[zihaomu]

Draft, stacked on #897. GitHub cannot base a PR on a fork branch, so this PR targets main (27740fe) and its diff currently also shows #897's three commits (29c98e7, 252528e, 01f3b47). The change proposed here is only the top commit cf44001. I will rebase onto main once #897 lands.

Follow-up to #897 (gfx1201 packed FP8 KV reads); related to #894 and #893.

#897 accelerates FP8 KV on RDNA4 only. On RDNA3 (gfx11), fp8_per_token_head KV is functionally supported but is much slower than BF16 KV: on a Radeon PRO W7900D (gfx1100) serving Qwen/Qwen3.5-0.8B, FP8 KV decode was 15% slower at 512 tokens and 55% slower at 64K, and 64K TTFT was 4.8x worse. This PR lists what RDNA3 needs and implements it: FP8 KV reaches TTFT parity with BF16 KV and decodes faster at long context (+20% at 64K), while keeping FP8's KV capacity (70K tokens: 0.41 GiB vs 0.80 GiB).

Scope: gfx11 + fp8_per_token_head only. RDNA4 (including #897's packed reader), CUDA, BF16 KV and INT8 KV keep their current code; evidence below.

Why the current FP8 path is slow on RDNA3

• gfx11 has no FP8 conversion instructions and no hardware FP32 to BF16 conversion. e4m3_native_cx() is false on HIP, so kv_tile runs e4m3_u8_to_f32 and then .to(tl.bfloat16), which LLVM lowers to an integer RNE sequence (v_add3_u32, v_cmp_o_f32, v_cndmask). The generated D256 FP8 decode stage-1 contains 3846 v_and, 3716 v_lshr, 1313 v_add3_u32 and 1312 v_cmp_o_f32 instructions, against 80 WMMA instructions.
• Triton places that elementwise decode in the dot-operand layout, so every WMMA operand replica decodes the same bytes again (the reason kv_tile_staged pins the decode on NVIDIA; _stage is a no-op on HIP).
• Prefill multiplies this: every query block re-decodes every prefix key tile. With a 2048-token chunk and 64-row query tiles, each FP8 element is decoded 32 query blocks x 4 GQA heads = 128 times, before operand replication.

Changes (top commit only)

1. Gate kv_f16_reader(fmt, device_index) in kv_rows.py: fp8_per_token_head, gcnArchName starts with gfx11, and FREETOKEN_FORCE_E4M3_EMU not set.
2. Decode (_decode_grouped_stage1_kernel, new constexpr F16_KV):
   • kv_tile_f16 turns each E4M3FN code into the fp16 bit pattern of value/256: ((c & 0x80) << 8) | ((c & 0x7F) << 7). This is exact for every finite code (E4M3 subnormals become fp16 subnormals) and uses integer ops only. The 256 is folded into the row scales; Q is cast to fp16, so QK and PV run as v_wmma_f32_16x16x16_f16 with FP32 accumulation.
   • A v_mov_b32 inline-asm pin keeps the decode in the load layout, so each byte is decoded once rather than per operand replica. The pipeliner cannot predicate this op, so the path uses one stage.
   • gfx11 launch: 16-key tiles, 1 stage, 2 warps; minimum split length 32; split cap at 8 programs per SM (max 256). W7900D reports 48 SMs, so batch-1 GQA4 gets 192 splits instead of 64. decode_kv_split_cap takes an optional kv_format, and the Triton backend sizes its decode scratch with the KV cache's format so the cap and the scratch agree, including Graph capture.
3. Prefill (extend_paged_attention): on gfx11 FP8, dequant_kv_rows decodes the batch's KV rows once into a transient BF16 buffer (row i = slot kv_indices[i]), then the unchanged BF16 extend kernel runs on identity indices. Each element is decoded once, not once per query block.
   • Transient memory: kv_indices.numel() x kv_heads x head_dim x 4 bytes (K+V BF16), about 134 MB per layer call for Qwen3.5-0.8B at 65,792 tokens. It is released after the call and does not change the KV pool size.
   • I also measured an in-kernel fp16 reader for prefill (same kv_tile_f16 inside _extend_attend_tile): 48.3 ms vs 35.1 ms per 2048-token chunk at a 62K prefix, so it was not kept.
4. Tests: test_rdna3_kv_f16_reader_preserves_every_finite_fp8_code compares all 256 codes bit-for-bit against torch.float8_e4m3fn / 256 in fp16, including signed zero and subnormals. test_rdna3_fp8_kv_paths_match_bf16_kernels_on_dequantized_rows checks:
   • dequant_kv_rows bit-for-bit against a PyTorch dequantization;
   • prefill bit-for-bit against the BF16 kernel on those rows;
   • decode (1/33/1800 keys, permuted slots) within 1e-2 of the BF16 kernel on the same rows.
     One existing backend test's FakeKVCache now subclasses the file's _PlainKVStorage because the backend reads kv_format when sizing decode scratch; production pools already define it (kv_format: KVFormat = NONE).

Numerics

• Prefill is bit-identical to running the BF16 kernel on dequantized rows. Relative to the previous FP8 path, the only extra rounding is code x scale to BF16 (relative error at most 2^-9, far below E4M3's 2^-4 step).
• Decode uses fp16 instead of BF16 matrix operands: Q is rounded to fp16 (finer mantissa, but values above 65504 would overflow; I have not seen such Q values, and Qwen3.5 applies q-norm) and p x scale_v x 256 is rounded to fp16. Greedy outputs therefore diverge from the previous FP8 path at some contexts (hash table below). Task-quality evaluation has not been run.

Isolation: RDNA4, CUDA and other formats

• kv_f16_reader is false off gfx11 (torch.version.hip is None short-circuits on CUDA), so host-side launch parameters, split caps and the prefill route are unchanged there, and F16_KV=False removes the new kernel branches at compile time.
• RDNA4 codegen evidence. I cross-compiled decode stage-1 for gfx1201 and gfx1200 (Triton target override, no launch) on 01f3b47 (perf(attention): add packed fp8 kv reads on gfx1201 #897) and on cf44001, for D128 GQA4. The instruction streams are identical in all eight comparisons:

Case	Instructions	v_cvt_pk_f32_fp8	SHA-256 prefix (both commits, both targets)
FP8 packed (KV_PACK=4, #897 path)	2861	192	861efefedf9f2fdf
FP8 generic reader	9871	0	b5de1d190d4f935f
INT8	8567	0	1be8768c41784f68
BF16 (none)	2395	0	9eb09bc2e2ab7ebf

• Not run on real RDNA4 or NVIDIA hardware for this change.

Test environment

• GPU: AMD Radeon PRO W7900D, gfx1100, 48 GiB, PCI 0000:A3:00.0 (one GPU visible via ROCR_VISIBLE_DEVICES; torch.cuda.get_device_properties(0) reports gfx1100, multi_processor_count=48).
• CPU: 2x AMD EPYC 9334; Linux 6.8.0-146-generic; amdgpu 6.16.13.
• Container image sha256:31cc953da2ee8345774f92c8f288dded6a19cbcd59c1ce96051970f48d627d18: ROCm 10.1.0, HIP 7.16.26385, Python 3.13.13, Torch 2.14.0+rocm10.1.0, Triton 3.8.0, Transformers 5.18.0.
• Checkpoint: Qwen/Qwen3.5-0.8B (BF16; model.safetensors SHA-256 04b1c301231dd422b8860db31311ab2721511346a32cb1e079c4c4e5f1fe4696). 24 layers, 6 full-attention layers; head_dim 256; 8 query / 2 KV heads (GQA4).

Tests (all on cf44001)

• pytest tests/kernels/test_triton_attention.py: 113 passed, 4 skipped. The four skips are test_rdna4_kv_decode_preserves_every_finite_fp8_code[*], which require gfx1201.
• pytest tests/engine/test_attention_backend_arch.py tests/kvcache/test_kv_cache_rebuild.py tests/kvcache/test_pool_sizing_surface.py tests/attention/test_capture_rows_per_req.py: 62 passed, 4 skipped (skips: sgl_kernel / flashinfer not installed).
• FREETOKEN_FORCE_E4M3_EMU=1 pytest tests/kernels/test_triton_attention.py -k test_quantized_kv_attention_matches_dequantized_reference with a separate Triton cache: 48 passed (the gfx11 path is disabled under forced emulation).

End-to-end A/B

Same GPU, same image, one service at a time, 1 warmup + 3 measured requests per point (medians), 128 output tokens, temperature=0, ignore_eos=true, --cache-type naive. Every request was checked for exact prompt tokens, 128 completion tokens and cached_tokens == 0. Prompts are exact-length synthetic documents (<|im_start|>user ... <|im_start|>assistant). TTFT is request start to first text event; decode tok/s = 127 / (finish event - first text event).

Baselines are 01f3b47 (#897 head = main + #897). #897 only touches decode stage-1 and its wrapper; for these gfx1100 D256 shapes the stage-1 instruction streams of main and 01f3b47 are identical for FP8, INT8 and BF16 (cross-compiled, SHA-256 prefixes 397ac1fad1cd6b35 / 1d0826d1a4170012 / fe517d008a142427), and the extend path is untouched, so these columns also represent main.

# BF16 KV baseline (01f3b47): --kv-cache-dtype auto; FP8 baseline (01f3b47) and this PR (cf44001): fp8_per_token_head
ROCR_VISIBLE_DEVICES=<gpu> python -m freetoken.cli serve --model Qwen/Qwen3.5-0.8B --dtype bfloat16 \
  --attention-backend triton --kv-cache-dtype {auto|fp8_per_token_head} --num-tokens 70000 \
  --max-seq-len-override 65792 --max-prefill-length 2048 --max-running-requests 1 --cuda-graph-max-bs 1 \
  --memory-ratio 0.7 --cache-type naive --text-model-only --enable-cache-report --enable-timings-report \
  --sampling-defaults none --reasoning-parser off --port <port>
python bench_service.py --origin http://127.0.0.1:<port> --prompts prompts --label <label> --repeats 3 --max-tokens 128 --output <label>.json

Context	TTFT BF16 / FP8 01f3b47 / FP8 cf44001	Decode tok/s BF16 / FP8 01f3b47 / FP8 cf44001	E2E BF16 / FP8 01f3b47 / FP8 cf44001	cf44001 vs BF16 (TTFT / decode / E2E)
512	0.036 / 0.032 / 0.032 s	198.2 / 168.5 / 198.4	0.678 / 0.786 / 0.671 s	-12.2% / +0.1% / -1.0%
8,192	0.402 / 0.669 / 0.390 s	186.4 / 147.4 / 191.2	1.088 / 1.529 / 1.054 s	-3.0% / +2.6% / -3.1%
32,768	2.203 / 7.734 / 2.191 s	160.3 / 92.5 / 176.7	2.993 / 9.107 / 2.909 s	-0.6% / +10.3% / -2.8%
65,536	6.129 / 29.553 / 6.150 s	135.6 / 61.1 / 162.9	7.065 / 31.633 / 6.931 s	+0.4% / +20.1% / -1.9%

The 512-token TTFT difference is about 4 ms and within that point's noise. KV allocation for 70,000 tokens: BF16 0.80 GiB, FP8 0.41 GiB (service log); service-reported VRAM BF16 2.719 GiB, FP8 2.318 GiB (01f3b47) / 2.316 GiB (cf44001).

Greedy output hashes (identical across the three repeats in every run): cf44001 matches BF16 at 512, matches the previous FP8 path at 8,192, and matches neither at 32,768 / 65,536.

Kernel microbenchmark

Qwen3.5-0.8B full-attention shapes (8 Q heads, 2 KV heads, D256), contiguous slots, per-layer time in us, median of 7 HIP-Graph replay batches. Decode uses the backend's split cap; prefill is one 2048-token chunk after the given prefix with k_extend/v_extend.

Kernel	Keys / prefix	BF16 01f3b47	FP8 01f3b47	FP8 cf44001	cf44001 max abs err vs BF16-on-dequantized
decode	512	22.6	87.3	21.4	4.2e-3
decode	2,048	42.2	157.4	26.5	2.4e-3
decode	8,192	57.7	288.2	47.7	9.8e-4
decode	32,768	184.4	1041.7	108.6	9.8e-4
decode	65,536	400.3	2079.9	173.0	6.7e-4
prefill	0	704.8	782.7	712.7	0
prefill	8,192	5003.6	36405.2	5020.4	0
prefill	32,768	18158.3	143474.7	18274.9	0
prefill	63,488	34847.1	301561.5	35170.9	0

Open questions for reviewers

• Gate width. The gate covers all gfx11* (gfx1100/1101/1102, RDNA3.5 gfx1150/1151), which share the missing conversions, but only gfx1100 was measured. Should it be narrowed to measured targets?
• BF16 decode on gfx11 benefits from the same launch direction: with microbenchmark overrides (16-key tiles, 1 stage, 4 warps, min split 32, up to 512 splits) BF16 decode measured 18 / 45 / 115 / 225 us at 512 / 8K / 32K / 64K. FP8 with this PR is 172 us at 64K, about par at 32K and slightly slower below. Retuning BF16 is out of scope here and would be a separate PR.
• All D256 decode/extend kernels on gfx1100, BF16 included, use 256 VGPRs and spill. That is pre-existing and not addressed here.

Scripts

bench_service.py (end-to-end SSE client)

from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import time
import urllib.request
from pathlib import Path


def get_json(url: str, timeout: float = 30) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return json.load(response)


def stream_completion(origin: str, model: str, prompt: str, max_tokens: int) -> dict:
    body = {
        "model": model,
        "prompt": prompt,
        "max_tokens": max_tokens,
        "temperature": 0,
        "ignore_eos": True,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    request = urllib.request.Request(
        f"{origin}/v1/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    t0 = time.perf_counter()
    first_text = None
    finish = None
    usage = None
    pieces: list[str] = []
    event_count = 0
    with urllib.request.urlopen(request, timeout=1800) as response:
        for raw in response:
            line = raw.strip()
            if not line.startswith(b"data:"):
                continue
            payload = line[5:].strip()
            if payload == b"[DONE]":
                break
            now = time.perf_counter()
            event = json.loads(payload)
            event_count += 1
            if event.get("usage"):
                usage = event["usage"]
            for choice in event.get("choices", []):
                text = choice.get("text") or ""
                if text:
                    if first_text is None:
                        first_text = now
                    pieces.append(text)
                if choice.get("finish_reason") is not None:
                    finish = now
    if usage is None or first_text is None or finish is None:
        raise RuntimeError(
            f"incomplete stream: usage={usage is not None}, "
            f"first_text={first_text}, finish={finish}, events={event_count}"
        )
    completion = int(usage["completion_tokens"])
    decode_s = finish - first_text
    return {
        "prompt_tokens": int(usage["prompt_tokens"]),
        "completion_tokens": completion,
        "cached_tokens": int(
            (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0) or 0
        ),
        "ttft_s": first_text - t0,
        "decode_s": decode_s,
        "decode_tok_s": (completion - 1) / decode_s,
        "e2e_s": finish - t0,
        "events": event_count,
        "text_sha256": hashlib.sha256("".join(pieces).encode()).hexdigest(),
        "text_chars": sum(map(len, pieces)),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--origin", required=True)
    parser.add_argument("--prompts", required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    model = get_json(f"{args.origin}/v1/models")["data"][0]["id"]
    manifest = json.loads((Path(args.prompts) / "manifest.json").read_text())
    rows = []
    for item in manifest:
        target = int(item["tokens"])
        prompt = (Path(args.prompts) / f"prompt-{target}.txt").read_text()
        warm = stream_completion(args.origin, model, prompt, args.max_tokens)
        if warm["prompt_tokens"] != target or warm["completion_tokens"] != args.max_tokens:
            raise RuntimeError(f"warmup usage mismatch at {target}: {warm}")
        print(f"warmup label={args.label} tokens={target} ttft={warm['ttft_s']:.3f}s decode={warm['decode_tok_s']:.2f}", flush=True)
        for repeat in range(1, args.repeats + 1):
            row = stream_completion(args.origin, model, prompt, args.max_tokens)
            if row["prompt_tokens"] != target:
                raise RuntimeError(f"prompt token mismatch: expected {target}, got {row}")
            if row["completion_tokens"] != args.max_tokens:
                raise RuntimeError(f"completion token mismatch: {row}")
            if row["cached_tokens"] != 0:
                raise RuntimeError(f"unexpected prefix cache hit: {row}")
            row.update(label=args.label, target_tokens=target, repeat=repeat)
            rows.append(row)
            print(json.dumps(row, sort_keys=True), flush=True)

    summaries = []
    for target in sorted({r["target_tokens"] for r in rows}):
        subset = [r for r in rows if r["target_tokens"] == target]
        summary = {
            "label": args.label,
            "target_tokens": target,
            "repeats": len(subset),
            "ttft_median_s": statistics.median(r["ttft_s"] for r in subset),
            "decode_tok_s_median": statistics.median(r["decode_tok_s"] for r in subset),
            "e2e_median_s": statistics.median(r["e2e_s"] for r in subset),
            "text_hashes": sorted({r["text_sha256"] for r in subset}),
        }
        summaries.append(summary)
        print("SUMMARY " + json.dumps(summary, sort_keys=True), flush=True)

    stats = get_json(f"{args.origin}/v1/stats")
    Path(args.output).write_text(
        json.dumps({"label": args.label, "model": model, "rows": rows, "summaries": summaries, "stats": stats}, indent=2)
    )


if __name__ == "__main__":
    main()

bench_kernels.py (kernel microbenchmark)

"""Kernel-level BF16 vs FP8 KV benchmark for Qwen3.5-0.8B full-attention shapes on one GPU."""
from __future__ import annotations

import argparse
import json
import statistics

import torch

from freetoken.kernel.triton.attention import decode_paged_attention, extend_paged_attention
from freetoken.kernel.triton.kv_rows import store_kv_rows
from freetoken.kvcache.format import FP8_PER_TOKEN_HEAD, NONE

QH, KVH, D = 8, 2, 256
CHUNK = 2048
SPLITS = 64


def make_cache(fmt, n, gen):
    k = torch.randn(n, KVH * D, generator=gen, device="cuda").mul_(2.0).to(torch.bfloat16)
    v = torch.randn(n, KVH * D, generator=gen, device="cuda").to(torch.bfloat16)
    if not fmt.quantized:
        return k.view(n, KVH, D), v.view(n, KVH, D), None, None
    kc = torch.empty(n, KVH, D, dtype=torch.uint8, device="cuda")
    vc = torch.empty_like(kc)
    ks = torch.empty(n, KVH, dtype=torch.float32, device="cuda")
    vs = torch.empty_like(ks)
    store_kv_rows(fmt, k, v, torch.arange(n, device="cuda"), kc, vc, ks, vs)
    return kc, vc, ks, vs


def dequant(codes, scale):
    return (codes.view(torch.float8_e4m3fn).float() * scale[..., None]).to(torch.bfloat16)


def time_graph(fn, iters):
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(3):
            fn()
    torch.cuda.current_stream().wait_stream(s)
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        fn()
    samples = []
    for _ in range(7):
        a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        a.record()
        for _ in range(iters):
            g.replay()
        b.record()
        torch.cuda.synchronize()
        samples.append(a.elapsed_time(b) / iters * 1000.0)
    return statistics.median(samples)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--fmt", choices=["bf16", "fp8"], required=True)
    ap.add_argument("--decode-ctx", default="512,8192,32768,65536")
    ap.add_argument("--prefix", default="0,8192,32768,63488")
    ap.add_argument("--output", required=True)
    ap.add_argument("--splits", type=int, default=None, help="max_kv_splits for decode (default: kernel cap)")
    ap.add_argument("--min-split-len", type=int, default=None)
    ap.add_argument("--dec-cfg", default=None, help="override kv_decode_config: block_n,stages,warps")
    args = ap.parse_args()
    fmt = FP8_PER_TOKEN_HEAD if args.fmt == "fp8" else NONE
    if args.dec_cfg:
        import freetoken.kernel.triton.attention as attn
        bn, st, wp = map(int, args.dec_cfg.split(","))
        attn.kv_decode_config = lambda *a: (1, bn, st, wp)
    try:
        from freetoken.kernel.triton.attention import decode_kv_split_cap
        cap = decode_kv_split_cap(1, QH, KVH, torch.cuda.current_device(), fmt)
    except TypeError:
        cap = SPLITS
    splits = args.splits or cap
    print(json.dumps({"scratch_splits": splits}), flush=True)
    if args.min_split_len:
        import freetoken.kernel.triton.attention as attn
        attn._MIN_SPLIT_LEN = args.min_split_len
    gen = torch.Generator(device="cuda").manual_seed(0)
    n = 65536 + CHUNK
    kc, vc, ks, vs = make_cache(fmt, n, gen)
    kv_args = dict(kv_format=fmt, k_scale=ks, v_scale=vs)
    ref_args = dict(kv_format=NONE, k_scale=None, v_scale=None)
    if fmt.quantized:
        kr, vr = dequant(kc, ks), dequant(vc, vs)
    rows = []

    for ctx in map(int, filter(None, args.decode_ctx.split(","))):
        q = torch.randn(1, QH, D, generator=gen, device="cuda").to(torch.bfloat16)
        indptr = torch.tensor([0, ctx], dtype=torch.int32, device="cuda")
        indices = torch.arange(ctx, dtype=torch.int32, device="cuda")
        qpos = torch.tensor([ctx - 1], dtype=torch.int32, device="cuda")
        logits = torch.empty(1, QH, splits, D, dtype=torch.float32, device="cuda")
        lse = torch.empty(1, QH, splits, dtype=torch.float32, device="cuda")
        out = torch.empty_like(q)

        def run(c=kc, v=vc, kw=kv_args, o=out):
            decode_paged_attention(q, c, v, indptr, indices, qpos, logits, lse, D ** -0.5, out=o,
                                   max_kv_splits=args.splits, **kw)

        us = time_graph(run, 50)
        err = None
        if fmt.quantized:
            ref = torch.empty_like(q)
            run(kr, vr, ref_args, ref)
            torch.cuda.synchronize()
            err = (out.float() - ref.float()).abs().max().item()
        rows.append({"kernel": "decode", "ctx": ctx, "us": us, "max_err_vs_dequant": err,
                     "splits": args.splits, "dec_cfg": args.dec_cfg})
        print(json.dumps(rows[-1]), flush=True)

    for prefix in map(int, filter(None, args.prefix.split(","))):
        total = prefix + CHUNK
        q = torch.randn(CHUNK, QH, D, generator=gen, device="cuda").to(torch.bfloat16)
        ke = torch.randn(CHUNK, KVH, D, generator=gen, device="cuda").to(torch.bfloat16)
        ve = torch.randn(CHUNK, KVH, D, generator=gen, device="cuda").to(torch.bfloat16)
        qo = torch.tensor([0, CHUNK], dtype=torch.int32, device="cuda")
        kvp = torch.tensor([0, total], dtype=torch.int32, device="cuda")
        idx = torch.arange(total, dtype=torch.int32, device="cuda")
        pl = torch.tensor([prefix], dtype=torch.int32, device="cuda")
        out = torch.empty_like(q)

        def run(c=kc, v=vc, kw=kv_args, o=out):
            extend_paged_attention(q, c, v, qo, kvp, idx, pl, CHUNK, D ** -0.5, out=o, k_extend=ke, v_extend=ve, **kw)

        us = time_graph(run, 5)
        err = None
        if fmt.quantized:
            ref = torch.empty_like(q)
            run(kr, vr, ref_args, ref)
            torch.cuda.synchronize()
            err = (out.float() - ref.float()).abs().max().item()
        rows.append({"kernel": "extend", "prefix": prefix, "us": us, "max_err_vs_dequant": err})
        print(json.dumps(rows[-1]), flush=True)

    with open(args.output, "w") as f:
        json.dump({"fmt": args.fmt, "rows": rows}, f, indent=2)


if __name__ == "__main__":
    main()

xcompile_gfx12.py (RDNA4 cross-compile check)

"""Cross-compile decode stage1 for an RDNA4 target and hash the instruction stream (no launch)."""
import hashlib
import re
import sys

import torch
import triton.language as tl
from triton.backends.compiler import GPUTarget
from triton.runtime import driver

arch = sys.argv[1]
driver.active.get_current_target = lambda: GPUTarget("hip", arch, 32)
import freetoken.kernel.triton.attention as A  # noqa: E402

fn = A._decode_grouped_stage1_kernel
has_f16 = "F16_KV" in fn.arg_names
H, KVH, D = 8, 2, 128
dev = "cuda"
q = torch.zeros(1, H, D, dtype=torch.bfloat16, device=dev)
idx = torch.arange(512, dtype=torch.int32, device=dev)
indptr = torch.tensor([0, 512], dtype=torch.int32, device=dev)
qpos = torch.tensor([511], dtype=torch.int32, device=dev)
logits = torch.zeros(1, H, 64, D, device=dev)
lse = torch.zeros(1, H, 64, device=dev)
cases = [
    ("fp8_per_token_head", torch.uint8, 4, 64, 1, 2),   # gfx1201 packed reader
    ("fp8_per_token_head", torch.uint8, 1, 32, 3, 4),   # generic fp8 reader
    ("int8_per_token_head", torch.int8, 1, 32, 3, 4),
    ("none", torch.bfloat16, 1, 32, 3, 4),
]
for fmt, dt, pack, bn, st, wp in cases:
    kc = torch.zeros(1024, KVH, D, dtype=dt, device=dev)
    vc = torch.zeros_like(kc)
    ks = torch.zeros(1024, KVH, device=dev)
    vs = torch.zeros_like(ks)
    args = [q, kc, vc, ks, vs, D ** -0.5, indptr, idx, qpos, logits, lse, 64,
            q.stride(0), q.stride(1), kc.stride(0), kc.stride(1), vc.stride(0), vc.stride(1),
            ks.stride(0), vs.stride(0), logits.stride(0), logits.stride(1), logits.stride(2),
            lse.stride(0), lse.stride(1), lse.stride(2)]
    kw = dict(KV_OFFSET_DTYPE=tl.int32, KV_FMT=fmt, KV_PACK=pack, GROUP=4, NUM_Q_HEADS=H, BLOCK_D=D, BLOCK_DV=D,
              BLOCK_N=bn, BLOCK_H=4, VALID_BLOCK_H=4, MIN_BLOCK_KV=32, MIN_SPLIT_LEN=128, MIN_SPLITS=8, D=D, DV=D,
              SLIDING_WINDOW=0, num_warps=wp, num_stages=st)
    if has_f16:
        kw["F16_KV"] = False
    k = fn.warmup(*args, grid=(1,), **kw)
    asm = k.asm["amdgcn"]
    target = re.search(r'amdgcn-amd-amdhsa--(gfx\w+)', asm).group(1)
    instr = [l.split("//")[0].strip() for l in asm.splitlines()
             if re.match(r"\s+(v|s|global|buffer|ds|scratch)_[a-z0-9_]+", l)]
    print(f"{fmt:20s} pack={pack} target={target} instrs={len(instr)} "
          f"cvt_pk_fp8={sum('v_cvt_pk_f32_fp8' in i for i in instr)} "
          f"sha={hashlib.sha256(chr(10).join(instr).encode()).hexdigest()[:16]}")