diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index f6a84692e..9a8cc9b0c 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,9 +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.e4m3_compat import FORCE_EMU +from freetoken.kernel.triton.kv_rows import kv_row_scale, kv_scale_args, kv_tile, kv_tile_rows, kv_tile_staged from freetoken.kvcache.format import NONE, KVFormat -from freetoken.utils.arch import is_rocm +from freetoken.utils.arch import get_rocm_gfx_arch, is_rocm _MIN_BLOCK_KV = 32 @@ -52,6 +53,19 @@ def _quantized_decode_launch(block_d: int, block_n: int, num_stages: int) -> tup return block_n, num_stages, 4 +def _use_gfx12_fp8_reader(kv_format: KVFormat, q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool: + return ( + kv_format.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 get_rocm_gfx_arch() == "gfx1201" + and k_cache.stride(-1) == v_cache.stride(-1) == 1 + # With only four-byte alignment, the packed load layout spills SGPRs. + and k_cache.data_ptr() % 16 == 0 and v_cache.data_ptr() % 16 == 0 + and all(s % 16 == 0 for s in (*k_cache.stride()[:2], *v_cache.stride()[:2])) + ) + + def _multi_row_decode_launch(m: int, block_d: int, quantized: bool, q_bytes: int, kv_bytes: int, smem: int) -> tuple[int, int, int]: """(block_n, num_stages, num_warps) of stage 1 for ``m`` > 1 query rows x heads per program. @@ -296,6 +310,7 @@ def _decode_grouped_stage1_kernel( SLIDING_WINDOW: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, KV_FMT: tl.constexpr, + KV_PACK: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -374,7 +389,12 @@ 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) + # Keep the address math inline: moving it into a helper changes CUDA codegen. + if KV_PACK == 4: + k = kv_tile_rows(k_ptr, slots, kv_head, stride_ks, stride_kh, mask_n, + D, BLOCK_D, 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, :] mask = mask_h[:, None] & mask_n[None, :] @@ -385,7 +405,11 @@ def _decode_grouped_stage1_kernel( mask = mask & ((logical_offs[None, :] + SLIDING_WINDOW) > row_pos[:, None]) scores = tl.where(mask, 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, 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) row_max = tl.max(scores, axis=1) @@ -616,6 +640,10 @@ def decode_paged_attention( rows * block_h, block_d, kv_format.quantized, elem, k_cache.element_size(), smem ) + kv_pack = 1 + if is_rocm() and rows == 1 and _use_gfx12_fp8_reader(kv_format, q, k_cache, v_cache): + kv_pack, block_n, num_stages, num_warps = 4, 64, 1, 4 + 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, @@ -647,6 +675,7 @@ 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, GROUP=group, NUM_Q_HEADS=num_q_heads, ROWS=rows, diff --git a/python/freetoken/kernel/triton/kv_rows.py b/python/freetoken/kernel/triton/kv_rows.py index afa19bfbe..eb9dd3991 100644 --- a/python/freetoken/kernel/triton/kv_rows.py +++ b/python/freetoken/kernel/triton/kv_rows.py @@ -15,6 +15,30 @@ from freetoken.kvcache.format import KVFormat +@triton.jit +def _kv_tile_gfx12_bf16(ptrs, mask, TRANSPOSE: tl.constexpr): + # Force VOP3 (e64): LLVM's word-0 packed intrinsic selects an invalid VOP1 on gfx12. + # Recheck plain FP8 casts when upgrading Triton (triton-lang/triton#11809). + 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 +55,25 @@ 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, TRANSPOSE: tl.constexpr, +): + """Read four E4M3FN codes per uint32 load and return a BF16 K or V tile.""" + tl.static_assert(D % 4 == 0) + base = (ptr + head * stride_head).to(tl.pointer_type(tl.uint32)) + row_stride = stride_slot // 4 + dims = tl.arange(0, BLOCK_D // 4) + if TRANSPOSE: + ptrs = base + slots[None, :] * row_stride + dims[:, None] + mask = row_mask[None, :] & (dims[:, None] < D // 4) + else: + ptrs = base + slots[:, None] * row_stride + dims[None, :] + mask = row_mask[:, None] & (dims[None, :] < D // 4) + return _kv_tile_gfx12_bf16(ptrs, mask, TRANSPOSE) + + @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 841bf4533..679efce9c 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -213,12 +213,15 @@ 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", "decode_rows", "extend", "split_extend"]) +@pytest.mark.parametrize(("path", "num_q_heads"), [ + ("paged", 4), ("decode", 4), ("decode", 8), ("decode", 16), + ("decode_rows", 4), ("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 = {"paged": 1, "decode": 1, "decode_rows": 3}.get(path, 2) row = num_kv_heads * head_dim @@ -260,6 +263,36 @@ 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.kv_rows import kv_tile_rows + from freetoken.utils.arch import get_rocm_gfx_arch + + if get_rocm_gfx_arch() != "gfx1201": + pytest.skip("packed KV conversion requires gfx1201") + + @triton.jit + def decode_codes(x, y, TRANSPOSE: tl.constexpr): + rows = tl.arange(0, 64) + values = kv_tile_rows(x, rows, 0, 128, 128, tl.full((64,), True, tl.int1), + 128, 128, 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, num_warps=4) + 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):