From c0fe37bb40d80cda4dcb7eea87d22a26c24f4899 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 08:30:25 -0700 Subject: [PATCH 01/12] [None][perf] Migrate MiniMax-M3 piecewise CUDA graphs to main Combine the symbolic FP8/FlashInfer fixes from #17216 and context producer capture from 6d282523 (#17473), adapting both to main after #18205. Preserve current attention-DP routing, MSA cache contracts, and MXFP8 decode tuning. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp | 42 ----- .../torch_compile_and_piecewise_cuda_graph.md | 13 ++ .../backends/sparse/minimax_m3/msa_backend.py | 4 + .../_torch/custom_ops/cpp_custom_ops.py | 45 ++++++ .../custom_ops/flashinfer_custom_ops.py | 37 +++++ .../_torch/models/modeling_minimaxm3.py | 62 ++++++- tensorrt_llm/_torch/modules/linear.py | 21 +-- .../_torch/pyexecutor/model_engine.py | 58 ++++++- .../defs/accuracy/test_llm_api_pytorch.py | 26 ++- .../integration/test_lists/test-db/l0_cpu.yml | 3 + .../test_lists/test-db/l0_dgx_b200.yml | 2 + .../test_fused_qk_norm_rope.py | 88 ++++++++++ .../attention/sparse/msa/test_msa_backend.py | 40 +++++ .../test_pytorch_model_engine_warmup.py | 152 +++++++++++++++++- .../unittest/_torch/models/test_minimax_m3.py | 116 +++++++++++++ .../_torch/modules/test_mxfp8_linear.py | 59 +++++-- 16 files changed, 694 insertions(+), 74 deletions(-) diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 933ed62432c2..0ae4ba1f52bc 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -170,18 +170,6 @@ torch::Tensor fused_qk_norm_rope_to_fp8(torch::Tensor const& qkv, // [num_tokens return out; } -// Meta (fake) implementation for torch.compile / tracing: only shape+dtype. -torch::Tensor fused_qk_norm_rope_to_fp8_meta(torch::Tensor const& qkv, int64_t num_heads_q, int64_t num_heads_k, - int64_t num_heads_v, int64_t head_dim, int64_t /*rotary_dim*/, double /*eps*/, torch::Tensor const& /*q_weight*/, - torch::Tensor const& /*k_weight*/, double /*base*/, bool /*is_neox*/, torch::Tensor const& /*position_ids*/, - double /*factor*/, double /*low*/, double /*high*/, double /*attention_factor*/, bool /*is_qk_norm*/, - bool /*use_gemma*/, bool /*use_mrope*/, int64_t /*mrope_section1*/, int64_t /*mrope_section2*/) -{ - int64_t num_tokens = qkv.size(0); - int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; - return torch::empty({num_tokens, total_heads * head_dim}, qkv.options().dtype(torch::kFloat8_e4m3fn)); -} - torch::Tensor minimaxM3Fp8QKNormRopeKVInsert(torch::Tensor const& qkv, torch::Tensor& kvCache, torch::Tensor const& outCacheLoc, int64_t numHeadsQ, int64_t numHeadsK, int64_t numHeadsV, int64_t headDim, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, double base, bool isNeox, @@ -251,14 +239,6 @@ torch::Tensor minimaxM3Fp8QKNormRopeKVInsert(torch::Tensor const& qkv, torch::Te return qOut; } -torch::Tensor minimaxM3Fp8QKNormRopeKVInsertMeta(torch::Tensor const& qkv, torch::Tensor& /*kvCache*/, - torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, int64_t /*numHeadsK*/, int64_t /*numHeadsV*/, - int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, torch::Tensor const& /*qWeight*/, - torch::Tensor const& /*kWeight*/, double /*base*/, bool /*isNeox*/, torch::Tensor const& /*positionIds*/) -{ - return torch::empty({qkv.size(0), numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn)); -} - std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsert(torch::Tensor const& packed, torch::Tensor& kvCache, torch::Tensor& indexKCache, torch::Tensor const& outCacheLoc, int64_t numHeadsQ, int64_t numHeadsKV, int64_t numHeadsIndex, int64_t headDim, int64_t rotaryDim, double eps, @@ -352,19 +332,6 @@ std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsert( return {qOut, indexQOut}; } -std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta(torch::Tensor const& packed, - torch::Tensor& /*kvCache*/, torch::Tensor& /*indexKCache*/, torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, - int64_t /*numHeadsKV*/, int64_t numHeadsIndex, int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, - torch::Tensor const& /*qWeight*/, torch::Tensor const& /*kWeight*/, torch::Tensor const& /*indexQWeight*/, - torch::Tensor const& /*indexKWeight*/, torch::Tensor const& /*rotaryCosSin*/, torch::Tensor const& /*positionIds*/) -{ - auto options = packed.options().dtype(at::ScalarType::Float8_e4m3fn); - return { - torch::empty({packed.size(0), numHeadsQ, headDim}, options), - torch::empty({packed.size(0), numHeadsIndex, headDim}, options), - }; -} - // Register the PyTorch operators TORCH_LIBRARY_FRAGMENT(trtllm, m) { @@ -390,7 +357,6 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor rotary_cos_sin, Tensor position_ids) -> (Tensor, Tensor)"); } -// Register the CUDA implementation TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("fused_qk_norm_rope", &fused_qk_norm_rope); @@ -399,14 +365,6 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert", &minimaxM3Fp8QKVIndexerNormRopeKVInsert); } -// Register the Meta implementation (shape/dtype inference for torch.compile). -TORCH_LIBRARY_IMPL(trtllm, Meta, m) -{ - m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8_meta); - m.impl("minimax_m3_fp8_qk_norm_rope_kv_insert", &minimaxM3Fp8QKNormRopeKVInsertMeta); - m.impl("minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert", &minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta); -} - } // namespace torch_ext TRTLLM_NAMESPACE_END diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 8b179030177f..fa957fc2fa08 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -79,6 +79,19 @@ validated before use. ### Piecewise CUDA Graph & Generation Only CUDA Graph +For decoder models, piecewise mode compiles capture-eligible context and mixed +batches. Generation-only batches and context batches above the capture ceiling +use the eager model, including any speculative-decoding epilogue. With attention +DP, eligibility is shared across ranks, including ranks without local context +requests. Torch compile without piecewise graphs retains its all-batch behavior. + +MiniMax-M3 with MSA supports FP8 KV and index-K caches in piecewise mode. With +`sparse_attention_config.fuse_qkv_index_projection: true`, projection, norm, +RoPE, FP8 conversion and cache insertion are captured together; sparse attention +remains in the eager boundary. Padded rows do not write to the caches. Automatic +MXFP8 dispatch uses the native backend for compiled context while preserving +decode-graph backend tuning. + Piecewise CUDA Graph only handles context-only and mixed context+generation iterations, while the generation-only CUDA Graph only handles pure generation iterations. Users need to specify the number of tokens to capture for each type of CUDA Graph separately in the extra config. Currently, the default value for `capture_num_tokens` is `[2**i for i in range(8)] + [i for i in range(256, 3073, 256)]`. However, this configuration should be tuned based on specific hardware, model, and parallel strategy. For guidance on tuning these values, see the [Performance Tuning](#performance-tuning) section below. ```yaml diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 3e45870d64a8..add137126681 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -965,6 +965,10 @@ def _build_msa_fields(self) -> None: self._msa_fields_ready = False if not self._msa_buffers_ready: return + # Captured producers execute the whole padded bucket, including on + # attention-DP ranks without local requests. Invalidate the tail before + # any early return so replay cannot write padding into stale KV slots. + self.msa_out_cache_loc.fill_(-1) request_ids = self.request_ids qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index d41ae7c5b066..042012c99c1f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -397,6 +397,51 @@ def minimax_m3_fp8_indexer_qk_norm_rope_fake( return qk.new_empty((qk.shape[0], num_heads_q, head_dim), dtype=torch.float8_e4m3fn) + @torch.library.register_fake("trtllm::fused_qk_norm_rope_to_fp8") + def _(qkv: torch.Tensor, num_heads_q: int, num_heads_k: int, + num_heads_v: int, head_dim: int, rotary_dim: int, eps: float, + q_weight: torch.Tensor, k_weight: torch.Tensor, base: float, + is_neox: bool, position_ids: torch.Tensor, factor: float, low: float, + high: float, attention_factor: float, is_qk_norm: bool, + use_gemma: bool, use_mrope: bool, mrope_section1: int, + mrope_section2: int) -> torch.Tensor: + del rotary_dim, eps, q_weight, k_weight, base, is_neox, position_ids + del factor, low, high, attention_factor, is_qk_norm, use_gemma + del use_mrope, mrope_section1, mrope_section2 + total_heads = num_heads_q + num_heads_k + num_heads_v + return qkv.new_empty((qkv.shape[0], total_heads * head_dim), + dtype=torch.float8_e4m3fn) + + @torch.library.register_fake( + "trtllm::minimax_m3_fp8_qk_norm_rope_kv_insert") + def _(qkv: torch.Tensor, kv_cache: torch.Tensor, + out_cache_loc: torch.Tensor, num_heads_q: int, num_heads_k: int, + num_heads_v: int, head_dim: int, rotary_dim: int, eps: float, + q_weight: torch.Tensor, k_weight: torch.Tensor, base: float, + is_neox: bool, position_ids: torch.Tensor) -> torch.Tensor: + del kv_cache, out_cache_loc, num_heads_k, num_heads_v, rotary_dim, eps + del q_weight, k_weight, base, is_neox, position_ids + return qkv.new_empty((qkv.shape[0], num_heads_q, head_dim), + dtype=torch.float8_e4m3fn) + + @torch.library.register_fake( + "trtllm::minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert") + def _(packed: torch.Tensor, kv_cache: torch.Tensor, + index_k_cache: torch.Tensor, out_cache_loc: torch.Tensor, + num_heads_q: int, num_heads_kv: int, num_heads_index: int, + head_dim: int, rotary_dim: int, eps: float, q_weight: torch.Tensor, + k_weight: torch.Tensor, index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + del kv_cache, index_k_cache, out_cache_loc, num_heads_kv, rotary_dim + del eps, q_weight, k_weight, index_q_weight, index_k_weight + del rotary_cos_sin, position_ids + num_tokens = packed.shape[0] + return (packed.new_empty((num_tokens, num_heads_q, head_dim), + dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, num_heads_index, head_dim), + dtype=torch.float8_e4m3fn)) + @torch.library.register_fake("trtllm::userbuffers_allreduce_finalize") def _(input, force_applying_finalize): return torch.empty_like(input) diff --git a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py index c9c85767d5c3..e1a0a3917bca 100644 --- a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import torch from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl @@ -126,3 +129,37 @@ def _( is_neox: bool = True, ): return + + # mm_mxfp8 is newer than the entry points above, so probe it separately + # rather than breaking this module's import on an older flashinfer build. + try: + from flashinfer import mm_mxfp8 + except ImportError: + mm_mxfp8 = None + + if mm_mxfp8 is not None: + + # Wrap this into a custom op so torch.compile traces one opaque node + # instead of inlining flashinfer's Python-level tactic lookup. + @torch.library.custom_op("trtllm::flashinfer_mm_mxfp8", mutates_args=()) + def flashinfer_mm_mxfp8(act: torch.Tensor, act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype) -> torch.Tensor: + # Argument order mirrors trtllm::mxfp8_mxfp8_gemm: weight arrives as + # [N, K] and mm_mxfp8 wants [K, N]. Both scale buffers are the 1D + # padded swizzled CUTLASS layout, hence use_8x4_sf_layout=False. + return mm_mxfp8(act, + weight.t(), + act_scale, + weight_scale, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend="cutlass") + + @flashinfer_mm_mxfp8.register_fake + def _(act: torch.Tensor, act_scale: torch.Tensor, weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype) -> torch.Tensor: + return act.new_empty((act.size(0), weight.size(0)), + dtype=output_dtype) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index fb443da45e70..a93d2f8468cf 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -838,6 +838,38 @@ def _minimax_m3_qkv_index_proj_fake( return hidden_states.new_empty((hidden_states.shape[0], sum(qkv_proj.local_output_sizes))) +@torch.library.custom_op("trtllm::minimax_m3_fused_sparse_qkv_producer", mutates_args=()) +def minimax_m3_fused_sparse_qkv_producer( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, +) -> List[torch.Tensor]: + """Capture projection, norm, RoPE and FP8 cache insertion together.""" + attn_metadata, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) + packed = attn_layer.qkv_proj(hidden_states) + result = attn_layer._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed, position_ids, attn_metadata + ) + if result is None: + raise RuntimeError("MiniMax-M3 piecewise graph requires the fused FP8 sparse QKV producer.") + return list(result) + + +@minimax_m3_fused_sparse_qkv_producer.register_fake +def _minimax_m3_fused_sparse_qkv_producer_fake( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, +) -> List[torch.Tensor]: + del position_ids + _, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) + num_tokens = hidden_states.shape[0] + return [ + hidden_states.new_empty((num_tokens, attn_layer.q_size), dtype=torch.float8_e4m3fn), + hidden_states.new_empty((num_tokens, attn_layer.index_q_size), dtype=torch.float8_e4m3fn), + ] + + @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",)) def minimax_m3_attn_custom_op_inplace( q: Optional[torch.Tensor], @@ -1929,6 +1961,18 @@ def _sparse_forward( packed_qkv = None packed_idx_qk = None if self.enable_fused_qkv_index_projection: + if ( + self.register_to_config + and is_torch_compiling() + and isinstance(self.attn, MiniMaxM3MsaSparseAttention) + and self._emit_fp8_main_qkv() + and self.attn.indexer_kv_dtype == "fp8" + ): + q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( + hidden_states, position_ids, self.layer_idx_str + ) + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): # Keep the projection in the captured segment while hiding # its shape-specializing MXFP8 internals behind a symbolic @@ -2003,10 +2047,19 @@ def _index_norm_rope(): idx_qk = ( packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) ) - fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) - if fp8_idx_q is not None: - # Index-K was inserted directly into the paged side cache. - return fp8_idx_q, None + graph_fp8_indexer = ( + self.register_to_config + and is_torch_compiling() + and isinstance(self.attn, MiniMaxM3MsaSparseAttention) + and self.attn.indexer_kv_dtype == "fp8" + ) + if not graph_fp8_indexer: + fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) + if fp8_idx_q is not None: + # Index-K was inserted directly into the paged side cache. + return fp8_idx_q, None + # During compilation keep norm/RoPE/FP8 conversion captured, but + # leave dynamic index-K cache insertion in the attention boundary. fused_idx = self._fused_qk_norm_rope( idx_qk, position_ids, @@ -2016,6 +2069,7 @@ def _index_norm_rope(): head_dim=self.sparse_index_dim, q_norm=self.index_q_norm, k_norm=self.index_k_norm, + out_fp8=graph_fp8_indexer, ) if fused_idx is not None: return self._split_index_qk(fused_idx) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index b3e29a427482..d8cbf963c468 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -37,7 +37,8 @@ from ...models.modeling_utils import QuantConfig from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, - replace_parameter_and_save_metadata, unswizzle_sf) + is_torch_compiling, replace_parameter_and_save_metadata, + unswizzle_sf) from .low_m_gemm import _should_apply_low_m_gemm, apply_low_m_gemm @@ -3340,9 +3341,13 @@ def _load_flashinfer(self, *, required: bool) -> bool: "quantization ops on Blackwell") return False try: - from flashinfer import autotune, mm_mxfp8 + from flashinfer import autotune if not callable(autotune): raise ImportError("flashinfer.autotune is unavailable") + flashinfer_mxfp8 = getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", + None) + if flashinfer_mxfp8 is None: + raise ImportError("trtllm::flashinfer_mm_mxfp8 is unavailable") except ImportError as error: if required: raise RuntimeError( @@ -3353,7 +3358,7 @@ def _load_flashinfer(self, *, required: bool) -> bool: "TensorRT-LLM GEMM backend.", key="flashinfer_mxfp8_unavailable") return False - self._flashinfer_mxfp8 = mm_mxfp8 + self._flashinfer_mxfp8 = flashinfer_mxfp8 return True def enable_flashinfer_auto(self) -> bool: @@ -3425,7 +3430,7 @@ def apply(self, module: Linear, input: torch.Tensor, if self.use_cutlass: input = input.contiguous() - if (self.tune_decode_graph_backends + if (self.tune_decode_graph_backends and not is_torch_compiling() and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()): # Tune only in the warmup-only pass (flashinfer_mxfp8_autotune). output = mxfp8_quantize_gemm_autotuned( @@ -3440,7 +3445,7 @@ def apply(self, module: Linear, input: torch.Tensor, # then the CUTLASS block-scaled e4m3xe4m3 GEMM. act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize(input, True) use_flashinfer = self.backend == "flashinfer" or ( - self.backend == "auto" and + self.backend == "auto" and not is_torch_compiling() and (_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() or (self._flashinfer_autotuned and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()))) @@ -3449,12 +3454,10 @@ def apply(self, module: Linear, input: torch.Tensor, assert flashinfer_mxfp8 is not None output = flashinfer_mxfp8( act_e4m3, - module.weight.t(), act_sf, + module.weight, module.weight_scale, - out_dtype=module.dtype, - use_8x4_sf_layout=False, - backend="cutlass", + module.dtype, ) else: # globalScale is the alpha multiplier; pure MXFP8xMXFP8 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 932e563b7f75..cb2a4c5aa1d7 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -74,7 +74,8 @@ from ..utils import (get_model_extra_attrs, get_per_request_prefill_cuda_graph_flag, set_per_request_prefill_cuda_graph_flag, - set_torch_compiling, with_model_extra_attrs) + set_torch_compiling, torch_compiling, + with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner from .config_utils import is_hybrid_linear from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, @@ -195,6 +196,34 @@ def _make_single_token_context_graph_batch( return graph_batch, promoted_context_request_ids +class _ContextOnlyCompiledModel(torch.nn.Module): + """Share parameters between captured context and eager generation paths. + + The prefill flag includes the all-rank attention-DP decision and capture + ceiling. A decode-only rank must still compile when another rank prefills. + """ + + def __init__(self, eager_model: torch.nn.Module, + compiled_model: torch.nn.Module) -> None: + super().__init__() + self.eager_model = eager_model + self.compiled_model = compiled_model + + def forward(self, *args: Any, **kwargs: Any) -> Any: + model = (self.compiled_model + if get_per_request_prefill_cuda_graph_flag() else + self.eager_model) + return model(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + # Model-specific epilogues (including M3 Eagle3) access embed_tokens + # and other transformer attributes after the wrapped forward returns. + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("eager_model"), name) + + class ModelEngine(ABC): @abstractmethod @@ -598,6 +627,7 @@ def __init__( self._torch_compile_enabled = torch_compile_enabled self._torch_compile_piecewise_cuda_graph = torch_compile_piecewise_cuda_graph + self._torch_compile_context_only = False prefill_cuda_graph_num_tokens = self.llm_args.prefill_capture_num_tokens if prefill_cuda_graph_num_tokens is None: @@ -644,10 +674,15 @@ def __init__( "apply_llm_torch_compile", None) if isinstance(self.model, DecoderModelForCausalLM): - self.model.model = torch.compile( - self.model.model, + eager_model = self.model.model + compiled_model = torch.compile( + eager_model, backend=self._torch_compile_backend, fullgraph=torch_compile_fullgraph) + self._torch_compile_context_only = self._torch_compile_piecewise_cuda_graph + self.model.model = ( + _ContextOnlyCompiledModel(eager_model, compiled_model) + if self._torch_compile_context_only else compiled_model) elif callable(apply_llm_torch_compile): # TODO: Move this contract to MultimodalModelMixin once # multimodal models consistently expose their LLM compile @@ -2132,8 +2167,16 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: native_mxfp8_methods = [ method for method in mxfp8_methods if method.needs_native_autotune ] + compile_all_batches = ( + self._torch_compile_enabled + and not getattr(self, "_torch_compile_context_only", False)) + if compile_all_batches: + # Compiled auto dispatch uses native; do not tune unused backends. + # Context-only compile retains the eager generation-graph policy. + for method in mxfp8_methods: + method.disable_flashinfer_auto() use_mxfp8_flashinfer_graph_default = ( - self.cuda_graph_runner.enabled + self.cuda_graph_runner.enabled and not compile_all_batches 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())) @@ -6257,7 +6300,12 @@ def model_forward(self, **kwargs): if reclaimer is not None and not self.is_warmup and isinstance(metadata, TrtllmAttentionMetadata) else contextlib.nullcontext()) - with reclaim_scope: + # Scope the entire top-level forward, including Eagle3's epilogue, so + # eager decode and over-ceiling prefill do not select compile-only ops. + compile_scope = ( + torch_compiling(get_per_request_prefill_cuda_graph_flag()) + if self._torch_compile_context_only else contextlib.nullcontext()) + with reclaim_scope, compile_scope: if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): return trace_func(self.model.forward)(**kwargs) else: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index fbc15c15f2e6..335464ee8647 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7295,6 +7295,22 @@ def test_nvfp4(self, use_msa): # NVFP4 checkpoint: MXFP8 base layers with NVFP4 routed experts # (MIXED_PRECISION checkpoint). The MSA path runs an FP8 KV cache; the # Triton path keeps the KV cache in BF16. + self._run_nvfp4(use_msa) + + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("fuse_qkv_index_projection", [False, True]) + def test_nvfp4_piecewise_cuda_graph( + self, fuse_qkv_index_projection: bool) -> None: + self._run_nvfp4(True, + piecewise=True, + fuse_qkv_index_projection=fuse_qkv_index_projection) + + def _run_nvfp4(self, + use_msa: bool, + *, + piecewise: bool = False, + fuse_qkv_index_projection: bool = False) -> None: tp_size = ep_size = 4 model_name = "nvidia/MiniMax-M3-NVFP4" model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" @@ -7303,7 +7319,8 @@ def test_nvfp4(self, use_msa): dtype="fp8" if use_msa else "auto") sparse_attention_config = MiniMaxM3SparseAttentionConfig( implementation="msa" if use_msa else "triton", - indexer_kv_dtype="fp8" if use_msa else "bf16") + indexer_kv_dtype="fp8" if use_msa else "bf16", + fuse_qkv_index_projection=fuse_qkv_index_projection) moe_config = MoeConfig(backend="CUTLASS") with LLM(model_path, tensor_parallel_size=tp_size, @@ -7311,6 +7328,13 @@ def test_nvfp4(self, use_msa): kv_cache_config=kv_cache_config, sparse_attention_config=sparse_attention_config, moe_config=moe_config, + prefill_cuda_graph_backend=(PrefillCudaGraphBackend.PIECEWISE + if piecewise else + PrefillCudaGraphBackend.DISABLED), + prefill_capture_num_tokens=[128, 512, 2048] + if piecewise else None, + torch_compile_config=TorchCompileConfig() + if piecewise else None, max_seq_len=4096, trust_remote_code=True) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 38247c9ad825..85df459a034c 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -58,6 +58,9 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_piecewise_attention_boundary_runs_horizontal_producer - unittest/_torch/models/test_minimax_m3.py::test_piecewise_projection_fake_preserves_padded_hidden_rows - unittest/_torch/models/test_minimax_m3.py::test_piecewise_fused_projection_preserves_input_token_dimension + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captured_producer_preserves_symbolic_shapes + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captures_horizontal_producer_before_attention + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_unfused_indexer_keeps_cache_write_eager - unittest/_torch/models/test_minimax_m3.py::test_msa_attention_core_routes_compact_q_to_attention_dispatcher - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 61a9b0881f18..fa456bd01d1b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -54,6 +54,8 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_piecewise_cuda_graph[fuse_qkv_index_projection=False] + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_piecewise_cuda_graph[fuse_qkv_index_projection=True] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=False-eval_mode=inferencex] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] diff --git a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py index e58a7dda4b36..b759661c3a9c 100644 --- a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py +++ b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py @@ -20,6 +20,94 @@ from tensorrt_llm._torch.modules.rms_norm import RMSNorm +@pytest.mark.cpu_only +@pytest.mark.parametrize("producer", ["norm_rope", "main_kv", "horizontal"]) +def test_fp8_producer_meta_keeps_dynamic_num_tokens(producer: str) -> None: + """All FP8 fake kernels must retain the symbolic token dimension.""" + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + with FakeTensorMode(shape_env=ShapeEnv()) as mode: + num_tokens = mode.shape_env.create_unbacked_symint() + qkv = torch.empty((num_tokens, 1280), dtype=torch.bfloat16) + positions = torch.empty((num_tokens,), dtype=torch.int32) + slots = torch.empty((num_tokens,), dtype=torch.int32) + weight = torch.empty(128, dtype=torch.bfloat16) + kv_cache = torch.empty((2, 2, 1, 128, 128), dtype=torch.float8_e4m3fn) + if producer == "norm_rope": + output = torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + 8, + 1, + 1, + 128, + 64, + 1e-5, + weight, + weight, + 10000.0, + True, + positions, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + outputs, tails = [output], [(1280,)] + elif producer == "main_kv": + output = torch.ops.trtllm.minimax_m3_fp8_qk_norm_rope_kv_insert( + qkv, + kv_cache, + slots, + 8, + 1, + 1, + 128, + 64, + 1e-5, + weight, + weight, + 10000.0, + True, + positions, + ) + outputs, tails = [output], [(8, 128)] + else: + packed = torch.empty((num_tokens, 1536), dtype=torch.bfloat16) + index_cache = torch.empty((2, 1, 128, 128), dtype=torch.float8_e4m3fn) + rope = torch.empty((16, 2, 32), dtype=torch.float32) + outputs = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + kv_cache, + index_cache, + slots, + 8, + 1, + 1, + 128, + 64, + 1e-5, + weight, + weight, + weight, + weight, + rope, + positions, + ) + tails = [(8, 128), (1, 128)] + + for output, tail in zip(outputs, tails): + assert isinstance(output.shape[0], torch.SymInt) + assert output.shape[0].node.expr == num_tokens.node.expr + assert output.shape[1:] == tail + assert output.dtype == torch.float8_e4m3fn + + @torch.inference_mode() def torch_ref_rms_norm_rope( qkv, diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index dc202f5e1658..53931d31a933 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -39,6 +39,46 @@ from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig +@pytest.mark.cpu_only +def test_msa_metadata_clears_padded_cache_slot_tail(monkeypatch: pytest.MonkeyPatch) -> None: + """A smaller replay must not reuse the previous step's live cache slots.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend + + monkeypatch.setattr(msa_backend, "maybe_pin_memory", lambda tensor: tensor) + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata._msa_buffers_ready = True + metadata.request_ids = [0] + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=4, + get_buffers=lambda layer_idx: torch.empty(1), + get_block_ids_per_seq=lambda request_ids: torch.tensor([[3]], dtype=torch.int32), + ) + metadata.msa_out_cache_loc = torch.full((4,), 99, dtype=torch.int32) + metadata.msa_block_table = torch.zeros((1, 1), dtype=torch.int32) + metadata.msa_seq_lens_cuda = torch.zeros(1, dtype=torch.int32) + metadata.msa_subpage_block_table = None + metadata._msa_runs_no_fmha = lambda: True + metadata._msa_kv_lens_may_change = lambda: False + original_ptr = metadata.msa_out_cache_loc.data_ptr() + + for count in (4, 2, 1): + metadata._msa_qo_lens_cpu = torch.tensor([count], dtype=torch.int32) + metadata._msa_kv_lens_cpu = metadata._msa_qo_lens_cpu.clone() + metadata._msa_qo_offset_cpu = torch.zeros(1, dtype=torch.int32) + metadata._build_msa_fields() + assert metadata.msa_out_cache_loc.tolist() == list(range(12, 12 + count)) + [-1] * ( + 4 - count + ) + assert metadata.msa_out_cache_loc.data_ptr() == original_ptr + assert metadata._msa_fields_ready + + metadata.request_ids = [] + metadata._msa_qo_lens_cpu = torch.empty(0, dtype=torch.int32) + metadata._build_msa_fields() + assert metadata.msa_out_cache_loc.tolist() == [-1] * 4 + + def test_msa_package_availability_installs_cutlass_compatibility_aliases(monkeypatch): from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( msa_package_available, 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 7dd860b72556..5e9693abb810 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -22,15 +22,20 @@ import torch import tensorrt_llm +import tensorrt_llm._torch.pyexecutor.model_engine as model_engine_module from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.model_engine import ( + PyTorchModelEngine, + _ContextOnlyCompiledModel, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.speculative.utils import update_draft_len +from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig from tensorrt_llm.llmapi.llm_args import ( DecodingBaseConfig, @@ -41,6 +46,146 @@ from tensorrt_llm.mapping import Mapping +@pytest.mark.cpu_only +@pytest.mark.parametrize("local_contexts", [0, 1]) +def test_context_only_compile_uses_all_rank_prefill_decision( + monkeypatch: pytest.MonkeyPatch, + local_contexts: int, +) -> None: + eager = torch.nn.Linear(4, 4) + compiled = torch.nn.Module() + compiled.shared = eager + expected = torch.ones((2, 4)) + eager.forward = Mock(return_value=expected) + compiled.forward = Mock(return_value=expected) + router = _ContextOnlyCompiledModel(eager, compiled) + assert router.weight is eager.weight + assert list(router.parameters()) == list(eager.parameters()) + # Local decode-only ranks participate when another attention-DP rank + # prefills; an over-ceiling local context batch uses eager instead. + metadata = SimpleNamespace(num_contexts=local_contexts) + for eligible in (True, False): + monkeypatch.setattr( + model_engine_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + ) + assert router(expected, attn_metadata=metadata) is expected + selected = compiled if eligible else eager + selected.forward.assert_called_once_with(expected, attn_metadata=metadata) + assert compiled.forward.call_count == eager.forward.call_count == 1 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("eligible", [False, True]) +@pytest.mark.parametrize("raises", [False, True]) +def test_context_only_compile_scopes_whole_model_forward( + monkeypatch: pytest.MonkeyPatch, + eligible: bool, + raises: bool, +) -> None: + observed = [] + + def forward(**kwargs: object) -> str: + observed.append(is_torch_compiling()) + # This represents work after the transformer, such as Eagle3 drafting. + if raises: + raise RuntimeError("epilogue failure") + observed.append(is_torch_compiling()) + return "done" + + engine = SimpleNamespace( + model=SimpleNamespace(model_config=SimpleNamespace(extra_attrs={}), forward=forward), + _torch_compile_backend=None, + _torch_compile_context_only=True, + _eager_workspace_reclaimer=None, + is_warmup=False, + ) + monkeypatch.setattr(model_engine_module, "get_model_extra_attrs", lambda: {}) + monkeypatch.setattr( + model_engine_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + ) + monkeypatch.setattr(model_engine_module, "is_trace_enabled", lambda name: False) + with torch_compiling(True): + if raises: + with pytest.raises(RuntimeError, match="epilogue failure"): + PyTorchModelEngine.model_forward(engine, attn_metadata=Mock()) + else: + assert PyTorchModelEngine.model_forward(engine, attn_metadata=Mock()) == "done" + assert is_torch_compiling() + assert observed == [eligible] * (1 if raises else 2) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("context_only", [False, True]) +@pytest.mark.parametrize("backend", [None, "auto", "flashinfer"]) +def test_compiled_mxfp8_warmup_backend_selection( + monkeypatch: pytest.MonkeyPatch, + context_only: bool, + backend: str | None, +) -> None: + """PCG retains eager decode tuning; all-batch compile settles auto on native.""" + import tensorrt_llm._torch.modules.linear as linear_module + + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + monkeypatch.delenv("TLLM_AUTOTUNER_CACHE_PATH", raising=False) + if backend is not None: + monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", backend) + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + monkeypatch.setattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", Mock(), raising=False) + flashinfer_tune = Mock(return_value=contextlib.nullcontext()) + monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(autotune=flashinfer_tune)) + method = MXFP8LinearMethod() + engine = SimpleNamespace( + llm_args=SimpleNamespace(enable_autotuner=True), + _torch_compile_enabled=True, + _torch_compile_context_only=context_only, + cuda_graph_runner=SimpleNamespace(enabled=True), + model=SimpleNamespace( + modules=lambda: [ + SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), + SimpleNamespace(quant_method=method), + ] + ), + mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False), + dist=object(), + kv_cache_manager_key="kv_cache", + max_num_tokens=16, + batch_size=16, + max_seq_len=2, + original_max_draft_len=0, + max_total_draft_tokens=0, + is_draft_model=False, + guided_decoder=None, + no_cuda_graph=lambda: contextlib.nullcontext(), + _create_warmup_request=Mock(return_value=object()), + _release_batch_context=lambda *args: contextlib.nullcontext(object()), + _should_run_warmup_batch=Mock(return_value=True), + _release_megamoe_profiling_scratch=Mock(), + forward=Mock(), + ) + cache = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) + resources = SimpleNamespace( + get_resource_manager=lambda key: cache if key == "kv_cache" else None + ) + tuner = Mock(profiling_cache={}) + monkeypatch.setattr(model_engine_module.AutoTuner, "get", lambda: tuner) + monkeypatch.setattr(model_engine_module, "autotune", lambda **kwargs: contextlib.nullcontext()) + monkeypatch.setattr(MXFP8GemmRunner, "sync_all_tactic_caches", Mock()) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(model_engine_module, "clear_memory_buffers", lambda: None) + + PyTorchModelEngine._run_autotuner_warmup(engine, resources) + + flashinfer_expected = context_only or backend == "flashinfer" + assert method.backend == ( + "flashinfer" if backend == "flashinfer" else "auto" if context_only else "trtllm" + ) + assert method._native_autotuned + assert method._flashinfer_autotuned == flashinfer_expected + assert flashinfer_tune.call_count == int(flashinfer_expected) + assert engine.forward.call_count == (4 if flashinfer_expected else 2) + + @pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig]) def test_warmup_overrides_dynamic_draft_length(config_cls): config = config_cls(max_draft_len=3, speculative_model="dummy", draft_len_schedule={1: 3, 4: 2}) @@ -424,6 +569,7 @@ def flashinfer_autotune(): engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=False), cuda_graph_runner=SimpleNamespace(enabled=True), + _torch_compile_enabled=False, model=SimpleNamespace( modules=lambda: [ SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), @@ -485,6 +631,7 @@ def flashinfer_autotune(): engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), cuda_graph_runner=SimpleNamespace(enabled=True), + _torch_compile_enabled=False, model=SimpleNamespace( modules=lambda: [ SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), @@ -598,6 +745,7 @@ def flashinfer_autotune(): engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), cuda_graph_runner=SimpleNamespace(enabled=True), + _torch_compile_enabled=False, model=SimpleNamespace( modules=lambda: [ SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), @@ -696,6 +844,7 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), cuda_graph_runner=SimpleNamespace(enabled=True), + _torch_compile_enabled=False, model=SimpleNamespace( modules=lambda: [ SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), @@ -762,6 +911,7 @@ def test_native_mxfp8_respects_disabled_global_autotuner(self): engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=False), cuda_graph_runner=SimpleNamespace(enabled=False), + _torch_compile_enabled=False, model=SimpleNamespace(modules=lambda: [SimpleNamespace(quant_method=method)]), ) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 59564a029292..0f016899f536 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -483,6 +483,7 @@ def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, ou layer = SimpleNamespace( enable_fused_qkv_index_projection=True, qkv_proj=lambda hidden_states: packed, + attn=object(), # Compatibility path, without the captured FP8 producer. register_to_config=True, num_heads=1, head_dim=3, @@ -517,6 +518,121 @@ def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, ou assert result.shape == (position_ids.shape[-1], 3) +@pytest.mark.cpu_only +def test_piecewise_captured_producer_preserves_symbolic_shapes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + layer = SimpleNamespace(q_size=1024, index_q_size=128) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (None, layer), + ) + with FakeTensorMode(shape_env=ShapeEnv()) as mode: + num_tokens = mode.shape_env.create_unbacked_symint() + hidden = torch.empty((num_tokens, 512), dtype=torch.bfloat16) + positions = torch.empty((1, num_tokens), dtype=torch.int32) + q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer(hidden, positions, "3") + for output, width in ((q, 1024), (idx_q, 128)): + assert output.shape[0].node.expr == num_tokens.node.expr + assert output.shape[1] == width + assert output.dtype == torch.float8_e4m3fn + + +@pytest.mark.cpu_only +def test_piecewise_captures_horizontal_producer_before_attention( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import Mock + + backend = object.__new__(MiniMaxM3MsaSparseAttention) + backend.indexer_kv_dtype = "fp8" + packed = torch.empty((4, 7), dtype=torch.bfloat16) + q = torch.empty((4, 3), dtype=torch.float8_e4m3fn) + idx_q = torch.empty((4, 1), dtype=torch.float8_e4m3fn) + output = torch.empty((4, 3), dtype=torch.bfloat16) + metadata = SimpleNamespace(num_tokens=2) + layer = SimpleNamespace( + enable_fused_qkv_index_projection=True, + register_to_config=True, + attn=backend, + _emit_fp8_main_qkv=lambda: True, + layer_idx_str="3", + qkv_proj=Mock(return_value=packed), + _fused_fp8_qkv_indexer_norm_rope_kv_insert=Mock(return_value=(q, idx_q)), + _forward_attention_core=Mock(return_value=output), + o_proj=lambda output, all_reduce_params: output, + ) + monkeypatch.setattr(modeling_minimaxm3, "is_torch_compiling", lambda: True) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (metadata, layer), + ) + hidden = torch.empty((4, 5), dtype=torch.bfloat16) + positions = torch.arange(4).reshape(1, 4) + result = MiniMaxM3Attention._sparse_forward(layer, positions, hidden, metadata) + assert result is output + layer.qkv_proj.assert_called_once_with(hidden) + layer._fused_fp8_qkv_indexer_norm_rope_kv_insert.assert_called_once_with( + packed, positions, metadata + ) + layer._forward_attention_core.assert_called_once_with(q, None, None, idx_q, None, metadata) + + +@pytest.mark.cpu_only +def test_piecewise_unfused_indexer_keeps_cache_write_eager(monkeypatch: pytest.MonkeyPatch) -> None: + from unittest.mock import Mock + + backend = object.__new__(MiniMaxM3MsaSparseAttention) + backend.indexer_kv_dtype = "fp8" + q, k, v, idx_q, idx_k = [torch.empty((4, 128), dtype=torch.float8_e4m3fn) for _ in range(5)] + metadata = SimpleNamespace(num_tokens=2) + layer = SimpleNamespace( + enable_fused_qkv_index_projection=False, + register_to_config=True, + attn=backend, + _emit_fp8_main_qkv=lambda: True, + qkv_proj=Mock(return_value=torch.empty((4, 384), dtype=torch.bfloat16)), + index_qk_proj=Mock(return_value=torch.empty((4, 256), dtype=torch.bfloat16)), + _fused_qk_norm_rope=Mock( + side_effect=[torch.cat((q, k, v), dim=-1), torch.cat((idx_q, idx_k), dim=-1)] + ), + _fused_fp8_index_qk_norm_rope=Mock(), + _split_main_qkv=lambda tensor: (q, k, v), + _split_index_qk=lambda tensor: (idx_q, idx_k), + num_heads=1, + num_key_value_heads=1, + head_dim=128, + sparse_num_index_heads=1, + sparse_index_dim=128, + q_norm=object(), + k_norm=object(), + index_q_norm=object(), + index_k_norm=object(), + ln_events=(None, None), + aux_stream=None, + _forward_attention_core=Mock(return_value=torch.empty((4, 128), dtype=torch.bfloat16)), + o_proj=lambda output, all_reduce_params: output, + ) + monkeypatch.setattr(modeling_minimaxm3, "is_torch_compiling", lambda: True) + monkeypatch.setattr( + modeling_minimaxm3, + "maybe_execute_in_parallel", + lambda first, second, *args, **kwargs: (first(), second()), + ) + MiniMaxM3Attention._sparse_forward( + layer, torch.arange(4).reshape(1, 4), torch.empty((4, 128), dtype=torch.bfloat16), metadata + ) + layer._fused_fp8_index_qk_norm_rope.assert_not_called() + assert layer._fused_qk_norm_rope.call_count == 2 + assert all(call.kwargs["out_fp8"] for call in layer._fused_qk_norm_rope.call_args_list) + layer._forward_attention_core.assert_called_once_with(q, k, v, idx_q, idx_k, metadata) + + @pytest.mark.cpu_only def test_msa_attention_core_routes_compact_q_to_attention_dispatcher() -> None: selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index 6ebb346d8610..f5e5e8ca6507 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -76,7 +76,10 @@ def test_mxfp8_dispatch_returns_mxfp8_method(monkeypatch): assert not method.use_native_autotuner -def _mock_mxfp8_ops(monkeypatch): +def _mock_mxfp8_ops( + monkeypatch: pytest.MonkeyPatch, + flashinfer_gemm: Mock | None = None, +) -> tuple[torch.Tensor, torch.Tensor, Mock, Mock, torch.Tensor, Mock, torch.Tensor]: quantized = torch.empty((2, 4), dtype=torch.float8_e4m3fn) activation_scale = torch.empty(512, dtype=torch.uint8) quantize = Mock(return_value=(quantized, activation_scale)) @@ -88,6 +91,7 @@ def _mock_mxfp8_ops(monkeypatch): mxfp8_quantize=quantize, mxfp8_mxfp8_gemm=native_gemm, mxfp8_mxfp8_gemm_autotuned=autotuned_gemm, + flashinfer_mm_mxfp8=flashinfer_gemm or Mock(), ) fake_torch = SimpleNamespace( ops=SimpleNamespace(trtllm=fake_trtllm_ops), @@ -107,7 +111,7 @@ def _mock_mxfp8_ops(monkeypatch): def test_mxfp8_flashinfer_call_contract(monkeypatch): - """The forced backend reuses TRT tensors and a zero-copy weight transpose.""" + """The forced backend passes native-layout tensors to the opaque op.""" monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", "flashinfer") monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) @@ -118,7 +122,7 @@ def test_mxfp8_flashinfer_call_contract(monkeypatch): "flashinfer", SimpleNamespace(mm_mxfp8=mm_mxfp8, autotune=Mock()), ) - quantized, activation_scale, quantize, _, _, _, _ = _mock_mxfp8_ops(monkeypatch) + quantized, activation_scale, quantize, _, _, _, _ = _mock_mxfp8_ops(monkeypatch, mm_mxfp8) weight = torch.empty((3, 4), dtype=torch.float8_e4m3fn) weight_scale = torch.empty(512, dtype=torch.uint8) @@ -133,15 +137,11 @@ def test_mxfp8_flashinfer_call_contract(monkeypatch): 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[1] is activation_scale + assert args[2] is weight assert args[3] is weight_scale - assert kwargs == { - "out_dtype": torch.bfloat16, - "use_8x4_sf_layout": False, - "backend": "cutlass", - } + assert args[4] == torch.bfloat16 + assert kwargs == {} def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): @@ -156,8 +156,9 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): SimpleNamespace(mm_mxfp8=mm_mxfp8, autotune=Mock()), ) _, _, _, native_gemm, native_output, autotuned_gemm, autotuned_output = _mock_mxfp8_ops( - monkeypatch + monkeypatch, mm_mxfp8 ) + monkeypatch.setattr(linear_module, "is_torch_compiling", lambda: False) module = SimpleNamespace( weight=torch.empty((3, 4), dtype=torch.float8_e4m3fn), @@ -187,6 +188,40 @@ def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): assert native_gemm.call_count == 2 +@pytest.mark.cpu_only +@pytest.mark.parametrize("backend", ["auto", "flashinfer"]) +def test_mxfp8_compile_skips_context_dispatch( + monkeypatch: pytest.MonkeyPatch, backend: str +) -> None: + """Compilation must not read Python ContextVars, even after decode tuning.""" + monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", backend) + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + monkeypatch.setattr(linear_module, "is_torch_compiling", lambda: True) + monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(autotune=Mock())) + flashinfer_output = torch.empty((2, 3), dtype=torch.bfloat16) + flashinfer_gemm = Mock(return_value=flashinfer_output) + _, _, _, native_gemm, native_output, _, _ = _mock_mxfp8_ops(monkeypatch, flashinfer_gemm) + for name in ( + "_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE", + "_FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE", + ): + monkeypatch.setattr( + linear_module, name, SimpleNamespace(get=Mock(side_effect=AssertionError)) + ) + method = MXFP8LinearMethod() + method.tune_decode_graph_backends = True + method.mark_flashinfer_autotuned() + module = SimpleNamespace( + weight=torch.empty((3, 4), dtype=torch.float8_e4m3fn), + weight_scale=torch.empty(512, dtype=torch.uint8), + dtype=torch.bfloat16, + ) + result = method.apply(module, torch.empty((2, 4), dtype=torch.bfloat16), bias=None) + assert result is (native_output if backend == "auto" else flashinfer_output) + assert native_gemm.call_count == (backend == "auto") + assert flashinfer_gemm.call_count == (backend == "flashinfer") + + 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) From ed2bd0c249f997286daeaa299652c8f4245e160f Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:25:25 -0700 Subject: [PATCH 02/12] [None][fix] Address MiniMax-M3 PCG review findings Expose both cache mutations to compilation and preserve producer outputs during in-place recovery. Keep checkpoint module names transparent for partial reloads. Add focused regressions for all five review findings and document touched function contracts. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../backends/sparse/minimax_m3/msa_backend.py | 11 ++ .../_torch/compilation/remove_copy_pass.py | 26 ++- tensorrt_llm/_torch/compilation/utils.py | 7 + .../_torch/custom_ops/cpp_custom_ops.py | 4 + .../custom_ops/flashinfer_custom_ops.py | 2 + .../_torch/models/modeling_minimaxm3.py | 55 ++++-- tensorrt_llm/_torch/modules/linear.py | 2 + .../_torch/pyexecutor/model_engine.py | 22 ++- .../defs/accuracy/test_llm_api_pytorch.py | 3 + .../integration/test_lists/test-db/l0_cpu.yml | 3 + .../attention/sparse/msa/test_msa_backend.py | 28 +++ .../compilation/test_remove_copy_pass.py | 57 +++++- .../test_pytorch_model_engine_warmup.py | 51 ++++++ .../unittest/_torch/models/test_minimax_m3.py | 162 +++++++++++++++++- .../_torch/modules/test_mxfp8_linear.py | 62 ++++++- 15 files changed, 471 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index add137126681..02610e5ad282 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -152,6 +152,9 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # Graph-stable buffers; consumers slice to the live count at the call # site. Filled once the current step's cache write is prepared. msa_out_cache_loc: Optional[torch.Tensor] = None + # Zero-copy pool views prepared outside Dynamo; PCG passes these explicitly + # to its mutable producer instead of hiding writes behind runtime metadata. + msa_layer_cache_tensors: Optional[dict[int, tuple[torch.Tensor, torch.Tensor]]] = None msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None @@ -399,6 +402,14 @@ def _create_msa_buffers(self) -> None: self._msa_buffers_ready = False if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): return + self.msa_layer_cache_tensors = { + layer_idx: ( + kv_cache_manager.get_buffers(layer_idx, kv_layout="HND"), + self.msa_idx_k_cache(layer_idx), + ) + for layer_idx in getattr(kv_cache_manager, "sparse_layer_ids", ()) + if layer_idx in kv_cache_manager.layer_offsets + } capture_graph = self.is_cuda_graph buffers = self.cuda_graph_buffers max_num_sequences = int(self.max_num_sequences) diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py index f9bb925a1e1c..ab253aea9ee3 100644 --- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py +++ b/tensorrt_llm/_torch/compilation/remove_copy_pass.py @@ -32,6 +32,7 @@ def remove_copy_for_mutates_args(graph: Graph): nodes_to_remove: list[Node] = [] def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): + """Restore in-place calls while preserving regular and mutated outputs.""" getitem_nodes = [ user for user in node.users if is_call_function(user, getitem) ] @@ -69,8 +70,30 @@ def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): kwargs[arg.name] = (None if base_index is None else all_bases[base_index]) + with graph.inserting_before(node): + inplace_node = graph.call_function(inplace_func, kwargs=kwargs) + num_returns = len(inplace_func._schema.returns) + if num_returns: + inplace_node.meta = node.meta.copy() + for key in ("val", "example_value"): + if key in node.meta: + values = node.meta[key] + inplace_node.meta[key] = (values[0] if num_returns == 1 else + values[:num_returns]) + for getitem_node in getitem_nodes: idx = getitem_node.args[1] + if idx < num_returns: + # Mutable producers can also return fresh tensors. Preserve + # those values while reconnecting the cache mutation outputs. + with graph.inserting_before(node): + replacement = (inplace_node + if num_returns == 1 else graph.call_function( + getitem, args=(inplace_node, idx))) + replacement.meta = getitem_node.meta.copy() + getitem_node.replace_all_uses_with(replacement) + nodes_to_remove.append(getitem_node) + continue if idx in tensor_list_replacements: mutated_arg, replacement = tensor_list_replacements[idx] else: @@ -82,9 +105,6 @@ def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): getitem_node.replace_all_uses_with(replacement) nodes_to_remove.append(getitem_node) - with graph.inserting_before(node): - graph.call_function(inplace_func, kwargs=kwargs) - nodes_to_remove.append(node) for node in graph.nodes: diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index 7fcd8464f898..2e98607b84b6 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/tensorrt_llm/_torch/compilation/utils.py @@ -64,6 +64,7 @@ def capture_piecewise_cuda_graph(enable: bool): def inplace_info(): + """Map functionalized mutation outputs to their original argument names.""" inplace_map = { torch.ops.trtllm.flashinfer_fused_add_rmsnorm.default: { 1: "input", @@ -222,6 +223,12 @@ def inplace_info(): "minimax_m3_attn_custom_op_inplace": { 1: "output" }, + # The ordinary outputs are compact Q/index-Q; the next + # two outputs of auto_functionalized are the mutated paged caches. + "minimax_m3_fused_sparse_qkv_producer": { + 2: "kv_cache", + 3: "index_k_cache" + }, "fused_sigmoid_mul_inplace": { 1: "attention_output" }, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 042012c99c1f..e8496115fa30 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -15,6 +15,7 @@ def _register_fake(): + """Register shape-only implementations for native operators during tracing.""" @torch.library.register_fake("trtllm::allreduce") def allreduce( @@ -405,6 +406,7 @@ def _(qkv: torch.Tensor, num_heads_q: int, num_heads_k: int, high: float, attention_factor: float, is_qk_norm: bool, use_gemma: bool, use_mrope: bool, mrope_section1: int, mrope_section2: int) -> torch.Tensor: + """Infer FP8 QKV output geometry while preserving symbolic token counts.""" del rotary_dim, eps, q_weight, k_weight, base, is_neox, position_ids del factor, low, high, attention_factor, is_qk_norm, use_gemma del use_mrope, mrope_section1, mrope_section2 @@ -419,6 +421,7 @@ def _(qkv: torch.Tensor, kv_cache: torch.Tensor, num_heads_v: int, head_dim: int, rotary_dim: int, eps: float, q_weight: torch.Tensor, k_weight: torch.Tensor, base: float, is_neox: bool, position_ids: torch.Tensor) -> torch.Tensor: + """Infer FP8 query geometry without performing the KV-cache write.""" del kv_cache, out_cache_loc, num_heads_k, num_heads_v, rotary_dim, eps del q_weight, k_weight, base, is_neox, position_ids return qkv.new_empty((qkv.shape[0], num_heads_q, head_dim), @@ -433,6 +436,7 @@ def _(packed: torch.Tensor, kv_cache: torch.Tensor, k_weight: torch.Tensor, index_q_weight: torch.Tensor, index_k_weight: torch.Tensor, rotary_cos_sin: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Infer main and index query shapes without mutating either cache.""" del kv_cache, index_k_cache, out_cache_loc, num_heads_kv, rotary_dim del eps, q_weight, k_weight, index_q_weight, index_k_weight del rotary_cos_sin, position_ids diff --git a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py index e1a0a3917bca..7412b4b0d519 100644 --- a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py @@ -146,6 +146,7 @@ def flashinfer_mm_mxfp8(act: torch.Tensor, act_scale: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, output_dtype: torch.dtype) -> torch.Tensor: + """Run CUTLASS MXFP8 GEMM with row-major weights and swizzled scales.""" # Argument order mirrors trtllm::mxfp8_mxfp8_gemm: weight arrives as # [N, K] and mm_mxfp8 wants [K, N]. Both scale buffers are the 1D # padded swizzled CUTLASS layout, hence use_8x4_sf_layout=False. @@ -161,5 +162,6 @@ def flashinfer_mm_mxfp8(act: torch.Tensor, act_scale: torch.Tensor, def _(act: torch.Tensor, act_scale: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, output_dtype: torch.dtype) -> torch.Tensor: + """Infer the GEMM output shape and dtype without invoking FlashInfer.""" return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index a93d2f8468cf..066b8dba2416 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -838,36 +838,49 @@ def _minimax_m3_qkv_index_proj_fake( return hidden_states.new_empty((hidden_states.shape[0], sum(qkv_proj.local_output_sizes))) -@torch.library.custom_op("trtllm::minimax_m3_fused_sparse_qkv_producer", mutates_args=()) +@torch.library.custom_op( + "trtllm::minimax_m3_fused_sparse_qkv_producer", + mutates_args=("kv_cache", "index_k_cache"), +) def minimax_m3_fused_sparse_qkv_producer( hidden_states: torch.Tensor, position_ids: Optional[torch.Tensor], + kv_cache: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, layer_idx: str, -) -> List[torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor]: """Capture projection, norm, RoPE and FP8 cache insertion together.""" attn_metadata, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) packed = attn_layer.qkv_proj(hidden_states) result = attn_layer._fused_fp8_qkv_indexer_norm_rope_kv_insert( - packed, position_ids, attn_metadata + packed, + position_ids, + attn_metadata, + cache_tensors=(kv_cache, index_k_cache, out_cache_loc), ) if result is None: raise RuntimeError("MiniMax-M3 piecewise graph requires the fused FP8 sparse QKV producer.") - return list(result) + return result @minimax_m3_fused_sparse_qkv_producer.register_fake def _minimax_m3_fused_sparse_qkv_producer_fake( hidden_states: torch.Tensor, position_ids: Optional[torch.Tensor], + kv_cache: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, layer_idx: str, -) -> List[torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor]: + """Infer FP8 query shapes while retaining the symbolic token dimension.""" del position_ids _, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) num_tokens = hidden_states.shape[0] - return [ + return ( hidden_states.new_empty((num_tokens, attn_layer.q_size), dtype=torch.float8_e4m3fn), hidden_states.new_empty((num_tokens, attn_layer.index_q_size), dtype=torch.float8_e4m3fn), - ] + ) @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",)) @@ -1324,6 +1337,8 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( packed: torch.Tensor, position_ids: Optional[torch.Tensor], attn_metadata: AttentionMetadata, + *, + cache_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: """Run the vLLM-style horizontal producer for every sparse batch. @@ -1369,12 +1384,15 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): return None - kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) - if kv_cache_manager is None: - return None - buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") - index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) - out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + if cache_tensors is None: + kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) + if kv_cache_manager is None: + return None + buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") + index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + else: + buffers, index_k_cache, out_cache_loc = cache_tensors num_tokens = int(packed.shape[0]) supported_main_cache = ( buffers is not None @@ -1968,8 +1986,16 @@ def _sparse_forward( and self._emit_fp8_main_qkv() and self.attn.indexer_kv_dtype == "fp8" ): + # Metadata stages these zero-copy views before compilation; + # tracing the cache manager's native pointer access is unsafe. + kv_cache, index_k_cache = attn_metadata.msa_layer_cache_tensors[self.layer_idx] q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( - hidden_states, position_ids, self.layer_idx_str + hidden_states, + position_ids, + kv_cache, + index_k_cache, + attn_metadata.msa_out_cache_loc, + self.layer_idx_str, ) o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) return self.o_proj(o, all_reduce_params=all_reduce_params) @@ -2044,6 +2070,7 @@ def _main_norm_rope(): return q, k, v def _index_norm_rope(): + """Project and normalize index queries, with RoPE and cache updates.""" idx_qk = ( packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) ) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index d8cbf963c468..f3b59efae565 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3334,6 +3334,7 @@ def needs_native_autotune(self) -> bool: and self.use_cutlass) def _load_flashinfer(self, *, required: bool) -> bool: + """Load the optional GEMM backend, raising only when explicitly required.""" if not self.use_cutlass: if required: raise RuntimeError( @@ -3424,6 +3425,7 @@ def create_weights(self, module: Linear, in_features: int, def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + """Apply MXFP8 linear projection with eager or capture-safe dispatch.""" original_shape = input.shape if input.dim() > 2: input = input.reshape(-1, input.shape[-1]) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index cb2a4c5aa1d7..5b227a7664fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -10,8 +10,8 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, - Union, cast) +from typing import (Any, Callable, Dict, Iterator, List, Optional, Sequence, + Tuple, Type, Union, cast) import torch import torch._dynamo.config @@ -205,17 +205,33 @@ class _ContextOnlyCompiledModel(torch.nn.Module): def __init__(self, eager_model: torch.nn.Module, compiled_model: torch.nn.Module) -> None: + """Keep eager and compiled entry points sharing the same model weights.""" super().__init__() self.eager_model = eager_model self.compiled_model = compiled_model + def named_modules( + self, + memo: Optional[set[torch.nn.Module]] = None, + prefix: str = "", + remove_duplicate: bool = True, + ) -> Iterator[Tuple[str, torch.nn.Module]]: + """Expose checkpoint-compatible module names for partial weight reloads.""" + # Weight reloads match checkpoint prefixes against this traversal. + # Neither routing wrapper owns parameters; expose the original tree + # once, including when the loader requests remove_duplicate=False. + yield from self.eager_model.named_modules(memo, prefix, + remove_duplicate) + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Use the compiled path only for globally eligible prefill batches.""" model = (self.compiled_model if get_per_request_prefill_cuda_graph_flag() else self.eager_model) return model(*args, **kwargs) def __getattr__(self, name: str) -> Any: + """Delegate model-specific attributes to the original eager model.""" # Model-specific epilogues (including M3 Eagle3) access embed_tokens # and other transformer attributes after the wrapped forward returns. try: @@ -362,6 +378,7 @@ def __init__( model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): + """Initialize model execution, cache management, and graph configuration.""" _configure_deep_gemm_pdl() self.forward_pass_callable = None @@ -6281,6 +6298,7 @@ def capture_postprocess_fn(inputs: Dict[str, Any]): return outputs def model_forward(self, **kwargs): + """Run the full model under the current batch's compile and graph flags.""" attrs = get_model_extra_attrs() assert attrs is not None, "Model extra attrs is not set" attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 335464ee8647..cd1305fd6d2e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7292,6 +7292,7 @@ def test_mxfp8_piecewise_cuda_graph(self, use_msa): @pytest.mark.skip_less_device_memory(140000) @parametrize_with_ids("use_msa", [False, True]) def test_nvfp4(self, use_msa): + """Check mixed-precision M3 accuracy with MSA or Triton attention.""" # NVFP4 checkpoint: MXFP8 base layers with NVFP4 routed experts # (MIXED_PRECISION checkpoint). The MSA path runs an FP8 KV cache; the # Triton path keeps the KV cache in BF16. @@ -7302,6 +7303,7 @@ def test_nvfp4(self, use_msa): @parametrize_with_ids("fuse_qkv_index_projection", [False, True]) def test_nvfp4_piecewise_cuda_graph( self, fuse_qkv_index_projection: bool) -> None: + """Check PCG accuracy with separate or fused QKV and index projections.""" self._run_nvfp4(True, piecewise=True, fuse_qkv_index_projection=fuse_qkv_index_projection) @@ -7311,6 +7313,7 @@ def _run_nvfp4(self, *, piecewise: bool = False, fuse_qkv_index_projection: bool = False) -> None: + """Run the shared four-GPU NVFP4 M3 accuracy workload.""" tp_size = ep_size = 4 model_name = "nvidia/MiniMax-M3-NVFP4" model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 85df459a034c..4b7666b6552d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -38,6 +38,7 @@ l0_cpu: # runs -m cpu_only, so this entry contributes only the 3 marked files under attention/. - unittest/_torch/attention - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py + - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs - unittest/_torch/distributed - unittest/_torch/executor - unittest/_torch/disaggregation @@ -60,6 +61,8 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_piecewise_fused_projection_preserves_input_token_dimension - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captured_producer_preserves_symbolic_shapes - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captures_horizontal_producer_before_attention + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captured_producer_rejects_unavailable_fusion + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captured_producer_declares_cache_mutations - unittest/_torch/models/test_minimax_m3.py::test_piecewise_unfused_indexer_keeps_cache_write_eager - unittest/_torch/models/test_minimax_m3.py::test_msa_attention_core_routes_compact_q_to_attention_dispatcher - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 53931d31a933..878fb782889f 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -312,6 +312,34 @@ def test_msa_buffers_include_graph_stable_block_table(): assert requested["msa_seq_lens_cuda"] == ((MAX_NUM_SEQUENCES,), torch.int32, True) +@pytest.mark.cpu_only +def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch) -> None: + """Stage zero-copy cache views only for sparse layers on the local rank.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend + + metadata = _buffer_metadata(sparse_layer_ids=[3, 4], layer_offsets={3: 0}) + main_cache = torch.zeros(2, 2, 1, 128, 128) + index_cache = torch.zeros(2, 1, 128, 128) + manager = metadata.kv_cache_manager + manager.get_buffers = lambda layer_idx, kv_layout: main_cache + manager.get_index_k_buffer = lambda layer_idx, kv_layout: index_cache + monkeypatch.setattr( + metadata, + "get_empty", + lambda buffers, shape, **kwargs: torch.empty(shape, dtype=kwargs["dtype"]), + ) + # No native pool in this CPU test; only zero-copy cache-view staging is under test. + monkeypatch.setattr(msa_backend, "uniform_subpages_per_slot", lambda manager: 0) + metadata._create_msa_buffers() + assert set(metadata.msa_layer_cache_tensors) == {3} + main, index = metadata.msa_layer_cache_tensors[3] + assert main is main_cache and index is index_cache + main.fill_(2) + index.fill_(3) + torch.testing.assert_close(main_cache, torch.full_like(main_cache, 2)) + torch.testing.assert_close(index_cache, torch.full_like(index_cache, 3)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( ("factors", "expected"), diff --git a/tests/unittest/_torch/compilation/test_remove_copy_pass.py b/tests/unittest/_torch/compilation/test_remove_copy_pass.py index 26c62db92dd2..4e7740792a19 100644 --- a/tests/unittest/_torch/compilation/test_remove_copy_pass.py +++ b/tests/unittest/_torch/compilation/test_remove_copy_pass.py @@ -17,7 +17,7 @@ import pytest import torch -from torch._higher_order_ops.auto_functionalize import auto_functionalized_v2 +from torch._higher_order_ops.auto_functionalize import auto_functionalized, auto_functionalized_v2 from torch.fx import Graph # Registers torch.ops.trtllm.mla_custom_op_inplace, used below. The op is a @@ -31,6 +31,61 @@ ) +@pytest.mark.cpu_only +@pytest.mark.parametrize("use_v2", [False, True]) +def test_remove_copy_preserves_minimax_producer_outputs(use_v2: bool) -> None: + """Preserve query outputs and cache aliases when removing functionalization.""" + import tensorrt_llm._torch.models.modeling_minimaxm3 # noqa: F401 + + graph = Graph() + hidden = graph.placeholder("hidden") + positions = graph.placeholder("positions") + main = graph.placeholder("main") + index = graph.placeholder("index") + slots = graph.placeholder("slots") + op = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer.default + kwargs = dict(hidden_states=hidden, position_ids=positions, out_cache_loc=slots, layer_idx="3") + if use_v2: + kwargs.update(_all_bases=(main, index), _kv_cache_base_index=0, _index_k_cache_base_index=1) + else: + kwargs.update(kv_cache=main, index_k_cache=index) + functionalized = graph.call_function( + auto_functionalized_v2 if use_v2 else auto_functionalized, args=(op,), kwargs=kwargs + ) + q = graph.call_function(getitem, args=(functionalized, 0)) + index_q = graph.call_function(getitem, args=(functionalized, 1)) + q.meta["val"] = torch.empty(4, 3) + index_q.meta["val"] = torch.empty(4, 1) + functionalized.meta["val"] = ( + q.meta["val"], + index_q.meta["val"], + torch.empty(8, 3), + torch.empty(8, 1), + ) + updated_main = graph.call_function(getitem, args=(functionalized, 2)) + updated_index = graph.call_function(getitem, args=(functionalized, 3)) + output = graph.output((q, index_q, updated_main, updated_index)) + + remove_copy_pass.remove_copy_for_mutates_args(graph) + + calls = [node for node in graph.nodes if node.target == op] + assert len(calls) == 1 + assert calls[0].kwargs["kv_cache"] is main + assert calls[0].kwargs["index_k_cache"] is index + assert calls[0].kwargs["out_cache_loc"] is slots + assert calls[0].meta["val"] == (q.meta["val"], index_q.meta["val"]) + new_q, new_index_q, new_main, new_index = output.args[0] + assert new_q.args == (calls[0], 0) + assert new_index_q.args == (calls[0], 1) + assert new_main is main and new_index is index + assert new_q.meta["val"] is q.meta["val"] + assert new_index_q.meta["val"] is index_q.meta["val"] + assert all( + node.target not in (auto_functionalized, auto_functionalized_v2) for node in graph.nodes + ) + graph.lint() + + def test_remove_copy_for_mutates_args_auto_functionalized_v2( monkeypatch: pytest.MonkeyPatch, ) -> None: 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 5e9693abb810..7878f7668de3 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -14,6 +14,7 @@ import os import sys import unittest +from collections import OrderedDict from dataclasses import dataclass from types import ModuleType, SimpleNamespace from unittest.mock import Mock, call, patch @@ -25,6 +26,8 @@ import tensorrt_llm._torch.pyexecutor.model_engine as model_engine_module from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner @@ -33,6 +36,7 @@ PyTorchModelEngine, _ContextOnlyCompiledModel, ) +from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.speculative.utils import update_draft_len from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling @@ -52,6 +56,7 @@ def test_context_only_compile_uses_all_rank_prefill_decision( monkeypatch: pytest.MonkeyPatch, local_contexts: int, ) -> None: + """Route using the global prefill decision, regardless of local contexts.""" eager = torch.nn.Linear(4, 4) compiled = torch.nn.Module() compiled.shared = eager @@ -74,6 +79,49 @@ def test_context_only_compile_uses_all_rank_prefill_decision( assert compiled.forward.call_count == eager.forward.call_count == 1 +@pytest.mark.cpu_only +def test_context_only_compile_preserves_partial_weight_reload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reload selected weights by original names into both model entry points.""" + import tensorrt_llm._torch.models.modeling_utils as modeling_utils + + eager = torch.nn.Sequential(OrderedDict(layers=torch.nn.Sequential(torch.nn.Linear(4, 4)))) + compiled = torch.compile(eager, backend="eager") + model = DecoderModelForCausalLM.__new__(DecoderModelForCausalLM) + torch.nn.Module.__init__(model) + model.config = SimpleNamespace(tie_word_embeddings=False) + model.model = _ContextOnlyCompiledModel(eager, compiled) + mapper = HfWeightMapper() + mapper._model = model + loader = ModelLoader.__new__(ModelLoader) + loader.weight_mapper = mapper + monkeypatch.setenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "True") + monkeypatch.setattr(modeling_utils, "local_mpi_rank", lambda: 0) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + weight = eager.layers[0].weight + old_bias = eager.layers[0].bias.detach().clone() + replacement = torch.full_like(weight, 3) + for remove_duplicate in (False, True): + names = dict(model.named_modules(remove_duplicate=remove_duplicate)) + assert names["model.layers.0"] is eager.layers[0] + assert not any("eager_model" in name or "compiled_model" in name for name in names) + loader.reload(model, {"model.layers.0.weight": replacement}, allow_partial_loading=True) + + assert eager.layers[0].weight is weight + assert compiled.layers[0].weight is weight + torch.testing.assert_close(weight, replacement) + torch.testing.assert_close(eager.layers[0].bias, old_bias) + inputs = torch.ones(2, 4) + for eligible in (False, True): + monkeypatch.setattr( + model_engine_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + ) + torch.testing.assert_close(model.model(inputs), inputs @ replacement.t() + old_bias) + + @pytest.mark.cpu_only @pytest.mark.parametrize("eligible", [False, True]) @pytest.mark.parametrize("raises", [False, True]) @@ -82,9 +130,11 @@ def test_context_only_compile_scopes_whole_model_forward( eligible: bool, raises: bool, ) -> None: + """Keep compile state active through model epilogues and restore it on exit.""" observed = [] def forward(**kwargs: object) -> str: + """Record compile state as a stand-in for a model-specific epilogue.""" observed.append(is_torch_compiling()) # This represents work after the transformer, such as Eagle3 drafting. if raises: @@ -899,6 +949,7 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): self.assertEqual(engine.forward.call_count, 1) def test_native_mxfp8_respects_disabled_global_autotuner(self): + """Avoid native MXFP8 warmup when the global autotuner is disabled.""" with ( patch( "tensorrt_llm._torch.modules.linear._mxfp8_cutlass_op_available", diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 0f016899f536..c3cff4a9190e 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -522,6 +522,7 @@ def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, ou def test_piecewise_captured_producer_preserves_symbolic_shapes( monkeypatch: pytest.MonkeyPatch, ) -> None: + """Retain an unbacked symbolic token count in both fake query outputs.""" from torch._subclasses.fake_tensor import FakeTensorMode from torch.fx.experimental.symbolic_shapes import ShapeEnv @@ -535,7 +536,12 @@ def test_piecewise_captured_producer_preserves_symbolic_shapes( num_tokens = mode.shape_env.create_unbacked_symint() hidden = torch.empty((num_tokens, 512), dtype=torch.bfloat16) positions = torch.empty((1, num_tokens), dtype=torch.int32) - q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer(hidden, positions, "3") + kv_cache = torch.empty((2, 2, 1, 128, 128), dtype=torch.float8_e4m3fn) + index_cache = torch.empty((2, 1, 128, 128), dtype=torch.float8_e4m3fn) + slots = torch.empty((4096,), dtype=torch.int32) + q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( + hidden, positions, kv_cache, index_cache, slots, "3" + ) for output, width in ((q, 1024), (idx_q, 128)): assert output.shape[0].node.expr == num_tokens.node.expr assert output.shape[1] == width @@ -546,6 +552,7 @@ def test_piecewise_captured_producer_preserves_symbolic_shapes( def test_piecewise_captures_horizontal_producer_before_attention( monkeypatch: pytest.MonkeyPatch, ) -> None: + """Run the captured producer before eager sparse attention consumes caches.""" from unittest.mock import Mock backend = object.__new__(MiniMaxM3MsaSparseAttention) @@ -554,12 +561,18 @@ def test_piecewise_captures_horizontal_producer_before_attention( q = torch.empty((4, 3), dtype=torch.float8_e4m3fn) idx_q = torch.empty((4, 1), dtype=torch.float8_e4m3fn) output = torch.empty((4, 3), dtype=torch.bfloat16) - metadata = SimpleNamespace(num_tokens=2) + kv_cache = torch.empty(8, dtype=torch.float8_e4m3fn) + index_cache = torch.empty_like(kv_cache) + slots = torch.tensor([0, 1, -1, -1], dtype=torch.int32) + metadata = SimpleNamespace( + num_tokens=2, msa_layer_cache_tensors={3: (kv_cache, index_cache)}, msa_out_cache_loc=slots + ) layer = SimpleNamespace( enable_fused_qkv_index_projection=True, register_to_config=True, attn=backend, _emit_fp8_main_qkv=lambda: True, + layer_idx=3, layer_idx_str="3", qkv_proj=Mock(return_value=packed), _fused_fp8_qkv_indexer_norm_rope_kv_insert=Mock(return_value=(q, idx_q)), @@ -578,13 +591,156 @@ def test_piecewise_captures_horizontal_producer_before_attention( assert result is output layer.qkv_proj.assert_called_once_with(hidden) layer._fused_fp8_qkv_indexer_norm_rope_kv_insert.assert_called_once_with( - packed, positions, metadata + packed, positions, metadata, cache_tensors=(kv_cache, index_cache, slots) ) layer._forward_attention_core.assert_called_once_with(q, None, None, idx_q, None, metadata) +@pytest.mark.cpu_only +def test_piecewise_captured_producer_rejects_unavailable_fusion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reject unsupported fused geometry without an eager fallback in capture.""" + from unittest.mock import Mock + + layer = SimpleNamespace( + qkv_proj=Mock(side_effect=lambda hidden: hidden.clone()), + _fused_fp8_qkv_indexer_norm_rope_kv_insert=Mock(return_value=None), + ) + monkeypatch.setattr( + modeling_minimaxm3, "_extract_minimax_m3_attention_extra_attrs", lambda _: (None, layer) + ) + with pytest.raises(RuntimeError) as exc: + torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( + torch.zeros(2, 4), + torch.arange(2), + torch.zeros(4), + torch.zeros(4), + torch.arange(2, dtype=torch.int32), + "3", + ) + assert str(exc.value) == ( + "MiniMax-M3 piecewise graph requires the fused FP8 sparse QKV producer." + ) + layer._fused_fp8_qkv_indexer_norm_rope_kv_insert.assert_called_once() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("restore_inplace", [False, True]) +def test_piecewise_captured_producer_declares_cache_mutations( + monkeypatch: pytest.MonkeyPatch, + restore_inplace: bool, +) -> None: + """Real custom-op schema/AOT checks with CPU cache writes in place of CUDA math.""" + from torch._dynamo.backends.common import aot_autograd + from torch._functorch.aot_autograd import make_boxed_func + from torch._higher_order_ops.auto_functionalize import ( + auto_functionalized, + auto_functionalized_v2, + ) + + from tensorrt_llm._torch.compilation.remove_copy_pass import remove_copy_for_mutates_args + + def producer( + packed: torch.Tensor, + positions: torch.Tensor, + metadata: object, + *, + cache_tensors: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Emulate native cache writes using only the explicit tensor arguments.""" + kv_cache, index_cache, slots = cache_tensors + # Deliberately use only the explicit caches, not the runtime metadata. + valid = slots >= 0 + indices = slots[valid].long() + kv_cache.index_copy_(0, indices, packed[valid, :3].float()) + index_cache.index_copy_(0, indices, packed[valid, :1].float() + 1) + return packed[:, :3].to(torch.float8_e4m3fn), packed[:, :1].to(torch.float8_e4m3fn) + + layer = SimpleNamespace( + q_size=3, + index_q_size=1, + qkv_proj=lambda hidden: hidden + 2, + _fused_fp8_qkv_indexer_norm_rope_kv_insert=producer, + ) + monkeypatch.setattr( + modeling_minimaxm3, "_extract_minimax_m3_attention_extra_attrs", lambda _: (None, layer) + ) + op = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer.default + assert { + arg.name for arg in op._schema.arguments if arg.alias_info and arg.alias_info.is_write + } == {"kv_cache", "index_k_cache"} + hidden = torch.arange(16, dtype=torch.float32).reshape(4, 4) + positions = torch.arange(4) + kv_cache = torch.zeros(8, 3) + index_cache = torch.zeros(8, 1) + slots = torch.tensor([1, 3, -1, -1], dtype=torch.int32) + args = (hidden, positions, kv_cache, index_cache, slots, "3") + assert all(result == "SUCCESS" for result in torch.library.opcheck(op, args).values()) + + def run( + hidden: torch.Tensor, + positions: torch.Tensor, + main: torch.Tensor, + index: torch.Tensor, + slots: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Expose query results and immediate observations of both live caches.""" + q, idx_q = op(hidden, positions, main, index, slots, "3") + return q.float(), idx_q.float(), main.clone(), index.clone() + + optimized_graphs = [] + + def optimize(gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]) -> object: + """Apply the production in-place recovery pass and inspect its aliases.""" + remove_copy_for_mutates_args(gm.graph) + gm.graph.lint() + gm.recompile() + producer_nodes = [node for node in gm.graph.nodes if node.target == op] + assert len(producer_nodes) == 1 + # No full-pool functionalization clones between the graph inputs and + # the producer: MSA's eager boundary must observe the original pools. + assert producer_nodes[0].kwargs["kv_cache"].op == "placeholder" + assert producer_nodes[0].kwargs["index_k_cache"].op == "placeholder" + assert all( + node.target not in (auto_functionalized, auto_functionalized_v2) + for node in gm.graph.nodes + ) + optimized_graphs.append(gm) + return make_boxed_func(gm.forward) + + # Match the production backend's functionalization version. + monkeypatch.setattr(torch._inductor.config, "enable_auto_functionalized_v2", False) + compiled = torch.compile( + run, + backend=aot_autograd(fw_compiler=optimize) if restore_inplace else "aot_eager", + fullgraph=True, + ) + for offset in (0, 4): + kv_cache.zero_() + index_cache.zero_() + q, idx_q, observed_main, observed_index = compiled( + hidden + offset, positions, kv_cache, index_cache, slots + ) + packed = hidden + offset + 2 + expected_main = torch.zeros_like(kv_cache) + expected_index = torch.zeros_like(index_cache) + expected_main[[1, 3]] = packed[:2, :3] + expected_index[[1, 3]] = packed[:2, :1] + 1 + torch.testing.assert_close(kv_cache, expected_main) + torch.testing.assert_close(index_cache, expected_index) + torch.testing.assert_close(observed_main, expected_main) + torch.testing.assert_close(observed_index, expected_index) + torch.testing.assert_close(q, packed[:, :3].to(torch.float8_e4m3fn).float()) + torch.testing.assert_close(idx_q, packed[:, :1].to(torch.float8_e4m3fn).float()) + torch.testing.assert_close(slots, torch.tensor([1, 3, -1, -1], dtype=torch.int32)) + if restore_inplace: + assert len(optimized_graphs) == 1 + + @pytest.mark.cpu_only def test_piecewise_unfused_indexer_keeps_cache_write_eager(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep index-cache mutation outside capture when projections are separate.""" from unittest.mock import Mock backend = object.__new__(MiniMaxM3MsaSparseAttention) diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index f5e5e8ca6507..edeced442d2e 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -20,6 +20,7 @@ import pytest import torch +import tensorrt_llm._torch.custom_ops.flashinfer_custom_ops as flashinfer_ops_module import tensorrt_llm._torch.custom_ops.torch_custom_ops as custom_ops_module import tensorrt_llm._torch.modules.linear as linear_module from tensorrt_llm._torch.autotuner import AutoTuner @@ -79,7 +80,10 @@ def test_mxfp8_dispatch_returns_mxfp8_method(monkeypatch): def _mock_mxfp8_ops( monkeypatch: pytest.MonkeyPatch, flashinfer_gemm: Mock | None = None, + *, + flashinfer_op_available: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, Mock, Mock, torch.Tensor, Mock, torch.Tensor]: + """Replace MXFP8 kernels with CPU doubles and optional FlashInfer registration.""" quantized = torch.empty((2, 4), dtype=torch.float8_e4m3fn) activation_scale = torch.empty(512, dtype=torch.uint8) quantize = Mock(return_value=(quantized, activation_scale)) @@ -91,8 +95,9 @@ def _mock_mxfp8_ops( mxfp8_quantize=quantize, mxfp8_mxfp8_gemm=native_gemm, mxfp8_mxfp8_gemm_autotuned=autotuned_gemm, - flashinfer_mm_mxfp8=flashinfer_gemm or Mock(), ) + if flashinfer_op_available: + fake_trtllm_ops.flashinfer_mm_mxfp8 = flashinfer_gemm or Mock() fake_torch = SimpleNamespace( ops=SimpleNamespace(trtllm=fake_trtllm_ops), ones=torch.ones, @@ -110,6 +115,60 @@ def _mock_mxfp8_ops( ) +@pytest.mark.cpu_only +def test_registered_flashinfer_mxfp8_wrapper_contract(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise the real registered wrapper, mocking only the FlashInfer kernel.""" + expected = torch.full((2, 3), 7, dtype=torch.bfloat16) + kernel = Mock(return_value=expected) + monkeypatch.setattr(flashinfer_ops_module, "mm_mxfp8", kernel) + act = torch.empty((2, 4), dtype=torch.float8_e4m3fn) + weight = torch.arange(12, dtype=torch.float32).reshape(3, 4).to(torch.float8_e4m3fn) + act_scale = torch.ones(512, dtype=torch.uint8) + weight_scale = torch.ones(512, dtype=torch.uint8) + + output = torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, torch.bfloat16 + ) + + torch.testing.assert_close(output, expected) + kernel.assert_called_once() + args = kernel.call_args.args + assert args[0] is act + assert args[1].shape == (4, 3) + assert args[1].stride() == weight.t().stride() + assert args[1].data_ptr() == weight.data_ptr() + torch.testing.assert_close(args[1].float(), weight.t().float()) + assert args[2] is act_scale and args[3] is weight_scale + assert kernel.call_args.kwargs == { + "out_dtype": torch.bfloat16, + "use_8x4_sf_layout": False, + "backend": "cutlass", + } + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("backend", ["auto", "flashinfer"]) +def test_mxfp8_missing_registered_flashinfer_op( + monkeypatch: pytest.MonkeyPatch, backend: str +) -> None: + """Allow automatic fallback but reject an unavailable explicitly chosen op.""" + monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", backend) + monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) + monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(autotune=Mock())) + _mock_mxfp8_ops(monkeypatch, flashinfer_op_available=False) + if backend == "flashinfer": + with pytest.raises( + RuntimeError, match="requires the pinned flashinfer-python package" + ) as exc: + MXFP8LinearMethod() + assert str(exc.value.__cause__) == "trtllm::flashinfer_mm_mxfp8 is unavailable" + else: + method = MXFP8LinearMethod() + assert method.backend == "trtllm" + assert not method.uses_flashinfer + assert method._flashinfer_mxfp8 is None + + def test_mxfp8_flashinfer_call_contract(monkeypatch): """The forced backend passes native-layout tensors to the opaque op.""" monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", "flashinfer") @@ -145,6 +204,7 @@ def test_mxfp8_flashinfer_call_contract(monkeypatch): def test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer(monkeypatch): + """Keep native eager GEMM while routing captured work to the opaque wrapper.""" monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) From efebec94ef9cce5e3a87ff64c9a30284ce2488f8 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:46:46 -0700 Subject: [PATCH 03/12] [None][fix] Preserve MiniMax-M3 PCG indexer cache contract Drop live FP8 index-K only after the attention boundary writes it to cache, retaining the BF16 handoff. Exercise the real boundary and indexer for unfused FP8, fused prewritten FP8, and BF16; assert HND cache-view requests. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3.py | 30 ++--- .../integration/test_lists/test-db/l0_cpu.yml | 2 +- .../attention/sparse/msa/test_msa_backend.py | 6 +- .../unittest/_torch/models/test_minimax_m3.py | 117 +++++++++++++----- 4 files changed, 108 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 066b8dba2416..127e3b254a08 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -897,11 +897,10 @@ def minimax_m3_attn_custom_op_inplace( ) -> None: """Run MiniMax-M3 cache and attention work behind a compile boundary. - The horizontal producer needs live paged-cache tensors and cache-slot - metadata, which are intentionally resolved inside this opaque attention - boundary rather than traced through Dynamo. Projection remains in the - captured segment; only the cache-writing producer and MSA attention stay - on the eager side of the existing piecewise boundary. + The captured horizontal producer can populate both caches before this + boundary. The packed-projection fallback runs that producer here instead, + while separate projections leave their cache writes here. Slice padded + inputs to live tokens before request-dependent cache and MSA work. """ attn_metadata, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) num_tokens = attn_metadata.num_tokens @@ -1889,25 +1888,28 @@ def _msa_attention_core( FMHA forward; this layer selects the top-k blocks (sparse only) and builds the forward_args the FMHA reads. - This layer owns the cache write: write_layer_caches stores the - new-token K/V (and, on the bf16 indexer path, index-K) in one launch - before the indexer's proxy pass reads the index-K cache. forward() - then receives k=v=None, which is the backend's contract for "K/V are - already resident", so neither FMHA phase writes them again. + Unless a producer has already populated the caches, write_layer_caches + stores new-token K/V and any live index-K before the indexer's proxy + pass reads the cache. FP8 index-K is then omitted from run_indexer; + BF16 retains its live tensor with idx_k_prewritten=True. forward() + receives k=v=None so neither FMHA phase writes them again. """ assert (k is None) == (v is None) if self.is_sparse_attention_layer: assert idx_q is not None - # On the FP8 indexer path idx_k is None: the fused producer already - # inserted E4M3 index-K into the side cache, so only K/V are written. + # Unfused PCG supplies live FP8 index-K; eager FP8 producers may + # already have inserted it and supply None instead. if k is not None: self.attn.write_layer_caches(k, v, idx_k, attn_metadata) + if self.attn.indexer_kv_dtype == "fp8": + # The FP8 indexer accepts only an already-populated cache. + idx_k = None else: # The horizontal producer has already written both caches. assert idx_k is None # Publish the selected blocks so the FMHA runs the sparse path. - # idx_k_prewritten: index-K is already in the cache (written above - # on bf16, or by the FP8 producer), so run_indexer must not write it. + # idx_k_prewritten: index-K is already in the cache (written here + # or by an FP8 producer), so run_indexer must not write it. kv_block_indexes = self.attn.run_indexer( idx_q, idx_k, attn_metadata, idx_k_prewritten=True ) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 4b7666b6552d..324d171210b4 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -64,7 +64,7 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captured_producer_rejects_unavailable_fusion - unittest/_torch/models/test_minimax_m3.py::test_piecewise_captured_producer_declares_cache_mutations - unittest/_torch/models/test_minimax_m3.py::test_piecewise_unfused_indexer_keeps_cache_write_eager - - unittest/_torch/models/test_minimax_m3.py::test_msa_attention_core_routes_compact_q_to_attention_dispatcher + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_attention_boundary_preserves_indexer_cache_contract - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip - unittest/_torch/models/checkpoints diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 878fb782889f..a712d1af21b9 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -321,8 +321,8 @@ def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch) -> main_cache = torch.zeros(2, 2, 1, 128, 128) index_cache = torch.zeros(2, 1, 128, 128) manager = metadata.kv_cache_manager - manager.get_buffers = lambda layer_idx, kv_layout: main_cache - manager.get_index_k_buffer = lambda layer_idx, kv_layout: index_cache + manager.get_buffers = Mock(return_value=main_cache) + manager.get_index_k_buffer = Mock(return_value=index_cache) monkeypatch.setattr( metadata, "get_empty", @@ -331,6 +331,8 @@ def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch) -> # No native pool in this CPU test; only zero-copy cache-view staging is under test. monkeypatch.setattr(msa_backend, "uniform_subpages_per_slot", lambda manager: 0) metadata._create_msa_buffers() + manager.get_buffers.assert_called_once_with(3, kv_layout="HND") + manager.get_index_k_buffer.assert_called_once_with(3, kv_layout="HND") assert set(metadata.msa_layer_cache_tensors) == {3} main, index = metadata.msa_layer_cache_tensors[3] assert main is main_cache and index is index_cache diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index c3cff4a9190e..01680035f09b 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -790,43 +790,100 @@ def test_piecewise_unfused_indexer_keeps_cache_write_eager(monkeypatch: pytest.M @pytest.mark.cpu_only -def test_msa_attention_core_routes_compact_q_to_attention_dispatcher() -> None: - selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) - - class FakeMsaBackend: - def __init__(self) -> None: - self.prepopulated_call = None - - def write_layer_caches(self, k, v, idx_k, metadata) -> None: - pytest.fail("The horizontal producer has already written both caches") - - def run_indexer(self, idx_q, idx_k, metadata, *, idx_k_prewritten): - assert idx_k is None - assert idx_k_prewritten - assert metadata is attn_metadata - return selected_blocks +@pytest.mark.parametrize( + ("indexer_dtype", "caches_prewritten"), + [("fp8", False), ("fp8", True), ("bf16", False)], + ids=["unfused-fp8", "fused-fp8", "bf16"], +) +def test_piecewise_attention_boundary_preserves_indexer_cache_contract( + monkeypatch: pytest.MonkeyPatch, indexer_dtype: str, caches_prewritten: bool +) -> None: + """Exercise the real MSA indexer after live-row slicing and cache insertion.""" + from unittest.mock import Mock - def forward(self, q, k, v, metadata, forward_args) -> None: - assert k is None and v is None - self.prepopulated_call = (q, metadata, forward_args) + dtype = torch.float8_e4m3fn if indexer_dtype == "fp8" else torch.bfloat16 + q, k, v, idx_q, idx_k = [torch.full((4, 128), value).to(dtype) for value in range(1, 6)] + index_cache = torch.zeros((4, 1, 1, 128), dtype=dtype) + expected_cache = torch.zeros_like(index_cache) + expected_cache[:2, 0, 0].copy_(idx_k[:2]) + if caches_prewritten: + index_cache.copy_(expected_cache) + selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) + attn_metadata = SimpleNamespace( + num_tokens=2, + msa_decode_span=None, + msa_idx_k_cache=Mock(return_value=index_cache), + msa_write_idx_k=Mock(), + msa_prefill_proxy_plan=None, + msa_prefill_n_valid_blocks=None, + msa_kv_indices=torch.tensor([0, 1], dtype=torch.int32), + msa_qo_lens_cpu=torch.tensor([2], dtype=torch.int32), + msa_kv_lens_cpu=torch.tensor([2], dtype=torch.int32), + msa_qo_offset_cpu=torch.tensor([0], dtype=torch.int32), + ) + + def write_caches( + live_k: torch.Tensor, + live_v: torch.Tensor, + live_idx_k: torch.Tensor, + metadata: SimpleNamespace, + ) -> None: + """Replace only CUDA scatter math, retaining the live cache-write inputs.""" + assert metadata is attn_metadata + torch.testing.assert_close(live_k.float(), k[:2].float()) + torch.testing.assert_close(live_v.float(), v[:2].float()) + torch.testing.assert_close(live_idx_k.float(), idx_k[:2].float()) + index_cache[:2, 0, 0].copy_(live_idx_k) + + def select_blocks( + live_idx_q: torch.Tensor, cache: torch.Tensor, **kwargs: object + ) -> torch.Tensor: + """Require the cache write to precede selection, without running CUDA scoring.""" + assert cache is index_cache + torch.testing.assert_close(cache.float(), expected_cache.float()) + torch.testing.assert_close(live_idx_q.flatten(1).float(), idx_q[:2].float()) + return selected_blocks + backend = object.__new__(MiniMaxM3MsaSparseAttention) + backend.layer_idx = 3 + backend.indexer_kv_dtype = indexer_dtype + backend.m3_config = SimpleNamespace(num_index_heads=1, sparse_index_dim=128) + backend.write_layer_caches = Mock(side_effect=write_caches) + backend.indexer = SimpleNamespace(select_blocks=Mock(side_effect=select_blocks)) + backend.forward = Mock() layer = MiniMaxM3Attention.__new__(MiniMaxM3Attention) - backend = FakeMsaBackend() layer.attn = backend layer.is_sparse_attention_layer = True - q = torch.randn(2, 8) - idx_q = torch.randn(2, 4) - attn_metadata = SimpleNamespace() - output = torch.empty_like(q) - - result = layer._msa_attention_core(q, None, None, idx_q, None, attn_metadata, output) + output = torch.empty((4, 128), dtype=torch.bfloat16) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (attn_metadata, layer), + ) + modeling_minimaxm3.minimax_m3_attn_custom_op_inplace( + q, + None if caches_prewritten else k, + None if caches_prewritten else v, + idx_q, + None if caches_prewritten else idx_k, + None, + None, + "3", + output, + ) - assert result is output - assert backend.prepopulated_call is not None - called_q, called_metadata, forward_args = backend.prepopulated_call - assert called_q is q + assert backend.write_layer_caches.call_count == (0 if caches_prewritten else 1) + attn_metadata.msa_write_idx_k.assert_not_called() + attn_metadata.msa_idx_k_cache.assert_called_once_with(3) + backend.indexer.select_blocks.assert_called_once() + backend.forward.assert_called_once() + called_q, called_k, called_v, called_metadata = backend.forward.call_args.args + assert called_k is None and called_v is None + torch.testing.assert_close(called_q.float(), q[:2].float()) assert called_metadata is attn_metadata - assert forward_args.output is output + forward_args = backend.forward.call_args.kwargs["forward_args"] + assert forward_args.output.shape == (2, 128) + assert forward_args.output.data_ptr() == output.data_ptr() assert forward_args.sparse_backend_args.topk_indices is selected_blocks From 2a5573240b4eb5ff6195f44cfb92a2575501b7f6 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:23:39 -0700 Subject: [PATCH 04/12] [None][test] Cover FP8 indexer cache handoff in MSA fixture Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../attention/sparse/msa/test_msa_backend.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index a712d1af21b9..e48d00915f71 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -1738,13 +1738,13 @@ def test_fused_scatter_matches_reference(src_dtype, cache_dtype, with_idx): torch.testing.assert_close(idx_pool, ref_idx_pool) -@pytest.mark.parametrize("sparse", [True, False]) -def test_msa_attention_core_owns_the_cache_write(sparse): +@pytest.mark.parametrize("sparse,indexer_dtype", [(True, "fp8"), (True, "bf16"), (False, "fp8")]) +def test_msa_attention_core_owns_the_cache_write(sparse: bool, indexer_dtype: str) -> None: """The model layer's MSA core must write the caches exactly once and in the right place: write_layer_caches runs before run_indexer (whose proxy pass reads the index-K cache), run_indexer is told index-K is already - resident, and forward() receives k=v=None so no FMHA phase writes K/V - again.""" + resident (with no live index-K for FP8), and forward() receives k=v=None + so no FMHA phase writes K/V again.""" from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3Attention num_tokens, width = 3, 128 @@ -1753,6 +1753,7 @@ def test_msa_attention_core_owns_the_cache_write(sparse): class FakeBackend: layer_idx = 7 + indexer_kv_dtype = indexer_dtype def write_layer_caches(self, k, v, idx_k, metadata): events.append(("write", k, v, idx_k, metadata)) @@ -1778,7 +1779,11 @@ def forward(self, q, k, v, metadata, forward_args=None): if sparse: assert names == ["write", "indexer", "forward"] _, indexer_q, indexer_k, indexer_metadata, prewritten = events[1] - assert indexer_q is idx_q and indexer_k is idx_k and indexer_metadata is metadata + assert indexer_q is idx_q and indexer_metadata is metadata + if indexer_dtype == "fp8": + assert indexer_k is None + else: + assert indexer_k is idx_k assert prewritten is True else: assert names == ["write", "forward"] From 3cff17fc9b099e21130529c4c0fd1ef788f85d8d Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:58:25 -0700 Subject: [PATCH 05/12] [None][docs] Drop PCG documentation additions from migration Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../torch_compile_and_piecewise_cuda_graph.md | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index fa957fc2fa08..8b179030177f 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -79,19 +79,6 @@ validated before use. ### Piecewise CUDA Graph & Generation Only CUDA Graph -For decoder models, piecewise mode compiles capture-eligible context and mixed -batches. Generation-only batches and context batches above the capture ceiling -use the eager model, including any speculative-decoding epilogue. With attention -DP, eligibility is shared across ranks, including ranks without local context -requests. Torch compile without piecewise graphs retains its all-batch behavior. - -MiniMax-M3 with MSA supports FP8 KV and index-K caches in piecewise mode. With -`sparse_attention_config.fuse_qkv_index_projection: true`, projection, norm, -RoPE, FP8 conversion and cache insertion are captured together; sparse attention -remains in the eager boundary. Padded rows do not write to the caches. Automatic -MXFP8 dispatch uses the native backend for compiled context while preserving -decode-graph backend tuning. - Piecewise CUDA Graph only handles context-only and mixed context+generation iterations, while the generation-only CUDA Graph only handles pure generation iterations. Users need to specify the number of tokens to capture for each type of CUDA Graph separately in the extra config. Currently, the default value for `capture_num_tokens` is `[2**i for i in range(8)] + [i for i in range(256, 3073, 256)]`. However, this configuration should be tuned based on specific hardware, model, and parallel strategy. For guidance on tuning these values, see the [Performance Tuning](#performance-tuning) section below. ```yaml From 50486b7c8c31c8e0c404a25740699734472b51cf Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:11:32 -0700 Subject: [PATCH 06/12] [None][test] Cover empty ADP ranks and Eagle3 piecewise prefill Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_pytorch.py | 43 +++- .../test_lists/test-db/l0_dgx_b200.yml | 3 + .../multi_gpu/test_minimax_m3_piecewise.py | 196 ++++++++++++++++++ 3 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index cd1305fd6d2e..8c13b255b272 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7356,6 +7356,34 @@ def _run_nvfp4(self, def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, overlap_scheduler, fuse_qkv_index_projection, eval_mode): + self._run_nvfp4_eagle3(tp_size, ep_size, attention_dp, + overlap_scheduler, fuse_qkv_index_projection, + eval_mode) + + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("fuse_qkv_index_projection", [False, True]) + def test_nvfp4_eagle3_piecewise_cuda_graph( + self, fuse_qkv_index_projection: bool) -> None: + """Check accuracy and draft acceptance from piecewise prefill to decode.""" + self._run_nvfp4_eagle3( + 4, + 4, + attention_dp=True, + overlap_scheduler=True, + fuse_qkv_index_projection=fuse_qkv_index_projection, + eval_mode="default", + piecewise=True) + + def _run_nvfp4_eagle3(self, + tp_size: int, + ep_size: int, + attention_dp: bool, + overlap_scheduler: bool, + fuse_qkv_index_projection: bool, + eval_mode: str, + *, + piecewise: bool = False) -> None: # One-model Eagle3 on the MSA backend with an FP8 KV cache and CUDA # graphs; the GQA drafter shares the target KV cache. MMLU + GSM8K, or # InferenceX GSM8K, plus a chat-GSM8K acceptance probe, since accuracy @@ -7384,6 +7412,18 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, else: max_seq_len = 4096 max_batch_size = 256 if attention_dp else 512 + piecewise_kwargs = {} + if piecewise: + # Cover the entire scheduler token budget so these evaluations + # exercise captured prefill, rather than falling back above the + # capture ceiling. Generation-only steps still take the eager + # model path, including the Eagle3 epilogue and decode graphs. + piecewise_kwargs = dict( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + prefill_capture_num_tokens=[128, 512, 2048, 4096], + torch_compile_config=TorchCompileConfig(), + max_num_tokens=4096, + ) with LLM( model_path, tensor_parallel_size=tp_size, @@ -7407,7 +7447,8 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, # Keep the whole acceptance probe: the default buffer holds # only the latest 1000 iterations. max_stats_len=5000, - trust_remote_code=True) as llm: + trust_remote_code=True, + **piecewise_kwargs) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION def drain_spec_stats(llm): diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index fa456bd01d1b..ffb715211573 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -16,6 +16,7 @@ l0_dgx_b200: orchestrator: mpi tests: - unittest/_torch/misc/test_autotuner.py::test_autotuner_distributed_strategy + - unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py::test_minimax_m3_piecewise_empty_adp_rank_preserves_caches - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-TRTLLM] # ------------- KV Cache V2 Scheduler IT (multi-GPU) --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_draft_tokens @@ -56,6 +57,8 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_piecewise_cuda_graph[fuse_qkv_index_projection=False] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_piecewise_cuda_graph[fuse_qkv_index_projection=True] + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3_piecewise_cuda_graph[fuse_qkv_index_projection=False] + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3_piecewise_cuda_graph[fuse_qkv_index_projection=True] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=False-eval_mode=inferencex] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] diff --git a/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py b/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py new file mode 100644 index 000000000000..6da7e43195b3 --- /dev/null +++ b/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Distributed PCG cache regressions without loading a full M3 checkpoint.""" + +import pickle +import sys +from types import SimpleNamespace + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from mpi4py.futures import MPIPoolExecutor + +import tensorrt_llm +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.compilation.piecewise_optimizer import PiecewiseRunner +from tensorrt_llm._torch.compilation.utils import capture_piecewise_cuda_graph +from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.pyexecutor.engine.runners.common import get_padding_params +from tensorrt_llm._torch.utils import ( + get_per_request_prefill_cuda_graph_flag, + model_extra_attrs, + piecewise_cuda_graph, + set_per_request_prefill_cuda_graph_flag, +) +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +# Match the other MPIPoolExecutor tests: its worker thread outlives the test. +pytestmark = pytest.mark.threadleak(enabled=False) + + +@torch.inference_mode() +def _run_empty_adp_rank(world_size: int) -> int: + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + torch.manual_seed(42 + rank) + mapping = Mapping( + world_size=world_size, rank=rank, tp_size=world_size, enable_attention_dp=True + ) + dist = Distributed.get(mapping) + bucket = 4 + num_heads_q, num_kv_heads, num_index_heads = 8, 2, 2 + width = (num_heads_q + 2 * num_kv_heads + num_index_heads + 1) * 128 + packed = torch.randn(bucket, width, dtype=torch.bfloat16, device="cuda") + weights = [torch.ones(128, dtype=torch.bfloat16, device="cuda") for _ in range(4)] + positions = torch.arange(bucket, dtype=torch.int32, device="cuda") + frequency = torch.outer( + positions.float(), + 5_000_000.0 ** (-torch.arange(0, 64, 2, dtype=torch.float32, device="cuda") / 64), + ) + rope_cache = torch.stack((frequency.cos(), frequency.sin()), dim=1).contiguous() + # HND views with interleaved pages, as in the producer's kernel tests. + main_backing = torch.full( + (6, 2, num_kv_heads, 128, 128), 1.0, dtype=torch.float8_e4m3fn, device="cuda" + ) + index_backing = torch.full((10, 1, 128, 128), 2.0, dtype=torch.float8_e4m3fn, device="cuda") + main_cache, index_cache = main_backing[::3], index_backing[::5] + + # Reuse the metadata fixture pattern from test_msa_backend. Only the cache + # allocator is stubbed; slot construction/clearing and GPU writes are real. + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata._msa_buffers_ready = True + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=128, + get_buffers=lambda layer_idx: main_cache, + get_block_ids_per_seq=lambda request_ids: torch.tensor([[1]], dtype=torch.int32), + ) + metadata.msa_out_cache_loc = torch.full((bucket,), -1, dtype=torch.int32, device="cuda") + metadata.msa_block_table = torch.zeros((1, 1), dtype=torch.int32, device="cuda") + metadata.msa_seq_lens_cuda = torch.zeros(1, dtype=torch.int32, device="cuda") + metadata.msa_subpage_block_table = None + metadata._msa_runs_no_fmha = lambda: True + metadata._msa_kv_lens_may_change = lambda: False + + def producer( + projection: torch.Tensor, + kv: torch.Tensor, + index_k: torch.Tensor, + slots: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + projection, + kv, + index_k, + slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + *weights, + rope_cache, + positions, + ) + + # Exercise the production PCG runner, including its actual capture/replay, + # rather than replacing the producer or its cache mutations with mocks. + graph = torch.fx.symbolic_trace(producer) + runner = PiecewiseRunner( + graph=graph, + name="empty_adp_rank_producer", + compile_time_num_tokens=bucket, + runtime_num_tokens_idx=None, + capture_num_tokens=[bucket], + graph_pool_handle=torch.cuda.graph_pool_handle(), + default_callable=graph.forward, + enable_inductor=False, + is_first_runner=True, + is_last_runner=True, + ) + previous_prefill_flag = get_per_request_prefill_cuda_graph_flag() + try: + with model_extra_attrs({}), piecewise_cuda_graph(True): + # First capture with both ranks live, then alternate which rank is + # empty. Each empty replay therefore has stale live slots to clear. + for empty_rank in (None, 0, 1): + live_tokens = 0 if rank == empty_rank else 3 + local_contexts = int(live_tokens > 0) + all_tokens = dist.tp_allgather_int64([live_tokens])[:, 0].tolist() + padded, eligible, padded_all_tokens = get_padding_params( + live_tokens, + local_contexts, + all_tokens, + dist=dist, + enable_attention_dp=True, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + prefill_cuda_graph_num_tokens=[bucket], + ) + assert eligible and padded == bucket + assert padded_all_tokens == [bucket] * world_size + if empty_rank is not None: + assert all_tokens[empty_rank] == 0 and max(all_tokens) == 3 + set_per_request_prefill_cuda_graph_flag(eligible) + + metadata.request_ids = [0] if live_tokens else [] + metadata._msa_qo_lens_cpu = torch.tensor( + [live_tokens] if live_tokens else [], dtype=torch.int32 + ) + metadata._msa_kv_lens_cpu = metadata._msa_qo_lens_cpu.clone() + metadata._msa_qo_offset_cpu = torch.zeros(local_contexts, dtype=torch.int32) + metadata._build_msa_fields() + assert metadata.msa_out_cache_loc.tolist() == list( + range(128, 128 + live_tokens) + ) + [-1] * (bucket - live_tokens) + packed.normal_() + before_main, before_index = main_backing.clone(), index_backing.clone() + args = (packed, main_cache, index_cache, metadata.msa_out_cache_loc) + if empty_rank is None: + with capture_piecewise_cuda_graph(True): + # Three warmups, then capture. Keep the capture outputs + # alive because the runner stores non-owning views. + for _ in range(4): + capture_output = runner(*args) + assert runner.entries[bucket].cuda_graph is not None + else: + runner(*args) + torch.cuda.synchronize() + assert all(torch.isfinite(output.float()).all() for output in capture_output) + + for cache, before, page in ( + (main_backing, before_main, 3), + (index_backing, before_index, 5), + ): + if live_tokens: + assert not torch.equal( + cache[page, ..., :live_tokens, :].view(torch.uint8), + before[page, ..., :live_tokens, :].view(torch.uint8), + ), "prefilling rank did not update its cache" + # Exclude only legal live-token writes; padding, other + # pages, and interleaved storage must remain unchanged. + before[page, ..., :live_tokens, :].copy_(cache[page, ..., :live_tokens, :]) + assert torch.equal(cache.view(torch.uint8), before.view(torch.uint8)), ( + f"rank {rank}: empty/padded cache storage was modified" + ) + finally: + set_per_request_prefill_cuda_graph_flag(previous_prefill_flag) + runner.clear_cuda_graphs() + return rank + + +@pytest.mark.skip_less_device(2) +def test_minimax_m3_piecewise_empty_adp_rank_preserves_caches() -> None: + """An empty ADP rank replays PCG without writing main K/V or index-K.""" + if torch.cuda.device_count() < 2: + pytest.skip("requires two CUDA devices") + with MPIPoolExecutor(max_workers=2) as executor: + ranks = list(executor.map(_run_empty_adp_rank, [2, 2])) + assert sorted(ranks) == [0, 1] From 57ba84769c74fd10e853764cfd9f1ea1e0759ba7 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:18:24 -0700 Subject: [PATCH 07/12] [None][fix] Scope piecewise eager fallback to MiniMax-M3 Keep other decoder models on their existing compiled fallback. Preserve explicit MXFP8 backend overrides, register shared weights once, and invalidate only unused MSA cache slots. Cover model opt-in, compile state, reload traversal and empty-rank metadata. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../backends/sparse/minimax_m3/msa_backend.py | 10 +- .../_torch/compilation/remove_copy_pass.py | 1 - .../_torch/custom_ops/cpp_custom_ops.py | 1 - .../_torch/models/modeling_minimaxm3.py | 2 + tensorrt_llm/_torch/models/modeling_utils.py | 7 +- .../_torch/pyexecutor/model_engine.py | 38 ++-- tensorrt_llm/_torch/utils.py | 6 + .../attention/sparse/msa/test_msa_backend.py | 16 +- .../test_pytorch_model_engine_warmup.py | 178 ++++++++++++++++-- 9 files changed, 208 insertions(+), 51 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 02610e5ad282..7692baa199d9 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -976,18 +976,16 @@ def _build_msa_fields(self) -> None: self._msa_fields_ready = False if not self._msa_buffers_ready: return - # Captured producers execute the whole padded bucket, including on - # attention-DP ranks without local requests. Invalidate the tail before - # any early return so replay cannot write padding into stale KV slots. - self.msa_out_cache_loc.fill_(-1) request_ids = self.request_ids qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu if request_ids is None or qo_lens_cpu is None: + self.msa_out_cache_loc.fill_(-1) return batch_size = int(qo_lens_cpu.shape[0]) if batch_size == 0: + self.msa_out_cache_loc.fill_(-1) return kv_cache_manager = self.kv_cache_manager @@ -1037,6 +1035,10 @@ def _build_msa_fields(self) -> None: ) self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) + # Captured producers also execute padded rows. Invalidate only the + # unwritten tail so they cannot reuse the previous step's live slots. + if total_new_tokens < self.msa_out_cache_loc.shape[0]: + self.msa_out_cache_loc[total_new_tokens:].fill_(-1) if kv_indices is not None: self.msa_kv_indices[: int(kv_indices.shape[0])].copy_(kv_indices, non_blocking=True) diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py index ab253aea9ee3..fd536e11b70c 100644 --- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py +++ b/tensorrt_llm/_torch/compilation/remove_copy_pass.py @@ -32,7 +32,6 @@ def remove_copy_for_mutates_args(graph: Graph): nodes_to_remove: list[Node] = [] def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): - """Restore in-place calls while preserving regular and mutated outputs.""" getitem_nodes = [ user for user in node.users if is_call_function(user, getitem) ] diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index e8496115fa30..1261612fcb5e 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -15,7 +15,6 @@ def _register_fake(): - """Register shape-only implementations for native operators during tracing.""" @torch.library.register_fake("trtllm::allreduce") def allreduce( diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 127e3b254a08..3f5a5e26158b 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -2690,6 +2690,8 @@ def _fold_gemma_boundary_norm_weights(weights): class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" + use_prefill_only_compile = True + @classmethod def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: """Prefer KV cache manager V2 for MiniMax-M3. diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0c16cb0972f3..ba922b255577 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -8,8 +8,8 @@ import os import time from dataclasses import dataclass -from typing import (Any, Callable, Dict, Generic, Iterator, List, Literal, - Optional, Tuple, Type, TypeVar, Union) +from typing import (Any, Callable, ClassVar, Dict, Generic, Iterator, List, + Literal, Optional, Tuple, Type, TypeVar, Union) import torch from torch import nn @@ -388,6 +388,9 @@ class DecoderModelForCausalLM(nn.Module, Generic[TModel, TConfig], metaclass=PostInitCaller): + # Opt in to original eager execution outside capture-eligible PCG batches. + use_prefill_only_compile: ClassVar[bool] = False + @staticmethod def _checkpoint_has_lm_head_scale(config: ModelConfig[TConfig]) -> bool: """Whether the checkpoint stores a quantized lm_head (a weight scale). diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5b227a7664fc..270f4aad489e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -196,8 +196,8 @@ def _make_single_token_context_graph_batch( return graph_batch, promoted_context_request_ids -class _ContextOnlyCompiledModel(torch.nn.Module): - """Share parameters between captured context and eager generation paths. +class _PrefillCompiledModel(torch.nn.Module): + """Share weights between eligible prefill/mixed and original eager paths. The prefill flag includes the all-rank attention-DP decision and capture ceiling. A decode-only rank must still compile when another rank prefills. @@ -208,7 +208,9 @@ def __init__(self, eager_model: torch.nn.Module, """Keep eager and compiled entry points sharing the same model weights.""" super().__init__() self.eager_model = eager_model - self.compiled_model = compiled_model + # The compiled callable references the same weights. Register only the + # eager tree so state_dict(), children() and _apply() visit it once. + object.__setattr__(self, "compiled_model", compiled_model) def named_modules( self, @@ -218,8 +220,7 @@ def named_modules( ) -> Iterator[Tuple[str, torch.nn.Module]]: """Expose checkpoint-compatible module names for partial weight reloads.""" # Weight reloads match checkpoint prefixes against this traversal. - # Neither routing wrapper owns parameters; expose the original tree - # once, including when the loader requests remove_duplicate=False. + # Hide the eager_model prefix even with remove_duplicate=False. yield from self.eager_model.named_modules(memo, prefix, remove_duplicate) @@ -232,8 +233,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: def __getattr__(self, name: str) -> Any: """Delegate model-specific attributes to the original eager model.""" - # Model-specific epilogues (including M3 Eagle3) access embed_tokens - # and other transformer attributes after the wrapped forward returns. + # Epilogues can access transformer attributes after forward returns. try: return super().__getattr__(name) except AttributeError: @@ -378,7 +378,6 @@ def __init__( model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): - """Initialize model execution, cache management, and graph configuration.""" _configure_deep_gemm_pdl() self.forward_pass_callable = None @@ -644,7 +643,7 @@ def __init__( self._torch_compile_enabled = torch_compile_enabled self._torch_compile_piecewise_cuda_graph = torch_compile_piecewise_cuda_graph - self._torch_compile_context_only = False + self._torch_compile_prefill_only = False prefill_cuda_graph_num_tokens = self.llm_args.prefill_capture_num_tokens if prefill_cuda_graph_num_tokens is None: @@ -696,10 +695,12 @@ def __init__( eager_model, backend=self._torch_compile_backend, fullgraph=torch_compile_fullgraph) - self._torch_compile_context_only = self._torch_compile_piecewise_cuda_graph + self._torch_compile_prefill_only = ( + self._torch_compile_piecewise_cuda_graph + and self.model.use_prefill_only_compile) self.model.model = ( - _ContextOnlyCompiledModel(eager_model, compiled_model) - if self._torch_compile_context_only else compiled_model) + _PrefillCompiledModel(eager_model, compiled_model) + if self._torch_compile_prefill_only else compiled_model) elif callable(apply_llm_torch_compile): # TODO: Move this contract to MultimodalModelMixin once # multimodal models consistently expose their LLM compile @@ -2184,12 +2185,12 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: native_mxfp8_methods = [ method for method in mxfp8_methods if method.needs_native_autotune ] - compile_all_batches = ( - self._torch_compile_enabled - and not getattr(self, "_torch_compile_context_only", False)) - if compile_all_batches: + compile_all_batches = (self._torch_compile_enabled + and not self._torch_compile_prefill_only) + if (compile_all_batches + and "TRTLLM_MXFP8_GEMM_BACKEND" not in os.environ): # Compiled auto dispatch uses native; do not tune unused backends. - # Context-only compile retains the eager generation-graph policy. + # Prefill-only compile retains the eager generation-graph policy. for method in mxfp8_methods: method.disable_flashinfer_auto() use_mxfp8_flashinfer_graph_default = ( @@ -6298,7 +6299,6 @@ def capture_postprocess_fn(inputs: Dict[str, Any]): return outputs def model_forward(self, **kwargs): - """Run the full model under the current batch's compile and graph flags.""" attrs = get_model_extra_attrs() assert attrs is not None, "Model extra attrs is not set" attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) @@ -6322,7 +6322,7 @@ def model_forward(self, **kwargs): # eager decode and over-ceiling prefill do not select compile-only ops. compile_scope = ( torch_compiling(get_per_request_prefill_cuda_graph_flag()) - if self._torch_compile_context_only else contextlib.nullcontext()) + if self._torch_compile_prefill_only else contextlib.nullcontext()) with reclaim_scope, compile_scope: if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): return trace_func(self.model.forward)(**kwargs) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index da6797268ffe..dd94124bda1d 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -100,6 +100,12 @@ def set_torch_compiling(enable: bool): def is_torch_compiling() -> bool: + """Return the runtime compile-dispatch policy, not Dynamo tracing state. + + Prefill-only models scope this process-global flag to each full forward: + eligible prefill/mixed batches use True, original eager fallbacks False. + Other compiled models retain the engine-wide value. + """ global is_torch_compiling_flag return is_torch_compiling_flag diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index e48d00915f71..e8251e29578b 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -40,7 +40,10 @@ @pytest.mark.cpu_only -def test_msa_metadata_clears_padded_cache_slot_tail(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("empty_state", ["empty_batch", "no_requests", "no_lengths"]) +def test_msa_metadata_clears_padded_cache_slot_tail( + monkeypatch: pytest.MonkeyPatch, empty_state: str +) -> None: """A smaller replay must not reuse the previous step's live cache slots.""" from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend @@ -62,7 +65,7 @@ def test_msa_metadata_clears_padded_cache_slot_tail(monkeypatch: pytest.MonkeyPa metadata._msa_kv_lens_may_change = lambda: False original_ptr = metadata.msa_out_cache_loc.data_ptr() - for count in (4, 2, 1): + for count in (4, 2, 0, 1): metadata._msa_qo_lens_cpu = torch.tensor([count], dtype=torch.int32) metadata._msa_kv_lens_cpu = metadata._msa_qo_lens_cpu.clone() metadata._msa_qo_offset_cpu = torch.zeros(1, dtype=torch.int32) @@ -73,8 +76,13 @@ def test_msa_metadata_clears_padded_cache_slot_tail(monkeypatch: pytest.MonkeyPa assert metadata.msa_out_cache_loc.data_ptr() == original_ptr assert metadata._msa_fields_ready - metadata.request_ids = [] - metadata._msa_qo_lens_cpu = torch.empty(0, dtype=torch.int32) + if empty_state == "empty_batch": + metadata.request_ids = [] + metadata._msa_qo_lens_cpu = torch.empty(0, dtype=torch.int32) + elif empty_state == "no_requests": + metadata.request_ids = None + else: + metadata._msa_qo_lens_cpu = None metadata._build_msa_fields() assert metadata.msa_out_cache_loc.tolist() == [-1] * 4 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 7878f7668de3..9e7361135d87 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -32,10 +32,7 @@ from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.model_engine import ( - PyTorchModelEngine, - _ContextOnlyCompiledModel, -) +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine, _PrefillCompiledModel from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.speculative.utils import update_draft_len @@ -45,14 +42,143 @@ DecodingBaseConfig, DraftTargetDecodingConfig, PARDDecodingConfig, + PrefillCudaGraphBackend, TorchLlmArgs, ) from tensorrt_llm.mapping import Mapping +@pytest.mark.cpu_only +@pytest.mark.parametrize("model_kind", ["default", "m3", "m3_vl", "gdn", "mamba"]) +@pytest.mark.parametrize("compile_enabled,piecewise", [(False, False), (True, False), (True, True)]) +def test_prefill_compile_initialization_is_model_opt_in( + monkeypatch: pytest.MonkeyPatch, + model_kind: str, + compile_enabled: bool, + piecewise: bool, +) -> None: + """Exercise the real constructor's routing gate without loading weights or CUDA.""" + from tensorrt_llm._torch.models.modeling_minimaxm3 import ( + MiniMaxM3ForCausalLM, + MiniMaxM3VLForConditionalGeneration, + ) + from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM + from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM + from tensorrt_llm._torch.pyexecutor import _util + + model_cls = { + "default": DecoderModelForCausalLM, + "m3": MiniMaxM3ForCausalLM, + "m3_vl": MiniMaxM3VLForConditionalGeneration, + "gdn": Qwen3NextForCausalLM, + "mamba": NemotronHForCausalLM, + }[model_kind] + model = model_cls.__new__(model_cls) + torch.nn.Module.__init__(model) + mapping = Mapping() + model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(torch_dtype=torch.float32, hidden_size=4), + mapping=mapping, + sparse_attention_config=None, + ) + eager = torch.nn.Linear(4, 4) + model.model = eager + compile_config = SimpleNamespace( + enable_fullgraph=True, enable_inductor=False, enable_userbuffers=False, max_num_streams=1 + ) + llm_args = SimpleNamespace( + encode_only=False, + mm_encoder_only=False, + get_runtime_sizes=lambda: (1, 16, 64, 4), + encoder_max_batch_size=None, + encoder_max_num_tokens=None, + enable_in_graph_sampling=False, + multimodal_config=SimpleNamespace(video_pruning_rate=None), + checkpoint_format="HF", + trust_remote_code=False, + disable_overlap_scheduler=True, + kv_cache_config=KvCacheConfig(), + enable_layerwise_nvtx_marker=False, + cuda_graph_config=None, + torch_compile_config=compile_config if compile_enabled else None, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE if piecewise else None, + prefill_capture_num_tokens=[4], + allreduce_strategy="AUTO", + attn_backend="TRTLLM", + sparse_attention_config=None, + ) + for name in ( + "should_enable_adp_dummy_fixes", + "should_enable_scheduler_aware_adp_dummy", + "should_enable_non_overlap_adp_forward_intent", + "should_enable_overlap_headroom", + "resolved_kv_cache_manager_is_v2", + ): + monkeypatch.setattr(_util, name, Mock(return_value=False)) + monkeypatch.setattr(_util, "compute_max_num_sequences", Mock(return_value=4)) + for name in ( + "_configure_deep_gemm_pdl", + "create_input_processor", + "setup_mm_encoder_attn_metadata", + ): + monkeypatch.setattr(model_engine_module, name, Mock()) + monkeypatch.setattr( + model_engine_module, "resolve_mrope_position_deltas_cache", lambda model: None + ) + monkeypatch.setattr( + model_engine_module, "is_hybrid_linear", lambda config: model_kind in ("gdn", "mamba") + ) + monkeypatch.setattr( + model_engine_module.MultimodalItemScheduler, "maybe_create", Mock(return_value=None) + ) + monkeypatch.setattr(PyTorchModelEngine, "_validate_breakable_cuda_graph_compatibility", Mock()) + monkeypatch.setattr(PyTorchModelEngine, "_init_model_capacity", Mock()) + monkeypatch.setattr(PyTorchModelEngine, "__del__", lambda self: None) + backend_factory = Mock() + backend_factory.Streams = list + monkeypatch.setattr(model_engine_module, "Backend", backend_factory) + compiled = torch.nn.Identity() + compile_model = Mock(return_value=compiled) + monkeypatch.setattr(torch, "compile", compile_model) + monkeypatch.setattr(torch._dynamo.config, "cache_size_limit", 16) + # Stop after the complete compilation block, before runtime cache allocation. + monkeypatch.setattr( + model_engine_module, + "get_attention_backend", + Mock(side_effect=RuntimeError("compile setup complete")), + ) + engine = PyTorchModelEngine.__new__(PyTorchModelEngine) + with torch_compiling(False), pytest.raises(RuntimeError, match="compile setup complete"): + PyTorchModelEngine.__init__( + engine, + model_path="dummy", + mapping=mapping, + model=model, + llm_args=llm_args, + checkpoint_loader=Mock(), + ) + + expected_prefill_only = compile_enabled and piecewise and model_kind in ("m3", "m3_vl") + assert engine._torch_compile_prefill_only is expected_prefill_only + if not compile_enabled: + compile_model.assert_not_called() + assert model.model is eager + else: + compile_model.assert_called_once_with( + eager, backend=engine._torch_compile_backend, fullgraph=True + ) + assert backend_factory.call_args.args[0] is False # Inductor stays disabled. + if expected_prefill_only: + assert isinstance(model.model, _PrefillCompiledModel) + assert model.model.eager_model is eager + assert model.model.compiled_model is compiled + else: + assert model.model is compiled + + @pytest.mark.cpu_only @pytest.mark.parametrize("local_contexts", [0, 1]) -def test_context_only_compile_uses_all_rank_prefill_decision( +def test_prefill_compile_uses_all_rank_prefill_decision( monkeypatch: pytest.MonkeyPatch, local_contexts: int, ) -> None: @@ -63,7 +189,7 @@ def test_context_only_compile_uses_all_rank_prefill_decision( expected = torch.ones((2, 4)) eager.forward = Mock(return_value=expected) compiled.forward = Mock(return_value=expected) - router = _ContextOnlyCompiledModel(eager, compiled) + router = _PrefillCompiledModel(eager, compiled) assert router.weight is eager.weight assert list(router.parameters()) == list(eager.parameters()) # Local decode-only ranks participate when another attention-DP rank @@ -80,7 +206,7 @@ def test_context_only_compile_uses_all_rank_prefill_decision( @pytest.mark.cpu_only -def test_context_only_compile_preserves_partial_weight_reload( +def test_prefill_compile_preserves_partial_weight_reload( monkeypatch: pytest.MonkeyPatch, ) -> None: """Reload selected weights by original names into both model entry points.""" @@ -90,8 +216,10 @@ def test_context_only_compile_preserves_partial_weight_reload( compiled = torch.compile(eager, backend="eager") model = DecoderModelForCausalLM.__new__(DecoderModelForCausalLM) torch.nn.Module.__init__(model) - model.config = SimpleNamespace(tie_word_embeddings=False) - model.model = _ContextOnlyCompiledModel(eager, compiled) + model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(tie_word_embeddings=False) + ) + model.model = _PrefillCompiledModel(eager, compiled) mapper = HfWeightMapper() mapper._model = model loader = ModelLoader.__new__(ModelLoader) @@ -108,6 +236,14 @@ def test_context_only_compile_preserves_partial_weight_reload( names = dict(model.named_modules(remove_duplicate=remove_duplicate)) assert names["model.layers.0"] is eager.layers[0] assert not any("eager_model" in name or "compiled_model" in name for name in names) + assert list(model.model.children()) == [eager] + assert list(model.model.state_dict()) == [ + "eager_model.layers.0.weight", + "eager_model.layers.0.bias", + ] + apply_tensor = Mock(side_effect=lambda tensor: tensor) + model.model._apply(apply_tensor) + assert apply_tensor.call_count == 2 loader.reload(model, {"model.layers.0.weight": replacement}, allow_partial_loading=True) assert eager.layers[0].weight is weight @@ -123,10 +259,12 @@ def test_context_only_compile_preserves_partial_weight_reload( @pytest.mark.cpu_only +@pytest.mark.parametrize("prefill_only", [False, True]) @pytest.mark.parametrize("eligible", [False, True]) @pytest.mark.parametrize("raises", [False, True]) -def test_context_only_compile_scopes_whole_model_forward( +def test_prefill_compile_scopes_whole_model_forward( monkeypatch: pytest.MonkeyPatch, + prefill_only: bool, eligible: bool, raises: bool, ) -> None: @@ -145,7 +283,7 @@ def forward(**kwargs: object) -> str: engine = SimpleNamespace( model=SimpleNamespace(model_config=SimpleNamespace(extra_attrs={}), forward=forward), _torch_compile_backend=None, - _torch_compile_context_only=True, + _torch_compile_prefill_only=prefill_only, _eager_workspace_reclaimer=None, is_warmup=False, ) @@ -161,15 +299,16 @@ def forward(**kwargs: object) -> str: else: assert PyTorchModelEngine.model_forward(engine, attn_metadata=Mock()) == "done" assert is_torch_compiling() - assert observed == [eligible] * (1 if raises else 2) + expected = eligible if prefill_only else True + assert observed == [expected] * (1 if raises else 2) @pytest.mark.cpu_only -@pytest.mark.parametrize("context_only", [False, True]) -@pytest.mark.parametrize("backend", [None, "auto", "flashinfer"]) +@pytest.mark.parametrize("prefill_only", [False, True]) +@pytest.mark.parametrize("backend", [None, "auto", "flashinfer", "trtllm"]) def test_compiled_mxfp8_warmup_backend_selection( monkeypatch: pytest.MonkeyPatch, - context_only: bool, + prefill_only: bool, backend: str | None, ) -> None: """PCG retains eager decode tuning; all-batch compile settles auto on native.""" @@ -187,7 +326,7 @@ def test_compiled_mxfp8_warmup_backend_selection( engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), _torch_compile_enabled=True, - _torch_compile_context_only=context_only, + _torch_compile_prefill_only=prefill_only, cuda_graph_runner=SimpleNamespace(enabled=True), model=SimpleNamespace( modules=lambda: [ @@ -226,10 +365,9 @@ def test_compiled_mxfp8_warmup_backend_selection( PyTorchModelEngine._run_autotuner_warmup(engine, resources) - flashinfer_expected = context_only or backend == "flashinfer" - assert method.backend == ( - "flashinfer" if backend == "flashinfer" else "auto" if context_only else "trtllm" - ) + expected_backend = backend or ("auto" if prefill_only else "trtllm") + flashinfer_expected = expected_backend in ("auto", "flashinfer") + assert method.backend == expected_backend assert method._native_autotuned assert method._flashinfer_autotuned == flashinfer_expected assert flashinfer_tune.call_count == int(flashinfer_expected) From 5119f0c208c00676b3bcbcdb4d24232d64f5cd78 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 21 Sep 2026 10:30:28 -0700 Subject: [PATCH 08/12] [None][fix] Clarify PCG fallback policy and skip unused MXFP8 tuning Rename the model policy to use_fx_for_pcg_fallback, defaulting to true and disabled only by MiniMax-M3. Preserve explicit auto backend settings while excluding compile-suppressed FlashInfer warmup. Exercise actual engine compile scopes and MXFP8 dispatch in the warmup regression. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3.py | 2 +- tensorrt_llm/_torch/models/modeling_utils.py | 4 +- .../_torch/pyexecutor/model_engine.py | 7 +- .../test_pytorch_model_engine_warmup.py | 71 +++++++++++++++---- 4 files changed, 67 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 3f5a5e26158b..d43f8064b421 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -2690,7 +2690,7 @@ def _fold_gemma_boundary_norm_weights(weights): class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" - use_prefill_only_compile = True + use_fx_for_pcg_fallback = False @classmethod def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index ba922b255577..f99a146c342a 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -388,8 +388,8 @@ class DecoderModelForCausalLM(nn.Module, Generic[TModel, TConfig], metaclass=PostInitCaller): - # Opt in to original eager execution outside capture-eligible PCG batches. - use_prefill_only_compile: ClassVar[bool] = False + # Keep FX optimizations for decode and prefill above the PCG capture ceiling. + use_fx_for_pcg_fallback: ClassVar[bool] = True @staticmethod def _checkpoint_has_lm_head_scale(config: ModelConfig[TConfig]) -> bool: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 270f4aad489e..cdf482c6a0d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -697,7 +697,7 @@ def __init__( fullgraph=torch_compile_fullgraph) self._torch_compile_prefill_only = ( self._torch_compile_piecewise_cuda_graph - and self.model.use_prefill_only_compile) + and not self.model.use_fx_for_pcg_fallback) self.model.model = ( _PrefillCompiledModel(eager_model, compiled_model) if self._torch_compile_prefill_only else compiled_model) @@ -2206,9 +2206,12 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: quant_method.enable_flashinfer_auto() quant_method.tune_decode_graph_backends = ( tune_with_cute_dsl and quant_method.uses_flashinfer) + # An explicit auto setting stays intact, but compiled auto dispatch + # uses native GEMM. Do not run or mark an unused FlashInfer tuning pass. flashinfer_mxfp8_methods = [ method for method in mxfp8_methods - if method.needs_flashinfer_autotune + if method.needs_flashinfer_autotune and ( + not compile_all_batches or method.backend == "flashinfer") ] # Every TP and PP rank must make the same backend decision before any 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 9e7361135d87..b930cdc44b4b 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -51,7 +51,7 @@ @pytest.mark.cpu_only @pytest.mark.parametrize("model_kind", ["default", "m3", "m3_vl", "gdn", "mamba"]) @pytest.mark.parametrize("compile_enabled,piecewise", [(False, False), (True, False), (True, True)]) -def test_prefill_compile_initialization_is_model_opt_in( +def test_pcg_fx_fallback_policy_is_model_specific( monkeypatch: pytest.MonkeyPatch, model_kind: str, compile_enabled: bool, @@ -73,6 +73,7 @@ def test_prefill_compile_initialization_is_model_opt_in( "gdn": Qwen3NextForCausalLM, "mamba": NemotronHForCausalLM, }[model_kind] + assert model_cls.use_fx_for_pcg_fallback is (model_kind not in ("m3", "m3_vl")) model = model_cls.__new__(model_cls) torch.nn.Module.__init__(model) mapping = Mapping() @@ -304,35 +305,54 @@ def forward(**kwargs: object) -> str: @pytest.mark.cpu_only -@pytest.mark.parametrize("prefill_only", [False, True]) +@pytest.mark.parametrize("compile_mode", ["eager", "all_batches", "prefill_only"]) @pytest.mark.parametrize("backend", [None, "auto", "flashinfer", "trtllm"]) def test_compiled_mxfp8_warmup_backend_selection( monkeypatch: pytest.MonkeyPatch, - prefill_only: bool, + compile_mode: str, backend: str | None, ) -> None: - """PCG retains eager decode tuning; all-batch compile settles auto on native.""" + """Tune only backends reached by real dispatch, preserving explicit choices.""" import tensorrt_llm._torch.modules.linear as linear_module + compile_enabled = compile_mode != "eager" + prefill_only = compile_mode == "prefill_only" + inputs = torch.zeros(2, 4) + output = torch.zeros(2, 3) + layer = SimpleNamespace( + weight=torch.zeros(3, 4), weight_scale=torch.ones(4), dtype=torch.float32 + ) + native_gemm = Mock(return_value=output) + flashinfer_gemm = Mock(return_value=output) monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) monkeypatch.delenv("TLLM_AUTOTUNER_CACHE_PATH", raising=False) if backend is not None: monkeypatch.setenv("TRTLLM_MXFP8_GEMM_BACKEND", backend) monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) - monkeypatch.setattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", Mock(), raising=False) + monkeypatch.setattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", flashinfer_gemm, raising=False) + monkeypatch.setattr( + torch.ops.trtllm, "mxfp8_quantize", Mock(return_value=(inputs, inputs)), raising=False + ) + for name in ("mxfp8_mxfp8_gemm", "mxfp8_mxfp8_gemm_autotuned"): + monkeypatch.setattr(torch.ops.trtllm, name, native_gemm, raising=False) flashinfer_tune = Mock(return_value=contextlib.nullcontext()) monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(autotune=flashinfer_tune)) method = MXFP8LinearMethod() engine = SimpleNamespace( llm_args=SimpleNamespace(enable_autotuner=True), - _torch_compile_enabled=True, + _torch_compile_enabled=compile_enabled, _torch_compile_prefill_only=prefill_only, + _torch_compile_backend=None, + _eager_workspace_reclaimer=None, + is_warmup=True, cuda_graph_runner=SimpleNamespace(enabled=True), model=SimpleNamespace( modules=lambda: [ SimpleNamespace(_use_flashinfer_mxfp8_decode_graph_default=True), SimpleNamespace(quant_method=method), - ] + ], + model_config=SimpleNamespace(extra_attrs={}), + forward=lambda **kwargs: method.apply(layer, inputs, None), ), mapping=SimpleNamespace(tp_size=1, has_pp=lambda: False), dist=object(), @@ -345,8 +365,10 @@ def test_compiled_mxfp8_warmup_backend_selection( is_draft_model=False, guided_decoder=None, no_cuda_graph=lambda: contextlib.nullcontext(), - _create_warmup_request=Mock(return_value=object()), - _release_batch_context=lambda *args: contextlib.nullcontext(object()), + _create_warmup_request=lambda resources, num_tokens, num_gen_requests: Mock( + num_gen_requests=num_gen_requests + ), + _release_batch_context=lambda batch, resources: contextlib.nullcontext(batch), _should_run_warmup_batch=Mock(return_value=True), _release_megamoe_profiling_scratch=Mock(), forward=Mock(), @@ -362,16 +384,41 @@ def test_compiled_mxfp8_warmup_backend_selection( monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(model_engine_module, "clear_memory_buffers", lambda: None) + monkeypatch.setattr(model_engine_module, "get_model_extra_attrs", lambda: {}) + monkeypatch.setattr(model_engine_module, "is_trace_enabled", lambda name: False) + + def forward(batch: Mock, **kwargs: object) -> torch.Tensor: + # Stand in for input preparation, but use the real engine compile scope + # and linear dispatch for each prefill/generation warmup batch. + monkeypatch.setattr( + model_engine_module, + "get_per_request_prefill_cuda_graph_flag", + lambda: batch.num_gen_requests == 0, + ) + return PyTorchModelEngine.model_forward(engine, attn_metadata=batch) - PyTorchModelEngine._run_autotuner_warmup(engine, resources) + engine.forward.side_effect = forward + with torch_compiling(compile_enabled): + PyTorchModelEngine._run_autotuner_warmup(engine, resources) + assert is_torch_compiling() is compile_enabled - expected_backend = backend or ("auto" if prefill_only else "trtllm") - flashinfer_expected = expected_backend in ("auto", "flashinfer") + expected_backend = backend or ("trtllm" if compile_mode == "all_batches" else "auto") + flashinfer_expected = expected_backend == "flashinfer" or ( + expected_backend == "auto" and compile_mode != "all_batches" + ) assert method.backend == expected_backend assert method._native_autotuned assert method._flashinfer_autotuned == flashinfer_expected assert flashinfer_tune.call_count == int(flashinfer_expected) assert engine.forward.call_count == (4 if flashinfer_expected else 2) + expected_flashinfer_calls = ( + (4 if expected_backend == "flashinfer" else (1 if prefill_only else 2)) + if flashinfer_expected + else 0 + ) + assert flashinfer_gemm.call_count == expected_flashinfer_calls + assert native_gemm.call_count == engine.forward.call_count - expected_flashinfer_calls + assert os.environ.get("TRTLLM_MXFP8_GEMM_BACKEND") == backend @pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig]) From 1ddef7a49c1d681bc70960a539a3ac48239b6986 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:03:17 -0700 Subject: [PATCH 09/12] [None][fix] Address MiniMax-M3 PCG review follow-ups Explain the opaque fused producer and M3 fallback policy. Use direct metadata access, strengthen PCG test fixtures, and move test helpers to module scope. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3.py | 7 +- .../test_fused_qk_norm_rope.py | 5 +- .../unittest/_torch/models/test_minimax_m3.py | 66 ++++++++++--------- 3 files changed, 44 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index d43f8064b421..a7722bccfc1e 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -838,6 +838,9 @@ def _minimax_m3_qkv_index_proj_fake( return hidden_states.new_empty((hidden_states.shape[0], sum(qkv_proj.local_output_sizes))) +# Projection and cache insertion inspect runtime shapes and layouts. Keep those +# checks opaque to avoid Dynamo specialization or graph breaks while PCG captures +# the fused kernels; explicit mutable cache inputs keep their writes visible. @torch.library.custom_op( "trtllm::minimax_m3_fused_sparse_qkv_producer", mutates_args=("kv_cache", "index_k_cache"), @@ -1384,7 +1387,7 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( return None if cache_tensors is None: - kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) + kv_cache_manager = attn_metadata.kv_cache_manager if kv_cache_manager is None: return None buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") @@ -2690,6 +2693,8 @@ def _fold_gemma_boundary_norm_weights(weights): class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" + # Preserve M3's hand-fused eager decode and above-ceiling prefill paths, + # including MXFP8 decode-graph backend tuning, instead of the FX fallback. use_fx_for_pcg_fallback = False @classmethod diff --git a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py index b759661c3a9c..b8db4ae76479 100644 --- a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py +++ b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_fused_qk_norm_rope.py @@ -14,6 +14,8 @@ # limitations under the License. import pytest import torch +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv from tensorrt_llm._torch.attention.backends.interface import RopeParams from tensorrt_llm._torch.attention.rotary_embedding import MRotaryEmbedding, RotaryEmbedding @@ -24,9 +26,6 @@ @pytest.mark.parametrize("producer", ["norm_rope", "main_kv", "horizontal"]) def test_fp8_producer_meta_keeps_dynamic_num_tokens(producer: str) -> None: """All FP8 fake kernels must retain the symbolic token dimension.""" - from torch._subclasses.fake_tensor import FakeTensorMode - from torch.fx.experimental.symbolic_shapes import ShapeEnv - with FakeTensorMode(shape_env=ShapeEnv()) as mode: num_tokens = mode.shape_env.create_unbacked_symint() qkv = torch.empty((num_tokens, 1280), dtype=torch.bfloat16) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 01680035f09b..32705c5708e8 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -23,23 +23,37 @@ import json import os from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from safetensors import safe_open from torch import nn +from torch._dynamo.backends.common import aot_autograd +from torch._functorch.aot_autograd import make_boxed_func +from torch._higher_order_ops.auto_functionalize import auto_functionalized, auto_functionalized_v2 +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv from transformers import AutoConfig from utils.llm_data import llm_models_root import tensorrt_llm._torch.models.modeling_minimaxm3 as modeling_minimaxm3 -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3MsaSparseAttention, + MiniMaxM3SparseRuntimeBackend, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( MiniMaxM3SparseConfig, MiniMaxM3SparseMetadataParams, MiniMaxM3SparseParams, index_head_range, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _group_max_reduce +from tensorrt_llm._torch.compilation.remove_copy_pass import remove_copy_for_mutates_args from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( MiniMaxM3HfWeightMapper, @@ -417,7 +431,8 @@ def _dispatch_attention_backend(self, q, k, v, idx_q, idx_k, attn_metadata, outp assert idx_q.shape == (attn_metadata.num_tokens, 1) output.copy_(q) - metadata = SimpleNamespace(num_tokens=2) + metadata = AttentionMetadata(max_num_requests=1, max_num_tokens=4) + metadata._num_tokens = 2 layer = FakeAttentionLayer() monkeypatch.setattr( modeling_minimaxm3, @@ -451,10 +466,11 @@ def test_piecewise_projection_fake_preserves_padded_hidden_rows(monkeypatch) -> nn.Module.__init__(projection) projection.local_output_sizes = (3, 4) layer = SimpleNamespace(qkv_proj=projection) + metadata = AttentionMetadata(max_num_requests=1, max_num_tokens=256) monkeypatch.setattr( modeling_minimaxm3, "_extract_minimax_m3_attention_extra_attrs", - lambda layer_idx: (SimpleNamespace(), layer), + lambda layer_idx: (metadata, layer), ) hidden_states = torch.randn(256, 5) position_ids = torch.arange(6).reshape(1, 6) @@ -471,6 +487,8 @@ def test_piecewise_fused_projection_preserves_input_token_dimension(monkeypatch) """Do not inherit a bucket-specialized token dimension from the GEMM output.""" packed = torch.randn(2, 7) captured = {} + metadata = AttentionMetadata(max_num_requests=1, max_num_tokens=6) + metadata._num_tokens = 6 def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, output): assert q is None and k is None and v is None @@ -483,7 +501,7 @@ def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, ou layer = SimpleNamespace( enable_fused_qkv_index_projection=True, qkv_proj=lambda hidden_states: packed, - attn=object(), # Compatibility path, without the captured FP8 producer. + attn=Mock(spec_set=MiniMaxM3SparseRuntimeBackend), # Non-MSA compatibility path. register_to_config=True, num_heads=1, head_dim=3, @@ -494,7 +512,7 @@ def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, ou monkeypatch.setattr( modeling_minimaxm3, "_extract_minimax_m3_attention_extra_attrs", - lambda layer_idx: (SimpleNamespace(), layer), + lambda layer_idx: (metadata, layer), ) monkeypatch.setattr(modeling_minimaxm3, "is_torch_compiling", lambda: True) monkeypatch.setattr( @@ -509,7 +527,7 @@ def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, ou layer, position_ids=position_ids, hidden_states=hidden_states, - attn_metadata=SimpleNamespace(), + attn_metadata=metadata, ) assert captured["packed"] is packed @@ -523,9 +541,6 @@ def test_piecewise_captured_producer_preserves_symbolic_shapes( monkeypatch: pytest.MonkeyPatch, ) -> None: """Retain an unbacked symbolic token count in both fake query outputs.""" - from torch._subclasses.fake_tensor import FakeTensorMode - from torch.fx.experimental.symbolic_shapes import ShapeEnv - layer = SimpleNamespace(q_size=1024, index_q_size=128) monkeypatch.setattr( modeling_minimaxm3, @@ -553,8 +568,6 @@ def test_piecewise_captures_horizontal_producer_before_attention( monkeypatch: pytest.MonkeyPatch, ) -> None: """Run the captured producer before eager sparse attention consumes caches.""" - from unittest.mock import Mock - backend = object.__new__(MiniMaxM3MsaSparseAttention) backend.indexer_kv_dtype = "fp8" packed = torch.empty((4, 7), dtype=torch.bfloat16) @@ -564,8 +577,12 @@ def test_piecewise_captures_horizontal_producer_before_attention( kv_cache = torch.empty(8, dtype=torch.float8_e4m3fn) index_cache = torch.empty_like(kv_cache) slots = torch.tensor([0, 1, -1, -1], dtype=torch.int32) - metadata = SimpleNamespace( - num_tokens=2, msa_layer_cache_tensors={3: (kv_cache, index_cache)}, msa_out_cache_loc=slots + metadata = Mock( + spec_set=MiniMaxM3MsaSparseAttentionMetadata, + num_tokens=2, + kv_cache_manager=None, + msa_layer_cache_tensors={3: (kv_cache, index_cache)}, + msa_out_cache_loc=slots, ) layer = SimpleNamespace( enable_fused_qkv_index_projection=True, @@ -601,8 +618,6 @@ def test_piecewise_captured_producer_rejects_unavailable_fusion( monkeypatch: pytest.MonkeyPatch, ) -> None: """Reject unsupported fused geometry without an eager fallback in capture.""" - from unittest.mock import Mock - layer = SimpleNamespace( qkv_proj=Mock(side_effect=lambda hidden: hidden.clone()), _fused_fp8_qkv_indexer_norm_rope_kv_insert=Mock(return_value=None), @@ -632,14 +647,6 @@ def test_piecewise_captured_producer_declares_cache_mutations( restore_inplace: bool, ) -> None: """Real custom-op schema/AOT checks with CPU cache writes in place of CUDA math.""" - from torch._dynamo.backends.common import aot_autograd - from torch._functorch.aot_autograd import make_boxed_func - from torch._higher_order_ops.auto_functionalize import ( - auto_functionalized, - auto_functionalized_v2, - ) - - from tensorrt_llm._torch.compilation.remove_copy_pass import remove_copy_for_mutates_args def producer( packed: torch.Tensor, @@ -741,12 +748,11 @@ def optimize(gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]) -> ob @pytest.mark.cpu_only def test_piecewise_unfused_indexer_keeps_cache_write_eager(monkeypatch: pytest.MonkeyPatch) -> None: """Keep index-cache mutation outside capture when projections are separate.""" - from unittest.mock import Mock - backend = object.__new__(MiniMaxM3MsaSparseAttention) backend.indexer_kv_dtype = "fp8" q, k, v, idx_q, idx_k = [torch.empty((4, 128), dtype=torch.float8_e4m3fn) for _ in range(5)] - metadata = SimpleNamespace(num_tokens=2) + metadata = AttentionMetadata(max_num_requests=1, max_num_tokens=4) + metadata._num_tokens = 2 layer = SimpleNamespace( enable_fused_qkv_index_projection=False, register_to_config=True, @@ -799,8 +805,6 @@ def test_piecewise_attention_boundary_preserves_indexer_cache_contract( monkeypatch: pytest.MonkeyPatch, indexer_dtype: str, caches_prewritten: bool ) -> None: """Exercise the real MSA indexer after live-row slicing and cache insertion.""" - from unittest.mock import Mock - dtype = torch.float8_e4m3fn if indexer_dtype == "fp8" else torch.bfloat16 q, k, v, idx_q, idx_k = [torch.full((4, 128), value).to(dtype) for value in range(1, 6)] index_cache = torch.zeros((4, 1, 1, 128), dtype=dtype) @@ -809,8 +813,10 @@ def test_piecewise_attention_boundary_preserves_indexer_cache_contract( if caches_prewritten: index_cache.copy_(expected_cache) selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) - attn_metadata = SimpleNamespace( + attn_metadata = Mock( + spec_set=MiniMaxM3MsaSparseAttentionMetadata, num_tokens=2, + kv_cache_manager=None, msa_decode_span=None, msa_idx_k_cache=Mock(return_value=index_cache), msa_write_idx_k=Mock(), @@ -826,7 +832,7 @@ def write_caches( live_k: torch.Tensor, live_v: torch.Tensor, live_idx_k: torch.Tensor, - metadata: SimpleNamespace, + metadata: AttentionMetadata, ) -> None: """Replace only CUDA scatter math, retaining the live cache-write inputs.""" assert metadata is attn_metadata From 943145f4c0f0e66a57dffb2831b616d43f7be1fd Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:22:56 -0700 Subject: [PATCH 10/12] [None][test] Bound MiniMax-M3 Eagle3 PCG test batches Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8c13b255b272..ae326fbea8ea 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7412,8 +7412,14 @@ def _run_nvfp4_eagle3(self, else: max_seq_len = 4096 max_batch_size = 256 if attention_dp else 512 + cuda_graph_max_batch_size = 64 if (inferencex or attention_dp) else 128 piecewise_kwargs = {} if piecewise: + # With ADP's 64 query heads, at most 32 requests keep MSA's + # short-query plans (up to 32 tokens/request) within 65536. + # The 2K/4K many-request warmups then use the long-prefill path. + max_batch_size = 32 + cuda_graph_max_batch_size = max_batch_size # Cover the entire scheduler token budget so these evaluations # exercise captured prefill, rather than falling back above the # capture ceiling. Generation-only steps still take the eager @@ -7439,7 +7445,7 @@ def _run_nvfp4_eagle3(self, speculative_config=spec_config, cuda_graph_config=CudaGraphConfig( enable_padding=True, - max_batch_size=64 if (inferencex or attention_dp) else 128, + max_batch_size=cuda_graph_max_batch_size, ), disable_overlap_scheduler=not overlap_scheduler, enable_attention_dp=attention_dp, From acfd0a945e8517718cbba735708695fb975460b4 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:20:04 -0700 Subject: [PATCH 11/12] [None][fix] Preserve Eagle3 captures at compiled graph exit Port the multistream graph-exit ordering fix from feature-branch PR #18066 (95ea460996e0). Preserve unreturned in-place mutations before graph output and wait for their streams before external consumers read them. Add the nine scheduling regressions to CPU CI; retain the Eagle3 PCG configuration and acceptance thresholds. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../multi_stream/auto_multi_stream.py | 28 ++++ .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../compilation/test_auto_multi_stream.py | 125 ++++++++++++++++++ 3 files changed, 154 insertions(+) create mode 100644 tests/unittest/_torch/compilation/test_auto_multi_stream.py diff --git a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py index cfc3ef93bec0..42f9bf1447ac 100644 --- a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py +++ b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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 time from dataclasses import dataclass, field from operator import getitem @@ -194,6 +208,20 @@ def flatten_args(args): elif isinstance(arg, torch.fx.Node) and arg.op != "placeholder": in_edges[arg] = self.nodes[arg] + if node.op == "output": + # An in-place op may mutate a graph input without returning a + # value (Eagle3 captures hidden states into a preallocated + # buffer with inplace_slice_copy), so the FX output does not + # reach that side effect. Make graph exit depend on the last + # mutation of every touched tensor: the scheduled graph then + # emits the mutation before `output` (a node emitted after + # `output` is dead code once the module is recompiled), and + # with live auxiliary streams the exit waits on the mutating + # stream before a graph-external consumer reads the buffer. + for mutated_arg, mutator in latest_inplace_stat.items(): + if isinstance(mutated_arg, torch.fx.Node): + in_edges[mutated_arg] = mutator + # For node without in edge, connect it to the entry if len(in_edges) == 0: in_edges[None] = self.entry_node diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 324d171210b4..71b610e1da49 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -38,6 +38,7 @@ l0_cpu: # runs -m cpu_only, so this entry contributes only the 3 marked files under attention/. - unittest/_torch/attention - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py + - unittest/_torch/compilation/test_auto_multi_stream.py - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs - unittest/_torch/distributed - unittest/_torch/executor diff --git a/tests/unittest/_torch/compilation/test_auto_multi_stream.py b/tests/unittest/_torch/compilation/test_auto_multi_stream.py new file mode 100644 index 000000000000..4a2487a435aa --- /dev/null +++ b/tests/unittest/_torch/compilation/test_auto_multi_stream.py @@ -0,0 +1,125 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""Scheduling of in-place side effects that the FX output does not reach. + +Eagle3 captures decoder hidden states into a preallocated buffer with +``inplace_slice_copy``; the drafter reads that buffer outside the compiled +graph. The multi-stream scheduler must emit every such mutation before ``output`` +(a node emitted after ``output`` is dead code once the module is recompiled) +and make the exit wait on the mutating stream. +""" + +from collections.abc import Callable, Iterable +from typing import TypeVar + +import pytest +import torch +from torch.fx import Graph, GraphModule, Node + +from tensorrt_llm._torch.compilation.multi_stream.auto_multi_stream import ( + MultiStreamDAG, + multi_stream_schedule, +) + +pytestmark = pytest.mark.cpu_only + +COPY = torch.ops.trtllm.inplace_slice_copy.default +T = TypeVar("T") + + +def _capture(graph: Graph, dest: Node, src: Node, layer: int) -> Node: + return graph.call_function( + COPY, kwargs={"dest": dest, "src": src, "dim1_start": layer, "dim1_end": layer + 1} + ) + + +def _decoder_stack( + n_layers: int, capture_layers: tuple[int, ...] +) -> tuple[GraphModule, Node, list[Node]]: + """Build a chain of layers with selected outputs copied into ``dest``.""" + graph = Graph() + dest = graph.placeholder("dest") + x = graph.placeholder("x") + hidden = x + captures = [] + for layer in range(n_layers): + mm = graph.call_function(torch.ops.aten.mm.default, args=(hidden, hidden)) + add = graph.call_function(torch.ops.aten.add.Tensor, args=(mm, hidden)) + hidden = graph.call_function(torch.ops.aten.mul.Tensor, args=(add, 2.0)) + if layer in capture_layers: + captures.append(_capture(graph, dest, hidden, layer)) + out = graph.call_function(torch.ops.aten.neg.default, args=(hidden,)) + graph.output((out,)) + return GraphModule({}, graph), dest, captures + + +def _index_of(nodes: Iterable[T], predicate: Callable[[T], bool]) -> int: + return next(i for i, node in enumerate(nodes) if predicate(node)) + + +def test_graph_exit_depends_on_unreturned_inplace_side_effect() -> None: + """Keep unreturned mutations and their cross-stream waits before graph exit.""" + graph = Graph() + dest = graph.placeholder("dest") + src = graph.placeholder("src") + returned = graph.call_function(torch.ops.aten.neg.default, args=(src,)) + mutation = _capture(graph, dest, src, 0) + output = graph.output(returned) + graph_module = GraphModule({}, graph) + + dag = MultiStreamDAG(graph_module) + assert dag.nodes[output].in_edges[dest] is dag.nodes[mutation] + + dag.assign_streams(max_num_streams=2) + scheduled = dag.create_new_graph() + scheduled.lint() + nodes = list(scheduled.nodes) + output_index = _index_of(nodes, lambda n: n.op == "output") + mutation_index = _index_of(nodes, lambda n: n.target is COPY) + assert mutation_index < output_index + + if dag.nodes[mutation].stream is not dag.nodes[output].stream: + event = dag.nodes[mutation].event + assert event is not None + assert any( + node.target is torch.ops.trtllm.wait_event and node.args == (event,) + for node in nodes[mutation_index:output_index] + ) + + +@pytest.mark.parametrize("max_num_streams", [2, 3]) +@pytest.mark.parametrize( + "n_layers,capture_layers", + [(6, (1, 3, 5)), (6, (1, 3, 4)), (12, (1, 5, 11)), (12, (1, 5, 8))], +) +def test_every_capture_precedes_output( + n_layers: int, capture_layers: tuple[int, ...], max_num_streams: int +) -> None: + """Keep every hidden-state capture ahead of return in the emitted Python.""" + graph_module, _, captures = _decoder_stack(n_layers, capture_layers) + multi_stream_schedule(graph_module, max_num_streams) + graph_module.graph.lint() + nodes = list(graph_module.graph.nodes) + output_index = _index_of(nodes, lambda n: n.op == "output") + copy_indices = [i for i, node in enumerate(nodes) if node.target is COPY] + assert len(copy_indices) == len(captures) + assert all(i < output_index for i in copy_indices), (copy_indices, output_index) + + # Every capture must survive into the generated code ahead of `return`. + graph_module.recompile() + lines = [line.strip() for line in graph_module.code.splitlines() if line.strip()] + return_index = _index_of(lines, lambda line: line.startswith("return")) + copy_lines = [i for i, line in enumerate(lines) if "inplace_slice_copy" in line] + assert len(copy_lines) == len(captures) + assert all(i < return_index for i in copy_lines), (copy_lines, return_index) From 82d066bbef636c5f6002fc2c677fbb9844d27896 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:27:18 -0700 Subject: [PATCH 12/12] [None][test] Complete standalone PCG regression setup Register the MoE operator needed by the scheduler cost model in its isolated regression tests. Initialize the eager workspace engine fixture with the default prefill-only compilation policy. Keep runtime behavior and Eagle3 acceptance gates unchanged. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tests/unittest/_torch/compilation/test_auto_multi_stream.py | 2 ++ tests/unittest/_torch/executor/test_eager_workspace_engine.py | 1 + 2 files changed, 3 insertions(+) diff --git a/tests/unittest/_torch/compilation/test_auto_multi_stream.py b/tests/unittest/_torch/compilation/test_auto_multi_stream.py index 4a2487a435aa..f4b6c3e4bd1a 100644 --- a/tests/unittest/_torch/compilation/test_auto_multi_stream.py +++ b/tests/unittest/_torch/compilation/test_auto_multi_stream.py @@ -27,6 +27,8 @@ import torch from torch.fx import Graph, GraphModule, Node +# Register the MoE op referenced by the scheduler's cost model. +import tensorrt_llm._torch.moe.fused_moe.interface # noqa: F401 from tensorrt_llm._torch.compilation.multi_stream.auto_multi_stream import ( MultiStreamDAG, multi_stream_schedule, diff --git a/tests/unittest/_torch/executor/test_eager_workspace_engine.py b/tests/unittest/_torch/executor/test_eager_workspace_engine.py index deb2eae4fb84..62244db5fd7b 100644 --- a/tests/unittest/_torch/executor/test_eager_workspace_engine.py +++ b/tests/unittest/_torch/executor/test_eager_workspace_engine.py @@ -173,6 +173,7 @@ def setUp(self) -> None: self.engine.mapping = SimpleNamespace(cp_size=1) self.engine.sparse_attention_config = None self.engine._torch_compile_backend = None + self.engine._torch_compile_prefill_only = False self.engine.breakable_cuda_graph_runner = None self.engine._is_warmup = False self.metadata = object.__new__(TrtllmAttentionMetadata)