Repository navigation
[Attention][DSA] Enable W4A16 DSA - #51724
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Documentation preview: https://vllm--51724.org.readthedocs.build/en/51724/ |
mergify
Bot
added
documentation
|
This pull request has merge conflicts that must be resolved before it can be |
|
Marking as |
d3703dd to
c249781
Compare
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #86016. |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #86016. |
|
/ci run |
|
✅ Triggered Buildkite CI #86110 for commit |
|
/ci retry |
|
✅ Queued 4 failed job(s) for retry in Buildkite CI #86110. |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #86110. |
LucasWilkinson
left a comment
LucasWilkinson
left a comment
There was a problem hiding this comment.
LGTM other then a couple nits
|
For the 2 models and 5 tasks, NVFP4 NoPE and FP8 RoPE is almost always better than NVFP4 NoPE and RoPE (only exception is DSV3.2 on Tau2). The accuracy difference between FP8 and NVFP4 NoPE and FP8 RoPE kv is < 1% except DSV3.2 on Tau2.
Note that nvfp4_nvfp4RoPE_ds_mla is not added in this PR. It is only used during experimentation. AIME25: 64 repeats |
|
/ci run |
|
✅ Triggered Buildkite CI #86391 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #86408 for commit |
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
vllm-project#52861 routed the DSA models to a fused norm+rope Triton kernel that writes the MLA KV cache itself and only supports fp8_ds_mla. vllm-project#51724 added nvfp4_ds_mla after that and the rebase missed it. Teach the fused kernel the nvfp4_ds_mla layout. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
The A.X-K2 line lets a dense EAGLE3/DFlash/DSpark draft ride a sparse-MLA target by defaulting the draft's KV cache dtype to "auto" whenever the target cache is fp8_ds_mla, since the draft's dense attention backend cannot hold the opaque DS-MLA record. Upstream vllm-project#51724/vllm-project#55538 added a second such record, nvfp4_ds_mla (SM100 only), which the planned B200 deployment will use; match on the "_ds_mla" suffix, as upstream's attention layer does, so the override covers both. Validation: ruff check/format and py_compile on vllm/config/vllm.py passed on macOS; no GPU runtime test was run on this host. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Seungmin <seungminheo@sk.com> (cherry picked from commit 6961a1a)
The A.X-K2 line lets a dense EAGLE3/DFlash/DSpark draft ride a sparse-MLA target by defaulting the draft's KV cache dtype to "auto" whenever the target cache is fp8_ds_mla, since the draft's dense attention backend cannot hold the opaque DS-MLA record. Upstream vllm-project#51724/vllm-project#55538 added a second such record, nvfp4_ds_mla (SM100 only), which the planned B200 deployment will use; match on the "_ds_mla" suffix, as upstream's attention layer does, so the override covers both. Validation: ruff check/format and py_compile on vllm/config/vllm.py passed on macOS; no GPU runtime test was run on this host. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Seungmin <seungminheo@sk.com> (cherry picked from commit 6961a1a)
5 participants
Corresponding FlashMLA PR: vllm-project/FlashMLA#18
concat_and_cache_nvfp4_ds_mla — quantize and store
cp_gather_and_upconvert_nvfp4_kv_cache — dequantize for prefill
Read path, mirroring cp_gather_and_upconvert_fp8_kv_cache.
Gathers a batch's scattered cache pages and upconverts them to a contiguous bf16 [total_tokens, 576] workspace for the bf16 prefill kernel.
Only runs at ≥32 query heads per rank, so at TP8 neither DSv3.2 nor GLM-5.2 reaches it — covered by unit tests, not their E2E runs.
Purpose
Enable W4A16 DSA
This PR is now in draft for testing purpose, the pin of FlashMLA needs to be changed before merging
Test Plan
Unittest
E2E with GLM 5.2 and DSV3.2
Test Result
passed.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.