[rebklee]

Summary

Mono decode reads only bf16 weights, so mono deployments, including the agentic recipe's C1
launch, exclude five dense projections from the ptpc_fp8 quant: KDA's in_proj and f_b (K1),
and o_proj and the shared experts (K2). C1 then carries 5.3 GiB more weight per rank and runs
bf16 GEMMs on every prefill.

This PR has K1 and K2 read the FP8 planes in place (E4M3 in AITER's (16, 16) preshuffle, an fp32
scale per row). FP8 widens exactly to bf16, so activations, MFMAs and summation order are
unchanged, and bf16 planes run as before. C1 now quantizes all five, like the other bands, which
don't run mono and are unchanged.

8x MI355X, the agentic C1 launch, against main 58ba50180; 3 launches per arm, ranges apart
unless noted:

• Decode is 5.1% faster. Median ITL 14.82 to 14.07 ms, output throughput +6%. Most of it is
  K1 (2445 to 1775 us per target forward, -27.4%), whose largest plane, in_proj, is now half
  the bytes.
• Prefill is 6 to 7% faster. TTFT 540 to 506 ms at 8192-token prompts and 1690 to 1571 ms at
  32768, because prefill runs the five as FP8 GEMMs.
• 5.3 GiB less weight per rank, 44% more KV blocks at C1 (7157 to 10325).
• No measurable accuracy change. GSM8K 0.9666 against main's 0.9674.
• FP8 reads are verified. Kernel tests match the dequantized product bit for bit, and on a
  served model all 2720 in_proj checks are within rounding error of an fp64 recompute.
• Cost: the draft's o_proj becomes FP8 too, making the draft forward 27 us slower; the
  decode numbers above include it.

Changes

• atom/mono/device/gemv.py (shared by K1, K2 and the DSA indexer) gains an FP8 (w8) path: in
  AITER's (16, 16) preshuffle a lane's 8 K columns are 8 contiguous bytes, so it loads 2 dwords
  instead of 4, widens them to bf16 and scales each reduced row before the bf16 rounding.
• K1 (kda.py) reads in_proj and f_b as FP8; K2 (layer.py, back.py, stages/moe.py)
  reads o_proj and the shared experts. The router and the latent projection stay bf16.
• runner.py adds w8 to the K1 / K2 build keys and passes the row scales as the last kernel
  arguments, so bf16 builds keep their argument layout.
• weights.py binds FP8 planes in place and refuses what the kernels can't read, such as a
  non-E4M3fn or unpadded plane, a bias, or bf16 and FP8 mixed in one build.
• model.py: _gated_core gains fuse_o_quant, so mono's forward_pre_o no longer fuses an
  FP8 o_proj's activation quant into the MLA gate; production forward is unchanged.
• Recipes: the agentic C1 launch drops the five exclusions (keeping fp32 KDA state, which mono
  needs on main), and Kimi-K3.md drops the exclusion advice.

Performance

8x MI355X, the agentic C1 launch, ATOM main 58ba50180 with and without this commit (without it,
the five are excluded, as the recipe does; both arms use fp32 KDA state). 3 interleaved launches
per arm; mean, with the range in parentheses.

	main (five projections bf16)	this PR (all FP8)	change
peak_torch, worst rank	202.01 GiB	196.69 GiB	-5.32 GiB
KV pool, worst rank	7157 blocks	10325 blocks (10317-10332)	+44.3%
median ITL, mean of waves 1 and 2	14.82 ms (14.70-14.93)	14.07 ms (14.04-14.10)	-5.1%
output throughput, waves 1 / 2	184 / 246 tok/s	196 / 260 tok/s	+6.1% / +5.6%
TTFT, 8192-token prompts (median)	540.4 ms	505.9 ms	-6.4%
TTFT, 32768-token prompts (median)	1689.6 ms	1571.2 ms	-7.0%
profiled step period	15518 us	14471 us	-6.7%
K1 per target forward	2445 us (2444-2447)	1775 us (1773-1777)	-27.4%
K2 per target forward	8142 us (7956-8376)	7760 us (7371-8232)	ranges overlap
draft model forward	951 us, 112 kernels	978 us, 117 kernels	+27 us

• Decode. K1 saves 670 us per target forward, the same to within 4 us in every launch. K2
  averages 382 us faster, but its launches spread over up to 860 us, so that part is unresolved.
  Net of the draft's 27 us, the step is 0.75 ms shorter in median ITL (1.05 ms profiled).
• Prefill. This PR is faster in 33 of 33 same-prompt pairs at each length, and its slowest
  prompt beats main's fastest.

Method

• rocm/atom-dev:nightly_202611011508-rocm10.0, ATOM main 58ba50180 mounted over the image's,
  with and without this commit. Every arm preloads the image's libgomp.so.1, which main's
  loader needed on this image before fix(loader): resolve OpenMP symbol through torch #2936; 58ba50180 includes fix(loader): resolve OpenMP symbol through torch #2936, and the preload is kept
  so every arm runs the same way.
• recipes/Agentic-Kimi-K3.md's C1 launch: TP8, DCP1, DSpark with 7 speculative tokens, synthetic
  acceptance length 3.84, mono decode on.
• Decode: atom.benchmarks.benchmark_serving, random dataset, 8192 in / 1024 out (range ratio
  0.8), --ignore-eos, 5 prompts at concurrency 1, seed 0; a warmup wave, then two timed waves.
  With speculative decoding each step streams one chunk, so the median ITL is the step period.
• Prefill: 12 random prompts of exactly 8192 tokens, then 12 of 32768 (range ratio 1, one seed
  per length, no shared prefixes), 8 output tokens, one request at a time. The benchmark's
  initial test request caches each set's first prompt, so TTFTs are over prompts 2 to 12.
• peak_torch (weights plus peak activations) and KV blocks are from the server log. K1, K2, the
  draft forward and the step period are medians over rank 0's graph replays in a torch-profiler
  window during a third wave; the profiler lengthens the step, so its period only compares arms.

Limitations

• Mono keeps activations bf16, while the production GEMM quantizes them per token, so the two
  differ on the same plane; mono's are closer to the exact product (see Testing).
• Only per-token E4M3fn in AITER's preshuffle is read; anything else, including e4m3fnuz on
  gfx942, is refused at bind. The (16, 16) layout isn't recorded on the tensor, so this relies on
  ATOM's linears shuffling only at that default.
• The draft's o_proj is now FP8, as at every other band (+27 us a draft forward; it affects
  acceptance, not output). A C1-only exclusion of layers.*.self_attn.o_proj, which matches only
  the draft's, would recover this, but would bring back a C1-only quant override, and its effect
  on acceptance is unmeasured.

Testing

gfx950 (MI355X), rocm/atom-dev:nightly_202611011508-rocm10.0.

• tests/models/test_kimi_k3_mono_dense_bind.py (new), CPU: bf16 and FP8 planes bind in
  place, unreadable formats and mixed bf16 / FP8 layers are refused, and the kernels' FP8 operand
  addresses match AITER's shuffle_weight(layout=(16, 16)).
• Same file, GPU (37 passed): gemv's bf16 and FP8 paths equal the dequantized product bit
  for bit at 1, 5 and 16 tokens and read all 254 finite E4M3 codes exactly; breaking the row
  scale or the dword order fails them.
• Served model: 2720 / 2720 checks of the KDA in_proj (C2, C8, and concurrency 5 and 8 at
  C4 settings, 5 to 64 rows) are within worst-case fp32 rounding of an fp64 recompute, with
  relative L2 error 1.7e-3 against 1.1e-2 to 1.3e-2 for the production GEMM. f_b is checked
  only by GSM8K.
• ISA (gfx950) against main: bf16 K1 / K2 builds keep their registers, spills and
  instructions (three gain one s_nop), the DSA indexer is byte-identical, and no FP8 build
  spills a VGPR.
• Kimi-K3 and mono tests: 275 pass on CPU and 247 on GPU. The image's site-packages tests
  package shadows ATOM's and must be masked; otherwise 2 fail, on main too.
• GSM8K (C2, mono on, no synthetic acceptance): 0.9666, against main's 0.9674 (0.9613 with mono
  off); McNemar p = 1.00. Draft acceptance 53.1% on both. Main's runs are on 606aceb11, with the
  same mono kernels.
• black and ruff clean.

python3 -m pytest -q tests/models/test_kimi_k3_mono_dense_bind.py

[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 2980 --add-label <label>

[valar-cop]

a few small things..

ci green.

[valar-cop]

use _need here too, bare assert breaks the file's own discipline..

[valar-cop]

these two builds unused? remove them or thread the scale args through..