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

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)
 

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)

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

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"
     ),

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)