[ROCm][Perf] Reuse graph-stable attention metadata across Kimi-K3 cache groups

[JohnQinAMD]

Purpose

Kimi-K3 on ROCm runs many KV-cache groups over one batch: at TP8, six MLA groups and
fourteen KDA groups. Every group currently builds its own attention metadata, but most
of it describes the batch — query offsets, sequence lengths, speculative offsets, the
split-K schedule — and is identical across groups. Only the page indices and state
indices differ, because only the block table differs.

This builds once per step and, for the remaining groups, redoes only the
block-table-dependent part.

Full CUDA graphs bake tensor addresses at capture, so metadata cannot simply be handed
to another group, and buffers two groups will share must be rebound before capture.
Four hooks, declared on AttentionMetadataBuilder with no-op defaults so backends that
don't opt in are unaffected:

• share_reusable_metadata_buffers() — before capture, point compatible groups at one
  allocation for the buffers that are read-only during execution.
• can_reuse_metadata() — per step; anything it rejects falls back to a normal build.
• update_block_table() — redo only page/state indices.
• get_metadata_reuse_key() — extra identity beyond (builder type, kv-cache spec).

AITER MLA and Kimi-K3 KDA implement them.

Test Plan

Unit tests, in the upstream image matching this branch's base:

docker run --rm --entrypoint bash --device /dev/kfd --device /dev/dri \
  --group-add video --ipc=host -v $PWD:/work \
  vllm/vllm-openai-rocm:nightly-d626108b1841888ec90aced33367149a6bbc7e4b -c '
  pip install -q tblib
  SP=$(python3 -c "import vllm,os;print(os.path.dirname(vllm.__file__))")
  for f in models/kimi_k3/amd/kda.py models/kimi_k3/amd/kda_metadata.py \
           models/kimi_k3/nvidia/kda_metadata.py v1/attention/backend.py \
           v1/attention/backends/mla/rocm_aiter_mla.py v1/worker/gpu/attn_utils.py; do
    cp /work/vllm/$f $SP/$f; done
  cd /work && python3 -m pytest -q tests/v1/worker/test_attn_utils.py \
    tests/v1/attention/test_rocm_aiter_mla_mtp_split.py \
    tests/models/kimi_k3/test_kda_metadata.py'

ruff check and ruff format --check (pinned 0.14.0) on the six changed files.

Accuracy on 8x MI355X, Kimi-K3 TP8 with a DSpark draft and fp8 KV. Control and candidate
are the same image with only these files swapped in (same approach as #52356); serve
flags follow #51253 / #50619:

vllm serve /model --served-model-name k3 --tensor-parallel-size 8 --trust-remote-code \
  --kv-cache-dtype fp8 --max-model-len 16384 --max-num-seqs 128 \
  --max-num-batched-tokens 16384 --gpu-memory-utilization 0.93 --block-size 128 \
  --no-enable-prefix-caching --mm-encoder-tp-mode data \
  --compilation-config '{"custom_ops":["+fused_rms_norm_gated"]}' \
  --speculative-config '{"method":"dspark","model":"/draft","num_speculative_tokens":2,"attention_backend":"TRITON_MLA"}'

lm_eval --model local-chat-completions \
  --model_args model=k3,base_url=http://127.0.0.1:8000/v1/chat/completions,num_concurrent=32,tokenized_requests=False \
  --tasks gsm8k --num_fewshot 5 --batch_size 32 --limit 200 \
  --apply_chat_template --gen_kwargs temperature=0

Test Result

Unit tests: 72 passed. Two of them are regression tests that fail without this
change: one covers mamba_cache_mode="none" (the default), where the state column is a
contiguous prefix rather than the align offset; the other covers the full-cudagraph
branch of update_block_table.

Lint: clean on all six files.

gsm8k, 200 questions, 5-shot greedy: 0.995 control, 0.995 candidate, both filters.

Metadata build cost (14 KDA groups, 8 requests, num_speculative_tokens=2),
fourteen builds versus one build plus thirteen reuses:

context	14x build	1 build + 13 reuse	saved
1,024	663.4 us	191.1 us	472.3 us	3.47x
16,384	674.8 us	199.4 us	475.4 us	3.38x
262,144	668.3 us	197.2 us	471.1 us	3.39x

Scope of the claim:

• The saving does not vary with context length. Both halves are launch-bound: MLA's
  page-index expansion is also flat at ~10.3 us per group from 1K to 262K tokens. It
  scales with group count and batch size instead.
• It is a fixed ~0.5 ms per decode step, so its relative value depends on step time. At
  num_speculative_tokens=2 a step is ~10 ms. Under a num_speculative_tokens=7
  draft-verify loop a step is ~163 ms and the same saving is ~0.3%.

No serving throughput number is claimed: a single control/candidate pair on this
hardware is within run-to-run drift, and separating them needs alternating arms with a
paired bootstrap.

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

[mergify]

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

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

[mergify]

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

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