diff --git a/LLama/Sampling/DefaultSamplingPipeline.cs b/LLama/Sampling/DefaultSamplingPipeline.cs
index 08fbf05d8..b8da759af 100644
--- a/LLama/Sampling/DefaultSamplingPipeline.cs
+++ b/LLama/Sampling/DefaultSamplingPipeline.cs
@@ -112,6 +112,11 @@ public float PresencePenalty
     /// Seed to use for random sampling
     /// </summary>
     public uint Seed { get; set; } = GetRandomSeed();
+    
+    /// <summary>
+    /// Selected grammar optimization mode
+    /// </summary>
+    public GrammarOptimizationMode GrammarOptimization { get; init; } = GrammarOptimizationMode.None;
 
     /// <summary>
     /// A chain with just the grammar
@@ -158,7 +163,6 @@ private SafeLLamaSamplerChainHandle CreateGrammarChain(SafeLLamaContextHandle co
 
         var chain = SafeLLamaSamplerChainHandle.Create(LLamaSamplerChainParams.Default());
         chain.AddGrammar(context.ModelHandle, Grammar.Gbnf, Grammar.Root);
-        chain.AddDistributionSampler(Seed);
         return chain;
     }
 
@@ -217,30 +221,75 @@ public override LLamaToken Sample(SafeLLamaContextHandle ctx, int index)
         // Rent some buffers to use later
         var rentedBufferVocabSize = ArrayPool<LLamaTokenData>.Shared.Rent(ctx.ModelHandle.Vocab.Count);
         var rentedBufferSingleItem = ArrayPool<LLamaTokenData>.Shared.Rent(1);
+        
         try
         {
-            using (LLamaTokenDataArrayNative.Create(LLamaTokenDataArray.Create(ctx.GetLogitsIth(index), rentedBufferVocabSize), out var nativeAll))
+            // Handle grammar optimization modes
+            if (GrammarOptimization != GrammarOptimizationMode.None)
             {
-                // Apply the chain without the grammar to select one token which may or may not be valid
-                Apply(ctx, ref nativeAll);
-                var candidateToken = nativeAll.Data[checked((int)nativeAll.Selected)].ID;
-
-                // Now create another token data array with just that one token
-                rentedBufferSingleItem[0] = new LLamaTokenData(candidateToken, 1, 0);
-                using (LLamaTokenDataArrayNative.Create(new LLamaTokenDataArray(rentedBufferSingleItem, true), out var nativeSingleCandidate))
+                // Basic optimization : Apply the grammar to the selected token and check if it's valid
+                using (LLamaTokenDataArrayNative.Create(LLamaTokenDataArray.Create(ctx.GetLogitsIth(index), rentedBufferVocabSize), out var nativeAll))
                 {
-                    // Apply the grammar to this single candidate.
-                    _grammarChain.Apply(ref nativeSingleCandidate);
-
-                    // Test if that single token was rejected by the grammar
-                    if (!float.IsNegativeInfinity(nativeSingleCandidate.Data[0].Logit))
+                    // Apply the chain without the grammar to select one token which may or may not be valid
+                    Apply(ctx, ref nativeAll);
+                    
+                    // Select the candidate token
+                    var candidateToken = nativeAll.Data[checked((int)nativeAll.Selected)].ID;
+
+                    // Now create another token data array with just that one token
+                    rentedBufferSingleItem[0] = new LLamaTokenData(candidateToken, 1, 0);
+                    using (LLamaTokenDataArrayNative.Create(new LLamaTokenDataArray(rentedBufferSingleItem, true), out var nativeSingleCandidate))
                     {
-                        Accept(candidateToken);
-                        return candidateToken;
+                        // Apply the grammar chain to the single candidate
+                        _grammarChain.Apply(ref nativeSingleCandidate);
+
+                        // Check if the token passes the grammar
+                        if (!float.IsNegativeInfinity(nativeSingleCandidate.Data[0].Logit))
+                        {
+                            Accept(candidateToken);
+                            return candidateToken;
+                        }
+                    }
+                    
+                    // Extended optimization : Apply the grammar to the TopK tokens and check if the selected token is valid
+                    if (GrammarOptimization == GrammarOptimizationMode.Extended)
+                    {
+                        // Calculate a safe TopK value
+                        int safeTopK = Math.Min(TopK, nativeAll.Data.Length);
+
+                        // Rent a buffer for the TopK candidates
+                        var rentedBufferTopK = ArrayPool<LLamaTokenData>.Shared.Rent(safeTopK);
+                        try
+                        {
+                            // Copy only the TopK tokens from the existing candidate pool to the new buffer
+                            nativeAll.Data.Slice(0, safeTopK).CopyTo(rentedBufferTopK.AsSpan(0, safeTopK));
+
+                            // Create a native array with the TopK tokens
+                            using (LLamaTokenDataArrayNative.Create(new LLamaTokenDataArray(rentedBufferTopK, true), out var nativeTopK))
+                            {
+                                // Apply the grammar chain to the TopK candidates
+                                _grammarChain.Apply(ref nativeTopK);
+                                
+                                // Select the candidate token
+                                var candidateTokenTopK = nativeTopK.Data[checked((int)nativeTopK.Selected)];
+                                
+                                // Check if the token passes the grammar
+                                if (!float.IsNegativeInfinity(candidateTokenTopK.Logit))
+                                {
+                                    // Accept and return the token
+                                    Accept(candidateTokenTopK.ID);
+                                    return candidateTokenTopK.ID;
+                                }
+                            }
+                        }
+                        finally
+                        {
+                            ArrayPool<LLamaTokenData>.Shared.Return(rentedBufferTopK);
+                        }
                     }
                 }
             }
-
+            
             // If we get here the grammar rejected the token
             using (LLamaTokenDataArrayNative.Create(LLamaTokenDataArray.Create(ctx.GetLogitsIth(index), rentedBufferVocabSize), out var nativeAll))
             {
@@ -262,4 +311,25 @@ public override LLamaToken Sample(SafeLLamaContextHandle ctx, int index)
             ArrayPool<LLamaTokenData>.Shared.Return(rentedBufferSingleItem);
         }
     }
+    
+    /// <summary>
+    /// Grammar Optimization Mode
+    /// </summary>
+    public enum GrammarOptimizationMode
+    {
+        /// <summary>
+        /// No grammar optimization, slow because it has to apply the grammar to the entire vocab.
+        /// </summary>
+        None,
+
+        /// <summary>
+        /// Attempts to return early by only applying the grammar to the selected token and checking if it's valid.
+        /// </summary>
+        Basic,
+
+        /// <summary>
+        /// Attempts to return early by applying the grammar to the top K tokens and checking if the selected token is valid.
+        /// </summary>
+        Extended
+    }
 }
\ No newline at end of file