--- /dev/null
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -0,0 +1,372 @@
+"""Paged attention decode kernel for RTX PRO 6000 (SM120 Blackwell).
+
+CUDA C++ kernel with online softmax, shared-memory tiling, and vectorized loads.
+"""
+
+import math
+from typing import List
+
+import torch
+import torch.nn as nn
+
+OP_TYPE = "attention"
+SUPPORTED_PRECISIONS = ["bf16"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+BATCH = 8
+NUM_HEADS = 32
+NUM_KV_HEADS = 8
+HEAD_DIM = 128
+SEQ_LEN = 1024
+PAGE_SIZE = 16
+
+# ---------------------------------------------------------------------------
+# CUDA kernel
+# ---------------------------------------------------------------------------
+
+_cuda_kernel_src = r"""
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+
+constexpr int PAGE_SIZE = 16;
+constexpr int MAX_D = 128;
+constexpr int WARP_SIZE = 32;
+constexpr int MAX_GROUP_SIZE = 8;  // max num_heads/num_kv_heads across shapes
+
+extern "C" __global__ void paged_attention_decode_kernel(
+    const __nv_bfloat16* __restrict__ query,       // (B, H, D)
+    const __nv_bfloat16* __restrict__ kv_cache,    // (num_blocks, P, Hkv, 2*D)
+    const int* __restrict__ block_table,           // (B, max_blocks)
+    const int* __restrict__ seq_lens,              // (B,)
+    __nv_bfloat16* __restrict__ output,            // (B, H, D)
+    int B,
+    int H,
+    int Hkv,
+    int D,
+    int max_blocks,
+    int kv_cache_stride_block,   // stride of block dim in bf16 elements
+    int kv_cache_stride_pos,     // stride of position dim
+    int kv_cache_stride_head     // stride of kv_head dim
+) {
+    int b = blockIdx.y;
+    int kv_head = blockIdx.x;
+    int group_size = H / Hkv;
+    int seq_len = seq_lens[b];
+    if (seq_len <= 0) return;
+
+    int num_pages = (seq_len + PAGE_SIZE - 1) / PAGE_SIZE;
+
+    __shared__ __nv_bfloat16 K_smem[PAGE_SIZE * MAX_D];
+    __shared__ __nv_bfloat16 V_smem[PAGE_SIZE * MAX_D];
+
+    int warp_id = threadIdx.x / WARP_SIZE;
+    int lane_id = threadIdx.x % WARP_SIZE;
+    int tid = threadIdx.x;
+    int num_threads = blockDim.x;
+
+    // Map each warp to a Q head in the group.
+    // With up to 8 warps (256 threads) and max group_size=8, each Q head
+    // gets one dedicated warp.  If group_size < 8, extra warps idle.
+    int q_head_in_group = warp_id;
+    bool active = (q_head_in_group < group_size);
+    int q_head = kv_head * group_size + q_head_in_group;
+
+    // --- Pre-load Q for this warp's head into registers ---
+    float q_reg[MAX_D / WARP_SIZE];
+    if (active) {
+        #pragma unroll
+        for (int i = 0; i < D / WARP_SIZE; i++) {
+            int d = lane_id + i * WARP_SIZE;
+            if (d < D) {
+                q_reg[i] = __bfloat162float(query[b * H * D + q_head * D + d]);
+            } else {
+                q_reg[i] = 0.0f;
+            }
+        }
+    }
+
+    // --- Online-softmax state (fp32, per thread's stripe of D) ---
+    float m_global = -1e30f;
+    float l_global = 0.0f;
+    float acc[MAX_D / WARP_SIZE];
+    if (active) {
+        #pragma unroll
+        for (int i = 0; i < D / WARP_SIZE; i++) {
+            acc[i] = 0.0f;
+        }
+    }
+
+    float scale = rsqrtf(static_cast<float>(D));
+
+    // --- Page loop ---
+    for (int page_idx = 0; page_idx < num_pages; page_idx++) {
+        int block_idx = block_table[b * max_blocks + page_idx];
+        int tokens_this_page;
+        if (page_idx == num_pages - 1) {
+            int rem = seq_len % PAGE_SIZE;
+            tokens_this_page = (rem == 0) ? PAGE_SIZE : rem;
+        } else {
+            tokens_this_page = PAGE_SIZE;
+        }
+
+        // ----- Cooperative load of K and V into shared memory -----
+        // All threads (including idle warps) participate so we fill smem
+        // quickly with coalesced accesses.
+        int total_k_elems = PAGE_SIZE * D;
+        int elems_per_thread = (total_k_elems + num_threads - 1) / num_threads;
+
+        // Load K
+        for (int i = 0; i < elems_per_thread; i++) {
+            int idx = tid + i * num_threads;
+            if (idx < total_k_elems) {
+                int pos = idx / D;
+                int d   = idx % D;
+                int gmem_offset = block_idx * kv_cache_stride_block
+                                + pos * kv_cache_stride_pos
+                                + kv_head * kv_cache_stride_head
+                                + d;
+                K_smem[idx] = kv_cache[gmem_offset];
+            }
+        }
+
+        // Load V
+        for (int i = 0; i < elems_per_thread; i++) {
+            int idx = tid + i * num_threads;
+            if (idx < total_k_elems) {
+                int pos = idx / D;
+                int d   = idx % D;
+                int gmem_offset = block_idx * kv_cache_stride_block
+                                + pos * kv_cache_stride_pos
+                                + kv_head * kv_cache_stride_head
+                                + D + d;
+                V_smem[idx] = kv_cache[gmem_offset];
+            }
+        }
+
+        __syncthreads();
+
+        // ----- Per-warp attention computation -----
+        if (active) {
+            // Compute scores[t] = dot(Q, K[t]) * scale for every valid
+            // position in this tile, using warp shuffles for the reduction.
+            float scores[PAGE_SIZE];
+
+            float m_local = m_global;
+            float max_score = m_global;
+
+            #pragma unroll
+            for (int t = 0; t < PAGE_SIZE; t++) {
+                float dot = 0.0f;
+                if (t < tokens_this_page) {
+                    #pragma unroll
+                    for (int i = 0; i < D / WARP_SIZE; i++) {
+                        int d = lane_id + i * WARP_SIZE;
+                        if (d < D) {
+                            dot += q_reg[i] * __bfloat162float(K_smem[t * D + d]);
+                        }
+                    }
+                }
+                // Full-warp reduction for

... (6117 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -33,7 +33,7 @@
 constexpr int WARP_SIZE = 32;
 constexpr int MAX_GROUP_SIZE = 8;  // max num_heads/num_kv_heads across shapes
 
-extern "C" __global__ void paged_attention_decode_kernel(
+__global__ void paged_attention_decode_kernel(
     const __nv_bfloat16* __restrict__ query,       // (B, H, D)
     const __nv_bfloat16* __restrict__ kv_cache,    // (num_blocks, P, Hkv, 2*D)
     const int* __restrict__ block_table,           // (B, max_blocks)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,13 +1,14 @@
 """Paged attention decode kernel for RTX PRO 6000 (SM120 Blackwell).
 
-CUDA C++ kernel with online softmax, shared-memory tiling, and vectorized loads.
+Triton-based decode kernel with online softmax and paged KV-cache gathering.
 """
 
 import math
-from typing import List
 
 import torch
 import torch.nn as nn
+import triton
+import triton.language as tl
 
 OP_TYPE = "attention"
 SUPPORTED_PRECISIONS = ["bf16"]
@@ -21,240 +22,194 @@
 PAGE_SIZE = 16
 
 # ---------------------------------------------------------------------------
-# CUDA kernel
-# ---------------------------------------------------------------------------
-
-_cuda_kernel_src = r"""
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
-
-constexpr int PAGE_SIZE = 16;
-constexpr int MAX_D = 128;
-constexpr int WARP_SIZE = 32;
-constexpr int MAX_GROUP_SIZE = 8;  // max num_heads/num_kv_heads across shapes
-
-__global__ void paged_attention_decode_kernel(
-    const __nv_bfloat16* __restrict__ query,       // (B, H, D)
-    const __nv_bfloat16* __restrict__ kv_cache,    // (num_blocks, P, Hkv, 2*D)
-    const int* __restrict__ block_table,           // (B, max_blocks)
-    const int* __restrict__ seq_lens,              // (B,)
-    __nv_bfloat16* __restrict__ output,            // (B, H, D)
-    int B,
-    int H,
-    int Hkv,
-    int D,
-    int max_blocks,
-    int kv_cache_stride_block,   // stride of block dim in bf16 elements
-    int kv_cache_stride_pos,     // stride of position dim
-    int kv_cache_stride_head     // stride of kv_head dim
-) {
-    int b = blockIdx.y;
-    int kv_head = blockIdx.x;
-    int group_size = H / Hkv;
-    int seq_len = seq_lens[b];
-    if (seq_len <= 0) return;
-
-    int num_pages = (seq_len + PAGE_SIZE - 1) / PAGE_SIZE;
-
-    __shared__ __nv_bfloat16 K_smem[PAGE_SIZE * MAX_D];
-    __shared__ __nv_bfloat16 V_smem[PAGE_SIZE * MAX_D];
-
-    int warp_id = threadIdx.x / WARP_SIZE;
-    int lane_id = threadIdx.x % WARP_SIZE;
-    int tid = threadIdx.x;
-    int num_threads = blockDim.x;
-
-    // Map each warp to a Q head in the group.
-    // With up to 8 warps (256 threads) and max group_size=8, each Q head
-    // gets one dedicated warp.  If group_size < 8, extra warps idle.
-    int q_head_in_group = warp_id;
-    bool active = (q_head_in_group < group_size);
-    int q_head = kv_head * group_size + q_head_in_group;
-
-    // --- Pre-load Q for this warp's head into registers ---
-    float q_reg[MAX_D / WARP_SIZE];
-    if (active) {
-        #pragma unroll
-        for (int i = 0; i < D / WARP_SIZE; i++) {
-            int d = lane_id + i * WARP_SIZE;
-            if (d < D) {
-                q_reg[i] = __bfloat162float(query[b * H * D + q_head * D + d]);
-            } else {
-                q_reg[i] = 0.0f;
-            }
-        }
-    }
-
-    // --- Online-softmax state (fp32, per thread's stripe of D) ---
-    float m_global = -1e30f;
-    float l_global = 0.0f;
-    float acc[MAX_D / WARP_SIZE];
-    if (active) {
-        #pragma unroll
-        for (int i = 0; i < D / WARP_SIZE; i++) {
-            acc[i] = 0.0f;
-        }
-    }
-
-    float scale = rsqrtf(static_cast<float>(D));
-
-    // --- Page loop ---
-    for (int page_idx = 0; page_idx < num_pages; page_idx++) {
-        int block_idx = block_table[b * max_blocks + page_idx];
-        int tokens_this_page;
-        if (page_idx == num_pages - 1) {
-            int rem = seq_len % PAGE_SIZE;
-            tokens_this_page = (rem == 0) ? PAGE_SIZE : rem;
-        } else {
-            tokens_this_page = PAGE_SIZE;
-        }
-
-        // ----- Cooperative load of K and V into shared memory -----
-        // All threads (including idle warps) participate so we fill smem
-        // quickly with coalesced accesses.
-        int total_k_elems = PAGE_SIZE * D;
-        int elems_per_thread = (total_k_elems + num_threads - 1) / num_threads;
-
-        // Load K
-        for (int i = 0; i < elems_per_thread; i++) {
-            int idx = tid + i * num_threads;
-            if (idx < total_k_elems) {
-                int pos = idx / D;
-                int d   = idx % D;
-                int gmem_offset = block_idx * kv_cache_stride_block
-                                + pos * kv_cache_stride_pos
-                                + kv_head * kv_cache_stride_head
-                                + d;
-                K_smem[idx] = kv_cache[gmem_offset];
-            }
-        }
-
-        // Load V
-        for (int i = 0; i < elems_per_thread; i++) {
-            int idx = tid + i * num_threads;
-            if (idx < total_k_elems) {
-                int pos = idx / D;
-                int d   = idx % D;
-                int gmem_offset = block_idx * kv_cache_stride_block
-                                + pos * kv_cache_stride_pos
-                                + kv_head * kv_cache_stride_head
-                                + D + d;
-                V_smem[idx] = kv_cache[gmem_offset];
-            }
-        }
-
-        __syncthreads();
-
-        // ----- Per-warp attention computation -----
-        if (active) {
-            // Compute scores[t] = dot(Q, K[t]) * scale for every valid
-            // position in this tile, using warp shuffles for the reduction.
-            float scores[PAGE_SIZE];
-
-            float m_local = m_global;
-            float max_score = m_global;
-
-            #pragma unroll
-            for (int t = 0; t < PAGE_SIZE; t++) {
-                float dot = 0.0f;
-                if (t < tokens_this_page) {
-                    #pragma unroll
-                    for (int i = 0; i < D / WARP_SIZE; i++) {
-                        int d = lane_id + i * WARP_SIZE;
-                        if (d < D) {
-                            dot += q_reg[i] * __bfloat162float(K_smem[t * D + d]);
-                        }
-                    }
-                }

... (12836 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -25,8 +25,7 @@
 # Triton kernel
 # ---------------------------------------------------------------------------
 
-# Autotune over tile sizes along the sequence-length dimension.
-# Decode is memory-bound so the primary goal is saturating DRAM bandwidth.
+
 @triton.autotune(
     configs=[
         triton.Config({"BLOCK_L": 16, "BLOCK_D": 64}, num_warps=4),
@@ -54,129 +53,105 @@
     group_size: tl.constexpr,
 ):
     """
-    Each program instance handles one (batch_element, kv_head) pair and
-    iterates over the sequence in BLOCK_L-sized tiles.
-
-    Within each tile, Q @ K^T is computed as a dot-product between the query
-    rows of the group and the K tile, using online softmax rescaling, then
-    the V tile is accumulated.
+    Each program handles one (batch, kv_head) pair.
+
+    Loads all Q heads in the GQA group together in one shot, then iterates
+    over KV-cache pages in BLOCK_L-sized tiles.  Online softmax (fp32)
+    rescaling keeps accumulator state per Q head.
     """
-    pid_b = tl.program_id(0)  # batch index
-    pid_kv = tl.program_id(1)  # KV head index
-
-    # Number of programs along the sequence-length dimension
+    pid_b = tl.program_id(0)
+    pid_kv = tl.program_id(1)
+
     seq_len = tl.load(seq_lens_ptr + pid_b)
     num_pages = (seq_len + P - 1) // P
 
-    # Pre-compute offsets that are constant across the tile loop
-    offs_d = tl.arange(0, BLOCK_D)  # dimension offsets (may exceed D)
-    offs_l = tl.arange(0, BLOCK_L)   # length offsets within a tile
-
-    # Online softmax state — one per Q head in the group (group_size rows)
-    # We store them as a [group_size, BLOCK_D] accumulator and per-row m/l.
-    # Use fp32 for accumulation.
+    offs_d = tl.arange(0, BLOCK_D)
+    offs_l = tl.arange(0, BLOCK_L)
+    offs_g = tl.arange(0, group_size)  # Q-heads within the GQA group
+
+    # --- Load all Q heads of this group at once ---
+    q_head_base = pid_kv * group_size
+    q_offs = (
+        pid_b * stride_q_b
+        + (q_head_base + offs_g[:, None]) * stride_q_h
+        + offs_d[None, :]
+    )
+    mask_q = offs_d[None, :] < D
+    q = tl.load(query_ptr + q_offs, mask=mask_q, other=0.0).to(tl.float32)
+    # q: (group_size, BLOCK_D)
+
+    # --- Online-softmax state (fp32) ---
     m = tl.full([group_size], float("-inf"), dtype=tl.float32)
     l_sum = tl.zeros([group_size], dtype=tl.float32)
-
-    # Accumulator: (group_size, BLOCK_D) — padded to BLOCK_D
     acc = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
 
     scale = 1.0 / tl.sqrt(D.to(tl.float32))
 
-    # Load Q for all heads in this group, shape (group_size, D)
-    # q_head = kv_head * group_size + g
-    kv_head = pid_kv
-    q_heads_base = kv_head * group_size * D  # offset within a batch's Q
-
-    # Load Q rows: (group_size, D) padded to BLOCK_D
-    q = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
-    for g in range(group_size):
-        q_head_idx = kv_head * group_size + g
-        q_offs = pid_b * stride_q_b + q_head_idx * stride_q_h + offs_d
-        mask_d = offs_d < D
-        q_g = tl.load(query_ptr + q_offs, mask=mask_d, other=0.0)
-        q[g, :] = q_g.to(tl.float32)
-
-    # Page loop
+    # --- Page loop ---
     for page_idx in range(num_pages):
         blk_idx = tl.load(block_table_ptr + pid_b * stride_bt_b + page_idx)
 
-        # Tokens in this page: PAGE_SIZE except possibly the last page
         if page_idx == num_pages - 1:
             rem = seq_len % P
             tokens_this_page = tl.where(rem == 0, P, rem)
         else:
             tokens_this_page = P
 
-        # Tile loop within a page (a page may be processed as one tile
-        # when BLOCK_L == P, but the autotuner may pick other sizes)
+        # Tile loop within a page
         for l_start in range(0, P, BLOCK_L):
             l_offs = l_start + offs_l
-            # Which positions in the page are valid?
-            # valid_l_mask: positions within this tile AND within the page token budget
             valid_l = l_offs < tokens_this_page
 
-            # --- Load K tile: (BLOCK_L, BLOCK_D) from this page, this KV head ---
-            # Global offset for K:
-            # kv_cache[blk_idx, pos, kv_head, d] for d in 0..D-1
-            # = blk_idx * stride_kv_blk + pos * stride_kv_pos + kv_head * stride_kv_h + d
+            # K / V base offset for this block + kv_head
             k_base = (
                 blk_idx.to(tl.int64) * stride_kv_blk
-                + kv_head * stride_kv_h
+                + pid_kv * stride_kv_h
             )
-            # We need K[l_offs[pos], d] — l_offs[pos] is the in-page position, d is dimension
-            # Reshape to 2D indexing: use l_offs for the pos dimension
-            k_offs = k_base + (l_offs[:, None] * stride_kv_pos) + offs_d[None, :]
+
+            # --- Load K tile ---
+            k_offs = k_base + l_offs[:, None] * stride_kv_pos + offs_d[None, :]
             mask_2d = valid_l[:, None] & (offs_d[None, :] < D)
             k_tile = tl.load(kv_cache_ptr + k_offs, mask=mask_2d, other=0.0)
-            k_tile_fp32 = k_tile.to(tl.float32)  # (BLOCK_L, BLOCK_D)
-
-            # --- Compute scores: (group_size, BLOCK_L) = Q @ K^T ---
-            # Q: (group_size, BLOCK_D), K: (BLOCK_L, BLOCK_D)
-            # scores[g, l] = sum_d Q[g, d] * K[l, d]
-            scores = tl.dot(q, tl.trans(k_tile_fp32))  # (group_size, BLOCK_L)
-            scores = scores * scale  # apply scaling
-
-            # Mask invalid positions to -inf for the softmax max-finding
+            k_f32 = k_tile.to(tl.float32)             # (BLOCK_L, BLOCK_D)
+
+            # --- QK^T ---
+            scores = tl.dot(q, tl.trans(k_f32))        # (group_size, BLOCK_L)
+            scores = scores * scale
+            # Mask invalid positions
             scores = tl.where(valid_l[None, :], scores, float("-inf"))
 
             # --- Online 

... (2833 more chars)

--- /dev/null
+++ b//workspace/problems/03_paged_attention/debug_kernel.py
@@ -0,0 +1,162 @@
+"""Debug script to compare ref vs sol for specific cases."""
+import sys
+sys.path.insert(0, '/workspace')
+import torch
+import reference
+import triton
+import triton.language as tl
+from src.eval.numeric_stress import numeric_stress_cases, numeric_stress_context, tolerance_for_case
+from src.eval.correctness import check_correctness
+import shapes, yaml
+from pathlib import Path
+
+meta = yaml.safe_load(Path('problem.yaml').read_text())
+device = torch.device('cuda:0')
+
+# Non-autotuned kernel with proper D handling
+@triton.jit
+def debug_kernel(
+    query_ptr, kv_cache_ptr, block_table_ptr, seq_lens_ptr, output_ptr,
+    B, H, Hkv, D, max_blocks, P,
+    stride_q_b, stride_q_h,
+    stride_kv_blk, stride_kv_pos, stride_kv_h,
+    stride_bt_b,
+    BLOCK_L: tl.constexpr, BLOCK_D: tl.constexpr, group_size: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_kv = tl.program_id(1)
+    seq_len = tl.load(seq_lens_ptr + pid_b)
+    num_pages = (seq_len + P - 1) // P
+    offs_l = tl.arange(0, BLOCK_L)
+    offs_g = tl.arange(0, group_size)
+    
+    q_head_base = pid_kv * group_size
+    
+    m = tl.full([group_size], float('-inf'), dtype=tl.float32)
+    l_sum = tl.zeros([group_size], dtype=tl.float32)
+    acc = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
+    scale = 1.0 / tl.sqrt(D.to(tl.float32))
+    
+    for page_idx in range(num_pages):
+        blk_idx = tl.load(block_table_ptr + pid_b * stride_bt_b + page_idx)
+        if page_idx == num_pages - 1:
+            rem = seq_len % P
+            tokens_this_page = tl.where(rem == 0, P, rem)
+        else:
+            tokens_this_page = P
+        
+        for l_start in range(0, P, BLOCK_L):
+            l_offs = l_start + offs_l
+            valid_l = l_offs < tokens_this_page
+            k_base = blk_idx.to(tl.int64) * stride_kv_blk + pid_kv * stride_kv_h
+            
+            # Process D in BLOCK_D chunks, accumulating the full dot product
+            scores_sum = tl.zeros([group_size, BLOCK_L], dtype=tl.float32)
+            acc_update = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
+            
+            for d_start in range(0, D, BLOCK_D):
+                offs_d = d_start + tl.arange(0, BLOCK_D)
+                mask_d = offs_d < D
+                
+                # Load Q chunk
+                q_offs = pid_b * stride_q_b + (q_head_base + offs_g[:, None]) * stride_q_h + offs_d[None, :]
+                q_chunk = tl.load(query_ptr + q_offs, mask=mask_d[None, :], other=0.0).to(tl.float32)
+                
+                # Load K chunk
+                k_offs = k_base + l_offs[:, None] * stride_kv_pos + offs_d[None, :]
+                mask_2d = valid_l[:, None] & mask_d[None, :]
+                k_chunk = tl.load(kv_cache_ptr + k_offs, mask=mask_2d, other=0.0).to(tl.float32)
+                
+                # Partial dot product
+                scores_sum += tl.dot(q_chunk, tl.trans(k_chunk))
+            
+            scores = scores_sum * scale
+            scores = tl.where(valid_l[None, :], scores, float('-inf'))
+            
+            m_new = tl.maximum(m, tl.max(scores, axis=1))
+            rescale = tl.exp(m - m_new)
+            acc = acc * rescale[:, None]
+            l_sum = l_sum * rescale
+            
+            p = tl.exp(scores - m_new[:, None])
+            p = tl.where(valid_l[None, :], p, 0.0)
+            l_sum = l_sum + tl.sum(p, axis=1)
+            
+            # Accumulate V (also in D chunks)
+            for d_start in range(0, D, BLOCK_D):
+                offs_d = d_start + tl.arange(0, BLOCK_D)
+                mask_d = offs_d < D
+                
+                v_offs = k_base + D + l_offs[:, None] * stride_kv_pos + offs_d[None, :]
+                mask_2d = valid_l[:, None] & mask_d[None, :]
+                v_chunk = tl.load(kv_cache_ptr + v_offs, mask=mask_2d, other=0.0).to(tl.float32)
+                
+                acc_update_d = tl.dot(p, v_chunk)  # (group_size, BLOCK_D)
+                # Scatter into acc — only valid D positions
+                acc_d = tl.where(mask_d[None, :], acc_update_d, 0.0)
+                acc += acc_d
+            
+            m = m_new
+    
+    # Write output (in D chunks)
+    for d_start in range(0, D, BLOCK_D):
+        offs_d = d_start + tl.arange(0, BLOCK_D)
+        mask_d = offs_d < D
+        
+        out_offs = pid_b * stride_q_b + (q_head_base + offs_g[:, None]) * stride_q_h + offs_d[None, :]
+        
+        # Extract the right portion of acc
+        acc_chunk = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
+        # Hmm, this is tricky — acc has shape (group_size, BLOCK_D) but we need
+        # to extract the d_start..d_start+BLOCK_D portion.
+        # Actually acc stores the full D-dim accumulation, but padded to BLOCK_D.
+        # Since BLOCK_D divides D (or we have multiple chunks), this gets complex.
+        
+        out_vals = acc / l_sum[:, None]
+        tl.store(output_ptr + out_offs, out_vals.to(tl.bfloat16), mask=mask_d[None, :])
+
+
+# Test nominal case
+shape = shapes.SHAPES[0]
+reference.BATCH = shape['batch']
+reference.NUM_HEADS = shape['num_heads']
+reference.NUM_KV_HEADS = shape['num_kv_heads']
+reference.HEAD_DIM = shape['head_dim']
+reference.SEQ_LEN = shape['seq_len']
+reference.PAGE_SIZE = shape['page_size']
+
+init_args = reference.get_init_inputs()
+ref_model = reference.Model(*init_args).to(device).eval()
+
+for case in numeric_stress_cases(meta.get('name', '')):
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+    base_inputs = [t.to(device) for t in reference.get_inputs()]
+    
+    with numeric_stress_context(ref_model, ref_model, base_inputs, case) as inputs:
+        q, kvc, bt, sl = inputs
+        B, H, D = q.shape
+        Hkv = kvc.shape[2]
+        P = kvc.shape[1]
+        group_size = H // Hkv
+        out = torch.zeros(B, H, D, dtype=t

... (930 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,7 +1,7 @@
 """Paged attention decode kernel for RTX PRO 6000 (SM120 Blackwell).
 
 Triton-based decode kernel with online softmax and paged KV-cache gathering.
-Uses one block per (batch, kv_head) pair, processes all GQA query heads together.
+One thread block per query head for maximum parallelism.
 """
 
 import math
@@ -23,16 +23,8 @@
 PAGE_SIZE = 16
 
 # ---------------------------------------------------------------------------
-# Triton kernel
-# ---------------------------------------------------------------------------
-#
-# Design:
-#   - Grid: (batch, num_kv_heads) — one program per GQA group per batch element.
-#   - Each program loads all Q heads in its GQA group at once and iterates
-#     over KV-cache pages one page (16 positions) at a time.
-#   - BLOCK_L = PAGE_SIZE = 16 (compile-time constant).
-#   - BLOCK_D = D (pinned at launch time so the full head_dim is processed).
-#   - Online softmax in fp32 avoids materialising the full scores matrix.
+# Triton kernel — one block per query head
+# ---------------------------------------------------------------------------
 
 
 @triton.jit
@@ -58,13 +50,16 @@
     group_size: tl.constexpr,
 ):
     """
-    Exact-match variant: BLOCK_D == D.
-
-    One program per (batch_element, kv_head).  KV-cache pages are loaded
-    one at a time (16 positions per page, entire D dimension at once).
+    Paged attention decode — one program per (batch, query_head).
+
+    BLOCK_D == D (pinned at launch).  Each program loads one page at a time
+    (BLOCK_L = PAGE_SIZE = 16 positions), computes dot products with its
+    single Q head, and maintains online-softmax state in fp32.
     """
-    pid_b = tl.program_id(0)
-    pid_kv = tl.program_id(1)
+    pid_q = tl.program_id(0)   # query head index
+    pid_b = tl.program_id(1)   # batch index
+
+    kv_head = pid_q // group_size
 
     seq_len = tl.load(seq_lens_ptr + pid_b)
     num_pages = (seq_len + P - 1) // P
@@ -72,25 +67,21 @@
     BLOCK_L: tl.constexpr = 16
     offs_l = tl.arange(0, BLOCK_L)
     offs_d = tl.arange(0, BLOCK_D)
-    offs_g = tl.arange(0, group_size)
-
-    q_head_base = pid_kv * group_size
-
-    # Load all Q heads at once: (group_size, D)
-    q_offs = (
-        pid_b * stride_q_b
-        + (q_head_base + offs_g[:, None]) * stride_q_h
-        + offs_d[None, :]
-    )
-    mask_q = offs_d[None, :] < D
+
+    # --- Load single Q head ---
+    q_offs = pid_b * stride_q_b + pid_q * stride_q_h + offs_d
+    mask_q = offs_d < D
     q = tl.load(query_ptr + q_offs, mask=mask_q, other=0.0).to(tl.float32)
-
-    m = tl.full([group_size], float("-inf"), dtype=tl.float32)
-    l_sum = tl.zeros([group_size], dtype=tl.float32)
-    acc = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
+    # q: (BLOCK_D,) with zeros beyond D
+
+    # --- Online-softmax state ---
+    m = float("-inf")
+    l_sum = 0.0
+    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
 
     scale = 1.0 / tl.sqrt(D.to(tl.float32))
 
+    # --- Page loop ---
     for page_idx in range(num_pages):
         blk_idx = tl.load(block_table_ptr + pid_b * stride_bt_b + page_idx)
 
@@ -101,42 +92,48 @@
             tokens_this_page = P
 
         valid_l = offs_l < tokens_this_page
-        k_base = blk_idx.to(tl.int64) * stride_kv_blk + pid_kv * stride_kv_h
-
-        # K tile: (BLOCK_L, D)
+        k_base = blk_idx.to(tl.int64) * stride_kv_blk + kv_head * stride_kv_h
+
+        # --- Load K tile: (BLOCK_L, BLOCK_D) ---
         k_offs = k_base + offs_l[:, None] * stride_kv_pos + offs_d[None, :]
         mask_kv = valid_l[:, None] & (offs_d[None, :] < D)
         k_tile = tl.load(kv_cache_ptr + k_offs, mask=mask_kv, other=0.0).to(tl.float32)
 
-        # Score = Q @ K^T : (group_size, BLOCK_L)
-        scores = tl.dot(q, tl.trans(k_tile)) * scale
-        scores = tl.where(valid_l[None, :], scores, float("-inf"))
-
-        # Online softmax
-        m_new = tl.maximum(m, tl.max(scores, axis=1))
+        # --- Scores: (BLOCK_L,) = q @ K^T ---
+        # tl.dot((1, BLOCK_D), (BLOCK_D, BLOCK_L)) → (1, BLOCK_L)
+        scores = tl.dot(
+            tl.reshape(q, (1, BLOCK_D)),
+            tl.trans(k_tile),
+        )
+        scores = tl.reshape(scores, (BLOCK_L,)) * scale
+        scores = tl.where(valid_l, scores, float("-inf"))
+
+        # --- Online softmax update ---
+        m_new = tl.maximum(m, tl.max(scores, axis=0))
+
         rescale = tl.exp(m - m_new)
-        acc = acc * rescale[:, None]
+        acc = acc * rescale
         l_sum = l_sum * rescale
 
-        p = tl.exp(scores - m_new[:, None])
-        p = tl.where(valid_l[None, :], p, 0.0)
-        l_sum = l_sum + tl.sum(p, axis=1)
-
-        # V tile: (BLOCK_L, D)
+        p = tl.exp(scores - m_new)
+        p = tl.where(valid_l, p, 0.0)
+        l_sum = l_sum + tl.sum(p, axis=0)
+
+        # --- Load V tile: (BLOCK_L, BLOCK_D) & accumulate ---
         v_offs = k_base + D + offs_l[:, None] * stride_kv_pos + offs_d[None, :]
         v_tile = tl.load(kv_cache_ptr + v_offs, mask=mask_kv, other=0.0).to(tl.float32)
-        acc = acc + tl.dot(p, v_tile)
+
+        # p: (BLOCK_L,), v_tile: (BLOCK_L, BLOCK_D)
+        # p @ V → (BLOCK_D,)
+        acc += tl.sum(p[:, None] * v_tile, axis=0)
 
         m = m_new
 
+    # --- Write output ---
     l_safe = tl.where(l_sum == 0.0, 1.0, l_sum)
-    out_vals = acc / l_safe[:, None]
-
-    out_offs = (
-        pid_b * stride_q_b
-        + (q_head_base + offs_g[:, None]) * stride_q_h
-        + offs_d[None, :]
-    )
+    out_vals = acc / l_safe
+
+    out_offs = pid_b * stride_q_b + pid_q * stride_q_h + offs_d
     tl.store(output_ptr + out_offs, out_vals.to(tl.bfloat16), mask=mask_q)
 
 
@@ -158,7 +155,9 @@
     group_size = H // Hkv
 
     out = torch.empty(B, H, D, dtype=torch.bfloat16, device=query.device)
-    grid = (B, Hkv)
+
+    # Grid: (num_heads, batch) — one block per query head
+   

... (76 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,7 +1,8 @@
 """Paged attention decode kernel for RTX PRO 6000 (SM120 Blackwell).
 
 Triton-based decode kernel with online softmax and paged KV-cache gathering.
-One thread block per query head for maximum parallelism.
+Splits work across (batch, kv_head, seq_chunk) for maximum SM occupancy,
+then reduces partial results with a second kernel.
 """
 
 import math
@@ -23,66 +24,81 @@
 PAGE_SIZE = 16
 
 # ---------------------------------------------------------------------------
-# Triton kernel — one block per query head
+# Pass 1: partial attention over a chunk of pages
 # ---------------------------------------------------------------------------
 
 
 @triton.jit
-def _paged_attention_decode_kernel(
+def _partial_decode_kernel(
     query_ptr,
     kv_cache_ptr,
     block_table_ptr,
     seq_lens_ptr,
-    output_ptr,
+    partial_out_ptr,     # (num_chunks, B, H, D + 2)  — last 2 elems = (m, l)
     B,
     H,
     Hkv,
     D,
     max_blocks,
     P,
+    num_pages_total,
+    pages_per_chunk,
     stride_q_b,
     stride_q_h,
     stride_kv_blk,
     stride_kv_pos,
     stride_kv_h,
     stride_bt_b,
+    stride_po_c,         # partial_out stride: chunk dim
+    stride_po_b,         # partial_out stride: batch dim
+    stride_po_h,         # partial_out stride: head dim
     BLOCK_D: tl.constexpr,
     group_size: tl.constexpr,
 ):
     """
-    Paged attention decode — one program per (batch, query_head).
-
-    BLOCK_D == D (pinned at launch).  Each program loads one page at a time
-    (BLOCK_L = PAGE_SIZE = 16 positions), computes dot products with its
-    single Q head, and maintains online-softmax state in fp32.
+    Compute partial attention for one chunk of pages.
+
+    Grid: (num_chunks, B, Hkv)
+    - pid_c = program_id(0): chunk index
+    - pid_b = program_id(1): batch index
+    - pid_kv = program_id(2): KV head index (handles group_size Q heads)
     """
-    pid_q = tl.program_id(0)   # query head index
-    pid_b = tl.program_id(1)   # batch index
-
-    kv_head = pid_q // group_size
+    pid_c = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_kv = tl.program_id(2)
 
     seq_len = tl.load(seq_lens_ptr + pid_b)
     num_pages = (seq_len + P - 1) // P
+
+    # Page range for this chunk
+    page_start = pid_c * pages_per_chunk
+    page_end = tl.minimum(page_start + pages_per_chunk, num_pages)
+    if page_start >= num_pages:
+        return
 
     BLOCK_L: tl.constexpr = 16
     offs_l = tl.arange(0, BLOCK_L)
     offs_d = tl.arange(0, BLOCK_D)
-
-    # --- Load single Q head ---
-    q_offs = pid_b * stride_q_b + pid_q * stride_q_h + offs_d
-    mask_q = offs_d < D
+    offs_g = tl.arange(0, group_size)
+
+    q_head_base = pid_kv * group_size
+
+    # Load all Q heads in this group
+    q_offs = (
+        pid_b * stride_q_b
+        + (q_head_base + offs_g[:, None]) * stride_q_h
+        + offs_d[None, :]
+    )
+    mask_q = offs_d[None, :] < D
     q = tl.load(query_ptr + q_offs, mask=mask_q, other=0.0).to(tl.float32)
-    # q: (BLOCK_D,) with zeros beyond D
-
-    # --- Online-softmax state ---
-    m = float("-inf")
-    l_sum = 0.0
-    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
+
+    m = tl.full([group_size], float("-inf"), dtype=tl.float32)
+    l_sum = tl.zeros([group_size], dtype=tl.float32)
+    acc = tl.zeros([group_size, BLOCK_D], dtype=tl.float32)
 
     scale = 1.0 / tl.sqrt(D.to(tl.float32))
 
-    # --- Page loop ---
-    for page_idx in range(num_pages):
+    for page_idx in range(page_start, page_end):
         blk_idx = tl.load(block_table_ptr + pid_b * stride_bt_b + page_idx)
 
         if page_idx == num_pages - 1:
@@ -92,49 +108,152 @@
             tokens_this_page = P
 
         valid_l = offs_l < tokens_this_page
-        k_base = blk_idx.to(tl.int64) * stride_kv_blk + kv_head * stride_kv_h
-
-        # --- Load K tile: (BLOCK_L, BLOCK_D) ---
+        k_base = blk_idx.to(tl.int64) * stride_kv_blk + pid_kv * stride_kv_h
+
+        # K tile
         k_offs = k_base + offs_l[:, None] * stride_kv_pos + offs_d[None, :]
         mask_kv = valid_l[:, None] & (offs_d[None, :] < D)
         k_tile = tl.load(kv_cache_ptr + k_offs, mask=mask_kv, other=0.0).to(tl.float32)
 
-        # --- Scores: (BLOCK_L,) = q @ K^T ---
-        # tl.dot((1, BLOCK_D), (BLOCK_D, BLOCK_L)) → (1, BLOCK_L)
-        scores = tl.dot(
-            tl.reshape(q, (1, BLOCK_D)),
-            tl.trans(k_tile),
-        )
-        scores = tl.reshape(scores, (BLOCK_L,)) * scale
-        scores = tl.where(valid_l, scores, float("-inf"))
-
-        # --- Online softmax update ---
-        m_new = tl.maximum(m, tl.max(scores, axis=0))
-
+        scores = tl.dot(q, tl.trans(k_tile)) * scale
+        scores = tl.where(valid_l[None, :], scores, float("-inf"))
+
+        m_new = tl.maximum(m, tl.max(scores, axis=1))
         rescale = tl.exp(m - m_new)
-        acc = acc * rescale
+        acc = acc * rescale[:, None]
         l_sum = l_sum * rescale
 
-        p = tl.exp(scores - m_new)
-        p = tl.where(valid_l, p, 0.0)
-        l_sum = l_sum + tl.sum(p, axis=0)
-
-        # --- Load V tile: (BLOCK_L, BLOCK_D) & accumulate ---
+        p = tl.exp(scores - m_new[:, None])
+        p = tl.where(valid_l[None, :], p, 0.0)
+        l_sum = l_sum + tl.sum(p, axis=1)
+
+        # V tile
         v_offs = k_base + D + offs_l[:, None] * stride_kv_pos + offs_d[None, :]
         v_tile = tl.load(kv_cache_ptr + v_offs, mask=mask_kv, other=0.0).to(tl.float32)
-
-        # p: (BLOCK_L,), v_tile: (BLOCK_L, BLOCK_D)
-        # p @ V → (BLOCK_D,)
-        acc += tl.sum(p[:, None] * v_tile, axis=0)
+        acc = acc + tl.dot(p, v_tile)
 
         m = m_new
 
-    # --- Write output ---
-    l_safe = tl.where(l_sum == 0.0, 1.0, l_sum)
-    out_vals = acc / l_safe
+    # Write partial results: (m[g], l_sum[g], acc[g, :D]) for each g in group
+    for g in range(group_si

... (5474 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -155,7 +155,7 @@
             + D + 1,
             l_sum[g],
         )
-        # Store acc
+            # Store acc (fp32, no conversion needed — partial_out is fp32)
         val_offs = offs_d
         mask_d = offs_d < D
         tl.store(
@@ -164,7 +164,7 @@
             + pid_b * stride_po_b
             + q_head * stride_po_h
             + val_offs,
-            acc[g, :].to(tl.bfloat16),
+            acc[g, :],
             mask=mask_d,
         )
 

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -284,10 +284,10 @@
     num_chunks = max(1, min(max_pages, target_blocks // base_blocks))
     pages_per_chunk = (max_pages + num_chunks - 1) // num_chunks
 
-    # Intermediate storage: (num_chunks, B, H, D + 2) bf16
+    # Intermediate storage: (num_chunks, B, H, D + 2) fp32 for precision
     partial = torch.zeros(
         num_chunks, B, H, D + 2,
-        dtype=torch.bfloat16,
+        dtype=torch.float32,
         device=query.device,
     )
 

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -234,7 +234,7 @@
         o_global = tl.where(chunk_valid, o_global * rescale_old, o_global)
         l_global = tl.where(chunk_valid, l_global * rescale_old, l_global)
 
-        # Load acc_c for this chunk
+        # Load acc_c for this chunk (stored as fp32)
         acc_c = tl.load(
             partial_in_ptr
             + c * stride_pi_c
@@ -243,7 +243,7 @@
             + offs_d,
             mask=mask_d,
             other=0.0,
-        ).to(tl.float32)
+        )
 
         o_global = tl.where(chunk_valid, o_global + rescale_c * l_c * acc_c, o_global)
         l_global = tl.where(chunk_valid, l_global + rescale_c * l_c, l_global)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,33 +1,19 @@
-    # Write partial results: (m[g], l_sum[g], acc[g, :D]) for each g in group
-    for g in range(group_size):
-        q_head = q_head_base + g
-        # Store m
-        tl.store(
-            partial_out_ptr
-            + pid_c * stride_po_c
-            + pid_b * stride_po_b
-            + q_head * stride_po_h
-            + D,
-            m[g],
-        )
-        # Store l_sum
-        tl.store(
-            partial_out_ptr
-            + pid_c * stride_po_c
-            + pid_b * stride_po_b
-            + q_head * stride_po_h
-            + D + 1,
-            l_sum[g],
-        )
-        # Store acc (fp32, no conversion needed — partial_out is fp32)
-        val_offs = offs_d
-        mask_d = offs_d < D
-        tl.store(
-            partial_out_ptr
-            + pid_c * stride_po_c
-            + pid_b * stride_po_b
-            + q_head * stride_po_h
-            + val_offs,
-            acc[g, :],
-            mask=mask_d,
-        )+    # Write partial results vectorised across all Q heads in the group.
+    base_off = pid_c * stride_po_c + pid_b * stride_po_b
+
+    # m: (group_size,) → store at [..., q_head, D]
+    m_offs = base_off + (q_head_base + offs_g) * stride_po_h + D
+    tl.store(partial_out_ptr + m_offs, m)
+
+    # l_sum: (group_size,) → store at [..., q_head, D+1]
+    l_offs = base_off + (q_head_base + offs_g) * stride_po_h + D + 1
+    tl.store(partial_out_ptr + l_offs, l_sum)
+
+    # acc: (group_size, BLOCK_D) → store at [..., q_head, 0:D]
+    acc_offs = (
+        base_off
+        + (q_head_base + offs_g[:, None]) * stride_po_h
+        + offs_d[None, :]
+    )
+    mask_acc = offs_d[None, :] < D
+    tl.store(partial_out_ptr + acc_offs, acc, mask=mask_acc)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,7 +1,7 @@
 """Paged attention decode kernel for RTX PRO 6000 (SM120 Blackwell).
 
 Triton-based decode kernel with online softmax and paged KV-cache gathering.
-Splits work across (batch, kv_head, seq_chunk) for maximum SM occupancy,
+Splits work across (batch, kv_head, seq_chunk) for SM occupancy,
 then reduces partial results with a second kernel.
 """
 
@@ -34,7 +34,7 @@
     kv_cache_ptr,
     block_table_ptr,
     seq_lens_ptr,
-    partial_out_ptr,     # (num_chunks, B, H, D + 2)  — last 2 elems = (m, l)
+    partial_out_ptr,
     B,
     H,
     Hkv,
@@ -49,19 +49,15 @@
     stride_kv_pos,
     stride_kv_h,
     stride_bt_b,
-    stride_po_c,         # partial_out stride: chunk dim
-    stride_po_b,         # partial_out stride: batch dim
-    stride_po_h,         # partial_out stride: head dim
+    stride_po_c,
+    stride_po_b,
+    stride_po_h,
     BLOCK_D: tl.constexpr,
     group_size: tl.constexpr,
 ):
-    """
-    Compute partial attention for one chunk of pages.
+    """Compute partial attention for one chunk of pages.
 
     Grid: (num_chunks, B, Hkv)
-    - pid_c = program_id(0): chunk index
-    - pid_b = program_id(1): batch index
-    - pid_kv = program_id(2): KV head index (handles group_size Q heads)
     """
     pid_c = tl.program_id(0)
     pid_b = tl.program_id(1)
@@ -70,7 +66,6 @@
     seq_len = tl.load(seq_lens_ptr + pid_b)
     num_pages = (seq_len + P - 1) // P
 
-    # Page range for this chunk
     page_start = pid_c * pages_per_chunk
     page_end = tl.minimum(page_start + pages_per_chunk, num_pages)
     if page_start >= num_pages:
@@ -134,39 +129,25 @@
 
         m = m_new
 
-    # Write partial results: (m[g], l_sum[g], acc[g, :D]) for each g in group
-    for g in range(group_size):
-        q_head = q_head_base + g
-        # Store m
-        tl.store(
-            partial_out_ptr
-            + pid_c * stride_po_c
-            + pid_b * stride_po_b
-            + q_head * stride_po_h
-            + D,
-            m[g],
-        )
-        # Store l_sum
-        tl.store(
-            partial_out_ptr
-            + pid_c * stride_po_c
-            + pid_b * stride_po_b
-            + q_head * stride_po_h
-            + D + 1,
-            l_sum[g],
-        )
-            # Store acc (fp32, no conversion needed — partial_out is fp32)
-        val_offs = offs_d
-        mask_d = offs_d < D
-        tl.store(
-            partial_out_ptr
-            + pid_c * stride_po_c
-            + pid_b * stride_po_b
-            + q_head * stride_po_h
-            + val_offs,
-            acc[g, :],
-            mask=mask_d,
-        )
+    # Write partial results — vectorised across all Q heads in the group.
+    base_off = pid_c * stride_po_c + pid_b * stride_po_b
+
+    # m: (group_size,) -> store at [c, b, q_head, D]
+    m_offs = base_off + (q_head_base + offs_g) * stride_po_h + D
+    tl.store(partial_out_ptr + m_offs, m)
+
+    # l_sum: (group_size,) -> store at [c, b, q_head, D+1]
+    ls_offs = base_off + (q_head_base + offs_g) * stride_po_h + D + 1
+    tl.store(partial_out_ptr + ls_offs, l_sum)
+
+    # acc: (group_size, BLOCK_D) -> store at [c, b, q_head, 0:D]
+    acc_offs = (
+        base_off
+        + (q_head_base + offs_g[:, None]) * stride_po_h
+        + offs_d[None, :]
+    )
+    mask_acc = offs_d[None, :] < D
+    tl.store(partial_out_ptr + acc_offs, acc, mask=mask_acc)
 
 
 # ---------------------------------------------------------------------------
@@ -176,8 +157,8 @@
 
 @triton.jit
 def _reduce_partial_kernel(
-    partial_in_ptr,      # (num_chunks, B, H, D + 2)
-    output_ptr,          # (B, H, D)
+    partial_in_ptr,
+    output_ptr,
     num_chunks,
     B,
     H,
@@ -189,8 +170,7 @@
     stride_q_h,
     BLOCK_D: tl.constexpr,
 ):
-    """
-    Reduce partial results across chunks into final output.
+    """Reduce partial results across chunks into final output.
 
     Grid: (H, B) — one program per query head.
     """
@@ -198,16 +178,14 @@
     pid_b = tl.program_id(1)
 
     offs_d = tl.arange(0, BLOCK_D)
-
-    # Global softmax reduction
+    mask_d = offs_d < D
+
     m_global = float("-inf")
     l_global = 0.0
     o_global = tl.zeros([BLOCK_D], dtype=tl.float32)
 
-    mask_d = offs_d < D
-
     for c in range(num_chunks):
-        # Load m_c
+        # Load m_c, l_c
         m_c = tl.load(
             partial_in_ptr
             + c * stride_pi_c
@@ -215,7 +193,6 @@
             + pid_q * stride_pi_h
             + D
         )
-        # Load l_c
         l_c = tl.load(
             partial_in_ptr
             + c * stride_pi_c
@@ -224,17 +201,19 @@
             + D + 1
         )
 
-        # Skip empty chunks (l_c == 0 and m_c == -inf)
         chunk_valid = l_c > 0.0
 
         m_new = tl.maximum(m_global, m_c)
+
+        # Only rescale if chunk is valid; otherwise keep current state.
+        # tl.where selects element-wise — both branches are evaluated.
         rescale_old = tl.exp(m_global - m_new)
         rescale_c = tl.exp(m_c - m_new)
 
         o_global = tl.where(chunk_valid, o_global * rescale_old, o_global)
         l_global = tl.where(chunk_valid, l_global * rescale_old, l_global)
 
-        # Load acc_c for this chunk (stored as fp32)
+        # Load acc_c for this chunk
         acc_c = tl.load(
             partial_in_ptr
             + c * stride_pi_c
@@ -245,8 +224,16 @@
             other=0.0,
         )
 
-        o_global = tl.where(chunk_valid, o_global + rescale_c * l_c * acc_c, o_global)
-        l_global = tl.where(chunk_valid, l_global + rescale_c * l_c, l_global)
+        o_global = tl.where(
+            chunk_valid,
+            o_global + rescale_c * l_c * acc_c,
+            o_global,
+        )
+        l_global = tl.where(
+            chunk_valid,
+            l_global + rescale_c * l_c,
+            l_global,
+        )
         m_global = tl.where(chunk_valid

... (1526 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -224,9 +224,10 @@
             other=0.0,
         )
 
+        # acc_c is already the exp-weighted sum: Σ exp(s-m_c)·V
         o_global = tl.where(
             chunk_valid,
-            o_global + rescale_c * l_c * acc_c,
+            o_global + rescale_c * acc_c,
             o_global,
         )
         l_global = tl.where(

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -277,52 +277,99 @@
         device=query.device,
     )
 
-    # Pass 1
-    grid1 = (num_chunks, B, Hkv)
-    _partial_decode_kernel[grid1](
-        query,
-        kv_cache,
-        block_table,
-        seq_lens,
-        partial,
-        B,
-        H,
-        Hkv,
-        D,
-        max_blocks,
-        P,
-        max_pages,
-        pages_per_chunk,
-        query.stride(0),
-        query.stride(1),
-        kv_cache.stride(0),
-        kv_cache.stride(1),
-        kv_cache.stride(2),
-        block_table.stride(0),
-        partial.stride(0),
-        partial.stride(1),
-        partial.stride(2),
-        BLOCK_D=D,
-        group_size=group_size,
-    )
-
-    # Pass 2
     out = torch.empty(B, H, D, dtype=torch.bfloat16, device=query.device)
-    grid2 = (H, B)
-    _reduce_partial_kernel[grid2](
-        partial,
-        out,
-        num_chunks,
-        B,
-        H,
-        D,
-        partial.stride(0),
-        partial.stride(1),
-        partial.stride(2),
-        out.stride(0),
-        out.stride(1),
-        BLOCK_D=D,
-    )
+
+    if num_chunks == 1:
+        # Fast path: single chunk, write output directly from pass 1.
+        # Use the partial tensor as output (ignoring the m/l tail).
+        # We still compute via the partial kernel, then copy acc/l to output.
+        grid1 = (1, B, Hkv)
+        _partial_decode_kernel[grid1](
+            query,
+            kv_cache,
+            block_table,
+            seq_lens,
+            partial,
+            B,
+            H,
+            Hkv,
+            D,
+            max_blocks,
+            P,
+            max_pages,
+            pages_per_chunk,
+            query.stride(0),
+            query.stride(1),
+            kv_cache.stride(0),
+            kv_cache.stride(1),
+            kv_cache.stride(2),
+            block_table.stride(0),
+            partial.stride(0),
+            partial.stride(1),
+            partial.stride(2),
+            BLOCK_D=D,
+            group_size=group_size,
+        )
+        # Single-chunk reduction pass
+        grid2 = (H, B)
+        _reduce_partial_kernel[grid2](
+            partial,
+            out,
+            num_chunks,
+            B,
+            H,
+            D,
+            partial.stride(0),
+            partial.stride(1),
+            partial.stride(2),
+            out.stride(0),
+            out.stride(1),
+            BLOCK_D=D,
+        )
+    else:
+        # Multi-chunk: partial + reduce
+        grid1 = (num_chunks, B, Hkv)
+        _partial_decode_kernel[grid1](
+            query,
+            kv_cache,
+            block_table,
+            seq_lens,
+            partial,
+            B,
+            H,
+            Hkv,
+            D,
+            max_blocks,
+            P,
+            max_pages,
+            pages_per_chunk,
+            query.stride(0),
+            query.stride(1),
+            kv_cache.stride(0),
+            kv_cache.stride(1),
+            kv_cache.stride(2),
+            block_table.stride(0),
+            partial.stride(0),
+            partial.stride(1),
+            partial.stride(2),
+            BLOCK_D=D,
+            group_size=group_size,
+        )
+        grid2 = (H, B)
+        _reduce_partial_kernel[grid2](
+            partial,
+            out,
+            num_chunks,
+            B,
+            H,
+            D,
+            partial.stride(0),
+            partial.stride(1),
+            partial.stride(2),
+            out.stride(0),
+            out.stride(1),
+            BLOCK_D=D,
+        )
 
     return out
 

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -41,7 +41,6 @@
     D,
     max_blocks,
     P,
-    num_pages_total,
     pages_per_chunk,
     stride_q_b,
     stride_q_h,

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -278,97 +278,50 @@
 
     out = torch.empty(B, H, D, dtype=torch.bfloat16, device=query.device)
 
-    if num_chunks == 1:
-        # Fast path: single chunk, write output directly from pass 1.
-        # Use the partial tensor as output (ignoring the m/l tail).
-        # We still compute via the partial kernel, then copy acc/l to output.
-        grid1 = (1, B, Hkv)
-        _partial_decode_kernel[grid1](
-            query,
-            kv_cache,
-            block_table,
-            seq_lens,
-            partial,
-            B,
-            H,
-            Hkv,
-            D,
-            max_blocks,
-            P,
-            max_pages,
-            pages_per_chunk,
-            query.stride(0),
-            query.stride(1),
-            kv_cache.stride(0),
-            kv_cache.stride(1),
-            kv_cache.stride(2),
-            block_table.stride(0),
-            partial.stride(0),
-            partial.stride(1),
-            partial.stride(2),
-            BLOCK_D=D,
-            group_size=group_size,
-        )
-        # Single-chunk reduction pass
-        grid2 = (H, B)
-        _reduce_partial_kernel[grid2](
-            partial,
-            out,
-            num_chunks,
-            B,
-            H,
-            D,
-            partial.stride(0),
-            partial.stride(1),
-            partial.stride(2),
-            out.stride(0),
-            out.stride(1),
-            BLOCK_D=D,
-        )
-    else:
-        # Multi-chunk: partial + reduce
-        grid1 = (num_chunks, B, Hkv)
-        _partial_decode_kernel[grid1](
-            query,
-            kv_cache,
-            block_table,
-            seq_lens,
-            partial,
-            B,
-            H,
-            Hkv,
-            D,
-            max_blocks,
-            P,
-            max_pages,
-            pages_per_chunk,
-            query.stride(0),
-            query.stride(1),
-            kv_cache.stride(0),
-            kv_cache.stride(1),
-            kv_cache.stride(2),
-            block_table.stride(0),
-            partial.stride(0),
-            partial.stride(1),
-            partial.stride(2),
-            BLOCK_D=D,
-            group_size=group_size,
-        )
-        grid2 = (H, B)
-        _reduce_partial_kernel[grid2](
-            partial,
-            out,
-            num_chunks,
-            B,
-            H,
-            D,
-            partial.stride(0),
-            partial.stride(1),
-            partial.stride(2),
-            out.stride(0),
-            out.stride(1),
-            BLOCK_D=D,
-        )
+    # Pass 1: partial attention over chunks
+    grid1 = (num_chunks, B, Hkv)
+    _partial_decode_kernel[grid1](
+        query,
+        kv_cache,
+        block_table,
+        seq_lens,
+        partial,
+        B,
+        H,
+        Hkv,
+        D,
+        max_blocks,
+        P,
+        pages_per_chunk,
+        query.stride(0),
+        query.stride(1),
+        kv_cache.stride(0),
+        kv_cache.stride(1),
+        kv_cache.stride(2),
+        block_table.stride(0),
+        partial.stride(0),
+        partial.stride(1),
+        partial.stride(2),
+        BLOCK_D=D,
+        group_size=group_size,
+    )
+
+    # Pass 2: reduce partial results across chunks
+    grid2 = (H, B)
+    _reduce_partial_kernel[grid2](
+        partial,
+        out,
+        num_chunks,
+        B,
+        H,
+        D,
+        partial.stride(0),
+        partial.stride(1),
+        partial.stride(2),
+        out.stride(0),
+        out.stride(1),
+        BLOCK_D=D,
+    )
 
     return out