[ROCm][Spec Decode] Add Aiter MLA decode support non-causal draft block

[JohnQinAMD]

Purpose

On ROCm a Kimi-K3 DSpark draft has to be pinned to TritonMLA:

--speculative-config '{"method":"dspark","model":"/draft",
                       "num_speculative_tokens":7,
                       "attention_backend":"TRITON_MLA"}'

The draft's query block is non-causal. On ROCm only TritonMLA sets supports_non_causal_multi_token_decode, so the pin.

TritonMLA and FlashInferMLA serve such a block by flattening it: one decode row per query token, every row reading the same committed KV. That is cheap for them because their kernels take a block table, so flattening is a repeat_interleave. AiterMLA's kernel takes a flat page-size-1 index list instead. Flattening there means rebuilding that list every forward -- roughly half a million entries at batch 8, block 8 and an 8k context.

Fix: don't flatten. Pass the mask, which the kernel already accepts.

get_mla_metadata_v1 has always taken is_causal. aiter's mla_decode_fwd takes causal as of v0.1.20. vLLM hardcoded the first and never passed the second:

get_mla_metadata_v1(..., self._num_attention_heads, 1, True, ...)  # <- the mask

Both now get the real value. Same one row per request, same index list.

Three pieces make that safe.

Selection. AiterMLA offers a non-causal block only if the installed mla_decode_fwd takes causal. Probed by signature, the same way this file already probes for fp8 scale support.

Routing. Gluon's mask is causal and cannot be turned off, so use_gluon_verify now returns False for any non-causal block. Such a block takes the padded asm decode instead. Padding the head count up is what gives it a kernel there -- the fp8 path already relies on that. This is also what lets a bf16 cache serve a non-causal block at all. Builder and impl call the same predicate, so they cannot disagree.

Build. One shape still reaches _build_decode and it raises: an fp8 2-token block. aiter's dispatch has no fold for that length, so the lookup lands on a causal-only kernel, and a missing kernel aborts the process rather than raising. ROCm/aiter#4872 adds the missing entry; once the pinned aiter carries it, this guard and its test can go. The bf16 fold has no such hole.

Outside ROCm, _build_decode gains a causal kwarg in MLACommonMetadataBuilder. The FlashAttn, FlashInfer, FlashMLA and Dots3Note overrides are signature-only updates. DCP is untouched; supports_non_causal_multi_token_dcp stays False.

vLLM pins aiter v0.1.19 (docker/Dockerfile.rocm_base). That build has no causal argument, so on the default image this PR is a no-op: supports_non_causal() returns False, selection keeps handing the draft group to TritonMLA, and the recipe keeps its pin. It starts working once the pin moves.

Test Plan

pytest tests/v1/attention/test_rocm_aiter_mla_non_causal.py    # 15, CPU only

They stub aiter.mla, so both aiter generations can be exercised on one machine. The two _build_decode guards are covered in test_rocm_aiter_mla_mtp_split.py. That file already drives the method directly.

End to end on 8x MI355X, /model = Kimi-K3 and /draft its DSpark draft. Same server twice; the only difference is whether the draft is pinned:

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 \
  --load-format auto --mm-encoder-tp-mode data \
  --safetensors-load-strategy prefetch \
  --compilation-config '{"custom_ops":["+fused_rms_norm_gated"]}' \
  --speculative-config "$SPEC" --port 8000

# A -- what the recipe does today
SPEC='{"method":"dspark","model":"/draft","num_speculative_tokens":7,"attention_backend":"TRITON_MLA"}'
# B -- this PR: no pin, the draft picks AiterMLA
SPEC='{"method":"dspark","model":"/draft","num_speculative_tokens":7}'

Drop --kv-cache-dtype fp8 and run again for the bf16 numbers. That run also needs #50619, for an unrelated main-branch cudagraph defect in the target's causal verify.

Both were served against a v0.1.19 install with v0.1.20's MLA files copied in. That install reproduces a full v0.1.20 tree's kernel numbers bit for bit.

for C in 1 8; do
  vllm bench serve --backend openai-chat --endpoint /v1/chat/completions \
    --model k3 --tokenizer /model --trust-remote-code --port 8000 \
    --dataset-name random --random-input-len 1024 --random-output-len 128 \
    --ignore-eos --max-concurrency $C --num-prompts $((C == 1 ? 16 : 48))
done

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 \
  --apply_chat_template --gen_kwargs temperature=0

Test Result

Unit tests: 15 + 46 passed.

Both runs land on FULL_AND_PIECEWISE with no cudagraph downgrade. Unpinning drops the second attention backend. So the gap below is the draft's backend and nothing else.

draft backend	conc.	output throughput	mean TPOT
TritonMLA (pinned)	1	68.44 tok/s	13.23 ms
AiterMLA (this PR)	1	76.32 tok/s	11.67 ms
TritonMLA (pinned)	8	294.57 tok/s	23.32 ms
AiterMLA (this PR)	8	324.44 tok/s	21.26 ms

That is 11.5% at concurrency 1 and 10.1% at 8.

A bf16 cache reaches 84.16 tok/s at concurrency 1 and 331.66 at 8, TPOT 10.45 ms and 20.76 ms.

Accuracy, full GSM8K test set (1319 questions), 5-shot, greedy, fp8 KV:

	flexible-extract	strict-match
TritonMLA (pinned), fp8	0.9644 ± 0.0051	0.9636 ± 0.0052
AiterMLA (this PR), fp8	0.9613 ± 0.0053	0.9606 ± 0.0054
AiterMLA (this PR), bf16	0.9644 ± 0.0051	0.9636 ± 0.0052

At the kernel, checked with ad-hoc scripts rather than in-tree tests:

With an fp8 cache and 16 heads, causal=False lands 7.1e-03 from a non-causal reference. Against a causal one it sits 6.3e-02 away at block 4 and 9.1e-02 at block 8. So it runs the mask that was asked for. Four identical query rows come back identical, spread 0.0.

With a bf16 cache the padded asm decode holds 3.5e-04 to 5.6e-04. That is across 8, 12, 16, 24 and 32 heads, and block lengths 1 through 16. fp8 covers 1 and 3 through 16.

---

Written with AI assistance (Claude Code); the commit carries Co-authored-by: Claude. I reviewed every changed line and ran the tests above myself.

[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

[ppalanga]

This branch is still conflicting with main, and I cannot push to JohnQinAMD/vllm-amd.

Rebased copy (current main): #55966

[JohnQinAMD]

close and replaced by #55966