--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -0,0 +1,193 @@
+"""KDA forward (chunk form) -- custom kernel implementation.
+
+Vectorized PyTorch implementation with chunk-parallel intra-chunk processing
+and sequential inter-chunk recurrence. No FLA imports.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+OP_TYPE = "linear_attention"
+SUPPORTED_PRECISIONS = ["bf16"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+
+def _chunk_kda_forward(
+    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, chunk form. No initial/final state. Returns o in v's dtype."""
+    dtype = v.dtype
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    BT = chunk_size
+    assert T % BT == 0
+    NT = T // BT
+
+    # Work in fp32
+    q = q.float() * scale
+    k = k.float()
+    v = v.float()
+    g = g.float()
+    beta = beta.float()
+
+    # Reshape to chunk form: (B, H, NT, BT, D)
+    q = q.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
+    k = k.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
+    v = v.reshape(B, T, H, V).permute(0, 2, 1, 3).reshape(B, H, NT, BT, V)
+    g = g.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
+    beta = beta.reshape(B, T, H).permute(0, 2, 1).reshape(B, H, NT, BT)
+
+    # In-chunk cumsum of g (the input g already has cumsum applied per the spec,
+    # but the reference does cumsum again on the reshaped g -- we follow the reference)
+    g = g.cumsum(-2)
+
+    # ---- Build A_kk (intra-chunk K-K interaction, lower-triangular w/ diag masked) ----
+    # g shape: (B, H, NT, BT, K)
+    # k shape: (B, H, NT, BT, K)
+    # For each position i, compute: A[..., i, c] = sum_d k[c,d] * exp(g[c,d] - g[i,d]) * k[i,d]
+    # This is: A[..., i, c] = (k * exp(g - g_i)).sum(-1) dot k_i, but per-element
+
+    # g_diff[b,h,n,c,i,d] = g[b,h,n,c,d] - g[b,h,n,i,d]
+    # A[b,h,n,i,c] = sum_d k[b,h,n,c,d] * exp(g_diff[b,h,n,c,i,d]) * k[b,h,n,i,d]
+    #               = sum_d k[b,h,n,c,d] * k[b,h,n,i,d] * exp(g[b,h,n,c,d] - g[b,h,n,i,d])
+
+    # Efficient vectorization: use outer product structure
+    # k_exp = k * exp(g)  -> (B, H, NT, BT, K)
+    # Then A[b,h,n,i,c] = sum_d k_exp[b,h,n,c,d] * k[b,h,n,i,d] * exp(-g[b,h,n,i,d])
+    #                    = sum_d (k[b,h,n,c,d]*exp(g[b,h,n,c,d])) * (k[b,h,n,i,d]*exp(-g[b,h,n,i,d]))
+
+    k_pos = k * g.exp()          # k * exp(g), shape (B,H,NT,BT,K)
+    k_neg = k * (-g).exp()       # k * exp(-g), shape (B,H,NT,BT,K)
+
+    # A[b,h,n,i,c] = sum_d k_pos[b,h,n,c,d] * k_neg[b,h,n,i,d]
+    # This is a batched outer product: A = k_pos @ k_neg^T
+    # k_pos: (B,H,NT,BT,K), k_neg: (B,H,NT,BT,K)
+    # A: (B,H,NT,BT,BT) = k_pos @ k_neg.transpose(-1,-2)
+    A = torch.einsum("bhnck,bhnik->bhnic", k_pos, k_neg)  # (B,H,NT,BT,BT)
+
+    # Apply beta and mask
+    A = A * beta[..., None]  # broadcast beta over last dim
+
+    # Mask upper triangular (diagonal and above) -> set to 0, then negate
+    mask_diag_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    A = -A.masked_fill(mask_diag_upper, 0)
+
+    # Newton iteration to solve cumulative system
+    # For i in 1..BT-1: A[..., i, :i] += sum_j(A[..., i, j] * A[..., j, :i]) for j < i
+    # This is essentially computing (I - L)^{-1} where L is strictly lower triangular
+    # We can do this with a series of rank-1 updates
+    for i in range(1, BT):
+        # A[..., i, :i] += (A[..., i, :, None] * A[..., :, :i]).sum(-2)
+        # But we need to be careful about which elements of A[..., i, :] are nonzero
+        # Since A is lower-triangular (upper masked to 0), A[..., i, j] is nonzero only for j < i
+        # So the sum is over j=0..i-1
+        A[..., i, :i] = A[..., i, :i] + (A[..., i, :i, None] * A[..., :i, :i].transpose(-2, -1)).sum(-2)
+
+    # Add identity and apply beta
+    A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta[..., None, :]
+
+    # Compute w = A @ (exp(g) * k) and u = A @ v
+    # g_exp_k = exp(g) * k: (B,H,NT,BT,K)
+    g_exp_k = g.exp() * k
+    # w = A @ g_exp_k: (B,H,NT,BT,BT) @ (B,H,NT,BT,K) -> (B,H,NT,BT,K)
+    w = torch.einsum("bhnic,bhncd->bhnic", A, g_exp_k)
+    # u = A @ v: (B,H,NT,BT,BT) @ (B,H,NT,BT,V) -> (B,H,NT,BT,V)
+    u = torch.einsum("bhnic,bhncd->bhnic", A, v)
+
+    # ---- Recurrent inter-chunk pass ----
+    S = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
+    o = torch.zeros(B, H, NT, BT, V, dtype=torch.float32, device=q.device)
+
+    # Strict upper triangular mask for Aqk
+    mask_strict_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
+
+    for i in range(NT):
+        q_i = q[:, :, i]       # (B, H, BT, K)
+        k_i = k[:, :, i]       # (B, H, BT, K)
+        g_i = g[:, :, i]       # (B, H, BT, K)
+        u_i = u[:, :, i]       # (B, H, BT, V)
+        w_i = w[:, :, i]       # (B, H, BT, K)
+
+        # Build Aqk: intra-chunk q-k attention with decay
+        # Aqk[b,h,j,c] = sum_d q[b,h,c,d] * exp(g[b,h,c,d] - g[b,h,j,d]) * k[b,h,j,d]
+        # Similar structure to A but with q instead of k for the "row" side
+        q_pos = q_i * g_i.exp()    # (B,H,BT,K)
+        k_neg_i = k_i * (-g_i).exp()  # (B,H,BT,K)
+        Aqk = torch.einsum("bhck,bhjk->bhjc", q_pos, k_neg_i)  # (B,H,BT,BT)
+
+        # Mask strictly upper triangular
+        Aqk = Aqk.masked_fill(mask_strict_upper, 0)
+
+        # v_i = u_i - w_i @ S
+        # w_i @ S: (B,H,BT,K) @ (B,H,K,V) -> (B,H,BT,V)
+        v_i = u_i - torch.einsum("bhck,bhkd->bhcd", w_i, S)
+
+        # o_i = (q_i * exp(g_i)) @ S + Aqk @ v_i
+        # (q_i * exp(g_i)) @ S: (B,H,BT,K) @ (B,H,K,V) -> (B,H,BT,V)
+        q_exp = q_i * g_i.exp()
+        o_i = torch.einsum("bhck,bhkd->bhcd", q_exp, S) + t

... (1988 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
@@ -13,7 +13,7 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-def _chunk_kda_forward(
+def _kda_fwd(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -168,7 +168,7 @@
         g: torch.Tensor,
         beta: torch.Tensor,
     ) -> torch.Tensor:
-        return _chunk_kda_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)
 
 
 B = 2

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -97,9 +97,9 @@
     # g_exp_k = exp(g) * k: (B,H,NT,BT,K)
     g_exp_k = g.exp() * k
     # w = A @ g_exp_k: (B,H,NT,BT,BT) @ (B,H,NT,BT,K) -> (B,H,NT,BT,K)
-    w = torch.einsum("bhnic,bhncd->bhnic", A, g_exp_k)
+    w = torch.einsum("bhnij,bhnjd->bhnid", A, g_exp_k)
     # u = A @ v: (B,H,NT,BT,BT) @ (B,H,NT,BT,V) -> (B,H,NT,BT,V)
-    u = torch.einsum("bhnic,bhncd->bhnic", A, v)
+    u = torch.einsum("bhnij,bhnjd->bhnid", A, v)
 
     # ---- Recurrent inter-chunk pass ----
     S = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -84,11 +84,7 @@
     # This is essentially computing (I - L)^{-1} where L is strictly lower triangular
     # We can do this with a series of rank-1 updates
     for i in range(1, BT):
-        # A[..., i, :i] += (A[..., i, :, None] * A[..., :, :i]).sum(-2)
-        # But we need to be careful about which elements of A[..., i, :] are nonzero
-        # Since A is lower-triangular (upper masked to 0), A[..., i, j] is nonzero only for j < i
-        # So the sum is over j=0..i-1
-        A[..., i, :i] = A[..., i, :i] + (A[..., i, :i, None] * A[..., :i, :i].transpose(-2, -1)).sum(-2)
+        A[..., i, :i] = A[..., i, :i] + (A[..., i, :, None] * A[..., :, :i]).sum(-2)
 
     # Add identity and apply beta
     A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta[..., None, :]

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -116,7 +116,7 @@
         # Similar structure to A but with q instead of k for the "row" side
         q_pos = q_i * g_i.exp()    # (B,H,BT,K)
         k_neg_i = k_i * (-g_i).exp()  # (B,H,BT,K)
-        Aqk = torch.einsum("bhck,bhjk->bhjc", q_pos, k_neg_i)  # (B,H,BT,BT)
+        Aqk = torch.einsum("bhcd,bhjd->bhjc", q_pos, k_neg_i)  # (B,H,BT,BT)
 
         # Mask strictly upper triangular
         Aqk = Aqk.masked_fill(mask_strict_upper, 0)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -116,7 +116,7 @@
         # Similar structure to A but with q instead of k for the "row" side
         q_pos = q_i * g_i.exp()    # (B,H,BT,K)
         k_neg_i = k_i * (-g_i).exp()  # (B,H,BT,K)
-        Aqk = torch.einsum("bhcd,bhjd->bhjc", q_pos, k_neg_i)  # (B,H,BT,BT)
+        Aqk = torch.einsum("bhcd,bhjd->bhcj", q_pos, k_neg_i)  # (B,H,BT,BT)
 
         # Mask strictly upper triangular
         Aqk = Aqk.masked_fill(mask_strict_upper, 0)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -44,45 +44,23 @@
     g = g.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
     beta = beta.reshape(B, T, H).permute(0, 2, 1).reshape(B, H, NT, BT)
 
-    # In-chunk cumsum of g (the input g already has cumsum applied per the spec,
-    # but the reference does cumsum again on the reshaped g -- we follow the reference)
+    # In-chunk cumsum of g
     g = g.cumsum(-2)
 
-    # ---- Build A_kk (intra-chunk K-K interaction, lower-triangular w/ diag masked) ----
-    # g shape: (B, H, NT, BT, K)
-    # k shape: (B, H, NT, BT, K)
-    # For each position i, compute: A[..., i, c] = sum_d k[c,d] * exp(g[c,d] - g[i,d]) * k[i,d]
-    # This is: A[..., i, c] = (k * exp(g - g_i)).sum(-1) dot k_i, but per-element
+    # ---- Build A_kk ----
+    # A[c, i] = sum_d k[c,d] * exp(g[c,d] - g[i,d]) * k[i,d]
+    # = sum_d (k[c,d]*exp(g[c,d])) * (k[i,d]*exp(-g[i,d]))
+    k_pos = k * g.exp()       # k * exp(g)
+    k_neg = k * (-g).exp()    # k * exp(-g)
+    # A[b,h,n,c,i] = sum_d k_pos[b,h,n,c,d] * k_neg[b,h,n,i,d]
+    A = torch.einsum("bhnck,bhnik->bhnci", k_pos, k_neg)
 
-    # g_diff[b,h,n,c,i,d] = g[b,h,n,c,d] - g[b,h,n,i,d]
-    # A[b,h,n,i,c] = sum_d k[b,h,n,c,d] * exp(g_diff[b,h,n,c,i,d]) * k[b,h,n,i,d]
-    #               = sum_d k[b,h,n,c,d] * k[b,h,n,i,d] * exp(g[b,h,n,c,d] - g[b,h,n,i,d])
-
-    # Efficient vectorization: use outer product structure
-    # k_exp = k * exp(g)  -> (B, H, NT, BT, K)
-    # Then A[b,h,n,i,c] = sum_d k_exp[b,h,n,c,d] * k[b,h,n,i,d] * exp(-g[b,h,n,i,d])
-    #                    = sum_d (k[b,h,n,c,d]*exp(g[b,h,n,c,d])) * (k[b,h,n,i,d]*exp(-g[b,h,n,i,d]))
-
-    k_pos = k * g.exp()          # k * exp(g), shape (B,H,NT,BT,K)
-    k_neg = k * (-g).exp()       # k * exp(-g), shape (B,H,NT,BT,K)
-
-    # A[b,h,n,i,c] = sum_d k_pos[b,h,n,c,d] * k_neg[b,h,n,i,d]
-    # This is a batched outer product: A = k_pos @ k_neg^T
-    # k_pos: (B,H,NT,BT,K), k_neg: (B,H,NT,BT,K)
-    # A: (B,H,NT,BT,BT) = k_pos @ k_neg.transpose(-1,-2)
-    A = torch.einsum("bhnck,bhnik->bhnic", k_pos, k_neg)  # (B,H,NT,BT,BT)
-
-    # Apply beta and mask
-    A = A * beta[..., None]  # broadcast beta over last dim
-
-    # Mask upper triangular (diagonal and above) -> set to 0, then negate
+    # Apply beta (broadcast over last dim) and mask upper triangular
     mask_diag_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    A = A * beta[..., None]
     A = -A.masked_fill(mask_diag_upper, 0)
 
-    # Newton iteration to solve cumulative system
-    # For i in 1..BT-1: A[..., i, :i] += sum_j(A[..., i, j] * A[..., j, :i]) for j < i
-    # This is essentially computing (I - L)^{-1} where L is strictly lower triangular
-    # We can do this with a series of rank-1 updates
+    # Newton iteration: forward substitution to solve (I + A)x = b
     for i in range(1, BT):
         A[..., i, :i] = A[..., i, :i] + (A[..., i, :, None] * A[..., :, :i]).sum(-2)
 
@@ -90,55 +68,41 @@
     A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta[..., None, :]
 
     # Compute w = A @ (exp(g) * k) and u = A @ v
-    # g_exp_k = exp(g) * k: (B,H,NT,BT,K)
     g_exp_k = g.exp() * k
-    # w = A @ g_exp_k: (B,H,NT,BT,BT) @ (B,H,NT,BT,K) -> (B,H,NT,BT,K)
     w = torch.einsum("bhnij,bhnjd->bhnid", A, g_exp_k)
-    # u = A @ v: (B,H,NT,BT,BT) @ (B,H,NT,BT,V) -> (B,H,NT,BT,V)
     u = torch.einsum("bhnij,bhnjd->bhnid", A, v)
 
     # ---- Recurrent inter-chunk pass ----
     S = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
     o = torch.zeros(B, H, NT, BT, V, dtype=torch.float32, device=q.device)
 
-    # Strict upper triangular mask for Aqk
     mask_strict_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
 
-    for i in range(NT):
-        q_i = q[:, :, i]       # (B, H, BT, K)
-        k_i = k[:, :, i]       # (B, H, BT, K)
-        g_i = g[:, :, i]       # (B, H, BT, K)
-        u_i = u[:, :, i]       # (B, H, BT, V)
-        w_i = w[:, :, i]       # (B, H, BT, K)
+    for ic in range(NT):
+        q_i = q[:, :, ic]
+        k_i = k[:, :, ic]
+        g_i = g[:, :, ic]
+        u_i = u[:, :, ic]
+        w_i = w[:, :, ic]
 
-        # Build Aqk: intra-chunk q-k attention with decay
-        # Aqk[b,h,j,c] = sum_d q[b,h,c,d] * exp(g[b,h,c,d] - g[b,h,j,d]) * k[b,h,j,d]
-        # Similar structure to A but with q instead of k for the "row" side
-        q_pos = q_i * g_i.exp()    # (B,H,BT,K)
-        k_neg_i = k_i * (-g_i).exp()  # (B,H,BT,K)
-        Aqk = torch.einsum("bhcd,bhjd->bhcj", q_pos, k_neg_i)  # (B,H,BT,BT)
-
-        # Mask strictly upper triangular
+        # Aqk[c, j] = sum_d q[c,d] * exp(g[c,d] - g[j,d]) * k[j,d]
+        q_pos_i = q_i * g_i.exp()
+        k_neg_i = k_i * (-g_i).exp()
+        Aqk = torch.einsum("bhcd,bhjd->bhcj", q_pos_i, k_neg_i)
         Aqk = Aqk.masked_fill(mask_strict_upper, 0)
 
         # v_i = u_i - w_i @ S
-        # w_i @ S: (B,H,BT,K) @ (B,H,K,V) -> (B,H,BT,V)
         v_i = u_i - torch.einsum("bhck,bhkd->bhcd", w_i, S)
 
         # o_i = (q_i * exp(g_i)) @ S + Aqk @ v_i
-        # (q_i * exp(g_i)) @ S: (B,H,BT,K) @ (B,H,K,V) -> (B,H,BT,V)
         q_exp = q_i * g_i.exp()
-        o_i = torch.einsum("bhck,bhkd->bhcd", q_exp, S) + torch.einsum("bhjc,bhcd->bhjd", Aqk, v_i)
-        o[:, :, i] = o_i
+        o_i = torch.einsum("bhck,bhkd->bhcd", q_exp, S) + torch.einsum("bhcj,bhcd->bhjd", Aqk, v_i)
+        o[:, :, ic] = o_i
 
         # Update state: S = S * exp(g_last) + k_decayed^T @ v_i
-        # g_last = g_i[:, :, -1, :]  -> (B, H, K)
-        g_last = g_i[:, :, -1, :]  # (B, H, K)
-        S = S * g_last.exp().unsqueeze(-1)  # (B, H, K, V)
-
-        # k_decayed[b,h,c,k] = exp(g_last[k] - g_i[b,h,c,k]) * k_i[b,h,c,k]
-        # k_decayed^T @ v_i: (B,H,K,BT) @ (B,H,BT,V) -> (B,H,K,V)
-  

... (336 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,7 +1,7 @@
 """KDA forward (chunk form) -- custom kernel implementation.
 
-Vectorized PyTorch implementation with chunk-parallel intra-chunk processing
-and sequential inter-chunk recurrence. No FLA imports.
+Vectorized PyTorch implementation with batched matmuls.
+No FLA imports.
 """
 from __future__ import annotations
 
@@ -37,76 +37,77 @@
     g = g.float()
     beta = beta.float()
 
-    # Reshape to chunk form: (B, H, NT, BT, D)
-    q = q.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
-    k = k.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
-    v = v.reshape(B, T, H, V).permute(0, 2, 1, 3).reshape(B, H, NT, BT, V)
-    g = g.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B, H, NT, BT, K)
-    beta = beta.reshape(B, T, H).permute(0, 2, 1).reshape(B, H, NT, BT)
+    # Reshape to chunk form: (B*H, NT, BT, D) for batched matmul
+    q = q.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B * H, NT, BT, K)
+    k = k.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B * H, NT, BT, K)
+    v = v.reshape(B, T, H, V).permute(0, 2, 1, 3).reshape(B * H, NT, BT, V)
+    g = g.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B * H, NT, BT, K)
+    beta = beta.reshape(B, T, H).permute(0, 2, 1).reshape(B * H, NT, BT)
 
     # In-chunk cumsum of g
     g = g.cumsum(-2)
 
     # ---- Build A_kk ----
-    # A[c, i] = sum_d k[c,d] * exp(g[c,d] - g[i,d]) * k[i,d]
-    # = sum_d (k[c,d]*exp(g[c,d])) * (k[i,d]*exp(-g[i,d]))
-    k_pos = k * g.exp()       # k * exp(g)
-    k_neg = k * (-g).exp()    # k * exp(-g)
-    # A[b,h,n,c,i] = sum_d k_pos[b,h,n,c,d] * k_neg[b,h,n,i,d]
-    A = torch.einsum("bhnck,bhnik->bhnci", k_pos, k_neg)
+    # A[c, i] = sum_d (k[c,d]*exp(g[c,d])) * (k[i,d]*exp(-g[i,d]))
+    k_pos = k * g.exp()       # (B*H, NT, BT, K)
+    k_neg = k * (-g).exp()    # (B*H, NT, BT, K)
 
-    # Apply beta (broadcast over last dim) and mask upper triangular
+    # A = k_pos @ k_neg^T: (B*H, NT, BT, K) @ (B*H, NT, K, BT) -> (B*H, NT, BT, BT)
+    A = torch.matmul(k_pos, k_neg.transpose(-1, -2))
+
+    # Apply beta and mask upper triangular
     mask_diag_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
-    A = A * beta[..., None]
+    A = A * beta.unsqueeze(-1)
     A = -A.masked_fill(mask_diag_upper, 0)
 
-    # Newton iteration: forward substitution to solve (I + A)x = b
+    # Newton iteration (forward substitution)
     for i in range(1, BT):
         A[..., i, :i] = A[..., i, :i] + (A[..., i, :, None] * A[..., :, :i]).sum(-2)
 
     # Add identity and apply beta
-    A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta[..., None, :]
+    A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta.unsqueeze(-2)
 
     # Compute w = A @ (exp(g) * k) and u = A @ v
     g_exp_k = g.exp() * k
-    w = torch.einsum("bhnij,bhnjd->bhnid", A, g_exp_k)
-    u = torch.einsum("bhnij,bhnjd->bhnid", A, v)
+    w = torch.matmul(A, g_exp_k)  # (B*H, NT, BT, K)
+    u = torch.matmul(A, v)         # (B*H, NT, BT, V)
 
     # ---- Recurrent inter-chunk pass ----
-    S = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
-    o = torch.zeros(B, H, NT, BT, V, dtype=torch.float32, device=q.device)
+    BH = B * H
+    S = torch.zeros(BH, K, V, dtype=torch.float32, device=q.device)
+    o = torch.zeros(BH, NT, BT, V, dtype=torch.float32, device=q.device)
 
     mask_strict_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
 
     for ic in range(NT):
-        q_i = q[:, :, ic]
-        k_i = k[:, :, ic]
-        g_i = g[:, :, ic]
-        u_i = u[:, :, ic]
-        w_i = w[:, :, ic]
+        q_i = q[:, ic]       # (BH, BT, K)
+        k_i = k[:, ic]       # (BH, BT, K)
+        g_i = g[:, ic]       # (BH, BT, K)
+        u_i = u[:, ic]       # (BH, BT, V)
+        w_i = w[:, ic]       # (BH, BT, K)
 
         # Aqk[c, j] = sum_d q[c,d] * exp(g[c,d] - g[j,d]) * k[j,d]
         q_pos_i = q_i * g_i.exp()
         k_neg_i = k_i * (-g_i).exp()
-        Aqk = torch.einsum("bhcd,bhjd->bhcj", q_pos_i, k_neg_i)
+        Aqk = torch.matmul(q_pos_i, k_neg_i.transpose(-1, -2))
         Aqk = Aqk.masked_fill(mask_strict_upper, 0)
 
         # v_i = u_i - w_i @ S
-        v_i = u_i - torch.einsum("bhck,bhkd->bhcd", w_i, S)
+        v_i = u_i - torch.matmul(w_i, S)
 
         # o_i = (q_i * exp(g_i)) @ S + Aqk @ v_i
         q_exp = q_i * g_i.exp()
-        o_i = torch.einsum("bhck,bhkd->bhcd", q_exp, S) + torch.einsum("bhcj,bhcd->bhjd", Aqk, v_i)
-        o[:, :, ic] = o_i
+        o_i = torch.matmul(q_exp, S) + torch.matmul(Aqk, v_i)
+        o[:, ic] = o_i
 
         # Update state: S = S * exp(g_last) + k_decayed^T @ v_i
-        g_last = g_i[:, :, -1, :]
-        S = S * g_last.exp().unsqueeze(-1)
-        k_decayed = (g_last.unsqueeze(-2) - g_i).exp() * k_i
-        S = S + torch.einsum("bhck,bhcd->bhkd", k_decayed, v_i)
+        g_last = g_i[:, -1, :]  # (BH, K)
+        S = S * g_last.unsqueeze(-1).exp()
+        k_decayed = (g_last.unsqueeze(-2) - g_i).exp() * k_i  # (BH, BT, K)
+        S = S + torch.matmul(k_decayed.transpose(-1, -2), v_i)
 
-    # Reshape back: (B, H, NT, BT, V) -> (B, T, H, V)
-    o = o.permute(0, 2, 3, 1, 4).reshape(B, T, H, V)
+    # Reshape back: (B*H, NT, BT, V) -> (B, T, H, V)
+    o = o.reshape(B, H, NT, BT, V).permute(0, 2, 3, 1, 4).reshape(B, T, H, V)
     return o.to(dtype)
 
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/02_kda_cutlass/solution.py
@@ -48,29 +48,23 @@
     g = g.cumsum(-2)
 
     # ---- Build A_kk ----
-    # A[c, i] = sum_d (k[c,d]*exp(g[c,d])) * (k[i,d]*exp(-g[i,d]))
-    k_pos = k * g.exp()       # (B*H, NT, BT, K)
-    k_neg = k * (-g).exp()    # (B*H, NT, BT, K)
-
-    # A = k_pos @ k_neg^T: (B*H, NT, BT, K) @ (B*H, NT, K, BT) -> (B*H, NT, BT, BT)
+    k_pos = k * g.exp()
+    k_neg = k * (-g).exp()
     A = torch.matmul(k_pos, k_neg.transpose(-1, -2))
 
-    # Apply beta and mask upper triangular
     mask_diag_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
     A = A * beta.unsqueeze(-1)
-    A = -A.masked_fill(mask_diag_upper, 0)
+    A_neg = -A.masked_fill(mask_diag_upper, 0)
 
-    # Newton iteration (forward substitution)
-    for i in range(1, BT):
-        A[..., i, :i] = A[..., i, :i] + (A[..., i, :, None] * A[..., :, :i]).sum(-2)
-
-    # Add identity and apply beta
-    A = (A + torch.eye(BT, dtype=torch.float32, device=q.device)) * beta.unsqueeze(-2)
+    # Newton iteration replaced with solve_triangular
+    I_mat = torch.eye(BT, dtype=torch.float32, device=q.device).expand_as(A_neg).contiguous()
+    A_inv = torch.linalg.solve_triangular(I_mat - A_neg, I_mat, upper=False, unitriangular=False)
+    A = A_inv * beta.unsqueeze(-2)
 
     # Compute w = A @ (exp(g) * k) and u = A @ v
     g_exp_k = g.exp() * k
-    w = torch.matmul(A, g_exp_k)  # (B*H, NT, BT, K)
-    u = torch.matmul(A, v)         # (B*H, NT, BT, V)
+    w = torch.matmul(A, g_exp_k)
+    u = torch.matmul(A, v)
 
     # ---- Recurrent inter-chunk pass ----
     BH = B * H
@@ -80,11 +74,11 @@
     mask_strict_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
 
     for ic in range(NT):
-        q_i = q[:, ic]       # (BH, BT, K)
-        k_i = k[:, ic]       # (BH, BT, K)
-        g_i = g[:, ic]       # (BH, BT, K)
-        u_i = u[:, ic]       # (BH, BT, V)
-        w_i = w[:, ic]       # (BH, BT, K)
+        q_i = q[:, ic]
+        k_i = k[:, ic]
+        g_i = g[:, ic]
+        u_i = u[:, ic]
+        w_i = w[:, ic]
 
         # Aqk[c, j] = sum_d q[c,d] * exp(g[c,d] - g[j,d]) * k[j,d]
         q_pos_i = q_i * g_i.exp()
@@ -100,13 +94,13 @@
         o_i = torch.matmul(q_exp, S) + torch.matmul(Aqk, v_i)
         o[:, ic] = o_i
 
-        # Update state: S = S * exp(g_last) + k_decayed^T @ v_i
-        g_last = g_i[:, -1, :]  # (BH, K)
+        # Update state
+        g_last = g_i[:, -1, :]
         S = S * g_last.unsqueeze(-1).exp()
-        k_decayed = (g_last.unsqueeze(-2) - g_i).exp() * k_i  # (BH, BT, K)
+        k_decayed = (g_last.unsqueeze(-2) - g_i).exp() * k_i
         S = S + torch.matmul(k_decayed.transpose(-1, -2), v_i)
 
-    # Reshape back: (B*H, NT, BT, V) -> (B, T, H, V)
+    # Reshape back
     o = o.reshape(B, H, NT, BT, V).permute(0, 2, 3, 1, 4).reshape(B, T, H, V)
     return o.to(dtype)
 

--- 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,16 +1,121 @@
-"""KDA forward (chunk form) -- custom kernel implementation.
-
-Vectorized PyTorch implementation with batched matmuls.
+"""KDA forward (chunk form) -- Triton kernel implementation.
+
+Triton kernel for the inter-chunk recurrence loop + PyTorch for intra-chunk.
 No FLA imports.
 """
 from __future__ import annotations
 
 import torch
 import torch.nn as nn
+import triton
+import triton.language as tl
 
 OP_TYPE = "linear_attention"
 SUPPORTED_PRECISIONS = ["bf16"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+
+@triton.jit
+def _inter_chunk_kernel(
+    # Input pointers
+    q_ptr, k_ptr, g_ptr, w_ptr, u_ptr, o_ptr,
+    # Strides for q, k, g: (BH, NT, BT, K)
+    stride_q_bh, stride_q_n, stride_q_t
+    , stride_q_k,
+    stride_k_bh, stride_k_n, stride_k_t, stride_k_k,
+    stride_g_bh, stride_g_n, stride_g_t, stride_g_k,
+    # Strides for w: (BH, NT, BT, K)
+    stride_w_bh, stride_w_n, stride_w_t, stride_w_k,
+    # Strides for u: (BH, NT, BT, V)
+    stride_u_bh, stride_u_n, stride_u_t, stride_u_v,
+    # Strides for o: (BH, NT, BT, V)
+    stride_o_bh, stride_o_n, stride_o_t, stride_o_v,
+    # Dimensions
+    NT, BT, K, V,
+    # Block sizes
+    BLOCK_K: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    # S accumulator in registers: (K, V) tiled as (BLOCK_K, BLOCK_V)
+    # We need to iterate over K tiles and V tiles
+    num_k_tiles = K // BLOCK_K
+    num_v_tiles = V // BLOCK_V
+
+    # Initialize S to zero in SRAM (we'll accumulate across tiles)
+    # For each (k_tile, v_tile) pair, maintain an accumulator
+    # Since BLOCK_K=128 and BLOCK_V=64, we have 2 V tiles
+
+    # Pointers for this (b, h) pair
+    q_base = q_ptr + pid * stride_q_bh
+    k_base = k_ptr + pid * stride_k_bh
+    g_base = g_ptr + pid * stride_g_bh
+    w_base = w_ptr + pid * stride_w_bh
+    u_base = u_ptr + pid * stride_u_bh
+    o_base = o_ptr + pid * stride_o_bh
+
+    # mask for strict upper triangular (BT x BT)
+    # We'll compute this on the fly
+
+    # Loop over chunks
+    for ic in range(NT):
+        # Pointers for this chunk
+        q_ic = q_base + ic * stride_q_n
+        k_ic = k_base + ic * stride_k_n
+        g_ic = g_base + ic * stride_g_n
+        w_ic = w_base + ic * stride_w_n
+        u_ic = u_base + ic * stride_u_n
+        o_ic = o_base + ic * stride_o_n
+
+        # For each output position j in BT:
+        #   v_i[j] = u_i[j] - sum_c w_i[j, c] * S[c, :]
+        #   o_i[j] = sum_c q_exp[j, c] * S[c, :] + sum_c Aqk[j, c] * v_i[c]
+        #   S update: S[c, :] += k_decayed[c] * v_i[c]
+
+        # We'll iterate over output positions j
+        for j in range(BT):
+            # Load q, k, g, w, u at position j
+            # q_pos[j] = q[j] * exp(g[j])
+            # k_neg[j] = k[j] * exp(-g[j])
+            # q_exp[j] = q[j] * exp(g[j])
+            # w_j = w[j]
+            # u_j = u[j]
+
+            # Load q[j], k[j], g[j]
+            q_offs = j * stride_q_t
+            k_offs = j * stride_k_t
+            g_offs = j * stride_g_t
+            w_offs = j * stride_w_t
+            u_offs = j * stride_u_t
+
+            # For each V tile
+            for v_tile in range(num_v_tiles):
+                v_start = v_tile * BLOCK_V
+                v_offs = v_start + tl.arange(0, BLOCK_V)
+
+                # Load u[j, v_start:v_start+BLOCK_V]
+                u_ptrs = u_ic + u_offs + v_offs * stride_u_v
+                u_val = tl.load(u_ptrs)
+
+                # Initialize output accumulator
+                out_acc = u_val  # start with u
+
+                # Subtract w_i @ S: sum_c w[j, c] * S[c, v_start:]
+                for c in range(BT):
+                    w_val = tl.load(w_ic + w_offs + c * stride_w_k)  # scalar
+                    # Load S[c, v_start:v_start+BLOCK_V] -- but S is in global mem initially
+                    # Actually S needs to be maintained in SRAM. This is the challenge.
+
+                    pass
+
+                # This approach won't work well because S needs to be in SRAM
+                # Let me restructure
+                pass
+            pass
+
+    # This approach is too complex. Let me use a simpler tiled matmul approach.
+    pass
 
 
 def _kda_fwd(
@@ -37,7 +142,7 @@
     g = g.float()
     beta = beta.float()
 
-    # Reshape to chunk form: (B*H, NT, BT, D) for batched matmul
+    # Reshape to chunk form: (B*H, NT, BT, D)
     q = q.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B * H, NT, BT, K)
     k = k.reshape(B, T, H, K).permute(0, 2, 1, 3).reshape(B * H, NT, BT, K)
     v = v.reshape(B, T, H, V).permute(0, 2, 1, 3).reshape(B * H, NT, BT, V)
@@ -56,7 +161,7 @@
     A = A * beta.unsqueeze(-1)
     A_neg = -A.masked_fill(mask_diag_upper, 0)
 
-    # Newton iteration replaced with solve_triangular
+    # Newton iteration -> solve_triangular
     I_mat = torch.eye(BT, dtype=torch.float32, device=q.device).expand_as(A_neg).contiguous()
     A_inv = torch.linalg.solve_triangular(I_mat - A_neg, I_mat, upper=False, unitriangular=False)
     A = A_inv * beta.unsqueeze(-2)
@@ -80,27 +185,22 @@
         u_i = u[:, ic]
         w_i = w[:, ic]
 
-        # Aqk[c, j] = sum_d q[c,d] * exp(g[c,d] - g[j,d]) * k[j,d]
         q_pos_i = q_i * g_i.exp()
         k_neg_i = k_i * (-g_i).exp()
         Aqk = torch.matmul(q_pos_i, k_neg_i.transpose(-1, -2))
         Aqk = Aqk.masked_fill(mask_strict_upper, 0)
 
-        # v_i = u_i - w_i @ S
         v_i = u_i - torch.matmul(w_i, S)
 
-        # o_i = (q_i * exp(g_i)) @ S + Aqk @ v_i
         q_exp = q_i * g_i.exp()
         o_i = torch.matmul(q_exp, S) + torch.matmul(Aqk, v_i)
         o[:, ic] = o_i
 
-        # Update state
         g_last = g_i[:, -1, :]
         S = S * g_last.unsqueeze(-1).exp()
         k_decayed = (g_last.unsqueeze(-2) - g_i).exp() * k_i
         S = S

... (177 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,121 +1,37 @@
-"""KDA forward (chunk form) -- Triton kernel implementation.
+"""KDA forward (chunk form) -- optimized PyTorch with fused operations.
 
-Triton kernel for the inter-chunk recurrence loop + PyTorch for intra-chunk.
+Uses concatenated matmuls and torch.compile to reduce kernel launch overhead.
 No FLA imports.
 """
 from __future__ import annotations
 
 import torch
 import torch.nn as nn
-import triton
-import triton.language as tl
 
 OP_TYPE = "linear_attention"
 SUPPORTED_PRECISIONS = ["bf16"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-@triton.jit
-def _inter_chunk_kernel(
-    # Input pointers
-    q_ptr, k_ptr, g_ptr, w_ptr, u_ptr, o_ptr,
-    # Strides for q, k, g: (BH, NT, BT, K)
-    stride_q_bh, stride_q_n, stride_q_t
-    , stride_q_k,
-    stride_k_bh, stride_k_n, stride_k_t, stride_k_k,
-    stride_g_bh, stride_g_n, stride_g_t, stride_g_k,
-    # Strides for w: (BH, NT, BT, K)
-    stride_w_bh, stride_w_n, stride_w_t, stride_w_k,
-    # Strides for u: (BH, NT, BT, V)
-    stride_u_bh, stride_u_n, stride_u_t, stride_u_v,
-    # Strides for o: (BH, NT, BT, V)
-    stride_o_bh, stride_o_n, stride_o_t, stride_o_v,
-    # Dimensions
-    NT, BT, K, V,
-    # Block sizes
-    BLOCK_K: tl.constexpr,
-    BLOCK_V: tl.constexpr,
-):
-    pid = tl.program_id(0)
+@torch.compile(mode="reduce-overhead")
+def _fused_inter_chunk(q_i, k_i, g_i, u_i, w_i, S, mask_strict_upper):
+    """Fused inter-chunk body. Compiled to reduce kernel launches."""
+    q_pos_i = q_i * g_i.exp()
+    k_neg_i = k_i * (-g_i).exp()
+    Aqk = torch.matmul(q_pos_i, k_neg_i.transpose(-1, -2))
+    Aqk = Aqk.masked_fill(mask_strict_upper, 0)
 
-    # S accumulator in registers: (K, V) tiled as (BLOCK_K, BLOCK_V)
-    # We need to iterate over K tiles and V tiles
-    num_k_tiles = K // BLOCK_K
-    num_v_tiles = V // BLOCK_V
+    v_i = u_i - torch.matmul(w_i, S)
 
-    # Initialize S to zero in SRAM (we'll accumulate across tiles)
-    # For each (k_tile, v_tile) pair, maintain an accumulator
-    # Since BLOCK_K=128 and BLOCK_V=64, we have 2 V tiles
+    q_exp = q_i * g_i.exp()
+    o_i = torch.matmul(q_exp, S) + torch.matmul(Aqk, v_i)
 
-    # Pointers for this (b, h) pair
-    q_base = q_ptr + pid * stride_q_bh
-    k_base = k_ptr + pid * stride_k_bh
-    g_base = g_ptr + pid * stride_g_bh
-    w_base = w_ptr + pid * stride_w_bh
-    u_base = u_ptr + pid * stride_u_bh
-    o_base = o_ptr + pid * stride_o_bh
+    g_last = g_i[:, -1, :]
+    S = S * g_last.unsqueeze(-1).exp()
+    k_decayed = (g_last.unsqueeze(-2) - g_i).exp() * k_i
+    S = S + torch.matmul(k_decayed.transpose(-1, -2), v_i)
 
-    # mask for strict upper triangular (BT x BT)
-    # We'll compute this on the fly
-
-    # Loop over chunks
-    for ic in range(NT):
-        # Pointers for this chunk
-        q_ic = q_base + ic * stride_q_n
-        k_ic = k_base + ic * stride_k_n
-        g_ic = g_base + ic * stride_g_n
-        w_ic = w_base + ic * stride_w_n
-        u_ic = u_base + ic * stride_u_n
-        o_ic = o_base + ic * stride_o_n
-
-        # For each output position j in BT:
-        #   v_i[j] = u_i[j] - sum_c w_i[j, c] * S[c, :]
-        #   o_i[j] = sum_c q_exp[j, c] * S[c, :] + sum_c Aqk[j, c] * v_i[c]
-        #   S update: S[c, :] += k_decayed[c] * v_i[c]
-
-        # We'll iterate over output positions j
-        for j in range(BT):
-            # Load q, k, g, w, u at position j
-            # q_pos[j] = q[j] * exp(g[j])
-            # k_neg[j] = k[j] * exp(-g[j])
-            # q_exp[j] = q[j] * exp(g[j])
-            # w_j = w[j]
-            # u_j = u[j]
-
-            # Load q[j], k[j], g[j]
-            q_offs = j * stride_q_t
-            k_offs = j * stride_k_t
-            g_offs = j * stride_g_t
-            w_offs = j * stride_w_t
-            u_offs = j * stride_u_t
-
-            # For each V tile
-            for v_tile in range(num_v_tiles):
-                v_start = v_tile * BLOCK_V
-                v_offs = v_start + tl.arange(0, BLOCK_V)
-
-                # Load u[j, v_start:v_start+BLOCK_V]
-                u_ptrs = u_ic + u_offs + v_offs * stride_u_v
-                u_val = tl.load(u_ptrs)
-
-                # Initialize output accumulator
-                out_acc = u_val  # start with u
-
-                # Subtract w_i @ S: sum_c w[j, c] * S[c, v_start:]
-                for c in range(BT):
-                    w_val = tl.load(w_ic + w_offs + c * stride_w_k)  # scalar
-                    # Load S[c, v_start:v_start+BLOCK_V] -- but S is in global mem initially
-                    # Actually S needs to be maintained in SRAM. This is the challenge.
-
-                    pass
-
-                # This approach won't work well because S needs to be in SRAM
-                # Let me restructure
-                pass
-            pass
-
-    # This approach is too complex. Let me use a simpler tiled matmul approach.
-    pass
+    return o_i, S
 
 
 def _kda_fwd(
@@ -179,28 +95,13 @@
     mask_strict_upper = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
 
     for ic in range(NT):
-        q_i = q[:, ic]
-        k_i = k[:, ic]
-        g_i = g[:, ic]
-        u_i = u[:, ic]
-        w_i = w[:, ic]
-
-        q_pos_i = q_i * g_i.exp()
-        k_neg_i = k_i * (-g_i).exp()
-        Aqk = torch.matmul(q_pos_i, k_neg_i.transpose(-1, -2))
-        Aqk = Aqk.masked_fill(mask_strict_upper, 0)
-
-        v_i = u_i - torch.matmul(w_i, S)
-
-        q_exp = q_i * g_i.exp()
-        o_i = torch.matmul(q_exp, S) + torch.matmul(Aqk, v_i)
+        o_i, S = _fused_inter_chunk(
+            q[:, ic], k[:, ic], g[:, ic], u[:, ic], w[:, ic],
+            S, mask_strict_upper,
+        )
         o[:, ic] = o_i
 
-        g_last = g_i[:, -1, :]
-        S = S * g_last.unsqueeze(-1).exp()
-        k_decayed 

... (234 more chars)