Repository navigation
[AMD][kimik3][ROCm][Perf] Fuse MLA q/kv RMSNorm in AMD Kimi-K3 MLA wrapper - #52494
Conversation
Keep the AITER fused_qk_rmsnorm path for eager Kimi-K3 out of shared mla.py by subclassing MultiHeadLatentAttentionWrapper under vllm/models/kimi_k3/amd, matching the vendor-local model layout. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@rbrugaro-amd this looks great! Just a small question: the ITL and TTFT for the conc 8 run in your results seem to move in the opposite direction as the rest of your results. Is that anomalous? |
|
Good catch @jamesETsmith — that’s an artifact of the table, not of C=8 fusion behavior. The posted C=8 baseline was a standalone idle rerun (3 warmups, fresh server) while optimized C=8 came from the zip sweep after C1/C2, so prefix-cache mix differed. Same-protocol C=8 (both sides from the original zip sweep) is flat: 150.6 vs 150.1 tok/s; ITL mean 36.96 vs 36.75; identical TTFT histograms. p50 TTFT/ITL at conc 8 is bimodal (cached prefix vs queued prefill) because that run used Updated table: all five concs from the original zip sweep (fusion off vs on). Idle-node C1/C2/C8 reruns dropped.
The remaining cross-conc sign flips are not a fusion story. This opt removes one kernel launch with the same FLOPs, and only on MLA layers (~1:3 vs KDA), so E2E is a small decode-step slice. A 68k-prefix / 350-decode agentic sweep is dominated by prefill + prefix-cache queueing, so those deltas are noise (and the C32 tok/s bump is the opposite of “launch savings show up at low conc”). To check the implementation under a tighter decode constraint I ran fixed 8k/1k, no prefix-cache sharing, C=1/4/8, fusion off vs on (ITL is the decode metric; TTFT still includes 8k prefill): aiperf profile --model moonshotai/Kimi-K3 \
--tokenizer /data/workloads-inference/models/Kimi-K3 \
--tokenizer-trust-remote-code \
--url http://127.0.0.1:8000 --api-key EMPTY \
--endpoint-type chat --streaming --use-server-token-count \
--synthetic-input-tokens-mean 8192 --synthetic-input-tokens-stddev 0 \
--output-tokens-mean 1024 --output-tokens-stddev 0 \
--extra-inputs ignore_eos:true \
--extra-inputs min_tokens:1024 --extra-inputs max_tokens:1024 \
--warmup-request-count 3 --sweep-type zip \
--concurrency 1,4,8 --request-count 10,20,40 \
--random-seed 42 --ui simple
Decode ITL p50/p90 at C=1 is ~0.5% — consistent with a launch-only MLA fusion being smaller than serving noise. Correctness of the path is the decode kernel census ( |
|
/ci run |
|
✅ Triggered Buildkite CI #85999 for commit |
rbrugaro-amd
changed the title
|
/ci run |
|
❌ @rbrugaro-amd, A reviewer with write access must run |
|
@dllehr-amd @tjtanaa could you please retry ci? the failed test is unrelated to PR changes |
tjtanaa
added
the
verified
|
/ci run |
|
✅ Triggered Buildkite CI #86193 for commit |
AndreasKaratzas
added
the
ready
|
✅ @rbrugaro-amd, CI is now available for this PR.
|
|
/ci run |
|
✅ Triggered Buildkite CI #86659 for commit |
dllehr-amd
left a comment
dllehr-amd
left a comment
There was a problem hiding this comment.
Thanks @rbrugaro-amd ! I think we can move forward with this!
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe change adds an AMD-specific Kimi-K3 MLA wrapper with optional fused RMSNorm and routes ChangesAMD Kimi-K3 MLA
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This change routes AMD Kimi-K3 MLA through a vendor-local wrapper that fuses eager RMSNorm launches while preserving the compiled execution path. No merge-blocking correctness or runtime risk remains identified. Sequence Diagram(s)sequenceDiagram
participant KimiMLAAttention
participant KimiK3MultiHeadLatentAttentionWrapper
participant fused_qkv_a_proj
participant mla_attn
participant o_proj
KimiMLAAttention->>KimiK3MultiHeadLatentAttentionWrapper: forward positions and hidden states
KimiK3MultiHeadLatentAttentionWrapper->>fused_qkv_a_proj: project q and kv when q_lora_rank is set
KimiK3MultiHeadLatentAttentionWrapper->>mla_attn: execute latent attention
mla_attn-->>KimiK3MultiHeadLatentAttentionWrapper: return attention output
KimiK3MultiHeadLatentAttentionWrapper->>o_proj: project gated output
o_proj-->>KimiMLAAttention: return attention result
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/ci run |
|
✅ Triggered Buildkite CI #87185 for commit |
…apper (vllm-project#52494) Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
…ITER On a stock build, Kimi-K3 MLA decode runs two kernels per layer per step: the fp8 KV-cache write (`vllm::concat_and_cache_mla_kernel`) and the `[ql_nope | q_pe]` concat plus that query's static fp8 quant, which the `mla_decode_concat_quant_fp8` custom op emits as a single `triton_poi_fused__to_copy_cat_clamp_mul_reciprocal_view` kernel. Collapse both into one AITER `fused_qk_rope_concat_and_cache_mla` launch. Measured on MI355X at TP=8, plain `vllm serve` on the stock upstream nightly with no speculative decoding and no compilation-config override: 288 -> 144 launches and 1351.8 -> 682.8 us per rank over 6 decode steps, i.e. 9.39 -> 4.74 us per MLA layer per decode step. End-to-end serving throughput is unchanged within measurement noise and this change does not claim a throughput win. K3 MLA is NoPE, while the kernel always applies RoPE and takes the cos/sin caches as required arguments. Feed it an identity rotation (cos=1, sin=0) so the rotation is a no-op -- the approach ROCm/ATOM uses for K3. The cache is a single row: the `_opt` decode kernel clamps `pos` into `[0, cos_cache.size(0))`, so real positions pass through untouched and K3's 1M max_position_embeddings does not cost ~268 MB to express a no-op. AITER's general decode kernel does not clamp, so the head counts that would select it are excluded. Kept entirely in `vllm/models/kimi_k3/amd/mla.py`, mirroring the q/kv RMSNorm fusion this wrapper already owns (vllm-project#52494); no shared MLA code is touched, so this carries no dependency on vllm-project#47757 and cannot affect any other model or backend. Everything that is fixed at load time -- platform, KV dtype, shapes, DCP/PCP/HiSparse, head padding and the W_UK variant -- is resolved once in `_fused_qk_prep_supported()`, so the only per-batch condition is that the batch is decode-only; anything else falls through to the stock `MLAAttention` path. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
5 participants
Purpose
Fuse the MLA layer's
q_a_layernorm/kv_a_layernormpair into a singleAITER
fused_qk_rmsnormlaunch for eager Kimi-K3 on ROCm, without puttingvendor checks in shared
vllm/model_executor/layers/mla.py.Kimi-K3 (no
@support_torch_compile) never hitsMLADualRMSNormFusionPass,so the two RMSNorms stay as separate AITER launches. NVIDIA already owns a
vendor-local MLA under
vllm/models/kimi_k3/nvidia/; this mirrors that layoutfor AMD:
vllm/models/kimi_k3/amd/mla.pywithKimiK3MultiHeadLatentAttentionWrappervllm/models/kimi_k3/amd/linear.pymla.pyuntouchedThe fused path uses
torch.ops.vllm.fused_mla_dual_rms_normwhen AITER isenabled and
not torch.compiler.is_compiling(), so compiled MLA models keepusing the fusion pass.
Supersedes the shared-
mla.pyapproach in #52215 and the compiler pass approach in #52190Trace evidence
MI355X TP=8, MTP k=2, ISL 128k / OSL 1k, conc 8, decode-phase torch profile
(rank 0, 6 engine steps):
aiter::fused_qk_rmsnorm_kernel<bf16,256,8,true,1>Paired per-MLA
add_rmsnormlaunches for the q/kv latents are goneConfig used
Image:
vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263(+ this PR mounted under
vllm/models/kimi_k3/amd/)Env:
GSM8k accuracy:
2026-08-16:06:25:52 INFO [_cli.run:375] Selected Tasks: [{'tag': ['math_word_problems'], 'task': 'gsm8k', 'dataset_path': 'openai/gsm8k', 'dataset_name': 'main', 'output_type': 'generate_until', 'training_split': 'train', 'fewshot_split': 'train', 'test_split': 'test', 'doc_to_text': 'Question: {{question}}\nEnd your response with the answer on the last line, formatted as: #### [number]\nAnswer:', 'doc_to_target': '{{answer}}', 'metric_list': [{'metric': 'exact_match', 'aggregation': 'mean', 'higher_is_better': True, 'ignore_case': True, 'ignore_punctuation': False, 'regexes_to_ignore': [',', '\\$', '(?s).*#### ', '\\.$']}], 'generation_kwargs': {'until': ['</s>', '<|im_end|>'], 'do_sample': False, 'temperature': 0.0}, 'repeats': 1, 'num_fewshot': 5, 'filter_list': [{'name': 'strict-match', 'filter': [{'function': 'regex', 'group_select': -1, 'regex_pattern': '#### (\\-?[0-9\\.\\,]+)'}, {'function': 'take_first'}]}, {'name': 'flexible-extract', 'filter': [{'function': 'regex', 'group_select': -1, 'regex_pattern': '(-?[$0-9.,]{2,})|(-?[0-9]+)'}, {'function': 'take_first'}]}], 'metadata': {'version': 3.0}}] 2026-08-16:06:37:46 INFO [loggers.evaluation_tracker:209] Saving results aggregated |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.9644|± |0.0051| | | |strict-match | 5|exact_match|↑ |0.9644|± |0.0051|Performance
(
--tokenizeris the local Kimi-K3 checkout; equivalent tomoonshotai/Kimi-K3with hub access. Idle-node baseline reruns used the same flags with
--concurrency/--request-countsubset to the point under test.)