Speculative One Model: FlashInfer sampling

Signed-off-by: Izzy Putterman <[email protected]>

Author: IzzyPutterman

tensorrt_llm/_torch/speculative/eagle3.py

@@ -367,6 +367,9 @@ def __init__(self, spec_config: "EagleDecodingConfig", mapping: Mapping):
         self.mapping = mapping
         self.guided_decoder: Optional[CapturableGuidedDecoder] = None
         self.force_num_accepted_tokens = get_force_num_accepted_tokens()
+        self.use_flashinfer = False
+        self.seed = 0
+        self.offset = 0
 
     # Skip torch.compile for now since current Torch is not compatible with Triton 3.4
     # @torch.compile(options={"max-autotune": True})
@@ -529,9 +532,17 @@ def _sample_tokens_for_batch(
             temperatures = spec_metadata.temperatures[:num_tokens]
             top_ks = spec_metadata.top_ks[:num_tokens]
             top_ps = spec_metadata.top_ps[:num_tokens]
+            if self.use_flashinfer:
+                self.seed += 1
 
             sampled_tokens = sampling_batch_spec_dec_one_model(
-                logits, temperatures, top_ks, top_ps)
+                logits,
+                temperatures,
+                top_ks,
+                top_ps,
+                use_flashinfer=self.use_flashinfer,
+                seed=self.seed,
+                offset=self.offset)
         else:
             sampled_tokens = torch.argmax(logits, dim=-1)
 

tensorrt_llm/_torch/speculative/one_model_sampler.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import torch
+from flashinfer.sampling import top_k_top_p_sampling_from_logits
 
 
 def forward_native(
@@ -78,6 +79,9 @@ def sampling_batch_spec_dec_one_model(
     temperatures: torch.Tensor,
     top_k: torch.Tensor,
     top_p: torch.Tensor,
+    use_flashinfer: bool = False,
+    seed: Optional[int] = None,
+    offset: Optional[int] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     CUDA-graph compatible sampling. Supports mixed sampling params.
@@ -87,5 +91,7 @@ def sampling_batch_spec_dec_one_model(
     sampling is opt-in for now.
     """
     logits = apply_temperature(logits, temperatures)
+    if use_flashinfer:
+        return top_k_top_p_sampling_from_logits(logits, top_k, top_p, seed=seed, offset=offset)
     random_sampled = forward_native(logits, top_k, top_p)
     return random_sampled