[None][perf] Add CuteDSL to MiniMax-M3 autotune (for MXFP8 linear GEMM+quant)

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(

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(