diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index 42e2b86965..5df08b0e40 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/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( diff --git a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py index c52dee413d..970b31c702 100644 --- a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py +++ b/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 ) diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index ed8ca2a784..a5598c719b 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/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