perf(gguf): sm_120 dispatch thresholds + upstream int8-MMA MMQ port (Ornith)

.gitignore

@@ -227,3 +227,6 @@ benchmarks/cross_framework
 
 # local e2e/bench artifacts (harnesses may run with repo cwd)
 /results/
+
+# Local scratch / agent worktrees (never committed)
+.claude/

benchmarks/README.md

@@ -27,5 +27,28 @@ batch size x miss rate.
 python benchmarks/bench_offload_cache_copy.py
 ```
 
+**`bench_ornith_attention.py`** — synthetic (no checkpoint): Ornith's exact attention
+geometry (16 query heads, 2 KV heads, head_dim 256 — the GQA shape `decode_launch_config`
+tunes packed-int4/Q4_0 decode for) through the production Triton kernels directly
+(`decode_paged_attention` / `paged_attention` / `extend_paged_attention`, no server).
+Sweeps decode context length x batch size x the `max_kv_splits` scratch ceiling, plus
+representative prefill (fresh chunk) and extend (cached prefix + new chunk) cases, over
+one or more `--kv-quant` pool formats (`int4`/`q4_0`, `q8_0`, `fp8_e4m3`, `bf16`). Every
+quantized case is checked against the same kernel fed the pool's dequantized values
+before it is timed — the correctness gate `test_ornith_q4_tuned_decode_matches_dequantized_oracle`
+pins at unit scale, exercised here at benchmark scale.
+
+```bash
+python benchmarks/bench_ornith_attention.py
+python benchmarks/bench_ornith_attention.py --decode-lengths 8192 32768 131072 200000 \
+    --kv-quant int4 q8_0 --batch-sizes 1 4 16 --json out.jsonl
+```
+
 For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw`
 instead — it writes the JSON profile the engine reads.
+
+`bench_decode_moe.py` also accepts `--max-context` (full-context `--max-seq-len-override`
++ `--num-tokens`), `--kv-cache-dtype`, `--prefill-chunk` (`--max-prefill-length`), and
+`--prefill-hit-d2d`, for reproducing the long-context configurations in `docs/models.md`
+(e.g. Ornith Q4_0 at 200K) through the real serving path — all optional, defaults
+unchanged when omitted.

benchmarks/bench_decode_moe.py

@@ -112,6 +112,31 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         default=1800,
         help="seconds to wait for the spawned server to become ready",
     )
+    p.add_argument(
+        "--max-context",
+        type=int,
+        default=None,
+        help="server --max-seq-len-override AND --num-tokens, for a full-context run "
+        "(e.g. a long-context Ornith/Laguna session); default keeps the prior "
+        "8192 + --decode sizing with the server's own --num-tokens default",
+    )
+    p.add_argument(
+        "--kv-cache-dtype",
+        default=None,
+        help="server --kv-cache-dtype (auto|q8_0|fp8_e4m3|int4|q4_0); default leaves the "
+        "server's own default (auto, unquantized) in place",
+    )
+    p.add_argument(
+        "--prefill-chunk",
+        type=int,
+        default=None,
+        help="server --max-prefill-length; default leaves the server's own default chunk size",
+    )
+    p.add_argument(
+        "--prefill-hit-d2d",
+        action="store_true",
+        help="pass --moe-prefill-hit-d2d to the server (off by default, matching the server default)",
+    )
     p.add_argument("--json", dest="json_out", default=None, help="append the result rows here")
     return p.parse_args(argv)
 
@@ -172,17 +197,28 @@ def free_port() -> int:
 
 
 def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]:
+    max_seq_len = args.max_context if args.max_context is not None else 8192 + args.decode
     cmd = [
         sys.executable, "-m", "freetoken.cli", "serve",
         "--model", args.model,
         "--host", "127.0.0.1", "--port", str(port),
         "--moe-backend", backend,
         "--max-running-requests", "1",
-        "--max-seq-len-override", str(8192 + args.decode),
+        "--max-seq-len-override", str(max_seq_len),
         "--memory-ratio", str(args.mem_ratio),
         "--cuda-graph-max-bs", "0" if args.no_graph else "1",
         "--moe-hybrid-max-fetch", str(args.hybrid_fetch),
     ]
+    # Every flag below is opt-in and omitted unless passed, so a bare invocation keeps
+    # the server's own defaults exactly as before this option set existed.
+    if args.max_context is not None:
+        cmd += ["--num-tokens", str(args.max_context)]
+    if args.kv_cache_dtype is not None:
+        cmd += ["--kv-cache-dtype", args.kv_cache_dtype]
+    if args.prefill_chunk is not None:
+        cmd += ["--max-prefill-length", str(args.prefill_chunk)]
+    if args.prefill_hit_d2d:
+        cmd.append("--moe-prefill-hit-d2d")
     if args.cache > 0:
         cmd += ["--moe-cache-size", str(args.cache)]
     elif args.cache_rate is not None:
@@ -309,6 +345,9 @@ def run_one(args: argparse.Namespace, backend: str) -> dict:
         f"[bench] model={args.model}\n"
         f"[bench] backend={backend} cache={args.cache or args.cache_rate or 'auto'} "
         f"mem_ratio={args.mem_ratio} decode={args.decode} graph={not args.no_graph}\n"
+        f"[bench] max_context={args.max_context or f'{8192 + args.decode} (default)'} "
+        f"kv_cache_dtype={args.kv_cache_dtype or 'auto (default)'} "
+        f"prefill_chunk={args.prefill_chunk or 'default'} prefill_hit_d2d={args.prefill_hit_d2d}\n"
         f"[bench] sampling={sampling} <- {sampling_src}\n"
         f"[bench] server log: {log_path}",
         flush=True,

benchmarks/bench_gguf_gemm.py

@@ -0,0 +1,171 @@
+"""GGUF quantized-matmul benchmark: int8-MMA MMQ vs DP4A MMQ vs dequant+cuBLAS vs MMVQ.
+
+Sweeps the dense (``fused_mul_mat_gguf`` seams) and grouped-MoE
+(``_moe_matmul`` seams) kernel families over batch size on synthetic
+random-but-safe packed weights (fp16 scale fields masked small, as in
+``tests/kernels/test_gguf_quant_types.py``), at Ornith-1.5-35B geometry by
+default. Every timed case is first cross-checked against the transient
+dequantized weights (the oracle): rel error must stay below --tol.
+
+The int8-MMA columns need the ``freetoken_gguf_mmq`` extension (sm_75+ build,
+sm_120-dispatched); they are skipped with a note where unavailable.
+
+Run:
+    CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_gguf_gemm.py
+    ... --dense-rows 4 16 32 512 8192 --moe-tokens 16 320 8192 --json out.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+
+
+def build_cases(args) -> list[dict]:
+    """Pure case list (unit-testable without CUDA)."""
+    cases = []
+    for rows in args.dense_rows:
+        for qtype_name in args.dense_types:
+            cases.append({"op": "dense", "rows": rows, "qtype": qtype_name,
+                          "out_features": args.dense_out, "in_features": args.hidden})
+    for tokens in args.moe_tokens:
+        cases.append({"op": "moe", "tokens": tokens, "experts": args.experts,
+                      "top_k": args.top_k, "hidden": args.hidden, "inter": args.inter})
+    return cases
+
+
+def parse_args(argv=None):
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--dense-rows", type=int, nargs="+", default=[4, 8, 16, 32, 128, 512, 2048, 8192])
+    p.add_argument("--dense-types", nargs="+", default=["q4_k", "q6_k"], choices=["q4_k", "q6_k"])
+    p.add_argument("--dense-out", type=int, default=8192)
+    p.add_argument("--moe-tokens", type=int, nargs="+", default=[16, 64, 320, 1024, 8192])
+    p.add_argument("--experts", type=int, default=256)
+    p.add_argument("--top-k", type=int, default=8)
+    p.add_argument("--hidden", type=int, default=2048)
+    p.add_argument("--inter", type=int, default=512)
+    p.add_argument("--iters", type=int, default=20)
+    p.add_argument("--tol", type=float, default=0.02)
+    p.add_argument("--json", type=str, default=None)
+    return p.parse_args(argv)
+
+
+def _med_ms(fn, iters):
+    import torch
+
+    for _ in range(3):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(iters):
+        s, e = torch.cuda.Event(True), torch.cuda.Event(True)
+        s.record()
+        fn()
+        e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    return statistics.median(times)
+
+
+def _packed(qtype, rows_of_blocks, seed):
+    import numpy as np
+    import torch
+    from freetoken.models.gguf.dequant import BLOCK_SHAPE
+
+    rng = np.random.default_rng(seed)
+    raw = rng.integers(0, 256, (rows_of_blocks, BLOCK_SHAPE[qtype][1]), dtype=np.uint8)
+    raw.view(np.uint16)[:] &= np.uint16(0x3BFF)
+    return torch.from_numpy(raw)
+
+
+def main(argv=None) -> int:
+    args = parse_args(argv)
+    import torch
+
+    from freetoken.kernel.gguf import (
+        ggml_dequantize,
+        ggml_moe_a8,
+        ggml_moe_get_block_size,
+        ggml_mul_mat_a8,
+        ggml_mul_mat_vec_a8,
+    )
+    from freetoken.models.gguf.dequant import BLOCK_SHAPE, GGML_Q4_K, GGML_Q6_K
+    from freetoken.moe.fused import moe_align_block_size
+
+    qtypes = {"q4_k": GGML_Q4_K, "q6_k": GGML_Q6_K}
+    try:
+        from freetoken.kernel.gguf import ggml_moe_a8_mma, ggml_mul_mat_a8_mma
+
+        mma_ok = True
+    except Exception as exc:  # noqa: BLE001 - report and continue without MMA
+        print(f"# int8-MMA extension unavailable: {exc}")
+        mma_ok = False
+
+    out_rows = []
+    torch.manual_seed(0)
+    for case in build_cases(args):
+        if case["op"] == "dense":
+            qtype = qtypes[case["qtype"]]
+            block, ts = BLOCK_SHAPE[qtype]
+            out_f, in_f, rows = case["out_features"], case["in_features"], case["rows"]
+            w = _packed(qtype, out_f * in_f // block, seed=qtype).reshape(out_f, -1).cuda()
+            dense = ggml_dequantize(w, qtype, out_f, in_f, torch.float16)
+            x = torch.randn(rows, in_f, dtype=torch.float16, device="cuda")
+            ref = (x.float() @ dense.float().T)
+            res = {"op": "dense", "qtype": case["qtype"], "rows": rows}
+
+            def check(name, y):
+                rel = ((y.float() - ref).norm() / ref.norm()).item()
+                assert rel < args.tol, (name, rel)
+
+            y = ggml_mul_mat_a8(w, x, qtype, out_f)
+            check("dp4a", y)
+            res["dp4a_ms"] = _med_ms(lambda: ggml_mul_mat_a8(w, x, qtype, out_f), args.iters)
+            res["dequant_ms"] = _med_ms(
+                lambda: x @ ggml_dequantize(w, qtype, out_f, in_f, torch.float16).T, args.iters
+            )
+            if rows <= 8:
+                y = ggml_mul_mat_vec_a8(w, x, qtype, out_f)
+                check("mmvq", y)
+                res["mmvq_ms"] = _med_ms(lambda: ggml_mul_mat_vec_a8(w, x, qtype, out_f), args.iters)
+            if mma_ok:
+                y = ggml_mul_mat_a8_mma(w, x.float(), qtype, out_f)
+                check("mma", y)
+                res["mma_ms"] = _med_ms(lambda: ggml_mul_mat_a8_mma(w, x, qtype, out_f), args.iters)
+        else:
+            experts, top_k = case["experts"], case["top_k"]
+            hidden, inter, tokens = case["hidden"], case["inter"], case["tokens"]
+            block, ts = BLOCK_SHAPE[GGML_Q4_K]
+            gu = _packed(GGML_Q4_K, experts * 2 * inter * hidden // block, seed=1).reshape(experts, -1).cuda()
+            x = torch.randn(tokens, hidden, dtype=torch.float16, device="cuda")
+            ids = torch.stack([torch.randperm(experts)[:top_k] for _ in range(tokens)]).int().cuda()
+            res = {"op": "moe", "tokens": tokens}
+
+            def dp4a_moe():
+                bs = ggml_moe_get_block_size(GGML_Q4_K)
+                s, ei, npad = moe_align_block_size(ids, bs, experts)
+                return ggml_moe_a8(x, gu, s, ei, npad, GGML_Q4_K, 2 * inter, top_k, tokens)
+
+            ref_moe = dp4a_moe().float()
+            res["dp4a_ms"] = _med_ms(dp4a_moe, args.iters)
+            if mma_ok:
+                y = ggml_moe_a8_mma(x, gu, ids, top_k, GGML_Q4_K, 2 * inter, tokens, gu.shape[1], True)
+                rel = ((y.float() - ref_moe).norm() / ref_moe.norm()).item()
+                assert rel < args.tol, ("moe-mma", rel)
+                res["mma_ms"] = _med_ms(
+                    lambda: ggml_moe_a8_mma(x, gu, ids, top_k, GGML_Q4_K, 2 * inter, tokens, gu.shape[1], True),
+                    args.iters,
+                )
+        print(res)
+        out_rows.append(res)
+
+    if args.json:
+        with open(args.json, "w") as f:
+            for row in out_rows:
+                f.write(json.dumps(row) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())