diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index ef2470519714..c2e308ba8049 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -16,9 +16,10 @@ import enum import os import threading +from contextlib import nullcontext from dataclasses import replace from functools import lru_cache -from typing import ClassVar, List, Mapping, Optional, Tuple, Union +from typing import Callable, ClassVar, List, Mapping, Optional, Tuple, Union import torch import triton # type: ignore[import] @@ -34,7 +35,7 @@ from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, DynamicTensorSpec, OptimizationProfile, TunableRunner, - TuningConfig) + TuningConfig, autotune) from ..cublaslt_utils import IS_CUBLASLT_AVAILABLE from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl @@ -48,7 +49,8 @@ from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets, fp4_scale_infer_shape, get_last_power_of_2_num_tokens_buckets, - last_positive_power_of_2) + get_power_of_2_num_tokens_buckets, + last_positive_power_of_2, next_positive_power_of_2) if IS_CUTLASS_DSL_AVAILABLE: from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ @@ -555,6 +557,9 @@ def _( _MXFP8_LARGE_M_BUCKETS = (8192, 16384, 32768) _MXFP8_LARGE_M_BANDS = ((6553, 8192), (13106, 16384), (19659, 32768)) _MXFP8_AUTOTUNED_OP = "trtllm::mxfp8_mxfp8_gemm_autotuned::gemm" +_MXFP8_QUANTIZE_AUTOTUNED_OP = "trtllm::mxfp8_quantize_autotuned::quantize" +_FLASHINFER_MXFP8_GEMM_AUTOTUNED_OP = ( + "trtllm::flashinfer_mxfp8_gemm_autotuned::gemm") def _map_to_mxfp8_large_m_bucket(num_tokens: int) -> int: @@ -670,6 +675,184 @@ def _( return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) +@lru_cache(maxsize=1) +def _get_flashinfer_mxfp8_cute_dsl_ops( +) -> Optional[Tuple[Callable, Callable, Callable]]: + if not IS_FLASHINFER_AVAILABLE: + return None + try: + from flashinfer import mm_mxfp8, mxfp8_quantize + from flashinfer.autotuner import autotune as flashinfer_autotune + from flashinfer.cute_dsl import is_cute_dsl_available + + if is_cute_dsl_available(): + return mxfp8_quantize, mm_mxfp8, flashinfer_autotune + except (ImportError, RuntimeError): + pass + return None + + +def is_flashinfer_mxfp8_cute_dsl_available() -> bool: + """Return whether both CuTeDSL MXFP8 stages can join backend tuning.""" + return _get_flashinfer_mxfp8_cute_dsl_ops() is not None + + +# Both runners include process-local CuTeDSL JIT state, so their winners stay +# in the in-process cache rather than being persisted. +class MXFP8QuantizeRunner(TunableRunner): + """Profile native and FlashInfer CuTeDSL activation quantization.""" + + TRTLLM = -1 + CUTE_DSL = 0 + + tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_power_of_2_num_tokens_buckets, next_positive_power_of_2), ), + exclude_from_cache=True) + + def __init__(self, input_dtype: torch.dtype) -> None: + self.input_dtype = input_dtype + ops = _get_flashinfer_mxfp8_cute_dsl_ops() + assert ops is not None + self.cute_dsl_quantize = ops[0] + + def unique_id(self) -> Tuple[torch.dtype]: + return (self.input_dtype, ) + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + return [self.TRTLLM, self.CUTE_DSL] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = TRTLLM, + ) -> Tuple[torch.Tensor, torch.Tensor]: + activation = inputs[0] + if tactic == self.CUTE_DSL: + return self.cute_dsl_quantize( + activation, + is_sf_swizzled_layout=True, + alignment=32, + enable_pdl=None, + backend="cute-dsl", + ) + return torch.ops.trtllm.mxfp8_quantize(activation, True) + + +class FlashInferMXFP8GemmRunner(TunableRunner): + """Profile FlashInfer CUTLASS and CuTeDSL MXFP8 GEMMs.""" + + CUTLASS = -1 + CUTE_DSL = 0 + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_power_of_2_num_tokens_buckets, + next_positive_power_of_2), ), + constraint_specs=(ConstraintSpec(1, 0, _mxfp8_scale_infer_shape), ), + exclude_from_cache=True, + ) + + def __init__(self, output_dtype: torch.dtype) -> None: + self.output_dtype = output_dtype + ops = _get_flashinfer_mxfp8_cute_dsl_ops() + assert ops is not None + _, self.cute_dsl_gemm, self.flashinfer_autotune = ops + + def unique_id(self) -> Tuple[torch.dtype]: + return (self.output_dtype, ) + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + return [self.CUTLASS, self.CUTE_DSL] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = CUTLASS, + ) -> torch.Tensor: + act, act_scale, weight, weight_scale = inputs + if tactic == self.CUTLASS: + return torch.ops.trtllm.flashinfer_mm_mxfp8(act, act_scale, weight, + weight_scale, + self.output_dtype) + + with self.flashinfer_autotune(tune_mode=False, skip_ops="mxfp8_gemm"): + return self.cute_dsl_gemm( + act, + weight.t(), + act_scale, + weight_scale, + out_dtype=self.output_dtype, + use_8x4_sf_layout=False, + backend="cute-dsl", + ) + + +def _choose_mxfp8_tactic(custom_op: str, + runner: TunableRunner, + inputs: List[torch.Tensor], + tune: bool, + flashinfer_autotune: Optional[Callable] = None) -> int: + """Profile each process-local shape bucket once, then reuse its winner.""" + tuner = AutoTuner.get() + should_tune = tune + if should_tune: + # Excluded ops always profile in tuning mode. Enter that mode only + # when this process has not selected a winner for the bucket yet. + is_cache_hit, *_ = tuner.profiling_cache.search_cache( + custom_op, + [runner], + tuple(input.size() for input in inputs), + runner.tuning_config, + ) + should_tune = not is_cache_hit + flashinfer_context = (flashinfer_autotune() + if should_tune and flashinfer_autotune is not None + else nullcontext()) + with autotune(tune_mode=should_tune, + skip_dynamic_tuning_buckets=True), flashinfer_context: + return tuner.choose_one( + custom_op, + [runner], + runner.tuning_config, + inputs, + )[1] + + +def mxfp8_quantize_autotuned( + activation: torch.Tensor, + tune: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Select an MXFP8 quantizer for the current generation-graph bucket.""" + runner = MXFP8QuantizeRunner(activation.dtype) + inputs = [activation] + tactic = _choose_mxfp8_tactic(_MXFP8_QUANTIZE_AUTOTUNED_OP, runner, inputs, + tune) + return runner(inputs, tactic=tactic) + + +def flashinfer_mxfp8_gemm_autotuned( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool = False, +) -> torch.Tensor: + """Select a FlashInfer MXFP8 GEMM for the current graph bucket.""" + runner = FlashInferMXFP8GemmRunner(output_dtype) + inputs = [act, act_scale, weight, weight_scale] + tactic = _choose_mxfp8_tactic( + _FLASHINFER_MXFP8_GEMM_AUTOTUNED_OP, + runner, + inputs, + tune, + runner.flashinfer_autotune, + ) + return runner(inputs, tactic=tactic) + + class FP4GemmRunner(TunableRunner): runner_dict = dict() tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index e5b3acfa28ef..cada2fb1f3ea 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -15,7 +15,9 @@ from torch.nn.parameter import Parameter import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils -from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + BufferKind, flashinfer_mxfp8_gemm_autotuned, + is_flashinfer_mxfp8_cute_dsl_available, mxfp8_quantize_autotuned) from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm._utils import is_device_integrated, mpi_disabled from tensorrt_llm.bindings import ipc_nvls_supported @@ -3048,6 +3050,8 @@ def _flashinfer_mxfp8_op(): "flashinfer_mxfp8_autotune_active", default=False) _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE = ContextVar( "flashinfer_mxfp8_decode_graph_capture_active", default=False) +_MXFP8_GRAPH_BACKEND_TUNING_ACTIVE = ContextVar( + "mxfp8_graph_backend_tuning_active", default=False) @contextmanager @@ -3064,13 +3068,15 @@ def flashinfer_mxfp8_autotune(): @contextmanager -def flashinfer_mxfp8_decode_graph_capture(): - """Enable auto-dispatched FlashInfer calls only for decode graph capture.""" - token = _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.set(True) +def flashinfer_mxfp8_decode_graph_capture(*, tune_backends: bool = False): + """Mark generation graph capture and optionally tune MXFP8 backends.""" + capture_token = _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.set(True) + tuning_token = _MXFP8_GRAPH_BACKEND_TUNING_ACTIVE.set(tune_backends) try: yield finally: - _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.reset(token) + _MXFP8_GRAPH_BACKEND_TUNING_ACTIVE.reset(tuning_token) + _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.reset(capture_token) class MXFP8LinearMethod(LinearMethodBase): @@ -3084,15 +3090,15 @@ class MXFP8LinearMethod(LinearMethodBase): MXFP8 activation quantize + block-scaled e4m3xe4m3 GEMM. - FlashInfer: reuse the CUTLASS-layout activations, weights, and scales via trtllm::flashinfer_mm_mxfp8, which wraps mm_mxfp8 so a compiled - graph gets one opaque node. MiniMax-M3 enables this path automatically - only while tuning or capturing decode CUDA graphs, and that automatic - selection is skipped under torch.compile. A pinned flashinfer backend - applies everywhere. + graph gets one opaque node. MiniMax-M3 can independently tune + quantization and GEMM backends while warming its decode CUDA graphs. + Automatic selection is skipped under torch.compile; a pinned + FlashInfer backend applies everywhere. ``TRTLLM_MXFP8_GEMM_BACKEND`` can explicitly select ``trtllm``, ``flashinfer``, or ``auto``; ``auto`` settles on ``trtllm`` when the model - is compiled. The reference layout is 2D [O,K/32]; both - compiled backends consume the same 1D padded swizzled scale layout. + is compiled. The reference layout is 2D [O,K/32]; both compiled backends + consume the same 1D padded swizzled scale layout. When the TensorRT-LLM autotuner is enabled, the native backend profiles its compiled tactics during startup. Learned tactics are registered in the native op so serving avoids the Python autotuner lookup; the generic @@ -3109,11 +3115,12 @@ def __init__(self): self.backend = os.environ.get("TRTLLM_MXFP8_GEMM_BACKEND", "trtllm") self.use_native_autotuner = True self._native_autotuned = False + self._flashinfer_autotuned = False + self._use_graph_backend_selection = False if self.backend not in ("trtllm", "flashinfer", "auto"): raise ValueError("TRTLLM_MXFP8_GEMM_BACKEND must be 'trtllm', " f"'flashinfer', or 'auto', got {self.backend!r}") self._flashinfer_mxfp8 = None - self._flashinfer_autotuned = False if self.backend == "flashinfer": self._load_flashinfer(required=True) elif self.backend == "auto" and not self._load_flashinfer( @@ -3126,7 +3133,13 @@ def uses_flashinfer(self) -> bool: @property def needs_flashinfer_autotune(self) -> bool: - return self.uses_flashinfer and self._flashinfer_mxfp8 is not None + return (self.uses_flashinfer and not self._use_graph_backend_selection + and self._flashinfer_mxfp8 is not None) + + @property + def uses_graph_backend_selection(self) -> bool: + """Whether graph capture should use the per-stage backend winners.""" + return self._use_graph_backend_selection @property def needs_native_autotune(self) -> bool: @@ -3163,6 +3176,20 @@ def enable_flashinfer_auto(self) -> bool: self.backend = "auto" return True + def configure_default_graph_dispatch(self, *, + enable_backend_tuning: bool) -> None: + """Configure model-default generation-graph dispatch. + + Leave explicit backend overrides untouched. When backend tuning is + unavailable, retain the eager-tuned FlashInfer graph path. + """ + if "TRTLLM_MXFP8_GEMM_BACKEND" in os.environ: + return + if not self.enable_flashinfer_auto(): + return + self._use_graph_backend_selection = ( + enable_backend_tuning and is_flashinfer_mxfp8_cute_dsl_available()) + def mark_flashinfer_autotuned(self) -> None: self._flashinfer_autotuned = True @@ -3177,6 +3204,7 @@ def disable_flashinfer_auto(self) -> None: if self.backend == "auto": self.backend = "trtllm" self._flashinfer_autotuned = False + self._use_graph_backend_selection = False @classmethod def _swizzled_scale_size(cls, out_features: int, in_features: int) -> int: @@ -3220,8 +3248,16 @@ def apply(self, module: Linear, input: torch.Tensor, if self.use_cutlass: # Dynamic MXFP8 activation quantization (swizzled SF layout), then # the CUTLASS block-scaled e4m3xe4m3 GEMM. - act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize( - input.contiguous(), True) + input = input.contiguous() + use_graph_backend_selection = ( + self._use_graph_backend_selection and not is_torch_compiling() + and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()) + if use_graph_backend_selection: + tune_backends = _MXFP8_GRAPH_BACKEND_TUNING_ACTIVE.get() + act_e4m3, act_sf = mxfp8_quantize_autotuned(input, + tune=tune_backends) + else: + act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize(input, True) # The automatic path switches per call on context state, which # Dynamo cannot trace and a compiled graph could not honor. A # pinned backend resolves before tracing, so only the automatic @@ -3231,7 +3267,16 @@ def apply(self, module: Linear, input: torch.Tensor, (_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() or (self._flashinfer_autotuned and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()))) - if use_flashinfer: + if use_graph_backend_selection: + output = flashinfer_mxfp8_gemm_autotuned( + act_e4m3, + act_sf, + module.weight, + module.weight_scale, + module.dtype, + tune=tune_backends, + ) + elif use_flashinfer: flashinfer_mxfp8 = self._flashinfer_mxfp8 assert flashinfer_mxfp8 is not None output = flashinfer_mxfp8( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c308a1847b6e..a78281431149 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1187,6 +1187,10 @@ def warmup(self, resource_manager: ResourceManager) -> None: self._run_cuda_graph_warmup(resource_manager) finally: self.cuda_graph_runner.is_warmup_only = False + if getattr(self, "_use_mxfp8_graph_backend_selection", False): + # Drop profiling intermediates before allocating capture graphs. + gc.collect() + torch.cuda.empty_cache() self.cuda_graph_runner.padding_dummy_requests = {} self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") @@ -1462,34 +1466,39 @@ def trtllm_gen_fmha_jit_warmup(): torch.cuda.synchronize() def _run_autotuner_warmup(self, resource_manager: ResourceManager): - """Runs a forward pass to populate the autotuner cache.""" + """Configure autotuners and run eager warmup when required.""" from ..modules.linear import (MXFP8LinearMethod, flashinfer_mxfp8_autotune) enable_trtllm_autotuner = self.llm_args.enable_autotuner - # The automatic FlashInfer dispatch keys off context state that Dynamo - # cannot trace, so a compiled model never reaches it (see - # MXFP8LinearMethod.apply); leave those layers on the native backend - # instead of tuning tactics that would go unused. - use_mxfp8_flashinfer_graph_default = ( - self.cuda_graph_runner.enabled and not self._torch_compile_enabled - and "TRTLLM_MXFP8_GEMM_BACKEND" not in os.environ and any( + # This is only the model-level opt-in. MXFP8LinearMethod owns explicit + # user overrides and backend capability checks. + mxfp8_decode_graph_default_requested = ( + self.cuda_graph_runner.enabled and any( getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) for module in self.model.modules())) + self._use_mxfp8_graph_backend_selection = False flashinfer_mxfp8_methods = [] native_mxfp8_methods = [] for module in self.model.modules(): quant_method = getattr(module, "quant_method", None) if not isinstance(quant_method, MXFP8LinearMethod): continue - if use_mxfp8_flashinfer_graph_default: - quant_method.enable_flashinfer_auto() - elif self._torch_compile_enabled: - # Same reasoning for a user-requested - # TRTLLM_MXFP8_GEMM_BACKEND=auto: settle on native here so - # warmup tunes the backend that execution will pick. + if self._torch_compile_enabled: + # Automatic dispatch keys off context state that Dynamo cannot + # trace. Settle auto on native so warmup tunes the path that + # execution will use; an explicitly pinned backend is retained. quant_method.disable_flashinfer_auto() + elif mxfp8_decode_graph_default_requested: + # Graph warmup has no pipeline-parallel cache handoff; retain + # the existing eager FlashInfer tuner when PP is enabled. + quant_method.configure_default_graph_dispatch( + enable_backend_tuning=not self.mapping.has_pp()) + if quant_method.uses_graph_backend_selection: + self._use_mxfp8_graph_backend_selection = True if quant_method.needs_flashinfer_autotune: + # Includes the model-default fallback when per-stage graph + # backend tuning is unavailable. flashinfer_mxfp8_methods.append(quant_method) if enable_trtllm_autotuner and quant_method.needs_native_autotune: native_mxfp8_methods.append(quant_method) @@ -1498,10 +1507,10 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods) enable_native_mxfp8_autotuner = bool(native_mxfp8_methods) + if (enable_trtllm_autotuner or self._use_mxfp8_graph_backend_selection): + AutoTuner.get().setup_distributed_state(self.mapping, self.dist) if not enable_trtllm_autotuner and not enable_flashinfer_mxfp8_autotuner: return - if enable_trtllm_autotuner: - AutoTuner.get().setup_distributed_state(self.mapping, self.dist) logger.info( f"Running autotuner warmup (TRT-LLM={enable_trtllm_autotuner}, " f"native MXFP8={enable_native_mxfp8_autotuner}, " @@ -1846,7 +1855,11 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): # The automatic MiniMax-M3 MXFP8 selection is decode-graph-only. # Keep piecewise context/prefill graph capture on the native backend. - with flashinfer_mxfp8_decode_graph_capture(): + tune_mxfp8_backends = ( + self.cuda_graph_runner.is_warmup_only + and getattr(self, "_use_mxfp8_graph_backend_selection", False)) + with flashinfer_mxfp8_decode_graph_capture( + tune_backends=tune_mxfp8_backends): self._capture_generation_cuda_graphs(resource_manager) # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index df37a09b9326..9a173979745a 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -216,9 +216,15 @@ def test_step_b_cleanup_skipped_with_helix_cp(self): calls.count("empty_cache"), 0, f"Helix CP should skip all warmup cleanup; got {calls}" ) - def test_flashinfer_mxfp8_autotunes_before_graph_capture(self): - """An auto-enabled M3 linear tunes even when TRT autotuning is disabled.""" + def _run_mxfp8_graph_default_warmup( + self, + *, + model_default=True, + cute_dsl_available=True, + pipeline_parallel=False, + ): calls = [] + tuner = SimpleNamespace(setup_distributed_state=Mock()) @contextlib.contextmanager def flashinfer_autotune(): @@ -228,7 +234,7 @@ def flashinfer_autotune(): flashinfer_module = ModuleType("flashinfer") flashinfer_autotuner_module = ModuleType("flashinfer.autotuner") - flashinfer_autotuner_module.autotune = Mock(side_effect=flashinfer_autotune) + flashinfer_autotuner_module.autotune = flashinfer_autotune flashinfer_module.autotuner = flashinfer_autotuner_module with ( @@ -247,27 +253,32 @@ def flashinfer_autotune(): "tensorrt_llm._torch.modules.linear._flashinfer_mxfp8_op", return_value=Mock(), ), + patch( + "tensorrt_llm._torch.modules.linear.is_flashinfer_mxfp8_cute_dsl_available", + return_value=cute_dsl_available, + ), patch.dict("os.environ", {}, clear=True), ): method = MXFP8LinearMethod() - self.assertEqual(method.backend, "trtllm") + modules = [SimpleNamespace(quant_method=method)] + if model_default: + modules.insert( + 0, + SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), + ) engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=False), cuda_graph_runner=SimpleNamespace(enabled=True), _torch_compile_enabled=False, - model=SimpleNamespace( - modules=lambda: [ - SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), - SimpleNamespace(quant_method=method), - ] - ), + model=SimpleNamespace(modules=lambda: modules), kv_cache_manager_key="kv_cache", max_num_tokens=16, batch_size=16, max_seq_len=2, original_max_draft_len=0, - mapping=SimpleNamespace(tp_size=1), + mapping=SimpleNamespace(tp_size=1, has_pp=lambda: pipeline_parallel), + dist=object(), is_draft_model=False, no_cuda_graph=lambda: contextlib.nullcontext(), _create_warmup_request=Mock(return_value=object()), @@ -277,23 +288,74 @@ def flashinfer_autotune(): ) kv_cache_manager = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resource_manager = SimpleNamespace( - get_resource_manager=lambda key: (kv_cache_manager if key == "kv_cache" else None) + get_resource_manager=lambda key: kv_cache_manager if key == "kv_cache" else None ) with ( + patch( + "tensorrt_llm._torch.pyexecutor.model_engine.AutoTuner.get", + return_value=tuner, + ), patch("torch.cuda.synchronize"), patch("torch.cuda.empty_cache"), patch("tensorrt_llm._torch.pyexecutor.model_engine.clear_memory_buffers"), ): PyTorchModelEngine._run_autotuner_warmup(engine, resource_manager) - self.assertEqual( - calls, - ["flashinfer_autotune_enter", "forward", "flashinfer_autotune_exit"], - ) - self.assertEqual(method.backend, "auto") - self.assertTrue(method._flashinfer_autotuned) - flashinfer_autotuner_module.autotune.assert_called_once_with() + return calls, method, engine, tuner + + def test_mxfp8_graph_default_routing(self): + """Use graph tuning only when the M3 default is fully eligible.""" + eager_calls = [ + "flashinfer_autotune_enter", + "forward", + "flashinfer_autotune_exit", + ] + cases = { + "eligible": ({}, [], True), + "missing_cutedsl": ({"cute_dsl_available": False}, eager_calls, False), + "pipeline_parallel": ({"pipeline_parallel": True}, eager_calls, False), + "other_model": ({"model_default": False}, [], False), + } + for name, (kwargs, expected_calls, graph_selection) in cases.items(): + with self.subTest(name=name): + calls, method, engine, tuner = self._run_mxfp8_graph_default_warmup(**kwargs) + + self.assertEqual(calls, expected_calls) + self.assertEqual(method.uses_graph_backend_selection, graph_selection) + self.assertEqual(engine._use_mxfp8_graph_backend_selection, graph_selection) + if graph_selection: + tuner.setup_distributed_state.assert_called_once_with( + engine.mapping, engine.dist + ) + else: + tuner.setup_distributed_state.assert_not_called() + + def test_graph_warmup_tunes_mxfp8_only_on_first_opted_in_pass(self): + tune_values = [] + + @contextlib.contextmanager + def graph_capture(*, tune_backends=False): + tune_values.append(tune_backends) + yield + + for warmup_only in (True, False): + engine = SimpleNamespace( + cuda_graph_runner=SimpleNamespace(enabled=True, is_warmup_only=warmup_only), + _torch_compile_piecewise_cuda_graph=False, + _use_mxfp8_graph_backend_selection=True, + _capture_generation_cuda_graphs=Mock(), + _capture_piecewise_cuda_graphs=Mock(), + ) + with patch( + "tensorrt_llm._torch.modules.linear.flashinfer_mxfp8_decode_graph_capture", + side_effect=graph_capture, + ): + PyTorchModelEngine._run_cuda_graph_warmup(engine, Mock()) + + engine._capture_generation_cuda_graphs.assert_called_once() + + self.assertEqual(tune_values, [True, False]) def _run_torch_compile_warmup_with_env(self, environ): """Run warmup for a compiled M3-style engine under the given environ.""" diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index a7827b8ac517..1092d59154b5 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -182,6 +182,52 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): assert native_gemm.call_count == 2 +def test_mxfp8_graph_backend_selection_is_an_explicit_opt_in(monkeypatch): + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + monkeypatch.setattr(linear_module, "is_flashinfer_mxfp8_cute_dsl_available", lambda: True) + + graph_output = torch.empty((2, 3), dtype=torch.bfloat16) + flashinfer_gemm = _mock_flashinfer_mxfp8_op(monkeypatch, graph_output) + quantized, activation_scale, _, native_gemm, native_output, _, _ = _mock_mxfp8_ops(monkeypatch) + graph_quantize = Mock(return_value=(quantized, activation_scale)) + graph_gemm = Mock(return_value=graph_output) + monkeypatch.setattr(linear_module, "mxfp8_quantize_autotuned", graph_quantize) + monkeypatch.setattr(linear_module, "flashinfer_mxfp8_gemm_autotuned", graph_gemm) + + module = SimpleNamespace( + weight=torch.empty((3, 4), dtype=torch.float8_e4m3fn), + weight_scale=torch.empty(512, dtype=torch.uint8), + dtype=torch.bfloat16, + ) + activation = torch.randn((2, 4), dtype=torch.bfloat16) + method = MXFP8LinearMethod() + method.configure_default_graph_dispatch(enable_backend_tuning=True) + assert method.uses_graph_backend_selection + assert not method.needs_flashinfer_autotune + + assert method.apply(module, activation, bias=None) is native_output + native_gemm.assert_called_once() + flashinfer_gemm.assert_not_called() + + with flashinfer_mxfp8_decode_graph_capture(tune_backends=True): + assert method.apply(module, activation, bias=None) is graph_output + graph_quantize.assert_called_once_with(activation, tune=True) + graph_gemm.assert_called_once_with( + quantized, + activation_scale, + module.weight, + module.weight_scale, + module.dtype, + tune=True, + ) + flashinfer_gemm.assert_not_called() + + # Leaving the decode-capture scope restores the eager/native path. + assert method.apply(module, activation, bias=None) is native_output + assert native_gemm.call_count == 2 + + def test_mxfp8_auto_stays_native_under_torch_compile(monkeypatch): """Dynamo cannot trace the context lookups that gate FlashInfer dispatch.""" monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False)