[None][perf] Use FlashInfer MXFP8 GEMM for MiniMax-M3 decode

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -81,6 +81,7 @@
 # and flash SDPA does not accept attn_mask.
 _DENSE_SDPA_BACKENDS = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
 
+
 # ---------------------------------------------------------------------------
 # Config normalization helpers
 # ---------------------------------------------------------------------------
@@ -1865,6 +1866,10 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"):
                 for layer_idx in range(config.num_hidden_layers)
             ]
         )
+        # The executor owns the authoritative CUDA-graph configuration. Mark
+        # this model as eligible here and let the executor enable the automatic
+        # FlashInfer MXFP8 path only when its decode graph runner is active.
+        self._use_flashinfer_mxfp8_decode_graph_default = True
         # Final norm is a plain (non-Gemma) RMSNorm for the same reason as the
         # layer-boundary norms (see MiniMaxM3DecoderLayer.__init__): it doubles
         # as the last layer's next_layer_layernorm, so the last MoE/MLP output

tensorrt_llm/_torch/modules/linear.py

@@ -4,6 +4,8 @@
 import math
 import os
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass
 from typing import ClassVar, Dict, List, Optional, Union
 
@@ -2956,19 +2958,52 @@ def _mxfp8_cutlass_op_available() -> bool:
                    ) and torch.cuda.get_device_capability()[0] >= 10
 
 
+_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE = ContextVar(
+    "flashinfer_mxfp8_autotune_active", default=False)
+_FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE = ContextVar(
+    "flashinfer_mxfp8_decode_graph_capture_active", default=False)
+
+
+@contextmanager
+def flashinfer_mxfp8_autotune():
+    """Tune FlashInfer MXFP8 tactics while enabling auto-dispatched calls."""
+    from flashinfer.autotuner import autotune
+
+    token = _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.set(True)
+    try:
+        with autotune():
+            yield
+    finally:
+        _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.reset(token)
+
+
+@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)
+    try:
+        yield
+    finally:
+        _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.reset(token)
+
+
 class MXFP8LinearMethod(LinearMethodBase):
     """MXFP8 weights (e4m3 + UE8M0 1x32) x dynamic MXFP8 activations (W8A8).
 
-    Two execution paths share a common loader:
+    Three execution paths share a common loader:
       - Reference (no CUTLASS op compiled): dequantize the weight to compute
         dtype and run F.linear. Slow but correct -- used to seed M1 tests and
         as a portable fallback.
       - CUTLASS (Blackwell sm100/103 + mxfp8_mxfp8_gemm op present): dynamic
         MXFP8 activation quantize + block-scaled e4m3xe4m3 GEMM.
-
-    The path is selected at create_weights time via _mxfp8_cutlass_op_available;
-    the weight_scale tensor's layout matches the chosen path (2D [O,K/32] for
-    the reference, 1D padded swizzled for CUTLASS) so apply() stays branchless.
+      - FlashInfer: reuse the CUTLASS-layout activations, weights, and scales
+        with ``mm_mxfp8``. MiniMax-M3 enables this path automatically only
+        while tuning or capturing decode CUDA graphs; eager execution remains
+        on the native TensorRT-LLM op.
+
+    ``TRTLLM_MXFP8_GEMM_BACKEND`` can explicitly select ``trtllm``,
+    ``flashinfer``, or ``auto``. The reference layout is 2D [O,K/32]; both
+    compiled backends consume the same 1D padded swizzled scale layout.
     """
     BLOCK_SIZE = 32
     # Swizzled-SF layout padding (matches W4A8MXFP4FP8: rows->128, cols/SFblock->4).
@@ -2978,6 +3013,64 @@ class MXFP8LinearMethod(LinearMethodBase):
     def __init__(self):
         super().__init__()
         self.use_cutlass = _mxfp8_cutlass_op_available()
+        self.backend = os.environ.get("TRTLLM_MXFP8_GEMM_BACKEND", "trtllm")
+        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(
+                required=False):
+            self.backend = "trtllm"
+
+    @property
+    def uses_flashinfer(self) -> bool:
+        return self.backend in ("flashinfer", "auto")
+
+    @property
+    def needs_flashinfer_autotune(self) -> bool:
+        return self.uses_flashinfer and self._flashinfer_mxfp8 is not None
+
+    def _load_flashinfer(self, *, required: bool) -> bool:
+        if not self.use_cutlass:
+            if required:
+                raise RuntimeError(
+                    "FlashInfer MXFP8 GEMM requires the TensorRT-LLM MXFP8 "
+                    "quantization ops on Blackwell")
+            return False
+        try:
+            from flashinfer import mm_mxfp8
+        except ImportError as error:
+            if required:
+                raise RuntimeError(
+                    "TRTLLM_MXFP8_GEMM_BACKEND=flashinfer requires the "
+                    "pinned flashinfer-python package") from error
+            logger.warning_once(
+                "FlashInfer MXFP8 is unavailable; using the native "
+                "TensorRT-LLM GEMM backend.",
+                key="flashinfer_mxfp8_unavailable")
+            return False
+        self._flashinfer_mxfp8 = mm_mxfp8
+        return True
+
+    def enable_flashinfer_auto(self) -> bool:
+        """Enable graph-only FlashInfer dispatch unless the user overrode it."""
+        if "TRTLLM_MXFP8_GEMM_BACKEND" in os.environ:
+            return self.backend == "auto"
+        if not self._load_flashinfer(required=False):
+            return False
+        self.backend = "auto"
+        return True
+
+    def mark_flashinfer_autotuned(self) -> None:
+        self._flashinfer_autotuned = True
+
+    def disable_flashinfer_auto(self) -> None:
+        if self.backend == "auto":
+            self.backend = "trtllm"
+            self._flashinfer_autotuned = False
 
     @classmethod
     def _swizzled_scale_size(cls, out_features: int, in_features: int) -> int:
@@ -3023,15 +3116,36 @@ def apply(self, module: Linear, input: torch.Tensor,
             # the CUTLASS block-scaled e4m3xe4m3 GEMM.
             act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize(
                 input.contiguous(), True)
-            # globalScale is the alpha multiplier; pure MXFP8xMXFP8 uses 1.0.
-            global_scale = torch.ones([1],
-                                      dtype=torch.float32,
-                                      device=input.device)
-            output = torch.ops.trtllm.mxfp8_mxfp8_gemm(act_e4m3, act_sf,
-                                                       module.weight,
-                                                       module.weight_scale,
-                                                       global_scale,
-                                                       module.dtype)
+            use_flashinfer = self.backend == "flashinfer" or (
+                self.backend == "auto" and
+                (_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() or
+                 (self._flashinfer_autotuned
+                  and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get())))
+            if use_flashinfer:
+                flashinfer_mxfp8 = self._flashinfer_mxfp8
+                assert flashinfer_mxfp8 is not None
+                output = flashinfer_mxfp8(
+                    act_e4m3,
+                    module.weight.t(),
+                    act_sf,
+                    module.weight_scale,
+                    out_dtype=module.dtype,
+                    use_8x4_sf_layout=False,
+                    backend="cutlass",
+                )
+            else:
+                # globalScale is the alpha multiplier; pure MXFP8xMXFP8 uses 1.0.
+                global_scale = torch.ones([1],
+                                          dtype=torch.float32,
+                                          device=input.device)
+                output = torch.ops.trtllm.mxfp8_mxfp8_gemm(
+                    act_e4m3,
+                    act_sf,
+                    module.weight,
+                    module.weight_scale,
+                    global_scale,
+                    module.dtype,
+                )
             if bias is not None:
                 output = output + bias
         else:

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1307,10 +1307,33 @@ def trtllm_gen_fmha_jit_warmup():
 
     def _run_autotuner_warmup(self, resource_manager: ResourceManager):
         """Runs a forward pass to populate the autotuner cache."""
-        if not self.llm_args.enable_autotuner:
+        from ..modules.linear import (MXFP8LinearMethod,
+                                      flashinfer_mxfp8_autotune)
+
+        enable_trtllm_autotuner = self.llm_args.enable_autotuner
+        use_mxfp8_flashinfer_graph_default = (
+            self.cuda_graph_runner.enabled
+            and "TRTLLM_MXFP8_GEMM_BACKEND" not in os.environ and any(
+                getattr(module, "_use_flashinfer_mxfp8_decode_graph_default",
+                        False) for module in self.model.modules()))
+        flashinfer_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()
+            if quant_method.needs_flashinfer_autotune:
+                flashinfer_mxfp8_methods.append(quant_method)
+        enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods)
+
+        if not enable_trtllm_autotuner and not enable_flashinfer_mxfp8_autotuner:
             return
-        AutoTuner.get().setup_distributed_state(self.mapping, self.dist)
-        logger.info("Running autotuner warmup...")
+        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"FlashInfer MXFP8={enable_flashinfer_mxfp8_autotuner})...")
         kv_cache_manager = resource_manager.get_resource_manager(
             self.kv_cache_manager_key)
         token_num_upper_bound = min(self.max_num_tokens,
@@ -1320,7 +1343,15 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
             max_num_draft_tokens=self.original_max_draft_len)
 
         cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None)
-        with self.no_cuda_graph(), autotune(cache_path=cache_path):
+        trtllm_autotune_context = (autotune(
+            cache_path=cache_path) if enable_trtllm_autotuner else
+                                   contextlib.nullcontext())
+        flashinfer_autotune_context = (flashinfer_mxfp8_autotune()
+                                       if enable_flashinfer_mxfp8_autotuner else
+                                       contextlib.nullcontext())
+        ran_forward = False
+        with self.no_cuda_graph(
+        ), trtllm_autotune_context, flashinfer_autotune_context:
             warmup_request = self._create_warmup_request(
                 resource_manager, curr_max_num_tokens, 0)
             with self._release_batch_context(warmup_request,
@@ -1342,21 +1373,41 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
                     self.forward(batch,
                                  new_tensors_device=None,
                                  resource_manager=resource_manager)
+                    ran_forward = True
 
-                    # pp_recv in AutoTuner choose_one will never be called if there is no tuning op during the forward pass.
-                    # So we need to make an extra call to consume the previous rank's pp_send to guarantee that the previous rank's pp_send is released.
-                    AutoTuner.get().cache_pp_recv()
-                    # Send the cache after the tuning process to the next PP rank
-                    AutoTuner.get().cache_pp_send()
-                    # Clean the pp flag to avoid deadlock with synchronous send/recv
-                    AutoTuner.get().clean_pp_flag()
+                    if enable_trtllm_autotuner:
+                        # pp_recv in AutoTuner choose_one will never be called if there is no tuning op during the forward pass.
+                        # So we need to make an extra call to consume the previous rank's pp_send to guarantee that the previous rank's pp_send is released.
+                        AutoTuner.get().cache_pp_recv()
+                        # Send the cache after the tuning process to the next PP rank
+                        AutoTuner.get().cache_pp_send()
+                        # Clean the pp flag to avoid deadlock with synchronous send/recv
+                        AutoTuner.get().clean_pp_flag()
 
                     torch.cuda.synchronize()
 
-        logger.info(
-            f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}"
-        )
-        AutoTuner.get().print_profiling_cache()
+        if enable_flashinfer_mxfp8_autotuner:
+            if ran_forward:
+                for method in flashinfer_mxfp8_methods:
+                    method.mark_flashinfer_autotuned()
+            else:
+                forced_flashinfer = any(method.backend == "flashinfer"
+                                        for method in flashinfer_mxfp8_methods)
+                for method in flashinfer_mxfp8_methods:
+                    method.disable_flashinfer_auto()
+                if forced_flashinfer:
+                    raise RuntimeError(
+                        "FlashInfer MXFP8 was explicitly requested but its autotuner "
+                        "warmup forward could not run")
+                logger.warning(
+                    "FlashInfer MXFP8 autotuning could not run; using the native "
+                    "TensorRT-LLM GEMM backend.")
+
+        if enable_trtllm_autotuner:
+            logger.info(
+                f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}"
+            )
+            AutoTuner.get().print_profiling_cache()
 
         # Clear workspace buffers allocated during the autotuner forward pass.
         # The autotuner runs a context-only forward with max_num_tokens, which
@@ -1471,7 +1522,12 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager):
                 or self._torch_compile_piecewise_cuda_graph):
             return
 
-        self._capture_generation_cuda_graphs(resource_manager)
+        from ..modules.linear import flashinfer_mxfp8_decode_graph_capture
+
+        # 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():
+            self._capture_generation_cuda_graphs(resource_manager)
         self._capture_piecewise_cuda_graphs(resource_manager)
 
     def _capture_generation_cuda_graphs(self,