iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core] Triton kernel for small-batch top-p only masking (#54651)

Browse Source 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Nick Hill
2026-09-03 07:52:47 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent bb363db9a5
commit bf95f58d10
3 changed files with 627 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/sample/test_topk_topp_sampler.py
+71
View File
@@ -771,6 +771,77 @@ class TestTritonTopkTopp:
if finite_in > 0:
assert kept > 0, f"Row {i}: no tokens kept"
def test_mixed_batch_grammar_rows_get_topp(self):
"""Grammar-masked rows in a mixed top-k batch must still get top-p.
Regression: with the split top-p pipeline active (batch <= 64), a
row with k < vocab but <= k finite logits was covered by neither the
monolithic kernel (standalone top-p disabled) nor the split pipeline
(skips rows with k < vocab), so its top-p was silently dropped.
"""
batch_size, vocab_size = 8, 32000
logits = torch.full((batch_size, vocab_size), float("-inf"))
# Rows 0-2: grammar rows (30 finite tokens); k=50 is a no-op for them.
for i in range(3):
idx = torch.randperm(vocab_size, generator=self.generator)[:30]
logits[i, idx] = torch.randn(30, generator=self.generator)
# Rows 3-4: p-only rows (k = vocab); rows 5-7: normal top-k+top-p.
logits[3:5] = torch.randn(2, vocab_size, generator=self.generator)
logits[5:] = torch.randn(3, vocab_size, generator=self.generator)
k = torch.tensor(
[50, 50, 50, vocab_size, vocab_size, 50, 50, 50], dtype=torch.int32
)
p = torch.tensor([0.5, 0.5, 0.5, 0.9, 0.9, 0.9, 0.9, 0.9])
self._compare_results(logits, k, p)
# Top-p must actually trim the grammar rows beyond the (no-op) top-k.
from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
kept_topk_only = (
apply_top_k_top_p_triton(logits.clone(), k, None)[:3] != float("-inf")
).sum(-1)
kept_topk_topp = (
apply_top_k_top_p_triton(logits.clone(), k, p)[:3] != float("-inf")
).sum(-1)
assert (kept_topk_topp < kept_topk_only).all(), (
f"top-p did not trim grammar rows: {kept_topk_only.tolist()} -> "
f"{kept_topk_topp.tolist()}"
)
def test_topp_tie_break_deterministic(self):
"""Boundary-duplicate trimming must be deterministic across calls.
Regression: the split top-p mask kernel handed out the tie budget in
cross-program atomic order, so which copies of a tied boundary value
survived could vary between identical calls.
"""
from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
batch_size, vocab_size = 4, 128256
# 1000 tied top tokens; the p=0.5 boundary lands inside the tie
# group (each tied token has mass e/(1000e+1000) ~ 0.073%).
logits = torch.full((batch_size, vocab_size), float("-inf"))
logits[:, :1000] = 1.0
logits[:, 1000:2000] = 0.0
p = torch.full((batch_size,), 0.5, dtype=torch.float32)
results = [apply_top_k_top_p_triton(logits.clone(), None, p) for _ in range(5)]
for r in results[1:]:
assert torch.equal(results[0], r), "non-deterministic tie break"
# The split pipeline now keeps the first `numkeep` copies in index
# order, exactly like the monolithic kernel and the PyTorch impl.
kept_idx = (results[0][0] != float("-inf")).nonzero().flatten()
assert kept_idx.min().item() == 0
assert kept_idx.max().item() == len(kept_idx) - 1
pytorch_kept = (
(apply_top_k_top_p_pytorch(logits.clone(), None, p)[0] != float("-inf"))
.sum()
.item()
)
assert len(kept_idx) == pytorch_kept
# =============================================================================
# FlashInfer top-k/top-p robustness tests
vllm/v1/sample/ops/topk_topp_sampler.py
+4 -12
View File
@@ -352,16 +352,11 @@ def apply_top_k_top_p(
if p is None and k is None:
return logits
if current_platform.is_cpu():
if HAS_TRITON:
return apply_top_k_top_p_triton(logits, k, p)
return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)
if HAS_TRITON and logits.shape[0] >= 8:
if HAS_TRITON:
return apply_top_k_top_p_triton(logits, k, p)
# Use pytorch sort implementation for small batch sizes.
return apply_top_k_top_p_pytorch(logits, k, p)
is_cpu = current_platform.is_cpu()
return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=is_cpu)
def apply_top_k_top_p_pytorch(
@@ -513,7 +508,4 @@ def flashinfer_sample(
def _to_tensor_scalar_tuple(x):
if isinstance(x, torch.Tensor):
return (x, 0)
else:
return (None, x)
return (x, 0) if isinstance(x, torch.Tensor) else (None, x)
vllm/v1/sample/ops/topk_topp_triton.py
+552 -1
View File
@@ -107,6 +107,7 @@ def _topk_topp_kernel(
BLOCK_SIZE_TRUNC: tl.constexpr,
TOPK_ENABLED: tl.constexpr,
TOPP_ENABLED: tl.constexpr,
SPLIT_COVERS_PONLY: tl.constexpr,
):
NUM_TILES: tl.constexpr = (VOCAB_SIZE + BLOCK_SIZE - 1) // BLOCK_SIZE
pid = tl.program_id(0)
@@ -594,8 +595,16 @@ def _topk_topp_kernel(
if TOPP_ENABLED and final_pivot == -float("inf"):
#### STANDALONE TOP-P SAMPLING ####
# When the split top-p pipeline co-runs (mixed top-k + top-p
# batches), it covers p-only rows (k >= VOCAB_SIZE); skip those
# here. Rows whose top-k was a no-op because they have <= k
# finite logits (e.g. grammar masks) are NOT covered by the
# split pipeline and must get standalone top-p here.
run_standalone = True
if SPLIT_COVERS_PONLY and TOPK_ENABLED:
run_standalone = tl.load(K + row_id) < VOCAB_SIZE
p = tl.load(P + row_id)
if p < 1.0:
if run_standalone and p < 1.0:
# Zeroth pass: Compute avg and std from a sample block
offs = tl.arange(0, BLOCK_SIZE)
mask_n = offs < VOCAB_SIZE
@@ -854,6 +863,529 @@ def _topk_topp_kernel(
tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n)
# ---------------------------------------------------------------------------
# Split-row pipeline for small batches.
#
# The monolithic kernel above assigns one program per row, so at small batch
# sizes nearly all SMs idle and every full-vocab sweep is latency-bound
# (~25us each on SM100). For standalone top-p (no top-k truncation) the pivot
# search can hit its iteration cap while sweeping ~the whole vocab per
# iteration, which is slower than the PyTorch sort path.
#
# The pipeline below instead splits each row across S programs and runs the
# pivot search as parallel partial reductions over a fixed number of rounds
# (no host synchronization, CUDA-graph safe). The search state is never
# stored: the per-round combine is a tiny pure function of the stored
# partials, so every program redundantly recomputes the bracket chain from
# them. This is deterministic (same inputs, same code) and far cheaper than
# a vocab sweep, and it saves a kernel launch per combine. The pipeline
# handles only "p-only" rows (p < 1 and no active top-k); rows with an
# active top-k are left to the monolithic kernel, whose truncation path is
# already fast for them.
# ---------------------------------------------------------------------------
_SPLIT_MAX_BATCH = 64
_SPLIT_MAX_SPLITS = 32 # power of two
_SPLIT_FANOUT = 8
_SPLIT_ROUNDS = 5
_TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {}
@triton.jit
def _topp_sb_stats_kernel(
LOGITS,
LOGITS_STRIDE_0,
STATS,
K,
P,
HAS_K: tl.constexpr,
VOCAB_SIZE: tl.constexpr,
S: tl.constexpr,
BLOCK: tl.constexpr,
):
# Per-slice online-softmax partials: (max, sum_exp, min, finite_count).
pid = tl.program_id(0)
row = pid // S
slice_id = pid % S
p = tl.load(P + row)
if p >= 1.0:
return
if HAS_K and tl.load(K + row) < VOCAB_SIZE:
return
slice_len = (VOCAB_SIZE + S - 1) // S
start = slice_id * slice_len
end = tl.minimum(start + slice_len, VOCAB_SIZE)
ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0
offs = tl.arange(0, BLOCK)
m = -float("inf")
exp_sum = 0.0
mn = float("inf")
cnt = tl.zeros((), dtype=tl.float32)
for i in range(start, end, BLOCK):
x = tl.load(ROW + i + offs, mask=i + offs < end, other=-float("inf"))
finite = x > -float("inf")
nm = tl.maximum(m, tl.max(x))
rescale = tl.where(m == -float("inf"), 0.0, tl.exp(m - nm))
e = tl.where(finite, tl.exp(x - nm), 0.0)
exp_sum = exp_sum * rescale + tl.sum(e)
m = nm
mn = tl.minimum(mn, tl.min(tl.where(finite, x, float("inf"))))
cnt += tl.sum(finite.to(tl.float32))
base = (row * S + slice_id) * 4
tl.store(STATS + base + 0, m)
tl.store(STATS + base + 1, exp_sum)
tl.store(STATS + base + 2, mn)
tl.store(STATS + base + 3, cnt)
@triton.jit
def _topp_sb_row_stats(STATS, row, S: tl.constexpr):
"""Combine per-slice partials into row (max, sum_exp, min, n_finite)."""
sidx = tl.arange(0, S)
sb = STATS + row * S * 4
m_s = tl.load(sb + sidx * 4 + 0)
exp_s = tl.load(sb + sidx * 4 + 1)
mn_s = tl.load(sb + sidx * 4 + 2)
c_s = tl.load(sb + sidx * 4 + 3)
M = tl.max(m_s)
w = tl.where(m_s == -float("inf"), 0.0, tl.exp(m_s - M))
Z = tl.sum(exp_s * w)
return M, Z, tl.min(mn_s), tl.sum(c_s)
@triton.jit
def _topp_sb_ladder(lo, hi, F: tl.constexpr):
"""Geometric ladder of F pivots inside (lo, hi)."""
fidx = tl.arange(0, F).to(tl.float32)
return lo * tl.exp(tl.log(hi / lo) * (fidx + 1.0) / (F + 1.0))
@triton.jit
def _topp_sb_combine(
PART_BASE,
Z,
p,
lo,
hi,
best,
best_mass,
best_minl,
best_nmin,
S: tl.constexpr,
F: tl.constexpr,
IS_LAST: tl.constexpr,
):
"""Combine one round's partials and advance the search state.
Returns (done, nomask, pivot, dup, numdup, numkeep, lo, hi, best,
best_mass, best_minl, best_nmin, sel_fidx, best_sel). `best*` tracks the
tightest fallback: the largest evaluated pivot whose above-mass is >= p,
which guarantees the top-p mass invariant if the search fails to
converge. Pivots are recomputed from (lo, hi) via the same ladder the
sweep used, so partials line up with pivots without storing them.
`sel_fidx` / `best_sel` report which ladder slot produced `dup` /
updated `best` this round (-1 if none), so callers can recover the
per-slice (minl, nmin) partials behind the boundary value for
deterministic duplicate handling.
"""
fidx = tl.arange(0, F)
sidx = tl.arange(0, S)
offs2 = sidx[:, None] * (F * 3) + fidx[None, :] * 3
mass_s = tl.load(PART_BASE + offs2 + 0)
minl_s = tl.load(PART_BASE + offs2 + 1)
nmin_s = tl.load(PART_BASE + offs2 + 2)
mass = tl.sum(mass_s, axis=0)
gmin = tl.min(minl_s, axis=0)
nmin = tl.sum(tl.where(tl.abs(minl_s - gmin[None, :]) < 1e-9, nmin_s, 0.0), axis=0)
pivs = _topp_sb_ladder(lo, hi, F)
ok = (mass >= p) & (mass - gmin * nmin < p)
ok_idx = tl.max(tl.where(ok, fidx, -1).to(tl.float32))
ge = mass >= p
ge_idx = tl.max(tl.where(ge, fidx, -1).to(tl.float32))
best_sel = -1.0
if ge_idx >= 0:
bmask = fidx == ge_idx.to(tl.int32)
bp = tl.sum(tl.where(bmask, pivs, 0.0))
if bp > best:
best = bp
best_mass = tl.sum(tl.where(bmask, mass, 0.0))
best_minl = tl.sum(tl.where(bmask, gmin, 0.0))
best_nmin = tl.sum(tl.where(bmask, nmin, 0.0))
best_sel = ge_idx
done = 0.0
nomask = 0.0
pivot = 0.0
sel_mass = 0.0
dup = 1.0
numdup = 1.0
numkeep = 1.0
sel_fidx = -1.0
if ok_idx >= 0:
bmask = fidx == ok_idx.to(tl.int32)
pivot = tl.sum(tl.where(bmask, pivs, 0.0))
sel_mass = tl.sum(tl.where(bmask, mass, 0.0))
dup = tl.sum(tl.where(bmask, gmin, 0.0))
numdup = tl.sum(tl.where(bmask, nmin, 0.0))
sel_fidx = ok_idx
done = 1.0
else:
new_lo = tl.maximum(lo, tl.max(tl.where(ge, pivs, -float("inf"))))
new_hi = tl.minimum(hi, tl.min(tl.where(~ge, pivs, float("inf"))))
if IS_LAST or not (new_hi > new_lo * (1.0 + 1e-7)):
# Out of rounds or the range collapsed: fall back to the
# tightest pivot with mass >= p, or keep the whole row if no
# pivot qualified.
if best > 0.0:
pivot = best
sel_mass = best_mass
dup = best_minl
numdup = best_nmin
else:
nomask = 1.0
done = 1.0
lo = new_lo
hi = new_hi
if (done != 0.0) and (nomask == 0.0):
# Mirror the monolithic guard: a pivot at/above the max logit would
# mask everything under the strict `>` comparison.
nomask = tl.where(pivot * Z >= 1.0, 1.0, 0.0)
# Duplicate (boundary-value) handling, same formula as the
# monolithic kernel: keep only some of the boundary duplicates.
numkeep = numdup - tl.cast(tl.maximum(sel_mass - p, 0.0) / dup, tl.uint32).to(
tl.float32
)
numkeep = tl.minimum(tl.maximum(numkeep, 1.0), numdup)
return (
done,
nomask,
pivot,
dup,
numdup,
numkeep,
lo,
hi,
best,
best_mass,
best_minl,
best_nmin,
sel_fidx,
best_sel,
)
@triton.jit
def _topp_sb_step_kernel(
LOGITS,
LOGITS_STRIDE_0,
STATS,
PARTS,
K,
P,
ROUND,
HAS_K: tl.constexpr,
S: tl.constexpr,
F: tl.constexpr,
NUM_ROUNDS: tl.constexpr,
VOCAB_SIZE: tl.constexpr,
BLOCK: tl.constexpr,
):
pid = tl.program_id(0)
row = pid // S
slice_id = pid % S
p = tl.load(P + row)
if p >= 1.0:
return
if HAS_K and tl.load(K + row) < VOCAB_SIZE:
return
M, Z, mn, num_finite = _topp_sb_row_stats(STATS, row, S)
if (Z <= 0.0) or (num_finite < 1.0):
return
max_prob = 1.0 / Z
min_prob = tl.maximum(tl.exp(mn - M) / Z, 1e-30)
if max_prob <= min_prob * (1.0 + 1e-6):
# (Nearly) uniform row: no pivot can differentiate; keep everything.
return
# Replay the combine chain for rounds 0..ROUND-1 (deterministic; each
# is a tiny pure function of the stored partials). Partials from rounds
# past a row's convergence are stale, so the chain stops at `done`.
lo = min_prob
hi = max_prob
best = -1.0
best_mass = 0.0
best_minl = 1.0
best_nmin = 0.0
done = 0.0
for j in range(ROUND):
if done == 0.0:
(
done,
nomask,
pivot,
dup,
numdup,
numkeep,
lo,
hi,
best,
best_mass,
best_minl,
best_nmin,
sel_fidx,
best_sel,
) = _topp_sb_combine(
PARTS + (row * NUM_ROUNDS + j) * S * F * 3,
Z,
p,
lo,
hi,
best,
best_mass,
best_minl,
best_nmin,
S=S,
F=F,
IS_LAST=False,
)
if done != 0.0:
return
# Sweep this slice against the ladder pivots derived from (lo, hi).
pivs = _topp_sb_ladder(lo, hi, F)
slice_len = (VOCAB_SIZE + S - 1) // S
start = slice_id * slice_len
end = tl.minimum(start + slice_len, VOCAB_SIZE)
ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0
offs = tl.arange(0, BLOCK)
fidx = tl.arange(0, F)
mass = tl.zeros([F], dtype=tl.float32)
minl = tl.full([F], 1.0, dtype=tl.float32)
nmin = tl.zeros([F], dtype=tl.float32)
for i in range(start, end, BLOCK):
x = tl.load(ROW + i + offs, mask=i + offs < end, other=-float("inf"))
pr = tl.exp(x - M) / Z
pr2 = tl.broadcast_to(pr[None, :], (F, BLOCK))
above = pr2 > pivs[:, None]
mass += tl.sum(tl.where(above, pr2, 0.0), axis=1)
cand = tl.where(above, pr2, float("inf"))
tmin = tl.min(cand, axis=1)
teq = above & (tl.abs(pr2 - tmin[:, None]) < 1e-9)
tcnt = tl.sum(teq.to(tl.float32), axis=1)
is_new = tmin < minl
is_same = tl.abs(tmin - minl) < 1e-9
nmin = tl.where(is_new, tcnt, nmin + tl.where(is_same, tcnt, 0.0))
minl = tl.minimum(minl, tmin)
pbase = PARTS + ((row * NUM_ROUNDS + ROUND) * S + slice_id) * F * 3
tl.store(pbase + fidx * 3 + 0, mass)
tl.store(pbase + fidx * 3 + 1, minl)
tl.store(pbase + fidx * 3 + 2, nmin)
@triton.jit
def _topp_sb_mask_kernel(
LOGITS,
LOGITS_STRIDE_0,
STATS,
PARTS,
K,
P,
HAS_K: tl.constexpr,
MASK_VALUE: tl.constexpr,
S: tl.constexpr,
F: tl.constexpr,
NUM_ROUNDS: tl.constexpr,
VOCAB_SIZE: tl.constexpr,
BLOCK: tl.constexpr,
):
pid = tl.program_id(0)
row = pid // S
slice_id = pid % S
p = tl.load(P + row)
if p >= 1.0:
return
if HAS_K and tl.load(K + row) < VOCAB_SIZE:
return
M, Z, mn, num_finite = _topp_sb_row_stats(STATS, row, S)
if (Z <= 0.0) or (num_finite < 1.0):
return
max_prob = 1.0 / Z
min_prob = tl.maximum(tl.exp(mn - M) / Z, 1e-30)
if max_prob <= min_prob * (1.0 + 1e-6):
return # (nearly) uniform row: keep everything
# Replay the full combine chain; the last round finalizes with fallback
# semantics. Rows that converged early stop the chain at that round.
# `fin` tracks the (round, ladder slot) whose per-slice partials define
# the boundary value, for deterministic duplicate handling below.
lo = min_prob
hi = max_prob
best = -1.0
best_mass = 0.0
best_minl = 1.0
best_nmin = 0.0
done = 0.0
nomask = 0.0
pivot = 0.0
dup = 1.0
numdup = 1.0
numkeep = 1.0
fin = -1.0
best_fin = -1.0
for j in range(NUM_ROUNDS):
if done == 0.0:
(
done,
nomask,
pivot,
dup,
numdup,
numkeep,
lo,
hi,
best,
best_mass,
best_minl,
best_nmin,
sel_fidx,
best_sel,
) = _topp_sb_combine(
PARTS + (row * NUM_ROUNDS + j) * S * F * 3,
Z,
p,
lo,
hi,
best,
best_mass,
best_minl,
best_nmin,
S=S,
F=F,
IS_LAST=j == NUM_ROUNDS - 1,
)
if sel_fidx >= 0.0:
fin = j * F + sel_fidx
if best_sel >= 0.0:
best_fin = j * F + best_sel
if nomask != 0.0:
return
if fin < 0.0:
fin = best_fin
logZ = tl.log(Z)
pivot_logit = tl.log(pivot) + logZ + M
# numdup/numkeep are exact small integers held in fp32.
ties = numkeep < numdup
num_preceding = 0.0
if ties:
# Deterministic, index-ordered duplicate budget: per-slice counts of
# the boundary value come from the finalizing round's partials (a
# slice holds copies of `dup` iff its minl matches `dup`, with the
# count in nmin), so each slice knows exactly how many copies were
# kept by lower-indexed slices.
fin_round = tl.cast(fin / F, tl.int32)
fin_slot = tl.cast(fin - tl.cast(fin_round, tl.float32) * F, tl.int32)
pbase = PARTS + ((row * NUM_ROUNDS + fin_round) * S * F + fin_slot) * 3
sidx = tl.arange(0, S)
minl_s = tl.load(pbase + sidx * (F * 3) + 1)
nmin_s = tl.load(pbase + sidx * (F * 3) + 2)
dupc_s = tl.where(tl.abs(minl_s - dup) < 1e-9, nmin_s, 0.0)
num_preceding = tl.sum(tl.where(sidx < slice_id, dupc_s, 0.0))
slice_len = (VOCAB_SIZE + S - 1) // S
start = slice_id * slice_len
end = tl.minimum(start + slice_len, VOCAB_SIZE)
ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0
offs = tl.arange(0, BLOCK)
for i in range(start, end, BLOCK):
mask_n = i + offs < end
x = tl.load(ROW + i + offs, mask=mask_n, other=-float("inf"))
keep = x > pivot_logit
if ties:
# Match the boundary value in probability space with exact
# equality: `dup` is a probability computed by the step kernels
# as exp(x - M) / Z, so recomputing it here with the same ops
# reproduces it bitwise on any platform. (The log/exp round-trip
# to logit space loses ~1e-6 to cancellation, breaking any
# absolute tolerance.)
dmask = tl.exp(x - M) / Z == dup
cum = tl.cumsum(dmask.to(tl.int32)).to(tl.float32) + num_preceding
keep_dup = dmask & (cum <= numkeep)
keep = keep & (~dmask | keep_dup)
tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n)
def _apply_topp_split(
logits: torch.Tensor,
k: torch.Tensor | None,
p: torch.Tensor,
mask_value: float,
num_sm: int,
) -> None:
"""Split-row top-p pipeline for small batches; masks logits in place."""
batch_size, vocab_size = logits.shape
splits = 1
while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:
splits *= 2
ws = _TRITON_SPLIT_CACHE.get(logits.device)
if ws is None:
ws = {
"stats": logits.new_empty((_SPLIT_MAX_BATCH * _SPLIT_MAX_SPLITS, 4)),
"parts": logits.new_empty(
(_SPLIT_MAX_BATCH, _SPLIT_ROUNDS, _SPLIT_MAX_SPLITS, _SPLIT_FANOUT, 3)
),
}
_TRITON_SPLIT_CACHE[logits.device] = ws
k_ptr = k if k is not None else logits # dummy pointer when HAS_K=False
has_k = k is not None
_topp_sb_stats_kernel[(batch_size * splits,)](
logits,
logits.stride(0),
ws["stats"],
k_ptr,
p,
HAS_K=has_k,
VOCAB_SIZE=vocab_size,
S=splits,
BLOCK=8192,
num_warps=8,
)
for round_i in range(_SPLIT_ROUNDS):
_topp_sb_step_kernel[(batch_size * splits,)](
logits,
logits.stride(0),
ws["stats"],
ws["parts"],
k_ptr,
p,
round_i,
HAS_K=has_k,
S=splits,
F=_SPLIT_FANOUT,
NUM_ROUNDS=_SPLIT_ROUNDS,
VOCAB_SIZE=vocab_size,
BLOCK=2048,
num_warps=8,
)
_topp_sb_mask_kernel[(batch_size * splits,)](
logits,
logits.stride(0),
ws["stats"],
ws["parts"],
k_ptr,
p,
HAS_K=has_k,
MASK_VALUE=mask_value,
S=splits,
F=_SPLIT_FANOUT,
NUM_ROUNDS=_SPLIT_ROUNDS,
VOCAB_SIZE=vocab_size,
BLOCK=8192,
num_warps=8,
)
def apply_top_k_top_p_triton(
logits: torch.Tensor,
k: torch.Tensor | None,
@@ -905,6 +1437,21 @@ def apply_top_k_top_p_triton(
p_ptr = logits # Dummy pointer (won't be read)
num_sm = num_compute_units(logits.device.index)
# At small batch sizes the monolithic kernel's standalone top-p path is
# latency-bound (one program per row, serial vocab sweeps), so use the
# split-row pipeline for rows without an active top-k. Rows with an
# active top-k still go through the monolithic kernel (its truncation
# path is already fast). The monolithic's standalone top-p branch then
# covers only rows whose top-k is a no-op (k < vocab but <= k finite
# logits, e.g. grammar masks), which the split pipeline skips.
use_split = (
topp_enabled and batch_size <= _SPLIT_MAX_BATCH and logits.device.type == "cuda"
)
if use_split and not topk_enabled:
_apply_topp_split(logits, None, p_ptr, mask_value, num_sm)
return logits
NUM_PROGRAMS = min(num_sm, batch_size)
# Cache per-Triton Program buffer on each device.
@@ -961,8 +1508,11 @@ def apply_top_k_top_p_triton(
BLOCK_SIZE_TRUNC=block_size_trunc,
TOPK_ENABLED=topk_enabled,
TOPP_ENABLED=topp_enabled,
SPLIT_COVERS_PONLY=use_split,
**launch_kwargs,
)
if use_split:
_apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm)
return logits
@@ -970,4 +1520,5 @@ def apply_top_k_top_p_triton(
def reset_buffer_cache():
_TRITON_BUFFER_CACHE.clear()
_TRITON_TABLE_CACHE.clear()
_TRITON_SPLIT_CACHE.clear()
torch.accelerator.empty_cache()