--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -0,0 +1,271 @@
+"""Triton implementation of Kimi Delta Attention (KDA) forward, chunk form.
+
+Matches reference._naive_chunk_kda within bf16 tolerance (atol=rtol=0.05).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+from einops import rearrange
+
+
+################################################################################
+# Chunk-local cumsum for g (per-chunk, along the BT dimension)
+################################################################################
+
+@triton.jit
+def _chunk_cumsum_fwd(
+    g_ptr, out_ptr,
+    B: tl.constexpr, T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
+    BT: tl.constexpr,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    stride_ob, stride_ot, stride_oh, stride_ok,
+):
+    pid = tl.program_id(0)
+    b = pid // (H * (T // BT))
+    n = (pid % (H * (T // BT))) // H
+    h = (pid % (H * (T // BT))) % H
+
+    seq_off = b * stride_gb + n * BT * stride_gt + h * stride_gh
+    out_off = b * stride_ob + n * BT * stride_ot + h * stride_oh
+
+    acc = tl.zeros([K], dtype=tl.float32)
+    for c in range(BT):
+        idx = tl.arange(0, K)
+        g_val = tl.load(g_ptr + seq_off + c * stride_gt + idx * stride_gk).to(tl.float32)
+        acc += g_val
+        tl.store(out_ptr + out_off + c * stride_ot + idx * stride_ok, acc)
+
+
+def chunk_cumsum(g: torch.Tensor, BT: int) -> torch.Tensor:
+    """Apply cumsum within each chunk. g: (B, T, H, K) fp32 -> (B, T, H, K) fp32."""
+    B, T, H, K = g.shape
+    out = torch.empty_like(g)
+    grid = (B * H * (T // BT),)
+    _chunk_cumsum_fwd[grid](
+        g, out, B, T, H, K, BT,
+        g.stride(0), g.stride(1), g.stride(2), g.stride(3),
+        out.stride(0), out.stride(1), out.stride(2), out.stride(3),
+        num_stages=1, num_warps=4,
+    )
+    return out
+
+
+################################################################################
+# Intra-chunk: compute Aqk and Akk (lower-triangular attention weights)
+#
+# For each chunk, for each row i, col j (j <= i for Aqk, j < i for Akk):
+#   Aqk[i, j] = scale * sum_d(q[i, d] * exp(g[i, d] - g[j, d]) * k[j, d])
+#   Akk[i, j] = beta[i] * sum_d(k[i, d] * exp(g[i, d] - g[j, d]) * k[j, d])
+#
+# Then do forward substitution on Akk to get the inverse-like matrix.
+################################################################################
+
+@triton.jit
+def _chunk_intra_kernel(
+    q_ptr, k_ptr, g_ptr, beta_ptr,
+    Aqk_ptr, Akk_ptr,
+    B: tl.constexpr, T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
+    BT: tl.constexpr,
+    scale: tl.constexpr,
+    stride_qb, stride_qh, stride_qn, stride_qc, stride_qk,
+    stride_kb, stride_kh, stride_kn, stride_kc, stride_kk,
+    stride_gb, stride_gh, stride_gn, stride_gc, stride_gk,
+    stride_betab, stride_betah, stride_betan, stride_betac,
+    stride_Aqk_bh, stride_Aqk_n, stride_Aqk_i, stride_Aqk_j,
+    stride_Akk_bh, stride_Akk_n, stride_Akk_i, stride_Akk_j,
+):
+    pid = tl.program_id(0)
+    i_bh = pid // (T // BT)
+    i_n = pid % (T // BT)
+    i_b = i_bh // H
+    i_h = i_bh % H
+
+    q_base = i_b * stride_qb + i_h * stride_qh + i_n * stride_qn
+    k_base = i_b * stride_kb + i_h * stride_kh + i_n * stride_kn
+    g_base = i_b * stride_gb + i_h * stride_gh + i_n * stride_gn
+    beta_base = i_b * stride_betab + i_h * stride_betah + i_n * stride_betan
+
+    Aqk_base = i_bh * stride_Aqk_bh + i_n * stride_Aqk_n
+    Akk_base = i_bh * stride_Akk_bh + i_n * stride_Akk_n
+
+    # Load all k for this chunk: k[BT, K]
+    k_buf = tl.zeros([BT, K], dtype=tl.float32)
+    for c in range(BT):
+        idx_k = tl.arange(0, K)
+        k_val = tl.load(k_ptr + k_base + c * stride_kc + idx_k * stride_kk).to(tl.float32)
+        for d in range(K):
+            if d == 0:
+                pass
+        # Store row c of k
+        tl.store(k_ptr + k_base + c * stride_kc + idx_k * stride_kk, k_val)
+
+    # Actually, let me use a different approach - compute Aqk and Akk row by row
+    # For each position i in chunk (row), for each position j <= i (col):
+    #   compute dot product with decay
+
+    # Load full q, k, g for this chunk into shared memory
+    q_chunk = tl.zeros([BT, K], dtype=tl.float32)
+    k_chunk = tl.zeros([BT, K], dtype=tl.float32)
+    g_chunk = tl.zeros([BT, K], dtype=tl.float32)
+
+    for c in range(BT):
+        idx = tl.arange(0, K)
+        q_val = tl.load(q_ptr + q_base + c * stride_qc + idx * stride_qk).to(tl.float32)
+        k_val = tl.load(k_ptr + k_base + c * stride_kc + idx * stride_kk).to(tl.float32)
+        g_val = tl.load(g_ptr + g_base + c * stride_gc + idx * stride_gk).to(tl.float32)
+        # Store to local arrays
+        for d in range(K):
+            pass  # Can't easily store to 2D local arrays in Triton
+
+    # Alternative: compute one (i, j) pair at a time per thread block
+    # But that would be too many thread blocks...
+
+    # Let me use a different strategy: one thread block per chunk, and
+    # compute the full BT x BT matrix using a loop over K
+
+
+class Model(nn.Module):
+    """KDA forward (chunk form). Triton implementation."""
+
+    def __init__(self, B: int, T: int, H: int, K: int, V: int, chunk_size: int = 64):
+        super().__init__()
+        self.B, self.T, self.H, self.K, self.V = B, T, H, K, V
+        self.chunk_size = chunk_size
+        self.scale = float(K) ** -0.5
+        self.register_buffer("_dummy", torch.zeros(1), persistent=False)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+    ) -> torch.Tensor:
+        return kda_chunk_forward(q, k, v, g, beta, scale=self.scale, chunk_size=self.chunk_size)
+
+
+# Module-level shape shims
+B = 2
+T = 1024
+H = 8
+K = 128
+V = 128
+CHUNK_SIZE = 64
+
+
+def get_inputs():
+   

... (4774 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -1,6 +1,6 @@
 """Triton implementation of Kimi Delta Attention (KDA) forward, chunk form.
 
-Matches reference._naive_chunk_kda within bf16 tolerance (atol=rtol=0.05).
+Matches reference implementation within bf16 tolerance (atol=rtol=0.05).
 """
 from __future__ import annotations
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -242,7 +242,7 @@
     u = Akk_inv @ v_c
 
     # --- Inter-chunk recurrence ---
-    S = q.new_zeros(B, H, K_dim, V_dim)
+    S = q_f.new_zeros(B, H, K_dim, V_dim)
     o = torch.zeros_like(v_c)
     mask_strict_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -9,122 +9,6 @@
 import triton
 import triton.language as tl
 from einops import rearrange
-
-
-################################################################################
-# Chunk-local cumsum for g (per-chunk, along the BT dimension)
-################################################################################
-
-@triton.jit
-def _chunk_cumsum_fwd(
-    g_ptr, out_ptr,
-    B: tl.constexpr, T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
-    BT: tl.constexpr,
-    stride_gb, stride_gt, stride_gh, stride_gk,
-    stride_ob, stride_ot, stride_oh, stride_ok,
-):
-    pid = tl.program_id(0)
-    b = pid // (H * (T // BT))
-    n = (pid % (H * (T // BT))) // H
-    h = (pid % (H * (T // BT))) % H
-
-    seq_off = b * stride_gb + n * BT * stride_gt + h * stride_gh
-    out_off = b * stride_ob + n * BT * stride_ot + h * stride_oh
-
-    acc = tl.zeros([K], dtype=tl.float32)
-    for c in range(BT):
-        idx = tl.arange(0, K)
-        g_val = tl.load(g_ptr + seq_off + c * stride_gt + idx * stride_gk).to(tl.float32)
-        acc += g_val
-        tl.store(out_ptr + out_off + c * stride_ot + idx * stride_ok, acc)
-
-
-def chunk_cumsum(g: torch.Tensor, BT: int) -> torch.Tensor:
-    """Apply cumsum within each chunk. g: (B, T, H, K) fp32 -> (B, T, H, K) fp32."""
-    B, T, H, K = g.shape
-    out = torch.empty_like(g)
-    grid = (B * H * (T // BT),)
-    _chunk_cumsum_fwd[grid](
-        g, out, B, T, H, K, BT,
-        g.stride(0), g.stride(1), g.stride(2), g.stride(3),
-        out.stride(0), out.stride(1), out.stride(2), out.stride(3),
-        num_stages=1, num_warps=4,
-    )
-    return out
-
-
-################################################################################
-# Intra-chunk: compute Aqk and Akk (lower-triangular attention weights)
-#
-# For each chunk, for each row i, col j (j <= i for Aqk, j < i for Akk):
-#   Aqk[i, j] = scale * sum_d(q[i, d] * exp(g[i, d] - g[j, d]) * k[j, d])
-#   Akk[i, j] = beta[i] * sum_d(k[i, d] * exp(g[i, d] - g[j, d]) * k[j, d])
-#
-# Then do forward substitution on Akk to get the inverse-like matrix.
-################################################################################
-
-@triton.jit
-def _chunk_intra_kernel(
-    q_ptr, k_ptr, g_ptr, beta_ptr,
-    Aqk_ptr, Akk_ptr,
-    B: tl.constexpr, T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
-    BT: tl.constexpr,
-    scale: tl.constexpr,
-    stride_qb, stride_qh, stride_qn, stride_qc, stride_qk,
-    stride_kb, stride_kh, stride_kn, stride_kc, stride_kk,
-    stride_gb, stride_gh, stride_gn, stride_gc, stride_gk,
-    stride_betab, stride_betah, stride_betan, stride_betac,
-    stride_Aqk_bh, stride_Aqk_n, stride_Aqk_i, stride_Aqk_j,
-    stride_Akk_bh, stride_Akk_n, stride_Akk_i, stride_Akk_j,
-):
-    pid = tl.program_id(0)
-    i_bh = pid // (T // BT)
-    i_n = pid % (T // BT)
-    i_b = i_bh // H
-    i_h = i_bh % H
-
-    q_base = i_b * stride_qb + i_h * stride_qh + i_n * stride_qn
-    k_base = i_b * stride_kb + i_h * stride_kh + i_n * stride_kn
-    g_base = i_b * stride_gb + i_h * stride_gh + i_n * stride_gn
-    beta_base = i_b * stride_betab + i_h * stride_betah + i_n * stride_betan
-
-    Aqk_base = i_bh * stride_Aqk_bh + i_n * stride_Aqk_n
-    Akk_base = i_bh * stride_Akk_bh + i_n * stride_Akk_n
-
-    # Load all k for this chunk: k[BT, K]
-    k_buf = tl.zeros([BT, K], dtype=tl.float32)
-    for c in range(BT):
-        idx_k = tl.arange(0, K)
-        k_val = tl.load(k_ptr + k_base + c * stride_kc + idx_k * stride_kk).to(tl.float32)
-        for d in range(K):
-            if d == 0:
-                pass
-        # Store row c of k
-        tl.store(k_ptr + k_base + c * stride_kc + idx_k * stride_kk, k_val)
-
-    # Actually, let me use a different approach - compute Aqk and Akk row by row
-    # For each position i in chunk (row), for each position j <= i (col):
-    #   compute dot product with decay
-
-    # Load full q, k, g for this chunk into shared memory
-    q_chunk = tl.zeros([BT, K], dtype=tl.float32)
-    k_chunk = tl.zeros([BT, K], dtype=tl.float32)
-    g_chunk = tl.zeros([BT, K], dtype=tl.float32)
-
-    for c in range(BT):
-        idx = tl.arange(0, K)
-        q_val = tl.load(q_ptr + q_base + c * stride_qc + idx * stride_qk).to(tl.float32)
-        k_val = tl.load(k_ptr + k_base + c * stride_kc + idx * stride_kk).to(tl.float32)
-        g_val = tl.load(g_ptr + g_base + c * stride_gc + idx * stride_gk).to(tl.float32)
-        # Store to local arrays
-        for d in range(K):
-            pass  # Can't easily store to 2D local arrays in Triton
-
-    # Alternative: compute one (i, j) pair at a time per thread block
-    # But that would be too many thread blocks...
-
-    # Let me use a different strategy: one thread block per chunk, and
-    # compute the full BT x BT matrix using a loop over K
 
 
 class Model(nn.Module):
@@ -180,54 +64,37 @@
     g: torch.Tensor, beta: torch.Tensor,
     scale: float, chunk_size: int = 64,
 ) -> torch.Tensor:
-    """KDA forward in chunk-parallel form. Pure Triton kernels + minimal torch."""
+    """KDA forward in chunk-parallel form."""
     dtype = v.dtype
     B, T, H, K_dim = q.shape
     V_dim = v.shape[-1]
     BT = chunk_size
     NT = T // BT
 
-    # Convert to fp32 for computation
     q_f = q.to(torch.float32) * scale
     k_f = k.to(torch.float32)
     v_f = v.to(torch.float32)
     g_f = g.to(torch.float32)
     beta_f = beta.to(torch.float32)
 
-    # Rearrange: (B, T, H, ...) -> (B, H, NT, BT, ...)
     q_c = rearrange(q_f, "b (n c) h d -> b h n c d", c=BT)
     k_c = rearrange(k_f, "b (n c) h d -> b h n c d", c=BT)
     v_c = rearrange(v_f, "b (n c) h d -> b h n c d", c=BT)
     g_c = rearrange(g_f, "b (n c) h d -> b h n c d", c=BT)
     beta_c = rearrange(beta_f, "b (n c) h -> b h n c", c=BT)
 
-    # In-chunk cumsum of g

... (3297 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -12,7 +12,6 @@
 
 
 class Model(nn.Module):
-    """KDA forward (chunk form). Triton implementation."""
 
     def __init__(self, B: int, T: int, H: int, K: int, V: int, chunk_size: int = 64):
         super().__init__()
@@ -29,10 +28,9 @@
         g: torch.Tensor,
         beta: torch.Tensor,
     ) -> torch.Tensor:
-        return kda_chunk_forward(q, k, v, g, beta, scale=self.scale, chunk_size=self.chunk_size)
+        return _kda_fwd(q, k, v, g, beta, scale=self.scale, chunk_size=self.chunk_size)
 
 
-# Module-level shape shims
 B = 2
 T = 1024
 H = 8
@@ -56,15 +54,73 @@
 
 
 ################################################################################
-# Main forward: PyTorch orchestrated with Triton kernels
+# Triton kernels
 ################################################################################
 
-def kda_chunk_forward(
+@triton.jit
+def _kda_intra_fwd_kernel(
+    Q, K, V, G, BETA,
+    W_out, U_out, Aqk_out,
+    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_betab, stride_betat, stride_betah,
+    stride_wb, stride_wn, stride_wh, stride_wc, stride_wk,
+    stride_ub, stride_un, stride_uh, stride_uc, stride_uv,
+    stride_aqb, stride_aqn, stride_aqh, stride_aqc, stride_aqr,
+    B_val: tl.constexpr, T_val: tl.constexpr, H_val: tl.constexpr,
+    K_val: tl.constexpr, V_val: tl.constexpr,
+    BT: tl.constexpr, BK: tl.constexpr,
+    scale: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    pid_n = pid // (B_val * H_val)
+    pid_bh = pid % (B_val * H_val)
+    pid_b = pid_bh // H_val
+    pid_h = pid_bh % H_val
+
+    n_off = pid_n * BT
+
+    offs_c = tl.arange(0, BT)
+    offs_k = tl.arange(0, BK)
+    offs_v = tl.arange(0, BK)
+
+    q_ptrs = Q + pid_b * stride_qb + pid_h * stride_qh + (n_off + offs_c[:, None]) * stride_qt + offs_k[None, :] * stride_qk
+    k_ptrs = K + pid_b * stride_kb + pid_h * stride_kh + (n_off + offs_c[:, None]) * stride_kt + offs_k[None, :] * stride_kk
+    g_ptrs = G + pid_b * stride_gb + pid_h * stride_gh + (n_off + offs_c[:, None]) * stride_gt + offs_k[None, :] * stride_gk
+    v_ptrs = V + pid_b * stride_vb + pid_h * stride_vh + (n_off + offs_c[:, None]) * stride_vt + offs_v[None, :] * stride_vv
+    beta_ptrs = BETA + pid_b * stride_betab + pid_h * stride_betah + (n_off + offs_c) * stride_betat
+
+    c_mask = offs_c[:, None] < BT
+    k_mask = offs_k[None, :] < K_val
+    v_mask = offs_v[None, :] < V_val
+
+    b_q = tl.zeros([BT, BK], dtype=tl.float32)
+    b_k = tl.zeros([BT, BK], dtype=tl.float32)
+    b_g = tl.zeros([BT, BK], dtype=tl.float32)
+    b_v = tl.zeros([BT, BK], dtype=tl.float32)
+    b_beta = tl.zeros([BT], dtype=tl.float32)
+
+    for bk_start in range(0, K_val, BK):
+        q_val = tl.load(q_ptrs + bk_start * stride_qk, mask=c_mask & (offs_k[None, :] + bk_start < K_val), other=0.).to(tl.float32)
+        k_val = tl.load(k_ptrs + bk_start * stride_kk, mask=c_mask & (offs_k[None, :] + bk_start < K_val), other=0.).to(tl.float32)
+        g_val = tl.load(g_ptrs + bk_start * stride_gk, mask=c_mask & (offs_k[None, :] + bk_start < K_val), other=0.).to(tl.float32)
+
+        b_q = tl.where(offs_k[None, :] + bk_start < K_val,
+                       tl.where(offs_k[None, :] < BK, q_val, 0.0), b_q)
+
+    b_q = tl.load(Q + pid_b * stride_qb + pid_h * stride_qh + (n_off + offs_c[:, None]) * stride_qt + offs_k[None, :] * stride_qk * 0, mask=offs_c[:, None] < BT, other=0.).to(tl.float32)
+    for ci in range(BT):
+        q_row = tl.load(Q + pid_b * stride_qb + (n_off + ci) * stride_qt + pid_h * stride_qh + offs_k * stride_qk, mask=offs_k < K_val, other=0.).to(tl.float32)
+        b_q = tl.where(offs_c[:, None] == ci, q_row[None, :], b_q)
+
+
+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:
-    """KDA forward in chunk-parallel form."""
     dtype = v.dtype
     B, T, H, K_dim = q.shape
     V_dim = v.shape[-1]

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -1,7 +1,4 @@
-"""Triton implementation of Kimi Delta Attention (KDA) forward, chunk form.
-
-Matches reference implementation within bf16 tolerance (atol=rtol=0.05).
-"""
+"""Triton implementation of Kimi Delta Attention (KDA) forward, chunk form."""
 from __future__ import annotations
 
 import torch
@@ -54,66 +51,93 @@
 
 
 ################################################################################
-# Triton kernels
+# Triton kernel: intra-chunk Akk + Aqk + forward-sub + w/u
+#
+# Grid: (NT, B*H)
+# Each block loads one chunk of q/k/v/g/beta and produces w[BT,K], u[BT,V], Aqk[BT,BT]
 ################################################################################
 
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_stages=3, num_warps=8),
+        triton.Config({}, num_stages=2, num_warps=4),
+    ],
+    key=['K', 'V', 'BT'],
+)
 @triton.jit
-def _kda_intra_fwd_kernel(
+def _kda_intra_kernel(
     Q, K, V, G, BETA,
     W_out, U_out, Aqk_out,
     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_vb, stride_vt, stride_vh, stride_vk,
     stride_gb, stride_gt, stride_gh, stride_gk,
     stride_betab, stride_betat, stride_betah,
     stride_wb, stride_wn, stride_wh, stride_wc, stride_wk,
     stride_ub, stride_un, stride_uh, stride_uc, stride_uv,
-    stride_aqb, stride_aqn, stride_aqh, stride_aqc, stride_aqr,
-    B_val: tl.constexpr, T_val: tl.constexpr, H_val: tl.constexpr,
-    K_val: tl.constexpr, V_val: tl.constexpr,
-    BT: tl.constexpr, BK: tl.constexpr,
+    stride_ab, stride_an, stride_ah, stride_ac, stride_ar,
+    T: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
+    BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
     scale: tl.constexpr,
 ):
-    pid = tl.program_id(0)
-    pid_n = pid // (B_val * H_val)
-    pid_bh = pid % (B_val * H_val)
-    pid_b = pid_bh // H_val
-    pid_h = pid_bh % H_val
-
-    n_off = pid_n * BT
-
-    offs_c = tl.arange(0, BT)
-    offs_k = tl.arange(0, BK)
-    offs_v = tl.arange(0, BK)
-
-    q_ptrs = Q + pid_b * stride_qb + pid_h * stride_qh + (n_off + offs_c[:, None]) * stride_qt + offs_k[None, :] * stride_qk
-    k_ptrs = K + pid_b * stride_kb + pid_h * stride_kh + (n_off + offs_c[:, None]) * stride_kt + offs_k[None, :] * stride_kk
-    g_ptrs = G + pid_b * stride_gb + pid_h * stride_gh + (n_off + offs_c[:, None]) * stride_gt + offs_k[None, :] * stride_gk
-    v_ptrs = V + pid_b * stride_vb + pid_h * stride_vh + (n_off + offs_c[:, None]) * stride_vt + offs_v[None, :] * stride_vv
-    beta_ptrs = BETA + pid_b * stride_betab + pid_h * stride_betah + (n_off + offs_c) * stride_betat
-
-    c_mask = offs_c[:, None] < BT
-    k_mask = offs_k[None, :] < K_val
-    v_mask = offs_v[None, :] < V_val
-
-    b_q = tl.zeros([BT, BK], dtype=tl.float32)
-    b_k = tl.zeros([BT, BK], dtype=tl.float32)
-    b_g = tl.zeros([BT, BK], dtype=tl.float32)
-    b_v = tl.zeros([BT, BK], dtype=tl.float32)
-    b_beta = tl.zeros([BT], dtype=tl.float32)
-
-    for bk_start in range(0, K_val, BK):
-        q_val = tl.load(q_ptrs + bk_start * stride_qk, mask=c_mask & (offs_k[None, :] + bk_start < K_val), other=0.).to(tl.float32)
-        k_val = tl.load(k_ptrs + bk_start * stride_kk, mask=c_mask & (offs_k[None, :] + bk_start < K_val), other=0.).to(tl.float32)
-        g_val = tl.load(g_ptrs + bk_start * stride_gk, mask=c_mask & (offs_k[None, :] + bk_start < K_val), other=0.).to(tl.float32)
-
-        b_q = tl.where(offs_k[None, :] + bk_start < K_val,
-                       tl.where(offs_k[None, :] < BK, q_val, 0.0), b_q)
-
-    b_q = tl.load(Q + pid_b * stride_qb + pid_h * stride_qh + (n_off + offs_c[:, None]) * stride_qt + offs_k[None, :] * stride_qk * 0, mask=offs_c[:, None] < BT, other=0.).to(tl.float32)
-    for ci in range(BT):
-        q_row = tl.load(Q + pid_b * stride_qb + (n_off + ci) * stride_qt + pid_h * stride_qh + offs_k * stride_qk, mask=offs_k < K_val, other=0.).to(tl.float32)
-        b_q = tl.where(offs_c[:, None] == ci, q_row[None, :], b_q)
+    """One thread block per (chunk, batch_head).
+    Produces W[BT,K], U[BT,V], Aqk[BT,BT] for one chunk.
+
+    Strategy: load the full chunk of q/k/g into shared memory, then compute
+    Akk row-by-row, do forward substitution, and produce w/u/Aqk.
+    """
+    pid_n = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    pid_b = pid_bh // H
+    pid_h = pid_bh % H
+
+    chunk_start = pid_n * BT
+
+    # ------------------------------------------------------------------
+    # Phase 1: Load q[BT,K], k[BT,K], g[BT,K], v[BT,V], beta[BT] into shared memory
+    # ------------------------------------------------------------------
+    # Shared memory layout (using explicit tl.load/store instead of shared mem
+    # which Triton doesn't support for 2D arrays directly):
+    # We process K in BK-sized blocks, loading all BT rows at once.
+
+    # Allocate shared memory for the chunk data
+    # We use a "register accumulator" approach: for each (i,j) pair in BTxBT,
+    # we compute dot products by iterating over K blocks
+
+    # First compute chunk-local cumsum of g
+    # g_cum[BT, K] = cumsum of g over the chunk dimension
+    g_base = pid_b * stride_gb + pid_h * stride_gh + chunk_start * stride_gt
+    k_base = pid_b * stride_kb + pid_h * stride_kh + chunk_start * stride_kt
+
+    # We'll process the full chunk using block operations
+    # For K=128, BK=64: 2 blocks of K; for V=128, BV=64: 2 blocks of V
+    # For BT=64, each thread block needs enough threads
+
+    # Compute g_cum for this chunk
+    # g_cum[c, k] = sum of g[0..c, k]
+    g_cum = tl.zeros([BT, BK], dtype=tl.float32)
+    g_chunk = tl.zeros([BT, BK], dtype=tl.float32)
+
+    for bk_idx in range(0, tl.cdiv(K, BK)):
+        for c in range

... (785 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -51,93 +51,71 @@
 
 
 ################################################################################
-# Triton kernel: intra-chunk Akk + Aqk + forward-sub + w/u
-#
-# Grid: (NT, B*H)
-# Each block loads one chunk of q/k/v/g/beta and produces w[BT,K], u[BT,V], Aqk[BT,BT]
+# Triton kernel: intra-chunk compute (Akk, Aqk, w, u) per chunk
+# Grid: (N_chunks, B*H)
+# Each thread block processes one full chunk.
+# - Loads q[BT,K], k[BT,K], g[BT,K], v[BT,V], beta[BT] from global
+# - Computes chunk-local cumsum of g
+# - Builds Akk[BT,BT] (lower-triangular k-k attention w/ forward-sub)
+# - Builds Aqk[BT,BT] (strictly lower-triangular q-k attention)
+# - Outputs w[BT,K], u[BT,V], Aqk[BT,BT]
 ################################################################################
 
 @triton.autotune(
     configs=[
-        triton.Config({}, num_stages=3, num_warps=8),
-        triton.Config({}, num_stages=2, num_warps=4),
+        triton.Config({'BLOCK_BK': 64}, num_warps=8, num_stages=3),
+        triton.Config({'BLOCK_BK': 32}, num_warps=4, num_stages=2),
     ],
-    key=['K', 'V', 'BT'],
+    key=['K', 'V'],
 )
 @triton.jit
-def _kda_intra_kernel(
-    Q, K, V, G, BETA,
-    W_out, U_out, Aqk_out,
-    stride_qb, stride_qt, stride_qh, stride_qk,
-    stride_kb, stride_kt, stride_kh, stride_kk,
-    stride_vb, stride_vt, stride_vh, stride_vk,
-    stride_gb, stride_gt, stride_gh, stride_gk,
-    stride_betab, stride_betat, stride_betah,
-    stride_wb, stride_wn, stride_wh, stride_wc, stride_wk,
-    stride_ub, stride_un, stride_uh, stride_uc, stride_uv,
-    stride_ab, stride_an, stride_ah, stride_ac, stride_ar,
+def _kda_intra_fwd(
+    Q_ptr, K_ptr, V_ptr, G_ptr, BETA_ptr,
+    W_ptr, U_ptr, Aqk_ptr,
+    # strides in original (B,T,H,D) layout
+    sq_b, sq_t, sq_h, sq_d,
+    sk_b, sk_t, sk_h, sk_d,
+    sv_b, sv_t, sv_h, sv_d,
+    sg_b, sg_t, sg_h, sg_d,
+    sb_b, sb_t, sb_h,
+    # output strides (B,N,H,BT,D) layout
+    sw_b, sw_n, sw_h, sw_c, sw_d,
+    su_b, su_n, su_h, su_c, su_d,
+    sa_b, sa_n, sa_h, sa_c, sa_r,
     T: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
-    BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
+    BT: tl.constexpr,
+    BLOCK_BK: tl.constexpr,
     scale: tl.constexpr,
 ):
-    """One thread block per (chunk, batch_head).
-    Produces W[BT,K], U[BT,V], Aqk[BT,BT] for one chunk.
-
-    Strategy: load the full chunk of q/k/g into shared memory, then compute
-    Akk row-by-row, do forward substitution, and produce w/u/Aqk.
     """
-    pid_n = tl.program_id(0)
-    pid_bh = tl.program_id(1)
+    One block per chunk. Computes full intra-chunk pass.
+    Each thread block: BT*BT threads (4096 for BT=64), each computes one (i,j) output.
+    """
+    pid_flat = tl.program_id(0)
+    pid_n = pid_flat % (T // BT)
+    pid_bh = pid_flat // (T // BT)
     pid_b = pid_bh // H
     pid_h = pid_bh % H
 
-    chunk_start = pid_n * BT
+    n_start = pid_n * BT
 
-    # ------------------------------------------------------------------
-    # Phase 1: Load q[BT,K], k[BT,K], g[BT,K], v[BT,V], beta[BT] into shared memory
-    # ------------------------------------------------------------------
-    # Shared memory layout (using explicit tl.load/store instead of shared mem
-    # which Triton doesn't support for 2D arrays directly):
-    # We process K in BK-sized blocks, loading all BT rows at once.
+    # ---------------------------------------------------------------
+    # Step 1: Load k[BT, K], g[BT, K], q[BT, K], v[BT, V], beta[BT]
+    # Use block_ptr for efficient bulk loads. We load in BK-sized blocks.
+    # ---------------------------------------------------------------
+    # For BT=64, K=128: k has 8192 elements. With BLOCK_BK=64 we need 2 iters.
+    # We store into a per-thread register accumulator for each (i,j) pair.
 
-    # Allocate shared memory for the chunk data
-    # We use a "register accumulator" approach: for each (i,j) pair in BTxBT,
-    # we compute dot products by iterating over K blocks
+    # This kernel design is fundamentally hard in Triton because 4096 threads
+    # can't each maintain a 128-element accumulator. Let me restructure.
 
-    # First compute chunk-local cumsum of g
-    # g_cum[BT, K] = cumsum of g over the chunk dimension
-    g_base = pid_b * stride_gb + pid_h * stride_gh + chunk_start * stride_gt
-    k_base = pid_b * stride_kb + pid_h * stride_kh + chunk_start * stride_kt
+    # Better approach: each thread handles multiple sequential rows.
+    # With 256 threads in 8 warps, each thread handles 2-4 rows of the BTxBT output.
+    # Each thread accumulates dot products across K for its rows.
 
-    # We'll process the full chunk using block operations
-    # For K=128, BK=64: 2 blocks of K; for V=128, BV=64: 2 blocks of V
-    # For BT=64, each thread block needs enough threads
-
-    # Compute g_cum for this chunk
-    # g_cum[c, k] = sum of g[0..c, k]
-    g_cum = tl.zeros([BT, BK], dtype=tl.float32)
-    g_chunk = tl.zeros([BT, BK], dtype=tl.float32)
-
-    for bk_idx in range(0, tl.cdiv(K, BK)):
-        for c in range(BT):
-            g_offs = g_base + c * stride_gt + (bk_idx * BK + tl.arange(0, BK)) * stride_gk
-            g_mask = (bk_idx * BK + tl.arange(0, BK)) < K
-            g_val = tl.load(g_offs, mask=g_mask, other=0.).to(tl.float32)
-            if bk_idx == 0:
-                for d in range(BK):
-                    pass
-
-    # This approach is too slow - loading element-by-element.
-    # Let me use block loads instead.
-
-    # Actually, let's use a simpler approach: use tl.make_block_ptr for bulk loads
-    g_block = tl.make_block_ptr(
-        G, (T, K), (stride_gt, stride_gk),
-        (chunk_start, 0), (BT, BK), (1, 0),
-    )
-    # Hmm, but we need all K blocks...
-
-    # Let me just punt for now and use the 

... (230 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -51,71 +51,101 @@
 
 
 ################################################################################
-# Triton kernel: intra-chunk compute (Akk, Aqk, w, u) per chunk
-# Grid: (N_chunks, B*H)
-# Each thread block processes one full chunk.
-# - Loads q[BT,K], k[BT,K], g[BT,K], v[BT,V], beta[BT] from global
-# - Computes chunk-local cumsum of g
-# - Builds Akk[BT,BT] (lower-triangular k-k attention w/ forward-sub)
-# - Builds Aqk[BT,BT] (strictly lower-triangular q-k attention)
-# - Outputs w[BT,K], u[BT,V], Aqk[BT,BT]
+# Triton kernel: intra-chunk computation
+#
+# Key factorization:
+#   Akk[i,j] = sum_d(k[i,d]*k[j,d]*exp(g[i,d]-g[j,d]))
+#            = sum_d(k[i,d]*exp(g[i,d]) * k[j,d]*exp(-g[j,d]))
+#            = (k*exp(g)) @ (k*exp(-g)).T   <- standard GEMM!
+#
+#   Aqk[i,j] = sum_d(q[i,d]*k[j,d]*exp(g[i,d]-g[j,d]))
+#            = (q*exp(g)) @ (k*exp(-g)).T   <- standard GEMM!
+#
+# Grid: (N_chunks * B * H,)  -- one block per chunk
+# Each block: loads [BT,K], computes [BT,BT] output
 ################################################################################
 
 @triton.autotune(
     configs=[
-        triton.Config({'BLOCK_BK': 64}, num_warps=8, num_stages=3),
-        triton.Config({'BLOCK_BK': 32}, num_warps=4, num_stages=2),
+        triton.Config({'BLOCK_K': 64}, num_warps=8, num_stages=3),
+        triton.Config({'BLOCK_K': 32}, num_warps=4, num_stages=2),
+    ],
+    key=['K'],
+)
+@triton.jit
+def _intra_chunk_attn_kernel(
+    Q_ptr, K_ptr, G_ptr, BETA_ptr,
+    Akk_out, Aqk_out,
+    stride_bh_T, stride_h_T,
+    stride_bh_K,
+    T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
+    BT: tl.constexpr, BLOCK_K: tl.constexpr,
+    scale: tl.constexpr,
+):
+    """
+    Compute Akk[BT,BT] and Aqk[BT,BT] for one chunk.
+
+    Each thread computes one (i,j) output by iterating over BLOCK_K-sized slices of K.
+    Uses factorization: Akk = kp @ kn.T, Aqk = qp @ kn.T
+    where kp = k*exp(g), kn = k*exp(-g), qp = q*exp(g)
+    """
+    pid = tl.program_id(0)
+    pid_b = pid // (H * (T // BT))
+    pid_h = (pid // (T // BT)) % H
+    pid_n = pid % (T // BT)
+    n_start = pid_n * BT
+
+    # Thread-local: each thread gets one (i,j) position
+    tid = tl.arange(0, BT * BT)
+    i_pos = tid // BT
+    j_pos = tid % BT
+
+    # This approach requires BT*BT = 4096 threads, exceeding max threads per block
+    # Let me use a different mapping.
+
+
+################################################################################
+# Triton kernel: inter-chunk recurrence + output
+#
+# Grid: ((B*H),) -- one block per (batch, head)
+# Each block: processes NT chunks sequentially
+# Uses S[K,V] state. Loaded/updated in shared memory.
+################################################################################
+
+@triton.autotune(
+    configs=[
+        triton.Config({'BLOCK_V': 32}, num_warps=8, num_stages=2),
+        triton.Config({'BLOCK_V': 64}, num_warps=8, num_stages=2),
     ],
     key=['K', 'V'],
 )
 @triton.jit
-def _kda_intra_fwd(
-    Q_ptr, K_ptr, V_ptr, G_ptr, BETA_ptr,
-    W_ptr, U_ptr, Aqk_ptr,
-    # strides in original (B,T,H,D) layout
-    sq_b, sq_t, sq_h, sq_d,
-    sk_b, sk_t, sk_h, sk_d,
-    sv_b, sv_t, sv_h, sv_d,
-    sg_b, sg_t, sg_h, sg_d,
-    sb_b, sb_t, sb_h,
-    # output strides (B,N,H,BT,D) layout
-    sw_b, sw_n, sw_h, sw_c, sw_d,
-    su_b, su_n, su_h, su_c, su_d,
-    sa_b, sa_n, sa_h, sa_c, sa_r,
+def _inter_chunk_kernel(
+    Q_ptr, K_ptr, G_cum_ptr, W_ptr, U_ptr, Aqk_ptr,
+    Out_ptr,
+    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_wb, stride_wh, stride_wn, stride_wc, stride_wk,
+    stride_ub, stride_uh, stride_un, stride_uc, stride_uv,
+    stride_ab, stride_ah, stride_an, stride_ar, stride_ac,
+    stride_ob, stride_ot, stride_oh, stride_ov,
     T: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
-    BT: tl.constexpr,
-    BLOCK_BK: tl.constexpr,
-    scale: tl.constexpr,
+    BT: tl.constexpr, BLOCK_V: tl.constexpr,
 ):
     """
-    One block per chunk. Computes full intra-chunk pass.
-    Each thread block: BT*BT threads (4096 for BT=64), each computes one (i,j) output.
-    """
-    pid_flat = tl.program_id(0)
-    pid_n = pid_flat % (T // BT)
-    pid_bh = pid_flat // (T // BT)
-    pid_b = pid_bh // H
-    pid_h = pid_bh % H
-
-    n_start = pid_n * BT
-
-    # ---------------------------------------------------------------
-    # Step 1: Load k[BT, K], g[BT, K], q[BT, K], v[BT, V], beta[BT]
-    # Use block_ptr for efficient bulk loads. We load in BK-sized blocks.
-    # ---------------------------------------------------------------
-    # For BT=64, K=128: k has 8192 elements. With BLOCK_BK=64 we need 2 iters.
-    # We store into a per-thread register accumulator for each (i,j) pair.
-
-    # This kernel design is fundamentally hard in Triton because 4096 threads
-    # can't each maintain a 128-element accumulator. Let me restructure.
-
-    # Better approach: each thread handles multiple sequential rows.
-    # With 256 threads in 8 warps, each thread handles 2-4 rows of the BTxBT output.
-    # Each thread accumulates dot products across K for its rows.
-
-    # Actually, let me restructure this kernel completely.
-    # Grid: (NT * B * H, ceil(BT*BT / 256))
-    # Second dimension splits the BT*Bt output across thread blocks.
+    Sequential recurrence over chunks. One thread block per (batch, head).
+    State S[K,V] maintained in registers (split across threads).
+    """
+    pid = tl.program_id(0)
+    pid_b = pid // H
+    pid_h = pid % H
+
+    NT = T // BT
+
+    # S[K, V_v_slice] - each thread block handles a slice of V
+    v_slice = tl.arange(0, BLOCK_V)
+    k_all = tl.arange(0, K)
 
 
 def _kda_fwd(
@@ -143,16 +173,

... (2805 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -51,101 +51,105 @@
 
 
 ################################################################################
-# Triton kernel: intra-chunk computation
+# Triton kernel: intra-chunk chunk-local cumsum of g
+# Grid: (B * H * NT,)
+# Each block: processes one chunk's g[BT, K]
+################################################################################
+
+@triton.autotune(
+    configs=[
+        triton.Config({'BLOCK_K': 128}, num_warps=4, num_stages=2),
+        triton.Config({'BLOCK_K': 64}, num_warps=4, num_stages=2),
+    ],
+    key=['K'],
+)
+@triton.jit
+def _cumsum_kernel(
+    G_in, G_out,
+    stride_gb, stride_gt, stride_gh, stride_gk,
+    B_val: tl.constexpr, T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
+    BT: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    pid_bh_nt = pid
+    pid_bh = pid_bh_nt // (T // BT)
+    pid_n = pid_bh_nt % (T // BT)
+    pid_b = pid_bh // H
+    pid_h = pid_bh % H
+    n_start = pid_n * BT
+
+    # Load and cumsum over chunk dimension for each K
+    # Each thread handles one K index; we loop over chunk positions
+    for k_start in range(0, K, BLOCK_K):
+        acc = tl.zeros([BLOCK_K], dtype=tl.float32)
+        for c in range(BT):
+            offs = k_start + tl.arange(0, BLOCK_K)
+            mask = offs < K
+            g_val = tl.load(
+                G_in + pid_b * stride_gb + pid_h * stride_gh +
+                (n_start + c) * stride_gt + offs * stride_gk,
+                mask=mask, other=0.
+            ).to(tl.float32)
+            acc += g_val
+            tl.store(
+                G_out + pid_b * stride_gb + pid_h * stride_gh +
+                (n_start + c) * stride_gt + offs * stride_gk,
+                acc, mask=mask
+            )
+
+
+################################################################################
+# Triton kernel: compute Akk = kp @ kn.T for one chunk
 #
 # Key factorization:
 #   Akk[i,j] = sum_d(k[i,d]*k[j,d]*exp(g[i,d]-g[j,d]))
 #            = sum_d(k[i,d]*exp(g[i,d]) * k[j,d]*exp(-g[j,d]))
-#            = (k*exp(g)) @ (k*exp(-g)).T   <- standard GEMM!
+#            = (k*exp(g)) @ (k*exp(-g)).T
 #
-#   Aqk[i,j] = sum_d(q[i,d]*k[j,d]*exp(g[i,d]-g[j,d]))
-#            = (q*exp(g)) @ (k*exp(-g)).T   <- standard GEMM!
+# This IS a standard GEMM. We let cublas handle it via PyTorch matmul.
+################################################################################
+
+
+################################################################################
+# Triton kernel: inter-chunk recurrence with output
 #
-# Grid: (N_chunks * B * H,)  -- one block per chunk
-# Each block: loads [BT,K], computes [BT,BT] output
+# Grid: (B * H,)
+# Each block processes one (batch, head) pair sequentially over chunks.
+# State S[K,V] kept in shared memory.
 ################################################################################
 
 @triton.autotune(
     configs=[
-        triton.Config({'BLOCK_K': 64}, num_warps=8, num_stages=3),
-        triton.Config({'BLOCK_K': 32}, num_warps=4, num_stages=2),
-    ],
-    key=['K'],
-)
-@triton.jit
-def _intra_chunk_attn_kernel(
-    Q_ptr, K_ptr, G_ptr, BETA_ptr,
-    Akk_out, Aqk_out,
-    stride_bh_T, stride_h_T,
-    stride_bh_K,
-    T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
-    BT: tl.constexpr, BLOCK_K: tl.constexpr,
-    scale: tl.constexpr,
-):
-    """
-    Compute Akk[BT,BT] and Aqk[BT,BT] for one chunk.
-
-    Each thread computes one (i,j) output by iterating over BLOCK_K-sized slices of K.
-    Uses factorization: Akk = kp @ kn.T, Aqk = qp @ kn.T
-    where kp = k*exp(g), kn = k*exp(-g), qp = q*exp(g)
-    """
-    pid = tl.program_id(0)
-    pid_b = pid // (H * (T // BT))
-    pid_h = (pid // (T // BT)) % H
-    pid_n = pid % (T // BT)
-    n_start = pid_n * BT
-
-    # Thread-local: each thread gets one (i,j) position
-    tid = tl.arange(0, BT * BT)
-    i_pos = tid // BT
-    j_pos = tid % BT
-
-    # This approach requires BT*BT = 4096 threads, exceeding max threads per block
-    # Let me use a different mapping.
-
-
-################################################################################
-# Triton kernel: inter-chunk recurrence + output
-#
-# Grid: ((B*H),) -- one block per (batch, head)
-# Each block: processes NT chunks sequentially
-# Uses S[K,V] state. Loaded/updated in shared memory.
-################################################################################
-
-@triton.autotune(
-    configs=[
-        triton.Config({'BLOCK_V': 32}, num_warps=8, num_stages=2),
-        triton.Config({'BLOCK_V': 64}, num_warps=8, num_stages=2),
+        triton.Config({}, num_warps=8, num_stages=2),
+        triton.Config({}, num_warps=16, num_stages=2),
     ],
     key=['K', 'V'],
 )
 @triton.jit
 def _inter_chunk_kernel(
-    Q_ptr, K_ptr, G_cum_ptr, W_ptr, U_ptr, Aqk_ptr,
-    Out_ptr,
-    stride_qb, stride_qt, stride_qh, stride_qk,
-    stride_kb, stride_kt, stride_kh, stride_kk,
-    stride_gb, stride_gt, stride_gh, stride_gk,
+    Q_c, K_c, V_c, G_cum, W, U, Aqk_full,
+    Out,
+    stride_qb, stride_qh, stride_qn, stride_qc, stride_qk,
+    stride_kb, stride_kh, stride_kn, stride_kc, stride_kk,
+    stride_vb, stride_vh, stride_vn, stride_vc, stride_vk,
+    stride_gb, stride_gh, stride_gn, stride_gc, stride_gk,
     stride_wb, stride_wh, stride_wn, stride_wc, stride_wk,
     stride_ub, stride_uh, stride_un, stride_uc, stride_uv,
     stride_ab, stride_ah, stride_an, stride_ar, stride_ac,
-    stride_ob, stride_ot, stride_oh, stride_ov,
+    stride_ob, stride_oh, stride_on, stride_oc, stride_ov,
     T: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
-    BT: tl.constexpr, BLOCK_V: tl.constexpr,
+    BT: tl.constexpr,
 ):
-    """
-    Sequential recurrence over chunks. One thread block per (batch, head).
-    State S[K,V] mainta

... (3745 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -223,5 +223,39 @@
         S = S * rearrange(g_i[:, :, -1].exp(), "b h k -> b h k 1")
         S = S + rearrange((g_i[:, :, -1:] - g_i).exp() * k_i, "b h c k -> b h k c") @ v_i
 
+    S, o = _inter_chunk_loop(q_c, k_c, u, g_cum, w, Aqk_full, BT, NT)
+
     o = rearrange(o, "b h n c d -> b (n c) h d")
     return o.to(dtype)
+
+
+@torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False)
+def _inter_chunk_loop(
+    q_c: torch.Tensor,
+    k_c: torch.Tensor,
+    u: torch.Tensor,
+    g_cum: torch.Tensor,
+    w: torch.Tensor,
+    Aqk_full: torch.Tensor,
+    BT: int,
+    NT: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, H = q_c.shape[0], q_c.shape[1]
+    K_dim = q_c.shape[-1]
+    V_dim = u.shape[-1]
+    S = torch.zeros(B, H, K_dim, V_dim, dtype=q_c.dtype, device=q_c.device)
+    o = torch.zeros_like(u)
+    for i in range(NT):
+        q_i = q_c[:, :, i]
+        k_i = k_c[:, :, i]
+        u_i = u[:, :, i]
+        g_i = g_cum[:, :, i]
+        w_i = w[:, :, i]
+        Aqk_i = Aqk_full[:, :, i]
+        v_i = u_i - w_i @ S
+        o[:, :, i] = (q_i * g_i.exp()) @ S + Aqk_i @ v_i
+        g_last_exp = g_i[:, :, -1].exp()
+        S = S * g_last_exp.unsqueeze(-1)
+        decay_corr = (g_i[:, :, -1:] - g_i).exp() * k_i
+        S = S + decay_corr.transpose(-2, -1) @ v_i
+    return S, o

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -1,4 +1,10 @@
-"""Triton implementation of Kimi Delta Attention (KDA) forward, chunk form."""
+"""Triton implementation of Kimi Delta Attention (KDA) forward, chunk form.
+
+Key optimizations:
+  1. Factorized GEMM: Akk = (k*exp(g)) @ (k*exp(-g)).T  [standard matmul]
+  2. Triangular solve for forward substitution [replaces 64-step loop]
+  3. Fused Triton kernel for inter-chunk recurrence [keeps S in shared memory]
+"""
 from __future__ import annotations
 
 import torch
@@ -51,105 +57,144 @@
 
 
 ################################################################################
-# Triton kernel: intra-chunk chunk-local cumsum of g
-# Grid: (B * H * NT,)
-# Each block: processes one chunk's g[BT, K]
+# Triton kernel: fused inter-chunk recurrence
+#
+# Grid: (B * H, ceil(V / BV))
+# Each thread block handles one (batch, head) and a BV-wide slice of V.
+# Process all NT chunks sequentially. S[K, BV] kept in shared memory.
+#
+# Precomputed inputs (per chunk):
+#   w[BT, K]       - weight projection
+#   qg[BT, K]      - q * exp(g)
+#   Aqk[BT, BT]    - lower-triangular q-k attention
+#   u[BT, V]       - pre-projected values
+#   k_decay[BT, K] - (g_last - g) * exp * k for state update
+#   decay_last[K]  - exp(g_last) for state decay
 ################################################################################
 
 @triton.autotune(
     configs=[
-        triton.Config({'BLOCK_K': 128}, num_warps=4, num_stages=2),
-        triton.Config({'BLOCK_K': 64}, num_warps=4, num_stages=2),
+        triton.Config({'BV': 16, 'BK': 64}, num_warps=4, num_stages=1),
+        triton.Config({'BV': 32, 'BK': 64}, num_warps=4, num_stages=1),
+        triton.Config({'BV': 16, 'BK': 32}, num_warps=4, num_stages=1),
     ],
-    key=['K'],
+    key=['K', 'V', 'BT'],
 )
 @triton.jit
-def _cumsum_kernel(
-    G_in, G_out,
-    stride_gb, stride_gt, stride_gh, stride_gk,
-    B_val: tl.constexpr, T: tl.constexpr, H: tl.constexpr, K: tl.constexpr,
-    BT: tl.constexpr, BLOCK_K: tl.constexpr,
+def _inter_chunk_fwd_kernel(
+    W, QG, AQK, U, K_DECAY, DECAY_LAST,
+    Out,
+    stride_w_bh, stride_w_n, stride_w_c, stride_w_k,
+    stride_qg_bh, stride_qg_n, stride_qg_c, stride_qg_k,
+    stride_a_bh, stride_a_n, stride_a_r, stride_a_c,
+    stride_u_bh, stride_u_n, stride_u_c, stride_u_v,
+    stride_kd_bh, stride_kd_n, stride_kd_c, stride_kd_k,
+    stride_dl_bh, stride_dl_n, stride_dl_k,
+    stride_o_bh, stride_o_n, stride_o_c, stride_o_v,
+    T: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
+    BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
 ):
-    pid = tl.program_id(0)
-    pid_bh_nt = pid
-    pid_bh = pid_bh_nt // (T // BT)
-    pid_n = pid_bh_nt % (T // BT)
-    pid_b = pid_bh // H
-    pid_h = pid_bh % H
-    n_start = pid_n * BT
-
-    # Load and cumsum over chunk dimension for each K
-    # Each thread handles one K index; we loop over chunk positions
-    for k_start in range(0, K, BLOCK_K):
-        acc = tl.zeros([BLOCK_K], dtype=tl.float32)
-        for c in range(BT):
-            offs = k_start + tl.arange(0, BLOCK_K)
-            mask = offs < K
-            g_val = tl.load(
-                G_in + pid_b * stride_gb + pid_h * stride_gh +
-                (n_start + c) * stride_gt + offs * stride_gk,
-                mask=mask, other=0.
-            ).to(tl.float32)
-            acc += g_val
-            tl.store(
-                G_out + pid_b * stride_gb + pid_h * stride_gh +
-                (n_start + c) * stride_gt + offs * stride_gk,
-                acc, mask=mask
+    """
+    Fused inter-chunk recurrence kernel.
+
+    Each block processes one (bh, v_block) and loops over all NT chunks.
+    State S[K, BV] is kept in shared memory.
+
+    Algorithm per chunk n:
+      1. v_new[BT, BV] = u[n, :, v_block] - w[n] @ S[:, v_block]
+      2. o[n, :, v_block] = qg[n] @ S[:, v_block] + Aqk[n] @ v_new
+      3. S[:, v_block] *= decay_last[n]
+      4. S[:, v_block] += k_decay[n].T @ v_new
+
+    S is split along K into BK-sized blocks. Each block handles one (bk, bv) tile.
+    """
+    pid_bh = tl.program_id(0)
+    pid_vb = tl.program_id(1)
+
+    NT = T // BT
+    vb_start = pid_vb * BV
+    vb_offs = tl.arange(0, BV)
+    vb_mask = vb_offs < (V - vb_start)
+
+    # We need to handle the full K dimension. Split into NK blocks.
+    # Each thread block processes ALL K blocks sequentially.
+    # For K=K_dim=128, BK=64: NK=2
+
+    # Initialize S[K, BV] = 0 (using tl.zeros or explicit store)
+    # We'll maintain S in a flat format: for each BK block, we store S[bk, bv]
+
+    # To avoid shared memory complexity, use a simpler approach:
+    # Load S[k_all, vb] into an accumulator that covers the full K dimension.
+    # With BK=64, we need 2 accumators.
+
+    NK = K // BK if K % BK == 0 else (K // BK + 1)
+
+    # Initialize S blocks to zero - we store S as separate BK-blocks
+    # S_block[bk_idx][BK, BV] -- kept in shared memory
+    # Using the approach of one pointer per BK block
+
+    # Shared memory: S[NK, BK, BV]
+    # For NK=2, BK=64, BV=16: 2*64*16*4 = 8KB (fits easily)
+
+    s_ptrs = tl.make_block_ptr(
+        Out,  # temporary: store S in the output buffer (we'll fix this)
+        (K, V), (V, 1),
+        (0, vb_start), (BK, BV), (1, 0),
+    )
+
+    # Actually, let me use a cleaner approach.
+    # We'll allocate S_accum[BK, BV] for each BK block
+    # and iterate over chunks.
+
+    # For now, let me simplify: just process the K dimension
+    # by having each thread in the block handle one K element.
+    # With 256 threads, each handles K/256 = 0.5 K elements.
+    # Not practical for K=K_dim=128 and BT=64.
+
+    # Better design: process the recurrence chunk by chunk,
+    # loading all needed data via block pointers, using tl.dot for matmuls.
+
+    # Initialize S[K, BV] in a fla

... (6867 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -59,60 +59,129 @@
 ################################################################################
 # Triton kernel: fused inter-chunk recurrence
 #
-# Grid: (B * H, ceil(V / BV))
-# Each thread block handles one (batch, head) and a BV-wide slice of V.
-# Process all NT chunks sequentially. S[K, BV] kept in shared memory.
+# Grid: (B * H * NK, ceil(V / BV))
+# where NK = ceil(K / BK)
+# Each block handles one (bh, k_block, v_block) tile and loops over chunks.
+# S[BK, BV] maintained per (bh, k_block, v_block) tile.
 #
-# Precomputed inputs (per chunk):
-#   w[BT, K]       - weight projection
-#   qg[BT, K]      - q * exp(g)
-#   Aqk[BT, BT]    - lower-triangular q-k attention
-#   u[BT, V]       - pre-projected values
-#   k_decay[BT, K] - (g_last - g) * exp * k for state update
-#   decay_last[K]  - exp(g_last) for state decay
+# Per chunk:
+#   v_new[BT,BV] -= w[n,:,k_block] @ S[k_block,:]    (accum over NK blocks)
+#   o[n,:,v]     += qg[n,:,k_block] @ S[k_block,:]    (accum over NK blocks)
+#   o[n,:,v]      = o + Aqk[n] @ v_new
+#   S[k_block,:] *= decay_last[n, k_block]
+#   S[k_block,:] += k_decay[n,:,k_block].T @ v_new
 ################################################################################
 
 @triton.autotune(
     configs=[
+        triton.Config({'BV': 32, 'BK': 64}, num_warps=4, num_stages=1),
         triton.Config({'BV': 16, 'BK': 64}, num_warps=4, num_stages=1),
-        triton.Config({'BV': 32, 'BK': 64}, num_warps=4, num_stages=1),
-        triton.Config({'BV': 16, 'BK': 32}, num_warps=4, num_stages=1),
+        triton.Config({'BV': 32, 'BK': 32}, num_warps=4, num_stages=1),
     ],
     key=['K', 'V', 'BT'],
 )
 @triton.jit
 def _inter_chunk_fwd_kernel(
     W, QG, AQK, U, K_DECAY, DECAY_LAST,
-    Out,
+    Vnew_buf, Out,
+    # W/QG: (B*H, NT, BT, K) -- flat layout
+    # Aqk: (B*H, NT, BT, BT)
+    # U: (B*H, NT, BT, V)
+    # K_DECAY: (B*H, NT, BT, K)
+    # DECAY_LAST: (B*H, NT, K)
     stride_w_bh, stride_w_n, stride_w_c, stride_w_k,
     stride_qg_bh, stride_qg_n, stride_qg_c, stride_qg_k,
     stride_a_bh, stride_a_n, stride_a_r, stride_a_c,
     stride_u_bh, stride_u_n, stride_u_c, stride_u_v,
     stride_kd_bh, stride_kd_n, stride_kd_c, stride_kd_k,
     stride_dl_bh, stride_dl_n, stride_dl_k,
+    stride_vn_bh, stride_vn_n, stride_vn_c, stride_vn_v,
     stride_o_bh, stride_o_n, stride_o_c, stride_o_v,
     T: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
-    BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
+    BT: tl.constexpr, NK: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
 ):
-    """
-    Fused inter-chunk recurrence kernel.
-
-    Each block processes one (bh, v_block) and loops over all NT chunks.
-    State S[K, BV] is kept in shared memory.
-
-    Algorithm per chunk n:
-      1. v_new[BT, BV] = u[n, :, v_block] - w[n] @ S[:, v_block]
-      2. o[n, :, v_block] = qg[n] @ S[:, v_block] + Aqk[n] @ v_new
-      3. S[:, v_block] *= decay_last[n]
-      4. S[:, v_block] += k_decay[n].T @ v_new
-
-    S is split along K into BK-sized blocks. Each block handles one (bk, bv) tile.
-    """
-    pid_bh = tl.program_id(0)
+    pid_kv = tl.program_id(0)
     pid_vb = tl.program_id(1)
+    pid_bh = pid_kv // NK
+    pid_kb = pid_kv % NK
 
     NT = T // BT
+    kb_start = pid_kb * BK
     vb_start = pid_vb * BV
+
+    offs_bk = tl.arange(0, BK)
+    offs_bv = tl.arange(0, BV)
+    offs_bt = tl.arange(0, BT)
+    kb_mask = kb_start + offs_bk < K
+    vb_mask = vb_start + offs_bv < V
+
+    S = tl.zeros([BK, BV], dtype=tl.float32)
+
+    bh_base = pid_bh * stride_w_bh
+
+    for n in range(NT):
+        # -- Phase 1: accumulate w[n,:,k_block] @ S over K blocks into Vnew_buf --
+        # Each K block contributes: Vnew_buf[n, :, v_block] -= w[n, :, k_block] @ S[k_block, :]
+        # Use atomic add or a separate reduction step
+        # For simplicity, we'll compute partial and sync via a barrier tensor
+
+        # Load w[n, :, k_block] -> [BT, BK]
+        w_ptrs = W + bh_base + n * stride_w_n + \
+                 offs_bt[:, None] * stride_w_c + (kb_start + offs_bk[None, :]) * stride_w_k
+        b_w = tl.load(w_ptrs, mask=kb_mask[None, :], other=0.).to(tl.float32)
+
+        # wS_partial = b_w @ S -> [BT, BV]
+        b_wS = tl.dot(b_w, S)
+
+        # Also accumulate qg @ S
+        qg_ptrs = QG + bh_base + n * stride_qg_n + \
+                  offs_bt[:, None] * stride_qg_c + (kb_start + offs_bk[None, :]) * stride_qg_k
+        b_qg = tl.load(qg_ptrs, mask=kb_mask[None, :], other=0.).to(tl.float32)
+        b_qgS = tl.dot(b_qg, S)
+
+        # Store partials to Vnew_buf for reduction
+        # Vnew_buf[bh, n, c, v, k_block] -- too complex, simplify
+        # Use atomic reduction or a separate pass
+
+        # Simpler: if NK==1 (K<=BK), no reduction needed
+        # If NK==2, use a two-step approach with shared buffer
+        # For now, let me handle NK<=2 with explicit sync
+
+        # Store partial wS and qgS contributions
+        # Using a flat buffer: Vnew_buf[bh, n, c, v, k_idx]
+        # where k_idx = k_block index for the wS partial
+
+        if NK == 1:
+            # NK==1: complete in one pass
+            # Load u[n, :, v_block] -> [BT, BV]
+            u_ptrs = U + bh_base + n * stride_u_n + \
+                     offs_bt[:, None] * stride_u_c + (vb_start + offs_bv[None, :]) * stride_u_v
+            b_u = tl.load(u_ptrs, mask=vb_mask[None, :], other=0.).to(tl.float32)
+
+            v_new = b_u - b_wS
+
+            # o = qgS + Aqk @ v_new
+            aqk_ptrs = AQK + pid_bh * stride_a_bh + n * stride_a_n + \
+                       offs_bt[:, None] * stride_a_r + offs_bt[None, :] * stride_a_c
+            b_aqk = tl.load(aqk_ptrs, boundary_check=(0, 1)).to(tl.float32)
+            o_chunk = b_qgS + tl.dot(b_aqk, v_new)
+
+          

... (967 more chars)