[k3] enable dual stream for shared expert and some fusions by gbyu-amd · Pull Request #1752 · ROCm/ATOM

gbyu-amd · 2026-07-30T17:12:32Z

Motivation

Depends on: ROCm/aiter#4467 and ROCm/aiter#4435

This pr:

Added dual stream for shared/routed expert overlapping
Added ptpc fp8 quant for attn and shared expert gemms, fuse the introduced per-token quant accordingly

Accuracy:

local-completions ({'model': '/workspace/shared/data/amd_int/models/Kimi-K3', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 64, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 5, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9568|±  |0.0056|
|     |       |strict-match    |     5|exact_match|↑  |0.9568|±  |0.0056|

Dual stream overlapping:

E2E perf:
Main:

Input Tokens	Output Tokens	Max Concurrency	Num Prompts	Median TTFT (ms)	Median TPOT (ms)	Output Token Throughput	Total Token Throughput
8192	1024	4	20	557.01	26.60	142.59	1270.64
8192	1024	8	40	548.81	29.98	248.98	2238.21
8192	1024	16	80	565.70	35.91	417.48	3713.71
8192	1024	32	160	596.72	47.92	624.52	5636.14
8192	1024	64	320	799.17	72.08	839.41	7506.28

This pr:

Input Tokens	Output Tokens	Max Concurrency	Num Prompts	Median TTFT (ms)	Median TPOT (ms)	Output Token Throughput	Total Token Throughput
8192	1024	4	20	512.26	25.66	147.14	1311.19
8192	1024	8	40	510.80	28.10	263.76	2371.03
8192	1024	16	80	523.14	37.75	401.82	3574.37
8192	1024	32	160	555.99	45.48	656.35	5923.39
8192	1024	64	320	759.50	65.76	923.87	8261.50

Remaining issue:
Overlap failed at CONC=16, caused some regression instead, which is strange:

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Port the deepseek-style norm+quant fusion onto the true-MLA base so the standalone activation-quant op emitted before online-quantized (ptpc_fp8) Linear layers is folded into the immediately-preceding RMSNorm. atom/model_ops/linear.py: - Skip post-load work for empty fused shells (weight.numel()==0), e.g. KDA b_proj/f_a_proj after their weights are concatenated into in_proj. - Pad a8w8 (per_Token fp8) preshuffle output up to the CK N-tile when the fused N is not tile-aligned, slicing the padding back off in forward. atom/models/kimi_k3.py: - Add _effective_layer_quant + _RMS_FUSABLE_QUANT_TYPES; route QK-norm through deepseek _fuse_rmsnorm_quant (degrades to bf16 when unquantized). - Thread (fp8, scale) tuples through KimiMLP, KimiFullAttention (q_scale into MLA) and KimiKDAAttention (scale carried across the splitting custom op). - Wire KimiDecoderLayer input_layernorm / post_attention_layernorm fusion, gated per layer so non-fusable schemes fall back to today's bf16 path. Every fusion is self-gated; validated by CPU unit tests. fp8 numerics still need one gsm8k run on GPU. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

QuantizationConfig only sets online_quant=True when quant_method is in a fixed allow-list. Kimi-K3's checkpoint declares quant_method "compressed-tensors", which was absent, so --online_quant_config was silently ignored: ptpc_fp8 never activated and every target Linear ran bf16 (gemm_a8w8_bpreshuffle was never called). Add "compressed-tensors" to the allow-list so online fp8 quant engages as intended. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The fusion implementation was renamed to _effective_layer_quant / _kda_packed_modules_mapping, but the tests still referenced the old _layer_effectively_consumes_per_token_fp8 / _kimi_packed_modules_mapping helpers, leaving the branch red (10 failures on HEAD). Update the tests to the current helpers and per-consumer fusion gating, and modernize a few Optional[X] annotations to X | None in kimi_k3.py (+ exec noqa in the linear alignment test). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Overlap shared-expert GEMMs (on alt_stream) with the routed-expert path on the main stream. Both TP all-reduces stay serial on the main stream to avoid concurrent collectives on the single TP communicator. GPU-validated on MI355 TP8 at gsm8k exact_match 0.9555 (no regression vs fp8 baseline). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

github-actions · 2026-07-30T17:12:58Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1752 --add-label <label>

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

gbyu-amd · 2026-08-03T11:09:24Z

kimi k3 test cases need ROCm/aiter#4467 merged as aiter fix.

gbyu-amd · 2026-08-04T10:46:35Z

another fix for ptpc fp8 gemm: ROCm/aiter#4546

* fix(plugin): gate dual-stream MoE by runtime graph mode Use the frontend's per-forward graph decision so FULL decode can retain overlap while PIECEWISE capture remains safely single-stream. Co-authored-by: Cursor <cursoragent@cursor.com> * style: sort forward context imports Match the repository's Ruff import ordering rules. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: xytpai <xytpai@foxmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>

gbyu-amd · 2026-08-05T11:59:12Z

@valarLip hi, please help take a look~ The failed CI cases should be unrelated

Rebased onto main; the branch's routed-MoE sub-batching knob (ATOM_K3_MOE_CHUNK) and the KimiSparseMoeBlock._forward_impl indirection are dropped -- main's dual-stream MoE (#1752) supersedes both. atom/model_ops/attentions/gdn_attn.py context_lens_tensor reads attn_metadata.num_cached_tokens instead of being hard-zeroed. Hard-zeroing restarts the linear-attention recurrence on every chunked-prefill chunk and silently corrupts the output. None (whole-prompt / first chunk) still yields all-zero, so the previous behaviour is preserved. Arch-independent bug fix. atom/model_ops/kimi_k3/attention_residual.py Clamp the candidate row index so the power-of-2-padded axis is not addressed out of bounds (addressing-only; masked lanes still resolve to other=0.0, result is bit-identical), and cap num_stages to 1 on gfx1250 only, where software-pipelining the H loop faults under concurrent load. atom/models/kimi_k3.py KDA prefill can run on fused_recurrent_kda via a new recurrent= argument to _run_kda, gated on ATOM_KDA_FORCE_RECURRENT (default off). chunk_kda NaNs on gfx1250 for prompts shorter than its chunk size, and its transpose_state_layout output can mismatch what the decode-time fused_recurrent_kda reader expects. safe_gate is popped because it is chunk_kda-only. The KDA output buffer is new_zeros, not new_empty. The fused decode kernel returns early for rows whose cache slot is PAD_SLOT_ID (CUDAGraph batch padding) and leaves them unwritten; uninitialized bf16 easily lands on a NaN/Inf pattern, which o_norm/o_proj propagate and the MoE router turns into out-of-range expert ids. atom/utils/envs.py ATOM_KDA_FORCE_RECURRENT, default 0. recipes/Kimi-K3.md gfx1250 / MI450 section: image, versions, server command, env-var table, GSM8K results. Validated on GFX1250, full 1319-question GSM8K, 5-shot x3: 0.9591 / 0.9621 / 0.9548 (flexible-extract).

gbyu-amd and others added 5 commits July 30, 2026 03:24

opt k3

695385d

zufayu requested a review from yhl-amd July 31, 2026 01:17

gbyu-amd and others added 4 commits July 31, 2026 10:12

Update atom/model_ops/linear.py

f18d0e6

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

rm local test files

b50e2bd

update recipe

198e9c2

make ar parallel on two streams

f4d810e

gbyu-amd requested review from XiaobingSuper, ZhangLirong-amd, valarLip and zejunchen-zejun July 31, 2026 08:22

small update

cd31755

zejunchen-zejun previously approved these changes Aug 3, 2026

View reviewed changes

change ar launch order to ensure better overlap

fc071e3

gbyu-amd dismissed zejunchen-zejun’s stale review via fc071e3 August 3, 2026 10:39

gbyu-amd force-pushed the guanbao/k3_rmsnorm_quant_fusion branch from 8f5b6aa to fc071e3 Compare August 3, 2026 10:39

Merge branch 'main' into guanbao/k3_rmsnorm_quant_fusion

73c99d7

gbyu-amd added 2 commits August 3, 2026 20:00

fix online quant with plugin

3b900ff

Merge branch 'main' into guanbao/k3_rmsnorm_quant_fusion

452a0ac

ZhangLirong-amd previously approved these changes Aug 4, 2026

View reviewed changes

wuhuikx mentioned this pull request Aug 4, 2026

fix(plugin): gate dual-stream MoE by runtime graph mode #1790

Closed

3 tasks

XiaobingSuper mentioned this pull request Aug 4, 2026

fix(plugin): gate dual-stream MoE by runtime graph mode #1791

Closed

4 tasks

XiaobingSuper and others added 2 commits August 4, 2026 19:37

Merge branch 'main' into guanbao/k3_rmsnorm_quant_fusion

5f22401

XiaobingSuper dismissed ZhangLirong-amd’s stale review via 540b1d1 August 4, 2026 11:38

correct rope_max_position

8487e5f

gbyu-amd and others added 6 commits August 6, 2026 09:58

Merge branch 'main' into guanbao/k3_rmsnorm_quant_fusion

09e932e

update linear.py

52aff26

clean comments

ebd7c55

update vllm plugin recipe

98dbb5d

fix format

c1f9c9f

Merge branch 'main' into guanbao/k3_rmsnorm_quant_fusion

1ffa09d

valarLip approved these changes Aug 6, 2026

View reviewed changes

valarLip merged commit d182e77 into main Aug 6, 2026
63 of 70 checks passed

valarLip deleted the guanbao/k3_rmsnorm_quant_fusion branch August 6, 2026 09:13

sajandhy mentioned this pull request Aug 16, 2026

feat(k3): Kimi-K3 vLLM plugin vision + DSpark draft support #1910

Open

7 tasks

Navigation Menu

Uh oh!

[k3] enable dual stream for shared expert and some fusions - #1752

[k3] enable dual stream for shared expert and some fusions#1752
valarLip merged 23 commits into
mainfrom
guanbao/k3_rmsnorm_quant_fusion

gbyu-amd commented Jul 30, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 30, 2026

Uh oh!

gbyu-amd commented Aug 3, 2026

Uh oh!

gbyu-amd commented Aug 4, 2026

Uh oh!

gbyu-amd commented Aug 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

gbyu-amd commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Submission Checklist

Uh oh!

github-actions Bot commented Jul 30, 2026

🏷️ CI Guide

Uh oh!

gbyu-amd commented Aug 3, 2026

Uh oh!

gbyu-amd commented Aug 4, 2026

Uh oh!

gbyu-amd commented Aug 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

gbyu-amd commented Jul 30, 2026 •

edited

Loading