perf(kernel): add opt-in hip fp8 decode kernels for r9700

benchmarks/ark/calibrate.py

@@ -0,0 +1,63 @@
+"""Recalibrate the exact source/environment; retain raw paired repeats and fallbacks."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import torch
+from freetoken.kernel.ark import initialize
+from freetoken.kernel.ark.common import policy
+
+from .cases import cases
+from .timing import assert_close, graph_pair, metrics, timing_pair
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--output", type=Path, required=True)
+    p.add_argument("--write-profile", type=Path)
+    a = p.parse_args()
+    torch.manual_seed(20261021)
+    torch.set_num_threads(8)
+    os.environ["FREETOKEN_USE_ARK"] = "1"
+    initialize()
+    policy.MODE = "native"
+    profile = {
+        "_environment": policy.current_environment(),
+        "_rules": {"eager_min_speedup": 1.05, "each_repeat_faster": True},
+        "_scope": "Eager decode only; Graph entries stay empty pending separate model qualification.",
+    }
+    rows = []
+    for op in ("dense_fp8_decode", "moe_decode_fp8"):
+        profile[op] = {"eager": {}, "graph": {}}
+        for case in cases(op):
+            expected = case.baseline().clone()
+            actual = case.candidate()
+            error = metrics(actual, expected)
+            assert_close(error, op)
+            row = {
+                "operator": op,
+                "case": case.name,
+                "metadata": case.metadata,
+                "correctness": error,
+                "eager": timing_pair(case.baseline, case.candidate, case.reset, False, 50),
+                "graph": graph_pair(case.baseline, case.candidate, case.reset, False, 100),
+            }
+            shape = policy.metadata_key(op, case.metadata)
+            timing = row["eager"]
+            paired = [x / y for x, y in zip(timing["samples_us"]["baseline"], timing["samples_us"]["candidate"])]
+            timing["enabled"] = timing["speedup"] >= 1.05 and min(paired) > 1
+            row["graph"]["enabled"] = False
+            if timing["enabled"]:
+                profile[op]["eager"][shape] = {"speedup": timing["speedup"], "min_paired_speedup": min(paired)}
+            rows.append(row)
+            a.output.parent.mkdir(parents=True, exist_ok=True)
+            a.output.write_text(json.dumps({"environment": profile["_environment"], "rows": rows}, indent=2) + "\n")
+            print(op, case.name, row["eager"]["speedup"], row["graph"]["speedup"], flush=True)
+    if a.write_profile:
+        a.write_profile.write_text(json.dumps(profile, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ark/cases.py

@@ -0,0 +1,75 @@
+"""Deterministic Qwen3 decode inputs, including the actual 20-GiB cache layout."""
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class Case:
+    name: str
+    baseline: object
+    candidate: object
+    metadata: dict
+
+    def reset(self):
+        pass
+
+
+def rand(*shape, scale=0.2, dtype=torch.bfloat16):
+    return (torch.randn(*shape, device="cuda") * scale).to(dtype)
+
+
+def quant(w):
+    *lead, n, k = w.shape
+    blocks = w.float().reshape(*lead, n // 128, 128, k // 128, 128)
+    scales = blocks.abs().amax(dim=(-3, -1), keepdim=True).clamp_min(1e-10) / 448
+    q = (blocks / scales).to(torch.float8_e4m3fn).reshape(*lead, n, k).contiguous()
+    return q, scales.squeeze(-3).squeeze(-1).bfloat16().contiguous()
+
+
+def dense_cases():
+    from freetoken.kernel.ark.dense_fp8_decode.adapter import (
+        block_fp8_linear as candidate,
+    )
+    from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear as baseline
+
+    for n, k in ((5120, 2048), (2048, 4096), (151936, 2048)):
+        w, scale = quant(rand(n, k, scale=k**-0.5))
+        x = rand(1, k, scale=0.5)
+        args = (x, w, scale)
+        yield Case(f"N{n}_K{k}", lambda a=args: baseline(*a), lambda a=args: candidate(*a), {"M": 1, "N": n, "K": k})
+
+
+def moe_cases():
+    from freetoken.kernel.ark.moe_decode_fp8.adapter import (
+        fused_experts_decode_fp8_block as candidate,
+    )
+    from freetoken.layers.quantization.moe.base import MoEConfig
+    from freetoken.layers.quantization.moe.fp8_block import TritonFp8BlockMoEKernel
+    from freetoken.moe.fused_fp8_block import fused_experts_decode_fp8_block as baseline
+
+    h, i, topk = 2048, 768, 8
+    for rows in (32, 4550):
+        specs = TritonFp8BlockMoEKernel().layout(MoEConfig(rows, h, i, topk))
+        banks = {name: torch.empty(rows, *s.shape, device="cuda", dtype=s.dtype) for name, s in specs.items()}
+        chosen = torch.linspace(1, rows - 1, 16, device="cuda").long()
+        gu, gs = quant(rand(16, 2 * i, h, scale=h**-0.5))
+        dn, ds = quant(rand(16, h, i, scale=i**-0.5))
+        banks["gate_up"].view(torch.uint8)[chosen] = gu.view(torch.uint8)
+        banks["down"].view(torch.uint8)[chosen] = dn.view(torch.uint8)
+        for name, values in (("gate_up_scale", gs), ("down_scale", ds)):
+            padded = torch.full((16, *banks[name].shape[1:]), 19.0, device="cuda", dtype=values.dtype)
+            padded[:, :, : values.shape[-1]] = values
+            banks[name][chosen] = padded
+        x = rand(1, h, scale=0.5)
+        ids = chosen[:topk].int().view(1, topk)
+        weights = torch.softmax(torch.randn(1, topk, device="cuda"), -1)
+        args = (x, banks["gate_up"], banks["gate_up_scale"], banks["down"], banks["down_scale"], weights, ids)
+        yield Case(
+            f"C{rows}_H{h}_I{i}", lambda a=args: baseline(*a), lambda a=args: candidate(*a), {"M": 1, "H": h, "I": i, "bank_rows": rows, "top_k": topk}
+        )
+
+
+def cases(op):
+    return moe_cases() if op == "moe_decode_fp8" else dense_cases()

benchmarks/ark/timing.py

@@ -0,0 +1,99 @@
+"""Correctness-gated, paired eager and HIP graph operator experiments."""
+
+import statistics
+import time
+
+import torch
+
+OPS = ["moe_decode_fp8", "dense_fp8_decode"]
+
+
+def metrics(actual, reference):
+    a, r = actual.float(), reference.float()
+    diff = (a - r).abs()
+    norm = diff.square().mean().sqrt() / (r.square().mean().sqrt() + 1e-12)
+    cos = torch.nn.functional.cosine_similarity(a.flatten(), r.flatten(), dim=0, eps=1e-30).item()
+    if r.abs().max().item() == 0:
+        cos = 1.0 if a.abs().max().item() == 0 else 0.0
+    return {
+        "normalized_rmse": norm.item(),
+        "cosine": cos,
+        "max_abs": diff.max().item(),
+        "max_ref": r.abs().max().item(),
+        "finite": bool(torch.isfinite(a).all().item()),
+    }
+
+
+def assert_close(m, op):
+    ok = m["finite"] and m["cosine"] > 0.999 and m["normalized_rmse"] < 0.03 and m["max_abs"] <= 0.04 * m["max_ref"] + 1e-3
+    if not ok:
+        raise AssertionError(m)
+
+
+def timing_pair(base, candidate, reset, restore, iters, repeats=5):
+    def invoke(fn):
+        if restore:
+            reset()
+        return fn()
+
+    for f in (base, candidate):
+        for _ in range(4):
+            invoke(f)
+    torch.cuda.synchronize()
+    samples = {"baseline": [], "candidate": []}
+    for r in range(repeats):
+        order = [("baseline", base), ("candidate", candidate)]
+        if r % 2:
+            order.reverse()
+        for name, fn in order:
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            for _ in range(iters):
+                invoke(fn)
+            torch.cuda.synchronize()
+            samples[name].append((time.perf_counter() - start) * 1e6 / iters)
+    return summarize(samples)
+
+
+def summarize(samples):
+    out = {k + "_us": statistics.median(v) for k, v in samples.items()}
+    out["speedup"] = out["baseline_us"] / out["candidate_us"]
+    out["samples_us"] = samples
+    return out
+
+
+def graph_pair(base, candidate, reset, restore, iters, repeats=5):
+    graphs = []
+    for fn in (base, candidate):
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for _ in range(4):
+                if restore:
+                    reset()
+                fn()
+        s.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        outputs = []
+        with torch.cuda.graph(graph, stream=s):
+            for _ in range(4):
+                if restore:
+                    reset()
+                out = fn()
+            outputs.append(out)
+        graphs.append((graph, outputs))
+    samples = {"baseline": [], "candidate": []}
+    for r in range(repeats):
+        order = [0, 1] if r % 2 == 0 else [1, 0]
+        for idx in order:
+            g = graphs[idx][0]
+            g.replay()
+            torch.cuda.synchronize()
+            start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+            start.record()
+            for _ in range(iters):
+                g.replay()
+            end.record()
+            end.synchronize()
+            samples["baseline" if idx == 0 else "candidate"].append(start.elapsed_time(end) * 1000 / (iters * 4))
+    return summarize(samples)

docs/ark-fp8-decode.md

@@ -0,0 +1,72 @@
+# Optional HIP FP8 decode kernels on R9700
+
+`FREETOKEN_USE_ARK=1` adds two native HIP candidates on Linux gfx1201:
+block-FP8 dense decode and block-FP8 routed MoE decode. The default is `0`.
+Both candidates retain independent FT fallbacks. The packaged profile enables
+only measured eager shapes; Graph entries remain empty.
+
+## Operator contracts
+
+- Dense: flattened M=1, BF16 activations, E4M3 weights and BF16 128x128 block
+  scales. W8A16 preserves leading dimensions and optional bias. M>1, including
+  batched decode, retains FT W8A8.
+- MoE: M=1, plain SiLU, BF16 activations, E4M3 weights, padded BF16 block
+  scales, contiguous INT32 slots and FP32 routing weights. W8A16 preserves
+  inputs and returns a new BF16 tensor. Prefill and unsupported configurations
+  retain FT.
+- The `LinearKernel` / `MoEKernel` subclasses inherit FT packing, layouts and
+  CPU formats and consume existing `ExpertView` banks. FTW remains compatible
+  with `fp8_block` / `triton` banks. CPU experts, hybrid routing, shared experts
+  and cache ownership remain with FT.
+- `supports_inactive_slots=False` preserves FT's safe slot-zero/zero-weight
+  routing. Explicit requests for the `triton` backend retain FT.
+
+## Build and selection
+
+Install in an existing ROCm environment with `uv pip install -e .`. The package
+includes both HIP translation units, the shared header, Torch bridge and shape
+profile. First enabled use compiles into `$FREETOKEN_ARK_CACHE_DIR`, defaulting
+to `$XDG_CACHE_HOME/freetoken/ark` (or `~/.cache/freetoken/ark`). Build keys include
+native source, compiler, build recipe and Torch/HIP versions. No external ARK
+workspace, generic `common` package or precompiled ARK binary is required.
+
+Selection checks GPU/architecture, visible device count, Torch/HIP/Triton,
+compiler, platform, ARK source and FT primitive fingerprints. Unmatched
+settings, shapes, dtypes, layouts and modes retain FT. This profile applies to
+the measured R9700 environment, not to every Radeon device or software stack.
+
+```bash
+export ROCR_VISIBLE_DEVICES=0
+export HIPCC=/opt/rocm/bin/hipcc
+export FREETOKEN_ARK_CACHE_DIR=/tmp/freetoken-ark
+FREETOKEN_USE_ARK=1 ft serve --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --attention-backend triton --moe-strategy offload --cuda-graph-max-bs 0
+```
+
+Set `FREETOKEN_USE_ARK=0` for the FT control. Default validated selection uses
+`common/validated_shapes.json`; `FREETOKEN_ARK_PROFILE` can name a separately
+calibrated profile. `FREETOKEN_ARK_POLICY=native` bypasses the measured table for
+operator validation and is not the production selection policy.
+
+## Validation and recalibration
+
+```bash
+python -m pytest tests/kernels/test_ark_fp8.py -q
+python -m benchmarks.ark.calibrate --output /tmp/ark-calibration.json --write-profile python/freetoken/kernel/ark/common/validated_shapes.json
+```
+
+The focused tests verify outputs, dtype/shape, input and alias contracts,
+repeated/zero/inactive routes, padded scales, bias, nondefault streams, updated
+inputs during direct Graph replay, FT fallback, bank packing, `Kernel.apply`
+and FTW selection. Native-hit tests observe the actual Torch operator call.
+Direct Graph API tests do not enable Graph in the model profile.
+
+The calibration benchmark covers only these two operators, records every
+paired eager/Graph repeat and checks numerical error against FT. Eager entries
+require at least 1.05x median speedup and every paired repeat to be faster.
+Graph entries stay empty; operator timing alone cannot qualify a model path.
+Changing source or runtime requires recalibration rather than removing checks.
+
+Model validation uses Qwen/Qwen3-30B-A3B-FP8 (local FTW), BF16 activations/KV,
+one R9700, offload, 20GiB expert cache, 512/32768 input tokens and 256 generated
+tokens. Updated A/B results and exact validation commands are in the PR body.
+This PR does not claim a Graph or sustained-CPU hybrid speedup.

pyproject.toml

@@ -111,6 +111,7 @@ version = {attr = "freetoken.version.__version__"}
 where = ["python"]
 
 [tool.setuptools.package-data]
+"freetoken.kernel.ark" = ["**/*.hip", "**/*.h", "**/*.cpp", "**/*.json", "NOTICE"]
 "*" = ["csrc/**/*", "moe/configs/**/*.json", "third_party/*/LICENSE", "third_party/*/VERSION", "third_party/*/freetoken.patch"]
 
 [tool.pytest.ini_options]

python/freetoken/engine/engine.py

@@ -378,6 +378,13 @@ def __init__(self, config: EngineConfig):
         torch.manual_seed(42)
         self.stream = torch.cuda.Stream()
         torch.cuda.set_stream(self.stream)
+        if os.name == "posix" and os.environ.get("FREETOKEN_USE_ARK", "0") == "1":
+            from freetoken.kernel.ark.common.policy import MODE, environment_matches
+
+            if MODE == "native" or environment_matches():
+                from freetoken.kernel.ark import initialize
+
+                initialize()
         self.dtype = config.dtype
         self.config = config  # retained for runtime cache rebuild (rebuild_runtime_cache)
         # KV pool family fixed at construction from the model config: its classmethods own the
@@ -814,7 +821,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                     f"expert banks hold experts of several formats ({sorted(set(banks.layer_formats))}) but the model "
                     f"binds {method.kind} / {method.kernel.name} on every MoE layer; reconvert the checkpoint"
                 )
-            if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name):
+            from freetoken.moe.legacy_format import bank_kernel_compatible
+
+            if banks.kind is not None and (
+                banks.kind != method.kind or not bank_kernel_compatible(method.kind, banks.kernel, method.kernel.name)
+            ):
                 raise ValueError(
                     f"expert banks were packed for {banks.kind} / {banks.kernel} but the model "
                     f"binds {method.kind} / {method.kernel.name}; reconvert or select that kernel"
@@ -1654,6 +1665,14 @@ def _adjust_ftw_quant_backend(model_path: str, quant_backend: QuantBackend) -> Q
     except KeyError:
         return quant_backend
     requested = quant_backend.select(LayerKind.MOE, kind)
+    if (
+        os.environ.get("FREETOKEN_USE_ARK", "0") == "1"
+        and kind is QuantKind.FP8_BLOCK and kernel == "triton"
+        and requested in ("auto", "ark_r9700")
+    ):
+        # Preserve auto selection for the layout-compatible candidate. Its own
+        # environment gate still selects FT when ARK has not been validated.
+        return quant_backend
     if requested == kernel:
         return quant_backend
     if requested != "auto":

python/freetoken/kernel/ark/NOTICE

@@ -0,0 +1,5 @@
+HIP kernels and adapter contracts adapted from Radeon Operator Library (ARK),
+revision 975e8cd9918ec1365d3f275d1259d215315f505a, and its FreeToken P0 experiments.
+The source project declares Apache-2.0 in its pyproject.toml.
+Original project: https://github.com/AMD-AIM/amd_radeon_kernels
+FreeToken integration retains source comments and the adapted C ABI.