From 9be573d223893e31591aac9328fe82d7a9853ab4 Mon Sep 17 00:00:00 2001 From: taking-lying-flat <1615405@qq.com> Date: Tue, 1 Sep 2026 20:02:08 +0800 Subject: [PATCH 1/4] fix(sampling): handle Triton top-k candidate limits --- python/freetoken/kernel/triton/sampling.py | 33 +++++++++-- tests/kernels/test_triton_sampling.py | 65 ++++++++++++++++++++++ 2 files changed, 93 insertions(+), 5 deletions(-) create mode 100644 tests/kernels/test_triton_sampling.py diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index 7345d65fc..caff156f8 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -555,13 +555,12 @@ def _topk_target(top_k, B, dev): return torch.full((B,), float(int(top_k)), device=dev, dtype=torch.float32) -def _topk_thr_ksum(probs, top_k): - """Return (threshold[B], kept_sum[B]) for a top-k keep: x >= threshold.""" +def _topk_thr_ksum(probs, target): + """Return threshold, buffered sum, and candidate count.""" B, V = probs.shape dev = probs.device G, CHUNK = _plan(B, V) grid = (B * G,) - target = _topk_target(top_k, B, dev) rmax = torch.zeros(B, device=dev, dtype=torch.float32) _rmax_pass[grid](probs, rmax, V, G, CHUNK, probs.stride(0), BLOCK_SIZE=2048, num_warps=8) buf = torch.empty(B * _CAP, device=dev, dtype=torch.float32) @@ -570,16 +569,40 @@ def _topk_thr_ksum(probs, top_k): thr = torch.empty(B, device=dev, dtype=torch.float32) ksum = torch.empty(B, device=dev, dtype=torch.float32) _refine_topk[(B,)](buf, cnt, rmax, target, thr, ksum, _CAP, _KR, _KBINS, 2048) - return thr, ksum + return thr, ksum, cnt + + +def _exact_top_k_renorm_probs(probs, target): + """Exact fallback for rows the fixed-size candidate buffer cannot represent.""" + _, V = probs.shape + target = target.to(torch.int64) + active_target = torch.where(target < V, target, 0) + max_top_k = max(1, int(active_target.max().item())) + values, indices = torch.topk(probs, max_top_k, dim=-1) + keep = torch.arange(max_top_k, device=probs.device).unsqueeze(0) < active_target.unsqueeze(1) + values = torch.where(keep, values, 0.0) + denom = values.sum(dim=-1, keepdim=True) + values = values / torch.where(denom > 0, denom, torch.ones_like(denom)) + out = torch.zeros_like(probs).scatter_(1, indices, values) + return torch.where((target >= V).unsqueeze(1), probs, out) def top_k_renorm_probs(probs, top_k): probs = probs.float() B, V = probs.shape + if not isinstance(top_k, torch.Tensor) and int(top_k) >= V: + return probs dev = probs.device + target = _topk_target(top_k, B, dev) + if (target >= V).any().item(): + return _exact_top_k_renorm_probs(probs, target) G, CHUNK = _plan(B, V) grid = (B * G,) - thr, ksum = _topk_thr_ksum(probs, top_k) + thr, _, cnt = _topk_thr_ksum(probs, target) + if ((cnt < target) | (cnt > _CAP)).any().item(): + return _exact_top_k_renorm_probs(probs, target) + ksum = torch.zeros(B, device=dev, dtype=torch.float32) + _ksum_pass[grid](probs, thr, ksum, V, G, CHUNK, probs.stride(0)) out = torch.empty_like(probs) _write_pass[grid](probs, out, thr, ksum, V, G, CHUNK, probs.stride(0)) return out diff --git a/tests/kernels/test_triton_sampling.py b/tests/kernels/test_triton_sampling.py new file mode 100644 index 000000000..24abd390f --- /dev/null +++ b/tests/kernels/test_triton_sampling.py @@ -0,0 +1,65 @@ +import pytest +import torch + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _reference_top_k_renorm(probs, top_k): + out = torch.zeros_like(probs) + vocab_size = probs.shape[1] + for row, k in enumerate(top_k.tolist()): + if k >= vocab_size: + out[row] = probs[row] + continue + values, indices = torch.topk(probs[row], k) + out[row, indices] = values / values.sum() + return out + + +def test_top_k_renorm_falls_back_when_too_few_candidates(): + from freetoken.kernel.triton.sampling import top_k_renorm_probs + + vocab_size = 10_000 + probs = torch.full((1, vocab_size), 0.8 / (vocab_size - 1), device="cuda") + probs[0, 0] = 0.2 + top_k = torch.tensor([50], dtype=torch.int32, device="cuda") + + actual = top_k_renorm_probs(probs, top_k) + expected = _reference_top_k_renorm(probs, top_k.cpu()) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual.sum(dim=-1), torch.ones(1, device="cuda")) + + +def test_top_k_renorm_falls_back_on_candidate_overflow(): + from freetoken.kernel.triton.sampling import top_k_renorm_probs + + vocab_size = 32_768 + probs = torch.linspace(1.0, 0.99, vocab_size, device="cuda").unsqueeze(0) + probs /= probs.sum(dim=-1, keepdim=True) + top_k = torch.tensor([50], dtype=torch.int32, device="cuda") + + actual = top_k_renorm_probs(probs, top_k) + expected = _reference_top_k_renorm(probs, top_k.cpu()) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual.sum(dim=-1), torch.ones(1, device="cuda")) + + +def test_top_k_renorm_mixed_batch_bypasses_vocab_size_row(): + from freetoken.kernel.triton.sampling import top_k_renorm_probs + + vocab_size = 10_000 + peaky = torch.full((vocab_size,), 0.8 / (vocab_size - 1), device="cuda") + peaky[0] = 0.2 + unfiltered = torch.linspace(1.0, 0.5, vocab_size, device="cuda") + unfiltered /= unfiltered.sum() + probs = torch.stack((peaky, unfiltered)) + top_k = torch.tensor([50, vocab_size], dtype=torch.int32, device="cuda") + + actual = top_k_renorm_probs(probs, top_k) + expected = _reference_top_k_renorm(probs, top_k.cpu()) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual.sum(dim=-1), torch.ones(2, device="cuda")) From de5b6cb947ad6a02b69008dcc4cedcb7d0484f77 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 2 Sep 2026 06:36:17 +0000 Subject: [PATCH 2/4] fix(kernel): exact single-launch triton top-k/top-p sampling --- python/freetoken/kernel/triton/sampling.py | 558 ++++++++++++++------- 1 file changed, 376 insertions(+), 182 deletions(-) diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index caff156f8..65237a68e 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -1,4 +1,4 @@ -"""Multi-CTA (split-vocab) Triton sampling ops (provenance: sampling<-vllm Qrita). +"""Multi-CTA (split-vocab) Triton sampling ops. Optional pure-triton drop-in for freetoken.kernel.sampling / flashinfer.sampling (softmax / top-k / top-p / combined + draw), self-contained. @@ -6,28 +6,33 @@ Design: * Every row is split across many CTAs (``_plan`` -> G column-chunks) so bs=1 uses the whole GPU, unlike a single-block-per-row kernel that is single-SM-bound. - * softmax is a multi-CTA online softmax; top-p uses a small fixed number of - histogram-bracket refinement passes instead of a ~48-iter bisection; the draw - is a multi-CTA inverse-CDF. - * The top-k path is adapted from vLLM's Qrita kernel - (v1/sample/ops/topk_topp_triton.py::_topk_topp_kernel): gather the small set of - "outlier" candidates (probs >= rmax*FRAC) into a compact per-row buffer in ONE - full-vocab pass, then run the k-th-value search on that tiny buffer (3 full-vocab - passes total vs ~6 for a pure histogram top-k). The outlier-pivot heuristic is - swapped for the probs domain (truncate at rmax*FRAC; softmax probs are not - Gaussian). Rows overflowing CAP silently drop the smallest gathered candidates; - the refine still finds the exact k-th since every value >= threshold is kept, and - an in-kernel guard keeps everything if fewer than k finite candidates are gathered. + * softmax is a multi-CTA online softmax; the draw is a multi-CTA inverse-CDF. + * top-k and top-p are each one cooperative kernel: the row's CTAs each bin their chunk over + the fp32 bit pattern (order-preserving for x >= 0), meet at a per-row spin barrier, and all + redo the refine so they share the bracket. Four rounds of 256 bins bring the 2**31 range + down to one bit pattern, so the threshold is exactly the k-th largest prob (top-k, counts) + or the value where the descending cumulative mass reaches p (top-p, exact per-bin mass); + the same kernel then renormalizes or draws. No candidate buffer, no data-dependent shape, + no host sync. Results are exact up to fp32 atomic summation order. + * If the cooperative launch is unavailable the module falls back to the multi-launch + search below (exact for top-k; top-p there is the older bin-center estimate). + * deterministic, generator and check_nan exist for flashinfer signature compatibility and are + ignored; seed and offset are honored. Given a seed, top-k draws reproduce; top-p may pick a + different token on rows whose cumulative mass sits within fp32 rounding of p. """ from __future__ import annotations +import logging + import torch import triton import triton.language as tl from freetoken.kernel.triton.autotune_cache import autotune_cache_kwargs +logger = logging.getLogger(__name__) + _NUM_SM = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count _MIN_CHUNK = 4096 # do not split a row finer than this @@ -145,7 +150,7 @@ def softmax(logits, temperature=None, enable_pdl=None): # multi-CTA top-p via histogram-bracket refinement. Every full-vocab pass is # split across all SMs; the sequential refine step is a tiny grid=(B,) kernel. # =========================================================================== -_PBINS = 64 # top-p uses count-hist + bin-center mass (64**4 ~ 1.7e7) +_PBINS = 64 # fallback top-p only: count-hist + bin-center mass (64**4 ~ 1.7e7) _PR = 4 _SR_CFGS = [ @@ -172,17 +177,21 @@ def _rmax_pass(probs_ptr, rmax_ptr, V, G, CHUNK, row_stride, BLOCK_SIZE: tl.cons tl.atomic_max(rmax_ptr + row, m) -@triton.autotune(configs=_SR_CFGS, key=["CHUNK"], reset_to_zero=["hist_ptr"], **autotune_cache_kwargs) +@triton.autotune(configs=_SR_CFGS, key=["CHUNK", "BINS", "BITS"], reset_to_zero=["hist_ptr"], **autotune_cache_kwargs) @triton.jit def _count_hist_pass( probs_ptr, lo_ptr, hi_ptr, hist_ptr, V, G, CHUNK, row_stride, - BINS: tl.constexpr, BLOCK_SIZE: tl.constexpr, + BINS: tl.constexpr, BITS: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): + # BITS: lo/hi are int32 bit patterns and bins are exact integer ranges; else float value bins pid = tl.program_id(0) row = pid // G lo = tl.load(lo_ptr + row) hi = tl.load(hi_ptr + row) - invw = BINS / tl.maximum(hi - lo, 1e-30) + if BITS: + w = (hi - lo + BINS - 1) // BINS + else: + invw = BINS / tl.maximum(hi - lo, 1e-30) base = row * row_stride start = (pid % G) * CHUNK end = tl.minimum(start + CHUNK, V) @@ -191,8 +200,13 @@ def _count_hist_pass( offs = s0 + tl.arange(0, BLOCK_SIZE) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=-1.0).to(tl.float32) - inrange = mask & (x >= lo) & (x < hi) - b = ((x - lo) * invw).to(tl.int32) + if BITS: + y = x.to(tl.int32, bitcast=True) + b = (y - lo) // w + else: + y = x + b = ((x - lo) * invw).to(tl.int32) + inrange = mask & (y >= lo) & (y < hi) # tl.histogram does NOT cleanly drop out-of-range indices; route every # out-of-bracket element to bin 0 and then subtract that count back out so # the histogram holds ONLY in-[lo,hi) counts (out-of-range is tracked via @@ -206,13 +220,16 @@ def _count_hist_pass( @triton.jit -def _refine_pass(lo_ptr, hi_ptr, above_ptr, hist_ptr, target_ptr, BINS: tl.constexpr): +def _refine_pass(lo_ptr, hi_ptr, above_ptr, hist_ptr, target_ptr, BINS: tl.constexpr, BITS: tl.constexpr): row = tl.program_id(0) lo = tl.load(lo_ptr + row) hi = tl.load(hi_ptr + row) above = tl.load(above_ptr + row) target = tl.load(target_ptr + row) - w = (hi - lo) / BINS + if BITS: + w = (hi - lo + BINS - 1) // BINS + else: + w = (hi - lo) / BINS jj = tl.arange(0, BINS) h = tl.load(hist_ptr + row * BINS + jj) prefix = tl.cumsum(h, 0) @@ -222,8 +239,11 @@ def _refine_pass(lo_ptr, hi_ptr, above_ptr, hist_ptr, target_ptr, BINS: tl.const j = tl.max(tl.where(ok, jj, -1)) prefix_j = tl.sum(tl.where(jj <= j, h, 0.0)) upd = j >= 0 + new_hi = lo + (j + 1) * w + if BITS: + new_hi = tl.minimum(new_hi, hi) tl.store(lo_ptr + row, tl.where(upd, lo + j * w, lo)) - tl.store(hi_ptr + row, tl.where(upd, lo + (j + 1) * w, hi)) + tl.store(hi_ptr + row, tl.where(upd, new_hi, hi)) tl.store(above_ptr + row, tl.where(upd, above + total - prefix_j, above)) # zero the row so the next iteration's atomic_add starts clean (reset_to_zero # only fires during autotuning, not on production calls) @@ -292,7 +312,7 @@ def _write_pass(probs_ptr, out_ptr, thr_ptr, ksum_ptr, V, G, CHUNK, row_stride, tl.store(out_ptr + base + offs, tl.where(x >= thr, x * inv_s, 0.0), mask=mask) -def _search(probs, target, mass, R, BINS): +def _search(probs, target, mass, R, BINS, bits=False): """Return per-row threshold: keep x >= thr, with count/mass(>=thr) ~ target.""" B, V = probs.shape dev = probs.device @@ -300,17 +320,21 @@ def _search(probs, target, mass, R, BINS): grid = (B * G,) rmax = torch.zeros(B, device=dev, dtype=torch.float32) _rmax_pass[grid](probs, rmax, V, G, CHUNK, probs.stride(0), BLOCK_SIZE=2048, num_warps=8) - lo = torch.zeros(B, device=dev, dtype=torch.float32) - hi = (rmax * 1.0000001).contiguous() + if bits: + lo = torch.zeros(B, device=dev, dtype=torch.int32) + hi = (rmax.view(torch.int32) + 1).contiguous() + else: + lo = torch.zeros(B, device=dev, dtype=torch.float32) + hi = (rmax * 1.0000001).contiguous() above = torch.zeros(B, device=dev, dtype=torch.float32) hist = torch.zeros(B * BINS, device=dev, dtype=torch.float32) for _ in range(R): - _count_hist_pass[grid](probs, lo, hi, hist, V, G, CHUNK, probs.stride(0), BINS) + _count_hist_pass[grid](probs, lo, hi, hist, V, G, CHUNK, probs.stride(0), BINS, bits) if mass: _refine_mass_pass[(B,)](lo, hi, above, hist, target, BINS) else: - _refine_pass[(B,)](lo, hi, above, hist, target, BINS) - return lo + _refine_pass[(B,)](lo, hi, above, hist, target, BINS, bits) + return lo.view(torch.float32) if bits else lo def _renorm(probs, thr): @@ -327,13 +351,7 @@ def _renorm(probs, thr): def top_p_renorm_probs(probs, top_p): probs = probs.float() - B, V = probs.shape - if isinstance(top_p, torch.Tensor): - target = top_p.float().to(probs.device).contiguous() - else: - target = torch.full((B,), float(top_p), device=probs.device, dtype=torch.float32) - thr = _search(probs, target, True, _PR, _PBINS) - return _renorm(probs, thr) + return _topp(probs, _topp_target(top_p, probs.size(0), probs.device), None, False) # --------------------------------------------------------------------------- @@ -358,38 +376,47 @@ def _draw_part(probs_ptr, thr_ptr, psum_ptr, V, G, CHUNK, row_stride, BLOCK_SIZE @triton.jit -def _draw_scan(psum_ptr, choff_ptr, u_ptr, target_ptr, G, G_POW2: tl.constexpr): +def _draw_scan(psum_ptr, choff_ptr, u_ptr, target_ptr, last_ptr, G, G_POW2: tl.constexpr): row = tl.program_id(0) goff = tl.arange(0, G_POW2) gmask = goff < G ps = tl.load(psum_ptr + row * G + goff, mask=gmask, other=0.0) tl.store(choff_ptr + row * G + goff, tl.cumsum(ps, 0) - ps, mask=gmask) tl.store(target_ptr + row, tl.load(u_ptr + row) * tl.sum(ps, 0)) + tl.store(last_ptr + row, tl.max(tl.where(gmask & (ps > 0), goff, -1), 0)) @triton.autotune(configs=_SR_CFGS, key=["CHUNK"], **autotune_cache_kwargs) @triton.jit -def _draw_find(probs_ptr, thr_ptr, choff_ptr, target_ptr, out_ptr, V, G, CHUNK, row_stride, +def _draw_find(probs_ptr, thr_ptr, choff_ptr, target_ptr, psum_ptr, last_ptr, out_ptr, V, G, CHUNK, row_stride, BLOCK_SIZE: tl.constexpr): pid = tl.program_id(0) row = pid // G thr = tl.load(thr_ptr + row) target = tl.load(target_ptr + row) acc = tl.load(choff_ptr + pid) + incl = acc + tl.load(psum_ptr + pid) base = row * row_stride start = (pid % G) * CHUNK end = tl.minimum(start + CHUNK, V) + last_kept = start * 0 - 1 for s0 in tl.range(start, end, BLOCK_SIZE): offs = s0 + tl.arange(0, BLOCK_SIZE) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - wv = tl.where((x >= thr) & mask, x, 0.0) + kept = (x >= thr) & mask + wv = tl.where(kept, x, 0.0) cval = acc + tl.cumsum(wv, 0) idx = tl.where(cval > target, offs, V) blk_min = tl.min(idx, 0) if (blk_min < V) and (acc <= target): tl.store(out_ptr + row, blk_min) acc += tl.sum(wv, 0) + last_kept = tl.maximum(last_kept, tl.max(tl.where(kept, offs, -1), 0)) + # see _keep_tail: the CTA owning an fp rounding gap writes its last kept token + is_last_mass = tl.load(last_ptr + row) == pid % G + if (acc <= target) and (last_kept >= 0) and ((incl > target) or is_last_mass): + tl.store(out_ptr + row, last_kept) _UGEN = {} @@ -417,11 +444,12 @@ def _draw(probs, thr, seed, offset): psum = torch.empty(B * G, device=dev, dtype=torch.float32) choff = torch.empty(B * G, device=dev, dtype=torch.float32) target = torch.empty(B, device=dev, dtype=torch.float32) - out = torch.empty(B, device=dev, dtype=torch.int32) + last = torch.empty(B, device=dev, dtype=torch.int32) + out = torch.zeros(B, device=dev, dtype=torch.int32) u = _gen_u(B, dev, seed, offset) _draw_part[grid](probs, thr, psum, V, G, CHUNK, probs.stride(0)) - _draw_scan[(B,)](psum, choff, u, target, G, _next_pow2(G)) - _draw_find[grid](probs, thr, choff, target, out, V, G, CHUNK, probs.stride(0)) + _draw_scan[(B,)](psum, choff, u, target, last, G, _next_pow2(G)) + _draw_find[grid](probs, thr, choff, target, psum, last, out, V, G, CHUNK, probs.stride(0)) return out @@ -442,178 +470,349 @@ def top_p_sampling_from_probs(probs, top_p, indices=None, deterministic=True, ge check_nan=False, seed=None, offset=None, return_valid=False): probs = probs.float() src = probs if indices is None else probs[indices].contiguous() - if isinstance(top_p, torch.Tensor): - target = top_p.float().to(src.device).contiguous() - else: - target = torch.full((src.size(0),), float(top_p), device=src.device, dtype=torch.float32) - thr = _search(src, target, True, _PR, _PBINS) - out = _draw(src, thr, seed, offset) + out = _topp(src, _topp_target(top_p, src.size(0), src.device), None, True, seed, offset) out = out.to(indices.dtype) if indices is not None else out return (out, torch.ones_like(out, dtype=torch.bool)) if return_valid else out # =========================================================================== -# top-k via Qrita outlier-gather + tiny-buffer refine (adapted from vLLM) -# Passes: rmax (full) + gather (full) + refine (tiny, buffer-only) + write (full) -# = 3 full-vocab passes. Keeps the multi-CTA vocab split so bs=1 uses the whole GPU. +# top-k: one cooperative kernel per call. Every CTA of a row histograms its column chunk +# over the fp32 bit pattern, the row's CTAs meet at a spin barrier, then each one redoes +# the tiny refine step so all of them hold the same bracket. Four rounds (exponent, then +# 8+8+7 mantissa bits) end on a single bit pattern, so thr is exactly the k-th largest +# prob. The same kernel then either renormalizes (DRAW=0) or draws a token (DRAW=1). # =========================================================================== -_FRAC = 0.05 # outlier gather pivot: keep probs >= rmax*FRAC -_CAP = 8192 # per-row candidate buffer capacity -_KBINS = 256 # bins per refine bracket iteration -_KR = 3 # refine iterations (256**3 ~ 1.7e7 resolution over [0, rmax]) +_KBINS = 256 +_INF_BITS = tl.constexpr(0x7F800000) +_FUSED_BLOCK = 2048 + + +def _topk_target(top_k, B, dev): + if isinstance(top_k, torch.Tensor): + target = top_k.float().to(dev).contiguous() + else: + target = torch.full((B,), float(int(top_k)), device=dev, dtype=torch.float32) + # k <= 0 would satisfy every bin and push the bracket past the max; k > V needs no clamp, the search leaves lo at 0 + return torch.clamp(target, min=1.0) + + +@triton.jit +def _row_barrier(bar_ptr, need): + # the row's CTAs must be co-resident (cooperative launch), or a lone CTA (G == 1) passes at once. + # every warp's preceding atomics must be issued before thread 0 announces arrival + tl.debug_barrier() + tl.atomic_add(bar_ptr, 1) + n = tl.atomic_add(bar_ptr, 0) + while n < need: + n = tl.atomic_add(bar_ptr, 0) + + +@triton.jit +def _bits_round( + probs_ptr, base, start, end, hist_ptr, bar_ptr, target, lo, above, need, + S: tl.constexpr, WIDTH: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, +): + jj = tl.arange(0, BINS) + acc = tl.zeros([BINS], tl.int32) + for s0 in tl.range(start, end, BLOCK): + offs = s0 + tl.arange(0, BLOCK) + mask = offs < end + y = tl.load(probs_ptr + base + offs, mask=mask, other=-1.0).to(tl.float32).to(tl.int32, bitcast=True) + d = y - lo + if WIDTH == 0: + inrange = mask & (y >= lo) & (y <= _INF_BITS) + else: + inrange = mask & (y >= lo) & (d < WIDTH) & (y <= _INF_BITS) + # every out-of-bracket lane (padding included) lands in bin 0 and is subtracted back out + b = tl.where(inrange, d >> S, 0) + h = tl.histogram(b, BINS) + acc += h - tl.where(jj == 0, tl.sum((~inrange).to(tl.int32)), 0) + tl.atomic_add(hist_ptr + jj, acc) + _row_barrier(bar_ptr, need) + h = tl.load(hist_ptr + jj, cache_modifier=".cg") + prefix = tl.cumsum(h, 0) + total = tl.sum(h, 0) + ok = (above + total - prefix + h).to(tl.float32) >= target + j = tl.max(tl.where(ok, jj, -1)) + prefix_j = tl.sum(tl.where(jj <= j, h, 0)) + upd = j >= 0 + lo = tl.where(upd, lo + (j << S), lo) + above = tl.where(upd, above + total - prefix_j, above) + return lo, above -@triton.autotune(configs=_SR_CFGS, key=["CHUNK"], reset_to_zero=["cnt_ptr"], **autotune_cache_kwargs) @triton.jit -def _gather_pass( - probs_ptr, rmax_ptr, buf_ptr, cnt_ptr, V, G, CHUNK, row_stride, - FRAC, CAP: tl.constexpr, BLOCK_SIZE: tl.constexpr, +def _topk_fused( + probs_ptr, target_ptr, lo0_ptr, hist_ptr, bar_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, + V, G, CHUNK, row_stride, + DRAW: tl.constexpr, G_POW2: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, ): pid = tl.program_id(0) row = pid // G - thr0 = tl.load(rmax_ptr + row) * FRAC + cta = pid % G base = row * row_stride - bufbase = row * CAP - start = (pid % G) * CHUNK + start = cta * CHUNK end = tl.minimum(start + CHUNK, V) - for s0 in tl.range(start, end, BLOCK_SIZE): - offs = s0 + tl.arange(0, BLOCK_SIZE) + target = tl.load(target_ptr + row) + hrow = hist_ptr + row * 4 * BINS + brow = bar_ptr + row + lo = tl.load(lo0_ptr + row) + above = lo + lo, above = _bits_round(probs_ptr, base, start, end, hrow, brow, target, lo, above, G, 23, 0, BINS, BLOCK) + lo, above = _bits_round(probs_ptr, base, start, end, hrow + BINS, brow, target, lo, above, 2 * G, 15, 1 << 23, BINS, BLOCK) + lo, above = _bits_round(probs_ptr, base, start, end, hrow + 2 * BINS, brow, target, lo, above, 3 * G, 7, 1 << 15, BINS, BLOCK) + lo, above = _bits_round(probs_ptr, base, start, end, hrow + 3 * BINS, brow, target, lo, above, 4 * G, 0, 1 << 7, BINS, BLOCK) + _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), brow, 5 * G, + ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, DRAW, G_POW2, BLOCK) + + +@triton.jit +def _keep_tail( + probs_ptr, base, start, end, pid, row, cta, thr, bar_ptr, need, + ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, + DRAW: tl.constexpr, G_POW2: tl.constexpr, BLOCK: tl.constexpr, +): + # keep x >= thr over this chunk: DRAW picks one token per row, else renormalize the kept mass to 1 + s = 0.0 + for s0 in tl.range(start, end, BLOCK): + offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - m = mask & (x >= thr0) - mi = m.to(tl.int32) - n = tl.sum(mi) - posbase = tl.atomic_add(cnt_ptr + row, n) - cpos = posbase + tl.cumsum(mi, 0) - 1 - wmask = m & (cpos < CAP) - tl.store(buf_ptr + bufbase + cpos, x, mask=wmask) + s += tl.sum(tl.where(x >= thr, x, 0.0), 0) + if DRAW: + tl.store(psum_ptr + pid, s) + _row_barrier(bar_ptr, need) + goff = tl.arange(0, G_POW2) + gmask = goff < G + ps = tl.load(psum_ptr + row * G + goff, mask=gmask, other=0.0, cache_modifier=".cg") + acc = tl.sum(tl.where(goff < cta, ps, 0.0), 0) + incl = tl.sum(tl.where(goff <= cta, ps, 0.0), 0) + tgt = tl.load(u_ptr + row) * tl.sum(ps, 0) + last_kept = start * 0 - 1 + for s0 in tl.range(start, end, BLOCK): + offs = s0 + tl.arange(0, BLOCK) + mask = offs < end + x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + kept = (x >= thr) & mask + wv = tl.where(kept, x, 0.0) + cval = acc + tl.cumsum(wv, 0) + idx = tl.where(cval > tgt, offs, V) + blk_min = tl.min(idx, 0) + if (blk_min < V) and (acc <= tgt): + tl.store(tok_ptr + row, blk_min) + acc += tl.sum(wv, 0) + last_kept = tl.maximum(last_kept, tl.max(tl.where(kept, offs, -1), 0)) + # fp rounding can leave tgt between this CTA's running sum and the next CTA's prefix, or past the total; + # the CTA that owns that gap (or the last one holding mass) writes its last kept token instead + is_last_mass = tl.sum(tl.where((goff > cta) & (ps > 0), 1, 0), 0) == 0 + if (acc <= tgt) and (last_kept >= 0) and ((incl > tgt) or is_last_mass): + tl.store(tok_ptr + row, last_kept) + else: + tl.atomic_add(ksum_ptr + row, s) + _row_barrier(bar_ptr, need) + inv = 1.0 / tl.atomic_add(ksum_ptr + row, 0.0) + for s0 in tl.range(start, end, BLOCK): + offs = s0 + tl.arange(0, BLOCK) + mask = offs < end + x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(out_ptr + base + offs, tl.where(x >= thr, x * inv, 0.0), mask=mask) + + +_PMBINS = 256 @triton.jit -def _refine_topk( - buf_ptr, cnt_ptr, rmax_ptr, target_ptr, thr_ptr, ksum_ptr, - CAP: tl.constexpr, R: tl.constexpr, BINS: tl.constexpr, BLK: tl.constexpr, +def _pmass_round( + probs_ptr, base, start, end, priv_ptr, mass_ptr, bar_ptr, target, lo, above, need, + S: tl.constexpr, WIDTH: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, ): - # single-CTA-per-row refine on the tiny buffer -> (threshold, kept-sum): - # histogram-bracket the k-th largest, then sum the kept mass (all kept values - # are in the buffer since threshold >= gather pivot). - row = tl.program_id(0) - cnt = tl.load(cnt_ptr + row) - cnt = tl.minimum(cnt, CAP) - target = tl.load(target_ptr + row) - base = row * CAP + # top-p round over the bit pattern: per-bin MASS (exact up to fp32 atomic order) via scatter-add into this + # CTA's private buffer, then one reduction into the row buffer, so the bin holding the p crossing is known jj = tl.arange(0, BINS) - lo = 0.0 - hi = tl.load(rmax_ptr + row) * 1.0000001 + for s0 in tl.range(start, end, BLOCK): + offs = s0 + tl.arange(0, BLOCK) + mask = offs < end + x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + y = x.to(tl.int32, bitcast=True) + d = y - lo + if WIDTH == 0: + inrange = mask & (y >= lo) & (y <= _INF_BITS) + else: + inrange = mask & (y >= lo) & (d < WIDTH) & (y <= _INF_BITS) + tl.atomic_add(priv_ptr + tl.where(inrange, d >> S, 0), x, mask=inrange) + # every warp's scatter-adds must land before any thread reads the private bins back + tl.debug_barrier() + tl.atomic_add(mass_ptr + jj, tl.load(priv_ptr + jj)) + _row_barrier(bar_ptr, need) + m = tl.load(mass_ptr + jj, cache_modifier=".cg") + prefix = tl.cumsum(m, 0) + total = tl.sum(m, 0) + ok = above + total - prefix + m >= target + # p above the total mass (fp rounding at p = 1): keep the whole bracket + j = tl.maximum(tl.max(tl.where(ok, jj, -1)), 0) + prefix_j = tl.sum(tl.where(jj <= j, m, 0.0)) + return lo + (j << S), above + total - prefix_j + + +@triton.jit +def _topp_fused( + probs_ptr, tp_ptr, tk_ptr, lo0_ptr, hist_ptr, priv_ptr, mass_ptr, bar_ptr, ksumk_ptr, ksum_ptr, psum_ptr, u_ptr, + out_ptr, tok_ptr, V, G, CHUNK, row_stride, + TOPK: tl.constexpr, DRAW: tl.constexpr, G_POW2: tl.constexpr, KBINS: tl.constexpr, PBINS: tl.constexpr, + BLOCK: tl.constexpr, +): + # top-p, optionally after an exact top-k stage: the top-k threshold becomes the lower edge of the top-p + # bracket and the p target is scaled by the kept top-k mass, so no renormalized copy is ever written + pid = tl.program_id(0) + row = pid // G + cta = pid % G + base = row * row_stride + start = cta * CHUNK + end = tl.minimum(start + CHUNK, V) + brow = bar_ptr + row + lo = tl.load(lo0_ptr + row) + if TOPK: + tk = tl.load(tk_ptr + row) + hk = hist_ptr + row * 4 * KBINS + above_i = lo + lo, above_i = _bits_round(probs_ptr, base, start, end, hk, brow, tk, lo, above_i, G, 23, 0, KBINS, BLOCK) + lo, above_i = _bits_round(probs_ptr, base, start, end, hk + KBINS, brow, tk, lo, above_i, 2 * G, 15, 1 << 23, KBINS, BLOCK) + lo, above_i = _bits_round(probs_ptr, base, start, end, hk + 2 * KBINS, brow, tk, lo, above_i, 3 * G, 7, 1 << 15, KBINS, BLOCK) + lo, above_i = _bits_round(probs_ptr, base, start, end, hk + 3 * KBINS, brow, tk, lo, above_i, 4 * G, 0, 1 << 7, KBINS, BLOCK) + thr_k = lo.to(tl.float32, bitcast=True) + s = 0.0 + for s0 in tl.range(start, end, BLOCK): + offs = s0 + tl.arange(0, BLOCK) + mask = offs < end + x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + s += tl.sum(tl.where(x >= thr_k, x, 0.0), 0) + tl.atomic_add(ksumk_ptr + row, s) + _row_barrier(brow, 5 * G) + target = tl.load(tp_ptr + row) * tl.atomic_add(ksumk_ptr + row, 0.0) + done = 5 + else: + target = tl.load(tp_ptr + row) + done = 0 + mp = mass_ptr + row * 4 * PBINS + pp = priv_ptr + pid * 4 * PBINS above = 0.0 - for _it in tl.static_range(R): - denom = tl.maximum(hi - lo, 1e-30) - w = denom / BINS - invw = BINS / denom - hc = tl.zeros([BINS], tl.int32) - for s0 in tl.range(0, cnt, BLK): - offs = s0 + tl.arange(0, BLK) - mask = offs < cnt - x = tl.load(buf_ptr + base + offs, mask=mask, other=-1.0) - inrange = mask & (x >= lo) & (x < hi) - b = ((x - lo) * invw).to(tl.int32) - b = tl.where(inrange, tl.maximum(0, tl.minimum(b, BINS - 1)), 0) - hcnt = tl.histogram(b, BINS) - noor = tl.sum((mask & (~inrange)).to(tl.int32)) - hc += hcnt - tl.where(jj == 0, noor, 0) - h = hc.to(tl.float32) - prefix = tl.cumsum(h, 0) - total = tl.sum(h, 0) - c_ge = above + total - prefix + h - ok = c_ge >= target - j = tl.max(tl.where(ok, jj, -1)) - prefix_j = tl.sum(tl.where(jj <= j, h, 0.0)) - upd = j >= 0 - new_lo = lo + j * w - new_hi = lo + (j + 1) * w - new_above = above + total - prefix_j - lo = tl.where(upd, new_lo, lo) - hi = tl.where(upd, new_hi, hi) - above = tl.where(upd, new_above, above) - thr = lo - # guard: if fewer finite candidates than k were gathered, keep everything. - if cnt < target: - thr = 0.0 - ks = 0.0 - for s0 in tl.range(0, cnt, BLK): - offs = s0 + tl.arange(0, BLK) - mask = offs < cnt - x = tl.load(buf_ptr + base + offs, mask=mask, other=0.0) - ks += tl.sum(tl.where(x >= thr, x, 0.0)) - tl.store(thr_ptr + row, thr) - tl.store(ksum_ptr + row, ks) + lo, above = _pmass_round(probs_ptr, base, start, end, pp, mp, brow, target, lo, above, (done + 1) * G, 23, 0, PBINS, BLOCK) + lo, above = _pmass_round(probs_ptr, base, start, end, pp + PBINS, mp + PBINS, brow, target, lo, above, (done + 2) * G, 15, 1 << 23, PBINS, BLOCK) + lo, above = _pmass_round(probs_ptr, base, start, end, pp + 2 * PBINS, mp + 2 * PBINS, brow, target, lo, above, (done + 3) * G, 7, 1 << 15, PBINS, BLOCK) + lo, above = _pmass_round(probs_ptr, base, start, end, pp + 3 * PBINS, mp + 3 * PBINS, brow, target, lo, above, (done + 4) * G, 0, 1 << 7, PBINS, BLOCK) + _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), brow, (done + 5) * G, + ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, DRAW, G_POW2, BLOCK) -def _topk_target(top_k, B, dev): - if isinstance(top_k, torch.Tensor): - return top_k.float().to(dev).contiguous() - return torch.full((B,), float(int(top_k)), device=dev, dtype=torch.float32) +_fused_ok = True +_COOP_CTAS_PER_SM = 2 # the fused kernels use ~80 regs/thread at 8 warps; 4/SM fails the cooperative launch + + +def _fused_plan(B, V): + # the cooperative launch needs the whole grid co-resident, so cap B*G by an occupancy budget instead of _plan's one CTA per SM + g_by_sm = max(1, (_COOP_CTAS_PER_SM * _NUM_SM) // B) + g_by_chunk = max(1, triton.cdiv(V, _MIN_CHUNK)) + G = min(g_by_sm, g_by_chunk) + return G, triton.cdiv(V, G) -def _topk_thr_ksum(probs, target): - """Return threshold, buffered sum, and candidate count.""" +def _fused_launch(probs, kernel, tk, tp, draw, seed, offset): B, V = probs.shape dev = probs.device - G, CHUNK = _plan(B, V) - grid = (B * G,) - rmax = torch.zeros(B, device=dev, dtype=torch.float32) - _rmax_pass[grid](probs, rmax, V, G, CHUNK, probs.stride(0), BLOCK_SIZE=2048, num_warps=8) - buf = torch.empty(B * _CAP, device=dev, dtype=torch.float32) - cnt = torch.zeros(B, device=dev, dtype=torch.int32) - _gather_pass[grid](probs, rmax, buf, cnt, V, G, CHUNK, probs.stride(0), _FRAC, _CAP) - thr = torch.empty(B, device=dev, dtype=torch.float32) - ksum = torch.empty(B, device=dev, dtype=torch.float32) - _refine_topk[(B,)](buf, cnt, rmax, target, thr, ksum, _CAP, _KR, _KBINS, 2048) - return thr, ksum, cnt - - -def _exact_top_k_renorm_probs(probs, target): - """Exact fallback for rows the fixed-size candidate buffer cannot represent.""" - _, V = probs.shape - target = target.to(torch.int64) - active_target = torch.where(target < V, target, 0) - max_top_k = max(1, int(active_target.max().item())) - values, indices = torch.topk(probs, max_top_k, dim=-1) - keep = torch.arange(max_top_k, device=probs.device).unsqueeze(0) < active_target.unsqueeze(1) - values = torch.where(keep, values, 0.0) - denom = values.sum(dim=-1, keepdim=True) - values = values / torch.where(denom > 0, denom, torch.ones_like(denom)) - out = torch.zeros_like(probs).scatter_(1, indices, values) - return torch.where((target >= V).unsqueeze(1), probs, out) + G, CHUNK = _fused_plan(B, V) + n_hist = 4 * _KBINS if (kernel is _topk_fused or tk is not None) else 0 + n_mass = 4 * _PMBINS if kernel is _topp_fused else 0 + # hist[B, n_hist] | mass[B, n_mass] | priv[B * G, n_mass] | bar[B] | lo0[B] | ksum[B] | ksum_k[B] | tok[B] + ws = torch.zeros(B * (n_hist + n_mass) + B * G * n_mass + 5 * B, device=dev, dtype=torch.int32) + hist = ws[:B * n_hist] + mass = ws[B * n_hist:B * (n_hist + n_mass)].view(torch.float32) + priv = ws[B * (n_hist + n_mass):B * (n_hist + n_mass) + B * G * n_mass].view(torch.float32) + tail = B * (n_hist + n_mass) + B * G * n_mass + bar, lo0 = ws[tail:tail + B], ws[tail + B:tail + 2 * B] + ksum = ws[tail + 2 * B:tail + 3 * B].view(torch.float32) + ksum_k = ws[tail + 3 * B:tail + 4 * B].view(torch.float32) + if draw: + psum = torch.empty(B * G, device=dev, dtype=torch.float32) + u = _gen_u(B, dev, seed, offset) + res = ws[tail + 4 * B:] + out, tok = probs, res + else: + psum, u = ksum, ksum + res = torch.empty_like(probs) + out, tok = res, lo0 + # a lone CTA per row in a single wave streams faster with more warps; with G > 1 the co-residency budget caps warps + wide = G == 1 and B <= _NUM_SM + common = dict(DRAW=draw, G_POW2=_next_pow2(G), BLOCK=8192 if wide else _FUSED_BLOCK, num_warps=32 if wide else 8, + launch_cooperative_grid=G > 1) + if kernel is _topk_fused: + _topk_fused[(B * G,)](probs, tk, lo0, hist, bar, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), + BINS=_KBINS, **common) + else: + _topp_fused[(B * G,)](probs, tp, tk if tk is not None else tp, lo0, hist, priv, mass, bar, ksum_k, ksum, psum, u, out, tok, + V, G, CHUNK, probs.stride(0), TOPK=tk is not None, KBINS=_KBINS, PBINS=_PMBINS, **common) + return res + + +def _topk_fused_launch(probs, target, draw, seed=None, offset=None): + return _fused_launch(probs, _topk_fused, target, None, draw, seed, offset) + + +def _topk_thr_search(probs, target): + # fallback when the cooperative launch is unavailable: rmax + 4 x (count-hist + refine) launches over the bit pattern + return _search(probs, target, False, 4, _KBINS, bits=True) + + +def _fused_or(fn_fused, fn_fallback): + global _fused_ok + if _fused_ok: + try: + return fn_fused() + except RuntimeError as exc: + # only a failed cooperative launch disqualifies this device; anything else (OOM, bad input) propagates + if "cooperative" not in str(exc) and "Triton Error" not in str(exc): + raise + _fused_ok = False + logger.warning("fused triton sampling unavailable (%s); using the multi-launch search", exc) + return fn_fallback() + + +def _topk(probs, target, draw, seed=None, offset=None): + if probs.size(0) == 0: + return torch.empty(0, device=probs.device, dtype=torch.int32) if draw else probs.clone() + + def fallback(): + thr = _topk_thr_search(probs, target) + return _draw(probs, thr, seed, offset) if draw else _renorm(probs, thr) + return _fused_or(lambda: _topk_fused_launch(probs, target, draw, seed, offset), fallback) + + +def _topp(probs, tp, tk, draw, seed=None, offset=None): + if probs.size(0) == 0: + return torch.empty(0, device=probs.device, dtype=torch.int32) if draw else probs.clone() + + def fallback(): + src = probs if tk is None else _renorm(probs, _topk_thr_search(probs, tk)) + thr = _search(src, tp, True, _PR, _PBINS) + return _draw(src, thr, seed, offset) if draw else _renorm(src, thr) + return _fused_or(lambda: _fused_launch(probs, _topp_fused, tk, tp, draw, seed, offset), fallback) + + +def _topp_target(top_p, B, dev): + if isinstance(top_p, torch.Tensor): + return top_p.float().to(dev).contiguous() + return torch.full((B,), float(top_p), device=dev, dtype=torch.float32) def top_k_renorm_probs(probs, top_k): probs = probs.float() - B, V = probs.shape - if not isinstance(top_k, torch.Tensor) and int(top_k) >= V: - return probs - dev = probs.device - target = _topk_target(top_k, B, dev) - if (target >= V).any().item(): - return _exact_top_k_renorm_probs(probs, target) - G, CHUNK = _plan(B, V) - grid = (B * G,) - thr, _, cnt = _topk_thr_ksum(probs, target) - if ((cnt < target) | (cnt > _CAP)).any().item(): - return _exact_top_k_renorm_probs(probs, target) - ksum = torch.zeros(B, device=dev, dtype=torch.float32) - _ksum_pass[grid](probs, thr, ksum, V, G, CHUNK, probs.stride(0)) - out = torch.empty_like(probs) - _write_pass[grid](probs, out, thr, ksum, V, G, CHUNK, probs.stride(0)) - return out + return _topk(probs, _topk_target(top_k, probs.size(0), probs.device), False) def top_k_sampling_from_probs(probs, top_k, indices=None, deterministic=True, generator=None, check_nan=False, seed=None, offset=None, return_valid=False): probs = probs.float() src = probs if indices is None else probs[indices].contiguous() - r = top_k_renorm_probs(src, top_k) - out = _draw(r, _zeros_thr(src.size(0), src.device), seed, offset) + out = _topk(src, _topk_target(top_k, src.size(0), src.device), True, seed, offset) out = out.to(indices.dtype) if indices is not None else out return (out, torch.ones_like(out, dtype=torch.bool)) if return_valid else out @@ -624,13 +823,8 @@ def top_k_top_p_sampling_from_probs(probs, top_k, top_p, indices=None, return_valid=False): probs = probs.float() src = probs if indices is None else probs[indices].contiguous() - r = top_k_renorm_probs(src, top_k) - if isinstance(top_p, torch.Tensor): - target = top_p.float().to(src.device).contiguous() - else: - target = torch.full((src.size(0),), float(top_p), device=src.device, dtype=torch.float32) - thr = _search(r, target, True, _PR, _PBINS) - out = _draw(r, thr, seed, offset) + B = src.size(0) + out = _topp(src, _topp_target(top_p, B, src.device), _topk_target(top_k, B, src.device), True, seed, offset) out = out.to(indices.dtype) if indices is not None else out return (out, torch.ones_like(out, dtype=torch.bool)) if return_valid else out From b51c580a8dedc3c502b874fb86e966c1c2e16fdb Mon Sep 17 00:00:00 2001 From: taking-lying-flat <1615405@qq.com> Date: Wed, 2 Sep 2026 21:32:23 +0800 Subject: [PATCH 3/4] fix(sampling): make fused filtering exact and portable --- python/freetoken/kernel/triton/sampling.py | 442 ++++++++------------- tests/kernels/test_triton_sampling.py | 65 --- 2 files changed, 156 insertions(+), 351 deletions(-) delete mode 100644 tests/kernels/test_triton_sampling.py diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index 65237a68e..2d3699fd6 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -7,15 +7,16 @@ * Every row is split across many CTAs (``_plan`` -> G column-chunks) so bs=1 uses the whole GPU, unlike a single-block-per-row kernel that is single-SM-bound. * softmax is a multi-CTA online softmax; the draw is a multi-CTA inverse-CDF. - * top-k and top-p are each one cooperative kernel: the row's CTAs each bin their chunk over + * top-k and top-p each use one Triton kernel: the row's CTAs bin their chunk over the fp32 bit pattern (order-preserving for x >= 0), meet at a per-row spin barrier, and all redo the refine so they share the bracket. Four rounds of 256 bins bring the 2**31 range down to one bit pattern, so the threshold is exactly the k-th largest prob (top-k, counts) - or the value where the descending cumulative mass reaches p (top-p, exact per-bin mass); - the same kernel then renormalizes or draws. No candidate buffer, no data-dependent shape, - no host sync. Results are exact up to fp32 atomic summation order. - * If the cooperative launch is unavailable the module falls back to the multi-launch - search below (exact for top-k; top-p there is the older bin-center estimate). + or the value where the descending cumulative mass reaches p (top-p, exact per-bin mass). + Boundary ties are clipped in token-id order, then the same kernel renormalizes or draws. + No candidate buffer, data-dependent shape, or host sync is needed. Results are exact up + to fp32 atomic summation order. + * If a cooperative launch is unavailable, the same exact kernel is retried with one CTA + per row; only parallelism changes. * deterministic, generator and check_nan exist for flashinfer signature compatibility and are ignored; seed and offset are honored. Given a seed, top-k draws reproduce; top-p may pick a different token on rows whose cumulative mass sits within fp32 rounding of p. @@ -24,6 +25,7 @@ from __future__ import annotations import logging +from functools import cache import torch import triton @@ -33,13 +35,17 @@ logger = logging.getLogger(__name__) -_NUM_SM = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count _MIN_CHUNK = 4096 # do not split a row finer than this -def _plan(B, V): +@cache +def _num_sm(device): + return torch.cuda.get_device_properties(device).multi_processor_count + + +def _plan(B, V, device): """Return (G, CHUNK): split each row into G column-chunks of size CHUNK.""" - g_by_sm = max(1, _NUM_SM // B) + g_by_sm = max(1, _num_sm(device) // B) g_by_chunk = max(1, triton.cdiv(V, _MIN_CHUNK)) G = min(g_by_sm, g_by_chunk) CHUNK = triton.cdiv(V, G) @@ -128,8 +134,10 @@ def _sm_finalize( def softmax(logits, temperature=None, enable_pdl=None): logits = logits.float() B, V = logits.shape + if B == 0: + return logits.clone() probs = torch.empty_like(logits) - G, CHUNK = _plan(B, V) + G, CHUNK = _plan(B, V, logits.device) if temperature is None: temperature = 1.0 if isinstance(temperature, torch.Tensor): @@ -146,13 +154,6 @@ def softmax(logits, temperature=None, enable_pdl=None): return probs -# =========================================================================== -# multi-CTA top-p via histogram-bracket refinement. Every full-vocab pass is -# split across all SMs; the sequential refine step is a tiny grid=(B,) kernel. -# =========================================================================== -_PBINS = 64 # fallback top-p only: count-hist + bin-center mass (64**4 ~ 1.7e7) -_PR = 4 - _SR_CFGS = [ triton.Config({"BLOCK_SIZE": bs}, num_warps=w, num_stages=s) for bs in (1024, 2048, 4096) @@ -161,194 +162,6 @@ def softmax(logits, temperature=None, enable_pdl=None): ] -@triton.jit -def _rmax_pass(probs_ptr, rmax_ptr, V, G, CHUNK, row_stride, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - row = pid // G - base = row * row_stride - start = (pid % G) * CHUNK - end = tl.minimum(start + CHUNK, V) - m = 0.0 - for s0 in tl.range(start, end, BLOCK_SIZE): - offs = s0 + tl.arange(0, BLOCK_SIZE) - mask = offs < end - x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - m = tl.maximum(m, tl.max(x, 0)) - tl.atomic_max(rmax_ptr + row, m) - - -@triton.autotune(configs=_SR_CFGS, key=["CHUNK", "BINS", "BITS"], reset_to_zero=["hist_ptr"], **autotune_cache_kwargs) -@triton.jit -def _count_hist_pass( - probs_ptr, lo_ptr, hi_ptr, hist_ptr, V, G, CHUNK, row_stride, - BINS: tl.constexpr, BITS: tl.constexpr, BLOCK_SIZE: tl.constexpr, -): - # BITS: lo/hi are int32 bit patterns and bins are exact integer ranges; else float value bins - pid = tl.program_id(0) - row = pid // G - lo = tl.load(lo_ptr + row) - hi = tl.load(hi_ptr + row) - if BITS: - w = (hi - lo + BINS - 1) // BINS - else: - invw = BINS / tl.maximum(hi - lo, 1e-30) - base = row * row_stride - start = (pid % G) * CHUNK - end = tl.minimum(start + CHUNK, V) - acc = tl.zeros([BINS], tl.int32) - for s0 in tl.range(start, end, BLOCK_SIZE): - offs = s0 + tl.arange(0, BLOCK_SIZE) - mask = offs < end - x = tl.load(probs_ptr + base + offs, mask=mask, other=-1.0).to(tl.float32) - if BITS: - y = x.to(tl.int32, bitcast=True) - b = (y - lo) // w - else: - y = x - b = ((x - lo) * invw).to(tl.int32) - inrange = mask & (y >= lo) & (y < hi) - # tl.histogram does NOT cleanly drop out-of-range indices; route every - # out-of-bracket element to bin 0 and then subtract that count back out so - # the histogram holds ONLY in-[lo,hi) counts (out-of-range is tracked via - # `above`/excluded, exactly like the one-hot path). - b = tl.where(inrange, tl.maximum(0, tl.minimum(b, BINS - 1)), 0) - hcnt = tl.histogram(b, BINS) - noor = tl.sum((mask & (~inrange)).to(tl.int32)) - hcnt = hcnt - tl.where(tl.arange(0, BINS) == 0, noor, 0) - acc += hcnt - tl.atomic_add(hist_ptr + row * BINS + tl.arange(0, BINS), acc.to(tl.float32)) - - -@triton.jit -def _refine_pass(lo_ptr, hi_ptr, above_ptr, hist_ptr, target_ptr, BINS: tl.constexpr, BITS: tl.constexpr): - row = tl.program_id(0) - lo = tl.load(lo_ptr + row) - hi = tl.load(hi_ptr + row) - above = tl.load(above_ptr + row) - target = tl.load(target_ptr + row) - if BITS: - w = (hi - lo + BINS - 1) // BINS - else: - w = (hi - lo) / BINS - jj = tl.arange(0, BINS) - h = tl.load(hist_ptr + row * BINS + jj) - prefix = tl.cumsum(h, 0) - total = tl.sum(h, 0) - c_ge_bottom = above + total - prefix + h - ok = c_ge_bottom >= target - j = tl.max(tl.where(ok, jj, -1)) - prefix_j = tl.sum(tl.where(jj <= j, h, 0.0)) - upd = j >= 0 - new_hi = lo + (j + 1) * w - if BITS: - new_hi = tl.minimum(new_hi, hi) - tl.store(lo_ptr + row, tl.where(upd, lo + j * w, lo)) - tl.store(hi_ptr + row, tl.where(upd, new_hi, hi)) - tl.store(above_ptr + row, tl.where(upd, above + total - prefix_j, above)) - # zero the row so the next iteration's atomic_add starts clean (reset_to_zero - # only fires during autotuning, not on production calls) - tl.store(hist_ptr + row * BINS + jj, 0.0) - - -@triton.jit -def _refine_mass_pass(lo_ptr, hi_ptr, above_ptr, hist_ptr, target_ptr, BINS: tl.constexpr): - # top-p refine: hist holds COUNTS; approximate per-bin MASS as count*bin_center - # (exact in the limit as the bracket narrows). target is the p mass threshold. - row = tl.program_id(0) - lo = tl.load(lo_ptr + row) - hi = tl.load(hi_ptr + row) - above = tl.load(above_ptr + row) - target = tl.load(target_ptr + row) - w = (hi - lo) / BINS - jj = tl.arange(0, BINS) - h = tl.load(hist_ptr + row * BINS + jj) - center = lo + (jj.to(tl.float32) + 0.5) * w - massbin = h * center - prefix = tl.cumsum(massbin, 0) - total = tl.sum(massbin, 0) - c_ge_bottom = above + total - prefix + massbin - ok = c_ge_bottom >= target - j = tl.max(tl.where(ok, jj, -1)) - prefix_j = tl.sum(tl.where(jj <= j, massbin, 0.0)) - upd = j >= 0 - tl.store(lo_ptr + row, tl.where(upd, lo + j * w, lo)) - tl.store(hi_ptr + row, tl.where(upd, lo + (j + 1) * w, hi)) - tl.store(above_ptr + row, tl.where(upd, above + total - prefix_j, above)) - tl.store(hist_ptr + row * BINS + jj, 0.0) - - -@triton.autotune(configs=_SR_CFGS, key=["CHUNK"], reset_to_zero=["ksum_ptr"], **autotune_cache_kwargs) -@triton.jit -def _ksum_pass(probs_ptr, thr_ptr, ksum_ptr, V, G, CHUNK, row_stride, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - row = pid // G - thr = tl.load(thr_ptr + row) - base = row * row_stride - start = (pid % G) * CHUNK - end = tl.minimum(start + CHUNK, V) - s = 0.0 - for s0 in tl.range(start, end, BLOCK_SIZE): - offs = s0 + tl.arange(0, BLOCK_SIZE) - mask = offs < end - x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - s += tl.sum(tl.where(x >= thr, x, 0.0), 0) - tl.atomic_add(ksum_ptr + row, s) - - -@triton.autotune(configs=_SR_CFGS, key=["CHUNK"], **autotune_cache_kwargs) -@triton.jit -def _write_pass(probs_ptr, out_ptr, thr_ptr, ksum_ptr, V, G, CHUNK, row_stride, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - row = pid // G - thr = tl.load(thr_ptr + row) - inv_s = 1.0 / tl.load(ksum_ptr + row) - base = row * row_stride - start = (pid % G) * CHUNK - end = tl.minimum(start + CHUNK, V) - for s0 in tl.range(start, end, BLOCK_SIZE): - offs = s0 + tl.arange(0, BLOCK_SIZE) - mask = offs < end - x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - tl.store(out_ptr + base + offs, tl.where(x >= thr, x * inv_s, 0.0), mask=mask) - - -def _search(probs, target, mass, R, BINS, bits=False): - """Return per-row threshold: keep x >= thr, with count/mass(>=thr) ~ target.""" - B, V = probs.shape - dev = probs.device - G, CHUNK = _plan(B, V) - grid = (B * G,) - rmax = torch.zeros(B, device=dev, dtype=torch.float32) - _rmax_pass[grid](probs, rmax, V, G, CHUNK, probs.stride(0), BLOCK_SIZE=2048, num_warps=8) - if bits: - lo = torch.zeros(B, device=dev, dtype=torch.int32) - hi = (rmax.view(torch.int32) + 1).contiguous() - else: - lo = torch.zeros(B, device=dev, dtype=torch.float32) - hi = (rmax * 1.0000001).contiguous() - above = torch.zeros(B, device=dev, dtype=torch.float32) - hist = torch.zeros(B * BINS, device=dev, dtype=torch.float32) - for _ in range(R): - _count_hist_pass[grid](probs, lo, hi, hist, V, G, CHUNK, probs.stride(0), BINS, bits) - if mass: - _refine_mass_pass[(B,)](lo, hi, above, hist, target, BINS) - else: - _refine_pass[(B,)](lo, hi, above, hist, target, BINS, bits) - return lo.view(torch.float32) if bits else lo - - -def _renorm(probs, thr): - B, V = probs.shape - dev = probs.device - G, CHUNK = _plan(B, V) - grid = (B * G,) - out = torch.empty_like(probs) - ksum = torch.zeros(B, device=dev, dtype=torch.float32) - _ksum_pass[grid](probs, thr, ksum, V, G, CHUNK, probs.stride(0)) - _write_pass[grid](probs, out, thr, ksum, V, G, CHUNK, probs.stride(0)) - return out - - def top_p_renorm_probs(probs, top_p): probs = probs.float() return _topp(probs, _topp_target(top_p, probs.size(0), probs.device), None, False) @@ -439,7 +252,9 @@ def _gen_u(B, device, seed, offset): def _draw(probs, thr, seed, offset): B, V = probs.shape dev = probs.device - G, CHUNK = _plan(B, V) + if B == 0: + return torch.empty(0, device=dev, dtype=torch.int32) + G, CHUNK = _plan(B, V, dev) grid = (B * G,) psum = torch.empty(B * G, device=dev, dtype=torch.float32) choff = torch.empty(B * G, device=dev, dtype=torch.float32) @@ -489,11 +304,8 @@ def top_p_sampling_from_probs(probs, top_p, indices=None, deterministic=True, ge def _topk_target(top_k, B, dev): if isinstance(top_k, torch.Tensor): - target = top_k.float().to(dev).contiguous() - else: - target = torch.full((B,), float(int(top_k)), device=dev, dtype=torch.float32) - # k <= 0 would satisfy every bin and push the bracket past the max; k > V needs no clamp, the search leaves lo at 0 - return torch.clamp(target, min=1.0) + return top_k.to(device=dev, dtype=torch.int32).contiguous() + return torch.full((B,), max(int(top_k), 1), device=dev, dtype=torch.int32) @triton.jit @@ -532,7 +344,7 @@ def _bits_round( h = tl.load(hist_ptr + jj, cache_modifier=".cg") prefix = tl.cumsum(h, 0) total = tl.sum(h, 0) - ok = (above + total - prefix + h).to(tl.float32) >= target + ok = above + total - prefix + h >= target j = tl.max(tl.where(ok, jj, -1)) prefix_j = tl.sum(tl.where(jj <= j, h, 0)) upd = j >= 0 @@ -541,9 +353,36 @@ def _bits_round( return lo, above +@triton.jit +def _tie_prefix(local, pid, row, cta, thr, tie_ptr, bar_ptr, need, target, above, G, + G_POW2: tl.constexpr, MASS_TARGET: tl.constexpr): + tl.store(tie_ptr + pid, local) + _row_barrier(bar_ptr, need) + + goff = tl.arange(0, G_POW2) + gmask = goff < G + counts = tl.load(tie_ptr + row * G + goff, mask=gmask, other=0, cache_modifier=".cg") + before = tl.sum(tl.where(goff < cta, counts, 0), 0) + total = tl.sum(counts, 0) + if MASS_TARGET: + remaining = tl.maximum(target - above, 0.0) + limit = tl.ceil(remaining / tl.maximum(thr, 1e-30)).to(tl.int32) + else: + limit = (target - above).to(tl.int32) + return before, tl.maximum(0, tl.minimum(limit, total)) + + +@triton.jit +def _keep_mask(x, mask, thr, tie_before, tie_limit, seen): + eq = mask & (x == thr) + rank = tie_before + seen + tl.cumsum(eq.to(tl.int32), 0) - 1 + keep = mask & ((x > thr) | (eq & (rank < tie_limit))) + return keep, seen + tl.sum(eq.to(tl.int32), 0) + + @triton.jit def _topk_fused( - probs_ptr, target_ptr, lo0_ptr, hist_ptr, bar_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, + probs_ptr, target_ptr, hist_ptr, tie_ptr, bar_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, CHUNK, row_stride, DRAW: tl.constexpr, G_POW2: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, ): @@ -553,35 +392,45 @@ def _topk_fused( base = row * row_stride start = cta * CHUNK end = tl.minimum(start + CHUNK, V) - target = tl.load(target_ptr + row) + target = tl.maximum(tl.load(target_ptr + row), 1) hrow = hist_ptr + row * 4 * BINS brow = bar_ptr + row - lo = tl.load(lo0_ptr + row) + lo = 0 above = lo lo, above = _bits_round(probs_ptr, base, start, end, hrow, brow, target, lo, above, G, 23, 0, BINS, BLOCK) lo, above = _bits_round(probs_ptr, base, start, end, hrow + BINS, brow, target, lo, above, 2 * G, 15, 1 << 23, BINS, BLOCK) lo, above = _bits_round(probs_ptr, base, start, end, hrow + 2 * BINS, brow, target, lo, above, 3 * G, 7, 1 << 15, BINS, BLOCK) lo, above = _bits_round(probs_ptr, base, start, end, hrow + 3 * BINS, brow, target, lo, above, 4 * G, 0, 1 << 7, BINS, BLOCK) - _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), brow, 5 * G, - ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, DRAW, G_POW2, BLOCK) + _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), above, target, + 0.0, 0, tie_ptr, brow, 5 * G, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, + False, False, DRAW, G_POW2, BLOCK) @triton.jit def _keep_tail( - probs_ptr, base, start, end, pid, row, cta, thr, bar_ptr, need, - ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, - DRAW: tl.constexpr, G_POW2: tl.constexpr, BLOCK: tl.constexpr, + probs_ptr, base, start, end, pid, row, cta, thr, above, target, floor_thr, floor_limit, + tie_ptr, bar_ptr, tie_need, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, + HAS_FLOOR: tl.constexpr, MASS_TARGET: tl.constexpr, DRAW: tl.constexpr, + G_POW2: tl.constexpr, BLOCK: tl.constexpr, ): - # keep x >= thr over this chunk: DRAW picks one token per row, else renormalize the kept mass to 1 s = 0.0 + local_ties = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - s += tl.sum(tl.where(x >= thr, x, 0.0), 0) + s += tl.sum(tl.where(mask & (x > thr), x, 0.0), 0) + local_ties += tl.sum((mask & (x == thr)).to(tl.int32), 0) + + tie_before, tie_limit = _tie_prefix(local_ties, pid, row, cta, thr, tie_ptr, bar_ptr, + tie_need, target, above, G, G_POW2, MASS_TARGET) + if HAS_FLOOR: + tie_limit = tl.where(thr == floor_thr, tl.minimum(tie_limit, floor_limit), tie_limit) + selected_ties = tl.maximum(0, tl.minimum(local_ties, tie_limit - tie_before)) + s += selected_ties.to(tl.float32) * thr if DRAW: tl.store(psum_ptr + pid, s) - _row_barrier(bar_ptr, need) + _row_barrier(bar_ptr, tie_need + G) goff = tl.arange(0, G_POW2) gmask = goff < G ps = tl.load(psum_ptr + row * G + goff, mask=gmask, other=0.0, cache_modifier=".cg") @@ -589,11 +438,12 @@ def _keep_tail( incl = tl.sum(tl.where(goff <= cta, ps, 0.0), 0) tgt = tl.load(u_ptr + row) * tl.sum(ps, 0) last_kept = start * 0 - 1 + seen = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - kept = (x >= thr) & mask + kept, seen = _keep_mask(x, mask, thr, tie_before, tie_limit, seen) wv = tl.where(kept, x, 0.0) cval = acc + tl.cumsum(wv, 0) idx = tl.where(cval > tgt, offs, V) @@ -601,7 +451,7 @@ def _keep_tail( if (blk_min < V) and (acc <= tgt): tl.store(tok_ptr + row, blk_min) acc += tl.sum(wv, 0) - last_kept = tl.maximum(last_kept, tl.max(tl.where(kept, offs, -1), 0)) + last_kept = tl.maximum(last_kept, tl.max(tl.where(kept & (x > 0), offs, -1), 0)) # fp rounding can leave tgt between this CTA's running sum and the next CTA's prefix, or past the total; # the CTA that owns that gap (or the last one holding mass) writes its last kept token instead is_last_mass = tl.sum(tl.where((goff > cta) & (ps > 0), 1, 0), 0) == 0 @@ -609,13 +459,15 @@ def _keep_tail( tl.store(tok_ptr + row, last_kept) else: tl.atomic_add(ksum_ptr + row, s) - _row_barrier(bar_ptr, need) + _row_barrier(bar_ptr, tie_need + G) inv = 1.0 / tl.atomic_add(ksum_ptr + row, 0.0) + seen = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - tl.store(out_ptr + base + offs, tl.where(x >= thr, x * inv, 0.0), mask=mask) + kept, seen = _keep_mask(x, mask, thr, tie_before, tie_limit, seen) + tl.store(out_ptr + base + offs, tl.where(kept, x * inv, 0.0), mask=mask) _PMBINS = 256 @@ -624,21 +476,26 @@ def _keep_tail( @triton.jit def _pmass_round( probs_ptr, base, start, end, priv_ptr, mass_ptr, bar_ptr, target, lo, above, need, + floor_thr, floor_before, floor_limit, TOPK: tl.constexpr, S: tl.constexpr, WIDTH: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, ): # top-p round over the bit pattern: per-bin MASS (exact up to fp32 atomic order) via scatter-add into this # CTA's private buffer, then one reduction into the row buffer, so the bin holding the p crossing is known jj = tl.arange(0, BINS) + seen = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + eligible = mask + if TOPK: + eligible, seen = _keep_mask(x, mask, floor_thr, floor_before, floor_limit, seen) y = x.to(tl.int32, bitcast=True) d = y - lo if WIDTH == 0: - inrange = mask & (y >= lo) & (y <= _INF_BITS) + inrange = eligible & (y >= lo) & (y <= _INF_BITS) else: - inrange = mask & (y >= lo) & (d < WIDTH) & (y <= _INF_BITS) + inrange = eligible & (y >= lo) & (d < WIDTH) & (y <= _INF_BITS) tl.atomic_add(priv_ptr + tl.where(inrange, d >> S, 0), x, mask=inrange) # every warp's scatter-adds must land before any thread reads the private bins back tl.debug_barrier() @@ -656,7 +513,7 @@ def _pmass_round( @triton.jit def _topp_fused( - probs_ptr, tp_ptr, tk_ptr, lo0_ptr, hist_ptr, priv_ptr, mass_ptr, bar_ptr, ksumk_ptr, ksum_ptr, psum_ptr, u_ptr, + probs_ptr, tp_ptr, tk_ptr, hist_ptr, priv_ptr, mass_ptr, tie_ptr, bar_ptr, ksumk_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, CHUNK, row_stride, TOPK: tl.constexpr, DRAW: tl.constexpr, G_POW2: tl.constexpr, KBINS: tl.constexpr, PBINS: tl.constexpr, BLOCK: tl.constexpr, @@ -670,9 +527,12 @@ def _topp_fused( start = cta * CHUNK end = tl.minimum(start + CHUNK, V) brow = bar_ptr + row - lo = tl.load(lo0_ptr + row) + lo = 0 + thr_k = 0.0 + tk_before = 0 + tk_limit = 0 if TOPK: - tk = tl.load(tk_ptr + row) + tk = tl.maximum(tl.load(tk_ptr + row), 1) hk = hist_ptr + row * 4 * KBINS above_i = lo lo, above_i = _bits_round(probs_ptr, base, start, end, hk, brow, tk, lo, above_i, G, 23, 0, KBINS, BLOCK) @@ -681,126 +541,136 @@ def _topp_fused( lo, above_i = _bits_round(probs_ptr, base, start, end, hk + 3 * KBINS, brow, tk, lo, above_i, 4 * G, 0, 1 << 7, KBINS, BLOCK) thr_k = lo.to(tl.float32, bitcast=True) s = 0.0 + local_ties = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - s += tl.sum(tl.where(x >= thr_k, x, 0.0), 0) + s += tl.sum(tl.where(mask & (x > thr_k), x, 0.0), 0) + local_ties += tl.sum((mask & (x == thr_k)).to(tl.int32), 0) + tk_before, tk_limit = _tie_prefix(local_ties, pid, row, cta, thr_k, tie_ptr, brow, + 5 * G, tk, above_i, G, G_POW2, False) + selected_ties = tl.maximum(0, tl.minimum(local_ties, tk_limit - tk_before)) + s += selected_ties.to(tl.float32) * thr_k tl.atomic_add(ksumk_ptr + row, s) - _row_barrier(brow, 5 * G) + _row_barrier(brow, 6 * G) target = tl.load(tp_ptr + row) * tl.atomic_add(ksumk_ptr + row, 0.0) - done = 5 + done = 6 else: target = tl.load(tp_ptr + row) done = 0 mp = mass_ptr + row * 4 * PBINS pp = priv_ptr + pid * 4 * PBINS above = 0.0 - lo, above = _pmass_round(probs_ptr, base, start, end, pp, mp, brow, target, lo, above, (done + 1) * G, 23, 0, PBINS, BLOCK) - lo, above = _pmass_round(probs_ptr, base, start, end, pp + PBINS, mp + PBINS, brow, target, lo, above, (done + 2) * G, 15, 1 << 23, PBINS, BLOCK) - lo, above = _pmass_round(probs_ptr, base, start, end, pp + 2 * PBINS, mp + 2 * PBINS, brow, target, lo, above, (done + 3) * G, 7, 1 << 15, PBINS, BLOCK) - lo, above = _pmass_round(probs_ptr, base, start, end, pp + 3 * PBINS, mp + 3 * PBINS, brow, target, lo, above, (done + 4) * G, 0, 1 << 7, PBINS, BLOCK) - _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), brow, (done + 5) * G, - ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, DRAW, G_POW2, BLOCK) - - -_fused_ok = True + lo, above = _pmass_round(probs_ptr, base, start, end, pp, mp, brow, target, lo, above, (done + 1) * G, + thr_k, tk_before, tk_limit, TOPK, 23, 0, PBINS, BLOCK) + lo, above = _pmass_round(probs_ptr, base, start, end, pp + PBINS, mp + PBINS, brow, target, lo, above, + (done + 2) * G, thr_k, tk_before, tk_limit, TOPK, 15, 1 << 23, PBINS, BLOCK) + lo, above = _pmass_round(probs_ptr, base, start, end, pp + 2 * PBINS, mp + 2 * PBINS, brow, target, lo, above, + (done + 3) * G, thr_k, tk_before, tk_limit, TOPK, 7, 1 << 15, PBINS, BLOCK) + lo, above = _pmass_round(probs_ptr, base, start, end, pp + 3 * PBINS, mp + 3 * PBINS, brow, target, lo, above, + (done + 4) * G, thr_k, tk_before, tk_limit, TOPK, 0, 1 << 7, PBINS, BLOCK) + _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), above, target, + thr_k, tk_limit, tie_ptr, brow, (done + 5) * G, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, + V, G, TOPK, True, DRAW, G_POW2, BLOCK) + + +_COOPERATIVE_DISABLED = set() _COOP_CTAS_PER_SM = 2 # the fused kernels use ~80 regs/thread at 8 warps; 4/SM fails the cooperative launch -def _fused_plan(B, V): +def _fused_plan(B, V, device, force_single=False): + if force_single: + return 1, V # the cooperative launch needs the whole grid co-resident, so cap B*G by an occupancy budget instead of _plan's one CTA per SM - g_by_sm = max(1, (_COOP_CTAS_PER_SM * _NUM_SM) // B) + g_by_sm = max(1, (_COOP_CTAS_PER_SM * _num_sm(device)) // B) g_by_chunk = max(1, triton.cdiv(V, _MIN_CHUNK)) G = min(g_by_sm, g_by_chunk) return G, triton.cdiv(V, G) -def _fused_launch(probs, kernel, tk, tp, draw, seed, offset): +def _fused_launch(probs, kernel, tk, tp, draw, seed, offset, force_single=False): B, V = probs.shape dev = probs.device - G, CHUNK = _fused_plan(B, V) + G, CHUNK = _fused_plan(B, V, dev, force_single) n_hist = 4 * _KBINS if (kernel is _topk_fused or tk is not None) else 0 n_mass = 4 * _PMBINS if kernel is _topp_fused else 0 - # hist[B, n_hist] | mass[B, n_mass] | priv[B * G, n_mass] | bar[B] | lo0[B] | ksum[B] | ksum_k[B] | tok[B] - ws = torch.zeros(B * (n_hist + n_mass) + B * G * n_mass + 5 * B, device=dev, dtype=torch.int32) + # hist[B, n_hist] | mass[B, n_mass] | priv[B * G, n_mass] | tie[B * G] | bar/ksum/ksum_k/tok[B] + ws = torch.zeros(B * (n_hist + n_mass) + B * G * (n_mass + 1) + 4 * B, device=dev, dtype=torch.int32) hist = ws[:B * n_hist] mass = ws[B * n_hist:B * (n_hist + n_mass)].view(torch.float32) priv = ws[B * (n_hist + n_mass):B * (n_hist + n_mass) + B * G * n_mass].view(torch.float32) - tail = B * (n_hist + n_mass) + B * G * n_mass - bar, lo0 = ws[tail:tail + B], ws[tail + B:tail + 2 * B] - ksum = ws[tail + 2 * B:tail + 3 * B].view(torch.float32) - ksum_k = ws[tail + 3 * B:tail + 4 * B].view(torch.float32) + tie_start = B * (n_hist + n_mass) + B * G * n_mass + tie = ws[tie_start:tie_start + B * G] + tail = tie_start + B * G + bar = ws[tail:tail + B] + ksum = ws[tail + B:tail + 2 * B].view(torch.float32) + ksum_k = ws[tail + 2 * B:tail + 3 * B].view(torch.float32) if draw: psum = torch.empty(B * G, device=dev, dtype=torch.float32) u = _gen_u(B, dev, seed, offset) - res = ws[tail + 4 * B:] + res = ws[tail + 3 * B:] out, tok = probs, res else: psum, u = ksum, ksum res = torch.empty_like(probs) - out, tok = res, lo0 + out, tok = res, bar # a lone CTA per row in a single wave streams faster with more warps; with G > 1 the co-residency budget caps warps - wide = G == 1 and B <= _NUM_SM + wide = G == 1 and B <= _num_sm(dev) common = dict(DRAW=draw, G_POW2=_next_pow2(G), BLOCK=8192 if wide else _FUSED_BLOCK, num_warps=32 if wide else 8, launch_cooperative_grid=G > 1) if kernel is _topk_fused: - _topk_fused[(B * G,)](probs, tk, lo0, hist, bar, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), + _topk_fused[(B * G,)](probs, tk, hist, tie, bar, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), BINS=_KBINS, **common) else: - _topp_fused[(B * G,)](probs, tp, tk if tk is not None else tp, lo0, hist, priv, mass, bar, ksum_k, ksum, psum, u, out, tok, + _topp_fused[(B * G,)](probs, tp, tk if tk is not None else tp, hist, priv, mass, tie, bar, ksum_k, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), TOPK=tk is not None, KBINS=_KBINS, PBINS=_PMBINS, **common) return res -def _topk_fused_launch(probs, target, draw, seed=None, offset=None): - return _fused_launch(probs, _topk_fused, target, None, draw, seed, offset) +def _cooperative_key(probs, kernel, tk, draw): + kind = "topk" if kernel is _topk_fused else "topk_topp" if tk is not None else "topp" + return probs.device, kind, draw -def _topk_thr_search(probs, target): - # fallback when the cooperative launch is unavailable: rmax + 4 x (count-hist + refine) launches over the bit pattern - return _search(probs, target, False, 4, _KBINS, bits=True) +def _is_cooperative_launch_error(exc): + message = str(exc).lower() + return "cooperative" in message or "too many resources requested for launch" in message -def _fused_or(fn_fused, fn_fallback): - global _fused_ok - if _fused_ok: - try: - return fn_fused() - except RuntimeError as exc: - # only a failed cooperative launch disqualifies this device; anything else (OOM, bad input) propagates - if "cooperative" not in str(exc) and "Triton Error" not in str(exc): - raise - _fused_ok = False - logger.warning("fused triton sampling unavailable (%s); using the multi-launch search", exc) - return fn_fallback() +def _exact_launch(probs, kernel, tk, tp, draw, seed, offset): + key = _cooperative_key(probs, kernel, tk, draw) + force_single = key in _COOPERATIVE_DISABLED + G, _ = _fused_plan(*probs.shape, probs.device, force_single) + try: + return _fused_launch(probs, kernel, tk, tp, draw, seed, offset, force_single) + except RuntimeError as exc: + if force_single or G == 1 or not _is_cooperative_launch_error(exc): + raise + _COOPERATIVE_DISABLED.add(key) + logger.warning("cooperative triton sampling unavailable on %s (%s); retrying with one CTA per row", + probs.device, exc) + return _fused_launch(probs, kernel, tk, tp, draw, seed, offset, force_single=True) def _topk(probs, target, draw, seed=None, offset=None): if probs.size(0) == 0: return torch.empty(0, device=probs.device, dtype=torch.int32) if draw else probs.clone() - def fallback(): - thr = _topk_thr_search(probs, target) - return _draw(probs, thr, seed, offset) if draw else _renorm(probs, thr) - return _fused_or(lambda: _topk_fused_launch(probs, target, draw, seed, offset), fallback) + return _exact_launch(probs, _topk_fused, target, None, draw, seed, offset) def _topp(probs, tp, tk, draw, seed=None, offset=None): if probs.size(0) == 0: return torch.empty(0, device=probs.device, dtype=torch.int32) if draw else probs.clone() - def fallback(): - src = probs if tk is None else _renorm(probs, _topk_thr_search(probs, tk)) - thr = _search(src, tp, True, _PR, _PBINS) - return _draw(src, thr, seed, offset) if draw else _renorm(src, thr) - return _fused_or(lambda: _fused_launch(probs, _topp_fused, tk, tp, draw, seed, offset), fallback) + return _exact_launch(probs, _topp_fused, tk, tp, draw, seed, offset) def _topp_target(top_p, B, dev): if isinstance(top_p, torch.Tensor): return top_p.float().to(dev).contiguous() - return torch.full((B,), float(top_p), device=dev, dtype=torch.float32) + return torch.full((B,), float(top_p), device=dev, dtype=torch.float32) def top_k_renorm_probs(probs, top_k): diff --git a/tests/kernels/test_triton_sampling.py b/tests/kernels/test_triton_sampling.py deleted file mode 100644 index 24abd390f..000000000 --- a/tests/kernels/test_triton_sampling.py +++ /dev/null @@ -1,65 +0,0 @@ -import pytest -import torch - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def _reference_top_k_renorm(probs, top_k): - out = torch.zeros_like(probs) - vocab_size = probs.shape[1] - for row, k in enumerate(top_k.tolist()): - if k >= vocab_size: - out[row] = probs[row] - continue - values, indices = torch.topk(probs[row], k) - out[row, indices] = values / values.sum() - return out - - -def test_top_k_renorm_falls_back_when_too_few_candidates(): - from freetoken.kernel.triton.sampling import top_k_renorm_probs - - vocab_size = 10_000 - probs = torch.full((1, vocab_size), 0.8 / (vocab_size - 1), device="cuda") - probs[0, 0] = 0.2 - top_k = torch.tensor([50], dtype=torch.int32, device="cuda") - - actual = top_k_renorm_probs(probs, top_k) - expected = _reference_top_k_renorm(probs, top_k.cpu()) - - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(actual.sum(dim=-1), torch.ones(1, device="cuda")) - - -def test_top_k_renorm_falls_back_on_candidate_overflow(): - from freetoken.kernel.triton.sampling import top_k_renorm_probs - - vocab_size = 32_768 - probs = torch.linspace(1.0, 0.99, vocab_size, device="cuda").unsqueeze(0) - probs /= probs.sum(dim=-1, keepdim=True) - top_k = torch.tensor([50], dtype=torch.int32, device="cuda") - - actual = top_k_renorm_probs(probs, top_k) - expected = _reference_top_k_renorm(probs, top_k.cpu()) - - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(actual.sum(dim=-1), torch.ones(1, device="cuda")) - - -def test_top_k_renorm_mixed_batch_bypasses_vocab_size_row(): - from freetoken.kernel.triton.sampling import top_k_renorm_probs - - vocab_size = 10_000 - peaky = torch.full((vocab_size,), 0.8 / (vocab_size - 1), device="cuda") - peaky[0] = 0.2 - unfiltered = torch.linspace(1.0, 0.5, vocab_size, device="cuda") - unfiltered /= unfiltered.sum() - probs = torch.stack((peaky, unfiltered)) - top_k = torch.tensor([50, vocab_size], dtype=torch.int32, device="cuda") - - actual = top_k_renorm_probs(probs, top_k) - expected = _reference_top_k_renorm(probs, top_k.cpu()) - - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(actual.sum(dim=-1), torch.ones(2, device="cuda")) From 65fedd442f2e8660e9eef5143099793438e749d1 Mon Sep 17 00:00:00 2001 From: taking-lying-flat <1615405@qq.com> Date: Thu, 3 Sep 2026 08:52:09 +0800 Subject: [PATCH 4/4] fix(sampling): preserve boundary ties --- python/freetoken/kernel/triton/sampling.py | 122 ++++++--------------- 1 file changed, 33 insertions(+), 89 deletions(-) diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index 2d3699fd6..a97071440 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -12,9 +12,9 @@ redo the refine so they share the bracket. Four rounds of 256 bins bring the 2**31 range down to one bit pattern, so the threshold is exactly the k-th largest prob (top-k, counts) or the value where the descending cumulative mass reaches p (top-p, exact per-bin mass). - Boundary ties are clipped in token-id order, then the same kernel renormalizes or draws. - No candidate buffer, data-dependent shape, or host sync is needed. Results are exact up - to fp32 atomic summation order. + Every boundary tie is kept, matching flashinfer, then the same kernel renormalizes or + draws. No candidate buffer, data-dependent shape, or host sync is needed. Results are + exact up to fp32 atomic summation order. * If a cooperative launch is unavailable, the same exact kernel is retried with one CTA per row; only parallelism changes. * deterministic, generator and check_nan exist for flashinfer signature compatibility and are @@ -353,36 +353,9 @@ def _bits_round( return lo, above -@triton.jit -def _tie_prefix(local, pid, row, cta, thr, tie_ptr, bar_ptr, need, target, above, G, - G_POW2: tl.constexpr, MASS_TARGET: tl.constexpr): - tl.store(tie_ptr + pid, local) - _row_barrier(bar_ptr, need) - - goff = tl.arange(0, G_POW2) - gmask = goff < G - counts = tl.load(tie_ptr + row * G + goff, mask=gmask, other=0, cache_modifier=".cg") - before = tl.sum(tl.where(goff < cta, counts, 0), 0) - total = tl.sum(counts, 0) - if MASS_TARGET: - remaining = tl.maximum(target - above, 0.0) - limit = tl.ceil(remaining / tl.maximum(thr, 1e-30)).to(tl.int32) - else: - limit = (target - above).to(tl.int32) - return before, tl.maximum(0, tl.minimum(limit, total)) - - -@triton.jit -def _keep_mask(x, mask, thr, tie_before, tie_limit, seen): - eq = mask & (x == thr) - rank = tie_before + seen + tl.cumsum(eq.to(tl.int32), 0) - 1 - keep = mask & ((x > thr) | (eq & (rank < tie_limit))) - return keep, seen + tl.sum(eq.to(tl.int32), 0) - - @triton.jit def _topk_fused( - probs_ptr, target_ptr, hist_ptr, tie_ptr, bar_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, + probs_ptr, target_ptr, hist_ptr, bar_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, CHUNK, row_stride, DRAW: tl.constexpr, G_POW2: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, ): @@ -401,36 +374,27 @@ def _topk_fused( lo, above = _bits_round(probs_ptr, base, start, end, hrow + BINS, brow, target, lo, above, 2 * G, 15, 1 << 23, BINS, BLOCK) lo, above = _bits_round(probs_ptr, base, start, end, hrow + 2 * BINS, brow, target, lo, above, 3 * G, 7, 1 << 15, BINS, BLOCK) lo, above = _bits_round(probs_ptr, base, start, end, hrow + 3 * BINS, brow, target, lo, above, 4 * G, 0, 1 << 7, BINS, BLOCK) - _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), above, target, - 0.0, 0, tie_ptr, brow, 5 * G, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, - False, False, DRAW, G_POW2, BLOCK) + _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), brow, 5 * G, + ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, DRAW, G_POW2, BLOCK) @triton.jit def _keep_tail( - probs_ptr, base, start, end, pid, row, cta, thr, above, target, floor_thr, floor_limit, - tie_ptr, bar_ptr, tie_need, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, - HAS_FLOOR: tl.constexpr, MASS_TARGET: tl.constexpr, DRAW: tl.constexpr, - G_POW2: tl.constexpr, BLOCK: tl.constexpr, + probs_ptr, base, start, end, pid, row, cta, thr, bar_ptr, need, + ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, + DRAW: tl.constexpr, G_POW2: tl.constexpr, BLOCK: tl.constexpr, ): + # Keep x >= thr over this chunk. This deliberately retains every boundary tie, + # matching flashinfer's top-k and top-p filtering semantics. s = 0.0 - local_ties = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - s += tl.sum(tl.where(mask & (x > thr), x, 0.0), 0) - local_ties += tl.sum((mask & (x == thr)).to(tl.int32), 0) - - tie_before, tie_limit = _tie_prefix(local_ties, pid, row, cta, thr, tie_ptr, bar_ptr, - tie_need, target, above, G, G_POW2, MASS_TARGET) - if HAS_FLOOR: - tie_limit = tl.where(thr == floor_thr, tl.minimum(tie_limit, floor_limit), tie_limit) - selected_ties = tl.maximum(0, tl.minimum(local_ties, tie_limit - tie_before)) - s += selected_ties.to(tl.float32) * thr + s += tl.sum(tl.where(mask & (x >= thr), x, 0.0), 0) if DRAW: tl.store(psum_ptr + pid, s) - _row_barrier(bar_ptr, tie_need + G) + _row_barrier(bar_ptr, need) goff = tl.arange(0, G_POW2) gmask = goff < G ps = tl.load(psum_ptr + row * G + goff, mask=gmask, other=0.0, cache_modifier=".cg") @@ -438,12 +402,11 @@ def _keep_tail( incl = tl.sum(tl.where(goff <= cta, ps, 0.0), 0) tgt = tl.load(u_ptr + row) * tl.sum(ps, 0) last_kept = start * 0 - 1 - seen = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - kept, seen = _keep_mask(x, mask, thr, tie_before, tie_limit, seen) + kept = mask & (x >= thr) wv = tl.where(kept, x, 0.0) cval = acc + tl.cumsum(wv, 0) idx = tl.where(cval > tgt, offs, V) @@ -459,15 +422,13 @@ def _keep_tail( tl.store(tok_ptr + row, last_kept) else: tl.atomic_add(ksum_ptr + row, s) - _row_barrier(bar_ptr, tie_need + G) + _row_barrier(bar_ptr, need) inv = 1.0 / tl.atomic_add(ksum_ptr + row, 0.0) - seen = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - kept, seen = _keep_mask(x, mask, thr, tie_before, tie_limit, seen) - tl.store(out_ptr + base + offs, tl.where(kept, x * inv, 0.0), mask=mask) + tl.store(out_ptr + base + offs, tl.where(x >= thr, x * inv, 0.0), mask=mask) _PMBINS = 256 @@ -476,26 +437,21 @@ def _keep_tail( @triton.jit def _pmass_round( probs_ptr, base, start, end, priv_ptr, mass_ptr, bar_ptr, target, lo, above, need, - floor_thr, floor_before, floor_limit, TOPK: tl.constexpr, S: tl.constexpr, WIDTH: tl.constexpr, BINS: tl.constexpr, BLOCK: tl.constexpr, ): # top-p round over the bit pattern: per-bin MASS (exact up to fp32 atomic order) via scatter-add into this # CTA's private buffer, then one reduction into the row buffer, so the bin holding the p crossing is known jj = tl.arange(0, BINS) - seen = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - eligible = mask - if TOPK: - eligible, seen = _keep_mask(x, mask, floor_thr, floor_before, floor_limit, seen) y = x.to(tl.int32, bitcast=True) d = y - lo if WIDTH == 0: - inrange = eligible & (y >= lo) & (y <= _INF_BITS) + inrange = mask & (y >= lo) & (y <= _INF_BITS) else: - inrange = eligible & (y >= lo) & (d < WIDTH) & (y <= _INF_BITS) + inrange = mask & (y >= lo) & (d < WIDTH) & (y <= _INF_BITS) tl.atomic_add(priv_ptr + tl.where(inrange, d >> S, 0), x, mask=inrange) # every warp's scatter-adds must land before any thread reads the private bins back tl.debug_barrier() @@ -513,7 +469,7 @@ def _pmass_round( @triton.jit def _topp_fused( - probs_ptr, tp_ptr, tk_ptr, hist_ptr, priv_ptr, mass_ptr, tie_ptr, bar_ptr, ksumk_ptr, ksum_ptr, psum_ptr, u_ptr, + probs_ptr, tp_ptr, tk_ptr, hist_ptr, priv_ptr, mass_ptr, bar_ptr, ksumk_ptr, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, CHUNK, row_stride, TOPK: tl.constexpr, DRAW: tl.constexpr, G_POW2: tl.constexpr, KBINS: tl.constexpr, PBINS: tl.constexpr, BLOCK: tl.constexpr, @@ -528,9 +484,6 @@ def _topp_fused( end = tl.minimum(start + CHUNK, V) brow = bar_ptr + row lo = 0 - thr_k = 0.0 - tk_before = 0 - tk_limit = 0 if TOPK: tk = tl.maximum(tl.load(tk_ptr + row), 1) hk = hist_ptr + row * 4 * KBINS @@ -541,21 +494,15 @@ def _topp_fused( lo, above_i = _bits_round(probs_ptr, base, start, end, hk + 3 * KBINS, brow, tk, lo, above_i, 4 * G, 0, 1 << 7, KBINS, BLOCK) thr_k = lo.to(tl.float32, bitcast=True) s = 0.0 - local_ties = 0 for s0 in tl.range(start, end, BLOCK): offs = s0 + tl.arange(0, BLOCK) mask = offs < end x = tl.load(probs_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) - s += tl.sum(tl.where(mask & (x > thr_k), x, 0.0), 0) - local_ties += tl.sum((mask & (x == thr_k)).to(tl.int32), 0) - tk_before, tk_limit = _tie_prefix(local_ties, pid, row, cta, thr_k, tie_ptr, brow, - 5 * G, tk, above_i, G, G_POW2, False) - selected_ties = tl.maximum(0, tl.minimum(local_ties, tk_limit - tk_before)) - s += selected_ties.to(tl.float32) * thr_k + s += tl.sum(tl.where(mask & (x >= thr_k), x, 0.0), 0) tl.atomic_add(ksumk_ptr + row, s) - _row_barrier(brow, 6 * G) + _row_barrier(brow, 5 * G) target = tl.load(tp_ptr + row) * tl.atomic_add(ksumk_ptr + row, 0.0) - done = 6 + done = 5 else: target = tl.load(tp_ptr + row) done = 0 @@ -563,16 +510,15 @@ def _topp_fused( pp = priv_ptr + pid * 4 * PBINS above = 0.0 lo, above = _pmass_round(probs_ptr, base, start, end, pp, mp, brow, target, lo, above, (done + 1) * G, - thr_k, tk_before, tk_limit, TOPK, 23, 0, PBINS, BLOCK) + 23, 0, PBINS, BLOCK) lo, above = _pmass_round(probs_ptr, base, start, end, pp + PBINS, mp + PBINS, brow, target, lo, above, - (done + 2) * G, thr_k, tk_before, tk_limit, TOPK, 15, 1 << 23, PBINS, BLOCK) + (done + 2) * G, 15, 1 << 23, PBINS, BLOCK) lo, above = _pmass_round(probs_ptr, base, start, end, pp + 2 * PBINS, mp + 2 * PBINS, brow, target, lo, above, - (done + 3) * G, thr_k, tk_before, tk_limit, TOPK, 7, 1 << 15, PBINS, BLOCK) + (done + 3) * G, 7, 1 << 15, PBINS, BLOCK) lo, above = _pmass_round(probs_ptr, base, start, end, pp + 3 * PBINS, mp + 3 * PBINS, brow, target, lo, above, - (done + 4) * G, thr_k, tk_before, tk_limit, TOPK, 0, 1 << 7, PBINS, BLOCK) - _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), above, target, - thr_k, tk_limit, tie_ptr, brow, (done + 5) * G, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, - V, G, TOPK, True, DRAW, G_POW2, BLOCK) + (done + 4) * G, 0, 1 << 7, PBINS, BLOCK) + _keep_tail(probs_ptr, base, start, end, pid, row, cta, lo.to(tl.float32, bitcast=True), brow, + (done + 5) * G, ksum_ptr, psum_ptr, u_ptr, out_ptr, tok_ptr, V, G, DRAW, G_POW2, BLOCK) _COOPERATIVE_DISABLED = set() @@ -595,14 +541,12 @@ def _fused_launch(probs, kernel, tk, tp, draw, seed, offset, force_single=False) G, CHUNK = _fused_plan(B, V, dev, force_single) n_hist = 4 * _KBINS if (kernel is _topk_fused or tk is not None) else 0 n_mass = 4 * _PMBINS if kernel is _topp_fused else 0 - # hist[B, n_hist] | mass[B, n_mass] | priv[B * G, n_mass] | tie[B * G] | bar/ksum/ksum_k/tok[B] - ws = torch.zeros(B * (n_hist + n_mass) + B * G * (n_mass + 1) + 4 * B, device=dev, dtype=torch.int32) + # hist[B, n_hist] | mass[B, n_mass] | priv[B * G, n_mass] | bar/ksum/ksum_k/tok[B] + ws = torch.zeros(B * (n_hist + n_mass) + B * G * n_mass + 4 * B, device=dev, dtype=torch.int32) hist = ws[:B * n_hist] mass = ws[B * n_hist:B * (n_hist + n_mass)].view(torch.float32) priv = ws[B * (n_hist + n_mass):B * (n_hist + n_mass) + B * G * n_mass].view(torch.float32) - tie_start = B * (n_hist + n_mass) + B * G * n_mass - tie = ws[tie_start:tie_start + B * G] - tail = tie_start + B * G + tail = B * (n_hist + n_mass) + B * G * n_mass bar = ws[tail:tail + B] ksum = ws[tail + B:tail + 2 * B].view(torch.float32) ksum_k = ws[tail + 2 * B:tail + 3 * B].view(torch.float32) @@ -620,10 +564,10 @@ def _fused_launch(probs, kernel, tk, tp, draw, seed, offset, force_single=False) common = dict(DRAW=draw, G_POW2=_next_pow2(G), BLOCK=8192 if wide else _FUSED_BLOCK, num_warps=32 if wide else 8, launch_cooperative_grid=G > 1) if kernel is _topk_fused: - _topk_fused[(B * G,)](probs, tk, hist, tie, bar, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), + _topk_fused[(B * G,)](probs, tk, hist, bar, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), BINS=_KBINS, **common) else: - _topp_fused[(B * G,)](probs, tp, tk if tk is not None else tp, hist, priv, mass, tie, bar, ksum_k, ksum, psum, u, out, tok, + _topp_fused[(B * G,)](probs, tp, tk if tk is not None else tp, hist, priv, mass, bar, ksum_k, ksum, psum, u, out, tok, V, G, CHUNK, probs.stride(0), TOPK=tk is not None, KBINS=_KBINS, PBINS=_PMBINS, **common) return res