From e39f0f4ff3a92a6c5941a396373be2c888fa90c6 Mon Sep 17 00:00:00 2001 From: cctry Date: Tue, 21 Apr 2026 22:54:37 -0700 Subject: [PATCH] Use libdevice tanh and support 2D-strided tensors in fused softcap kernel (#23157) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- python/sglang/srt/layers/logits_processor.py | 35 +++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 7151ef1f3..0f13dc9eb 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -21,6 +21,7 @@ import torch import triton import triton.language as tl from torch import nn +from triton.language.extra import libdevice from sglang.srt.distributed import ( get_tensor_model_parallel_world_size, @@ -1109,39 +1110,49 @@ class LogitsProcessor(nn.Module): def fused_softcap_kernel( full_logits_ptr, softcapping_value, - n_elements, + ncols, + row_stride, BLOCK_SIZE: tl.constexpr, ): + row = tl.program_id(1).to(tl.int64) pid = tl.program_id(0).to(tl.int64) block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements + mask = offsets < ncols # Load values - x = tl.load(full_logits_ptr + offsets, mask=mask) + row_ptr = full_logits_ptr + row * row_stride + x = tl.load(row_ptr + offsets, mask=mask) # Perform operations in-place x = x / softcapping_value - - # Manual tanh implementation using exp - exp2x = tl.exp(2 * x) - x = (exp2x - 1) / (exp2x + 1) - + x = libdevice.tanh(x) x = x * softcapping_value # Store result - tl.store(full_logits_ptr + offsets, x, mask=mask) + tl.store(row_ptr + offsets, x, mask=mask) def fused_softcap(full_logits, final_logit_softcapping): - n_elements = full_logits.numel() + if full_logits.is_contiguous(): + nrows, ncols = 1, full_logits.numel() + row_stride = ncols + else: + assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor" + assert ( + full_logits.stride(1) == 1 + ), "non-contiguous softcap requires contiguous columns" + nrows, ncols = full_logits.shape + row_stride = full_logits.stride(0) + BLOCK_SIZE = 1024 - grid = ((n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE, 1, 1) + grid = ((ncols + BLOCK_SIZE - 1) // BLOCK_SIZE, nrows) fused_softcap_kernel[grid]( full_logits_ptr=full_logits, softcapping_value=final_logit_softcapping, - n_elements=n_elements, + ncols=ncols, + row_stride=row_stride, BLOCK_SIZE=BLOCK_SIZE, ) return full_logits