fix(kda): preserve SM90 CP transitions in FP32

USAGE.md

@@ -213,6 +213,11 @@ Pass `use_intracard_cp` to the Hopper prefill:
 
 Works with **any sequence length** (non-CHUNK-aligned is handled internally) and **dense or varlen** input. The auto decision uses the device SM count plus two tunables in `cula/ops/kda/sm90/cp/plan.py` (`CULA_KDA_CP_RERUN_RATIO`, `CULA_KDA_CP_AUTO_MIN_SEG_TILES`).
 
+The SM90 CP pre-scan keeps its segment transition matrix canonically in FP32
+registers and maintains a BF16 mirror only for Tensor Core products. This
+prevents transition error from compounding when many slow-decay segments are
+composed.
+
 ```python
 o, final_state = flashkda_prefill(
     q=q, k=k, v=v, g=g, beta=beta, A_log=A_log, dt_bias=dt_bias,

cula/ops/kda/sm90/cp/pre_scan.py

@@ -91,6 +91,9 @@ def pre_scan_kernel(
     sINV = smem.allocate_tensor(cutlass.BFloat16, cc_stage_layout, 128)
     sState = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128)
     sM = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128)
+    # sState intentionally matches the serial FlashKDA BF16 recurrence.  sM
+    # is only its Tensor Core operand mirror; the CP-only affine transform is
+    # kept canonically in a distributed FP32 register fragment below.
     sGt = smem.allocate_tensor(cutlass.Float32, cute.make_layout((D, 1, STAGES), stride=(1, D, D)), 128)
     sBeta = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout((CHUNK, 1, STAGES), stride=(1, 64, 64)), 128)
     sMbar = smem.allocate_tensor(cutlass.Int64, cute.make_layout((STAGES,)), 16)
@@ -154,12 +157,10 @@ def pre_scan_kernel(
     gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout)
     gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout)
 
-    # sState=0, sM=I
+    # sState=0.  The distributed FP32 M accumulator initializes sM below.
     if tidx < D:
         for e in cutlass.range_constexpr(D):
             sState[tidx, e] = cutlass.BFloat16(0.0)
-            sM[tidx, e] = cutlass.BFloat16(0.0)
-        sM[tidx, tidx] = cutlass.BFloat16(1.0)
     cute.arch.barrier()
 
     mma_atom = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16))
@@ -223,10 +224,26 @@ def pre_scan_kernel(
     tCrKrA_state_blk_cv = smem_thr_copy_A_state.retile(tCrKrA_state_blk)
     tCrUpd_blk = thr_mma_state.make_fragment_C(tiled_mma_state.partition_shape_C((CHUNK, D)))
     sState_blk_tile = cute.flat_divide(sState, (CHUNK, D))
-    sM_blk_tile = cute.flat_divide(sM, (CHUNK, D))
     coord_state_blk = cute.make_identity_tensor((CHUNK, D))
     tCcState_blk = thr_mma_state.partition_C(coord_state_blk)
 
+    # Four compute warps jointly own the complete transition matrix.  Keeping
+    # this fragment in registers preserves the original shared-memory
+    # footprint and avoids round-tripping the canonical value through SMEM.
+    tCrKrA_m = thr_mma_state.make_fragment_A(thr_mma_state.partition_A(sKr_T_view_s0))
+    tCrKrA_m_cv = smem_thr_copy_A_state.retile(tCrKrA_m)
+    tCrMAcc = thr_mma_state.make_fragment_C(tiled_mma_state.partition_shape_C((D, D)))
+    tCcM = thr_mma_state.partition_C(cute.make_identity_tensor((D, D)))
+    tCsM = thr_mma_state.partition_C(sM)
+    if warp_idx < LOAD_WARP_IDX:
+        for i in cutlass.range_constexpr(cute.size(tCrMAcc)):
+            ii: cutlass.Constexpr[int] = i
+            row, col = tCcM[ii]
+            value = cutlass.Float32(1.0) if row == col else cutlass.Float32(0.0)
+            tCrMAcc[ii] = value
+            tCsM[ii] = cutlass.BFloat16(value)
+    cute.arch.barrier()
+
     tCrU_T = thr_mma.make_fragment_B(thr_mma.partition_B(sKr_T_ref))
 
     sINV_s0 = sINV[(None, None, 0)]
@@ -406,30 +423,21 @@ def pre_scan_kernel(
                 ii: cutlass.Constexpr[int] = i
                 tCrU_T_post_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_post_u32[ii]))
 
-            # State update': M = M*gt + kr^T @ U
-            for mi in cutlass.range_constexpr(M_BLOCKS):
-                sKr_T_blk_s = sKr_T_blk_tile_s[None, None, mi, 0]
-                cute.copy(
-                    smem_tiled_copy_A_state,
-                    smem_thr_copy_A_state.partition_S(sKr_T_blk_s),
-                    tCrKrA_state_blk_cv,
-                )
-                tCrUpd_blk.fill(0.0)
-                cute.gemm(tiled_mma_state, tCrUpd_blk, tCrKrA_state_blk, tCrU_T_post, tCrUpd_blk)
-
-                sM_blk = sM_blk_tile[None, None, mi, 0]
-                tCsM_blk = thr_mma_state.partition_C(sM_blk)
-                m_frag_blk = cute.make_fragment_like(tCsM_blk, cutlass.BFloat16)
-                gt_frag_blk_m = cute.make_fragment_like(tCsM_blk, cutlass.Float32)
-                m_off2: cutlass.Constexpr[int] = mi * CHUNK
-                for i in cutlass.range_constexpr(cute.size(m_frag_blk)):
-                    ii: cutlass.Constexpr[int] = i
-                    m_frag_blk[ii] = tCsM_blk[ii]
-                    gt_frag_blk_m[ii] = sGt_s[m_off2 + tCcState_blk[ii][0]]
-                for i in cutlass.range_constexpr(cute.size(tCrUpd_blk)):
-                    ii: cutlass.Constexpr[int] = i
-                    old = cutlass.Float32(m_frag_blk[ii]) * gt_frag_blk_m[ii]
-                    tCsM_blk[ii] = cutlass.BFloat16(old + tCrUpd_blk[ii])
+            # State update': M = M*gt + kr^T @ U.  Accumulate into the
+            # persistent FP32 register fragment and publish only a BF16 image
+            # for the next chunk's kd@M Tensor Core operand.
+            cute.copy(
+                smem_tiled_copy_A_state,
+                smem_thr_copy_A_state.partition_S(sKr_T_s),
+                tCrKrA_m_cv,
+            )
+            for i in cutlass.range_constexpr(cute.size(tCrMAcc)):
+                ii: cutlass.Constexpr[int] = i
+                tCrMAcc[ii] = tCrMAcc[ii] * sGt_s[tCcM[ii][0]]
+            cute.gemm(tiled_mma_state, tCrMAcc, tCrKrA_m, tCrU_T_post, tCrMAcc)
+            for i in cutlass.range_constexpr(cute.size(tCrMAcc)):
+                ii: cutlass.Constexpr[int] = i
+                tCsM[ii] = cutlass.BFloat16(tCrMAcc[ii])
 
             cute.arch.barrier(barrier_id=1, number_of_threads=128)
             cute.arch.fence_view_async_shared()
@@ -448,7 +456,11 @@ def pre_scan_kernel(
     if tidx < D:
         for d_out in cutlass.range_constexpr(D):
             b_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32(sState[tidx, d_out])
-            m_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32(sM[tidx, d_out])
+    if warp_idx < LOAD_WARP_IDX:
+        for i in cutlass.range_constexpr(cute.size(tCrMAcc)):
+            ii: cutlass.Constexpr[int] = i
+            row, col = tCcM[ii]
+            m_state_g[state_base_f + cutlass.Int32(col * D + row)] = tCrMAcc[ii]
 
 
 @cute.jit

tests/test_kda_sm90_intracard_cp.py

@@ -67,6 +67,11 @@ def _rel_rmse(a, b):
     return (a - b).pow(2).mean().sqrt().item() / max(b.pow(2).mean().sqrt().item(), 1e-6)
 
 
+def _rel_l2(a, b):
+    a, b = a.float(), b.float()
+    return ((a - b).norm() / b.norm().clamp_min(1e-30)).item()
+
+
 def _assert_cp_matches(actual, ref, name):
     rrmse = _rel_rmse(actual, ref)
     assert rrmse < TOL_RMSE, f"{name}: rel_rmse {rrmse:.2e} >= {TOL_RMSE}"
@@ -168,6 +173,121 @@ def test_trivial_manual_plan_falls_back_to_serial():
     assert torch.equal(fin_fb, fin_ref)
 
 
+@needs_cuda
+@pytest.mark.kda_slow
+def test_slow_decay_64k_cp_relative_l2_matches_fla_scale():
+    """Regression for #126 using the issue's public synthetic reproducer."""
+    torch.manual_seed(0)
+    h = 1
+    shape = (1, 65536, h, D)
+    q = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    g = torch.full(shape, -10.0, dtype=torch.bfloat16, device="cuda")
+    beta = torch.zeros(shape[:-1], dtype=torch.bfloat16, device="cuda")
+    A_log = torch.zeros(h, dtype=torch.float32, device="cuda")
+    dt_bias = torch.zeros(h, D, dtype=torch.float32, device="cuda")
+
+    def run(*, split):
+        out = torch.empty_like(v)
+        state = torch.empty(1, h, D, D, dtype=torch.float32, device="cuda")
+        if split == 1:
+            flash_kda_fwd(q, k, v, g, beta, SCALE, out, A_log, dt_bias, LB, final_state=state)
+        else:
+            intracard_prefill(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                SCALE,
+                out,
+                A_log,
+                dt_bias,
+                LB,
+                final_state=state,
+                s_split=split,
+                allow_fallback=False,
+            )
+        return out, state
+
+    out_serial, state_serial = run(split=1)
+    for split in (2, 4, 8, 16):
+        out_cp, state_cp = run(split=split)
+
+        if split == 2:
+            # There is only one carry handoff, so the CP path is exactly the
+            # same BF16 recurrence as serial FlashKDA.
+            assert torch.equal(out_cp, out_serial)
+            assert torch.equal(state_cp, state_serial)
+        else:
+            # This public reproducer has an almost-zero segment transition;
+            # it mostly measures the different BF16 state-rounding points.
+            # Keep the divergence bounded while the transition-sensitive
+            # regression below directly covers the #126 amplification.
+            assert _rel_l2(out_cp, out_serial) < 8e-3, f"output split={split}"
+            assert _rel_l2(state_cp, state_serial) < 8e-3, f"state split={split}"
+
+
+@needs_cuda
+@pytest.mark.kda_slow
+def test_slow_decay_16k_cp_transition_matches_fp64():
+    """The CP-only affine transition must not accumulate in BF16 (#126)."""
+    from fla.ops.kda import fused_recurrent_kda
+
+    torch.manual_seed(123)
+    h = 1
+    shape = (1, 16384, h, D)
+    q = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    g = torch.full(shape, -10.0, dtype=torch.bfloat16, device="cuda")
+    beta = torch.full(shape[:-1], -10.0, dtype=torch.bfloat16, device="cuda")
+    A_log = torch.zeros(h, dtype=torch.float32, device="cuda")
+    dt_bias = torch.zeros(h, D, dtype=torch.float32, device="cuda")
+    h0 = torch.randn(1, h, D, D, dtype=torch.float32, device="cuda") * 0.1
+
+    out = torch.empty_like(v)
+    state = torch.empty_like(h0)
+    intracard_prefill(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        SCALE,
+        out,
+        A_log,
+        dt_bias,
+        LB,
+        initial_state=h0,
+        final_state=state,
+        s_split=16,
+        allow_fallback=False,
+    )
+
+    g64 = LB * torch.sigmoid(torch.exp(A_log.double()).view(1, 1, h, 1) * (g.double() + dt_bias.double().view(1, 1, h, D)))
+    beta64 = torch.sigmoid(beta.double())
+    ref_out, ref_state = fused_recurrent_kda(
+        q.double(),
+        k.double(),
+        v.double(),
+        g64,
+        beta64,
+        scale=SCALE,
+        initial_state=h0.double(),
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=True,
+        use_gate_in_kernel=False,
+        transpose_state_layout=True,
+    )
+
+    # BF16 sM gives about 1.45e-1 / 5.60e-1 on this case.  Keeping the
+    # canonical transition in FP32 reduces that to about 5.5e-2 / 1.4e-2.
+    assert _rel_l2(out, ref_out) < 8e-2
+    assert _rel_l2(state, ref_state) < 3e-2
+
+
 @needs_cuda
 @pytest.mark.kda_fast
 @pytest.mark.parametrize("s_split", [2, 4, 7])