From 29c98e77f9fd5b19ec287f2e83d425b543e43885 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Thu, 12 Nov 2026 11:08:46 +0800 Subject: [PATCH 1/3] perf(attention): add packed fp8 kv reads on gfx1201 --- python/freetoken/kernel/triton/attention.py | 42 +++++++++++++++++++-- python/freetoken/kernel/triton/kv_rows.py | 23 +++++++++++ tests/kernels/test_triton_attention.py | 42 +++++++++++++++++++-- 3 files changed, 101 insertions(+), 6 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index b1a4455a4..3c06a7ec8 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,7 +6,8 @@ import triton import triton.language as tl -from freetoken.kernel.triton.kv_rows import kv_row_scale, kv_scale_args, kv_tile, kv_tile_staged +from freetoken.kernel.triton.e4m3_compat import FORCE_EMU +from freetoken.kernel.triton.kv_rows import kv_row_scale, kv_scale_args, kv_tile, kv_tile_gfx12_bf16, kv_tile_staged from freetoken.kvcache.format import NONE, KVFormat @@ -65,6 +66,13 @@ def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) +@functools.lru_cache(maxsize=None) +def _gfx1201(device_index: int) -> bool: + if torch.version.hip is None: + return False + return getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "").split(":")[0] == "gfx1201" + + def decode_kv_split_cap(batch: int, num_q_heads: int, num_kv_heads: int, device_index: int) -> int: """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS].""" head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads))) @@ -262,6 +270,7 @@ def _decode_grouped_stage1_kernel( SLIDING_WINDOW: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, KV_FMT: tl.constexpr, + PACKED_GFX12_KV: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -316,12 +325,26 @@ def _decode_grouped_stage1_kernel( logical_offs = effective_start + rel_offs slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) - k = kv_tile(k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None], KV_FMT) + if PACKED_GFX12_KV: + packed_d = tl.arange(0, BLOCK_D // 4) + k = kv_tile_gfx12_bf16( + (k_ptr + kv_head * stride_kh).to(tl.pointer_type(tl.uint32)) + + slots[None, :] * (stride_ks // 4) + packed_d[:, None], + mask_n[None, :] & (packed_d[:, None] < D // 4), True) + else: + k = kv_tile(k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None], KV_FMT) scale_k = kv_row_scale(k_scale_ptr + slots * stride_kss + kv_head, mask_n, KV_FMT) scores = tl.dot(q, k) * sm_scale * scale_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - v = kv_tile(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], KV_FMT) + if PACKED_GFX12_KV: + packed_dv = tl.arange(0, BLOCK_DV // 4) + v = kv_tile_gfx12_bf16( + (v_ptr + kv_head * stride_vh).to(tl.pointer_type(tl.uint32)) + + slots[:, None] * (stride_vs // 4) + packed_dv[None, :], + mask_n[:, None] & (packed_dv[None, :] < DV // 4), False) + else: + v = kv_tile(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], KV_FMT) scale_v = kv_row_scale(v_scale_ptr + slots * stride_vss + kv_head, mask_n, KV_FMT) m_new = tl.maximum(tl.max(scores, axis=1), m_i) @@ -515,6 +538,18 @@ def decode_paged_attention( if kv_format.quantized: block_n, num_stages, num_warps = _quantized_decode_launch(block_d, block_n, num_stages) + packed_gfx12_kv = ( + kv_format.name == "fp8_per_token_head" and not FORCE_EMU + and q.dtype == torch.bfloat16 and head_dim == 128 + and batch == 1 and group in (4, 8) and _gfx1201(q.device.index) + and k_cache.stride(-1) == v_cache.stride(-1) == 1 + and k_cache.data_ptr() % 4 == 0 and v_cache.data_ptr() % 4 == 0 + and all(s % 4 == 0 for s in (*k_cache.stride()[:2], *v_cache.stride()[:2])) + ) + if packed_gfx12_kv: + # Limit live conversion data while expanding packed KV tiles. + block_n, num_stages, num_warps = 64, 1, 2 + k_scale_arg, v_scale_arg, stride_kss, stride_vss = kv_scale_args(kv_format, k_cache, v_cache, k_scale, v_scale, q) _decode_grouped_stage1_kernel[grid]( q, @@ -545,6 +580,7 @@ def decode_paged_attention( attn_lse.stride(2), KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), KV_FMT=kv_format.name, + PACKED_GFX12_KV=packed_gfx12_kv, GROUP=group, NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, diff --git a/python/freetoken/kernel/triton/kv_rows.py b/python/freetoken/kernel/triton/kv_rows.py index afa19bfbe..f7f55f0a7 100644 --- a/python/freetoken/kernel/triton/kv_rows.py +++ b/python/freetoken/kernel/triton/kv_rows.py @@ -15,6 +15,29 @@ from freetoken.kvcache.format import KVFormat +@triton.jit +def kv_tile_gfx12_bf16(ptrs, mask, TRANSPOSE: tl.constexpr): + # Each uint32 holds four E4M3FN codes. Force e64 for both word selectors. + codes = tl.load(ptrs, mask=mask, other=0, cache_modifier=".cg") + lo = tl.inline_asm_elementwise( + "v_cvt_pk_f32_fp8_e64 $0, $1 op_sel:[0,0]", "=v,v", [codes], + dtype=tl.uint64, is_pure=True, pack=1) + hi = tl.inline_asm_elementwise( + "v_cvt_pk_f32_fp8_e64 $0, $1 op_sel:[1,0]", "=v,v", [codes], + dtype=tl.uint64, is_pure=True, pack=1) + # Every finite E4M3FN value is exactly representable in BF16. + a = (lo >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + b = (lo >> 48).to(tl.uint16).to(tl.bfloat16, bitcast=True) + c = (hi >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + d = (hi >> 48).to(tl.uint16).to(tl.bfloat16, bitcast=True) + if TRANSPOSE: + a, b, c, d = tl.trans(a), tl.trans(b), tl.trans(c), tl.trans(d) + values = tl.interleave(tl.interleave(a, c), tl.interleave(b, d)) + if TRANSPOSE: + values = tl.trans(values) + return values + + @triton.jit def kv_tile(ptrs, mask, FMT: tl.constexpr): """One K or V tile in the dtype ``tl.dot`` takes.""" diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index ad1120181..974e38a65 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -206,12 +206,14 @@ def test_store_kv_rows_round_trips_within_the_format_error_bound(fmt_name: str): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") @pytest.mark.parametrize("fmt_name", KV_QUANT_FORMATS) -@pytest.mark.parametrize("path", ["paged", "decode", "extend", "split_extend"]) +@pytest.mark.parametrize(("path", "num_q_heads"), [ + ("paged", 4), ("decode", 4), ("decode", 8), ("decode", 16), ("extend", 4), ("split_extend", 4), +]) @pytest.mark.parametrize("head_dim", [64, 128, 256, 512]) -def test_quantized_kv_attention_matches_dequantized_reference(fmt_name: str, path: str, head_dim: int): +def test_quantized_kv_attention_matches_dequantized_reference(fmt_name: str, path: str, num_q_heads: int, head_dim: int): torch.manual_seed(11) device = torch.device("cuda") - num_kv_heads, num_q_heads = 2, 4 + num_kv_heads = 2 count, capacity = 66, 200 q_len = 2 if "extend" in path else 1 row = num_kv_heads * head_dim @@ -253,6 +255,40 @@ def test_quantized_kv_attention_matches_dequantized_reference(fmt_name: str, pat torch.testing.assert_close(actual.float(), expected, atol=2e-2 * expected.abs().max().item(), rtol=2e-2) +@pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is None, reason="needs ROCm") +@pytest.mark.parametrize("transpose", [False, True]) +def test_rdna4_kv_decode_preserves_every_finite_fp8_code(transpose: bool): + import triton + import triton.language as tl + + from freetoken.kernel.triton.attention import _gfx1201 + from freetoken.kernel.triton.kv_rows import kv_tile_gfx12_bf16 + + if not _gfx1201(torch.cuda.current_device()): + pytest.skip("packed KV conversion requires gfx1201") + + @triton.jit + def decode_codes(x, y, TRANSPOSE: tl.constexpr): + rows = tl.arange(0, 64) + words = tl.arange(0, 32) + if TRANSPOSE: + ptrs = x + rows[None, :] * 32 + words[:, None] + else: + ptrs = x + rows[:, None] * 32 + words[None, :] + values = kv_tile_gfx12_bf16(ptrs, tl.full(ptrs.shape, True, tl.int1), TRANSPOSE) + if TRANSPOSE: + values = tl.trans(values) + tl.store(y + rows[:, None] * 128 + tl.arange(0, 128)[None, :], values) + + codes = torch.arange(256, dtype=torch.uint8).repeat(32) + expected = codes.view(torch.float8_e4m3fn).bfloat16().cuda() + actual = torch.empty_like(expected) + decode_codes[(1,)](codes.cuda().view(torch.uint32), actual, transpose, num_warps=2) + finite = expected.isfinite() + # Integer comparison also checks signed zero and the byte order in each word. + assert torch.equal(actual.view(torch.int16)[finite], expected.view(torch.int16)[finite]) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") @pytest.mark.parametrize("path", ["paged", "decode", "extend", "split_extend"]) def test_quantized_kv_attention_high_physical_slots(path: str): From 252528e6ed43ea210659d1f29cac72fe4adc2851 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Thu, 12 Nov 2026 13:48:26 +0800 Subject: [PATCH 2/3] refactor(attention): encapsulate packed fp8 kv readers --- python/freetoken/kernel/triton/attention.py | 48 ++++------------- python/freetoken/kernel/triton/kv_rows.py | 60 +++++++++++++++++++++ tests/kernels/test_triton_attention.py | 18 +++---- 3 files changed, 76 insertions(+), 50 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 3c06a7ec8..b0b0cc678 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,8 +6,7 @@ import triton import triton.language as tl -from freetoken.kernel.triton.e4m3_compat import FORCE_EMU -from freetoken.kernel.triton.kv_rows import kv_row_scale, kv_scale_args, kv_tile, kv_tile_gfx12_bf16, kv_tile_staged +from freetoken.kernel.triton.kv_rows import kv_decode_config, kv_row_scale, kv_scale_args, kv_tile, kv_tile_rows, kv_tile_staged from freetoken.kvcache.format import NONE, KVFormat @@ -66,13 +65,6 @@ def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) -@functools.lru_cache(maxsize=None) -def _gfx1201(device_index: int) -> bool: - if torch.version.hip is None: - return False - return getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "").split(":")[0] == "gfx1201" - - def decode_kv_split_cap(batch: int, num_q_heads: int, num_kv_heads: int, device_index: int) -> int: """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS].""" head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads))) @@ -270,7 +262,7 @@ def _decode_grouped_stage1_kernel( SLIDING_WINDOW: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, KV_FMT: tl.constexpr, - PACKED_GFX12_KV: tl.constexpr, + KV_PACK: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -311,8 +303,6 @@ def _decode_grouped_stage1_kernel( acc = tl.zeros((BLOCK_H, BLOCK_DV), dtype=tl.float32) q_offsets = batch_id * stride_qt + q_heads[:, None] * stride_qh + offs_d[None, :] - k_base_offsets = kv_head * stride_kh + offs_d[:, None] - v_base_offsets = kv_head * stride_vh + offs_dv[None, :] if split_end > split_start: q = tl.load(q_ptr + q_offsets, mask=mask_h[:, None] & mask_d[None, :], other=0.0) @@ -325,26 +315,14 @@ def _decode_grouped_stage1_kernel( logical_offs = effective_start + rel_offs slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) - if PACKED_GFX12_KV: - packed_d = tl.arange(0, BLOCK_D // 4) - k = kv_tile_gfx12_bf16( - (k_ptr + kv_head * stride_kh).to(tl.pointer_type(tl.uint32)) - + slots[None, :] * (stride_ks // 4) + packed_d[:, None], - mask_n[None, :] & (packed_d[:, None] < D // 4), True) - else: - k = kv_tile(k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None], KV_FMT) + k = kv_tile_rows(k_ptr, slots, kv_head, stride_ks, stride_kh, mask_n, + D, BLOCK_D, KV_FMT, KV_PACK, True) scale_k = kv_row_scale(k_scale_ptr + slots * stride_kss + kv_head, mask_n, KV_FMT) scores = tl.dot(q, k) * sm_scale * scale_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - if PACKED_GFX12_KV: - packed_dv = tl.arange(0, BLOCK_DV // 4) - v = kv_tile_gfx12_bf16( - (v_ptr + kv_head * stride_vh).to(tl.pointer_type(tl.uint32)) - + slots[:, None] * (stride_vs // 4) + packed_dv[None, :], - mask_n[:, None] & (packed_dv[None, :] < DV // 4), False) - else: - v = kv_tile(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], KV_FMT) + v = kv_tile_rows(v_ptr, slots, kv_head, stride_vs, stride_vh, mask_n, + DV, BLOCK_DV, KV_FMT, KV_PACK, False) scale_v = kv_row_scale(v_scale_ptr + slots * stride_vss + kv_head, mask_n, KV_FMT) m_new = tl.maximum(tl.max(scores, axis=1), m_i) @@ -538,17 +516,9 @@ def decode_paged_attention( if kv_format.quantized: block_n, num_stages, num_warps = _quantized_decode_launch(block_d, block_n, num_stages) - packed_gfx12_kv = ( - kv_format.name == "fp8_per_token_head" and not FORCE_EMU - and q.dtype == torch.bfloat16 and head_dim == 128 - and batch == 1 and group in (4, 8) and _gfx1201(q.device.index) - and k_cache.stride(-1) == v_cache.stride(-1) == 1 - and k_cache.data_ptr() % 4 == 0 and v_cache.data_ptr() % 4 == 0 - and all(s % 4 == 0 for s in (*k_cache.stride()[:2], *v_cache.stride()[:2])) + kv_pack, block_n, num_stages, num_warps = kv_decode_config( + kv_format, q, k_cache, v_cache, block_n, num_stages, num_warps, ) - if packed_gfx12_kv: - # Limit live conversion data while expanding packed KV tiles. - block_n, num_stages, num_warps = 64, 1, 2 k_scale_arg, v_scale_arg, stride_kss, stride_vss = kv_scale_args(kv_format, k_cache, v_cache, k_scale, v_scale, q) _decode_grouped_stage1_kernel[grid]( @@ -580,7 +550,7 @@ def decode_paged_attention( attn_lse.stride(2), KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), KV_FMT=kv_format.name, - PACKED_GFX12_KV=packed_gfx12_kv, + KV_PACK=kv_pack, GROUP=group, NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, diff --git a/python/freetoken/kernel/triton/kv_rows.py b/python/freetoken/kernel/triton/kv_rows.py index f7f55f0a7..a659495b9 100644 --- a/python/freetoken/kernel/triton/kv_rows.py +++ b/python/freetoken/kernel/triton/kv_rows.py @@ -1,11 +1,14 @@ from __future__ import annotations +import functools + import torch import triton import triton.language as tl from triton.language import target_info from freetoken.kernel.triton.e4m3_compat import ( + FORCE_EMU, e4m3_f32_to_u8, e4m3_native_cx, e4m3_u8_to_f32, @@ -15,6 +18,35 @@ from freetoken.kvcache.format import KVFormat +@functools.lru_cache(maxsize=None) +def _gfx1201(device_index: int) -> bool: + if torch.version.hip is None: + return False + return getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "").split(":")[0] == "gfx1201" + + +def kv_decode_config( + fmt: KVFormat, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_n: int, + num_stages: int, + num_warps: int, +) -> tuple[int, int, int, int]: + """Elements per packed load and launch geometry for the qualified KV reader.""" + if ( + fmt.name == "fp8_per_token_head" and not FORCE_EMU + and q.dtype == torch.bfloat16 and q.shape[2] == 128 + and q.shape[0] == 1 and q.shape[1] // k_cache.shape[1] in (4, 8) and _gfx1201(q.device.index) + and k_cache.stride(-1) == v_cache.stride(-1) == 1 + and k_cache.data_ptr() % 4 == 0 and v_cache.data_ptr() % 4 == 0 + and all(s % 4 == 0 for s in (*k_cache.stride()[:2], *v_cache.stride()[:2])) + ): + return 4, 64, 1, 2 + return 1, block_n, num_stages, num_warps + + @triton.jit def kv_tile_gfx12_bf16(ptrs, mask, TRANSPOSE: tl.constexpr): # Each uint32 holds four E4M3FN codes. Force e64 for both word selectors. @@ -54,6 +86,34 @@ def kv_tile(ptrs, mask, FMT: tl.constexpr): return tile +@triton.jit +def kv_tile_rows( + ptr, slots, head, stride_slot, stride_head, row_mask, + D: tl.constexpr, BLOCK_D: tl.constexpr, FMT: tl.constexpr, + PACK: tl.constexpr, TRANSPOSE: tl.constexpr, +): + """Load row-addressed K/V tiles through the selected reader, in dot-operand order.""" + tl.static_assert(PACK == 1 or PACK == 4) + if PACK == 4: + tl.static_assert(FMT == "fp8_per_token_head" and D % 4 == 0) + base = (ptr + head * stride_head).to(tl.pointer_type(tl.uint32)) + row_stride = stride_slot // 4 + else: + base = ptr + head * stride_head + row_stride = stride_slot + dims = tl.arange(0, BLOCK_D // PACK) + if TRANSPOSE: + ptrs = base + slots[None, :] * row_stride + dims[:, None] + mask = row_mask[None, :] & (dims[:, None] < D // PACK) + else: + ptrs = base + slots[:, None] * row_stride + dims[None, :] + mask = row_mask[:, None] & (dims[None, :] < D // PACK) + if PACK == 4: + return kv_tile_gfx12_bf16(ptrs, mask, TRANSPOSE) + else: + return kv_tile(ptrs, mask, FMT) + + @triton.jit def _stage(x): if target_info.is_hip(): diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 974e38a65..5d7c1cf05 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -257,25 +257,21 @@ def test_quantized_kv_attention_matches_dequantized_reference(fmt_name: str, pat @pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is None, reason="needs ROCm") @pytest.mark.parametrize("transpose", [False, True]) -def test_rdna4_kv_decode_preserves_every_finite_fp8_code(transpose: bool): +@pytest.mark.parametrize("pack", [1, 4]) +def test_rdna4_kv_decode_preserves_every_finite_fp8_code(transpose: bool, pack: int): import triton import triton.language as tl - from freetoken.kernel.triton.attention import _gfx1201 - from freetoken.kernel.triton.kv_rows import kv_tile_gfx12_bf16 + from freetoken.kernel.triton.kv_rows import _gfx1201, kv_tile_rows if not _gfx1201(torch.cuda.current_device()): pytest.skip("packed KV conversion requires gfx1201") @triton.jit - def decode_codes(x, y, TRANSPOSE: tl.constexpr): + def decode_codes(x, y, TRANSPOSE: tl.constexpr, PACK: tl.constexpr): rows = tl.arange(0, 64) - words = tl.arange(0, 32) - if TRANSPOSE: - ptrs = x + rows[None, :] * 32 + words[:, None] - else: - ptrs = x + rows[:, None] * 32 + words[None, :] - values = kv_tile_gfx12_bf16(ptrs, tl.full(ptrs.shape, True, tl.int1), TRANSPOSE) + values = kv_tile_rows(x, rows, 0, 128, 128, tl.full((64,), True, tl.int1), + 128, 128, "fp8_per_token_head", PACK, TRANSPOSE) if TRANSPOSE: values = tl.trans(values) tl.store(y + rows[:, None] * 128 + tl.arange(0, 128)[None, :], values) @@ -283,7 +279,7 @@ def decode_codes(x, y, TRANSPOSE: tl.constexpr): codes = torch.arange(256, dtype=torch.uint8).repeat(32) expected = codes.view(torch.float8_e4m3fn).bfloat16().cuda() actual = torch.empty_like(expected) - decode_codes[(1,)](codes.cuda().view(torch.uint32), actual, transpose, num_warps=2) + decode_codes[(1,)](codes.cuda(), actual, transpose, pack, num_warps=2) finite = expected.isfinite() # Integer comparison also checks signed zero and the byte order in each word. assert torch.equal(actual.view(torch.int16)[finite], expected.view(torch.int16)[finite]) From 01f3b472c6a3b96e279d6cd9da1e96e13ec23078 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Thu, 12 Nov 2026 15:44:22 +0800 Subject: [PATCH 3/3] fix(attention): isolate packed fp8 kv reads to amd --- python/freetoken/kernel/triton/attention.py | 25 +++++++++++++++------ python/freetoken/kernel/triton/kv_rows.py | 4 ++-- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index b0b0cc678..c9d737a50 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -303,6 +303,8 @@ def _decode_grouped_stage1_kernel( acc = tl.zeros((BLOCK_H, BLOCK_DV), dtype=tl.float32) q_offsets = batch_id * stride_qt + q_heads[:, None] * stride_qh + offs_d[None, :] + k_base_offsets = kv_head * stride_kh + offs_d[:, None] + v_base_offsets = kv_head * stride_vh + offs_dv[None, :] if split_end > split_start: q = tl.load(q_ptr + q_offsets, mask=mask_h[:, None] & mask_d[None, :], other=0.0) @@ -315,14 +317,21 @@ def _decode_grouped_stage1_kernel( logical_offs = effective_start + rel_offs slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) - k = kv_tile_rows(k_ptr, slots, kv_head, stride_ks, stride_kh, mask_n, - D, BLOCK_D, KV_FMT, KV_PACK, True) + # Preserve the merged reader's address expressions on non-packed targets. + if KV_PACK == 4: + k = kv_tile_rows(k_ptr, slots, kv_head, stride_ks, stride_kh, mask_n, + D, BLOCK_D, KV_FMT, KV_PACK, True) + else: + k = kv_tile(k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None], KV_FMT) scale_k = kv_row_scale(k_scale_ptr + slots * stride_kss + kv_head, mask_n, KV_FMT) scores = tl.dot(q, k) * sm_scale * scale_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - v = kv_tile_rows(v_ptr, slots, kv_head, stride_vs, stride_vh, mask_n, - DV, BLOCK_DV, KV_FMT, KV_PACK, False) + if KV_PACK == 4: + v = kv_tile_rows(v_ptr, slots, kv_head, stride_vs, stride_vh, mask_n, + DV, BLOCK_DV, KV_FMT, KV_PACK, False) + else: + v = kv_tile(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], KV_FMT) scale_v = kv_row_scale(v_scale_ptr + slots * stride_vss + kv_head, mask_n, KV_FMT) m_new = tl.maximum(tl.max(scores, axis=1), m_i) @@ -516,9 +525,11 @@ def decode_paged_attention( if kv_format.quantized: block_n, num_stages, num_warps = _quantized_decode_launch(block_d, block_n, num_stages) - kv_pack, block_n, num_stages, num_warps = kv_decode_config( - kv_format, q, k_cache, v_cache, block_n, num_stages, num_warps, - ) + kv_pack = 1 + if torch.version.hip is not None: + kv_pack, block_n, num_stages, num_warps = kv_decode_config( + kv_format, q, k_cache, v_cache, block_n, num_stages, num_warps, + ) k_scale_arg, v_scale_arg, stride_kss, stride_vss = kv_scale_args(kv_format, k_cache, v_cache, k_scale, v_scale, q) _decode_grouped_stage1_kernel[grid]( diff --git a/python/freetoken/kernel/triton/kv_rows.py b/python/freetoken/kernel/triton/kv_rows.py index a659495b9..220481cd3 100644 --- a/python/freetoken/kernel/triton/kv_rows.py +++ b/python/freetoken/kernel/triton/kv_rows.py @@ -34,7 +34,7 @@ def kv_decode_config( num_stages: int, num_warps: int, ) -> tuple[int, int, int, int]: - """Elements per packed load and launch geometry for the qualified KV reader.""" + """Elements per packed load and launch geometry for qualified gfx1201 decode.""" if ( fmt.name == "fp8_per_token_head" and not FORCE_EMU and q.dtype == torch.bfloat16 and q.shape[2] == 128 @@ -92,7 +92,7 @@ def kv_tile_rows( D: tl.constexpr, BLOCK_D: tl.constexpr, FMT: tl.constexpr, PACK: tl.constexpr, TRANSPOSE: tl.constexpr, ): - """Load row-addressed K/V tiles through the selected reader, in dot-operand order.""" + """Row-addressed K/V reader used by the qualified gfx1201 decode path.""" tl.static_assert(PACK == 1 or PACK == 4) if PACK == 4: tl.static_assert(FMT == "fp8_per_token_head" and D % 4 == 0)