[Test][UMBP] Restore a Mamba2 hybrid end to end through UMBP by AndreasKaratzas · Pull Request #64656 · vllm-project/vllm

AndreasKaratzas · 2026-10-30T04:47:24Z

Add a Granite 4.0-H hybrid restore test to test_mori_token_e2e.py, run through embedded UMBP.
Require bitwise-equal tokens and logprobs against a local GPU prefix hit, plus an exact external-hit count.
Cover a prompt ending mid-block and one ending exactly on a block boundary.

No hybrid model restores through UMBP end to end today. The Nemotron-H GSM8K row never gets a prefix hit because its 1088-token hybrid block is longer than the prompts, so the Mamba boundary-state handoff is only covered with fake runtimes. Mamba hybrids have no batch-invariant kernels, so a cold run can't be the reference, and prompts go one at a time because async load timing changes the batch layout and batched hybrid outputs aren't bitwise stable. On MI350X with MORI v1.2.3.post1 built with UMBP, pytest -v -s -rs v1/kv_connector/umbp_integration/test_mori_token_e2e.py -k hybrid passed 10/10 in a loop and the whole file gave 3 passed, but every case fails once the restored Mamba pages are zeroed after the load. The MI355 UMBP integration group already runs this directory, so CI is unchanged. No open UMBP PR (#61795-#61798, #63002, #59276) restores a hybrid end to end, though #63830 and #63670 rework Mamba align mode and may need this test adjusted if they land first.

AI assistance (Claude Code) was used to draft this change. Every line was reviewed by the submitter.

Add embedded runtime contract tests, MORI GPU integration, token-level restore coverage, reusable test helpers, user documentation, and AMD CI wiring. Signed-off-by: CharlesWu <Charles.Wu@amd.com> Co-authored-by: Yichao Zhu <Yichao.Zhu@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit d61f813) (cherry picked from commit 967175e) (cherry picked from commit 68cef07)

Verify completed stores do not retain publication-key history. Use native residency and byte checks for eviction coverage. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com> (cherry picked from commit 16c79a7) (cherry picked from commit c4b05e1)

Move the scheduler factory fix into the embedded test slice so its lookup instance matches the workers at this PR head. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

The TP2 x PP2 native roundtrip launches four ranks. Use the existing mi355_4 pool and request four GPUs for this integration step. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

Install gRPC and Protobuf dependencies in both ROCm image bases, request BUILD_UMBP=ON, and reject silently disabled builds through the CMake cache. Check UMBPClient and four visible GPUs before running the integration suite so a missing native backend cannot turn module skips into a green job. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

Keep the added embedded unit and GPU integration tests independent of Buildkite's long TMPDIR, including multiprocess workers and token restore. Share the existing short directory fixture across connector test suites. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

Remove the module-wide ROCm/MORI skip from the CPU contract suite so the four regressions merged with the embedded implementation continue to run. All 52 cases pass without MORI; native GPU tests retain their own guards. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

Run the integration suite with the image's pytest instead of a separate virtualenv. Fail collection when MORI lacks UMBPClient and fail the TP2 x PP2 round-trip when fewer than four GPUs are visible, so a misconfigured image or agent still cannot pass through skips. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

- Docker: install gRPC/Protobuf build deps only in build_mori; final gets the runtime libs and a UMBPClient import check; drop no-op BUILD_UMBP. - CI: single-GPU integration group in test_areas and test-amd.yaml, with complete source dependencies. - Tests: drop tautological and fake-only tests, use a fake-clock fixture for timeouts, move the hybrid-tail restore into the unit suite with a non-periodic fill, rewrite the e2e restore with vllm_runner and an external-hit assertion, and gate on ROCm + UMBP via pytestmark. - Accuracy: add a ROCm-gated UMBP GSM8K row that asserts the reload. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

…eduler - e2e: under VLLM_BATCH_INVARIANT=1 a UMBP restore must reproduce the cold run's tokens and logprobs, for a batch ending mid-block and for a restore that extends a shorter GPU prefix hit, with exact external-hit counts. - unit: an object evicted between lookup and load must fail or recompute per kv_load_failure_policy through the real Scheduler, for one and two KV cache groups, and never leave a stale prefix-cache hit. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify · 2026-10-30T09:41:09Z

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

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

Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

charxwu

Tested on MI350X (gfx950) with vLLM built from 687c0ca (MORI 1.2.3.post1 with UMBP, VLLM_WORKER_MULTIPROC_METHOD=spawn):

pytest -v -s -rs v1/kv_connector/umbp_integration: 5 passed.
... test_mori_token_e2e.py -k hybrid: 10/10 passed, about 95 s each.
Zeroing the restored Mamba blocks after the UMBP load fails both cases at the first generated token (70 vs 198, 319 vs 198), while the cached-token and external-hit asserts still pass.

LGTM once rebased.

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.

AndreasKaratzas · 2026-11-03T00:07:37Z

Thanks @charxwu for testing it, including the zeroed-Mamba check!

AndreasKaratzas · 2026-11-03T00:13:29Z

/ci run

github-actions · 2026-11-03T00:13:51Z

✅ Triggered Buildkite CI #99386 for commit cd4bf55f7ed7.

tjtanaa

Lgtm. Nice. We have such a small model for testing.

AndreasKaratzas · 2026-11-04T06:41:07Z

/ci run

github-actions · 2026-11-04T06:41:35Z

✅ Triggered Buildkite CI #99869 for commit 7cf9f1be45d5.

ℹ️ This PR is 2 commits behind upstream main. Merged upstream main into this branch as 7cf9f1be45d5 before starting CI. Pull the branch before pushing more changes.

peymanr · 2026-11-04T07:00:03Z

PR Pundit review

(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)

#		Finding	Where
1	🟠	Hybrid test cannot catch an align-mode checkpointing bug consistent across both paths: this PR uses the GPU hit, not the cold run, as oracle	`tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py:178`

1. 🟠 Should fix: Hybrid test cannot catch an align-mode checkpointing bug consistent across both paths: this PR uses the GPU hit, not the cold run, as oracle

tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py:178

test_hybrid_restore_matches_local_prefix_hit discards the cold llm.generate output and only compares the UMBP restore with a GPU prefix hit. Both paths resume from the same stored Mamba boundary state, so the test checks that UMBP round-trips the state pages, not that the align-mode checkpoint itself is correct. That is the documented intent: Mamba hybrids have no batch-invariant kernels, so an exact cold comparison would not hold. Checkpoint correctness against a no-cache run is covered by the test_apc_* tests in test_hybrid.py. Optionally, keep the cold output and add a tolerance check (check_logprobs_close) as a cheap sanity guard.

Who is affected: Only test coverage: CI runs of the ROCm+MORI UMBP integration test on the granite-4.0-h-350m hybrid.

Checked against the code by an independent pass.

_{Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at 7cf9f1be; findings checked and posted by @peymanr; points already raised on this PR are left out.}

mergify · 2026-11-04T07:14:26Z

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

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

charxwu and others added 14 commits October 29, 2026 09:05

[Test][UMBP] Create embedded GPU scheduler through its runtime

49b195a

Move the scheduler factory fix into the embedded test slice so its lookup instance matches the workers at this PR head. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

[CI][UMBP] Allocate four GPUs for embedded integration tests

a827544

The TP2 x PP2 native roundtrip launches four ranks. Use the existing mi355_4 pool and request four GPUs for this integration step. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>

Merge branch 'main' into review/umbp-split-3-embedded

18b84a4

Merge branch 'main' into review/umbp-split-3-embedded

98af942

Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Merge branch 'main' into review/umbp-split-3-embedded

4ab20f5

[Test][UMBP] Restore a Mamba2 hybrid end to end through UMBP

2c8b981

Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify Bot added ci/build rocm Related to AMD ROCm kv-connector labels Oct 30, 2026

github-project-automation Bot added this to AMD Oct 30, 2026

github-project-automation Bot moved this to Todo in AMD Oct 30, 2026

mergify Bot added the needs-rebase label Oct 30, 2026

Merge branch 'main' into akaratza_umbp_hybrid_restore

687c0ca

Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>

mergify Bot removed the needs-rebase label Oct 30, 2026

charxwu reviewed Nov 2, 2026

View reviewed changes

AndreasKaratzas marked this pull request as ready for review November 3, 2026 00:02

AndreasKaratzas requested review from ApostaC, ivanium and orozery as code owners November 3, 2026 00:02

claude Bot reviewed Nov 3, 2026

View reviewed changes

AndreasKaratzas requested a review from wzhao18 as a code owner November 3, 2026 00:02

Merge branch 'main' into akaratza_umbp_hybrid_restore

cd4bf55

tjtanaa approved these changes Nov 4, 2026

View reviewed changes

tjtanaa added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 4, 2026

tjtanaa enabled auto-merge (squash) November 4, 2026 05:29

Merge branch 'main' into akaratza_umbp_hybrid_restore

ef8704c

Merge branch 'main' into akaratza_umbp_hybrid_restore

7cf9f1b

mergify Bot added the needs-rebase label Nov 4, 2026

isytwu mentioned this pull request Nov 4, 2026

[Bugfix][UMBP] Store the attention page of Mamba hybrid boundaries isytwu/vllm#1

Open

6 tasks done

Navigation Menu

Repository navigation

[Test][UMBP] Restore a Mamba2 hybrid end to end through UMBP - #64656

[Test][UMBP] Restore a Mamba2 hybrid end to end through UMBP#64656
AndreasKaratzas wants to merge 18 commits into
vllm-project:mainfrom
AndreasKaratzas:akaratza_umbp_hybrid_restore

AndreasKaratzas commented Oct 30, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Oct 30, 2026

Uh oh!

charxwu left a comment

Uh oh!

claude Bot left a comment

Uh oh!

AndreasKaratzas commented Nov 3, 2026

Uh oh!

AndreasKaratzas commented Nov 3, 2026

Uh oh!

github-actions Bot commented Nov 3, 2026

Uh oh!

tjtanaa left a comment

Uh oh!

AndreasKaratzas commented Nov 4, 2026

Uh oh!

github-actions Bot commented Nov 4, 2026

Uh oh!

peymanr commented Nov 4, 2026

Uh oh!

mergify Bot commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

AndreasKaratzas commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mergify Bot commented Oct 30, 2026

Uh oh!

charxwu left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

AndreasKaratzas commented Nov 3, 2026

Uh oh!

AndreasKaratzas commented Nov 3, 2026

Uh oh!

github-actions Bot commented Nov 3, 2026

Uh oh!

tjtanaa left a comment

Choose a reason for hiding this comment

Uh oh!

AndreasKaratzas commented Nov 4, 2026

Uh oh!

github-actions Bot commented Nov 4, 2026

Uh oh!

peymanr commented Nov 4, 2026

PR Pundit review

1. 🟠 Should fix: Hybrid test cannot catch an align-mode checkpointing bug consistent across both paths: this PR uses the GPU hit, not the cold run, as oracle

Uh oh!

mergify Bot commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

AndreasKaratzas commented Oct 30, 2026 •

edited

Loading