diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 55bbd1ec0c..fdf7041692 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -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 diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 03716cb3c1..d714976d85 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -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) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index d8f8c8956d..7c1f080fcd 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -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()