[isytwu]

Internal review draft. Remove this note before submitting upstream. The Kimi K3 c70 runs are still in progress, so their results are pending. Links to sibling PRs point at this fork and need the upstream numbers later. Confirm the AI-assistance statement and checklist after review.

Overview

UMBP stores only the Mamba state page core hands off at a prompt-tail boundary. When that boundary falls inside an attention block, resuming there also needs the attention block's page, which UMBP never stores, so a Mamba hybrid restores from an older boundary or not at all. This PR stores that page with each handed-off boundary.

Claims

• A Mamba hybrid restores from UMBP at the same boundary as a local GPU prefix hit when prefix_match_unit is finer than the attention block. Granite 4.0-H 350M: 0 → 832 restored tokens, with identical tokens and logprobs.
• With EAGLE's block drop, the page is keyed one lookup unit past the boundary, where core registers its own partial tail. The request-finished path now keys its page the same way.
• Kimi K3 1P1D at concurrency 70: Pending: the Kimi K3 c70 runs are in progress.

Validation

MI355X, vllm/vllm-openai-rocm:nightly-40941f2c (MORI 1.2.3.post1 with UMBP) with this branch's umbp/ directory.

Unit tests: pytest -q tests/v1/kv_connector/unit/test_umbp_*.py gives 234 passed. Three of the four new cases fail on main. The non-EAGLE finished-path case already passes there and guards the keying rule.

New end-to-end regression test, one GPU:

pytest -v -s -rs tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py -k hybrid_restores

Granite 4.0-H 350M, prefix_match_unit=16, 400-token blocks, 840-token prompt; the tail checkpoint sits at 832:

	GPU prefix hit	Restored from UMBP	Result
main	832	0	AssertionError: assert 0 == 832
this PR	832	832	passed, outputs identical

Kimi K3 1P1D (TP8/DCP8 each, DSpark, prefix_match_unit 128, embedded UMBP 1799 GB on P, agentic trace at concurrency 70, 1 h):

Pending: the Kimi K3 c70 runs are in progress.

Details

Request flow on the prefill instance:

1. Turn 1 computes its prompt. Core caches the Mamba state at the prompt-tail boundary b, inside attention block k, and hands that state page to UMBP (boundary_state_offloads).
2. UMBP stores attention blocks 0..k-1 and the state at b. Attention block k, holding tokens [k*B, b), is never stored. The finished path stores it only when the request stops exactly at b.
3. Turn 2 arrives after the GPU cache dropped turn 1. The lookup coordinator needs every group to hit at the same boundary. Mamba hits at b but attention stops at k*B, so the coordinator settles on an older boundary, or on none when only the tail checkpoint is retained (prefix_cache_retention_interval=0, the default).
4. With this PR, step 2 also stores attention block k keyed at b, and turn 2 resumes at b.

Key position: the coordinator matches an EAGLE group one lookup unit past a candidate and then drops that unit. For groups whose manager has use_eagle, the page is keyed at b + unit; otherwise at b. Reading this from the coordinator keeps it correct for DFlash/DSpark, which skip the drop, and if vllm-project#65448 changes when the drop applies. A block-aligned end is a full block and is stored with its prefix.

Cost: one more attention page per handed-off boundary. On Kimi K3 that is about a third of a KDA checkpoint.

Not a duplicate: no open PR stores attention pages at hybrid boundaries. vllm-project#64656 adds a hybrid restore test whose prompts end on block boundaries, without a fine prefix_match_unit, where this change has no effect. Both tests use the same model, so whichever lands second can share the fixture.

AI assistance was used for this change and description. The submitter reviewed every line and ran the tests above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.