From 29c98e77f9fd5b19ec287f2e83d425b543e43885 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Thu, 12 Nov 2026 11:08:46 +0800 Subject: [PATCH 1/4] 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/4] 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/4] 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) From cf4400137375175f5cccaf2f8d0659252a5f9193 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Thu, 12 Nov 2026 19:36:06 +0800 Subject: [PATCH 4/4] perf(attention): fast fp8 kv path for rdna3 (gfx11) Decode reads E4M3 codes as fp16 bit patterns (value/256, exact) pinned to the load layout, folds 256 into the row scales, and uses 16-key single-stage tiles with shorter/more splits. Prefill dequantizes the batch rows to bf16 once and reuses the bf16 extend kernel instead of re-decoding every key tile per query block. --- python/freetoken/attention/triton.py | 5 +- python/freetoken/kernel/triton/attention.py | 47 ++++++++++-- python/freetoken/kernel/triton/kv_rows.py | 69 ++++++++++++++++++ tests/kernels/test_triton_attention.py | 80 ++++++++++++++++++++- 4 files changed, 192 insertions(+), 9 deletions(-) diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 179c57b5c..ad9b6a3c8 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -86,7 +86,10 @@ def __init__(self, config: ModelConfig): def _split_cap(self, bs: int) -> int: """The most decode KV splits any layer group takes at this batch size.""" - return max(decode_kv_split_cap(bs, max(1, self.num_q_heads), heads, self.device.index) for heads in self.kv_heads) + return max( + decode_kv_split_cap(bs, max(1, self.num_q_heads), heads, self.device.index, self.kvcache.kv_format) + for heads in self.kv_heads + ) def _ensure_decode_scratch( self, diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c9d737a50..d45c83b33 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,7 +6,10 @@ import triton import triton.language as tl -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.kernel.triton.kv_rows import ( + dequant_kv_rows, kv_decode_config, kv_f16_reader, kv_row_scale, kv_scale_args, kv_tile, kv_tile_f16, + kv_tile_rows, kv_tile_staged, +) from freetoken.kvcache.format import NONE, KVFormat @@ -20,6 +23,9 @@ _MAX_KV_SPLITS = 64 # stage 2 merges at most this many fp32 partial elements per step _STAGE2_TILE_ELEMS = 8192 +# the RDNA3 fp8 decode reads 16-key tiles with one stage: short splits and more programs hide its latency +_F16_KV_MIN_SPLIT_LEN = 32 +_F16_KV_MAX_SPLITS = 256 def _kv_offset_dtype(k_cache: torch.Tensor, v_cache: torch.Tensor): @@ -65,10 +71,14 @@ def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) -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].""" +def decode_kv_split_cap( + batch: int, num_q_heads: int, num_kv_heads: int, device_index: int, kv_format: KVFormat = NONE, +) -> int: + """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS] + (eight, within _F16_KV_MAX_SPLITS, for the RDNA3 fp8 reader).""" head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads))) - return max(_MIN_SPLITS, min(_MAX_KV_SPLITS, triton.cdiv(4 * _num_sms(device_index), batch * head_blocks))) + per_sm, max_splits = (8, _F16_KV_MAX_SPLITS) if kv_f16_reader(kv_format, device_index) else (4, _MAX_KV_SPLITS) + return max(_MIN_SPLITS, min(max_splits, triton.cdiv(per_sm * _num_sms(device_index), batch * head_blocks))) def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]: @@ -263,6 +273,7 @@ def _decode_grouped_stage1_kernel( KV_OFFSET_DTYPE: tl.constexpr, KV_FMT: tl.constexpr, KV_PACK: tl.constexpr, + F16_KV: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -310,6 +321,8 @@ def _decode_grouped_stage1_kernel( q = tl.load(q_ptr + q_offsets, mask=mask_h[:, None] & mask_d[None, :], other=0.0) if KV_FMT == "none": q = q.to(k_ptr.dtype.element_ty) + elif F16_KV: + q = q.to(tl.float16) for rel_start in tl.range(split_start, split_end, BLOCK_N): rel_offs = rel_start + tl.arange(0, BLOCK_N) @@ -321,18 +334,26 @@ def _decode_grouped_stage1_kernel( 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) + elif F16_KV: + k = kv_tile_f16(k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None]) 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) + if F16_KV: + scale_k = scale_k * 256.0 scores = tl.dot(q, k) * sm_scale * scale_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) 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) + elif F16_KV: + v = kv_tile_f16(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :]) 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) + if F16_KV: + scale_v = scale_v * 256.0 m_new = tl.maximum(tl.max(scores, axis=1), m_i) alpha = tl.exp(m_i - m_new) @@ -481,8 +502,9 @@ def decode_paged_attention( if sliding_window is not None and sliding_window <= 0: # the kernels treat a window <= 0 as no window sliding_window = None + f16_kv = kv_f16_reader(kv_format, q.device.index) if max_kv_splits is None: - max_kv_splits = decode_kv_split_cap(batch, num_q_heads, num_kv_heads, q.device.index) + max_kv_splits = decode_kv_split_cap(batch, num_q_heads, num_kv_heads, q.device.index, kv_format) if sliding_window: # no request attends more than the window, so programs past its splits would only exit max_kv_splits = min(max_kv_splits, _max_request_splits(sliding_window)) @@ -530,6 +552,10 @@ def decode_paged_attention( 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 f16_kv: + # a single stage: the pipeliner cannot predicate kv_tile_f16's asm pin + block_n, num_stages, num_warps = 16, 1, 2 + min_split_len = _F16_KV_MIN_SPLIT_LEN if f16_kv else _MIN_SPLIT_LEN 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]( @@ -562,6 +588,7 @@ def decode_paged_attention( KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), KV_FMT=kv_format.name, KV_PACK=kv_pack, + F16_KV=f16_kv, GROUP=group, NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, @@ -570,7 +597,7 @@ def decode_paged_attention( BLOCK_H=block_h, VALID_BLOCK_H=valid_block_h, MIN_BLOCK_KV=_MIN_BLOCK_KV, - MIN_SPLIT_LEN=_MIN_SPLIT_LEN, + MIN_SPLIT_LEN=min_split_len, MIN_SPLITS=_MIN_SPLITS, D=head_dim, DV=head_dim, @@ -596,7 +623,7 @@ def decode_paged_attention( o.stride(0), o.stride(1), MIN_BLOCK_KV=_MIN_BLOCK_KV, - MIN_SPLIT_LEN=_MIN_SPLIT_LEN, + MIN_SPLIT_LEN=min_split_len, MIN_SPLITS=_MIN_SPLITS, BLOCK_DV=block_dv, DV=head_dim, @@ -995,6 +1022,12 @@ def extend_paged_attention( if block_ends is not None: assert block_ends.is_cuda and block_ends.dtype == torch.int32 and block_ends.numel() == num_q_tokens + if kv_f16_reader(kv_format, q.device.index): + # every query block re-decodes each key tile, which RDNA3 does far slower than it streams bf16: + # dequantize the batch's rows once and run the bf16 kernel on them + k_cache, v_cache = dequant_kv_rows(k_cache, v_cache, k_scale, v_scale, kv_indices, kv_indptr) + kv_indices = torch.arange(kv_indices.numel(), dtype=torch.int32, device=q.device) + kv_format, k_scale, v_scale = NONE, None, None o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q block_ends_arg = block_ends if block_ends is not None else qo_indptr diff --git a/python/freetoken/kernel/triton/kv_rows.py b/python/freetoken/kernel/triton/kv_rows.py index 220481cd3..4b99d7adb 100644 --- a/python/freetoken/kernel/triton/kv_rows.py +++ b/python/freetoken/kernel/triton/kv_rows.py @@ -25,6 +25,19 @@ def _gfx1201(device_index: int) -> bool: return getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "").split(":")[0] == "gfx1201" +@functools.lru_cache(maxsize=None) +def _gfx11(device_index: int) -> bool: + if torch.version.hip is None: + return False + return getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "").startswith("gfx11") + + +def kv_f16_reader(fmt: KVFormat, device_index: int) -> bool: + """Read E4M3 K/V as fp16 bit patterns on RDNA3, which has no fp8 converts and lowers every + fp32 -> bf16 cast to integer RNE code; the generic decode then dominates attention time.""" + return fmt.name == "fp8_per_token_head" and not FORCE_EMU and _gfx11(device_index) + + def kv_decode_config( fmt: KVFormat, q: torch.Tensor, @@ -86,6 +99,17 @@ def kv_tile(ptrs, mask, FMT: tl.constexpr): return tile +@triton.jit +def kv_tile_f16(ptrs, mask): + """E4M3FN codes as fp16 holding ``value / 256``: exact for every finite code (subnormals + land on fp16 subnormals), integer ops only. Callers fold the 256 into the row scales.""" + codes = tl.load(ptrs, mask=mask, other=0).to(tl.uint16) + bits = ((codes & 0x80) << 8) | ((codes & 0x7F) << 7) + # pin the decode to the load layout so it runs once per element, not once per dot-operand replica + bits = tl.inline_asm_elementwise("v_mov_b32 $0, $1", "=v,v", [bits], dtype=tl.uint16, is_pure=False, pack=2) + return bits.to(tl.float16, bitcast=True) + + @triton.jit def kv_tile_rows( ptr, slots, head, stride_slot, stride_head, row_mask, @@ -268,6 +292,51 @@ def store_kv_rows( ) +@triton.jit +def _kv_dequant_rows_kernel( + k_src, v_src, k_scale, v_scale, k_dst, v_dst, indices, total, + stride_ks, stride_vs, stride_kss, stride_vss, + D: tl.constexpr, ROW: tl.constexpr, BLOCK: tl.constexpr, +): + pos = tl.program_id(0).to(tl.int64) + if pos < tl.load(total): + slot = tl.load(indices + pos).to(tl.int64) + offs = tl.arange(0, BLOCK) + mask = offs < ROW + head = offs // D + k = kv_tile_f16(k_src + slot * stride_ks + offs, mask).to(tl.float32) + ks = tl.load(k_scale + slot * stride_kss + head, mask=mask, other=0.0) * 256.0 + tl.store(k_dst + pos * ROW + offs, (k * ks).to(tl.bfloat16), mask=mask) + v = kv_tile_f16(v_src + slot * stride_vs + offs, mask).to(tl.float32) + vs = tl.load(v_scale + slot * stride_vss + head, mask=mask, other=0.0) * 256.0 + tl.store(v_dst + pos * ROW + offs, (v * vs).to(tl.bfloat16), mask=mask) + + +def dequant_kv_rows( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + indices: torch.Tensor, + indptr: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """BF16 K/V whose row ``i`` dequantizes E4M3 slot ``indices[i]``, for ``i < indptr[-1]``.""" + n = indices.numel() + _, heads, head_dim = k_cache.shape + assert k_cache.stride(1) == v_cache.stride(1) == head_dim and k_cache.stride(2) == v_cache.stride(2) == 1 + assert k_scale.stride(1) == v_scale.stride(1) == 1 + k = torch.empty(n, heads, head_dim, dtype=torch.bfloat16, device=k_cache.device) + v = torch.empty_like(k) + if n: + _kv_dequant_rows_kernel[(n,)]( + k_cache, v_cache, k_scale, v_scale, k, v, indices, indptr[-1:], + k_cache.stride(0), v_cache.stride(0), k_scale.stride(0), v_scale.stride(0), + D=head_dim, ROW=heads * head_dim, BLOCK=triton.next_power_of_2(heads * head_dim), + num_warps=2, + ) + return k, v + + @triton.jit def _latent_store_kernel( diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 5d7c1cf05..68020952a 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -285,6 +285,84 @@ def decode_codes(x, y, TRANSPOSE: tl.constexpr, PACK: tl.constexpr): assert torch.equal(actual.view(torch.int16)[finite], expected.view(torch.int16)[finite]) +def _require_gfx11(): + from freetoken.kernel.triton.kv_rows import _gfx11 + + if not _gfx11(torch.cuda.current_device()): + pytest.skip("the fp16-bit-pattern fp8 KV reader is RDNA3 only") + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is None, reason="needs ROCm") +def test_rdna3_kv_f16_reader_preserves_every_finite_fp8_code(): + import triton + import triton.language as tl + + from freetoken.kernel.triton.kv_rows import kv_tile_f16 + + _require_gfx11() + + @triton.jit + def decode_codes(x, y): + offs = tl.arange(0, 256) + tl.store(y + offs, kv_tile_f16(x + offs, tl.full((256,), True, tl.int1))) + + codes = torch.arange(256, dtype=torch.uint8) + expected = (codes.view(torch.float8_e4m3fn).float() / 256).half().cuda() + actual = torch.empty_like(expected) + decode_codes[(1,)](codes.cuda(), actual, num_warps=1) + finite = expected.isfinite() + # integer comparison also checks signed zero and the fp16 subnormals of the e4m3 subnormals + assert torch.equal(actual.view(torch.int16)[finite], expected.view(torch.int16)[finite]) + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is None, reason="needs ROCm") +def test_rdna3_fp8_kv_paths_match_bf16_kernels_on_dequantized_rows(): + from freetoken.kernel.triton.attention import decode_kv_split_cap, decode_paged_attention, extend_paged_attention + from freetoken.kernel.triton.kv_rows import dequant_kv_rows, store_kv_rows + + _require_gfx11() + fmt = KV_FORMATS["fp8_per_token_head"] + torch.manual_seed(0) + slots, heads, kv_heads, head_dim, prefix, chunk = 4096, 8, 2, 256, 1500, 300 + k = torch.randn(slots, kv_heads * head_dim, device="cuda").bfloat16() + v = torch.randn(slots, kv_heads * head_dim, device="cuda").bfloat16() + k_cache = torch.empty(slots, kv_heads, head_dim, dtype=torch.uint8, device="cuda") + v_cache = torch.empty_like(k_cache) + k_scale = torch.empty(slots, kv_heads, device="cuda") + v_scale = torch.empty_like(k_scale) + store_kv_rows(fmt, k, v, torch.arange(slots, device="cuda"), k_cache, v_cache, k_scale, v_scale) + rows = torch.randperm(slots, device="cuda")[: prefix + chunk].int() + indptr = torch.tensor([0, prefix + chunk], dtype=torch.int32, device="cuda") + + k_ref = (k_cache.view(torch.float8_e4m3fn).float() * k_scale[..., None]).bfloat16() + v_ref = (v_cache.view(torch.float8_e4m3fn).float() * v_scale[..., None]).bfloat16() + k_rows, v_rows = dequant_kv_rows(k_cache, v_cache, k_scale, v_scale, rows, indptr) + assert torch.equal(k_rows, k_ref[rows.long()]) and torch.equal(v_rows, v_ref[rows.long()]) + + # prefill runs the bf16 kernel on the dequantized rows, so it matches that kernel exactly + q = torch.randn(chunk, heads, head_dim, device="cuda").bfloat16() + k_ext = k_ref[rows[prefix:].long()].contiguous() + v_ext = v_ref[rows[prefix:].long()].contiguous() + extend_args = (torch.tensor([0, chunk], dtype=torch.int32, device="cuda"), indptr, rows, + torch.tensor([prefix], dtype=torch.int32, device="cuda"), chunk, head_dim ** -0.5) + actual = extend_paged_attention(q, k_cache, v_cache, *extend_args, k_extend=k_ext, v_extend=v_ext, + kv_format=fmt, k_scale=k_scale, v_scale=v_scale) + expected = extend_paged_attention(q, k_ref, v_ref, *extend_args, k_extend=k_ext, v_extend=v_ext) + assert torch.equal(actual, expected) + + # decode reorders only fp16 vs bf16 operand rounding + for keys in (1, 33, prefix + chunk): + q = torch.randn(1, heads, head_dim, device="cuda").bfloat16() + splits = decode_kv_split_cap(1, heads, kv_heads, q.device.index, fmt) + scratch = (torch.empty(1, heads, splits, head_dim, device="cuda"), torch.empty(1, heads, splits, device="cuda")) + decode_args = (torch.tensor([0, keys], dtype=torch.int32, device="cuda"), rows[:keys].contiguous(), + torch.tensor([keys - 1], dtype=torch.int32, device="cuda"), *scratch, head_dim ** -0.5) + actual = decode_paged_attention(q, k_cache, v_cache, *decode_args, kv_format=fmt, k_scale=k_scale, + v_scale=v_scale) + expected = decode_paged_attention(q, k_ref, v_ref, *decode_args) + torch.testing.assert_close(actual.float(), expected.float(), atol=1e-2, rtol=1e-2) + + @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): @@ -1125,7 +1203,7 @@ def v_cache(self, layer_id): def test_triton_backend_replay_metadata_uses_capture_buffers(monkeypatch): from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata - class FakeKVCache: + class FakeKVCache(_PlainKVStorage): def __init__(self, device: torch.device): self.device = device