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 b1a4455a4..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_row_scale, kv_scale_args, kv_tile, 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]: @@ -262,6 +272,8 @@ def _decode_grouped_stage1_kernel( SLIDING_WINDOW: tl.constexpr, 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) @@ -309,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) @@ -316,13 +330,30 @@ 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) + # 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) + 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")) - v = kv_tile(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], KV_FMT) + 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) @@ -471,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)) @@ -515,6 +547,16 @@ 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 = 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, + ) + 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]( q, @@ -545,6 +587,8 @@ def decode_paged_attention( attn_lse.stride(2), 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, @@ -553,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, @@ -579,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, @@ -978,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 afa19bfbe..4b99d7adb 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,71 @@ 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" + + +@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, + 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 qualified gfx1201 decode.""" + 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. + 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.""" @@ -31,6 +99,45 @@ 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, + D: tl.constexpr, BLOCK_D: tl.constexpr, FMT: tl.constexpr, + PACK: tl.constexpr, TRANSPOSE: tl.constexpr, +): + """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) + 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(): @@ -185,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 ad1120181..68020952a 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,114 @@ 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]) +@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.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, PACK: tl.constexpr): + rows = tl.arange(0, 64) + 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) + + 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(), 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]) + + +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): @@ -1093,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