perf(attention): fast fp8 kv decode and prefill on rdna3 (gfx11)

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,

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,20 +321,39 @@ 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)
             mask_n = rel_offs < split_end
             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

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(