[rebklee]

Summary

ATOM_GDN_SSM_DTYPE=fp16 halves Kimi-K3's KDA recurrent state pool, and it is the agentic
recipe's default in every band. Mono decode refused it, because K1 read and wrote the state in
place as fp32, so the recipe's C1 launch (the band that runs mono) sets the pool back to fp32 as a
precondition. This PR has K1 take the pool's dtype as a build constant (state16): every state K1
loads is widened to fp32, the recurrence runs in fp32 as before, and every state it stores is
rounded to fp16, which is how the production Triton kernels treat the same pool. An fp32 pool
builds the same instructions as before; bf16 is still refused. The recipe's C1 launch drops the
override.

On 8x MI355X with the recipe's C1 launch, this gives the worst rank 45% more KV blocks (7158 to
10407) with no measurable change in the median decode step.

Motivation

The fp32 precondition doubles the state pool wherever mono is armed, and keeps mono out of any
deployment that runs the pool in fp16, as the agentic recipe does by default. The production
path already serves this state in fp16 in every band.

Changes

• kda.py:
  • KdaPreBuild.state16, a build-key field (kernel names gain _h0 / _h1).
  • state_at(c, slot): one address function for every state access (a request's initial state,
    each token's snapshot slot, and under ReplaySSM the checkpoint load and the flush).
  • For fp16, state_loads widens each state to fp32 and store_states rounds it back to fp16,
    to nearest even (v_cvt_pk_f16_f32).
• runner.py:
  • _state_dtype() requires one recurrent dtype across every mono layer, as _state_len() does
    for the conv window; it goes into every K1 key.
  • _k1 checks each step's pool against the build: a pool of the other dtype would be read at
    the wrong element size.
  • _recurrent_state() refuses anything but fp32 or fp16, which would be reinterpreted rather
    than converted.
• config.py: config_refusal admits fp16 as well as fp32.
• recipes/Agentic-Kimi-K3.md: the C1 launch no longer sets ATOM_GDN_SSM_DTYPE=fp32.

Performance

8x MI355X, the recipe's C1 launch on main, against the same launch with this PR (the fp32
override gone):

	main (fp32 state)	this PR (fp16 state)
KV pool, worst rank	7158 blocks	10407 blocks (+45%)
weights, worst rank	202.01 GB	202.01 GB
median ITL, waves 1 / 2	14.71 / 14.78 ms	14.87 / 14.92 ms
K1 per target forward (changed by this PR)	2443 us	2422 us
K2 per target forward (unchanged)	7799 us	7668 us

The fp16 state pool is half the size, so with the same weights the worst rank fits 45% more KV
blocks. The median decode step does not measurably change: median ITL is 0.15 ms (1%) longer in
both waves, within the step's run-to-run variation, which K2 dominates (K2, which this PR does
not touch, is 131 us faster in this PR's trace, and moves as much between launches of the same
build). K1, the only kernel this PR changes, is 21 us (0.9%) faster per forward; its time varies
by under 10 us between steps and launches, so that gain is likely real.

The extra blocks are capacity, not speed: at 128 tokens a block, the worst rank's pool grows from
916k to 1.33M tokens, and no request in this benchmark needs more than 72 blocks.

Method

• rocm/atom-dev:nightly_202611011508-rocm10.0 with ATOM main 606aceb11 mounted over the
  image's, with and without this commit. Every arm preloads the image's
  /usr/lib/x86_64-linux-gnu/libgomp.so.1: main's weight loader (feat(mono): shared DSA indexer mono launch; Kimi-K3 mono to 64 rows, plain decode and ReplaySSM #2871) looks up
  omp_set_num_threads among the process's global symbols, this image's torch loads libgomp with
  local scope, and without the preload no rank loads.
• recipes/Agentic-Kimi-K3.md's C1 launch: TP8, DCP1, DSpark with 7 speculative tokens, synthetic
  acceptance length 3.84, mono decode on, the five dense projections excluded from ptpc_fp8;
  ATOM_GDN_SSM_DTYPE=fp32 on main, the default fp16 with this PR.
• atom.benchmarks.benchmark_serving, random dataset, input 8192 and output 1024 tokens with
  range ratio 0.8, --ignore-eos, 5 prompts at concurrency 1, seed 0: a warmup wave, then two
  timed waves. With speculative decoding a request gets one streamed chunk per step, so the median
  ITL is the decode step period.
• Weights (peak_torch) and KV blocks are from the server logs. K1 and K2 times are from a 3 s
  torch-profiler window 20 s into a third wave: medians over the target forward's graph replays
  in rank 0's trace (165 and 167).

Limitations

• Mono's decode on an fp16 pool is not bit-identical to the fp32 build: each stored state is
  rounded to fp16, as production's is. GSM8K (below) is the accuracy check.
• The recipe runs mono at C1 only, where ReplaySSM is off, so the ReplaySSM checkpoint load and
  flush are checked only by the fp32-twin probe below, not by a recipe band.
• A bf16 pool is still refused.
• Range is no new constraint: production keeps the same state in fp16. (With k l2-normalized,
  beta in (0, 1) and a decay of at most 1, a token adds at most v_j^2 to the squared norm of
  value channel j's state column.)

Testing

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

• New CPU tests in test_kimi_k3_mono_state_pools.py, test_kimi_k3_mono_config_refusal.py
  and test_kimi_k3_mono_mla_k2.py: the config gate admits fp32 and fp16 and refuses bf16; an
  all-fp16 pool set reads as fp16; every K1 key carries the dtype; layers that disagree are
  refused; a step refuses a pool of the other dtype.
• Kimi-K3 and mono tests in the image: 262 passed on CPU (251 on main) and 221 on GPU with
  none skipped. The image's site-packages tests package (from cufile_python) shadows ATOM's
  and was masked; unmasked, 2 tests fail, on main too.
• ISA, compiled for gfx950: with an fp32 pool, all 27 K1 and K2 builds compared (1 to 64
  rows, spec verify, plain decode and ReplaySSM) are main's instructions; each of the 17 fp16 K1
  builds has its fp32 build's VGPR, SGPR, spill and LDS counts.
• Against an fp32 twin, on hardware: in an eager served run at C2, before each fp16 K1 launch,
  the same build with state16=False runs on copies of everything K1 writes, from the pool
  widened exactly. The core, prefix, block bank, conv window and ReplaySSM records must be
  bit-equal, and the fp16 pool must equal the twin's fp32 pool rounded to fp16. 81,600 of 81,600
  launches pass (150 steps x 68 KDA layers x 8 ranks) in each of three runs: spec verify, plain
  decode, and spec verify with ReplaySSM forced on (78,880 of its launches flushed state).
• GSM8K, C2, mono on, against main with mono on (fp32 state): 0.9651 against 0.9674, 21 / 18
  questions right on one side only (McNemar p = 0.75); main with mono off scores 0.9613. 1319
  questions, 5-shot chat, greedy, no synthetic acceptance.
• black and ruff clean.

python3 -m pytest -q tests/models/test_kimi_k3_mono_state_pools.py \
  tests/models/test_kimi_k3_mono_config_refusal.py tests/models/test_kimi_k3_mono_mla_k2.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 2945 --add-label <label>

[valar-cop]

some changes needed..

e2e or accuracy of mono on fp16 pool with ReplaySSM on, and at DCP>1?

ci green.

[valar-cop]

with mono on by default this flips mesh cells to mono without measuring them, please pin ATOM_MONO_ENABLE=0 there or measure one..

[valar-cop]

this removes the only doc of ATOM_GDN_SSM_DTYPE, add a row in docs/environment_variables.md?