[XiaobingSuper]

Summary

Add a production Kimi-K3 Monokernel path for gfx1250:

• single-launch resident wave32/WMMA K1 for 1–48 rows
• fuse the attention-residual seam, hidden projections, convolution and KDA recurrence
• reuse projection weights across token tiles and pair projection row groups
• keep K2 on the existing production path, including tuned FP8/BF16 GEMMs, latent-MoE kernels, collective selection and communication/compute overlap
• avoid binding, dequantizing or compiling the slower experimental gfx1250 K2 kernels
• architecture-gated dispatch; the existing gfx950 path is unchanged

The path is limited to TP8, 256-CU gfx1250 and non-ReplaySSM deployments. Full-attention layers, DSA indexers and every post-K1 operation remain on their ordinary production paths.

Only the two projections consumed inside K1 need to remain BF16:

• *self_attn.in_proj*
• *self_attn.f_b_proj*

o_proj and both shared-expert projections use the normal online-quant configuration.

Performance

Single-GPU operator microbenchmarks compare the one-launch K1 with the repaired six-launch wave32 staged implementation using the same corrected KDA backend:

• S1/L1: 66.7% lower latency (95.2 us vs 286.1 us)
• S16/L1: 65.8% (96.1 us vs 281.3 us)
• S16/L16: 65.0% (95.1 us vs 271.9 us)
• S48/L1: 46.3% (161.3 us vs 300.6 us)
• S48/L16: 41.2% (160.4 us vs 272.7 us)

The projection pipeline computes both 16-row output groups from one x staging pass and delays the second weight group until the first MFMA chain has issued. S48 plain decode uses a wave32 head task with CTA-local norm reduction, removing cross-CTA QK/ONRM handoffs while keeping K1 at one launch.

An earlier two-node run used an experimental custom K2 and regressed E2E performance. That K2 path is no longer selected: commit 4004b393c hands K1's core directly to the model's production o_proj and latent MoE. A matched TP8 E2E rerun of the revised path is pending.

Validation

• deterministic nonzero S1 reference: exact prefix, core_out and conv-state match; recurrent-state maximum absolute error 7.45e-9
• S48/L16 repeated-invocation smoke, reset/non-reset block writes and GPU 0/1 smoke passed
• all S1/S16/S32/S48 legal qlen builds: zero scratch, zero VGPR/SGPR spills, LDS <= 320 KiB
• production K2 handoff contracts cover normal and block-reset seams
• complete targeted Kimi Mono suite after the production-K2 change: 474 passed
• Python compilation, diagnostics and git diff --check: clean

The earlier two-node production run passed accuracy (GSM8K 0.985/200, MLPerf inline accuracy 64.11%). A checkpoint rerun is still required for the revised performance number.

Related: https://github.com/wuhuikx/LLM-Booster/issues/366

[github-actions]

🏷️ 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 2894 --add-label <label>

[gbyu-amd]

Trying #2894 end to end: Kimi-K3, TP8 across two MI455 trays (gfx1250)

We merged #2894 at e07d31b6 into our Kimi-K3 TP8 image and tried to serve it. Mono never got to run
a kernel: it is refused at three successive points, and configuration alone cannot get past the last
one. Details below, in the order we hit them.

Setup

• ATOM: main 16ebac2e + feat(tp): run one tensor-parallel group across nodes #2647 (TP across nodes), [DSpark][DCP] Allow speculative decode with DCP on gfx1250's PS1 MLA kernel #2760, [DCP] Prefetch the next MLA layer's cached-prefix KV all-gathers on a high-priority side stream (opt-in) #2737, fix(mla): plan persistent decode over the rows q/o have #2557, perf(kimi-k3): stream attn-res candidate rows through LDS with TDM on gfx1250 #2814, [MLA] Opt-in FP8 prefill attention on gfx1250: ASM or FlyDSL #2899, feat(kimi-k3): record prefill as token graphs #2942, feat(moe): route FlyDSL MXFP4 decode through aiter HERD on gfx1250 (Kimi-K3 image branch) #2900,
  feat(kimi-k3): fuse o_proj's E8M0 activation quant into the MLA out gate #2994, feat(kimi-k3): fuse down_proj's E8M0 activation quant into SiTUv2 #2995, perf(kimi-k3): quantize the shared qkv_a / g_proj input once #2996; then feat(kimi-k3): add gfx1250 monokernel #2894 at e07d31b6, merged without conflicts (same git patch-id as the
  PR's own diff).
• aiter: main 8f4fe9e5 (with [HIP] [FlyDSL] [GFX1250] Kimi-K3 DSpark PTPC/blockscale and f_b_proj GEMM rows aiter#7142 and [Config] [gfx1250] Tune MiniMax-M3 TP4 BF16, indexer-CP QKV and MoE token 1024 aiter#7209) + [Config] Tune Kimi-K3 gfx1250 gemm_a4w4 aiter#7138,
  [ASM] [HIP] [FlyDSL] [gfx1250] asm mha fp8 fwd: add D192x128 per-tensor varlen kernel aiter#7375, [FlyDSL] optimize dsv4.1, k3, m3 moe perf for gfx1250 aiter#7296, [FlyDSL][MoE] Kimi-K3 HERD on gfx1250: fused_moe(herd_topk) and K+1 sigmoid top-k aiter#7313, [HIP] perf(custom_ar): gfx1250 LL two-shot all-reduce for 1 MiB+ messages aiter#7432.
• moonshotai/Kimi-K3 @ f831ab66, TP8 over two trays of 4 x MI455 (gfx1250, 256 CUs):
  -tp 8 --tensor-parallel-size-local 4, --decode-context-parallel-size 8, MoE TP (no expert
  parallelism), --cudagraph-mode FULL, FP8 KV cache, DSpark with 4 draft tokens at concurrency 8 (a
  full decode step is 40 rows), ATOM_MONO_ENABLE=1, and ATOM_ENABLE_REPLAYSSM=0 because the gfx1250
  runner refuses ReplaySSM.

1. A TP group that spans nodes is refused before the gfx1250 runner is chosen

common_refusals (atom/mono/runtime/deployment.py) has carried this check since a98c0dd:

(
    not getattr(atom_config, "is_multinode_tp", False),
    "TP group spans nodes (HIP IPC peer memory)",
),

so every rank logs Kimi-K3 mono decode off: TP group spans nodes (HIP IPC peer memory). The reason
holds for the gfx950 runner, but runner_gfx1250.py keeps no peer memory: it overrides __init__, so
owned_peer_buffer and StepMailboxes are never built, and K2's collectives go through
get_tp_group(). With the local patch below (a multinode_gfx argument next to supported_gfx, which
K3 sets to ("gfx1250",)), all eight ranks log Kimi-K3 mono decode armed (TP8), and
tests/models/test_kimi_k3_mono*.py plus tests/test_mono*.py give the same results as without it.

diff --git a/atom/models/kimi_k3/mono/config.py b/atom/models/kimi_k3/mono/config.py
--- a/atom/models/kimi_k3/mono/config.py
+++ b/atom/models/kimi_k3/mono/config.py
@@ -253,6 +253,9 @@
             SUPPORTED_TP,
             supported_gfx=("gfx950", "gfx1250"),
             refuse_decode_cp=False,
+            # runner_gfx1250 keeps no peer memory: its TP collectives are RCCL's,
+            # which reach the ranks of a TP group on the other node as well
+            multinode_gfx=("gfx1250",),
         )
         + resident_grid_refusals(atom_config)
     ) + [
diff --git a/atom/mono/runtime/deployment.py b/atom/mono/runtime/deployment.py
--- a/atom/mono/runtime/deployment.py
+++ b/atom/mono/runtime/deployment.py
@@ -88,6 +88,7 @@
     refuse_tr_hybrid: bool = True,
     supported_gfx: tuple[str, ...] = ("gfx950",),
     refuse_decode_cp: bool = True,
+    multinode_gfx: tuple[str, ...] = (),
 ) -> list[tuple[bool, str]]:
     """``(holds, why not)`` of the conditions every mono model needs: ``dims``
     (the model's ``Dims``) shards at this TP, which is one of ``supported_tp``,
@@ -113,7 +114,12 @@
     (its attention layers stay on the production path, which handles the
     split) or are written for that split and refuse any other themselves
     (the shared DSA indexer, ``dsa_indexer.bind``) passes
-    ``refuse_decode_cp=False``."""
+    ``refuse_decode_cp=False``.
+
+    ``multinode_gfx``: the architectures on which the model's runner serves a
+    TP group that spans nodes. Mono's peer buffers are HIP IPC mappings, which
+    cannot reach a GPU on another node, so a TP group across nodes is refused
+    unless the runner for this architecture keeps no peer memory at all."""
     tp = atom_config.tensor_parallel_size
     dpa = dp_attention_ranks(atom_config)
     if dpa not in (0, *supported_dpa):
@@ -147,7 +153,7 @@
         # another node. Read like `compilation_config` below: a config without
         # the attribute (plugin modes) is a single-node one.
         (
-            not getattr(atom_config, "is_multinode_tp", False),
+            gfx in multinode_gfx or not getattr(atom_config, "is_multinode_tp", False),
             "TP group spans nodes (HIP IPC peer memory)",
         ),
         (not (atom_config.enable_tbo or atom_config.enable_tbo_decode), "TBO"),

2. The five bf16 projections (already documented; FYI)

Our launch quantizes the checkpoint's BF16 layers online, so the bind refused with
layer 1 in_proj is torch.float8_e4m3fn. Excluding *self_attn.in_proj*, *self_attn.f_b_proj*,
*self_attn.o_proj*, *shared_experts.gate_up_proj* and *shared_experts.down_proj* from the online
quant, as recipes/Kimi-K3.md says, removes that refusal, and the bind goes on to the next check.

3. Blocking: the gate/up layout of the routed experts

On gfx1250 our launch sets ATOM_MOE_GU_ITLV=1 for the production MoE, and #2894 refuses it twice in
atom/models/kimi_k3/mono/weights.py: in the shared _experts (line 290), which is the one we hit,

Kimi-K3 mono decode off: layer 1 experts: gate/up interleaved (ATOM_MOE_GU_ITLV); the scale readers follow shuffle_scale's plain layout

and again in the gfx1250 bind_k2_expert_scales (line 366):
gate/up interleaving is not supported by Kimi mono.

With ATOM_MOE_GU_ITLV=0 the server does not get as far as the bind: the first prefill, which runs the
production MoE, fails in aiter (aiter/fused_moe.py, _fused_moe_impl):

ValueError: gfx1250 fused_moe requires gu_interleave weight layout (gate_mode='interleave'), got 'separated'

So on gfx1250, with today's aiter, the production MoE needs the interleaved layout and #2894's K2
refuses it, for the same expert weights. Which layout did the TP8 operator A/B run with, and is the
interleaved layout planned for the gfx1250 K2 expert stages?

What this does not cover

No mono kernel ran, because the runner never got past its bind. So nothing here says anything about
#2894's kernels, their speed or their accuracy.

[XiaobingSuper]

@gbyu-amd Thanks — your diagnosis is correct. The earlier TP8 A/B used the non-interleaved operator setup; it did not prove that the current production prefill layout could bind and hit the gfx1250 mono path. The production-readiness claim was too strong.

I pushed 520e5b6df with both blockers fixed:

1. cross-node TP admission is now architecture-local: Kimi gfx1250 is allowed because its runner owns no HIP-IPC peer buffer and uses standard RCCL; gfx950 and other peer-memory mono paths remain refused;
2. K2 now consumes production GUGU w13 and n32k4 scales zero-copy. The gate/up N32 WMMA tiles read physical interleaved rows, reconstruct adjacent gate/up pairs in registers, and keep the existing SiTU/output packing. There is no runtime repack. The separated GGUU path remains supported for compatibility.

I also fixed the gfx1250 scale-size check: native n32k4 w2 has no gfx950-style K-group padding.

Validation on MI455/gfx1250:

• production moe_shuffle_weight(..., is_guinterleave=True, gate_up=True) and moe_shuffle_scale(..., is_guinterleave=True, gate_up=True) compared against their separated counterparts through the K2 UG device stage: packed intermediate and scales are bit-identical (GUGU_EQUIVALENCE_OK);
• all GGUU/GUGU K2 finish widths compile with zero scratch and zero VGPR/SGPR spills;
• focused contracts pass, and the broader Kimi mono suite is 556 passed, 1 skipped (excluding two unrelated container import-teardown tests).

The remaining missing proof is your actual two-node TP8 model run. Please retry with the same production ATOM_MOE_GU_ITLV=1 recipe; no GU_ITLV=0 workaround should be used.

[gbyu-amd]

Retried with 520e5b6, thanks for the quick fix. The setup is the same two-node TP8 one as in my previous comment, now also with #2942's latest commits, plus 520e5b6. The launch keeps the production layout:

• ATOM_MOE_GU_ITLV=1, ATOM_MONO_ENABLE=1, ATOM_ENABLE_REPLAYSSM=0
• the five projections excluded from online quant
• DSpark with 4 draft tokens at concurrency 8, so decode steps are 5–45 rows

The bind now passes. All eight ranks log Kimi-K3 mono decode armed (TP8), and the packed gfx1250 runner binds. The GUGU w13, the n32k4 w2 scale size and the cross-node admission all hold on the real checkpoint.

One merge note: our image already carries #2647's multinode refusal in common_refusals (a98c0dd). 520e5b6's new check would sit next to that one, and the old check would still refuse. We put multinode_gfx on the existing check instead.

Blocking: no width builds K1, so mono never launches a kernel. During CUDA-graph capture, every rank fails the K1 build for every width:

  File ".../kernels/attention/kda_builder_gfx1250.py", line 357, in validate
ValueError: local K1 state_len must cover qlen plus the three-entry causal-convolution history
Kimi-K3 gfx1250 packed mono: building the width-770 kernels failed

It fails for all nine widths the run needs: rows 5, 10, …, 45 at qlen 5 (width keys 90 … 770). WidthBuilds catches the error, so the server comes up and serves those widths on the production path.

The cause is a one-slot mismatch with the pool ATOM allocates:

• K1LocalBuildSpec.validate (kda_builder_gfx1250.py:356) requires state_len >= qlen + 3.
• The runner's state_len is the KDA conv pool's window: K3MonoDecodeRunner._state_len (runner.py:223) reads conv_state.size(-1). ATOM sizes that window as conv_kernel_size - 1 + num_spec (_state_shape, gdn_attn.py:397).
• That window is qlen + 2 under speculation (7 for 4 drafts) and 3 for plain decode, so it is one slot short in both cases.
• The gfx1250 K1 tests build with state_len=qlen + 3 (test_kimi_k3_mono_k1_build_gfx1250.py:32) and state_len=4 at qlen 1 (test_kimi_k3_mono_k1_local_runner_gfx1250.py:37), so they never meet the production pool.

The conv itself looks fine with ATOM's window:

• The resident K1 reads the 3-entry history at acc_idx .. acc_idx + 2, and state_len - qlen "old" entries from acc_idx + 1 (kda_resident_gfx1250.py).
• It then writes back the old entries followed by the step's qlen inputs (kda_conv_compute, kda_gfx1250.py).
• That is the convention the gfx950 K1 uses (conv_state_store, kda.py), and the gfx950 K1 runs on ATOM's kernel - 1 + num_spec pool.

With state_len = qlen + 2, every read stays in the window; the last one is slot acc_idx + 2 ≤ qlen + 1. So qlen + 3 looks stricter than the kernel needs. Is there another reason for it?

We are trying state_len >= qlen + 2 locally and will report what we find. As before, no mono kernel ran, so nothing here measures #2894's kernels, their accuracy or their speed.

[gbyu-amd]

Follow-up: with the K1 bound relaxed to qlen + 2, mono serves the real model end to end, and accuracy holds. In our two-node TP8 run it is about 3.7x slower per output token than the production path.

Setup. ATOM is the image from my first comment, plus #2942's latest commits, 520e5b6 and the one-line K1 change below. The launch is:

• the recipe's default online quant (per_block_fp8, kv_b_proj ptpc_fp8), with the five projections excluded
• ATOM_MOE_GU_ITLV=1, ATOM_MONO_ENABLE=1, ATOM_ENABLE_REPLAYSSM=0
• DSpark with 4 draft tokens at concurrency 8, so decode steps are 5–45 rows

--- a/atom/models/kimi_k3/mono/kernels/attention/kda_builder_gfx1250.py
+++ b/atom/models/kimi_k3/mono/kernels/attention/kda_builder_gfx1250.py
@@ -353,10 +353,13 @@ class K1LocalBuildSpec:
-        if self.state_len < self.qlen + 3:
+        # ATOM's conv pool keeps the last `kernel - 1 + num_spec` inputs
+        # (`_state_shape`): the history read at `acc_idx <= qlen - 1` ends at
+        # `acc_idx + 2`, so qlen + 2 entries serve every accepted count.
+        if self.state_len < self.qlen + 2:
             raise ValueError(
-                "local K1 state_len must cover qlen plus the three-entry "
-                "causal-convolution history"
+                "local K1 state_len must cover qlen plus the two history "
+                "entries the conv window keeps"
             )

Mono runs, and the output is right.

• All nine widths build on all eight ranks (packed mono serving rows=5 … 45 qlen=5). The only decline reason logged is prefill.
• The accuracy gate passes.
• GSM8K, 5-shot, first 200 problems, at num_concurrent=8 so that every decode step is a mono width: strict-match 0.985 ± 0.0086, flexible-extract 0.985.
• The MLPerf Endpoints c8 run finishes with 0 failed turns and 64.11% inline accuracy.

But it is much slower. Both points are MLPerf Endpoints at concurrency 8 with a 10-minute issue cap, on the same harness and the same image:

c8	tok/s/GPU	system tok/s	tok/s/user	TTFT p50 / p90	TPOT p50 / p90 / p99	turns
mono off (recipe default online quant)	67	533	108.1	0.20 / 0.34 s	8.7 / 12.4 / 15.8 ms	795
mono on	25	202	30.7	0.29 / 0.55 s	32.1 / 41.9 / 50.9 ms	259

The two runs differ in mono, and in the five projections being BF16 in the mono run. Draft acceptance is unchanged: 40.8% vs 42.1%, 2.63 vs 2.68 tokens per forward. So each decode step got slower, from about 23 ms to about 86 ms (TPOT x tokens per forward).

We have not profiled it yet. One candidate is K2's collectives. _build_k2 hands ExternalCollectiveK2Runner an RCCL all-gather (use_custom=False, runner_gfx1250.py:302) and a torch.distributed.all_reduce (runner_gfx1250.py:305) for every layer. Across two trays these are small cross-node RCCL calls, while the production path's all-reduces use aiter's custom all-reduce (with the LL two-shot of ROCm/aiter#7432 in this image). Two such calls per layer at a few hundred microseconds each would account for most of the gap.

Was the operator A/B measured with these same collectives? And what did one S32 or S48 step break down into? A per-kernel breakdown would tell us where to look.

[XiaobingSuper]

@gbyu-amd Thanks — both findings are correct.

I pushed 06a7fbaad:

• K1 now accepts ATOM production conv windows at state_len >= qlen + 2; rows 5..45 / qlen 5 / state_len 7 all compile with zero scratch and zero spills.
• gfx1250 K2 no longer calls torch.distributed.all_reduce directly. It now uses tp_group.all_reduce(..., ca_fp8_quant=False), matching the production dispatcher so LL/custom AR, symmetric memory, or RCCL fallback is selected by the same communicator policy.

The earlier isolated K1 gains remain valid, but the previous full TP8 operator-path performance claim was not measured with this production cross-node collective stack and should not be used. The real model numbers in your comment are the current E2E baseline.

I am preparing a conc8/MTP4 rerun with the production online-quant baseline unchanged. Mono ON alone will carry the five required BF16 exclusions; Mono OFF will remain the original recipe.

[XiaobingSuper]

gfx1250 now keeps K2 on the production path (commit 4004b393c). The resident single-launch K1 hands its core to the model’s existing o_proj and latent MoE, preserving K3 tuned GEMMs, expert kernels, collectives, and overlap. The runner no longer binds/dequantizes/compiles custom K2. Only in_proj and f_b_proj need Mono BF16 exclusions; o_proj and shared-expert projections should use the baseline online-quant recipe. Targeted suite: 474 passed. Please use this head for the next TP8 A/B.