diff --git a/benchmarks/ark/__init__.py b/benchmarks/ark/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/ark/calibrate.py b/benchmarks/ark/calibrate.py new file mode 100644 index 000000000..da93da84a --- /dev/null +++ b/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() diff --git a/benchmarks/ark/cases.py b/benchmarks/ark/cases.py new file mode 100644 index 000000000..adfa04b92 --- /dev/null +++ b/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() diff --git a/benchmarks/ark/timing.py b/benchmarks/ark/timing.py new file mode 100644 index 000000000..b4865d552 --- /dev/null +++ b/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) diff --git a/docs/ark-fp8-decode.md b/docs/ark-fp8-decode.md new file mode 100644 index 000000000..817bcfa92 --- /dev/null +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index ba3c3a769..c35288da6 100644 --- a/pyproject.toml +++ b/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] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 7982095ac..de20541a7 100644 --- a/python/freetoken/engine/engine.py +++ b/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": diff --git a/python/freetoken/kernel/ark/NOTICE b/python/freetoken/kernel/ark/NOTICE new file mode 100644 index 000000000..2eda14b59 --- /dev/null +++ b/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. diff --git a/python/freetoken/kernel/ark/__init__.py b/python/freetoken/kernel/ark/__init__.py new file mode 100644 index 000000000..ddd9dc23f --- /dev/null +++ b/python/freetoken/kernel/ark/__init__.py @@ -0,0 +1,38 @@ +"""Optional measured gfx1201 FP8 decode kernels, disabled by default.""" + +import os + + +def enabled(): + return os.name == "posix" and os.environ.get("FREETOKEN_USE_ARK", "0") == "1" + + +def initialize(): + """Build/load outside HIP Graph capture, into a writable source-keyed cache.""" + from .common.native import load + + load() + + +def candidate_kernels(method, candidates): + if not enabled(): + return candidates + from .common.policy import MODE, current_environment, environment_matches + + if current_environment().get("arch") != "gfx1201": + return candidates + if MODE != "native" and not environment_matches(): + return candidates + from freetoken.layers.quantization.registry import LayerKind + from freetoken.layers.quantization.scheme import QuantKind + + from .kernels import ArkFp8LinearKernel, ArkFp8MoEKernel + + cls = { + (QuantKind.FP8_BLOCK, LayerKind.MOE): ArkFp8MoEKernel, + (QuantKind.FP8_BLOCK, LayerKind.LINEAR): ArkFp8LinearKernel, + }.get((method.kind, method.layer_kind)) + if cls is None or cls().unusable_reason(method.cfg): + return candidates + initialize() + return (cls, *candidates) diff --git a/python/freetoken/kernel/ark/common/__init__.py b/python/freetoken/kernel/ark/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/freetoken/kernel/ark/common/bindings.cpp b/python/freetoken/kernel/ark/common/bindings.cpp new file mode 100644 index 000000000..7373e522e --- /dev/null +++ b/python/freetoken/kernel/ark/common/bindings.cpp @@ -0,0 +1,75 @@ +#include +#include +#include +#include +#include + +extern "C" { +int radeon_dense_gemv_fp8(const void*, const void*, const void*, void*, int, int, int, void*); +int radeon_moe_decode_gemv_fp8(const void*, const void*, const void*, const void*, const void*, const void*, const void*, void*, void*, int, int, int, int, int, int, void*); +} + +using at::Tensor; + +namespace { +void check(const Tensor& t, const Tensor& x, at::ScalarType dtype) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous() && t.scalar_type() == dtype, + "ARK requires contiguous tensors of the declared dtype on one GPU"); +} + +void checked(int rc) { TORCH_CHECK(rc == 0, "ARK native HIP return code ", rc); } +void* stream() { return (void*)at::cuda::getCurrentCUDAStream().stream(); } + +Tensor dense_decode(const Tensor& x, const Tensor& w, const Tensor& scale) { + c10::cuda::CUDAGuard guard(x.device()); + check(x, x, at::kBFloat16); + check(w, x, at::kFloat8_e4m3fn); + check(scale, x, at::kBFloat16); + TORCH_CHECK(x.dim() == 2 && x.size(0) == 1 && w.dim() == 2 && scale.dim() == 2, + "ARK dense decode requires M=1 and 2D weight/scales"); + int K = x.size(1), N = w.size(0); + TORCH_CHECK(K > 0 && N > 0 && K % 128 == 0 && w.size(1) == K, + "ARK dense decode dimensions/configuration"); + TORCH_CHECK(scale.size(0) >= (N + 127) / 128 && scale.size(1) >= K / 128, + "ARK dense scale storage too small"); + auto out = at::empty({1, N}, x.options()); + checked(radeon_dense_gemv_fp8(x.data_ptr(), w.data_ptr(), scale.data_ptr(), out.data_ptr(), + N, K, scale.stride(0), stream())); + return out; +} + +Tensor moe_decode(const Tensor& x, const Tensor& w1, const Tensor& w2, const Tensor& weights, + const Tensor& ids, const Tensor& s1, const Tensor& s2) { + c10::cuda::CUDAGuard guard(x.device()); + check(x, x, at::kBFloat16); + for (const auto& w : {w1, w2}) check(w, x, at::kFloat8_e4m3fn); + for (const auto& s : {s1, s2}) check(s, x, at::kBFloat16); + check(ids, x, at::kInt); + check(weights, x, at::kFloat); + TORCH_CHECK(x.dim() == 2 && x.size(0) == 1 && w1.dim() == 3 && w2.dim() == 3 && + s1.dim() == 3 && s2.dim() == 3 && ids.dim() == 2 && weights.sizes() == ids.sizes(), + "ARK MoE decode requires M=1 and packed expert banks"); + int H = x.size(1), I = w2.size(2), K = ids.size(1), E = w1.size(0); + TORCH_CHECK(H > 0 && I > 0 && H % 128 == 0 && I % 128 == 0 && E > 0 && K > 0 && K <= 16 && + ids.size(0) == 1 && w1.size(1) == 2 * I && w1.size(2) == H && + w2.size(0) == E && w2.size(1) == H, "ARK MoE decode dimensions"); + TORCH_CHECK(s1.size(0) == E && s1.size(1) == 2 * I / 128 && s1.size(2) >= H / 128 && + s2.size(0) == E && s2.size(1) == H / 128 && s2.size(2) >= I / 128, + "ARK MoE scale storage too small"); + auto intermediate = at::empty({1, K, I}, x.options().dtype(at::kFloat)); + auto out = at::empty({1, H}, x.options()); + checked(radeon_moe_decode_gemv_fp8(x.data_ptr(), w1.data_ptr(), s1.data_ptr(), w2.data_ptr(), + s2.data_ptr(), ids.data_ptr(), weights.data_ptr(), intermediate.data_ptr(), out.data_ptr(), + 1, K, H, I, s1.stride(1), s2.stride(1), stream())); + return out; +} +} // namespace + +TORCH_LIBRARY(freetoken_ark, m) { + m.def("dense_decode(Tensor x, Tensor w, Tensor scale) -> Tensor"); + m.def("moe_decode(Tensor x, Tensor w1, Tensor w2, Tensor weights, Tensor ids, Tensor s1, Tensor s2) -> Tensor"); +} +TORCH_LIBRARY_IMPL(freetoken_ark, CUDA, m) { + m.impl("dense_decode", TORCH_FN(dense_decode)); + m.impl("moe_decode", TORCH_FN(moe_decode)); +} diff --git a/python/freetoken/kernel/ark/common/contracts.py b/python/freetoken/kernel/ark/common/contracts.py new file mode 100644 index 000000000..1ec2e7a22 --- /dev/null +++ b/python/freetoken/kernel/ark/common/contracts.py @@ -0,0 +1,63 @@ +import torch + + +def gfx1201(tensor): + return ( + bool(torch.version.hip) + and tensor.is_cuda + and getattr(torch.cuda.get_device_properties(tensor.device), "gcnArchName", "").startswith("gfx1201") + ) + + +def plain_silu(activation, apply_router_weight_on_input, act_alpha, act_limit): + return activation == "silu" and not apply_router_weight_on_input and act_alpha == 1.0 and act_limit == float("inf") + + +def moe_supported(x, w1, w2, tw, ids): + dtype = torch.float8_e4m3fn + return ( + gfx1201(x) + and x.dtype == torch.bfloat16 + and x.ndim == 2 + and x.is_contiguous() + and w1.ndim == 3 + and w2.ndim == 3 + and w1.dtype == dtype + and w2.dtype == dtype + and w1.is_contiguous() + and w2.is_contiguous() + and w1.shape[0] == w2.shape[0] + and w1.shape[1] == 2 * w2.shape[2] + and w1.shape[2] == x.shape[1] + and w2.shape[1] == x.shape[1] + and ids.ndim == 2 + and ids.shape == tw.shape + and ids.shape[0] == x.shape[0] + and ids.dtype == torch.int32 + and tw.dtype == torch.float32 + and ids.is_contiguous() + and tw.is_contiguous() + and x.shape[0] == 1 + and w1.shape[0] > 0 + and 0 < ids.shape[1] + and w2.shape[2] > 0 + and x.shape[1] % 128 == 0 + and w2.shape[2] % 128 == 0 + and ids.shape[1] <= 16 + and all(t.device == x.device for t in (w1, w2, tw, ids)) + ) + + +def scales_supported(s1, s2, w1, w2): + return ( + s1.dtype == s2.dtype == torch.bfloat16 + and s1.is_contiguous() + and s2.is_contiguous() + and s1.ndim == s2.ndim == 3 + and s1.shape[:2] == (w1.shape[0], w1.shape[1] // 128) + and s2.shape[:2] == (w2.shape[0], w2.shape[1] // 128) + and s1.shape[2] >= w1.shape[2] // 128 + and s2.shape[2] >= w2.shape[2] // 128 + and s1.device == w1.device + and s2.device == w2.device + ) diff --git a/python/freetoken/kernel/ark/common/include/ft_decode_impl.h b/python/freetoken/kernel/ark/common/include/ft_decode_impl.h new file mode 100644 index 000000000..5568ad458 --- /dev/null +++ b/python/freetoken/kernel/ark/common/include/ft_decode_impl.h @@ -0,0 +1,79 @@ +// ARK MoE decode with FreeToken intermediate rounding, without atomics or a memset. +#pragma once +#include +#include +#include +namespace ark_ft { +using bf16=__hip_bfloat16; +__device__ __forceinline__ float sum(float x){for(int s=16;s;s>>=1)x+=__shfl_xor(x,s);return x;} +__device__ __forceinline__ float e4m3(unsigned char x){ + #if defined(__gfx1201__) + return __builtin_amdgcn_cvt_f32_fp8((int)x,0); +#else + unsigned short u=((unsigned)(x&128)<<8)|((unsigned)(x&127)<<7);__half v;__builtin_memcpy(&v,&u,2);return __half2float(v)*256.f; +#endif +} +template __device__ __forceinline__ float product(float a,float b){ + if constexpr(FP8)return a*b;else return (float)(bf16)(a*b); +} +template +__global__ __launch_bounds__(WARPS*32) void gate_up(const bf16* x,const void* weights,const bf16* scales, + const int* ids,const float* tw,float* h,int T,int K,int H,int I,int ss){ + int row=blockIdx.x*WARPS+(threadIdx.x>>5),lane=threadIdx.x&31; + if(row>=T*K*I)return; + int j=row%I,route=row/I,t=route/K,e=ids[route]; + if(e<0 || tw[route]==0.f){if(lane==0)h[row]=0.f;return;} + float ag=0.f,au=0.f; + for(int kb=0;kb>(8*z))&255)*(float)xx[z];pu+=e4m3((uw>>(8*z))&255)*(float)xx[z];} + ag+=pg*(float)scales[((size_t)e*(2*I/128)+j/128)*ss+kb];au+=pu*(float)scales[((size_t)e*(2*I/128)+(I+j)/128)*ss+kb]; + }else{ + auto w=(const bf16*)weights; + float2 gr=*(const float2*)(w+((size_t)e*2*I+j)*H+k),ur=*(const float2*)(w+((size_t)e*2*I+I+j)*H+k); + auto gg=(const bf16*)&gr;auto uu=(const bf16*)&ur; + #pragma unroll + for(int z=0;z<4;++z){ag+=product((float)gg[z],(float)xx[z]);au+=product((float)uu[z],(float)xx[z]);} + } + } + ag=(float)(bf16)sum(ag);au=(float)(bf16)sum(au); + if(lane==0)h[row]=(float)(bf16)((ag/(1.f+expf(-ag)))*au); +} +template +__global__ __launch_bounds__(WARPS*32) void down_reduce(const float* h,const void* weights,const bf16* scales, + const int* ids,const float* tw,bf16* out,int T,int K,int H,int I,int ss){ + int row=blockIdx.x*WARPS+(threadIdx.x>>5),lane=threadIdx.x&31;if(row>=T*H)return; + int t=row/H,n=row%H;float total=0.f; + for(int route=0;route>(8*z))&255)*hv[z]; + acc+=p*(float)scales[((size_t)e*(H/128)+n/128)*ss+kb]; + }else{ + auto w=(const bf16*)weights;float2 wr=*(const float2*)(w+((size_t)e*H+n)*I+k);auto wv=(const bf16*)≀ + #pragma unroll + for(int z=0;z<4;++z)acc+=product((float)wv[z],hv[z]); + } + } + total+=(float)(bf16)(sum(acc)*weight); + } + if(lane==0)out[row]=(bf16)total; +} +template +void launch(const void* x,const void* w1,const void* s1,const void* w2,const void* s2,const void* ids,const void* tw, + void* h,void* out,int T,int K,int H,int I,int ss1,int ss2,void* stream){ + auto s=(hipStream_t)stream; + gate_up<<>>((const bf16*)x,w1,(const bf16*)s1,(const int*)ids,(const float*)tw,(float*)h,T,K,H,I,ss1); + down_reduce<<>>((const float*)h,w2,(const bf16*)s2,(const int*)ids,(const float*)tw,(bf16*)out,T,K,H,I,ss2); +} +} diff --git a/python/freetoken/kernel/ark/common/native.py b/python/freetoken/kernel/ark/common/native.py new file mode 100644 index 000000000..c27d58d8d --- /dev/null +++ b/python/freetoken/kernel/ark/common/native.py @@ -0,0 +1,122 @@ +"""Source-only HIP build, keyed by source, compiler, Torch and architecture.""" + +import fcntl +import hashlib +import json +import os +import shutil +import subprocess +import threading +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +_LOADED = False +_LOCK = threading.Lock() + + +def source_digest(native_only=False): + suffixes = (".hip", ".h", ".cpp") if native_only else (".hip", ".h", ".cpp", ".py") + paths = sorted(p for p in ROOT.rglob("*") if p.suffix in suffixes and "__pycache__" not in p.parts) + h = hashlib.sha256() + for p in paths: + h.update(str(p.relative_to(ROOT)).encode()) + h.update(p.read_bytes()) + return h.hexdigest() + + +def _compiler(): + from torch.utils.cpp_extension import ROCM_HOME + + if os.environ.get("HIPCC"): + return os.environ["HIPCC"] + for root in (ROCM_HOME, os.environ.get("ROCM_PATH"), "/opt/rocm"): + if root and (Path(root) / "bin/hipcc").is_file(): + return str(Path(root) / "bin/hipcc") + return shutil.which("hipcc") or "/opt/rocm/bin/hipcc" + + +def cache_dir(): + import torch + + compiler = _compiler() + version = subprocess.check_output([compiler, "--version"]) + recipe = Path(__file__).read_bytes() + key = hashlib.sha256( + source_digest(native_only=True).encode() + recipe + str(torch.__version__).encode() + str(torch.version.hip).encode() + version + ).hexdigest()[:20] + root = Path(os.environ.get("FREETOKEN_ARK_CACHE_DIR", Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "freetoken/ark")) + return root / key + + +def library_path(): + return cache_dir() / "libfreetoken_ark.so" + + +def load(): + global _LOADED + if _LOADED: + return + with _LOCK: + if _LOADED: + return + import torch + + if not torch.version.hip or not torch.cuda.is_available(): + raise RuntimeError("ARK requires ROCm and a gfx1201 GPU") + if not torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.startswith("gfx1201"): + raise RuntimeError("ARK kernels are compiled for gfx1201 only") + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Call freetoken.kernel.ark.initialize() before HIP Graph capture") + build = cache_dir() + build.mkdir(parents=True, exist_ok=True) + with (build / "build.lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + lib = build / "libfreetoken_ark.so" + if not lib.exists(): + compiler = _compiler() + sources = sorted(ROOT.glob("*/*.hip")) + flags = ["-O3", "-std=c++17", "--offload-arch=gfx1201", "-fPIC", "-I" + str(ROOT / "common/include")] + objects = [] + for source in sources: + obj = build / (source.parent.name + "_" + source.stem + ".o") + command = [compiler, *flags, "-c", str(source), "-o", str(obj)] + proc = subprocess.run(command, check=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + (build / (obj.stem + ".log")).write_text(proc.stdout) + if proc.returncode: + raise RuntimeError(f"ARK HIP build failed: {command}\n{proc.stdout}") + objects.append(obj) + subprocess.run([compiler, "-shared", *map(str, objects), "-o", str(lib)], check=True) + (build / "manifest.json").write_text( + json.dumps( + { + "source_sha256": source_digest(), + "torch": str(torch.__version__), + "hip": str(torch.version.hip), + "compiler": subprocess.check_output([compiler, "--version"], text=True), + "flags": flags, + "library_sha256": hashlib.sha256(lib.read_bytes()).hexdigest(), + }, + indent=2, + ) + ) + # hipify may generate a sibling source; keep installed packages read-only. + binding = build / "bindings.cpp" + shutil.copyfile(ROOT / "common/bindings.cpp", binding) + from torch.utils.cpp_extension import load as cpp_load + + cpp_load( + name="freetoken_ark_ops_" + build.name, + sources=[str(binding)], + extra_cflags=["-O3"], + extra_ldflags=[f"-L{build}", "-lfreetoken_ark", f"-Wl,-rpath,{build}"], + build_directory=str(build), + with_cuda=True, + is_python_module=False, + verbose=os.environ.get("FREETOKEN_ARK_BUILD_VERBOSE") == "1", + ) + _LOADED = True + + +if __name__ == "__main__": + load() + print(library_path()) diff --git a/python/freetoken/kernel/ark/common/policy.py b/python/freetoken/kernel/ark/common/policy.py new file mode 100644 index 000000000..c20be89aa --- /dev/null +++ b/python/freetoken/kernel/ark/common/policy.py @@ -0,0 +1,78 @@ +"""Fail-closed environment and measured-shape selection, separate for eager/Graph.""" + +import hashlib +import json +import os +import platform +import subprocess +from functools import lru_cache +from pathlib import Path + +import torch + +MODE = os.environ.get("FREETOKEN_ARK_POLICY", "validated") +if MODE not in ("native", "validated", "baseline"): + raise ValueError("FREETOKEN_ARK_POLICY must be native, validated or baseline") +PROFILE_PATH = Path(os.environ.get("FREETOKEN_ARK_PROFILE", Path(__file__).with_name("validated_shapes.json"))) +_profiles = json.loads(PROFILE_PATH.read_text()) if PROFILE_PATH.exists() else {} + + +@lru_cache(maxsize=1) +def current_environment(): + import triton + + if not torch.version.hip or not torch.cuda.is_available(): + return {} + from .native import _compiler, source_digest + + p = torch.cuda.get_device_properties(torch.cuda.current_device()) + try: + compiler = subprocess.check_output([_compiler(), "--version"]) + except (OSError, subprocess.CalledProcessError): + return {} + package = Path(__file__).resolve().parents[3] + baseline_paths = ("moe/fused_fp8_block.py", "kernel/triton/fp8_block_linear.py") + baseline = hashlib.sha256(b"".join((package / path).read_bytes() for path in baseline_paths)).hexdigest() + return { + "torch": str(torch.__version__), + "triton": triton.__version__, + "hip": str(torch.version.hip), + "gpu_name": p.name, + "arch": p.gcnArchName.split(":")[0], + "source_sha256": source_digest(), + "device_count": torch.cuda.device_count(), + "platform": platform.platform(), + "compiler_sha256": hashlib.sha256(compiler).hexdigest(), + "baseline_sha256": baseline, + } + + +@lru_cache(maxsize=1) +def environment_matches(): + expected = _profiles.get("_environment", {}) + actual = current_environment() + return bool(actual) and all(expected.get(k) == v for k, v in actual.items()) + + +def key(*values): + return ",".join(str(int(v)) for v in values) + + +def allowed(operator, shape_key, *, capturing=None): + if os.environ.get("FREETOKEN_USE_ARK", "0") != "1" or MODE == "baseline": + return False + if MODE == "native": + return True + if not environment_matches(): + return False + if capturing is None: + capturing = torch.cuda.is_current_stream_capturing() + return shape_key in _profiles.get(operator, {}).get("graph" if capturing else "eager", {}) + + +def metadata_key(op, m): + if op == "moe_decode_fp8": + return key(m["M"], m["H"], m["I"], m["bank_rows"], m["top_k"]) + if op == "dense_fp8_decode": + return key(m["M"], m["N"], m["K"]) + raise ValueError(f"unsupported ARK operator: {op}") diff --git a/python/freetoken/kernel/ark/common/validated_shapes.json b/python/freetoken/kernel/ark/common/validated_shapes.json new file mode 100644 index 000000000..84c0ffe7a --- /dev/null +++ b/python/freetoken/kernel/ark/common/validated_shapes.json @@ -0,0 +1,49 @@ +{ + "_environment": { + "torch": "2.14.0+rocm10.1.0", + "triton": "3.8.0", + "hip": "7.16.26385", + "gpu_name": "AMD Radeon AI PRO R9700", + "arch": "gfx1201", + "source_sha256": "27b45edd3f1fd2aa950ff6a6f1c714020e78f525f8c022b5ff46f727f896cf6b", + "device_count": 1, + "platform": "Linux-7.0.0-38-generic-x86_64-with-glibc2.39", + "compiler_sha256": "449f6779a4041275d18d801246d0c2c09ffb453a48093d4bebbbd82803954205", + "baseline_sha256": "9815bbdde6aa7c596020f7c1ec5c6e5ea3b363697877a6a052c833e34d74b3af" + }, + "_rules": { + "eager_min_speedup": 1.05, + "each_repeat_faster": true + }, + "_scope": "Eager decode only; Graph entries stay empty pending separate model qualification.", + "dense_fp8_decode": { + "eager": { + "1,5120,2048": { + "speedup": 2.8881421083394208, + "min_paired_speedup": 2.258698800667121 + }, + "1,2048,4096": { + "speedup": 2.7925056711471514, + "min_paired_speedup": 2.7456354563614886 + }, + "1,151936,2048": { + "speedup": 1.1631544640714144, + "min_paired_speedup": 1.1601689528473773 + } + }, + "graph": {} + }, + "moe_decode_fp8": { + "eager": { + "1,2048,768,32,8": { + "speedup": 2.149956471331526, + "min_paired_speedup": 2.097327574739232 + }, + "1,2048,768,4550,8": { + "speedup": 2.6720033324525043, + "min_paired_speedup": 2.6453511021618126 + } + }, + "graph": {} + } +} diff --git a/python/freetoken/kernel/ark/dense_fp8_decode/__init__.py b/python/freetoken/kernel/ark/dense_fp8_decode/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py b/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py new file mode 100644 index 000000000..5b3addb1e --- /dev/null +++ b/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py @@ -0,0 +1,39 @@ +"""M=1 FreeToken block_fp8_linear, including bias and leading dimensions.""" + +import torch + +from freetoken.kernel.ark.common.contracts import gfx1201 +from freetoken.kernel.ark.common.native import load +from freetoken.kernel.ark.common.policy import allowed, key + + +def block_fp8_linear(x, weight, weight_scale, bias=None): + k = x.shape[-1] + if not ( + weight.ndim == 2 + and allowed("dense_fp8_decode", key(x.numel() // k, weight.shape[0], k)) + and gfx1201(x) + and x.dtype == torch.bfloat16 + and x.numel() == k + and x.is_contiguous() + and weight.dtype == torch.float8_e4m3fn + and weight.ndim == 2 + and weight.is_contiguous() + and weight.shape[1] == k + and k % 128 == 0 + and weight_scale.dtype == torch.bfloat16 + and weight_scale.ndim == 2 + and weight_scale.is_contiguous() + and weight_scale.shape[0] * 128 >= weight.shape[0] + and weight_scale.shape[1] >= k // 128 + and weight.device == x.device + and weight_scale.device == x.device + ): + from freetoken.kernel.triton.fp8_block_linear import ( + block_fp8_linear as baseline, + ) + + return baseline(x, weight, weight_scale, bias) + load() + out = torch.ops.freetoken_ark.dense_decode(x.reshape(1, k), weight, weight_scale).reshape(*x.shape[:-1], weight.shape[0]) + return out + bias.to(out.dtype) if bias is not None else out diff --git a/python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip b/python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip new file mode 100644 index 000000000..c9712f92c --- /dev/null +++ b/python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip @@ -0,0 +1,79 @@ +// Radeon Operator Library — dense fp8 block-scale GEMV (HIP, W8A16, decode M=1). +// +// For the DeepSeek-V3-style fp8-block dense linears (GDN in_proj/out_proj, lm_head): weight [N,K] fp8-e4m3 +// with a per-128x128-block bf16 scale (weight_scale_inv [N/128, K/128]); activation x [1,K] bf16 stays +// bf16 (weight-only fp8). One wave per output row reduces over K in 128-element blocks, each block's partial +// scaled by weight_scale_inv[n/128][kb] before the single final warp reduce (the block scale is a scalar, +// so it factors out of the reduction). Vectorized: each lane reads a uint (4 fp8 weights) + a float2 of +// 4 bf16 activations per step (K%128==0, N any). Mirrors moe_decode_gemv_fp8 stage2 minus expert/h/atomic. + +#include +#include +#include + +namespace radeon_ops { +namespace hip_native { + +using bf16 = __hip_bfloat16; + +__device__ __forceinline__ float dg_wred(float x) { + for (int o = warpSize / 2; o > 0; o >>= 1) x += __shfl_xor(x, o); + return x; +} + +__device__ __forceinline__ float dg_e4m3_to_f32(unsigned char v) { +#if defined(__gfx1201__) + return __builtin_amdgcn_cvt_f32_fp8((int)v,0); +#else + unsigned short bits = (unsigned short)((unsigned)(v & 0x80) << 8) | + (unsigned short)((unsigned)(v & 0x7F) << 7); + __half hf; + __builtin_memcpy(&hf, &bits, sizeof(hf)); + return __half2float(hf) * 256.0f; +#endif +} + +template +__global__ void __launch_bounds__(ROWS * 32) dense_gemv_fp8_rows(const bf16* __restrict__ x, + const unsigned char* __restrict__ w, const bf16* __restrict__ wscale, + bf16* __restrict__ out, int N, int K, int wscale_stride) { + int n = blockIdx.x*ROWS+(threadIdx.x>>5); + if (n >= N) return; + int lane = threadIdx.x & 31; + const unsigned char* wr = w + (size_t)n * K; + const bf16* sc = wscale + (size_t)(n / 128) * wscale_stride; // 128 output rows share one scale row + int nkb = K / 128; + float acc = 0.f; // per-lane running sum (this warp's blocks) + for (int kb = 0; kb < nkb; ++kb) { + int idx = kb * 128 + lane * 4; // 32 lanes x 4 fp8 = 128 weights/block + unsigned int wv = *(const unsigned int*)(wr + idx); + float2 xv = *(const float2*)(x + idx); + const bf16* xp = (const bf16*)&xv; + float p = 0.f; + #pragma unroll + for (int b = 0; b < 4; ++b) p += dg_e4m3_to_f32((wv >> (8 * b)) & 0xFF) * __bfloat162float(xp[b]); + acc += p * __bfloat162float(sc[kb]); + } + #pragma unroll + for (int o = 16; o > 0; o >>= 1) acc += __shfl_xor(acc, o); // reduce this warp's lanes + if (lane == 0) out[n] = (bf16)acc; +} + +void launch_dense_gemv_fp8(const bf16* x, const unsigned char* w, const bf16* wscale, bf16* out, + int N, int K, int wscale_stride, hipStream_t stream) { + hipLaunchKernelGGL((dense_gemv_fp8_rows<4>), dim3((N+3)/4), dim3(128), 0, stream, + x, w, wscale, out, N, K, wscale_stride); +} + +} // namespace hip_native +} // namespace radeon_ops + +// Dense fp8 block-scale GEMV (W8A16, M=1). x [1,K] bf16; w [N,K] fp8-e4m3; wscale [N/128,K/128] bf16 +// (weight_scale_inv; wscale_stride = its dim-0/per-output-block stride); out [1,N] bf16. K % 128 == 0. +extern "C" int radeon_dense_gemv_fp8(const void* x, const void* w, const void* wscale, void* out, + int N, int K, int wscale_stride, void* stream) { + radeon_ops::hip_native::launch_dense_gemv_fp8( + (const __hip_bfloat16*)x, (const unsigned char*)w, (const __hip_bfloat16*)wscale, + (__hip_bfloat16*)out, N, K, wscale_stride, (hipStream_t)stream); + return (int)hipGetLastError(); +} diff --git a/python/freetoken/kernel/ark/kernels.py b/python/freetoken/kernel/ark/kernels.py new file mode 100644 index 000000000..606e6caea --- /dev/null +++ b/python/freetoken/kernel/ark/kernels.py @@ -0,0 +1,40 @@ +"""Reuse FreeToken packing, expert views and CPU formats; replace decode only.""" + +from freetoken.layers.quantization.linear.fp8_block import TritonFp8BlockLinearKernel +from freetoken.layers.quantization.moe.base import limit_or_inf +from freetoken.layers.quantization.moe.fp8_block import TritonFp8BlockMoEKernel + + +class ArkFp8MoEKernel(TritonFp8BlockMoEKernel): + name = "ark_r9700" + # FT fallbacks require the caller's safe slot-zero/zero-weight CPU routes. + supports_inactive_slots = False + + def apply(self, layer, x, topk_weights, topk_ids, view, *, is_prefill): + if is_prefill: + return super().apply(layer, x, topk_weights, topk_ids, view, is_prefill=True) + from .moe_decode_fp8.adapter import fused_experts_decode_fp8_block + + t = view.tensors + return fused_experts_decode_fp8_block( + x, + t["gate_up"], + t["gate_up_scale"], + t["down"], + t["down_scale"], + topk_weights, + topk_ids, + layer.activation, + layer.apply_router_weight_on_input, + float(layer.alpha), + limit_or_inf(layer), + ) + + +class ArkFp8LinearKernel(TritonFp8BlockLinearKernel): + name = "ark_r9700" + + def apply(self, layer, x): + from .dense_fp8_decode.adapter import block_fp8_linear + + return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias) diff --git a/python/freetoken/kernel/ark/moe_decode_fp8/__init__.py b/python/freetoken/kernel/ark/moe_decode_fp8/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py b/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py new file mode 100644 index 000000000..c5ff52741 --- /dev/null +++ b/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py @@ -0,0 +1,57 @@ +"""Drop-in for freetoken.moe.fused_fp8_block.fused_experts_decode_fp8_block.""" + +import torch + +from freetoken.kernel.ark.common.contracts import ( + moe_supported, + plain_silu, + scales_supported, +) +from freetoken.kernel.ark.common.native import load +from freetoken.kernel.ark.common.policy import allowed, key + + +def fused_experts_decode_fp8_block( + hidden_states, + gate_up, + gate_up_scale, + down, + down_scale, + topk_weights, + topk_ids, + activation="silu", + apply_router_weight_on_input=False, + act_alpha=1.0, + act_limit=float("inf"), +): + if not ( + hidden_states.ndim == 2 + and gate_up.ndim == 3 + and down.ndim == 3 + and topk_ids.ndim == 2 + and allowed( + "moe_decode_fp8", key(hidden_states.shape[0], hidden_states.shape[1], down.shape[2], gate_up.shape[0], topk_ids.shape[1]) + ) + and plain_silu(activation, apply_router_weight_on_input, act_alpha, act_limit) + and moe_supported(hidden_states, gate_up, down, topk_weights, topk_ids) + and scales_supported(gate_up_scale, down_scale, gate_up, down) + ): + from freetoken.moe.fused_fp8_block import ( + fused_experts_decode_fp8_block as baseline, + ) + + return baseline( + hidden_states, + gate_up, + gate_up_scale, + down, + down_scale, + topk_weights, + topk_ids, + activation, + apply_router_weight_on_input, + act_alpha, + act_limit, + ) + load() + return torch.ops.freetoken_ark.moe_decode(hidden_states, gate_up, down, topk_weights, topk_ids, gate_up_scale, down_scale) diff --git a/python/freetoken/kernel/ark/moe_decode_fp8/moe_decode_gemv_fp8.hip b/python/freetoken/kernel/ark/moe_decode_fp8/moe_decode_gemv_fp8.hip new file mode 100644 index 000000000..004da33f2 --- /dev/null +++ b/python/freetoken/kernel/ark/moe_decode_fp8/moe_decode_gemv_fp8.hip @@ -0,0 +1,6 @@ +#include "ft_decode_impl.h" +extern "C" int radeon_moe_decode_gemv_fp8(const void* x,const void* w1,const void* s1,const void* w2,const void* s2, + const void* ids,const void* tw,void* h,void* out,int T,int K,int H,int I,int ss1,int ss2,void* stream){ + ark_ft::launch(x,w1,s1,w2,s2,ids,tw,h,out,T,K,H,I,ss1,ss2,stream); + return (int)hipGetLastError(); +} diff --git a/python/freetoken/layers/quantization/method.py b/python/freetoken/layers/quantization/method.py index 2a079539f..a11eed147 100644 --- a/python/freetoken/layers/quantization/method.py +++ b/python/freetoken/layers/quantization/method.py @@ -2,6 +2,8 @@ from __future__ import annotations +import os + from abc import ABC from typing import Any, ClassVar, Sequence @@ -55,7 +57,13 @@ class QuantMethod(ABC): def __init__(self, cfg: Any, requested: str = "auto"): self.cfg = cfg self.requested = requested - self.kernel = select_kernel(self.candidates, requested, cfg) + candidates = self.candidates + + if os.environ.get("FREETOKEN_USE_ARK", "0") == "1": + from freetoken.kernel.ark import candidate_kernels + + candidates = candidate_kernels(self, candidates) + self.kernel = select_kernel(candidates, requested, cfg) @property def scheme(self): diff --git a/python/freetoken/moe/legacy_format.py b/python/freetoken/moe/legacy_format.py index 8f347e5c3..a3c2048e9 100644 --- a/python/freetoken/moe/legacy_format.py +++ b/python/freetoken/moe/legacy_format.py @@ -29,6 +29,7 @@ def legacy_bank_names(quant_format: str) -> dict[str, str]: LEGACY_FORMAT = { (QuantKind.NONE, "fused"): "bf16", (QuantKind.FP8_BLOCK, "triton"): "fp8_block", + (QuantKind.FP8_BLOCK, "ark_r9700"): "fp8_block", (QuantKind.NVFP4, "triton"): "nvfp4", # same banks as triton, plus the resident input scales; a "nvfp4" file reads back as triton (QuantKind.NVFP4, "w4a4"): "nvfp4", @@ -48,3 +49,8 @@ def legacy_format_for(kind: QuantKind, kernel: str) -> str: def kind_kernel_for(legacy_format: str) -> tuple[QuantKind, str]: return _KIND_KERNEL[legacy_format] + + +def bank_kernel_compatible(kind: QuantKind, stored: str, selected: str) -> bool: + """ARK inherits the block-FP8 Triton layout; other kernels keep strict identity.""" + return stored == selected or (kind is QuantKind.FP8_BLOCK and stored == "triton" and selected == "ark_r9700") diff --git a/tests/kernels/test_ark_fp8.py b/tests/kernels/test_ark_fp8.py new file mode 100644 index 000000000..b5193187e --- /dev/null +++ b/tests/kernels/test_ark_fp8.py @@ -0,0 +1,292 @@ +"""Native decode, fallback, bank and stream contracts against independent FT paths.""" + +import os +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +from benchmarks.ark.cases import quant, rand +from benchmarks.ark.timing import assert_close, metrics + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not torch.version.hip, + reason="requires ROCm gfx1201", +) + + +@pytest.fixture(autouse=True) +def native(monkeypatch): + if not torch.cuda.get_device_properties(0).gcnArchName.startswith("gfx1201"): + pytest.skip("requires gfx1201") + from freetoken.kernel.ark import initialize + from freetoken.kernel.ark.common import policy + + monkeypatch.setenv("FREETOKEN_USE_ARK", "1") + monkeypatch.setattr(policy, "MODE", "native") + initialize() + torch.manual_seed(20261021) + + +def compare(a, b): + assert a.shape == b.shape and a.dtype == b.dtype and a.device == b.device + assert_close(metrics(a, b), "decode") + + +def moe_inputs(m=1): + x = rand(m, 256) + w1, s1 = quant(rand(4, 256, 256, scale=1 / 16)) + w2, s2 = quant(rand(4, 256, 128, scale=128**-0.5)) + s1 = torch.cat((s1, torch.full((4, 2, 3), 19.0, device="cuda", dtype=s1.dtype)), -1) + s2 = torch.cat((s2, torch.full((4, 2, 2), 23.0, device="cuda", dtype=s2.dtype)), -1) + weights = torch.tensor([[0.2, 0.3, 0.0, 0.5]], device="cuda").expand(m, 4).contiguous() + ids = torch.tensor([[1, 1, 0, 3]], device="cuda", dtype=torch.int32).expand(m, 4).contiguous() + return x, w1, s1, w2, s2, weights, ids + + +@pytest.mark.parametrize("value", [None, 0.0, 1e-12]) +@pytest.mark.parametrize("bias", [False, True]) +def test_dense_bias_padding_alias(value, bias, monkeypatch): + from freetoken.kernel.ark.dense_fp8_decode import adapter + from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear as baseline + + x = rand(1, 1, 384) if value is None else torch.full((1, 1, 384), value, device="cuda", dtype=torch.bfloat16) + initial = x.clone() + w, scale = quant(rand(384, 384)) + scale = torch.cat((scale, torch.full((3, 2), 17.0, device="cuda", dtype=scale.dtype)), -1) + b = rand(384, dtype=torch.float32) if bias else None + hits = [] + native = torch.ops.freetoken_ark.dense_decode + + def dispatch(*args): + hits.append(True) + return native(*args) + + monkeypatch.setattr(torch.ops.freetoken_ark, "dense_decode", dispatch) + actual = adapter.block_fp8_linear(x, w, scale, b) + compare(actual, baseline(x, w, scale, b)) + assert torch.equal(x, initial) and actual.data_ptr() != x.data_ptr() + assert hits == [True] + + +@pytest.mark.parametrize("kind", ["batch", "strided", "disabled", "unknown_shape"]) +def test_dense_fallback(kind, monkeypatch): + from freetoken.kernel.ark.common import policy + from freetoken.kernel.ark.dense_fp8_decode import adapter + from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear as baseline + + x = rand(2 if kind == "batch" else 1, 256) + if kind == "strided": + x = rand(1, 512)[:, ::2] + w, scale = quant(rand(256, 256)) + if kind == "disabled": + monkeypatch.setenv("FREETOKEN_USE_ARK", "0") + if kind == "unknown_shape": + monkeypatch.setattr(policy, "MODE", "validated") + monkeypatch.setattr(policy, "environment_matches", lambda: True) + monkeypatch.setattr(policy, "_profiles", {}) + def reject_native(*args): + raise AssertionError("unsupported dense input reached ARK") + + monkeypatch.setattr(torch.ops.freetoken_ark, "dense_decode", reject_native) + compare(adapter.block_fp8_linear(x, w, scale), baseline(x, w, scale)) + + +def test_moe_routes_zero_inactive_and_input_unchanged(): + from freetoken.kernel.ark.moe_decode_fp8.adapter import ( + fused_experts_decode_fp8_block as candidate, + ) + from freetoken.moe.fused_fp8_block import fused_experts_decode_fp8_block as baseline + + args = moe_inputs() + original = [v.clone() for v in args] + expected = baseline(*args) + actual = candidate(*args) + compare(actual, expected) + assert actual.data_ptr() != args[0].data_ptr() + for a, b in zip(args, original): + assert torch.equal(a, b) + args[-1][0, 2] = -1 + compare(candidate(*args), expected) + args[-2].zero_() + assert candidate(*args).count_nonzero().item() == 0 + + +@pytest.mark.parametrize("kind", ["batch", "gelu", "disabled", "graph_excluded"]) +def test_moe_fallback(kind, monkeypatch): + from freetoken.kernel.ark.common import policy + from freetoken.kernel.ark.moe_decode_fp8 import adapter + from freetoken.moe.fused_fp8_block import fused_experts_decode_fp8_block as baseline + + args = moe_inputs(2 if kind == "batch" else 1) + kwargs = {"activation": "gelu"} if kind == "gelu" else {} + if kind == "disabled": + monkeypatch.setenv("FREETOKEN_USE_ARK", "0") + if kind == "graph_excluded": + monkeypatch.setattr(policy, "MODE", "validated") + monkeypatch.setattr(policy, "environment_matches", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + def reject_native(*args): + raise AssertionError("unsupported MoE input reached ARK") + + monkeypatch.setattr(torch.ops.freetoken_ark, "moe_decode", reject_native) + compare(adapter.fused_experts_decode_fp8_block(*args, **kwargs), baseline(*args, **kwargs)) + + +@pytest.mark.parametrize("op", ["dense", "moe"]) +def test_nondefault_stream_graph_and_updated_inputs(op): + if op == "dense": + 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, + ) + + w, scale = quant(rand(384, 256)) + args = (rand(1, 256), w, scale) + else: + from freetoken.kernel.ark.moe_decode_fp8.adapter import ( + fused_experts_decode_fp8_block as candidate, + ) + from freetoken.moe.fused_fp8_block import ( + fused_experts_decode_fp8_block as baseline, + ) + + args = moe_inputs() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + candidate(*args) + baseline(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = candidate(*args) + for step in range(3): + args[0].copy_(rand(*args[0].shape)) + if op == "moe": + args[-1].copy_(torch.tensor([[3, 1, 1, 0]], device="cuda", dtype=torch.int32)) + args[-2].copy_(torch.tensor([[0.0, 0.3, 0.3, 0.4]], device="cuda")) + expected = baseline(*args).clone() + graph.replay() + compare(actual, expected) + stream.synchronize() + + +@pytest.mark.parametrize("resident", [False, True]) +def test_real_method_bank_pack_apply_prefill(resident, monkeypatch): + from freetoken.kernel.ark.kernels import ArkFp8LinearKernel, ArkFp8MoEKernel + from freetoken.layers.quantization.linear.base import LinearConfig + from freetoken.layers.quantization.linear.fp8_block import ( + Fp8BlockLinearMethod, + TritonFp8BlockLinearKernel, + ) + from freetoken.layers.quantization.moe.base import ExpertView, MoEConfig + from freetoken.layers.quantization.moe.fp8_block import ( + Fp8BlockMoEMethod, + TritonFp8BlockMoEKernel, + ) + from freetoken.layers.quantization.scheme import fp8_block_scheme + + scheme = fp8_block_scheme("bf16") + cfg = MoEConfig(4, 256, 128, 4, scheme=scheme) + candidate, baseline = ArkFp8MoEKernel(), TritonFp8BlockMoEKernel() + assert candidate.layout(cfg) == baseline.layout(cfg) + assert candidate.cpu_format == baseline.cpu_format and not candidate.supports_inactive_slots + x, gu, gs, dn, ds, weights, ids = moe_inputs() + pieces = {"gate_up": gu, "gate_up_scale": gs[:, :, :2], "down": dn, "down_scale": ds[:, :, :1]} + specs = candidate.layout(cfg) + bank = {k: torch.zeros(4, *v.shape, device="cuda", dtype=v.dtype) for k, v in specs.items()} + ref_bank = {k: torch.zeros_like(v) for k, v in bank.items()} + assert candidate.pack(pieces, cfg, bank) == baseline.pack(pieces, cfg, ref_bank) + for name, value in bank.items(): + assert torch.equal(value, ref_bank[name]) + view = ExpertView(bank, n=4) + layer = SimpleNamespace( + activation="silu", + alpha=1.0, + limit=None, + num_experts=4, + apply_router_weight_on_input=False, + quant_method=SimpleNamespace(cfg=SimpleNamespace(strategy="resident" if resident else "offload")), + ) + for prefill in (False, True): + m = 33 if prefill else 1 + x = rand(m, 256) + initial = x.clone() + tw = weights.expand(m, 4).contiguous() + idx = ids.expand(m, 4).contiguous() + compare(candidate.apply(layer, x, tw, idx, view, is_prefill=prefill), baseline.apply(layer, x, tw, idx, view, is_prefill=prefill)) + assert torch.equal(x, initial) + assert isinstance(Fp8BlockMoEMethod(cfg).kernel, ArkFp8MoEKernel) + linear_cfg = LinearConfig(256, 384, scheme=scheme) + assert isinstance(Fp8BlockLinearMethod(linear_cfg).kernel, ArkFp8LinearKernel) + assert type(Fp8BlockMoEMethod(cfg, requested="triton").kernel) is TritonFp8BlockMoEKernel + monkeypatch.setenv("FREETOKEN_USE_ARK", "0") + assert type(Fp8BlockMoEMethod(cfg).kernel) is TritonFp8BlockMoEKernel + assert type(Fp8BlockLinearMethod(linear_cfg).kernel) is TritonFp8BlockLinearKernel + + +def test_ftw_identity_and_auto_selection(monkeypatch): + from freetoken.checkpoint import ftw + from freetoken.engine.engine import _adjust_ftw_quant_backend + from freetoken.layers.quantization import LayerKind, QuantBackend, QuantKind + from freetoken.moe.legacy_format import ( + bank_kernel_compatible, + kind_kernel_for, + legacy_format_for, + ) + + monkeypatch.setattr(ftw, "ftw_quant_format", lambda _: "fp8_block") + kind = QuantKind.FP8_BLOCK + auto = QuantBackend() + assert _adjust_ftw_quant_backend("/unused", auto).select(LayerKind.MOE, kind) == "auto" + assert legacy_format_for(kind, "ark_r9700") == "fp8_block" + assert kind_kernel_for("fp8_block") == (kind, "triton") + assert bank_kernel_compatible(kind, "triton", "ark_r9700") + assert not bank_kernel_compatible(QuantKind.NONE, "fused", "ark_r9700") + assert not bank_kernel_compatible(kind, "unknown", "ark_r9700") + monkeypatch.setenv("FREETOKEN_USE_ARK", "0") + assert _adjust_ftw_quant_backend("/unused", auto).select(LayerKind.MOE, kind) == "triton" + + +def test_policy_mismatch_default_and_no_builder_import(monkeypatch): + from freetoken.kernel.ark.common import policy + + monkeypatch.setattr(policy, "MODE", "validated") + monkeypatch.setattr(policy, "_profiles", {"_environment": {"torch": "mismatch"}}) + policy.environment_matches.cache_clear() + assert not policy.environment_matches() + assert not policy.allowed("dense_fp8_decode", "1,5120,2048") + policy.environment_matches.cache_clear() + env = dict(os.environ, FREETOKEN_USE_ARK="0") + script = "import sys; import freetoken.kernel.ark as a; assert not a.enabled(); assert 'freetoken.kernel.ark.common.native' not in sys.modules" + subprocess.run([sys.executable, "-c", script], env=env, check=True) + + +def test_finite_fp8_codes(): + from freetoken.kernel.ark.dense_fp8_decode.adapter import block_fp8_linear + + codes = torch.arange(256, device="cuda", dtype=torch.int32).to(torch.uint8) + codes[127] = codes[255] = 0 + storage = torch.zeros(256, 128, device="cuda", dtype=torch.uint8) + storage[:, 0] = codes + x = torch.zeros(1, 128, device="cuda", dtype=torch.bfloat16) + x[0, 0] = 1 + actual = block_fp8_linear(x, storage.view(torch.float8_e4m3fn), torch.ones(2, 1, device="cuda", dtype=torch.bfloat16)) + assert torch.equal(actual, codes.view(torch.float8_e4m3fn).bfloat16().reshape(1, 256)) + + +def test_packaged_policy_keeps_graph_fallback(monkeypatch): + from freetoken.kernel.ark.common import policy + + monkeypatch.setattr(policy, "MODE", "validated") + monkeypatch.setattr(policy, "environment_matches", lambda: True) + for op in ("dense_fp8_decode", "moe_decode_fp8"): + assert not policy._profiles[op]["graph"] + assert policy._profiles[op]["eager"] + for shape in policy._profiles[op]["eager"]: + assert policy.allowed(op, shape, capturing=False) + assert not policy.allowed(op, shape, capturing=True)