[Perf][Attention][MiniMax-M3] Select CUTLASS MSA by EAGLE3 decode shape (#55838)

Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py

@@ -65,30 +65,40 @@ SM_SCALE = HEAD_DIM**-0.5
         "decode_query_len",
         "num_q_heads",
         "num_kv_heads",
+        "is_sm103",
         "expected",
     ),
     [
-        pytest.param(8, 4, 64, 4, False, id="tp1-below-min-batch"),
-        pytest.param(16, 4, 64, 4, True, id="tp1-supported"),
-        pytest.param(16, 4, 16, 1, True, id="tp4-min-batch"),
-        pytest.param(24, 4, 16, 1, True, id="tp4-intermediate-batch"),
-        pytest.param(32, 4, 16, 1, True, id="tp4-supported"),
-        pytest.param(16, 1, 64, 4, True, id="tp1-query-len-1"),
-        pytest.param(16, 1, 16, 1, True, id="tp4-query-len-1"),
-        pytest.param(16, 2, 64, 4, True, id="tp1-query-len-2"),
-        pytest.param(16, 2, 16, 1, True, id="tp4-query-len-2"),
-        pytest.param(16, 32, 64, 4, True, id="query-len-upper-bound"),
-        pytest.param(16, 0, 64, 4, False, id="query-len-zero"),
-        pytest.param(16, 33, 64, 4, False, id="query-len-above-bound"),
+        pytest.param(15, 1, 64, 4, True, False, id="qlen1-below-row-floor"),
+        pytest.param(16, 1, 64, 4, True, True, id="qlen1-at-row-floor"),
+        pytest.param(10, 3, 32, 2, True, False, id="qlen3-below-row-floor"),
+        pytest.param(11, 3, 32, 2, True, True, id="qlen3-at-row-floor"),
+        pytest.param(3, 4, 64, 4, True, False, id="qlen4-below-row-floor"),
+        pytest.param(4, 4, 64, 4, True, True, id="qlen4-at-row-floor"),
+        pytest.param(7, 8, 16, 1, True, False, id="qlen8-below-row-floor"),
+        pytest.param(8, 8, 16, 1, True, True, id="qlen8-at-row-floor"),
+        pytest.param(15, 9, 64, 4, True, False, id="unmeasured-query-length"),
+        pytest.param(15, 4, 8, 1, True, False, id="unmeasured-heads"),
+        pytest.param(15, 4, 64, 4, False, False, id="sm100-keeps-old-floor"),
+        pytest.param(16, 32, 64, 4, True, True, id="query-len-upper-bound"),
+        pytest.param(16, 0, 64, 4, True, False, id="query-len-zero"),
+        pytest.param(16, 33, 64, 4, True, False, id="query-len-above-bound"),
     ],
 )
 def test_msa_cutlass_decode_static_dispatch(
     batch_size: int,
     decode_query_len: int,
-    expected: bool,
     num_q_heads: int,
     num_kv_heads: int,
+    is_sm103: bool,
+    expected: bool,
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability",
+        lambda capability: is_sm103 and capability == (10, 3),
+    )
     assert (
         should_prepare_decode_metadata(
             batch_size,
@@ -291,6 +301,7 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(
     builder.kv_cache_dtype = "fp8_e4m3"
     builder.decode_backend = "cutlass"
     builder.msa_cutlass_plan_cache = object()
+    builder.msa_cutlass_min_batch_size = None
 
     metadata = builder.build(
         0,
@@ -486,6 +497,10 @@ def _make_topk(
         pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"),
         pytest.param(16, 1, 8, 1, True, id="tp4-query-len-1"),
         pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"),
+        # Measured B300 EAGLE3 dispatch boundaries.
+        pytest.param(64, 4, 2, 4, True, id="tp1-batch-4-query-len-4"),
+        pytest.param(32, 2, 4, 4, True, id="tp2-batch-8-query-len-4"),
+        pytest.param(16, 1, 8, 4, True, id="tp4-batch-16-query-len-4"),
     ],
 )
 def test_msa_cutlass_decode_matches_triton_with_interleaved_cache(

vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py

@@ -23,9 +23,10 @@ _TOPK = 16
 # fixed planner allocation used by the MSA decode kernel.
 _MAX_QUERY_HEAD_ROWS = 65536
 _MAX_DECODE_QUERY_LEN = 32
-# Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4.
-# NVFP4 has no Triton fallback, so it takes CUTLASS at every batch size.
-_MIN_CUTLASS_BATCH_SIZE = 16
+_DEFAULT_MIN_CUTLASS_BATCH_SIZE = 16
+_B300_EAGLE3_QUERY_HEAD_ROWS = 1024
+_B300_EAGLE3_MAX_QUERY_LEN = 8
+_B300_EAGLE3_HEAD_GEOMETRIES = {(64, 4), (32, 2), (16, 1)}
 
 
 def is_nvfp4_kv_cache(kv_cache_dtype: str) -> bool:
@@ -235,6 +236,22 @@ def _supported_head_geometry(num_q_heads: int, num_kv_heads: int) -> bool:
     )
 
 
+def _min_cutlass_batch_size(
+    decode_query_len: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+) -> int:
+    if (
+        not current_platform.is_device_capability((10, 3))
+        or not 1 <= decode_query_len <= _B300_EAGLE3_MAX_QUERY_LEN
+        or (num_q_heads, num_kv_heads) not in _B300_EAGLE3_HEAD_GEOMETRIES
+    ):
+        return _DEFAULT_MIN_CUTLASS_BATCH_SIZE
+    # Use CUTLASS once batch_size * decode_query_len * num_q_heads >= 1024.
+    query_head_rows = decode_query_len * num_q_heads
+    return (_B300_EAGLE3_QUERY_HEAD_ROWS + query_head_rows - 1) // query_head_rows
+
+
 def supports_cutlass_sparse_decode(
     *,
     decode_backend: MiniMaxM3MSADecodeBackend,
@@ -267,9 +284,14 @@ def should_prepare_decode_metadata(
     kv_cache_dtype: str,
     page_size: int,
     topk_blocks: int,
+    min_batch_size: int | None = None,
 ) -> bool:
     """Return whether a graph shape can use the CUTLASS decode path."""
     total_q = batch_size * decode_query_len
+    if min_batch_size is None:
+        min_batch_size = _min_cutlass_batch_size(
+            decode_query_len, num_q_heads, num_kv_heads
+        )
     return (
         supports_cutlass_sparse_decode(
             decode_backend=decode_backend,
@@ -280,7 +302,7 @@ def should_prepare_decode_metadata(
             topk_blocks=topk_blocks,
         )
         and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN
-        and (batch_size >= _MIN_CUTLASS_BATCH_SIZE or is_nvfp4_kv_cache(kv_cache_dtype))
+        and (batch_size >= min_batch_size or is_nvfp4_kv_cache(kv_cache_dtype))
         and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
     )
 

vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

@@ -8,6 +8,7 @@ path for regular decode and speculative verification. NVFP4 KV caches read the
 vLLM packed pages directly in both prefill and CUTLASS decode.
 """
 
+import os
 from dataclasses import dataclass, replace
 from typing import ClassVar
 
@@ -49,6 +50,11 @@ from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
 logger = init_logger(__name__)
 
 
+def _get_msa_cutlass_min_batch_size() -> int | None:
+    value = os.getenv("VLLM_MINIMAX_M3_MSA_CUTLASS_MIN_BATCH")
+    return int(value) if value is not None else None
+
+
 def _dequantize_query(
     query_fp8: torch.Tensor, q_scale: float, out: torch.Tensor
 ) -> torch.Tensor:
@@ -133,6 +139,7 @@ class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
         self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
         self.decode_backend = vllm_config.attention_config.minimax_m3_msa_decode_backend
         self.msa_cutlass_plan_cache = MSACutlassDecodePlanCache()
+        self.msa_cutlass_min_batch_size = _get_msa_cutlass_min_batch_size()
 
     def build(
         self,
@@ -159,6 +166,7 @@ class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
             kv_cache_dtype=self.kv_cache_dtype,
             page_size=SPARSE_BLOCK_SIZE,
             topk_blocks=self.topk_blocks,
+            min_batch_size=self.msa_cutlass_min_batch_size,
         ):
             seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
             assert seq_lens_cpu is not None