feat: 'disable_flashinfer_sampling' config option

Signed-off-by: ixlmar <[email protected]>

Author: ixlmar

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -824,7 +824,7 @@ def create_torch_sampler_args(mapping: Mapping, *, max_seq_len: int,
                               max_batch_size: int,
                               speculative_config: SpeculativeConfig,
                               max_beam_width: int,
-                              disable_flash_infer_sampling: bool):
+                              disable_flashinfer_sampling: bool):
     max_num_sequences = max_batch_size * mapping.pp_size
     max_draft_len = (0 if speculative_config is None else
                      speculative_config.max_draft_len)
@@ -837,7 +837,7 @@ def create_torch_sampler_args(mapping: Mapping, *, max_seq_len: int,
         max_total_draft_tokens=max_total_draft_tokens,
         max_num_sequences=max_num_sequences,
         max_beam_width=max_beam_width,
-        disable_flash_infer_sampling=disable_flash_infer_sampling,
+        disable_flashinfer_sampling=disable_flashinfer_sampling,
     )
 
 
@@ -853,15 +853,15 @@ def instantiate_sampler(
     speculative_config: SpeculativeConfig,
     decoding_config: trtllm.DecodingConfig,
     kv_cache_config: KvCacheConfig,
-    disable_flash_infer_sampling: bool,
+    disable_flashinfer_sampling: bool,
 ):
     sampler_args = create_torch_sampler_args(
         mapping,
         max_seq_len=engine.max_seq_len,
         max_batch_size=max_batch_size,
         speculative_config=speculative_config,
         max_beam_width=max_beam_width,
-        disable_flash_infer_sampling=disable_flash_infer_sampling,
+        disable_flashinfer_sampling=disable_flashinfer_sampling,
     )
     decoding_mode = get_decoding_mode(decoding_config=decoding_config,
                                       max_beam_width=max_beam_width)

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -516,7 +516,7 @@ def drafting_loop_wrapper(model):
             speculative_config=spec_config,
             decoding_config=decoding_config,
             kv_cache_config=kv_cache_config,
-            disable_flash_infer_sampling=llm_args._disable_flash_infer_sampling,
+            disable_flashinfer_sampling=llm_args._disable_flashinfer_sampling,
         )
         logger.info(f"Using Sampler: {type(sampler).__name__}")
 

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -615,7 +615,7 @@ class Args:
         max_num_sequences: int
         max_beam_width: int
         max_total_draft_tokens: int
-        disable_flash_infer_sampling: bool = False
+        disable_flashinfer_sampling: bool = False
 
     def __init__(self, args: Args):
         self.max_seq_len = args.max_seq_len
@@ -651,7 +651,7 @@ def __init__(self, args: Args):
             }
 
         self._grouped_sampler_cls: Type[GroupedStrategySampler]
-        if IS_FLASHINFER_AVAILABLE and not args.disable_flash_infer_sampling:
+        if IS_FLASHINFER_AVAILABLE and not args.disable_flashinfer_sampling:
             from .sampling_utils_flashinfer import FlashInferGroupedStrategySampler
 
             self._grouped_sampler_cls = FlashInferGroupedStrategySampler

tensorrt_llm/llmapi/llm_args.py

@@ -2775,8 +2775,10 @@ class TorchLlmArgs(BaseLlmArgs):
     # PrivateVars
     _quant_config: Optional[QuantConfig] = PrivateAttr(default=None)
 
-    _disable_flash_infer_sampling: bool = PrivateAttr(default=True)
-    """Unless this is set to False, FlashInfer.sampling is not used, even if available."""
+    disable_flashinfer_sampling: bool = Field(
+        default=True,
+        description="Disable the use of FlashInfer.sampling.",
+        status="prototype")
 
     @property
     def quant_config(self) -> QuantConfig:

tests/unittest/_torch/sampler/test_torch_sampler.py

@@ -1077,7 +1077,7 @@ def _build_sampler(
                 max_beam_width=1,  # currently the only supported value
                 max_num_sequences=num_seq_slots,
                 max_total_draft_tokens=max_draft_len,
-                disable_flash_infer_sampling=(not use_flashinfer),
+                disable_flashinfer_sampling=(not use_flashinfer),
             )
         )