From 6de2d8e605dfc3c6cf0d7370b590b6f29d50109b Mon Sep 17 00:00:00 2001 From: zihaomu Date: Wed, 21 Oct 2026 13:24:06 +0800 Subject: [PATCH 1/2] perf(kernel): add opt-in hip fp8 decode kernels for r9700 --- benchmarks/ark/__init__.py | 0 benchmarks/ark/audit.py | 77 +++++ benchmarks/ark/audit_quality.py | 41 +++ benchmarks/ark/calibrate.py | 73 ++++ benchmarks/ark/cases.py | 71 ++++ benchmarks/ark/http_client.py | 132 ++++++++ benchmarks/ark/model_ab.py | 259 ++++++++++++++ benchmarks/ark/probe.py | 87 +++++ benchmarks/ark/quality_prompts.py | 44 +++ benchmarks/ark/timing.py | 99 ++++++ docs/ark-fp8-decode.md | 200 +++++++++++ .../ark-fp8-decode-20261021/calibration.json | 320 ++++++++++++++++++ .../graph-candidate-profile.json | 61 ++++ .../ark-fp8-decode-20261021/model-samples.csv | 25 ++ .../ark-fp8-decode-20261021/runtime.json | 48 +++ .../ark-fp8-decode-20261021/validation.json | 115 +++++++ pyproject.toml | 1 + python/freetoken/engine/engine.py | 21 +- python/freetoken/kernel/ark/NOTICE | 5 + python/freetoken/kernel/ark/__init__.py | 37 ++ .../freetoken/kernel/ark/common/__init__.py | 0 .../freetoken/kernel/ark/common/bindings.cpp | 75 ++++ .../freetoken/kernel/ark/common/contracts.py | 63 ++++ .../ark/common/include/ft_decode_impl.h | 79 +++++ python/freetoken/kernel/ark/common/native.py | 122 +++++++ python/freetoken/kernel/ark/common/policy.py | 77 +++++ python/freetoken/kernel/ark/common/trace.py | 43 +++ .../kernel/ark/common/validated_shapes.json | 56 +++ .../kernel/ark/dense_fp8_decode/__init__.py | 0 .../kernel/ark/dense_fp8_decode/adapter.py | 41 +++ .../ark/dense_fp8_decode/dense_gemv_fp8.hip | 139 ++++++++ python/freetoken/kernel/ark/kernels.py | 40 +++ .../kernel/ark/moe_decode_fp8/__init__.py | 0 .../kernel/ark/moe_decode_fp8/adapter.py | 57 ++++ .../moe_decode_fp8/moe_decode_gemv_fp8.hip | 6 + .../freetoken/layers/quantization/method.py | 10 +- python/freetoken/moe/legacy_format.py | 6 + tests/kernels/test_ark_fp8.py | 264 +++++++++++++++ 38 files changed, 2792 insertions(+), 2 deletions(-) create mode 100644 benchmarks/ark/__init__.py create mode 100644 benchmarks/ark/audit.py create mode 100644 benchmarks/ark/audit_quality.py create mode 100644 benchmarks/ark/calibrate.py create mode 100644 benchmarks/ark/cases.py create mode 100644 benchmarks/ark/http_client.py create mode 100644 benchmarks/ark/model_ab.py create mode 100644 benchmarks/ark/probe.py create mode 100644 benchmarks/ark/quality_prompts.py create mode 100644 benchmarks/ark/timing.py create mode 100644 docs/ark-fp8-decode.md create mode 100644 docs/benchmarks/ark-fp8-decode-20261021/calibration.json create mode 100644 docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json create mode 100644 docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv create mode 100644 docs/benchmarks/ark-fp8-decode-20261021/runtime.json create mode 100644 docs/benchmarks/ark-fp8-decode-20261021/validation.json create mode 100644 python/freetoken/kernel/ark/NOTICE create mode 100644 python/freetoken/kernel/ark/__init__.py create mode 100644 python/freetoken/kernel/ark/common/__init__.py create mode 100644 python/freetoken/kernel/ark/common/bindings.cpp create mode 100644 python/freetoken/kernel/ark/common/contracts.py create mode 100644 python/freetoken/kernel/ark/common/include/ft_decode_impl.h create mode 100644 python/freetoken/kernel/ark/common/native.py create mode 100644 python/freetoken/kernel/ark/common/policy.py create mode 100644 python/freetoken/kernel/ark/common/trace.py create mode 100644 python/freetoken/kernel/ark/common/validated_shapes.json create mode 100644 python/freetoken/kernel/ark/dense_fp8_decode/__init__.py create mode 100644 python/freetoken/kernel/ark/dense_fp8_decode/adapter.py create mode 100644 python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip create mode 100644 python/freetoken/kernel/ark/kernels.py create mode 100644 python/freetoken/kernel/ark/moe_decode_fp8/__init__.py create mode 100644 python/freetoken/kernel/ark/moe_decode_fp8/adapter.py create mode 100644 python/freetoken/kernel/ark/moe_decode_fp8/moe_decode_gemv_fp8.hip create mode 100644 tests/kernels/test_ark_fp8.py diff --git a/benchmarks/ark/__init__.py b/benchmarks/ark/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/benchmarks/ark/audit.py b/benchmarks/ark/audit.py new file mode 100644 index 0000000000..ad7ec6ad52 --- /dev/null +++ b/benchmarks/ark/audit.py @@ -0,0 +1,77 @@ +"""Independently recompute formal SSE timing and paired summaries from raw files.""" + +import argparse +import hashlib +import json +from pathlib import Path +import statistics + + +def audit(root): + rows = [] + for path in sorted(root.glob("round-*/measured-*.json")): + row = json.loads(path.read_text()) + assert row["success"] and row["done"] + assert row["usage"]["prompt_tokens"] == row["context"] + assert row["usage"]["completion_tokens"] == 256 + assert row["request"]["prompt"] == (root / f"prompt-{row['context']}.txt").read_text() + assert row["request"]["temperature"] == 0 and row["request"]["ignore_eos"] + assert (row["usage"].get("prompt_tokens_details") or {}).get("cached_tokens", 0) == 0 + events = [x for x in row["events"] if any(c.get("text") for c in x["data"].get("choices", []))] + span = events[-1]["elapsed_s"] - events[0]["elapsed_s"] + assert abs(255 / span - row["decode_tok_s"]) < 1e-9 + assert events[0]["elapsed_s"] == row["ttft_s"] + text = "".join(c.get("text", "") for x in events for c in x["data"]["choices"]) + assert hashlib.sha256(text.encode()).hexdigest() == row["text_sha256"] + command = json.loads((path.parent / "command.json").read_text()) + assert command["env"]["FREETOKEN_USE_ARK"] == str(row["enabled"]) + assert command["env"]["FREETOKEN_ARK_POLICY"] == "validated" + assert "FREETOKEN_ARK_TRACE" not in command["env"] + assert "benchmarks.ark.probe" not in command["argv"] + assert json.loads((path.parent / "complete.json").read_text())["success"] + assert json.loads((path.parent / "shutdown.json").read_text())["vram"] < 256 * 2**20 + rows.append({k: v for k, v in row.items() if k not in ("events", "text", "request")}) + paired = [] + for context in (512, 32768): + pairs = [] + for rid in sorted({r["round"] for r in rows}): + pair = {r["enabled"]: r for r in rows if r["round"] == rid and r["context"] == context} + assert set(pair) == {0, 1} + baseline, candidate = pair[0], pair[1] + pairs.append( + dict( + round=rid, + speedup=candidate["decode_tok_s"] / baseline["decode_tok_s"], + ft_decode_tok_s=baseline["decode_tok_s"], + ark_decode_tok_s=candidate["decode_tok_s"], + outputs_equal=baseline["text_sha256"] == candidate["text_sha256"], + ) + ) + if pairs: + ft = statistics.median(r["ft_decode_tok_s"] for r in pairs) + ark = statistics.median(r["ark_decode_tok_s"] for r in pairs) + paired.append( + dict( + context=context, + rounds=pairs, + ft_median=ft, + ark_median=ark, + median_change_percent=100 * (ark / ft - 1), + wins=sum(r["speedup"] > 1 for r in pairs), + n=len(pairs), + ) + ) + return dict(passed=True, samples=rows, paired=paired) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("root", type=Path) + a = p.parse_args() + result = audit(a.root) + (a.root / "audit.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result["paired"], indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ark/audit_quality.py b/benchmarks/ark/audit_quality.py new file mode 100644 index 0000000000..74ad49d8c5 --- /dev/null +++ b/benchmarks/ark/audit_quality.py @@ -0,0 +1,41 @@ +"""Score the first answer in saved non-thinking generation checks.""" + +import argparse +import json +from pathlib import Path + +from .quality_prompts import CASES, first_answer + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("root", type=Path) + a = p.parse_args() + groups = {} + for enabled in (0, 1): + folder = a.root / f"round-1-ark{enabled}" + rows = [] + for name, _, predicate in CASES: + paths = list(folder.glob(f"quality-{name}-*.json")) + assert len(paths) == 1, (name, paths) + raw = json.loads(paths[0].read_text()) + assert raw["success"] and raw["usage"]["completion_tokens"] == 128 + assert raw["request"]["prompt"].endswith("\n\n\n\n") + answer = first_answer(raw["text"]) + rows.append(dict(name=name, passed=bool(predicate(answer)), answer=answer)) + groups[str(enabled)] = rows + (folder / "quality-audited.json").write_text(json.dumps(rows, indent=2) + "\n") + old = {r["name"] for r in groups["0"] if not r["passed"]} + new = {r["name"] for r in groups["1"] if not r["passed"]} + result = dict( + groups=groups, + new_failures=sorted(new - old), + note="8 bounded functional checks, first answer only, non-thinking template, not a standardized accuracy benchmark", + ) + (a.root / "quality-audit.json").write_text(json.dumps(result, indent=2) + "\n") + print("baseline", len(groups["0"]) - len(old), "ARK", len(groups["1"]) - len(new), "new failures", sorted(new - old)) + assert not new - old, result + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ark/calibrate.py b/benchmarks/ark/calibrate.py new file mode 100644 index 0000000000..74f2ddb2f2 --- /dev/null +++ b/benchmarks/ark/calibrate.py @@ -0,0 +1,73 @@ +"""Recalibrate the exact source/environment; retain raw paired repeats and fallbacks.""" + +import argparse +import json +import os +from pathlib import Path + +import torch + +from .cases import cases +from .timing import metrics, assert_close, timing_pair, graph_pair +from freetoken.kernel.ark import initialize +from freetoken.kernel.ark.common import policy + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--write-profile", type=Path) + p.add_argument( + "--graph-candidates", action="store_true", help="emit dense Graph microbenchmark candidates for separate model evaluation" + ) + 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, "graph_min_speedup": 1.03, "each_repeat_faster": True}, + "_quality_exclusions": {"moe_decode_fp8": "All Graph shapes remain FT; Qwen3 matched-prefix KL previously exceeded 0.05 nats."}, + } + profile["_rules"]["graph_requires_model_benefit"] = True + profile["_model_exclusions"] = { + "dense_fp8_decode.graph": "Current Qwen3 candidate model A/B did not establish a Graph benefit; default Graph remains FT." + } + 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) + for mode, threshold in (("eager", 1.05), ("graph", 1.03)): + timing = row[mode] + paired = [x / y for x, y in zip(timing["samples_us"]["baseline"], timing["samples_us"]["candidate"])] + eligible = timing["speedup"] >= threshold and min(paired) > 1 + if mode == "graph" and (op == "moe_decode_fp8" or not a.graph_candidates): + eligible = False + row[mode]["enabled"] = eligible + if eligible: + profile[op][mode][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 0000000000..f49b7c1ef2 --- /dev/null +++ b/benchmarks/ark/cases.py @@ -0,0 +1,71 @@ +"""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.triton.fp8_block_linear import block_fp8_linear as baseline + from freetoken.kernel.ark.dense_fp8_decode.adapter import block_fp8_linear as candidate + + 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), dict(M=1, N=n, K=k)) + + +def moe_cases(): + 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 + from freetoken.kernel.ark.moe_decode_fp8.adapter import fused_experts_decode_fp8_block as candidate + + 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), dict(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/http_client.py b/benchmarks/ark/http_client.py new file mode 100644 index 0000000000..81bbacf5bb --- /dev/null +++ b/benchmarks/ark/http_client.py @@ -0,0 +1,132 @@ +"""SSE client and resource sampling, using the same timing definitions as prior reports.""" + +import hashlib, json, time, traceback, urllib.error, urllib.request +from pathlib import Path + +ROOT = Path("/artifacts") +ORIGIN = "http://127.0.0.1:19673" +DEVICE = None +MODEL = "/models/Qwen3-30B-A3B-FP8-FTW" + + +def save(path, data): + temp = path.with_suffix(path.suffix + ".tmp") + temp.write_text(json.dumps(data, indent=2) + "\n") + temp.replace(path) + + +def get(route, timeout=3): + with urllib.request.urlopen(ORIGIN + route, timeout=timeout) as response: + return json.load(response) + + +def tree(pid): + result, pending = set(), [pid] + while pending: + p = pending.pop() + if p in result: + continue + result.add(p) + try: + pending.extend(int(x) for x in Path(f"/proc/{p}/task/{p}/children").read_text().split()) + except OSError: + pass + return result + + +def snapshot(pid): + data = {"time": time.time(), "rss_sum_bytes": 0, "processes": []} + for p in tree(pid): + try: + status = dict(line.split(":", 1) for line in Path(f"/proc/{p}/status").read_text().splitlines() if ":" in line) + rss = int(status.get("VmRSS", "0 kB").split()[0]) * 1024 + data["rss_sum_bytes"] += rss + data["processes"].append( + { + "pid": p, + "name": status["Name"].strip(), + "rss_bytes": rss, + "cpus": status["Cpus_allowed_list"].strip(), + "mems": status["Mems_allowed_list"].strip(), + } + ) + except (OSError, KeyError): + pass + for key in ["mem_info_vram_used", "gpu_busy_percent"]: + try: + data[key] = int((DEVICE / key).read_text()) + except (OSError, ValueError): + data[key] = None + return data + + +def monitor(pid, stop, path): + with path.open("w") as output: + while not stop.is_set(): + output.write(json.dumps(snapshot(pid)) + "\n") + output.flush() + stop.wait(0.25) + + +def request(directory, context, phase, output_tokens=256, prompt=None): + prompt = prompt if prompt is not None else (ROOT / f"prompt-{context}.txt").read_text() + payload = { + "model": MODEL, + "prompt": prompt, + "max_tokens": output_tokens, + "temperature": 0, + "ignore_eos": True, + "stream": True, + "stream_options": {"include_usage": True}, + } + req = urllib.request.Request( + ORIGIN + "/v1/completions", data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"} + ) + row = {"context": context, "phase": phase, "expected_output_tokens": output_tokens, "request": payload, "events": []} + started = time.perf_counter() + try: + with urllib.request.urlopen(req, timeout=180) as response: + for line in response: + now = time.perf_counter() + if now - started > 180: + raise TimeoutError("request exceeded 180s") + if not line.startswith(b"data: "): + continue + value = line[6:].strip() + if value == b"[DONE]": + row["done"] = True + break + event = json.loads(value) + row["events"].append({"elapsed_s": now - started, "data": event}) + if "error" in event: + raise RuntimeError(event["error"]) + row["e2e_s"] = time.perf_counter() - started + row["usage"] = next(x["data"]["usage"] for x in reversed(row["events"]) if "usage" in x["data"]) + chunks = [x for x in row["events"] if any(c.get("text") for c in x["data"].get("choices", []))] + row["text"] = "".join(c.get("text", "") for x in chunks for c in x["data"]["choices"]) + row["text_sha256"] = hashlib.sha256(row["text"].encode()).hexdigest() + row["token_text_chunks"] = len(chunks) + row["ttft_s"] = chunks[0]["elapsed_s"] + row["decode_span_s"] = chunks[-1]["elapsed_s"] - chunks[0]["elapsed_s"] + row["decode_tok_s"] = (output_tokens - 1) / row["decode_span_s"] + # The official completion streamer emits at most one nonempty event per generated token. + row["decode_timing_kind"] = ( + "one-event-per-token verified" + if len(chunks) == row["usage"]["completion_tokens"] + else "client legacy estimate; empty-token events omitted" + ) + row["e2e_output_tok_s"] = row["usage"]["completion_tokens"] / row["e2e_s"] + row["prefill_effective_tok_s"] = context / row["ttft_s"] + assert row.get("done"), "missing SSE DONE" + assert row["usage"]["prompt_tokens"] == context, row["usage"] + assert row["usage"]["completion_tokens"] == output_tokens, row["usage"] + details = row["usage"].get("prompt_tokens_details") or {} + assert details.get("cached_tokens", 0) == 0, row["usage"] + assert row["text"].strip(), "empty output" + row["success"] = True + except Exception: + row.update(success=False, error=traceback.format_exc(), elapsed_s=time.perf_counter() - started) + raise + finally: + save(directory / f"{phase}-{context}.json", row) + return row diff --git a/benchmarks/ark/model_ab.py b/benchmarks/ark/model_ab.py new file mode 100644 index 0000000000..b5cc7c084e --- /dev/null +++ b/benchmarks/ark/model_ab.py @@ -0,0 +1,259 @@ +"""Fresh-process off/on Qwen3 A/B; instrumentation runs separately via --probe.""" + +import argparse +from contextlib import contextmanager +import hashlib +import json +import os +from pathlib import Path +import signal +import statistics +import subprocess +import sys +import threading +import time + +from . import http_client as client + + +@contextmanager +def service(args, directory, enabled): + directory.mkdir(parents=True, exist_ok=False) + if int((client.DEVICE / "mem_info_vram_used").read_text()) > 256 * 2**20: + raise RuntimeError("Target GPU occupied; no service launched") + module = "benchmarks.ark.probe" if args.probe else "freetoken.cli" + argv = [ + "numactl", + "--interleave=" + args.memory_nodes, + sys.executable, + "-m", + module, + "serve", + "--model-path", + args.model_path, + "--dtype", + "bfloat16", + "--attention-backend", + "triton", + "--moe-strategy", + "offload", + "--moe-cpu-threads", + "17", + "--moe-cache-mem", + "20GiB", + "--moe-cache-scope", + "unified", + "--moe-cache-split", + "equal", + "--num-tokens", + "33058", + "--page-size", + "1", + "--max-seq-len-override", + "33024", + "--max-prefill-length", + "8192", + "--max-running-requests", + "1", + "--cuda-graph-max-bs", + str(args.graph), + "--max-output-tokens", + "256", + "--memory-ratio", + "0.9", + "--expert-load", + "serial", + "--cache-type", + "naive", + "--enable-cache-report", + "--sampling-defaults", + "none", + "--reasoning-parser", + "off", + "--host", + "127.0.0.1", + "--port", + "19673", + ] + env = dict( + os.environ, + FREETOKEN_USE_ARK=str(enabled), + FREETOKEN_ARK_POLICY="validated", + OMP_NUM_THREADS="1", + MKL_NUM_THREADS="1", + OPENBLAS_NUM_THREADS="1", + PYTHONUNBUFFERED="1", + ) + for key in ("PYTHONPATH", "FREETOKEN_ARK_TRACE", "FREETOKEN_ARK_PROFILE"): + env.pop(key, None) + if args.profile: + env["FREETOKEN_ARK_PROFILE"] = str(args.profile.resolve()) + if args.probe: + env["FREETOKEN_ARK_TRACE"] = str(directory / "dispatch-{pid}.json") + env["ARK_MODEL_PROBE"] = str(directory) + client.save(directory / "command.json", dict(argv=argv, env={k: v for k, v in env.items() if k.startswith(("FREETOKEN_", "ROCR_"))})) + stop = threading.Event() + with (directory / "server.log").open("w") as out: + proc = subprocess.Popen(argv, env=env, stdout=out, stderr=subprocess.STDOUT, start_new_session=True) + monitor = threading.Thread(target=client.monitor, args=(proc.pid, stop, directory / "resources.jsonl"), daemon=True) + monitor.start() + start = time.monotonic() + try: + while True: + if proc.poll() is not None: + raise RuntimeError(f"server exited {proc.returncode}: {directory}") + if time.monotonic() - start > 600: + raise TimeoutError("startup timeout") + try: + health = client.get("/health") + except OSError: + time.sleep(1) + continue + if health.get("status") == "error": + raise RuntimeError(health) + if health.get("status") == "ok": + break + time.sleep(1) + print("READY", directory.name, flush=True) + yield + client.save(directory / "complete.json", dict(success=True, stats=client.get("/v1/stats"))) + finally: + owned = client.tree(proc.pid) + if proc.poll() is None: + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(30) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait(10) + for pid in owned - {proc.pid}: + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + stop.set() + monitor.join(5) + deadline = time.monotonic() + 180 + while int((client.DEVICE / "mem_info_vram_used").read_text()) > 256 * 2**20: + if time.monotonic() > deadline: + raise RuntimeError("GPU did not return to idle") + time.sleep(0.5) + client.save( + directory / "shutdown.json", dict(returncode=proc.returncode, vram=int((client.DEVICE / "mem_info_vram_used").read_text())) + ) + + +def make_prompts(root, model): + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(model) + head = "<|im_start|>user\nRead the following archive.\n" + tail = ( + "\nExplain how inspectors compare transport, energy and weather records. Give a detailed answer.<|im_end|>\n<|im_start|>assistant\n" + ) + body = ( + "The archive tracks transport, energy, weather and equipment inspections. Inspectors compare observations from different stations.\n" + * 2000 + ) + head_ids, tail_ids = (tok.encode(x, add_special_tokens=False) for x in (head, tail)) + body_ids = tok.encode(body, add_special_tokens=False) + for n in (512, 32768): + path = root / f"prompt-{n}.txt" + if not path.exists(): + budget = n - len(head_ids) - len(tail_ids) + for _ in range(32): + text = tok.decode(head_ids + body_ids[:budget] + tail_ids) + count = len(tok.encode(text, add_special_tokens=False)) + if count == n: + break + budget += n - count + else: + raise ValueError("prompt token roundtrip failed") + path.write_text(text) + assert len(tok.encode(path.read_text(), add_special_tokens=False)) == n + client.save(root / "prompts.json", {str(n): hashlib.sha256((root / f"prompt-{n}.txt").read_bytes()).hexdigest() for n in (512, 32768)}) + + +def summarize(root): + samples = [] + for path in sorted(root.glob("round-*/measured-*.json")): + row = json.loads(path.read_text()) + assert row["success"] + samples.append({k: v for k, v in row.items() if k not in ("events", "text", "request")}) + groups = [] + for n in (512, 32768): + for enabled in (0, 1): + rows = [r for r in samples if r["context"] == n and r["enabled"] == enabled] + if rows: + groups.append( + dict( + context=n, + enabled=enabled, + n=len(rows), + decode_tok_s=statistics.median(r["decode_tok_s"] for r in rows), + ttft_s=statistics.median(r["ttft_s"] for r in rows), + e2e_s=statistics.median(r["e2e_s"] for r in rows), + ) + ) + client.save(root / "summary.json", dict(samples=samples, groups=groups)) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model-path", required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--gpu-pci", required=True) + p.add_argument("--memory-nodes", required=True) + p.add_argument("--rounds", type=int, default=3) + p.add_argument("--start-round", type=int, default=1) + p.add_argument("--graph", type=int, default=0) + p.add_argument("--probe", action="store_true") + p.add_argument("--quality-only", action="store_true", help="run corrected short generation checks without context probes") + p.add_argument("--profile", type=Path, help="explicit candidate profile for separate experiments") + a = p.parse_args() + if a.quality_only and not a.probe: + p.error("--quality-only requires --probe") + a.output = a.output.resolve() + a.output.mkdir(parents=True, exist_ok=True) + client.ROOT = a.output + client.MODEL = a.model_path + client.DEVICE = Path("/sys/bus/pci/devices") / a.gpu_pci + make_prompts(a.output, a.model_path) + for r in range(a.start_round, a.start_round + a.rounds): + for enabled in (0, 1) if r % 2 else (1, 0): + directory = a.output / f"round-{r}-ark{enabled}" + with service(a, directory, enabled): + for n in () if a.quality_only else ((512, 32768) if r % 2 else (32768, 512)): + if not a.probe: + client.request(directory, n, "warmup") + row = client.request(directory, n, "probe" if a.probe else "measured", output_tokens=8 if a.probe else 256) + row.update(round=r, enabled=enabled) + client.save(directory / f"{row['phase']}-{n}.json", row) + if a.probe: + counts = [row for path in directory.glob("dispatch-*.json") for row in json.loads(path.read_text())["counts"]] + client.save(directory / f"hits-{n}.json", counts) + if not a.probe: + summarize(a.output) + print("RESULT", directory.name, n, row["decode_tok_s"], row["ttft_s"], flush=True) + if a.probe: + from transformers import AutoTokenizer + from .quality_prompts import CASES, first_answer + + tokenizer = AutoTokenizer.from_pretrained(a.model_path) + quality = [] + for name, question, predicate in CASES: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + n = len(tokenizer.encode(prompt, add_special_tokens=False)) + result = client.request(directory, n, "quality-" + name, output_tokens=128, prompt=prompt) + answer = first_answer(result["text"]) + quality.append(dict(name=name, answer=answer, passed=bool(predicate(answer)))) + client.save(directory / "quality.json", quality) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ark/probe.py b/benchmarks/ark/probe.py new file mode 100644 index 0000000000..f734efa5c9 --- /dev/null +++ b/benchmarks/ark/probe.py @@ -0,0 +1,87 @@ +"""Untimed diagnostic server: real dispatch, GPU symbols and same-input shadows.""" + +import json +import os +from pathlib import Path + +import torch + +from freetoken.engine.engine import Engine +from freetoken.kernel.ark.common.trace import flush +from freetoken.kernel.ark.kernels import ArkFp8LinearKernel, ArkFp8MoEKernel +from freetoken.layers.quantization.linear.fp8_block import TritonFp8BlockLinearKernel +from freetoken.layers.quantization.moe.fp8_block import TritonFp8BlockMoEKernel +from .timing import metrics, assert_close + +ROOT = Path(os.environ["ARK_MODEL_PROBE"]) +SEEN = set() +ROWS = [] +ACTIVE = False +PROFILED = False + + +def shadow(op, key, actual, expected): + error = metrics(actual, expected) + assert_close(error, op) + ROWS.append(dict(operator=op, key=str(key), **error)) + (ROOT / "shadows.json").write_text(json.dumps(ROWS, indent=2) + "\n") + + +_linear = ArkFp8LinearKernel.apply + + +def linear(self, layer, x): + result = _linear(self, layer, x) + key = ("dense_fp8_decode", id(layer)) + if ACTIVE and x.numel() == x.shape[-1] and key not in SEEN: + SEEN.add(key) + shadow(key[0], (tuple(x.shape), tuple(layer.weight.shape)), result, TritonFp8BlockLinearKernel().apply(layer, x)) + return result + + +_moe = ArkFp8MoEKernel.apply + + +def moe(self, layer, x, weights, ids, view, *, is_prefill): + result = _moe(self, layer, x, weights, ids, view, is_prefill=is_prefill) + key = ("moe_decode_fp8", layer.layer_id) + if ACTIVE and not is_prefill and x.shape[0] == 1 and key not in SEEN: + SEEN.add(key) + shadow( + key[0], + (layer.layer_id, tuple(x.shape), view.tensors["gate_up"].shape[0]), + result, + TritonFp8BlockMoEKernel().apply(layer, x, weights, ids, view, is_prefill=False), + ) + return result + + +ArkFp8LinearKernel.apply = linear +ArkFp8MoEKernel.apply = moe +_forward = Engine.forward_batch + + +def forward(self, batch, args): + global ACTIVE, PROFILED + ACTIVE = True + try: + if not batch.is_prefill and not PROFILED: + PROFILED = True + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]) as trace: + result = _forward(self, batch, args) + torch.cuda.synchronize() + trace.export_chrome_trace(str(ROOT / "decode-trace.json")) + else: + result = _forward(self, batch, args) + flush() + return result + finally: + ACTIVE = False + + +Engine.forward_batch = forward + +if __name__ == "__main__": + from freetoken.cli import main + + main() diff --git a/benchmarks/ark/quality_prompts.py b/benchmarks/ark/quality_prompts.py new file mode 100644 index 0000000000..097a48626e --- /dev/null +++ b/benchmarks/ark/quality_prompts.py @@ -0,0 +1,44 @@ +"""Small explicit generation gates; not a standardized accuracy benchmark.""" + +import json +import re + + +def first_answer(text): + for marker in ("<|im_end|>", "<|endoftext|>", "<|im_start|>"): + text = text.split(marker, 1)[0] + if "" in text and "" not in text: + return "" + if "" in text: + text = text.split("", 1)[1] + return text.strip() + + +def json_ok(text): + match = re.search(r"\{[^{}]*\}", text) + if not match: + return False + try: + return json.loads(match.group()) == {"name": "test", "count": 3} + except json.JSONDecodeError: + return False + + +CASES = [ + ( + "moe_cn", + "请用两句话解释什么是混合专家（MoE）模型。", + lambda t: "专家" in t and any(w in t for w in ("路由", "激活", "选择", "门控", "分配", "决定")), + ), + ("addition", "What is 137 + 286? Reply with the number only.", lambda t: "423" in t), + ("division", "What is 144 divided by 12? Reply with the number only.", lambda t: "12" in t), + ( + "python", + "Write a Python function named add that returns the sum of two arguments. Output code only.", + lambda t: "def add" in t and "return" in t and "+" in t, + ), + ("sky", "In one sentence, explain why the sky appears blue.", lambda t: "blue" in t.lower() and "scatter" in t.lower()), + ("json", 'Return only a JSON object with keys name and count, using name "test" and count 3.', json_ok), + ("capital_cn", "中国的首都是哪里？只回答城市名称。", lambda t: "北京" in t), + ("decimal", "Which is larger, 0.9 or 0.11? Answer with just the larger number.", lambda t: "0.9" in t), +] diff --git a/benchmarks/ark/timing.py b/benchmarks/ark/timing.py new file mode 100644 index 0000000000..584f23f143 --- /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 0000000000..7e91f12752 --- /dev/null +++ b/docs/ark-fp8-decode.md @@ -0,0 +1,200 @@ +# Optional ARK FP8 decode kernels on Radeon + +`FREETOKEN_USE_ARK=1` adds native HIP candidates for block-FP8 dense and MoE +**decode** on Linux gfx1201. The default is `0`. Other operations retain FreeToken. + +The packaged profile is specific to R9700, Torch 2.14.0+rocm10.1.0, HIP +7.16.26385, Triton 3.8.0, the compiler, host platform and source fingerprints. +An unmatched environment or unlisted shape uses FT. This is not an all-Radeon +support declaration. Recalibrate on another software stack; do not strip the +environment check. + +## Integration + +`ArkFp8MoEKernel` and `ArkFp8LinearKernel` inherit the existing Triton kernels' +packing, layouts and CPU formats. MoE consumes the current `ExpertView` directly, +including resident, materialized and cache-slot banks. The FTW format remains +`fp8_block`; reverse lookup still returns the original `triton` kernel. CPU expert +execution, hybrid routing, shared experts, cache management and upstream locked +expert fetching are owned by FT. + +- MoE: BF16 activations, E4M3 weights, BF16 padded 128x128 block scales, + contiguous INT32 slots/FP32 routing weights, M=1 and plain SiLU. It is W8A16. + Other activations, router-weight placement, dtypes, shapes and prefill use FT. + Inputs are preserved and the BF16 result is a new tensor. +- Dense: flattened M=1 uses W8A16, preserving leading dimensions and optional + bias. M>1, including batched decode, remains the original W8A8 FT path. + Unsupported layouts, scales and shapes use FT. +- `supports_inactive_slots=False` preserves FT's safe slot-zero/zero-weight + contract on hybrid routes. Native negative slots are covered only by direct + boundary tests; the default integration does not advertise that capability. +- The original primitives are independent fallback targets. No FT function is + overwritten by the production adapters. +- Every MoE Graph shape remains FT. A prior Qwen3 matched-prefix model probe + exceeded its 0.05-nat KL gate with native MoE Graph; operator tolerance alone + was insufficient. Native Graph correctness tests are separate from permission + to enable it in a model. Dense Graph also remains FT: its three-round model A/B did not demonstrate a benefit. + Native Graph correctness tests are retained for future qualification. + +## Build and use + +In an existing ROCm FreeToken environment, use `uv pip install -e .` (or install +the wheel). HIP/C++ sources and the profile are included in the package. +The first enabled use builds the two HIP translation units and Torch bridge in +`$FREETOKEN_ARK_CACHE_DIR`, or `$XDG_CACHE_HOME/freetoken/ark`. Build keys include +native source, build recipe, compiler version and Torch/HIP versions. The exact +ARK Python/native and baseline source fingerprints also gate the profile. +No experiment directory, external `common` package or prebuilt ARK library is +required. The engine loads the library before Graph capture. + +```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 +FREETOKEN_USE_ARK=0 ft serve --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --attention-backend triton --moe-strategy offload --cuda-graph-max-bs 0 +``` + +Expose one R9700 (the profile requires one visible GPU). + +Existing explicit FT backend requests, e.g. `moe.fp8_block=triton`, keep selecting +FT. `FREETOKEN_ARK_TRACE=/tmp/dispatch-{pid}.json` writes operator, shape, backend, +mode and Python call counts. Graph capture counts are not replay counts; inspect +a GPU trace for replay execution. Formal performance runs omit this tracing. + +## Reproduce validation and profile selection + +```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 +python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-eager --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 3 +python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-graph-candidate --profile docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 3 --graph 1 +python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-probe --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 1 --probe +``` + +Use an idle GPU and appropriate CPU/memory affinity for your host. The calibration +command intentionally forces native candidates for measurement. Serving uses +`FREETOKEN_ARK_POLICY=validated` (default); `native` is a calibration/debug override, +not a validated serving policy. `baseline` forces FT fallbacks. An external +`FREETOKEN_ARK_PROFILE` can hold a separately measured table. The generator keeps all Graph entries disabled by default. `--graph-candidates` +can emit dense-only microbenchmark candidates for a separate model evaluation; +it does not qualify them for production. MoE Graph always remains disabled. + +The tests cover numerical error, dtype/shape/alias and input preservation, +repeated/zero/inactive routes, padded scale strides, finite E4M3 codes, bias, +leading dimensions, non-default streams, Graph replay with changed inputs/routes, +actual method selection/packing/Kernel.apply, FTW compatibility, explicit FT +selection, disabled mode and environment mismatch. + +The model harness uses fresh off/on services, alternating off/on order across rounds, +warmup per context, no prefix reuse, 512/32768 input tokens and fixed 256 output +tokens. Its server command, SSE events, usage, per-round metrics and resources +are saved. Decode speed is 255 divided by the span between first and last nonempty +text events; if events differ from tokens, the record labels this a client timing +estimate. TTFT includes request/frontend work. These are concurrency-one client +measurements, not aggregate throughput. + +`--probe` is separate and untimed: it records real GPU kernel symbols, dispatch +counts, same-input per-layer comparisons and eight bounded generation checks. +It is a diagnostic wrapper; production code does not monkeypatch FT methods. +Numerical tolerance and bounded generation checks are not a standardized model +accuracy certification. + +Source provenance: [AMD-AIM/amd_radeon_kernels](https://github.com/AMD-AIM/amd_radeon_kernels) +revision `975e8cd9918ec1365d3f275d1259d215315f505a`; see the packaged ARK `NOTICE`. +Related: [AMD roadmap #541](https://github.com/FlashML-org/FreeToken/issues/541). + +## This calibration (2026-10-21) + +Fresh build against upstream `fe5a5ec9c5644e266527de5afb9f1156b9bdcabb`. +GPU `23:00.0` is R9700 32 GiB, PCIe 32 GT/s x16. CPU affinity selects one +EPYC 9334 socket (`0-31,64-95`) of a dual-socket host, memory interleave `2-3`. +Ubuntu 24.04.1, kernel 7.0.0-38, amdgpu 7.1.9.31600000; remaining environment +fingerprints are in [runtime.json](benchmarks/ark-fp8-decode-20261021/runtime.json). + +Each entry retains five alternating-order paired repeats. Eager timings include +Python dispatch, checks, allocation and the complete GPU operation. Graph timings +include the captured operation, using GPU events over replay; Python capture cost +is excluded. No activation quantization is omitted: these M=1 paths are W8A16. + +| Operator / shape | Eager FT / native, us | Eager speedup | Graph FT / native, us | Deployed Graph selected? | +|---|---:|---:|---:|---| +| dense_fp8_decode / N5120_K2048 | 42.737 / 14.336 | 2.981x | 15.091 / 12.952 | FT (model veto) | +| dense_fp8_decode / N2048_K4096 | 43.586 / 16.359 | 2.664x | 13.675 / 16.654 | FT | +| dense_fp8_decode / N151936_K2048 | 581.642 / 500.559 | 1.162x | 580.391 / 499.153 | FT (model veto) | +| moe_decode_fp8 / C32_H2048_I768 | 105.468 / 45.588 | 2.314x | 53.397 / 44.586 | FT | +| moe_decode_fp8 / C4550_H2048_I768 | 104.505 / 43.965 | 2.377x | 54.333 / 40.295 | FT | + +All five cases passed the numerical gate (NRMSE < 0.03, cosine > 0.999 and a +bounded maximum error). Observed NRMSE was 0 for QKV, output projection and both +MoE bank sizes; the lm-head-shaped case was 0.00000534. The tests include bias and +boundary cases separately. The lm-head shape is an operator benchmark, not a +claim that this model's lm head actually selected native FP8. + +[Raw calibration](benchmarks/ark-fp8-decode-20261021/calibration.json) includes +all repeats and the negative Graph result. Its `enabled` fields are the initial +operator-level candidates; subsequent model evaluation vetoed all Graph entries. +The packaged profile is the deployed decision table. No speedups from the original +experimental report are reused. A different R9700 was used because the prior +measurement GPU was busy; the new off/on model comparison uses this same card. + +## Model results and real-call evidence + +Three-round medians, fixed 256 output tokens, same model/prompt/settings per pair. +Baseline uses the original FT compute path with `FREETOKEN_USE_ARK=0`. +Formal timing excludes instrumentation. The Graph rows evaluate the **rejected +candidate profile**, not the deployed all-FT Graph policy. + +| Mode | Input tokens | FT / ARK tok/s | Change | Paired wins | FT / ARK TTFT, s | +|---|---:|---:|---:|---:|---:| +| eager | 512 | 29.454 / 34.838 | +18.28% | 3/3 | 0.6517 / 0.6866 | +| eager | 32768 | 28.366 / 36.337 | +28.10% | 3/3 | 10.8668 / 10.8626 | +| Graph candidate (rejected) | 512 | 100.180 / 100.099 | -0.08% | 1/3 | 0.7239 / 0.7942 | +| Graph candidate (rejected) | 32768 | 65.221 / 64.977 | -0.37% | 0/3 | 10.9913 / 10.9634 | + +**The recommended enabled scope is eager. Both deployed Graph tables are empty.** +512-token Graph samples vary between about 53 and 101 tok/s; 32K candidate +Graph was slightly slower in all three pairs. No Graph model gain is claimed. +The initial operator-only candidate profile is preserved in +[graph-candidate-profile.json](benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json). +The model harness's `--profile` flag reproduces that candidate without changing +the packaged production table. + +- Final focused and related tests: **41 passed, 42 skipped** (23 ARK cases; + skips mostly require other checkpoint configurations/weights). +- Real model eager trace contains `dense_gemv_fp8_rows<4>`, `gate_up` + and `down_reduce`. In each context's eight-token diagnostic request, + seven decode steps hit native dense **672** times and native MoE **336** times. + Native shapes are dense `(M,N,K)=(1,5120,2048)/(1,2048,4096)` and MoE + `(M,H,I,bank_rows,top_k)=(1,2048,768,4550,8)`. +- **144** real-input per-layer shadows passed, maximum NRMSE **0.000734784**. +- Deployed Graph replay probe: **zero ARK GPU kernels**, and eight-token outputs match + FT exactly for both input contexts. The server's eager warmup still hits native + kernels; these are not Graph capture/replay hits. Full fallback is not counted as ARK gain. +- Corrected bounded generation suite: FT **7/8**, ARK **7/8**, no new failure. + Both incorrectly answer `0.11` for the decimal comparison. Initial 64-token + thinking-mode probes were truncated and excluded; the scored run uses the + checkpoint's non-thinking template, 128 output tokens and the first answer. +- Standalone wheel passed with `python -I`, the repo removed from `sys.path`, + new native/Triton caches and validated policy. All 17 ARK packaged files match + this tree; the enabled dense kernel was hit and NRMSE was **0**. + +[Per-round model samples](benchmarks/ark-fp8-decode-20261021/model-samples.csv) +and [validation evidence](benchmarks/ark-fp8-decode-20261021/validation.json) +retain positive/negative outcomes and exact native symbol names. Complete raw +SSE, resource records, traces and layer/answer comparisons are produced by the +commands above. Validate raw timing with `python -m benchmarks.ark.audit PATH`; +for the generation checks use: + +```bash +python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-quality --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 1 --probe --quality-only +python -m benchmarks.ark.audit_quality /tmp/ark-quality +``` + +The CPU is shared; all slow samples are retained. FT and native generate different +continuations, changing subsequent expert-cache traffic. These are configuration +A/B measurements, not fixed-token kernel attribution or a universal speed guarantee. +Prefill compute is unchanged, so no TTFT acceleration is claimed. Only this FP8 +Qwen3 offload model was measured end to end. Full-model resident/hybrid performance, +other checkpoints, standardized accuracy and NVIDIA execution remain unverified; +operator/bank tests do not substitute for those results. diff --git a/docs/benchmarks/ark-fp8-decode-20261021/calibration.json b/docs/benchmarks/ark-fp8-decode-20261021/calibration.json new file mode 100644 index 0000000000..b7105bd9b9 --- /dev/null +++ b/docs/benchmarks/ark-fp8-decode-20261021/calibration.json @@ -0,0 +1,320 @@ +{ + "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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", + "device_count": 1, + "platform": "Linux-7.0.0-38-generic-x86_64-with-glibc2.39", + "compiler_sha256": "449f6779a4041275d18d801246d0c2c09ffb453a48093d4bebbbd82803954205", + "baseline_sha256": "9815bbdde6aa7c596020f7c1ec5c6e5ea3b363697877a6a052c833e34d74b3af" + }, + "rows": [ + { + "operator": "dense_fp8_decode", + "case": "N5120_K2048", + "metadata": { + "M": 1, + "N": 5120, + "K": 2048 + }, + "correctness": { + "normalized_rmse": 0.0, + "cosine": 0.9999998211860657, + "max_abs": 0.0, + "max_ref": 1.8515625, + "finite": true + }, + "eager": { + "baseline_us": 42.73695987649262, + "candidate_us": 14.336119929794222, + "speedup": 2.981068802840718, + "samples_us": { + "baseline": [ + 43.779159896075726, + 42.98155981814489, + 42.46317985234782, + 42.73695987649262, + 41.98615992208943 + ], + "candidate": [ + 14.82991996454075, + 14.136340178083628, + 14.336119929794222, + 14.085119764786214, + 14.75892000598833 + ] + }, + "enabled": true + }, + "graph": { + "baseline_us": 15.090957880020142, + "candidate_us": 12.952004671096802, + "speedup": 1.1651445674426404, + "samples_us": { + "baseline": [ + 15.374047756195068, + 15.15425443649292, + 15.090957880020142, + 13.583890199661255, + 13.605790138244629 + ], + "candidate": [ + 13.392394781112671, + 13.354895114898682, + 12.952004671096802, + 12.55481481552124, + 12.526215314865112 + ] + }, + "enabled": true + } + }, + { + "operator": "dense_fp8_decode", + "case": "N2048_K4096", + "metadata": { + "M": 1, + "N": 2048, + "K": 4096 + }, + "correctness": { + "normalized_rmse": 0.0, + "cosine": 1.0, + "max_abs": 0.0, + "max_ref": 1.7734375, + "finite": true + }, + "eager": { + "baseline_us": 43.58615988167003, + "candidate_us": 16.359120199922472, + "speedup": 2.664333983063257, + "samples_us": { + "baseline": [ + 43.6879598419182, + 43.26476017013192, + 42.61715977918357, + 43.58615988167003, + 43.689539888873696 + ], + "candidate": [ + 16.41471986658871, + 15.847500180825591, + 16.84950024355203, + 16.359120199922472, + 16.238120151683688 + ] + }, + "enabled": true + }, + "graph": { + "baseline_us": 13.674784898757935, + "candidate_us": 16.653517484664917, + "speedup": 0.8211349290833067, + "samples_us": { + "baseline": [ + 13.933980464935303, + 13.843982219696045, + 13.674784898757935, + 12.418814897537231, + 12.446115016937256 + ], + "candidate": [ + 16.684716939926147, + 17.051807641983032, + 16.653517484664917, + 16.02343201637268, + 16.268327236175537 + ] + }, + "enabled": false + } + }, + { + "operator": "dense_fp8_decode", + "case": "N151936_K2048", + "metadata": { + "M": 1, + "N": 151936, + "K": 2048 + }, + "correctness": { + "normalized_rmse": 5.340938969311537e-06, + "cosine": 1.0, + "max_abs": 0.0009765625, + "max_ref": 2.171875, + "finite": true + }, + "eager": { + "baseline_us": 581.6417402820662, + "candidate_us": 500.55860017891973, + "speedup": 1.1619853101598177, + "samples_us": { + "baseline": [ + 582.4773397762328, + 581.5167599939741, + 581.6417402820662, + 581.8621598882601, + 580.6271600886248 + ], + "candidate": [ + 500.55860017891973, + 500.20219990983605, + 501.82079983642325, + 500.08219986921176, + 500.90499978978187 + ] + }, + "enabled": true + }, + "graph": { + "baseline_us": 580.3911972045898, + "candidate_us": 499.1529083251953, + "speedup": 1.1627523100125228, + "samples_us": { + "baseline": [ + 579.3439483642578, + 588.4440231323242, + 580.3911972045898, + 580.3480911254883, + 589.8081588745117 + ], + "candidate": [ + 499.1529083251953, + 499.552001953125, + 499.11949157714844, + 499.1862487792969, + 499.14161682128906 + ] + }, + "enabled": true + } + }, + { + "operator": "moe_decode_fp8", + "case": "C32_H2048_I768", + "metadata": { + "M": 1, + "H": 2048, + "I": 768, + "bank_rows": 32, + "top_k": 8 + }, + "correctness": { + "normalized_rmse": 0.0, + "cosine": 0.9999999403953552, + "max_abs": 0.0, + "max_ref": 0.23046875, + "finite": true + }, + "eager": { + "baseline_us": 105.46781995799392, + "candidate_us": 45.58776010526344, + "speedup": 2.3135117784788224, + "samples_us": { + "baseline": [ + 105.63260002527386, + 107.13299998315051, + 100.50904005765915, + 105.46781995799392, + 100.46084003988653 + ], + "candidate": [ + 48.84213994955644, + 47.465919924434274, + 45.58776010526344, + 43.56513993116096, + 43.68355992482975 + ] + }, + "enabled": true + }, + "graph": { + "baseline_us": 53.39667797088623, + "candidate_us": 44.586381912231445, + "speedup": 1.197600605404536, + "samples_us": { + "baseline": [ + 57.948174476623535, + 52.832889556884766, + 53.39667797088623, + 53.42008113861084, + 53.18868160247803 + ], + "candidate": [ + 44.586381912231445, + 44.8009729385376, + 44.05508995056152, + 44.63818073272705, + 44.17378902435303 + ] + }, + "enabled": false + } + }, + { + "operator": "moe_decode_fp8", + "case": "C4550_H2048_I768", + "metadata": { + "M": 1, + "H": 2048, + "I": 768, + "bank_rows": 4550, + "top_k": 8 + }, + "correctness": { + "normalized_rmse": 0.0, + "cosine": 0.9999999403953552, + "max_abs": 0.0, + "max_ref": 0.216796875, + "finite": true + }, + "eager": { + "baseline_us": 104.50481990119442, + "candidate_us": 43.96515985717997, + "speedup": 2.3769916961675213, + "samples_us": { + "baseline": [ + 104.50481990119442, + 106.70819989172742, + 102.2930201725103, + 106.53300007106736, + 100.65804002806544 + ], + "candidate": [ + 44.85156008740887, + 43.23495988501236, + 43.75496006105095, + 43.96515985717997, + 44.09874003613368 + ] + }, + "enabled": true + }, + "graph": { + "baseline_us": 54.33335304260254, + "candidate_us": 40.294880867004395, + "speedup": 1.348393440395889, + "samples_us": { + "baseline": [ + 54.33335304260254, + 337.1195983886719, + 54.96233940124512, + 47.28722095489502, + 47.352614402770996 + ], + "candidate": [ + 40.70597171783447, + 39.989686012268066, + 40.99296569824219, + 40.294880867004395, + 39.54289674758911 + ] + }, + "enabled": false + } + } + ] +} diff --git a/docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json b/docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json new file mode 100644 index 0000000000..a746d2cc8d --- /dev/null +++ b/docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json @@ -0,0 +1,61 @@ +{ + "_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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", + "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, + "graph_min_speedup": 1.03, + "each_repeat_faster": true + }, + "_quality_exclusions": { + "moe_decode_fp8": "All Graph shapes remain FT; Qwen3 matched-prefix KL previously exceeded 0.05 nats." + }, + "dense_fp8_decode": { + "eager": { + "1,5120,2048": { + "speedup": 2.981068802840718, + "min_paired_speedup": 2.8447989354948624 + }, + "1,2048,4096": { + "speedup": 2.664333983063257, + "min_paired_speedup": 2.5292833118591935 + }, + "1,151936,2048": { + "speedup": 1.1619853101598177, + "min_paired_speedup": 1.1590626384391836 + } + }, + "graph": { + "1,5120,2048": { + "speedup": 1.1651445674426404, + "min_paired_speedup": 1.0819665920414685 + }, + "1,151936,2048": { + "speedup": 1.1627523100125228, + "min_paired_speedup": 1.1606542578468129 + } + } + }, + "moe_decode_fp8": { + "eager": { + "1,2048,768,32,8": { + "speedup": 2.3135117784788224, + "min_paired_speedup": 2.162734887012934 + }, + "1,2048,768,4550,8": { + "speedup": 2.3769916961675213, + "min_paired_speedup": 2.2825604528743484 + } + }, + "graph": {} + } +} diff --git a/docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv b/docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv new file mode 100644 index 0000000000..c80d87970a --- /dev/null +++ b/docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv @@ -0,0 +1,25 @@ +mode,round,enabled,context,decode_tok_s,ttft_s,e2e_s,token_text_chunks,text_sha256 +eager,1,0,32768,28.366052932428786,10.775955886012525,19.766115287013236,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb +eager,1,0,512,27.465098474060973,0.8498681829951238,10.134956226000213,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 +eager,1,1,32768,36.337070449318986,10.830030515004182,17.849467851003283,256,5c20deafdd17c988f9da13d473d60069d0b6eb9e854090b5884ead328e29181e +eager,1,1,512,30.28456095314772,0.9327965800039237,9.353914507999434,256,e809cdcb447a0670cefb38e065aaba2cc539f7b132e95fba4900045330fcc002 +eager,2,0,32768,31.473914371452555,10.866772359004244,18.970433360009338,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb +eager,2,0,512,31.61883942692908,0.6422161329974188,8.707675377998385,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 +eager,2,1,32768,35.719651721432726,13.993129288995988,21.133178561998648,256,5c20deafdd17c988f9da13d473d60069d0b6eb9e854090b5884ead328e29181e +eager,2,1,512,37.675944950872434,0.6477360019925982,7.416634577995865,256,e809cdcb447a0670cefb38e065aaba2cc539f7b132e95fba4900045330fcc002 +eager,3,0,32768,20.53616960530507,10.875875447993167,23.294661559993983,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb +eager,3,0,512,29.4541884296002,0.6517329810012598,9.309694625990232,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 +eager,3,1,32768,36.72163966356929,10.862550713005476,17.808452309996937,256,5c20deafdd17c988f9da13d473d60069d0b6eb9e854090b5884ead328e29181e +eager,3,1,512,34.83756625429643,0.6866280950052897,8.006714333008858,256,e809cdcb447a0670cefb38e065aaba2cc539f7b132e95fba4900045330fcc002 +graph_candidate_rejected,1,0,32768,65.22113636555925,10.602495545012061,14.513973229011754,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb +graph_candidate_rejected,1,0,512,56.33858678141364,0.7210316830023658,5.247661571993376,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 +graph_candidate_rejected,1,1,32768,64.9769935118381,10.927924178002286,14.853244415004156,256,1c1e3c7df3e702de994c988f2e2c8ac57f299713b9418c51de8e42756022cdc2 +graph_candidate_rejected,1,1,512,100.51418361740437,1.3024535959993955,3.840363162002177,256,555a5f1f33054486d0637db560b98685ec71a7d776fdd9736add0d9511f913a4 +graph_candidate_rejected,2,0,32768,65.32587621602467,14.445713420995162,18.350080314994557,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb +graph_candidate_rejected,2,0,512,100.88825964400858,0.7239236070017796,3.2523054859921103,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 +graph_candidate_rejected,2,1,32768,65.06508241655588,10.963409739997587,14.884385861005285,256,1c1e3c7df3e702de994c988f2e2c8ac57f299713b9418c51de8e42756022cdc2 +graph_candidate_rejected,2,1,512,100.09894710247522,0.7882882770063588,3.3363208860100713,256,555a5f1f33054486d0637db560b98685ec71a7d776fdd9736add0d9511f913a4 +graph_candidate_rejected,3,0,32768,65.14931377043827,10.991303474002052,14.9065101879969,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb +graph_candidate_rejected,3,0,512,100.17950068652583,1.1008289820019854,3.647286549996352,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 +graph_candidate_rejected,3,1,32768,64.84847037983543,11.109217436009203,15.042871676007053,256,1c1e3c7df3e702de994c988f2e2c8ac57f299713b9418c51de8e42756022cdc2 +graph_candidate_rejected,3,1,512,52.61622377409661,0.7942161740065785,5.641266263002763,256,555a5f1f33054486d0637db560b98685ec71a7d776fdd9736add0d9511f913a4 diff --git a/docs/benchmarks/ark-fp8-decode-20261021/runtime.json b/docs/benchmarks/ark-fp8-decode-20261021/runtime.json new file mode 100644 index 0000000000..1b1ce2296c --- /dev/null +++ b/docs/benchmarks/ark-fp8-decode-20261021/runtime.json @@ -0,0 +1,48 @@ +{ + "base": "fe5a5ec9c5644e266527de5afb9f1156b9bdcabb", + "branch": "perf/ark-fp8-decode", + "gpu": { + "name": "AMD Radeon AI PRO R9700", + "pci": "0000:23:00.0", + "uuid": "GPU-98103ef4c22a3ccc", + "render": "renderD131", + "memory_gib": 32, + "link": "32 GT/s x16" + }, + "cpu": { + "model": "AMD EPYC 9334 32-Core Processor", + "host_sockets": 2, + "allowed_cpus": "0-31,64-95", + "memory_nodes": "2-3", + "exclusive": false + }, + "os": "Ubuntu 24.04.1 LTS", + "amdgpu": "7.1.9.31600000", + "model": "Qwen/Qwen3-30B-A3B-FP8", + "checkpoint_format": "FTW", + "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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", + "device_count": 1, + "platform": "Linux-7.0.0-38-generic-x86_64-with-glibc2.39", + "compiler_sha256": "449f6779a4041275d18d801246d0c2c09ffb453a48093d4bebbbd82803954205", + "baseline_sha256": "9815bbdde6aa7c596020f7c1ec5c6e5ea3b363697877a6a052c833e34d74b3af" + }, + "native_cache_key": "b63ee64d5f192756348a", + "weight_dtype": "E4M3 FP8", + "activation_dtype": "BF16", + "kv_dtype": "BF16", + "expert_cache_gib": 20, + "expert_slots": 4550, + "input_contexts": [ + 512, + 32768 + ], + "output_tokens": 256, + "concurrency": 1, + "prefill_chunk": 8192 +} diff --git a/docs/benchmarks/ark-fp8-decode-20261021/validation.json b/docs/benchmarks/ark-fp8-decode-20261021/validation.json new file mode 100644 index 0000000000..a0d98fbad3 --- /dev/null +++ b/docs/benchmarks/ark-fp8-decode-20261021/validation.json @@ -0,0 +1,115 @@ +{ + "tests": { + "passed": 41, + "skipped": 42, + "ark_cases": 23 + }, + "formal_requests": 24, + "numerical_shadow": { + "comparisons": 144, + "max_normalized_rmse": 0.0007347841165028512 + }, + "native_gpu_kernel_names": [ + "void ark_ft::down_reduce(float const*, void const*, __hip_bfloat16 const*, int const*, float const*, __hip_bfloat16*, int, int, int, int, int) [clone .kd]", + "void ark_ft::gate_up(__hip_bfloat16 const*, void const*, __hip_bfloat16 const*, int const*, float const*, float*, int, int, int, int, int) [clone .kd]", + "void radeon_ops::hip_native::dense_gemv_fp8_rows<4>(__hip_bfloat16 const*, unsigned char const*, __hip_bfloat16 const*, __hip_bfloat16*, int, int, int) [clone .kd]" + ], + "native_eager_calls_per_context": { + "512": { + "dense_fp8_decode": 672, + "moe_decode_fp8": 336 + }, + "32768": { + "dense_fp8_decode": 672, + "moe_decode_fp8": 336 + } + }, + "graph": { + "deployed_graph_capture_native_calls": 0, + "probe_outputs_match_ft": [ + 512, + 32768 + ], + "candidate_speedups_rejected": [ + { + "context": 512, + "rounds": [ + { + "round": 1, + "speedup": 1.7841090691072228, + "ft_decode_tok_s": 56.33858678141364, + "ark_decode_tok_s": 100.51418361740437, + "outputs_equal": false + }, + { + "round": 2, + "speedup": 0.9921763687438111, + "ft_decode_tok_s": 100.88825964400858, + "ark_decode_tok_s": 100.09894710247522, + "outputs_equal": false + }, + { + "round": 3, + "speedup": 0.5252194651951735, + "ft_decode_tok_s": 100.17950068652583, + "ark_decode_tok_s": 52.61622377409661, + "outputs_equal": false + } + ], + "ft_median": 100.17950068652583, + "ark_median": 100.09894710247522, + "median_change_percent": -0.08040924889680667, + "wins": 1, + "n": 3 + }, + { + "context": 32768, + "rounds": [ + { + "round": 1, + "speedup": 0.9962566911997248, + "ft_decode_tok_s": 65.22113636555925, + "ark_decode_tok_s": 64.9769935118381, + "outputs_equal": false + }, + { + "round": 2, + "speedup": 0.9960078025037679, + "ft_decode_tok_s": 65.32587621602467, + "ark_decode_tok_s": 65.06508241655588, + "outputs_equal": false + }, + { + "round": 3, + "speedup": 0.995382247744575, + "ft_decode_tok_s": 65.14931377043827, + "ark_decode_tok_s": 64.84847037983543, + "outputs_equal": false + } + ], + "ft_median": 65.22113636555925, + "ark_median": 64.9769935118381, + "median_change_percent": -0.374330880027518, + "wins": 0, + "n": 3 + } + ] + }, + "generation": { + "ft_passed": 7, + "ark_passed": 7, + "total": 8, + "new_failures": [], + "shared_failure": "decimal: both answered 0.11 instead of 0.9", + "method": "non-thinking tokenizer template, 128 generated tokens, score first answer only; initial 64-token thinking probes excluded; not standardized accuracy" + }, + "standalone_wheel": { + "passed": true, + "normalized_rmse": 0.0, + "sha256": "01a49d978b14496514141cc4645c1372c12542898ad7f6a2dd7ef30876f4d9dc", + "ark_packaged_files": 17, + "new_native_cache": true, + "validated_policy": true, + "isolated_python": true + } +} diff --git a/pyproject.toml b/pyproject.toml index ba3c3a7693..c35288da6b 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 7982095ac4..de20541a70 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 0000000000..2eda14b59b --- /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 0000000000..75654b704b --- /dev/null +++ b/python/freetoken/kernel/ark/__init__.py @@ -0,0 +1,37 @@ +"""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 0000000000..e69de29bb2 diff --git a/python/freetoken/kernel/ark/common/bindings.cpp b/python/freetoken/kernel/ark/common/bindings.cpp new file mode 100644 index 0000000000..ef81bdbe05 --- /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, 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, int64_t nw) { + 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 && nw == 1, + "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), nw, 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, int nw) -> 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 0000000000..1ec2e7a222 --- /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 0000000000..5568ad4588 --- /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 0000000000..2864c34ca4 --- /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 +from pathlib import Path +import shutil +import subprocess +import threading + +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, 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 0000000000..947073cdc5 --- /dev/null +++ b/python/freetoken/kernel/ark/common/policy.py @@ -0,0 +1,77 @@ +"""Fail-closed environment and measured-shape selection, separate for eager/Graph.""" + +import json +import os +import hashlib +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 source_digest, _compiler + + 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/trace.py b/python/freetoken/kernel/ark/common/trace.py new file mode 100644 index 0000000000..833054d21c --- /dev/null +++ b/python/freetoken/kernel/ark/common/trace.py @@ -0,0 +1,43 @@ +"""Optional Python dispatch census; Graph captures are distinct from replays.""" + +import atexit +from collections import Counter +import json +import os +from pathlib import Path +import torch + +_PATH = os.environ.get("FREETOKEN_ARK_TRACE") +_COUNTS = Counter() +_TOTAL = 0 + + +def record(op, shape, backend): + global _TOTAL + if _PATH: + mode = "capture" if torch.cuda.is_current_stream_capturing() else "eager" + item = (op, shape, backend, mode) + first = item not in _COUNTS + _COUNTS[item] += 1 + _TOTAL += 1 + if first or _TOTAL % 1000 == 0: + flush() + + +def flush(): + if _PATH: + path = Path(_PATH.replace("{pid}", str(os.getpid()))) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "pid": os.getpid(), + "counts": [dict(operator=o, shape=s, backend=b, phase=m, calls=n) for (o, s, b, m), n in sorted(_COUNTS.items())], + "note": "Python eager calls and Graph captures only; replay GPU execution needs profiler evidence.", + }, + indent=2, + ) + ) + + +atexit.register(flush) 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 0000000000..ecc0c403d2 --- /dev/null +++ b/python/freetoken/kernel/ark/common/validated_shapes.json @@ -0,0 +1,56 @@ +{ + "_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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", + "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, + "graph_min_speedup": 1.03, + "each_repeat_faster": true, + "graph_requires_model_benefit": true + }, + "_quality_exclusions": { + "moe_decode_fp8": "All Graph shapes remain FT; Qwen3 matched-prefix KL previously exceeded 0.05 nats." + }, + "dense_fp8_decode": { + "eager": { + "1,5120,2048": { + "speedup": 2.981068802840718, + "min_paired_speedup": 2.8447989354948624 + }, + "1,2048,4096": { + "speedup": 2.664333983063257, + "min_paired_speedup": 2.5292833118591935 + }, + "1,151936,2048": { + "speedup": 1.1619853101598177, + "min_paired_speedup": 1.1590626384391836 + } + }, + "graph": {} + }, + "moe_decode_fp8": { + "eager": { + "1,2048,768,32,8": { + "speedup": 2.3135117784788224, + "min_paired_speedup": 2.162734887012934 + }, + "1,2048,768,4550,8": { + "speedup": 2.3769916961675213, + "min_paired_speedup": 2.2825604528743484 + } + }, + "graph": {} + }, + "_model_exclusions": { + "dense_fp8_decode.graph": "Three-round Qwen3 candidate model A/B: 512 -0.0804%, 32768 -0.3743%; keep FT until model benefit is demonstrated." + } +} 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 0000000000..e69de29bb2 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 0000000000..68e2f2aaa1 --- /dev/null +++ b/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py @@ -0,0 +1,41 @@ +"""M=1 FreeToken block_fp8_linear, including bias and leading dimensions.""" + +import torch +from freetoken.kernel.ark.common.trace import record +from freetoken.kernel.ark.common.policy import allowed, key +from freetoken.kernel.ark.common.contracts import gfx1201 +from freetoken.kernel.ark.common.native import load + +NW = 1 + + +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 + ): + record("dense_fp8_decode", key(x.numel() // k, weight.shape[0], k), "freetoken") + from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear as baseline + + return baseline(x, weight, weight_scale, bias) + record("dense_fp8_decode", key(x.numel() // k, weight.shape[0], k), "ark") + load() + out = torch.ops.freetoken_ark.dense_decode(x.reshape(1, k), weight, weight_scale, NW).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 0000000000..689bf6bb46 --- /dev/null +++ b/python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip @@ -0,0 +1,139 @@ +// 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 +} + +// Split-K: NW warps per output row, each striding over a 1/NW share of the K-blocks (shorter serial +// reduction chain at large K), combined across warps via LDS. NW=1 is the plain wave-per-output path. +template +__global__ void __launch_bounds__(NW * 32) dense_gemv_fp8_k(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; + if (n >= N) return; + int warp = threadIdx.x >> 5, 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 = warp; kb < nkb; kb += NW) { + 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 (NW == 1) { + if (lane == 0) out[n] = (bf16)acc; + return; + } + __shared__ float red[NW]; + if (lane == 0) red[warp] = acc; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + #pragma unroll + for (int i = 0; i < NW; ++i) s += red[i]; + out[n] = (bf16)s; + } +} + +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 warp = 0, 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 = warp; kb < nkb; kb += 1) { + 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 (1 == 1) { + if (lane == 0) out[n] = (bf16)acc; + return; + } + __shared__ float red[1]; + if (lane == 0) red[warp] = acc; + __syncthreads(); + if (threadIdx.x == 0) { + float s = 0.f; + #pragma unroll + for (int i = 0; i < 1; ++i) s += red[i]; + out[n] = (bf16)s; + } +} + +void launch_dense_gemv_fp8(const bf16* x, const unsigned char* w, const bf16* wscale, bf16* out, + int N, int K, int wscale_stride, int nw, hipStream_t stream) { + (void)K; + // auto split-K by N: small N underfills the GPU at NW=1, so split each output across more warps; large + // N already has plenty of output waves -> NW=1. NOTE this op WINS at moderate N (e.g. GDN out_proj + // K=4096 N=2048 ~2x both arches) but LOSES to triton's tuned big-N GEMV at large N (e.g. N>=16384 on + // gfx1201) -- the caller should gate on N and keep triton for large-N projections (lm_head, in_proj). + if (nw <= 0) nw = (N <= 2048) ? 4 : (N <= 8192 ? 2 : 1); + dim3 grid((unsigned)N); + if (nw >= 4) hipLaunchKernelGGL((dense_gemv_fp8_k<4>), grid, dim3(128), 0, stream, x, w, wscale, out, N, K, wscale_stride); + else if (nw == 2) hipLaunchKernelGGL((dense_gemv_fp8_k<2>), grid, dim3(64), 0, stream, x, w, wscale, out, N, K, wscale_stride); + else 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, int nw, 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, nw, (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 0000000000..606e6caeaa --- /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 0000000000..e69de29bb2 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 0000000000..2c03d2887a --- /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.trace import record +from freetoken.kernel.ark.common.policy import allowed, key +from freetoken.kernel.ark.common.contracts import moe_supported, scales_supported, plain_silu +from freetoken.kernel.ark.common.native import load + + +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) + ): + record( + "moe_decode_fp8", + key(hidden_states.shape[0], hidden_states.shape[1], down.shape[2], gate_up.shape[0], topk_ids.shape[1]), + "freetoken", + ) + 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, + ) + record("moe_decode_fp8", key(hidden_states.shape[0], hidden_states.shape[1], down.shape[2], gate_up.shape[0], topk_ids.shape[1]), "ark") + 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 0000000000..004da33f2c --- /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 2a079539f1..a11eed1476 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 8f347e5c30..a3c2048e97 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 0000000000..65a51cda1b --- /dev/null +++ b/tests/kernels/test_ark_fp8.py @@ -0,0 +1,264 @@ +"""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 rand, quant +from benchmarks.ark.timing import metrics, assert_close + +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.common import policy + from freetoken.kernel.ark import initialize + + 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 = [] + monkeypatch.setattr(adapter, "record", lambda *a: hits.append(a)) + 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[-1][-1] == "ark" + + +@pytest.mark.parametrize("kind", ["batch", "strided", "disabled", "unknown_shape"]) +def test_dense_fallback(kind, monkeypatch): + from freetoken.kernel.ark.dense_fp8_decode import adapter + from freetoken.kernel.ark.common import policy + 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", {}) + hits = [] + monkeypatch.setattr(adapter, "record", lambda *a: hits.append(a)) + compare(adapter.block_fp8_linear(x, w, scale), baseline(x, w, scale)) + assert hits[-1][-1] == "freetoken" + + +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.moe_decode_fp8 import adapter + from freetoken.kernel.ark.common import policy + 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) + hits = [] + monkeypatch.setattr(adapter, "record", lambda *a: hits.append(a)) + compare(adapter.fused_experts_decode_fp8_block(*args, **kwargs), baseline(*args, **kwargs)) + assert hits[-1][-1] == "freetoken" + + +@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 ArkFp8MoEKernel, ArkFp8LinearKernel + from freetoken.layers.quantization.moe.base import MoEConfig, ExpertView + from freetoken.layers.quantization.moe.fp8_block import Fp8BlockMoEMethod, TritonFp8BlockMoEKernel + from freetoken.layers.quantization.linear.base import LinearConfig + from freetoken.layers.quantization.linear.fp8_block import Fp8BlockLinearMethod, TritonFp8BlockLinearKernel + 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 = dict(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 k in bank: + assert torch.equal(bank[k], ref_bank[k]) + 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.engine.engine import _adjust_ftw_quant_backend + from freetoken.layers.quantization import QuantBackend, LayerKind, QuantKind + from freetoken.moe.legacy_format import legacy_format_for, kind_kernel_for, bank_kernel_compatible + import freetoken.checkpoint.ftw as ftw + + 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) From ca5a5eb9837c3faf79b2464e40e86111ede148f1 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Wed, 21 Oct 2026 16:13:53 +0800 Subject: [PATCH 2/2] refactor(kernel): narrow ark fp8 decode integration scope --- benchmarks/ark/audit.py | 77 ----- benchmarks/ark/audit_quality.py | 41 --- benchmarks/ark/calibrate.py | 32 +- benchmarks/ark/cases.py | 12 +- benchmarks/ark/http_client.py | 132 -------- benchmarks/ark/model_ab.py | 259 -------------- benchmarks/ark/probe.py | 87 ----- benchmarks/ark/quality_prompts.py | 44 --- benchmarks/ark/timing.py | 2 +- docs/ark-fp8-decode.md | 236 +++---------- .../ark-fp8-decode-20261021/calibration.json | 320 ------------------ .../graph-candidate-profile.json | 61 ---- .../ark-fp8-decode-20261021/model-samples.csv | 25 -- .../ark-fp8-decode-20261021/runtime.json | 48 --- .../ark-fp8-decode-20261021/validation.json | 115 ------- python/freetoken/kernel/ark/__init__.py | 1 + .../freetoken/kernel/ark/common/bindings.cpp | 10 +- python/freetoken/kernel/ark/common/native.py | 4 +- python/freetoken/kernel/ark/common/policy.py | 5 +- python/freetoken/kernel/ark/common/trace.py | 43 --- .../kernel/ark/common/validated_shapes.json | 33 +- .../kernel/ark/dense_fp8_decode/adapter.py | 14 +- .../ark/dense_fp8_decode/dense_gemv_fp8.hip | 76 +---- .../kernel/ark/moe_decode_fp8/adapter.py | 18 +- tests/kernels/test_ark_fp8.py | 84 +++-- 25 files changed, 177 insertions(+), 1602 deletions(-) delete mode 100644 benchmarks/ark/audit.py delete mode 100644 benchmarks/ark/audit_quality.py delete mode 100644 benchmarks/ark/http_client.py delete mode 100644 benchmarks/ark/model_ab.py delete mode 100644 benchmarks/ark/probe.py delete mode 100644 benchmarks/ark/quality_prompts.py delete mode 100644 docs/benchmarks/ark-fp8-decode-20261021/calibration.json delete mode 100644 docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json delete mode 100644 docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv delete mode 100644 docs/benchmarks/ark-fp8-decode-20261021/runtime.json delete mode 100644 docs/benchmarks/ark-fp8-decode-20261021/validation.json delete mode 100644 python/freetoken/kernel/ark/common/trace.py diff --git a/benchmarks/ark/audit.py b/benchmarks/ark/audit.py deleted file mode 100644 index ad7ec6ad52..0000000000 --- a/benchmarks/ark/audit.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Independently recompute formal SSE timing and paired summaries from raw files.""" - -import argparse -import hashlib -import json -from pathlib import Path -import statistics - - -def audit(root): - rows = [] - for path in sorted(root.glob("round-*/measured-*.json")): - row = json.loads(path.read_text()) - assert row["success"] and row["done"] - assert row["usage"]["prompt_tokens"] == row["context"] - assert row["usage"]["completion_tokens"] == 256 - assert row["request"]["prompt"] == (root / f"prompt-{row['context']}.txt").read_text() - assert row["request"]["temperature"] == 0 and row["request"]["ignore_eos"] - assert (row["usage"].get("prompt_tokens_details") or {}).get("cached_tokens", 0) == 0 - events = [x for x in row["events"] if any(c.get("text") for c in x["data"].get("choices", []))] - span = events[-1]["elapsed_s"] - events[0]["elapsed_s"] - assert abs(255 / span - row["decode_tok_s"]) < 1e-9 - assert events[0]["elapsed_s"] == row["ttft_s"] - text = "".join(c.get("text", "") for x in events for c in x["data"]["choices"]) - assert hashlib.sha256(text.encode()).hexdigest() == row["text_sha256"] - command = json.loads((path.parent / "command.json").read_text()) - assert command["env"]["FREETOKEN_USE_ARK"] == str(row["enabled"]) - assert command["env"]["FREETOKEN_ARK_POLICY"] == "validated" - assert "FREETOKEN_ARK_TRACE" not in command["env"] - assert "benchmarks.ark.probe" not in command["argv"] - assert json.loads((path.parent / "complete.json").read_text())["success"] - assert json.loads((path.parent / "shutdown.json").read_text())["vram"] < 256 * 2**20 - rows.append({k: v for k, v in row.items() if k not in ("events", "text", "request")}) - paired = [] - for context in (512, 32768): - pairs = [] - for rid in sorted({r["round"] for r in rows}): - pair = {r["enabled"]: r for r in rows if r["round"] == rid and r["context"] == context} - assert set(pair) == {0, 1} - baseline, candidate = pair[0], pair[1] - pairs.append( - dict( - round=rid, - speedup=candidate["decode_tok_s"] / baseline["decode_tok_s"], - ft_decode_tok_s=baseline["decode_tok_s"], - ark_decode_tok_s=candidate["decode_tok_s"], - outputs_equal=baseline["text_sha256"] == candidate["text_sha256"], - ) - ) - if pairs: - ft = statistics.median(r["ft_decode_tok_s"] for r in pairs) - ark = statistics.median(r["ark_decode_tok_s"] for r in pairs) - paired.append( - dict( - context=context, - rounds=pairs, - ft_median=ft, - ark_median=ark, - median_change_percent=100 * (ark / ft - 1), - wins=sum(r["speedup"] > 1 for r in pairs), - n=len(pairs), - ) - ) - return dict(passed=True, samples=rows, paired=paired) - - -def main(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("root", type=Path) - a = p.parse_args() - result = audit(a.root) - (a.root / "audit.json").write_text(json.dumps(result, indent=2) + "\n") - print(json.dumps(result["paired"], indent=2)) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/ark/audit_quality.py b/benchmarks/ark/audit_quality.py deleted file mode 100644 index 74ad49d8c5..0000000000 --- a/benchmarks/ark/audit_quality.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Score the first answer in saved non-thinking generation checks.""" - -import argparse -import json -from pathlib import Path - -from .quality_prompts import CASES, first_answer - - -def main(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("root", type=Path) - a = p.parse_args() - groups = {} - for enabled in (0, 1): - folder = a.root / f"round-1-ark{enabled}" - rows = [] - for name, _, predicate in CASES: - paths = list(folder.glob(f"quality-{name}-*.json")) - assert len(paths) == 1, (name, paths) - raw = json.loads(paths[0].read_text()) - assert raw["success"] and raw["usage"]["completion_tokens"] == 128 - assert raw["request"]["prompt"].endswith("\n\n\n\n") - answer = first_answer(raw["text"]) - rows.append(dict(name=name, passed=bool(predicate(answer)), answer=answer)) - groups[str(enabled)] = rows - (folder / "quality-audited.json").write_text(json.dumps(rows, indent=2) + "\n") - old = {r["name"] for r in groups["0"] if not r["passed"]} - new = {r["name"] for r in groups["1"] if not r["passed"]} - result = dict( - groups=groups, - new_failures=sorted(new - old), - note="8 bounded functional checks, first answer only, non-thinking template, not a standardized accuracy benchmark", - ) - (a.root / "quality-audit.json").write_text(json.dumps(result, indent=2) + "\n") - print("baseline", len(groups["0"]) - len(old), "ARK", len(groups["1"]) - len(new), "new failures", sorted(new - old)) - assert not new - old, result - - -if __name__ == "__main__": - main() diff --git a/benchmarks/ark/calibrate.py b/benchmarks/ark/calibrate.py index 74f2ddb2f2..da93da84ae 100644 --- a/benchmarks/ark/calibrate.py +++ b/benchmarks/ark/calibrate.py @@ -6,20 +6,17 @@ from pathlib import Path import torch - -from .cases import cases -from .timing import metrics, assert_close, timing_pair, graph_pair 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) - p.add_argument( - "--graph-candidates", action="store_true", help="emit dense Graph microbenchmark candidates for separate model evaluation" - ) a = p.parse_args() torch.manual_seed(20261021) torch.set_num_threads(8) @@ -28,12 +25,8 @@ def main(): policy.MODE = "native" profile = { "_environment": policy.current_environment(), - "_rules": {"eager_min_speedup": 1.05, "graph_min_speedup": 1.03, "each_repeat_faster": True}, - "_quality_exclusions": {"moe_decode_fp8": "All Graph shapes remain FT; Qwen3 matched-prefix KL previously exceeded 0.05 nats."}, - } - profile["_rules"]["graph_requires_model_benefit"] = True - profile["_model_exclusions"] = { - "dense_fp8_decode.graph": "Current Qwen3 candidate model A/B did not establish a Graph benefit; default Graph remains FT." + "_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"): @@ -52,15 +45,12 @@ def main(): "graph": graph_pair(case.baseline, case.candidate, case.reset, False, 100), } shape = policy.metadata_key(op, case.metadata) - for mode, threshold in (("eager", 1.05), ("graph", 1.03)): - timing = row[mode] - paired = [x / y for x, y in zip(timing["samples_us"]["baseline"], timing["samples_us"]["candidate"])] - eligible = timing["speedup"] >= threshold and min(paired) > 1 - if mode == "graph" and (op == "moe_decode_fp8" or not a.graph_candidates): - eligible = False - row[mode]["enabled"] = eligible - if eligible: - profile[op][mode][shape] = {"speedup": timing["speedup"], "min_paired_speedup": min(paired)} + 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") diff --git a/benchmarks/ark/cases.py b/benchmarks/ark/cases.py index f49b7c1ef2..adfa04b92b 100644 --- a/benchmarks/ark/cases.py +++ b/benchmarks/ark/cases.py @@ -29,21 +29,25 @@ def quant(w): 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 - from freetoken.kernel.ark.dense_fp8_decode.adapter import block_fp8_linear as candidate 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), dict(M=1, N=n, K=k)) + 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 - from freetoken.kernel.ark.moe_decode_fp8.adapter import fused_experts_decode_fp8_block as candidate h, i, topk = 2048, 768, 8 for rows in (32, 4550): @@ -63,7 +67,7 @@ def moe_cases(): 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), dict(M=1, H=h, I=i, bank_rows=rows, top_k=topk) + 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} ) diff --git a/benchmarks/ark/http_client.py b/benchmarks/ark/http_client.py deleted file mode 100644 index 81bbacf5bb..0000000000 --- a/benchmarks/ark/http_client.py +++ /dev/null @@ -1,132 +0,0 @@ -"""SSE client and resource sampling, using the same timing definitions as prior reports.""" - -import hashlib, json, time, traceback, urllib.error, urllib.request -from pathlib import Path - -ROOT = Path("/artifacts") -ORIGIN = "http://127.0.0.1:19673" -DEVICE = None -MODEL = "/models/Qwen3-30B-A3B-FP8-FTW" - - -def save(path, data): - temp = path.with_suffix(path.suffix + ".tmp") - temp.write_text(json.dumps(data, indent=2) + "\n") - temp.replace(path) - - -def get(route, timeout=3): - with urllib.request.urlopen(ORIGIN + route, timeout=timeout) as response: - return json.load(response) - - -def tree(pid): - result, pending = set(), [pid] - while pending: - p = pending.pop() - if p in result: - continue - result.add(p) - try: - pending.extend(int(x) for x in Path(f"/proc/{p}/task/{p}/children").read_text().split()) - except OSError: - pass - return result - - -def snapshot(pid): - data = {"time": time.time(), "rss_sum_bytes": 0, "processes": []} - for p in tree(pid): - try: - status = dict(line.split(":", 1) for line in Path(f"/proc/{p}/status").read_text().splitlines() if ":" in line) - rss = int(status.get("VmRSS", "0 kB").split()[0]) * 1024 - data["rss_sum_bytes"] += rss - data["processes"].append( - { - "pid": p, - "name": status["Name"].strip(), - "rss_bytes": rss, - "cpus": status["Cpus_allowed_list"].strip(), - "mems": status["Mems_allowed_list"].strip(), - } - ) - except (OSError, KeyError): - pass - for key in ["mem_info_vram_used", "gpu_busy_percent"]: - try: - data[key] = int((DEVICE / key).read_text()) - except (OSError, ValueError): - data[key] = None - return data - - -def monitor(pid, stop, path): - with path.open("w") as output: - while not stop.is_set(): - output.write(json.dumps(snapshot(pid)) + "\n") - output.flush() - stop.wait(0.25) - - -def request(directory, context, phase, output_tokens=256, prompt=None): - prompt = prompt if prompt is not None else (ROOT / f"prompt-{context}.txt").read_text() - payload = { - "model": MODEL, - "prompt": prompt, - "max_tokens": output_tokens, - "temperature": 0, - "ignore_eos": True, - "stream": True, - "stream_options": {"include_usage": True}, - } - req = urllib.request.Request( - ORIGIN + "/v1/completions", data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"} - ) - row = {"context": context, "phase": phase, "expected_output_tokens": output_tokens, "request": payload, "events": []} - started = time.perf_counter() - try: - with urllib.request.urlopen(req, timeout=180) as response: - for line in response: - now = time.perf_counter() - if now - started > 180: - raise TimeoutError("request exceeded 180s") - if not line.startswith(b"data: "): - continue - value = line[6:].strip() - if value == b"[DONE]": - row["done"] = True - break - event = json.loads(value) - row["events"].append({"elapsed_s": now - started, "data": event}) - if "error" in event: - raise RuntimeError(event["error"]) - row["e2e_s"] = time.perf_counter() - started - row["usage"] = next(x["data"]["usage"] for x in reversed(row["events"]) if "usage" in x["data"]) - chunks = [x for x in row["events"] if any(c.get("text") for c in x["data"].get("choices", []))] - row["text"] = "".join(c.get("text", "") for x in chunks for c in x["data"]["choices"]) - row["text_sha256"] = hashlib.sha256(row["text"].encode()).hexdigest() - row["token_text_chunks"] = len(chunks) - row["ttft_s"] = chunks[0]["elapsed_s"] - row["decode_span_s"] = chunks[-1]["elapsed_s"] - chunks[0]["elapsed_s"] - row["decode_tok_s"] = (output_tokens - 1) / row["decode_span_s"] - # The official completion streamer emits at most one nonempty event per generated token. - row["decode_timing_kind"] = ( - "one-event-per-token verified" - if len(chunks) == row["usage"]["completion_tokens"] - else "client legacy estimate; empty-token events omitted" - ) - row["e2e_output_tok_s"] = row["usage"]["completion_tokens"] / row["e2e_s"] - row["prefill_effective_tok_s"] = context / row["ttft_s"] - assert row.get("done"), "missing SSE DONE" - assert row["usage"]["prompt_tokens"] == context, row["usage"] - assert row["usage"]["completion_tokens"] == output_tokens, row["usage"] - details = row["usage"].get("prompt_tokens_details") or {} - assert details.get("cached_tokens", 0) == 0, row["usage"] - assert row["text"].strip(), "empty output" - row["success"] = True - except Exception: - row.update(success=False, error=traceback.format_exc(), elapsed_s=time.perf_counter() - started) - raise - finally: - save(directory / f"{phase}-{context}.json", row) - return row diff --git a/benchmarks/ark/model_ab.py b/benchmarks/ark/model_ab.py deleted file mode 100644 index b5cc7c084e..0000000000 --- a/benchmarks/ark/model_ab.py +++ /dev/null @@ -1,259 +0,0 @@ -"""Fresh-process off/on Qwen3 A/B; instrumentation runs separately via --probe.""" - -import argparse -from contextlib import contextmanager -import hashlib -import json -import os -from pathlib import Path -import signal -import statistics -import subprocess -import sys -import threading -import time - -from . import http_client as client - - -@contextmanager -def service(args, directory, enabled): - directory.mkdir(parents=True, exist_ok=False) - if int((client.DEVICE / "mem_info_vram_used").read_text()) > 256 * 2**20: - raise RuntimeError("Target GPU occupied; no service launched") - module = "benchmarks.ark.probe" if args.probe else "freetoken.cli" - argv = [ - "numactl", - "--interleave=" + args.memory_nodes, - sys.executable, - "-m", - module, - "serve", - "--model-path", - args.model_path, - "--dtype", - "bfloat16", - "--attention-backend", - "triton", - "--moe-strategy", - "offload", - "--moe-cpu-threads", - "17", - "--moe-cache-mem", - "20GiB", - "--moe-cache-scope", - "unified", - "--moe-cache-split", - "equal", - "--num-tokens", - "33058", - "--page-size", - "1", - "--max-seq-len-override", - "33024", - "--max-prefill-length", - "8192", - "--max-running-requests", - "1", - "--cuda-graph-max-bs", - str(args.graph), - "--max-output-tokens", - "256", - "--memory-ratio", - "0.9", - "--expert-load", - "serial", - "--cache-type", - "naive", - "--enable-cache-report", - "--sampling-defaults", - "none", - "--reasoning-parser", - "off", - "--host", - "127.0.0.1", - "--port", - "19673", - ] - env = dict( - os.environ, - FREETOKEN_USE_ARK=str(enabled), - FREETOKEN_ARK_POLICY="validated", - OMP_NUM_THREADS="1", - MKL_NUM_THREADS="1", - OPENBLAS_NUM_THREADS="1", - PYTHONUNBUFFERED="1", - ) - for key in ("PYTHONPATH", "FREETOKEN_ARK_TRACE", "FREETOKEN_ARK_PROFILE"): - env.pop(key, None) - if args.profile: - env["FREETOKEN_ARK_PROFILE"] = str(args.profile.resolve()) - if args.probe: - env["FREETOKEN_ARK_TRACE"] = str(directory / "dispatch-{pid}.json") - env["ARK_MODEL_PROBE"] = str(directory) - client.save(directory / "command.json", dict(argv=argv, env={k: v for k, v in env.items() if k.startswith(("FREETOKEN_", "ROCR_"))})) - stop = threading.Event() - with (directory / "server.log").open("w") as out: - proc = subprocess.Popen(argv, env=env, stdout=out, stderr=subprocess.STDOUT, start_new_session=True) - monitor = threading.Thread(target=client.monitor, args=(proc.pid, stop, directory / "resources.jsonl"), daemon=True) - monitor.start() - start = time.monotonic() - try: - while True: - if proc.poll() is not None: - raise RuntimeError(f"server exited {proc.returncode}: {directory}") - if time.monotonic() - start > 600: - raise TimeoutError("startup timeout") - try: - health = client.get("/health") - except OSError: - time.sleep(1) - continue - if health.get("status") == "error": - raise RuntimeError(health) - if health.get("status") == "ok": - break - time.sleep(1) - print("READY", directory.name, flush=True) - yield - client.save(directory / "complete.json", dict(success=True, stats=client.get("/v1/stats"))) - finally: - owned = client.tree(proc.pid) - if proc.poll() is None: - os.killpg(proc.pid, signal.SIGTERM) - try: - proc.wait(30) - except subprocess.TimeoutExpired: - os.killpg(proc.pid, signal.SIGKILL) - proc.wait(10) - for pid in owned - {proc.pid}: - try: - os.kill(pid, signal.SIGKILL) - except ProcessLookupError: - pass - stop.set() - monitor.join(5) - deadline = time.monotonic() + 180 - while int((client.DEVICE / "mem_info_vram_used").read_text()) > 256 * 2**20: - if time.monotonic() > deadline: - raise RuntimeError("GPU did not return to idle") - time.sleep(0.5) - client.save( - directory / "shutdown.json", dict(returncode=proc.returncode, vram=int((client.DEVICE / "mem_info_vram_used").read_text())) - ) - - -def make_prompts(root, model): - from transformers import AutoTokenizer - - tok = AutoTokenizer.from_pretrained(model) - head = "<|im_start|>user\nRead the following archive.\n" - tail = ( - "\nExplain how inspectors compare transport, energy and weather records. Give a detailed answer.<|im_end|>\n<|im_start|>assistant\n" - ) - body = ( - "The archive tracks transport, energy, weather and equipment inspections. Inspectors compare observations from different stations.\n" - * 2000 - ) - head_ids, tail_ids = (tok.encode(x, add_special_tokens=False) for x in (head, tail)) - body_ids = tok.encode(body, add_special_tokens=False) - for n in (512, 32768): - path = root / f"prompt-{n}.txt" - if not path.exists(): - budget = n - len(head_ids) - len(tail_ids) - for _ in range(32): - text = tok.decode(head_ids + body_ids[:budget] + tail_ids) - count = len(tok.encode(text, add_special_tokens=False)) - if count == n: - break - budget += n - count - else: - raise ValueError("prompt token roundtrip failed") - path.write_text(text) - assert len(tok.encode(path.read_text(), add_special_tokens=False)) == n - client.save(root / "prompts.json", {str(n): hashlib.sha256((root / f"prompt-{n}.txt").read_bytes()).hexdigest() for n in (512, 32768)}) - - -def summarize(root): - samples = [] - for path in sorted(root.glob("round-*/measured-*.json")): - row = json.loads(path.read_text()) - assert row["success"] - samples.append({k: v for k, v in row.items() if k not in ("events", "text", "request")}) - groups = [] - for n in (512, 32768): - for enabled in (0, 1): - rows = [r for r in samples if r["context"] == n and r["enabled"] == enabled] - if rows: - groups.append( - dict( - context=n, - enabled=enabled, - n=len(rows), - decode_tok_s=statistics.median(r["decode_tok_s"] for r in rows), - ttft_s=statistics.median(r["ttft_s"] for r in rows), - e2e_s=statistics.median(r["e2e_s"] for r in rows), - ) - ) - client.save(root / "summary.json", dict(samples=samples, groups=groups)) - - -def main(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--model-path", required=True) - p.add_argument("--output", type=Path, required=True) - p.add_argument("--gpu-pci", required=True) - p.add_argument("--memory-nodes", required=True) - p.add_argument("--rounds", type=int, default=3) - p.add_argument("--start-round", type=int, default=1) - p.add_argument("--graph", type=int, default=0) - p.add_argument("--probe", action="store_true") - p.add_argument("--quality-only", action="store_true", help="run corrected short generation checks without context probes") - p.add_argument("--profile", type=Path, help="explicit candidate profile for separate experiments") - a = p.parse_args() - if a.quality_only and not a.probe: - p.error("--quality-only requires --probe") - a.output = a.output.resolve() - a.output.mkdir(parents=True, exist_ok=True) - client.ROOT = a.output - client.MODEL = a.model_path - client.DEVICE = Path("/sys/bus/pci/devices") / a.gpu_pci - make_prompts(a.output, a.model_path) - for r in range(a.start_round, a.start_round + a.rounds): - for enabled in (0, 1) if r % 2 else (1, 0): - directory = a.output / f"round-{r}-ark{enabled}" - with service(a, directory, enabled): - for n in () if a.quality_only else ((512, 32768) if r % 2 else (32768, 512)): - if not a.probe: - client.request(directory, n, "warmup") - row = client.request(directory, n, "probe" if a.probe else "measured", output_tokens=8 if a.probe else 256) - row.update(round=r, enabled=enabled) - client.save(directory / f"{row['phase']}-{n}.json", row) - if a.probe: - counts = [row for path in directory.glob("dispatch-*.json") for row in json.loads(path.read_text())["counts"]] - client.save(directory / f"hits-{n}.json", counts) - if not a.probe: - summarize(a.output) - print("RESULT", directory.name, n, row["decode_tok_s"], row["ttft_s"], flush=True) - if a.probe: - from transformers import AutoTokenizer - from .quality_prompts import CASES, first_answer - - tokenizer = AutoTokenizer.from_pretrained(a.model_path) - quality = [] - for name, question, predicate in CASES: - prompt = tokenizer.apply_chat_template( - [{"role": "user", "content": question}], - tokenize=False, - add_generation_prompt=True, - enable_thinking=False, - ) - n = len(tokenizer.encode(prompt, add_special_tokens=False)) - result = client.request(directory, n, "quality-" + name, output_tokens=128, prompt=prompt) - answer = first_answer(result["text"]) - quality.append(dict(name=name, answer=answer, passed=bool(predicate(answer)))) - client.save(directory / "quality.json", quality) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/ark/probe.py b/benchmarks/ark/probe.py deleted file mode 100644 index f734efa5c9..0000000000 --- a/benchmarks/ark/probe.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Untimed diagnostic server: real dispatch, GPU symbols and same-input shadows.""" - -import json -import os -from pathlib import Path - -import torch - -from freetoken.engine.engine import Engine -from freetoken.kernel.ark.common.trace import flush -from freetoken.kernel.ark.kernels import ArkFp8LinearKernel, ArkFp8MoEKernel -from freetoken.layers.quantization.linear.fp8_block import TritonFp8BlockLinearKernel -from freetoken.layers.quantization.moe.fp8_block import TritonFp8BlockMoEKernel -from .timing import metrics, assert_close - -ROOT = Path(os.environ["ARK_MODEL_PROBE"]) -SEEN = set() -ROWS = [] -ACTIVE = False -PROFILED = False - - -def shadow(op, key, actual, expected): - error = metrics(actual, expected) - assert_close(error, op) - ROWS.append(dict(operator=op, key=str(key), **error)) - (ROOT / "shadows.json").write_text(json.dumps(ROWS, indent=2) + "\n") - - -_linear = ArkFp8LinearKernel.apply - - -def linear(self, layer, x): - result = _linear(self, layer, x) - key = ("dense_fp8_decode", id(layer)) - if ACTIVE and x.numel() == x.shape[-1] and key not in SEEN: - SEEN.add(key) - shadow(key[0], (tuple(x.shape), tuple(layer.weight.shape)), result, TritonFp8BlockLinearKernel().apply(layer, x)) - return result - - -_moe = ArkFp8MoEKernel.apply - - -def moe(self, layer, x, weights, ids, view, *, is_prefill): - result = _moe(self, layer, x, weights, ids, view, is_prefill=is_prefill) - key = ("moe_decode_fp8", layer.layer_id) - if ACTIVE and not is_prefill and x.shape[0] == 1 and key not in SEEN: - SEEN.add(key) - shadow( - key[0], - (layer.layer_id, tuple(x.shape), view.tensors["gate_up"].shape[0]), - result, - TritonFp8BlockMoEKernel().apply(layer, x, weights, ids, view, is_prefill=False), - ) - return result - - -ArkFp8LinearKernel.apply = linear -ArkFp8MoEKernel.apply = moe -_forward = Engine.forward_batch - - -def forward(self, batch, args): - global ACTIVE, PROFILED - ACTIVE = True - try: - if not batch.is_prefill and not PROFILED: - PROFILED = True - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]) as trace: - result = _forward(self, batch, args) - torch.cuda.synchronize() - trace.export_chrome_trace(str(ROOT / "decode-trace.json")) - else: - result = _forward(self, batch, args) - flush() - return result - finally: - ACTIVE = False - - -Engine.forward_batch = forward - -if __name__ == "__main__": - from freetoken.cli import main - - main() diff --git a/benchmarks/ark/quality_prompts.py b/benchmarks/ark/quality_prompts.py deleted file mode 100644 index 097a48626e..0000000000 --- a/benchmarks/ark/quality_prompts.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Small explicit generation gates; not a standardized accuracy benchmark.""" - -import json -import re - - -def first_answer(text): - for marker in ("<|im_end|>", "<|endoftext|>", "<|im_start|>"): - text = text.split(marker, 1)[0] - if "" in text and "" not in text: - return "" - if "" in text: - text = text.split("", 1)[1] - return text.strip() - - -def json_ok(text): - match = re.search(r"\{[^{}]*\}", text) - if not match: - return False - try: - return json.loads(match.group()) == {"name": "test", "count": 3} - except json.JSONDecodeError: - return False - - -CASES = [ - ( - "moe_cn", - "请用两句话解释什么是混合专家（MoE）模型。", - lambda t: "专家" in t and any(w in t for w in ("路由", "激活", "选择", "门控", "分配", "决定")), - ), - ("addition", "What is 137 + 286? Reply with the number only.", lambda t: "423" in t), - ("division", "What is 144 divided by 12? Reply with the number only.", lambda t: "12" in t), - ( - "python", - "Write a Python function named add that returns the sum of two arguments. Output code only.", - lambda t: "def add" in t and "return" in t and "+" in t, - ), - ("sky", "In one sentence, explain why the sky appears blue.", lambda t: "blue" in t.lower() and "scatter" in t.lower()), - ("json", 'Return only a JSON object with keys name and count, using name "test" and count 3.', json_ok), - ("capital_cn", "中国的首都是哪里？只回答城市名称。", lambda t: "北京" in t), - ("decimal", "Which is larger, 0.9 or 0.11? Answer with just the larger number.", lambda t: "0.9" in t), -] diff --git a/benchmarks/ark/timing.py b/benchmarks/ark/timing.py index 584f23f143..b4865d5524 100644 --- a/benchmarks/ark/timing.py +++ b/benchmarks/ark/timing.py @@ -2,8 +2,8 @@ import statistics import time -import torch +import torch OPS = ["moe_decode_fp8", "dense_fp8_decode"] diff --git a/docs/ark-fp8-decode.md b/docs/ark-fp8-decode.md index 7e91f12752..817bcfa921 100644 --- a/docs/ark-fp8-decode.md +++ b/docs/ark-fp8-decode.md @@ -1,200 +1,72 @@ -# Optional ARK FP8 decode kernels on Radeon - -`FREETOKEN_USE_ARK=1` adds native HIP candidates for block-FP8 dense and MoE -**decode** on Linux gfx1201. The default is `0`. Other operations retain FreeToken. - -The packaged profile is specific to R9700, Torch 2.14.0+rocm10.1.0, HIP -7.16.26385, Triton 3.8.0, the compiler, host platform and source fingerprints. -An unmatched environment or unlisted shape uses FT. This is not an all-Radeon -support declaration. Recalibrate on another software stack; do not strip the -environment check. - -## Integration - -`ArkFp8MoEKernel` and `ArkFp8LinearKernel` inherit the existing Triton kernels' -packing, layouts and CPU formats. MoE consumes the current `ExpertView` directly, -including resident, materialized and cache-slot banks. The FTW format remains -`fp8_block`; reverse lookup still returns the original `triton` kernel. CPU expert -execution, hybrid routing, shared experts, cache management and upstream locked -expert fetching are owned by FT. - -- MoE: BF16 activations, E4M3 weights, BF16 padded 128x128 block scales, - contiguous INT32 slots/FP32 routing weights, M=1 and plain SiLU. It is W8A16. - Other activations, router-weight placement, dtypes, shapes and prefill use FT. - Inputs are preserved and the BF16 result is a new tensor. -- Dense: flattened M=1 uses W8A16, preserving leading dimensions and optional - bias. M>1, including batched decode, remains the original W8A8 FT path. - Unsupported layouts, scales and shapes use FT. +# 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 - contract on hybrid routes. Native negative slots are covered only by direct - boundary tests; the default integration does not advertise that capability. -- The original primitives are independent fallback targets. No FT function is - overwritten by the production adapters. -- Every MoE Graph shape remains FT. A prior Qwen3 matched-prefix model probe - exceeded its 0.05-nat KL gate with native MoE Graph; operator tolerance alone - was insufficient. Native Graph correctness tests are separate from permission - to enable it in a model. Dense Graph also remains FT: its three-round model A/B did not demonstrate a benefit. - Native Graph correctness tests are retained for future qualification. + routing. Explicit requests for the `triton` backend retain FT. + +## Build and selection -## Build and use +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. -In an existing ROCm FreeToken environment, use `uv pip install -e .` (or install -the wheel). HIP/C++ sources and the profile are included in the package. -The first enabled use builds the two HIP translation units and Torch bridge in -`$FREETOKEN_ARK_CACHE_DIR`, or `$XDG_CACHE_HOME/freetoken/ark`. Build keys include -native source, build recipe, compiler version and Torch/HIP versions. The exact -ARK Python/native and baseline source fingerprints also gate the profile. -No experiment directory, external `common` package or prebuilt ARK library is -required. The engine loads the library before Graph capture. +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 -FREETOKEN_USE_ARK=0 ft serve --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --attention-backend triton --moe-strategy offload --cuda-graph-max-bs 0 ``` -Expose one R9700 (the profile requires one visible GPU). - -Existing explicit FT backend requests, e.g. `moe.fp8_block=triton`, keep selecting -FT. `FREETOKEN_ARK_TRACE=/tmp/dispatch-{pid}.json` writes operator, shape, backend, -mode and Python call counts. Graph capture counts are not replay counts; inspect -a GPU trace for replay execution. Formal performance runs omit this tracing. +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. -## Reproduce validation and profile selection +## 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 -python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-eager --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 3 -python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-graph-candidate --profile docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 3 --graph 1 -python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-probe --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 1 --probe -``` - -Use an idle GPU and appropriate CPU/memory affinity for your host. The calibration -command intentionally forces native candidates for measurement. Serving uses -`FREETOKEN_ARK_POLICY=validated` (default); `native` is a calibration/debug override, -not a validated serving policy. `baseline` forces FT fallbacks. An external -`FREETOKEN_ARK_PROFILE` can hold a separately measured table. The generator keeps all Graph entries disabled by default. `--graph-candidates` -can emit dense-only microbenchmark candidates for a separate model evaluation; -it does not qualify them for production. MoE Graph always remains disabled. - -The tests cover numerical error, dtype/shape/alias and input preservation, -repeated/zero/inactive routes, padded scale strides, finite E4M3 codes, bias, -leading dimensions, non-default streams, Graph replay with changed inputs/routes, -actual method selection/packing/Kernel.apply, FTW compatibility, explicit FT -selection, disabled mode and environment mismatch. - -The model harness uses fresh off/on services, alternating off/on order across rounds, -warmup per context, no prefix reuse, 512/32768 input tokens and fixed 256 output -tokens. Its server command, SSE events, usage, per-round metrics and resources -are saved. Decode speed is 255 divided by the span between first and last nonempty -text events; if events differ from tokens, the record labels this a client timing -estimate. TTFT includes request/frontend work. These are concurrency-one client -measurements, not aggregate throughput. - -`--probe` is separate and untimed: it records real GPU kernel symbols, dispatch -counts, same-input per-layer comparisons and eight bounded generation checks. -It is a diagnostic wrapper; production code does not monkeypatch FT methods. -Numerical tolerance and bounded generation checks are not a standardized model -accuracy certification. - -Source provenance: [AMD-AIM/amd_radeon_kernels](https://github.com/AMD-AIM/amd_radeon_kernels) -revision `975e8cd9918ec1365d3f275d1259d215315f505a`; see the packaged ARK `NOTICE`. -Related: [AMD roadmap #541](https://github.com/FlashML-org/FreeToken/issues/541). - -## This calibration (2026-10-21) - -Fresh build against upstream `fe5a5ec9c5644e266527de5afb9f1156b9bdcabb`. -GPU `23:00.0` is R9700 32 GiB, PCIe 32 GT/s x16. CPU affinity selects one -EPYC 9334 socket (`0-31,64-95`) of a dual-socket host, memory interleave `2-3`. -Ubuntu 24.04.1, kernel 7.0.0-38, amdgpu 7.1.9.31600000; remaining environment -fingerprints are in [runtime.json](benchmarks/ark-fp8-decode-20261021/runtime.json). - -Each entry retains five alternating-order paired repeats. Eager timings include -Python dispatch, checks, allocation and the complete GPU operation. Graph timings -include the captured operation, using GPU events over replay; Python capture cost -is excluded. No activation quantization is omitted: these M=1 paths are W8A16. - -| Operator / shape | Eager FT / native, us | Eager speedup | Graph FT / native, us | Deployed Graph selected? | -|---|---:|---:|---:|---| -| dense_fp8_decode / N5120_K2048 | 42.737 / 14.336 | 2.981x | 15.091 / 12.952 | FT (model veto) | -| dense_fp8_decode / N2048_K4096 | 43.586 / 16.359 | 2.664x | 13.675 / 16.654 | FT | -| dense_fp8_decode / N151936_K2048 | 581.642 / 500.559 | 1.162x | 580.391 / 499.153 | FT (model veto) | -| moe_decode_fp8 / C32_H2048_I768 | 105.468 / 45.588 | 2.314x | 53.397 / 44.586 | FT | -| moe_decode_fp8 / C4550_H2048_I768 | 104.505 / 43.965 | 2.377x | 54.333 / 40.295 | FT | - -All five cases passed the numerical gate (NRMSE < 0.03, cosine > 0.999 and a -bounded maximum error). Observed NRMSE was 0 for QKV, output projection and both -MoE bank sizes; the lm-head-shaped case was 0.00000534. The tests include bias and -boundary cases separately. The lm-head shape is an operator benchmark, not a -claim that this model's lm head actually selected native FP8. - -[Raw calibration](benchmarks/ark-fp8-decode-20261021/calibration.json) includes -all repeats and the negative Graph result. Its `enabled` fields are the initial -operator-level candidates; subsequent model evaluation vetoed all Graph entries. -The packaged profile is the deployed decision table. No speedups from the original -experimental report are reused. A different R9700 was used because the prior -measurement GPU was busy; the new off/on model comparison uses this same card. - -## Model results and real-call evidence - -Three-round medians, fixed 256 output tokens, same model/prompt/settings per pair. -Baseline uses the original FT compute path with `FREETOKEN_USE_ARK=0`. -Formal timing excludes instrumentation. The Graph rows evaluate the **rejected -candidate profile**, not the deployed all-FT Graph policy. - -| Mode | Input tokens | FT / ARK tok/s | Change | Paired wins | FT / ARK TTFT, s | -|---|---:|---:|---:|---:|---:| -| eager | 512 | 29.454 / 34.838 | +18.28% | 3/3 | 0.6517 / 0.6866 | -| eager | 32768 | 28.366 / 36.337 | +28.10% | 3/3 | 10.8668 / 10.8626 | -| Graph candidate (rejected) | 512 | 100.180 / 100.099 | -0.08% | 1/3 | 0.7239 / 0.7942 | -| Graph candidate (rejected) | 32768 | 65.221 / 64.977 | -0.37% | 0/3 | 10.9913 / 10.9634 | - -**The recommended enabled scope is eager. Both deployed Graph tables are empty.** -512-token Graph samples vary between about 53 and 101 tok/s; 32K candidate -Graph was slightly slower in all three pairs. No Graph model gain is claimed. -The initial operator-only candidate profile is preserved in -[graph-candidate-profile.json](benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json). -The model harness's `--profile` flag reproduces that candidate without changing -the packaged production table. - -- Final focused and related tests: **41 passed, 42 skipped** (23 ARK cases; - skips mostly require other checkpoint configurations/weights). -- Real model eager trace contains `dense_gemv_fp8_rows<4>`, `gate_up` - and `down_reduce`. In each context's eight-token diagnostic request, - seven decode steps hit native dense **672** times and native MoE **336** times. - Native shapes are dense `(M,N,K)=(1,5120,2048)/(1,2048,4096)` and MoE - `(M,H,I,bank_rows,top_k)=(1,2048,768,4550,8)`. -- **144** real-input per-layer shadows passed, maximum NRMSE **0.000734784**. -- Deployed Graph replay probe: **zero ARK GPU kernels**, and eight-token outputs match - FT exactly for both input contexts. The server's eager warmup still hits native - kernels; these are not Graph capture/replay hits. Full fallback is not counted as ARK gain. -- Corrected bounded generation suite: FT **7/8**, ARK **7/8**, no new failure. - Both incorrectly answer `0.11` for the decimal comparison. Initial 64-token - thinking-mode probes were truncated and excluded; the scored run uses the - checkpoint's non-thinking template, 128 output tokens and the first answer. -- Standalone wheel passed with `python -I`, the repo removed from `sys.path`, - new native/Triton caches and validated policy. All 17 ARK packaged files match - this tree; the enabled dense kernel was hit and NRMSE was **0**. - -[Per-round model samples](benchmarks/ark-fp8-decode-20261021/model-samples.csv) -and [validation evidence](benchmarks/ark-fp8-decode-20261021/validation.json) -retain positive/negative outcomes and exact native symbol names. Complete raw -SSE, resource records, traces and layer/answer comparisons are produced by the -commands above. Validate raw timing with `python -m benchmarks.ark.audit PATH`; -for the generation checks use: - -```bash -python -m benchmarks.ark.model_ab --model-path /path/to/Qwen3-30B-A3B-FP8-FTW --output /tmp/ark-quality --gpu-pci 0000:23:00.0 --memory-nodes 2-3 --rounds 1 --probe --quality-only -python -m benchmarks.ark.audit_quality /tmp/ark-quality ``` -The CPU is shared; all slow samples are retained. FT and native generate different -continuations, changing subsequent expert-cache traffic. These are configuration -A/B measurements, not fixed-token kernel attribution or a universal speed guarantee. -Prefill compute is unchanged, so no TTFT acceleration is claimed. Only this FP8 -Qwen3 offload model was measured end to end. Full-model resident/hybrid performance, -other checkpoints, standardized accuracy and NVIDIA execution remain unverified; -operator/bank tests do not substitute for those results. +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/docs/benchmarks/ark-fp8-decode-20261021/calibration.json b/docs/benchmarks/ark-fp8-decode-20261021/calibration.json deleted file mode 100644 index b7105bd9b9..0000000000 --- a/docs/benchmarks/ark-fp8-decode-20261021/calibration.json +++ /dev/null @@ -1,320 +0,0 @@ -{ - "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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", - "device_count": 1, - "platform": "Linux-7.0.0-38-generic-x86_64-with-glibc2.39", - "compiler_sha256": "449f6779a4041275d18d801246d0c2c09ffb453a48093d4bebbbd82803954205", - "baseline_sha256": "9815bbdde6aa7c596020f7c1ec5c6e5ea3b363697877a6a052c833e34d74b3af" - }, - "rows": [ - { - "operator": "dense_fp8_decode", - "case": "N5120_K2048", - "metadata": { - "M": 1, - "N": 5120, - "K": 2048 - }, - "correctness": { - "normalized_rmse": 0.0, - "cosine": 0.9999998211860657, - "max_abs": 0.0, - "max_ref": 1.8515625, - "finite": true - }, - "eager": { - "baseline_us": 42.73695987649262, - "candidate_us": 14.336119929794222, - "speedup": 2.981068802840718, - "samples_us": { - "baseline": [ - 43.779159896075726, - 42.98155981814489, - 42.46317985234782, - 42.73695987649262, - 41.98615992208943 - ], - "candidate": [ - 14.82991996454075, - 14.136340178083628, - 14.336119929794222, - 14.085119764786214, - 14.75892000598833 - ] - }, - "enabled": true - }, - "graph": { - "baseline_us": 15.090957880020142, - "candidate_us": 12.952004671096802, - "speedup": 1.1651445674426404, - "samples_us": { - "baseline": [ - 15.374047756195068, - 15.15425443649292, - 15.090957880020142, - 13.583890199661255, - 13.605790138244629 - ], - "candidate": [ - 13.392394781112671, - 13.354895114898682, - 12.952004671096802, - 12.55481481552124, - 12.526215314865112 - ] - }, - "enabled": true - } - }, - { - "operator": "dense_fp8_decode", - "case": "N2048_K4096", - "metadata": { - "M": 1, - "N": 2048, - "K": 4096 - }, - "correctness": { - "normalized_rmse": 0.0, - "cosine": 1.0, - "max_abs": 0.0, - "max_ref": 1.7734375, - "finite": true - }, - "eager": { - "baseline_us": 43.58615988167003, - "candidate_us": 16.359120199922472, - "speedup": 2.664333983063257, - "samples_us": { - "baseline": [ - 43.6879598419182, - 43.26476017013192, - 42.61715977918357, - 43.58615988167003, - 43.689539888873696 - ], - "candidate": [ - 16.41471986658871, - 15.847500180825591, - 16.84950024355203, - 16.359120199922472, - 16.238120151683688 - ] - }, - "enabled": true - }, - "graph": { - "baseline_us": 13.674784898757935, - "candidate_us": 16.653517484664917, - "speedup": 0.8211349290833067, - "samples_us": { - "baseline": [ - 13.933980464935303, - 13.843982219696045, - 13.674784898757935, - 12.418814897537231, - 12.446115016937256 - ], - "candidate": [ - 16.684716939926147, - 17.051807641983032, - 16.653517484664917, - 16.02343201637268, - 16.268327236175537 - ] - }, - "enabled": false - } - }, - { - "operator": "dense_fp8_decode", - "case": "N151936_K2048", - "metadata": { - "M": 1, - "N": 151936, - "K": 2048 - }, - "correctness": { - "normalized_rmse": 5.340938969311537e-06, - "cosine": 1.0, - "max_abs": 0.0009765625, - "max_ref": 2.171875, - "finite": true - }, - "eager": { - "baseline_us": 581.6417402820662, - "candidate_us": 500.55860017891973, - "speedup": 1.1619853101598177, - "samples_us": { - "baseline": [ - 582.4773397762328, - 581.5167599939741, - 581.6417402820662, - 581.8621598882601, - 580.6271600886248 - ], - "candidate": [ - 500.55860017891973, - 500.20219990983605, - 501.82079983642325, - 500.08219986921176, - 500.90499978978187 - ] - }, - "enabled": true - }, - "graph": { - "baseline_us": 580.3911972045898, - "candidate_us": 499.1529083251953, - "speedup": 1.1627523100125228, - "samples_us": { - "baseline": [ - 579.3439483642578, - 588.4440231323242, - 580.3911972045898, - 580.3480911254883, - 589.8081588745117 - ], - "candidate": [ - 499.1529083251953, - 499.552001953125, - 499.11949157714844, - 499.1862487792969, - 499.14161682128906 - ] - }, - "enabled": true - } - }, - { - "operator": "moe_decode_fp8", - "case": "C32_H2048_I768", - "metadata": { - "M": 1, - "H": 2048, - "I": 768, - "bank_rows": 32, - "top_k": 8 - }, - "correctness": { - "normalized_rmse": 0.0, - "cosine": 0.9999999403953552, - "max_abs": 0.0, - "max_ref": 0.23046875, - "finite": true - }, - "eager": { - "baseline_us": 105.46781995799392, - "candidate_us": 45.58776010526344, - "speedup": 2.3135117784788224, - "samples_us": { - "baseline": [ - 105.63260002527386, - 107.13299998315051, - 100.50904005765915, - 105.46781995799392, - 100.46084003988653 - ], - "candidate": [ - 48.84213994955644, - 47.465919924434274, - 45.58776010526344, - 43.56513993116096, - 43.68355992482975 - ] - }, - "enabled": true - }, - "graph": { - "baseline_us": 53.39667797088623, - "candidate_us": 44.586381912231445, - "speedup": 1.197600605404536, - "samples_us": { - "baseline": [ - 57.948174476623535, - 52.832889556884766, - 53.39667797088623, - 53.42008113861084, - 53.18868160247803 - ], - "candidate": [ - 44.586381912231445, - 44.8009729385376, - 44.05508995056152, - 44.63818073272705, - 44.17378902435303 - ] - }, - "enabled": false - } - }, - { - "operator": "moe_decode_fp8", - "case": "C4550_H2048_I768", - "metadata": { - "M": 1, - "H": 2048, - "I": 768, - "bank_rows": 4550, - "top_k": 8 - }, - "correctness": { - "normalized_rmse": 0.0, - "cosine": 0.9999999403953552, - "max_abs": 0.0, - "max_ref": 0.216796875, - "finite": true - }, - "eager": { - "baseline_us": 104.50481990119442, - "candidate_us": 43.96515985717997, - "speedup": 2.3769916961675213, - "samples_us": { - "baseline": [ - 104.50481990119442, - 106.70819989172742, - 102.2930201725103, - 106.53300007106736, - 100.65804002806544 - ], - "candidate": [ - 44.85156008740887, - 43.23495988501236, - 43.75496006105095, - 43.96515985717997, - 44.09874003613368 - ] - }, - "enabled": true - }, - "graph": { - "baseline_us": 54.33335304260254, - "candidate_us": 40.294880867004395, - "speedup": 1.348393440395889, - "samples_us": { - "baseline": [ - 54.33335304260254, - 337.1195983886719, - 54.96233940124512, - 47.28722095489502, - 47.352614402770996 - ], - "candidate": [ - 40.70597171783447, - 39.989686012268066, - 40.99296569824219, - 40.294880867004395, - 39.54289674758911 - ] - }, - "enabled": false - } - } - ] -} diff --git a/docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json b/docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json deleted file mode 100644 index a746d2cc8d..0000000000 --- a/docs/benchmarks/ark-fp8-decode-20261021/graph-candidate-profile.json +++ /dev/null @@ -1,61 +0,0 @@ -{ - "_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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", - "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, - "graph_min_speedup": 1.03, - "each_repeat_faster": true - }, - "_quality_exclusions": { - "moe_decode_fp8": "All Graph shapes remain FT; Qwen3 matched-prefix KL previously exceeded 0.05 nats." - }, - "dense_fp8_decode": { - "eager": { - "1,5120,2048": { - "speedup": 2.981068802840718, - "min_paired_speedup": 2.8447989354948624 - }, - "1,2048,4096": { - "speedup": 2.664333983063257, - "min_paired_speedup": 2.5292833118591935 - }, - "1,151936,2048": { - "speedup": 1.1619853101598177, - "min_paired_speedup": 1.1590626384391836 - } - }, - "graph": { - "1,5120,2048": { - "speedup": 1.1651445674426404, - "min_paired_speedup": 1.0819665920414685 - }, - "1,151936,2048": { - "speedup": 1.1627523100125228, - "min_paired_speedup": 1.1606542578468129 - } - } - }, - "moe_decode_fp8": { - "eager": { - "1,2048,768,32,8": { - "speedup": 2.3135117784788224, - "min_paired_speedup": 2.162734887012934 - }, - "1,2048,768,4550,8": { - "speedup": 2.3769916961675213, - "min_paired_speedup": 2.2825604528743484 - } - }, - "graph": {} - } -} diff --git a/docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv b/docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv deleted file mode 100644 index c80d87970a..0000000000 --- a/docs/benchmarks/ark-fp8-decode-20261021/model-samples.csv +++ /dev/null @@ -1,25 +0,0 @@ -mode,round,enabled,context,decode_tok_s,ttft_s,e2e_s,token_text_chunks,text_sha256 -eager,1,0,32768,28.366052932428786,10.775955886012525,19.766115287013236,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb -eager,1,0,512,27.465098474060973,0.8498681829951238,10.134956226000213,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 -eager,1,1,32768,36.337070449318986,10.830030515004182,17.849467851003283,256,5c20deafdd17c988f9da13d473d60069d0b6eb9e854090b5884ead328e29181e -eager,1,1,512,30.28456095314772,0.9327965800039237,9.353914507999434,256,e809cdcb447a0670cefb38e065aaba2cc539f7b132e95fba4900045330fcc002 -eager,2,0,32768,31.473914371452555,10.866772359004244,18.970433360009338,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb -eager,2,0,512,31.61883942692908,0.6422161329974188,8.707675377998385,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 -eager,2,1,32768,35.719651721432726,13.993129288995988,21.133178561998648,256,5c20deafdd17c988f9da13d473d60069d0b6eb9e854090b5884ead328e29181e -eager,2,1,512,37.675944950872434,0.6477360019925982,7.416634577995865,256,e809cdcb447a0670cefb38e065aaba2cc539f7b132e95fba4900045330fcc002 -eager,3,0,32768,20.53616960530507,10.875875447993167,23.294661559993983,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb -eager,3,0,512,29.4541884296002,0.6517329810012598,9.309694625990232,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 -eager,3,1,32768,36.72163966356929,10.862550713005476,17.808452309996937,256,5c20deafdd17c988f9da13d473d60069d0b6eb9e854090b5884ead328e29181e -eager,3,1,512,34.83756625429643,0.6866280950052897,8.006714333008858,256,e809cdcb447a0670cefb38e065aaba2cc539f7b132e95fba4900045330fcc002 -graph_candidate_rejected,1,0,32768,65.22113636555925,10.602495545012061,14.513973229011754,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb -graph_candidate_rejected,1,0,512,56.33858678141364,0.7210316830023658,5.247661571993376,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 -graph_candidate_rejected,1,1,32768,64.9769935118381,10.927924178002286,14.853244415004156,256,1c1e3c7df3e702de994c988f2e2c8ac57f299713b9418c51de8e42756022cdc2 -graph_candidate_rejected,1,1,512,100.51418361740437,1.3024535959993955,3.840363162002177,256,555a5f1f33054486d0637db560b98685ec71a7d776fdd9736add0d9511f913a4 -graph_candidate_rejected,2,0,32768,65.32587621602467,14.445713420995162,18.350080314994557,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb -graph_candidate_rejected,2,0,512,100.88825964400858,0.7239236070017796,3.2523054859921103,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 -graph_candidate_rejected,2,1,32768,65.06508241655588,10.963409739997587,14.884385861005285,256,1c1e3c7df3e702de994c988f2e2c8ac57f299713b9418c51de8e42756022cdc2 -graph_candidate_rejected,2,1,512,100.09894710247522,0.7882882770063588,3.3363208860100713,256,555a5f1f33054486d0637db560b98685ec71a7d776fdd9736add0d9511f913a4 -graph_candidate_rejected,3,0,32768,65.14931377043827,10.991303474002052,14.9065101879969,256,098971317aa58a592a9c863251350d8d3c9c7e53c7a96bfcab43d808d3d579eb -graph_candidate_rejected,3,0,512,100.17950068652583,1.1008289820019854,3.647286549996352,256,5e35ae8c573a8ad2dacaf4ed9f12390a0695dde01d58cd4d8191dc951ab1e579 -graph_candidate_rejected,3,1,32768,64.84847037983543,11.109217436009203,15.042871676007053,256,1c1e3c7df3e702de994c988f2e2c8ac57f299713b9418c51de8e42756022cdc2 -graph_candidate_rejected,3,1,512,52.61622377409661,0.7942161740065785,5.641266263002763,256,555a5f1f33054486d0637db560b98685ec71a7d776fdd9736add0d9511f913a4 diff --git a/docs/benchmarks/ark-fp8-decode-20261021/runtime.json b/docs/benchmarks/ark-fp8-decode-20261021/runtime.json deleted file mode 100644 index 1b1ce2296c..0000000000 --- a/docs/benchmarks/ark-fp8-decode-20261021/runtime.json +++ /dev/null @@ -1,48 +0,0 @@ -{ - "base": "fe5a5ec9c5644e266527de5afb9f1156b9bdcabb", - "branch": "perf/ark-fp8-decode", - "gpu": { - "name": "AMD Radeon AI PRO R9700", - "pci": "0000:23:00.0", - "uuid": "GPU-98103ef4c22a3ccc", - "render": "renderD131", - "memory_gib": 32, - "link": "32 GT/s x16" - }, - "cpu": { - "model": "AMD EPYC 9334 32-Core Processor", - "host_sockets": 2, - "allowed_cpus": "0-31,64-95", - "memory_nodes": "2-3", - "exclusive": false - }, - "os": "Ubuntu 24.04.1 LTS", - "amdgpu": "7.1.9.31600000", - "model": "Qwen/Qwen3-30B-A3B-FP8", - "checkpoint_format": "FTW", - "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": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", - "device_count": 1, - "platform": "Linux-7.0.0-38-generic-x86_64-with-glibc2.39", - "compiler_sha256": "449f6779a4041275d18d801246d0c2c09ffb453a48093d4bebbbd82803954205", - "baseline_sha256": "9815bbdde6aa7c596020f7c1ec5c6e5ea3b363697877a6a052c833e34d74b3af" - }, - "native_cache_key": "b63ee64d5f192756348a", - "weight_dtype": "E4M3 FP8", - "activation_dtype": "BF16", - "kv_dtype": "BF16", - "expert_cache_gib": 20, - "expert_slots": 4550, - "input_contexts": [ - 512, - 32768 - ], - "output_tokens": 256, - "concurrency": 1, - "prefill_chunk": 8192 -} diff --git a/docs/benchmarks/ark-fp8-decode-20261021/validation.json b/docs/benchmarks/ark-fp8-decode-20261021/validation.json deleted file mode 100644 index a0d98fbad3..0000000000 --- a/docs/benchmarks/ark-fp8-decode-20261021/validation.json +++ /dev/null @@ -1,115 +0,0 @@ -{ - "tests": { - "passed": 41, - "skipped": 42, - "ark_cases": 23 - }, - "formal_requests": 24, - "numerical_shadow": { - "comparisons": 144, - "max_normalized_rmse": 0.0007347841165028512 - }, - "native_gpu_kernel_names": [ - "void ark_ft::down_reduce(float const*, void const*, __hip_bfloat16 const*, int const*, float const*, __hip_bfloat16*, int, int, int, int, int) [clone .kd]", - "void ark_ft::gate_up(__hip_bfloat16 const*, void const*, __hip_bfloat16 const*, int const*, float const*, float*, int, int, int, int, int) [clone .kd]", - "void radeon_ops::hip_native::dense_gemv_fp8_rows<4>(__hip_bfloat16 const*, unsigned char const*, __hip_bfloat16 const*, __hip_bfloat16*, int, int, int) [clone .kd]" - ], - "native_eager_calls_per_context": { - "512": { - "dense_fp8_decode": 672, - "moe_decode_fp8": 336 - }, - "32768": { - "dense_fp8_decode": 672, - "moe_decode_fp8": 336 - } - }, - "graph": { - "deployed_graph_capture_native_calls": 0, - "probe_outputs_match_ft": [ - 512, - 32768 - ], - "candidate_speedups_rejected": [ - { - "context": 512, - "rounds": [ - { - "round": 1, - "speedup": 1.7841090691072228, - "ft_decode_tok_s": 56.33858678141364, - "ark_decode_tok_s": 100.51418361740437, - "outputs_equal": false - }, - { - "round": 2, - "speedup": 0.9921763687438111, - "ft_decode_tok_s": 100.88825964400858, - "ark_decode_tok_s": 100.09894710247522, - "outputs_equal": false - }, - { - "round": 3, - "speedup": 0.5252194651951735, - "ft_decode_tok_s": 100.17950068652583, - "ark_decode_tok_s": 52.61622377409661, - "outputs_equal": false - } - ], - "ft_median": 100.17950068652583, - "ark_median": 100.09894710247522, - "median_change_percent": -0.08040924889680667, - "wins": 1, - "n": 3 - }, - { - "context": 32768, - "rounds": [ - { - "round": 1, - "speedup": 0.9962566911997248, - "ft_decode_tok_s": 65.22113636555925, - "ark_decode_tok_s": 64.9769935118381, - "outputs_equal": false - }, - { - "round": 2, - "speedup": 0.9960078025037679, - "ft_decode_tok_s": 65.32587621602467, - "ark_decode_tok_s": 65.06508241655588, - "outputs_equal": false - }, - { - "round": 3, - "speedup": 0.995382247744575, - "ft_decode_tok_s": 65.14931377043827, - "ark_decode_tok_s": 64.84847037983543, - "outputs_equal": false - } - ], - "ft_median": 65.22113636555925, - "ark_median": 64.9769935118381, - "median_change_percent": -0.374330880027518, - "wins": 0, - "n": 3 - } - ] - }, - "generation": { - "ft_passed": 7, - "ark_passed": 7, - "total": 8, - "new_failures": [], - "shared_failure": "decimal: both answered 0.11 instead of 0.9", - "method": "non-thinking tokenizer template, 128 generated tokens, score first answer only; initial 64-token thinking probes excluded; not standardized accuracy" - }, - "standalone_wheel": { - "passed": true, - "normalized_rmse": 0.0, - "sha256": "01a49d978b14496514141cc4645c1372c12542898ad7f6a2dd7ef30876f4d9dc", - "ark_packaged_files": 17, - "new_native_cache": true, - "validated_policy": true, - "isolated_python": true - } -} diff --git a/python/freetoken/kernel/ark/__init__.py b/python/freetoken/kernel/ark/__init__.py index 75654b704b..ddd9dc23f2 100644 --- a/python/freetoken/kernel/ark/__init__.py +++ b/python/freetoken/kernel/ark/__init__.py @@ -25,6 +25,7 @@ def candidate_kernels(method, candidates): return candidates from freetoken.layers.quantization.registry import LayerKind from freetoken.layers.quantization.scheme import QuantKind + from .kernels import ArkFp8LinearKernel, ArkFp8MoEKernel cls = { diff --git a/python/freetoken/kernel/ark/common/bindings.cpp b/python/freetoken/kernel/ark/common/bindings.cpp index ef81bdbe05..7373e522e6 100644 --- a/python/freetoken/kernel/ark/common/bindings.cpp +++ b/python/freetoken/kernel/ark/common/bindings.cpp @@ -5,7 +5,7 @@ #include extern "C" { -int radeon_dense_gemv_fp8(const void*, const void*, const void*, void*, int, int, int, int, void*); +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*); } @@ -20,7 +20,7 @@ void check(const Tensor& t, const Tensor& x, at::ScalarType dtype) { 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, int64_t nw) { +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); @@ -28,13 +28,13 @@ Tensor dense_decode(const Tensor& x, const Tensor& w, const Tensor& scale, int64 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 && nw == 1, + 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), nw, stream())); + N, K, scale.stride(0), stream())); return out; } @@ -66,7 +66,7 @@ Tensor moe_decode(const Tensor& x, const Tensor& w1, const Tensor& w2, const Ten } // namespace TORCH_LIBRARY(freetoken_ark, m) { - m.def("dense_decode(Tensor x, Tensor w, Tensor scale, int nw) -> Tensor"); + 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) { diff --git a/python/freetoken/kernel/ark/common/native.py b/python/freetoken/kernel/ark/common/native.py index 2864c34ca4..c27d58d8da 100644 --- a/python/freetoken/kernel/ark/common/native.py +++ b/python/freetoken/kernel/ark/common/native.py @@ -4,10 +4,10 @@ import hashlib import json import os -from pathlib import Path import shutil import subprocess import threading +from pathlib import Path ROOT = Path(__file__).resolve().parents[1] _LOADED = False @@ -80,7 +80,7 @@ def load(): 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, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + 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}") diff --git a/python/freetoken/kernel/ark/common/policy.py b/python/freetoken/kernel/ark/common/policy.py index 947073cdc5..c20be89aa9 100644 --- a/python/freetoken/kernel/ark/common/policy.py +++ b/python/freetoken/kernel/ark/common/policy.py @@ -1,12 +1,13 @@ """Fail-closed environment and measured-shape selection, separate for eager/Graph.""" +import hashlib import json import os -import hashlib import platform import subprocess from functools import lru_cache from pathlib import Path + import torch MODE = os.environ.get("FREETOKEN_ARK_POLICY", "validated") @@ -22,7 +23,7 @@ def current_environment(): if not torch.version.hip or not torch.cuda.is_available(): return {} - from .native import source_digest, _compiler + from .native import _compiler, source_digest p = torch.cuda.get_device_properties(torch.cuda.current_device()) try: diff --git a/python/freetoken/kernel/ark/common/trace.py b/python/freetoken/kernel/ark/common/trace.py deleted file mode 100644 index 833054d21c..0000000000 --- a/python/freetoken/kernel/ark/common/trace.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Optional Python dispatch census; Graph captures are distinct from replays.""" - -import atexit -from collections import Counter -import json -import os -from pathlib import Path -import torch - -_PATH = os.environ.get("FREETOKEN_ARK_TRACE") -_COUNTS = Counter() -_TOTAL = 0 - - -def record(op, shape, backend): - global _TOTAL - if _PATH: - mode = "capture" if torch.cuda.is_current_stream_capturing() else "eager" - item = (op, shape, backend, mode) - first = item not in _COUNTS - _COUNTS[item] += 1 - _TOTAL += 1 - if first or _TOTAL % 1000 == 0: - flush() - - -def flush(): - if _PATH: - path = Path(_PATH.replace("{pid}", str(os.getpid()))) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text( - json.dumps( - { - "pid": os.getpid(), - "counts": [dict(operator=o, shape=s, backend=b, phase=m, calls=n) for (o, s, b, m), n in sorted(_COUNTS.items())], - "note": "Python eager calls and Graph captures only; replay GPU execution needs profiler evidence.", - }, - indent=2, - ) - ) - - -atexit.register(flush) diff --git a/python/freetoken/kernel/ark/common/validated_shapes.json b/python/freetoken/kernel/ark/common/validated_shapes.json index ecc0c403d2..84c0ffe7a8 100644 --- a/python/freetoken/kernel/ark/common/validated_shapes.json +++ b/python/freetoken/kernel/ark/common/validated_shapes.json @@ -5,7 +5,7 @@ "hip": "7.16.26385", "gpu_name": "AMD Radeon AI PRO R9700", "arch": "gfx1201", - "source_sha256": "44b97b10bff974118c6e6aacf0f5e18d1bd3ea89d974ea73101f621334ce3a87", + "source_sha256": "27b45edd3f1fd2aa950ff6a6f1c714020e78f525f8c022b5ff46f727f896cf6b", "device_count": 1, "platform": "Linux-7.0.0-38-generic-x86_64-with-glibc2.39", "compiler_sha256": "449f6779a4041275d18d801246d0c2c09ffb453a48093d4bebbbd82803954205", @@ -13,26 +13,22 @@ }, "_rules": { "eager_min_speedup": 1.05, - "graph_min_speedup": 1.03, - "each_repeat_faster": true, - "graph_requires_model_benefit": true - }, - "_quality_exclusions": { - "moe_decode_fp8": "All Graph shapes remain FT; Qwen3 matched-prefix KL previously exceeded 0.05 nats." + "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.981068802840718, - "min_paired_speedup": 2.8447989354948624 + "speedup": 2.8881421083394208, + "min_paired_speedup": 2.258698800667121 }, "1,2048,4096": { - "speedup": 2.664333983063257, - "min_paired_speedup": 2.5292833118591935 + "speedup": 2.7925056711471514, + "min_paired_speedup": 2.7456354563614886 }, "1,151936,2048": { - "speedup": 1.1619853101598177, - "min_paired_speedup": 1.1590626384391836 + "speedup": 1.1631544640714144, + "min_paired_speedup": 1.1601689528473773 } }, "graph": {} @@ -40,17 +36,14 @@ "moe_decode_fp8": { "eager": { "1,2048,768,32,8": { - "speedup": 2.3135117784788224, - "min_paired_speedup": 2.162734887012934 + "speedup": 2.149956471331526, + "min_paired_speedup": 2.097327574739232 }, "1,2048,768,4550,8": { - "speedup": 2.3769916961675213, - "min_paired_speedup": 2.2825604528743484 + "speedup": 2.6720033324525043, + "min_paired_speedup": 2.6453511021618126 } }, "graph": {} - }, - "_model_exclusions": { - "dense_fp8_decode.graph": "Three-round Qwen3 candidate model A/B: 512 -0.0804%, 32768 -0.3743%; keep FT until model benefit is demonstrated." } } diff --git a/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py b/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py index 68e2f2aaa1..5b3addb1e1 100644 --- a/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py +++ b/python/freetoken/kernel/ark/dense_fp8_decode/adapter.py @@ -1,12 +1,10 @@ """M=1 FreeToken block_fp8_linear, including bias and leading dimensions.""" import torch -from freetoken.kernel.ark.common.trace import record -from freetoken.kernel.ark.common.policy import allowed, key + from freetoken.kernel.ark.common.contracts import gfx1201 from freetoken.kernel.ark.common.native import load - -NW = 1 +from freetoken.kernel.ark.common.policy import allowed, key def block_fp8_linear(x, weight, weight_scale, bias=None): @@ -31,11 +29,11 @@ def block_fp8_linear(x, weight, weight_scale, bias=None): and weight.device == x.device and weight_scale.device == x.device ): - record("dense_fp8_decode", key(x.numel() // k, weight.shape[0], k), "freetoken") - from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear as baseline + from freetoken.kernel.triton.fp8_block_linear import ( + block_fp8_linear as baseline, + ) return baseline(x, weight, weight_scale, bias) - record("dense_fp8_decode", key(x.numel() // k, weight.shape[0], k), "ark") load() - out = torch.ops.freetoken_ark.dense_decode(x.reshape(1, k), weight, weight_scale, NW).reshape(*x.shape[:-1], weight.shape[0]) + 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 index 689bf6bb46..c9712f92ca 100644 --- a/python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip +++ b/python/freetoken/kernel/ark/dense_fp8_decode/dense_gemv_fp8.hip @@ -33,58 +33,18 @@ __device__ __forceinline__ float dg_e4m3_to_f32(unsigned char v) { #endif } -// Split-K: NW warps per output row, each striding over a 1/NW share of the K-blocks (shorter serial -// reduction chain at large K), combined across warps via LDS. NW=1 is the plain wave-per-output path. -template -__global__ void __launch_bounds__(NW * 32) dense_gemv_fp8_k(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; - if (n >= N) return; - int warp = threadIdx.x >> 5, 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 = warp; kb < nkb; kb += NW) { - 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 (NW == 1) { - if (lane == 0) out[n] = (bf16)acc; - return; - } - __shared__ float red[NW]; - if (lane == 0) red[warp] = acc; - __syncthreads(); - if (threadIdx.x == 0) { - float s = 0.f; - #pragma unroll - for (int i = 0; i < NW; ++i) s += red[i]; - out[n] = (bf16)s; - } -} - 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 warp = 0, lane = threadIdx.x & 31; + 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 = warp; kb < nkb; kb += 1) { + 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); @@ -96,33 +56,13 @@ __global__ void __launch_bounds__(ROWS * 32) dense_gemv_fp8_rows(const bf16* __r } #pragma unroll for (int o = 16; o > 0; o >>= 1) acc += __shfl_xor(acc, o); // reduce this warp's lanes - if (1 == 1) { - if (lane == 0) out[n] = (bf16)acc; - return; - } - __shared__ float red[1]; - if (lane == 0) red[warp] = acc; - __syncthreads(); - if (threadIdx.x == 0) { - float s = 0.f; - #pragma unroll - for (int i = 0; i < 1; ++i) s += red[i]; - out[n] = (bf16)s; - } + 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, int nw, hipStream_t stream) { - (void)K; - // auto split-K by N: small N underfills the GPU at NW=1, so split each output across more warps; large - // N already has plenty of output waves -> NW=1. NOTE this op WINS at moderate N (e.g. GDN out_proj - // K=4096 N=2048 ~2x both arches) but LOSES to triton's tuned big-N GEMV at large N (e.g. N>=16384 on - // gfx1201) -- the caller should gate on N and keep triton for large-N projections (lm_head, in_proj). - if (nw <= 0) nw = (N <= 2048) ? 4 : (N <= 8192 ? 2 : 1); - dim3 grid((unsigned)N); - if (nw >= 4) hipLaunchKernelGGL((dense_gemv_fp8_k<4>), grid, dim3(128), 0, stream, x, w, wscale, out, N, K, wscale_stride); - else if (nw == 2) hipLaunchKernelGGL((dense_gemv_fp8_k<2>), grid, dim3(64), 0, stream, x, w, wscale, out, N, K, wscale_stride); - else hipLaunchKernelGGL((dense_gemv_fp8_rows<4>), dim3((N+3)/4), dim3(128), 0, stream, x, w, wscale, out, N, K, wscale_stride); + 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 @@ -131,9 +71,9 @@ void launch_dense_gemv_fp8(const bf16* x, const unsigned char* w, const bf16* ws // 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, int nw, void* stream) { + 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, nw, (hipStream_t)stream); + (__hip_bfloat16*)out, N, K, wscale_stride, (hipStream_t)stream); return (int)hipGetLastError(); } diff --git a/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py b/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py index 2c03d2887a..c5ff527412 100644 --- a/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py +++ b/python/freetoken/kernel/ark/moe_decode_fp8/adapter.py @@ -1,10 +1,14 @@ """Drop-in for freetoken.moe.fused_fp8_block.fused_experts_decode_fp8_block.""" import torch -from freetoken.kernel.ark.common.trace import record -from freetoken.kernel.ark.common.policy import allowed, key -from freetoken.kernel.ark.common.contracts import moe_supported, scales_supported, plain_silu + +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( @@ -32,12 +36,9 @@ def fused_experts_decode_fp8_block( and moe_supported(hidden_states, gate_up, down, topk_weights, topk_ids) and scales_supported(gate_up_scale, down_scale, gate_up, down) ): - record( - "moe_decode_fp8", - key(hidden_states.shape[0], hidden_states.shape[1], down.shape[2], gate_up.shape[0], topk_ids.shape[1]), - "freetoken", + from freetoken.moe.fused_fp8_block import ( + fused_experts_decode_fp8_block as baseline, ) - from freetoken.moe.fused_fp8_block import fused_experts_decode_fp8_block as baseline return baseline( hidden_states, @@ -52,6 +53,5 @@ def fused_experts_decode_fp8_block( act_alpha, act_limit, ) - record("moe_decode_fp8", key(hidden_states.shape[0], hidden_states.shape[1], down.shape[2], gate_up.shape[0], topk_ids.shape[1]), "ark") 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/tests/kernels/test_ark_fp8.py b/tests/kernels/test_ark_fp8.py index 65a51cda1b..b5193187e8 100644 --- a/tests/kernels/test_ark_fp8.py +++ b/tests/kernels/test_ark_fp8.py @@ -8,8 +8,8 @@ import pytest import torch -from benchmarks.ark.cases import rand, quant -from benchmarks.ark.timing import metrics, assert_close +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, @@ -21,8 +21,8 @@ def native(monkeypatch): if not torch.cuda.get_device_properties(0).gcnArchName.startswith("gfx1201"): pytest.skip("requires gfx1201") - from freetoken.kernel.ark.common import policy from freetoken.kernel.ark import initialize + from freetoken.kernel.ark.common import policy monkeypatch.setenv("FREETOKEN_USE_ARK", "1") monkeypatch.setattr(policy, "MODE", "native") @@ -58,17 +58,23 @@ def test_dense_bias_padding_alias(value, bias, monkeypatch): 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 = [] - monkeypatch.setattr(adapter, "record", lambda *a: hits.append(a)) + 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[-1][-1] == "ark" + assert hits == [True] @pytest.mark.parametrize("kind", ["batch", "strided", "disabled", "unknown_shape"]) def test_dense_fallback(kind, monkeypatch): - from freetoken.kernel.ark.dense_fp8_decode import adapter 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) @@ -81,14 +87,17 @@ def test_dense_fallback(kind, monkeypatch): monkeypatch.setattr(policy, "MODE", "validated") monkeypatch.setattr(policy, "environment_matches", lambda: True) monkeypatch.setattr(policy, "_profiles", {}) - hits = [] - monkeypatch.setattr(adapter, "record", lambda *a: hits.append(a)) + 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)) - assert hits[-1][-1] == "freetoken" 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.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() @@ -107,8 +116,8 @@ def test_moe_routes_zero_inactive_and_input_unchanged(): @pytest.mark.parametrize("kind", ["batch", "gelu", "disabled", "graph_excluded"]) def test_moe_fallback(kind, monkeypatch): - from freetoken.kernel.ark.moe_decode_fp8 import adapter 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) @@ -119,23 +128,32 @@ def test_moe_fallback(kind, monkeypatch): monkeypatch.setattr(policy, "MODE", "validated") monkeypatch.setattr(policy, "environment_matches", lambda: True) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - hits = [] - monkeypatch.setattr(adapter, "record", lambda *a: hits.append(a)) + 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)) - assert hits[-1][-1] == "freetoken" @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 + 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 + 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() @@ -159,11 +177,17 @@ def test_nondefault_stream_graph_and_updated_inputs(op): @pytest.mark.parametrize("resident", [False, True]) def test_real_method_bank_pack_apply_prefill(resident, monkeypatch): - from freetoken.kernel.ark.kernels import ArkFp8MoEKernel, ArkFp8LinearKernel - from freetoken.layers.quantization.moe.base import MoEConfig, ExpertView - from freetoken.layers.quantization.moe.fp8_block import Fp8BlockMoEMethod, TritonFp8BlockMoEKernel + 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.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") @@ -172,13 +196,13 @@ def test_real_method_bank_pack_apply_prefill(resident, monkeypatch): 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 = dict(gate_up=gu, gate_up_scale=gs[:, :, :2], down=dn, down_scale=ds[:, :, :1]) + 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 k in bank: - assert torch.equal(bank[k], ref_bank[k]) + for name, value in bank.items(): + assert torch.equal(value, ref_bank[name]) view = ExpertView(bank, n=4) layer = SimpleNamespace( activation="silu", @@ -206,10 +230,14 @@ def test_real_method_bank_pack_apply_prefill(resident, monkeypatch): 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 QuantBackend, LayerKind, QuantKind - from freetoken.moe.legacy_format import legacy_format_for, kind_kernel_for, bank_kernel_compatible - import freetoken.checkpoint.ftw as ftw + 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