From 9234e40aed293c80134d8d0cad34729db2689396 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Wed, 19 Aug 2026 17:48:03 -0700 Subject: [PATCH] [sampling] Fix int32 offset overflow in top-k renorm Triton kernels (#35571) Co-authored-by: Xiaozhu Meng --- python/sglang/kernels/ops/sampling/renorm_triton.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/sampling/renorm_triton.py b/python/sglang/kernels/ops/sampling/renorm_triton.py index a50956688..38714fc4e 100644 --- a/python/sglang/kernels/ops/sampling/renorm_triton.py +++ b/python/sglang/kernels/ops/sampling/renorm_triton.py @@ -25,7 +25,7 @@ def _mask_and_partial_sum_kernel( chunk = tl.program_id(1) offsets = chunk * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < vocab_size - row_offsets = row * vocab_size + offsets + row_offsets = row.to(tl.int64) * vocab_size + offsets probs = tl.load(probs_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) pivot = tl.load(pivots_ptr + row) @@ -43,7 +43,7 @@ def _normalize_kernel( vocab_size: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < numel row = offsets // vocab_size values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32)