From 554b88ef791280af515fdc3b3752250b098784c3 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 22 Jul 2026 10:27:45 -0700 Subject: [PATCH 1/6] [None][perf] Use FlashInfer MXFP8 GEMM for MiniMax-M3 decode Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3.py | 5 + tensorrt_llm/_torch/modules/linear.py | 147 ++++++++++++++-- .../_torch/pyexecutor/model_engine.py | 101 +++++++++-- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_b300.yml | 1 + .../test_pytorch_model_engine_warmup.py | 78 ++++++++- .../_torch/modules/test_mxfp8_linear.py | 164 +++++++++++++++++- 7 files changed, 463 insertions(+), 34 deletions(-) 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 744661d415e4..60c5c4a9cf2a 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 e71c09e2ca6a..04a142f20c63 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1555,10 +1555,33 @@ def _release_megamoe_profiling_scratch(): 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, @@ -1568,7 +1591,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, @@ -1590,21 +1621,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() self._release_megamoe_profiling_scratch() @@ -1860,7 +1911,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 bcacc3e90bbc..074ef1c18ce0 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 4e7347edd644..101cf51af250 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..5378ec1dcf14 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,79 @@ 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, + mapping=SimpleNamespace(tp_size=1), + is_draft_model=False, + 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", "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..a247ffd6e407 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,99 @@ 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)) + 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)) + _, _, _, 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 +206,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) From 3c1a752d33d7ca1c7a72374672ed127f686bf365 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 13:57:20 -0700 Subject: [PATCH 2/6] [None][perf] Autotune large-M MXFP8 GEMM tactics Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/thop/mxfp8Gemm.cpp | 218 ++++++++++++++++-- .../_torch/custom_ops/torch_custom_ops.py | 118 ++++++++++ tensorrt_llm/_torch/modules/linear.py | 23 +- .../_torch/pyexecutor/model_engine.py | 16 ++ .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_b300.yml | 1 + .../test_pytorch_model_engine_warmup.py | 107 +++++++++ .../_torch/modules/test_mxfp8_linear.py | 140 ++++++++++- .../thop/parallel/test_mxfp8_mxfp8_gemm.py | 154 +++++++++++++ 9 files changed, 748 insertions(+), 30 deletions(-) create mode 100644 tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py diff --git a/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp b/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp index c6b9d8f20902..4a2594a0f7e1 100644 --- a/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp +++ b/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp @@ -25,6 +25,11 @@ #include #include +#include +#include +#include +#include +#include #include namespace tkc = tensorrt_llm::cutlass_extensions; @@ -39,11 +44,103 @@ namespace torch_ext namespace { +constexpr int64_t kMxfp8LargeMMin = 6553; +constexpr int64_t kMxfp8M8kBucket = 8192; +constexpr int64_t kMxfp8M16kMin = 13106; +constexpr int64_t kMxfp8M16kBucket = 16384; +constexpr int64_t kMxfp8M32kMin = 19659; +constexpr int64_t kMxfp8M32kBucket = 32768; +constexpr int64_t kMxfp8TacticCacheMiss = -2; + +int getMxfp8SmVersion() +{ + // PyExecutor binds one GPU architecture per rank. + static int const smVersion = tensorrt_llm::common::getSMVersion(); + return smVersion; +} + +int64_t getMxfp8TuningBucket(int64_t const m) +{ + if (m < kMxfp8LargeMMin) + { + return m; + } + if (m <= kMxfp8M8kBucket) + { + return kMxfp8M8kBucket; + } + if (m >= kMxfp8M16kMin && m <= kMxfp8M16kBucket) + { + return kMxfp8M16kBucket; + } + if (m >= kMxfp8M32kMin && m <= kMxfp8M32kBucket) + { + return kMxfp8M32kBucket; + } + return m; +} + +using Mxfp8TacticCacheKey = std::tuple; + +struct Mxfp8TacticCacheEntry +{ + tkc::CutlassGemmConfig config; + int64_t tactic; +}; + +using Mxfp8TacticCache = std::map; + +Mxfp8TacticCache& getMxfp8TacticCache() +{ + static Mxfp8TacticCache cache; + return cache; +} + +std::shared_mutex& getMxfp8TacticCacheMutex() +{ + static std::shared_mutex mutex; + return mutex; +} + +Mxfp8TacticCacheKey makeMxfp8TacticCacheKey( + int64_t const m, int64_t const n, int64_t const k, at::ScalarType const outputDtype) +{ + return {getMxfp8SmVersion(), outputDtype, getMxfp8TuningBucket(m), n, k}; +} + +std::optional findMxfp8TacticCacheEntry( + int64_t const m, int64_t const n, int64_t const k, at::ScalarType const outputDtype) +{ + std::shared_lock lock(getMxfp8TacticCacheMutex()); + auto const& cache = getMxfp8TacticCache(); + auto const iterator = cache.find(makeMxfp8TacticCacheKey(m, n, k, outputDtype)); + if (iterator == cache.end()) + { + return std::nullopt; + } + return iterator->second; +} + +void cacheMxfp8Tactic(int64_t const m, int64_t const n, int64_t const k, at::ScalarType const outputDtype, + tkc::CutlassGemmConfig const& config, int64_t const tactic) +{ + std::unique_lock lock(getMxfp8TacticCacheMutex()); + getMxfp8TacticCache().insert_or_assign( + makeMxfp8TacticCacheKey(m, n, k, outputDtype), Mxfp8TacticCacheEntry{config, tactic}); +} + +void clearMxfp8CachedTactics() +{ + std::unique_lock lock(getMxfp8TacticCacheMutex()); + getMxfp8TacticCache().clear(); +} + tkc::CutlassGemmConfig getDefaultMxfp8GemmConfig() { // Reuse the same default tile/cluster as MXFP8xMXFP4 -- the B operand is // 2x wider in MXFP8xMXFP8, but the same 4x4 cluster/256x256 tile shape is - // a reasonable starting point on B200. + // a reasonable Blackwell fallback before startup tuning populates the + // native tactic cache. return tkc::CutlassGemmConfig(tkc::CutlassTileConfigSM100::CtaShape128x256x256B, tkc::MainloopScheduleType::AUTO, tkc::EpilogueScheduleType::AUTO, tkc::ClusterShape::ClusterShape_4x4x1); } @@ -63,25 +160,9 @@ void runMxfp8Gemm(at::Tensor& out, at::Tensor const& act, at::Tensor const& weig reinterpret_cast(workspace.data_ptr()), wsBytes, at::cuda::getCurrentCUDAStream(act.get_device())); } -} // namespace - -// MXFP8 (e4m3 + UE8M0 1x32 block scales) x MXFP8 (e4m3 + UE8M0 1x32 block -// scales) GEMM on Blackwell sm_100/103. -// -// Operands (matching the CUTLASS block-scaled tensor-op convention): -// act: [M, K] Float8_e4m3fn, row-major. -// actScale: 1D uint8 (UE8M0), swizzled layout produced by -// torch.ops.trtllm.mxfp8_quantize(input, swizzedLayout=True). -// weight: [N, K] Float8_e4m3fn, expected to be column-major in memory. -// The caller is responsible for ensuring the weight tensor is -// contiguous in the column-major sense that CUTLASS expects. -// weightScale: 1D uint8 (UE8M0), swizzled layout produced by -// torch.ops.trtllm.block_scale_interleave(scale). -// globalScale: [1] float -- alpha multiplier baked into the epilogue. -// For pure MXFP8xMXFP8 this is usually [1.0]. -// out_dtype: fp16 / bf16 / fp32 output element type. -at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight, - at::Tensor const& weightScale, at::Tensor const& globalScale, std::optional out_dtype) +at::Tensor mxfp8Mxfp8GemmImpl(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight, + at::Tensor const& weightScale, at::Tensor const& globalScale, std::optional outDtype, + tkc::CutlassGemmConfig const* gemmConfig, bool const useTacticCache) { CHECK_INPUT(act, torch::kFloat8_e4m3fn); CHECK_INPUT(weight, torch::kFloat8_e4m3fn); @@ -105,14 +186,17 @@ at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, a constexpr int kAlignmentN = 32; TORCH_CHECK(n % kAlignmentN == 0, "N (", n, ") must be divisible by ", kAlignmentN); - auto chosen_dtype = out_dtype.value_or(torch::kBFloat16); - TORCH_CHECK(chosen_dtype == torch::kFloat || chosen_dtype == torch::kHalf || chosen_dtype == torch::kBFloat16, + auto const chosenDtype = outDtype.value_or(torch::kBFloat16); + TORCH_CHECK(chosenDtype == torch::kFloat || chosenDtype == torch::kHalf || chosenDtype == torch::kBFloat16, "out_dtype must be one of fp16/bf16/fp32 (default bf16)."); - at::Tensor out = at::detail::empty_cuda({m, n}, chosen_dtype, act.device(), std::nullopt); + at::Tensor out = at::detail::empty_cuda({m, n}, chosenDtype, act.device(), std::nullopt); - auto const config = getDefaultMxfp8GemmConfig(); - switch (chosen_dtype) + auto const cachedEntry = useTacticCache ? findMxfp8TacticCacheEntry(m, n, k, chosenDtype) : std::nullopt; + auto const config = gemmConfig != nullptr + ? *gemmConfig + : (cachedEntry.has_value() ? cachedEntry->config : getDefaultMxfp8GemmConfig()); + switch (chosenDtype) { case at::ScalarType::Half: runMxfp8Gemm(out, act, weight, actScale, weightScale, globalScale, m, n, k, config); @@ -132,12 +216,96 @@ at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, a return out; } +} // namespace + +// MXFP8 (e4m3 + UE8M0 1x32 block scales) x MXFP8 (e4m3 + UE8M0 1x32 block +// scales) GEMM on Blackwell sm_100/103. +// +// Operands (matching the CUTLASS block-scaled tensor-op convention): +// act: [M, K] Float8_e4m3fn, row-major. +// actScale: 1D uint8 (UE8M0), swizzled layout produced by +// torch.ops.trtllm.mxfp8_quantize(input, swizzedLayout=True). +// weight: [N, K] Float8_e4m3fn, expected to be column-major in memory. +// The caller is responsible for ensuring the weight tensor is +// contiguous in the column-major sense that CUTLASS expects. +// weightScale: 1D uint8 (UE8M0), swizzled layout produced by +// torch.ops.trtllm.block_scale_interleave(scale). +// globalScale: [1] float -- alpha multiplier baked into the epilogue. +// For pure MXFP8xMXFP8 this is usually [1.0]. +// out_dtype: fp16 / bf16 / fp32 output element type. +at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight, + at::Tensor const& weightScale, at::Tensor const& globalScale, std::optional outDtype) +{ + return mxfp8Mxfp8GemmImpl(act, actScale, weight, weightScale, globalScale, outDtype, /*gemmConfig=*/nullptr, + /*useTacticCache=*/true); +} + +class MXFP8GemmRunner : public torch::CustomClassHolder +{ +public: + explicit MXFP8GemmRunner(at::ScalarType outputDtype) + : mOutputDtype(outputDtype) + { + TORCH_CHECK(outputDtype == torch::kFloat || outputDtype == torch::kHalf || outputDtype == torch::kBFloat16, + "output_dtype must be one of fp16/bf16/fp32."); + mConfigs = CutlassFp4GemmRunner{}.getConfigs(); + } + + at::Tensor runGemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight, + at::Tensor const& weightScale, at::Tensor const& globalScale, int64_t configIdx) const + { + auto const config = configIdx == -1 ? getDefaultMxfp8GemmConfig() : getConfig(configIdx); + return mxfp8Mxfp8GemmImpl( + act, actScale, weight, weightScale, globalScale, mOutputDtype, &config, /*useTacticCache=*/false); + } + + void registerTactic(int64_t const m, int64_t const n, int64_t const k, int64_t const configIdx) const + { + tkc::CutlassGemmConfig const config = configIdx == -1 ? getDefaultMxfp8GemmConfig() : getConfig(configIdx); + cacheMxfp8Tactic(m, n, k, mOutputDtype, config, configIdx); + } + + int64_t getCachedTactic(int64_t const m, int64_t const n, int64_t const k) const + { + auto const entry = findMxfp8TacticCacheEntry(m, n, k, mOutputDtype); + return entry.has_value() ? entry->tactic : kMxfp8TacticCacheMiss; + } + + void clearTacticCache() const + { + clearMxfp8CachedTactics(); + } + + int64_t getNumConfigs() const + { + return static_cast(mConfigs.size()); + } + +private: + tkc::CutlassGemmConfig const& getConfig(int64_t const configIdx) const + { + TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs()); + return mConfigs.at(configIdx); + } + + at::ScalarType mOutputDtype; + std::vector mConfigs; +}; + } // namespace torch_ext TRTLLM_NAMESPACE_END TORCH_LIBRARY_FRAGMENT(trtllm, m) { + m.class_("MXFP8GemmRunner") + .def(torch::init()) + .def("run_gemm", &tensorrt_llm::torch_ext::MXFP8GemmRunner::runGemm) + .def("get_num_configs", &tensorrt_llm::torch_ext::MXFP8GemmRunner::getNumConfigs) + .def("register_tactic", &tensorrt_llm::torch_ext::MXFP8GemmRunner::registerTactic) + .def("get_cached_tactic", &tensorrt_llm::torch_ext::MXFP8GemmRunner::getCachedTactic) + .def("clear_tactic_cache", &tensorrt_llm::torch_ext::MXFP8GemmRunner::clearTacticCache); + m.def( "mxfp8_mxfp8_gemm(Tensor act, Tensor actScale, Tensor weight, Tensor weightScale, " "Tensor globalScale, ScalarType? out_dtype=None) -> Tensor"); diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f3b0dc0476d3..ef2470519714 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -552,6 +552,124 @@ def _( return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) +_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" + + +def _map_to_mxfp8_large_m_bucket(num_tokens: int) -> int: + for (lower_bound, upper_bound), bucket in zip(_MXFP8_LARGE_M_BANDS, + _MXFP8_LARGE_M_BUCKETS): + if lower_bound <= num_tokens <= upper_bound: + return bucket + return num_tokens + + +def _get_mxfp8_large_m_tuning_buckets(max_num_tokens: int) -> tuple[int, ...]: + mapped_max = _map_to_mxfp8_large_m_bucket(max_num_tokens) + return tuple(bucket for bucket in _MXFP8_LARGE_M_BUCKETS + if bucket <= mapped_max) + + +def _mxfp8_scale_infer_shape(input_shapes: List[List[int]]) -> int: + _, scale_shape = fp4_utils.get_fp4_shape(input_shapes[0], sf_vec_size=32) + return scale_shape + + +class MXFP8GemmRunner(TunableRunner): + runner_dict = dict() + tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, _get_mxfp8_large_m_tuning_buckets, + _map_to_mxfp8_large_m_bucket), ), + constraint_specs=(ConstraintSpec( + 1, 0, _mxfp8_scale_infer_shape), ), + use_cuda_graph=False) + + def __init__(self, output_dtype: torch.dtype): + self.output_dtype = output_dtype + self.sm_version = get_sm_version() + instance_key = (output_dtype, self.sm_version) + if instance_key not in MXFP8GemmRunner.runner_dict: + MXFP8GemmRunner.runner_dict[ + instance_key] = torch.classes.trtllm.MXFP8GemmRunner( + output_dtype) + self.mxfp8_gemm_runner = MXFP8GemmRunner.runner_dict[instance_key] + + def unique_id(self): + return (self.output_dtype, self.sm_version) + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + return [-1, *range(self.mxfp8_gemm_runner.get_num_configs())] + + def sync_tactic_cache(self, tuner: AutoTuner) -> None: + runner_name = self.__class__.__name__ + unique_id = str(self.unique_id()) + cache = tuner.profiling_cache.get_specific_custom_op( + _MXFP8_AUTOTUNED_OP) + for cache_key, (_runner_id, tactic, _min_time) in cache.items(): + _, cached_runner_name, cached_unique_id, profile = cache_key + if cached_runner_name != runner_name or cached_unique_id != unique_id: + continue + m, k = profile[0] + n, weight_k = profile[2] + if k != weight_k: + raise ValueError( + f"MXFP8 autotuner cache has mismatched K dimensions: " + f"activation K={k}, weight K={weight_k}") + self.mxfp8_gemm_runner.register_tactic(m, n, k, tactic) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + ) -> torch.Tensor: + act, act_scale, weight, weight_scale, global_scale = inputs + return self.mxfp8_gemm_runner.run_gemm( + act, + act_scale, + weight, + weight_scale, + global_scale, + tactic, + ) + + +@torch.library.custom_op("trtllm::mxfp8_mxfp8_gemm_autotuned", mutates_args=()) +def mxfp8_mxfp8_gemm_autotuned( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + global_scale: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + tuner = AutoTuner.get() + runner = MXFP8GemmRunner(output_dtype) + inputs = [act, act_scale, weight, weight_scale, global_scale] + _, best_tactic = tuner.choose_one( + _MXFP8_AUTOTUNED_OP, + [runner], + MXFP8GemmRunner.tuning_config, + inputs, + ) + if tuner.is_tuning_mode: + runner.sync_tactic_cache(tuner) + return runner(inputs=inputs, tactic=best_tactic) + + +@mxfp8_mxfp8_gemm_autotuned.register_fake +def _( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + global_scale: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) + + 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 60c5c4a9cf2a..6c2307b4625c 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3082,6 +3082,10 @@ class MXFP8LinearMethod(LinearMethodBase): ``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. + 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 + CUTLASS configuration remains the cache-miss fallback. """ BLOCK_SIZE = 32 # Swizzled-SF layout padding (matches W4A8MXFP4FP8: rows->128, cols/SFblock->4). @@ -3092,6 +3096,8 @@ def __init__(self) -> None: super().__init__() self.use_cutlass = _mxfp8_cutlass_op_available() self.backend = os.environ.get("TRTLLM_MXFP8_GEMM_BACKEND", "trtllm") + self.use_native_autotuner = True + self._native_autotuned = 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}") @@ -3111,6 +3117,11 @@ def uses_flashinfer(self) -> bool: def needs_flashinfer_autotune(self) -> bool: return self.uses_flashinfer and self._flashinfer_mxfp8 is not None + @property + def needs_native_autotune(self) -> bool: + return (self.use_native_autotuner and not self._native_autotuned + and self.use_cutlass and self.backend == "trtllm") + def _load_flashinfer(self, *, required: bool) -> bool: if not self.use_cutlass: if required: @@ -3147,6 +3158,13 @@ def enable_flashinfer_auto(self) -> bool: def mark_flashinfer_autotuned(self) -> None: self._flashinfer_autotuned = True + def mark_native_autotuned(self) -> None: + self._native_autotuned = True + + def disable_native_autotune(self) -> None: + self.use_native_autotuner = False + self._native_autotuned = False + def disable_flashinfer_auto(self) -> None: if self.backend == "auto": self.backend = "trtllm" @@ -3218,7 +3236,10 @@ def apply(self, module: Linear, input: torch.Tensor, global_scale = torch.ones([1], dtype=torch.float32, device=input.device) - output = torch.ops.trtllm.mxfp8_mxfp8_gemm( + gemm = (torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if self.needs_native_autotune else + torch.ops.trtllm.mxfp8_mxfp8_gemm) + output = gemm( act_e4m3, act_sf, module.weight, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 04a142f20c63..bbc6c3f8707a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1565,6 +1565,7 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) for module in self.model.modules())) 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): @@ -1573,7 +1574,12 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): quant_method.enable_flashinfer_auto() if quant_method.needs_flashinfer_autotune: flashinfer_mxfp8_methods.append(quant_method) + if enable_trtllm_autotuner and quant_method.needs_native_autotune: + native_mxfp8_methods.append(quant_method) + elif not enable_trtllm_autotuner: + quant_method.disable_native_autotune() enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods) + enable_native_mxfp8_autotuner = bool(native_mxfp8_methods) if not enable_trtllm_autotuner and not enable_flashinfer_mxfp8_autotuner: return @@ -1581,6 +1587,7 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): 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}, " f"FlashInfer MXFP8={enable_flashinfer_mxfp8_autotuner})...") kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) @@ -1651,6 +1658,15 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): "FlashInfer MXFP8 autotuning could not run; using the native " "TensorRT-LLM GEMM backend.") + if enable_native_mxfp8_autotuner: + if ran_forward: + for method in native_mxfp8_methods: + method.mark_native_autotuned() + else: + logger.warning( + "Native MXFP8 autotuning had no runnable warmup batch; " + "leaving tuning pending for a later warmup.") + if enable_trtllm_autotuner: logger.info( f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 074ef1c18ce0..234b8cf5a946 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -109,6 +109,7 @@ l0_b200: - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules + - unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 101cf51af250..eae48e24cfd8 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -33,6 +33,7 @@ l0_b300: - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules + - unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py 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 5378ec1dcf14..212f5787953d 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -289,6 +289,113 @@ def flashinfer_autotune(): self.assertTrue(method._flashinfer_autotuned) flashinfer_module.autotune.assert_called_once_with() + def test_native_mxfp8_retries_after_missing_warmup_batch(self): + """MXFP8 tuning remains pending until a warmup forward can run.""" + calls = [] + + @contextlib.contextmanager + def trtllm_autotune(**kwargs): + self.assertIsNone(kwargs["cache_path"]) + calls.append("autotune_enter") + yield + calls.append("autotune_exit") + + tuner = SimpleNamespace( + setup_distributed_state=Mock(), + cache_pp_recv=Mock(), + cache_pp_send=Mock(), + clean_pp_flag=Mock(), + profiling_cache={}, + print_profiling_cache=Mock(), + ) + + with patch( + "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", + return_value=True, + ): + method = MXFP8LinearMethod() + self.assertTrue(method.needs_native_autotune) + + engine = SimpleNamespace( + llm_args=SimpleNamespace(enable_autotuner=True), + cuda_graph_runner=SimpleNamespace(enabled=False), + model=SimpleNamespace(modules=lambda: [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, + mapping=SimpleNamespace(tp_size=1), + dist=object(), + is_draft_model=False, + no_cuda_graph=lambda: contextlib.nullcontext(), + _create_warmup_request=Mock(return_value=object()), + _release_batch_context=Mock( + side_effect=[ + contextlib.nullcontext(None), + contextlib.nullcontext(object()), + ] + ), + _assert_all_tp_ranks_have_warmup_batch=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( + "tensorrt_llm._torch.pyexecutor.model_engine.AutoTuner.get", + return_value=tuner, + ), + patch( + "tensorrt_llm._torch.pyexecutor.model_engine.autotune", + side_effect=trtllm_autotune, + ), + 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, ["autotune_enter", "autotune_exit"]) + self.assertFalse(method._native_autotuned) + self.assertTrue(method.needs_native_autotune) + + PyTorchModelEngine._run_autotuner_warmup(engine, resource_manager) + + self.assertEqual( + calls, + [ + "autotune_enter", + "autotune_exit", + "autotune_enter", + "forward", + "autotune_exit", + ], + ) + self.assertTrue(method._native_autotuned) + self.assertFalse(method.needs_native_autotune) + self.assertEqual(tuner.setup_distributed_state.call_count, 2) + tuner.setup_distributed_state.assert_called_with(engine.mapping, engine.dist) + + def test_native_mxfp8_respects_disabled_global_autotuner(self): + with patch( + "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", + return_value=True, + ): + method = MXFP8LinearMethod() + engine = SimpleNamespace( + llm_args=SimpleNamespace(enable_autotuner=False), + cuda_graph_runner=SimpleNamespace(enabled=False), + model=SimpleNamespace(modules=lambda: [SimpleNamespace(quant_method=method)]), + ) + + PyTorchModelEngine._run_autotuner_warmup(engine, Mock()) + + self.assertFalse(method.use_native_autotuner) + self.assertFalse(method.needs_native_autotune) + 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 a247ffd6e407..c0e5e3714d27 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -21,6 +21,12 @@ import torch import tensorrt_llm._torch.modules.linear as linear_module +from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + MXFP8GemmRunner, + _get_mxfp8_large_m_tuning_buckets, + _map_to_mxfp8_large_m_bucket, +) from tensorrt_llm._torch.modules.linear import ( Linear, MXFP8LinearMethod, @@ -63,6 +69,7 @@ def test_mxfp8_dispatch_returns_mxfp8_method(monkeypatch): method = get_quant_method(qc) assert isinstance(method, MXFP8LinearMethod) assert method.backend == "trtllm" + assert method.use_native_autotuner def _mock_mxfp8_ops(monkeypatch): @@ -71,9 +78,23 @@ def _mock_mxfp8_ops(monkeypatch): 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) + autotuned_output = torch.empty((2, 3), dtype=torch.bfloat16) + autotuned_gemm = Mock(return_value=autotuned_output) + fake_trtllm_ops = SimpleNamespace( + mxfp8_quantize=quantize, + mxfp8_mxfp8_gemm=native_gemm, + mxfp8_mxfp8_gemm_autotuned=autotuned_gemm, + ) monkeypatch.setattr(linear_module.torch, "ops", SimpleNamespace(trtllm=fake_trtllm_ops)) - return quantized, activation_scale, quantize, native_gemm, native_output + return ( + quantized, + activation_scale, + quantize, + native_gemm, + native_output, + autotuned_gemm, + autotuned_output, + ) def test_mxfp8_flashinfer_call_contract(monkeypatch): @@ -84,7 +105,7 @@ def test_mxfp8_flashinfer_call_contract(monkeypatch): expected = torch.empty((2, 3), dtype=torch.bfloat16) mm_mxfp8 = Mock(return_value=expected) monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(mm_mxfp8=mm_mxfp8)) - quantized, activation_scale, quantize, _, _ = _mock_mxfp8_ops(monkeypatch) + 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) @@ -117,7 +138,7 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): 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)) - _, _, _, native_gemm, native_output = _mock_mxfp8_ops(monkeypatch) + _, _, _, native_gemm, native_output, _, _ = _mock_mxfp8_ops(monkeypatch) module = SimpleNamespace( weight=torch.empty((3, 4), dtype=torch.float8_e4m3fn), @@ -142,6 +163,117 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): assert native_gemm.call_count == 2 +@pytest.mark.parametrize( + "num_tokens,expected", + [ + (1, 1), + (6552, 6552), + (6553, 8192), + (8192, 8192), + (8193, 8193), + (13105, 13105), + (13106, 16384), + (16384, 16384), + (16385, 16385), + (19658, 19658), + (19659, 32768), + (32768, 32768), + (32769, 32769), + ], +) +def test_mxfp8_large_m_bucket_mapping(num_tokens, expected): + assert _map_to_mxfp8_large_m_bucket(num_tokens) == expected + + +@pytest.mark.parametrize( + "max_num_tokens,expected", + [ + (4096, ()), + (6599, (8192,)), + (14906, (8192, 16384)), + (29765, (8192, 16384, 32768)), + ], +) +def test_mxfp8_large_m_tuning_buckets(max_num_tokens, expected): + assert _get_mxfp8_large_m_tuning_buckets(max_num_tokens) == expected + + +def test_mxfp8_large_m_cache_profile_maps_act_and_constrains_scale(): + AutoTuner._find_nearest_profile.cache_clear() + input_shapes = ( + torch.Size((6599, 6144)), + torch.Size((1277952,)), + torch.Size((9216, 6144)), + torch.Size((1769472,)), + torch.Size((1,)), + ) + profile = AutoTuner._find_nearest_profile( + input_shapes, + MXFP8GemmRunner.tuning_config.dynamic_tensor_specs, + MXFP8GemmRunner.tuning_config.constraint_specs, + ) + assert profile == ( + (8192, 6144), + (-1,), + (9216, 6144), + (1769472,), + (1,), + ) + + +def test_mxfp8_native_autotuner_dispatch(monkeypatch): + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + _, _, _, native_gemm, native_output, autotuned_gemm, autotuned_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.use_native_autotuner + assert method.needs_native_autotune + assert method.apply(module, activation, bias=None) is autotuned_output + autotuned_gemm.assert_called_once() + native_gemm.assert_not_called() + + method.mark_native_autotuned() + assert not method.needs_native_autotune + assert method.apply(module, activation, bias=None) is native_output + autotuned_gemm.assert_called_once() + native_gemm.assert_called_once() + + +def test_mxfp8_native_autotuner_syncs_profiles(): + runner = object.__new__(MXFP8GemmRunner) + runner.output_dtype = torch.bfloat16 + runner.sm_version = 100 + runner.mxfp8_gemm_runner = Mock() + profile = ( + (8192, 6144), + (-1,), + (9216, 6144), + (1769472,), + (1,), + ) + cache_key = ( + "trtllm::mxfp8_mxfp8_gemm_autotuned::gemm", + "MXFP8GemmRunner", + str(runner.unique_id()), + profile, + ) + profiling_cache = Mock() + profiling_cache.get_specific_custom_op.return_value = {cache_key: (0, 17, 0.25)} + + runner.sync_tactic_cache(SimpleNamespace(profiling_cache=profiling_cache)) + + runner.mxfp8_gemm_runner.register_tactic.assert_called_once_with(8192, 9216, 6144, 17) + + def test_mxfp8_rejects_unknown_backend(monkeypatch): monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", "unknown") with pytest.raises(ValueError, match="TRTLLM_MXFP8_GEMM_BACKEND"): diff --git a/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py b/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py new file mode 100644 index 000000000000..dae039f9ead3 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F +from utils.util import getSMVersion + +import tensorrt_llm._torch.custom_ops.torch_custom_ops # noqa: F401 + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="MXFP8 GEMM is supported on SM100 and SM103 only. Current SM is %d." % getSMVersion(), +) +@pytest.mark.parametrize( + "m,n,k", + [ + (6599, 9216, 6144), + (6599, 6144, 3072), + (14906, 6144, 8192), + (14906, 6144, 6144), + (29765, 24576, 6144), + (29765, 6144, 12288), + (8193, 9216, 6144), + ], +) +def test_mxfp8_mxfp8_gemm_large_m(m: int, n: int, k: int): + """The generic fallback agrees with BF16 GEMM for representative shapes.""" + torch.manual_seed(42) + mat_a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + mat_b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + + fp8_a, a_block_sf = torch.ops.trtllm.mxfp8_quantize(mat_a, True) + fp8_b, b_block_sf = torch.ops.trtllm.mxfp8_quantize(mat_b, True) + global_scale = torch.ones((1,), device="cuda", dtype=torch.float32) + + output = torch.ops.trtllm.mxfp8_mxfp8_gemm( + fp8_a, + a_block_sf, + fp8_b, + b_block_sf, + global_scale, + torch.bfloat16, + ) + output_ref = mat_a @ mat_b.t() + + assert F.cosine_similarity(output.flatten(), output_ref.flatten(), dim=0).item() > 0.98 + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="MXFP8 tactic runner requires SM100 or SM103. Current SM is %d." % getSMVersion(), +) +def test_mxfp8_mxfp8_runner_tactics(): + """Every exposed tactic and the generic fallback produce aligned output.""" + torch.manual_seed(42) + m, n, k = 128, 256, 512 + mat_a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + mat_b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + fp8_a, a_block_sf = torch.ops.trtllm.mxfp8_quantize(mat_a, True) + fp8_b, b_block_sf = torch.ops.trtllm.mxfp8_quantize(mat_b, True) + global_scale = torch.ones((1,), device="cuda", dtype=torch.float32) + output_ref = mat_a @ mat_b.t() + + runner = torch.classes.trtllm.MXFP8GemmRunner(torch.bfloat16) + expected_tactics = 20 if getSMVersion() == 100 else 10 + assert runner.get_num_configs() == expected_tactics + + for tactic in [-1, *range(expected_tactics)]: + output = runner.run_gemm( + fp8_a, + a_block_sf, + fp8_b, + b_block_sf, + global_scale, + tactic, + ) + similarity = F.cosine_similarity(output.flatten(), output_ref.flatten(), dim=0) + assert similarity.item() > 0.98 + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="MXFP8 tactic cache requires SM100 or SM103. Current SM is %d." % getSMVersion(), +) +def test_mxfp8_mxfp8_native_tactic_cache(): + """The direct op uses cached tactics and preserves the generic fallback.""" + torch.manual_seed(42) + m, n, k = 128, 256, 512 + mat_a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + mat_b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + fp8_a, a_block_sf = torch.ops.trtllm.mxfp8_quantize(mat_a, True) + fp8_b, b_block_sf = torch.ops.trtllm.mxfp8_quantize(mat_b, True) + global_scale = torch.ones((1,), device="cuda", dtype=torch.float32) + runner = torch.classes.trtllm.MXFP8GemmRunner(torch.bfloat16) + + try: + runner.clear_tactic_cache() + assert runner.get_cached_tactic(m, n, k) == -2 + + runner.register_tactic(m, n, k, 0) + assert runner.get_cached_tactic(m, n, k) == 0 + cached_output = torch.ops.trtllm.mxfp8_mxfp8_gemm( + fp8_a, + a_block_sf, + fp8_b, + b_block_sf, + global_scale, + torch.bfloat16, + ) + explicit_output = runner.run_gemm( + fp8_a, + a_block_sf, + fp8_b, + b_block_sf, + global_scale, + 0, + ) + torch.testing.assert_close(cached_output, explicit_output, rtol=0, atol=0) + + runner.register_tactic(m, n, k, -1) + assert runner.get_cached_tactic(m, n, k) == -1 + cached_fallback_output = torch.ops.trtllm.mxfp8_mxfp8_gemm( + fp8_a, + a_block_sf, + fp8_b, + b_block_sf, + global_scale, + torch.bfloat16, + ) + explicit_fallback_output = runner.run_gemm( + fp8_a, + a_block_sf, + fp8_b, + b_block_sf, + global_scale, + -1, + ) + torch.testing.assert_close(cached_fallback_output, explicit_fallback_output, rtol=0, atol=0) + finally: + runner.clear_tactic_cache() From c0318042a781bb7b6b07669661825a7bed8f407a Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:41:31 -0700 Subject: [PATCH 3/6] [None][fix] Harden MXFP8 autotuner warmup Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/thop/mxfp8Gemm.cpp | 3 +- .../_torch/custom_ops/torch_custom_ops.py | 18 +- tensorrt_llm/_torch/modules/linear.py | 2 +- .../_torch/pyexecutor/model_engine.py | 162 +++++++----- .../test_pytorch_model_engine_warmup.py | 249 +++++++++++++++--- .../_torch/modules/test_mxfp8_linear.py | 72 ++++- .../thop/parallel/test_mxfp8_mxfp8_gemm.py | 6 +- 7 files changed, 394 insertions(+), 118 deletions(-) diff --git a/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp b/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp index 4a2594a0f7e1..3b597ac6dda4 100644 --- a/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp +++ b/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp @@ -284,7 +284,8 @@ class MXFP8GemmRunner : public torch::CustomClassHolder private: tkc::CutlassGemmConfig const& getConfig(int64_t const configIdx) const { - TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs()); + TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs(), "MXFP8 config index ", configIdx, + " is out of range [0, ", getNumConfigs(), ")."); return mConfigs.at(configIdx); } diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 97138bebf552..3bd4ea11f345 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -558,6 +558,7 @@ def _( def _map_to_mxfp8_large_m_bucket(num_tokens: int) -> int: + """Map known large-M bands to stable native autotuning profiles.""" for (lower_bound, upper_bound), bucket in zip(_MXFP8_LARGE_M_BANDS, _MXFP8_LARGE_M_BUCKETS): if lower_bound <= num_tokens <= upper_bound: @@ -566,18 +567,22 @@ def _map_to_mxfp8_large_m_bucket(num_tokens: int) -> int: def _get_mxfp8_large_m_tuning_buckets(max_num_tokens: int) -> tuple[int, ...]: + """Return large-M profiles reachable by the configured token limit.""" mapped_max = _map_to_mxfp8_large_m_bucket(max_num_tokens) return tuple(bucket for bucket in _MXFP8_LARGE_M_BUCKETS if bucket <= mapped_max) def _mxfp8_scale_infer_shape(input_shapes: List[List[int]]) -> int: + """Infer the swizzled MXFP8 activation-scale storage size.""" _, scale_shape = fp4_utils.get_fp4_shape(input_shapes[0], sf_vec_size=32) return scale_shape class MXFP8GemmRunner(TunableRunner): runner_dict = dict() + synced_cache_keys: ClassVar[dict[tuple[torch.dtype, int, int], set[tuple]]] + synced_cache_keys = {} tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, _get_mxfp8_large_m_tuning_buckets, _map_to_mxfp8_large_m_bucket), ), @@ -595,19 +600,30 @@ def __init__(self, output_dtype: torch.dtype): output_dtype) self.mxfp8_gemm_runner = MXFP8GemmRunner.runner_dict[instance_key] - def unique_id(self): + def unique_id(self) -> tuple[torch.dtype, int]: + """Return the native tactic-cache identity for this runner.""" return (self.output_dtype, self.sm_version) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: + """Return the generic fallback followed by every compiled tactic.""" return [-1, *range(self.mxfp8_gemm_runner.get_num_configs())] def sync_tactic_cache(self, tuner: AutoTuner) -> None: + """Register newly profiled tactics in the native serving cache.""" runner_name = self.__class__.__name__ unique_id = str(self.unique_id()) cache = tuner.profiling_cache.get_specific_custom_op( _MXFP8_AUTOTUNED_OP) + sync_key = (*self.unique_id(), id(tuner.profiling_cache)) + synced_cache_keys = self.synced_cache_keys.setdefault(sync_key, set()) for cache_key, (_runner_id, tactic, _min_time) in cache.items(): + if cache_key in synced_cache_keys: + continue + # Each cache entry is immutable once profiled, so remember both + # matching and non-matching entries to avoid rescanning the full + # shared custom-op cache on every synchronization. + synced_cache_keys.add(cache_key) _, cached_runner_name, cached_unique_id, profile = cache_key if cached_runner_name != runner_name or cached_unique_id != unique_id: continue diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 6fa0ce1f275e..4bcb345b22ba 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3120,7 +3120,7 @@ def needs_flashinfer_autotune(self) -> bool: @property def needs_native_autotune(self) -> bool: return (self.use_native_autotuner and not self._native_autotuned - and self.use_cutlass and self.backend == "trtllm") + and self.use_cutlass) def _load_flashinfer(self, *, required: bool) -> bool: if not self.use_cutlass: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index bbc6c3f8707a..2fa89789da50 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1559,32 +1559,59 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): flashinfer_mxfp8_autotune) enable_trtllm_autotuner = self.llm_args.enable_autotuner + mxfp8_methods = [] + for module in self.model.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, MXFP8LinearMethod): + mxfp8_methods.append(quant_method) + if not enable_trtllm_autotuner: + for method in mxfp8_methods: + method.disable_native_autotune() + method.disable_flashinfer_auto() + return + + # Native and FlashInfer tuning are independent. Capture native + # eligibility before enabling graph-only FlashInfer dispatch. + native_mxfp8_methods = [ + method for method in mxfp8_methods if method.needs_native_autotune + ] 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 = [] - 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: + if use_mxfp8_flashinfer_graph_default: + for quant_method in mxfp8_methods: quant_method.enable_flashinfer_auto() - if quant_method.needs_flashinfer_autotune: - flashinfer_mxfp8_methods.append(quant_method) - if enable_trtllm_autotuner and quant_method.needs_native_autotune: - native_mxfp8_methods.append(quant_method) - elif not enable_trtllm_autotuner: - quant_method.disable_native_autotune() + flashinfer_mxfp8_methods = [ + method for method in mxfp8_methods + if method.needs_flashinfer_autotune + ] + + # Every TP rank must make the same backend decision before any rank + # returns or enters a tuning forward with TP collectives. + if self.mapping.tp_size > 1: + local_flashinfer_enabled = int(bool(flashinfer_mxfp8_methods)) + all_flashinfer_enabled = list( + self.dist.tp_allgather(local_flashinfer_enabled)) + if any(all_flashinfer_enabled) and not all(all_flashinfer_enabled): + forced_flashinfer = any(method.backend == "flashinfer" + for method in mxfp8_methods) + for method in mxfp8_methods: + method.disable_flashinfer_auto() + flashinfer_mxfp8_methods = [] + if forced_flashinfer: + raise RuntimeError( + "FlashInfer MXFP8 was explicitly requested but is not " + "available on every TP rank") + logger.warning( + "FlashInfer MXFP8 availability differs across TP ranks; " + "using the native TensorRT-LLM GEMM backend on every rank.") + enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods) enable_native_mxfp8_autotuner = bool(native_mxfp8_methods) - 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) + 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}, " @@ -1597,52 +1624,56 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): token_num_upper_bound=token_num_upper_bound, max_num_draft_tokens=self.original_max_draft_len) - cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) - 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, - resource_manager) as batch: - if batch is None and self.mapping.tp_size <= 1: - pass # Single rank, safe to skip - else: - self._assert_all_tp_ranks_have_warmup_batch( - batch, curr_max_num_tokens) - if batch is not None: - # Reset the flag is_first_draft for the draft model. - # This is necessary for overlap scheduler. - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance( - spec_resource_manager, Eagle3ResourceManager): - spec_resource_manager.is_first_draft = True + def run_autotuner_pass(autotune_context: Any, + synchronize_trtllm_cache: bool) -> bool: + """Run one isolated tuning pass with a fresh synthetic batch.""" + ran_forward = False + with self.no_cuda_graph(), autotune_context: + warmup_request = self._create_warmup_request( + resource_manager, curr_max_num_tokens, 0) + with self._release_batch_context(warmup_request, + resource_manager) as batch: + if batch is None and self.mapping.tp_size <= 1: + pass # Single rank, safe to skip + else: + self._assert_all_tp_ranks_have_warmup_batch( + batch, curr_max_num_tokens) + if batch is not None: + # Reset the flag is_first_draft for the draft model. + # This is necessary for overlap scheduler. + spec_resource_manager = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if self.is_draft_model and isinstance( + spec_resource_manager, Eagle3ResourceManager): + spec_resource_manager.is_first_draft = True - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) - ran_forward = True + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) + ran_forward = True - 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() + if synchronize_trtllm_cache: + # 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() + torch.cuda.synchronize() + return ran_forward + + cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) + ran_native_forward = run_autotuner_pass(autotune(cache_path=cache_path), + synchronize_trtllm_cache=True) + ran_flashinfer_forward = False + if enable_flashinfer_mxfp8_autotuner: + ran_flashinfer_forward = run_autotuner_pass( + flashinfer_mxfp8_autotune(), synchronize_trtllm_cache=False) if enable_flashinfer_mxfp8_autotuner: - if ran_forward: + if ran_flashinfer_forward: for method in flashinfer_mxfp8_methods: method.mark_flashinfer_autotuned() else: @@ -1659,19 +1690,20 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): "TensorRT-LLM GEMM backend.") if enable_native_mxfp8_autotuner: - if ran_forward: + if ran_native_forward: for method in native_mxfp8_methods: method.mark_native_autotuned() else: + for method in native_mxfp8_methods: + method.disable_native_autotune() logger.warning( "Native MXFP8 autotuning had no runnable warmup batch; " - "leaving tuning pending for a later warmup.") + "using the default native GEMM tactic.") - if enable_trtllm_autotuner: - logger.info( - f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}" - ) - AutoTuner.get().print_profiling_cache() + logger.info( + f"[Autotuner] Cache size after warmup is {len(AutoTuner.get().profiling_cache)}" + ) + AutoTuner.get().print_profiling_cache() self._release_megamoe_profiling_scratch() 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 212f5787953d..185d66f9d854 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -8,6 +8,7 @@ """ import contextlib +import os import sys import unittest from dataclasses import dataclass @@ -216,8 +217,8 @@ 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 test_flashinfer_mxfp8_respects_disabled_global_autotuner(self): + """The global autotuner switch also disables automatic FlashInfer tuning.""" calls = [] @contextlib.contextmanager @@ -241,8 +242,10 @@ def flashinfer_autotune(): "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", return_value=True, ), - patch.dict("os.environ", {}, clear=True), + patch.dict(os.environ, {}, clear=False), ): + os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) + os.environ.pop("TLLM_AUTOTUNER_CACHE_PATH", None) method = MXFP8LinearMethod() self.assertEqual(method.backend, "trtllm") @@ -255,16 +258,83 @@ def flashinfer_autotune(): SimpleNamespace(quant_method=method), ] ), + ) + PyTorchModelEngine._run_autotuner_warmup(engine, Mock()) + + self.assertEqual(calls, []) + self.assertEqual(method.backend, "trtllm") + self.assertFalse(method.use_native_autotuner) + self.assertFalse(method._flashinfer_autotuned) + flashinfer_module.autotune.assert_not_called() + + def test_mxfp8_native_and_flashinfer_use_separate_warmup_passes(self): + """Native and FlashInfer backends each receive an isolated tuning forward.""" + calls = [] + + @contextlib.contextmanager + def trtllm_autotune(**kwargs): + self.assertIsNone(kwargs["cache_path"]) + calls.append("trtllm_autotune_enter") + yield + calls.append("trtllm_autotune_exit") + + @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) + + tuner = SimpleNamespace( + setup_distributed_state=Mock(), + cache_pp_recv=Mock(), + cache_pp_send=Mock(), + clean_pp_flag=Mock(), + profiling_cache={}, + print_profiling_cache=Mock(), + ) + + 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=False), + ): + os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) + os.environ.pop("TLLM_AUTOTUNER_CACHE_PATH", None) + method = MXFP8LinearMethod() + self.assertTrue(method.needs_native_autotune) + + engine = SimpleNamespace( + llm_args=SimpleNamespace(enable_autotuner=True), + 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, mapping=SimpleNamespace(tp_size=1), + dist=object(), is_draft_model=False, no_cuda_graph=lambda: contextlib.nullcontext(), _create_warmup_request=Mock(return_value=object()), - _release_batch_context=Mock(return_value=contextlib.nullcontext(object())), + _release_batch_context=Mock( + side_effect=[ + contextlib.nullcontext(object()), + 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")), @@ -275,6 +345,14 @@ def flashinfer_autotune(): ) with ( + patch( + "tensorrt_llm._torch.pyexecutor.model_engine.AutoTuner.get", + return_value=tuner, + ), + patch( + "tensorrt_llm._torch.pyexecutor.model_engine.autotune", + side_effect=trtllm_autotune, + ), patch("torch.cuda.synchronize"), patch("torch.cuda.empty_cache"), patch("tensorrt_llm._torch.pyexecutor.model_engine.clear_memory_buffers"), @@ -283,14 +361,24 @@ def flashinfer_autotune(): self.assertEqual( calls, - ["flashinfer_autotune_enter", "forward", "flashinfer_autotune_exit"], + [ + "trtllm_autotune_enter", + "forward", + "trtllm_autotune_exit", + "flashinfer_autotune_enter", + "forward", + "flashinfer_autotune_exit", + ], ) + self.assertTrue(method._native_autotuned) + self.assertFalse(method.needs_native_autotune) self.assertEqual(method.backend, "auto") self.assertTrue(method._flashinfer_autotuned) - flashinfer_module.autotune.assert_called_once_with() + self.assertEqual(tuner.setup_distributed_state.call_count, 1) + tuner.setup_distributed_state.assert_called_with(engine.mapping, engine.dist) - def test_native_mxfp8_retries_after_missing_warmup_batch(self): - """MXFP8 tuning remains pending until a warmup forward can run.""" + def test_native_mxfp8_falls_back_after_missing_warmup_batch(self): + """A missing startup batch latches native MXFP8 to the default tactic.""" calls = [] @contextlib.contextmanager @@ -300,26 +388,41 @@ def trtllm_autotune(**kwargs): yield calls.append("autotune_exit") + @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) + tuner = SimpleNamespace( setup_distributed_state=Mock(), - cache_pp_recv=Mock(), - cache_pp_send=Mock(), - clean_pp_flag=Mock(), profiling_cache={}, print_profiling_cache=Mock(), ) - with patch( - "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", - return_value=True, + 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=False), ): + os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) method = MXFP8LinearMethod() - self.assertTrue(method.needs_native_autotune) - engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), - cuda_graph_runner=SimpleNamespace(enabled=False), - model=SimpleNamespace(modules=lambda: [SimpleNamespace(quant_method=method)]), + 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, @@ -333,11 +436,12 @@ def trtllm_autotune(**kwargs): _release_batch_context=Mock( side_effect=[ contextlib.nullcontext(None), - contextlib.nullcontext(object()), + contextlib.nullcontext(None), ] ), _assert_all_tp_ranks_have_warmup_batch=Mock(), - forward=Mock(side_effect=lambda *args, **kwargs: calls.append("forward")), + _release_megamoe_profiling_scratch=Mock(), + forward=Mock(), ) kv_cache_manager = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resource_manager = SimpleNamespace( @@ -353,37 +457,112 @@ def trtllm_autotune(**kwargs): "tensorrt_llm._torch.pyexecutor.model_engine.autotune", side_effect=trtllm_autotune, ), - 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, ["autotune_enter", "autotune_exit"]) - self.assertFalse(method._native_autotuned) - self.assertTrue(method.needs_native_autotune) - - PyTorchModelEngine._run_autotuner_warmup(engine, resource_manager) self.assertEqual( calls, [ "autotune_enter", "autotune_exit", - "autotune_enter", - "forward", - "autotune_exit", + "flashinfer_autotune_enter", + "flashinfer_autotune_exit", ], ) - self.assertTrue(method._native_autotuned) + self.assertFalse(method._native_autotuned) + self.assertFalse(method.use_native_autotuner) self.assertFalse(method.needs_native_autotune) - self.assertEqual(tuner.setup_distributed_state.call_count, 2) - tuner.setup_distributed_state.assert_called_with(engine.mapping, engine.dist) + self.assertEqual(method.backend, "trtllm") + self.assertFalse(method._flashinfer_autotuned) + engine.forward.assert_not_called() + + def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): + """TP ranks agree on native fallback before entering the tuning forward.""" + flashinfer_module = ModuleType("flashinfer") + flashinfer_module.mm_mxfp8 = Mock() + flashinfer_module.autotune = Mock(return_value=contextlib.nullcontext()) + tuner = SimpleNamespace( + setup_distributed_state=Mock(), + cache_pp_recv=Mock(), + cache_pp_send=Mock(), + clean_pp_flag=Mock(), + profiling_cache={}, + print_profiling_cache=Mock(), + ) + dist = SimpleNamespace(tp_allgather=Mock(return_value=[1, 0])) + + 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=False), + ): + os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) + method = MXFP8LinearMethod() + engine = SimpleNamespace( + llm_args=SimpleNamespace(enable_autotuner=True), + 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, + mapping=SimpleNamespace(tp_size=2), + dist=dist, + is_draft_model=False, + 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(), + ) + 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( + "tensorrt_llm._torch.pyexecutor.model_engine.AutoTuner.get", + return_value=tuner, + ), + patch( + "tensorrt_llm._torch.pyexecutor.model_engine.autotune", + return_value=contextlib.nullcontext(), + ), + 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) + + dist.tp_allgather.assert_called_once_with(1) + self.assertEqual(method.backend, "trtllm") + self.assertTrue(method._native_autotuned) + self.assertFalse(method._flashinfer_autotuned) + flashinfer_module.autotune.assert_not_called() + engine.forward.assert_called_once() def test_native_mxfp8_respects_disabled_global_autotuner(self): - with patch( - "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", - return_value=True, + with ( + patch( + "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", + return_value=True, + ), + patch.dict(os.environ, {}, clear=False), ): + os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) method = MXFP8LinearMethod() engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=False), diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index c0e5e3714d27..c91d200d5337 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -85,7 +85,12 @@ def _mock_mxfp8_ops(monkeypatch): mxfp8_mxfp8_gemm=native_gemm, mxfp8_mxfp8_gemm_autotuned=autotuned_gemm, ) - monkeypatch.setattr(linear_module.torch, "ops", SimpleNamespace(trtllm=fake_trtllm_ops)) + fake_torch = SimpleNamespace( + ops=SimpleNamespace(trtllm=fake_trtllm_ops), + ones=torch.ones, + float32=torch.float32, + ) + monkeypatch.setattr(linear_module, "torch", fake_torch) return ( quantized, activation_scale, @@ -104,7 +109,11 @@ def test_mxfp8_flashinfer_call_contract(monkeypatch): expected = torch.empty((2, 3), dtype=torch.bfloat16) mm_mxfp8 = Mock(return_value=expected) - monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(mm_mxfp8=mm_mxfp8)) + 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) @@ -137,8 +146,14 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): 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)) - _, _, _, native_gemm, native_output, _, _ = _mock_mxfp8_ops(monkeypatch) + monkeypatch.setitem( + sys.modules, + "flashinfer", + SimpleNamespace(mm_mxfp8=mm_mxfp8, autotune=Mock()), + ) + _, _, _, native_gemm, native_output, autotuned_gemm, autotuned_output = _mock_mxfp8_ops( + monkeypatch + ) module = SimpleNamespace( weight=torch.empty((3, 4), dtype=torch.float8_e4m3fn), @@ -149,10 +164,12 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): method = MXFP8LinearMethod() assert method.enable_flashinfer_auto() - assert method.apply(module, activation, bias=None) is native_output - native_gemm.assert_called_once() + assert method.apply(module, activation, bias=None) is autotuned_output + autotuned_gemm.assert_called_once() + native_gemm.assert_not_called() mm_mxfp8.assert_not_called() + method.mark_native_autotuned() method.mark_flashinfer_autotuned() with flashinfer_mxfp8_decode_graph_capture(): assert method.apply(module, activation, bias=None) is flashinfer_output @@ -160,7 +177,37 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): # 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 + native_gemm.assert_called_once() + + +def test_mxfp8_auto_fallback_does_not_rearm_native_autotuning(monkeypatch): + """Falling back after native warmup keeps serving on the plain native op.""" + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + monkeypatch.setitem( + sys.modules, + "flashinfer", + SimpleNamespace(mm_mxfp8=Mock(), autotune=Mock()), + ) + _, _, _, native_gemm, native_output, autotuned_gemm, _ = _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() + method.mark_native_autotuned() + assert method.enable_flashinfer_auto() + + method.disable_flashinfer_auto() + + assert method.backend == "trtllm" + assert not method.needs_native_autotune + assert method.apply(module, activation, bias=None) is native_output + native_gemm.assert_called_once() + autotuned_gemm.assert_not_called() @pytest.mark.parametrize( @@ -269,6 +316,7 @@ def test_mxfp8_native_autotuner_syncs_profiles(): profiling_cache = Mock() profiling_cache.get_specific_custom_op.return_value = {cache_key: (0, 17, 0.25)} + runner.sync_tactic_cache(SimpleNamespace(profiling_cache=profiling_cache)) runner.sync_tactic_cache(SimpleNamespace(profiling_cache=profiling_cache)) runner.mxfp8_gemm_runner.register_tactic.assert_called_once_with(8192, 9216, 6144, 17) @@ -372,7 +420,7 @@ def test_mxfp8_flashinfer_decode_graph_matches_native(monkeypatch, batch_size): dtype=torch.bfloat16, quant_config=quant_config, ).cuda() - flashinfer = Linear( + flashinfer_linear = Linear( in_features=in_f, out_features=out_f, bias=False, @@ -381,13 +429,13 @@ def test_mxfp8_flashinfer_decode_graph_matches_native(monkeypatch, batch_size): ).cuda() weights = [{"weight": weight_e4m3, "weight_scale_inv": weight_scale}] native.load_weights(weights) - flashinfer.load_weights(weights) + flashinfer_linear.load_weights(weights) native_output = native(x) - method = flashinfer.quant_method + method = flashinfer_linear.quant_method assert isinstance(method, MXFP8LinearMethod) assert method.enable_flashinfer_auto() with flashinfer_mxfp8_autotune(): - warmup_output = flashinfer(warmup_x) + warmup_output = flashinfer_linear(warmup_x) method.mark_flashinfer_autotuned() torch.testing.assert_close(warmup_output, native(warmup_x), rtol=2e-2, atol=2e-2) @@ -398,7 +446,7 @@ def test_mxfp8_flashinfer_decode_graph_matches_native(monkeypatch, batch_size): torch.cuda.synchronize() with torch.cuda.graph(graph): with flashinfer_mxfp8_decode_graph_capture(): - graph_output = flashinfer(static_x) + graph_output = flashinfer_linear(static_x) assert flashinfer_gemm.call_count == 1 graph.replay() torch.testing.assert_close(graph_output, native_output, rtol=2e-2, atol=2e-2) diff --git a/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py b/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py index dae039f9ead3..167e6203edf8 100644 --- a/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py @@ -76,10 +76,10 @@ def test_mxfp8_mxfp8_runner_tactics(): output_ref = mat_a @ mat_b.t() runner = torch.classes.trtllm.MXFP8GemmRunner(torch.bfloat16) - expected_tactics = 20 if getSMVersion() == 100 else 10 - assert runner.get_num_configs() == expected_tactics + num_tactics = runner.get_num_configs() + assert num_tactics > 0 - for tactic in [-1, *range(expected_tactics)]: + for tactic in [-1, *range(num_tactics)]: output = runner.run_gemm( fp8_a, a_block_sf, From daed63fb58b87cb79806ef2a45784fd00d57ffc4 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:35:29 -0700 Subject: [PATCH 4/6] [None][fix] Address MXFP8 autotuner review feedback Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- ...ployment-guide-for-minimax-m3-on-trtllm.md | 13 +++++++ .../_torch/custom_ops/torch_custom_ops.py | 2 ++ tensorrt_llm/_torch/modules/linear.py | 9 ++++- .../_torch/pyexecutor/model_engine.py | 6 ++++ .../test_lists/test-db/l0_b200.yml | 1 - .../test_lists/test-db/l0_b300.yml | 1 - .../test_pytorch_model_engine_warmup.py | 2 +- .../_torch/modules/test_mxfp8_linear.py | 23 ++++++++---- .../thop/parallel/test_mxfp8_mxfp8_gemm.py | 36 +++++++++++++++++++ 9 files changed, 82 insertions(+), 11 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md index e35678191c86..1c9c0c693cea 100644 --- a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md @@ -101,6 +101,19 @@ If you don't have access to the source code locally, you can manually create the The configuration uses Data-Expert Parallelism (DEP): `enable_attention_dp: true` runs the attention layers data-parallel across ranks while the MoE experts run expert-parallel, which favors high-throughput / large-batch serving on MiniMax-M3. +For MXFP8 checkpoints, TensorRT LLM selects the GEMM backend automatically. +`TRTLLM_MXFP8_GEMM_BACKEND` is an advanced override for debugging and +performance experiments: + +* `trtllm` uses the native TensorRT LLM GEMM for eager execution and CUDA graphs. +* `flashinfer` forces FlashInfer for both eager execution and CUDA graphs; it + requires the pinned `flashinfer-python` package and Blackwell MXFP8 support. +* `auto` keeps eager execution on the native GEMM and uses FlashInfer in + captured decode CUDA graphs after startup tuning. + +Leave the variable unset for normal deployments. An explicit value disables +MiniMax-M3's automatic backend selection and uses the requested policy. + ### Launch the TensorRT LLM Server MiniMax-M3 is launched through the `trtllm-llmapi-launch` wrapper, which sets up the multi-rank (MPI/Slurm) environment that the parallel server requires. The wrapper is run once per rank by Slurm (`srun`), with one task (rank) per GPU. The example below launches the server across 2 nodes (`-N 2`), 4 GPUs per node (`--ntasks-per-node 4`, 8 ranks total), using the curated YAML to drive parallelism, batching, and the MiniMax-M3 sparse-attention backend: diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 3bd4ea11f345..a2a03f43ef0f 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -552,6 +552,8 @@ def _( return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) +# The 8K-input workload produces one-, two-, and three/four-request context +# batches in these bands; their endpoints are validated on SM100 and SM103. _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" diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 4bcb345b22ba..63137c6997e3 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3096,7 +3096,11 @@ def __init__(self) -> None: super().__init__() self.use_cutlass = _mxfp8_cutlass_op_available() self.backend = os.environ.get("TRTLLM_MXFP8_GEMM_BACKEND", "trtllm") - self.use_native_autotuner = True + # Only PyTorchModelEngine owns the startup tuning lifecycle. Keep + # standalone modules and engine paths that skip warmup (for example, + # Helix CP) on the direct native op instead of leaving Python + # AutoTuner dispatch armed indefinitely. + self.use_native_autotuner = False self._native_autotuned = False if self.backend not in ("trtllm", "flashinfer", "auto"): raise ValueError("TRTLLM_MXFP8_GEMM_BACKEND must be 'trtllm', " @@ -3161,6 +3165,9 @@ def mark_flashinfer_autotuned(self) -> None: def mark_native_autotuned(self) -> None: self._native_autotuned = True + def enable_native_autotune(self) -> None: + self.use_native_autotuner = True + def disable_native_autotune(self) -> None: self.use_native_autotuner = False self._native_autotuned = False diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2fa89789da50..c53beef9890a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1570,6 +1570,12 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): method.disable_flashinfer_auto() return + # This engine owns startup warmup, so it explicitly opts its MXFP8 + # methods into native tuning. Standalone modules and engine paths that + # skip this warmup remain on the direct native op. + for method in mxfp8_methods: + method.enable_native_autotune() + # Native and FlashInfer tuning are independent. Capture native # eligibility before enabling graph-only FlashInfer dispatch. native_mxfp8_methods = [ diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 3f14d87a7d45..ae83e384989a 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -109,7 +109,6 @@ l0_b200: - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules - - unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index eae48e24cfd8..101cf51af250 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -33,7 +33,6 @@ l0_b300: - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules - - unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py 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 185d66f9d854..a3fe352fa3e5 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -308,7 +308,7 @@ def flashinfer_autotune(): os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) os.environ.pop("TLLM_AUTOTUNER_CACHE_PATH", None) method = MXFP8LinearMethod() - self.assertTrue(method.needs_native_autotune) + self.assertFalse(method.needs_native_autotune) engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index c91d200d5337..47528535b61e 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -69,7 +69,7 @@ def test_mxfp8_dispatch_returns_mxfp8_method(monkeypatch): method = get_quant_method(qc) assert isinstance(method, MXFP8LinearMethod) assert method.backend == "trtllm" - assert method.use_native_autotuner + assert not method.use_native_autotuner def _mock_mxfp8_ops(monkeypatch): @@ -164,11 +164,14 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): method = MXFP8LinearMethod() assert method.enable_flashinfer_auto() - assert method.apply(module, activation, bias=None) is autotuned_output - autotuned_gemm.assert_called_once() - native_gemm.assert_not_called() + assert method.apply(module, activation, bias=None) is native_output + native_gemm.assert_called_once() + autotuned_gemm.assert_not_called() mm_mxfp8.assert_not_called() + method.enable_native_autotune() + assert method.apply(module, activation, bias=None) is autotuned_output + autotuned_gemm.assert_called_once() method.mark_native_autotuned() method.mark_flashinfer_autotuned() with flashinfer_mxfp8_decode_graph_capture(): @@ -177,7 +180,7 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): # Leaving the decode-capture scope restores the eager/native path. assert method.apply(module, activation, bias=None) is native_output - native_gemm.assert_called_once() + assert native_gemm.call_count == 2 def test_mxfp8_auto_fallback_does_not_rearm_native_autotuning(monkeypatch): @@ -282,17 +285,23 @@ def test_mxfp8_native_autotuner_dispatch(monkeypatch): activation = torch.randn((2, 4), dtype=torch.bfloat16) method = MXFP8LinearMethod() + assert not method.use_native_autotuner + assert not method.needs_native_autotune + assert method.apply(module, activation, bias=None) is native_output + native_gemm.assert_called_once() + autotuned_gemm.assert_not_called() + + method.enable_native_autotune() assert method.use_native_autotuner assert method.needs_native_autotune assert method.apply(module, activation, bias=None) is autotuned_output autotuned_gemm.assert_called_once() - native_gemm.assert_not_called() method.mark_native_autotuned() assert not method.needs_native_autotune assert method.apply(module, activation, bias=None) is native_output autotuned_gemm.assert_called_once() - native_gemm.assert_called_once() + assert native_gemm.call_count == 2 def test_mxfp8_native_autotuner_syncs_profiles(): diff --git a/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py b/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py index 167e6203edf8..d6b489dad7b1 100644 --- a/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py @@ -19,6 +19,7 @@ from utils.util import getSMVersion import tensorrt_llm._torch.custom_ops.torch_custom_ops # noqa: F401 +from tensorrt_llm._torch.custom_ops.torch_custom_ops import _map_to_mxfp8_large_m_bucket @pytest.mark.skipif( @@ -152,3 +153,38 @@ def test_mxfp8_mxfp8_native_tactic_cache(): torch.testing.assert_close(cached_fallback_output, explicit_fallback_output, rtol=0, atol=0) finally: runner.clear_tactic_cache() + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="MXFP8 tactic cache requires SM100 or SM103. Current SM is %d." % getSMVersion(), +) +@pytest.mark.parametrize( + "lower_bound,upper_bound,bucket", + [ + (6553, 8192, 8192), + (13106, 16384, 16384), + (19659, 32768, 32768), + ], +) +def test_mxfp8_native_tactic_cache_large_m_bucket_boundaries( + lower_bound: int, upper_bound: int, bucket: int +): + """C++ cache bucketing stays aligned with the Python tuning profiles.""" + n, k = 9216, 6144 + runner = torch.classes.trtllm.MXFP8GemmRunner(torch.bfloat16) + + try: + runner.clear_tactic_cache() + runner.register_tactic(bucket, n, k, -1) + + assert _map_to_mxfp8_large_m_bucket(lower_bound) == bucket + assert _map_to_mxfp8_large_m_bucket(upper_bound) == bucket + assert runner.get_cached_tactic(lower_bound, n, k) == -1 + assert runner.get_cached_tactic(upper_bound, n, k) == -1 + assert _map_to_mxfp8_large_m_bucket(lower_bound - 1) == lower_bound - 1 + assert _map_to_mxfp8_large_m_bucket(upper_bound + 1) == upper_bound + 1 + assert runner.get_cached_tactic(lower_bound - 1, n, k) == -2 + assert runner.get_cached_tactic(upper_bound + 1, n, k) == -2 + finally: + runner.clear_tactic_cache() From a7816a4f952168de6f6ca27dbe87224d5d71c926 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:46:10 -0700 Subject: [PATCH 5/6] [None][fix] Harden MXFP8 autotuner cache and rank handling Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/thop/mxfp8Gemm.cpp | 15 ++++++ ...ployment-guide-for-minimax-m3-on-trtllm.md | 5 +- tensorrt_llm/_torch/autotuner.py | 7 +++ .../_torch/custom_ops/torch_custom_ops.py | 50 +++++++++++++++---- .../_torch/pyexecutor/model_engine.py | 21 +++++--- .../test_pytorch_model_engine_warmup.py | 10 ++-- .../_torch/modules/test_mxfp8_linear.py | 25 ++++++++-- 7 files changed, 108 insertions(+), 25 deletions(-) diff --git a/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp b/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp index 3b597ac6dda4..76835ab1fa2e 100644 --- a/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp +++ b/cpp/tensorrt_llm/thop/mxfp8Gemm.cpp @@ -240,9 +240,12 @@ at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, a /*useTacticCache=*/true); } +//! Profiles native MXFP8 GEMM tactics and registers selected tactics in the serving cache. class MXFP8GemmRunner : public torch::CustomClassHolder { public: + //! Constructs a runner for the requested output element type. + //! \param outputDtype Output type; supported values are FP16, BF16, and FP32. explicit MXFP8GemmRunner(at::ScalarType outputDtype) : mOutputDtype(outputDtype) { @@ -251,6 +254,14 @@ class MXFP8GemmRunner : public torch::CustomClassHolder mConfigs = CutlassFp4GemmRunner{}.getConfigs(); } + //! Runs one MXFP8 GEMM with a selected compiled tactic. + //! \param act Row-major MXFP8 activation tensor with shape [M, K]. + //! \param actScale Swizzled UE8M0 activation scales. + //! \param weight MXFP8 weight tensor with logical shape [N, K]. + //! \param weightScale Swizzled UE8M0 weight scales. + //! \param globalScale FP32 scalar tensor applied by the epilogue. + //! \param configIdx Compiled tactic index, or -1 for the generic fallback. + //! \return Output tensor with shape [M, N]. at::Tensor runGemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight, at::Tensor const& weightScale, at::Tensor const& globalScale, int64_t configIdx) const { @@ -259,23 +270,27 @@ class MXFP8GemmRunner : public torch::CustomClassHolder act, actScale, weight, weightScale, globalScale, mOutputDtype, &config, /*useTacticCache=*/false); } + //! Registers a compiled tactic for a serving shape. void registerTactic(int64_t const m, int64_t const n, int64_t const k, int64_t const configIdx) const { tkc::CutlassGemmConfig const config = configIdx == -1 ? getDefaultMxfp8GemmConfig() : getConfig(configIdx); cacheMxfp8Tactic(m, n, k, mOutputDtype, config, configIdx); } + //! Returns the registered tactic for a serving shape, or the cache-miss sentinel. int64_t getCachedTactic(int64_t const m, int64_t const n, int64_t const k) const { auto const entry = findMxfp8TacticCacheEntry(m, n, k, mOutputDtype); return entry.has_value() ? entry->tactic : kMxfp8TacticCacheMiss; } + //! Removes every registered MXFP8 serving tactic. void clearTacticCache() const { clearMxfp8CachedTactics(); } + //! Returns the number of compiled native tactics available for profiling. int64_t getNumConfigs() const { return static_cast(mConfigs.size()); diff --git a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md index 1c9c0c693cea..aae19a4b73bd 100644 --- a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md @@ -106,8 +106,9 @@ For MXFP8 checkpoints, TensorRT LLM selects the GEMM backend automatically. performance experiments: * `trtllm` uses the native TensorRT LLM GEMM for eager execution and CUDA graphs. -* `flashinfer` forces FlashInfer for both eager execution and CUDA graphs; it - requires the pinned `flashinfer-python` package and Blackwell MXFP8 support. +* `flashinfer` forces FlashInfer for eligible captured decode CUDA graphs; eager, + context/prefill, and piecewise CUDA-graph execution remain on the native GEMM. + It requires the pinned `flashinfer-python` package and Blackwell MXFP8 support. * `auto` keeps eager execution on the native GEMM and uses FlashInfer in captured decode CUDA graphs after startup tuning. diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 934487fce8dd..2c08b487d6d2 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -436,6 +436,7 @@ class AutoTunerProfilingCache: def __init__(self): self.cache: Dict[Tuple, Tuple] = dict() + self._generation = 0 # Track which ops use which distributed strategy # Maps custom_op name -> DistributedTuningStrategy @@ -466,6 +467,12 @@ def clear(self) -> None: self.cache.clear() self.independent_op.clear() self.excluded_op.clear() + self._generation += 1 + + @property + def generation(self) -> int: + """Return a counter that changes whenever this cache is cleared.""" + return self._generation def fallback_entry(self) -> Tuple: # runner_id = 0, tactic = -1 diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index a2a03f43ef0f..fa0fffdb7877 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -19,6 +19,7 @@ from dataclasses import replace from functools import lru_cache from typing import ClassVar, List, Mapping, Optional, Tuple, Union +from weakref import WeakKeyDictionary import torch import triton # type: ignore[import] @@ -32,9 +33,9 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.utils import fp8_quantize -from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, - DynamicTensorSpec, OptimizationProfile, TunableRunner, - TuningConfig) +from ..autotuner import (AutoTuner, AutoTunerProfilingCache, ConstraintSpec, + DistributedTuningStrategy, DynamicTensorSpec, + OptimizationProfile, TunableRunner, TuningConfig) 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 @@ -582,9 +583,17 @@ def _mxfp8_scale_infer_shape(input_shapes: List[List[int]]) -> int: class MXFP8GemmRunner(TunableRunner): + """Autotunable native MXFP8 GEMM runner with a serving tactic cache. + + Args: + output_dtype: Output element type. Supported types are FP16, BF16, and + FP32. + """ + runner_dict = dict() - synced_cache_keys: ClassVar[dict[tuple[torch.dtype, int, int], set[tuple]]] - synced_cache_keys = {} + synced_cache_keys: ClassVar[WeakKeyDictionary[AutoTunerProfilingCache, dict[ + tuple[torch.dtype, int], tuple[int, + set[tuple]]]]] = WeakKeyDictionary() tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, _get_mxfp8_large_m_tuning_buckets, _map_to_mxfp8_large_m_bucket), ), @@ -592,7 +601,7 @@ class MXFP8GemmRunner(TunableRunner): 1, 0, _mxfp8_scale_infer_shape), ), use_cuda_graph=False) - def __init__(self, output_dtype: torch.dtype): + def __init__(self, output_dtype: torch.dtype) -> None: self.output_dtype = output_dtype self.sm_version = get_sm_version() instance_key = (output_dtype, self.sm_version) @@ -615,10 +624,16 @@ def sync_tactic_cache(self, tuner: AutoTuner) -> None: """Register newly profiled tactics in the native serving cache.""" runner_name = self.__class__.__name__ unique_id = str(self.unique_id()) - cache = tuner.profiling_cache.get_specific_custom_op( - _MXFP8_AUTOTUNED_OP) - sync_key = (*self.unique_id(), id(tuner.profiling_cache)) - synced_cache_keys = self.synced_cache_keys.setdefault(sync_key, set()) + profiling_cache = tuner.profiling_cache + cache = profiling_cache.get_specific_custom_op(_MXFP8_AUTOTUNED_OP) + runner_sync_state = self.synced_cache_keys.setdefault( + profiling_cache, {}) + generation, synced_cache_keys = runner_sync_state.get( + self.unique_id(), (-1, set())) + if generation != profiling_cache.generation: + generation, synced_cache_keys = profiling_cache.generation, set() + runner_sync_state[self.unique_id()] = (generation, + synced_cache_keys) for cache_key, (_runner_id, tactic, _min_time) in cache.items(): if cache_key in synced_cache_keys: continue @@ -662,6 +677,21 @@ def mxfp8_mxfp8_gemm_autotuned( global_scale: torch.Tensor, output_dtype: torch.dtype, ) -> torch.Tensor: + """Run an autotuned native MXFP8-by-MXFP8 matrix multiplication. + + Args: + act: Row-major MXFP8 activation tensor with shape ``[M, K]``. + act_scale: Swizzled UE8M0 activation scales. + weight: MXFP8 weight tensor with logical shape ``[N, K]`` and the + column-major storage expected by CUTLASS. + weight_scale: Swizzled UE8M0 weight scales. + global_scale: FP32 scalar tensor applied by the GEMM epilogue. + output_dtype: Output element type. Supported types are FP16, BF16, and + FP32. + + Returns: + Output tensor with shape ``[M, N]`` and dtype ``output_dtype``. + """ tuner = AutoTuner.get() runner = MXFP8GemmRunner(output_dtype) inputs = [act, act_scale, weight, weight_scale, global_scale] diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 93371acec441..fe1c3cd8ce2f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1613,12 +1613,19 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): if method.needs_flashinfer_autotune ] - # Every TP rank must make the same backend decision before any rank - # returns or enters a tuning forward with TP collectives. - if self.mapping.tp_size > 1: + # Every TP and PP rank must make the same backend decision before any + # rank returns or enters a tuning forward with model collectives. + if self.mapping.tp_size > 1 or self.mapping.has_pp(): local_flashinfer_enabled = int(bool(flashinfer_mxfp8_methods)) - all_flashinfer_enabled = list( - self.dist.tp_allgather(local_flashinfer_enabled)) + all_flashinfer_enabled = [local_flashinfer_enabled] + if self.mapping.tp_size > 1: + all_flashinfer_enabled = list( + self.dist.tp_allgather(local_flashinfer_enabled)) + if self.mapping.has_pp(): + all_flashinfer_enabled = [ + enabled for stage_flags in self.dist.pp_allgather( + all_flashinfer_enabled) for enabled in stage_flags + ] if any(all_flashinfer_enabled) and not all(all_flashinfer_enabled): forced_flashinfer = any(method.backend == "flashinfer" for method in mxfp8_methods) @@ -1628,9 +1635,9 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): if forced_flashinfer: raise RuntimeError( "FlashInfer MXFP8 was explicitly requested but is not " - "available on every TP rank") + "available on every TP/PP rank") logger.warning( - "FlashInfer MXFP8 availability differs across TP ranks; " + "FlashInfer MXFP8 availability differs across TP/PP ranks; " "using the native TensorRT-LLM GEMM backend on every rank.") enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods) 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 003f72a3936a..f3d518b04e63 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -489,7 +489,7 @@ def flashinfer_autotune(): engine.forward.assert_not_called() def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): - """TP ranks agree on native fallback before entering the tuning forward.""" + """TP and PP ranks agree on fallback before the tuning forward.""" flashinfer_module = ModuleType("flashinfer") flashinfer_module.mm_mxfp8 = Mock() flashinfer_module.autotune = Mock(return_value=contextlib.nullcontext()) @@ -501,7 +501,10 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): profiling_cache={}, print_profiling_cache=Mock(), ) - dist = SimpleNamespace(tp_allgather=Mock(return_value=[1, 0])) + dist = SimpleNamespace( + tp_allgather=Mock(return_value=[1, 1]), + pp_allgather=Mock(return_value=[[1, 1], [1, 0]]), + ) with ( patch.dict(sys.modules, {"flashinfer": flashinfer_module}), @@ -527,7 +530,7 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): batch_size=16, max_seq_len=2, original_max_draft_len=0, - mapping=SimpleNamespace(tp_size=2, has_pp=lambda: False), + mapping=SimpleNamespace(tp_size=2, has_pp=lambda: True), dist=dist, is_draft_model=False, guided_decoder=None, @@ -560,6 +563,7 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): PyTorchModelEngine._run_autotuner_warmup(engine, resource_manager) dist.tp_allgather.assert_called_once_with(1) + dist.pp_allgather.assert_called_once_with([1, 1]) self.assertEqual(method.backend, "trtllm") self.assertTrue(method._native_autotuned) self.assertFalse(method._flashinfer_autotuned) diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index 47528535b61e..de7a60052e17 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gc import sys +import weakref from types import SimpleNamespace from unittest.mock import Mock @@ -305,6 +307,7 @@ def test_mxfp8_native_autotuner_dispatch(monkeypatch): def test_mxfp8_native_autotuner_syncs_profiles(): + initial_cache_count = len(MXFP8GemmRunner.synced_cache_keys) runner = object.__new__(MXFP8GemmRunner) runner.output_dtype = torch.bfloat16 runner.sm_version = 100 @@ -322,13 +325,29 @@ def test_mxfp8_native_autotuner_syncs_profiles(): str(runner.unique_id()), profile, ) - profiling_cache = Mock() + profiling_cache = Mock(generation=0) profiling_cache.get_specific_custom_op.return_value = {cache_key: (0, 17, 0.25)} - runner.sync_tactic_cache(SimpleNamespace(profiling_cache=profiling_cache)) - runner.sync_tactic_cache(SimpleNamespace(profiling_cache=profiling_cache)) + tuner = Mock(profiling_cache=profiling_cache) + runner.sync_tactic_cache(tuner) + runner.sync_tactic_cache(tuner) runner.mxfp8_gemm_runner.register_tactic.assert_called_once_with(8192, 9216, 6144, 17) + profiling_cache.generation = 1 + runner.sync_tactic_cache(tuner) + assert runner.mxfp8_gemm_runner.register_tactic.call_count == 2 + + replacement_cache = Mock(generation=0) + replacement_cache.get_specific_custom_op.return_value = {cache_key: (0, 17, 0.25)} + tuner.profiling_cache = replacement_cache + runner.sync_tactic_cache(tuner) + assert runner.mxfp8_gemm_runner.register_tactic.call_count == 3 + + cache_ref = weakref.ref(profiling_cache) + del profiling_cache + gc.collect() + assert cache_ref() is None + assert len(MXFP8GemmRunner.synced_cache_keys) == initial_cache_count + 1 def test_mxfp8_rejects_unknown_backend(monkeypatch): From 46f99c83337d9904fc88b8f233e08be504043a4c Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:55:29 -0700 Subject: [PATCH 6/6] [None][fix] Invalidate MXFP8 tactic sync after cache load Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/autotuner.py | 2 + .../_torch/modules/test_mxfp8_linear.py | 57 ++++++++++++++++++- 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 2c08b487d6d2..90bdf05e6cf8 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -746,6 +746,8 @@ def load_cache(self, file_path: Union[str, Path], rank: int) -> None: f"[AutoTuner] Loaded {len(rank_cache)} rank-specific cache entries for rank {rank}" ) + self._generation += 1 + logger.info( f"[AutoTuner] Successfully loaded cache from {file_path} using JSON format (total {len(self.cache)} entries)" ) diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index de7a60052e17..b4580d7ee502 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -14,6 +14,7 @@ # limitations under the License. import gc +import json import sys import weakref from types import SimpleNamespace @@ -23,7 +24,7 @@ import torch import tensorrt_llm._torch.modules.linear as linear_module -from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.autotuner import AutoTuner, AutoTunerProfilingCache from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( MXFP8GemmRunner, _get_mxfp8_large_m_tuning_buckets, @@ -350,6 +351,60 @@ def test_mxfp8_native_autotuner_syncs_profiles(): assert len(MXFP8GemmRunner.synced_cache_keys) == initial_cache_count + 1 +def test_mxfp8_native_autotuner_resyncs_loaded_tactic(tmp_path): + runner = object.__new__(MXFP8GemmRunner) + runner.output_dtype = torch.bfloat16 + runner.sm_version = 100 + runner.mxfp8_gemm_runner = Mock() + profile = ( + (8192, 6144), + (-1,), + (9216, 6144), + (1769472,), + (1,), + ) + cache_key = ( + "trtllm::mxfp8_mxfp8_gemm_autotuned::gemm", + "MXFP8GemmRunner", + str(runner.unique_id()), + profile, + ) + profiling_cache = object.__new__(AutoTunerProfilingCache) + profiling_cache.cache = {cache_key: (0, 17, 0.25)} + profiling_cache.independent_op = set() + profiling_cache.excluded_op = set() + profiling_cache._generation = 0 + tuner = Mock(profiling_cache=profiling_cache) + + runner.sync_tactic_cache(tuner) + runner.mxfp8_gemm_runner.register_tactic.assert_called_once_with(8192, 9216, 6144, 17) + + cache_path = tmp_path / "mxfp8-cache.json" + cache_path.write_text( + json.dumps( + { + "metadata": { + "lib_version": "test", + "creation_timestamp": 0, + "device_name": "test", + "device_capability": [10, 0], + }, + "rank_0": profiling_cache._serialize_cache_data({cache_key: (0, 23, 0.20)}), + } + ) + ) + profiling_cache.load_cache(cache_path, rank=0) + runner.sync_tactic_cache(tuner) + + assert profiling_cache.generation == 1 + assert runner.mxfp8_gemm_runner.register_tactic.call_args_list[-1].args == ( + 8192, + 9216, + 6144, + 23, + ) + + def test_mxfp8_rejects_unknown_backend(monkeypatch): monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", "unknown") with pytest.raises(ValueError, match="TRTLLM_MXFP8_GEMM_BACKEND"):