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)