Skip to content

fix(attention): map grouped decode heads to KV groups #177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
taking-lying-flat wants to merge 1 commit into
base: main
Choose a base branch
from
+21 −10
Open

fix(attention): map grouped decode heads to KV groups #177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions python/freetoken/kernel/triton/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
13 changes: 9 additions & 4 deletions tests/kernels/test_triton_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down