[BugFix] Cache Hit Rate with KIMI + RHAI Dspark - #52804
Open
robertgshaw2-redhat wants to merge 4 commits into
Open
[BugFix] Cache Hit Rate with KIMI + RHAI Dspark#52804robertgshaw2-redhat wants to merge 4 commits into
robertgshaw2-redhat wants to merge 4 commits into
Conversation
…age sizes For a hybrid target (e.g. MLA + linear attention, block size aligned to the mamba state page, 1536 tokens) plus a draft model whose layers have a different per-token KV footprint (e.g. an SWA DSpark drafter), unify_kv_cache_spec_page_size scaled the draft layers' block size by the byte ratio to make pages exactly equal. The resulting block size (e.g. 3456 tokens) is generally neither a multiple nor a divisor of the other groups' block sizes, so scheduler_block_size — the LCM of all group block sizes and the prefix-cache hit alignment — exploded (e.g. lcm(1536, 3456) = 13824 tokens). Every prefix-cache hit then rounds down to a 13824-token boundary, costing ~7K reusable tokens per request at 70K ISL (~10-15pp hit rate). Fix: when byte-exact scaling would break token commensurability and the layer's backend tolerates padded pages (indexes_kv_by_block_stride), pick the largest block size that divides the max-page groups' block-size LCM instead and pad the page to the shared page size. The scheduler block size then stays at the target's block size (1536), matching the behavior of an architecture-matched (e.g. MLA-native) drafter. Backends without padded-page support keep the byte-exact behavior; a new warning in resolve_kv_cache_block_sizes flags incommensurate group block sizes so the hit-granularity cost is visible. Co-authored-by: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDEqyZ6WAXtzkp5S3W9hdQ Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
robertgshaw2-redhat
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
heheda12345,
ivanium,
njhill,
orozery and
ywang96
as code owners
August 18, 2026 17:12
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
robertgshaw2-redhat
added
the
ready
Collaborator
Author
|
/ci run |
Collaborator
Author
|
/ci run |
|
✅ Triggered Buildkite CI #84423 for commit |
|
✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/84423 |
Contributor
|
Hi @robertgshaw2-redhat, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
Collaborator
Author
|
raised an IMA |
…che views The padded-page strided view in _reshape_attention_kv_cache strided every kernel block by the full manager block's padded page. With kernel-block splitting (spec block_size > kernel block size, e.g. a 1536-token padded SWA group viewed as 3x 512-token kernel blocks), the view overshot the allocation by the split factor and failed at init: RuntimeError: setStorage: ... requiring a storage size of 12760678400 are out of bounds for storage of size 4253810688 Previously padded pages only ever coexisted with split == 1 (skip-quant SW layers pick their block size to fit the shared page), so this path was never exercised. Distribute the padding evenly across the kernel blocks so they stay uniformly strided: each kernel block gets padded_page / split bytes, and manager block boundaries still land on padded-page multiples, keeping raw-tensor consumers (connectors, offload) that copy whole manager pages correct. Co-authored-by: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDEqyZ6WAXtzkp5S3W9hdQ Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
…/robertgshaw2-redhat/vllm into claude/pr-problem-solving-mkp5tz
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 22, 2026
… NaN class, mirrors carried dense fix); ubatch metadata slice propagates causal/encoder/dcp/is_prefilling/replayssm fields; ReplaySSM dtype-count loud-fail in gdn base; padded-page kernel-block split support in create_kv_cache_views/compute_layout_strides (completes vllm-project#52804 semantics); CPU rejection-sampler shims brought to current Triton conventions; TQ scratch shutdown release; speculator manager teardown widened; import/style repairs — 24 findings refuted with evidence Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 22, 2026
…: V1-runner GDN spec-split recovery fallback in build() (rebase regression — raw-count callers silently took the non-spec branch); mamba_attn cached-spec cursor hardening (capture-gated commit/reset, dummy-row NULL_BLOCK_ID guard); TQ continuation-prefill full-dequant launch missing Centroids_ptr positional (latent TypeError on every TQ preset) + value-codebook loud-fail on store and decode sides; decode_write_pos_d zeros-init for capture safety; TQ shared-scratch shutdown release completed (reset_shared_scratch wired, SoA buf_holder to shared holder); per-layer kv_quant_mode at TQ spec build (VLLM_TQ_LAYER_BITS); dflash/dspark fused-KV dtype from norm weight + reload invalidation; commensurate padded-block exact page-fit check; create_kv_cache_views dense-dtype demotion fix (innermost stride==1 is required, not disqualifying); relaxed_thinking async-sched auto-disable + CPU shim + monkeypatch; V1 normalize gate ReplaySSM-spec parity; thinking-budget in_think_mask persistent staging pair; hybrid-mamba APC warm leg; cudagraph profiling teardown hook reset + test coherence; marlin fp32-reduce zero-init test; vllm-project#52804 kernel-block-split test — cross-file not_fixed patches applied (3), remaining not_fixed entries carry proofs-of-not-our-path Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 22, 2026
…re-pads an already-padded layer to the shared page (boot AssertionError class); commensurate per-token-bytes zero guard; AR speculator draft-prefill dispatches on num_tokens_padded (restores upstream vllm-project#47352, reverted by carried vllm-project#48244 pick); gemma4_dspark + laguna_dflash fused-KV dtype derived from norm weight + bias cast (fp16 drafters); unpadded() propagates mm_req_doc_ranges; V1 runner zeroes padded-row num_prompt_tokens_cpu (stale-length dummy-row class); fused_recurrent + fused_sigmoid_gating INPLACE_FINAL_STATE load masked to the row (SM-fault class, mirrors vllm-project#50021's sibling bound); profile_cudagraph_memory empty-sample guard; TQ spec head_size_v; packed-codec (kvarn_/turboquant_) branch in _validate_cache_dtype logging; max_page_block_lcm includes padded-to-max layers (PR vllm-project#52804 intent) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 28, 2026
…KV page — hand-ported around the house commensurate-block path, LCM fed the scaled blocks, kernel-split guard added Hand-port (not a cherry-pick): our unify_kv_cache_spec_page_size carries the house commensurate-block-size path (vllm-project#52804 kernel block splitting) that the PR cannot see. Carried from the PR: - Branch choice and scaling ratio for non-MLA attention now come from the natural page (unpadded_page_size_bytes), not the possibly-stale page_size_padded; MLA keeps its aligned page as the base because __post_init__ reapplies alignment on every replace. Factored into _unify_branch_page_size_bytes() so the LCM predicate and the loop cannot drift apart. - The pad branch scales block_size by the whole part of the ratio before recording page_size_padded (_scale_block_size_before_padding), so a draft head next to a quantized primary stops claiming ~51x the blocks it needs. - The divisible branch drops stale padding when it scales (page_size_padded =None) instead of the previous "re-pad to the shared page": with the ratio taken from the natural page the scaled page lands exactly on the maximum. MLA/non-attention now take a plain replace(block_size=...) as upstream does, so alignment padding is never left stale at max_page_size. Deviations from the PR: - HOUSE INTERACTION: max_page_block_lcm (~:1145) is computed BEFORE the loop from exactly the layers this now rescales. Pad-branch scaling is resolved first into scaled_pad_specs and the LCM is built from those SCALED block sizes, so _commensurate_padded_block_size never caps candidates against a granularity no group ends up using. The LCM membership predicate also moved to the natural page so it matches the loop's branch test. - Kernel-split guard the PR has no equivalent of: the scale ratio must divide max_page_size (max_page_size % ratio == 0), mirroring the existing check in _commensurate_padded_block_size, because compute_layout_strides distributes the padding across spec.block_size // kernel_block_size kernel blocks and asserts even divisibility at boot. Worst case (backend splits down to the pre-scale block); a backend that supports the scaled block never splits. When the whole ratio fails the guard, the largest smaller ratio that passes is used instead of dropping the scale entirely, so the PR's own nvfp4 case (whole ratio 51 against a 3,354,624 B page) still scales by 48 rather than falling back to no scale at all. - The scaled candidate's page is re-checked against max_page_size per ratio (specs whose states are floored from tokens_per_state are not linear in block_size). Tests: the PR's three cases ported (pre-padded case uses new_kv_cache_spec's page_size_padded kwarg instead of dataclasses.replace) plus a house-only case, test_unify_kv_cache_page_size_padded_scale_feeds_the_commensurate_lcm, that fails on a stale (pre-scale) LCM: the divisible layer would drop from block 32 to a padded block 16. Verified in the vllm-openai container against the source tree: tests/v1/core/test_kv_cache_utils.py 69 passed vs 65 on the base branch, identical 23-item pre-existing failure set (missing flash-attn extensions in the stock image).
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
1 participant
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.
Keep KV cache group block sizes commensurate when unifying page sizes
For a hybrid target (e.g. MLA + linear attention, block size aligned to the mamba state page, 1536 tokens) plus a draft model whose layers have a different per-token KV footprint (e.g. an SWA DSpark drafter), unify_kv_cache_spec_page_size scaled the draft layers' block size by the byte ratio to make pages exactly equal. The resulting block size (e.g. 3456 tokens) is generally neither a multiple nor a divisor of the other groups' block sizes, so scheduler_block_size — the LCM of all group block sizes and the prefix-cache hit alignment — exploded (e.g. lcm(1536, 3456) = 13824 tokens). Every prefix-cache hit then rounds down to a 13824-token boundary, costing ~7K reusable tokens per request at 70K ISL (~10-15pp hit rate).
Fix: when byte-exact scaling would break token commensurability and the layer's backend tolerates padded pages (indexes_kv_by_block_stride), pick the largest block size that divides the max-page groups' block-size LCM instead and pad the page to the shared page size. The scheduler block size then stays at the target's block size (1536), matching the behavior of an architecture-matched (e.g. MLA-native) drafter. Backends without padded-page support keep the byte-exact behavior; a new warning in resolve_kv_cache_block_sizes flags incommensurate group block sizes so the hit-granularity cost is visible.
Claude-Session: https://claude.ai/code/session_01TDEqyZ6WAXtzkp5S3W9hdQ
Purpose
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.