feat(BACKEND-ROCM): select the attention backend in the runner

[tbrasser]

SelectAttentionBackendName had zero production callers on main. Every hit
outside registry.{h,cpp} was a comment, so the whole engine-level attention
registry was dead code. This is the change that reaches it.

GPUModelRunner::initialize_kv_cache now resolves a backend per KV-cache kind
inside the view loop, stores the name, logs it under VT_ATTN_SELECT_LOG, and
validates each group's view against the resolved backend's get_kv_cache_shape.

Per group, not per runner

An earlier revision resolved one backend for the whole runner from a defaulted —
therefore always dense — AttnSelectorConfig. That is wrong for MLA models, and
they do reach this code: runner.cpp:547-548 treats kMlaAttention as the
full-attention group, so DeepSeek-V2/V4, MiniCPM3, GLM4-MoE-Lite and Kimi-K3 would
have resolved FLASH_ATTN instead of TRITON_MLA. Resolution is now lazy and
keyed on the kind, with an mla_layer_mask parallel to fa_dims; MLA validates
the fused 3-dim shape and is tolerant, dense is loud.

Upstream resolves per layer (gpu_model_runner.py:6994-7099, via
layers[layer_name].get_attn_backend() at :7036, with the name chosen in
selector.py:185). Per kind is the coarser mirror that this tree's group
structure supports.

The shape check can now fail

The previous VT_CHECK compared the engine's numbers against an echo of the same
numbers, so it could not fail for any reachable backend. test_attn_backend_registry.cpp
now registers a deliberately mis-shaped scratch backend and asserts the throw,
with three positive controls including FLASH_ATTN checked against the MLA view.

The block-size contract, and why it is a user-visible change

Making get_kv_cache_shape reachable at engine init also makes its
multiple-of-16 requirement reachable. Two shipped paths were relying on that
never being checked: examples/bench/bench_core.h set block_size to an
arbitrary max_prompt + output_len + 4, and server_main.cpp took
--block-size through a bare std::stoi with no validation. Both are fixed —
the bench rounds up, the server rejects at startup with a clear message rather
than throwing during engine init — and the two fixtures that carried
kBlockSize = 8 are now 16.

build-test-cpu-arm64 confirms the two previously red suites
(test_kimi_linear_paged, test_bench) now pass.

Maintainer changes on top

• VT_ATTN_SELECT_LOG added to scripts/env-doc-allowlist.txt beside its
  sibling VT_KV_ALLOC_LOG; check-env-doc refused the change without it.
• The contract is now documented where a user meets it. docs/USAGE.md still
  described --block-size as an unconstrained KV block size, and
  include/vllm.h still told embedders <= 0 => 32 with no constraint — so an
  embedder calling vllm_engine_load with block_size = 8 would newly throw with
  nothing in the header to explain it. Both now state the requirement.
• Merged onto current main rather than cherry-picked; the branch predates main
  enough that replaying its diff conflicts, while the merge is clean.

Owed

CheckKvCacheShape is well tested; its production call site is not. Deleting
runner.cpp's install of it leaves the gate green, because no test drives the
runner with a non-multiple-of-16 block size and none registers a mis-shaped
backend for the device the runner resolves. A case in test_runner.cpp that
builds the runner with kBlockSize = 8 and asserts initialize_kv_cache throws
would close it and double as the executable statement of the contract above.

This edits the runner while .agents/backend-matrix.md:239 records BACKEND-ROCM
as landing with "ZERO selector/model/runner edit". Reaching a dead seam is worth
doing, but it belongs to its own row and issue rather than riding #41.

Issue: #41

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: ClaudeCode:claude-opus-5 [ClaudeCode]

[localai-bot]

Reviewed alongside #1056. This one is more important than the ROCm framing
suggests, and it needs another pass before it can land.

Why it matters more than you may realise. SelectAttentionBackendName has
zero production callers on main today — I checked every hit in src/ and
include/ outside registry.{h,cpp} and all six are comments. The entire
engine-level attention-backend registry is dead code. This PR is the first thing
that reaches it. That is exactly what our "Nothing lands dead" rule exists to
get, so the instinct here is right.

Order: #1056 must land first. These are not stacked — both branch from
0f8580e26 and touch disjoint files — but this one calls the selector
unconditionally as the first statement of initialize_kv_cache, and ROCm's
priority list is empty on main, so ROCm would go from "runs via the hardcoded
NHD layout" to "cannot initialise a KV cache at all". Your PR bodies say this
correctly.

Note first that no CI has run on either PR — fork PRs need maintainer
approval and nobody had granted it. Our fault, being fixed. It matters here
because the first finding is one a build-test-cpu run would have caught in
minutes.

1. There is probably a red test you have not seen.
tests/vllm/models/test_kimi_linear_paged.cpp:250 sets kBlockSize = 8 and
feeds it into GPUModelRunner at :311-312 and :421-426. After this change
initialize_kv_cache calls get_kv_cache_shape(…, 8, …), which throws
std::invalid_argument("Block size must be a multiple of 16.")
(src/vllm/v1/attention/backend.cpp:73-80). You hit exactly this in
test_runner.cpp and fixed it by moving kBlockSize 8 → 16 — worth grepping for
everyone else who constructs a runner before assuming that was the only one.

2. The % 16 enforcement is engine-wide, not ROCm-only, and it breaks a
shipped path. examples/bench/bench_core.h:559 does
params.block_size = seq_budget; where
seq_budget = max_prompt + cfg.output_len + 4 — an arbitrary integer,
essentially never a multiple of 16. The synthetic-model bench will now throw at
engine init on every device, not just ROCm. server_main.cpp:437 also takes
--block-size through a bare std::stoi with no validation, so --block-size 8
goes from working to throwing. Enforcing the constraint may well be right —
upstream does — but then it is its own change with its own note and its own
migration, not a side effect of a ROCm PR.

3. One backend per runner diverges from upstream, and the anchor does not
hold. gpu_model_runner.py:289-293 at our pin is async output-copy code, not
backend resolution. Upstream resolves per layer: initialize_attn_backend
at vllm/v1/worker/gpu_model_runner.py:6994-7099, calling
layers[layer_name].get_attn_backend() at :7036 and grouping into
AttentionGroupKey(attn_backend, layer_kv_cache_spec, num_heads_q), with the
name chosen in selector.py:185 from that layer's own selector config. You
resolve one backend for the whole runner with a defaulted — therefore always
dense — AttnSelectorConfig{}. That is wrong for MLA models, and they do reach
this code: runner.cpp:547-548 treats kMlaAttention as the full-attention
group, so DeepSeek-V2/V4 and MiniCPM3 would resolve FLASH_ATTN rather than
TRITON_MLA. It does not crash today only because of finding 4, and it cannot
be fixed by passing use_mla=true — TritonMLABackend returns a 3-dim shape
your new VT_CHECK would reject outright.

4. The VT_CHECK cannot currently fail. At runner.cpp:948-953 both
FLASH_ATTN and ROCM_ATTN return exactly the NHD tuple the engine just built
from its own fa_dims, so it compares the engine's numbers against an echo of
them. The only new behaviour it actually produces is the % 16 throw in finding
2. The guarantee you describe — a future backend with a different layout failing
loudly — has no test behind it. Worth registering a deliberately mis-shaped
scratch backend and asserting the throw;
test_attn_backend_registry.cpp:433 already has the TEST_ONLY_ATTN idiom to
copy.

5. Small one. include/vllm/v1/worker/gpu/runner.h:227-229 says the name is
"Empty only if no full-attention group exists (a pure-GDN / pooling model caches
no paged KV)" — but you assign it as the first statement of
initialize_kv_cache, before any group scan, so it is never empty, and a
pure-GDN or pooling model that caches no paged KV still pays the selection and
still throws if its platform has no registered backend. Moving the resolution
inside the full_attn_group_id_ >= 0 region makes the comment true and removes
a trap for the next device backend that starts life with an empty priority list
— which is how both Vulkan and ROCm started.

One scope point. .agents/backend-matrix.md:239 records BACKEND-ROCM as
landing with "ZERO selector/model/runner edit", and #1056's own body cites
that contract as the reason for the split. This PR then edits the runner under
the same banner. The work is worth doing — reaching a dead seam is exactly what
we want — but it is a different concern from the ROCm row and deserves its own
row, issue and spec rather than riding #41. Happy to help set that up.

Good news on the mechanics: this one is properly reached —
initialize_kv_cache is production (src/vllm/v1/worker/gpu/runner.cpp:413,454)
and deleting your block in a scratch copy reds test_runner.cpp:454. Your
commits and PR body both pass the trailer and style checkers cleanly.

[localai-bot]

CI has now run on this for the first time, and it confirms both predictions from
my review — so you have concrete evidence rather than my inference:

99% tests passed, 2 tests failed out of 495

	154 - test_kimi_linear_paged (Failed)
	341 - test_bench (Failed)

Block size must be a multiple of 16 appears 10 times in the log.

• test_kimi_linear_paged is finding 1: kBlockSize = 8 at
  tests/vllm/models/test_kimi_linear_paged.cpp:250, fed into GPUModelRunner
  at :311-312 and :421-426. Same shape as the one you already fixed in
  test_runner.cpp by moving 8 → 16.
• test_bench is finding 2, and it is the one worth pausing on:
  examples/bench/bench_core.h:559 sets params.block_size = seq_budget where
  seq_budget = max_prompt + output_len + 4. That is a shipped path on every
  device, not a test fixture and not ROCm-specific. It regressed because this
  PR makes a previously-unreached constraint reachable.

Both are green on main, so these are attributable to the change rather than
pre-existing.

That second one is the reason I would rather see the % 16 enforcement land as
its own change with its own note. server_main.cpp:437 also takes
--block-size through a bare std::stoi with no validation, so --block-size 8
goes from working to throwing for anyone running the server today. Upstream does
enforce the constraint, so enforcing it is probably right — it just deserves to
be a deliberate, announced change rather than a side effect.

Two notes on the CI itself, neither your fault:

• This is the first run either of your PRs has ever had. Fork-PR workflows
  here queue for maintainer approval, and roughly 166 runs were sitting unapproved
  across the repo — yours among them, and going back weeks for other contributors
  too. I approved yours. You have been working without gate feedback this whole
  time, which is on us.
• windows-msvc-cpu / windows-msvc-vulkan are red here but not yours. That
  lane had a compile queue that was cleared only today (windows-msvc-cpu/vulkan red on every PR based on c7cb59fbb: C4244 double->float in ltx2_video.cpp's StreamState<->Ltx2LatentState positions copy (#964) #968 landed as fix(ltx2): cast positions double/float to avoid MSVC C4244 (#968) #983,
  main does not compile under MSVC: #1054 removed a lambda capture that was not redundant #1068 as fix(qwen3.5): the capture #1054 dropped was not redundant, and MSVC says so #1069); anything based before that still sees the old break, and even
  after it the lane fails at test_openai_api_server crashes on Windows with STATUS_STACK_BUFFER_OVERRUN (0xC0000409), unmasked by the #512 fix #584's runtime crash. Ignore those two.

Six checks are still running; I will flag anything else that turns up.

[tbrasser]

Thanks — all five findings are addressed in the new commits, and the CI results you posted were the confirmation we needed.

1. test_kimi_linear_paged — fixture kBlockSize 8 → 16 (the vLLM gate block size), same shape as the test_runner fix. I also swept for other runner constructions with non-%16 block sizes; the only other odd one is test_kv_cache_coordinator (kBlockSize = 2), which never builds a GPUModelRunner, so it's unaffected.

2. The %16 contract — made deliberate at both shipped entry points. bench_core.h now rounds the synthetic unified block up to a multiple of 16 ((seq_budget + 15) / 16 * 16), and server_main --block-size validates (positive, %16) with a clear error instead of a bare stoi. Both are called out in the PR body as an announced change — the enforcement itself was already in FLASH_ATTN's get_kv_cache_shape on main; this PR just makes it reachable, and these two fixes keep every shipped path feeding a valid block size. If you'd rather have it as a standalone PR with its own migration note, say the word and I'll split it.

3. Per-group resolution, not per runner. Resolution now happens per attention group, lazily per kind, inside the full_attn_group_id_ >= 0 region. Dense groups resolve loudly (a platform with no registered dense backend fails at init — only for models that actually need one). MLA groups resolve TRITON_MLA on CUDA — its 3-dim get_kv_cache_shape is exactly the fused (num_blocks, block_size, head_size) cache deepseek_v2.cpp views (:576-578) — and stay op-driven on devices with no registered MLA backend (CPU, ROCm today), because the engine's MLA execution (TritonMLAImpl) is not registry-gated and a loud throw there would regress working MLA paths. That's the piece "cannot be fixed by passing use_mla=true": on CUDA it can be fixed properly, and on CPU/ROCm the honest answer is an empty name + recorded op-driven execution, not a throw.

4. The check is now real and tested. It moved to vllm::v1::CheckKvCacheShape (registry.h/cpp) with the per-group expected view — NHD 5-dim for dense, fused MLA 3-dim for MLA. A new registry test registers a deliberately mis-shaped scratch backend (upstream's K/V-outermost shape) and asserts the throw on both views, plus positive controls: FLASH_ATTN passes the dense view, TRITON_MLA passes the MLA view, and FLASH_ATTN against the wrong expected view throws — so the comparison is not an echo of the engine's own numbers.

5. Comment/guard fixed. Resolution moved inside the guarded region; the runner.h "empty only if no full-attention group exists" comment is now true, and a pure-GDN / pooling model pays no selection.

Scope — agreed, and it's now called out in the body: the runner work is a different concern from the ROCm row and is proposed to move to its own row/issue/spec (BACKEND-ATTN-SELECTION-RUNNER, spec rocm-attn-backend.md §7). Happy to set that up with you. Merge order stays #1056 → #1065.

Verification (CPU tier, clean -Werror build): test_attn_backend_registry 17/17 (61 assertions, incl. the new mis-shaped case), test_runner 19/19 (543), test_kimi_linear_paged 8/8 (206), test_bench 11/11 (80), test_llm_engine 24/24 (493), plus test_prepare_inputs / test_mla_attention_block green — the two tests that were red in your CI run are green locally. The gfx1151 e2e evidence from the first pass stands (selection resolves ROCM_ATTN, NHD geometry validated, identical generation); I can re-run the battery on the GPU with VT_ATTN_SELECT_LOG=1 if you want the per-group log shape confirmed on silicon.