[AndreasKaratzas]

• Split each layer's page into per-head runs in data.py, using the allocator's compute_layout_strides.
• Register the same runs in the embedded.py bulk path, falling back to per-range copies for a group whose stride doesn't match.
• Fail at startup instead of corrupting KV when a page is split and kv_cache_layout is unset.
• Round-trip all 6 layouts through the bulk and range paths in test_umbp_embedded.py.

UMBP took cache.stride(0) as the block stride and copied one contiguous page per layer, both on the range path in data.py and on the bulk path in _register_load_layouts, as noted in the review of #61794. Under LHBNC, the default with ROCM_AITER_FA and VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1, K and V are head planes spanning every block, so the engine fails at init with KV region bytes (65536) exceed its stride (32768), and under BHLNC heads of different layers interleave inside a block row, so UMBP silently restores the wrong KV. pytest -q v1/kv_connector/unit/test_umbp_embedded.py -k every_kv_cache_layout gives 4 failures without the fix and 12 passed with it. On MI350X with Qwen3-0.6B, 16 greedy tokens after a 192-token UMBP hit now match recompute under BHLNC (TRITON_ATTN with VLLM_KV_CACHE_LAYOUT=BHLNC) and under AITER shuffle-KV, and stay identical under the default layouts; the batch-invariant tests in test_mori_token_e2e.py likewise fail without the fix and pass with it under VLLM_KV_CACHE_LAYOUT=BHLNC and LHBNC. Batch invariance rules out ROCM_AITER_FA, so those tests run on TRITON_ATTN, and AITER shuffle-KV is covered only by the greedy check above. Other layouts and MLA keep the single-span copy, and refusing LHBNC and BHLNC instead would make UMBP unusable with AITER shuffle-KV, since forcing LBHNC trips the shuffle decode path. LHBNC and BHLNC now issue one bulk range per head instead of one per layer, which I haven't benchmarked yet, and no open PR touches UMBP layout handling.

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

[mergify]

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

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

/ci run

[github-actions]

✅ Triggered Buildkite CI #98771 for commit f217b63ac798.

[charxwu]

Thanks for the fix, @AndreasKaratzas! I re-ran this PR on MI350X (gfx950) with vLLM built from f217b63 (0.33.0rc1.dev401+gf217b63ac), Qwen3-0.6B. As the baseline I used the same build with umbp/data.py and umbp/runtime/embedded.py taken from 913fccd, so the two runs differ only by this PR. The fix checks out on my side:

• test_every_kv_cache_layout_restores_blocks_exactly: 4 failures (LHBNC/BHLNC × bulk/ranges) without the fix, 12 passed with it, and all 242 UMBP unit tests pass with the fix.
• test_mori_token_e2e.py: with VLLM_KV_CACHE_LAYOUT=BHLNC both tests fail without the fix (restored tokens differ from token 0) and pass with it. With LHBNC, init fails without the fix (KV region bytes (65536) exceed its stride (8192)) and both tests pass with it.

One small thing about the coverage note in the description, unless I'm missing a setup step: I don't think the batch-invariant tests in test_mori_token_e2e.py can reach ROCM_AITER_FA. That means they don't exercise the AITER shuffle-KV (LHBNC) case, even with VLLM_ROCM_USE_AITER=1 VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1.

As far as I can tell:

• The umbp_llm fixture forces VLLM_BATCH_INVARIANT=1 (test_mori_token_e2e.py#L31-L32).
• The selector passes that on as use_batch_invariant (selector.py#L171).
• validate_configuration then marks every backend whose supports_batch_invariance() is false as invalid (backend.py#L345-L346). False is the base default (backend.py#L213-L215), and AiterFlashAttentionBackend doesn't override it (rocm_aiter_fa.py#L886).
• On ROCm, auto-selection then falls back to TRITON_ATTN (rocm.py#L891-L912), and an explicit --attention-backend ROCM_AITER_FA raises (rocm.py#L788-L831).
• LHBNC is only picked inside AITER FA when shuffle is on (rocm_aiter_fa.py#L973-L984), so the shuffle flag ends up having no effect in these tests.

On the build above, running test_mori_token_e2e.py with those two env vars logs the following, once per test:

Found incompatible backend(s) [ROCM_AITER_FA, ROCM_AITER_UNIFIED_ATTN, TURBOQUANT, ROCM_OCTAVE, ULTRAQUANT] with AttentionType.DECODER. Overriding with TRITON_ATTN out of potential backends: ['TRITON_ATTN'].
Using LBNHC KV cache layout.

Both tests pass with and without the fix. Forcing the backend with VLLM_BATCH_INVARIANT=1 gives ValueError: Selected backend AttentionBackendEnum.ROCM_AITER_FA is not valid for this configuration. Reason: ['batch invariance not supported'], while the same call with VLLM_BATCH_INVARIANT=0 starts on ROCM_AITER_FA. So I read the "pass with it under both" part as showing the fix doesn't break the TRITON_ATTN path, rather than as coverage of AITER shuffle-KV.

For AITER shuffle-KV itself, I ran a non-batch-invariant check modeled on your description, with attention_backend="ROCM_AITER_FA", VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1. It used 3 prompts of 192 tokens plus a partial block, and 16 greedy tokens with logprobs=5. The layout resolves to LHBNC there. Without the fix the engine fails at init with KV region bytes (65536) exceed its stride (32768), matching the description. With the fix, all 576 prefix tokens are restored from UMBP, and the output matches a local prefix-cache hit token-for-token with identical logprobs on all 3 prompts. Since AITER FA isn't batch invariant, I compared against a local prefix-cache hit rather than recompute. That way both runs take the same cached-prefix path and differ only in where the KV bytes come from.

[AndreasKaratzas]

Thanks for rerunning this and for the AITER shuffle-KV check! You're right: the fixture forces VLLM_BATCH_INVARIANT=1, which rules out ROCM_AITER_FA, so those e2e tests only exercise BHLNC/LHBNC through TRITON_ATTN. I've corrected that sentence in the description.

[ivanium]

Thanks for the fix!

[ivanium]

/ci run

[github-actions]

✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/98771