--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/test_triton.py
@@ -0,0 +1,15 @@
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def k(a, b, c, n):
+    pid = tl.program_id(0)
+    tl.store(c + pid, a[pid] + b[pid])
+
+a = torch.zeros(64, device='cuda')
+b = torch.ones(64, device='cuda')
+c = torch.zeros(64, device='cuda')
+k[(64,)](a, b, c, 64)
+print(c[:8])
+print('ok')

--- a//workspace/problems/02_kda_cutlass/test_triton.py
+++ b//workspace/problems/02_kda_cutlass/test_triton.py
@@ -3,9 +3,11 @@
 import triton.language as tl
 
 @triton.jit
-def k(a, b, c, n):
+def k(a_ptr, b_ptr, c_ptr, n):
     pid = tl.program_id(0)
-    tl.store(c + pid, a[pid] + b[pid])
+    a = tl.load(a_ptr + pid)
+    b = tl.load(b_ptr + pid)
+    tl.store(c_ptr + pid, a + b)
 
 a = torch.zeros(64, device='cuda')
 b = torch.ones(64, device='cuda')

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -0,0 +1,223 @@
+"""Kimi Delta Attention (chunk form) forward, Triton implementation."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kda_fwd_kernel(
+    Q_ptr, K_ptr, V_ptr, G_ptr, BETA_ptr, O_ptr,
+    scale,
+    stride_qb, stride_qt, stride_qh, stride_qk,
+    stride_kb, stride_kt, stride_kh, stride_kk,
+    stride_vb, stride_vt, stride_vh, stride_vv,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_bb, stride_bt, stride_bh,
+    stride_ob, stride_ot, stride_oh, stride_ov,
+    H: tl.constexpr,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    K_C: tl.constexpr,
+    V_C: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    b = pid // H
+    h = pid % H
+
+    i_idx = tl.arange(0, BT)
+    j_idx = tl.arange(0, BT)
+    k_idx = tl.arange(0, K_C)
+    v_idx = tl.arange(0, V_C)
+    mask_lt = i_idx[:, None] > j_idx[None, :]   # strict lower triangular
+    eye_mask = i_idx[:, None] == j_idx[None, :]
+
+    # State S: [K_C, V_C] zero-initialized
+    S = tl.zeros((K_C, V_C), dtype=tl.float32)
+
+    for chunk_i in range(NT):
+        t_start = chunk_i * BT
+
+        # Load q: [BT, K_C]
+        q_off = (
+            Q_ptr
+            + b * stride_qb
+            + h * stride_qh
+            + (t_start + i_idx)[:, None] * stride_qt
+            + k_idx[None, :] * stride_qk
+        )
+        q = tl.load(q_off).to(tl.float32) * scale
+
+        # Load k: [BT, K_C]
+        k_off = (
+            K_ptr
+            + b * stride_kb
+            + h * stride_kh
+            + (t_start + i_idx)[:, None] * stride_kt
+            + k_idx[None, :] * stride_kk
+        )
+        k = tl.load(k_off).to(tl.float32)
+
+        # Load v: [BT, V_C]
+        v_off = (
+            V_ptr
+            + b * stride_vb
+            + h * stride_vh
+            + (t_start + i_idx)[:, None] * stride_vt
+            + v_idx[None, :] * stride_vv
+        )
+        v = tl.load(v_off).to(tl.float32)
+
+        # Load g: [BT, K_C]
+        g_off = (
+            G_ptr
+            + b * stride_gb
+            + h * stride_gh
+            + (t_start + i_idx)[:, None] * stride_gt
+            + k_idx[None, :] * stride_gk
+        )
+        g = tl.load(g_off)
+
+        # Load beta: [BT]
+        beta_off = (
+            BETA_ptr
+            + b * stride_bb
+            + h * stride_bh
+            + (t_start + i_idx) * stride_bt
+        )
+        beta = tl.load(beta_off).to(tl.float32)
+
+        # in-chunk cumsum of g
+        g = tl.cumsum(g, axis=0)
+
+        g_exp = tl.exp(g)             # [BT, K_C]
+        g_neg_exp = tl.exp(-g)        # [BT, K_C]
+        g_last = g[BT - 1, :]         # [K_C]
+        g_last_exp = tl.exp(g_last)   # [K_C]
+
+        K_ng = k * g_neg_exp          # [BT, K_C]
+        T = k * g_exp                 # [BT, K_C]
+
+        # A = K_ng @ T.T   [BT, BT]
+        A = tl.dot(K_ng, tl.trans(T))
+        A = tl.where(mask_lt, -A, 0.0)
+
+        # Neumann doubling: S_inv = (I - A)^{-1} = I + A + A^2 + ... + A^63
+        S_inv = tl.where(eye_mask, 1.0, 0.0) + A
+        M2 = tl.dot(A, A)
+        S_inv = S_inv + tl.dot(M2, S_inv)
+        M4 = tl.dot(M2, M2)
+        S_inv = S_inv + tl.dot(M4, S_inv)
+        M8 = tl.dot(M4, M4)
+        S_inv = S_inv + tl.dot(M8, S_inv)
+        M16 = tl.dot(M8, M8)
+        S_inv = S_inv + tl.dot(M16, S_inv)
+        M32 = tl.dot(M16, M16)
+        S_inv = S_inv + tl.dot(M32, S_inv)
+
+        # Apply beta on columns
+        S_inv = S_inv * beta[None, :]
+
+        # w = S_inv @ T
+        w = tl.dot(S_inv, tl.trans(T))
+        # u = S_inv @ v
+        u = tl.dot(S_inv, tl.trans(v))
+
+        # Inter-chunk QK with decay (strict lower triangular)
+        # Aqk[i, j] = sum_d Q[i, d] * exp(g[j, d] - g[i, d]) * k[j, d]
+        # = (Q * exp(-g)) @ (k * exp(g)).T
+        Qg = q * g_neg_exp
+        Aqk = tl.dot(Qg, tl.trans(T))
+        Aqk = tl.where(mask_lt, Aqk, 0.0)
+
+        # o = (q * exp(g)) @ S + Aqk @ (u - w @ S)
+        qg = q * g_exp
+        wS = tl.dot(w, S)                       # [BT, V_C]
+        v_i = u - wS
+        part1 = tl.dot(qg, S)                   # [BT, V_C]
+        part2 = tl.dot(Aqk, v_i)                # [BT, V_C]
+        o = part1 + part2
+
+        # S update
+        # S_new = S * g_last_exp + (k * exp(g_last - g))^T @ v_i
+        K_g = k * tl.exp(g_last - g)            # [BT, K_C]
+        delta = tl.dot(tl.trans(K_g), v_i)      # [K_C, V_C]
+        S = S * g_last_exp[:, None] + delta
+
+        # Store o
+        o_off = (
+            O_ptr
+            + b * stride_ob
+            + h * stride_oh
+            + (t_start + i_idx)[:, None] * stride_ot
+            + v_idx[None, :] * stride_ov
+        )
+        tl.store(o_off, o.to(tl.bfloat16))
+
+
+def kda_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    chunk_size: int = 64,
+) -> torch.Tensor:
+    B, T, H, K_dim = q.shape
+    V_dim = v.shape[-1]
+    NT = T // chunk_size
+    o = torch.empty(B, T, H, V_dim, dtype=v.dtype, device=q.device)
+
+    grid = (B * H,)
+    _kda_fwd_kernel[grid](
+        q, k, v, g, beta, o,
+        scale,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+        g.stride(0), g.stride(1), g.stride(2), g.stride(3),
+        beta.stride(0), beta.stride(1), beta.stride(2),
+        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+        H=H, NT=NT, BT=chunk_size, K_C=K_dim, V_C=V_dim,
+        num_warps=4, num_stages=2,
+    )
+    return o
+
+
+class Model(nn.Module):
+    """KDA forward (chunk form)."""
+
+    def __init__(self, B: int, T: int, H: int, K: int, V: int, chunk_size: int = 64):
+        super().__in

... (994 more chars)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -95,7 +95,8 @@
 
         g_exp = tl.exp(g)             # [BT, K_C]
         g_neg_exp = tl.exp(-g)        # [BT, K_C]
-        g_last = g[BT - 1, :]         # [K_C]
+        # g_last: last row of g, shape [K_C]
+        g_last = tl.sum(tl.where(i_idx[:, None] == BT - 1, g, 0.0), axis=0)
         g_last_exp = tl.exp(g_last)   # [K_C]
 
         K_ng = k * g_neg_exp          # [BT, K_C]

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/test_triton2.py
@@ -0,0 +1,19 @@
+import triton
+import triton.language as tl
+import torch
+
+@triton.jit
+def test(a_ptr, b_ptr, c_ptr, BT: tl.constexpr, K: tl.constexpr):
+    i = tl.arange(0, BT)
+    A = tl.zeros((BT, BT), dtype=tl.float32) + 1.0
+    T = tl.zeros((BT, K), dtype=tl.float32) + 2.0
+    beta = tl.zeros((BT,), dtype=tl.float32) + 0.5
+    tl.static_print('A shape:', A.shape)
+    tl.static_print('T shape:', T.shape)
+    S_inv = A * beta[:, None]
+    tl.static_print('S_inv shape:', S_inv.shape)
+    w = tl.dot(S_inv, tl.trans(T))
+    tl.static_print('w shape:', w.shape)
+
+a = torch.zeros(1, device='cuda')
+test[(1,)](a, a, a, 64, 128)

--- a//workspace/problems/02_kda_cutlass/test_triton2.py
+++ b//workspace/problems/02_kda_cutlass/test_triton2.py
@@ -14,6 +14,7 @@
     tl.static_print('S_inv shape:', S_inv.shape)
     w = tl.dot(S_inv, tl.trans(T))
     tl.static_print('w shape:', w.shape)
+    tl.static_print('T trans shape:', tl.trans(T).shape)
 
 a = torch.zeros(1, device='cuda')
 test[(1,)](a, a, a, 64, 128)

--- a//workspace/problems/02_kda_cutlass/test_triton2.py
+++ b//workspace/problems/02_kda_cutlass/test_triton2.py
@@ -14,7 +14,10 @@
     tl.static_print('S_inv shape:', S_inv.shape)
     w = tl.dot(S_inv, tl.trans(T))
     tl.static_print('w shape:', w.shape)
-    tl.static_print('T trans shape:', tl.trans(T).shape)
+    Tt = tl.trans(T)
+    tl.static_print('T trans shape:', Tt.shape)
+    w2 = tl.dot(S_inv, Tt)
+    tl.static_print('w2 shape:', w2.shape)
 
 a = torch.zeros(1, device='cuda')
 test[(1,)](a, a, a, 64, 128)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -123,9 +123,9 @@
         S_inv = S_inv * beta[None, :]
 
         # w = S_inv @ T
-        w = tl.dot(S_inv, tl.trans(T))
+        w = tl.dot(S_inv, T)
         # u = S_inv @ v
-        u = tl.dot(S_inv, tl.trans(v))
+        u = tl.dot(S_inv, v)
 
         # Inter-chunk QK with decay (strict lower triangular)
         # Aqk[i, j] = sum_d Q[i, d] * exp(g[j, d] - g[i, d]) * k[j, d]

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -184,7 +184,7 @@
         beta.stride(0), beta.stride(1), beta.stride(2),
         o.stride(0), o.stride(1), o.stride(2), o.stride(3),
         H=H, NT=NT, BT=chunk_size, K_C=K_dim, V_C=V_dim,
-        num_warps=4, num_stages=2,
+        num_warps=4, num_stages=1,
     )
     return o
 

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -184,7 +184,7 @@
         beta.stride(0), beta.stride(1), beta.stride(2),
         o.stride(0), o.stride(1), o.stride(2), o.stride(3),
         H=H, NT=NT, BT=chunk_size, K_C=K_dim, V_C=V_dim,
-        num_warps=4, num_stages=1,
+        num_warps=2, num_stages=1,
     )
     return o
 

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -1,4 +1,10 @@
-"""Kimi Delta Attention (chunk form) forward, Triton implementation."""
+"""Kimi Delta Attention (chunk form) forward, Triton implementation.
+
+Multi-kernel design:
+  kernel 1: per-chunk w, u compute (parallel over B*H*NT)
+  kernel 2: per-chunk Aqk compute (parallel over B*H*NT)
+  kernel 3: inter-chunk output pass (one program per B*H, sequential over NT)
+"""
 from __future__ import annotations
 
 import torch
@@ -8,15 +14,232 @@
 import triton.language as tl
 
 
+# -------------------------------------------------------------------------
+# Kernel 1: compute w and u per chunk.
+#   A = (k * exp(-g)) @ (k * exp(g)).T
+#   A is masked to strictly lower triangular, then negated.
+#   Apply Neumann doubling: S_inv = (I - A)^{-1} = I + A + A^2 + ... + A^63
+#   w = (S_inv * beta) @ (k * exp(g))
+#   u = (S_inv * beta) @ v
+# -------------------------------------------------------------------------
 @triton.jit
-def _kda_fwd_kernel(
-    Q_ptr, K_ptr, V_ptr, G_ptr, BETA_ptr, O_ptr,
+def _kda_wu_kernel(
+    K_ptr, V_ptr, G_ptr, BETA_ptr,
+    W_ptr, U_ptr,
+    stride_kb, stride_kt, stride_kh, stride_kk,
+    stride_vb, stride_vt, stride_vh, stride_vv,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_bb, stride_bt, stride_bh,
+    stride_wb, stride_wt, stride_wh, stride_wk,
+    stride_ub, stride_ut, stride_uh, stride_uv,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    K_C: tl.constexpr,
+    V_C: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    NT_H = tl.num_programs(0) // H if False else 0
+    pid_bh = tl.program_id(0)
+    # Grid: (B*H, NT) -- actually we use (B*H*NT,) below
+    chunk_idx = tl.program_id(1) if tl.num_programs(0) == 0 else 0  # placeholder
+    b = pid_bh // 0  # placeholder
+    pass
+
+
+# We'll just use a simpler grid: program_id(0) encodes (B, H, NT_chunk).
+# First do kernel 1 properly.
+
+@triton.jit
+def _wu_kernel(
+    K_ptr, V_ptr, G_ptr, BETA_ptr,
+    W_ptr, U_ptr,
+    stride_kb, stride_kt, stride_kh, stride_kk,
+    stride_vb, stride_vt, stride_vh, stride_vv,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_bb, stride_bt, stride_bh,
+    stride_wb, stride_wt, stride_wh, stride_wk,
+    stride_ub, stride_ut, stride_uh, stride_uv,
+    H: tl.constexpr,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    K_C: tl.constexpr,
+    V_C: tl.constexpr,
+):
+    pid_bh = tl.program_id(0)
+    chunk_idx = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    i_idx = tl.arange(0, BT)
+    j_idx = tl.arange(0, BT)
+    k_idx = tl.arange(0, K_C)
+    v_idx = tl.arange(0, V_C)
+    mask_lt = i_idx[:, None] > j_idx[None, :]
+    eye_mask = i_idx[:, None] == j_idx[None, :]
+
+    t_start = chunk_idx * BT
+
+    # Load k, v, g, beta
+    k_off = (
+        K_ptr + b * stride_kb + h * stride_kh
+        + (t_start + i_idx)[:, None] * stride_kt
+        + k_idx[None, :] * stride_kk
+    )
+    k = tl.load(k_off).to(tl.float32)
+
+    v_off = (
+        V_ptr + b * stride_vb + h * stride_vh
+        + (t_start + i_idx)[:, None] * stride_vt
+        + v_idx[None, :] * stride_vv
+    )
+    v = tl.load(v_off).to(tl.float32)
+
+    g_off = (
+        G_ptr + b * stride_gb + h * stride_gh
+        + (t_start + i_idx)[:, None] * stride_gt
+        + k_idx[None, :] * stride_gk
+    )
+    g = tl.load(g_off)
+    g = tl.cumsum(g, axis=0)
+
+    beta_off = (
+        BETA_ptr + b * stride_bb + h * stride_bh
+        + (t_start + i_idx) * stride_bt
+    )
+    beta = tl.load(beta_off).to(tl.float32)
+
+    g_exp = tl.exp(g)
+    g_neg_exp = tl.exp(-g)
+
+    K_ng = k * g_neg_exp
+    T = k * g_exp
+
+    # A = K_ng @ T.T   [BT, BT]
+    A = tl.dot(K_ng, tl.trans(T))
+    A = tl.where(mask_lt, -A, 0.0)
+
+    # Neumann doubling
+    S_inv = tl.where(eye_mask, 1.0, 0.0) + A
+    M2 = tl.dot(A, A)
+    S_inv = S_inv + tl.dot(M2, S_inv)
+    M4 = tl.dot(M2, M2)
+    S_inv = S_inv + tl.dot(M4, S_inv)
+    M8 = tl.dot(M4, M4)
+    S_inv = S_inv + tl.dot(M8, S_inv)
+    M16 = tl.dot(M8, M8)
+    S_inv = S_inv + tl.dot(M16, S_inv)
+    M32 = tl.dot(M16, M16)
+    S_inv = S_inv + tl.dot(M32, S_inv)
+
+    S_inv = S_inv * beta[None, :]
+
+    w = tl.dot(S_inv, T)
+    u = tl.dot(S_inv, v)
+
+    w_off = (
+        W_ptr + b * stride_wb + h * stride_wh
+        + (t_start + i_idx)[:, None] * stride_wt
+        + k_idx[None, :] * stride_wk
+    )
+    tl.store(w_off, w)
+
+    u_off = (
+        U_ptr + b * stride_ub + h * stride_uh
+        + (t_start + i_idx)[:, None] * stride_ut
+        + v_idx[None, :] * stride_uv
+    )
+    tl.store(u_off, u)
+
+
+# -------------------------------------------------------------------------
+# Kernel 2: compute Aqk per chunk.
+#   Aqk[i, j] = sum_d Q[i, d] * exp(g[j, d] - g[i, d]) * k[j, d]
+#             = (Q * exp(-g)) @ (k * exp(g)).T  (strict lower triangular)
+# -------------------------------------------------------------------------
+@triton.jit
+def _aqk_kernel(
+    Q_ptr, K_ptr, G_ptr,
+    AQK_ptr,
+    scale,
+    stride_qb, stride_qt, stride_qh, stride_qk,
+    stride_kb, stride_kt, stride_kh, stride_kk,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_ab, stride_at, stride_ah, stride_aq,
+    H: tl.constexpr,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    K_C: tl.constexpr,
+):
+    pid_bh = tl.program_id(0)
+    chunk_idx = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    i_idx = tl.arange(0, BT)
+    j_idx = tl.arange(0, BT)
+    k_idx = tl.arange(0, K_C)
+    mask_lt = i_idx[:, None] > j_idx[None, :]
+
+    t_start = chunk_idx * BT
+
+    q_off = (
+        Q_ptr + b * stride_qb + h * stride_qh
+        + (t_start + i_idx)[:, None] * stride_qt
+        + k_idx[None, :] * stride_qk
+    )
+    q = tl.load(q_off).to(tl.float32) * scale
+
+    k_off = (
+        K_ptr + b * stride_kb + h * stride_kh
+        + (t_start + i_idx)[:, None] * stride_kt
+ 

... (9847 more chars)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -221,11 +221,12 @@
 
 
 # -------------------------------------------------------------------------
-# Kernel 3: inter-chunk output pass.
+# Kernel 3: inter-chunk output pass with V-tiling.
 #   For each chunk:
-#     v_i = u - w @ S
-#     o = (q * exp(g)) @ S + Aqk @ v_i
-#     S = S * exp(g_last) + (k * exp(g_last - g)).T @ v_i
+#     For each V block:
+#       v_i = u[:, v_block] - w @ S[:, v_block]
+#       o[:, v_block] = (q * exp(g)) @ S[:, v_block] + Aqk @ v_i
+#       S[:, v_block] = S[:, v_block] * exp(g_last) + (k * exp(g_last - g)).T @ v_i
 # -------------------------------------------------------------------------
 @triton.jit
 def _o_kernel(
@@ -246,6 +247,7 @@
     BT: tl.constexpr,
     K_C: tl.constexpr,
     V_C: tl.constexpr,
+    BV: tl.constexpr,  # V-tile size
 ):
     pid_bh = tl.program_id(0)
     b = pid_bh // H
@@ -254,13 +256,15 @@
     i_idx = tl.arange(0, BT)
     j_idx = tl.arange(0, BT)
     k_idx = tl.arange(0, K_C)
-    v_idx = tl.arange(0, V_C)
+    v_idx_full = tl.arange(0, V_C)
+    v_idx_local = tl.arange(0, BV)
 
     S = tl.zeros((K_C, V_C), dtype=tl.float32)
 
     for chunk_i in range(NT):
         t_start = chunk_i * BT
 
+        # Load q, k, g, w (only depend on K, not V)
         q_off = (
             Q_ptr + b * stride_qb + h * stride_qh
             + (t_start + i_idx)[:, None] * stride_qt
@@ -290,13 +294,6 @@
         )
         w = tl.load(w_off)
 
-        u_off = (
-            U_ptr + b * stride_ub + h * stride_uh
-            + (t_start + i_idx)[:, None] * stride_ut
-            + v_idx[None, :] * stride_uv
-        )
-        u = tl.load(u_off)
-
         aqk_off = (
             AQK_ptr + b * stride_ab + h * stride_ah
             + (t_start + i_idx)[:, None] * stride_at
@@ -307,29 +304,47 @@
         g_exp = tl.exp(g)
         g_last = tl.sum(tl.where(i_idx[:, None] == BT - 1, g, 0.0), axis=0)
         g_last_exp = tl.exp(g_last)
-
-        # v_i = u - w @ S
-        wS = tl.dot(w, S)
-        v_i = u - wS
-
-        # o = (q * g_exp) @ S + Aqk @ v_i
         qg = q * g_exp
-        part1 = tl.dot(qg, S)
-        part2 = tl.dot(Aqk, v_i)
-        o = part1 + part2
-
-        # S update
-        K_g = k * tl.exp(g_last - g)
-        delta = tl.dot(tl.trans(K_g), v_i)
-        S = S * g_last_exp[:, None] + delta
+        K_g = k * tl.exp(g_last - g)  # [BT, K_C]
+
+        for v_block in tl.static_range(V_C // BV):
+            v_block_start = v_block * BV
+            v_idx = v_block_start + v_idx_local
+
+            # u partial
+            u_off = (
+                U_ptr + b * stride_ub + h * stride_uh
+                + (t_start + i_idx)[:, None] * stride_ut
+                + v_idx[None, :] * stride_uv
+            )
+            u_blk = tl.load(u_off)
+
+            S_blk = tl.reshape(
+                tl.load(S_ptr + ...) if False else S, (K_C, V_C)
+            )  # placeholder
+            # Just use the S tile: S[:, v_block_start:v_block_start+BV]
+            # In Triton, slicing on a tensor is tricky. We'll do it differently.
+
+            # v_i = u_blk - w @ S_blk  -- we'll compute S_blk via load or use mask
+            # Since S is in registers, we can index it via tl.arange
+            # But Triton doesn't support arbitrary indexing. So we maintain
+            # a separate S for this block.
+
+            # Actually, let's just use a different approach: keep S as one tensor
+            # and use mask to extract blocks. But that wastes compute.
+            #
+            # Cleaner: have S stored in registers as a [K_C, V_C] tensor but
+            # process V blocks separately by storing/loading S to/from shared mem.
+            #
+            # Simplest: store S to scratch in global memory each chunk. But that
+            # adds bandwidth.
+            #
+            # For now, let's just process all V in one go but with reduced sizes.
+
+            pass
 
         # Store o
-        o_off = (
-            O_ptr + b * stride_ob + h * stride_oh
-            + (t_start + i_idx)[:, None] * stride_ot
-            + v_idx[None, :] * stride_ov
-        )
-        tl.store(o_off, o.to(tl.bfloat16))
+        pass
 
 
 def kda_fwd(

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -1,9 +1,10 @@
 """Kimi Delta Attention (chunk form) forward, Triton implementation.
 
 Multi-kernel design:
-  kernel 1: per-chunk w, u compute (parallel over B*H*NT)
-  kernel 2: per-chunk Aqk compute (parallel over B*H*NT)
-  kernel 3: inter-chunk output pass (one program per B*H, sequential over NT)
+  kernel 1 (wu_kernel): per-chunk w, u compute (parallel over B*H*NT)
+  kernel 2 (aqk_kernel): per-chunk Aqk compute (parallel over B*H*NT)
+  kernel 3 (o_kernel): inter-chunk output pass with V-tiling
+                       (one program per B*H, sequential over NT)
 """
 from __future__ import annotations
 
@@ -15,40 +16,13 @@
 
 
 # -------------------------------------------------------------------------
-# Kernel 1: compute w and u per chunk.
+# Kernel 1: per-chunk w, u compute.
 #   A = (k * exp(-g)) @ (k * exp(g)).T
 #   A is masked to strictly lower triangular, then negated.
 #   Apply Neumann doubling: S_inv = (I - A)^{-1} = I + A + A^2 + ... + A^63
 #   w = (S_inv * beta) @ (k * exp(g))
 #   u = (S_inv * beta) @ v
 # -------------------------------------------------------------------------
-@triton.jit
-def _kda_wu_kernel(
-    K_ptr, V_ptr, G_ptr, BETA_ptr,
-    W_ptr, U_ptr,
-    stride_kb, stride_kt, stride_kh, stride_kk,
-    stride_vb, stride_vt, stride_vh, stride_vv,
-    stride_gb, stride_gt, stride_gh, stride_gk,
-    stride_bb, stride_bt, stride_bh,
-    stride_wb, stride_wt, stride_wh, stride_wk,
-    stride_ub, stride_ut, stride_uh, stride_uv,
-    NT: tl.constexpr,
-    BT: tl.constexpr,
-    K_C: tl.constexpr,
-    V_C: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    NT_H = tl.num_programs(0) // H if False else 0
-    pid_bh = tl.program_id(0)
-    # Grid: (B*H, NT) -- actually we use (B*H*NT,) below
-    chunk_idx = tl.program_id(1) if tl.num_programs(0) == 0 else 0  # placeholder
-    b = pid_bh // 0  # placeholder
-    pass
-
-
-# We'll just use a simpler grid: program_id(0) encodes (B, H, NT_chunk).
-# First do kernel 1 properly.
-
 @triton.jit
 def _wu_kernel(
     K_ptr, V_ptr, G_ptr, BETA_ptr,
@@ -79,7 +53,6 @@
 
     t_start = chunk_idx * BT
 
-    # Load k, v, g, beta
     k_off = (
         K_ptr + b * stride_kb + h * stride_kh
         + (t_start + i_idx)[:, None] * stride_kt
@@ -114,11 +87,9 @@
     K_ng = k * g_neg_exp
     T = k * g_exp
 
-    # A = K_ng @ T.T   [BT, BT]
     A = tl.dot(K_ng, tl.trans(T))
     A = tl.where(mask_lt, -A, 0.0)
 
-    # Neumann doubling
     S_inv = tl.where(eye_mask, 1.0, 0.0) + A
     M2 = tl.dot(A, A)
     S_inv = S_inv + tl.dot(M2, S_inv)
@@ -152,7 +123,7 @@
 
 
 # -------------------------------------------------------------------------
-# Kernel 2: compute Aqk per chunk.
+# Kernel 2: per-chunk Aqk compute.
 #   Aqk[i, j] = sum_d Q[i, d] * exp(g[j, d] - g[i, d]) * k[j, d]
 #             = (Q * exp(-g)) @ (k * exp(g)).T  (strict lower triangular)
 # -------------------------------------------------------------------------
@@ -222,11 +193,11 @@
 
 # -------------------------------------------------------------------------
 # Kernel 3: inter-chunk output pass with V-tiling.
-#   For each chunk:
-#     For each V block:
-#       v_i = u[:, v_block] - w @ S[:, v_block]
-#       o[:, v_block] = (q * exp(g)) @ S[:, v_block] + Aqk @ v_i
-#       S[:, v_block] = S[:, v_block] * exp(g_last) + (k * exp(g_last - g)).T @ v_i
+#   We keep S as [K_C, V_C] in registers and use static_range to unroll
+#   the V-tile loop. For each chunk, for each V block:
+#     v_i_blk = u_blk - w @ S_blk
+#     o_blk = qg @ S_blk + Aqk @ v_i_blk
+#     S_blk = S_blk * g_last_exp + K_g.T @ v_i_blk
 # -------------------------------------------------------------------------
 @triton.jit
 def _o_kernel(
@@ -247,7 +218,6 @@
     BT: tl.constexpr,
     K_C: tl.constexpr,
     V_C: tl.constexpr,
-    BV: tl.constexpr,  # V-tile size
 ):
     pid_bh = tl.program_id(0)
     b = pid_bh // H
@@ -256,15 +226,14 @@
     i_idx = tl.arange(0, BT)
     j_idx = tl.arange(0, BT)
     k_idx = tl.arange(0, K_C)
-    v_idx_full = tl.arange(0, V_C)
-    v_idx_local = tl.arange(0, BV)
+    v_idx = tl.arange(0, V_C)
 
     S = tl.zeros((K_C, V_C), dtype=tl.float32)
 
     for chunk_i in range(NT):
         t_start = chunk_i * BT
 
-        # Load q, k, g, w (only depend on K, not V)
+        # Load q, k, g (V-independent)
         q_off = (
             Q_ptr + b * stride_qb + h * stride_qh
             + (t_start + i_idx)[:, None] * stride_qt
@@ -287,6 +256,7 @@
         g = tl.load(g_off)
         g = tl.cumsum(g, axis=0)
 
+        # Load w (V-independent)
         w_off = (
             W_ptr + b * stride_wb + h * stride_wh
             + (t_start + i_idx)[:, None] * stride_wt
@@ -294,6 +264,7 @@
         )
         w = tl.load(w_off)
 
+        # Load Aqk
         aqk_off = (
             AQK_ptr + b * stride_ab + h * stride_ah
             + (t_start + i_idx)[:, None] * stride_at
@@ -307,44 +278,33 @@
         qg = q * g_exp
         K_g = k * tl.exp(g_last - g)  # [BT, K_C]
 
-        for v_block in tl.static_range(V_C // BV):
-            v_block_start = v_block * BV
-            v_idx = v_block_start + v_idx_local
-
-            # u partial
-            u_off = (
-                U_ptr + b * stride_ub + h * stride_uh
-                + (t_start + i_idx)[:, None] * stride_ut
-                + v_idx[None, :] * stride_uv
-            )
-            u_blk = tl.load(u_off)
-
-            S_blk = tl.reshape(
-                tl.load(S_ptr + ...) if False else S, (K_C, V_C)
-            )  # placeholder
-            # Just use the S tile: S[:, v_block_start:v_block_start+BV]
-            # In Triton, slicing on a tensor is tricky. We'll do it differently.
-
-            # v_i = u_blk - w @ S_blk  -- we'll compute S_blk via load or use mask
-            # Since S is in registers, we can index it via tl.arange
-            # But Triton doesn't support arbitrar

... (2534 more chars)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -1,10 +1,10 @@
 """Kimi Delta Attention (chunk form) forward, Triton implementation.
 
-Multi-kernel design:
+Multi-kernel design with V-tiling in the output pass:
   kernel 1 (wu_kernel): per-chunk w, u compute (parallel over B*H*NT)
   kernel 2 (aqk_kernel): per-chunk Aqk compute (parallel over B*H*NT)
-  kernel 3 (o_kernel): inter-chunk output pass with V-tiling
-                       (one program per B*H, sequential over NT)
+  kernel 3 (o_kernel): inter-chunk output pass; V is split into 2 tiles of
+                       BV = V_C/2, with one S register per tile.
 """
 from __future__ import annotations
 
@@ -17,11 +17,6 @@
 
 # -------------------------------------------------------------------------
 # Kernel 1: per-chunk w, u compute.
-#   A = (k * exp(-g)) @ (k * exp(g)).T
-#   A is masked to strictly lower triangular, then negated.
-#   Apply Neumann doubling: S_inv = (I - A)^{-1} = I + A + A^2 + ... + A^63
-#   w = (S_inv * beta) @ (k * exp(g))
-#   u = (S_inv * beta) @ v
 # -------------------------------------------------------------------------
 @triton.jit
 def _wu_kernel(
@@ -124,8 +119,6 @@
 
 # -------------------------------------------------------------------------
 # Kernel 2: per-chunk Aqk compute.
-#   Aqk[i, j] = sum_d Q[i, d] * exp(g[j, d] - g[i, d]) * k[j, d]
-#             = (Q * exp(-g)) @ (k * exp(g)).T  (strict lower triangular)
 # -------------------------------------------------------------------------
 @triton.jit
 def _aqk_kernel(
@@ -193,11 +186,7 @@
 
 # -------------------------------------------------------------------------
 # Kernel 3: inter-chunk output pass with V-tiling.
-#   We keep S as [K_C, V_C] in registers and use static_range to unroll
-#   the V-tile loop. For each chunk, for each V block:
-#     v_i_blk = u_blk - w @ S_blk
-#     o_blk = qg @ S_blk + Aqk @ v_i_blk
-#     S_blk = S_blk * g_last_exp + K_g.T @ v_i_blk
+#   V is split into 2 tiles (BV = V_C / 2). We keep one S register per tile.
 # -------------------------------------------------------------------------
 @triton.jit
 def _o_kernel(
@@ -218,6 +207,7 @@
     BT: tl.constexpr,
     K_C: tl.constexpr,
     V_C: tl.constexpr,
+    BV: tl.constexpr,
 ):
     pid_bh = tl.program_id(0)
     b = pid_bh // H
@@ -226,14 +216,15 @@
     i_idx = tl.arange(0, BT)
     j_idx = tl.arange(0, BT)
     k_idx = tl.arange(0, K_C)
-    v_idx = tl.arange(0, V_C)
-
-    S = tl.zeros((K_C, V_C), dtype=tl.float32)
+    bv_idx = tl.arange(0, BV)
+
+    # Two S registers (one per V tile)
+    S0 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S1 = tl.zeros((K_C, BV), dtype=tl.float32)
 
     for chunk_i in range(NT):
         t_start = chunk_i * BT
 
-        # Load q, k, g (V-independent)
         q_off = (
             Q_ptr + b * stride_qb + h * stride_qh
             + (t_start + i_idx)[:, None] * stride_qt
@@ -256,7 +247,6 @@
         g = tl.load(g_off)
         g = tl.cumsum(g, axis=0)
 
-        # Load w (V-independent)
         w_off = (
             W_ptr + b * stride_wb + h * stride_wh
             + (t_start + i_idx)[:, None] * stride_wt
@@ -264,7 +254,6 @@
         )
         w = tl.load(w_off)
 
-        # Load Aqk
         aqk_off = (
             AQK_ptr + b * stride_ab + h * stride_ah
             + (t_start + i_idx)[:, None] * stride_at
@@ -278,33 +267,49 @@
         qg = q * g_exp
         K_g = k * tl.exp(g_last - g)  # [BT, K_C]
 
-        # Single matmul block: wS = w @ S
-        wS = tl.dot(w, S)
-        # v_i = u - w @ S, u must be loaded in full
-        u_off = (
+        # v_block 0
+        v0_idx = 0 + bv_idx
+        u0_off = (
             U_ptr + b * stride_ub + h * stride_uh
             + (t_start + i_idx)[:, None] * stride_ut
-            + v_idx[None, :] * stride_uv
-        )
-        u = tl.load(u_off)
-        v_i = u - wS
-
-        # o = qg @ S + Aqk @ v_i
-        part1 = tl.dot(qg, S)
-        part2 = tl.dot(Aqk, v_i)
-        o = part1 + part2
-
-        # S update
-        delta = tl.dot(tl.trans(K_g), v_i)
-        S = S * g_last_exp[:, None] + delta
-
-        # Store o
-        o_off = (
+            + v0_idx[None, :] * stride_uv
+        )
+        u0 = tl.load(u0_off)
+        wS0 = tl.dot(w, S0)
+        v_i0 = u0 - wS0
+        part1_0 = tl.dot(qg, S0)
+        part2_0 = tl.dot(Aqk, v_i0)
+        o0 = part1_0 + part2_0
+        delta0 = tl.dot(tl.trans(K_g), v_i0)
+        S0 = S0 * g_last_exp[:, None] + delta0
+        o0_off = (
             O_ptr + b * stride_ob + h * stride_oh
             + (t_start + i_idx)[:, None] * stride_ot
-            + v_idx[None, :] * stride_ov
-        )
-        tl.store(o_off, o.to(tl.bfloat16))
+            + v0_idx[None, :] * stride_ov
+        )
+        tl.store(o0_off, o0.to(tl.bfloat16))
+
+        # v_block 1
+        v1_idx = BV + bv_idx
+        u1_off = (
+            U_ptr + b * stride_ub + h * stride_uh
+            + (t_start + i_idx)[:, None] * stride_ut
+            + v1_idx[None, :] * stride_uv
+        )
+        u1 = tl.load(u1_off)
+        wS1 = tl.dot(w, S1)
+        v_i1 = u1 - wS1
+        part1_1 = tl.dot(qg, S1)
+        part2_1 = tl.dot(Aqk, v_i1)
+        o1 = part1_1 + part2_1
+        delta1 = tl.dot(tl.trans(K_g), v_i1)
+        S1 = S1 * g_last_exp[:, None] + delta1
+        o1_off = (
+            O_ptr + b * stride_ob + h * stride_oh
+            + (t_start + i_idx)[:, None] * stride_ot
+            + v1_idx[None, :] * stride_ov
+        )
+        tl.store(o1_off, o1.to(tl.bfloat16))
 
 
 def kda_fwd(
@@ -319,6 +324,7 @@
     B, T, H, K_dim = q.shape
     V_dim = v.shape[-1]
     NT = T // chunk_size
+    BV = V_dim // 2
     device = q.device
     dtype = v.dtype
 
@@ -360,7 +366,7 @@
         u_buf.stride(0), u_buf.stride(1), u_buf.stride(2), u_buf.stride(3),
         aqk_buf.stride(0), aqk_buf.stride(1), aqk_buf.stride(2), aqk_buf.stride(3),
         o.stride(0), o.stride(

... (210 more chars)

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_wu.py
@@ -0,0 +1,73 @@
+"""Debug: verify w, u computation matches naive reference."""
+import torch
+import solution
+import reference
+from einops import rearrange
+
+device = 'cuda:0'
+torch.manual_seed(0)
+B, T, H, K, V = 2, 1024, 8, 128, 128
+chunk_size = 64
+NT = T // chunk_size
+
+q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16, device=device))
+
+# Compute w, u using the naive reference
+def naive_wu(q, k, v, g, beta, scale, chunk_size=64):
+    BT = chunk_size
+    q, k, v, g, beta = (x.to(torch.float32) for x in (q, k, v, g, beta))
+    q = q * scale
+    q = rearrange(q, "b (n c) h d -> b h n c d", c=BT)
+    k = rearrange(k, "b (n c) h d -> b h n c d", c=BT)
+    v = rearrange(v, "b (n c) h d -> b h n c d", c=BT)
+    g = rearrange(g, "b (n c) h d -> b h n c d", c=BT)
+    beta = rearrange(beta, "b (n c) h -> b h n c", c=BT)
+    g = g.cumsum(-2)
+    mask_diag_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    A = torch.zeros(*q.shape[:-1], BT, dtype=torch.float32, device=q.device)
+    for i in range(BT):
+        k_i = k[..., i, :]
+        g_i = g[..., i:i + 1, :]
+        A[..., i] = torch.einsum("... c d, ... d -> ... c", k * (g - g_i).exp(), k_i)
+    A = A * beta[..., None]
+    A = -A.masked_fill(mask_diag_upper, 0)
+    for i in range(1, BT):
+        A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2)
+    A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta[..., None, :]
+    w = A @ (g.exp() * k)
+    u = A @ v
+    return w, u  # [B, H, NT, BT, K] / [B, H, NT, BT, V]
+
+scale = K ** -0.5
+w_ref, u_ref = naive_wu(q, k, v, g, beta, scale, chunk_size)
+print('w_ref shape:', w_ref.shape)
+print('u_ref shape:', u_ref.shape)
+
+# Compute w, u using our triton kernel
+w_buf = torch.empty(B, T, H, K, dtype=torch.float32, device=device)
+u_buf = torch.empty(B, T, H, V, dtype=torch.float32, device=device)
+aqk_buf = torch.empty(B, T, H, chunk_size, dtype=torch.float32, device=device)
+solution._wu_kernel[(B * H, NT)](
+    k, v, g, beta, w_buf, u_buf,
+    k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+    v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+    g.stride(0), g.stride(1), g.stride(2), g.stride(3),
+    beta.stride(0), beta.stride(1), beta.stride(2),
+    w_buf.stride(0), w_buf.stride(1), w_buf.stride(2), w_buf.stride(3),
+    u_buf.stride(0), u_buf.stride(1), u_buf.stride(2), u_buf.stride(3),
+    H=H, NT=NT, BT=chunk_size, K_C=K, V_C=V,
+    num_warps=4, num_stages=1,
+)
+
+# w_buf is [B, T, H, K], reshape to [B, H, NT, BT, K]
+w_sol = w_buf.permute(0, 2, 1, 3).reshape(B, H, NT, chunk_size, K)
+u_sol = u_buf.permute(0, 2, 1, 3).reshape(B, H, NT, chunk_size, V)
+
+print('w diff:', (w_ref - w_sol).abs().max().item())
+print('u diff:', (u_ref - u_sol).abs().max().item())
+print('w_ref max:', w_ref.abs().max().item())
+print('u_ref max:', u_ref.abs().max().item())

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -86,21 +86,21 @@
     A = tl.where(mask_lt, -A, 0.0)
 
     S_inv = tl.where(eye_mask, 1.0, 0.0) + A
-    M2 = tl.dot(A, A)
-    S_inv = S_inv + tl.dot(M2, S_inv)
-    M4 = tl.dot(M2, M2)
-    S_inv = S_inv + tl.dot(M4, S_inv)
-    M8 = tl.dot(M4, M4)
-    S_inv = S_inv + tl.dot(M8, S_inv)
-    M16 = tl.dot(M8, M8)
-    S_inv = S_inv + tl.dot(M16, S_inv)
-    M32 = tl.dot(M16, M16)
-    S_inv = S_inv + tl.dot(M32, S_inv)
+    M2 = tl.dot(A, A, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M2, S_inv, input_precision="ieee")
+    M4 = tl.dot(M2, M2, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M4, S_inv, input_precision="ieee")
+    M8 = tl.dot(M4, M4, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M8, S_inv, input_precision="ieee")
+    M16 = tl.dot(M8, M8, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M16, S_inv, input_precision="ieee")
+    M32 = tl.dot(M16, M16, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M32, S_inv, input_precision="ieee")
 
     S_inv = S_inv * beta[None, :]
 
-    w = tl.dot(S_inv, T)
-    u = tl.dot(S_inv, v)
+    w = tl.dot(S_inv, T, input_precision="ieee")
+    u = tl.dot(S_inv, v, input_precision="ieee")
 
     w_off = (
         W_ptr + b * stride_wb + h * stride_wh

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -82,7 +82,7 @@
     K_ng = k * g_neg_exp
     T = k * g_exp
 
-    A = tl.dot(K_ng, tl.trans(T))
+    A = tl.dot(K_ng, tl.trans(T), input_precision="ieee")
     A = tl.where(mask_lt, -A, 0.0)
 
     S_inv = tl.where(eye_mask, 1.0, 0.0) + A

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -173,7 +173,7 @@
     T = k * g_exp
     Qg = q * g_neg_exp
 
-    Aqk = tl.dot(Qg, tl.trans(T))
+    Aqk = tl.dot(Qg, tl.trans(T), input_precision="ieee")
     Aqk = tl.where(mask_lt, Aqk, 0.0)
 
     aqk_off = (

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -275,12 +275,12 @@
             + v0_idx[None, :] * stride_uv
         )
         u0 = tl.load(u0_off)
-        wS0 = tl.dot(w, S0)
+        wS0 = tl.dot(w, S0, input_precision="ieee")
         v_i0 = u0 - wS0
-        part1_0 = tl.dot(qg, S0)
-        part2_0 = tl.dot(Aqk, v_i0)
+        part1_0 = tl.dot(qg, S0, input_precision="ieee")
+        part2_0 = tl.dot(Aqk, v_i0, input_precision="ieee")
         o0 = part1_0 + part2_0
-        delta0 = tl.dot(tl.trans(K_g), v_i0)
+        delta0 = tl.dot(tl.trans(K_g), v_i0, input_precision="ieee")
         S0 = S0 * g_last_exp[:, None] + delta0
         o0_off = (
             O_ptr + b * stride_ob + h * stride_oh
@@ -297,12 +297,12 @@
             + v1_idx[None, :] * stride_uv
         )
         u1 = tl.load(u1_off)
-        wS1 = tl.dot(w, S1)
+        wS1 = tl.dot(w, S1, input_precision="ieee")
         v_i1 = u1 - wS1
-        part1_1 = tl.dot(qg, S1)
-        part2_1 = tl.dot(Aqk, v_i1)
+        part1_1 = tl.dot(qg, S1, input_precision="ieee")
+        part2_1 = tl.dot(Aqk, v_i1, input_precision="ieee")
         o1 = part1_1 + part2_1
-        delta1 = tl.dot(tl.trans(K_g), v_i1)
+        delta1 = tl.dot(tl.trans(K_g), v_i1, input_precision="ieee")
         S1 = S1 * g_last_exp[:, None] + delta1
         o1_off = (
             O_ptr + b * stride_ob + h * stride_oh

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_a.py
@@ -0,0 +1,98 @@
+"""Debug: compare A computation only."""
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _a_kernel(
+    K_ptr, G_ptr, A_ptr,
+    stride_kb, stride_kt, stride_kh, stride_kk,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_ab, stride_at, stride_ah, stride_aq1, stride_aq2,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    K_C: tl.constexpr,
+):
+    pid_bh = tl.program_id(0)
+    chunk_idx = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    i_idx = tl.arange(0, BT)
+    j_idx = tl.arange(0, BT)
+    k_idx = tl.arange(0, K_C)
+    mask_lt = i_idx[:, None] > j_idx[None, :]
+
+    t_start = chunk_idx * BT
+
+    k_off = (
+        K_ptr + b * stride_kb + h * stride_kh
+        + (t_start + i_idx)[:, None] * stride_kt
+        + k_idx[None, :] * stride_kk
+    )
+    k = tl.load(k_off).to(tl.float32)
+
+    g_off = (
+        G_ptr + b * stride_gb + h * stride_gh
+        + (t_start + i_idx)[:, None] * stride_gt
+        + k_idx[None, :] * stride_gk
+    )
+    g = tl.load(g_off)
+    g = tl.cumsum(g, axis=0)
+
+    g_exp = tl.exp(g)
+    g_neg_exp = tl.exp(-g)
+    K_ng = k * g_neg_exp
+    T = k * g_exp
+
+    A = tl.dot(K_ng, tl.trans(T), input_precision="ieee")
+    A = tl.where(mask_lt, -A, 0.0)
+
+    a_off = (
+        A_ptr + b * stride_ab + h * stride_ah
+        + (t_start + i_idx)[:, None] * stride_at
+        + j_idx[None, :] * stride_aq2
+    )
+    tl.store(a_off, A)
+
+
+device = 'cuda:0'
+torch.manual_seed(0)
+B, T, H, K = 2, 1024, 8, 128
+chunk_size = 64
+NT = T // chunk_size
+
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+
+# Reference computation
+def ref_A(k, g, chunk_size=64):
+    BT = chunk_size
+    k, g = k.to(torch.float32), g.to(torch.float32)
+    k = k.reshape(B, NT, BT, H, K).permute(0, 3, 1, 2, 4).contiguous()  # [B, H, NT, BT, K]
+    g = g.reshape(B, NT, BT, H, K).permute(0, 3, 1, 2, 4).contiguous()  # [B, H, NT, BT, K]
+    g = g.cumsum(-2)
+    A = torch.zeros(B, H, NT, BT, BT, dtype=torch.float32, device=k.device)
+    for i in range(BT):
+        k_i = k[..., i, :]
+        g_i = g[..., i:i + 1, :]
+        A[..., i] = torch.einsum("... c d, ... d -> ... c", k * (g - g_i).exp(), k_i)
+    return A  # [B, H, NT, BT, BT]
+
+A_ref = ref_A(k, g, chunk_size)
+A_ref = A_ref.permute(0, 2, 1, 3, 4).reshape(B, T, H, chunk_size)  # to [B, T, H, BT]
+
+A_sol = torch.empty(B, T, H, chunk_size, dtype=torch.float32, device=device)
+_a_kernel[(B * H, NT)](
+    k, g, A_sol,
+    k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+    g.stride(0), g.stride(1), g.stride(2), g.stride(3),
+    A_sol.stride(0), A_sol.stride(1), A_sol.stride(2), A_sol.stride(3), 0,
+    H=H, BT=chunk_size, K_C=K,
+    num_warps=4, num_stages=1,
+)
+
+print('A_ref shape:', A_ref.shape, 'A_sol shape:', A_sol.shape)
+print('A diff:', (A_ref - A_sol).abs().max().item())
+print('A_ref max:', A_ref.abs().max().item())
+print('A_sol max:', A_sol.abs().max().item())

--- a//workspace/problems/02_kda_cutlass/debug_a.py
+++ b//workspace/problems/02_kda_cutlass/debug_a.py
@@ -96,3 +96,21 @@
 print('A diff:', (A_ref - A_sol).abs().max().item())
 print('A_ref max:', A_ref.abs().max().item())
 print('A_sol max:', A_sol.abs().max().item())
+
+# Check non-zero entries only
+mask_lower = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=-1)
+mask_lower = mask_lower.expand_as(A_ref)
+A_ref_lower = A_ref[mask_lower]
+A_sol_lower = A_sol[mask_lower]
+print('Lower triangular:')
+print('  A_ref[:5]:', A_ref_lower[:5])
+print('  A_sol[:5]:', A_sol_lower[:5])
+print('  A_ref mean abs:', A_ref_lower.abs().mean().item())
+print('  A_sol mean abs:', A_sol_lower.abs().mean().item())
+# A_sol should equal A_ref (since we apply beta later)
+# But wait, my A_sol applies the negation and mask but no beta
+# So A_sol[i, j] = -sum_d k[i, d] * k[j, d] * exp(g[j, d] - g[i, d])  for j < i
+# A_ref[i, j] = -beta[j] * sum_d k[i, d] * k[j, d] * exp(g[j, d] - g[i, d])  for j < i
+# So A_sol = A_ref / beta[j]  (i.e., A_sol > A_ref since beta < 1)
+ratio = A_sol_lower / A_ref_lower
+print('  ratio mean:', ratio.mean().item(), 'min:', ratio.min().item(), 'max:', ratio.max().item())

--- a//workspace/problems/02_kda_cutlass/debug_a.py
+++ b//workspace/problems/02_kda_cutlass/debug_a.py
@@ -99,9 +99,17 @@
 
 # Check non-zero entries only
 mask_lower = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=-1)
-mask_lower = mask_lower.expand_as(A_ref)
-A_ref_lower = A_ref[mask_lower]
-A_sol_lower = A_sol[mask_lower]
+# A_ref is [B, T, H, BT] where T = NT*BT, so the last dim is BT (the column index of A in chunk view)
+# But T is the time dim, so [B, T, H, BT] = [B, NT*BT, H, BT] reshaped
+# A_ref_lower should be: A_ref[:, t, h, j] for j < chunk_pos_in_chunk(t)
+# Since each row of A has only lower-triangular entries, this is hard to extract.
+# Let me just check the ratio across the full thing
+ratio = A_sol / A_ref
+print('  ratio mean:', ratio[~torch.isnan(ratio) & ~torch.isinf(ratio)].mean().item())
+# Where A_ref is 0, A_sol should also be 0
+print('  both zero:', ((A_ref == 0) & (A_sol == 0)).sum().item(), '/', A_ref.numel())
+print('  ref zero, sol not:', ((A_ref == 0) & (A_sol != 0)).sum().item())
+print('  ref not, sol zero:', ((A_ref != 0) & (A_sol == 0)).sum().item())
 print('Lower triangular:')
 print('  A_ref[:5]:', A_ref_lower[:5])
 print('  A_sol[:5]:', A_sol_lower[:5])

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -84,6 +84,8 @@
 
     A = tl.dot(K_ng, tl.trans(T), input_precision="ieee")
     A = tl.where(mask_lt, -A, 0.0)
+    # Reference multiplies A by beta on rows before Neumann, then by beta on cols.
+    A = A * beta[:, None]
 
     S_inv = tl.where(eye_mask, 1.0, 0.0) + A
     M2 = tl.dot(A, A, input_precision="ieee")

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -99,6 +99,7 @@
     M32 = tl.dot(M16, M16, input_precision="ieee")
     S_inv = S_inv + tl.dot(M32, S_inv, input_precision="ieee")
 
+    # Column-beta application
     S_inv = S_inv * beta[None, :]
 
     w = tl.dot(S_inv, T, input_precision="ieee")

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_small.py
@@ -0,0 +1,133 @@
+"""Debug with tiny dimensions."""
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _wu_small_kernel(
+    K_ptr, V_ptr, G_ptr, BETA_ptr,
+    W_ptr, U_ptr,
+    stride_kb, stride_kt, stride_kh, stride_kk,
+    stride_vb, stride_vt, stride_vh, stride_vv,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_bb, stride_bt, stride_bh,
+    stride_wb, stride_wt, stride_wh, stride_wk,
+    stride_ub, stride_ut, stride_uh, stride_uv,
+    H: tl.constexpr,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    K_C: tl.constexpr,
+    V_C: tl.constexpr,
+):
+    pid_bh = tl.program_id(0)
+    chunk_idx = tl.program_id(1)
+    b = pid_bh // H
+    h = pid_bh % H
+
+    i_idx = tl.arange(0, BT)
+    j_idx = tl.arange(0, BT)
+    k_idx = tl.arange(0, K_C)
+    v_idx = tl.arange(0, V_C)
+    mask_lt = i_idx[:, None] > j_idx[None, :]
+    eye_mask = i_idx[:, None] == j_idx[None, :]
+
+    t_start = chunk_idx * BT
+
+    k = tl.load(K_ptr + b * stride_kb + h * stride_kh + (t_start + i_idx)[:, None] * stride_kt + k_idx[None, :] * stride_kk).to(tl.float32)
+    v = tl.load(V_ptr + b * stride_vb + h * stride_vh + (t_start + i_idx)[:, None] * stride_vt + v_idx[None, :] * stride_vv).to(tl.float32)
+    g = tl.load(G_ptr + b * stride_gb + h * stride_gh + (t_start + i_idx)[:, None] * stride_gt + k_idx[None, :] * stride_gk)
+    g = tl.cumsum(g, axis=0)
+    beta = tl.load(BETA_ptr + b * stride_bb + h * stride_bh + (t_start + i_idx) * stride_bt).to(tl.float32)
+
+    g_exp = tl.exp(g)
+    g_neg_exp = tl.exp(-g)
+    K_ng = k * g_neg_exp
+    T = k * g_exp
+
+    A = tl.dot(K_ng, tl.trans(T), input_precision="ieee")
+    A = tl.where(mask_lt, -A, 0.0)
+    A = A * beta[:, None]
+
+    S_inv = tl.where(eye_mask, 1.0, 0.0) + A
+    M2 = tl.dot(A, A, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M2, S_inv, input_precision="ieee")
+    M4 = tl.dot(M2, M2, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M4, S_inv, input_precision="ieee")
+    M8 = tl.dot(M4, M4, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M8, S_inv, input_precision="ieee")
+    M16 = tl.dot(M8, M8, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M16, S_inv, input_precision="ieee")
+    M32 = tl.dot(M16, M16, input_precision="ieee")
+    S_inv = S_inv + tl.dot(M32, S_inv, input_precision="ieee")
+
+    S_inv = S_inv * beta[None, :]
+
+    w = tl.dot(S_inv, T, input_precision="ieee")
+    u = tl.dot(S_inv, v, input_precision="ieee")
+
+    tl.store(W_ptr + b * stride_wb + h * stride_wh + (t_start + i_idx)[:, None] * stride_wt + k_idx[None, :] * stride_wk, w)
+    tl.store(U_ptr + b * stride_ub + h * stride_uh + (t_start + i_idx)[:, None] * stride_ut + v_idx[None, :] * stride_uv, u)
+
+
+# Use smaller dimensions
+device = 'cuda:0'
+torch.manual_seed(0)
+B, T, H, K, V = 1, 64, 1, 8, 8
+chunk_size = 64
+
+q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16, device=device))
+
+# Reference
+from einops import rearrange
+def naive_wu(q, k, v, g, beta, scale, chunk_size=64):
+    BT = chunk_size
+    k, v, g, beta = (x.to(torch.float32) for x in (k, v, g, beta))
+    k = rearrange(k, "b (n c) h d -> b h n c d", c=BT)
+    v = rearrange(v, "b (n c) h d -> b h n c d", c=BT)
+    g = rearrange(g, "b (n c) h d -> b h n c d", c=BT)
+    beta = rearrange(beta, "b (n c) h -> b h n c", c=BT)
+    g = g.cumsum(-2)
+    mask_diag_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=k.device), diagonal=0)
+    A = torch.zeros(*k.shape[:-1], BT, dtype=torch.float32, device=k.device)
+    for i in range(BT):
+        k_i = k[..., i, :]
+        g_i = g[..., i:i + 1, :]
+        A[..., i] = torch.einsum("... c d, ... d -> ... c", k * (g - g_i).exp(), k_i)
+    A = A * beta[..., None]
+    A = -A.masked_fill(mask_diag_upper, 0)
+    for i in range(1, BT):
+        A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2)
+    A = (A + torch.eye(BT, dtype=torch.float32, device=k.device)) * beta[..., None, :]
+    w = A @ (g.exp() * k)
+    u = A @ v
+    return w, u
+
+NT = T // chunk_size
+w_ref, u_ref = naive_wu(q, k, v, g, beta, K ** -0.5, chunk_size)
+
+# Triton
+w_buf = torch.empty(B, T, H, K, dtype=torch.float32, device=device)
+u_buf = torch.empty(B, T, H, V, dtype=torch.float32, device=device)
+_wu_small_kernel[(B * H, NT)](
+    k, v, g, beta, w_buf, u_buf,
+    k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+    v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+    g.stride(0), g.stride(1), g.stride(2), g.stride(3),
+    beta.stride(0), beta.stride(1), beta.stride(2),
+    w_buf.stride(0), w_buf.stride(1), w_buf.stride(2), w_buf.stride(3),
+    u_buf.stride(0), u_buf.stride(1), u_buf.stride(2), u_buf.stride(3),
+    H=H, NT=NT, BT=chunk_size, K_C=K, V_C=V,
+    num_warps=2, num_stages=1,
+)
+
+w_sol = w_buf.permute(0, 2, 1, 3).reshape(B, H, NT, chunk_size, K)
+u_sol = u_buf.permute(0, 2, 1, 3).reshape(B, H, NT, chunk_size, V)
+
+print('w diff:', (w_ref - w_sol).abs().max().item())
+print('u diff:', (u_ref - u_sol).abs().max().item())
+print('w_ref:', w_ref.flatten()[:10])
+print('w_sol:', w_sol.flatten()[:10])

--- a//workspace/problems/02_kda_cutlass/debug_small.py
+++ b//workspace/problems/02_kda_cutlass/debug_small.py
@@ -72,7 +72,7 @@
 # Use smaller dimensions
 device = 'cuda:0'
 torch.manual_seed(0)
-B, T, H, K, V = 1, 64, 1, 8, 8
+B, T, H, K, V = 1, 64, 1, 128, 128
 chunk_size = 64
 
 q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1

--- a//workspace/problems/02_kda_cutlass/debug_small.py
+++ b//workspace/problems/02_kda_cutlass/debug_small.py
@@ -129,5 +129,5 @@
 
 print('w diff:', (w_ref - w_sol).abs().max().item())
 print('u diff:', (u_ref - u_sol).abs().max().item())
-print('w_ref:', w_ref.flatten()[:10])
-print('w_sol:', w_sol.flatten()[:10])
+print('w_ref[0,0,0,0,:5]:', w_ref[0,0,0,0,:5])
+print('w_sol[0,0,0,0,:5]:', w_sol[0,0,0,0,:5])

--- a//workspace/problems/02_kda_cutlass/debug_small.py
+++ b//workspace/problems/02_kda_cutlass/debug_small.py
@@ -131,3 +131,13 @@
 print('u diff:', (u_ref - u_sol).abs().max().item())
 print('w_ref[0,0,0,0,:5]:', w_ref[0,0,0,0,:5])
 print('w_sol[0,0,0,0,:5]:', w_sol[0,0,0,0,:5])
+print('w_ref max abs:', w_ref.abs().max().item())
+print('w_sol max abs:', w_sol.abs().max().item())
+print()
+# Look at the row with max diff
+diff_per_entry = (w_ref - w_sol).abs()
+idx = diff_per_entry.argmax()
+b_idx, h_idx, n_idx, i_idx, j_idx = torch.unravel_index(idx, w_ref.shape)
+print(f'Max diff at (b={b_idx},h={h_idx},n={n_idx},i={i_idx},j={j_idx})')
+print(f'  w_ref = {w_ref[b_idx, h_idx, n_idx, i_idx, j_idx].item()}')
+print(f'  w_sol = {w_sol[b_idx, h_idx, n_idx, i_idx, j_idx].item()}')

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -0,0 +1,65 @@
+"""Verify the math by computing (I - A)^{-1} directly."""
+import torch
+from einops import rearrange
+
+device = 'cuda:0'
+torch.manual_seed(0)
+B, T, H, K, V = 1, 64, 1, 128, 128
+chunk_size = 64
+
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16, device=device))
+
+# Reference A computation
+k, g, beta = (x.to(torch.float32) for x in (k, g, beta))
+k_re = rearrange(k, "b (n c) h d -> b h n c d", c=chunk_size)
+g_re = rearrange(g, "b (n c) h d -> b h n c d", c=chunk_size)
+beta_re = rearrange(beta, "b (n c) h -> b h n c", c=chunk_size)
+g_re = g_re.cumsum(-2)
+
+A = torch.zeros(*k_re.shape[:-1], chunk_size, dtype=torch.float32, device=device)
+for i in range(chunk_size):
+    k_i = k_re[..., i, :]
+    g_i = g_re[..., i:i + 1, :]
+    A[..., i] = torch.einsum("... c d, ... d -> ... c", k_re * (g_re - g_i).exp(), k_i)
+A = A * beta_re[..., None]
+A = -A.masked_fill(torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=0), 0)
+
+# A is now [B, H, NT, BT, BT] -- take the first chunk
+A = A[0, 0, 0]  # [BT, BT]
+print('A diagonal (should be 0):', A.diag()[:5])
+print('A norm:', A.norm().item())
+
+# Direct inverse: (I - A)^{-1}
+I = torch.eye(chunk_size, device=device, dtype=torch.float32)
+A_inv = torch.linalg.solve(I - A, I)
+print('A_inv (I - A)^{-1}:')
+print('  diag[:5]:', A_inv.diag()[:5])
+print('  [62, 7]:', A_inv[62, 7].item())
+
+# Now compute the in-place Neumann reference
+A_inplace = A.clone()
+for i in range(1, chunk_size):
+    A_inplace[i, :i] = A_inplace[i, :i] + (A_inplace[i, :, None] * A_inplace[:, :i]).sum(-2)
+# Add I and multiply by beta
+A_inplace_full = (A_inplace + I) * beta_re[0, 0, 0, :, None]
+print('A_inplace_full[62, 7]:', A_inplace_full[62, 7].item())
+print('A_inv * beta[j]:', (A_inv * beta_re[0, 0, 0, :].unsqueeze(0))[62, 7].item())
+
+# Doubling approach
+S_inv = I + A
+A2 = A @ A
+S_inv = S_inv + A2 @ S_inv
+A4 = A2 @ A2
+S_inv = S_inv + A4 @ S_inv
+A8 = A4 @ A4
+S_inv = S_inv + A8 @ S_inv
+A16 = A8 @ A8
+S_inv = S_inv + A16 @ S_inv
+A32 = A16 @ A16
+S_inv = S_inv + A32 @ S_inv
+
+S_inv_with_beta = S_inv * beta_re[0, 0, 0, :].unsqueeze(0)
+print('Doubling S_inv * beta[62, 7]:', S_inv_with_beta[62, 7].item())
+print('A_inv * beta[j]:', (A_inv * beta_re[0, 0, 0, :].unsqueeze(0))[62, 7].item())

--- a//workspace/problems/02_kda_cutlass/debug_pyt.py
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -33,10 +33,11 @@
 
 # Direct inverse: (I - A)^{-1}
 I = torch.eye(chunk_size, device=device, dtype=torch.float32)
-A_inv = torch.linalg.solve(I - A, I)
+A_inv = torch.linalg.inv(I - A)
 print('A_inv (I - A)^{-1}:')
 print('  diag[:5]:', A_inv.diag()[:5])
 print('  [62, 7]:', A_inv[62, 7].item())
+print('  max abs:', A_inv.abs().max().item())
 
 # Now compute the in-place Neumann reference
 A_inplace = A.clone()

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_doubling.py
@@ -0,0 +1,76 @@
+"""Verify the Triton doubling gives the same result as PyTorch."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _doubling_kernel(
+    A_ptr, S_ptr,
+    BT: tl.constexpr,
+):
+    """Given A (BT, BT), compute S = (I - A)^{-1} using doubling."""
+    i_idx = tl.arange(0, BT)
+    j_idx = tl.arange(0, BT)
+    eye_mask = i_idx[:, None] == j_idx[None, :]
+
+    a_off = i_idx[:, None] * BT + j_idx[None, :]
+    A = tl.load(A_ptr + a_off)
+
+    S = tl.where(eye_mask, 1.0, 0.0) + A
+    M2 = tl.dot(A, A, input_precision="ieee")
+    S = S + tl.dot(M2, S, input_precision="ieee")
+    M4 = tl.dot(M2, M2, input_precision="ieee")
+    S = S + tl.dot(M4, S, input_precision="ieee")
+    M8 = tl.dot(M4, M4, input_precision="ieee")
+    S = S + tl.dot(M8, S, input_precision="ieee")
+    M16 = tl.dot(M8, M8, input_precision="ieee")
+    S = S + tl.dot(M16, S, input_precision="ieee")
+    M32 = tl.dot(M16, M16, input_precision="ieee")
+    S = S + tl.dot(M32, S, input_precision="ieee")
+
+    s_off = i_idx[:, None] * BT + j_idx[None, :]
+    tl.store(S_ptr + s_off, S)
+
+
+device = 'cuda:0'
+torch.manual_seed(0)
+BT = 64
+K = 128
+
+# Build the same A as in the KDA
+chunk_size = 64
+B, T, H = 1, 64, 1
+
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16, device=device))
+
+from einops import rearrange
+k, g, beta = (x.to(torch.float32) for x in (k, g, beta))
+k_re = rearrange(k, "b (n c) h d -> b h n c d", c=chunk_size)
+g_re = rearrange(g, "b (n c) h d -> b h n c d", c=chunk_size)
+beta_re = rearrange(beta, "b (n c) h -> b h n c", c=chunk_size)
+g_re = g_re.cumsum(-2)
+
+A = torch.zeros(*k_re.shape[:-1], chunk_size, dtype=torch.float32, device=device)
+for i in range(chunk_size):
+    k_i = k_re[..., i, :]
+    g_i = g_re[..., i:i + 1, :]
+    A[..., i] = torch.einsum("... c d, ... d -> ... c", k_re * (g_re - g_i).exp(), k_i)
+A = A * beta_re[..., None]
+A = -A.masked_fill(torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=0), 0)
+
+# A shape: [B, H, NT, BT, BT] = [1, 1, 1, 64, 64]
+A = A[0, 0, 0].contiguous()
+S = torch.zeros(BT, BT, device=device, dtype=torch.float32)
+_doubling_kernel[(1,)](A, S, BT=BT, num_warps=2, num_stages=1)
+
+# Compare with PyTorch
+I = torch.eye(BT, device=device, dtype=torch.float32)
+A_inv_pyt = torch.linalg.inv(I - A)
+print('S[62, 7]:', S[62, 7].item())
+print('A_inv[62, 7]:', A_inv_pyt[62, 7].item())
+print('Diff:', (S - A_inv_pyt).abs().max().item())
+print('S max abs:', S.abs().max().item())
+print('A_inv max abs:', A_inv_pyt.abs().max().item())

--- a//workspace/problems/02_kda_cutlass/debug_pyt.py
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -39,6 +39,15 @@
 print('  [62, 7]:', A_inv[62, 7].item())
 print('  max abs:', A_inv.abs().max().item())
 
+# Compute in-place Neumann in a different way
+A2 = A.clone()
+for i in range(1, chunk_size):
+    A2[i, :i] = A2[i, :i] + (A2[i, :, None] * A2[:, :i]).sum(-2)
+# Now A2 = (I - A)^{-1} - I (in the strictly lower triangular)
+# Add I and multiply by beta on columns
+A2_full = (A2 + I) * beta_re[0, 0, 0, :, None]
+print('A2_full[62, 7]:', A2_full[62, 7].item())
+
 # Now compute the in-place Neumann reference
 A_inplace = A.clone()
 for i in range(1, chunk_size):

--- a//workspace/problems/02_kda_cutlass/debug_pyt.py
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -39,6 +39,15 @@
 print('  [62, 7]:', A_inv[62, 7].item())
 print('  max abs:', A_inv.abs().max().item())
 
+# In-place: compute f = (I - A)^{-1} - I
+A_inplace = A.clone()
+for i in range(1, chunk_size):
+    A_inplace[i, :i] = A_inplace[i, :i] + (A_inplace[i, :, None] * A_inplace[:, :i]).sum(-2)
+print('A_inplace (just lower triangular, after in-place):')
+print('  [62, 7]:', A_inplace[62, 7].item())
+print('  diff from A_inv[62, 7]:', abs(A_inplace[62, 7].item() - A_inv[62, 7].item()))
+print('  max diff from A_inv[62, :62]:', (A_inplace[62, :62] - A_inv[62, :62]).abs().max().item())
+
 # Compute in-place Neumann in a different way
 A2 = A.clone()
 for i in range(1, chunk_size):

--- a//workspace/problems/02_kda_cutlass/debug_pyt.py
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -65,6 +65,8 @@
 A_inplace_full = (A_inplace + I) * beta_re[0, 0, 0, :, None]
 print('A_inplace_full[62, 7]:', A_inplace_full[62, 7].item())
 print('A_inv * beta[j]:', (A_inv * beta_re[0, 0, 0, :].unsqueeze(0))[62, 7].item())
+print('beta[7]:', beta_re[0, 0, 0, 7].item())
+print('A_inplace[62, 7] * beta[7]:', A_inplace[62, 7].item() * beta_re[0, 0, 0, 7].item())
 
 # Doubling approach
 S_inv = I + A

--- a//workspace/problems/02_kda_cutlass/debug_pyt.py
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -63,10 +63,13 @@
     A_inplace[i, :i] = A_inplace[i, :i] + (A_inplace[i, :, None] * A_inplace[:, :i]).sum(-2)
 # Add I and multiply by beta
 A_inplace_full = (A_inplace + I) * beta_re[0, 0, 0, :, None]
+print('A_inplace_full shape:', A_inplace_full.shape)
 print('A_inplace_full[62, 7]:', A_inplace_full[62, 7].item())
-print('A_inv * beta[j]:', (A_inv * beta_re[0, 0, 0, :].unsqueeze(0))[62, 7].item())
-print('beta[7]:', beta_re[0, 0, 0, 7].item())
-print('A_inplace[62, 7] * beta[7]:', A_inplace[62, 7].item() * beta_re[0, 0, 0, 7].item())
+print('A_inplace_full[62, 62]:', A_inplace_full[62, 62].item())
+print('A_inplace[62, 7] =', A_inplace[62, 7].item())
+print('I[62, 7] =', I[62, 7].item())
+print('(A_inplace + I)[62, 7] =', (A_inplace + I)[62, 7].item())
+print('beta_re[0, 0, 0, 7, None] =', beta_re[0, 0, 0, 7, None].item())
 
 # Doubling approach
 S_inv = I + A

--- a//workspace/problems/02_kda_cutlass/debug_pyt.py
+++ b//workspace/problems/02_kda_cutlass/debug_pyt.py
@@ -91,6 +91,7 @@
 A32 = A16 @ A16
 S_inv = S_inv + A32 @ S_inv
 
-S_inv_with_beta = S_inv * beta_re[0, 0, 0, :].unsqueeze(0)
+# Column-beta (multiply j-th column by beta[j])
+S_inv_with_beta = S_inv * beta_re[0, 0, 0, None, :]
 print('Doubling S_inv * beta[62, 7]:', S_inv_with_beta[62, 7].item())
-print('A_inv * beta[j]:', (A_inv * beta_re[0, 0, 0, :].unsqueeze(0))[62, 7].item())
+print('A_inv * beta[j]:', (A_inv * beta_re[0, 0, 0, None, :])[62, 7].item())

--- a//workspace/problems/02_kda_cutlass/debug_wu.py
+++ b//workspace/problems/02_kda_cutlass/debug_wu.py
@@ -1,3 +1,13 @@
 print('w_ref shape:', w_ref.shape, 'u_ref shape:', u_ref.shape)
 print('w diff:', (w_ref - w_sol).abs().max().item())
-print('u diff:', (u_ref - u_sol).abs().max().item())+print('u diff:', (u_ref - u_sol).abs().max().item())
+print('w_ref max abs:', w_ref.abs().max().item())
+print('w_sol max abs:', w_sol.abs().max().item())
+
+# Look at the (i, j, d) with max diff in w
+diff_w = (w_ref - w_sol).abs()
+idx = diff_w.argmax()
+b_idx, h_idx, n_idx, i_idx, j_idx = torch.unravel_index(idx, w_ref.shape)
+print(f'Max diff in w at (b={b_idx},h={h_idx},n={n_idx},i={i_idx},j={j_idx})')
+print(f'  w_ref = {w_ref[b_idx, h_idx, n_idx, i_idx, j_idx].item()}')
+print(f'  w_sol = {w_sol[b_idx, h_idx, n_idx, i_idx, j_idx].item()}')

--- a//workspace/problems/02_kda_cutlass/debug_wu.py
+++ b//workspace/problems/02_kda_cutlass/debug_wu.py
@@ -71,3 +71,11 @@
 print('u diff:', (u_ref - u_sol).abs().max().item())
 print('w_ref max:', w_ref.abs().max().item())
 print('u_ref max:', u_ref.abs().max().item())
+
+# Look at max diff location
+diff_w = (w_ref - w_sol).abs()
+idx = diff_w.argmax()
+b_idx, h_idx, n_idx, i_idx, j_idx = torch.unravel_index(idx, w_ref.shape)
+print(f'Max diff in w at (b={b_idx},h={h_idx},n={n_idx},i={i_idx},j={j_idx})')
+print(f'  w_ref = {w_ref[b_idx, h_idx, n_idx, i_idx, j_idx].item()}')
+print(f'  w_sol = {w_sol[b_idx, h_idx, n_idx, i_idx, j_idx].item()}')

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_a_doubling.py
@@ -0,0 +1,85 @@
+"""Debug: check the A norm and condition for the wu kernel."""
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+
+device = 'cuda:0'
+torch.manual_seed(0)
+B, T, H, K, V = 2, 1024, 8, 128, 128
+chunk_size = 64
+NT = T // chunk_size
+
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16, device=device))
+
+# Compute A in PyTorch for the chunk with max diff: n=13 (b=1, h=0)
+n_test = 13
+b_test, h_test = 1, 0
+
+k_t = k.to(torch.float32)
+g_t = g.to(torch.float32)
+beta_t = beta.to(torch.float32)
+
+k_re = rearrange(k_t, "b (n c) h d -> b h n c d", c=chunk_size)
+g_re = rearrange(g_t, "b (n c) h d -> b h n c d", c=chunk_size)
+beta_re = rearrange(beta_t, "b (n c) h -> b h n c", c=chunk_size)
+g_re = g_re.cumsum(-2)
+
+A = torch.zeros(*k_re.shape[:-1], chunk_size, dtype=torch.float32, device=device)
+for i in range(chunk_size):
+    k_i = k_re[..., i, :]
+    g_i = g_re[..., i:i + 1, :]
+    A[..., i] = torch.einsum("... c d, ... d -> ... c", k_re * (g_re - g_i).exp(), k_i)
+A = A * beta_re[..., None]
+A = -A.masked_fill(torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=0), 0)
+
+# A shape: [B, H, NT, BT, BT]
+A_chunk = A[b_test, h_test, n_test]
+print('A_chunk norm:', A_chunk.norm().item())
+print('A_chunk abs max:', A_chunk.abs().max().item())
+print('A_chunk[i=63, j<62] sample:', A_chunk[63, :62].abs().max().item())
+
+# (I - A)^{-1} - the true inverse
+I = torch.eye(chunk_size, device=device, dtype=torch.float32)
+A_inv = torch.linalg.inv(I - A_chunk)
+print('A_inv[63, 51]:', A_inv[63, 51].item())
+print('A_inv max abs:', A_inv.abs().max().item())
+
+# Doubling
+S_inv = I + A_chunk
+M2 = A_chunk @ A_chunk
+S_inv = S_inv + M2 @ S_inv
+M4 = M2 @ M2
+S_inv = S_inv + M4 @ S_inv
+M8 = M4 @ M4
+S_inv = S_inv + M8 @ S_inv
+M16 = M8 @ M8
+S_inv = S_inv + M16 @ S_inv
+M32 = M16 @ M16
+S_inv = S_inv + M32 @ S_inv
+print('S_inv[63, 51]:', S_inv[63, 51].item())
+print('S_inv max abs:', S_inv.abs().max().item())
+
+# In-place
+A_inplace = A_chunk.clone()
+for i in range(1, chunk_size):
+    A_inplace[i, :i] = A_inplace[i, :i] + (A_inplace[i, :, None] * A_inplace[:, :i]).sum(-2)
+print('A_inplace[63, 51]:', A_inplace[63, 51].item())
+
+# Column-beta
+beta_v = beta_re[b_test, h_test, n_test]
+print('S_inv[63, 51] * beta[51]:', (S_inv[63, 51] * beta_v[51]).item())
+print('A_inplace[63, 51] * beta[51]:', (A_inplace[63, 51] * beta_v[51]).item())
+
+# Now w_ref[i, j] = sum_k A_ref[i, k] * T[k, j]
+# A_ref = (I - A)^{-1} * diag(beta on columns) = S_inv * beta on columns
+A_ref = S_inv * beta_v.unsqueeze(0)
+T = (k_re * g_re.exp())[b_test, h_test, n_test]  # [BT, K]
+w_pyt = A_ref @ T
+print('w_pyt[63, 51]:', w_pyt[63, 51].item())
+
+A_ref_inplace = A_inplace * beta_v.unsqueeze(0)
+w_pyt_inplace = A_ref_inplace @ T
+print('w_pyt_inplace[63, 51]:', w_pyt_inplace[63, 51].item())

--- /dev/null
+++ b//workspace/problems/02_kda_cutlass/debug_wu_v2.py
@@ -0,0 +1,82 @@
+"""Check: does the wu kernel match a PyTorch reference that uses the doubling algorithm?"""
+import torch
+from einops import rearrange
+import solution
+
+device = 'cuda:0'
+torch.manual_seed(0)
+B, T, H, K, V = 2, 1024, 8, 128, 128
+chunk_size = 64
+NT = T // chunk_size
+
+k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device=device) * 0.1
+v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device=device) * 0.1
+g = torch.randn(B, T, H, K, dtype=torch.float32, device=device) * 0.1 - 0.05
+beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16, device=device))
+
+
+def pyt_wu_doubling(k, v, g, beta, chunk_size=64):
+    """Compute w, u using the doubling algorithm."""
+    BT = chunk_size
+    k, v, g, beta = (x.to(torch.float32) for x in (k, v, g, beta))
+    k_re = rearrange(k, "b (n c) h d -> b h n c d", c=BT)
+    v_re = rearrange(v, "b (n c) h d -> b h n c d", c=BT)
+    g_re = rearrange(g, "b (n c) h d -> b h n c d", c=BT)
+    beta_re = rearrange(beta, "b (n c) h -> b h n c", c=BT)
+    g_re = g_re.cumsum(-2)
+    g_exp = g_re.exp()
+    T = k_re * g_exp  # [B, H, NT, BT, K]
+
+    # A = -K_ng @ T.T, mask strict lower
+    K_ng = k_re * (-g_re).exp()
+    A_full = torch.einsum("...id,...jd->...ij", K_ng, T)  # [B, H, NT, BT, BT]
+    A = -torch.tril(A_full, diagonal=-1)  # strict lower triangular
+    A = A * beta_re[..., None]  # row-beta
+
+    # Doubling
+    I = torch.eye(BT, device=k.device, dtype=torch.float32)
+    S_inv = I + A
+    A2 = A @ A
+    S_inv = S_inv + A2 @ S_inv
+    A4 = A2 @ A2
+    S_inv = S_inv + A4 @ S_inv
+    A8 = A4 @ A4
+    S_inv = S_inv + A8 @ S_inv
+    A16 = A8 @ A8
+    S_inv = S_inv + A16 @ S_inv
+    A32 = A16 @ A16
+    S_inv = S_inv + A32 @ S_inv
+
+    S_inv = S_inv * beta_re[..., None, :]  # column-beta
+
+    w = S_inv @ T
+    u = S_inv @ v_re
+    return w, u  # [B, H, NT, BT, K], [B, H, NT, BT, V]
+
+
+w_pyt, u_pyt = pyt_wu_doubling(k, v, g, beta, chunk_size)
+print('w_pyt shape:', w_pyt.shape)
+
+# Run my triton wu kernel
+w_buf = torch.empty(B, T, H, K, dtype=torch.float32, device=device)
+u_buf = torch.empty(B, T, H, V, dtype=torch.float32, device=device)
+solution._wu_kernel[(B * H, NT)](
+    k, v, g, beta, w_buf, u_buf,
+    k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+    v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+    g.stride(0), g.stride(1), g.stride(2), g.stride(3),
+    beta.stride(0), beta.stride(1), beta.stride(2),
+    w_buf.stride(0), w_buf.stride(1), w_buf.stride(2), w_buf.stride(3),
+    u_buf.stride(0), u_buf.stride(1), u_buf.stride(2), u_buf.stride(3),
+    H=H, NT=NT, BT=chunk_size, K_C=K, V_C=V,
+    num_warps=4, num_stages=1,
+)
+
+# Reshape Triton output
+w_sol = w_buf.permute(0, 2, 1, 3).reshape(B, H, NT, chunk_size, K)
+u_sol = u_buf.permute(0, 2, 1, 3).reshape(B, H, NT, chunk_size, V)
+
+print('w diff (Triton doubling vs PyTorch doubling):', (w_pyt - w_sol).abs().max().item())
+print('u diff:', (u_pyt - u_sol).abs().max().item())
+print('w_pyt max abs:', w_pyt.abs().max().item())
+print('w_sol max abs:', w_sol.abs().max().item())

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -221,9 +221,12 @@
     k_idx = tl.arange(0, K_C)
     bv_idx = tl.arange(0, BV)
 
-    # Two S registers (one per V tile)
+    # V_C / BV S registers (one per V tile)
+    # V_C=128, BV=32 -> 4 S registers
     S0 = tl.zeros((K_C, BV), dtype=tl.float32)
     S1 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S2 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S3 = tl.zeros((K_C, BV), dtype=tl.float32)
 
     for chunk_i in range(NT):
         t_start = chunk_i * BT

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -317,6 +317,50 @@
         )
         tl.store(o1_off, o1.to(tl.bfloat16))
 
+        # v_block 2
+        v2_idx = 2 * BV + bv_idx
+        u2_off = (
+            U_ptr + b * stride_ub + h * stride_uh
+            + (t_start + i_idx)[:, None] * stride_ut
+            + v2_idx[None, :] * stride_uv
+        )
+        u2 = tl.load(u2_off)
+        wS2 = tl.dot(w, S2, input_precision="ieee")
+        v_i2 = u2 - wS2
+        part1_2 = tl.dot(qg, S2, input_precision="ieee")
+        part2_2 = tl.dot(Aqk, v_i2, input_precision="ieee")
+        o2 = part1_2 + part2_2
+        delta2 = tl.dot(tl.trans(K_g), v_i2, input_precision="ieee")
+        S2 = S2 * g_last_exp[:, None] + delta2
+        o2_off = (
+            O_ptr + b * stride_ob + h * stride_oh
+            + (t_start + i_idx)[:, None] * stride_ot
+            + v2_idx[None, :] * stride_ov
+        )
+        tl.store(o2_off, o2.to(tl.bfloat16))
+
+        # v_block 3
+        v3_idx = 3 * BV + bv_idx
+        u3_off = (
+            U_ptr + b * stride_ub + h * stride_uh
+            + (t_start + i_idx)[:, None] * stride_ut
+            + v3_idx[None, :] * stride_uv
+        )
+        u3 = tl.load(u3_off)
+        wS3 = tl.dot(w, S3, input_precision="ieee")
+        v_i3 = u3 - wS3
+        part1_3 = tl.dot(qg, S3, input_precision="ieee")
+        part2_3 = tl.dot(Aqk, v_i3, input_precision="ieee")
+        o3 = part1_3 + part2_3
+        delta3 = tl.dot(tl.trans(K_g), v_i3, input_precision="ieee")
+        S3 = S3 * g_last_exp[:, None] + delta3
+        o3_off = (
+            O_ptr + b * stride_ob + h * stride_oh
+            + (t_start + i_idx)[:, None] * stride_ot
+            + v3_idx[None, :] * stride_ov
+        )
+        tl.store(o3_off, o3.to(tl.bfloat16))
+
 
 def kda_fwd(
     q: torch.Tensor,

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -374,7 +374,7 @@
     B, T, H, K_dim = q.shape
     V_dim = v.shape[-1]
     NT = T // chunk_size
-    BV = V_dim // 2
+    BV = V_dim // 4
     device = q.device
     dtype = v.dtype
 

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -3,8 +3,8 @@
 Multi-kernel design with V-tiling in the output pass:
   kernel 1 (wu_kernel): per-chunk w, u compute (parallel over B*H*NT)
   kernel 2 (aqk_kernel): per-chunk Aqk compute (parallel over B*H*NT)
-  kernel 3 (o_kernel): inter-chunk output pass; V is split into 2 tiles of
-                       BV = V_C/2, with one S register per tile.
+  kernel 3 (o_kernel): inter-chunk output pass; K is split into 2 tiles of
+                       BK = K_C/2, with one S register per K tile.
 """
 from __future__ import annotations
 
@@ -84,7 +84,6 @@
 
     A = tl.dot(K_ng, tl.trans(T), input_precision="ieee")
     A = tl.where(mask_lt, -A, 0.0)
-    # Reference multiplies A by beta on rows before Neumann, then by beta on cols.
     A = A * beta[:, None]
 
     S_inv = tl.where(eye_mask, 1.0, 0.0) + A
@@ -99,7 +98,6 @@
     M32 = tl.dot(M16, M16, input_precision="ieee")
     S_inv = S_inv + tl.dot(M32, S_inv, input_precision="ieee")
 
-    # Column-beta application
     S_inv = S_inv * beta[None, :]
 
     w = tl.dot(S_inv, T, input_precision="ieee")
@@ -188,8 +186,8 @@
 
 
 # -------------------------------------------------------------------------
-# Kernel 3: inter-chunk output pass with V-tiling.
-#   V is split into 2 tiles (BV = V_C / 2). We keep one S register per tile.
+# Kernel 3: inter-chunk output pass with K-tiling.
+#   K is split into 2 tiles (BK = K_C / 2). We keep one S register per K tile.
 # -------------------------------------------------------------------------
 @triton.jit
 def _o_kernel(
@@ -210,7 +208,7 @@
     BT: tl.constexpr,
     K_C: tl.constexpr,
     V_C: tl.constexpr,
-    BV: tl.constexpr,
+    BK: tl.constexpr,
 ):
     pid_bh = tl.program_id(0)
     b = pid_bh // H
@@ -219,14 +217,11 @@
     i_idx = tl.arange(0, BT)
     j_idx = tl.arange(0, BT)
     k_idx = tl.arange(0, K_C)
-    bv_idx = tl.arange(0, BV)
-
-    # V_C / BV S registers (one per V tile)
-    # V_C=128, BV=32 -> 4 S registers
-    S0 = tl.zeros((K_C, BV), dtype=tl.float32)
-    S1 = tl.zeros((K_C, BV), dtype=tl.float32)
-    S2 = tl.zeros((K_C, BV), dtype=tl.float32)
-    S3 = tl.zeros((K_C, BV), dtype=tl.float32)
+    v_idx = tl.arange(0, V_C)
+    bk_idx = tl.arange(0, BK)
+
+    S0 = tl.zeros((BK, V_C), dtype=tl.float32)
+    S1 = tl.zeros((BK, V_C), dtype=tl.float32)
 
     for chunk_i in range(NT):
         t_start = chunk_i * BT
@@ -253,12 +248,27 @@
         g = tl.load(g_off)
         g = tl.cumsum(g, axis=0)
 
+        # Load full v, w
         w_off = (
             W_ptr + b * stride_wb + h * stride_wh
             + (t_start + i_idx)[:, None] * stride_wt
             + k_idx[None, :] * stride_wk
         )
         w = tl.load(w_off)
+
+        v_off = (
+            V_ptr + b * stride_vb + h * stride_vh
+            + (t_start + i_idx)[:, None] * stride_vt
+            + v_idx[None, :] * stride_vv
+        )
+        v = tl.load(v_off)
+
+        u_off = (
+            U_ptr + b * stride_ub + h * stride_uh
+            + (t_start + i_idx)[:, None] * stride_ut
+            + v_idx[None, :] * stride_uv
+        )
+        u = tl.load(u_off)
 
         aqk_off = (
             AQK_ptr + b * stride_ab + h * stride_ah
@@ -273,93 +283,52 @@
         qg = q * g_exp
         K_g = k * tl.exp(g_last - g)  # [BT, K_C]
 
-        # v_block 0
-        v0_idx = 0 + bv_idx
-        u0_off = (
-            U_ptr + b * stride_ub + h * stride_uh
-            + (t_start + i_idx)[:, None] * stride_ut
-            + v0_idx[None, :] * stride_uv
-        )
-        u0 = tl.load(u0_off)
-        wS0 = tl.dot(w, S0, input_precision="ieee")
-        v_i0 = u0 - wS0
-        part1_0 = tl.dot(qg, S0, input_precision="ieee")
-        part2_0 = tl.dot(Aqk, v_i0, input_precision="ieee")
-        o0 = part1_0 + part2_0
-        delta0 = tl.dot(tl.trans(K_g), v_i0, input_precision="ieee")
-        S0 = S0 * g_last_exp[:, None] + delta0
-        o0_off = (
+        # K-block 0
+        k0_idx = 0 + bk_idx
+        w0 = w[:, k0_idx]  # actually can't slice
+        # Hmm, we need to slice. Let me use a different approach.
+
+        # Process both k blocks
+        # k_block 0: rows of S0 are k[0:BK]
+        S = S0
+        w_k = w  # full w for now
+        K_g_k = K_g[:, k0_idx]  # [BT, BK]
+        # Need to slice w[:, 0:BK]. Use mask.
+
+        # Actually, this is getting complex. Let me use a simpler approach: full V but small K
+        # Just use S as a single tensor for now
+
+        # K-gate times full S:
+        # w @ S = w @ [S0; S1] = [w @ S0, w @ S1]
+        # We can compute them separately
+
+        # v_i = u - w @ S
+        wS = tl.dot(w, tl.cat(S0, S1, can_reorder=False), input_precision="ieee")
+        v_i = u - wS
+
+        # o = qg @ S + Aqk @ v_i
+        part1 = tl.dot(qg, tl.cat(S0, S1, can_reorder=False), input_precision="ieee")
+        part2 = tl.dot(Aqk, v_i, input_precision="ieee")
+        o = part1 + part2
+
+        # S update
+        delta = tl.dot(tl.trans(K_g), v_i, input_precision="ieee")
+        # Split delta into S0 and S1
+        # delta is [K_C, V_C] = [128, 128]
+        # We need to split into [BK, V_C] = [64, 128] and [BK, V_C] = [64, 128]
+        # Use mask
+
+        # Hmm, slicing in Triton is hard
+        # Let me just keep S as a single tensor
+        S0 = S0 * g_last_exp[k0_idx, None] + delta  # placeholder
+
+        # Store o
+        o_off = (
             O_ptr + b * stride_ob + h * stride_oh
             + (t_start + i_idx)[:, None] * stride_ot
-            + v0_idx[None, :] * stride_ov
-        )
-        tl.store(o0_off, o0.to(tl.bfloat16))
-
-        # v_block 1
-        v1_idx = BV + bv_idx
-        u1_off = (
-            U_ptr + b * stride_ub + h * stride_uh
-            + (t_start + i_idx)[:, None] * stride_ut
-            + v1_idx[None, :] * stride_uv
-        )
-        u1 = tl.load(u1_off)
-        wS1 = tl.dot(w, S1, input_pr

... (2945 more chars)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -3,8 +3,8 @@
 Multi-kernel design with V-tiling in the output pass:
   kernel 1 (wu_kernel): per-chunk w, u compute (parallel over B*H*NT)
   kernel 2 (aqk_kernel): per-chunk Aqk compute (parallel over B*H*NT)
-  kernel 3 (o_kernel): inter-chunk output pass; K is split into 2 tiles of
-                       BK = K_C/2, with one S register per K tile.
+  kernel 3 (o_kernel): inter-chunk output pass; V is split into 2 tiles of
+                       BV = V_C/2, with one S register per tile.
 """
 from __future__ import annotations
 
@@ -186,8 +186,8 @@
 
 
 # -------------------------------------------------------------------------
-# Kernel 3: inter-chunk output pass with K-tiling.
-#   K is split into 2 tiles (BK = K_C / 2). We keep one S register per K tile.
+# Kernel 3: inter-chunk output pass with V-tiling.
+#   V is split into 2 tiles (BV = V_C / 2). We keep one S register per tile.
 # -------------------------------------------------------------------------
 @triton.jit
 def _o_kernel(
@@ -208,7 +208,7 @@
     BT: tl.constexpr,
     K_C: tl.constexpr,
     V_C: tl.constexpr,
-    BK: tl.constexpr,
+    BV: tl.constexpr,
 ):
     pid_bh = tl.program_id(0)
     b = pid_bh // H
@@ -217,11 +217,10 @@
     i_idx = tl.arange(0, BT)
     j_idx = tl.arange(0, BT)
     k_idx = tl.arange(0, K_C)
-    v_idx = tl.arange(0, V_C)
-    bk_idx = tl.arange(0, BK)
-
-    S0 = tl.zeros((BK, V_C), dtype=tl.float32)
-    S1 = tl.zeros((BK, V_C), dtype=tl.float32)
+    bv_idx = tl.arange(0, BV)
+
+    S0 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S1 = tl.zeros((K_C, BV), dtype=tl.float32)
 
     for chunk_i in range(NT):
         t_start = chunk_i * BT
@@ -248,27 +247,12 @@
         g = tl.load(g_off)
         g = tl.cumsum(g, axis=0)
 
-        # Load full v, w
         w_off = (
             W_ptr + b * stride_wb + h * stride_wh
             + (t_start + i_idx)[:, None] * stride_wt
             + k_idx[None, :] * stride_wk
         )
         w = tl.load(w_off)
-
-        v_off = (
-            V_ptr + b * stride_vb + h * stride_vh
-            + (t_start + i_idx)[:, None] * stride_vt
-            + v_idx[None, :] * stride_vv
-        )
-        v = tl.load(v_off)
-
-        u_off = (
-            U_ptr + b * stride_ub + h * stride_uh
-            + (t_start + i_idx)[:, None] * stride_ut
-            + v_idx[None, :] * stride_uv
-        )
-        u = tl.load(u_off)
 
         aqk_off = (
             AQK_ptr + b * stride_ab + h * stride_ah
@@ -283,52 +267,49 @@
         qg = q * g_exp
         K_g = k * tl.exp(g_last - g)  # [BT, K_C]
 
-        # K-block 0
-        k0_idx = 0 + bk_idx
-        w0 = w[:, k0_idx]  # actually can't slice
-        # Hmm, we need to slice. Let me use a different approach.
-
-        # Process both k blocks
-        # k_block 0: rows of S0 are k[0:BK]
-        S = S0
-        w_k = w  # full w for now
-        K_g_k = K_g[:, k0_idx]  # [BT, BK]
-        # Need to slice w[:, 0:BK]. Use mask.
-
-        # Actually, this is getting complex. Let me use a simpler approach: full V but small K
-        # Just use S as a single tensor for now
-
-        # K-gate times full S:
-        # w @ S = w @ [S0; S1] = [w @ S0, w @ S1]
-        # We can compute them separately
-
-        # v_i = u - w @ S
-        wS = tl.dot(w, tl.cat(S0, S1, can_reorder=False), input_precision="ieee")
-        v_i = u - wS
-
-        # o = qg @ S + Aqk @ v_i
-        part1 = tl.dot(qg, tl.cat(S0, S1, can_reorder=False), input_precision="ieee")
-        part2 = tl.dot(Aqk, v_i, input_precision="ieee")
-        o = part1 + part2
-
-        # S update
-        delta = tl.dot(tl.trans(K_g), v_i, input_precision="ieee")
-        # Split delta into S0 and S1
-        # delta is [K_C, V_C] = [128, 128]
-        # We need to split into [BK, V_C] = [64, 128] and [BK, V_C] = [64, 128]
-        # Use mask
-
-        # Hmm, slicing in Triton is hard
-        # Let me just keep S as a single tensor
-        S0 = S0 * g_last_exp[k0_idx, None] + delta  # placeholder
-
-        # Store o
-        o_off = (
+        # v_block 0
+        v0_idx = 0 + bv_idx
+        u0_off = (
+            U_ptr + b * stride_ub + h * stride_uh
+            + (t_start + i_idx)[:, None] * stride_ut
+            + v0_idx[None, :] * stride_uv
+        )
+        u0 = tl.load(u0_off)
+        wS0 = tl.dot(w, S0, input_precision="ieee")
+        v_i0 = u0 - wS0
+        part1_0 = tl.dot(qg, S0, input_precision="ieee")
+        part2_0 = tl.dot(Aqk, v_i0, input_precision="ieee")
+        o0 = part1_0 + part2_0
+        delta0 = tl.dot(tl.trans(K_g), v_i0, input_precision="ieee")
+        S0 = S0 * g_last_exp[:, None] + delta0
+        o0_off = (
             O_ptr + b * stride_ob + h * stride_oh
             + (t_start + i_idx)[:, None] * stride_ot
-            + v_idx[None, :] * stride_ov
-        )
-        tl.store(o_off, o.to(tl.bfloat16))
+            + v0_idx[None, :] * stride_ov
+        )
+        tl.store(o0_off, o0.to(tl.bfloat16))
+
+        # v_block 1
+        v1_idx = BV + bv_idx
+        u1_off = (
+            U_ptr + b * stride_ub + h * stride_uh
+            + (t_start + i_idx)[:, None] * stride_ut
+            + v1_idx[None, :] * stride_uv
+        )
+        u1 = tl.load(u1_off)
+        wS1 = tl.dot(w, S1, input_precision="ieee")
+        v_i1 = u1 - wS1
+        part1_1 = tl.dot(qg, S1, input_precision="ieee")
+        part2_1 = tl.dot(Aqk, v_i1, input_precision="ieee")
+        o1 = part1_1 + part2_1
+        delta1 = tl.dot(tl.trans(K_g), v_i1, input_precision="ieee")
+        S1 = S1 * g_last_exp[:, None] + delta1
+        o1_off = (
+            O_ptr + b * stride_ob + h * stride_oh
+            + (t_start + i_idx)[:, None] * stride_ot
+            + v1_idx[None, :] * stride_ov
+        )
+        tl.store(o1_off, o1.to(tl.bfloat16))
 
 
 def kda_fwd(
@@ -343,7 +324,7 @@
     B, T, H, K_dim = q.shape
     V_d

... (561 more chars)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -1,6 +1,10 @@
     # V_C / BV S registers (one per V tile)
-    # V_C=128, BV=32 -> 4 S registers
+    # V_C=128, BV=16 -> 8 S registers
     S0 = tl.zeros((K_C, BV), dtype=tl.float32)
     S1 = tl.zeros((K_C, BV), dtype=tl.float32)
     S2 = tl.zeros((K_C, BV), dtype=tl.float32)
-    S3 = tl.zeros((K_C, BV), dtype=tl.float32)+    S3 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S4 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S5 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S6 = tl.zeros((K_C, BV), dtype=tl.float32)
+    S7 = tl.zeros((K_C, BV), dtype=tl.float32)

--- a//workspace/problems/02_kda_cutlass/solution.py
+++ b//workspace/problems/02_kda_cutlass/solution.py
@@ -367,7 +367,7 @@
         aqk_buf.stride(0), aqk_buf.stride(1), aqk_buf.stride(2), aqk_buf.stride(3),
         o.stride(0), o.stride(1), o.stride(2), o.stride(3),
         H=H, NT=NT, BT=chunk_size, K_C=K_dim, V_C=V_dim, BV=BV,
-        num_warps=2, num_stages=1,
+        num_warps=1, num_stages=1,
     )
     return o