[Misc][Platform] Check aligned KV block sizes against every attention backend

[divyvasal]

Purpose

Follow-up to #49845. Phase 1 there picks a KV cache block size every attention
backend accepts. The hybrid and heterogeneous-KV alignment steps that run
afterwards still round by the first backend's kernel block only, as
@Sanjays2402 pointed out in review of #49845, so they can land on a size a
sibling backend rejects. Today that surfaces late, as the opaque
select_common_block_size() error at worker setup.

This PR does not change how alignment rounds. It validates the result: if
alignment changed the block size, every backend's supports_block_size() must
accept it, otherwise it raises immediately and names the rejecting backends. I
don't know of a model configuration that hits this today, so this is hardening
rather than a fix for an observed failure.

Why not round by every backend

An earlier revision rounded by the LCM of each backend's smallest kernel block.
That regresses valid configurations: with a backend supporting [32, 48] and an
exact-size [48] backend, phase 1 correctly picks 48, and the LCM rounding
bumps it to 96, which the exact-size backend rejects. Choosing an aligned size
every backend accepts would need a candidate search inside the mamba and indexer
alignment logic, which is a larger change. Thanks @yewentao256 and
@jayzuccarelli for working through this.

The check is per backend rather than a common-size check across all backends.
Kernel sizes are chosen per KV cache group at runtime, so a global check can
reject configurations that work, and it would bypass backends' exact-size
overrides of supports_block_size().

Changes

• update_block_size_for_backend() records the block size before alignment and
  calls a new _check_aligned_block_size() afterwards.
• CpuPlatform.update_block_size_for_backend() overrides the base method and
  aligns directly, so it calls the same check.
• XPUPlatform.update_block_size_for_backend() rounds the block size up to a
  multiple of 64 for GDN after super() has run the check, so it checks that
  rounding too.
• Alignment itself and phase 1 are unchanged.

Test Plan

Added to tests/v1/worker/test_gpu_model_runner.py:

• test_alignment_rejected_by_a_sibling_backend_raises, run for both the base
  and CPU platforms: an alignment that lands on 32 beside an exact-size [16]
  backend raises, naming that backend.
• test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises: XPU's GDN
  rounding from 16 to 64 beside an exact-size [16] backend raises. It skips
  without vllm_xpu_kernels, like the existing XPU tests.

Test Result

Ran the three tests locally in a CPU environment with the XPU kernel modules
stubbed: all pass, and the XPU test fails with the XPU check removed. Full
pre-commit is clean.

🤖 Generated with Claude Code

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

[jayzuccarelli]

I worked on this path in #48313, so I read it against the phase 1 code from #49845. Two things:

1. _check_aligned_block_size validates with each backend's supports_block_size, which only asks whether the size is a multiple of one of that backend's kernel sizes. The failure it pre-empts comes from select_common_block_size, which needs one kernel size every backend accepts that also divides the manager block. Those can disagree: with exact sizes [128] and [64], 128 passes the new check and select_common_block_size(128, ...) still raises No common block size for 128. Validating with select_common_block_size itself (re-raising with the backend list) would make the early and late errors agree. Phase 1's _preferred_block_size_for_backends has the same gap, so the helper could be shared.

2. In _align_heterogeneous_kv_block_size, smallest_kernel_block now feeds both block_alignment and required_page. The cross-backend LCM makes sense for the rounding step, but required_page is the padded spec's page at its own smallest kernel block, which comes from the primary backend. With FLASHINFER_MLA_SPARSE ([32, 64]) plus the CUDA indexer ([64]) the LCM is 64 instead of 32, so required_page doubles and the block size goes 64 to 128 at a 2x padded/primary ratio, and 128 to 256 at ~3.6x (nvfp4 primary, bf16 skip layers). Nothing fails, since select_common_block_size still picks 64, but it changes behavior for DSA models with skip layers. Keeping required_page on backend_cls's own smallest block and using the granularity only for block_alignment would avoid that.

[divyvasal]

@jayzuccarelli Thanks, both are right and both are now addressed in 3b737ef.

On (2): you are correct that this was a behavior change I introduced. The padded spec is packed against the primary backend alone, as the function's own docstring says, so widening its page to the cross-backend LCM inflated it exactly as you describe. required_page now uses backend_classes[0]'s own smallest kernel block again, and the granularity is used only for block_alignment. Your FLASHINFER_MLA_SPARSE plus indexer example is what made it obvious.

On (1): also correct, and the mismatch is sharper than I realized. supports_block_size reads a fixed kernel size as "any multiple of it", and says so in a comment, while select_common_block_size requires exact equality for a plain int. So the two disagree precisely on exact-size backends. _check_aligned_block_size now validates through a shared _block_size_is_usable helper that calls select_common_block_size and re-raises with the backend list, so the early and late errors agree. Added a test for your [128] plus [64] case, which passes the permissive predicate and fails the real one.

On sharing it with phase 1: I agree the same gap is there, but I left _preferred_block_size_for_backends alone in this PR on purpose. Switching it changes merged behavior rather than just tightening new code: with exact [32] and [48] it currently returns 96, and under select_common_block_size it would raise instead. That also flips the ([[32], [48]], 96) case in the existing test. I think that is the correct end state, but it deserves its own PR and its own discussion rather than riding along in a hardening change. Happy to open it if you agree with the direction.

For transparency: I cannot run vLLM's test suite in my environment, so CI is the check on the new tests. Lint is clean at the pinned ruff version.

[jayzuccarelli]

Thanks for turning these around so fast. 3b737ef reads right to me on (2).

One correction to my own point (1), which I should have caught: at runtime select_common_block_size isn't called across all backends. prepare_kernel_block_sizes (v1/worker/utils.py) calls it per KV cache group with only that group's backends, so backends in different groups can each pick their own kernel size. _block_size_is_usable checks all non-SSM backends together, which is stricter than what actually runs: exact [32] and [48] in different groups can work at 96 (kernel 32 in one group, 48 in the other), but the new check rejects it. Grouping isn't known at update_block_size_for_backend time, so neither predicate is exact: supports_block_size can miss a same-group mismatch, and the global check can reject a cross-group config that works. Today's backends always share a size with their siblings (64 for sparse MLA plus the indexer), so neither case is live, but a false raise blocks a working model while a missed one only delays an error. I'd lean toward the permissive check here, or a warning instead of a raise, unless the grouping can be known this early.

Same reason on phase 1: I'd hold off on switching it. Your [32]/[48] at 96 case is exactly the one that can work across groups.

[yewentao256]

Thanks for the work!

Please do the following:

• shrink the diff, idealy <100 LOC

Also take a look at these AI generated comments:

I think `_check_aligned_block_size()` is validating at the wrong abstraction level.

The new code calls `select_common_block_size(block_size, backend_classes)` across all non-SSM backends globally, but at runtime kernel block sizes are selected per KV cache group. Different groups can therefore legitimately use different physical kernel block sizes.

For example, if one group uses a backend that supports only 32 and another uses one that supports only 48, a manager block size of 96 is valid: the two groups can use kernel block sizes 32 and 48 respectively. A global `select_common_block_size(96, [backend32, backend48])` would reject this because there is no single kernel block size common to both backends.

Conversely, `select_common_block_size()` only looks at `get_supported_kernel_block_sizes()` and bypasses backend-specific `supports_block_size()` overrides. So a backend that explicitly requires the manager block itself to be exactly 16 could incorrectly accept a manager block of 64 simply because 16 divides 64.

It seems safer for this early validation step to keep checking `supports_block_size(block_size)` independently for each backend, and leave common physical kernel-size selection to `prepare_kernel_block_sizes()`, where the actual KV cache grouping is known.

The new exact-backend tests also seem inconsistent with the current implementation.

`_mock_backend(..., exact=True)` enforces exactness through `supports_block_size()`, but `_check_aligned_block_size()` no longer calls that method. For example, an exact `[16]` backend with manager block size 64 would still pass `select_common_block_size()` because 16 divides 64, so the expected `ValueError` should not be raised with the current code.

I think these tests are actually highlighting why the early check should continue to respect `supports_block_size()`.

[divyvasal]

@yewentao256 @jayzuccarelli Thanks, both of you. You're right, and my last push was wrong in two ways. Checking with select_common_block_size across all backends rejects configurations that work once kernel sizes are picked per KV cache group, and it bypasses backends' exact-size supports_block_size() overrides. That second part also meant the exact-size tests in that push expected an error the code no longer raised.

I've rebuilt it smaller, now +72/-11:

• Alignment rounds by the LCM of every backend's smallest kernel block.
• If alignment changed the block size, each backend's supports_block_size() must accept it, otherwise it raises early.
• The heterogeneous path's padded page is still sized from the primary backend, keeping the fix from 3b737ef.
• Phase 1 is untouched.

[jayzuccarelli]

a6ab16d reads right to me. The padded page is still sized from the primary backend, and removing the global select_common_block_size check resolves the per-group point. Two non-blocking notes:

1. CpuPlatform.update_block_size_for_backend overrides the base method and now aligns with the LCM across all backends, but doesn't run the new supports_block_size rejection, so CPU keeps the late failure. Might be worth sharing the check or leaving a note.
2. The early check is per backend, so two backends in the same KV cache group can each accept the aligned size without sharing a kernel size (e.g. MultipleOf(16) and exact [24] at 48). That still fails as late as it does on main, so it's not a regression. [Misc] Name each backend and its kernel block sizes in block-size errors #58557 makes that late error name each backend and its kernel sizes, so it at least reads clearly.

[divyvasal]

@jayzuccarelli Thanks. On (1), agreed: the check is now a shared _check_aligned_block_size that both the base method and CpuPlatform.update_block_size_for_backend call, and the rejection test runs against both. On (2), agreed it's out of scope here since that same-group case fails no later than on main, and #58557 naming each backend and its sizes is the right fix for it. Now +88/-11.

[yewentao256]

Thanks for the work!

[yewentao256]

Please take a look at this AI generated comment:

Is taking the minimum advertised size from each backend sufficient here? `get_supported_kernel_block_sizes()` represents alternative supported sizes, and `supports_block_size()` may also implement exact-size behavior.

For example, if backend A supports `[32, 48]` and backend B supports exactly `[48]`, Phase 1 can validly choose `48`. This helper would compute `lcm(32, 48) = 96`, potentially bumping an already-valid block size to `96`, which B then rejects.

Could we avoid collapsing each backend's supported set to its minimum, and instead choose an aligned candidate that is explicitly accepted by all backends?

[divyvasal]

@yewentao256 Thanks, the AI comment is right, and it's a regression my LCM rounding introduced. With [32, 48] and an exact [48] backend, phase 1 correctly picks 48, and the LCM of each backend's minimum bumps it to 96, which the exact backend rejects. Choosing a size every backend accepts would need a candidate search inside the mamba and indexer alignment, so instead I've dropped the rounding change entirely. Alignment is now identical to main, and the PR only validates its result: if alignment changed the block size, every backend's supports_block_size() must accept it, otherwise it raises early and names the rejecting backends. Now +59/-0.

[yewentao256]

Shall we support xpu as well?

[divyvasal]

Yes, done in d333dfd. XPUPlatform.update_block_size_for_backend calls super() first, so the check ran before its GDN step rounds the block size up to a multiple of 64. That rounding was unchecked, so it now runs the same check on its result. Added test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises, which skips without vllm_xpu_kernels, like the existing XPU tests. Also rebased onto current main. Still +99/-0.

[yewentao256]

LGTM, thanks for the work!

[yewentao256]

/ci run

[github-actions]

❌ This PR is 150 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[github-actions]

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

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /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.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #92487 for commit be424011f8f9.