chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Migrate renorm kernels from sgl-kernel to FlashInfer JIT (#18854)

Browse Source
This commit is contained in:
Johnsonms
2026-03-06 06:53:28 -08:00
committed by GitHub
parent 6d22c9f369
commit 2d266c73ea
3 changed files with 458 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
sgl-kernel/python/sgl_kernel/sampling.py
View File

@@ -3,6 +3,13 @@ from typing import Optional, Union
import torch
from sgl_kernel.utils import _to_tensor_scalar_tuple
try:
import flashinfer.sampling as _flashinfer_sampling
_has_flashinfer = True
except ImportError:
_has_flashinfer = False
def _top_k_renorm_probs_internal(
probs: torch.Tensor,
@@ -46,7 +53,10 @@ def top_k_renorm_probs(
This combination of ``top_k_renorm_probs`` and ``sampling_from_probs`` should be equivalent to
``top_k_sampling_from_probs``.
"""
return _top_k_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_k))
if probs.device.type == "musa" or not _has_flashinfer:
return _top_k_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_k))
else:
return _flashinfer_sampling.top_k_renorm_probs(probs, top_k)
top_k_renorm_prob = top_k_renorm_probs
@@ -96,7 +106,10 @@ def top_p_renorm_probs(
``top_p_sampling_from_probs``.
"""
return _top_p_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_p))
if probs.device.type == "musa" or not _has_flashinfer:
return _top_p_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_p))
else:
return _flashinfer_sampling.top_p_renorm_probs(probs, top_p)
top_p_renorm_prob = top_p_renorm_probs
@@ -169,4 +182,7 @@ def top_k_mask_logits(
--------
top_k_renorm_probs
"""
return _top_k_mask_logits_internal(logits, *_to_tensor_scalar_tuple(top_k))
if logits.device.type == "musa" or not _has_flashinfer:
return _top_k_mask_logits_internal(logits, *_to_tensor_scalar_tuple(top_k))
else:
return _flashinfer_sampling.top_k_mask_logits(logits, top_k)