[ROCm] Return Kimi-K3 MLA output directly

[LiuYinfeng01]

Summary

• Return the AMD Kimi-K3 MLA projection result directly instead of allocating a second hidden-state tensor and copying into it.
• Preserve the caller-owned output buffer used by KDA layers.
• Add ownership-focused coverage for both dispatch paths.

Kimi-K3 has 24 MLA layers. For BF16 hidden states of width 7168, this removes 24 copies of num_tokens * 7168 * 2 bytes per model forward. The returned values are unchanged; only tensor ownership changes.

This is intentionally separate from #50664, which owns output-gate fusion and explicitly leaves output-buffer ownership out of scope.

Test plan

• Fixed-image unit test: 2 passed in rocm/vllm-dev:gfx950_kimi_k3_20260727.
• Fixed image plus [GFX950] Relocate MLA Gluon kernel and unify decode dispatch ROCm/aiter#4450 overlay: Kimi-K3 TP8 server boot and short serving smoke passed.
• GSM8K subset: baseline 100/100, candidate 100/100; zero invalid responses in both arms.
• Short random serving A/B, ISL 8192 / OSL 512 / concurrency 4 / 8 requests: all requests succeeded.
• git diff --check and IDE diagnostics passed.

The one-shot serving measurement was neutral within noise: median TPOT was 25.53 ms baseline and 25.75 ms candidate. This PR therefore does not claim a measured end-to-end speedup from that short run; its deterministic effect is removal of the redundant MLA output allocation/copy.

Tool assistance

Cursor assisted with implementation, tests, benchmarking, and drafting this description.

Made with Cursor

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

[Fangzhou-Ai]

Thanks for the direct-return cleanup. A new fixed-cohort pure-decode trace gives a useful estimate for why the short one-shot serving A/B can look neutral even though the redundant work is deterministic.

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 plus the 24 MLA copies removed here. The complete 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 24-copy MLA share is approximately:

• C16: 0.13 ms/step, about 0.39% of the pure-GPU period
• C24: 0.15 ms/step, about 0.40% of the pure-GPU period

So the realistic standalone expectation at these shapes is below 0.2 ms/step, which is small enough to be hidden by run-to-run end-to-end noise. This is a trace-derived removable-cost estimate, not a candidate A/B measurement. The change still removes 24 deterministic copies and should combine with #50592, which removes the other 69 post-projection copies.

[LiuYinfeng01]

Thanks for the direct-return cleanup. A new fixed-cohort pure-decode trace gives a useful estimate for why the short one-shot serving A/B can look neutral even though the redundant work is deterministic.

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 plus the 24 MLA copies removed here. The complete 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 24-copy MLA share is approximately:

• C16: 0.13 ms/step, about 0.39% of the pure-GPU period
• C24: 0.15 ms/step, about 0.40% of the pure-GPU period

So the realistic standalone expectation at these shapes is below 0.2 ms/step, which is small enough to be hidden by run-to-run end-to-end noise. This is a trace-derived removable-cost estimate, not a candidate A/B measurement. The change still removes 24 deterministic copies and should combine with #50592, which removes the other 69 post-projection copies.

@Fangzhou-Ai Thank you for the detailed analysis and recommendation. I’ll fold the functionality of #50847 into #50592 so that a single PR removes all 93 KDA and MLA post-projection copies. I really appreciate your profiling work and helpful guidance.

[hongxiayang]

should one of the #50592 and #50847 be closed as one supersede another?

[LiuYinfeng01]

should one of the #50592 and #50847 be closed as one supersede another?

I will closed this PR #50847 and I have moved this PR codes to PR #50592.

[LiuYinfeng01]

Superseded by #50592, which now includes this MLA direct-return change together with the complementary KDA post-projection direct-return path and combined fixed-cohort A/B evidence.