Merge pull request #1248 from krisbiradar/add-support-for-gemma-3n

Add support for gemma 3n

Author: martindevans

LLama.Unittest/LLamaContextTests.cs

@@ -13,10 +13,9 @@ public LLamaContextTests()
         {
             var @params = new ModelParams(Constants.GenerativeModelPath2)
             {
-                ContextSize = 128,
+                ContextSize = 512,
                 BatchSize = 8,
                 UBatchSize = 8,
-                SeqMax = 1,
                 VocabOnly = false,
                 GpuLayerCount = Constants.CIGpuLayerCount,
             };
@@ -33,7 +32,7 @@ public void Dispose()
         [Fact]
         public void CheckProperties()
         {
-            Assert.Equal(128u, _context.ContextSize);
+            Assert.Equal(_context.NativeHandle.MaxSeq * 256, _context.ContextSize);
             Assert.Equal(960, _context.EmbeddingSize);
             Assert.Equal(49152, _context.Vocab.Count);
         }

LLama.Unittest/LLamaContextWithCustomLoggerTests.cs

@@ -30,7 +30,7 @@ public LLamaContextWithCustomLoggerTests()
         {
             var @params = new ModelParams(Constants.GenerativeModelPath2)
             {
-                ContextSize = 128,
+                ContextSize = 512,
                 GpuLayerCount = Constants.CIGpuLayerCount,
             };
 
@@ -55,7 +55,7 @@ public void Dispose()
         [Fact]
         public void CheckProperties()
         {
-            Assert.Equal(128u, _context.ContextSize);
+            Assert.Equal(_context.NativeHandle.MaxSeq * 256, _context.ContextSize);
             Assert.Equal(960, _context.EmbeddingSize);
             Assert.Equal(49152, _context.Vocab.Count);
         }

LLama.Unittest/LLamaRerankerTests.cs

@@ -20,7 +20,6 @@ public LLamaRerankerTests(ITestOutputHelper testOutputHelper)
             ContextSize = 0,
             PoolingType = LLamaPoolingType.Rank,
             GpuLayerCount = Constants.CIGpuLayerCount,
-
         };
         using var weights = LLamaWeights.LoadFromFile(@params);
         _reranker = new LLamaReranker(weights, @params);

LLama.Unittest/SamplingTests.cs

@@ -104,7 +104,7 @@ public void BatchedSampling()
                 }
             }
 
-            // Add " repeat" and test whether next tokens will be "this phrase forever.".
+            // Add " repeat" and test whether next tokens will be "this phrase forever."
             for (int i = 0; i < 4; i++)
             {
                 for (int b = 0; b < batch_count; b++)

LLama.Web/Common/ModelOptions.cs

@@ -102,7 +102,7 @@ public class ModelOptions
         public bool NoKqvOffload { get; set; }
 
         /// <inheritdoc />
-        public bool FlashAttention { get; set; }
+        public bool? FlashAttention { get; set; }
 
         /// <inheritdoc />
         public Encoding Encoding { get; set; } = Encoding.UTF8;

LLama/Abstractions/IContextParams.cs

@@ -106,8 +106,8 @@ public interface IContextParams
     /// <summary>
     /// Whether to use flash attention
     /// </summary>
-    bool FlashAttention { get; }
-
+    bool? FlashAttention { get; }
+    
     /// <summary>
     /// defragment the KV cache if holes/size &gt; defrag_threshold, Set to &lt;= 0 to disable (default)
     /// </summary>

LLama/Common/ModelParams.cs

@@ -1,3 +1,4 @@
+using System;
 using LLama.Abstractions;
 using System.Text;
 using System.Text.Json.Serialization;
@@ -95,12 +96,12 @@ public record ModelParams
 
         /// <inheritdoc />
         public bool NoKqvOffload { get; set; }
-
+        
         /// <inheritdoc />
-
-        public bool FlashAttention { get; set; }
+        public bool? FlashAttention { get; set; }
 
         /// <inheritdoc />
+        [Obsolete]
         public float? DefragThreshold { get; set; }
 
         /// <inheritdoc />

LLama/Extensions/IContextParamsExtensions.cs

@@ -37,7 +37,7 @@ public static void ToLlamaContextParams(this IContextParams @params, out LLamaCo
             result.yarn_beta_slow = @params.YarnBetaSlow ?? 1f;
             result.yarn_orig_ctx = @params.YarnOriginalContext ?? 0;
             result.rope_scaling_type = @params.YarnScalingType ?? RopeScalingType.Unspecified;
-
+            
             result.defrag_threshold = @params.DefragThreshold ?? -1;
 
             result.cb_eval = IntPtr.Zero;
@@ -49,9 +49,16 @@ public static void ToLlamaContextParams(this IContextParams @params, out LLamaCo
             result.type_k = @params.TypeK ?? GGMLType.GGML_TYPE_F16;
             result.type_v = @params.TypeV ?? GGMLType.GGML_TYPE_F16;
             result.offload_kqv = !@params.NoKqvOffload;
-            result.flash_attention = @params.FlashAttention;
             result.llama_pooling_type = @params.PoolingType;
             result.attention_type = @params.AttentionType;
+            result.llama_flash_attn_type = @params.FlashAttention switch
+            {
+                true => LLamaFlashAttentionType.LLAMA_FLASH_ATTENTION_TYPE_ENABLED,
+                false => LLamaFlashAttentionType.LLAMA_FLASH_ATTENTION_TYPE_DISABLED,
+                null => LLamaFlashAttentionType.LLAMA_FLASH_ATTENTION_TYPE_AUTO
+            };
+            result.kv_unified = true;
+            result.n_seq_max = (uint)Math.Min(Math.Max(10,result.n_ctx/8),256);
 
             result.n_threads = Threads(@params.Threads);
             result.n_threads_batch = Threads(@params.BatchThreads);

LLama/LLamaSharp.csproj

@@ -57,7 +57,7 @@
   </ItemGroup>
 
   <PropertyGroup>
-    <BinaryReleaseId>11dd5a44eb180e</BinaryReleaseId>
+    <BinaryReleaseId>86587da</BinaryReleaseId>
   </PropertyGroup>
 
   <PropertyGroup>

LLama/Native/LLamaContextParams.cs

@@ -64,6 +64,11 @@ public struct LLamaContextParams
         /// Attention type to use for embeddings
         /// </summary>
         public LLamaAttentionType attention_type;
+
+        /// <summary>
+        /// when to enable Flash Attention
+        /// </summary>
+        public LLamaFlashAttentionType llama_flash_attn_type;
 
         /// <summary>
         /// RoPE base frequency, 0 = from model