diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py
index c2358d84..208bb65c 100644
--- a/python/freetoken/kernel/triton/attention.py
+++ b/python/freetoken/kernel/triton/attention.py
@@ -187,11 +187,13 @@ def _decode_grouped_stage1_kernel(
     # VALID_BLOCK_H == min(cap, GROUP) is the number of query heads actually handled per program;
     # BLOCK_H is the power-of-two tile size for the head axis (tl.arange requires a power of two),
     # so a non-power-of-two GQA group (e.g. 24/4 == 6) rounds the tile up and masks the extra
-    # lanes. Each kv head spans cdiv(GROUP, VALID_BLOCK_H) head blocks.
-    kv_head = head_block_id // tl.cdiv(GROUP, VALID_BLOCK_H)
-    q_heads = head_block_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
-    mask_h = q_heads < (head_block_id + 1) * VALID_BLOCK_H
-    mask_h = mask_h & (q_heads < NUM_Q_HEADS)
+    # lanes. Each kv head spans cdiv(GROUP, VALID_BLOCK_H) head blocks. Keep those
+    # blocks within the kv group so a partial tail block cannot spill into the next one.
+    blocks_per_kv = tl.cdiv(GROUP, VALID_BLOCK_H)
+    kv_head = head_block_id // blocks_per_kv
+    block_in_kv = head_block_id % blocks_per_kv
+    q_heads = kv_head * GROUP + block_in_kv * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = (q_heads < (kv_head + 1) * GROUP) & (q_heads < NUM_Q_HEADS)
 
     offs_d = tl.arange(0, BLOCK_D)
     offs_dv = tl.arange(0, BLOCK_DV)
@@ -400,7 +402,11 @@ def decode_paged_attention(
     block_dv = triton.next_power_of_2(head_dim)
 
     _decode_grouped_stage1_kernel[
-        (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits)
+        (
+            batch,
+            num_kv_heads * triton.cdiv(group, valid_block_h),
+            max_kv_splits,
+        )
     ](
         q,
         k_cache,
diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py
index 6f4afca9..431b5a1b 100644
--- a/tests/kernels/test_triton_attention.py
+++ b/tests/kernels/test_triton_attention.py
@@ -315,11 +315,16 @@ def test_decode_triton_attention_matches_reference(
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA")
-@pytest.mark.parametrize(("num_q_heads", "num_kv_heads"), [(24, 4), (20, 4), (28, 4)])
+@pytest.mark.parametrize(
+    ("num_q_heads", "num_kv_heads"),
+    [(24, 4), (20, 4), (28, 4), (48, 2), (40, 2), (80, 4)],
+)
 def test_decode_triton_attention_non_pow2_group(num_q_heads: int, num_kv_heads: int):
-    """GQA groups that are not a power of two (e.g. Qwen3.6-27B's 24/4 == 6). The grouped
-    decode tiles the head axis to a power of two (tl.arange constraint) and masks the extra
-    lanes; the result must still match the reference."""
+    """Non-power-of-two GQA groups, including groups that span multiple 16-head tiles.
+
+    The grouped decode must mask a group's partial tail tile without assigning its lanes to
+    the next KV head.
+    """
     from freetoken.kernel.triton.attention import decode_paged_attention
 
     torch.manual_seed(3)