--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -0,0 +1,336 @@
+"""Custom Triton paged attention decode kernel.
+
+Layout:
+  query:      (batch, num_heads, head_dim)             bf16
+  kv_cache:   (num_blocks, page_size, num_kv_heads, head_dim * 2)  bf16
+              last dim packs [K | V] so a single gather pulls both
+  block_table: (batch, max_blocks)                     int32
+  seq_lens:   (batch,)                                 int32
+
+Output: bf16 (batch, num_heads, head_dim)
+
+Strategy:
+  One program per (batch, kv_head). Each program:
+    1. Loads block_table[b] pointers for this batch.
+    2. Iterates over pages, loading K and V blocks from the page cache.
+    3. For each head in the GQA group, computes QK^T dot product.
+    4. Accumulates online softmax numerator and denominator.
+  We vectorize the head_dim dimension.
+"""
+import math
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _paged_decode_kernel(
+    # Output: (batch, num_heads, head_dim)
+    Out,
+    # Inputs
+    Q,
+    KVC,
+    BlockTable,
+    SeqLens,
+    # Strides
+    stride_qb, stride_qh, stride_qd,
+    stride_kc_b, stride_kc_s, stride_kc_h, stride_kc_d,
+    stride_bt_b, stride_bt_s,
+    stride_ob, stride_oh, stride_od,
+    # Shapes
+    BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, MAX_SEQ_LEN, PAGE_SIZE,
+    HEAD_DIM_PADDED: tl.constexpr,
+    NUM_PAGES: tl.constexpr,
+    PAGE_SIZE_CONST: tl.constexpr,
+    BLOCK_SIZE_D: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_hkv = tl.program_id(1)
+
+    # Compute seq_len for this batch element
+    seq_len = tl.load(SeqLens + pid_b)
+    num_pages_needed = (seq_len + PAGE_SIZE_CONST - 1) // PAGE_SIZE_CONST
+
+    # Load block table for this batch
+    # block_table[pid_b, :] -> page indices
+    # Load all page indices
+    offs_bt = tl.arange(0, NUM_PAGES)
+    page_indices = tl.load(
+        BlockTable + pid_b * stride_bt_b + offs_bt,
+        mask=offs_bt < num_pages_needed,
+        other=0,
+    )
+
+    # Compute Q pointer for the group of heads
+    # Q[pid_b, pid_hkv * group_size : (pid_hkv+1) * group_size, :]
+    q_base = Q + pid_b * stride_qb + pid_hkv * GROUP_SIZE * stride_qh
+
+    # Output pointer base
+    out_base = Out + pid_b * stride_ob + pid_hkv * GROUP_SIZE * stride_oh
+
+    # Initialize accumulators for each head in the group
+    # We'll process HEAD_DIM_PADDED elements at a time
+    num_d_blocks = HEAD_DIM // BLOCK_SIZE_D
+
+    # Running max and denominator for online softmax (per head in group)
+    # acc_num[group][d]: accumulated numerator
+    # acc_denom[group]: accumulated denominator
+    acc_num = tl.zeros((GROUP_SIZE, BLOCK_SIZE_D), dtype=tl.float32)
+    acc_denom = tl.zeros((GROUP_SIZE,), dtype=tl.float32)
+    acc_max = tl.full((GROUP_SIZE,), -1e30, dtype=tl.float32)
+
+    # Load all Q vectors for the group
+    # q_vecs[group][d_block]
+    offs_d_block = tl.arange(0, BLOCK_SIZE_D)
+    q_ptrs_base = tl.make_block_ptr(
+        base=Q + pid_b * stride_qb,
+        shape=(NUM_HEADS, HEAD_DIM),
+        strides=(stride_qh, stride_qd),
+        offsets=(0, 0),
+        block_shape=(GROUP_SIZE, BLOCK_SIZE_D),
+        order=(1, 0),
+    )
+
+    # For each head group member, load Q
+    q_group_ptrs = tl.make_block_ptr(
+        base=Q + pid_b * stride_qb,
+        shape=(NUM_HEADS, HEAD_DIM),
+        strides=(stride_qh, stride_qd),
+        offsets=(pid_hkv * GROUP_SIZE, 0),
+        block_shape=(GROUP_SIZE, BLOCK_SIZE_D),
+        order=(1, 0),
+    )
+    q_vec = tl.load(q_group_ptrs, boundary_check=(0, 1), padding_option="zero")
+    q_vec = q_vec.to(tl.float32)
+
+    # Scale factor
+    scale = 1.0 / tl.sqrt(HEAD_DIM)
+
+    # Iterate over pages
+    for page_idx in range(NUM_PAGES):
+        if page_idx >= num_pages_needed:
+            break
+
+        page_id = tl.load(BlockTable + pid_b * stride_bt_b + page_idx)
+
+        # Compute K and V pointers for this page
+        # kv_cache[page_id, :, pid_hkv, :] -> (page_size, head_dim * 2)
+        page_base = KVC + page_id * stride_kc_b + pid_hkv * stride_kc_h
+
+        # Load K block: (page_size, head_dim)
+        # Then load V block: (page_size, head_dim) starting at head_dim offset
+
+        # Use block pointers for K
+        k_block_ptrs = tl.make_block_ptr(
+            base=page_base,
+            shape=(PAGE_SIZE_CONST, HEAD_DIM),
+            strides=(stride_kc_s, stride_kc_d),
+            offsets=(0, 0),
+            block_shape=(PAGE_SIZE_CONST, BLOCK_SIZE_D),
+            order=(1, 0),
+        )
+        k_block = tl.load(k_block_ptrs, boundary_check=(0, 1), padding_option="zero")
+        k_block = k_block.to(tl.float32)
+
+        # V block
+        v_block_ptrs = tl.make_block_ptr(
+            base=page_base + HEAD_DIM,
+            shape=(PAGE_SIZE_CONST, HEAD_DIM),
+            strides=(stride_kc_s, stride_kc_d),
+            offsets=(0, 0),
+            block_shape=(PAGE_SIZE_CONST, BLOCK_SIZE_D),
+            order=(1, 0),
+        )
+        v_block = tl.load(v_block_ptrs, boundary_check=(0, 1), padding_option="zero")
+        v_block = v_block.to(tl.float32)
+
+        # For each token position in the page:
+        for s in range(PAGE_SIZE_CONST):
+            token_global = page_idx * PAGE_SIZE_CONST + s
+            if token_global >= seq_len:
+                break
+
+            # k_vec: (1, BLOCK_SIZE_D)
+            k_vec = tl.broadcast_to(k_block[s:s+1, :], (GROUP_SIZE, BLOCK_SIZE_D))
+
+            # QK^T for the group: (GROUP_SIZE,) dot product
+            qk = tl.sum(q_vec * k_vec, axis=1) * scale
+
+            # Online softmax update
+            old_max = acc_max
+            new_max = tl.maximum(old_max, qk)
+            exp_diff = tl.exp(qk - new_max)
+            rescale = tl.exp(old_max - new_max)
+
+            # V contribution
+            v_vec = tl.broadc

... (5541 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -201,10 +201,9 @@
     stride_ob, stride_oh, stride_od,
     # Shapes
     BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, PAGE_SIZE,
-    HEAD_DIM_PACKED: tl.constexpr,   # next power of 2
-    BLOCK_D: tl.constexpr,
+    HEAD_DIM_PADDED: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
-):
+    IS_VAR_LEN: tl.constexpr,
     """Version 2: one program per (batch, head), not per kv_head group.
     Simpler indexing, more parallelism.
     """

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -218,7 +218,7 @@
     num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
 
     # Offsets in d
-    d_offs = tl.arange(0, BLOCK_D)
+    d_offs = tl.arange(0, HEAD_DIM_PADDED)
     d_mask = d_offs < HEAD_DIM
 
     # Load Q

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -226,7 +226,8 @@
     q = tl.load(Q + q_offs, mask=d_mask, other=0.0).to(tl.float32)
 
     # Scale
-    scale = 1.0 / tl.sqrt(HEAD_DIM.to(tl.float32) if hasattr(HEAD_DIM, '__float__') else float(HEAD_DIM))
+    SCALE: tl.constexpr = 1.0 / (HEAD_DIM ** 0.5)
+    scale: tl.constexpr = SCALE
 
     # Accumulators
     acc_num = tl.zeros([BLOCK_D], dtype=tl.float32)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,22 +1,4 @@
-"""Custom Triton paged attention decode kernel.
-
-Layout:
-  query:      (batch, num_heads, head_dim)             bf16
-  kv_cache:   (num_blocks, page_size, num_kv_heads, head_dim * 2)  bf16
-              last dim packs [K | V] so a single gather pulls both
-  block_table: (batch, max_blocks)                     int32
-  seq_lens:   (batch,)                                 int32
-
-Output: bf16 (batch, num_heads, head_dim)
-
-Strategy:
-  One program per (batch, kv_head). Each program:
-    1. Loads block_table[b] pointers for this batch.
-    2. Iterates over pages, loading K and V blocks from the page cache.
-    3. For each head in the GQA group, computes QK^T dot product.
-    4. Accumulates online softmax numerator and denominator.
-  We vectorize the head_dim dimension.
-"""
+"""Custom Triton paged attention decode kernel."""
 import math
 import torch
 import torch.nn as nn
@@ -26,245 +8,61 @@
 
 @triton.jit
 def _paged_decode_kernel(
-    # Output: (batch, num_heads, head_dim)
     Out,
-    # Inputs
     Q,
-    KVC,
+    K,
+    V,
     BlockTable,
     SeqLens,
     # Strides
     stride_qb, stride_qh, stride_qd,
-    stride_kc_b, stride_kc_s, stride_kc_h, stride_kc_d,
-    stride_bt_b, stride_bt_s,
-    stride_ob, stride_oh, stride_od,
-    # Shapes
-    BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, MAX_SEQ_LEN, PAGE_SIZE,
-    HEAD_DIM_PADDED: tl.constexpr,
-    NUM_PAGES: tl.constexpr,
-    PAGE_SIZE_CONST: tl.constexpr,
-    BLOCK_SIZE_D: tl.constexpr,
-    GROUP_SIZE: tl.constexpr,
-):
-    pid_b = tl.program_id(0)
-    pid_hkv = tl.program_id(1)
-
-    # Compute seq_len for this batch element
-    seq_len = tl.load(SeqLens + pid_b)
-    num_pages_needed = (seq_len + PAGE_SIZE_CONST - 1) // PAGE_SIZE_CONST
-
-    # Load block table for this batch
-    # block_table[pid_b, :] -> page indices
-    # Load all page indices
-    offs_bt = tl.arange(0, NUM_PAGES)
-    page_indices = tl.load(
-        BlockTable + pid_b * stride_bt_b + offs_bt,
-        mask=offs_bt < num_pages_needed,
-        other=0,
-    )
-
-    # Compute Q pointer for the group of heads
-    # Q[pid_b, pid_hkv * group_size : (pid_hkv+1) * group_size, :]
-    q_base = Q + pid_b * stride_qb + pid_hkv * GROUP_SIZE * stride_qh
-
-    # Output pointer base
-    out_base = Out + pid_b * stride_ob + pid_hkv * GROUP_SIZE * stride_oh
-
-    # Initialize accumulators for each head in the group
-    # We'll process HEAD_DIM_PADDED elements at a time
-    num_d_blocks = HEAD_DIM // BLOCK_SIZE_D
-
-    # Running max and denominator for online softmax (per head in group)
-    # acc_num[group][d]: accumulated numerator
-    # acc_denom[group]: accumulated denominator
-    acc_num = tl.zeros((GROUP_SIZE, BLOCK_SIZE_D), dtype=tl.float32)
-    acc_denom = tl.zeros((GROUP_SIZE,), dtype=tl.float32)
-    acc_max = tl.full((GROUP_SIZE,), -1e30, dtype=tl.float32)
-
-    # Load all Q vectors for the group
-    # q_vecs[group][d_block]
-    offs_d_block = tl.arange(0, BLOCK_SIZE_D)
-    q_ptrs_base = tl.make_block_ptr(
-        base=Q + pid_b * stride_qb,
-        shape=(NUM_HEADS, HEAD_DIM),
-        strides=(stride_qh, stride_qd),
-        offsets=(0, 0),
-        block_shape=(GROUP_SIZE, BLOCK_SIZE_D),
-        order=(1, 0),
-    )
-
-    # For each head group member, load Q
-    q_group_ptrs = tl.make_block_ptr(
-        base=Q + pid_b * stride_qb,
-        shape=(NUM_HEADS, HEAD_DIM),
-        strides=(stride_qh, stride_qd),
-        offsets=(pid_hkv * GROUP_SIZE, 0),
-        block_shape=(GROUP_SIZE, BLOCK_SIZE_D),
-        order=(1, 0),
-    )
-    q_vec = tl.load(q_group_ptrs, boundary_check=(0, 1), padding_option="zero")
-    q_vec = q_vec.to(tl.float32)
-
-    # Scale factor
-    scale = 1.0 / tl.sqrt(HEAD_DIM)
-
-    # Iterate over pages
-    for page_idx in range(NUM_PAGES):
-        if page_idx >= num_pages_needed:
-            break
-
-        page_id = tl.load(BlockTable + pid_b * stride_bt_b + page_idx)
-
-        # Compute K and V pointers for this page
-        # kv_cache[page_id, :, pid_hkv, :] -> (page_size, head_dim * 2)
-        page_base = KVC + page_id * stride_kc_b + pid_hkv * stride_kc_h
-
-        # Load K block: (page_size, head_dim)
-        # Then load V block: (page_size, head_dim) starting at head_dim offset
-
-        # Use block pointers for K
-        k_block_ptrs = tl.make_block_ptr(
-            base=page_base,
-            shape=(PAGE_SIZE_CONST, HEAD_DIM),
-            strides=(stride_kc_s, stride_kc_d),
-            offsets=(0, 0),
-            block_shape=(PAGE_SIZE_CONST, BLOCK_SIZE_D),
-            order=(1, 0),
-        )
-        k_block = tl.load(k_block_ptrs, boundary_check=(0, 1), padding_option="zero")
-        k_block = k_block.to(tl.float32)
-
-        # V block
-        v_block_ptrs = tl.make_block_ptr(
-            base=page_base + HEAD_DIM,
-            shape=(PAGE_SIZE_CONST, HEAD_DIM),
-            strides=(stride_kc_s, stride_kc_d),
-            offsets=(0, 0),
-            block_shape=(PAGE_SIZE_CONST, BLOCK_SIZE_D),
-            order=(1, 0),
-        )
-        v_block = tl.load(v_block_ptrs, boundary_check=(0, 1), padding_option="zero")
-        v_block = v_block.to(tl.float32)
-
-        # For each token position in the page:
-        for s in range(PAGE_SIZE_CONST):
-            token_global = page_idx * PAGE_SIZE_CONST + s
-            if token_global >= seq_len:
-                break
-
-            # k_vec: (1, BLOCK_SIZE_D)
-            k_vec = tl.broadcast_to(k_block[s:s+1, :], (GROUP_SIZE, BLOCK_SIZE_D))
-
-            # QK^T for the group: (GROUP_SIZE,) dot product
-            qk = tl.sum(q_vec * k_vec, axis=1) * scale
-
-            # Online softmax update
-            old_max = acc_max
-            new_max = tl.maximum(old_max, qk)
-            exp_diff = tl.exp(qk - new_max)
-

... (6485 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -10,15 +10,13 @@
 def _paged_decode_kernel(
     Out,
     Q,
-    K,
-    V,
+    KV,
     BlockTable,
     SeqLens,
     # Strides
     stride_qb, stride_qh, stride_qd,
-    stride_kb, stride_ks, stride_kh, stride_kd,
-    stride_vb, stride_vs, stride_vh, stride_vd,
-    stride_bt_b, stride_bt_s,
+    stride_kvb, stride_kvs, stride_kvh, stride_kvd,
+    stride_bt_b,
     stride_ob, stride_oh, stride_od,
     # Shapes
     BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, PAGE_SIZE,
@@ -50,17 +48,19 @@
     for page_i in tl.range(num_pages):
         page_id = tl.load(BlockTable + bt_base + page_i)
 
+        page_base = page_id * stride_kvb + pid_hkv * stride_kvh
+
         for s in tl.range(PAGE_SIZE):
             global_tok = page_i * PAGE_SIZE + s
             if global_tok >= seq_len:
                 break
 
-            k_ptr = K + page_id * stride_kb + s * stride_ks + pid_hkv * stride_kh + d_offs * stride_kd
+            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
             k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-            qk = tl.sum(q * k)
+            qk = tl.sum(q * k) * scale
 
-            v_ptr = V + page_id * stride_vb + s * stride_vs + pid_hkv * stride_vh + d_offs * stride_vd
+            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
             v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
             old_max = acc_max
@@ -81,7 +81,7 @@
 class Model(nn.Module):
     def __init__(self, batch, num_heads, num_kv_heads, head_dim, seq_len, page_size):
         super().__init__()
-        assert num_heads % num_kv_heads == 0, "num_heads must be a multiple of num_kv_heads (GQA)"
+        assert num_heads % num_kv_heads == 0
         self.batch = batch
         self.num_heads = num_heads
         self.num_kv_heads = num_kv_heads
@@ -97,28 +97,21 @@
         Hkv = self.num_kv_heads
         D = self.head_dim
         P = self.page_size
-        G = self.group_size
-        D_pad = triton.next_power_of_2(D)
-
-        # kv_cache: (num_blocks, page_size, num_kv_heads, head_dim * 2)
-        # We need separate K and V views
-        kv_cache_4d = kv_cache.view(B * -1 or -1, P, Hkv, D * 2)
-        K = kv_cache[..., :D].contiguous()
-        V = kv_cache[..., D:].contiguous()
 
         out = torch.empty(B, H, D, dtype=query.dtype, device=query.device)
 
+        assert kv_cache.stride(3) == 1, "kv_cache must have innermost stride 1"
+
         grid = (B * H,)
         _paged_decode_kernel[grid](
-            out, query, K, V, block_table, seq_lens,
+            out, query, kv_cache, block_table, seq_lens,
             query.stride(0), query.stride(1), query.stride(2),
-            K.stride(0), K.stride(1), K.stride(2), K.stride(3),
-            V.stride(0), V.stride(1), V.stride(2), V.stride(3),
-            block_table.stride(0), block_table.stride(1),
+            kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3),
+            block_table.stride(0),
             out.stride(0), out.stride(1), out.stride(2),
             B, H, Hkv, D, P,
-            D_pad,
-            G,
+            triton.next_power_of_2(D),
+            self.group_size,
         )
 
         return out

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -37,7 +37,7 @@
     q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
     q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-    scale = 1.0 / tl.sqrt(float(HEAD_DIM))
+    scale = 1.0 / (HEAD_DIM ** 0.5)
 
     acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
     acc_denom = 0.0

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -4,6 +4,10 @@
 import torch.nn as nn
 import triton
 import triton.language as tl
+
+
+def _make_scale(head_dim: int) -> float:
+    return 1.0 / math.sqrt(head_dim)
 
 
 @triton.jit
@@ -19,9 +23,13 @@
     stride_bt_b,
     stride_ob, stride_oh, stride_od,
     # Shapes
-    BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, PAGE_SIZE,
+    BATCH, NUM_HEADS, NUM_KV_HEADS,
+    PAGE_SIZE,
+    # constexpr
+    HEAD_DIM: tl.constexpr,
     HEAD_DIM_PACKED: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
+    SCALE: tl.constexpr,
 ):
     pid_b = tl.program_id(0) // NUM_HEADS
     pid_h = tl.program_id(0) % NUM_HEADS
@@ -36,8 +44,6 @@
 
     q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
     q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
-
-    scale = 1.0 / (HEAD_DIM ** 0.5)
 
     acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
     acc_denom = 0.0
@@ -58,7 +64,7 @@
             k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
             k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-            qk = tl.sum(q * k) * scale
+            qk = tl.sum(q * k) * SCALE
 
             v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
             v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
@@ -97,6 +103,8 @@
         Hkv = self.num_kv_heads
         D = self.head_dim
         P = self.page_size
+        D_pad = triton.next_power_of_2(D)
+        scale = _make_scale(D)
 
         out = torch.empty(B, H, D, dtype=query.dtype, device=query.device)
 
@@ -109,9 +117,10 @@
             kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3),
             block_table.stride(0),
             out.stride(0), out.stride(1), out.stride(2),
-            B, H, Hkv, D, P,
-            triton.next_power_of_2(D),
+            B, H, Hkv, P,
+            D, D_pad,
             self.group_size,
+            scale,
         )
 
         return out

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -30,6 +30,7 @@
     HEAD_DIM_PACKED: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     SCALE: tl.constexpr,
+    MAX_PAGES: tl.constexpr,
 ):
     pid_b = tl.program_id(0) // NUM_HEADS
     pid_h = tl.program_id(0) % NUM_HEADS
@@ -51,32 +52,33 @@
 
     bt_base = pid_b * stride_bt_b
 
-    for page_i in tl.range(num_pages):
+    for page_i in tl.range(MAX_PAGES):
+        if page_i >= num_pages:
+            continue
+
         page_id = tl.load(BlockTable + bt_base + page_i)
 
         page_base = page_id * stride_kvb + pid_hkv * stride_kvh
 
         for s in tl.range(PAGE_SIZE):
             global_tok = page_i * PAGE_SIZE + s
-            if global_tok >= seq_len:
-                break
+            if global_tok < seq_len:
+                k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+                k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
-            k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
+                qk = tl.sum(q * k) * SCALE
 
-            qk = tl.sum(q * k) * SCALE
+                v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+                v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
-            v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
+                old_max = acc_max
+                new_max = tl.maximum(old_max, qk)
+                exp_val = tl.exp(qk - new_max)
+                rescale_old = tl.exp(old_max - new_max)
 
-            old_max = acc_max
-            new_max = tl.maximum(old_max, qk)
-            exp_val = tl.exp(qk - new_max)
-            rescale_old = tl.exp(old_max - new_max)
-
-            acc_num = acc_num * rescale_old + v * exp_val
-            acc_denom = acc_denom * rescale_old + exp_val
-            acc_max = new_max
+                acc_num = acc_num * rescale_old + v * exp_val
+                acc_denom = acc_denom * rescale_old + exp_val
+                acc_max = new_max
 
     out = acc_num / acc_denom
 
@@ -95,6 +97,7 @@
         self.seq_len = seq_len
         self.page_size = page_size
         self.group_size = num_heads // num_kv_heads
+        self.max_pages = (seq_len + page_size - 1) // page_size
         self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
 
     def forward(self, query, kv_cache, block_table, seq_lens):
@@ -105,6 +108,7 @@
         P = self.page_size
         D_pad = triton.next_power_of_2(D)
         scale = _make_scale(D)
+        max_pages = self.max_pages
 
         out = torch.empty(B, H, D, dtype=query.dtype, device=query.device)
 
@@ -121,6 +125,7 @@
             D, D_pad,
             self.group_size,
             scale,
+            max_pages,
         )
 
         return out

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution_triton_v3.py
@@ -0,0 +1,83 @@
+"""Custom Triton paged attention decode kernel."""
+import math
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+def _make_scale(head_dim: int) -> float:
+    return 1.0 / math.sqrt(head_dim)
+
+
+@triton.jit
+def _paged_decode_kernel(
+    Out,
+    Q,
+    KV,
+    BlockTable,
+    SeqLens,
+    # Strides
+    stride_qb, stride_qh, stride_qd,
+    stride_kvb, stride_kvs, stride_kvh, stride_kvd,
+    stride_bt_b,
+    stride_ob, stride_oh, stride_od,
+    # constexpr
+    HEAD_DIM: tl.constexpr,
+    HEAD_DIM_PACKED: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    SCALE: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    pid_hkv = pid_h // num_kv_heads
+
+    seq_len = tl.load(SeqLens + pid_b)
+
+    d_offs = tl.arange(0, HEAD_DIM_PACKED)
+    d_mask = d_offs < HEAD_DIM
+
+    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
+    q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
+
+    acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
+    acc_denom = 0.0
+    acc_max = -1e30
+
+    bt_base = pid_b * stride_bt_b
+
+    # Compute total tokens for iteration
+    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
+    total_tokens = num_pages * PAGE_SIZE
+
+    for t in tl.range(total_tokens):
+        valid = t < seq_len
+        if valid:
+            page_i = t // PAGE_SIZE
+            s = t % PAGE_SIZE
+
+            page_id = tl.load(BlockTable + bt_base + page_i)
+            page_base = page_id * stride_kvb + pid_hkv * stride_kvh
+
+            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+            k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
+
+            qk = tl.sum(q * k) * SCALE
+
+            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+            v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
+
+            old_max = acc_max
+            new_max = tl.maximum(old_max, qk)
+            exp_val = tl.exp(qk - new_max)
+            rescale_old = tl.exp(old_max - new_max)
+
+            acc_num = acc_num * rescale_old + v * exp_val
+            acc_denom = acc_denom * rescale_old + exp_val
+            acc_max = new_max
+
+    out = acc_num / acc_denom
+
+    out_ptr = Out + pid_b * stride_ob + pid_h * stride_oh + d_offs * stride_od
+    tl.store(out_ptr, out.to(Out.type.element_ty), mask=d_mask)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,4 +1,13 @@
-"""Custom Triton paged attention decode kernel."""
+"""Custom Triton paged attention decode kernel.
+
+Strategy:
+- One program per (batch, query_head)
+- Each program loads page pointers from block_table for its batch
+- Loads one page of K and V at a time into registers/SRAM
+- Computes qk dot product, online softmax accumulation
+- Uses tl.dot for QK^T and weighted sum
+- GQA: pid_hkv = pid_h // group_size
+"""
 import math
 import torch
 import torch.nn as nn
@@ -20,29 +29,28 @@
     # Strides
     stride_qb, stride_qh, stride_qd,
     stride_kvb, stride_kvs, stride_kvh, stride_kvd,
-    stride_bt_b,
+    stride_bt_b, stride_bt_s,
     stride_ob, stride_oh, stride_od,
-    # Shapes
-    BATCH, NUM_HEADS, NUM_KV_HEADS,
-    PAGE_SIZE,
-    # constexpr
+    # constexpr params
+    NUM_KV_HEADS: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
     HEAD_DIM: tl.constexpr,
     HEAD_DIM_PACKED: tl.constexpr,
-    GROUP_SIZE: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
     SCALE: tl.constexpr,
-    MAX_PAGES: tl.constexpr,
+    MAX_SEQ_LEN: tl.constexpr,
 ):
-    pid_b = tl.program_id(0) // NUM_HEADS
-    pid_h = tl.program_id(0) % NUM_HEADS
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
 
     pid_hkv = pid_h // GROUP_SIZE
 
     seq_len = tl.load(SeqLens + pid_b)
-    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
 
     d_offs = tl.arange(0, HEAD_DIM_PACKED)
     d_mask = d_offs < HEAD_DIM
 
+    # Load Q: (HEAD_DIM,)
     q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
     q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
@@ -52,33 +60,37 @@
 
     bt_base = pid_b * stride_bt_b
 
-    for page_i in tl.range(MAX_PAGES):
-        if page_i >= num_pages:
-            continue
+    max_tokens = MAX_SEQ_LEN
 
-        page_id = tl.load(BlockTable + bt_base + page_i)
+    for t in tl.range(max_tokens):
+        valid = t < seq_len
+        page_i = t // PAGE_SIZE
+        s = t % PAGE_SIZE
+
+        page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
 
         page_base = page_id * stride_kvb + pid_hkv * stride_kvh
 
-        for s in tl.range(PAGE_SIZE):
-            global_tok = page_i * PAGE_SIZE + s
-            if global_tok < seq_len:
-                k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
-                k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
+        k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+        k = tl.load(k_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
 
-                qk = tl.sum(q * k) * SCALE
+        qk = tl.sum(q * k) * SCALE
 
-                v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
-                v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
+        v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+        v = tl.load(v_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
 
-                old_max = acc_max
-                new_max = tl.maximum(old_max, qk)
-                exp_val = tl.exp(qk - new_max)
-                rescale_old = tl.exp(old_max - new_max)
+        # Safe logit: if !valid, qk = 0 but we also don't want it to affect softmax
+        # Set qk to -inf for invalid tokens so exp(qk-max)=0
+        qk = tl.where(valid, qk, -1e30)
 
-                acc_num = acc_num * rescale_old + v * exp_val
-                acc_denom = acc_denom * rescale_old + exp_val
-                acc_max = new_max
+        old_max = acc_max
+        new_max = tl.maximum(old_max, qk)
+        exp_val = tl.exp(qk - new_max)
+        rescale_old = tl.exp(old_max - new_max)
+
+        acc_num = acc_num * rescale_old + v * exp_val
+        acc_denom = acc_denom * rescale_old + exp_val
+        acc_max = new_max
 
     out = acc_num / acc_denom
 
@@ -97,7 +109,6 @@
         self.seq_len = seq_len
         self.page_size = page_size
         self.group_size = num_heads // num_kv_heads
-        self.max_pages = (seq_len + page_size - 1) // page_size
         self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
 
     def forward(self, query, kv_cache, block_table, seq_lens):
@@ -108,24 +119,23 @@
         P = self.page_size
         D_pad = triton.next_power_of_2(D)
         scale = _make_scale(D)
-        max_pages = self.max_pages
 
         out = torch.empty(B, H, D, dtype=query.dtype, device=query.device)
 
         assert kv_cache.stride(3) == 1, "kv_cache must have innermost stride 1"
 
-        grid = (B * H,)
+        grid = (B, H)
         _paged_decode_kernel[grid](
             out, query, kv_cache, block_table, seq_lens,
             query.stride(0), query.stride(1), query.stride(2),
             kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3),
-            block_table.stride(0),
+            block_table.stride(0), block_table.stride(1),
             out.stride(0), out.stride(1), out.stride(2),
-            B, H, Hkv, P,
+            Hkv, self.group_size,
             D, D_pad,
-            self.group_size,
+            P,
             scale,
-            max_pages,
+            self.seq_len,
         )
 
         return out

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,12 +1,15 @@
-"""Custom Triton paged attention decode kernel.
+"""Optimized Triton paged attention decode kernel.
 
 Strategy:
-- One program per (batch, query_head)
-- Each program loads page pointers from block_table for its batch
-- Loads one page of K and V at a time into registers/SRAM
-- Computes qk dot product, online softmax accumulation
-- Uses tl.dot for QK^T and weighted sum
-- GQA: pid_hkv = pid_h // group_size
+- One program per (batch, kv_head)
+- Each program:
+  1. Pre-loads all page indices from block_table for this batch
+  2. Loads entire pages of K and V into SRAM at once (PAGE_SIZE × head_dim blocks)
+  3. For each head in the GQA group, computes QK^T dot products
+  4. Maintains online softmax accumulators per head
+
+This processes all heads in a GQA group together, reducing redundant K/V loads.
+We load K/V pages as full 2D blocks using tl.load, and compute Q@K^T via row-wise dot.
 """
 import math
 import torch
@@ -32,6 +35,7 @@
     stride_bt_b, stride_bt_s,
     stride_ob, stride_oh, stride_od,
     # constexpr params
+    NUM_HEADS: tl.constexpr,
     NUM_KV_HEADS: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     HEAD_DIM: tl.constexpr,
@@ -39,63 +43,99 @@
     PAGE_SIZE: tl.constexpr,
     SCALE: tl.constexpr,
     MAX_SEQ_LEN: tl.constexpr,
+    BLOCK_BLOCKSIZE: tl.constexpr,
 ):
     pid_b = tl.program_id(0)
-    pid_h = tl.program_id(1)
+    pid_hkv = tl.program_id(1)
 
-    pid_hkv = pid_h // GROUP_SIZE
+    pid_h_start = pid_hkv * GROUP_SIZE
 
+    # Number of pages for this batch element
     seq_len = tl.load(SeqLens + pid_b)
+    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
 
-    d_offs = tl.arange(0, HEAD_DIM_PACKED)
-    d_mask = d_offs < HEAD_DIM
-
-    # Load Q: (HEAD_DIM,)
-    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
-    q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
-
-    acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
-    acc_denom = 0.0
-    acc_max = -1e30
+    # --- Pre-load page indices for this batch (at most BLOCK_BLOCKSIZE pages) ---
+    page_ids = tl.zeros([BLOCK_BLOCKSIZE], dtype=tl.int32)
+    page_valid = tl.zeros([BLOCK_BLOCKSIZE], dtype=tl.int1)
 
     bt_base = pid_b * stride_bt_b
+    bt_offs = tl.arange(0, BLOCK_BLOCKSIZE)
+    page_mask = bt_offs < num_pages
+    page_ids = tl.load(BlockTable + bt_base + bt_offs * stride_bt_s, mask=page_mask, other=0)
+    page_valid = page_mask
 
-    max_tokens = MAX_SEQ_LEN
+    # --- Load Q vectors for all heads in this group ---
+    # Q pointers: (GROUP_SIZE, HEAD_DIM)
+    q_offs_h = tl.arange(0, GROUP_SIZE)
+    q_offs_d = tl.arange(0, HEAD_DIM_PACKED)
+    q_mask_h = q_offs_h < GROUP_SIZE
+    q_mask_d = q_offs_d < HEAD_DIM
 
-    for t in tl.range(max_tokens):
-        valid = t < seq_len
-        page_i = t // PAGE_SIZE
-        s = t % PAGE_SIZE
+    # q_vals shape: (GROUP_SIZE, HEAD_DIM_PACKED)
+    q_ptrs = Q + (pid_b * stride_qb + (pid_h_start + q_offs_h[:, None]) * stride_qh
+                  + q_offs_d[None, :] * stride_qd)
+    q_mask = q_mask_h[:, None] & q_mask_d[None, :]
+    q_vals = tl.load(q_ptrs, mask=q_mask, other=0.0).to(tl.float32)
 
-        page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
+    # --- Accumulators per head in group ---
+    acc_num = tl.zeros([GROUP_SIZE, HEAD_DIM_PACKED], dtype=tl.float32)
+    acc_denom = tl.zeros([GROUP_SIZE], dtype=tl.float32)
+    acc_max = tl.full([GROUP_SIZE], -1e30, dtype=tl.float32)
 
+    # --- Process each page ---
+    max_pages = (MAX_SEQ_LEN + PAGE_SIZE - 1) // PAGE_SIZE
+
+    for p_idx in tl.range(max_pages):
+        p_valid = p_idx < num_pages
+
+        page_id = tl.load(BlockTable + bt_base + p_idx * stride_bt_s)
         page_base = page_id * stride_kvb + pid_hkv * stride_kvh
 
-        k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
-        k = tl.load(k_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
+        for s in tl.range(PAGE_SIZE):
+            global_tok = p_idx * PAGE_SIZE + s
+            tok_valid = (global_tok < seq_len) & p_valid
 
-        qk = tl.sum(q * k) * SCALE
+            # Load K: (HEAD_DIM,)
+            k_offs = d_offs = tl.arange(0, HEAD_DIM_PACKED)
+            k_mask = k_offs < HEAD_DIM
 
-        v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
-        v = tl.load(v_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
+            k_ptr = KV + page_base + s * stride_kvs + k_offs * stride_kvd
+            k = tl.load(k_ptr, mask=k_mask & tok_valid, other=0.0).to(tl.float32)
 
-        # Safe logit: if !valid, qk = 0 but we also don't want it to affect softmax
-        # Set qk to -inf for invalid tokens so exp(qk-max)=0
-        qk = tl.where(valid, qk, -1e30)
+            # QK dot: broadcast k to (GROUP_SIZE, HEAD_DIM), dot with q_vals
+            # qk_vals shape: (GROUP_SIZE,)
+            qk = tl.sum(q_vals * k[None, :], axis=1) * SCALE
 
-        old_max = acc_max
-        new_max = tl.maximum(old_max, qk)
-        exp_val = tl.exp(qk - new_max)
-        rescale_old = tl.exp(old_max - new_max)
+            # Load V: (HEAD_DIM,)
+            v_ptr = KV + page_base + s * stride_kvs + (k_offs + HEAD_DIM) * stride_kvd
+            v = tl.load(v_ptr, mask=k_mask & tok_valid, other=0.0).to(tl.float32)
 
-        acc_num = acc_num * rescale_old + v * exp_val
-        acc_denom = acc_denom * rescale_old + exp_val
-        acc_max = new_max
+            # Logit masking: -inf for invalid
+            qk = tl.where(tok_valid, qk, -1e30)
 
-    out = acc_num / acc_denom
+            # Online softmax update
+            old_max = acc_max
+            new_max = tl.maximum(old_max, qk)
+            exp_val = tl.exp(qk - new_max)
+            rescale = tl.exp(old_max - new_max)
 
-    out_ptr = Out + pid_b

... (1665 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,15 +1,13 @@
 """Optimized Triton paged attention decode kernel.
 
 Strategy:
-- One program per (batch, kv_head)
-- Each program:
-  1. Pre-loads all page indices from block_table for this batch
-  2. Loads entire pages of K and V into SRAM at once (PAGE_SIZE × head_dim blocks)
-  3. For each head in the GQA group, computes QK^T dot products
-  4. Maintains online softmax accumulators per head
+- One program per (batch, query_head)
+- Load entire page blocks at once using 2D block pointers
+- Process pages sequentially, accumulating online softmax
+- GQA: pid_hkv = pid_h // group_size
 
-This processes all heads in a GQA group together, reducing redundant K/V loads.
-We load K/V pages as full 2D blocks using tl.load, and compute Q@K^T via row-wise dot.
+Optimization: instead of loading token-by-token, we use
+tl.load with block pointers to fetch entire rows of K/V at once.
 """
 import math
 import torch
@@ -35,7 +33,6 @@
     stride_bt_b, stride_bt_s,
     stride_ob, stride_oh, stride_od,
     # constexpr params
-    NUM_HEADS: tl.constexpr,
     NUM_KV_HEADS: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     HEAD_DIM: tl.constexpr,
@@ -43,99 +40,63 @@
     PAGE_SIZE: tl.constexpr,
     SCALE: tl.constexpr,
     MAX_SEQ_LEN: tl.constexpr,
-    BLOCK_BLOCKSIZE: tl.constexpr,
 ):
     pid_b = tl.program_id(0)
-    pid_hkv = tl.program_id(1)
+    pid_h = tl.program_id(1)
 
-    pid_h_start = pid_hkv * GROUP_SIZE
+    pid_hkv = pid_h // GROUP_SIZE
 
-    # Number of pages for this batch element
     seq_len = tl.load(SeqLens + pid_b)
-    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
 
-    # --- Pre-load page indices for this batch (at most BLOCK_BLOCKSIZE pages) ---
-    page_ids = tl.zeros([BLOCK_BLOCKSIZE], dtype=tl.int32)
-    page_valid = tl.zeros([BLOCK_BLOCKSIZE], dtype=tl.int1)
+    d_offs = tl.arange(0, HEAD_DIM_PACKED)
+    d_mask = d_offs < HEAD_DIM
+
+    # Load Q: (HEAD_DIM,)
+    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
+    q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
+
+    acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
+    acc_denom = 0.0
+    acc_max = -1e30
 
     bt_base = pid_b * stride_bt_b
-    bt_offs = tl.arange(0, BLOCK_BLOCKSIZE)
-    page_mask = bt_offs < num_pages
-    page_ids = tl.load(BlockTable + bt_base + bt_offs * stride_bt_s, mask=page_mask, other=0)
-    page_valid = page_mask
 
-    # --- Load Q vectors for all heads in this group ---
-    # Q pointers: (GROUP_SIZE, HEAD_DIM)
-    q_offs_h = tl.arange(0, GROUP_SIZE)
-    q_offs_d = tl.arange(0, HEAD_DIM_PACKED)
-    q_mask_h = q_offs_h < GROUP_SIZE
-    q_mask_d = q_offs_d < HEAD_DIM
+    max_tokens = MAX_SEQ_LEN
 
-    # q_vals shape: (GROUP_SIZE, HEAD_DIM_PACKED)
-    q_ptrs = Q + (pid_b * stride_qb + (pid_h_start + q_offs_h[:, None]) * stride_qh
-                  + q_offs_d[None, :] * stride_qd)
-    q_mask = q_mask_h[:, None] & q_mask_d[None, :]
-    q_vals = tl.load(q_ptrs, mask=q_mask, other=0.0).to(tl.float32)
+    # Pre-compute page offsets to avoid repeated division/mod
+    for t in tl.range(max_tokens):
+        valid = t < seq_len
 
-    # --- Accumulators per head in group ---
-    acc_num = tl.zeros([GROUP_SIZE, HEAD_DIM_PACKED], dtype=tl.float32)
-    acc_denom = tl.zeros([GROUP_SIZE], dtype=tl.float32)
-    acc_max = tl.full([GROUP_SIZE], -1e30, dtype=tl.float32)
+        page_i = t // PAGE_SIZE
+        s = t % PAGE_SIZE
 
-    # --- Process each page ---
-    max_pages = (MAX_SEQ_LEN + PAGE_SIZE - 1) // PAGE_SIZE
+        page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
 
-    for p_idx in tl.range(max_pages):
-        p_valid = p_idx < num_pages
-
-        page_id = tl.load(BlockTable + bt_base + p_idx * stride_bt_s)
         page_base = page_id * stride_kvb + pid_hkv * stride_kvh
 
-        for s in tl.range(PAGE_SIZE):
-            global_tok = p_idx * PAGE_SIZE + s
-            tok_valid = (global_tok < seq_len) & p_valid
+        k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+        k = tl.load(k_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
 
-            # Load K: (HEAD_DIM,)
-            k_offs = d_offs = tl.arange(0, HEAD_DIM_PACKED)
-            k_mask = k_offs < HEAD_DIM
+        qk = tl.sum(q * k) * SCALE
 
-            k_ptr = KV + page_base + s * stride_kvs + k_offs * stride_kvd
-            k = tl.load(k_ptr, mask=k_mask & tok_valid, other=0.0).to(tl.float32)
+        v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+        v = tl.load(v_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
 
-            # QK dot: broadcast k to (GROUP_SIZE, HEAD_DIM), dot with q_vals
-            # qk_vals shape: (GROUP_SIZE,)
-            qk = tl.sum(q_vals * k[None, :], axis=1) * SCALE
+        qk = tl.where(valid, qk, -1e30)
 
-            # Load V: (HEAD_DIM,)
-            v_ptr = KV + page_base + s * stride_kvs + (k_offs + HEAD_DIM) * stride_kvd
-            v = tl.load(v_ptr, mask=k_mask & tok_valid, other=0.0).to(tl.float32)
+        old_max = acc_max
+        new_max = tl.maximum(old_max, qk)
+        exp_val = tl.exp(qk - new_max)
+        rescale_old = tl.exp(old_max - new_max)
 
-            # Logit masking: -inf for invalid
-            qk = tl.where(tok_valid, qk, -1e30)
+        acc_num = acc_num * rescale_old + v * exp_val
+        acc_denom = acc_denom * rescale_old + exp_val
+        acc_max = new_max
 
-            # Online softmax update
-            old_max = acc_max
-            new_max = tl.maximum(old_max, qk)
-            exp_val = tl.exp(qk - new_max)
-            rescale = tl.exp(old_max - new_max)
+    out = acc_num / acc_denom
 
-            acc_num = acc_num * rescale[:, None] + v[None, :] * exp_val[:, None]
-            acc_denom = acc_denom * rescale + exp_val
-   

... (1550 more chars)

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/sol_cuda.py
@@ -0,0 +1,242 @@
+"""CUDA C++ paged attention decode kernel using torch.compile/inline PTX.
+
+Strategy:
+- Use torch.compile with custom Triton kernel that loads full page tiles
+- Or use cuda kernel via load_inline
+- One thread block per (batch, kv_head)
+- Each thread handles subset of heads in GQA group + subset of head_dim
+- Load full page K/V into shared memory, process page at a time
+"""
+import math
+import torch
+import torch.nn as nn
+
+
+# We'll write a CUDA kernel and use load_inline
+_kernel_source = """
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <math.h>
+#include <float.h>
+
+#define WARP_SIZE 32
+
+// Paged attention decode kernel.
+// Configuration: (batch, kv_head, head_dim/VEC_SIZE)
+// One warp per (batch, kv_head), threads in warp share K/V loads and compute per-head dot products.
+
+__device__ inline float vec_dot(const nv_bfloat16* a, const nv_bfloat16* b, int d) {
+    float sum = 0.0f;
+    for (int i = 0; i < d; ++i) {
+        float ai = __bfloat162float(a[i]);
+        float bi = __bfloat162float(b[i]);
+        sum += ai * bi;
+    }
+    return sum;
+}
+
+__global__ void paged_decode_kernel(
+    nv_bfloat16* __restrict__ Out,
+    const nv_bfloat16* __restrict__ Q,
+    const nv_bfloat16* __restrict__ KV,
+    const int32_t* __restrict__ BlockTable,
+    const int32_t* __restrict__ SeqLens,
+    int batch,
+    int num_heads,
+    int num_kv_heads,
+    int head_dim,
+    int page_size,
+    float scale,
+    int stride_qb,
+    int stride_qh,
+    int stride_qd,
+    int stride_kvb,
+    int stride_kvs,
+    int stride_kvh,
+    int stride_kvd,
+    int stride_bt_b,
+    int stride_bt_s,
+    int stride_ob,
+    int stride_oh,
+    int stride_od
+) {
+    int pid_b = blockIdx.x % batch;
+    int pid_hkv = blockIdx.y;
+    int group_start = threadIdx.x;
+    int group_size = num_heads / num_kv_heads;
+
+    int block_idx_b = blockIdx.x;
+    pid_b = block_idx_b % batch;
+    pid_hkv = blockIdx.y;
+    group_start = threadIdx.x;
+
+    // Each thread in warp handles one head in the GQA group
+    // But only threads with group_start < group_size do work
+
+    // Load seq_len for this batch
+    int seq_len = SeqLens[pid_b];
+    int num_pages = (seq_len + page_size - 1) / page_size;
+
+    // Base pointer for Q
+    int h_base = pid_hkv * group_size;
+
+    // Shared memory for K and V page: (page_size, head_dim) * 2
+    extern __shared__ char smem[];
+    nv_bfloat16* sK = (nv_bfloat16*)smem;
+    nv_bfloat16* sV = sK + page_size * head_dim;
+
+    // Each thread gets its Q vector if within group
+    float q_vec[256]; // max head_dim
+    bool thread_active = (group_start < group_size);
+
+    if (thread_active) {
+        int h = h_base + group_start;
+        const nv_bfloat16* q_ptr = Q + pid_b * stride_qb + h * stride_qh;
+        for (int d = 0; d < head_dim; ++d) {
+            q_vec[d] = (float)__bfloat162float(q_ptr[d]);
+        }
+    }
+
+    // Accumulators per thread
+    float acc_num[256];
+    float acc_denom;
+    float acc_max;
+    for (int d = 0; d < head_dim; ++d) acc_num[d] = 0.0f;
+    acc_denom = 1e-30f;
+    acc_max = -1e30f;
+
+    // Process pages
+    for (int page_idx = 0; page_idx < num_pages; ++page_idx) {
+        int block_id = BlockTable[pid_b * stride_bt_b + page_idx * stride_bt_s];
+
+        // Load K and V page into shared memory
+        // Parallelize loading: each thread loads some rows
+        int tokens_per_thread = (page_size + blockDim.x - 1) / blockDim.x;
+        for (int t = 0; t < tokens_per_thread && (threadIdx.x + t) < page_size; ++t) {
+            int tok = threadIdx.x * tokens_per_thread + t;
+            // K row
+            const nv_bfloat16* k_src = KV + block_id * stride_kvb + tok * stride_kvs;
+            nv_bfloat16* k_dst = sK + tok * head_dim;
+            // V row
+            const nv_bfloat16* v_src = k_src + head_dim; // K and V are contiguous
+            nv_bfloat16* v_dst = sV + tok * head_dim;
+
+            for (int d = 0; d < head_dim; ++d) {
+                k_dst[d] = k_src[d];
+                v_dst[d] = v_src[d];
+            }
+        }
+        __syncthreads();
+
+        // Each active thread computes QK^T for its head and accumulates
+        if (thread_active) {
+            int num_tokens_in_page = (page_idx + 1) * page_size <= seq_len ? page_size : seq_len - page_idx * page_size;
+
+            for (int tok = 0; tok < num_tokens_in_page; ++tok) {
+                const nv_bfloat16* k_ptr = sK + tok * head_dim;
+                float qk = 0.0f;
+                for (int d = 0; d < head_dim; ++d) {
+                    qk += q_vec[d] * (float)__bfloat162float(k_ptr[d]);
+                }
+                qk *= scale;
+
+                const nv_bfloat16* v_ptr = sV + tok * head_dim;
+
+                float old_max = acc_max;
+                float new_max = fmaxf(old_max, qk);
+                float exp_val = __expf(qk - new_max);
+                float rescale = __expf(old_max - new_max);
+
+                for (int d = 0; d < head_dim; ++d) {
+                    acc_num[d] = acc_num[d] * rescale + exp_val * (float)__bfloat162float(v_ptr[d]);
+                }
+                acc_denom = acc_denom * rescale + exp_val;
+                acc_max = new_max;
+            }
+        }
+        __syncthreads();
+    }
+
+    // Compute final output
+    if (thread_active) {
+        int h = h_base + group_start;
+        float inv_denom = 1.0f / acc_denom;
+        nv_bfloat16* out_ptr = Out + pid_b * stride_ob + h * stride_oh;
+        for (int d = 0; d < head_dim; ++d) {
+            out_ptr[d] = __float2bfloat16(acc_num[d] * inv_denom);
+        }
+    }
+}
+"""
+
+
+def _make_scale(head_dim: int) -> float:
+    return 1.0 / math.sqrt(head_dim)
+
+
+def _compile_kernel(head_dim: int, page_size: int) -> object:
+    """Compile k

... (2098 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -2,12 +2,15 @@
 
 Strategy:
 - One program per (batch, query_head)
-- Load entire page blocks at once using 2D block pointers
-- Process pages sequentially, accumulating online softmax
-- GQA: pid_hkv = pid_h // group_size
+- Load entire page of K and V from KV cache using block pointers into registers
+- Compute QK^T for all tokens in the page at once
+- Use tl.dot for Q @ K_page^T to get scores for all tokens in the page
+- Online softmax across all pages
+- Store result
 
-Optimization: instead of loading token-by-token, we use
-tl.load with block pointers to fetch entire rows of K/V at once.
+Key optimization: Process PAGE_SIZE tokens per iteration instead of 1.
+Use tl.dot to do (1, head_dim) @ (head_dim, PAGE_SIZE) = (1, PAGE_SIZE) scores.
+Then broadcast V page (PAGE_SIZE, head_dim) by softmax weights.
 """
 import math
 import torch
@@ -63,36 +66,52 @@
 
     max_tokens = MAX_SEQ_LEN
 
-    # Pre-compute page offsets to avoid repeated division/mod
-    for t in tl.range(max_tokens):
-        valid = t < seq_len
-
-        page_i = t // PAGE_SIZE
-        s = t % PAGE_SIZE
-
+    # Pre-compute page offsets for this batch, then process page-sized chunks
+    for page_i in tl.range((MAX_SEQ_LEN + PAGE_SIZE - 1) // PAGE_SIZE):
         page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
 
         page_base = page_id * stride_kvb + pid_hkv * stride_kvh
 
-        k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
-        k = tl.load(k_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
+        # Process each token in this page
+        for s in tl.range(PAGE_SIZE):
+            global_tok = page_i * PAGE_SIZE + s
+            valid = global_tok < seq_len
 
-        qk = tl.sum(q * k) * SCALE
+            # If invalid, all accumulators stay the same (soft no-op via masking)
+            # Load K and V
+            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+            k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-        v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
-        v = tl.load(v_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
+            qk = tl.sum(q * k) * SCALE
+            qk = tl.where(valid, qk, -1e30)
 
-        qk = tl.where(valid, qk, -1e30)
+            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+            v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
 
-        old_max = acc_max
-        new_max = tl.maximum(old_max, qk)
-        exp_val = tl.exp(qk - new_max)
-        rescale_old = tl.exp(old_max - new_max)
+            # Softmax update: if invalid, qk=-inf so exp=0, no effect
+            # But we still need to handle the exp(-inf - new_max) = 0 correctly
+            new_max = tl.maximum(acc_max, qk)
 
-        acc_num = acc_num * rescale_old + v * exp_val
-        acc_denom = acc_denom * rescale_old + exp_val
-        acc_max = new_max
+            # When qk = -1e30 and acc_max is finite, new_max = acc_max
+            # exp_val = exp(-1e30 - acc_max) ~ 0
+            # rescale = exp(acc_max - acc_max) = 1.0
+            # This is correct: valid tokens get updated, invalid contribute nothing
 
+            # Edge: if only invalid tokens and acc_max = -1e30 initially
+            # Then new_max = -1e30, exp_val = 1, this is bad
+            # Fix: check if any valid token exists
+            is_first_valid = tl.cast(acc_max == -1e30, tl.float32)
+
+            # For invalid token when no valid tokens seen yet, set exp_val=0
+            exp_val = tl.where(valid | (is_first_valid == 1.0), tl.exp(qk - new_max), 0.0)
+            rescale = tl.exp(tl.where(is_first_valid == 1.0, 0.0, acc_max - new_max))
+
+            acc_num = acc_num * rescale + v * exp_val
+            acc_denom = acc_denom * rescale + exp_val
+            acc_max = tl.where(valid, new_max, acc_max)
+
+    # Avoid division by zero
+    acc_denom = tl.maximum(acc_denom, 1e-30)
     out = acc_num / acc_denom
 
     out_ptr = Out + pid_b * stride_ob + pid_h * stride_oh + d_offs * stride_od

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,16 +1,10 @@
 """Optimized Triton paged attention decode kernel.
 
 Strategy:
-- One program per (batch, query_head)
-- Load entire page of K and V from KV cache using block pointers into registers
-- Compute QK^T for all tokens in the page at once
-- Use tl.dot for Q @ K_page^T to get scores for all tokens in the page
-- Online softmax across all pages
-- Store result
-
-Key optimization: Process PAGE_SIZE tokens per iteration instead of 1.
-Use tl.dot to do (1, head_dim) @ (head_dim, PAGE_SIZE) = (1, PAGE_SIZE) scores.
-Then broadcast V page (PAGE_SIZE, head_dim) by softmax weights.
+- One program per (batch, kv_head). The program processes all heads in the GQA group.
+- Load entire page K and V into registers/SRAM once, share across group heads.
+- Use vectorized loads for head_dim dimension.
+- Online softmax accumulation per head.
 """
 import math
 import torch
@@ -36,86 +30,116 @@
     stride_bt_b, stride_bt_s,
     stride_ob, stride_oh, stride_od,
     # constexpr params
+    NUM_HEADS: tl.constexpr,
     NUM_KV_HEADS: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     HEAD_DIM: tl.constexpr,
-    HEAD_DIM_PACKED: tl.constexpr,
+    HEAD_DIM_PADDED: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
+    PAGE_SIZE_PADDED: tl.constexpr,
     SCALE: tl.constexpr,
     MAX_SEQ_LEN: tl.constexpr,
+    MAX_PAGES: tl.constexpr,
+    USE_BLOCK_PTRS: tl.constexpr,
 ):
     pid_b = tl.program_id(0)
-    pid_h = tl.program_id(1)
-
-    pid_hkv = pid_h // GROUP_SIZE
+    pid_hkv = tl.program_id(1)
 
     seq_len = tl.load(SeqLens + pid_b)
 
-    d_offs = tl.arange(0, HEAD_DIM_PACKED)
-    d_mask = d_offs < HEAD_DIM
+    # --- Pre-load page table for this batch ---
+    bt_base = pid_b * stride_bt_b
+    bt_offs = tl.arange(0, MAX_PAGES)
+    page_ids = tl.load(BlockTable + bt_base + bt_offs * stride_bt_s,
+                       mask=bt_offs < MAX_PAGES, other=0)
 
-    # Load Q: (HEAD_DIM,)
-    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
-    q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
+    # --- Load Q for all heads in the GQA group ---
+    # Shape: (GROUP_SIZE, HEAD_DIM_PADDED)
+    h_offs = tl.arange(0, GROUP_SIZE)
+    h_base = pid_hkv * GROUP_SIZE
+    d_offs = tl.arange(0, HEAD_DIM_PADDED)
+    q_mask = (h_offs < GROUP_SIZE)[:, None] & (d_offs < HEAD_DIM)[None, :]
 
-    acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
-    acc_denom = 0.0
-    acc_max = -1e30
+    # Q is (batch, num_heads, head_dim)
+    # q_ptrs: Q[pid_b, h_base:h_base+GROUP_SIZE, :]
+    q_row_base = pid_b * stride_qb + (h_base + h_offs[:, None]) * stride_qh
+    q_ptrs = q_row_base + d_offs[None, :] * stride_qd
+    q_vals = tl.load(q_ptrs, mask=q_mask, other=0.0).to(tl.float32)
 
-    bt_base = pid_b * stride_bt_b
+    # --- Accumulators: one per head in the group ---
+    acc_num = tl.zeros([GROUP_SIZE, HEAD_DIM_PADDED], dtype=tl.float32)
+    acc_denom = tl.zeros([GROUP_SIZE], dtype=tl.float32)
+    acc_max = tl.full([GROUP_SIZE], -1e30, dtype=tl.float32)
 
+    # --- Process each page ---
+    page_base_offsets = page_ids * stride_kvb + pid_hkv * stride_kvh
+
+    # We'll do token-by-token but with vectorized head_dim loads
     max_tokens = MAX_SEQ_LEN
 
-    # Pre-compute page offsets for this batch, then process page-sized chunks
-    for page_i in tl.range((MAX_SEQ_LEN + PAGE_SIZE - 1) // PAGE_SIZE):
+    for t in tl.range(max_tokens):
+        valid = t < seq_len
+        page_i = t // PAGE_SIZE
+        s = t % PAGE_SIZE
+
+        # Get page base offset for this token
+        pbo = tl.load(page_base_offsets + page_i) if MAX_PAGES > 1 else tl.load(page_base_offsets)
+
+        # Actually let's compute offset without gather
         page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
+        base = page_id * stride_kvb + pid_hkv * stride_kvh + s * stride_kvs
 
-        page_base = page_id * stride_kvb + pid_hkv * stride_kvh
+        # Load K: (HEAD_DIM,)
+        k_offs = d_offs
+        k_ptr = KV + base + k_offs * stride_kvd
+        k = tl.load(k_ptr, mask=(d_offs < HEAD_DIM) & valid, other=0.0).to(tl.float32)
 
-        # Process each token in this page
-        for s in tl.range(PAGE_SIZE):
-            global_tok = page_i * PAGE_SIZE + s
-            valid = global_tok < seq_len
+        # QK^T: (GROUP_SIZE,) = sum(q_vals * k[None, :], axis=1) * SCALE
+        qk = tl.sum(q_vals * k[None, :], axis=1) * SCALE
+        qk = tl.where(valid, qk, -1e30)
 
-            # If invalid, all accumulators stay the same (soft no-op via masking)
-            # Load K and V
-            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
-            k = tl.load(k_ptr, mask=d_mask, other=0.0).to(tl.float32)
+        # Load V: (HEAD_DIM,)
+        v_ptr = KV + base + (k_offs + HEAD_DIM) * stride_kvd
+        v = tl.load(v_ptr, mask=(d_offs < HEAD_DIM) & valid, other=0.0).to(tl.float32)
 
-            qk = tl.sum(q * k) * SCALE
-            qk = tl.where(valid, qk, -1e30)
+        # Online softmax update
+        # For invalid tokens: qk=-inf, so exp(qk-new_max) ~ 0, doesn't affect accumulators
+        old_max = acc_max
+        new_max = tl.maximum(old_max, qk)
 
-            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
-            v = tl.load(v_ptr, mask=d_mask, other=0.0).to(tl.float32)
+        # Careful: when all tokens are invalid, acc_max stays -1e30
+        # exp(-1e30 - (-1e30)) = 1, which would give wrong result
+        # Fix: only update when valid
+        any_valid = tl.reduce(acc_max > -2e29, 0, combine_fn=tl.maximum)
+        # Simpler: just mask out invalid updates
 
-            # Softmax update: if invalid, qk=-inf so exp=0, no effect
-            # But we still need to handle the exp(-inf - new_max) = 0 correctly
-            new

... (3681 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,10 +1,10 @@
-"""Optimized Triton paged attention decode kernel.
+"""CUDA C++ paged attention decode kernel via torch.compile/extension.
 
 Strategy:
-- One program per (batch, kv_head). The program processes all heads in the GQA group.
-- Load entire page K and V into registers/SRAM once, share across group heads.
-- Use vectorized loads for head_dim dimension.
-- Online softmax accumulation per head.
+- Write raw CUDA kernel and use torch.compile/inline PTX for best performance
+- One thread block per (batch, kv_head)
+- Load K/V pages into shared memory
+- Compute attention for all heads in GQA group sharing the K/V load
 """
 import math
 import torch
@@ -13,6 +13,7 @@
 import triton.language as tl
 
 
+# Keep the working Triton kernel as solution, we'll optimize later
 def _make_scale(head_dim: int) -> float:
     return 1.0 / math.sqrt(head_dim)
 
@@ -30,116 +31,69 @@
     stride_bt_b, stride_bt_s,
     stride_ob, stride_oh, stride_od,
     # constexpr params
-    NUM_HEADS: tl.constexpr,
     NUM_KV_HEADS: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     HEAD_DIM: tl.constexpr,
-    HEAD_DIM_PADDED: tl.constexpr,
+    HEAD_DIM_PACKED: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
-    PAGE_SIZE_PADDED: tl.constexpr,
     SCALE: tl.constexpr,
     MAX_SEQ_LEN: tl.constexpr,
-    MAX_PAGES: tl.constexpr,
-    USE_BLOCK_PTRS: tl.constexpr,
 ):
     pid_b = tl.program_id(0)
-    pid_hkv = tl.program_id(1)
+    pid_h = tl.program_id(1)
+
+    pid_hkv = pid_h // GROUP_SIZE
 
     seq_len = tl.load(SeqLens + pid_b)
 
-    # --- Pre-load page table for this batch ---
+    d_offs = tl.arange(0, HEAD_DIM_PACKED)
+    d_mask = d_offs < HEAD_DIM
+
+    # Load Q: (HEAD_DIM,)
+    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
+    q = tl.load(q_ptr, mask=d_mask, other=0.0).to(tl.float32)
+
+    acc_num = tl.zeros([HEAD_DIM_PACKED], dtype=tl.float32)
+    acc_denom = 0.0
+    acc_max = -1e30
+
     bt_base = pid_b * stride_bt_b
-    bt_offs = tl.arange(0, MAX_PAGES)
-    page_ids = tl.load(BlockTable + bt_base + bt_offs * stride_bt_s,
-                       mask=bt_offs < MAX_PAGES, other=0)
 
-    # --- Load Q for all heads in the GQA group ---
-    # Shape: (GROUP_SIZE, HEAD_DIM_PADDED)
-    h_offs = tl.arange(0, GROUP_SIZE)
-    h_base = pid_hkv * GROUP_SIZE
-    d_offs = tl.arange(0, HEAD_DIM_PADDED)
-    q_mask = (h_offs < GROUP_SIZE)[:, None] & (d_offs < HEAD_DIM)[None, :]
-
-    # Q is (batch, num_heads, head_dim)
-    # q_ptrs: Q[pid_b, h_base:h_base+GROUP_SIZE, :]
-    q_row_base = pid_b * stride_qb + (h_base + h_offs[:, None]) * stride_qh
-    q_ptrs = q_row_base + d_offs[None, :] * stride_qd
-    q_vals = tl.load(q_ptrs, mask=q_mask, other=0.0).to(tl.float32)
-
-    # --- Accumulators: one per head in the group ---
-    acc_num = tl.zeros([GROUP_SIZE, HEAD_DIM_PADDED], dtype=tl.float32)
-    acc_denom = tl.zeros([GROUP_SIZE], dtype=tl.float32)
-    acc_max = tl.full([GROUP_SIZE], -1e30, dtype=tl.float32)
-
-    # --- Process each page ---
-    page_base_offsets = page_ids * stride_kvb + pid_hkv * stride_kvh
-
-    # We'll do token-by-token but with vectorized head_dim loads
     max_tokens = MAX_SEQ_LEN
 
     for t in tl.range(max_tokens):
         valid = t < seq_len
+
         page_i = t // PAGE_SIZE
         s = t % PAGE_SIZE
 
-        # Get page base offset for this token
-        pbo = tl.load(page_base_offsets + page_i) if MAX_PAGES > 1 else tl.load(page_base_offsets)
+        page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
 
-        # Actually let's compute offset without gather
-        page_id = tl.load(BlockTable + bt_base + page_i * stride_bt_s)
-        base = page_id * stride_kvb + pid_hkv * stride_kvh + s * stride_kvs
+        page_base = page_id * stride_kvb + pid_hkv * stride_kvh
 
-        # Load K: (HEAD_DIM,)
-        k_offs = d_offs
-        k_ptr = KV + base + k_offs * stride_kvd
-        k = tl.load(k_ptr, mask=(d_offs < HEAD_DIM) & valid, other=0.0).to(tl.float32)
+        k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+        k = tl.load(k_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
 
-        # QK^T: (GROUP_SIZE,) = sum(q_vals * k[None, :], axis=1) * SCALE
-        qk = tl.sum(q_vals * k[None, :], axis=1) * SCALE
+        qk = tl.sum(q * k) * SCALE
+
+        v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+        v = tl.load(v_ptr, mask=d_mask & valid, other=0.0).to(tl.float32)
+
         qk = tl.where(valid, qk, -1e30)
 
-        # Load V: (HEAD_DIM,)
-        v_ptr = KV + base + (k_offs + HEAD_DIM) * stride_kvd
-        v = tl.load(v_ptr, mask=(d_offs < HEAD_DIM) & valid, other=0.0).to(tl.float32)
-
-        # Online softmax update
-        # For invalid tokens: qk=-inf, so exp(qk-new_max) ~ 0, doesn't affect accumulators
         old_max = acc_max
         new_max = tl.maximum(old_max, qk)
+        exp_val = tl.exp(qk - new_max)
+        rescale_old = tl.exp(old_max - new_max)
 
-        # Careful: when all tokens are invalid, acc_max stays -1e30
-        # exp(-1e30 - (-1e30)) = 1, which would give wrong result
-        # Fix: only update when valid
-        any_valid = tl.reduce(acc_max > -2e29, 0, combine_fn=tl.maximum)
-        # Simpler: just mask out invalid updates
-
-        # When valid=False: qk = -1e30 (already masked above)
-        # new_max doesn't change because old_max >= any actual score >> -1e30
-        # Actually old_max = -1e30 initially when no valid tokens seen
-        # Then qk = -1e30 for invalid -> new_max = -1e30
-        # exp(qk - new_max) = exp(0) = 1.0 -> WRONG
-
-        # Solution: set invalid qk values so they are ignored
-        is_first = acc_max <= -1e29  # no valid tokens seen yet
-        safe_qk = tl.where(valid, qk, -1e30)
-        new_max = tl.maximu

... (2212 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,99 +1,202 @@
-"""CUDA C++ paged attention decode kernel via torch.compile/extension.
+"""High-performance CUDA paged attention decode kernel.
 
 Strategy:
-- Write raw CUDA kernel and use torch.compile/inline PTX for best performance
 - One thread block per (batch, kv_head)
-- Load K/V pages into shared memory
-- Compute attention for all heads in GQA group sharing the K/V load
+- Each thread handles one head in the GQA group
+- Load each page of K/V into shared memory once
+- All threads in block compute QK^T using their own Q and shared K/V
+- Online softmax accumulation per thread
+- Vectorized memory access for K/V loads
+
+This reduces KV cache bandwidth by sharing K/V loads across a GQA group of heads.
 """
 import math
+import os
 import torch
 import torch.nn as nn
-import triton
-import triton.language as tl
-
-
-# Keep the working Triton kernel as solution, we'll optimize later
+
+
+_CUDA_SOURCE = """
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <math.h>
+
+// Paged attention decode kernel.
+// Grid: (batch, num_kv_heads), Block: (group_size, 1, 1)
+// Each thread handles one query head in the GQA group.
+// Shared memory holds the current page's K and V.
+
+extern "C" __global__
+void __launch_bounds__(1024) paged_decode_kernel(
+    nv_bfloat16* __restrict__ Out,
+    const nv_bfloat16* __restrict__ Q,
+    const nv_bfloat16* __restrict__ KV,
+    const int* __restrict__ BlockTable,
+    const int* __restrict__ SeqLens,
+    int batch,
+    int num_heads,
+    int num_kv_heads,
+    int head_dim,
+    int page_size,
+    float scale,
+    int stride_qb, int stride_qh, int stride_qd,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
+    int stride_bt_b, int stride_bt_s,
+    int stride_ob, int stride_oh, int stride_od
+) {
+    // One block per (batch, kv_head)
+    extern __shared__ char smem_ptr[];
+    const int pid_b = blockIdx.x;
+    const int pid_hkv = blockIdx.y;
+    const int group_size = num_heads / num_kv_heads;
+
+    // Each thread handles one head
+    const int tid = threadIdx.x;
+    const int h = pid_hkv * group_size + tid;
+    const bool is_valid_thread = (tid < group_size);
+
+    // seq_len
+    const int seq_len = SeqLens[pid_b];
+    const int num_pages = (seq_len + page_size - 1) / page_size;
+
+    // Shared memory: K page + V page
+    // K: page_size x head_dim, V: page_size x head_dim
+    nv_bfloat16* __restrict__ smem_k = (nv_bfloat16*)smem_ptr;
+    nv_bfloat16* __restrict__ smem_v = smem_k + page_size * head_dim;
+
+    // Load Q for this thread
+    float q_vec[256]; // max head_dim = 128
+    if (is_valid_thread) {
+        const nv_bfloat16* __restrict__ q_base = Q + pid_b * stride_qb + h * stride_qh;
+        #pragma unroll
+        for (int d = 0; d < 128; d += 8) {
+            #pragma unroll
+            for (int i = 0; i < 8; ++i) {
+                if ((d + i) < head_dim)
+                    q_vec[d + i] = __bfloat162float(q_base[d + i]);
+            }
+        }
+    }
+
+    // Accumulators
+    float acc_num[256];
+    float acc_denom = 1e-30f;
+    float acc_max = -1e30f;
+    #pragma unroll
+    for (int d = 0; d < 128; ++d) acc_num[d] = 0.0f;
+
+    // Process each page
+    for (int page_idx = 0; page_idx < num_pages; ++page_idx) {
+        const int block_id = BlockTable[pid_b * stride_bt_b + page_idx * stride_bt_s];
+
+        // Load K and V page into shared memory (coalesced across threads)
+        const int tokens_per_thread = (page_size + group_size - 1) / group_size;
+        #pragma unroll
+        for (int t_local = 0; t_local < 16; ++t_local) {  // max page_size=16 tokens_per_thread=1 at most
+            const int tok_idx = tid * tokens_per_thread + t_local;
+            if (tok_idx < page_size) {
+                const nv_bfloat16* __restrict__ k_src = KV + block_id * stride_kvb + tok_idx * stride_kvs;
+                nv_bfloat16* __restrict__ k_dst = smem_k + tok_idx * head_dim;
+                // K and V are contiguous: V is right after K
+                const nv_bfloat16* __restrict__ v_src = k_src + head_dim;
+                nv_bfloat16* __restrict__ v_dst = smem_v + tok_idx * head_dim;
+                #pragma unroll
+                for (int d = 0; d < 128; d += 8) {
+                    #pragma unroll
+                    for (int i = 0; i < 8; ++i) {
+                        if ((d + i) < head_dim) {
+                            k_dst[d + i] = k_src[d + i];
+                            v_dst[d + i] = v_src[d + i];
+                        }
+                    }
+                }
+            }
+        }
+        __syncthreads();
+
+        // Each thread computes attention for its head using shared K/V
+        if (is_valid_thread) {
+            const int tokens_in_page = (page_idx + 1) * page_size <= seq_len
+                ? page_size : seq_len - page_idx * page_size;
+
+            for (int tok = 0; tok < page_size; ++tok) {
+                if (tok >= tokens_in_page) break;
+
+                // Load K row from shared memory
+                const nv_bfloat16* __restrict__ k_row = smem_k + tok * head_dim;
+                float qk = 0.0f;
+                #pragma unroll
+                for (int d = 0; d < 128; d += 8) {
+                    #pragma unroll
+                    for (int i = 0; i < 8; ++i) {
+                        if ((d + i) < head_dim) {
+                            qk += q_vec[d + i] * __bfloat162float(k_row[d + i]);
+                        }
+                    }
+                }
+                qk *= scale;
+
+                // Load V row from shared memory
+                const nv_bfloat16* __restrict__ v_row = smem_v + tok * head_dim;
+
+                // Online softmax
+                const float old_max = acc_max;
+             

... (5483 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -185,6 +185,7 @@
             "-Xptxas=-v",
             "--use_fast_math",
         ],
+        extra_include_paths=[],
         verbose=False,
     )
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -175,17 +175,31 @@
 def _compile_kernel():
     from torch.utils.cpp_extension import load_inline
     # Include a shape-dependent key for caching
+    CPP_SRC = '''
+#include <torch/extension.h>
+void paged_decode_kernel(
+    void *Out, const void *Q, const void *KV,
+    const int *BlockTable, const int *SeqLens,
+    int batch, int num_heads, int num_kv_heads,
+    int head_dim, int page_size, float scale,
+    int stride_qb, int stride_qh, int stride_qd,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
+    int stride_bt_b, int stride_bt_s,
+    int stride_ob, int stride_oh, int stride_od);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("paged_decode_kernel", &paged_decode_kernel, "paged_decode_kernel");
+}
+'''
     return load_inline(
         name="paged_decode_v2",
-        cpp_sources='',
+        cpp_sources=CPP_SRC,
         cuda_sources=_CUDA_SOURCE,
-        functions=["paged_decode_kernel"],
         extra_cuda_cflags=[
             "-O3",
             "-Xptxas=-v",
             "--use_fast_math",
         ],
-        extra_include_paths=[],
         verbose=False,
     )
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,217 +1,16 @@
-"""High-performance CUDA paged attention decode kernel.
+"""Paged attention decode kernel using optimized PyTorch tensor operations.
 
 Strategy:
-- One thread block per (batch, kv_head)
-- Each thread handles one head in the GQA group
-- Load each page of K/V into shared memory once
-- All threads in block compute QK^T using their own Q and shared K/V
-- Online softmax accumulation per thread
-- Vectorized memory access for K/V loads
-
-This reduces KV cache bandwidth by sharing K/V loads across a GQA group of heads.
+- Gather all K/V pages at once using advanced indexing
+- Reshape to (batch, seq_len, num_kv_heads, head_dim) for K and V
+- Use broadcasted einsum for QK^T and attention
+- Repeat interleave for GQA
+- This is essentially the reference but without Python loops over batch
+- All operations done in fused tensor form on GPU
 """
 import math
-import os
 import torch
 import torch.nn as nn
-
-
-_CUDA_SOURCE = """
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <math.h>
-
-// Paged attention decode kernel.
-// Grid: (batch, num_kv_heads), Block: (group_size, 1, 1)
-// Each thread handles one query head in the GQA group.
-// Shared memory holds the current page's K and V.
-
-extern "C" __global__
-void __launch_bounds__(1024) paged_decode_kernel(
-    nv_bfloat16* __restrict__ Out,
-    const nv_bfloat16* __restrict__ Q,
-    const nv_bfloat16* __restrict__ KV,
-    const int* __restrict__ BlockTable,
-    const int* __restrict__ SeqLens,
-    int batch,
-    int num_heads,
-    int num_kv_heads,
-    int head_dim,
-    int page_size,
-    float scale,
-    int stride_qb, int stride_qh, int stride_qd,
-    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
-    int stride_bt_b, int stride_bt_s,
-    int stride_ob, int stride_oh, int stride_od
-) {
-    // One block per (batch, kv_head)
-    extern __shared__ char smem_ptr[];
-    const int pid_b = blockIdx.x;
-    const int pid_hkv = blockIdx.y;
-    const int group_size = num_heads / num_kv_heads;
-
-    // Each thread handles one head
-    const int tid = threadIdx.x;
-    const int h = pid_hkv * group_size + tid;
-    const bool is_valid_thread = (tid < group_size);
-
-    // seq_len
-    const int seq_len = SeqLens[pid_b];
-    const int num_pages = (seq_len + page_size - 1) / page_size;
-
-    // Shared memory: K page + V page
-    // K: page_size x head_dim, V: page_size x head_dim
-    nv_bfloat16* __restrict__ smem_k = (nv_bfloat16*)smem_ptr;
-    nv_bfloat16* __restrict__ smem_v = smem_k + page_size * head_dim;
-
-    // Load Q for this thread
-    float q_vec[256]; // max head_dim = 128
-    if (is_valid_thread) {
-        const nv_bfloat16* __restrict__ q_base = Q + pid_b * stride_qb + h * stride_qh;
-        #pragma unroll
-        for (int d = 0; d < 128; d += 8) {
-            #pragma unroll
-            for (int i = 0; i < 8; ++i) {
-                if ((d + i) < head_dim)
-                    q_vec[d + i] = __bfloat162float(q_base[d + i]);
-            }
-        }
-    }
-
-    // Accumulators
-    float acc_num[256];
-    float acc_denom = 1e-30f;
-    float acc_max = -1e30f;
-    #pragma unroll
-    for (int d = 0; d < 128; ++d) acc_num[d] = 0.0f;
-
-    // Process each page
-    for (int page_idx = 0; page_idx < num_pages; ++page_idx) {
-        const int block_id = BlockTable[pid_b * stride_bt_b + page_idx * stride_bt_s];
-
-        // Load K and V page into shared memory (coalesced across threads)
-        const int tokens_per_thread = (page_size + group_size - 1) / group_size;
-        #pragma unroll
-        for (int t_local = 0; t_local < 16; ++t_local) {  // max page_size=16 tokens_per_thread=1 at most
-            const int tok_idx = tid * tokens_per_thread + t_local;
-            if (tok_idx < page_size) {
-                const nv_bfloat16* __restrict__ k_src = KV + block_id * stride_kvb + tok_idx * stride_kvs;
-                nv_bfloat16* __restrict__ k_dst = smem_k + tok_idx * head_dim;
-                // K and V are contiguous: V is right after K
-                const nv_bfloat16* __restrict__ v_src = k_src + head_dim;
-                nv_bfloat16* __restrict__ v_dst = smem_v + tok_idx * head_dim;
-                #pragma unroll
-                for (int d = 0; d < 128; d += 8) {
-                    #pragma unroll
-                    for (int i = 0; i < 8; ++i) {
-                        if ((d + i) < head_dim) {
-                            k_dst[d + i] = k_src[d + i];
-                            v_dst[d + i] = v_src[d + i];
-                        }
-                    }
-                }
-            }
-        }
-        __syncthreads();
-
-        // Each thread computes attention for its head using shared K/V
-        if (is_valid_thread) {
-            const int tokens_in_page = (page_idx + 1) * page_size <= seq_len
-                ? page_size : seq_len - page_idx * page_size;
-
-            for (int tok = 0; tok < page_size; ++tok) {
-                if (tok >= tokens_in_page) break;
-
-                // Load K row from shared memory
-                const nv_bfloat16* __restrict__ k_row = smem_k + tok * head_dim;
-                float qk = 0.0f;
-                #pragma unroll
-                for (int d = 0; d < 128; d += 8) {
-                    #pragma unroll
-                    for (int i = 0; i < 8; ++i) {
-                        if ((d + i) < head_dim) {
-                            qk += q_vec[d + i] * __bfloat162float(k_row[d + i]);
-                        }
-                    }
-                }
-                qk *= scale;
-
-                // Load V row from shared memory
-                const nv_bfloat16* __restrict__ v_row = smem_v + tok * head_dim;
-
-                // Online softmax
-                const float old_m

... (6939 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,16 +1,179 @@
-"""Paged attention decode kernel using optimized PyTorch tensor operations.
+"""High-performance CUDA paged attention decode kernel via torch.utils.cpp_extension.
 
 Strategy:
-- Gather all K/V pages at once using advanced indexing
-- Reshape to (batch, seq_len, num_kv_heads, head_dim) for K and V
-- Use broadcasted einsum for QK^T and attention
-- Repeat interleave for GQA
-- This is essentially the reference but without Python loops over batch
-- All operations done in fused tensor form on GPU
+- One thread block per (batch, kv_head)
+- Each thread handles one head in the GQA group
+- Load each page of K/V into shared memory once
+- All threads compute QK^T using shared K/V (saves KV bandwidth)
+- Online softmax per thread
 """
 import math
 import torch
 import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+def _make_scale(head_dim: int) -> float:
+    return 1.0 / math.sqrt(head_dim)
+
+_CUDA_SRC = r"""
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+__global__ void __launch_bounds__(512) paged_decode_kernel(
+    nv_bfloat16* __restrict__ Out,
+    const nv_bfloat16* __restrict__ Q,
+    const nv_bfloat16* __restrict__ KV,
+    const int* __restrict__ BlockTable,
+    const int* __restrict__ SeqLens,
+    int batch, int num_heads, int num_kv_heads, int head_dim,
+    int page_size, float scale,
+    int stride_qb, int stride_qh, int stride_qd,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
+    int stride_bt_b, int stride_bt_s,
+    int stride_ob, int stride_oh, int stride_od
+) {
+    int pid_b = blockIdx.x;
+    int pid_hkv = blockIdx.y;
+    int group_size = num_heads / num_kv_heads;
+    int tid = threadIdx.x;
+
+    if (tid >= group_size) return;
+
+    int h = pid_hkv * group_size + tid;
+    int seq_len = SeqLens[pid_b];
+    int num_pages = (seq_len + page_size - 1) / page_size;
+
+    // Shared memory for current page K + V
+    extern __shared__ char smem[];
+    nv_bfloat16* __restrict__ sK = (nv_bfloat16*)smem;
+    nv_bfloat16* __restrict__ sV = sK + page_size * head_dim;
+
+    // Load Q
+    float q[128];
+    const nv_bfloat16* qp = Q + pid_b * stride_qb + h * stride_qh;
+    for (int d = 0; d < head_dim; ++d)
+        q[d] = __bfloat162float(qp[d]);
+
+    // Accumulators
+    float acc_num[128] = {0};
+    float acc_denom = 1e-30f;
+    float acc_max = -1e30f;
+
+    for (int pg = 0; pg < num_pages; ++pg) {
+        int block_id = BlockTable[pid_b * stride_bt_b + pg * stride_bt_s];
+
+        // Cooperative load of K/V page
+        const nv_bfloat16* k_page_src = KV + block_id * stride_kvb;
+        const int tokens_to_load = page_size;
+        const int items_thread = (tokens_to_load * head_dim + group_size - 1) / group_size;
+
+        for (int i_idx = 0; i_idx < items_thread; ++i_idx) {
+            int flat = tid * items_thread + i_idx;
+            if (flat < tokens_to_load * head_dim) {
+                int row = flat / head_dim;
+                int col = flat % head_dim;
+                // K
+                sK[row * head_dim + col] = k_page_src[row * stride_kvs + col * stride_kvd];
+                // V (contiguous after K in last dim)
+                sV[row * head_dim + col] = k_page_src[row * stride_kvs + head_dim + col * stride_kvd];
+            }
+        }
+        __syncthreads();
+
+        // Compute attention for tokens in this page
+        int tokens_in_page = (pg + 1) * page_size <= seq_len ? page_size : seq_len - pg * page_size;
+        for (int t = 0; t < tokens_in_page; ++t) {
+            float qk = 0.0f;
+            const nv_bfloat16* kr = sK + t * head_dim;
+            #pragma unroll
+            for (int d = 0; d < 128; ++d) {
+                if (d < head_dim) qk += q[d] * __bfloat162float(kr[d]);
+            }
+            qk *= scale;
+
+            float old_max = acc_max;
+            float new_max = fmaxf(old_max, qk);
+            float exp_v = __expf(qk - new_max);
+            float rescale = __expf(old_max - new_max);
+
+            const nv_bfloat16* vr = sV + t * head_dim;
+            for (int d = 0; d < head_dim; ++d) {
+                acc_num[d] = acc_num[d] * rescale + exp_v * __bfloat162float(vr[d]);
+            }
+            acc_denom = acc_denom * rescale + exp_v;
+            acc_max = new_max;
+        }
+        __syncthreads();
+    }
+
+    // Write output
+    nv_bfloat16* __restrict__ op = Out + pid_b * stride_ob + h * stride_oh;
+    float inv = 1.0f / acc_denom;
+    for (int d = 0; d < head_dim; ++d)
+        op[d] = __float2bfloat16(acc_num[d] * inv);
+}
+"""
+
+_CPP_SRC = r"""
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+void paged_decode_forward(
+    torch::Tensor out,
+    torch::Tensor query,
+    torch::Tensor kv_cache,
+    torch::Tensor block_table,
+    torch::Tensor seq_lens,
+    float scale) {
+
+    int B = query.size(0);
+    int H = query.size(1);
+    int Hkv = kv_cache.size(2);
+    int D = query.size(2);
+    int P = kv_cache.size(1);  // page_size
+    int G = H / Hkv;
+
+    int smem_size = 2 * P * D * 2;  // K + V in bf16
+
+    int gridX = B;
+    int gridY = Hkv;
+    int blockZ = G;
+
+    paged_decode_kernel<<<dim3(gridX, gridY, 1), dim3(blockZ, 1, 1), smem_size>>>(
+        (nv_bfloat16*)out.data_ptr(),
+        (const nv_bfloat16*)query.data_ptr(),
+        (const nv_bfloat16*)kv_cache.data_ptr(),
+        block_table.data_ptr<int>(),
+        seq_lens.data_ptr<int>(),
+        B, H, Hkv, D, P, scale,
+        query.stride(0), query.stride(1), query.stride(2),
+        kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3),
+        block_table.stride(0), block_table.stride(1),
+        out.stride(0), out.stride(1), out.stride(2)
+    );
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION

... (4813 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,182 +1,154 @@
-"""High-performance CUDA paged attention decode kernel via torch.utils.cpp_extension.
+"""High-performance CUDA paged attention decode kernel.
 
 Strategy:
 - One thread block per (batch, kv_head)
-- Each thread handles one head in the GQA group
-- Load each page of K/V into shared memory once
-- All threads compute QK^T using shared K/V (saves KV bandwidth)
-- Online softmax per thread
+- Block has num_warps * 32 threads: each warp is one query head in GQA group
+- Load K/V page into shared memory cooperatively
+- Each warp computes QK^T using warp shuffle / register loads
+- Online softmax per warp with per-warp accumulator
 """
 import math
 import torch
 import torch.nn as nn
-from torch.utils.cpp_extension import load_inline
+import triton
+import triton.language as tl
+
 
 def _make_scale(head_dim: int) -> float:
     return 1.0 / math.sqrt(head_dim)
 
-_CUDA_SRC = r"""
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
 
-__global__ void __launch_bounds__(512) paged_decode_kernel(
-    nv_bfloat16* __restrict__ Out,
-    const nv_bfloat16* __restrict__ Q,
-    const nv_bfloat16* __restrict__ KV,
-    const int* __restrict__ BlockTable,
-    const int* __restrict__ SeqLens,
-    int batch, int num_heads, int num_kv_heads, int head_dim,
-    int page_size, float scale,
-    int stride_qb, int stride_qh, int stride_qd,
-    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
-    int stride_bt_b, int stride_bt_s,
-    int stride_ob, int stride_oh, int stride_od
-) {
-    int pid_b = blockIdx.x;
-    int pid_hkv = blockIdx.y;
-    int group_size = num_heads / num_kv_heads;
-    int tid = threadIdx.x;
+# The key optimization for paged decode: load full K/V pages into registers/sram,
+# process PAGE_SIZE tokens per page iteration, share loads across GQA group.
+#
+# With PAGE_SIZE=16 and HEAD_DIM=128:
+# - Per page K load: 16 * 128 * 2 = 4096 bytes
+# - Per page V load: 16 * 128 * 2 = 4096 bytes
+# - Process all 16 tokens at once via tl.dot or vectorized loads
+#
+# For head_dim=128, one thread can hold the entire K/V row in registers.
 
-    if (tid >= group_size) return;
 
-    int h = pid_hkv * group_size + tid;
-    int seq_len = SeqLens[pid_b];
-    int num_pages = (seq_len + page_size - 1) / page_size;
+@triton.jit
+def _paged_decode_kernel(
+    Out,
+    Q,
+    KV,
+    BlockTable,
+    SeqLens,
+    # strides
+    stride_qb,
+    stride_qh,
+    stride_qd,
+    stride_kvb,
+    stride_kvs,
+    stride_kvh,
+    stride_kvd,
+    stride_ob,
+    stride_oh,
+    stride_od,
+    stride_bt_b,
+    # constexpr
+    HEAD_DIM: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    SCALE: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Decode kernel with PAGE_SIZE tokens processed per page iteration.
+    
+    Each program handles one (batch, head) pair.
+    Uses tl.dot to compute (1, HEAD_DIM) @ (HEAD_DIM, BLOCK_N) = (1, BLOCK_N).
+    """
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
 
-    // Shared memory for current page K + V
-    extern __shared__ char smem[];
-    nv_bfloat16* __restrict__ sK = (nv_bfloat16*)smem;
-    nv_bfloat16* __restrict__ sV = sK + page_size * head_dim;
+    pid_hkv = pid_h // GROUP_SIZE
 
-    // Load Q
-    float q[128];
-    const nv_bfloat16* qp = Q + pid_b * stride_qb + h * stride_qh;
-    for (int d = 0; d < head_dim; ++d)
-        q[d] = __bfloat162float(qp[d]);
+    seq_len = tl.load(SeqLens + pid_b)
 
-    // Accumulators
-    float acc_num[128] = {0};
-    float acc_denom = 1e-30f;
-    float acc_max = -1e30f;
+    d_offs = tl.arange(0, HEAD_DIM)
+    n_offs = tl.arange(0, BLOCK_N)
+    n_mask = n_offs < PAGE_SIZE
 
-    for (int pg = 0; pg < num_pages; ++pg) {
-        int block_id = BlockTable[pid_b * stride_bt_b + pg * stride_bt_s];
+    # Load Q
+    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
+    q = tl.load(q_ptr).to(tl.float32)
 
-        // Cooperative load of K/V page
-        const nv_bfloat16* k_page_src = KV + block_id * stride_kvb;
-        const int tokens_to_load = page_size;
-        const int items_thread = (tokens_to_load * head_dim + group_size - 1) / group_size;
+    bt_base = pid_b * stride_bt_b
 
-        for (int i_idx = 0; i_idx < items_thread; ++i_idx) {
-            int flat = tid * items_thread + i_idx;
-            if (flat < tokens_to_load * head_dim) {
-                int row = flat / head_dim;
-                int col = flat % head_dim;
-                // K
-                sK[row * head_dim + col] = k_page_src[row * stride_kvs + col * stride_kvd];
-                // V (contiguous after K in last dim)
-                sV[row * head_dim + col] = k_page_src[row * stride_kvs + head_dim + col * stride_kvd];
-            }
-        }
-        __syncthreads();
+    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
 
-        // Compute attention for tokens in this page
-        int tokens_in_page = (pg + 1) * page_size <= seq_len ? page_size : seq_len - pg * page_size;
-        for (int t = 0; t < tokens_in_page; ++t) {
-            float qk = 0.0f;
-            const nv_bfloat16* kr = sK + t * head_dim;
-            #pragma unroll
-            for (int d = 0; d < 128; ++d) {
-                if (d < head_dim) qk += q[d] * __bfloat162float(kr[d]);
-            }
-            qk *= scale;
+    acc_num = tl.zeros([HEAD_DIM], dtype=tl.float32)
+    acc_denom = 0.0
+    acc_max = -1e30
 
-            float old_max = acc_max;
-            float new_max = fmaxf(old_max, qk);
-            float exp_v = __expf(qk - new_max);
-            float rescale = __expf(old_max - new_max);
+    for page_idx in tl.range(num_pages):
+        page_id = tl.load(BlockTable + bt_base + page_idx)
+        page_base = page_id * stride_kvb + pid_hkv

... (6348 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,13 +1,13 @@
-"""High-performance CUDA paged attention decode kernel.
+"""Optimized Triton paged attention decode kernel.
 
-Strategy:
-- One thread block per (batch, kv_head)
-- Block has num_warps * 32 threads: each warp is one query head in GQA group
-- Load K/V page into shared memory cooperatively
-- Each warp computes QK^T using warp shuffle / register loads
-- Online softmax per warp with per-warp accumulator
+Key optimizations over naive:
+1. Each program handles one (batch, head) but pre-caches page_id at start of each page
+2. Use triton's constexpr PAGE_SIZE loop for full unrolling (16 iterations max)
+3. Reduce memory reads: block_table only readed once per page, not per token
+4. Cache page_base in register (avoids recomputing ptr arithmetic per token)
 """
 import math
+
 import torch
 import torch.nn as nn
 import triton
@@ -18,17 +18,6 @@
     return 1.0 / math.sqrt(head_dim)
 
 
-# The key optimization for paged decode: load full K/V pages into registers/sram,
-# process PAGE_SIZE tokens per page iteration, share loads across GQA group.
-#
-# With PAGE_SIZE=16 and HEAD_DIM=128:
-# - Per page K load: 16 * 128 * 2 = 4096 bytes
-# - Per page V load: 16 * 128 * 2 = 4096 bytes
-# - Process all 16 tokens at once via tl.dot or vectorized loads
-#
-# For head_dim=128, one thread can hold the entire K/V row in registers.
-
-
 @triton.jit
 def _paged_decode_kernel(
     Out,
@@ -37,29 +26,16 @@
     BlockTable,
     SeqLens,
     # strides
-    stride_qb,
-    stride_qh,
-    stride_qd,
-    stride_kvb,
-    stride_kvs,
-    stride_kvh,
-    stride_kvd,
-    stride_ob,
-    stride_oh,
-    stride_od,
+    stride_qb, stride_qh, stride_qd,
+    stride_kvb, stride_kvs, stride_kvh, stride_kvd,
+    stride_ob, stride_oh, stride_od,
     stride_bt_b,
     # constexpr
     HEAD_DIM: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     SCALE: tl.constexpr,
-    BLOCK_N: tl.constexpr,
 ):
-    """Decode kernel with PAGE_SIZE tokens processed per page iteration.
-    
-    Each program handles one (batch, head) pair.
-    Uses tl.dot to compute (1, HEAD_DIM) @ (HEAD_DIM, BLOCK_N) = (1, BLOCK_N).
-    """
     pid_b = tl.program_id(0)
     pid_h = tl.program_id(1)
 
@@ -68,10 +44,8 @@
     seq_len = tl.load(SeqLens + pid_b)
 
     d_offs = tl.arange(0, HEAD_DIM)
-    n_offs = tl.arange(0, BLOCK_N)
-    n_mask = n_offs < PAGE_SIZE
 
-    # Load Q
+    # Load Q: (D,)
     q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
     q = tl.load(q_ptr).to(tl.float32)
 
@@ -83,60 +57,40 @@
     acc_denom = 0.0
     acc_max = -1e30
 
+    # Loop over pages - each iteration loads page_id and processes PAGE_SIZE tokens
     for page_idx in tl.range(num_pages):
+        # Load page_id ONCE per page (not per token)
         page_id = tl.load(BlockTable + bt_base + page_idx)
+        # Cache page_base in register
         page_base = page_id * stride_kvb + pid_hkv * stride_kvh
-        global_start = page_idx * PAGE_SIZE
 
-        # Load K page: (PAGE_SIZE, HEAD_DIM)
-        # k_ptrs: (BLOCK_N, HEAD_DIM)
-        k_base = KV + page_base
-        k_ptrs = k_base + (n_offs[:, None] * stride_kvs + d_offs[None, :] * stride_kvd)
-        k_mask = n_mask[:, None]
-        k_page = tl.load(k_ptrs, mask=k_mask, other=0.0).to(tl.float32)
+        # Unroll PAGE_SIZE tokens
+        for s in tl.range(PAGE_SIZE):
+            global_tok = page_idx * PAGE_SIZE + s
+            valid = global_tok < seq_len
 
-        # Load V page: (PAGE_SIZE, HEAD_DIM)
-        v_base = KV + page_base + HEAD_DIM * stride_kvd
-        v_ptrs = v_base + (n_offs[:, None] * stride_kvs + d_offs[None, :] * stride_kvd)
-        v_page = tl.load(v_ptrs, mask=k_mask, other=0.0).to(tl.float32)
+            # K pointer
+            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
+            k = tl.load(k_ptr, mask=valid, other=0.0).to(tl.float32)
 
-        # Compute QK^T: (1, BLOCK_N) = q[1, D] @ K[D, BLOCK_N].T
-        # Actually q is [D,] and k_page is [BLOCK_N, D]
-        # We need q @ K^T = sum(q[d] * K[n, d], d) for each n
-        scores = tl.sum(q[None, :] * k_page, axis=1) * SCALE  # (BLOCK_N,)
+            # QK dot
+            qk = tl.sum(q * k) * SCALE
 
-        # Mask for tokens beyond seq_len
-        valid_mask = (global_start + n_offs) < seq_len
-        scores = tl.where(valid_mask, scores, -1e30)
+            # V pointer (K and V contiguous in last dim)
+            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
+            v = tl.load(v_ptr, mask=valid, other=0.0).to(tl.float32)
 
-        # Online softmax across the BLOCK_N tokens
-        # Use tl.max and tl.exp for parallel reduction
-        old_max = acc_max
-        new_max = tl.maximum(old_max, scores)  # (BLOCK_N,) per-element max
-
-        # Online softmax must be done sequentially or carefully
-        # We need the running max across all BLOCK_N elements
-        # This is tricky: we can't just do elementwise online softmax
-
-        # Instead: compute max, then compute exp, then sum for this block
-        # Update online softmax:
-        # For the block, we process BLOCK_N tokens. The max across all of them,
-        # then update acc with exp-weighted sum.
-
-        # Simplest correct approach: for small BLOCK_N, do sequential
-        for s in tl.range(PAGE_SIZE):
-            qk = scores[s]
-            if s >= PAGE_SIZE:
-                tl.static_assert(True)  # should not happen, but triton needs this
+            # Online softmax: qk=0 for invalid tokens, exp(0-max) handled
+            # When valid=False: qk=0 from masked load, but we still have nonzero result
+            # Need to set qk=-inf for invalid tokens to avoid affecting softmax
+            qk = tl.where(valid, qk, -1e30)
 
    

... (1666 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,104 +1,395 @@
-"""Optimized Triton paged attention decode kernel.
-
-Key optimizations over naive:
-1. Each program handles one (batch, head) but pre-caches page_id at start of each page
-2. Use triton's constexpr PAGE_SIZE loop for full unrolling (16 iterations max)
-3. Reduce memory reads: block_table only readed once per page, not per token
-4. Cache page_base in register (avoids recomputing ptr arithmetic per token)
-"""
+"""High-performance CUDA C++ paged attention decode kernel."""
 import math
-
 import torch
 import torch.nn as nn
-import triton
-import triton.language as tl
-
 
 def _make_scale(head_dim: int) -> float:
     return 1.0 / math.sqrt(head_dim)
 
-
-@triton.jit
-def _paged_decode_kernel(
-    Out,
-    Q,
-    KV,
-    BlockTable,
-    SeqLens,
-    # strides
-    stride_qb, stride_qh, stride_qd,
-    stride_kvb, stride_kvs, stride_kvh, stride_kvd,
-    stride_ob, stride_oh, stride_od,
-    stride_bt_b,
-    # constexpr
-    HEAD_DIM: tl.constexpr,
-    PAGE_SIZE: tl.constexpr,
-    GROUP_SIZE: tl.constexpr,
-    SCALE: tl.constexpr,
-):
-    pid_b = tl.program_id(0)
-    pid_h = tl.program_id(1)
-
-    pid_hkv = pid_h // GROUP_SIZE
-
-    seq_len = tl.load(SeqLens + pid_b)
-
-    d_offs = tl.arange(0, HEAD_DIM)
-
-    # Load Q: (D,)
-    q_ptr = Q + pid_b * stride_qb + pid_h * stride_qh + d_offs * stride_qd
-    q = tl.load(q_ptr).to(tl.float32)
-
-    bt_base = pid_b * stride_bt_b
-
-    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
-
-    acc_num = tl.zeros([HEAD_DIM], dtype=tl.float32)
-    acc_denom = 0.0
-    acc_max = -1e30
-
-    # Loop over pages - each iteration loads page_id and processes PAGE_SIZE tokens
-    for page_idx in tl.range(num_pages):
-        # Load page_id ONCE per page (not per token)
-        page_id = tl.load(BlockTable + bt_base + page_idx)
-        # Cache page_base in register
-        page_base = page_id * stride_kvb + pid_hkv * stride_kvh
-
-        # Unroll PAGE_SIZE tokens
-        for s in tl.range(PAGE_SIZE):
-            global_tok = page_idx * PAGE_SIZE + s
-            valid = global_tok < seq_len
-
-            # K pointer
-            k_ptr = KV + page_base + s * stride_kvs + d_offs * stride_kvd
-            k = tl.load(k_ptr, mask=valid, other=0.0).to(tl.float32)
-
-            # QK dot
-            qk = tl.sum(q * k) * SCALE
-
-            # V pointer (K and V contiguous in last dim)
-            v_ptr = KV + page_base + s * stride_kvs + (d_offs + HEAD_DIM) * stride_kvd
-            v = tl.load(v_ptr, mask=valid, other=0.0).to(tl.float32)
-
-            # Online softmax: qk=0 for invalid tokens, exp(0-max) handled
-            # When valid=False: qk=0 from masked load, but we still have nonzero result
-            # Need to set qk=-inf for invalid tokens to avoid affecting softmax
-            qk = tl.where(valid, qk, -1e30)
-
-            old_max = acc_max
-            new_max = tl.maximum(old_max, qk)
-            exp_val = tl.exp(qk - new_max)
-            rescale = tl.exp(old_max - new_max)
-
-            acc_num = acc_num * rescale + v * exp_val
-            acc_denom = acc_denom * rescale + exp_val
-            acc_max = new_max
-
-    out = acc_num / acc_denom
-
-    out_ptr = Out + pid_b * stride_ob + pid_h * stride_oh + d_offs * stride_od
-    tl.store(out_ptr, out.to(Out.type.element_ty))
-
+# The CUDA source. One thread block per (batch, kv_head). Each thread handles
+# one query head in the GQA group. Shared memory holds K/V page for current iteration.
+
+_CUDA_SRC = r"""
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <math.h>
+#include <stdio.h>
+
+using bfloat = __nv_bfloat16;
+
+__forceinline__ __device__ float bf16_to_float(bfloat v) {
+    return __bfloat162float(v);
+}
+
+__forceinline__ __device__ bfloat float_to_bfloat(float v) {
+    return __float2bfloat16(v);
+}
+
+__global__ void paged_decode_kernel(
+    bfloat* __restrict__ Out,
+    const bfloat* __restrict__ Q,
+    const bfloat* __restrict__ KV,
+    const int* __restrict__ BlockTable,
+    const int* __restrict__ SeqLens,
+    int B, int H, int Hkv, int D, int P, float scale,
+    int stride_qb, int stride_qh, int stride_qd,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
+    int stride_bt_b, int stride_bt_s,
+    int stride_ob, int stride_oh, int stride_od
+) {
+    // grid: (B, Hkv), block: (group_size)
+    extern __shared__ char smem_raw[];
+    bfloat* __restrict__ smem = (bfloat*)smem_raw;
+
+    int pid_b = blockIdx.x;
+    int pid_hkv = blockIdx.y;
+    int group_size = H / Hkv;
+
+    int h_in_group = threadIdx.x;
+    if (h_in_group >= group_size) return;
+
+    int seq_len = SeqLens[pid_b];
+    int num_pages = (seq_len + P - 1) / P;
+
+    // Shared memory layout:
+    // [0: P*D) = K page
+    // [P*D: 2*P*D) = V page
+    int KV_stride = P * D;
+    bfloat* __restrict__ sK = smem;
+    bfloat* __restrict__ sV = smem + KV_stride;
+
+    // Each thread gets its Q vector
+    float q[128];
+    const bfloat* q_base = Q + pid_b * stride_qb + (pid_hkv * group_size + h_in_group) * stride_qh;
+    for (int d = 0; d < D; d++) {
+        q[d] = bf16_to_float(q_base[d]);
+    }
+
+    // Accumulators per thread
+    float acc_num[128] = {0};
+    float acc_denom = 1e-30f;
+    float acc_max = -1e30f;
+
+    // Process each page
+    for (int pg = 0; pg < num_pages; pg++) {
+        int block_id = BlockTable[pid_b * stride_bt_b + pg * stride_bt_s];
+
+        // Load K page into shared memory cooperatively
+        // Total elements: P * D. Divide among group_size threads.
+        int total_k_items = P * D;
+        int items_per_thread = (total_k_items + group_size - 1) / group_size;
+
+        for (int i = 0; i < items_per_thread; i++) {
+            int idx = h_in_group * items_per_thread + i;
+            if (idx >= total_k_i

... (11741 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -312,7 +312,7 @@
     int stride_ob, int stride_oh, int stride_od
 ) {
     int group_size = H / Hkv;
-    int smem_size = 2 * P * D * sizeoof(bfloat);
+    int smem_size = 2 * P * D * sizeof(bfloat);
 
     dim3 grid(B, Hkv, 1);
     dim3 block(group_size, 1, 1);

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -2,30 +2,39 @@
 import math
 import torch
 import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
 
 def _make_scale(head_dim: int) -> float:
     return 1.0 / math.sqrt(head_dim)
 
-# The CUDA source. One thread block per (batch, kv_head). Each thread handles
-# one query head in the GQA group. Shared memory holds K/V page for current iteration.
-
-_CUDA_SRC = r"""
+_module = None
+
+def _compile():
+    CUDA_SRC = r"""
 #include <cuda_runtime.h>
 #include <cuda_bf16.h>
 #include <math.h>
-#include <stdio.h>
 
 using bfloat = __nv_bfloat16;
 
-__forceinline__ __device__ float bf16_to_float(bfloat v) {
-    return __bfloat162float(v);
-}
-
-__forceinline__ __device__ bfloat float_to_bfloat(float v) {
-    return __float2bfloat16(v);
-}
-
-__global__ void paged_decode_kernel(
+/*
+ * Design: one thread block per (batch, head).
+ * - 128 threads per block (4 warps, each warp handles 32 elements of head_dim)
+ * - Load one K/V page tile into shared memory cooperatively
+ * - Process seq_len tokens, accumulate online softmax
+ * - GQA: all heads sharing a KV head will compete for shared memory bandwidth,
+ *       but each computes independently.
+ *
+ * Shared memory layout:
+ *   smem[0:P*D]         = K page (PAGE_SIZE x HEAD_DIM)
+ *   smem[P*D:2*P*D]     = V page (PAGE_SIZE x HEAD_DIM)
+ *
+ * Each thread handles D/128 elements of the head_dim vector.
+ */
+#define MAX_HEAD_DIM 128
+#define BLOCK_THREADS 128
+
+__global__ void __launch_bounds__(BLOCK_THREADS) paged_decode_kernel(
     bfloat* __restrict__ Out,
     const bfloat* __restrict__ Q,
     const bfloat* __restrict__ KV,
@@ -37,269 +46,125 @@
     int stride_bt_b, int stride_bt_s,
     int stride_ob, int stride_oh, int stride_od
 ) {
-    // grid: (B, Hkv), block: (group_size)
+    int pid_b = blockIdx.x;
+    int pid_h = blockIdx.y;
+    int tid = threadIdx.x;
+
+    int seq_len = SeqLens[pid_b];
+    int num_pages = (seq_len + P - 1) / P;
+
+    // Shared memory
     extern __shared__ char smem_raw[];
     bfloat* __restrict__ smem = (bfloat*)smem_raw;
-
-    int pid_b = blockIdx.x;
-    int pid_hkv = blockIdx.y;
-    int group_size = H / Hkv;
-
-    int h_in_group = threadIdx.x;
-    if (h_in_group >= group_size) return;
-
-    int seq_len = SeqLens[pid_b];
-    int num_pages = (seq_len + P - 1) / P;
-
-    // Shared memory layout:
-    // [0: P*D) = K page
-    // [P*D: 2*P*D) = V page
-    int KV_stride = P * D;
     bfloat* __restrict__ sK = smem;
-    bfloat* __restrict__ sV = smem + KV_stride;
-
-    // Each thread gets its Q vector
-    float q[128];
-    const bfloat* q_base = Q + pid_b * stride_qb + (pid_hkv * group_size + h_in_group) * stride_qh;
-    for (int d = 0; d < D; d++) {
-        q[d] = bf16_to_float(q_base[d]);
+    bfloat* __restrict__ sV = smem + P * D;
+
+    // Each thread owns D/BLOCK_THREADS elements
+    int items_per_thread = (D + BLOCK_THREADS - 1) / BLOCK_THREADS;
+    int start_d = tid * items_per_thread;
+    int end_d = min(start_d + items_per_thread, D);
+    int my_items = end_d - start_d;
+
+    // Load Q elements for this thread's range
+    float q_local[8]; // MAX_HEAD_DIM / BLOCK_THREADS = 1 when 128 threads and D=128
+    if (tid < D) {
+        const bfloat* q_base = Q + pid_b * stride_qb + pid_h * stride_qh;
+        for (int j = 0; j < my_items; j++)
+            q_local[j] = __bfloat162float(q_base[start_d + j]);
     }
 
-    // Accumulators per thread
-    float acc_num[128] = {0};
+    // Accumulators for this thread's range
+    float acc_num_local[8] = {0};
     float acc_denom = 1e-30f;
     float acc_max = -1e30f;
 
     // Process each page
-    for (int pg = 0; pg < num_pages; pg++) {
-        int block_id = BlockTable[pid_b * stride_bt_b + pg * stride_bt_s];
-
-        // Load K page into shared memory cooperatively
-        // Total elements: P * D. Divide among group_size threads.
-        int total_k_items = P * D;
-        int items_per_thread = (total_k_items + group_size - 1) / group_size;
-
-        for (int i = 0; i < items_per_thread; i++) {
-            int idx = h_in_group * items_per_thread + i;
-            if (idx >= total_k_items) break;
-
+    for (int pg_idx = 0; pg_idx < num_pages; pg_idx++) {
+        int page_id = BlockTable[pid_b * stride_bt_b + pg_idx * stride_bt_s];
+
+        // Load K page into smem cooperatively
+        // Each thread loads P*D/BLOCK_THREADS elements
+        int total_kv = P * D;
+        int load_per_thread = (total_kv + BLOCK_THREADS - 1) / BLOCK_THREADS;
+        for (int i = 0; i < load_per_thread; i++) {
+            int idx = tid * load_per_thread + i;
+            if (idx >= total_kv) break;
             int row = idx / D;
             int col = idx % D;
-
-            // K is at KV[block_id, row, pid_hkv, col]
-            // offset = block_id * stride_kvb + row * stride_kvs + pid_hkv * stride_kvh + col * stride_kvd
-            int k_off = block_id * stride_kvb + row * stride_kvs + pid_hkv * stride_kvh + col * stride_kvd;
-            // V is at offset + D (contiguous in last dim)
-            int v_off = k_off + D;
-
-            sK[row * D + col] = KV[k_off];
-            sV[row * D + col] = KV[v_off];
+            int base_off = page_id * stride_kvb + row * stride_kvs + (KV + 0) == 0 ? 0 : 0;
+            int k_off = page_id * stride_kvb + row * stride_kvs + (KV + 0 - KV) * stride_kvd + col * stride_kvd;
+            // Fix: proper offset computation
+            int kv_row_offset = page_id * stride_kvb + row * stride_kvs;
+            int k_idx = kv_row_offset + col * stride_kvd;
+            int v_idx = k_idx + D * stride_kvd;
+            sK[row * D + col] = KV[k_idx];
+            sV[row * D + col] = KV[v_idx];
         }
         __syncthreads();
 
-        // Each thread computes attention for its head
-        // Process ea

... (10907 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,4 +1,12 @@
-"""High-performance CUDA C++ paged attention decode kernel."""
+"""Optimized CUDA C++ paged attention decode kernel.
+
+Design:
+- One thread block per (batch, kv_head)
+- 128 threads cooperate: load K/V page into shared memory, then each thread 
+  computes its query head's attention output
+- Shared memory holds one K/V page (PAGE_SIZE x HEAD_DIM each)
+- All threads share the K/V load, reducing bandwidth
+"""
 import math
 import torch
 import torch.nn as nn
@@ -18,173 +26,121 @@
 using bfloat = __nv_bfloat16;
 
 /*
- * Design: one thread block per (batch, head).
- * - 128 threads per block (4 warps, each warp handles 32 elements of head_dim)
- * - Load one K/V page tile into shared memory cooperatively
- * - Process seq_len tokens, accumulate online softmax
- * - GQA: all heads sharing a KV head will compete for shared memory bandwidth,
- *       but each computes independently.
+ * Block layout: 128 threads per (batch, kv_head).
+ * First GROUP_SIZE threads each handle one query head's output.
+ * All 128 threads cooperate to load K/V into shared memory.
  *
- * Shared memory layout:
- *   smem[0:P*D]         = K page (PAGE_SIZE x HEAD_DIM)
- *   smem[P*D:2*P*D]     = V page (PAGE_SIZE x HEAD_DIM)
- *
- * Each thread handles D/128 elements of the head_dim vector.
+ * Shared memory: sK[PAGE_SIZE x HEAD_DIM], sV[PAGE_SIZE x HEAD_DIM]
+ * Each thread loads a portion of the page.
  */
-#define MAX_HEAD_DIM 128
-#define BLOCK_THREADS 128
-
-__global__ void __launch_bounds__(BLOCK_THREADS) paged_decode_kernel(
+__global__ void __launch_bounds__(128) paged_decode_kernel(
     bfloat* __restrict__ Out,
     const bfloat* __restrict__ Q,
     const bfloat* __restrict__ KV,
     const int* __restrict__ BlockTable,
     const int* __restrict__ SeqLens,
-    int B, int H, int Hkv, int D, int P, float scale,
-    int stride_qb, int stride_qh, int stride_qd,
-    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
-    int stride_bt_b, int stride_bt_s,
-    int stride_ob, int stride_oh, int stride_od
+    int B, int H, int Hkv, int D, int P, float scale
 ) {
+    int group_size = H / Hkv;
+
     int pid_b = blockIdx.x;
-    int pid_h = blockIdx.y;
-    int tid = threadIdx.x;
+    int pid_hkv = blockIdx.y;
+
+    int tid = threadIdx.x;  // 0..127
+    int num_threads = 128;
+
+    // Each thread loads a portion of K/V page
+    // Total elements per page: P * D (K) + P * D (V) = 2 * P * D
+    int total_kv = 2 * P * D;
+    int load_per_thread = (total_kv + num_threads - 1) / num_threads;
+
+    // Shared memory
+    extern __shared__ char smem_raw[];
+    bfloat* __restrict__ sK = (bfloat*)smem_raw;
+    bfloat* __restrict__ sV = sK + P * D;
 
     int seq_len = SeqLens[pid_b];
     int num_pages = (seq_len + P - 1) / P;
 
-    // Shared memory
-    extern __shared__ char smem_raw[];
-    bfloat* __restrict__ smem = (bfloat*)smem_raw;
-    bfloat* __restrict__ sK = smem;
-    bfloat* __restrict__ sV = smem + P * D;
-
-    // Each thread owns D/BLOCK_THREADS elements
-    int items_per_thread = (D + BLOCK_THREADS - 1) / BLOCK_THREADS;
-    int start_d = tid * items_per_thread;
-    int end_d = min(start_d + items_per_thread, D);
-    int my_items = end_d - start_d;
-
-    // Load Q elements for this thread's range
-    float q_local[8]; // MAX_HEAD_DIM / BLOCK_THREADS = 1 when 128 threads and D=128
-    if (tid < D) {
-        const bfloat* q_base = Q + pid_b * stride_qb + pid_h * stride_qh;
-        for (int j = 0; j < my_items; j++)
-            q_local[j] = __bfloat162float(q_base[start_d + j]);
+    // Thread tid < group_size handles query head pid_hkv*group_size+tid
+    bool is_worker = (tid < group_size);
+
+    // Load Q for this thread's head
+    float q_local[128];
+    if (is_worker) {
+        int h = pid_hkv * group_size + tid;
+        const bfloat* q_base = Q + pid_b * H * D + h * D;  // contiguous layout
+        for (int j = 0; j < D; j++)
+            q_local[j] = __bfloat162float(q_base[j]);
     }
 
-    // Accumulators for this thread's range
-    float acc_num_local[8] = {0};
+    // Each thread accumulates for its own head's range
+    float acc_num[128];
+    for (int j = 0; j < 128; j++) acc_num[j] = 0.0f;
     float acc_denom = 1e-30f;
     float acc_max = -1e30f;
 
-    // Process each page
-    for (int pg_idx = 0; pg_idx < num_pages; pg_idx++) {
-        int page_id = BlockTable[pid_b * stride_bt_b + pg_idx * stride_bt_s];
-
-        // Load K page into smem cooperatively
-        // Each thread loads P*D/BLOCK_THREADS elements
-        int total_kv = P * D;
-        int load_per_thread = (total_kv + BLOCK_THREADS - 1) / BLOCK_THREADS;
+    // Process pages
+    for (int pg = 0; pg < num_pages; pg++) {
+        int page_id = BlockTable[pid_b * group_size + pg];  // max_blocks per batch
+
+        // Cooperative load of K/V page into shared memory
         for (int i = 0; i < load_per_thread; i++) {
             int idx = tid * load_per_thread + i;
             if (idx >= total_kv) break;
-            int row = idx / D;
-            int col = idx % D;
-            int base_off = page_id * stride_kvb + row * stride_kvs + (KV + 0) == 0 ? 0 : 0;
-            int k_off = page_id * stride_kvb + row * stride_kvs + (KV + 0 - KV) * stride_kvd + col * stride_kvd;
-            // Fix: proper offset computation
-            int kv_row_offset = page_id * stride_kvb + row * stride_kvs;
-            int k_idx = kv_row_offset + col * stride_kvd;
-            int v_idx = k_idx + D * stride_kvd;
-            sK[row * D + col] = KV[k_idx];
-            sV[row * D + col] = KV[v_idx];
+
+            int slot = idx % D;
+            int tmp = idx / D;
+            int row = tmp % P;
+            int kv_part = tmp / P;  // 0=K, 1=V
+
+            int base_off = page_id * P * Hkv * 2 * D + row * H

... (6484 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,12 +1,4 @@
-"""Optimized CUDA C++ paged attention decode kernel.
-
-Design:
-- One thread block per (batch, kv_head)
-- 128 threads cooperate: load K/V page into shared memory, then each thread 
-  computes its query head's attention output
-- Shared memory holds one K/V page (PAGE_SIZE x HEAD_DIM each)
-- All threads share the K/V load, reducing bandwidth
-"""
+"""Optimized CUDA C++ paged attention decode kernel."""
 import math
 import torch
 import torch.nn as nn
@@ -26,96 +18,96 @@
 using bfloat = __nv_bfloat16;
 
 /*
- * Block layout: 128 threads per (batch, kv_head).
- * First GROUP_SIZE threads each handle one query head's output.
- * All 128 threads cooperate to load K/V into shared memory.
- *
- * Shared memory: sK[PAGE_SIZE x HEAD_DIM], sV[PAGE_SIZE x HEAD_DIM]
- * Each thread loads a portion of the page.
+ * One thread block per (batch, kv_head).
+ * 128 threads: all load K/V page cooperatively, first GROUP_SIZE threads
+ * each compute one query head's output.
+ * 
+ * KV cache: (num_blocks, page_size, num_kv_heads, 2*head_dim) bf16
+ * contiguous in last dim: [K_0,...,K_{D-1}, V_0,...,V_{D-1}]
  */
-__global__ void __launch_bounds__(128) paged_decode_kernel(
+
+__global__ void paged_decode_kernel(
     bfloat* __restrict__ Out,
     const bfloat* __restrict__ Q,
     const bfloat* __restrict__ KV,
     const int* __restrict__ BlockTable,
     const int* __restrict__ SeqLens,
-    int B, int H, int Hkv, int D, int P, float scale
+    int B, int H, int Hkv, int D, int P, float scale,
+    int max_blocks_per_batch,
+    bfloat* __restrict__ smem_KV  // externally allocated shared memory
 ) {
     int group_size = H / Hkv;
-
     int pid_b = blockIdx.x;
     int pid_hkv = blockIdx.y;
-
-    int tid = threadIdx.x;  // 0..127
-    int num_threads = 128;
-
-    // Each thread loads a portion of K/V page
-    // Total elements per page: P * D (K) + P * D (V) = 2 * P * D
+    int tid = threadIdx.x;
+    int num_threads = blockDim.x;
+
+    int seq_len = SeqLens[pid_b];
+    int num_pages = (seq_len + P - 1) / P;
+
+    // Shared memory pointers
+    bfloat* __restrict__ sK = smem_KV;
+    bfloat* __restrict__ sV = smem_KV + P * D;
+
+    // Each thread loads portion of K/V page (combined K+V = 2*P*D elements)
     int total_kv = 2 * P * D;
     int load_per_thread = (total_kv + num_threads - 1) / num_threads;
 
-    // Shared memory
-    extern __shared__ char smem_raw[];
-    bfloat* __restrict__ sK = (bfloat*)smem_raw;
-    bfloat* __restrict__ sV = sK + P * D;
-
-    int seq_len = SeqLens[pid_b];
-    int num_pages = (seq_len + P - 1) / P;
-
-    // Thread tid < group_size handles query head pid_hkv*group_size+tid
+    // Worker threads compute attention (one per query head in GQA group)
     bool is_worker = (tid < group_size);
 
-    // Load Q for this thread's head
-    float q_local[128];
+    // Load Q for worker threads
+    float q_vec[128];
     if (is_worker) {
         int h = pid_hkv * group_size + tid;
-        const bfloat* q_base = Q + pid_b * H * D + h * D;  // contiguous layout
+        const bfloat* q_base = Q + pid_b * stride_qb + h * stride_qh;
         for (int j = 0; j < D; j++)
-            q_local[j] = __bfloat162float(q_base[j]);
-    }
-
-    // Each thread accumulates for its own head's range
+            q_vec[j] = __bfloat162float(q_base[j * stride_qd]);
+    }
+
+    // Local accumulators
     float acc_num[128];
-    for (int j = 0; j < 128; j++) acc_num[j] = 0.0f;
+    for (int j = 0; j < D; j++) acc_num[j] = 0.0f;
     float acc_denom = 1e-30f;
     float acc_max = -1e30f;
 
     // Process pages
     for (int pg = 0; pg < num_pages; pg++) {
-        int page_id = BlockTable[pid_b * group_size + pg];  // max_blocks per batch
-
-        // Cooperative load of K/V page into shared memory
+        int page_id = BlockTable[pid_b * max_blocks_per_batch + pg];
+
+        // Load K/V page cooperatively into shared memory
         for (int i = 0; i < load_per_thread; i++) {
             int idx = tid * load_per_thread + i;
             if (idx >= total_kv) break;
 
-            int slot = idx % D;
+            int col = idx % D;
             int tmp = idx / D;
             int row = tmp % P;
-            int kv_part = tmp / P;  // 0=K, 1=V
-
-            int base_off = page_id * P * Hkv * 2 * D + row * Hkv * 2 * D + pid_hkv * 2 * D;
-            // KV cache: [page_id, row, pid_hkv, slot] for K, [page_id, row, pid_hkv, slot] for V
-
-            if (kv_part == 0) {
-                sK[row * D + slot] = KV[base_off + slot];
-            } else {
-                sV[row * D + slot] = KV[base_off + D + slot];
-            }
-        }
-        __syncthreads();
-
-        // Each processing thread computes QK^T for its head using shared K/V
+            int kv_half = tmp / P;  // 0=K, 1=V
+
+            // KV offset: page_id * stride_kvb + row * stride_kvs + pid_hkv * stride_kvh + col * stride_kvd
+            // K is at col, V is at col + D
+            int base = page_id * stride_kvb + row * stride_kvs + pid_hkv * stride_kvh;
+            int k_idx = base + col * stride_kvd;
+            int v_idx = base + (D + col) * stride_kvd;
+
+            // Store into shared memory: row-major, D columns
+            if (kv_half == 0)
+                sK[row * D + col] = KV[k_idx];
+            else
+                sV[row * D + col] = KV[v_idx];
+        }
+        __syncthreads();
+
         if (is_worker) {
-            int tokens_in_page = min((pg + 1) * P, seq_len) - pg * P;
-            if (tokens_in_page < 0) tokens_in_page = 0;
-
-            for (int t = 0; t < tokens_in_page; t++) {
-                // QK^T
+            int tokens_in = min((pg + 1) * P, seq_len) - pg * P;
+            if (tokens_in < 0) tokens_in = 0;
+
+            for (int t = 0; t < tokens_in; t++) {
                 flo

... (12611 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,4 +1,15 @@
-"""Optimized CUDA C++ paged attention decode kernel."""
+"""Optimized CUDA C++ paged attention decode kernel.
+
+Architecture:
+- One block per (batch, kv_head), 128 threads
+- All 128 threads cooperatively load K/V page into shared memory 
+- First GROUP_SIZE threads compute attention for their query heads
+- Correct batch indexing for BlockTable (stride=1 within each batch's rows)
+
+The key performance insight: decode is memory-bound since KV cache must be streamed.
+By loading each page once into shared memory and sharing across the GQA group,
+we avoid redundant loads that would happen with per-head programs.
+"""
 import math
 import torch
 import torch.nn as nn
@@ -17,274 +28,6 @@
 
 using bfloat = __nv_bfloat16;
 
-/*
- * One thread block per (batch, kv_head).
- * 128 threads: all load K/V page cooperatively, first GROUP_SIZE threads
- * each compute one query head's output.
- * 
- * KV cache: (num_blocks, page_size, num_kv_heads, 2*head_dim) bf16
- * contiguous in last dim: [K_0,...,K_{D-1}, V_0,...,V_{D-1}]
- */
-
-__global__ void paged_decode_kernel(
-    bfloat* __restrict__ Out,
-    const bfloat* __restrict__ Q,
-    const bfloat* __restrict__ KV,
-    const int* __restrict__ BlockTable,
-    const int* __restrict__ SeqLens,
-    int B, int H, int Hkv, int D, int P, float scale,
-    int max_blocks_per_batch,
-    bfloat* __restrict__ smem_KV  // externally allocated shared memory
-) {
-    int group_size = H / Hkv;
-    int pid_b = blockIdx.x;
-    int pid_hkv = blockIdx.y;
-    int tid = threadIdx.x;
-    int num_threads = blockDim.x;
-
-    int seq_len = SeqLens[pid_b];
-    int num_pages = (seq_len + P - 1) / P;
-
-    // Shared memory pointers
-    bfloat* __restrict__ sK = smem_KV;
-    bfloat* __restrict__ sV = smem_KV + P * D;
-
-    // Each thread loads portion of K/V page (combined K+V = 2*P*D elements)
-    int total_kv = 2 * P * D;
-    int load_per_thread = (total_kv + num_threads - 1) / num_threads;
-
-    // Worker threads compute attention (one per query head in GQA group)
-    bool is_worker = (tid < group_size);
-
-    // Load Q for worker threads
-    float q_vec[128];
-    if (is_worker) {
-        int h = pid_hkv * group_size + tid;
-        const bfloat* q_base = Q + pid_b * stride_qb + h * stride_qh;
-        for (int j = 0; j < D; j++)
-            q_vec[j] = __bfloat162float(q_base[j * stride_qd]);
-    }
-
-    // Local accumulators
-    float acc_num[128];
-    for (int j = 0; j < D; j++) acc_num[j] = 0.0f;
-    float acc_denom = 1e-30f;
-    float acc_max = -1e30f;
-
-    // Process pages
-    for (int pg = 0; pg < num_pages; pg++) {
-        int page_id = BlockTable[pid_b * max_blocks_per_batch + pg];
-
-        // Load K/V page cooperatively into shared memory
-        for (int i = 0; i < load_per_thread; i++) {
-            int idx = tid * load_per_thread + i;
-            if (idx >= total_kv) break;
-
-            int col = idx % D;
-            int tmp = idx / D;
-            int row = tmp % P;
-            int kv_half = tmp / P;  // 0=K, 1=V
-
-            // KV offset: page_id * stride_kvb + row * stride_kvs + pid_hkv * stride_kvh + col * stride_kvd
-            // K is at col, V is at col + D
-            int base = page_id * stride_kvb + row * stride_kvs + pid_hkv * stride_kvh;
-            int k_idx = base + col * stride_kvd;
-            int v_idx = base + (D + col) * stride_kvd;
-
-            // Store into shared memory: row-major, D columns
-            if (kv_half == 0)
-                sK[row * D + col] = KV[k_idx];
-            else
-                sV[row * D + col] = KV[v_idx];
-        }
-        __syncthreads();
-
-        if (is_worker) {
-            int tokens_in = min((pg + 1) * P, seq_len) - pg * P;
-            if (tokens_in < 0) tokens_in = 0;
-
-            for (int t = 0; t < tokens_in; t++) {
-                float qk = 0.0f;
-                const bfloat* kr = sK + t * D;
-                for (int j = 0; j < D; j++)
-                    qk += q_vec[j] * __bfloat162float(kr[j]);
-                qk *= scale;
-
-                float old_max = acc_max;
-                float new_max = fmaxf(old_max, qk);
-                float exp_v = __expf(qk - new_max);
-                float rescale = __expf(old_max - new_max);
-
-                const bfloat* vr = sV + t * D;
-                for (int j = 0; j < D; j++)
-                    acc_num[j] = acc_num[j] * rescale + exp_v * __bfloat162float(vr[j]);
-                acc_denom = acc_denom * rescale + exp_v;
-                acc_max = new_max;
-            }
-        }
-        __syncthreads();
-    }
-
-    // Write output
-    if (is_worker) {
-        int h = pid_hkv * group_size + tid;
-        bfloat* out_ptr = Out + pid_b * stride_ob + h * stride_oh;
-        float inv = 1.0f / acc_denom;
-        for (int j = 0; j < D; j++)
-            out_ptr[j * stride_od] = __float2bfloat16(acc_num[j] * inv);
-    }
-}
-
-// We need to pass strides from Python, so we wrap the kernel launch
-extern "C" void launch_paged_decode(
-    bfloat* Out,
-    const bfloat* Q,
-    const bfloat* KV,
-    const int* BlockTable,
-    const int* SeqLens,
-    int B, int H, int Hkv, int D, int P, float scale,
-    int max_blocks_per_batch,
-    int smem_size_bytes,
-    int stride_qb, int stride_qh, int stride_qd,
-    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvd,
-    int stride_ob, int stride_oh, int stride_od
-) {
-    // We can't pass shared memory size dynamically from a C function
-    // without modifying the kernel. Let me just use a kernel template
-    // or pass smem as a kernel parameter.
-    // Actually, the kernel above already uses a separate smem_KV pointer.
-    // We need to allocate it via cudaMalloc or use dynamic shared memory.
-}
-"""
-    # The smem approach requires k

... (13157 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,15 +1,4 @@
-"""Optimized CUDA C++ paged attention decode kernel.
-
-Architecture:
-- One block per (batch, kv_head), 128 threads
-- All 128 threads cooperatively load K/V page into shared memory 
-- First GROUP_SIZE threads compute attention for their query heads
-- Correct batch indexing for BlockTable (stride=1 within each batch's rows)
-
-The key performance insight: decode is memory-bound since KV cache must be streamed.
-By loading each page once into shared memory and sharing across the GQA group,
-we avoid redundant loads that would happen with per-head programs.
-"""
+"""Optimized CUDA C++ paged attention decode kernel."""
 import math
 import torch
 import torch.nn as nn
@@ -37,8 +26,10 @@
     int B, int H, int Hkv, int D, int P, float scale,
     int stride_qb, int stride_qo,
     int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
-    int stride_ob, int stride_oo
+    int stride_ob, int stride_oo,
+    int stride_bt_row
 ) {
+    // Grid: (B, Hkv), Block: (128 threads)
     int BLOCK_THREADS = blockDim.x;
     int GROUP_SIZE = H / Hkv;
 
@@ -55,7 +46,7 @@
 
     bool is_worker = (tid < GROUP_SIZE);
 
-    // Load Q for this thread's head
+    // Load Q
     float q[256];
     if (is_worker) {
         int h = pid_hkv * GROUP_SIZE + tid;
@@ -70,49 +61,191 @@
     float acc_denom = 1e-30f;
     float acc_max = -1e30f;
 
-    // Cooperative K/V load
+    // K/V page load distribution
     int total_items = P * D;
     int per_thread = (total_items + BLOCK_THREADS - 1) / BLOCK_THREADS;
 
-    // BlockTable is (B, max_blocks), row-major. Entry [b, pg] is at:
-    //   BlockTable + b * max_blocks + pg
-    // But we need max_blocks stride from Python.
-    // Since each row is contiguous, stride_bt is 1 within a row.
-    // We need the batch stride. The kernel receives BlockTable already offset to this batch
-    // via the caller passing (BlockTable + pid_b * stride_bt_b).
-    // Actually, let's compute it: BlockTable is (B, max_blocks), so stride_bt_b = max_blocks.
-    // Let's just pass max_blocks from the kernel launch wrapper.
+    for (int pg = 0; pg < num_pages; pg++) {
+        int page_id = BlockTable[pid_b * stride_bt_row + pg];
+
+        // Cooperative load of K/V page into shared memory
+        for (int i = 0; i < per_thread; i++) {
+            int idx = tid * per_thread + i;
+            if (idx >= total_items) break;
+            int row = idx / D;
+            int col = idx % D;
+
+            int base = page_id * stride_kvb + row * stride_kvs + col * stride_kvo;
+            sK[row * D + col] = KV[base];
+            sV[row * D + col] = KV[base + D];
+        }
+        __syncthreads();
+
+        // Compute attention
+        if (is_worker) {
+            int tok_end = min((pg + 1) * P, seq_len);
+            int tok_start = pg * P;
+            for (int t = tok_start; t < tok_end; t++) {
+                int local_t = t - tok_start;
+                float qk = 0.0f;
+                const bfloat* kr = sK + local_t * D;
+                for (int j = 0; j < D; j++)
+                    qk += q[j] * __bfloat162float(kr[j]);
+                qk *= scale;
+
+                float old_max = acc_max;
+                float new_max = fmaxf(old_max, qk);
+                float exp_v = __expf(qk - new_max);
+                float rescale = __expf(old_max - new_max);
+
+                const bfloat* vr = sV + local_t * D;
+                for (int j = 0; j < D; j++)
+                    acc_num[j] = acc_num[j] * rescale + exp_v * __bfloat162float(vr[j]);
+                acc_denom = acc_denom * rescale + exp_v;
+                acc_max = new_max;
+            }
+        }
+        __syncthreads();
+    }
+
+    // Write output
+    if (is_worker) {
+        int h = pid_hkv * GROUP_SIZE + tid;
+        bfloat* ob = Out + pid_b * stride_ob + h * stride_oo;
+        float inv = 1.0f / acc_denom;
+        for (int j = 0; j < D; j++)
+            ob[j] = __float2bfloat16(acc_num[j] * inv);
+    }
+}
+
+// Wrapper: called from PyTorch
+void launch_paged_decode(
+    bfloat* Out, const bfloat* Q, const bfloat* KV,
+    const int* BlockTable, const int* SeqLens,
+    int B, int H, int Hkv, int D, int P, float scale,
+    int stride_bt_row,
+    int stride_qb, int stride_qo,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
+    int stride_ob, int stride_oo
+) {
+    int smem_bytes = 2 * P * D * sizeof(bfloat);
+    dim3 grid(B, Hkv);
+    dim3 block(128, 1, 1);
+    cudaStream_t stream = 0;
+
+    // Use the non-extern-shared version by passing through the wrapper
+    // We need to call the kernel with dynamic smem
+    // In CUDA: kernel<<<grid, block, smem_bytes, stream>>>
     
-    // For now, assume BlockTable passed is already offset to this batch
-    // (caller does: BlockTable + pid_b * stride_bt_b)
-    // Wait, we can't do that easily. Let me rework: use max_blocks from blockTable shape
-    // passed via kernel params.
+    // Since the kernel uses extern __shared__ and we launch from here,
+    // we can pass the smem size:
+    paged_decode_kernel<<<grid, block, smem_bytes, stream>>>(
+        Out, Q, KV, BlockTable, SeqLens,
+        B, H, Hkv, D, P, scale,
+        stride_qb, stride_qo,
+        stride_kvb, stride_kvs, stride_kvh, stride_kvo,
+        stride_ob, stride_oo, stride_bt_row);
+}
+"""
+
+    CPP_SRC = r"""
+#include <torch/extension.h>
+#include <cuda_bf16.h>
+
+void launch_paged_decode(
+    __nv_bfloat16* Out, const __nv_bfloat16* Q, const __nv_bfloat16* KV,
+    const int* BlockTable, const int* SeqLens,
+    int B, int H, int Hkv, int D, int P, float scale,
+    int stride_bt_row,
+    int stride_qb, int stride_qo,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
+    int stride_ob, int stride_oo);
+
+void paged_attention_

... (4398 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,4 +1,11 @@
-"""Optimized CUDA C++ paged attention decode kernel."""
+"""Paged attention decode kernel - per-head design optimized for bandwidth.
+
+Architecture: one block per (batch, query_head), 128 threads.
+- Each thread loads and computes a portion of the attention (D / BLOCK_DIM elements)
+- No shared memory - each thread loads its own K/V from global memory
+- Use __ldg (read-only data cache) for K/V loads
+- Use vectorized loads (float4 = 8 bf16) where possible
+"""
 import math
 import torch
 import torch.nn as nn
@@ -17,7 +24,17 @@
 
 using bfloat = __nv_bfloat16;
 
-__global__ void paged_decode_kernel(
+/*
+ * Design: one thread block per (batch, head), 128 threads.
+ * Each thread handles D/128 elements (1 element when D=128).
+ * No shared memory: each thread loads from global and accumulates.
+ * 
+ * For D=128, each thread handles exactly 1 element of the 128-dim vector.
+ * Thread d loads KV[d], computes partial QK, and accumulates acc_num[d].
+ * We need a reduction across threads for qk.
+ */
+
+__global__ void __launch_bounds__(128) paged_decode_kernel(
     bfloat* __restrict__ Out,
     const bfloat* __restrict__ Q,
     const bfloat* __restrict__ KV,
@@ -29,122 +46,85 @@
     int stride_ob, int stride_oo,
     int stride_bt_row
 ) {
-    // Grid: (B, Hkv), Block: (128 threads)
-    int BLOCK_THREADS = blockDim.x;
-    int GROUP_SIZE = H / Hkv;
-
+    int BLOCK_DIM = 128;
     int pid_b = blockIdx.x;
-    int pid_hkv = blockIdx.y;
+    int pid_h = blockIdx.y;
     int tid = threadIdx.x;
 
+    int pid_hkv = pid_h / (H / Hkv);
     int seq_len = SeqLens[pid_b];
     int num_pages = (seq_len + P - 1) / P;
 
-    extern __shared__ bfloat smem[];
-    bfloat* __restrict__ sK = smem;
-    bfloat* __restrict__ sV = smem + P * D;
-
-    bool is_worker = (tid < GROUP_SIZE);
-
-    // Load Q
-    float q[256];
-    if (is_worker) {
-        int h = pid_hkv * GROUP_SIZE + tid;
-        const bfloat* qb = Q + pid_b * stride_qb + h * stride_qo;
-        for (int j = 0; j < D; j++)
-            q[j] = __bfloat162float(qb[j]);
-    }
+    // Each thread works on element tid of the D-dimensional vector
+    // For D < 128, some threads are idle
+    if (tid >= D) return;
+
+    // Load this thread's Q element
+    float q_val = __bfloat162float(Q[pid_b * stride_qb + pid_h * stride_qo + tid]);
 
     // Accumulators
-    float acc_num[256];
-    for (int j = 0; j < D; j++) acc_num[j] = 0.0f;
+    float acc_num_val = 0.0f;
     float acc_denom = 1e-30f;
     float acc_max = -1e30f;
 
-    // K/V page load distribution
-    int total_items = P * D;
-    int per_thread = (total_items + BLOCK_THREADS - 1) / BLOCK_THREADS;
-
+    // Process pages and tokens
     for (int pg = 0; pg < num_pages; pg++) {
         int page_id = BlockTable[pid_b * stride_bt_row + pg];
 
-        // Cooperative load of K/V page into shared memory
-        for (int i = 0; i < per_thread; i++) {
-            int idx = tid * per_thread + i;
-            if (idx >= total_items) break;
-            int row = idx / D;
-            int col = idx % D;
-
-            int base = page_id * stride_kvb + row * stride_kvs + col * stride_kvo;
-            sK[row * D + col] = KV[base];
-            sV[row * D + col] = KV[base + D];
+        int tok_end = min((pg + 1) * P, seq_len);
+        int tok_start = pg * P;
+
+        for (int t = tok_start; t < tok_end; t++) {
+            int local_t = t - tok_start;
+
+            // Load K[tid] for this token
+            int k_base = page_id * stride_kvb + local_t * stride_kvs + pid_hkv * stride_kvh;
+            float k_val = __bfloat162float(KV[k_base + tid * stride_kvo]);
+
+            // Each thread computes its partial Q * K contribution
+            float partial = q_val * k_val;
+
+            // Warp-wide reduction to compute full qk dot product
+            // Use __shfl_xor for butterfly reduction across 128 threads
+            // 128 threads = 4 warps (32 each)
+
+            // First: reduce within warp
+            #pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1)
+                partial += __shfl_xor_sync(0xffffffff, partial, offset);
+
+            // Now each warp has partial sum of its 32 elements
+            // Warp 0 holds sum of elements [0..31], warp 1: [32..63], etc.
+            // Need to get all partial sums to all threads
+
+            __shared__ float warp_sum[4];  // 128/32 = 4 warps max
+            if (tid % 32 == 0)
+                warp_sum[tid / 32] = partial;
+            __syncthreads();
+
+            float qk = 0.0f;
+            for (int w = 0; w < (D + 31) / 32 && w * 32 < D; w++)
+                qk += warp_sum[w];
+            qk *= scale;
+
+            // Now each thread loads V[tid] and does its softmax update
+            float v_val = __bfloat162float(KV[k_base + D + tid * stride_kvo]);
+
+            // All threads need consistent qk, max, exp, rescale
+            float old_max = acc_max;
+            float new_max = fmaxf(old_max, qk);
+            float exp_v = __expf(qk - new_max);
+            float rescale = __expf(old_max - new_max);
+
+            // Update: acc_num[tid] += exp_v * v_val, acc_denom += exp_v
+            acc_num_val = acc_num_val * rescale + exp_v * v_val;
+            acc_denom = acc_denom * rescale + exp_v;
+            acc_max = new_max;
         }
-        __syncthreads();
-
-        // Compute attention
-        if (is_worker) {
-            int tok_end = min((pg + 1) * P, seq_len);
-            int tok_start = pg * P;
-            for (int t = tok_start; t < tok_end; t++) {
-                int local_t = t - tok_start;
-                float qk = 0.0f;
-                const bfloat* kr = sK + local_t * D;
-                for (int j = 0; j < D; j++)
-                    qk += q[j] * __bfloat162float(kr[

... (3838 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -133,13 +133,23 @@
 #include <cuda_bf16.h>
 
 void launch_paged_decode(
-    __nv_bfloat16* Out, const __nv_bfloat16* Q, const __nv_bfloat16* KV,
+    bfloat* Out, const bfloat* Q, const bfloat* KV,
     const int* BlockTable, const int* SeqLens,
     int B, int H, int Hkv, int D, int P, float scale,
     int stride_bt_row,
     int stride_qb, int stride_qo,
     int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
-    int stride_ob, int stride_oo);
+    int stride_ob, int stride_oo
+) {
+    dim3 grid(B, H);
+    dim3 block(128, 1, 1);
+    paged_decode_kernel<<<grid, block>>>(
+        Out, Q, KV, BlockTable, SeqLens,
+        B, H, Hkv, D, P, scale,
+        stride_qb, stride_qo,
+        stride_kvb, stride_kvs, stride_kvh, stride_kvo,
+        stride_ob, stride_oo, stride_bt_row);
+}
 
 void paged_attention_forward(
     torch::Tensor out,

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -132,8 +132,18 @@
 #include <torch/extension.h>
 #include <cuda_bf16.h>
 
+// Kernel declaration (defined in .cu)
+__global__ void paged_decode_kernel(
+    __nv_bfloat16* Out, const __nv_bfloat16* Q, const __nv_bfloat16* KV,
+    const int* BlockTable, const int* SeqLens,
+    int B, int H, int Hkv, int D, int P, float scale,
+    int stride_qb, int stride_qo,
+    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
+    int stride_ob, int stride_oo,
+    int stride_bt_row);
+
 void launch_paged_decode(
-    bfloat* Out, const bfloat* Q, const bfloat* KV,
+    __nv_bfloat16* Out, const __nv_bfloat16* Q, const __nv_bfloat16* KV,
     const int* BlockTable, const int* SeqLens,
     int B, int H, int Hkv, int D, int P, float scale,
     int stride_bt_row,

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/03_paged_attention/solution.py
@@ -1,10 +1,7 @@
-"""Paged attention decode kernel - per-head design optimized for bandwidth.
+"""Per-head CUDA paged attention decode kernel.
 
-Architecture: one block per (batch, query_head), 128 threads.
-- Each thread loads and computes a portion of the attention (D / BLOCK_DIM elements)
-- No shared memory - each thread loads its own K/V from global memory
-- Use __ldg (read-only data cache) for K/V loads
-- Use vectorized loads (float4 = 8 bf16) where possible
+Architecture: one thread block per (batch, query_head), 128 threads.
+Each thread handles one element of D, reduces qk across threads.
 """
 import math
 import torch
@@ -17,22 +14,15 @@
 _module = None
 
 def _compile():
+    # Put ALL the code in the CUDA source. No C++ wrapper needed.
+    # We define the kernel AND the C++ interface function in the .cu file.
     CUDA_SRC = r"""
 #include <cuda_runtime.h>
 #include <cuda_bf16.h>
 #include <math.h>
+#include <torch/extension.h>
 
 using bfloat = __nv_bfloat16;
-
-/*
- * Design: one thread block per (batch, head), 128 threads.
- * Each thread handles D/128 elements (1 element when D=128).
- * No shared memory: each thread loads from global and accumulates.
- * 
- * For D=128, each thread handles exactly 1 element of the 128-dim vector.
- * Thread d loads KV[d], computes partial QK, and accumulates acc_num[d].
- * We need a reduction across threads for qk.
- */
 
 __global__ void __launch_bounds__(128) paged_decode_kernel(
     bfloat* __restrict__ Out,
@@ -46,7 +36,6 @@
     int stride_ob, int stride_oo,
     int stride_bt_row
 ) {
-    int BLOCK_DIM = 128;
     int pid_b = blockIdx.x;
     int pid_h = blockIdx.y;
     int tid = threadIdx.x;
@@ -55,11 +44,10 @@
     int seq_len = SeqLens[pid_b];
     int num_pages = (seq_len + P - 1) / P;
 
-    // Each thread works on element tid of the D-dimensional vector
-    // For D < 128, some threads are idle
+    // Early exit
     if (tid >= D) return;
 
-    // Load this thread's Q element
+    // Load Q element
     float q_val = __bfloat162float(Q[pid_b * stride_qb + pid_h * stride_qo + tid]);
 
     // Accumulators
@@ -67,7 +55,10 @@
     float acc_denom = 1e-30f;
     float acc_max = -1e30f;
 
-    // Process pages and tokens
+    // Warp count
+    int num_warps = (D + 31) / 32;
+    int my_warp = tid / 32;
+
     for (int pg = 0; pg < num_pages; pg++) {
         int page_id = BlockTable[pid_b * stride_bt_row + pg];
 
@@ -77,46 +68,39 @@
         for (int t = tok_start; t < tok_end; t++) {
             int local_t = t - tok_start;
 
-            // Load K[tid] for this token
+            // Load K[tid]
             int k_base = page_id * stride_kvb + local_t * stride_kvs + pid_hkv * stride_kvh;
             float k_val = __bfloat162float(KV[k_base + tid * stride_kvo]);
 
-            // Each thread computes its partial Q * K contribution
+            // Partial dot product
             float partial = q_val * k_val;
 
-            // Warp-wide reduction to compute full qk dot product
-            // Use __shfl_xor for butterfly reduction across 128 threads
-            // 128 threads = 4 warps (32 each)
-
-            // First: reduce within warp
+            // Intra-warp reduction
             #pragma unroll
             for (int offset = 16; offset > 0; offset >>= 1)
                 partial += __shfl_xor_sync(0xffffffff, partial, offset);
 
-            // Now each warp has partial sum of its 32 elements
-            // Warp 0 holds sum of elements [0..31], warp 1: [32..63], etc.
-            // Need to get all partial sums to all threads
-
-            __shared__ float warp_sum[4];  // 128/32 = 4 warps max
+            // Each warp leader writes its sum to shared memory
+            __shared__ float warp_sum[4];
             if (tid % 32 == 0)
                 warp_sum[tid / 32] = partial;
             __syncthreads();
 
+            // Compute full qk reduction
             float qk = 0.0f;
-            for (int w = 0; w < (D + 31) / 32 && w * 32 < D; w++)
+            for (int w = 0; w < num_warps; w++)
                 qk += warp_sum[w];
             qk *= scale;
 
-            // Now each thread loads V[tid] and does its softmax update
+            // Load V[tid]
             float v_val = __bfloat162float(KV[k_base + D + tid * stride_kvo]);
 
-            // All threads need consistent qk, max, exp, rescale
+            // Online softmax (consistent across all threads)
             float old_max = acc_max;
             float new_max = fmaxf(old_max, qk);
             float exp_v = __expf(qk - new_max);
             float rescale = __expf(old_max - new_max);
 
-            // Update: acc_num[tid] += exp_v * v_val, acc_denom += exp_v
             acc_num_val = acc_num_val * rescale + exp_v * v_val;
             acc_denom = acc_denom * rescale + exp_v;
             acc_max = new_max;
@@ -126,41 +110,8 @@
     // Write output
     Out[pid_b * stride_ob + pid_h * stride_oo + tid] = __float2bfloat16(acc_num_val / acc_denom);
 }
-"""
 
-    CPP_SRC = r"""
-#include <torch/extension.h>
-#include <cuda_bf16.h>
-
-// Kernel declaration (defined in .cu)
-__global__ void paged_decode_kernel(
-    __nv_bfloat16* Out, const __nv_bfloat16* Q, const __nv_bfloat16* KV,
-    const int* BlockTable, const int* SeqLens,
-    int B, int H, int Hkv, int D, int P, float scale,
-    int stride_qb, int stride_qo,
-    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
-    int stride_ob, int stride_oo,
-    int stride_bt_row);
-
-void launch_paged_decode(
-    __nv_bfloat16* Out, const __nv_bfloat16* Q, const __nv_bfloat16* KV,
-    const int* BlockTable, const int* SeqLens,
-    int B, int H, int Hkv, int D, int P, float scale,
-    int stride_bt_row,
-    int stride_qb, int stride_qo,
-    int stride_kvb, int stride_kvs, int stride_kvh, int stride_kvo,
-    i

... (1591 more chars)