[None][perf] Extend MiniMax-M3 piecewise CUDA graphs coverage

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<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor> minimaxM3Fp8QKVIndexerNormRopeKVInsert(
     return {qOut, indexQOut};
 }
 
-std::tuple<torch::Tensor, torch::Tensor> 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

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)
@@ -970,9 +981,11 @@ def _build_msa_fields(self) -> None:
         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
@@ -1022,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)
 

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

tensorrt_llm/_torch/compilation/remove_copy_pass.py

@@ -69,8 +69,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 +104,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:

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"
         },

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -397,6 +397,54 @@ 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:
+        """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
+        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:
+        """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),
+                             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]:
+        """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
+        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)

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,39 @@ 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:
+            """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.
+            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:
+            """Infer the GEMM output shape and dtype without invoking FlashInfer."""
+            return act.new_empty((act.size(0), weight.size(0)),
+                                 dtype=output_dtype)