diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 28effa5ea..bfb59f7df 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -19,6 +19,15 @@ _STAGE2_TILE_ELEMS = 8192 +def _kv_offset_dtype(k_cache: torch.Tensor, v_cache: torch.Tensor): + # Bound physical offsets, not request lengths. + max_offset = max( + sum((size - 1) * stride for size, stride in zip(cache.shape, cache.stride())) + for cache in (k_cache, v_cache) + ) + return tl.int32 if max_offset < 1 << 31 else tl.int64 + + def _stage2_launch(max_kv_splits: int, block_dv: int) -> tuple[int, int]: """(splits merged per step, num_warps) for stage 2; the tile is at most the split cap rounded up to a power of two.""" tile = min(triton.next_power_of_2(max_kv_splits), max(4, _STAGE2_TILE_ELEMS // block_dv), 64) @@ -97,6 +106,7 @@ def _paged_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + KV_OFFSET_DTYPE: tl.constexpr, ): q_tok = tl.program_id(0) q_head = tl.program_id(1) @@ -135,7 +145,7 @@ def _paged_attention_kernel( skip_tile = tl.max(mask_n.to(tl.int32), axis=0) == 0 if not skip_tile: - slots = tl.load(indices_ptr + kv_start + offs_n, mask=offs_n < kv_len, other=0) + slots = tl.load(indices_ptr + kv_start + offs_n, mask=offs_n < kv_len, other=0).to(KV_OFFSET_DTYPE) k = tl.load( k_ptr + slots[:, None] * stride_ks @@ -217,6 +227,7 @@ def _decode_grouped_stage1_kernel( D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, + KV_OFFSET_DTYPE: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -268,7 +279,7 @@ def _decode_grouped_stage1_kernel( rel_offs = rel_start + tl.arange(0, BLOCK_N) mask_n = rel_offs < split_end logical_offs = effective_start + rel_offs - slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0) + slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) k = tl.load( k_ptr + slots[None, :] * stride_ks + k_base_offsets, @@ -492,6 +503,7 @@ def decode_paged_attention( attn_lse.stride(0), attn_lse.stride(1), attn_lse.stride(2), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), GROUP=group, NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, @@ -568,6 +580,7 @@ def _extend_attention_kernel( SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, HAS_BLOCKS: tl.constexpr, + KV_OFFSET_DTYPE: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -626,7 +639,7 @@ def _extend_attention_kernel( skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 if not skip_tile: - slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) + slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) k = tl.load( k_ptr + slots[None, :] * stride_ks @@ -745,6 +758,7 @@ def _extend_attention_split_kernel( SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, HAS_BLOCKS: tl.constexpr, + KV_OFFSET_DTYPE: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -793,7 +807,7 @@ def _extend_attention_split_kernel( prefix_len + block_start, prefix_end // BLOCK_N * BLOCK_N, prefix_end, BLOCK_M, BLOCK_N, SLIDING_WINDOW ) for start_n in tl.range(u0, u1, BLOCK_N): - slots = tl.load(kv_indices_ptr + kv_start + start_n + offs_n) + slots = tl.load(kv_indices_ptr + kv_start + start_n + offs_n).to(KV_OFFSET_DTYPE) m_i, l_i, acc = _extend_attend_tile( q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, offs_d, offs_dv, mask_d, mask_dv, None, None, False, @@ -805,7 +819,7 @@ def _extend_attention_split_kernel( final_mask = mask_n[None, :] if SLIDING_WINDOW > 0: final_mask = final_mask & ((kv_offsets[None, :] + SLIDING_WINDOW) > q_abs_pos[:, None]) - slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) + slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) m_i, l_i, acc = _extend_attend_tile( q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True, @@ -942,6 +956,7 @@ def extend_paged_attention( v_cache.stride(1), o.stride(0), o.stride(1), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), GROUP=num_q_heads // num_kv_heads, D=head_dim, BLOCK_D=block_d, @@ -976,6 +991,7 @@ def extend_paged_attention( v_cache.stride(1), o.stride(0), o.stride(1), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), GROUP=num_q_heads // num_kv_heads, D=head_dim, BLOCK_D=block_d, @@ -1048,6 +1064,7 @@ def paged_attention( v_cache.stride(1), o.stride(0), o.stride(1), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), GROUP=num_q_heads // num_kv_heads, D=head_dim, BLOCK_D=block_d, diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 60a5e2747..d86467b5c 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -49,6 +49,68 @@ def _reference_paged_attention( return torch.stack(outs, dim=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs a GPU") +@pytest.mark.parametrize("path", ["paged", "decode", "extend", "split_extend"]) +@pytest.mark.parametrize("offset_bits", [31, 32]) +def test_triton_attention_high_physical_slots(path: str, offset_bits: int): + from freetoken.kernel.triton.attention import ( + decode_paged_attention, + extend_paged_attention, + paged_attention, + ) + + torch.manual_seed(1347) + device = torch.device("cuda") + head_dim, num_kv_heads, num_q_heads = 64, 2, 4 + stride = num_kv_heads * head_dim + low, count = 128, 66 + high = (1 << 32) // stride + low + required = (high + count) * stride * 2 + if torch.cuda.mem_get_info()[0] < required + (1 << 30): + # Release inactive cached blocks before deciding this case cannot fit. + torch.cuda.empty_cache() + if torch.cuda.mem_get_info()[0] < required + (1 << 30): + pytest.skip("High-offset attention regression needs 9 GiB of free GPU memory") + + # Share K/V storage to keep the real allocation near 8 GiB. Int32 overflow + # aliases initialized low rows, so a regression fails numerically, not by OOB. + cache = torch.empty(high + count, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16) + cache[low : low + count].fill_(-2) + cache[high : high + count].normal_(mean=3, std=2) + if offset_bits == 31: + # Keep the old negative offset inside allocated guard storage too. + guard = (1 << 31) // stride + cache = cache[guard:] + high -= guard + q = torch.randn(1, num_q_heads, head_dim, device=device, dtype=cache.dtype) + indices = torch.arange(high, high + count, device=device, dtype=torch.int32) + indptr = torch.tensor([0, count], device=device, dtype=torch.int32) + q_to_req = torch.zeros(1, device=device, dtype=torch.int32) + q_positions = torch.tensor([count - 1], device=device, dtype=torch.int64) + sm_scale = head_dim**-0.5 + expected = _reference_paged_attention( + q, cache, cache, indptr, indices, q_to_req, q_positions, sm_scale, None, + ) + + if path == "paged": + actual = paged_attention(q, cache, cache, indptr, indices, q_to_req, q_positions, sm_scale) + elif path == "decode": + mid = torch.empty(1, num_q_heads, 64, head_dim, device=device, dtype=torch.float32) + lse = torch.empty(1, num_q_heads, 64, device=device, dtype=torch.float32) + actual = decode_paged_attention(q, cache, cache, indptr, indices, q_positions, mid, lse, sm_scale) + else: + qo_indptr = torch.tensor([0, 1], device=device, dtype=torch.int32) + prefix_lens = torch.tensor([count - 1], device=device, dtype=torch.int32) + # 65 cached tokens cover both the full-tile and masked split-prefix loads. + extend = cache[high + count - 1 : high + count] if path == "split_extend" else None + actual = extend_paged_attention( + q, cache, cache, qo_indptr, indptr, indices, prefix_lens, 1, sm_scale, + k_extend=extend, v_extend=extend, + ) + + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + def test_triton_backend_passes_attention_sinks_to_paged_kernel(monkeypatch): from freetoken.attention import AttentionSpec from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata