diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 8d3e0ddcac1e..e2d4364076d7 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -84,6 +84,7 @@ # and flash SDPA does not accept attn_mask. _DENSE_SDPA_BACKENDS = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH] + # --------------------------------------------------------------------------- # Config normalization helpers # --------------------------------------------------------------------------- @@ -1802,6 +1803,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 diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 53f91b6a4a06..052c065ff2df 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -4,6 +4,9 @@ import math import os from abc import ABC, abstractmethod +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from typing import ClassVar, Dict, List, Optional, Union @@ -3033,28 +3036,121 @@ 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() -> Iterator[None]: + """Tune FlashInfer MXFP8 tactics while enabling auto-dispatched calls.""" + from flashinfer 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() -> Iterator[None]: + """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). _SF_ROW_PAD = 128 _SF_COL_PAD = 4 - def __init__(self): + def __init__(self) -> None: 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 autotune, mm_mxfp8 + if not callable(autotune): + raise ImportError("flashinfer.autotune is unavailable") + 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: @@ -3100,15 +3196,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: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 12e98c96a5d4..75293c2f53a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1574,10 +1574,33 @@ def _release_megamoe_profiling_scratch(): def _run_autotuner_warmup(self, resource_manager: ResourceManager): """Runs forward passes 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, @@ -1593,8 +1616,15 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): warmup_configs.append((1 + self.max_total_draft_tokens, 1)) cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) - with self.no_cuda_graph(), autotune(cache_path=cache_path): - ran_forward = False + 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: for num_tokens, num_gen_requests in warmup_configs: warmup_request = self._create_warmup_request( resource_manager, num_tokens, num_gen_requests) @@ -1620,7 +1650,7 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): torch.cuda.synchronize() ran_forward = True - if ran_forward: + if ran_forward and 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() @@ -1629,10 +1659,28 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): # Clean the pp flag to avoid deadlock with synchronous send/recv AutoTuner.get().clean_pp_flag() - 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() self._release_megamoe_profiling_scratch() @@ -1888,7 +1936,25 @@ 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 (MXFP8LinearMethod, + flashinfer_mxfp8_autotune, + flashinfer_mxfp8_decode_graph_capture) + + # The automatic MiniMax-M3 MXFP8 selection is decode-graph-only. + # Tune every generation graph shape during the warmup-only pass. Keep + # piecewise context/prefill graph capture on the native backend. + flashinfer_methods = [ + quant_method for module in self.model.modules() + if isinstance((quant_method := getattr(module, "quant_method", None) + ), MXFP8LinearMethod) + and quant_method.needs_flashinfer_autotune + ] + flashinfer_autotune_context = ( + flashinfer_mxfp8_autotune() if self.cuda_graph_runner.is_warmup_only + and flashinfer_methods else contextlib.nullcontext()) + with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture( + ): + 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. if not self.cuda_graph_runner.is_warmup_only: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 32c237390afa..e2702762cb82 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -104,6 +104,7 @@ l0_b200: - unittest/_torch/modules/test_fp4_num_tokens_slice.py - unittest/_torch/modules/test_awq_quantization.py - unittest/_torch/modules/test_triton_linear.py + - unittest/_torch/modules/test_mxfp8_linear.py - unittest/_torch/modules/test_group_rmn_norm.py - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 95938a31cd98..b52eebd0897c 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -28,6 +28,7 @@ l0_b300: - unittest/_torch/modules/test_fused_activation_quant.py - unittest/_torch/modules/test_awq_quantization.py - unittest/_torch/modules/test_triton_linear.py + - unittest/_torch/modules/test_mxfp8_linear.py - unittest/_torch/modules/test_group_rmn_norm.py - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba 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 9724ad3d0a5e..9498a847b710 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -8,14 +8,17 @@ """ import contextlib +import sys import unittest from dataclasses import dataclass -from unittest.mock import patch +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch import torch import tensorrt_llm from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine from tensorrt_llm._torch.pyexecutor.resource_manager import ( KVCacheManager, @@ -213,6 +216,86 @@ 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.""" + calls = [] + + @contextlib.contextmanager + def flashinfer_autotune(): + calls.append("flashinfer_autotune_enter") + yield + calls.append("flashinfer_autotune_exit") + + flashinfer_module = ModuleType("flashinfer") + flashinfer_module.mm_mxfp8 = Mock() + flashinfer_module.autotune = Mock(side_effect=flashinfer_autotune) + + with ( + patch.dict( + sys.modules, + { + "flashinfer": flashinfer_module, + }, + ), + patch( + "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", + return_value=True, + ), + patch.dict("os.environ", {}, clear=True), + ): + method = MXFP8LinearMethod() + self.assertEqual(method.backend, "trtllm") + + engine = SimpleNamespace( + llm_args=SimpleNamespace(enable_autotuner=False), + cuda_graph_runner=SimpleNamespace(enabled=True), + model=SimpleNamespace( + modules=lambda: [ + SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), + SimpleNamespace(quant_method=method), + ] + ), + kv_cache_manager_key="kv_cache", + max_num_tokens=16, + batch_size=16, + max_seq_len=2, + original_max_draft_len=0, + max_total_draft_tokens=0, + mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False), + is_draft_model=False, + guided_decoder=None, + no_cuda_graph=lambda: contextlib.nullcontext(), + _create_warmup_request=Mock(return_value=object()), + _release_batch_context=Mock(return_value=contextlib.nullcontext(object())), + _assert_all_tp_ranks_have_warmup_batch=Mock(), + _release_megamoe_profiling_scratch=Mock(), + forward=Mock(side_effect=lambda *args, **kwargs: calls.append("forward")), + ) + 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) + ) + + with ( + 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", + "forward", + "flashinfer_autotune_exit", + ], + ) + self.assertEqual(method.backend, "auto") + self.assertTrue(method._flashinfer_autotuned) + flashinfer_module.autotune.assert_called_once_with() + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index 19fb731effa2..fa7bb426d513 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -13,10 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +import sys +from types import SimpleNamespace +from unittest.mock import Mock + import pytest import torch -from tensorrt_llm._torch.modules.linear import Linear, MXFP8LinearMethod, get_quant_method +import tensorrt_llm._torch.modules.linear as linear_module +from tensorrt_llm._torch.modules.linear import ( + Linear, + MXFP8LinearMethod, + flashinfer_mxfp8_autotune, + flashinfer_mxfp8_decode_graph_capture, + get_quant_method, +) from tensorrt_llm._torch.modules.mxfp8_utils import dequant_mxfp8_weight, quant_bf16_to_mxfp8 from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -42,14 +53,103 @@ def test_quant_dequant_roundtrip_is_close(): assert rel < 0.1, f"relative error too high: {rel}" -def test_mxfp8_dispatch_returns_mxfp8_method(): +def test_mxfp8_dispatch_returns_mxfp8_method(monkeypatch): """get_quant_method must dispatch QuantAlgo.MXFP8 to MXFP8LinearMethod. This is a pure dispatch check; no CUDA required. """ + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) qc = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) method = get_quant_method(qc) assert isinstance(method, MXFP8LinearMethod) + assert method.backend == "trtllm" + + +def _mock_mxfp8_ops(monkeypatch): + quantized = torch.empty((2, 4), dtype=torch.float8_e4m3fn) + activation_scale = torch.empty(512, dtype=torch.uint8) + quantize = Mock(return_value=(quantized, activation_scale)) + native_output = torch.empty((2, 3), dtype=torch.bfloat16) + native_gemm = Mock(return_value=native_output) + fake_trtllm_ops = SimpleNamespace(mxfp8_quantize=quantize, mxfp8_mxfp8_gemm=native_gemm) + monkeypatch.setattr(linear_module.torch, "ops", SimpleNamespace(trtllm=fake_trtllm_ops)) + return quantized, activation_scale, quantize, native_gemm, native_output + + +def test_mxfp8_flashinfer_call_contract(monkeypatch): + """The forced backend reuses TRT tensors and a zero-copy weight transpose.""" + monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", "flashinfer") + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + + expected = torch.empty((2, 3), dtype=torch.bfloat16) + mm_mxfp8 = Mock(return_value=expected) + monkeypatch.setitem( + sys.modules, "flashinfer", SimpleNamespace(mm_mxfp8=mm_mxfp8, autotune=Mock()) + ) + quantized, activation_scale, quantize, _, _ = _mock_mxfp8_ops(monkeypatch) + + weight = torch.empty((3, 4), dtype=torch.float8_e4m3fn) + weight_scale = torch.empty(512, dtype=torch.uint8) + module = SimpleNamespace(weight=weight, weight_scale=weight_scale, dtype=torch.bfloat16) + activation = torch.randn((2, 4), dtype=torch.bfloat16) + + method = MXFP8LinearMethod() + output = method.apply(module, activation, bias=None) + + assert output is expected + quantize.assert_called_once_with(activation, True) + args = mm_mxfp8.call_args.args + kwargs = mm_mxfp8.call_args.kwargs + assert args[0] is quantized + assert args[1].shape == (4, 3) + assert args[1].data_ptr() == weight.data_ptr() + assert args[2] is activation_scale + assert args[3] is weight_scale + assert kwargs == { + "out_dtype": torch.bfloat16, + "use_8x4_sf_layout": False, + "backend": "cutlass", + } + + +def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + + flashinfer_output = torch.empty((2, 3), dtype=torch.bfloat16) + mm_mxfp8 = Mock(return_value=flashinfer_output) + monkeypatch.setitem( + sys.modules, "flashinfer", SimpleNamespace(mm_mxfp8=mm_mxfp8, autotune=Mock()) + ) + _, _, _, native_gemm, native_output = _mock_mxfp8_ops(monkeypatch) + + 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() + assert method.enable_flashinfer_auto() + + assert method.apply(module, activation, bias=None) is native_output + native_gemm.assert_called_once() + mm_mxfp8.assert_not_called() + + method.mark_flashinfer_autotuned() + with flashinfer_mxfp8_decode_graph_capture(): + assert method.apply(module, activation, bias=None) is flashinfer_output + mm_mxfp8.assert_called_once() + + # 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_rejects_unknown_backend(monkeypatch): + monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", "unknown") + with pytest.raises(ValueError, match="TRTLLM_MXFP8_GEMM_BACKEND"): + MXFP8LinearMethod() @pytest.mark.skipif(not torch.cuda.is_available(), reason="MXFP8 Linear load path requires CUDA") @@ -110,3 +210,67 @@ def test_mxfp8_linear_cutlass_matches_reference(): ref = (x.float() @ w_deq.t()).to(torch.bfloat16) rel = (got.float() - ref.float()).norm() / ref.float().norm().clamp_min(1e-6) assert rel < 0.05, f"CUTLASS vs reference rel err {rel}" + + +@pytest.mark.skipif( + not _mxfp8_cutlass_op_available(), reason="MXFP8xMXFP8 GEMM op not compiled or sm < 100" +) +@pytest.mark.parametrize("batch_size", (1, 8, 16, 32)) +def test_mxfp8_flashinfer_decode_graph_matches_native(monkeypatch, batch_size): + """FlashInfer must consume TRT-LLM's swizzled scales like the native op. + + Tune a large-M warmup shape, then replay several decode graph shapes. This + protects the decode-only path from a silent scale-layout or tactic-cache + miss during graph capture. + """ + try: + import flashinfer # noqa: F401 + except ImportError: + pytest.skip("FlashInfer is not installed") + + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + torch.manual_seed(0) + out_f, in_f = 256, 512 + weight = torch.randn(out_f, in_f, dtype=torch.bfloat16) + weight_e4m3, weight_scale = quant_bf16_to_mxfp8(weight, 32) + warmup_x = torch.randn(128, in_f, dtype=torch.bfloat16, device="cuda") + x = torch.randn(batch_size, in_f, dtype=torch.bfloat16, device="cuda") + quant_config = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) + + native = Linear( + in_features=in_f, + out_features=out_f, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + ).cuda() + flashinfer = Linear( + in_features=in_f, + out_features=out_f, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + ).cuda() + weights = [{"weight": weight_e4m3, "weight_scale_inv": weight_scale}] + native.load_weights(weights) + flashinfer.load_weights(weights) + native_output = native(x) + method = flashinfer.quant_method + assert isinstance(method, MXFP8LinearMethod) + assert method.enable_flashinfer_auto() + with flashinfer_mxfp8_autotune(): + warmup_output = flashinfer(warmup_x) + method.mark_flashinfer_autotuned() + torch.testing.assert_close(warmup_output, native(warmup_x), rtol=2e-2, atol=2e-2) + + flashinfer_gemm = Mock(wraps=method._flashinfer_mxfp8) + method._flashinfer_mxfp8 = flashinfer_gemm + static_x = x.clone() + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + with flashinfer_mxfp8_decode_graph_capture(): + graph_output = flashinfer(static_x) + assert flashinfer_gemm.call_count == 1 + graph.replay() + torch.testing.assert_close(graph_output, native_output, rtol=2e-2, atol=2e-2)