Conversation
mgoin
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
heheda12345,
ivanium,
njhill,
orozery,
robertgshaw2-redhat and
ywang96
as code owners
July 15, 2026 19:10
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mgoin
force-pushed
the
fix-kv-alloc-for-sparse-MLA-with-SWA-draft
branch
from
July 22, 2026 20:17
d5bc7bf to
61818c9
Compare
mgoin
changed the title
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: mgoin <mgoin64@gmail.com>
mgoin
force-pushed
the
fix-kv-alloc-for-sparse-MLA-with-SWA-draft
branch
from
July 22, 2026 20:24
61818c9 to
b005aff
Compare
ivanium
approved these changes
Jul 23, 2026
ivanium
enabled auto-merge (squash)
July 23, 2026 18:29
4 tasks
njhill
added a commit
that referenced
this pull request
Aug 6, 2026
The KV-layout standardization removed indexes_kv_by_block_stride and with it the gate on page-size padding, so unify_kv_cache_spec_page_size now pads any attention layer with a non-divisible page. That silently re-broke the sparse-MLA-target + SWA-draft case from #48776: instead of raising and falling back to a single full-allocation group, the sparse indexer's page is padded and the layers split into hybrid groups. MLA kernels address custom packed page interiors (the 656-byte fp8_ds_mla layout, the sparse indexer k-cache) rather than the view's block stride, so a padded MLA page would be misread; #48776's fallback exists precisely to avoid that. Refuse padding for MLAAttentionSpec so the NotImplementedError fires again and the full-allocation fallback is taken. test_mla_with_incompatible_swa_uses_one_full_allocation_group passes again; the full test_kv_cache_utils.py file passes on this commit (71/71). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ThBqMtrBGwXZnMYHky55on Signed-off-by: Nick Hill <nickhill123@gmail.com>
LucasWilkinson
pushed a commit
that referenced
this pull request
Aug 6, 2026
The KV-layout standardization removed indexes_kv_by_block_stride and with it the gate on page-size padding, so unify_kv_cache_spec_page_size now pads any attention layer with a non-divisible page. That silently re-broke the sparse-MLA-target + SWA-draft case from #48776: instead of raising and falling back to a single full-allocation group, the sparse indexer's page is padded and the layers split into hybrid groups. MLA kernels address custom packed page interiors (the 656-byte fp8_ds_mla layout, the sparse indexer k-cache) rather than the view's block stride, so a padded MLA page would be misread; #48776's fallback exists precisely to avoid that. Refuse padding for MLAAttentionSpec so the NotImplementedError fires again and the full-allocation fallback is taken. test_mla_with_incompatible_swa_uses_one_full_allocation_group passes again; the full test_kv_cache_utils.py file passes on this commit (71/71). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ThBqMtrBGwXZnMYHky55on Signed-off-by: Nick Hill <nickhill123@gmail.com>
LucasWilkinson
pushed a commit
that referenced
this pull request
Aug 6, 2026
The KV-layout standardization removed indexes_kv_by_block_stride and with it the gate on page-size padding, so unify_kv_cache_spec_page_size now pads any attention layer with a non-divisible page. That silently re-broke the sparse-MLA-target + SWA-draft case from #48776: instead of raising and falling back to a single full-allocation group, the sparse indexer's page is padded and the layers split into hybrid groups. MLA kernels address custom packed page interiors (the 656-byte fp8_ds_mla layout, the sparse indexer k-cache) rather than the view's block stride, so a padded MLA page would be misread; #48776's fallback exists precisely to avoid that. Refuse padding for MLAAttentionSpec so the NotImplementedError fires again and the full-allocation fallback is taken. test_mla_with_incompatible_swa_uses_one_full_allocation_group passes again; the full test_kv_cache_utils.py file passes on this commit (71/71). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ThBqMtrBGwXZnMYHky55on Signed-off-by: Nick Hill <nickhill123@gmail.com>
realliujiaxu
pushed a commit
to vllm-project/vllm-ascend
that referenced
this pull request
Aug 17, 2026
### What this PR does / why we need it? The following PRs are not cherry-picked to vllm v0.26.0 branch, so patches should be added on vllm-ascend releases/v0.26.0rc branch, if we want to use some kinds of DSpark (e.g. https://huggingface.co/RedHatAI/GLM-5.2-speculator.dspark). - vllm-project/vllm#48524: solving fc shape error when num_target_layers != num_hidden_layers - vllm-project/vllm#48776: solving KV Cache planning when target model is sparse-MLA and draft model is SWA - vllm-project/vllm#48639: solving `sample_from_anchor` wrongly loaded ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - GLM-5.2 W4A8 + DSpark acceptance rate ``` INFO 08-14 15:40:22 [metrics.py:120] SpecDecoding metrics: Mean acceptance length: 2.82, Accepted throughput: 73.39 tokens/s, Drafted throughput: 282.88 tokens/s, Accepted: 790 tokens, Drafted: 3045 tokens, Per-position acceptance rate: 0.729, 0.457, 0.278, 0.161, 0.083, 0.064, 0.044, Avg Draft acceptance rate: 25.9% INFO 08-14 15:40:32 [metrics.py:120] SpecDecoding metrics: Mean acceptance length: 4.38, Accepted throughput: 146.65 tokens/s, Drafted throughput: 303.29 tokens/s, Accepted: 1469 tokens, Drafted: 3038 tokens, Per-position acceptance rate: 0.836, 0.652, 0.537, 0.408, 0.353, 0.318, 0.281, Avg Draft acceptance rate: 48.4% INFO 08-14 15:40:42 [metrics.py:120] SpecDecoding metrics: Mean acceptance length: 3.48, Accepted throughput: 59.89 tokens/s, Drafted throughput: 169.09 tokens/s, Accepted: 600 tokens, Drafted: 1694 tokens, Per-position acceptance rate: 0.719, 0.537, 0.372, 0.273, 0.219, 0.186, 0.174, Avg Draft acceptance rate: 35.4% ``` - vLLM version: v0.26.0 - vLLM main: vllm-project/vllm@d02df74 Signed-off-by: Wangbingjie <wangbj1207@126.com>
Leetrytry
pushed a commit
to Leetrytry/vllm-ascend
that referenced
this pull request
Sep 11, 2026
### What this PR does / why we need it? The following PRs are not cherry-picked to vllm v0.26.0 branch, so patches should be added on vllm-ascend releases/v0.26.0rc branch, if we want to use some kinds of DSpark (e.g. https://huggingface.co/RedHatAI/GLM-5.2-speculator.dspark). - vllm-project/vllm#48524: solving fc shape error when num_target_layers != num_hidden_layers - vllm-project/vllm#48776: solving KV Cache planning when target model is sparse-MLA and draft model is SWA - vllm-project/vllm#48639: solving `sample_from_anchor` wrongly loaded ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - GLM-5.2 W4A8 + DSpark acceptance rate ``` INFO 08-14 15:40:22 [metrics.py:120] SpecDecoding metrics: Mean acceptance length: 2.82, Accepted throughput: 73.39 tokens/s, Drafted throughput: 282.88 tokens/s, Accepted: 790 tokens, Drafted: 3045 tokens, Per-position acceptance rate: 0.729, 0.457, 0.278, 0.161, 0.083, 0.064, 0.044, Avg Draft acceptance rate: 25.9% INFO 08-14 15:40:32 [metrics.py:120] SpecDecoding metrics: Mean acceptance length: 4.38, Accepted throughput: 146.65 tokens/s, Drafted throughput: 303.29 tokens/s, Accepted: 1469 tokens, Drafted: 3038 tokens, Per-position acceptance rate: 0.836, 0.652, 0.537, 0.408, 0.353, 0.318, 0.281, Avg Draft acceptance rate: 48.4% INFO 08-14 15:40:42 [metrics.py:120] SpecDecoding metrics: Mean acceptance length: 3.48, Accepted throughput: 59.89 tokens/s, Drafted throughput: 169.09 tokens/s, Accepted: 600 tokens, Drafted: 1694 tokens, Per-position acceptance rate: 0.719, 0.537, 0.372, 0.273, 0.219, 0.186, 0.174, Avg Draft acceptance rate: 35.4% ``` - vLLM version: v0.26.0 - vLLM main: vllm-project/vllm@d02df74 Signed-off-by: Wangbingjie <wangbj1207@126.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Serving a sparse-MLA target such as
nvidia/GLM-5.2-NVFP4with a regular sliding-window DSpark draft could fail during KV-cache planning because the target MLA/indexer pages cannot be unified with the draft page.Design
Keep the existing grouping path unchanged when page sizes can be unified. If page-size unification fails for the narrow MLA + regular-SWA case, promote only the draft's cache-allocation spec to
FullAttentionSpecand reconcile it to the target's unique block size (64 for GLM). Draft attention computation remains sliding-window.This produces one
UniformTypeKVCacheSpecsgroup with a shared block table and the existing contiguous per-layer tensors. It avoids adding a new allocator layout or changing packing, zeroing, or memory accounting. Existing page padding is enlarged when necessary after block-size promotion.The tradeoff is deliberate: the single draft layer allocates KV slots like full attention, so this does not retain SWA's physical-memory saving. With today's global block pool, however, it avoids the additive block-ID demand of separate target and draft groups.
Validation
.venv/bin/python -m pytest tests/v1/core/test_kv_cache_utils.py \ tests/v1/core/test_contiguous_kv_packing.py -q \ -k 'unify_hybrid_kv_cache_specs or mla_draft_prefers_standard_layout or mla_with_incompatible_swa or contiguous'End-to-end checks:
RedHatAI/GLM-5.2-speculator.dspark: 100/100 requests, 236.19 output tok/s, 40.44% draft-token acceptance, mean accepted length 3.83.shanjiaz/glm-dspark-0(fully SWA): the final path selected 64-token blocks, initialized successfully with 504,217 KV-cache tokens and 25.21x maximum concurrency at 20k context.Related work
This is not duplicate work: #49472 optimizes the normal page-unification path, while this fallback runs only after that path fails. #48993 concerns the DeepSeek-V4 packed layout, which this PR does not change. #47926 covers broader restored-prefix DSpark correctness; this PR only validates ordinary local prefix-cache reuse.
Follow-up
If draft KV memory becomes material, explore a real per-group/variable-size pool that preserves SWA reclamation without charging separate groups against one global block-ID pool. That is intentionally outside this bug fix.
AI assistance
AI assistance was used for code review, test refinement, related-PR analysis, and benchmark execution/analysis. The human submitter reviewed the final two-file change and validation evidence.