[ROCm][Kimi-K3] Fuse gated RMSNorm with per-token FP8 for KDA o_proj by rebklee · Pull Request #54254 · vllm-project/vllm

rebklee · 2026-08-28T19:12:02Z

Purpose

Depends on #54248 (K3 o_proj becomes Fp8PtpcOnlineLinearMethod). Until this lands, this path is a no-op.

Kimi-K3 KDA is gated RMSNorm then o_proj. With PTPC that is FusedRMSNormGated plus standalone dynamic_per_token_scaled_quant` in the linear. Fuse them in one Triton kernel with per-token amax.

Wrap (fp8, scale) as QuantizedActivation(kFp8DynamicTokenSym) when o_proj.input_quant_key matches. Skip o_norm on that path so we do not double-norm.

fused_kda_decode (#50654) already applied gated RMSNorm on pure non-spec decode. Quantize that output with per_token_fp8_quant so we do not double-norm.

Profile

8×MI355X, TP8, Kimi-K3 + DSpark, 50 decode steps. Overlay is --quantization-config PTPC o_proj only.

Kimi-K3 mix in this window: 69 KDA / 24 MLA → 3450 KDA launches (69 × 50).

Unfused

Gated RMSNorm and per-token quant are still two launches.

Kernel	Calls	CUDA avg	CUDA total
`layer_norm_gated_fwd_kernel.kd`	3450	4.375 µs	15.094 ms
`dynamic_per_token_scaled_quant`	4650	4.086 µs	18.998 ms

The 4650 quant launches are 3450 KDA + 1200 MLA. KDA-only epilogue ≈ 8.5 µs (4.4 + 4.1).

Fused (this PR)

Gated RMSNorm and per-token quant are one launch.

Kernel	Calls	CUDA avg	CUDA total
`_rmsnorm_gated_fp8_per_token_kernel.kd`	3450	4.221 µs	14.561 ms
`layer_norm_gated_fwd_kernel.kd`	0	—	—
`dynamic_per_token_scaled_quant`	0	—	—

KDA epilogue 8.5 → 4.2 µs (−4.3 µs / KDA layer, ~0.30 ms/step × 69 KDA).

Accuracy

8×MI355X, TP8, DSpark MTP (num_speculative_tokens=2). Same recipe as the fused PTPC KDA o_proj path (this change). Full GSM8K, 1,319 examples, lm_eval 5-shot (local-completions). No unfused PTPC pair in this run.

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9613	±	0.0053
		strict-match	5	exact_match	↑	0.9621	±	0.0053

VLLM_ROCM_USE_AITER=1 \
VLLM_USE_BREAKABLE_CUDAGRAPH=0 \
vllm serve moonshotai/Kimi-K3 \
  --served-model-name moonshotai/Kimi-K3 \
  --host 0.0.0.0 --port 8888 \
  --trust-remote-code \
  --tensor-parallel-size 8 \
  --language-model-only \
  --reasoning-parser kimi_k3 \
  --max-model-len 1048576 \
  --kv-cache-dtype fp8 \
  --attention-config '{"mla_prefill_backend":"FLASH_ATTN"}' \
  --quantization-config '{"linear":"fp8_per_channel","ignore":["re:(?!.*self_attn\\.o_proj$).*"]}' \
  --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","num_speculative_tokens":2,"method":"dspark","attention_backend":"TRITON_MLA","kv_cache_dtype":"auto","draft_sample_method":"probabilistic","rejection_sample_method":"block"}'

pip install --quiet 'lm_eval[api]==0.4.12'
lm_eval --model local-completions \
  --model_args model=moonshotai/Kimi-K3,base_url=http://127.0.0.1:8888/v1/completions,tokenized_requests=False,trust_remote_code=True,num_concurrent=8,max_gen_toks=8192,timeout=3600 \
  --tasks gsm8k \
  --num_fewshot 5 \
  --seed 1234

Tests

tests/models/kimi_k3/test_amd_kda_oproj_ptpc.py — fused vs sequential; per_token_fp8_quant vs already-normed (fused_kda_decode path).
Overlay self-check (CPU torch fallback, no GPU): rmsnorm_gated_fp8_per_token matches sequential gated RMSNorm + per-token amax (scale exact, dequant atol/rtol 1e-3). All self.o_norm(...) sites stay behind if not _o_proj_is_ptpc_fp8(...) so PTPC does not double-norm.

FP8ScaledMMLinearKernel.apply_weights already consumes a matching QuantizedActivation, but the AITER PTPC kernels never exposed input_quant_key(), so expose_input_quant_key left the layer unmarked. Return the key from the AITER PTPC kernels and expose it on Fp8PtpcOnlineLinearMethod. Do not override input_quant_key on FP8ScaledMMLinearKernel: torch fallbacks call get_output_padding() via get_current_vllm_config(), which is unset in profile_run. Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>

When o_proj.input_quant_key is kFp8DynamicTokenSym, fuse gated RMSNorm and per-token FP8 in one Triton kernel so the linear skips scaled_fp8_quant. Skip o_norm on that path so we do not double-norm. Signed-off-by: Rebecca Lee <rebecca.lee@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

mergify · 2026-09-03T18:56:00Z

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

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

…orm-fp8-oproj

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.

dllehr-amd · 2026-09-09T17:50:00Z

I think this may be a dup of #52968

dllehr-amd · 2026-09-09T17:51:33Z

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused gated RMSNorm + per-token FP8 for Kimi-K3 KDA ``o_proj``.


can we look at putting this into aiter also if it's not already there?

mergify · 2026-09-10T08:57:13Z

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

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

simondanielsson

Nice work!

simondanielsson · 2026-09-10T16:37:13Z

+    return (xn.to(dtype) * weight.to(dtype)) * torch.sigmoid(gate)
+
+
+def rmsnorm_gated_fp8_per_token(


Suggestion: should we UT this? I saw the description mentioned a tests/models/kimi_k3/test_amd_kda_oproj_ptpc.py

simondanielsson · 2026-09-10T16:45:22Z

+            )
+            if used_fused_decode:
+                x_2d = rearrange(core_attn_out, "1 n h d -> n (h d)")
+                q, scale = per_token_fp8_quant(x_2d, fp8_dtype)


Question: is there a reason for now using rocm_aiter_ops.per_token_quant here?

simondanielsson · 2026-09-10T16:48:49Z

+    t, heads, h = x.shape
+    fp8_max = float(torch.finfo(quant_dtype).max)
+    if not _HAS_TRITON or t == 0 or h > 8192:
+        normed = _rmsnorm_gated_torch(x, weight, gate, eps).reshape(t, heads * h)


Question: Is the fused torch variant better than unfused default impl?

simondanielsson · 2026-09-10T17:10:20Z

+                if isinstance(attn_metadata_raw, dict)
+                else None
+            )
+            used_fused_decode = bool(


Question: is this safe with breakable cudagraphs? or might this condition be misaligned with that inside _forward

Signed-off-by: Rebecca Lee <rebecca.lee@amd.com> # Conflicts: # tests/fusion/test_quant_activation_contract.py # vllm/models/kimi_k3/amd/kda.py

This branch carried its own copy of vllm-project#54248: the three aiter.py input_quant_key() methods, the online/fp8.py QuantizedActivation plumbing, and the contract test's AITER kernel registrations. Those files now come from vllm-project#54248 alone, and this branch keeps only the KDA gated-RMSNorm fusion, inert until vllm-project#54248 lands -- the same dependency that vllm-project#55043 already has. Signed-off-by: Rebecca Lee <rebecca.lee@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

rebklee and others added 2 commits August 28, 2026 12:43

mergify Bot added quantization kimi k3 rocm Related to AMD ROCm labels Aug 28, 2026

github-project-automation Bot added this to AMD Aug 28, 2026

github-project-automation Bot moved this to Todo in AMD Aug 28, 2026

okorzh-amd mentioned this pull request Aug 31, 2026

[ROCm] Expose kFp8DynamicTokenSym on AITER PTPC linears #54248

Merged

2 tasks

rebklee mentioned this pull request Sep 3, 2026

[ROCm][Kimi-K3] Fuse MLA sigmoid-mul with per-token FP8 for o_proj #55043

Open

Merge branch 'main' into rebklee/k3-kda-rmsnorm-fp8-oproj

15af463

mergify Bot added the needs-rebase label Sep 3, 2026

Merge remote-tracking branch 'upstream/main' into rebklee/k3-kda-rmsn…

0567a0c

…orm-fp8-oproj

mergify Bot removed the needs-rebase label Sep 9, 2026

Merge branch 'main' into rebklee/k3-kda-rmsnorm-fp8-oproj

20f2d7f

rebklee marked this pull request as ready for review September 9, 2026 16:34

rebklee requested review from AndreasKaratzas, dllehr-amd, hongxiayang, mgoin, pavanimajety, robertgshaw2-redhat, tjtanaa, tlrmchlsmth, yewentao256 and zyongye as code owners September 9, 2026 16:34

claude Bot reviewed Sep 9, 2026

View reviewed changes

yewentao256 mentioned this pull request Sep 9, 2026

[Feature]: Kimi K3 Performance Optimization #50587

Open

46 tasks

dllehr-amd reviewed Sep 9, 2026

View reviewed changes

mergify Bot added the needs-rebase label Sep 10, 2026

simondanielsson reviewed Sep 10, 2026

View reviewed changes

Merge branch 'main' into rebklee/k3-kda-rmsnorm-fp8-oproj

5da6e3b

Signed-off-by: Rebecca Lee <rebecca.lee@amd.com> # Conflicts: # tests/fusion/test_quant_activation_contract.py # vllm/models/kimi_k3/amd/kda.py

mergify Bot removed the needs-rebase label Sep 16, 2026

Navigation Menu

Uh oh!

Uh oh!

[ROCm][Kimi-K3] Fuse gated RMSNorm with per-token FP8 for KDA o_proj - #54254

[ROCm][Kimi-K3] Fuse gated RMSNorm with per-token FP8 for KDA o_proj#54254
rebklee wants to merge 7 commits into
vllm-project:mainfrom
rebklee:rebklee/k3-kda-rmsnorm-fp8-oproj

rebklee commented Aug 28, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Sep 3, 2026

Uh oh!

claude Bot left a comment

Uh oh!

dllehr-amd commented Sep 9, 2026

Uh oh!

dllehr-amd Sep 9, 2026

Uh oh!

mergify Bot commented Sep 10, 2026

Uh oh!

simondanielsson left a comment

Uh oh!

simondanielsson Sep 10, 2026

Uh oh!

simondanielsson Sep 10, 2026

Uh oh!

simondanielsson Sep 10, 2026

Uh oh!

simondanielsson Sep 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return (xn.to(dtype) * weight.to(dtype)) * torch.sigmoid(gate)


		def rmsnorm_gated_fp8_per_token(

Uh oh!

Uh oh!

Conversation

rebklee commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Profile

Unfused

Fused (this PR)

Accuracy

Tests

Uh oh!

mergify Bot commented Sep 3, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

dllehr-amd commented Sep 9, 2026

Uh oh!

dllehr-amd Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Sep 10, 2026

Uh oh!

simondanielsson left a comment

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rebklee commented Aug 28, 2026 •

edited

Loading