Repository navigation
Conversation
…e attention Under DCP, compact this rank's top-k slots into a prefix and pass their count to FlashMLA as topk_length, with one batch entry per token, instead of running the kernel on the full row with the other ranks' slots masked. Assisted-by: Claude Signed-off-by: LoongPei <3136347099@qq.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
The masked call came from #46514, written when the SM90 V3.2 kernel did not take Kernel. vllm-project/FlashMLA#27 and #28 cherry-pick cleanly onto vLLM's current pin
End to end. Nightly
One run per point. The prefill gain shows up the same way in all four prefill-bound measurements (both prefill rows, the needle and the cold prompts); the decode differences at 8 and 32 are single samples. That is in line with your H20 numbers, slightly smaller on prefill. Chunk size. In the per-token layout each call allocates split accumulators of Separately, and not caused by this PR: on the current pin the masked mixed-batch call allocates the same accumulators again (2 GiB at 8192 tokens, 8 GiB at 32768) because the FlashMLA V4.1 sync dropped the AI assistance was used for this work. |
|
For anyone who wants to try this on Hopper before it lands: |
Split the per-token FlashMLA call so that its fp32 split-KV accumulators stay within 256 MiB (2048 tokens at 64 heads), instead of planning up to 8192 tokens per call, which needed about 1 GiB. Assisted-by: Claude Signed-off-by: LoongPei <3136347099@qq.com>
|
Thanks a lot for testing this on H200 NVL with GLM-5.3, and for the kernel check. Filling the tail of the compacted rows with other valid slots is a stronger test than mine. Agreed on the chunk size. fb94721 bounds the split-KV accumulators by bytes instead: each call now takes as many tokens as fit in 256 MiB of fp32 accumulators, which is 2048 tokens at 64 heads and 1024 at 128 heads. That is about 270 MiB per call including the On 8x H20 (TP8 + DCP4, V3.2) the cost is small. In a 16K-token mixed batch the attention call is 1.4-1.8% slower than with 8192-token chunks (3.5-4.6% with 1024). Fresh prefill goes from 65.17 / 79.10 s to 65.44 / 79.30 s, which is -28.0% / -24.2% against stock instead of -28.3% / -24.4%; I updated the table in the description. Decode-only steps fit in one call either way. Unit tests and pre-commit pass. Thanks also for #59305 and vllm-project/FlashMLA#29. |
This branch has not been deployed
2 participants
Purpose
Fixes #58980. Under DCP, each rank owns about 2048/dcp slots of the indexer's 2048-wide top-k row.
_forward_fp8_kv_mixed_batchkept the other ranks' slots as-1and ran FlashMLA on the full width, so ~(dcp-1)/dcp of every rank's sparse attention went to masked slots, for decode rows and, under DCP, prefill rows alike.This PR:
compact_valid_to_front, already used by the trtllm-gen path) and passes their count astopk_length.topk_lengthis per batch entry (the layout the DeepSeek-V4 decode path uses). Each call gets a freshFlashMLASchedMetabecause the lengths, and so the plan, differ per layer. Tokens are split across calls so that FlashMLA's fp32 split-KV accumulators stay within 256 MiB per call (2048 tokens at 64 heads).topk_length == 0instead of an all--1scan._fp8_flash_mla_kerneloverride is no longer on the DCP path.topk_lengthsupport ([SM90] Sparse decode: honor topk_length for DeepSeek-V3.2 FlashMLA#27) and the faster planner (Faster decoding schedule planner with identical plans FlashMLA#28), and adds the two new kernel files tocmake/external_projects/flashmla.cmake.The attended token set is exactly the same as before; only the kernel's accumulation order changes.
This does not duplicate an open PR: the masked call comes from #46514, written when the SM90 V3.2 kernel did not accept
topk_length, and no open PR changes it. This stays a draft until vllm-project/FlashMLA#27 and vllm-project/FlashMLA#28 merge; theGIT_TAGbump will then point at the merge commit.Test Plan
pytest tests/v1/attention/test_sparse_mla_backends.py tests/v1/attention/test_indexer_dcp_localize.py -k "flashmla or FlashMLA or dcp or fp8 or hyv4"test_fp8_mixed_batch_dcp_neutralizes_empty_rowsis updated to the new filter return value.test_fp8_dcp_topk_length_matches_masked_rows(GPU) checks the new path against the masked call for DCP 2 and 4 with KV interleave 1 and 16. It covers rows with fewer than top-k candidates, rows the rank owns nothing of, and batches split across several plans.test_hyv4_fp8_per_token_kernel_passes_sinkchecks that HY V4's sink reaches the new call.tests/evals/gsm8k/gsm8k_eval.py(1319 questions, 5-shot), without and with MTP (num_speculative_tokens=3).Setup: 8x H20, vLLM main (a9eafde), DeepSeek-V3.2,
--tensor-parallel-size 8 --decode-context-parallel-size 4 --enable-expert-parallel --kv-cache-dtype fp8_ds_mla --max-num-batched-tokens 16384 --max-num-seqs 64, CUDA graphs on.Test Result
Unit: 219 passed, 0 failed on H20 (FlashMLA with vllm-project/FlashMLA#27).
pre-commit run --fileson the four changed files: all hooks pass.Decode step time (ms):
"after" uses vllm-project/FlashMLA#27 only; the last column adds vllm-project/FlashMLA#28. All columns come from one session. With the first alone, planning every layer (FlashMLA's planner takes ~16 us per call, and V3.2 has 61 layers) takes back what the kernel saves at 8 concurrent requests; the second cuts the planner to ~4 us.
Prefill wall time of fresh prompts (s, mean of two runs):
"after" is with the 256 MiB accumulator bound. Planning up to 8192 tokens per call measured 65.17 and 79.10 s; in a 16K-token batch the bound makes the attention call 1.4-1.8% slower on H20.
Accuracy:
The mean MTP acceptance length is 3.11 before and 3.12 after. 38 of the 40 passkey completions are byte-identical before and after. The one miss answers "The pass key is" with "the piece of information you were supposed to find and memorize." instead of the digits. The per-document NLL differs by at most 0.0017 nats per token, also beyond the first 4K tokens.
GSM8K without MTP, passkey and NLL ran on v0.30.0 with the same backend code; everything else ran on main. The accuracy runs used a FlashMLA build with both FlashMLA PRs; the planner PR's plans are identical, so the outputs are the same as with vllm-project/FlashMLA#27 alone.
AI assistance (Claude) was used for parts of the implementation, the tests and the benchmark scripts. I reviewed every changed line and ran the tests and evaluations above.