minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[DSA] Q8KV8 FP8 Sparse Prefill on GLM-5.2 & DeepSeek-V3.2: Q8-Path & Shared-Path Optimizations (#31888)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Ho-Ren (Jack) Chuang
2026-07-30 15:15:11 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 4f51dad1da
commit e4a40a71f8
30 changed files with 3837 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py
Executable
+262
View File
@@ -0,0 +1,262 @@
#!/usr/bin/env python3
"""Microbenchmark: Q8KV8 sparse-prefill KV gather overhaul.
Compares the legacy gather (``gather_dequant_requant_fp8_paged_legacy``:
fresh ``torch.zeros`` destination + one program per (token, 128-elem
slice)) against the new gather (``gather_dequant_requant_fp8_paged``:
no pre-zeroing needed, fused pad-row zero-fill, TOKENS_PER_PROG tokens
per program with 16B-vectorized access), in three flavors:
legacy torch.zeros alloc + legacy kernel (baseline)
new_alloc torch.empty alloc + vectorized kernel (vectorized copy only)
new_cached persistent grow-only buffer + vec kernel (changes 1 + 2,
= production path)
For each scenario it checks BIT-EXACT equality of the fp8 output bytes
(``torch.equal`` on ``uint8`` views; the requant scale is the identity
scalar 1.0 on this path, so the buffer is the entire output), including
the `topk` zero landing-pad rows, then reports us/call and effective
TB/s.
Shapes model GLM / DeepSeek-V3.2 DSA prefill on one rank: d = 576 fp8
out (512 nope + 64 rope), 656 B/token paged cache rows (512 nope fp8 +
16 B f32 group scales + 128 B bf16 rope), page_size 64, topk 2048.
NOTE: gather traffic scales with kv_len (= number of gathered KV rows
= len(page_table_1_flattened)), NOT with s_q; s_q below only labels the
chunk that a scenario represents. The GLM-5.2 il=64k profile point
(112.0 us/call, ~0.7 TB/s effective) corresponds to the
(s_q=4096, kv_len=65536) row.
Usage (single GPU, < 2 min):
python benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py [--device cuda:0]
[--iters 200] [--warmup 20]
Also runs correctness-only edge cases: ragged tail (kv_len % 4 != 0),
extra_rows=0, num_tokens=0, and a cached-buffer SHRINK reuse (big call
then small call) that proves stale bytes from the earlier, larger call
cannot leak into the smaller call's pad rows.
"""
import argparse
import importlib.util
import sys
from pathlib import Path
import torch
_REPO_ROOT = Path(__file__).resolve().parents[3]
_MOD_PATH = _REPO_ROOT / "python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py"
# Import the module straight from its file so the benchmark stays
# standalone (no sglang package import side effects; needs only
# torch + triton).
_spec = importlib.util.spec_from_file_location("dequant_k_cache", _MOD_PATH)
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)
gather_new = _mod.gather_dequant_requant_fp8_paged
gather_legacy = _mod.gather_dequant_requant_fp8_paged_legacy
PAGE_SIZE = 64
DIM_QUANT = 656 # 512 nope fp8 + 16 scale bytes + 128 rope bytes
OUT_DIM = 576 # 512 nope + 64 rope, fp8
TOPK = 2048
# (s_q label, kv_len = gathered rows). kv_len drives the bytes moved.
SCENARIOS = [
(512, 8192),
(2048, 32768),
(4096, 65536),
]
def build_paged_kv_pool(pool_tokens: int, device: str) -> torch.Tensor:
"""Synthetic paged fp8 KV cache: [pool_tokens, 1, 656] fp8_e4m3fn."""
g = torch.Generator(device=device).manual_seed(0)
nope = (torch.randn(pool_tokens, 512, generator=g, device=device) * 2.0).to(
torch.float8_e4m3fn
)
# Positive, realistically small per-group dequant scales.
scales = (torch.rand(pool_tokens, 4, generator=g, device=device) * 0.05 + 1e-3).to(
torch.float32
)
rope = torch.randn(pool_tokens, 64, generator=g, device=device).to(torch.bfloat16)
raw = torch.empty(pool_tokens, DIM_QUANT, dtype=torch.uint8, device=device)
raw[:, :512] = nope.view(torch.uint8)
raw[:, 512:528] = scales.view(torch.uint8)
raw[:, 528:] = rope.view(torch.uint8)
return raw.view(torch.float8_e4m3fn).view(pool_tokens, 1, DIM_QUANT)
def build_page_table_flattened(
kv_lens, pool_tokens: int, device: str, seed: int = 1
) -> torch.Tensor:
"""Realistic page_table_1_flattened: per request, random distinct
64-token pages, tokens contiguous within a page (production paged
layout), requests concatenated."""
n_pool_pages = pool_tokens // PAGE_SIZE
g = torch.Generator(device="cpu").manual_seed(seed)
parts = []
for kv_len in kv_lens:
n_pages = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE
assert n_pages <= n_pool_pages, "pool too small for scenario"
pages = torch.randperm(n_pool_pages, generator=g)[:n_pages]
toks = (pages[:, None] * PAGE_SIZE + torch.arange(PAGE_SIZE)[None, :]).reshape(
-1
)[:kv_len]
parts.append(toks)
return torch.cat(parts).to(torch.int32).to(device)
class CachedGather:
"""Mimics the dsa_backend production path: persistent grow-only fp8
destination buffer, gather zero-fills only the pad tail in-kernel."""
def __init__(self):
self.buf = None
def __call__(self, pool, pt, extra_rows):
total = pt.shape[0] + extra_rows
if self.buf is None or self.buf.shape[0] < total:
self.buf = torch.empty(
(total, OUT_DIM), dtype=torch.float8_e4m3fn, device=pool.device
)
return gather_new(pool, pt, extra_rows=extra_rows, out=self.buf[:total])
def assert_bit_exact(ref: torch.Tensor, got: torch.Tensor, what: str):
assert ref.shape == got.shape, f"{what}: shape {got.shape} != {ref.shape}"
ok = torch.equal(
ref.contiguous().view(torch.uint8), got.contiguous().view(torch.uint8)
)
assert ok, f"{what}: fp8 bytes NOT bit-exact"
def bench_us(fn, warmup: int, iters: int) -> float:
for _ in range(warmup):
fn()
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iters):
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) * 1000.0 / iters # ms -> us
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--device", default="cuda:0")
ap.add_argument("--iters", type=int, default=200)
ap.add_argument("--warmup", type=int, default=20)
args = ap.parse_args()
if not torch.cuda.is_available():
print("CUDA not available; this microbench needs 1 GPU.", file=sys.stderr)
sys.exit(1)
torch.cuda.set_device(args.device)
dev = args.device
pool_tokens = 131072 # 2x the largest kv_len; 131072*656 B ~ 86 MB
pool = build_paged_kv_pool(pool_tokens, dev)
print(f"device={dev} ({torch.cuda.get_device_name(dev)})")
print(f"pool: {pool_tokens} tokens x {DIM_QUANT} B, topk={TOPK}")
print()
# ------------------------------------------------------------------
# Correctness edge cases (not timed)
# ------------------------------------------------------------------
print("== correctness edge cases ==")
cached_edge = CachedGather()
for kv_len, extra in [(1234, TOPK), (8192, 0), (0, TOPK), (63, 17)]:
pt = build_page_table_flattened([kv_len], pool_tokens, dev, seed=7)
ref = gather_legacy(pool, pt, extra_rows=extra)
got_alloc = gather_new(pool, pt, extra_rows=extra)
got_cached = cached_edge(pool, pt, extra_rows=extra)
assert_bit_exact(ref, got_alloc, f"kv_len={kv_len},extra={extra} new_alloc")
assert_bit_exact(ref, got_cached, f"kv_len={kv_len},extra={extra} new_cached")
# Pad rows must be exactly zero bytes.
if extra > 0:
pad = got_cached[kv_len:].view(torch.uint8)
assert int(pad.max()) == 0 if pad.numel() else True
print(f" kv_len={kv_len:6d} extra_rows={extra:5d}: bit-exact OK")
# Cached-buffer SHRINK reuse: big call dirties the buffer, then a
# smaller call must still produce zero pad rows (stale-data test for
# the grow-only buffer + tail-only zeroing invariant).
cached_shrink = CachedGather()
pt_big = build_page_table_flattened([65536], pool_tokens, dev, seed=11)
cached_shrink(pool, pt_big, TOPK)
pt_small = build_page_table_flattened([4096], pool_tokens, dev, seed=13)
ref_small = gather_legacy(pool, pt_small, extra_rows=TOPK)
got_small = cached_shrink(pool, pt_small, TOPK)
assert_bit_exact(ref_small, got_small, "shrink-reuse (65536 -> 4096)")
print(" shrink-reuse 65536 -> 4096 rows: pad rows clean, bit-exact OK")
print()
# ------------------------------------------------------------------
# Timed scenarios
# ------------------------------------------------------------------
print("== timing ==")
print(
"metrics: us/call = mean wall time per gather call incl. any alloc/"
"zero-fill (LOWER = faster); eff TB/s = payload (656+4 B/token read"
" + 576 B/row written incl. pad) / time (HIGHER = faster);"
" speedup = legacy_us / variant_us (>1 = faster than legacy)."
)
header = (
f"{'s_q':>5} {'kv_len':>7} {'variant':>10} {'us/call':>9} "
f"{'eff TB/s':>9} {'speedup':>8}"
)
print(header)
print("-" * len(header))
for s_q, kv_len in SCENARIOS:
pt = build_page_table_flattened([kv_len], pool_tokens, dev, seed=s_q)
total_rows = kv_len + TOPK
payload_bytes = kv_len * (DIM_QUANT + 4) + total_rows * OUT_DIM
cached = CachedGather()
variants = [
("legacy", lambda: gather_legacy(pool, pt, extra_rows=TOPK)),
("new_alloc", lambda: gather_new(pool, pt, extra_rows=TOPK)),
("new_cached", lambda: cached(pool, pt, TOPK)),
]
# Bit-exactness at the benchmarked shape before timing.
ref = gather_legacy(pool, pt, extra_rows=TOPK)
for name, fn in variants[1:]:
assert_bit_exact(ref, fn(), f"s_q={s_q} {name}")
del ref
legacy_us = None
for name, fn in variants:
us = bench_us(fn, args.warmup, args.iters)
tbps = payload_bytes / (us * 1e-6) / 1e12
if name == "legacy":
legacy_us = us
speedup = "1.00x"
else:
speedup = f"{legacy_us / us:.2f}x"
print(
f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} "
f"{tbps:>9.3f} {speedup:>8}"
)
print()
print(
"note: legacy additionally writes a full-buffer zero fill "
f"({total_rows * OUT_DIM / 1e6:.1f} MB at the largest shape) that is "
"NOT counted in its payload bytes; its true HW bandwidth is higher "
"than the eff TB/s shown, which is exactly why us/call is the "
"decision metric."
)
print("ALL CHECKS PASSED")
if __name__ == "__main__":
main()
benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py
Executable
+486
View File
@@ -0,0 +1,486 @@
#!/usr/bin/env python3
"""Microbenchmark: Q8KV8 sparse-prefill q-prep — old path vs born-fp8 fused path.
Old path (production default):
1. q_nope_out = torch.bmm(q_nope.transpose(0, 1), w_kc).transpose(0, 1)
(cublas bf16 bmm, writes bf16 [H, T, N] to DRAM)
2. concat_and_cast_q_fp8_pad(q_fp8, q_nope_out, q_rope, H)
(Triton: re-reads the bf16 bmm output + q_rope, writes fp8 [T, H, N+R])
New path (SGLANG_ENABLE_DSA_Q8KV8_BORN_FP8_Q):
absorbed_bmm_concat_cast_q_fp8(q_fp8, q_nope, w_kc, q_rope, H)
(one Triton kernel: bmm + concat + fp8 cast; the bf16 q_nope_out and the
standalone concat-cast launch disappear)
Non-power-of-2 K (GLM 192) K-dimension codegen variants (A/B'd here; all keep
the identical fp32 -> bf16 -> fp8 epilogue, see cache_ops.py):
loop : split-K loop, BLOCK_K=64 x 3 (the original K=192 path)
two_dot : preload a once as 128+64 tiles, two chained tl.dot, no K-loop
three_dot : preload a once as 3 x 64 tiles, three chained tl.dot
(same fp32 add order as `loop`, loads hoisted)
pad : single tl.dot at BLOCK_K=256 with zero-masked k tail
single_k : single tl.dot at BLOCK_K=192 -- documents the Triton
non-power-of-2 tl.arange limitation (compile fails <= 3.5.x)
Power-of-2 K (DeepSeek 128) collapses every variant to the same single-dot
fast path, so only one NEW row is shown there.
Shapes (both models: N = kv_lora_rank = 512, R = qk_rope_head_dim = 64;
K = qk_nope_head_dim differs per model):
GLM-5.2: heads = 64, K = 192 (w_kc [64, 192, 512]; DP attention,
per-rank full heads; K=192 exercises the kernel's split-K path)
DS-V3.2: heads = 128, K = 128 (power-of-2 K, preload-once fast path)
Metric conventions:
* time is reported in microseconds per call (us/call) — LOWER = FASTER.
* bandwidth is analytic-bytes / time in GB/s — HIGHER = BETTER.
* "speedup x" = old_time / new_time — >1.0 means the NEW path is faster.
Correctness:
* rope half must be BIT-EXACT (same bf16 source, same Triton conversion).
* nope half: same rounding stages (fp32 accum -> bf16 -> fp8) but a
different GEMM accumulation order than cublas -> near- but not
guaranteed bit-exact. We report the bitwise-match fraction, the max
dequantized |diff|, and which path lands closer to an fp64 reference.
Usage (single GPU):
python benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py # both model shapes
python benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py --tokens 8192 --iters 300
python benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py --variants two_dot,pad
python benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py --sweep # + tile/warp sweep
python benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py --rounding-study
"""
import argparse
import torch
from sglang.kernels.ops.kvcache.cache_ops import (
absorbed_bmm_concat_cast_q_fp8,
concat_and_cast_q_fp8_pad,
)
N_LORA = 512 # kv_lora_rank (post-absorb q_nope dim; "d_nope" at the kernel)
R_ROPE = 64 # qk_rope_head_dim
def make_inputs(
num_tokens: int,
num_heads: int,
k_nope: int,
device,
seed: int,
magnitude: float,
):
g = torch.Generator(device=device).manual_seed(seed)
# Production layout: q = q_b_proj output [T, H, K+R] bf16; q_nope/q_rope are
# strided views of it (rope applied in-place on the q_rope slice).
q = (
torch.randn(
(num_tokens, num_heads, k_nope + R_ROPE),
generator=g,
device=device,
dtype=torch.float32,
)
* magnitude
).to(torch.bfloat16)
q_nope = q[..., :k_nope]
q_rope = q[..., k_nope:]
# Production w_kc layout: [H, K, N] with strides (K*N, 1, K) (N-major), the
# result of w_kc.transpose(1, 2).contiguous().transpose(1, 2) at load.
w_base = (
torch.randn(
(num_heads, N_LORA, k_nope),
generator=g,
device=device,
dtype=torch.float32,
)
/ (k_nope**0.5)
).to(torch.bfloat16)
w_kc = w_base.transpose(1, 2)
return q, q_nope, q_rope, w_kc
def old_path(q_fp8, q_nope, w_kc, q_rope, num_heads):
q_nope_out = torch.bmm(q_nope.transpose(0, 1), w_kc).transpose(0, 1)
concat_and_cast_q_fp8_pad(q_fp8, q_nope_out, q_rope, num_heads)
def old_path_bmm_only(q_nope, w_kc):
return torch.bmm(q_nope.transpose(0, 1), w_kc)
def new_path(q_fp8, q_nope, w_kc, q_rope, num_heads, **kw):
absorbed_bmm_concat_cast_q_fp8(q_fp8, q_nope, w_kc, q_rope, num_heads, **kw)
# Non-power-of-2-K variants, in bench order (power-of-2 K collapses to "auto").
ALL_VARIANTS = ["loop", "two_dot", "three_dot", "pad", "single_k"]
# (block_m, block_n, num_warps, num_stages) sweep grid; num_stages 0 = Triton
# default. N=512 is a multiple of every block_n here; block_m stays power of 2.
SWEEP_TILES = [
(64, 128, 4, 0), # kernel default
(64, 128, 8, 0),
(64, 128, 4, 2),
(64, 128, 4, 4),
(128, 128, 4, 0),
(128, 128, 8, 0),
(64, 256, 8, 0),
(128, 256, 8, 0),
(32, 128, 4, 0),
(64, 64, 4, 0),
(128, 64, 8, 0),
]
def time_fn(fn, iters: int, warmup: int) -> float:
"""Median wall time of fn() in microseconds per call (lower = faster)."""
for _ in range(warmup):
fn()
torch.cuda.synchronize()
times = []
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
chunk = 10
for _ in range(max(1, iters // chunk)):
start.record()
for _ in range(chunk):
fn()
end.record()
torch.cuda.synchronize()
times.append(start.elapsed_time(end) * 1e3 / chunk) # ms -> us
times.sort()
return times[len(times) // 2]
def analytic_bytes(num_tokens: int, num_heads: int, k_nope: int):
"""(old_bytes, new_bytes) of DRAM traffic per call, analytic lower bound."""
t, h = num_tokens, num_heads
a_read = t * h * k_nope * 2 # q_nope bf16
w_read = h * k_nope * N_LORA * 2 # w_kc bf16
nope_bf16 = t * h * N_LORA * 2 # bmm bf16 out (written then re-read)
rope_read = t * h * R_ROPE * 2 # q_rope bf16
fp8_write = t * h * (N_LORA + R_ROPE) # q_fp8 out
old = (a_read + w_read + nope_bf16) + (nope_bf16 + rope_read + fp8_write)
new = a_read + w_read + rope_read + fp8_write
return old, new
def make_check_ctx(num_tokens, num_heads, k_nope, device, seed, magnitude):
"""Fresh inputs + old-path fp8 output + fp64 bmm reference (once/config)."""
q, q_nope, q_rope, w_kc = make_inputs(
num_tokens, num_heads, k_nope, device, seed, magnitude
)
q_fp8_old = torch.zeros(
(num_tokens, num_heads, N_LORA + R_ROPE),
dtype=torch.float8_e4m3fn,
device=device,
)
old_path(q_fp8_old, q_nope, w_kc, q_rope, num_heads)
torch.cuda.synchronize()
# fp64 reference: which path's fp8 lands closer to the exact bmm?
ref = torch.bmm(
q_nope.transpose(0, 1).to(torch.float64), w_kc.to(torch.float64)
).transpose(0, 1)
err_old = (q_fp8_old[..., :N_LORA].to(torch.float64) - ref).abs()
return {
"q_nope": q_nope,
"q_rope": q_rope,
"w_kc": w_kc,
"num_heads": num_heads,
"q_fp8_old": q_fp8_old,
"ref": ref,
"meanerr_old": err_old.mean().item(),
"maxerr_old": err_old.max().item(),
}
def check_variant(ctx, **new_kwargs):
"""Correctness of one new-path variant vs the old path + fp64 reference."""
q_fp8_old = ctx["q_fp8_old"]
q_fp8_new = torch.zeros_like(q_fp8_old)
new_path(
q_fp8_new,
ctx["q_nope"],
ctx["w_kc"],
ctx["q_rope"],
ctx["num_heads"],
**new_kwargs,
)
torch.cuda.synchronize()
rope_old = q_fp8_old[..., N_LORA:].view(torch.uint8)
rope_new = q_fp8_new[..., N_LORA:].view(torch.uint8)
rope_bitexact = bool(torch.equal(rope_old, rope_new))
nope_old = q_fp8_old[..., :N_LORA]
nope_new = q_fp8_new[..., :N_LORA]
match = (
(nope_old.view(torch.uint8) == nope_new.view(torch.uint8)).float().mean().item()
)
diff = (nope_old.to(torch.float32) - nope_new.to(torch.float32)).abs()
max_diff = diff.max().item()
err_new = (nope_new.to(torch.float64) - ctx["ref"]).abs()
return {
"rope_bitexact": rope_bitexact,
"nope_bitwise_match_frac": match,
"nope_max_dequant_absdiff": max_diff,
"nope_meanerr_old_vs_fp64": ctx["meanerr_old"],
"nope_meanerr_new_vs_fp64": err_new.mean().item(),
"nope_maxerr_old_vs_fp64": ctx["maxerr_old"],
"nope_maxerr_new_vs_fp64": err_new.max().item(),
}
def rounding_study(device, seed):
"""sweep summary: fp8(bf16(x)) double round vs fp8(x) single round.
(Informational only — the born-fp8 kernel deliberately keeps the
fp32->bf16->fp8 double round to match the default path's rounding stages.)
"""
g = torch.Generator(device=device).manual_seed(seed)
x = torch.randn((1 << 22,), generator=g, device=device, dtype=torch.float32) * 8.0
double = x.to(torch.bfloat16).to(torch.float8_e4m3fn)
single = x.to(torch.float8_e4m3fn)
mismatch = (
(double.view(torch.uint8) != single.view(torch.uint8)).float().mean().item()
)
err_double = (double.to(torch.float32) - x).abs()
err_single = (single.to(torch.float32) - x).abs()
print("\n=== rounding study: fp32->bf16->fp8 (double) vs fp32->fp8 (single) ===")
print(f"elements : {x.numel()}")
print(f"byte-mismatch fraction : {mismatch:.3e} (fraction, lower = closer)")
print(
f"mean |err| vs fp32 (double) : {err_double.mean().item():.6e} (lower = more accurate)"
)
print(
f"mean |err| vs fp32 (single) : {err_single.mean().item():.6e} (lower = more accurate)"
)
winner = (
"single (direct fp32->fp8)"
if err_single.mean() <= err_double.mean()
else "double (via bf16)"
)
print(f"more accurate on average : {winner}")
print(
"NOTE: the born-fp8 kernel keeps the DOUBLE round on purpose to match "
"the default path's rounding stages."
)
def run_config(
name,
num_tokens,
num_heads,
k_nope,
iters,
warmup,
device,
seed,
magnitude,
variants,
sweep,
):
print(f"\n=== {name}: tokens={num_tokens} heads={num_heads} K={k_nope} ===")
print(
f" (K={k_nope} nope-in, N={N_LORA} nope-out, R={R_ROPE} rope; "
"us/call LOWER = FASTER; GB/s HIGHER = BETTER; speedup >1 = new faster)"
)
q, q_nope, q_rope, w_kc = make_inputs(
num_tokens, num_heads, k_nope, device, seed, magnitude
)
q_fp8 = torch.zeros(
(num_tokens, num_heads, N_LORA + R_ROPE),
dtype=torch.float8_e4m3fn,
device=device,
)
t_old = time_fn(
lambda: old_path(q_fp8, q_nope, w_kc, q_rope, num_heads), iters, warmup
)
t_bmm = time_fn(lambda: old_path_bmm_only(q_nope, w_kc), iters, warmup)
# standalone concat-cast (reads a fresh bf16 bmm out, like production)
q_nope_out = torch.bmm(q_nope.transpose(0, 1), w_kc).transpose(0, 1)
t_cast = time_fn(
lambda: concat_and_cast_q_fp8_pad(q_fp8, q_nope_out, q_rope, num_heads),
iters,
warmup,
)
b_old, b_new = analytic_bytes(num_tokens, num_heads, k_nope)
print(
f"OLD bmm (cublas bf16) : {t_bmm:10.1f} us/call"
f" (component of OLD total)"
)
print(f"OLD concat_and_cast_q_fp8_pad: {t_cast:10.1f} us/call (component)")
print(
f"OLD total (bmm + concat-cast): {t_old:10.1f} us/call"
f" ({b_old / 1e6:8.1f} MB analytic, {b_old / t_old / 1e3:7.0f} GB/s)"
)
# Power-of-2 K collapses every variant to the same single-dot codegen.
pow2 = k_nope & (k_nope - 1) == 0
run_variants = ["auto"] if pow2 else variants
ctx = make_check_ctx(num_tokens, num_heads, k_nope, device, seed + 1, magnitude)
results = {}
for v in run_variants:
kw = {"variant": v}
try:
t_new = time_fn(
lambda: new_path(q_fp8, q_nope, w_kc, q_rope, num_heads, **kw),
iters,
warmup,
)
except Exception as e:
msg = (str(e).splitlines() or [type(e).__name__])[0]
print(f"NEW {v:<24}: COMPILE/RUN FAIL — {msg[:100]}")
continue
c = check_variant(ctx, **kw)
results[v] = t_new
faster = "NEW FASTER" if t_new < t_old else "OLD FASTER"
print(
f"NEW {v:<24}: {t_new:10.1f} us/call"
f" ({b_new / 1e6:8.1f} MB analytic, {b_new / t_new / 1e3:7.0f} GB/s,"
f" speedup {t_old / t_new:5.2f}x vs OLD, {faster})"
)
rope = "PASS (bitwise identical)" if c["rope_bitexact"] else "FAIL (BUG)"
print(
f" rope bit-exact: {rope}; nope bitwise match vs OLD"
f" {c['nope_bitwise_match_frac'] * 100:9.4f}% (100% = bit-exact);"
f" max |dequant diff| {c['nope_max_dequant_absdiff']:.4f}"
)
print(
f" nope |err| vs fp64 ref: mean old {c['nope_meanerr_old_vs_fp64']:.3e}"
f" / new {c['nope_meanerr_new_vs_fp64']:.3e}; max old"
f" {c['nope_maxerr_old_vs_fp64']:.3e} / new"
f" {c['nope_maxerr_new_vs_fp64']:.3e} (lower = more accurate)"
)
if results:
best = min(results, key=results.get)
print(
f"BEST variant : {best} @ {results[best]:.1f} us/call"
f" (speedup {t_old / results[best]:.2f}x vs OLD total)"
)
if sweep and results:
print(
f"\n--- tile sweep: {name} (us/call LOWER = FASTER;"
" stages=0 -> Triton default) ---"
)
rows = []
for v in results:
for bm, bn, nw, ns in SWEEP_TILES:
kw = dict(
variant=v, block_m=bm, block_n=bn, num_warps=nw, num_stages=ns
)
try:
t = time_fn(
lambda: new_path(q_fp8, q_nope, w_kc, q_rope, num_heads, **kw),
max(iters // 2, 20),
warmup,
)
except Exception as e:
msg = (str(e).splitlines() or [type(e).__name__])[0]
print(
f" {v:<10} bm={bm:<3} bn={bn:<3} warps={nw} stages={ns}:"
f" FAIL — {msg[:70]}"
)
continue
rows.append((t, v, bm, bn, nw, ns))
print(
f" {v:<10} bm={bm:<3} bn={bn:<3} warps={nw} stages={ns}:"
f" {t:8.1f} us/call ({b_new / t / 1e3:5.0f} GB/s,"
f" {t_old / t:5.2f}x vs OLD)"
)
rows.sort()
print(" -- top 5 (fastest first) --")
for t, v, bm, bn, nw, ns in rows[:5]:
print(
f" {v:<10} bm={bm:<3} bn={bn:<3} warps={nw} stages={ns}:"
f" {t:8.1f} us/call ({t_old / t:5.2f}x vs OLD)"
)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--tokens", type=int, default=4096, help="s_q per call")
parser.add_argument("--iters", type=int, default=200)
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--seed", type=int, default=1234)
parser.add_argument(
"--magnitude",
type=float,
default=1.0,
help="input scale multiplier (q amax stress)",
)
parser.add_argument(
"--heads",
type=int,
default=None,
help="run a single head count instead of the GLM(64,K192)+DS(128,K128) pair",
)
parser.add_argument(
"--k-nope",
type=int,
default=128,
help="qk_nope_head_dim for --heads runs (ignored for the default pair)",
)
parser.add_argument(
"--variants",
type=str,
default="all",
help=(
"comma list of non-power-of-2-K variants to bench "
f"(default: all = {','.join(ALL_VARIANTS)}); power-of-2-K configs "
"always run the single collapsed 'auto' variant"
),
)
parser.add_argument(
"--sweep",
action="store_true",
help="also sweep (block_m, block_n, num_warps, num_stages) per variant",
)
parser.add_argument("--rounding-study", action="store_true")
args = parser.parse_args()
if args.variants == "all":
variants = ALL_VARIANTS
else:
variants = [v.strip() for v in args.variants.split(",") if v.strip()]
unknown = set(variants) - set(ALL_VARIANTS) - {"auto"}
assert not unknown, f"unknown variants: {sorted(unknown)}"
assert torch.cuda.is_available(), "CUDA GPU required"
device = torch.device("cuda")
name = torch.cuda.get_device_name(device)
print(f"device: {name}; torch {torch.__version__}")
if args.heads is not None:
configs = [(f"custom h={args.heads} K={args.k_nope}", args.heads, args.k_nope)]
else:
configs = [
("GLM-5.2 (h=64, K=192)", 64, 192),
("DS-V3.2 (h=128, K=128)", 128, 128),
]
for cfg_name, heads, k_nope in configs:
run_config(
cfg_name,
args.tokens,
heads,
k_nope,
args.iters,
args.warmup,
device,
args.seed,
args.magnitude,
variants,
args.sweep,
)
if args.rounding_study:
rounding_study(device, args.seed)
if __name__ == "__main__":
main()