From 2eb491e18b1ce5b8a25e4c88f32dbb5d028d525c Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Sat, 10 Oct 2026 22:40:45 -0700 Subject: [PATCH] [Perf][MiniMax-M3] MSA prefill attends with the E4M3 query; bump fmha_sm100 to 968061a (#60684) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- cmake/external_projects/fmha_sm100.cmake | 11 +- ...st_minimax_m3_msa_cutlass_sparse_decode.py | 179 +++++++++++++++--- vllm/models/minimax_m3/nvidia/model.py | 2 +- .../minimax_m3/nvidia/sparse_attention_msa.py | 16 +- 4 files changed, 176 insertions(+), 32 deletions(-) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index f6bb0b8ace..6b84b35c90 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG be4a52aea420b6bef73be9af1a649859964790cb + GIT_TAG 968061ab16cb14991a5c7d4c0e8eabf531be0dcf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -36,11 +36,14 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100") install(FILES "${FMHA_SM100_PY_ROOT}/__init__.py" + "${FMHA_SM100_PY_ROOT}/_jit_cache.py" "${FMHA_SM100_PY_ROOT}/api.py" "${FMHA_SM100_PY_ROOT}/bench_utils.py" "${FMHA_SM100_PY_ROOT}/jit.py" + "${FMHA_SM100_PY_ROOT}/msa_warmup.py" "${FMHA_SM100_PY_ROOT}/nvfp4_kv.py" "${FMHA_SM100_PY_ROOT}/q8kv4_decode_adapter.py" + "${FMHA_SM100_PY_ROOT}/q8kv4_prefill_adapter.py" "${FMHA_SM100_PY_ROOT}/sparse.py" "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py" DESTINATION vllm/third_party/fmha_sm100 @@ -65,6 +68,12 @@ install(DIRECTORY "${FMHA_SM100_PY_ROOT}/kvouter/" PATTERN "__pycache__" EXCLUDE PATTERN "*.pyc" EXCLUDE) +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/prefill_q8kv4/" + DESTINATION vllm/third_party/fmha_sm100/prefill_q8kv4 + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE) + install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/" DESTINATION vllm/third_party/fmha_sm100/cute COMPONENT fmha_sm100 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 0007f20c5b..42e2b86965 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 @@ -782,6 +782,46 @@ def _make_nvfp4_kv_cache( return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16) +def _e4m3_staged_kv_reference( + kv_reference: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor +) -> torch.Tensor: + """The dequantized NVFP4 cache as an E4M3-query kernel reads it: each + code * block_scale rounded to E4M3, then the global scale.""" + sides = [] + for side, global_scale in zip( + kv_reference.float().split(HEAD_DIM, dim=-1), (k_scale, v_scale) + ): + sides.append( + (side / global_scale).to(torch.float8_e4m3fn).float() * global_scale + ) + return torch.cat(sides, dim=-1).to(torch.bfloat16) + + +def _assert_close_fp8_prefill(actual: torch.Tensor, expected: torch.Tensor) -> None: + """The fp8-q prefill rounds each split's probabilities to E4M3, and the splits + follow the SM count, so a rare element moves past the BF16 tolerance.""" + close = torch.isclose(actual, expected, atol=0.05, rtol=0.05) + assert close.float().mean().item() > 1 - 1e-4 + torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.05) + + +def _record_query_dtypes( + monkeypatch: pytest.MonkeyPatch, kernel_name: str +) -> list[torch.dtype]: + """Record the query dtype of every ``fmha_sm100.sparse.`` call.""" + from vllm.third_party.fmha_sm100 import sparse as msa_sparse + + kernel = getattr(msa_sparse, kernel_name) + query_dtypes: list[torch.dtype] = [] + + def record(q, *args, **kwargs): + query_dtypes.append(q.dtype) + return kernel(q, *args, **kwargs) + + monkeypatch.setattr(msa_sparse, kernel_name, record) + return query_dtypes + + def test_nvfp4_kv_cache_views_address_head_slots() -> None: """Slot 2h (K) and 2h + 1 (V) hold head h's data then scales, so each head is one contiguous run of the page: a TP rank's page is a byte slice.""" @@ -898,27 +938,14 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache( torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) -@pytest.mark.parametrize( - ("num_q_heads", "num_kv_heads"), - [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], -) -def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( - num_q_heads: int, +def _make_prefill_metadata( + q_lens: list[int], + seq_lens_list: list[int], + block_table: torch.Tensor, num_kv_heads: int, -) -> None: - """MSA prefill reads NVFP4 pages as their exact dequantized values, - including chunked prefill over cached context, and takes q from its fp8 - copy (the only one the fused insert writes on this path).""" - torch.manual_seed(0) - q_lens = [300, 128, 77] - seq_lens_list = [300, 640, 1100] - block_table, num_pages = _make_paged_layout(seq_lens_list) - k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda") - v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda") - kv_cache, kv_reference = _make_nvfp4_kv_cache( - num_pages, num_kv_heads, "nvfp4", k_scale, v_scale - ) - +) -> tuple[MiniMaxM3SparseMetadata, torch.Tensor]: + """Prefill-only metadata (chunks over cached context) and a top-k selecting + every causally visible page (<= TOPK per query), ending with the local one.""" total_q = sum(q_lens) seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") context_lens = seq_lens - torch.tensor(q_lens, dtype=torch.int32, device="cuda") @@ -948,7 +975,6 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( prefill=prefill, ) - # Every causally visible page (<= TOPK per query), ending with the local one. topk = torch.full( (total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device="cuda" ) @@ -959,6 +985,35 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( assert visible <= TOPK topk[token, :, :visible] = torch.arange(visible, device="cuda") token += 1 + return metadata, topk + + +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( + monkeypatch: pytest.MonkeyPatch, + num_q_heads: int, + num_kv_heads: int, +) -> None: + """MSA prefill reads NVFP4 pages as their dequantized values, staged to E4M3 + for the E4M3 q, including chunked prefill over cached context; it attends + with the fp8 q (the only one the fused insert writes on this path).""" + torch.manual_seed(0) + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] + block_table, num_pages = _make_paged_layout(seq_lens_list) + k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda") + v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda") + kv_cache, kv_reference = _make_nvfp4_kv_cache( + num_pages, num_kv_heads, "nvfp4", k_scale, v_scale + ) + + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + total_q = sum(q_lens) layer_name = "model.layers.0.self_attn.attn" q_scale = 0.5 @@ -976,6 +1031,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( / q_scale ).to(torch.float8_e4m3fn) query = query_fp8.to(torch.bfloat16) * q_scale + kv_reference = _e4m3_staged_kv_reference(kv_reference, k_scale, v_scale) common = dict( num_heads=num_q_heads, head_size=HEAD_DIM, @@ -984,6 +1040,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( topk_blocks=TOPK, sparse_block_size=BLOCK_SIZE, ) + query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_nvfp4_kv_func") forward_context = ForwardContext( no_compile_layers={}, attn_metadata={layer_name: metadata}, @@ -1001,7 +1058,8 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( query_fp8=query_fp8, ) current_platform.synchronize() - torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) + assert query_dtypes == [torch.float8_e4m3fn] + _assert_close_fp8_prefill(actual, expected) def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: @@ -1079,3 +1137,80 @@ def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: ) current_platform.synchronize() torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_reads_query_fp8_over_fp8_cache( + monkeypatch: pytest.MonkeyPatch, + num_q_heads: int, + num_kv_heads: int, +) -> None: + """With CUTLASS decode on, FP8 prefill (incl. chunks over cached context) + attends with the fp8 q the fused insert writes, not a BF16 copy, and + matches Triton on the dequantized q.""" + torch.manual_seed(0) + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] + block_table, num_pages = _make_paged_layout(seq_lens_list) + kv_cache = ( + torch.randn( + num_pages, + num_kv_heads, + BLOCK_SIZE, + 2 * HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.5 + ).to(torch.float8_e4m3fn) + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + layer_name = "model.layers.0.self_attn.attn" + q_scale = 0.5 + layer = SimpleNamespace( + layer_name=layer_name, + topk_indices_buffer=topk, + _q_scale_float=q_scale, + ) + query_fp8 = ( + torch.randn( + sum(q_lens), num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + / q_scale + ).to(torch.float8_e4m3fn) + query = query_fp8.to(torch.bfloat16) * q_scale + common = dict( + num_heads=num_q_heads, + head_size=HEAD_DIM, + scale=SM_SCALE, + num_kv_heads=num_kv_heads, + kv_cache_dtype="fp8", + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + ) + impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass") + assert impl.use_cutlass_decode + query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_func") + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={layer_name: metadata}, + slot_mapping={}, + ) + with override_forward_context(forward_context): + expected = MiniMaxM3SparseTritonImpl(**common).forward( + layer, query, kv_cache, torch.empty_like(query) + ) + actual = impl.forward( + layer, + torch.full_like(query, float("nan")), + kv_cache, + torch.zeros_like(query), + query_fp8=query_fp8, + ) + current_platform.synchronize() + assert query_dtypes == [torch.float8_e4m3fn] + _assert_close_fp8_prefill(actual, expected) diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 06c28a7a38..d67e9ea19d 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -614,7 +614,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] q = qkv.new_empty((num_tokens, self.q_size)) # With query_fp8, q is emitted only in fp8; the attend dequantizes it - # into ``q`` for the tokens that run on bf16-query kernels. + # into ``q`` for the tokens that run on Triton decode. query_fp8 = self._allocate_query_fp8(qkv) # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); # the fused kernel emits fp8 directly when this buffer is e4m3. diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index 5aa23b4ac8..ed8ca2a784 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -266,7 +266,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): q = query[:num_tokens].view(-1, self.num_heads, hd) out = output[:num_tokens].view(-1, self.num_heads, hd) # Given query_fp8, it is the only q the fused insert wrote: CUTLASS - # decode reads it directly, and the bf16-query kernels read it + # decode and prefill read it directly, and Triton decode reads it # dequantized into ``query``. q_fp8 = ( None @@ -339,11 +339,11 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): # [H, prefill, MK] transposed view; build_k2q_csr consumes the # strided view directly (topK stays innermost-contiguous). prefill_topk = topk[nd:num_tokens].transpose(0, 1) - qp = ( - q[nd:] - if q_fp8 is None - else _dequantize_query(q_fp8[nd:], q_scale_float, q[nd:]) - ) + # An E4M3 query attends as is, its scale folded into softmax_scale. + if q_fp8 is None: + qp, softmax_scale = q[nd:], self.scale + else: + qp, softmax_scale = q_fp8[nd:], self.scale * q_scale_float k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr( prefill_topk, p.cu_seqlens_q, @@ -375,7 +375,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): max_seqlen_k=p.max_seq_len, blk_kv=SPARSE_BLOCK_SIZE, causal=True, - softmax_scale=self.scale, + softmax_scale=softmax_scale, page_table=p.block_table, seqused_k=p.seq_lens, schedule=schedule, @@ -393,7 +393,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): topK=self.topk_blocks, blk_kv=SPARSE_BLOCK_SIZE, causal=True, - softmax_scale=self.scale, + softmax_scale=softmax_scale, cu_seqlens_q=p.cu_seqlens_q, cu_seqlens_k=p.cu_seqlens_k, max_seqlen_q=p.max_query_len,