perf(attention): add packed fp8 kv reads on gfx1201

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,

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():

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):