[Bugfix][KV Offload] Do not let a recurrent group's unhashed block truncate the load boundary

[yifjiang]

update_state_after_alloc scans the allocated block list for the first block that is not is_null and block_hash is None and treats that index as the start of the freshly allocated region.

This PR also renames that value from num_locally_computed_gpu_blocks to load_start_gpu_block_idx. It is not a count of computed blocks — it is the index in group_blocks where the load region begins, used only as a slice bound. The old name is actively misleading: in test_partial_lookup_returns_exact_boundary_and_group_load_keys it is 1 while num_locally_computed_tokens is 0. The _gpu_block_idx suffix follows the sibling store path (_build_store_jobs), which already calls its equivalent start_gpu_block_idx and feeds the same block_indices field; the load_ prefix keeps the two unambiguous now that both are visible in one file. The store path is left untouched.

The assertion that fires is scheduler.py#L1034-L1037, reached from Scheduler.schedule.

The scan started at index 0 — inside the region already reported as locally computed, where every block is computed by definition. For a recurrent (Mamba / GDN) group that is not merely redundant but wrong. Such a group has no per-token KV at all: it carries a fixed-size state, which is why get_sliding_window_size_in_chunks() returns 1 for MambaSpec ("Mamba depends on a single state"). Most positions in its block table point at the shared null sentinel, and the one real block holding the state is legitimately not full-and-cached — non-null, unhashed, sitting below the computed mark.

legend   [H] hashed   [·] null placeholder   [U] non-null, UNHASHED   [F] fresh

recurrent group (GDN/Mamba, window = 1 chunk), 21 chunks computed + 1 loaded:

  idx      0    1    2    3    4   ...   20   21
         [H]  [H]  [·]  [U]  [·]  ...   [·]  [F]
                         ^
BEFORE   scan starts at 0, stops here          boundary = 3
         |<-------- claimed computed: 21 ------------>|
         assert 21 <= 3   ✗  AssertionError
         load region would be blocks 3..21  ->  19 wrong destinations, wrong keys

AFTER    scan starts at first_fresh = 21 ------------>^
         boundary = 21,  assert 21 <= 21  ✓
         load region = block 21 only  ->  matches the 1-chunk hit

For a full-attention group this is a no-op: every token has KV, every block below the boundary is full-and-cached, so the old scan skipped them anyway. The existing assert now holds by construction (nlcgb >= cdiv(nlct, tpb) ⟹ nlct <= nlcgb * tpb); it is kept as a guard against other causes rather than deleted.

Reproduction

4-node GB200, Qwen3.5-family hybrid GDN+GQA (4 KV cache groups, 3 recurrent), MTP, TP=4 x DP=4, 200 GiB host tier, --num-gpu-blocks-override 259. Under sustained load the leader crashed after 47,142 requests:

  File "vllm/v1/core/sched/scheduler.py", line 963, in schedule
    self.connector.update_state_after_alloc(
  File ".../kv_connector/v1/offloading/scheduler.py", line 832, in update_state_after_alloc
    num_locally_computed_tokens
AssertionError

leader exit=1 | other three exit=137 (gang teardown)

The captured 3h22m window on the crashing rank makes the trigger unambiguous. It saw 22 lookups that hit offloaded tokens, but only 4 that were true partial hits (num_locally_computed_tokens > 0) — and all four arrived in the last 36 seconds:

time	external tokens	locally computed	outcome
00:57 – 04:18	18 lookups	0	fine — assert reads 0 <= X
04:18:45	4256	8512	ok
04:18:52	2128	4256	ok
04:19:07	2128	10640	ok
04:19:21	4256	10640	AssertionError, same second

Why this is easy to miss. It needs num_external_tokens > 0 and num_locally_computed_tokens > 0. Tests that flood to full eviction before re-sending produce num_computed_tokens == 0 — 18 of the 22 lookups above — where the assert reads 0 <= X and cannot fire. The rank sat clean for 3h21m under exactly that pattern before the first partial hits appeared.

Reproduced on a build that carries the #49146 clamp

The assert is not masked by #49146's num_allocated_chunks clamp (94ed0bf4e0, in main). A build composed as #49146 clamp + #52771, without this fix was put under sustained production traffic and asserted three times across two independent clusters, with a fault-free lifetime of roughly 20 minutes / ~850 requests:

vllm/v1/core/sched/scheduler.py:963                     schedule
vllm/.../kv_connector/v1/offloading/scheduler.py:840    update_state_after_alloc
AssertionError

Both tracebacks are this same assertion. The line numbers above come from the deployed vLLM-0.26.0-based image, not from main (where it is L1034-L1037): within that tree it sits at 832 without the #49146 clamp and at 840 with it, the clamp inserting exactly 8 lines above.

Across four builds under the same workload and config, the correlation is with this fix, not the clamp and not #52771:

build	#49146 clamp	#52771	this fix	fault-free lifetime
A	✓	—	—	22–70 min, 3 asserts
B	—	✓	—	10 / 30 / 40 / 88 min, ≥4 asserts
C	✓	✓	—	~20 min, 3 asserts
D	✓	—	✓	336 min, zero asserts

Every build lacking this fix asserts — including two that carry the clamp (A and C); the only build carrying this fix is the only one that holds. Memory was never implicated (no OOMKilled; peak well under limit).

A separate note on reproducing it: a synthetic driver that holds the hit shape constant does not surface this. On a smaller same-family deployment I drove 12,000 consecutive requests at a 100% true-partial-hit rate with no assert — but every request had an identical num_locally_computed_tokens, so it probed one boundary alignment repeatedly. The production crashes arise under organically varied prefix lengths, which is what walks the boundary across block offsets.

Why the boundary is exact, and why the null-skip stays

The caller block-aligns the local hit before handing it to the connector (scheduler.py), so num_locally_computed_tokens % tokens_per_block == 0 and first_fresh_gpu_block is exact. The existing assert then reduces to boundary >= first_fresh_gpu_block — this fix makes it hold by construction rather than detecting the violation after the fact.

The null-skipping loop is retained and is not redundant. With num_locally_computed_tokens == 0 a group's block table can begin with the shared null sentinel, and the load must start at the first real block. test_partial_lookup_returns_exact_boundary_and_group_load_keys (from #50507) pins this on a Mamba hybrid: replacing the loop with a plain = first_fresh_gpu_block makes it emit [31, 32, 0, 41] — writing loaded KV into null block 0, which is shared across every request. Verified by running it both ways.

Relationship to #52735 / #52771

This is independent of the drafter-annotation defect, and of the abort fix. Three arms, all at current main (which already carries the #49146 num_allocated_chunks clamp, commit 94ed0bf4e0), running the regression test below:

arm	#49146 clamp	#52771	this fix	offloading_connector/ suite
A	✓ (in main)	—	—	regression test fails
B	✓	✓	—	1 failed, 254 passed
C	✓	✓	✓	255 passed, 0 failed

Arms B and C were run back-to-back in a single GPU allocation, toggling only this hunk, so the environment is identical; the one failure in B is this PR's regression test and nothing else. Every test that passes without the fix still passes with it.

Arm B is the decisive one: with #52771 (44a3045ec1) cherry-picked on top of main, the all-groups drafter fallback removed and no group marked as a drafter — so the volatile-tail pop never runs — the boundary assert still fires. The assert is not gated on is_eagle_group anywhere.

#52771 is also complementary: by restoring real loads it raises how often num_external_tokens > 0, and therefore how often this path is reached. It is the right fix for the annotation defect and should land; this is the crash underneath it.

(This PR previously also carried a competing fix for #52735 — removing the eagle query-widen, required_window += 1 and the volatile-tail pop. I have dropped all of it in favour of #52771, which was filed first, targets the root cause rather than the symptom, and additionally fixes a store-side hole mine did not. What remains here is only the boundary fix.)

Tests

test_recurrent_group_unhashed_block_does_not_truncate_load_boundary plants exactly that block — non-null, unhashed, below the boundary — and asserts the destination blocks are correct (41 not in loaded, 42 in loaded), not merely that nothing raised. Verified to fail against unpatched main and pass with this change.

Full tests/v1/kv_connector/unit/offloading_connector/ suite green on an aarch64 GB200 node.

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

[yifjiang]

validated: the boundary assert fires with the competing fix applied and none of this PR's eagle changes present

Earlier in this PR I argued from code that the update_state_after_alloc boundary bug is independent of the eagle changes. That is now measured rather than argued.

I built an image carrying only #52771 (the alternative fix for #52735) on our production baseline — so the all-groups drafter fallback is removed, no group is marked a drafter, the volatile-tail pop never applies, and none of this PR's eagle commits are present. Under sustained partial-hit load on a 4-node GB200 deployment it crashed after 47,142 requests:

scheduler.py:963  self.connector.update_state_after_alloc(
offloading/scheduler.py       num_locally_computed_tokens
AssertionError

leader exit=1 Error (1 assert); other three exit=137 (gang teardown)
partial hits just before: 12768, 6384, 8512, 6384, 4256 offloaded tokens

Every alternative explanation is eliminated by construction: no drafter annotation, no pop, no eagle change from this PR. 518af3b is necessary on its own.

It also explains why this took so long to reproduce on dev. The assert needs num_external_tokens > 0 and num_locally_computed_tokens > 0 — a partial hit. Probe-style tests that flood to full eviction before re-sending produce ncomp 0, where the assert reads 0 <= X and cannot fire. Only sustained partial-overlap traffic reaches it; the same arm had been clean and idle for 3h47m immediately beforehand.

On overlap with #52771: that PR was filed first and is the better fix for the annotation defect — it targets the root cause and additionally fixes a store-side finish hole that this PR does not. I have offered there to close the eagle half of this PR in favour of theirs and send the boundary fix standalone. Maintainers should feel free to take that route; the boundary fix is the part that must survive either way.

[kamb-code]

Independent verification of this PR, CPU-only, at head 469b23f8 against merge-base 41f179b5.

1. The regression test is load-bearing. test_recurrent_group_unhashed_block_does_not_truncate_load_boundary passes on this branch, and fails with AssertionError on exactly the num_locally_computed_tokens <= num_locally_computed_gpu_blocks * tokens_per_block assert when only vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py is reverted to main. So it pins this fix rather than passing incidentally.
2. No regressions: full tests/v1/kv_connector/unit/offloading_connector/ suite on this branch — 244 passed, 2 skipped.
3. The two PRs compose. Since I said on [Bugfix] OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode #52771 that it should not merge alone, I checked the combination rather than assuming it. [Bugfix] OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode #52771's source changes apply cleanly on top of this branch, and with both applied (both test sets merged onto one tree) this PR's regression test and [Bugfix] OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode #52771's three all pass together — 4/4 — with the full offloading-connector suite at 247 passed, 2 skipped.

The reasoning in your diff comment matches what I read in the code: the scan starting at index 0 walks the already-computed prefix, and a recurrent group legitimately holds non-null unhashed blocks there because it keeps a fixed-size state rather than per-token KV. Starting at cdiv(num_locally_computed_tokens, tokens_per_block) makes the assert hold by construction.

For reviewers: this fix and #52771 are complementary and have now been tested together. #52771 makes partial hits ordinary, which is what makes this boundary reachable, so this one should land first or alongside it.

Disclosure: AI-assisted verification (Claude Code); results reviewed by me.

[mergify]

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

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

[orozery]

File "vllm/v1/core/sched/scheduler.py", line 963, in schedule
self.connector.update_state_after_alloc(
File ".../kv_connector/v1/offloading/scheduler.py", line 832, in update_state_after_alloc
num_locally_computed_tokens
AssertionError

leader exit=1 | other three exit=137 (gang teardown)

which vllm version (git sha) is this?

[yifjiang]

ffd46bfab (2026-07-24).

One thing to watch: offloading/scheduler.py:832 is not an upstream line number — that file is patched. Upstream at this commit the same assert is line 816.

The patch is #52771 and nothing else — none of this PR's commits — verified against the file extracted from that image. That's the point of the run: the boundary assert still fired, after 47,142 partial-hit requests.

The image does carry unrelated vendor vLLM patches (Qwen3.5 text support, DCP flashinfer a2a, some kernels, modelopt mtp-exclude); none of them touch the offload connector.

Line-number key for the builds in this thread, in case they come up: 816 upstream · 832 this build · 840 and 849 are later builds where I'd applied #49146 (unrelated to this bug — it shifts the file by 8 lines).

[orozery]

This should be replaced by:

assert num_locally_computed_tokens % tokens_per_block == 0

[yifjiang]

Applied. You're right that the old one is tautological now — the scan starts at cdiv(num_locally_computed_tokens, tokens_per_block), so load_start_gpu_block_idx * tokens_per_block >= num_locally_computed_tokens holds by construction. Asserting the alignment states the invariant that actually makes the cdiv exact.

One thing I could not verify locally, so flagging rather than sitting on it: tokens_per_block here is per group (spec.tokens_per_block is a list), and the code just above deliberately handles groups having different chunk sizes —

full_attn_tokens_per_chunk: set[int] = set()
...
# Only apply the optimization if there's a single consistent
# full-attention alignment size.

with the comment that "load hits are always aligned to this boundary" referring to the full-attention group. So the new assert now checks alignment against each group's block size, including recurrent ones. If a Mamba/SWA group's tokens_per_block doesn't divide the full-attention alignment, this would fire on exactly the hybrid deployment the PR is about.

I could not run it to check: this branch is now main-based and my only environment with the offloading tests installed is on an older vLLM (Medium isn't importable from vllm.v1.kv_offload.base there), so the suite won't load. If you're confident hits are aligned to every group's block size then this is strictly better than what I had, and CI's hybrid cases should cover it — I just didn't want to ship an assert I hadn't been able to exercise.

[yifjiang]

Correction to my note above — I did manage to exercise it, and the assert never fires.

I said the suite wouldn't load because my environment is on an older vLLM. That was me giving up too early: vllm/ in git is pure Python, so overlaying main's tree onto the older install (keeping its compiled .so) makes main's scheduler importable, and the offloading tests then run.

Same overlay, offloading/scheduler.py swapped underneath, tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py:

pristine main   82 failed / 106 passed
this PR         81 failed / 107 passed

regressions vs pristine:  none
fixed by this PR:         test_recurrent_group_unhashed_block_does_not_truncate_load_boundary
`num_locally_computed_tokens % tokens_per_block` assertion fired: 0 times

The 81 shared failures are the version mismatch between the overlaid Python and the older compiled extensions — identical on both sides, so the delta is attributable.

Zero firings across all 188 cases, including the TestEagle, sliding-window and mixed full-attn/SWA tests — the hybrid shapes my concern was about. So the per-group tokens_per_block worry doesn't materialise, and your version is strictly better than the tautological one it replaced. Nothing outstanding from my side.

[yifjiang]

Ran it e2e as well, on the real model. Summary: no violations, but the interesting case wasn't reachable — worth stating both halves.

Built an image from v22 that evaluates the new predicate against the live per-group block sizes and logs instead of asserting, so a violation would be observable rather than engine-fatal, and the ok line proves the path actually executed (otherwise "no violations" is indistinguishable from "never ran" — which is exactly what the first attempt turned out to be).

Deployment: 4-node GB200, TieringOffloadingSpec, 300 GiB host pool, OakHaven (Qwen3.5-Max, 3 Mamba + 1 full-attention KV group), 5/5 Ready.

Getting the path to execute took a rank-pinned probe — host pools are private per DP rank, so an unpinned re-send lands on a cold rank and never reads back. 41 unpinned requests of 160K tokens produced 0 loads; pinning to one rank produced a real read-back (8.04 GB CPU_to_GPU, warm re-send 0.8 s vs 13.7 s cold).

Result across 8 evaluations, one per KV group per load:

MODCHK ok  nlct=0  tokens_per_block=2128     x8
MODCHK VIOLATION                              x0

Two things worth having on the record:

1. All four KV groups report the same tokens_per_block (2128) — including the three recurrent ones. So the per-group divergence I raised above does not arise on this model; the assert sees one block size, not four.
2. Every e2e evaluation had nlct = 0, so they only exercise the trivial case. I tried to force a partial local hit (lighter eviction, then a prefix-superset request of 208K tokens which clearly reused the prefix at 5.6 s) and num_locally_computed_tokens was still 0 at the point of the check.

So the non-trivial nlct > 0 case is covered by the unit test (test_recurrent_group_unhashed_block_does_not_truncate_load_boundary, which fails without this fix), and e2e adds the real-model group sizes rather than a second test of the arithmetic. I did not manufacture a synthetic non-zero nlct in production, and I'd rather say that than imply broader coverage than I have.

Nothing outstanding from my side — happy to go further if you want the nlct > 0 path exercised in a live deployment too.

[orozery]

Shorter:

Suggested change

	# ``load_start_gpu_block_idx``: the index in ``group_blocks`` where the
	# load region begins -- a slice bound, not a count of computed blocks.
	# Scan from the computed boundary, not 0: ``num_locally_computed_tokens``
	# is block-aligned by the caller, so lower blocks are computed by
	# definition, and a recurrent (Mamba / GDN) group legitimately holds a
	# non-null *unhashed* block there -- a fixed-size state, not per-token
	# KV. Taking that as the start drags the boundary below the computed
	# mark and trips the assert. Skip nulls: block 0 is the shared sentinel
	# and must never be a load destination.
	# Scan from the computed boundary, not 0: sparse groups (Mamba,
	# SWA) legitimately hold non-null unhashed blocks below it.
	# Skip nulls (sentinel / out-of-retention padding).

[yifjiang]

Applied — also rebased onto current main (was 476 behind, needs-rebase). Only conflict was an import in the test file; kept both sides (KVCacheBlockCopy + make_block_hash_with_group_id).

[orozery]

/ci run

[github-actions]

✅ @yifjiang, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #87047 for commit bb9351d3c090.