feat(k3-mono): run on an fp16 KDA recurrent state pool

atom/models/kimi_k3/mono/config.py

@@ -27,11 +27,12 @@
 * **TP sharding** -- ``SUPPORTED_TP``; ``Dims`` refuses any size that does not
   divide every width.
 * **weight layouts** -- ``weights.py`` checks each tensor it binds.
-* **KV dtype** -- the conv state must be bf16 and the recurrent state fp32.
-  ``ATOM_GDN_SSM_DTYPE`` is checked here, but that is the *request*; the pools
-  that were actually allocated are checked in ``runner._conv_state`` and
-  ``runner._recurrent_state`` -- at build for the first mono layer, and again
-  where each step packs them.
+* **KV dtype** -- the conv state must be bf16 and the recurrent state fp32 or
+  fp16 (computed in fp32 either way: an fp16 pool is widened on load and each
+  state written back rounded to fp16). ``ATOM_GDN_SSM_DTYPE`` is checked here, but
+  that is the *request*; the pools that were actually allocated are checked in
+  ``runner._conv_state`` and ``runner._recurrent_state`` -- at build for every
+  mono layer, and again, against the build, where each step packs them.
 """
 
 from dataclasses import dataclass
@@ -249,10 +250,10 @@ def config_refusal(atom_config) -> str | None:
         common_refusals(atom_config, Dims, SUPPORTED_TP, refuse_decode_cp=False)
         + resident_grid_refusals(atom_config)
     ) + [
-        # the KDA recurrence reads and writes fp32 state in place; a narrower
-        # pool would be read as fp32 rather than converted
+        # the KDA recurrence reads and writes the state in place as fp32 or as
+        # fp16; any other pool would be reinterpreted rather than converted
         (
-            envs.ATOM_GDN_SSM_DTYPE == "fp32",
+            envs.ATOM_GDN_SSM_DTYPE in ("fp32", "fp16"),
             f"ATOM_GDN_SSM_DTYPE={envs.ATOM_GDN_SSM_DTYPE}",
         ),
         # the model's own widths, which the kernels bake

atom/models/kimi_k3/mono/kernels/attention/kda.py

@@ -262,6 +262,9 @@ class KdaPreBuild:
     timeline: bool = field(default=False, metadata={"sym": "t"})
     # a block-write layer's block: the updated prefix stored there
     write_idx: int = field(default=-1, metadata={"sym": "wi"})
+    # the recurrent state pool is fp16, else fp32: widened on load, every state
+    # written back rounded to fp16 (``state_loads`` / ``store_states``)
+    state16: bool = field(default=False, metadata={"sym": "h"})
     # ReplaySSM (``fla_ops.replayssm``): the record buffer's length (its
     # ``cache_len``), 0 without. A request's state is its checkpoint plus the
     # (k, u, g) records committed since (``write_pos``), rebuilt here; each
@@ -1029,12 +1032,30 @@ def state_late(s: int) -> bool:
     return s >= 8
 
 
-def state_loads(k):
-    """A task's initial state rows ``k["rows"]`` into ``k["st"]``."""
-    k["st"] = [
-        fx.Vector(bo.buffer_load(rsrc(k["st_base"]), off, vec_width=4, dtype=T.f32))
-        for off in k["head_offs"]
-    ]
+def state_at(c, slot):
+    """Where ``slot``'s recurrent state starts: 4 B an element, 2 an fp16 pool's."""
+    a = c["args"]
+    return a["rstate"] + fx.Int64(slot) * fx.Int64(a["rs_seq"]) * (
+        2 if c["state16"] else 4
+    )
+
+
+def state_loads(c, k):
+    """A task's initial state rows ``k["rows"]`` into ``k["st"]``, fp32."""
+    if const_expr(c["state16"]):
+        k["st"] = [
+            fx.Vector(
+                bo.buffer_load(rsrc(k["st_base"]), off // 2, vec_width=2, dtype=T.i32)
+            )
+            .bitcast(fx.Float16)
+            .to(fx.Float32)
+            for off in k["head_offs"]
+        ]
+    else:
+        k["st"] = [
+            fx.Vector(bo.buffer_load(rsrc(k["st_base"]), off, vec_width=4, dtype=T.f32))
+            for off in k["head_offs"]
+        ]
 
 
 def fb_loads(c, k):
@@ -1097,11 +1118,9 @@ def kda_task_inputs(c, req, h, m0, channels, rows):
     k0 = (tid % 32) * 4
     k["rows"] = rows
     k["head_offs"] = [(h * HD + r) * HD + k0 for r in rows]
-    k["st_base"] = (
-        a["rstate"] + fx.Int64(fx.max(k["slot0"], 0)) * fx.Int64(a["rs_seq"]) * 4
-    )
+    k["st_base"] = state_at(c, fx.max(k["slot0"], 0))
     if const_expr(not state_late(c["S"])):
-        state_loads(k)
+        state_loads(c, k)
     # f_b: wave w's rows 16 w .. of head h, a lane 64 B of its row
     fr = h * HD + wave * 16 + lane % 16
     fb = (fr * HD + (lane // 16) * 32) // 2
@@ -1156,7 +1175,7 @@ def kda_gates(c, k, req, h):
     if const_expr(L > 1):
         fb_loads(c, k)
     if const_expr(state_late(c["S"])):
-        state_loads(k)
+        state_loads(c, k)
     if const_expr(not c["kda_head"]):
         gpu.barrier()
     gates_compute(c, k, h)
@@ -1257,12 +1276,14 @@ def kda_recur(c, k, req, h, ch):
 
 def store_states(c, k, slot, stv):
     """The task's state rows ``stv`` (a chunk's each) into ``slot``."""
-    a = c["args"]
-    dst = rsrc(a["rstate"] + fx.Int64(slot) * fx.Int64(a["rs_seq"]) * 4)
+    dst = rsrc(state_at(c, slot))
     for m in range_constexpr(len(stv)):
-        bo.buffer_store(
-            fx.Vector.from_elements(stv[m], fx.Float32), dst, k["head_offs"][m]
-        )
+        st = fx.Vector.from_elements(stv[m], fx.Float32)
+        if const_expr(c["state16"]):
+            st16 = st.to(fx.Float16).bitcast(fx.Int32)
+            bo.buffer_store(st16, dst, k["head_offs"][m] // 2)
+        else:
+            bo.buffer_store(st, dst, k["head_offs"][m])
 
 
 @traced
@@ -1731,6 +1752,7 @@ def kda_pre(
             "lb": key.lower_bound,
             "kda0": key.kda0,
             "write_idx": key.write_idx,
+            "state16": key.state16,
             "put": mb.put,
             "poll": mb.poll,
             "red": lds.red.ptr,

atom/models/kimi_k3/mono/runner.py

@@ -204,6 +204,7 @@ def _bind(self, model, drafter=None, metadata_builder=None) -> None:
         first = next(p for p in self.plans if p.mono)
         self.device = first.kda.w_in.device
         self.state_len = self._state_len()
+        self.state16 = self._state_dtype() == torch.float16
         self.replay_cap = self._replay_cap()
         self._replay_tables = {}  # layer id -> (K1's ReplaySSM table, its pointers)
         # the DSA layers' indexers: static checks every rank shares, so no
@@ -265,6 +266,21 @@ def _state_len(self) -> int:
             )
         return lens.pop()
 
+    def _state_dtype(self) -> torch.dtype:
+        """The recurrent state pools' dtype, a build constant (``state16``) for the
+        same reason ``state_len`` is one."""
+        dtypes = {
+            self._recurrent_state(plan.layer_id).dtype
+            for plan in self.plans
+            if plan.mono
+        }
+        if len(dtypes) != 1:
+            raise MonoUnsupported(
+                "mono layers disagree on the recurrent state dtype: "
+                f"{sorted(map(str, dtypes))}"
+            )
+        return dtypes.pop()
+
     def _replay_cap(self) -> int:
         """ReplaySSM's record buffer length (its ``cache_len``), 0 without:
         read off the buffers allocated, every mono layer's checked as
@@ -302,17 +318,18 @@ def _recurrent_state(self, layer_id: int) -> torch.Tensor:
 
         `config_refusal` checks `ATOM_GDN_SSM_DTYPE`, which is the *request*;
         this is the pool that was actually allocated. The recurrence reads and
-        writes it in place as fp32 (see `_allocate`), so a narrower pool is
+        writes it in place as fp32, or as fp16 under ``state16`` (widened to
+        fp32 on load, rounded to nearest on every store), so any other dtype is
         reinterpreted rather than converted -- wrong numbers, no error. The
         conv state next door has been checked all along; this one was not,
         although the module docstring claimed both were.
         """
         cache = get_forward_context().kv_cache_data[f"layer_{layer_id}"]
         state = cache.v_cache
-        if state.dtype != torch.float32:
+        if state.dtype not in (torch.float32, torch.float16):
             raise MonoUnsupported(
                 f"layer {layer_id}: recurrent state {state.dtype}, the "
-                "recurrence reads it in place as float32"
+                "recurrence reads it in place as float32 or float16"
             )
         return state
 
@@ -395,6 +412,7 @@ def _k1_keys(self, rows: int, qlen: int) -> dict[int, k1.KdaPreBuild]:
                 onorm_eps=p.kda.onorm_eps,
                 lower_bound=p.kda.lower_bound,
                 write_idx=p.kda.write_idx,
+                state16=self.state16,
                 replay_cap=self.replay_cap,
                 replay_window=window,
                 **k1.gemv_staging(rows),
@@ -618,6 +636,11 @@ def _k1(self, build, plan, prefix, delta, blocks, metadata, taps=None):
                 f"built {key.state_len}"
             )
         rstate = self._recurrent_state(plan.layer_id)
+        if (rstate.dtype == torch.float16) != key.state16:
+            raise MonoUnsupported(
+                f"layer {plan.layer_id}: recurrent state {rstate.dtype} != the "
+                f"built {torch.float16 if key.state16 else torch.float32}"
+            )
         slots, accepted = state_slots(metadata, self.ones)
         core = prefix.new_empty(prefix.size(0), self.dims.proj)
         p = torch.Tensor.data_ptr

recipes/Agentic-Kimi-K3.md

@@ -187,14 +187,10 @@ case "${CONC}" in
     SPEC_DECODE_ACCEPTANCE_LENGTH=3.84
     if [[ "${CONC}" == "1" ]]; then
       ATOM_MONO_ENABLE=1
-      # Both overrides are C1-only, because mono is C1-only. They are
-      # preconditions, not tuning: without them mono refuses at build time and
-      # C1 silently serves the paged path, which is what `ATOM_MONO_ENABLE=1`
+      # The override is C1-only, because mono is C1-only. It is a
+      # precondition, not tuning: without it mono refuses at build time and C1
+      # silently serves the paged path, which is what `ATOM_MONO_ENABLE=1`
       # above claims it does not do.
-      #
-      # The mono KDA kernel reads the SSM state pool as fp32; the band default
-      # above requests fp16, which `kimi_k3/mono/config.py` refuses outright.
-      ATOM_GDN_SSM_DTYPE=fp32
       # The five projections mono loads as plain bf16 (`kimi_k3/mono/weights.py`:
       # `_plain_bf16` on in_proj, f_b_proj, o_proj and the two shared-expert
       # projections). Online ptpc_fp8 would quantize them and the bf16 load

tests/models/test_kimi_k3_mono_config_refusal.py

@@ -165,3 +165,16 @@ def test_the_query_cap_bounds_the_speculative_width(servable, draft_tokens, refu
         method="dspark", num_speculative_tokens=draft_tokens
     )
     assert (config_refusal(atom_config) is not None) is refused
+
+
+@pytest.mark.parametrize(
+    ("ssm", "refused"), [("fp32", False), ("fp16", False), ("bf16", True)]
+)
+def test_the_recurrent_state_dtype_gate(servable, monkeypatch, ssm, refused):
+    # K1 reads the pool in place as fp32 or fp16; a bf16 one would be
+    # reinterpreted rather than converted
+    monkeypatch.setattr(k3_config.envs, "ATOM_GDN_SSM_DTYPE", ssm)
+    reason = config_refusal(_atom_config())
+    assert (reason is not None) is refused
+    if refused:
+        assert f"ATOM_GDN_SSM_DTYPE={ssm}" in reason

tests/models/test_kimi_k3_mono_mla_k2.py

@@ -265,6 +265,7 @@ def test_k2_builds_cover_full_attention_layers_and_k1_builds_do_not(runner_modul
     runner = object.__new__(runner_module.K3MonoDecodeRunner)
     runner.debug = False
     runner.state_len = 3
+    runner.state16 = False
     runner.replay_cap = 0  # no ReplaySSM record buffers
     runner.plans = [
         _plan(1, kda=_stub_kda(), moe=_stub_moe()),
@@ -278,6 +279,18 @@ def test_k2_builds_cover_full_attention_layers_and_k1_builds_do_not(runner_modul
     assert keys[1] == keys[3]
 
 
+@pytest.mark.parametrize("fp16", [False, True])
+def test_k1_keys_carry_the_recurrent_state_dtype(runner_module, fp16):
+    # K1 reads the pool at its element size, so the dtype is part of the key
+    runner = object.__new__(runner_module.K3MonoDecodeRunner)
+    runner.debug = False
+    runner.state_len = 3
+    runner.state16 = fp16
+    runner.replay_cap = 0
+    runner.plans = [_plan(1, kda=_stub_kda(), moe=_stub_moe())]
+    assert [key.state16 for key in runner._k1_keys(8, 8).values()] == [fp16]
+
+
 # ----------------------------------------------------------- forward_pre_o
 def test_forward_pre_o_is_forward_without_o_proj():
     try:

tests/models/test_kimi_k3_mono_state_pools.py

@@ -24,6 +24,7 @@
 launched at all.
 """
 
+from collections import namedtuple
 from types import SimpleNamespace
 
 import pytest
@@ -104,3 +105,48 @@ def test_a_non_mono_layers_pool_is_not_checked(monkeypatch):
     caches[3] = _pools(WINDOW + 7, conv=torch.float16)
     runner = _runner(monkeypatch, caches)
     assert runner._state_len() == WINDOW
+
+
+def test_an_fp16_recurrent_pool_is_admitted(monkeypatch):
+    # the recurrence widens fp16 to fp32 and rounds each state it writes back;
+    # a bf16 pool is refused (`test_a_later_mono_layers_bad_pool_is_refused_at_bind`)
+    runner = _runner(monkeypatch, {lid: _pools(rec=torch.float16) for lid in range(4)})
+    assert runner._state_len() == WINDOW
+    assert runner._state_dtype() == torch.float16
+
+
+def test_mono_layers_that_disagree_on_the_recurrent_dtype_are_refused(monkeypatch):
+    # like the window, the dtype is one build constant baked into every K1 key
+    caches = {lid: _pools() for lid in range(4)}
+    caches[2] = _pools(rec=torch.float16)
+    runner = _runner(monkeypatch, caches)
+    with pytest.raises(MonoUnsupported, match="disagree on the recurrent state dtype"):
+        runner._state_dtype()
+
+
+class _PastTheChecks(Exception):
+    pass
+
+
+def _past_the_checks(*_):
+    raise _PastTheChecks
+
+
+@pytest.mark.parametrize("state16", [False, True])
+@pytest.mark.parametrize("rec", [torch.float32, torch.float16])
+def test_a_step_takes_only_the_pool_its_build_was_made_for(monkeypatch, rec, state16):
+    # K1 reads the pool at the element size its key was built for, so `_k1`
+    # checks the dtype against the key, as it checks the window; `state_slots`
+    # is the first thing it does past the checks
+    runner = _runner(monkeypatch, {lid: _pools(rec=rec) for lid in range(4)})
+    runner.ones = None
+    monkeypatch.setattr(runner_mod, "state_slots", _past_the_checks)
+    key = namedtuple("Key", "state_len state16")(WINDOW, state16)
+    build = SimpleNamespace(k1_keys={1: key}, k1_launchers={key: None})
+    plan = SimpleNamespace(layer_id=1, kda=None)
+    if (rec == torch.float16) == state16:
+        with pytest.raises(_PastTheChecks):
+            runner._k1(build, plan, None, None, None, None)
+    else:
+        with pytest.raises(MonoUnsupported, match="recurrent state .* != the built"):
+            runner._k1(build, plan, None, None, None, None)