A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 23:15 UTC
[Bugfix] Fix native Triton top-k/top-p kernel assumes contiguous logi… (#42739)
Signed-off-by: xiaogang.zhou <xiaogang.zhou@bytedance.com> Co-authored-by: xiaogang.zhou <xiaogang.zhou@bytedance.com>
Xiaogang Zhou committed
May 23, 2026 at 13:56 UTC
d19db10974587335ec3a37e0a424abb57430574e
2 files changed
+62
-5
tests/v1/sample/test_topk_topp_sampler.py
+50
@@ -320,6 +320,56 @@ class TestTritonTopkTopp:
320
321
self._compare_results(logits, k, p)
322
323
+ @pytest.mark.parametrize(
324
+ "mode",
325
+ ["topk_only", "topp_only", "topk_and_topp"],
326
+ )
327
+ def test_noncontiguous_logits_match_contiguous(self, mode: str):
328
+ """Non-contiguous logits views should behave like contiguous inputs."""
329
+ from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
330
+
331
+ device = torch.device(DEVICE_TYPE)
332
+ batch_size, vocab_size, pad = 16, 4096, 8
333
+ backing = torch.full(
334
+ (batch_size, vocab_size + pad),
335
+ -1000.0,
336
+ device=device,
337
+ dtype=torch.float32,
338
+ )
339
+ base = torch.linspace(
340
+ 10.0, -10.0, vocab_size, device=device, dtype=torch.float32
341
+ )
342
+ source = base[None, :] + (
343
+ torch.arange(batch_size, device=device, dtype=torch.float32)[:, None]
344
+ / 1000.0
345
+ )
346
+
347
+ logits = backing[:, :vocab_size]
348
+ logits.copy_(source)
349
+ contig_logits = source.clone()
350
+ pytorch_logits = source.clone()
351
+
352
+ assert logits.shape == (batch_size, vocab_size)
353
+ assert logits.stride() == (vocab_size + pad, 1)
354
+ assert not logits.is_contiguous()
355
+
356
+ k: torch.Tensor | None = None
357
+ p: torch.Tensor | None = None
358
+ if mode in ("topk_only", "topk_and_topp"):
359
+ k = torch.full((batch_size,), 154, device=device, dtype=torch.int32)
360
+ if mode in ("topp_only", "topk_and_topp"):
361
+ p = torch.full((batch_size,), 0.95, device=device, dtype=torch.float32)
362
+
363
+ noncontig_out = apply_top_k_top_p_triton(logits, k, p)
364
+ contig_out = apply_top_k_top_p_triton(contig_logits, k, p)
365
+ pytorch_out = apply_top_k_top_p_pytorch(pytorch_logits, k, p)
366
+
367
+ assert noncontig_out.data_ptr() == logits.data_ptr()
368
+ assert not noncontig_out.is_contiguous()
369
+ assert torch.equal(logits, noncontig_out)
370
+ assert torch.equal(torch.isfinite(noncontig_out), torch.isfinite(contig_out))
371
+ assert torch.equal(torch.isfinite(noncontig_out), torch.isfinite(pytorch_out))
372
+
373
# -----------------------------------------------------------------
374
# Tests for -inf logits (e.g. from grammar / structured output masks)
375
# -----------------------------------------------------------------
vllm/v1/sample/ops/topk_topp_triton.py
+12
-5
@@ -93,6 +93,7 @@ def _update_min_larger_stats(data, above_mask, min_larger, num_min_larger, senti
93
@triton.jit
94
def _topk_topp_kernel(
95
LOGITS,
96
+ LOGITS_STRIDE_0,
97
BUFFER,
98
PERCENTILE_TO_STD_TABLE,
99
NORMAL_CDF_TO_SIGMA_TABLE,
@@ -110,7 +111,7 @@ def _topk_topp_kernel(
111
pid = tl.program_id(0)
112
num_programs = tl.num_programs(0)
113
for row_id in tl.range(pid, BATCH_SIZE, num_programs):
113
- LOGITS_ROW = LOGITS + row_id * VOCAB_SIZE
114
+ LOGITS_ROW = LOGITS + row_id * LOGITS_STRIDE_0
115
BUFFER_ROW = BUFFER + pid * VOCAB_SIZE
116
117
final_pivot = -float("inf")
@@ -975,26 +976,31 @@ def apply_top_k_top_p_triton(
976
to the remaining k values (by probability).
977
978
Args:
978
- logits: [batch_size, vocab_size] float32 tensor, modified in-place
979
+ logits: [batch_size, vocab_size] float32 tensor. The returned tensor
980
+ may alias this input or be a new contiguous tensor for unsupported
981
+ layouts.
982
k: [batch_size] int32 tensor of top-k values per row, or None to disable top-k
983
p: [batch_size] float32 tensor of top-p values per row (0 to 1),
984
or None to disable top-p
985
mask_value: Value for masked positions (default: -inf)
986
987
Returns:
985
- The logits tensor (modified in-place)
988
+ The masked logits tensor. It may or may not be modified in-place.
989
"""
990
assert logits.ndim == 2
991
assert logits.dtype == torch.float32
989
-
992
batch_size, vocab_size = logits.shape
991
-
993
topk_enabled = k is not None
994
topp_enabled = p is not None
995
996
if batch_size == 0 or not (topk_enabled or topp_enabled):
997
return logits
998
999
+ # The Triton kernel supports arbitrary row strides, but it still assumes
1000
+ # the vocab dimension is laid out contiguously within each row.
1001
+ if logits.stride(1) != 1:
1002
+ logits = logits.contiguous()
1003
+
1004
if k is not None:
1005
assert k.ndim == 1 and k.shape[0] == batch_size
1006
k_ptr = k.to(torch.int32)
@@ -1034,6 +1040,7 @@ def apply_top_k_top_p_triton(
1040
1041
_topk_topp_kernel[(NUM_PROGRAMS,)](
1042
logits,
1043
+ logits.stride(0),
1044
buffer,
1045
percentile_to_std_table,
1046
normal_cdf_to_sigma_table,