iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][MiniMax-M3] MSA prefill attends with the E4M3 query; bump fmha_sm100 to 968061a (#60684)

Browse Source 
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yongye Zhu
2026-10-11 13:40:45 +08:00
committed by GitHub
co-authored by Claude Opus 5.5
parent f8471164bf
commit 2eb491e18b
4 changed files with 176 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/external_projects/fmha_sm100.cmake
+10 -1
View File
@@ -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
tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
+157 -22
View File
@@ -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.<kernel_name>`` 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)
vllm/models/minimax_m3/nvidia/model.py
+1 -1
View File
@@ -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.
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+8 -8
View File
@@ -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,