[kernel] Add opt-in deterministic CuTe-DSL DSA top-k backend (--dsa-topk-backend cutedsl)#31115
[kernel] Add opt-in deterministic CuTe-DSL DSA top-k backend (--dsa-topk-backend cutedsl)#31115zkyue wants to merge 3 commits into
Conversation
--dsa-topk-backend cutedsl: bitwise-deterministic exact top-k selection (canonical order: score desc, index asc on bitwise-equal scores; smallest-index tie subset at the k-th boundary; threshold-bin overflow falls back to full-row rescans instead of dropping candidates). Complements the deterministic-inference effort: the default radix kernel emits arrival-order results and can return a different index set per run under exact score ties, which realistic DSA indexer logits (relu zero plateau) hit routinely. Selection-only for now: like the torch backend it requires SGLANG_DSA_FUSE_TOPK=false.
zkyue
requested review from
1am9trash,
BBuf,
DarkSharpness,
Fridge003,
HaiShaw,
HydraQYH,
YAMY1234,
celve,
hubertlu-tw,
kkHuang-amd,
rainj-me and
yuan-luo
as code owners
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
…tests Host wrapper: replace asserts with actionable errors and enforce the supports() envelope (topk <= 4096, L < 2**24, B*L and B*topk < 2**31) so out-of-envelope calls fail loudly instead of silently mis-tie-breaking (L >= 2**24 wraps nidx24) or mis-addressing; move lengths/row_starts to score.device (int32) and validate their size, matching the torch backend; refuse first-touch allocation of the cached default-row_starts buffer under CUDA graph capture and make its zero-fill stream-safe; serialize cute.compile behind a lock, key the compile cache by device, and compile under score.device; validate SGLANG_DSA_TOPK_CUTEDSL_SURVCAP; document NaN ordering (bitwise total order, not torch.topk NaN-maximal) and the trusted-metadata window clamps. Tests: verify the advertised smallest-index tie subset at the k-th boundary (previously only order, count, and score multiset were checked); stable data seeds (zlib.crc32 instead of PYTHONHASHSEED-randomized hash()); skip cleanly when nvidia-cutlass-dsl is unavailable; cover an all--inf row with a nonzero window.
|
Pushed 48bce24 — pre-review hardening pass over the wrapper and tests:
All 12 registered tests pass on B200, and the determinism A/B harness stays bitwise-stable across runs on tie-heavy inputs. |
|
Intersting observation. I never expect there will be that many 0 in the indexer output (this is a degeneration case since we cannot distuiguish 0 values). |
|
Right — and the zeros are structural rather than a fixture artifact: the DSA lightning-indexer score is a sum of ReLU-ed per-head dot products, so every non-positive contribution lands on exactly 0.0. In realistic captures a sizable fraction of each row sits on that 0.0 plateau, and whenever a row's k-th boundary falls inside the plateau the tie set is huge — that is exactly where the default kernel's selection varies run-to-run (and where the #17747 overflow-drop bites hardest). Bitwise determinism therefore needs a total order among equal scores; this backend breaks ties by smallest index (a second radix pass over ~idx), which is what keeps replays byte-identical. Happy to add a short note about the tie/0.0-plateau semantics to the backend docs if that would help. |
|
@zkyue Could you please share some realistic data on this (e.g. score distribution for the indexer output in online serving)? I wonder whether our topk_v2 kernel (JIT kernel) is safe under this setting. |
|
Working on exactly that — I will follow up here with per-row score-distribution stats (fraction of exact-0.0 entries, tie-set size at the k-th boundary) from a reproducible setup, plus a small harness so you can measure the same on any live deployment, and a focused check of how topk_v2 behaves as the boundary tie-set grows. Should have numbers up shortly. |
|
Thanks for taking a look — here is the data we promised. Short answer to "is topk_v2 safe under this setting": in our tests it is memory-safe on arbitrarily large tie sets (no crashes, no illegal addresses, no Provenance, first. We cannot share score captures from our internal serving deployments, so everything here is synthetic and reproducible: the "realistic relu logits" used in this PR's tests are Why exact-0.0 plateaus are structural. In the DSA/lightning-indexer scoring as implemented in sglang ( When does the plateau reach the top-k boundary? It is a threshold phenomenon. The k-th boundary sits inside the 0.0 plateau iff
Either the tie set at the boundary is 1, or it is thousands of tokens — nothing in between. When it is thousands, which zeros pad the tail of the top-k is underdetermined, and those tokens still participate in sparse attention (an indexer score of exactly 0.0 does not mean zero attention weight), so the model output can differ between identical runs. topk_v2 behavior vs boundary tie-set size (
Reading of this: on pure exact-tie plateaus the kernel is well-behaved in the correctness sense — every run returns a valid top-k (dropped candidates are ties, so the score multiset stays right) — and the failure is reproducibility only: a different token set nearly every run once the boundary is inside the plateau. The realistic ReLU shape (plateau plus a spray of small positives) additionally produced a rare row where a strictly-positive candidate was replaced by a 0.0 one, and the narrow-range regime from #17747 drops true candidates on every row. No illegal behavior anywhere. The cutedsl backend in this PR does not change which regime traffic lands in — it only makes the tie-break defined (score desc, index asc: the smallest-index subset of the boundary tie set), so the selected set is a pure function of the input, run to run and across batch compositions. Both scripts are below; the first one is the interesting one to run on real serving traffic — capture the dsa_topk_score_stats.py — per-row tie/plateau stats for captured or synthetic indexer scores# dsa_topk_score_stats.py — per-row tie/plateau statistics for DSA indexer scores.
#
# Answers: "how big is the tie set at the top-k boundary of the scores that
# feed fast_topk_v2 / the DSA topk backend?"
#
# Two modes:
# 1) Real capture: python dsa_topk_score_stats.py --input scores.pt [--lengths lengths.pt]
# scores.pt = torch.save'd fp32 tensor (B, L), the exact `score` argument of
# DSATopKBackend.topk_func. Capture in a live deployment with a 2-line hook, e.g.
# in dsa_topk_backend.py right before the topk call:
# torch.save(score.float().cpu(), f"/tmp/idx_scores_{time.time_ns()}.pt")
# 2) Synthetic model of the indexer (no weights needed):
# python dsa_topk_score_stats.py --synthetic
# score[s] = ks[s] * sum_h w_h * relu(x[h,s]) with x[h,s] = per-head logit,
# modeled as correlated Gaussians: x[h,s] = mu + sqrt(rho)*c[s] + sqrt(1-rho)*z[h,s].
# rho = cross-head correlation of "is this key relevant", mu = mean logit.
# Exact-0.0 plateau fraction = P(all H head-logits <= 0); rho/mu sweep spans
# "no plateau" -> "row is mostly exact 0.0" so no single number is assumed.
import argparse
import torch
def row_stats(row: torch.Tensor, topks=(512, 1024, 2048)) -> dict:
"""row: 1-D fp32 scores (only the valid window). Returns tie stats per topk."""
n = row.numel()
out = {"len": n, "frac_zero": (row == 0.0).float().mean().item()}
vals = torch.sort(row, descending=True).values
for k in topks:
if n <= k:
out[k] = dict(boundary=None, tie=0, above=n, straddle=False, zero_bound=False)
continue
v = vals[k - 1].item() # k-th largest = boundary score
tie = (row == v).sum().item() # tie set containing the boundary
above = (row > v).sum().item() # rows strictly above: always selected
out[k] = dict(
boundary=v,
tie=tie,
above=above,
straddle=(tie > 1 and above + tie > k), # selection within tie set is underdetermined
zero_bound=(v == 0.0),
)
return out
def synthetic_rows(L, H=64, rho=0.9, mu=-1.0, n_rows=64, seed=0, device="cpu"):
"""Documented synthetic indexer: per-head logits are equicorrelated Gaussians."""
g = torch.Generator(device=device).manual_seed(seed)
c = torch.randn(n_rows, 1, L, device=device, generator=g) # shared relevance factor
z = torch.randn(n_rows, H, L, device=device, generator=g) # per-head noise
x = mu + (rho**0.5) * c + ((1 - rho) ** 0.5) * z
w = torch.rand(n_rows, H, 1, device=device, generator=g) / H # nonneg head weights
score = (w * torch.relu(x)).sum(dim=1) # (n_rows, L)
ks = torch.exp2(torch.randint(-2, 3, (1, L), device=device, generator=g).float())
return score * ks # e8m0 per-key scale: preserves exact 0.0
def report(rows_iter, topks):
import statistics as st
agg = {k: {"tie": [], "straddle": 0, "zero_bound": 0, "n": 0} for k in topks}
fz = []
for row in rows_iter:
s = row_stats(row.float(), topks)
fz.append(s["frac_zero"])
for k in topks:
if s[k]["boundary"] is None:
continue
a = agg[k]
a["tie"].append(s[k]["tie"])
a["straddle"] += s[k]["straddle"]
a["zero_bound"] += s[k]["zero_bound"]
a["n"] += 1
print(f" frac(score==0.0): mean {st.mean(fz):.3f} min {min(fz):.3f} max {max(fz):.3f}")
for k in topks:
a = agg[k]
if not a["n"]:
print(f" topk={k}: all rows shorter than k")
continue
t = sorted(a["tie"])
print(
f" topk={k}: boundary-tie-set size med {t[len(t)//2]} p90 {t[int(len(t)*0.9)]} "
f"max {t[-1]} | straddling rows {a['straddle']}/{a['n']} "
f"| boundary==0.0 rows {a['zero_bound']}/{a['n']}"
)
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--input", help=".pt tensor (B,L) of captured indexer scores")
p.add_argument("--lengths", help="optional .pt int tensor (B) of valid lengths")
p.add_argument("--synthetic", action="store_true")
p.add_argument("--topks", default="512,1024,2048")
args = p.parse_args()
topks = tuple(int(x) for x in args.topks.split(","))
if args.input:
sc = torch.load(args.input, map_location="cpu").float()
ln = torch.load(args.lengths, map_location="cpu") if args.lengths else None
rows = [sc[i, : (ln[i].item() if ln is not None else sc.shape[1])] for i in range(sc.shape[0])]
print(f"captured scores: {sc.shape[0]} rows")
report(rows, topks)
elif args.synthetic:
for L in (2048, 8192, 32768, 131072):
for rho, mu in ((0.0, -1.0), (0.5, -1.0), (0.9, -1.0), (0.9, -2.0), (0.99, -1.0)):
n_rows = max(8, min(64, 2**22 // L))
print(f"L={L} H=64 rho={rho} mu={mu} ({n_rows} rows)")
report(synthetic_rows(L, rho=rho, mu=mu, n_rows=n_rows), topks)
else:
p.error("need --input or --synthetic")topk_v2_tie_safety.py — the replay/tie-sweep used for the table above# topk_v2_tie_safety.py — sgl_kernel.fast_topk_v2 behavior vs boundary tie-set size.
#
# For each synthetic config we build B identical-distribution rows where the
# tie set containing the k-th boundary has a controlled size, replay the
# kernel N times (with a noisy sibling stream perturbing SM scheduling), and
# check per rep:
# - bitwise output hash + per-batch sorted-index-set hash (selection stability)
# - contract: -1 short-fill count (rows have >= k valid candidates, so any -1
# is a dropped candidate), duplicate indices, out-of-range indices
# - selected-score multiset vs torch.topk (tie-tolerant correctness):
# mismatch = candidates at/above the true boundary were dropped and
# replaced by strictly-worse ones (issue #17747 mode)
#
# Families:
# plateau: m distinct positive scores, rest exactly 0.0 (ReLU indexer shape).
# If m < k the boundary sits in the zero plateau: tie = L - m.
# tie@v: T scores exactly v=1.0 straddling the boundary, distinct values
# above and below (generic quantization-tie shape).
import argparse
import hashlib
import sys
import torch
from sgl_kernel import fast_topk_v2
def make_plateau(B, L, m, device, seed):
g = torch.Generator(device=device).manual_seed(seed)
sc = torch.zeros(B, L, device=device)
if m:
vals = 0.5 + torch.arange(m, device=device) * 1e-3 # distinct positives
for b in range(B):
pos = torch.randperm(L, device=device, generator=g)[:m]
sc[b, pos] = vals[torch.randperm(m, device=device, generator=g)]
return sc
def make_tie(B, L, k, T, device, seed):
g = torch.Generator(device=device).manual_seed(seed)
above = max(0, k - T // 2) # boundary lands inside the tie set
sc = torch.empty(B, L, device=device)
for b in range(B):
pos = torch.randperm(L, device=device, generator=g)
sc[b, pos[:above]] = 2.0 + torch.arange(above, device=device) * 1e-4
sc[b, pos[above : above + T]] = 1.0
nb = L - above - T
sc[b, pos[above + T :]] = 0.01 + torch.arange(nb, device=device) * 1e-7
return sc
def make_relu_pr(B, L, device, seed):
"""The (public) PR fixture: relu of shifted gaussian, ~10% nonzero.
At L=16384 nonzeros ~1638 < 2048 so the boundary sits in the 0.0 plateau."""
g = torch.Generator(device=device).manual_seed(seed)
return torch.relu(torch.randn(B, L, device=device, generator=g) - 1.28) * 11.3
def make_narrow(B, L, device, seed):
"""Distinct values in a narrow range (all within ~1e-6 of 1.0): the radix
threshold-bin staging-cap overflow regime (issue #17747)."""
g = torch.Generator(device=device).manual_seed(seed)
return 1.0 + torch.randn(B, L, device=device, generator=g) * 1e-7
def h(t):
return hashlib.sha1(t.cpu().numpy().tobytes()).hexdigest()[:12]
def check(sc, idx, k):
"""Returns (short_fill_rows, dup_rows, oob_rows, wrong_set_rows)."""
B, L = sc.shape
short = dup = oob = wrong = 0
ref = torch.sort(sc, dim=-1, descending=True).values[:, :k]
for b in range(B):
row = idx[b]
valid = row[row >= 0]
if (row >= L).any() or (row < -1).any():
oob += 1
continue
if valid.numel() < k:
short += 1
if valid.numel() != valid.unique().numel():
dup += 1
got = torch.sort(sc[b, valid.long()], descending=True).values
if got.numel() != k or not torch.equal(got, ref[b]):
wrong += 1
return short, dup, oob, wrong
def run(name, sc, k, reps, noise_mat, results):
B, L = sc.shape
ln = torch.full((B,), L, dtype=torch.int32, device=sc.device)
ns = torch.cuda.Stream()
outs, sets = set(), set()
worst = (0, 0, 0, 0)
err = ""
try:
for r in range(reps):
with torch.cuda.stream(ns):
for _ in range(4):
noise_mat @ noise_mat
idx = fast_topk_v2(sc, ln, k)
torch.cuda.synchronize()
outs.add(h(idx))
sets.add(h(torch.sort(idx, dim=-1).values))
worst = tuple(max(a, b) for a, b in zip(worst, check(sc, idx, k)))
except RuntimeError as e:
torch.cuda.synchronize()
err = str(e).splitlines()[0][:60]
results.append((name, B, L, k, len(outs), len(sets), *worst, err))
print(results[-1], flush=True)
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--gpu", type=int, default=0)
p.add_argument("--reps", type=int, default=25)
args = p.parse_args()
torch.cuda.set_device(args.gpu)
dev = "cuda"
noise_mat = torch.randn(2048, 2048, device=dev)
results = []
print("name B L k distinct_out distinct_set short dup oob wrong_set err", flush=True)
# fast_topk_v2 supports only topk=2048 (asserted in sgl_kernel wrapper;
# kernel is specialized for the DSv3.2 production config)
for B, L in ((64, 32768), (4, 32768), (8, 131072)):
for k in (2048,):
for m in (0, 256, 1024, 4095, 8192):
tie = L - m if m < k else 1
run(f"plateau_m{m}_tie{tie}", make_plateau(B, L, m, dev, 7), k, args.reps, noise_mat, results)
for T in (64, 512, 4096, 16384):
run(f"tie1.0_T{T}", make_tie(B, L, k, T, dev, 11), k, args.reps, noise_mat, results)
# PR-public relu fixture: L=16384 -> ~1638 nonzeros < 2048, boundary in plateau;
# L=32768 -> ~3277 nonzeros > 2048, boundary among distinct nonzeros (control)
for B, L in ((64, 16384), (256, 16384), (64, 32768)):
run(f"relu_pr_L{L}", make_relu_pr(B, L, dev, 13), 2048, args.reps, noise_mat, results)
# narrow-range distinct values: #17747 staging-cap overflow regime
for B, L in ((64, 32768), (8, 131072)):
run(f"narrow_L{L}", make_narrow(B, L, dev, 17), 2048, args.reps, noise_mat, results)
import csv
with open("topk_v2_tie_safety.csv", "w") as f:
w = csv.writer(f)
w.writerow(["name", "B", "L", "k", "distinct_out", "distinct_set", "short", "dup", "oob", "wrong_set", "err"])
w.writerows(results)
print("done", len(results), "configs -> topk_v2_tie_safety.csv") |
2 participants
Motivation
The deterministic-inference effort (#15041) currently has a gap at DSA indexer top-k selection:
sgl-kernelfast_topk_v2radix kernel is very fast, but by design it emits winners in shared-memory-atomic arrival order, resolves the k-th-boundary tie set by arrival order, and (as documented in Reduce topk kernel shared memory from 128KB to 32KB for better occupancy #17747) stages threshold-bin candidates in a fixed 4K-entry smem window where overflow candidates are gracefully skipped. These are reasonable throughput tradeoffs, but they mean the output — and under exact score ties even the selected token set — can differ run to run for identical input.sum_h relu(q·k) * w) contain a large exact-0.0 tie plateau wherever every head's dot product is negative. On such realistic inputs we measured the default kernel returning a different top-k index set on every invocation (see A/B below). Downstream, a different gather set breaks bitwise reproducibility of sparse attention even when everything else is deterministic.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC=1), which routes through the unfused masking wrapper.This PR adds an opt-in deterministic top-k backend so a deterministic DSA stack does not need to leave the default kernel family:
--dsa-topk-backend cutedsl.Modifications
python/sglang/jit_kernel/cutedsl_topk.py(new): CuTe-DSL SIMT top-k selector, one CTA per row, 512 threads. Coarse 8-bit histogram pass → MSB-first radix refinement to the exact threshold key → smallest-index tie resolution at the k-th boundary via an index radix-select → canonical descending bitonic sort of unique (score, index) composites. Properties:DSATopKBackend.topk_func:(score fp32 (B, L), lengths i32, topk, row_starts)→(B, topk) i32local indices, front-packed, −1 padded; exact-infexclusion matching thetorchbackend.topk ≤ 4096(production DSA topk = 2048), any row width < 2^24, SM90+, CUDA-graph-safe after one eager warmup per config (replay verified bitwise-equal to eager).python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py:DSATopKBackend.CUTEDSLenum value +topk_funcdispatch.python/sglang/srt/server_args.py:cutedsladded to--dsa-topk-backendchoices. Like the existingtorchbackend,cutedslcurrently requiresSGLANG_DSA_FUSE_TOPK=false(it implements selection, not the fused page-table transforms); help text updated accordingly.test/registered/jit/test_cutedsl_topk.py(new): contract tests (tie-robust exact-set parity vs atorch.topkoracle across randn / realistic-relu / quantized-tie /-infinputs, with and withoutrow_starts, edge lengths 0/topk/full-width), bitwise-determinism repeat test, CUDA-graph capture/replay test, empty-batch test.Accuracy Tests
Tie-robust checker (per row: index range/uniqueness/count, selected-score multiset bitwise-equal to the
torch.topkreference — a unique answer even under exact fp32 ties, unlike index comparison): all green for the new kernel across randn, realistic-relu plateau, quantized ties, all-equal rows (L=32k), narrow-range distributions (L=32k), sprinkled/all-infrows, 128k-wide decode rows,topk∈ {512, 1024, 2048, 4096}, with and withoutrow_starts.Determinism A/B on B200 (30 repetitions per input, identical input tensor, busy sibling stream to perturb scheduling; "sets" compares row-wise sorted index sets):
Speed Tests and Profiling
B200, topk=2048, paired both orders, µs/call. Caveat: the host was shared (74–89% background GPU load from other jobs) — values are min-of-rounds as the best uncontended estimate; ratios were stable across three runs.
flashinfer-detis the existing deterministic path measured as sglang dispatches it (_topk_unfusedmasking +flashinfer.top_k(deterministic=True)).Honest summary: the default kernel remains 2.3–4.8× faster on real-work shapes (more on short rows, where it writes identity indices without reading scores) — this PR does not propose changing the default. As the deterministic backend,
cutedslis faster than the current flashinfer deterministic path on 9/11 shapes and at parity on the rest, while additionally being exact under tie/overflow regimes.Known follow-up headroom: a 1024-thread variant (largest remaining lever for the decode gap), fp16 coarse first pass, host-launch pre-binding.
Related
#15041 (deterministic inference), #17747 / #16858 (top-k smem/perf background), flashinfer#2215 (deterministic top-k precedent).
Checklist
CI States
Latest PR Test (Base): ❌ Run #29315598173
Latest PR Test (Extra): ❌ Run #29315597939