@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bf95f58d10

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

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nick Hill committed Sep 3, 2026 at 07:52 UTC bf95f58d10891293aa6fde4e21f184d12ec0ba34

3 files changed +627 -13

tests/v1/sample/test_topk_topp_sampler.py

+71

@@ -771,6 +771,77 @@ class TestTritonTopkTopp:

771 if finite_in > 0:

772 assert kept > 0, f"Row {i}: no tokens kept"

773

774 + def test_mixed_batch_grammar_rows_get_topp(self):

775 + """Grammar-masked rows in a mixed top-k batch must still get top-p.

776 +

777 + Regression: with the split top-p pipeline active (batch <= 64), a

778 + row with k < vocab but <= k finite logits was covered by neither the

779 + monolithic kernel (standalone top-p disabled) nor the split pipeline

780 + (skips rows with k < vocab), so its top-p was silently dropped.

781 + """

782 + batch_size, vocab_size = 8, 32000

783 + logits = torch.full((batch_size, vocab_size), float("-inf"))

784 + # Rows 0-2: grammar rows (30 finite tokens); k=50 is a no-op for them.

785 + for i in range(3):

786 + idx = torch.randperm(vocab_size, generator=self.generator)[:30]

787 + logits[i, idx] = torch.randn(30, generator=self.generator)

788 + # Rows 3-4: p-only rows (k = vocab); rows 5-7: normal top-k+top-p.

789 + logits[3:5] = torch.randn(2, vocab_size, generator=self.generator)

790 + logits[5:] = torch.randn(3, vocab_size, generator=self.generator)

791 + k = torch.tensor(

792 + [50, 50, 50, vocab_size, vocab_size, 50, 50, 50], dtype=torch.int32

793 + )

794 + p = torch.tensor([0.5, 0.5, 0.5, 0.9, 0.9, 0.9, 0.9, 0.9])

795 +

796 + self._compare_results(logits, k, p)

797 +

798 + # Top-p must actually trim the grammar rows beyond the (no-op) top-k.

799 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton

800 +

801 + kept_topk_only = (

802 + apply_top_k_top_p_triton(logits.clone(), k, None)[:3] != float("-inf")

803 + ).sum(-1)

804 + kept_topk_topp = (

805 + apply_top_k_top_p_triton(logits.clone(), k, p)[:3] != float("-inf")

806 + ).sum(-1)

807 + assert (kept_topk_topp < kept_topk_only).all(), (

808 + f"top-p did not trim grammar rows: {kept_topk_only.tolist()} -> "

809 + f"{kept_topk_topp.tolist()}"

810 + )

811 +

812 + def test_topp_tie_break_deterministic(self):

813 + """Boundary-duplicate trimming must be deterministic across calls.

814 +

815 + Regression: the split top-p mask kernel handed out the tie budget in

816 + cross-program atomic order, so which copies of a tied boundary value

817 + survived could vary between identical calls.

818 + """

819 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton

820 +

821 + batch_size, vocab_size = 4, 128256

822 + # 1000 tied top tokens; the p=0.5 boundary lands inside the tie

823 + # group (each tied token has mass e/(1000e+1000) ~ 0.073%).

824 + logits = torch.full((batch_size, vocab_size), float("-inf"))

825 + logits[:, :1000] = 1.0

826 + logits[:, 1000:2000] = 0.0

827 + p = torch.full((batch_size,), 0.5, dtype=torch.float32)

828 +

829 + results = [apply_top_k_top_p_triton(logits.clone(), None, p) for _ in range(5)]

830 + for r in results[1:]:

831 + assert torch.equal(results[0], r), "non-deterministic tie break"

832 +

833 + # The split pipeline now keeps the first `numkeep` copies in index

834 + # order, exactly like the monolithic kernel and the PyTorch impl.

835 + kept_idx = (results[0][0] != float("-inf")).nonzero().flatten()

836 + assert kept_idx.min().item() == 0

837 + assert kept_idx.max().item() == len(kept_idx) - 1

838 + pytorch_kept = (

839 + (apply_top_k_top_p_pytorch(logits.clone(), None, p)[0] != float("-inf"))

840 + .sum()

841 + .item()

842 + )

843 + assert len(kept_idx) == pytorch_kept

844 +

845

846 # =============================================================================

847 # FlashInfer top-k/top-p robustness tests

vllm/v1/sample/ops/topk_topp_sampler.py

+4 -12

@@ -352,16 +352,11 @@ def apply_top_k_top_p(

352 if p is None and k is None:

353 return logits

354

355 - if current_platform.is_cpu():

356 - if HAS_TRITON:

357 - return apply_top_k_top_p_triton(logits, k, p)

358 - return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)

359 -

360 - if HAS_TRITON and logits.shape[0] >= 8:

355 + if HAS_TRITON:

356 return apply_top_k_top_p_triton(logits, k, p)

357

363 - # Use pytorch sort implementation for small batch sizes.

364 - return apply_top_k_top_p_pytorch(logits, k, p)

358 + is_cpu = current_platform.is_cpu()

359 + return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=is_cpu)

360

361

362 def apply_top_k_top_p_pytorch(

@@ -513,7 +508,4 @@ def flashinfer_sample(

508

509

510 def _to_tensor_scalar_tuple(x):

516 - if isinstance(x, torch.Tensor):

517 - return (x, 0)

518 - else:

519 - return (None, x)

511 + return (x, 0) if isinstance(x, torch.Tensor) else (None, x)

vllm/v1/sample/ops/topk_topp_triton.py

+552 -1

@@ -107,6 +107,7 @@ def _topk_topp_kernel(

107 BLOCK_SIZE_TRUNC: tl.constexpr,

108 TOPK_ENABLED: tl.constexpr,

109 TOPP_ENABLED: tl.constexpr,

110 + SPLIT_COVERS_PONLY: tl.constexpr,

111 ):

112 NUM_TILES: tl.constexpr = (VOCAB_SIZE + BLOCK_SIZE - 1) // BLOCK_SIZE

113 pid = tl.program_id(0)

@@ -594,8 +595,16 @@ def _topk_topp_kernel(

595

596 if TOPP_ENABLED and final_pivot == -float("inf"):

597 #### STANDALONE TOP-P SAMPLING ####

598 + # When the split top-p pipeline co-runs (mixed top-k + top-p

599 + # batches), it covers p-only rows (k >= VOCAB_SIZE); skip those

600 + # here. Rows whose top-k was a no-op because they have <= k

601 + # finite logits (e.g. grammar masks) are NOT covered by the

602 + # split pipeline and must get standalone top-p here.

603 + run_standalone = True

604 + if SPLIT_COVERS_PONLY and TOPK_ENABLED:

605 + run_standalone = tl.load(K + row_id) < VOCAB_SIZE

606 p = tl.load(P + row_id)

598 - if p < 1.0:

607 + if run_standalone and p < 1.0:

608 # Zeroth pass: Compute avg and std from a sample block

609 offs = tl.arange(0, BLOCK_SIZE)

610 mask_n = offs < VOCAB_SIZE

@@ -854,6 +863,529 @@ def _topk_topp_kernel(

863 tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n)

864

865

866 +# ---------------------------------------------------------------------------

867 +# Split-row pipeline for small batches.

868 +#

869 +# The monolithic kernel above assigns one program per row, so at small batch

870 +# sizes nearly all SMs idle and every full-vocab sweep is latency-bound

871 +# (~25us each on SM100). For standalone top-p (no top-k truncation) the pivot

872 +# search can hit its iteration cap while sweeping ~the whole vocab per

873 +# iteration, which is slower than the PyTorch sort path.

874 +#

875 +# The pipeline below instead splits each row across S programs and runs the

876 +# pivot search as parallel partial reductions over a fixed number of rounds

877 +# (no host synchronization, CUDA-graph safe). The search state is never

878 +# stored: the per-round combine is a tiny pure function of the stored

879 +# partials, so every program redundantly recomputes the bracket chain from

880 +# them. This is deterministic (same inputs, same code) and far cheaper than

881 +# a vocab sweep, and it saves a kernel launch per combine. The pipeline

882 +# handles only "p-only" rows (p < 1 and no active top-k); rows with an

883 +# active top-k are left to the monolithic kernel, whose truncation path is

884 +# already fast for them.

885 +# ---------------------------------------------------------------------------

886 +

887 +_SPLIT_MAX_BATCH = 64

888 +_SPLIT_MAX_SPLITS = 32 # power of two

889 +_SPLIT_FANOUT = 8

890 +_SPLIT_ROUNDS = 5

891 +

892 +_TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {}

893 +

894 +

895 +@triton.jit

896 +def _topp_sb_stats_kernel(

897 + LOGITS,

898 + LOGITS_STRIDE_0,

899 + STATS,

900 + K,

901 + P,

902 + HAS_K: tl.constexpr,

903 + VOCAB_SIZE: tl.constexpr,

904 + S: tl.constexpr,

905 + BLOCK: tl.constexpr,

906 +):

907 + # Per-slice online-softmax partials: (max, sum_exp, min, finite_count).

908 + pid = tl.program_id(0)

909 + row = pid // S

910 + slice_id = pid % S

911 + p = tl.load(P + row)

912 + if p >= 1.0:

913 + return

914 + if HAS_K and tl.load(K + row) < VOCAB_SIZE:

915 + return

916 + slice_len = (VOCAB_SIZE + S - 1) // S

917 + start = slice_id * slice_len

918 + end = tl.minimum(start + slice_len, VOCAB_SIZE)

919 + ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0

920 + offs = tl.arange(0, BLOCK)

921 + m = -float("inf")

922 + exp_sum = 0.0

923 + mn = float("inf")

924 + cnt = tl.zeros((), dtype=tl.float32)

925 + for i in range(start, end, BLOCK):

926 + x = tl.load(ROW + i + offs, mask=i + offs < end, other=-float("inf"))

927 + finite = x > -float("inf")

928 + nm = tl.maximum(m, tl.max(x))

929 + rescale = tl.where(m == -float("inf"), 0.0, tl.exp(m - nm))

930 + e = tl.where(finite, tl.exp(x - nm), 0.0)

931 + exp_sum = exp_sum * rescale + tl.sum(e)

932 + m = nm

933 + mn = tl.minimum(mn, tl.min(tl.where(finite, x, float("inf"))))

934 + cnt += tl.sum(finite.to(tl.float32))

935 + base = (row * S + slice_id) * 4

936 + tl.store(STATS + base + 0, m)

937 + tl.store(STATS + base + 1, exp_sum)

938 + tl.store(STATS + base + 2, mn)

939 + tl.store(STATS + base + 3, cnt)

940 +

941 +

942 +@triton.jit

943 +def _topp_sb_row_stats(STATS, row, S: tl.constexpr):

944 + """Combine per-slice partials into row (max, sum_exp, min, n_finite)."""

945 + sidx = tl.arange(0, S)

946 + sb = STATS + row * S * 4

947 + m_s = tl.load(sb + sidx * 4 + 0)

948 + exp_s = tl.load(sb + sidx * 4 + 1)

949 + mn_s = tl.load(sb + sidx * 4 + 2)

950 + c_s = tl.load(sb + sidx * 4 + 3)

951 + M = tl.max(m_s)

952 + w = tl.where(m_s == -float("inf"), 0.0, tl.exp(m_s - M))

953 + Z = tl.sum(exp_s * w)

954 + return M, Z, tl.min(mn_s), tl.sum(c_s)

955 +

956 +

957 +@triton.jit

958 +def _topp_sb_ladder(lo, hi, F: tl.constexpr):

959 + """Geometric ladder of F pivots inside (lo, hi)."""

960 + fidx = tl.arange(0, F).to(tl.float32)

961 + return lo * tl.exp(tl.log(hi / lo) * (fidx + 1.0) / (F + 1.0))

962 +

963 +

964 +@triton.jit

965 +def _topp_sb_combine(

966 + PART_BASE,

967 + Z,

968 + p,

969 + lo,

970 + hi,

971 + best,

972 + best_mass,

973 + best_minl,

974 + best_nmin,

975 + S: tl.constexpr,

976 + F: tl.constexpr,

977 + IS_LAST: tl.constexpr,

978 +):

979 + """Combine one round's partials and advance the search state.

980 +

981 + Returns (done, nomask, pivot, dup, numdup, numkeep, lo, hi, best,

982 + best_mass, best_minl, best_nmin, sel_fidx, best_sel). `best*` tracks the

983 + tightest fallback: the largest evaluated pivot whose above-mass is >= p,

984 + which guarantees the top-p mass invariant if the search fails to

985 + converge. Pivots are recomputed from (lo, hi) via the same ladder the

986 + sweep used, so partials line up with pivots without storing them.

987 + `sel_fidx` / `best_sel` report which ladder slot produced `dup` /

988 + updated `best` this round (-1 if none), so callers can recover the

989 + per-slice (minl, nmin) partials behind the boundary value for

990 + deterministic duplicate handling.

991 + """

992 + fidx = tl.arange(0, F)

993 + sidx = tl.arange(0, S)

994 + offs2 = sidx[:, None] * (F * 3) + fidx[None, :] * 3

995 + mass_s = tl.load(PART_BASE + offs2 + 0)

996 + minl_s = tl.load(PART_BASE + offs2 + 1)

997 + nmin_s = tl.load(PART_BASE + offs2 + 2)

998 + mass = tl.sum(mass_s, axis=0)

999 + gmin = tl.min(minl_s, axis=0)

1000 + nmin = tl.sum(tl.where(tl.abs(minl_s - gmin[None, :]) < 1e-9, nmin_s, 0.0), axis=0)

1001 + pivs = _topp_sb_ladder(lo, hi, F)

1002 +

1003 + ok = (mass >= p) & (mass - gmin * nmin < p)

1004 + ok_idx = tl.max(tl.where(ok, fidx, -1).to(tl.float32))

1005 + ge = mass >= p

1006 + ge_idx = tl.max(tl.where(ge, fidx, -1).to(tl.float32))

1007 +

1008 + best_sel = -1.0

1009 + if ge_idx >= 0:

1010 + bmask = fidx == ge_idx.to(tl.int32)

1011 + bp = tl.sum(tl.where(bmask, pivs, 0.0))

1012 + if bp > best:

1013 + best = bp

1014 + best_mass = tl.sum(tl.where(bmask, mass, 0.0))

1015 + best_minl = tl.sum(tl.where(bmask, gmin, 0.0))

1016 + best_nmin = tl.sum(tl.where(bmask, nmin, 0.0))

1017 + best_sel = ge_idx

1018 +

1019 + done = 0.0

1020 + nomask = 0.0

1021 + pivot = 0.0

1022 + sel_mass = 0.0

1023 + dup = 1.0

1024 + numdup = 1.0

1025 + numkeep = 1.0

1026 + sel_fidx = -1.0

1027 + if ok_idx >= 0:

1028 + bmask = fidx == ok_idx.to(tl.int32)

1029 + pivot = tl.sum(tl.where(bmask, pivs, 0.0))

1030 + sel_mass = tl.sum(tl.where(bmask, mass, 0.0))

1031 + dup = tl.sum(tl.where(bmask, gmin, 0.0))

1032 + numdup = tl.sum(tl.where(bmask, nmin, 0.0))

1033 + sel_fidx = ok_idx

1034 + done = 1.0

1035 + else:

1036 + new_lo = tl.maximum(lo, tl.max(tl.where(ge, pivs, -float("inf"))))

1037 + new_hi = tl.minimum(hi, tl.min(tl.where(~ge, pivs, float("inf"))))

1038 + if IS_LAST or not (new_hi > new_lo * (1.0 + 1e-7)):

1039 + # Out of rounds or the range collapsed: fall back to the

1040 + # tightest pivot with mass >= p, or keep the whole row if no

1041 + # pivot qualified.

1042 + if best > 0.0:

1043 + pivot = best

1044 + sel_mass = best_mass

1045 + dup = best_minl

1046 + numdup = best_nmin

1047 + else:

1048 + nomask = 1.0

1049 + done = 1.0

1050 + lo = new_lo

1051 + hi = new_hi

1052 + if (done != 0.0) and (nomask == 0.0):

1053 + # Mirror the monolithic guard: a pivot at/above the max logit would

1054 + # mask everything under the strict `>` comparison.

1055 + nomask = tl.where(pivot * Z >= 1.0, 1.0, 0.0)

1056 + # Duplicate (boundary-value) handling, same formula as the

1057 + # monolithic kernel: keep only some of the boundary duplicates.

1058 + numkeep = numdup - tl.cast(tl.maximum(sel_mass - p, 0.0) / dup, tl.uint32).to(

1059 + tl.float32

1060 + )

1061 + numkeep = tl.minimum(tl.maximum(numkeep, 1.0), numdup)

1062 + return (

1063 + done,

1064 + nomask,

1065 + pivot,

1066 + dup,

1067 + numdup,

1068 + numkeep,

1069 + lo,

1070 + hi,

1071 + best,

1072 + best_mass,

1073 + best_minl,

1074 + best_nmin,

1075 + sel_fidx,

1076 + best_sel,

1077 + )

1078 +

1079 +

1080 +@triton.jit

1081 +def _topp_sb_step_kernel(

1082 + LOGITS,

1083 + LOGITS_STRIDE_0,

1084 + STATS,

1085 + PARTS,

1086 + K,

1087 + P,

1088 + ROUND,

1089 + HAS_K: tl.constexpr,

1090 + S: tl.constexpr,

1091 + F: tl.constexpr,

1092 + NUM_ROUNDS: tl.constexpr,

1093 + VOCAB_SIZE: tl.constexpr,

1094 + BLOCK: tl.constexpr,

1095 +):

1096 + pid = tl.program_id(0)

1097 + row = pid // S

1098 + slice_id = pid % S

1099 + p = tl.load(P + row)

1100 + if p >= 1.0:

1101 + return

1102 + if HAS_K and tl.load(K + row) < VOCAB_SIZE:

1103 + return

1104 + M, Z, mn, num_finite = _topp_sb_row_stats(STATS, row, S)

1105 + if (Z <= 0.0) or (num_finite < 1.0):

1106 + return

1107 + max_prob = 1.0 / Z

1108 + min_prob = tl.maximum(tl.exp(mn - M) / Z, 1e-30)

1109 + if max_prob <= min_prob * (1.0 + 1e-6):

1110 + # (Nearly) uniform row: no pivot can differentiate; keep everything.

1111 + return

1112 + # Replay the combine chain for rounds 0..ROUND-1 (deterministic; each

1113 + # is a tiny pure function of the stored partials). Partials from rounds

1114 + # past a row's convergence are stale, so the chain stops at `done`.

1115 + lo = min_prob

1116 + hi = max_prob

1117 + best = -1.0

1118 + best_mass = 0.0

1119 + best_minl = 1.0

1120 + best_nmin = 0.0

1121 + done = 0.0

1122 + for j in range(ROUND):

1123 + if done == 0.0:

1124 + (

1125 + done,

1126 + nomask,

1127 + pivot,

1128 + dup,

1129 + numdup,

1130 + numkeep,

1131 + lo,

1132 + hi,

1133 + best,

1134 + best_mass,

1135 + best_minl,

1136 + best_nmin,

1137 + sel_fidx,

1138 + best_sel,

1139 + ) = _topp_sb_combine(

1140 + PARTS + (row * NUM_ROUNDS + j) * S * F * 3,

1141 + Z,

1142 + p,

1143 + lo,

1144 + hi,

1145 + best,

1146 + best_mass,

1147 + best_minl,

1148 + best_nmin,

1149 + S=S,

1150 + F=F,

1151 + IS_LAST=False,

1152 + )

1153 + if done != 0.0:

1154 + return

1155 + # Sweep this slice against the ladder pivots derived from (lo, hi).

1156 + pivs = _topp_sb_ladder(lo, hi, F)

1157 + slice_len = (VOCAB_SIZE + S - 1) // S

1158 + start = slice_id * slice_len

1159 + end = tl.minimum(start + slice_len, VOCAB_SIZE)

1160 + ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0

1161 + offs = tl.arange(0, BLOCK)

1162 + fidx = tl.arange(0, F)

1163 + mass = tl.zeros([F], dtype=tl.float32)

1164 + minl = tl.full([F], 1.0, dtype=tl.float32)

1165 + nmin = tl.zeros([F], dtype=tl.float32)

1166 + for i in range(start, end, BLOCK):

1167 + x = tl.load(ROW + i + offs, mask=i + offs < end, other=-float("inf"))

1168 + pr = tl.exp(x - M) / Z

1169 + pr2 = tl.broadcast_to(pr[None, :], (F, BLOCK))

1170 + above = pr2 > pivs[:, None]

1171 + mass += tl.sum(tl.where(above, pr2, 0.0), axis=1)

1172 + cand = tl.where(above, pr2, float("inf"))

1173 + tmin = tl.min(cand, axis=1)

1174 + teq = above & (tl.abs(pr2 - tmin[:, None]) < 1e-9)

1175 + tcnt = tl.sum(teq.to(tl.float32), axis=1)

1176 + is_new = tmin < minl

1177 + is_same = tl.abs(tmin - minl) < 1e-9

1178 + nmin = tl.where(is_new, tcnt, nmin + tl.where(is_same, tcnt, 0.0))

1179 + minl = tl.minimum(minl, tmin)

1180 + pbase = PARTS + ((row * NUM_ROUNDS + ROUND) * S + slice_id) * F * 3

1181 + tl.store(pbase + fidx * 3 + 0, mass)

1182 + tl.store(pbase + fidx * 3 + 1, minl)

1183 + tl.store(pbase + fidx * 3 + 2, nmin)

1184 +

1185 +

1186 +@triton.jit

1187 +def _topp_sb_mask_kernel(

1188 + LOGITS,

1189 + LOGITS_STRIDE_0,

1190 + STATS,

1191 + PARTS,

1192 + K,

1193 + P,

1194 + HAS_K: tl.constexpr,

1195 + MASK_VALUE: tl.constexpr,

1196 + S: tl.constexpr,

1197 + F: tl.constexpr,

1198 + NUM_ROUNDS: tl.constexpr,

1199 + VOCAB_SIZE: tl.constexpr,

1200 + BLOCK: tl.constexpr,

1201 +):

1202 + pid = tl.program_id(0)

1203 + row = pid // S

1204 + slice_id = pid % S

1205 + p = tl.load(P + row)

1206 + if p >= 1.0:

1207 + return

1208 + if HAS_K and tl.load(K + row) < VOCAB_SIZE:

1209 + return

1210 + M, Z, mn, num_finite = _topp_sb_row_stats(STATS, row, S)

1211 + if (Z <= 0.0) or (num_finite < 1.0):

1212 + return

1213 + max_prob = 1.0 / Z

1214 + min_prob = tl.maximum(tl.exp(mn - M) / Z, 1e-30)

1215 + if max_prob <= min_prob * (1.0 + 1e-6):

1216 + return # (nearly) uniform row: keep everything

1217 + # Replay the full combine chain; the last round finalizes with fallback

1218 + # semantics. Rows that converged early stop the chain at that round.

1219 + # `fin` tracks the (round, ladder slot) whose per-slice partials define

1220 + # the boundary value, for deterministic duplicate handling below.

1221 + lo = min_prob

1222 + hi = max_prob

1223 + best = -1.0

1224 + best_mass = 0.0

1225 + best_minl = 1.0

1226 + best_nmin = 0.0

1227 + done = 0.0

1228 + nomask = 0.0

1229 + pivot = 0.0

1230 + dup = 1.0

1231 + numdup = 1.0

1232 + numkeep = 1.0

1233 + fin = -1.0

1234 + best_fin = -1.0

1235 + for j in range(NUM_ROUNDS):

1236 + if done == 0.0:

1237 + (

1238 + done,

1239 + nomask,

1240 + pivot,

1241 + dup,

1242 + numdup,

1243 + numkeep,

1244 + lo,

1245 + hi,

1246 + best,

1247 + best_mass,

1248 + best_minl,

1249 + best_nmin,

1250 + sel_fidx,

1251 + best_sel,

1252 + ) = _topp_sb_combine(

1253 + PARTS + (row * NUM_ROUNDS + j) * S * F * 3,

1254 + Z,

1255 + p,

1256 + lo,

1257 + hi,

1258 + best,

1259 + best_mass,

1260 + best_minl,

1261 + best_nmin,

1262 + S=S,

1263 + F=F,

1264 + IS_LAST=j == NUM_ROUNDS - 1,

1265 + )

1266 + if sel_fidx >= 0.0:

1267 + fin = j * F + sel_fidx

1268 + if best_sel >= 0.0:

1269 + best_fin = j * F + best_sel

1270 + if nomask != 0.0:

1271 + return

1272 + if fin < 0.0:

1273 + fin = best_fin

1274 + logZ = tl.log(Z)

1275 + pivot_logit = tl.log(pivot) + logZ + M

1276 + # numdup/numkeep are exact small integers held in fp32.

1277 + ties = numkeep < numdup

1278 + num_preceding = 0.0

1279 + if ties:

1280 + # Deterministic, index-ordered duplicate budget: per-slice counts of

1281 + # the boundary value come from the finalizing round's partials (a

1282 + # slice holds copies of `dup` iff its minl matches `dup`, with the

1283 + # count in nmin), so each slice knows exactly how many copies were

1284 + # kept by lower-indexed slices.

1285 + fin_round = tl.cast(fin / F, tl.int32)

1286 + fin_slot = tl.cast(fin - tl.cast(fin_round, tl.float32) * F, tl.int32)

1287 + pbase = PARTS + ((row * NUM_ROUNDS + fin_round) * S * F + fin_slot) * 3

1288 + sidx = tl.arange(0, S)

1289 + minl_s = tl.load(pbase + sidx * (F * 3) + 1)

1290 + nmin_s = tl.load(pbase + sidx * (F * 3) + 2)

1291 + dupc_s = tl.where(tl.abs(minl_s - dup) < 1e-9, nmin_s, 0.0)

1292 + num_preceding = tl.sum(tl.where(sidx < slice_id, dupc_s, 0.0))

1293 + slice_len = (VOCAB_SIZE + S - 1) // S

1294 + start = slice_id * slice_len

1295 + end = tl.minimum(start + slice_len, VOCAB_SIZE)

1296 + ROW = LOGITS + row.to(tl.int64) * LOGITS_STRIDE_0

1297 + offs = tl.arange(0, BLOCK)

1298 + for i in range(start, end, BLOCK):

1299 + mask_n = i + offs < end

1300 + x = tl.load(ROW + i + offs, mask=mask_n, other=-float("inf"))

1301 + keep = x > pivot_logit

1302 + if ties:

1303 + # Match the boundary value in probability space with exact

1304 + # equality: `dup` is a probability computed by the step kernels

1305 + # as exp(x - M) / Z, so recomputing it here with the same ops

1306 + # reproduces it bitwise on any platform. (The log/exp round-trip

1307 + # to logit space loses ~1e-6 to cancellation, breaking any

1308 + # absolute tolerance.)

1309 + dmask = tl.exp(x - M) / Z == dup

1310 + cum = tl.cumsum(dmask.to(tl.int32)).to(tl.float32) + num_preceding

1311 + keep_dup = dmask & (cum <= numkeep)

1312 + keep = keep & (~dmask | keep_dup)

1313 + tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n)

1314 +

1315 +

1316 +def _apply_topp_split(

1317 + logits: torch.Tensor,

1318 + k: torch.Tensor | None,

1319 + p: torch.Tensor,

1320 + mask_value: float,

1321 + num_sm: int,

1322 +) -> None:

1323 + """Split-row top-p pipeline for small batches; masks logits in place."""

1324 + batch_size, vocab_size = logits.shape

1325 + splits = 1

1326 + while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:

1327 + splits *= 2

1328 +

1329 + ws = _TRITON_SPLIT_CACHE.get(logits.device)

1330 + if ws is None:

1331 + ws = {

1332 + "stats": logits.new_empty((_SPLIT_MAX_BATCH * _SPLIT_MAX_SPLITS, 4)),

1333 + "parts": logits.new_empty(

1334 + (_SPLIT_MAX_BATCH, _SPLIT_ROUNDS, _SPLIT_MAX_SPLITS, _SPLIT_FANOUT, 3)

1335 + ),

1336 + }

1337 + _TRITON_SPLIT_CACHE[logits.device] = ws

1338 +

1339 + k_ptr = k if k is not None else logits # dummy pointer when HAS_K=False

1340 + has_k = k is not None

1341 +

1342 + _topp_sb_stats_kernel[(batch_size * splits,)](

1343 + logits,

1344 + logits.stride(0),

1345 + ws["stats"],

1346 + k_ptr,

1347 + p,

1348 + HAS_K=has_k,

1349 + VOCAB_SIZE=vocab_size,

1350 + S=splits,

1351 + BLOCK=8192,

1352 + num_warps=8,

1353 + )

1354 + for round_i in range(_SPLIT_ROUNDS):

1355 + _topp_sb_step_kernel[(batch_size * splits,)](

1356 + logits,

1357 + logits.stride(0),

1358 + ws["stats"],

1359 + ws["parts"],

1360 + k_ptr,

1361 + p,

1362 + round_i,

1363 + HAS_K=has_k,

1364 + S=splits,

1365 + F=_SPLIT_FANOUT,

1366 + NUM_ROUNDS=_SPLIT_ROUNDS,

1367 + VOCAB_SIZE=vocab_size,

1368 + BLOCK=2048,

1369 + num_warps=8,

1370 + )

1371 + _topp_sb_mask_kernel[(batch_size * splits,)](

1372 + logits,

1373 + logits.stride(0),

1374 + ws["stats"],

1375 + ws["parts"],

1376 + k_ptr,

1377 + p,

1378 + HAS_K=has_k,

1379 + MASK_VALUE=mask_value,

1380 + S=splits,

1381 + F=_SPLIT_FANOUT,

1382 + NUM_ROUNDS=_SPLIT_ROUNDS,

1383 + VOCAB_SIZE=vocab_size,

1384 + BLOCK=8192,

1385 + num_warps=8,

1386 + )

1387 +

1388 +

1389 def apply_top_k_top_p_triton(

1390 logits: torch.Tensor,

1391 k: torch.Tensor | None,

@@ -905,6 +1437,21 @@ def apply_top_k_top_p_triton(

1437 p_ptr = logits # Dummy pointer (won't be read)

1438

1439 num_sm = num_compute_units(logits.device.index)

1440 +

1441 + # At small batch sizes the monolithic kernel's standalone top-p path is

1442 + # latency-bound (one program per row, serial vocab sweeps), so use the

1443 + # split-row pipeline for rows without an active top-k. Rows with an

1444 + # active top-k still go through the monolithic kernel (its truncation

1445 + # path is already fast). The monolithic's standalone top-p branch then

1446 + # covers only rows whose top-k is a no-op (k < vocab but <= k finite

1447 + # logits, e.g. grammar masks), which the split pipeline skips.

1448 + use_split = (

1449 + topp_enabled and batch_size <= _SPLIT_MAX_BATCH and logits.device.type == "cuda"

1450 + )

1451 + if use_split and not topk_enabled:

1452 + _apply_topp_split(logits, None, p_ptr, mask_value, num_sm)

1453 + return logits

1454 +

1455 NUM_PROGRAMS = min(num_sm, batch_size)

1456

1457 # Cache per-Triton Program buffer on each device.

@@ -961,8 +1508,11 @@ def apply_top_k_top_p_triton(

1508 BLOCK_SIZE_TRUNC=block_size_trunc,

1509 TOPK_ENABLED=topk_enabled,

1510 TOPP_ENABLED=topp_enabled,

1511 + SPLIT_COVERS_PONLY=use_split,

1512 **launch_kwargs,

1513 )

1514 + if use_split:

1515 + _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm)

1516

1517 return logits

1518

@@ -970,4 +1520,5 @@ def apply_top_k_top_p_triton(

1520 def reset_buffer_cache():

1521 _TRITON_BUFFER_CACHE.clear()

1522 _TRITON_TABLE_CACHE.clear()

1523 + _TRITON_SPLIT_CACHE.clear()

1524 torch.accelerator.empty_cache()