From f57903f36c092556312fc42878860f147bc1ad4d Mon Sep 17 00:00:00 2001
From: Thom Hadfield <100659456+festeraeb@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:19:49 +0000
Subject: [PATCH] fix(attention): size the decode KV tile to the shared-memory
 budget

Assisted-by: Claude (Anthropic)
---
 python/freetoken/kernel/triton/attention.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py
index d19457cffc..e49346dd5d 100644
--- a/python/freetoken/kernel/triton/attention.py
+++ b/python/freetoken/kernel/triton/attention.py
@@ -454,9 +454,15 @@ def decode_paged_attention(
     # a grid that fits in one wave runs about one program per SM, which then needs wider tiles to keep enough loads in flight
     block_n = 64 if grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index) else _MIN_BLOCK_KV
     num_stages = 3
-    # stage 1 keeps one K tile and one V tile in shared memory
-    if 2 * block_n * block_d * k_cache.element_size() > 0.8 * _optin_smem_bytes(q.device.index):
+    # stage 1 keeps the Q tile and one K and one V tile in shared memory
+    smem = _optin_smem_bytes(q.device.index)
+    elem = k_cache.element_size()
+    if 2 * block_n * block_d * elem > 0.8 * smem:
         block_n = _MIN_BLOCK_KV
+    # 64 KiB parts (sm_75) cannot fit a head_dim 512 Q tile next to 32-row K/V tiles:
+    # halve the KV tile down to tl.dot's 16-row minimum until the three tiles fit.
+    while smem and block_n > 16 and (block_h * block_d + 2 * block_n * block_d) * elem > smem:
+        block_n //= 2
     if sliding_window:
         split_len = triton.cdiv(triton.cdiv(sliding_window, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV
         if split_len <= block_n: