[None][perf] Extend packed MiniMax-M3 producer to mixed and decode batches

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -1316,21 +1316,20 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert(
         position_ids: Optional[torch.Tensor],
         attn_metadata: AttentionMetadata,
     ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
-        """Run the vLLM-style horizontal sparse producer for pure prefill."""
+        """Run the vLLM-style horizontal producer for every sparse batch.
+
+        The CUDA kernel is token-major and batch-type agnostic: per-token
+        positions and cache slots cover pure prefill, mixed aggregate batches,
+        and CUDA-graph decode.
+        """
         if (
             not self.enable_fused_qkv_index_projection
             or not isinstance(self.attn, MiniMaxM3MsaSparseAttention)
             or not self._emit_fp8_main_qkv()
             or self.attn.indexer_kv_dtype != "fp8"
         ):
             return None
-        if (
-            is_torch_compiling()
-            or getattr(attn_metadata, "is_cuda_graph", False)
-            or int(getattr(attn_metadata, "num_generations", 0)) != 0
-            or int(getattr(attn_metadata, "num_contexts", 0)) == 0
-            or int(packed.shape[0]) != int(attn_metadata.num_tokens)
-        ):
+        if is_torch_compiling() or int(packed.shape[0]) != int(attn_metadata.num_tokens):
             return None
         if (
             packed.dtype != torch.bfloat16

tensorrt_llm/llmapi/llm_args.py

@@ -738,10 +738,10 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig):
         default=False,
         description=
         "Fuse Q/K/V and index-Q/index-K into one quantized projection. Index-Q "
-        "is sharded with the KV heads and index-K is replicated. This prototype "
-        "currently targets disaggregated prefill workers; leave it disabled on "
-        "decode or mixed workers. The MiniMax-M3-specific path requires the MSA "
-        "implementation.",
+        "is sharded with the KV heads and index-K is replicated. MSA batches "
+        "also use a horizontal norm/RoPE/cache-insertion producer for prefill, "
+        "mixed, and CUDA-graph decode execution. The MiniMax-M3-specific path "
+        "requires the MSA implementation.",
         status="prototype",
     )
     num_attention_heads: Optional[int] = Field(

tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py

@@ -63,6 +63,9 @@ def test_minimax_m3_horizontal_producer_matches_separate_producers(num_tokens):
     slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 37) % (
         (num_pages - 1) * 128
     )
+    # Keep the parity reference slots valid: the legacy separate main-K/V
+    # producer does not support negative slots. Negative-slot handling is
+    # exercised below using horizontal eager execution versus graph replay.
     rope_cache = _rope_cache(max(256, num_tokens))
 
     main_width = (num_heads_q + 2 * num_kv_heads) * 128
@@ -121,8 +124,9 @@ def test_minimax_m3_horizontal_producer_matches_separate_producers(num_tokens):
         position_ids,
     )
 
-    pages = slots.long() // 128
-    within = slots.long() % 128
+    valid = slots >= 0
+    pages = slots[valid].long() // 128
+    within = slots[valid].long() % 128
     assert torch.equal(q.view(torch.uint8), q_reference.view(torch.uint8))
     # The horizontal producer follows vLLM's CUDA contract and converts its
     # normalized/RoPE FP32 registers directly to E4M3. The existing separate
@@ -145,3 +149,96 @@ def test_minimax_m3_horizontal_producer_matches_separate_producers(num_tokens):
         rtol=0.13,
         atol=0.05,
     )
+
+    # Aggregate decode captures this producer in a CUDA graph. The operator
+    # allocates compact Q/index-Q outputs while writing graph-stable paged
+    # caches through a graph-stable slot mapping, so exercise both capture and
+    # replay for decode-sized (1) and larger mixed/prefill token counts.
+    graph_packed = packed.clone()
+    graph_positions = position_ids.clone()
+    graph_slots = slots.clone()
+    graph_main_cache = _main_cache(num_pages, num_kv_heads)
+    graph_index_cache = _index_cache(num_pages)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        graph_q, graph_index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert(
+            graph_packed,
+            graph_main_cache,
+            graph_index_cache,
+            graph_slots,
+            num_heads_q,
+            num_kv_heads,
+            num_index_heads,
+            128,
+            64,
+            1e-5,
+            q_weight,
+            k_weight,
+            index_q_weight,
+            index_k_weight,
+            rope_cache,
+            graph_positions,
+        )
+
+    # Replay with different projection values, nonuniform positions, and new
+    # cache destinations. This proves replay reads the refreshed graph buffers
+    # rather than retaining capture-time values or slots.
+    replay_packed = torch.randn_like(packed)
+    replay_positions = (
+        torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 7 + 3
+    ) % rope_cache.shape[0]
+    replay_slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 53 + 11) % (
+        (num_pages - 1) * 128
+    )
+    if num_tokens > 1:
+        replay_slots[-1] = -1
+    graph_packed.copy_(replay_packed)
+    graph_positions.copy_(replay_positions)
+    graph_slots.copy_(replay_slots)
+    graph_main_cache.zero_()
+    graph_index_cache.zero_()
+
+    replay_main_cache = _main_cache(num_pages, num_kv_heads)
+    replay_index_cache = _index_cache(num_pages)
+    replay_q, replay_index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert(
+        replay_packed,
+        replay_main_cache,
+        replay_index_cache,
+        replay_slots,
+        num_heads_q,
+        num_kv_heads,
+        num_index_heads,
+        128,
+        64,
+        1e-5,
+        q_weight,
+        k_weight,
+        index_q_weight,
+        index_k_weight,
+        rope_cache,
+        replay_positions,
+    )
+    graph.replay()
+    torch.cuda.synchronize()
+
+    replay_valid = replay_slots >= 0
+    replay_pages = replay_slots[replay_valid].long() // 128
+    replay_within = replay_slots[replay_valid].long() % 128
+    assert torch.equal(graph_q.view(torch.uint8), replay_q.view(torch.uint8))
+    torch.testing.assert_close(
+        graph_index_q.float(),
+        replay_index_q.float(),
+        rtol=0.0,
+        atol=0.0,
+    )
+    assert torch.equal(
+        graph_main_cache[replay_pages, :, :, replay_within, :].view(torch.uint8),
+        replay_main_cache[replay_pages, :, :, replay_within, :].view(torch.uint8),
+    )
+    torch.testing.assert_close(
+        graph_index_cache[replay_pages, :, replay_within, :].float(),
+        replay_index_cache[replay_pages, :, replay_within, :].float(),
+        rtol=0.0,
+        atol=0.0,
+    )