[Kimi-K3][AMD] Return KDA and MLA projection outputs directly

[LiuYinfeng01]

Purpose

Kimi-K3's AMD attention paths copied each output projection into a caller-owned buffer even though the projection already returns a tensor with the required shape, dtype, layout, and lifetime.

This PR removes that redundant post-projection allocation and copy from both attention families:

• 69 KDA layers
• 24 MLA layers
• 93 post-attention-projection copies per model step in total

The functionality from #50847 is now folded into this PR. #50654 is separate and complementary: it optimizes the earlier KDA recurrent/conv/RMSNorm path, while this PR removes the later copy after the output projection.

Implementation

• Return KDA output-projection storage directly from both the shared KDA implementation and the dedicated full-rank Kimi-K3 implementation selected by current main.
• Return MLA output-projection storage directly.
• Use one return-value contract in KimiDecoderLayer for KDA and MLA.
• Preserve attention arithmetic, state updates, normalization, projection, tensor-parallel communication, and HIP-graph replay behavior.

Validation

Targeted ownership tests:  3 passed
Ruff 0.14.0 check:         passed
Ruff 0.14.0 format check:  passed
git diff --check:          passed
DCO:                       passed

The tests cover KDA dispatch ownership, the concrete full-rank Kimi-K3 KDA forward path, and MLA dispatch ownership. In each case, the returned tensor is the exact projection tensor and has the same data_ptr().

Fixed-cohort decode A/B

Setup

GPU:                 8 x AMD Instinct MI355X (gfx950)
Parallelism:         TP8
Model:               Kimi-K3, BF16 activations / MXFP4 weights
Decode:              non-speculative
Prefix caching:      enabled
Prompt:              63,911-token shared prefix + 4,089-token suffix
Output:              1,024 tokens per request, ignore EOS
Concurrency:         C16 and C24 fixed cohorts
HIP graph:           FULL, capture sizes 1/16/24
Sampling:            greedy, seed 42
Runs:                1 warmup + 3 measured runs per concurrency and variant
Success rate:        240/240 measured requests

Both variants used image localhost/kimi-pp8-upstream-clean (vLLM beca88e59, AITER 99733dc00). The baseline was the unmodified image. The candidate mounted only this PR's AMD runtime files from 43d81ba3e. The relevant AMD Kimi-K3 files in the image are identical to the rebased main baseline before this PR.

Results

Means over three measured runs. TPOT and ITL are end-to-end serving metrics, not the reviewer's pure-GPU trace period.

Cohort	Metric	Baseline	Combined PR	Delta
C16	Mean TPOT	73.141 ms	72.428 ms	-0.713 ms (-0.97%)
C16	Median ITL	70.488 ms	70.186 ms	-0.302 ms
C16	Output throughput	193.575 tok/s	195.555 tok/s	+1.02%
C24	Mean TPOT	74.958 ms	74.457 ms	-0.501 ms (-0.67%)
C24	Median ITL	71.311 ms	70.916 ms	-0.396 ms
C24	Output throughput	278.245 tok/s	280.452 tok/s	+0.79%

Mean-TPOT standard deviation across the three runs was 0.184 ms / 0.348 ms for baseline C16/C24 and 0.019 ms / 0.122 ms for candidate C16/C24.

Interpretation

The reviewer's trace found that the complete 93-copy group costs approximately 0.495 ms at C16 and 0.563 ms at C24. The combined A/B removes all 93 copies and measures a 0.713 ms C16 and 0.501 ms C24 TPOT reduction. C24 closely matches the trace-derived removable cost. C16 is modestly larger, which is plausible because returning projection storage directly also changes allocator and lifetime behavior; the baseline's run-to-run variance is also higher.

The absolute TPOT values should not be compared directly with the reviewer's 32.82/35.98 ms pure-GPU periods because TPOT includes serving and scheduling overhead. The A/B delta is the relevant same-harness comparison.

Scope summary

#50592: remove 69 KDA + 24 MLA post-projection copies
#50847: functionality absorbed by this PR
#50654: separate earlier KDA recurrent-output/fusion optimization

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

[LiuYinfeng01]

Hi @tjtanaa,

Could you please take a look at this PR when you have a chance? Your feedback would be greatly appreciated.

Thank you!

[LiuYinfeng01]

@hongxiayang
Could you help review this PR? Thank you !

I completed a full Kimi-K3 GSM8K A/B for PR #50592 on 8x MI355X using
the same gfx950_kimi_k3_20260727 image, weights, TP8 configuration,
AITER environment variables, prompts, and request concurrency for both
variants.

Neither baseline nor the KDA direct-return variant produced garbled text.
Deterministic raw-completion and chat probes had zero UTF-8 replacement
characters, control characters, or empty responses. The full logged
1,319-sample runs also had zero garbled responses.

Results:

vLLM tests/evals/gsm8k/gsm8k_eval.py, 1,319 questions, 5-shot,
temperature 0, seed 42, max tokens 256:

baseline: 93.8590% (1,238/1,319), invalid rate 0.0758%

KDA direct return: 95.5269% (1,260/1,319), invalid rate 0.0758%

delta: +1.6679 percentage points

lm_eval 0.4.12, gsm8k, 5-shot,
exact_match,flexible-extract:

baseline: 96.7400% (1,276/1,319)

KDA direct return: 96.0576% (1,267/1,319)

delta: -0.6823 percentage points

The paired lm_eval breakdown was 1,259 both correct, 17 baseline-only
correct, 8 optimized-only correct, and 35 both wrong. The exact paired
McNemar p-value is 0.1078, so the observed difference is not statistically
significant at p < 0.05. The independent vLLM evaluator moved in the opposite
direction. There were no failed requests or increases in malformed output.

Environment qualification: PR #50592 head is
bbe70a0dd864d756322ab6e3c7ab30f9f70568e9 and its direct parent is
94e9ef0. The fixed Docker image contains
the older vLLM g5f76ae224, so mounting the complete current parent/head source
would mix incompatible internal APIs with the image binaries. For this
Docker-specific A/B, baseline was the unmodified image and the optimized
variant mounted only this PR's two runtime hunks onto the image-matching source:
kimi_gdn_linear_attn.py (+2/-3) and AMD linear.py (+11/-4).

All aggregate JSON, per-sample logs, raw smoke responses, server logs,
container inspect output, paired comparison data, and reproduction scripts are
saved under /home/yinfeliu/kimi-k3-kda-gsm8k-ab/.

[LiuYinfeng01]

@hongxiayang
The recipe is :

#!/usr/bin/env bash
set -euo pipefail

VARIANT=${1:?usage: run_server.sh baseline|optimized}
PODMAN=(podman --root /data/podman-k3-storage --runroot /run/podman-k3)
IMAGE=docker.io/rocm/vllm-dev:gfx950_kimi_k3_20260727
MODEL_HOST=/data/models/Kimi-K3
ROOT=/home/yinfeliu/kimi-k3-kda-gsm8k-ab
NAME="k3-kda-gsm8k-${VARIANT}"

case "$VARIANT" in
baseline)
overlay_args=()
;;
optimized)
overlay_args=(
-v "$ROOT/image-overlay-optimized/kimi_gdn_linear_attn.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py:ro"
-v "$ROOT/image-overlay-optimized/linear.py:/usr/local/lib/python3.12/dist-packages/vllm/models/kimi_k3/amd/linear.py:ro"
)
;;
*)
echo "variant must be baseline or optimized" >&2
exit 2
;;
esac

"${PODMAN[@]}" rm -f "$NAME" 2>/dev/null || true
"${PODMAN[@]}" run -d
--name "$NAME"
--ipc=host
--network=host
--device=/dev/kfd
--device=/dev/dri
--security-opt seccomp=unconfined
--group-add video
-v "$MODEL_HOST:/data/Kimi-K3/Kimi-K3:ro"
-v "$ROOT/results/$VARIANT:/results"
"${overlay_args[@]}"
-e VLLM_ROCM_USE_AITER=1
-e VLLM_ROCM_USE_AITER_MOE=1
-e VLLM_ROCM_USE_AITER_MLA=1
-e AITER_ROCM_ARCH=gfx950
-e AITER_SITUV2_A8W4=1
-e AITER_BF16_FP8_MOE_BOUND=0
-e SAFETENSORS_FAST_GPU=1
--entrypoint /bin/bash
"$IMAGE"
-lc 'exec vllm serve /data/Kimi-K3/Kimi-K3
--trust-remote-code
--host 127.0.0.1
--port 8005
--gpu-memory-utilization 0.95
--tensor-parallel-size 8
--mm-encoder-tp-mode data
--max-model-len 16384
--max-num-batched-tokens 8192
--max-num-seqs 64
--no-enable-prefix-caching'

[Fangzhou-Ai]

Thanks for this optimization. I mapped this PR's copy site against a new fixed-cohort pure-decode trace so we have a concurrency-relevant estimate in addition to the batch-one result in the description.

Profile setup: 8 x MI355X (gfx950), TP8, BF16/MXFP4, non-speculative decode, prefix caching, 63,911-token shared prefix + 4,089-token suffix, full HIP-graph replay, and +fused_rms_norm_gated. The baseline was vLLM 25e1f81e1b8def93d7228a38b83348880babbe7b with AITER 61549c5c2623316d24ce54dcf35438173cb4d95b. The fixed-cohort pure-GPU step periods were 32.82 ms at C16 and 35.98 ms at C24.

The trace contains 93 post-attention-projection __amd_rocclr_copyBuffer launches per step: 69 KDA copies targeted by this PR and 24 MLA copies targeted by #50847. The complete 93-copy group costs about 0.495 ms at C16 and 0.563 ms at C24. Since all copies have the same [batch, 7168] BF16 shape, the 69-copy KDA share is approximately:

• C16: 0.37 ms/step, about 1.1% of the pure-GPU period
• C24: 0.42 ms/step, about 1.2% of the pure-GPU period

I would therefore expect an isolated fixed-cohort gain around 0.35-0.42 ms/step for this PR at C16/C24. This is a trace-derived removable-cost estimate, not a candidate A/B result; allocator/lifetime effects can make end-to-end behavior differ.

For scope clarity, the 24 MLA projection copies are covered by #50847, while the separate 69 KDA recurrent-output placement copies are covered on the gfx950 pure-decode path by #50654. Those costs should not be attributed to this PR.

[mergify]

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

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

[LiuYinfeng01]

Thanks for this optimization. I mapped this PR's copy site against a new fixed-cohort pure-decode trace so we have a concurrency-relevant estimate in addition to the batch-one result in the description.

Profile setup: 8 x MI355X (gfx950), TP8, BF16/MXFP4, non-speculative decode, prefix caching, 63,911-token shared prefix + 4,089-token suffix, full HIP-graph replay, and +fused_rms_norm_gated. The baseline was vLLM 25e1f81e1b8def93d7228a38b83348880babbe7b with AITER 61549c5c2623316d24ce54dcf35438173cb4d95b. The fixed-cohort pure-GPU step periods were 32.82 ms at C16 and 35.98 ms at C24.

The trace contains 93 post-attention-projection __amd_rocclr_copyBuffer launches per step: 69 KDA copies targeted by this PR and 24 MLA copies targeted by #50847. The complete 93-copy group costs about 0.495 ms at C16 and 0.563 ms at C24. Since all copies have the same [batch, 7168] BF16 shape, the 69-copy KDA share is approximately:

• C16: 0.37 ms/step, about 1.1% of the pure-GPU period
• C24: 0.42 ms/step, about 1.2% of the pure-GPU period

I would therefore expect an isolated fixed-cohort gain around 0.35-0.42 ms/step for this PR at C16/C24. This is a trace-derived removable-cost estimate, not a candidate A/B result; allocator/lifetime effects can make end-to-end behavior differ.

For scope clarity, the 24 MLA projection copies are covered by #50847, while the separate 69 KDA recurrent-output placement copies are covered on the gfx950 pure-decode path by #50654. Those costs should not be attributed to this PR.

Hi. @Fangzhou-Ai Thanks for clarifying the scope. I will fold the functionality of #50847 into #50592, so a single PR will remove both the KDA and MLA post-projection copies.

This does not overlap with #50654: #50654 removes the earlier KDA recurrent-output placement copy and fuses the decode operations, while #50592 removes the separate copy after the output projection. The two optimizations are complementary and can be applied together.

[github-actions]

✅ Triggered Buildkite CI #83945 for commit 6a5cde61b9cb.

[LiuYinfeng01]

/ci retry

[github-actions]

❌ @LiuYinfeng01, A reviewer with write access must run /ci run, approve the PR, or add the ready label first.

[Fangzhou-Ai]

@AndreasKaratzas seems this PR is ready to be merged?

[AndreasKaratzas]

/ci run

[AndreasKaratzas]

b200 batch invariance seems to be broken. I merged main to re eval it cause it was not failing in latest nightly.

[mergify]

Hi @LiuYinfeng01, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[AndreasKaratzas]

/ci retry

[github-actions]

✅ Triggered Buildkite CI #84061 for commit 6edede507e8a, running 1 failed step(s) from Buildkite CI #83945.

[AndreasKaratzas]

This also affects NVIDIA’s low-rank Kimi-Linear path. vllm/models/kimi_k3/nvidia/model.py imports this shared class when use_full_rank_gate=False, but still sets _self_attn_writes_output = True and calls it with output=output. Since this forward() no longer accepts output, that path will fail with an unexpected-keyword TypeError.

Could you update the NVIDIA caller to the new return-value contract too? The clean fix would be to make _run_self_attn() always return self.self_attn(...), since all three NVIDIA attention implementations now return tensors. vllm/models/kimi_k3/nvidia/kda.py itself does not need changing.

[LiuYinfeng01]

This also affects NVIDIA’s low-rank Kimi-Linear path. vllm/models/kimi_k3/nvidia/model.py imports this shared class when use_full_rank_gate=False, but still sets _self_attn_writes_output = True and calls it with output=output. Since this forward() no longer accepts output, that path will fail with an unexpected-keyword TypeError.

Could you update the NVIDIA caller to the new return-value contract too? The clean fix would be to make _run_self_attn() always return self.self_attn(...), since all three NVIDIA attention implementations now return tensors. vllm/models/kimi_k3/nvidia/kda.py itself does not need changing.

Hi. Andreas. I’ve fixed the NVIDIA path and added a test for it. Thank you very much.

[AndreasKaratzas]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84163 for commit 498a91bdeff2.

NVIDIA path covered