From b82994e95c7f3d347a4a427d0ac8063d0d270a3e Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Thu, 3 Sep 2026 06:05:26 -0500 Subject: [PATCH] perf(kimi-k3): make the KDA temporal state dtype configurable KDA decode is state-bandwidth bound: 69 layers each stream a [12, 128, 128] fp32 recurrent state per token, and the fused gating kernel already runs at 71-90% of achievable bandwidth. Halving the state's element size is the lever that is left. ATOM_KDA_SSM_DTYPE ("fp32" | "fp16" | "bf16", default "fp32") picks the storage dtype of the temporal state pool for the KDA families (kimi_linear and glm5_next_text). The default is unchanged, so main behaves identically. 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. 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. Prefill needs ROCm/aiter#5249. 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 and k are L2-normalized in-kernel so the state is O(1) and bf16's range buys nothing, while its three fewer mantissa bits cost roughly 8x the error. Measured on MI355X, Kimi-K3 TP8, per recipes/Kimi-K3.md: GSM8K, 1319 questions, 5-shot, greedy (strict = flexible): fp32 0.9659 fp16 0.9591 +- 0.0055 Serving, 256 in / 1024 out, --ignore-eos: fp32 fp16 c=32 tok/s 1144.71 1164.95 +1.8% TPOT ms 26.91 26.55 -1.3% ITL ms 33.86 32.13 -5.1% c=64 tok/s 1902.21 1913.54 +0.6% TPOT ms 32.34 31.75 -1.8% ITL ms 40.83 39.48 -3.3% State pool: 56.17 -> 29.04 MB per slot, 3.35 -> 1.73 GB for 64 slots, which the KV pool takes over (53.58 -> 55.61 GB, +1236 blocks). 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. Co-Authored-By: Claude Opus 5 --- atom/model_ops/attentions/gdn_attn.py | 29 +++++++++++++------ .../model_ops/attentions/kimi_mla_gdn_attn.py | 9 +++--- .../model_ops/fla_ops/fused_sigmoid_gating.py | 6 +++- atom/utils/envs.py | 5 ++++ tests/test_kda_layout_id.py | 5 ++-- 5 files changed, 38 insertions(+), 16 deletions(-) diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index ea84f7e6d2..dcec82c936 100644 --- a/atom/model_ops/attentions/gdn_attn.py +++ b/atom/model_ops/attentions/gdn_attn.py @@ -42,6 +42,12 @@ # both kinds cannot number them into the same space. LINEAR_STATE_ROWS = "linear_state" +_GDN_SSM_DTYPES = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + class GDNAttentionBackend(AiterBackend): @staticmethod @@ -316,17 +322,21 @@ def _state_shape( return conv_state_shape, temporal_state_shape def _state_dtypes(self) -> tuple[torch.dtype, torch.dtype]: - # KDA recurrence accumulates in fp32 and aiter's chunk_kimi_delta_attn - # reads the state back verbatim, so the temporal state must be fp32 for - # every KDA model (Kimi-Linear and GLM-5.3-Flash). + # The KDA recurrence accumulates in fp32 whatever the pool stores, so + # the KDA models pick their storage dtype; everyone else keeps the + # state at the model dtype. if getattr(self.model_runner.config.hf_config, "model_type", None) in ( "kimi_linear", "glm5_next_text", ): - return ( - self.model_runner.config.torch_dtype, - torch.float32, - ) + requested = envs.ATOM_GDN_SSM_DTYPE + temporal_dtype = _GDN_SSM_DTYPES.get(requested) + if temporal_dtype is None: + raise ValueError( + f"ATOM_GDN_SSM_DTYPE={requested!r} is not one of " + f"{sorted(_GDN_SSM_DTYPES)}." + ) + return (self.model_runner.config.torch_dtype, temporal_dtype) return ( self.model_runner.config.torch_dtype, self.model_runner.config.torch_dtype, @@ -420,8 +430,9 @@ def state_transfer(self) -> StateTransfer: 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 fp32 v side is the one pool that breaks it, and - it overrides (`_KimiMLAGDNCommon.state_transfer`). + agreeing; kimi_linear's temporal side is `ATOM_GDN_SSM_DTYPE` and need + not, so it overrides unconditionally + (`_KimiMLAGDNCommon.state_transfer`). """ return StateTransfer.fork(1, readable_midstep=False) diff --git a/atom/model_ops/attentions/kimi_mla_gdn_attn.py b/atom/model_ops/attentions/kimi_mla_gdn_attn.py index 18c834b5db..b351b2a55d 100644 --- a/atom/model_ops/attentions/kimi_mla_gdn_attn.py +++ b/atom/model_ops/attentions/kimi_mla_gdn_attn.py @@ -95,10 +95,11 @@ def state_transfer(self) -> StateTransfer: own reasons, so the two agree meanwhile. Dtype-safe by construction, which a checkpoint cut from `h` would not - be here — `_state_dtypes` gives kimi_linear an fp32 v side. An image is - copied slot to slot with no kernel output in between, so that fp32 side - round-trips exactly. Both dtypes are named in the layout id, so a build - that changed either cannot read another's images. + be here — `_state_dtypes` gives kimi_linear a v side of its own dtype + (`ATOM_GDN_SSM_DTYPE`). An image is copied slot to slot with no kernel + output in between, so it round-trips exactly whatever that dtype is. + Both dtypes are named in the layout id, so a build that changed either + cannot read another's images. """ if not self._uses_paged_checkpoints(): return StateTransfer.fork(1) diff --git a/atom/model_ops/fla_ops/fused_sigmoid_gating.py b/atom/model_ops/fla_ops/fused_sigmoid_gating.py index c143fc7703..526c9dd0dd 100644 --- a/atom/model_ops/fla_ops/fused_sigmoid_gating.py +++ b/atom/model_ops/fla_ops/fused_sigmoid_gating.py @@ -230,7 +230,11 @@ def fused_sigmoid_gating_delta_rule_update( B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) + # 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) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" num_stages = 3 diff --git a/atom/utils/envs.py b/atom/utils/envs.py index bc229f77b9..1b16dafda9 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -256,6 +256,11 @@ # serial one at verify windows >= ~12 tokens (measured on gfx950), so # "auto" keeps practical MTP windows on the serial route. "ATOM_REPLAYSSM_ROUTE": lambda: os.getenv("ATOM_REPLAYSSM_ROUTE", "auto").lower(), + # "fp32" | "fp16" | "bf16". Storage dtype of the KDA temporal state pool, + # whose per-token traffic dominates KDA decode; the recurrence itself + # always accumulates in fp32. fp16 over bf16 when narrowing: the state is + # O(1), so bf16's range buys nothing and its short mantissa costs accuracy. + "ATOM_GDN_SSM_DTYPE": lambda: os.getenv("ATOM_GDN_SSM_DTYPE", "fp32").lower(), "ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_RMSNORM_QUANT": lambda: ( os.getenv("ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_RMSNORM_QUANT", "1") == "1" ), diff --git a/tests/test_kda_layout_id.py b/tests/test_kda_layout_id.py index 5ab584cc36..3a8264eafa 100644 --- a/tests/test_kda_layout_id.py +++ b/tests/test_kda_layout_id.py @@ -112,9 +112,10 @@ def test_the_conv_dtype(self): assert layout_of() != layout_of(dt_k=torch.float16) def test_the_ssm_dtype(self): - """The fp32 v side is the reason a PAGE copy round-trips exactly. A - build that narrowed it must not read this one's images.""" + """The v side is `ATOM_GDN_SSM_DTYPE`. fp16 and bf16 are the same size, + so the id is the only thing telling those two apart.""" assert layout_of() != layout_of(dt_v=torch.bfloat16) + assert layout_of(dt_v=torch.float16) != layout_of(dt_v=torch.bfloat16) def test_the_layer_count(self): assert layout_of() != layout_of(layers=68)