--- /dev/null
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -0,0 +1,212 @@
+"""Paged attention decode kernel via Triton."""
+import math
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+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
+
+
+@triton.jit
+def _paged_decode_kernel(
+    Q_ptr,
+    KV_ptr,
+    BT_ptr,
+    SL_ptr,
+    O_ptr,
+    stride_qb,
+    stride_qh,
+    stride_qd,
+    stride_kv_block,
+    stride_kv_page,
+    stride_kv_h,
+    stride_kv_d,
+    stride_bt_b,
+    stride_bt_p,
+    stride_ob,
+    stride_oh,
+    stride_od,
+    num_heads,
+    num_kv_heads,
+    head_dim,
+    page_size,
+    group_size,
+    scale,
+    BLOCK_D: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+    head_idx = tl.program_id(1)
+
+    kv_head_idx = head_idx // group_size
+    seq_len = tl.load(SL_ptr + batch_idx)
+
+    offs_d = tl.arange(0, BLOCK_D)
+    d_mask = offs_d < head_dim
+
+    q_ptr = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh
+    q = tl.load(q_ptr + offs_d * stride_qd, mask=d_mask, other=0.0).to(tl.float32)
+
+    m_i = tl.full([], -float("inf"), tl.float32)
+    l_i = tl.full([], 0.0, tl.float32)
+    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
+
+    num_pages = (seq_len + page_size - 1) // page_size
+
+    for page_idx in range(num_pages):
+        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
+
+        page_start = page_idx * page_size
+        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
+
+        for t in range(BLOCK_T):
+            if t >= tokens_in_page:
+                break
+
+            token_idx = page_start + t
+            kv_base = (
+                block_id * stride_kv_block
+                + t * stride_kv_page
+                + kv_head_idx * stride_kv_h
+            )
+
+            k = tl.load(
+                KV_ptr + kv_base + offs_d * stride_kv_d,
+                mask=d_mask,
+                other=0.0,
+            ).to(tl.float32)
+            v = tl.load(
+                KV_ptr + kv_base + (head_dim + offs_d) * stride_kv_d,
+                mask=d_mask,
+                other=0.0,
+            ).to(tl.float32)
+
+            score = tl.sum(q * k, axis=0) * scale
+            m_ij = tl.maximum(m_i, score)
+            alpha = tl.exp(m_i - m_ij)
+            beta = tl.exp(score - m_ij)
+            l_i = l_i * alpha + beta
+            acc = acc * alpha + beta * v
+            m_i = m_ij
+
+    out = (acc / l_i).to(tl.bfloat16)
+    o_ptr = O_ptr + batch_idx * stride_ob + head_idx * stride_oh
+    tl.store(o_ptr + offs_d * stride_od, out, mask=d_mask)
+
+
+def paged_attention_decode(
+    query: torch.Tensor,
+    kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    B, H, D = query.shape
+    _, P, Hkv, D2 = kv_cache.shape
+    assert D2 == 2 * D
+    group_size = H // Hkv
+
+    out = torch.empty_like(query)
+
+    BLOCK_D = triton.next_power_of_2(D)
+    BLOCK_T = P  # page_size
+
+    grid = (B, H)
+    _paged_decode_kernel[grid](
+        query,
+        kv_cache,
+        block_table,
+        seq_lens,
+        out,
+        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),
+        num_heads=H,
+        num_kv_heads=Hkv,
+        head_dim=D,
+        page_size=P,
+        group_size=group_size,
+        scale=scale,
+        BLOCK_D=BLOCK_D,
+        BLOCK_T=BLOCK_T,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out
+
+
+class Model(nn.Module):
+    def __init__(
+        self,
+        batch: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        seq_len: int,
+        page_size: int,
+    ):
+        super().__init__()
+        assert num_heads % num_kv_heads == 0
+        self.batch = batch
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.seq_len = seq_len
+        self.page_size = page_size
+        self.group_size = num_heads // num_kv_heads
+        self.scale = 1.0 / math.sqrt(head_dim)
+        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        kv_cache: torch.Tensor,
+        block_table: torch.Tensor,
+        seq_lens: torch.Tensor,
+    ) -> torch.Tensor:
+        return paged_attention_decode(query, kv_cache, block_table, seq_lens, self.scale)
+
+
+def get_inputs():
+    B = BATCH
+    H = NUM_HEADS
+    Hkv = NUM_KV_HEADS
+    D = HEAD_DIM
+    L = SEQ_LEN
+    P = PAGE_SIZE
+
+    pages_per_seq = (L + P - 1) // P
+    total_pages = max(B * pages_per_seq + 8, 64)
+
+    query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1
+    kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1
+
+    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
+    block_table = perm.contiguous()
+    seq_lens = torch.full((B,), L, dtype=torch.int32)
+
+    return [query, kv_cache, block_table, seq_lens]
+
+
+def get_init_inputs():
+    return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE]

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,80 +1,5 @@
-"""Paged attention decode kernel via Triton."""
-import math
-
-import torch
-import torch.nn as nn
-import triton
-import triton.language as tl
-
-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
-
-
-@triton.jit
-def _paged_decode_kernel(
-    Q_ptr,
-    KV_ptr,
-    BT_ptr,
-    SL_ptr,
-    O_ptr,
-    stride_qb,
-    stride_qh,
-    stride_qd,
-    stride_kv_block,
-    stride_kv_page,
-    stride_kv_h,
-    stride_kv_d,
-    stride_bt_b,
-    stride_bt_p,
-    stride_ob,
-    stride_oh,
-    stride_od,
-    num_heads,
-    num_kv_heads,
-    head_dim,
-    page_size,
-    group_size,
-    scale,
-    BLOCK_D: tl.constexpr,
-    BLOCK_T: tl.constexpr,
-):
-    batch_idx = tl.program_id(0)
-    head_idx = tl.program_id(1)
-
-    kv_head_idx = head_idx // group_size
-    seq_len = tl.load(SL_ptr + batch_idx)
-
-    offs_d = tl.arange(0, BLOCK_D)
-    d_mask = offs_d < head_dim
-
-    q_ptr = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh
-    q = tl.load(q_ptr + offs_d * stride_qd, mask=d_mask, other=0.0).to(tl.float32)
-
-    m_i = tl.full([], -float("inf"), tl.float32)
-    l_i = tl.full([], 0.0, tl.float32)
-    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
-
-    num_pages = (seq_len + page_size - 1) // page_size
-
-    for page_idx in range(num_pages):
-        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
-
-        page_start = page_idx * page_size
-        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
-
         for t in range(BLOCK_T):
-            if t >= tokens_in_page:
-                break
-
-            token_idx = page_start + t
+            t_valid = t < tokens_in_page
             kv_base = (
                 block_id * stride_kv_block
                 + t * stride_kv_page
@@ -83,130 +8,20 @@
 
             k = tl.load(
                 KV_ptr + kv_base + offs_d * stride_kv_d,
-                mask=d_mask,
+                mask=d_mask & t_valid,
                 other=0.0,
             ).to(tl.float32)
             v = tl.load(
                 KV_ptr + kv_base + (head_dim + offs_d) * stride_kv_d,
-                mask=d_mask,
+                mask=d_mask & t_valid,
                 other=0.0,
             ).to(tl.float32)
 
             score = tl.sum(q * k, axis=0) * scale
+            score = tl.where(t_valid, score, -float("inf"))
             m_ij = tl.maximum(m_i, score)
             alpha = tl.exp(m_i - m_ij)
             beta = tl.exp(score - m_ij)
             l_i = l_i * alpha + beta
             acc = acc * alpha + beta * v
-            m_i = m_ij
-
-    out = (acc / l_i).to(tl.bfloat16)
-    o_ptr = O_ptr + batch_idx * stride_ob + head_idx * stride_oh
-    tl.store(o_ptr + offs_d * stride_od, out, mask=d_mask)
-
-
-def paged_attention_decode(
-    query: torch.Tensor,
-    kv_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    seq_lens: torch.Tensor,
-    scale: float,
-) -> torch.Tensor:
-    B, H, D = query.shape
-    _, P, Hkv, D2 = kv_cache.shape
-    assert D2 == 2 * D
-    group_size = H // Hkv
-
-    out = torch.empty_like(query)
-
-    BLOCK_D = triton.next_power_of_2(D)
-    BLOCK_T = P  # page_size
-
-    grid = (B, H)
-    _paged_decode_kernel[grid](
-        query,
-        kv_cache,
-        block_table,
-        seq_lens,
-        out,
-        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),
-        num_heads=H,
-        num_kv_heads=Hkv,
-        head_dim=D,
-        page_size=P,
-        group_size=group_size,
-        scale=scale,
-        BLOCK_D=BLOCK_D,
-        BLOCK_T=BLOCK_T,
-        num_warps=4,
-        num_stages=2,
-    )
-    return out
-
-
-class Model(nn.Module):
-    def __init__(
-        self,
-        batch: int,
-        num_heads: int,
-        num_kv_heads: int,
-        head_dim: int,
-        seq_len: int,
-        page_size: int,
-    ):
-        super().__init__()
-        assert num_heads % num_kv_heads == 0
-        self.batch = batch
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_dim = head_dim
-        self.seq_len = seq_len
-        self.page_size = page_size
-        self.group_size = num_heads // num_kv_heads
-        self.scale = 1.0 / math.sqrt(head_dim)
-        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        kv_cache: torch.Tensor,
-        block_table: torch.Tensor,
-        seq_lens: torch.Tensor,
-    ) -> torch.Tensor:
-        return paged_attention_decode(query, kv_cache, block_table, seq_lens, self.scale)
-
-
-def get_inputs():
-    B = BATCH
-    H = NUM_HEADS
-    Hkv = NUM_KV_HEADS
-    D = HEAD_DIM
-    L = SEQ_LEN
-    P = PAGE_SIZE
-
-    pages_per_seq = (L + P - 1) // P
-    total_pages = max(B * pages_per_seq + 8, 64)
-
-    query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1
-    kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1
-
-    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
-    block_table = perm.contiguous()
-    seq_lens = torch.full((B,), L, dtype=torch.int32)
-
-    return [query, kv_cache, block_table, seq_lens]
-
-
-def get_init_inputs():
-    return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE]
+            m_i = m_ij

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,27 +1,234 @@
-        for t in range(BLOCK_T):
-            t_valid = t < tokens_in_page
-            kv_base = (
-                block_id * stride_kv_block
-                + t * stride_kv_page
-                + kv_head_idx * stride_kv_h
-            )
-
-            k = tl.load(
+"""Paged attention decode kernel via Triton (GQA-aware, page-batched)."""
+import math
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+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
+
+
+@triton.jit
+def _paged_decode_gqa_kernel(
+    Q_ptr,
+    KV_ptr,
+    BT_ptr,
+    SL_ptr,
+    O_ptr,
+    stride_qb,
+    stride_qh,
+    stride_qd,
+    stride_kv_block,
+    stride_kv_page,
+    stride_kv_h,
+    stride_kv_d,
+    stride_bt_b,
+    stride_bt_p,
+    stride_ob,
+    stride_oh,
+    stride_od,
+    num_heads,
+    head_dim,
+    page_size,
+    group_size,
+    scale,
+    BLOCK_G: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+    kv_head_idx = tl.program_id(1)
+
+    seq_len = tl.load(SL_ptr + batch_idx)
+    num_pages = (seq_len + page_size - 1) // page_size
+
+    offs_g = tl.arange(0, BLOCK_G)
+    offs_t = tl.arange(0, BLOCK_T)
+    offs_d = tl.arange(0, BLOCK_D)
+    g_mask = (kv_head_idx * group_size + offs_g) < num_heads
+    d_mask = offs_d < head_dim
+
+    q = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
+    head_base = kv_head_idx * group_size
+    for g in tl.static_range(BLOCK_G):
+        head_idx = head_base + g
+        q_ptr = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh
+        q_row = tl.load(
+            q_ptr + offs_d * stride_qd,
+            mask=g_mask[g] & d_mask,
+            other=0.0,
+        ).to(tl.float32)
+        q = tl.where(g_mask[g], tl.expand_dims(q_row, 0), q)  # need proper assignment
+
+    # Load Q rows explicitly
+    q_rows = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
+    for g in tl.static_range(BLOCK_G):
+        head_idx = head_base + g
+        q_ptr = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh
+        q_row = tl.load(q_ptr + offs_d * stride_qd, mask=d_mask, other=0.0).to(tl.float32)
+        q_rows = tl.where(tl.expand_dims(g_mask[g], 1), tl.expand_dims(q_row, 0), q_rows)
+    q = q_rows
+
+    m_i = tl.full([BLOCK_G], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
+
+    for page_idx in range(num_pages):
+        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
+        page_start = page_idx * page_size
+        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
+
+        t_mask = offs_t[:, None] < tokens_in_page
+        kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
+
+        k = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32)
+        v = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32)
+        for t in tl.static_range(BLOCK_T):
+            kv_base = kv_h_base + t * stride_kv_page
+            k_row = tl.load(
                 KV_ptr + kv_base + offs_d * stride_kv_d,
-                mask=d_mask & t_valid,
+                mask=d_mask,
                 other=0.0,
             ).to(tl.float32)
-            v = tl.load(
+            v_row = tl.load(
                 KV_ptr + kv_base + (head_dim + offs_d) * stride_kv_d,
-                mask=d_mask & t_valid,
+                mask=d_mask,
                 other=0.0,
             ).to(tl.float32)
-
-            score = tl.sum(q * k, axis=0) * scale
-            score = tl.where(t_valid, score, -float("inf"))
-            m_ij = tl.maximum(m_i, score)
-            alpha = tl.exp(m_i - m_ij)
-            beta = tl.exp(score - m_ij)
-            l_i = l_i * alpha + beta
-            acc = acc * alpha + beta * v
-            m_i = m_ij+            k = tl.where(tl.expand_dims(offs_t == t, 1), tl.expand_dims(k_row, 0), k)
+            v = tl.where(tl.expand_dims(offs_t == t, 1), tl.expand_dims(v_row, 0), v)
+
+        qk = tl.dot(q, tl.trans(k)) * scale
+        qk = tl.where(t_mask, qk, -float("inf"))
+
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        alpha = tl.exp(m_i - m_ij)
+        p = tl.exp(qk - tl.expand_dims(m_ij, 1))
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * tl.expand_dims(alpha, 1) + tl.dot(p.to(tl.bfloat16), v)
+        m_i = m_ij
+
+    out_vals = acc / tl.expand_dims(l_i, 1)
+    for g in tl.static_range(BLOCK_G):
+        head_idx = head_base + g
+        o_ptr = O_ptr + batch_idx * stride_ob + head_idx * stride_oh
+        o_row = tl.load(out_vals + g)  # invalid
+        tl.store(o_ptr + offs_d * stride_od, out_vals[g, :].to(tl.bfloat16), mask=g_mask[g] & d_mask)
+
+
+def paged_attention_decode(
+    query: torch.Tensor,
+    kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    B, H, D = query.shape
+    _, P, Hkv, D2 = kv_cache.shape
+    assert D2 == 2 * D
+    group_size = H // Hkv
+
+    out = torch.empty_like(query)
+
+    BLOCK_G = triton.next_power_of_2(group_size)
+    BLOCK_T = P
+    BLOCK_D = triton.next_power_of_2(D)
+
+    grid = (B, Hkv)
+    _paged_decode_gqa_kernel[grid](
+        query,
+        kv_cache,
+        block_table,
+        seq_lens,
+        out,
+        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),
+        num_heads=H,
+ 

... (1955 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -55,29 +55,18 @@
     offs_g = tl.arange(0, BLOCK_G)
     offs_t = tl.arange(0, BLOCK_T)
     offs_d = tl.arange(0, BLOCK_D)
-    g_mask = (kv_head_idx * group_size + offs_g) < num_heads
+
+    head_ids = kv_head_idx * group_size + offs_g
+    g_mask = head_ids < num_heads
     d_mask = offs_d < head_dim
 
-    q = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
-    head_base = kv_head_idx * group_size
-    for g in tl.static_range(BLOCK_G):
-        head_idx = head_base + g
-        q_ptr = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh
-        q_row = tl.load(
-            q_ptr + offs_d * stride_qd,
-            mask=g_mask[g] & d_mask,
-            other=0.0,
-        ).to(tl.float32)
-        q = tl.where(g_mask[g], tl.expand_dims(q_row, 0), q)  # need proper assignment
-
-    # Load Q rows explicitly
-    q_rows = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
-    for g in tl.static_range(BLOCK_G):
-        head_idx = head_base + g
-        q_ptr = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh
-        q_row = tl.load(q_ptr + offs_d * stride_qd, mask=d_mask, other=0.0).to(tl.float32)
-        q_rows = tl.where(tl.expand_dims(g_mask[g], 1), tl.expand_dims(q_row, 0), q_rows)
-    q = q_rows
+    q_ptrs = (
+        Q_ptr
+        + batch_idx * stride_qb
+        + head_ids[:, None] * stride_qh
+        + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=g_mask[:, None] & d_mask[None, :], other=0.0).to(tl.float32)
 
     m_i = tl.full([BLOCK_G], -float("inf"), tl.float32)
     l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
@@ -88,56 +77,51 @@
         page_start = page_idx * page_size
         tokens_in_page = tl.minimum(page_size, seq_len - page_start)
 
-        t_mask = offs_t[:, None] < tokens_in_page
         kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
-
-        k = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32)
-        v = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32)
-        for t in tl.static_range(BLOCK_T):
-            kv_base = kv_h_base + t * stride_kv_page
-            k_row = tl.load(
-                KV_ptr + kv_base + offs_d * stride_kv_d,
-                mask=d_mask,
-                other=0.0,
-            ).to(tl.float32)
-            v_row = tl.load(
-                KV_ptr + kv_base + (head_dim + offs_d) * stride_kv_d,
-                mask=d_mask,
-                other=0.0,
-            ).to(tl.float32)
-            k = tl.where(tl.expand_dims(offs_t == t, 1), tl.expand_dims(k_row, 0), k)
-            v = tl.where(tl.expand_dims(offs_t == t, 1), tl.expand_dims(v_row, 0), v)
+        t_valid = offs_t[:, None] < tokens_in_page
+
+        k_ptrs = (
+            KV_ptr
+            + kv_h_base
+            + offs_t[:, None] * stride_kv_page
+            + offs_d[None, :] * stride_kv_d
+        )
+        v_ptrs = (
+            KV_ptr
+            + kv_h_base
+            + offs_t[:, None] * stride_kv_page
+            + (head_dim + offs_d)[None, :] * stride_kv_d
+        )
+        kv_mask = t_valid & d_mask[None, :]
+
+        k = tl.load(k_ptrs, mask=kv_mask, other=0.0).to(tl.float32)
+        v = tl.load(v_ptrs, mask=kv_mask, other=0.0).to(tl.float32)
 
         qk = tl.dot(q, tl.trans(k)) * scale
-        qk = tl.where(t_mask, qk, -float("inf"))
+        qk = tl.where(t_valid, qk, -float("inf"))
 
         m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
         alpha = tl.exp(m_i - m_ij)
-        p = tl.exp(qk - tl.expand_dims(m_ij, 1))
+        p = tl.exp(qk - m_ij[:, None])
         l_i = l_i * alpha + tl.sum(p, axis=1)
-        acc = acc * tl.expand_dims(alpha, 1) + tl.dot(p.to(tl.bfloat16), v)
+        acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v)
         m_i = m_ij
 
-    out_vals = acc / tl.expand_dims(l_i, 1)
-    for g in tl.static_range(BLOCK_G):
-        head_idx = head_base + g
-        o_ptr = O_ptr + batch_idx * stride_ob + head_idx * stride_oh
-        o_row = tl.load(out_vals + g)  # invalid
-        tl.store(o_ptr + offs_d * stride_od, out_vals[g, :].to(tl.bfloat16), mask=g_mask[g] & d_mask)
-
-
-def paged_attention_decode(
-    query: torch.Tensor,
-    kv_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    seq_lens: torch.Tensor,
-    scale: float,
-) -> torch.Tensor:
+    out_vals = acc / l_i[:, None]
+    o_ptrs = (
+        O_ptr
+        + batch_idx * stride_ob
+        + head_ids[:, None] * stride_oh
+        + offs_d[None, :] * stride_od
+    )
+    tl.store(o_ptrs, out_vals.to(tl.bfloat16), mask=g_mask[:, None] & d_mask[None, :])
+
+
+def _launch(query, kv_cache, block_table, seq_lens, scale, num_warps, num_stages):
     B, H, D = query.shape
     _, P, Hkv, D2 = kv_cache.shape
     assert D2 == 2 * D
     group_size = H // Hkv
-
     out = torch.empty_like(query)
 
     BLOCK_G = triton.next_power_of_2(group_size)
@@ -171,10 +155,23 @@
         BLOCK_G=BLOCK_G,
         BLOCK_T=BLOCK_T,
         BLOCK_D=BLOCK_D,
-        num_warps=4 if D <= 64 else 8,
-        num_stages=2,
+        num_warps=num_warps,
+        num_stages=num_stages,
     )
     return out
+
+
+def paged_attention_decode(
+    query: torch.Tensor,
+    kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    D = query.shape[-1]
+    if D <= 64:
+        return _launch(query, kv_cache, block_table, seq_lens, scale, num_warps=4, num_stages=2)
+    return _launch(query, kv_cache, block_table, seq_lens, scale, num_warps=8, num_stages=3)
 
 
 class Model(nn.Module):

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,231 +1,5 @@
-"""Paged attention decode kernel via Triton (GQA-aware, page-batched)."""
-import math
-
-import torch
-import torch.nn as nn
-import triton
-import triton.language as tl
-
-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
-
-
-@triton.jit
-def _paged_decode_gqa_kernel(
-    Q_ptr,
-    KV_ptr,
-    BT_ptr,
-    SL_ptr,
-    O_ptr,
-    stride_qb,
-    stride_qh,
-    stride_qd,
-    stride_kv_block,
-    stride_kv_page,
-    stride_kv_h,
-    stride_kv_d,
-    stride_bt_b,
-    stride_bt_p,
-    stride_ob,
-    stride_oh,
-    stride_od,
-    num_heads,
-    head_dim,
-    page_size,
-    group_size,
-    scale,
-    BLOCK_G: tl.constexpr,
-    BLOCK_T: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-):
-    batch_idx = tl.program_id(0)
-    kv_head_idx = tl.program_id(1)
-
-    seq_len = tl.load(SL_ptr + batch_idx)
-    num_pages = (seq_len + page_size - 1) // page_size
-
-    offs_g = tl.arange(0, BLOCK_G)
-    offs_t = tl.arange(0, BLOCK_T)
-    offs_d = tl.arange(0, BLOCK_D)
-
-    head_ids = kv_head_idx * group_size + offs_g
-    g_mask = head_ids < num_heads
-    d_mask = offs_d < head_dim
-
-    q_ptrs = (
-        Q_ptr
-        + batch_idx * stride_qb
-        + head_ids[:, None] * stride_qh
-        + offs_d[None, :] * stride_qd
-    )
-    q = tl.load(q_ptrs, mask=g_mask[:, None] & d_mask[None, :], other=0.0).to(tl.float32)
-
-    m_i = tl.full([BLOCK_G], -float("inf"), tl.float32)
-    l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
-    acc = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
-
-    for page_idx in range(num_pages):
-        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
         page_start = page_idx * page_size
         tokens_in_page = tl.minimum(page_size, seq_len - page_start)
 
         kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
-        t_valid = offs_t[:, None] < tokens_in_page
-
-        k_ptrs = (
-            KV_ptr
-            + kv_h_base
-            + offs_t[:, None] * stride_kv_page
-            + offs_d[None, :] * stride_kv_d
-        )
-        v_ptrs = (
-            KV_ptr
-            + kv_h_base
-            + offs_t[:, None] * stride_kv_page
-            + (head_dim + offs_d)[None, :] * stride_kv_d
-        )
-        kv_mask = t_valid & d_mask[None, :]
-
-        k = tl.load(k_ptrs, mask=kv_mask, other=0.0).to(tl.float32)
-        v = tl.load(v_ptrs, mask=kv_mask, other=0.0).to(tl.float32)
-
-        qk = tl.dot(q, tl.trans(k)) * scale
-        qk = tl.where(t_valid, qk, -float("inf"))
-
-        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
-        alpha = tl.exp(m_i - m_ij)
-        p = tl.exp(qk - m_ij[:, None])
-        l_i = l_i * alpha + tl.sum(p, axis=1)
-        acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v)
-        m_i = m_ij
-
-    out_vals = acc / l_i[:, None]
-    o_ptrs = (
-        O_ptr
-        + batch_idx * stride_ob
-        + head_ids[:, None] * stride_oh
-        + offs_d[None, :] * stride_od
-    )
-    tl.store(o_ptrs, out_vals.to(tl.bfloat16), mask=g_mask[:, None] & d_mask[None, :])
-
-
-def _launch(query, kv_cache, block_table, seq_lens, scale, num_warps, num_stages):
-    B, H, D = query.shape
-    _, P, Hkv, D2 = kv_cache.shape
-    assert D2 == 2 * D
-    group_size = H // Hkv
-    out = torch.empty_like(query)
-
-    BLOCK_G = triton.next_power_of_2(group_size)
-    BLOCK_T = P
-    BLOCK_D = triton.next_power_of_2(D)
-
-    grid = (B, Hkv)
-    _paged_decode_gqa_kernel[grid](
-        query,
-        kv_cache,
-        block_table,
-        seq_lens,
-        out,
-        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),
-        num_heads=H,
-        head_dim=D,
-        page_size=P,
-        group_size=group_size,
-        scale=scale,
-        BLOCK_G=BLOCK_G,
-        BLOCK_T=BLOCK_T,
-        BLOCK_D=BLOCK_D,
-        num_warps=num_warps,
-        num_stages=num_stages,
-    )
-    return out
-
-
-def paged_attention_decode(
-    query: torch.Tensor,
-    kv_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    seq_lens: torch.Tensor,
-    scale: float,
-) -> torch.Tensor:
-    D = query.shape[-1]
-    if D <= 64:
-        return _launch(query, kv_cache, block_table, seq_lens, scale, num_warps=4, num_stages=2)
-    return _launch(query, kv_cache, block_table, seq_lens, scale, num_warps=8, num_stages=3)
-
-
-class Model(nn.Module):
-    def __init__(
-        self,
-        batch: int,
-        num_heads: int,
-        num_kv_heads: int,
-        head_dim: int,
-        seq_len: int,
-        page_size: int,
-    ):
-        super().__init__()
-        assert num_heads % num_kv_heads == 0
-        self.batch = batch
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_dim = head_dim
-        self.seq_len = seq_len
-        self.page_size = page_size
-        self.group_size = num_heads // num_kv_heads
-        self.scale = 1.0 / math.sqrt(head_dim)
-        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        kv_cache: torch.Tensor,
-        block_table: torch.Tensor,
-        seq_lens: torch.Tensor,
-    ) -> torch.Tensor:
-        return paged_attention_decode(query, kv_cache, block_table, seq_lens, self.scale)
-
-
-def get_inputs():
-    B = BATCH
-    H = NUM_HEADS
-    Hkv = NUM_KV_HEADS
-    D = HEAD_DIM
-    L = SEQ_LEN
-    P =

... (645 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,5 +1,3 @@
-        page_start = page_idx * page_size
-        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
-
         kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
-        t_valid = offs_t[None, :] < tokens_in_page+        t_valid_row = offs_t[None, :] < tokens_in_page
+        t_valid_2d = offs_t[:, None] < tokens_in_page

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,3 +1 @@
-        kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
-        t_valid_row = offs_t[None, :] < tokens_in_page
-        t_valid_2d = offs_t[:, None] < tokens_in_page+        kv_mask = t_valid_2d & d_mask[None, :]

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1 +1 @@
-        kv_mask = t_valid_2d & d_mask[None, :]+        qk = tl.where(t_valid_row, qk, -float("inf"))

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1 +1 @@
-        qk = tl.where(t_valid_row, qk, -float("inf"))+        acc = acc * alpha[:, None] + tl.dot(p, v)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1 +1,247 @@
-        acc = acc * alpha[:, None] + tl.dot(p, v)+"""Paged attention decode kernel via Triton (GQA-aware, page-batched)."""
+import math
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+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
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64, "PAGE_UNROLL": 1}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64, "PAGE_UNROLL": 2}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 32, "BLOCK_D": 64, "PAGE_UNROLL": 2}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128, "PAGE_UNROLL": 1}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128, "PAGE_UNROLL": 2}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 32, "BLOCK_D": 128, "PAGE_UNROLL": 2}, num_warps=8, num_stages=4),
+    ],
+    key=["num_heads", "head_dim", "page_size", "group_size"],
+)
+@triton.jit
+def _paged_decode_gqa_kernel(
+    Q_ptr,
+    KV_ptr,
+    BT_ptr,
+    SL_ptr,
+    O_ptr,
+    stride_qb,
+    stride_qh,
+    stride_qd,
+    stride_kv_block,
+    stride_kv_page,
+    stride_kv_h,
+    stride_kv_d,
+    stride_bt_b,
+    stride_bt_p,
+    stride_ob,
+    stride_oh,
+    stride_od,
+    num_heads,
+    head_dim,
+    page_size,
+    group_size,
+    scale,
+    BLOCK_G: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    PAGE_UNROLL: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+    kv_head_idx = tl.program_id(1)
+
+    seq_len = tl.load(SL_ptr + batch_idx)
+    num_pages = (seq_len + page_size - 1) // page_size
+    tokens_per_iter = BLOCK_T // PAGE_UNROLL
+
+    offs_g = tl.arange(0, BLOCK_G)
+    offs_t = tl.arange(0, BLOCK_T)
+    offs_d = tl.arange(0, BLOCK_D)
+
+    head_ids = kv_head_idx * group_size + offs_g
+    g_mask = head_ids < num_heads
+    d_mask = offs_d < head_dim
+
+    q_ptrs = (
+        Q_ptr
+        + batch_idx * stride_qb
+        + head_ids[:, None] * stride_qh
+        + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=g_mask[:, None] & d_mask[None, :], other=0.0)
+
+    m_i = tl.full([BLOCK_G], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
+
+    for page_base in range(0, num_pages, PAGE_UNROLL):
+        k = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.bfloat16)
+        v = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.bfloat16)
+
+        for pu in tl.static_range(PAGE_UNROLL):
+            page_idx = page_base + pu
+            page_active = page_idx < num_pages
+
+            block_id = tl.load(
+                BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p,
+                mask=page_active,
+                other=0,
+            ).to(tl.int64)
+
+            page_start = page_idx * page_size
+            tokens_in_page = tl.minimum(page_size, seq_len - page_start)
+            kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
+
+            local_t = offs_t - pu * tokens_per_iter
+            t_in_page = (local_t >= 0) & (local_t < tokens_per_iter)
+            t_valid = t_in_page & (local_t < tokens_in_page) & page_active
+
+            k_ptrs = (
+                KV_ptr
+                + kv_h_base
+                + local_t[:, None] * stride_kv_page
+                + offs_d[None, :] * stride_kv_d
+            )
+            v_ptrs = (
+                KV_ptr
+                + kv_h_base
+                + local_t[:, None] * stride_kv_page
+                + (head_dim + offs_d)[None, :] * stride_kv_d
+            )
+            kv_mask = t_valid[:, None] & d_mask[None, :]
+
+            k_page = tl.load(k_ptrs, mask=kv_mask, other=0.0)
+            v_page = tl.load(v_ptrs, mask=kv_mask, other=0.0)
+
+            slot = pu * tokens_per_iter + local_t
+            k = tl.where(slot[:, None] == offs_t[:, None], k_page, k)
+            v = tl.where(slot[:, None] == offs_t[:, None], v_page, v)
+
+        t_valid_row = offs_t[None, :] < tl.minimum(seq_len - page_base * tokens_per_iter, BLOCK_T)
+        qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale
+        qk = tl.where(t_valid_row, qk, -float("inf"))
+
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        alpha = tl.exp(m_i - m_ij)
+        p = tl.exp(qk - m_ij[:, None])
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v).to(tl.float32)
+        m_i = m_ij
+
+    out_vals = acc / l_i[:, None]
+    o_ptrs = (
+        O_ptr
+        + batch_idx * stride_ob
+        + head_ids[:, None] * stride_oh
+        + offs_d[None, :] * stride_od
+    )
+    tl.store(o_ptrs, out_vals.to(tl.bfloat16), mask=g_mask[:, None] & d_mask[None, :])
+
+
+def paged_attention_decode(
+    query: torch.Tensor,
+    kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    B, H, D = query.shape
+    _, P, Hkv, D2 = kv_cache.shape
+    assert D2 == 2 * D
+    group_size = H // Hkv
+    out = torch.empty_like(query)
+
+    grid = (B, Hkv)
+    _paged_decode_gqa_kernel[grid](
+        query,
+        kv_cache,
+        block_table,
+        seq_lens,
+        out,
+        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),
+        o

... (1852 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -20,12 +20,12 @@
 
 @triton.autotune(
     configs=[
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64, "PAGE_UNROLL": 1}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64, "PAGE_UNROLL": 2}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 32, "BLOCK_D": 64, "PAGE_UNROLL": 2}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128, "PAGE_UNROLL": 1}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128, "PAGE_UNROLL": 2}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 32, "BLOCK_D": 128, "PAGE_UNROLL": 2}, num_warps=8, num_stages=4),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=4),
     ],
     key=["num_heads", "head_dim", "page_size", "group_size"],
 )
@@ -56,14 +56,12 @@
     BLOCK_G: tl.constexpr,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
-    PAGE_UNROLL: tl.constexpr,
 ):
     batch_idx = tl.program_id(0)
     kv_head_idx = tl.program_id(1)
 
     seq_len = tl.load(SL_ptr + batch_idx)
     num_pages = (seq_len + page_size - 1) // page_size
-    tokens_per_iter = BLOCK_T // PAGE_UNROLL
 
     offs_g = tl.arange(0, BLOCK_G)
     offs_t = tl.arange(0, BLOCK_T)
@@ -85,50 +83,32 @@
     l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
     acc = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
 
-    for page_base in range(0, num_pages, PAGE_UNROLL):
-        k = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.bfloat16)
-        v = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.bfloat16)
-
-        for pu in tl.static_range(PAGE_UNROLL):
-            page_idx = page_base + pu
-            page_active = page_idx < num_pages
-
-            block_id = tl.load(
-                BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p,
-                mask=page_active,
-                other=0,
-            ).to(tl.int64)
-
-            page_start = page_idx * page_size
-            tokens_in_page = tl.minimum(page_size, seq_len - page_start)
-            kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
-
-            local_t = offs_t - pu * tokens_per_iter
-            t_in_page = (local_t >= 0) & (local_t < tokens_per_iter)
-            t_valid = t_in_page & (local_t < tokens_in_page) & page_active
-
-            k_ptrs = (
-                KV_ptr
-                + kv_h_base
-                + local_t[:, None] * stride_kv_page
-                + offs_d[None, :] * stride_kv_d
-            )
-            v_ptrs = (
-                KV_ptr
-                + kv_h_base
-                + local_t[:, None] * stride_kv_page
-                + (head_dim + offs_d)[None, :] * stride_kv_d
-            )
-            kv_mask = t_valid[:, None] & d_mask[None, :]
-
-            k_page = tl.load(k_ptrs, mask=kv_mask, other=0.0)
-            v_page = tl.load(v_ptrs, mask=kv_mask, other=0.0)
-
-            slot = pu * tokens_per_iter + local_t
-            k = tl.where(slot[:, None] == offs_t[:, None], k_page, k)
-            v = tl.where(slot[:, None] == offs_t[:, None], v_page, v)
-
-        t_valid_row = offs_t[None, :] < tl.minimum(seq_len - page_base * tokens_per_iter, BLOCK_T)
+    for page_idx in range(num_pages):
+        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
+        page_start = page_idx * page_size
+        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
+
+        kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
+        t_valid_row = offs_t[None, :] < tokens_in_page
+        t_valid_2d = offs_t[:, None] < tokens_in_page
+
+        k_ptrs = (
+            KV_ptr
+            + kv_h_base
+            + offs_t[:, None] * stride_kv_page
+            + offs_d[None, :] * stride_kv_d
+        )
+        v_ptrs = (
+            KV_ptr
+            + kv_h_base
+            + offs_t[:, None] * stride_kv_page
+            + (head_dim + offs_d)[None, :] * stride_kv_d
+        )
+        kv_mask = t_valid_2d & d_mask[None, :]
+
+        k = tl.load(k_ptrs, mask=kv_mask, other=0.0)
+        v = tl.load(v_ptrs, mask=kv_mask, other=0.0)
+
         qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale
         qk = tl.where(t_valid_row, qk, -float("inf"))
 
@@ -161,6 +141,10 @@
     assert D2 == 2 * D
     group_size = H // Hkv
     out = torch.empty_like(query)
+
+    BLOCK_G = triton.next_power_of_2(group_size)
+    BLOCK_T = P
+    BLOCK_D = triton.next_power_of_2(D)
 
     grid = (B, Hkv)
     _paged_decode_gqa_kernel[grid](
@@ -186,6 +170,9 @@
         page_size=P,
         group_size=group_size,
         scale=scale,
+        BLOCK_G=BLOCK_G,
+        BLOCK_T=BLOCK_T,
+        BLOCK_D=BLOCK_D,
     )
     return out
 

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,151 +1,3 @@
-"""Paged attention decode kernel via Triton (GQA-aware, page-batched)."""
-import math
-
-import torch
-import torch.nn as nn
-import triton
-import triton.language as tl
-
-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
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=4),
-    ],
-    key=["num_heads", "head_dim", "page_size", "group_size"],
-)
-@triton.jit
-def _paged_decode_gqa_kernel(
-    Q_ptr,
-    KV_ptr,
-    BT_ptr,
-    SL_ptr,
-    O_ptr,
-    stride_qb,
-    stride_qh,
-    stride_qd,
-    stride_kv_block,
-    stride_kv_page,
-    stride_kv_h,
-    stride_kv_d,
-    stride_bt_b,
-    stride_bt_p,
-    stride_ob,
-    stride_oh,
-    stride_od,
-    num_heads,
-    head_dim,
-    page_size,
-    group_size,
-    scale,
-    BLOCK_G: tl.constexpr,
-    BLOCK_T: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-):
-    batch_idx = tl.program_id(0)
-    kv_head_idx = tl.program_id(1)
-
-    seq_len = tl.load(SL_ptr + batch_idx)
-    num_pages = (seq_len + page_size - 1) // page_size
-
-    offs_g = tl.arange(0, BLOCK_G)
-    offs_t = tl.arange(0, BLOCK_T)
-    offs_d = tl.arange(0, BLOCK_D)
-
-    head_ids = kv_head_idx * group_size + offs_g
-    g_mask = head_ids < num_heads
-    d_mask = offs_d < head_dim
-
-    q_ptrs = (
-        Q_ptr
-        + batch_idx * stride_qb
-        + head_ids[:, None] * stride_qh
-        + offs_d[None, :] * stride_qd
-    )
-    q = tl.load(q_ptrs, mask=g_mask[:, None] & d_mask[None, :], other=0.0)
-
-    m_i = tl.full([BLOCK_G], -float("inf"), tl.float32)
-    l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
-    acc = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
-
-    for page_idx in range(num_pages):
-        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
-        page_start = page_idx * page_size
-        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
-
-        kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
-        t_valid_row = offs_t[None, :] < tokens_in_page
-        t_valid_2d = offs_t[:, None] < tokens_in_page
-
-        k_ptrs = (
-            KV_ptr
-            + kv_h_base
-            + offs_t[:, None] * stride_kv_page
-            + offs_d[None, :] * stride_kv_d
-        )
-        v_ptrs = (
-            KV_ptr
-            + kv_h_base
-            + offs_t[:, None] * stride_kv_page
-            + (head_dim + offs_d)[None, :] * stride_kv_d
-        )
-        kv_mask = t_valid_2d & d_mask[None, :]
-
-        k = tl.load(k_ptrs, mask=kv_mask, other=0.0)
-        v = tl.load(v_ptrs, mask=kv_mask, other=0.0)
-
-        qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale
-        qk = tl.where(t_valid_row, qk, -float("inf"))
-
-        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
-        alpha = tl.exp(m_i - m_ij)
-        p = tl.exp(qk - m_ij[:, None])
-        l_i = l_i * alpha + tl.sum(p, axis=1)
-        acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v).to(tl.float32)
-        m_i = m_ij
-
-    out_vals = acc / l_i[:, None]
-    o_ptrs = (
-        O_ptr
-        + batch_idx * stride_ob
-        + head_ids[:, None] * stride_oh
-        + offs_d[None, :] * stride_od
-    )
-    tl.store(o_ptrs, out_vals.to(tl.bfloat16), mask=g_mask[:, None] & d_mask[None, :])
-
-
-def paged_attention_decode(
-    query: torch.Tensor,
-    kv_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    seq_lens: torch.Tensor,
-    scale: float,
-) -> torch.Tensor:
-    B, H, D = query.shape
-    _, P, Hkv, D2 = kv_cache.shape
-    assert D2 == 2 * D
-    group_size = H // Hkv
-    out = torch.empty_like(query)
-
-    BLOCK_G = triton.next_power_of_2(group_size)
-    BLOCK_T = P
-    BLOCK_D = triton.next_power_of_2(D)
-
     grid = (B, Hkv)
     _paged_decode_gqa_kernel[grid](
         query,
@@ -170,65 +22,4 @@
         page_size=P,
         group_size=group_size,
         scale=scale,
-        BLOCK_G=BLOCK_G,
-        BLOCK_T=BLOCK_T,
-        BLOCK_D=BLOCK_D,
-    )
-    return out
-
-
-class Model(nn.Module):
-    def __init__(
-        self,
-        batch: int,
-        num_heads: int,
-        num_kv_heads: int,
-        head_dim: int,
-        seq_len: int,
-        page_size: int,
-    ):
-        super().__init__()
-        assert num_heads % num_kv_heads == 0
-        self.batch = batch
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_dim = head_dim
-        self.seq_len = seq_len
-        self.page_size = page_size
-        self.group_size = num_heads // num_kv_heads
-        self.scale = 1.0 / math.sqrt(head_dim)
-        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        kv_cache: torch.Tensor,
-        block_table: torch.Tensor,
-        seq_lens: torch.Tensor,
-    ) -> torch.Tensor:
-        return paged_attention_decode(query, kv_cache, block_table, seq_lens, self.scale)
-
-
-def get_inputs():
-    B = BATCH
-    H = NUM_HEADS
-    Hkv = NUM_KV_HEADS
-    D = HEAD_DIM
-    L = SEQ_LEN
-    P = PAGE_SIZE
-
-    pages_per_seq = (L + P - 1) // P
-    total_pages

... (533 more chars)

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,25 +1,9 @@
-    grid = (B, Hkv)
-    _paged_decode_gqa_kernel[grid](
-        query,
-        kv_cache,
-        block_table,
-        seq_lens,
-        out,
-        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),
-        num_heads=H,
-        head_dim=D,
-        page_size=P,
-        group_size=group_size,
-        scale=scale,
-    )+    configs=[
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=4),
+    ],

--- a//workspace/problems/03_paged_attention/solution.py
+++ b//workspace/problems/03_paged_attention/solution.py
@@ -1,9 +1,230 @@
-    configs=[
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_G": 4, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_G": 8, "BLOCK_T": 16, "BLOCK_D": 128}, num_warps=8, num_stages=4),
-    ],+"""Paged attention decode kernel via Triton (GQA-aware, page-batched)."""
+import math
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+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
+
+
+@triton.jit
+def _paged_decode_gqa_kernel(
+    Q_ptr,
+    KV_ptr,
+    BT_ptr,
+    SL_ptr,
+    O_ptr,
+    stride_qb,
+    stride_qh,
+    stride_qd,
+    stride_kv_block,
+    stride_kv_page,
+    stride_kv_h,
+    stride_kv_d,
+    stride_bt_b,
+    stride_bt_p,
+    stride_ob,
+    stride_oh,
+    stride_od,
+    num_heads,
+    head_dim,
+    page_size,
+    group_size,
+    scale,
+    BLOCK_G: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+    kv_head_idx = tl.program_id(1)
+
+    seq_len = tl.load(SL_ptr + batch_idx)
+    num_pages = (seq_len + page_size - 1) // page_size
+
+    offs_g = tl.arange(0, BLOCK_G)
+    offs_t = tl.arange(0, BLOCK_T)
+    offs_d = tl.arange(0, BLOCK_D)
+
+    head_ids = kv_head_idx * group_size + offs_g
+    g_mask = head_ids < num_heads
+    d_mask = offs_d < head_dim
+
+    q_ptrs = (
+        Q_ptr
+        + batch_idx * stride_qb
+        + head_ids[:, None] * stride_qh
+        + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=g_mask[:, None] & d_mask[None, :], other=0.0)
+
+    m_i = tl.full([BLOCK_G], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_G, BLOCK_D], dtype=tl.float32)
+
+    for page_idx in range(num_pages):
+        block_id = tl.load(BT_ptr + batch_idx * stride_bt_b + page_idx * stride_bt_p).to(tl.int64)
+        page_start = page_idx * page_size
+        tokens_in_page = tl.minimum(page_size, seq_len - page_start)
+
+        kv_h_base = block_id * stride_kv_block + kv_head_idx * stride_kv_h
+        t_valid_row = offs_t[None, :] < tokens_in_page
+        t_valid_2d = offs_t[:, None] < tokens_in_page
+
+        k_ptrs = (
+            KV_ptr
+            + kv_h_base
+            + offs_t[:, None] * stride_kv_page
+            + offs_d[None, :] * stride_kv_d
+        )
+        v_ptrs = (
+            KV_ptr
+            + kv_h_base
+            + offs_t[:, None] * stride_kv_page
+            + (head_dim + offs_d)[None, :] * stride_kv_d
+        )
+        kv_mask = t_valid_2d & d_mask[None, :]
+
+        k = tl.load(k_ptrs, mask=kv_mask, other=0.0)
+        v = tl.load(v_ptrs, mask=kv_mask, other=0.0)
+
+        qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale
+        qk = tl.where(t_valid_row, qk, -float("inf"))
+
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        alpha = tl.exp(m_i - m_ij)
+        p = tl.exp(qk - m_ij[:, None])
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v).to(tl.float32)
+        m_i = m_ij
+
+    out_vals = acc / l_i[:, None]
+    o_ptrs = (
+        O_ptr
+        + batch_idx * stride_ob
+        + head_ids[:, None] * stride_oh
+        + offs_d[None, :] * stride_od
+    )
+    tl.store(o_ptrs, out_vals.to(tl.bfloat16), mask=g_mask[:, None] & d_mask[None, :])
+
+
+def paged_attention_decode(
+    query: torch.Tensor,
+    kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    B, H, D = query.shape
+    _, P, Hkv, D2 = kv_cache.shape
+    assert D2 == 2 * D
+    group_size = H // Hkv
+    out = torch.empty_like(query)
+
+    BLOCK_G = triton.next_power_of_2(group_size)
+    BLOCK_T = P
+    BLOCK_D = triton.next_power_of_2(D)
+
+    if D <= 64:
+        num_warps, num_stages = 4, 3
+    else:
+        num_warps, num_stages = 8, 3
+
+    grid = (B, Hkv)
+    _paged_decode_gqa_kernel[grid](
+        query,
+        kv_cache,
+        block_table,
+        seq_lens,
+        out,
+        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),
+        num_heads=H,
+        head_dim=D,
+        page_size=P,
+        group_size=group_size,
+        scale=scale,
+        BLOCK_G=BLOCK_G,
+        BLOCK_T=BLOCK_T,
+        BLOCK_D=BLOCK_D,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return out
+
+
+class Model(nn.Module):
+    def __init__(
+        self,
+        batch: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        seq_len: int,
+        page_size: int,
+    ):
+        super().__init__()
+        assert num_heads % num_kv_heads == 0
+        self.batch = batch
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.seq_len = seq_len
+        self.page_size = page_size
+        self.group_size = num_heads // num_kv_heads
+        se

... (1133 more chars)