perf(kimi-k3): make the KDA temporal state dtype configurable

[XiaobingSuper]

What

ATOM_GDN_SSM_DTYPE ("fp32" | "fp16" | "bf16", default "fp32") picks the storage dtype of the KDA temporal (recurrent) state pool, for kimi_linear and glm5_next_text. The default is unchanged, so main behaves identically.

KDA decode is state-bandwidth bound — 69 layers each stream a [12, 128, 128] fp32 state per token, and the fused gating kernel already runs at 71-90% of achievable bandwidth. Halving the element size is the lever that is left.

The dtype is decided in one place, GDNStateMixin._state_dtypes(). Pool sizing, per-request allocation, the checkpoint plane shapes and the checkpoint layout id all derive their bytes from it, so nothing else needed to change; the layout id already names both dtypes, so a build that flips the variable cannot read another's checkpoint images.

No cast is introduced anywhere. The recurrence accumulates in fp32 whatever the pool stores, and every path — prefill, decode, spec-decode, ReplaySSM — reads and writes the state through the destination pointer's element type.

One other change: the gating kernel's BV cap now follows the state's element size. It was tuned for a 4-byte state; a 2-byte one wants twice the V per block to keep the same bytes in flight (HV=12, K=V=128, N=64: 18.2 → 16.1 us at BV=64), while fp32 is slightly worse there (24.6 → 25.2 us).

fp16 rather than bf16 for the narrow setting: q/k are L2-normalized in-kernel so the state is O(1) and bf16's range buys nothing, while its 3 fewer mantissa bits cost roughly 8× the error.

Depends on

ROCm/aiter#5249 — chunk_kimi_delta_attn (the prefill kernel) previously required an fp32 initial_state.

Accuracy

GSM8K, 1319 questions, 5-shot, greedy. strict-match and flexible-extract agree in both arms.

	exact_match
fp32 (default)	0.9659
fp16	0.9591 ± 0.0055

Both are inside the range recipes/Kimi-K3.md records as verified (0.9538–0.9591).

Performance

MI355X, Kimi-K3 TP8, per recipes/Kimi-K3.md. Serving benchmark, 256 in / 1024 out, --ignore-eos.

		fp32	fp16
c=32	output tok/s	1144.71	1164.95	+1.8%
	mean TPOT	26.91 ms	26.55 ms	−1.3%
	mean ITL	33.86 ms	32.13 ms	−5.1%
c=64	output tok/s	1902.21	1913.54	+0.6%
	mean TPOT	32.34 ms	31.75 ms	−1.8%
	mean ITL	40.83 ms	39.48 ms	−3.3%

Memory

	fp32	fp16
state per slot	56.17 MB	29.04 MB
state pool (64 slots)	3.35 GB	1.73 GB
KV pool	53.58 GB	55.61 GB (+1236 blocks)

Note on the default

This PR only adds the switch and the data. Flipping the default is the accuracy owner's call: GSM8K's short generations do not exercise long-context state accumulation, and the evidence that the error does not grow with sequence length is offline numerics rather than an end-to-end run.

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

[Copilot]

🟡 Changes recommended

There’s a concrete correctness issue in fused_sigmoid_gating_delta_rule_update where initial_state is treated as optional but dereferenced unconditionally.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds an environment-variable switch to control the storage dtype of the KDA temporal (recurrent) state pool for KDA models, enabling a lower-bandwidth decode path while keeping the default behavior unchanged.

Changes:

• Introduce ATOM_KDA_SSM_DTYPE (default fp32) and plumb it into GDNStateMixin._state_dtypes() for kimi_linear / glm5_next_text.
• Adjust the fused sigmoid gating kernel’s BV cap based on the state element size.
• Add/extend tests to validate accepted dtype values, default behavior, pool sizing, and layout-id differentiation.

File summaries

File	Description
tests/test_kda_ssm_dtype.py	New unit tests covering accepted env values, default, rejection of unknown values, and pool sizing impact.
tests/test_kda_layout_id.py	Extends layout-id assertions to distinguish fp16 vs bf16 temporal dtype (same size, different semantics).
atom/utils/envs.py	Adds ATOM_KDA_SSM_DTYPE env var definition and documentation.
atom/model_ops/fla_ops/fused_sigmoid_gating.py	Makes BV tuning depend on the state’s element size for bandwidth-bound KDA decode.
atom/model_ops/attentions/kimi_mla_gdn_attn.py	Updates paged-checkpoint transfer docstring to reflect configurable temporal dtype.
atom/model_ops/attentions/gdn_attn.py	Adds dtype mapping + env parsing/validation and uses it for KDA temporal-state dtype selection.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

It changes a performance-critical recurrent-state dtype and kernel tiling behavior that can affect numerical stability and runtime characteristics and warrants final human review.

Review details

Suppressed comments (1)

atom/model_ops/attentions/gdn_attn.py:446

• The comment says kimi_linear's temporal side “breaks it at any setting”, but with ATOM_KDA_SSM_DTYPE set to match config.torch_dtype (e.g. bf16), the two dtypes can still agree. This is documentation-only, but it’s misleading about when the midstep-exactness argument fails.

        Exact, not approximate, when it is turned back on: `h` is `k.new_empty`
        and `_state_dtypes` returns `config.torch_dtype`, so slicing `h` rounds
        exactly where a shortened forward would. That rests on the two dtypes
        agreeing; kimi_linear's temporal side is dtype-configurable
        (ATOM_KDA_SSM_DTYPE) and breaks it at any setting, so it overrides
        (`_KimiMLAGDNCommon.state_transfer`).

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

There are a few correctness/robustness issues to address (notably an API contract mismatch around initial_state being effectively required despite a None default, plus missing test coverage for the new env var).

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

atom/model_ops/fla_ops/fused_sigmoid_gating.py:237

• initial_state is annotated with a default of None, but this function unconditionally dereferences it (element_size(), dtype, stride). If a caller relies on the default it will raise an AttributeError; consider validating early and raising a clear ValueError (or make the parameter required).

    # State-bandwidth bound: a 2-byte state wants twice the V per block to keep
    # the same bytes in flight (HV=12, K=V=128, N=64: 18.2 -> 16.1 us, while
    # fp32 goes 24.6 -> 25.2).
    bv_cap = 64 if initial_state.element_size() <= 2 else 32
    BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), bv_cap)

atom/model_ops/attentions/gdn_attn.py:444

• Docstring grammar: “and need not” is missing the implied verb/object, which makes the sentence hard to parse. Suggest clarifying that the temporal dtype need not agree with the model dtype.

        agreeing; kimi_linear's temporal side is `ATOM_KDA_SSM_DTYPE` and need
        not, so it overrides unconditionally
        (`_KimiMLAGDNCommon.state_transfer`).

• Files reviewed: 5/5 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🟢 Approval recommended

The functional changes are localized and default-preserving, with validation and tests updated; only a minor docstring grammar nit was found.

Review details

• Files reviewed: 5/5 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🟢 Approval recommended

The functional changes are localized, validated with an updated layout-id test, and the only remaining feedback is a minor docstring grammar nit.

Review details

Suppressed comments (1)

atom/model_ops/attentions/gdn_attn.py:435

• Docstring reads "and need not" without an object; adding "agree" makes the sentence grammatical and clearer about why the override is unconditional.

        agreeing; kimi_linear's temporal side is `ATOM_GDN_SSM_DTYPE` and need
        not, so it overrides unconditionally
        (`_KimiMLAGDNCommon.state_transfer`).

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

[valarLip]

LGTM