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

[kernel slimming] Clean many useless sgl-kernel deprecated kernels (#20277)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-14 16:45:54 +08:00
committed by GitHub
parent 75a7879fd4
commit 25e38216b6
26 changed files with 60 additions and 1483 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
sgl-kernel/python/sgl_kernel/__init__.py
View File

@@ -18,8 +18,6 @@ from sgl_kernel.attention import (
)
from sgl_kernel.cutlass_moe import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data
from sgl_kernel.elementwise import (
FusedSetKVBufferArg,
apply_rope_with_cos_sin_cache_inplace,
concat_mla_absorb_q,
concat_mla_k,
copy_to_gpu_no_ce,
@@ -41,7 +39,6 @@ from sgl_kernel.expert_specialization import (
from sgl_kernel.gemm import (
awq_dequantize,
bmm_fp8,
cutlass_scaled_fp4_mm,
dsv3_fused_a_gemm,
dsv3_router_gemm,
fp8_blockwise_scaled_mm,
@@ -51,15 +48,11 @@ from sgl_kernel.gemm import (
int8_scaled_mm,
qserve_w4a8_per_chn_gemm,
qserve_w4a8_per_group_gemm,
scaled_fp4_grouped_quant,
scaled_fp4_quant,
sgl_per_tensor_quant_fp8,
sgl_per_token_group_quant_8bit,
sgl_per_token_group_quant_fp8,
sgl_per_token_group_quant_int8,
sgl_per_token_quant_fp8,
shuffle_rows,
silu_and_mul_scaled_fp4_grouped_quant,
)
from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
from sgl_kernel.kvcacheio import (
@@ -75,7 +68,7 @@ from sgl_kernel.mamba import (
causal_conv1d_update_cpu,
chunk_gated_delta_rule_cpu,
)
from sgl_kernel.memory import set_kv_buffer_kernel, weak_ref_tensor
from sgl_kernel.memory import weak_ref_tensor
from sgl_kernel.moe import (
apply_shuffle_mul_sum,
fp8_blockwise_scaled_grouped_mm,

116
sgl-kernel/python/sgl_kernel/elementwise.py
View File

@@ -1,4 +1,3 @@
from dataclasses import dataclass
from typing import Optional
import torch
@@ -332,121 +331,6 @@ if torch.version.hip is not None:
return out
@dataclass
class FusedSetKVBufferArg:
"""
value : Optional[torch.Tensor]
Value tensor, shape: ``(nnz, num_v_heads * head_size)``.
k_buffer : Optional[torch.Tensor]
Buffer for keys, shape: ``(nnz, num_k_heads * head_size)``.
v_buffer : Optional[torch.Tensor]
Buffer for values, shape: ``(nnz, num_v_heads * head_size)``.
k_scale : Optional[float]
Scale factor for keys.
v_scale : Optional[float]
Scale factor for values.
cache_loc : Optional[torch.Tensor]
Cache location tensor, used for indexing kv cache.
"""
value: torch.Tensor
k_buffer: torch.Tensor
v_buffer: torch.Tensor
k_scale: Optional[float]
v_scale: Optional[float]
cache_loc: torch.Tensor
def _view_3d(x, head_size):
return x.view(x.shape[0], -1, head_size)
def apply_rope_with_cos_sin_cache_inplace(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox: bool = True,
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
enable_pdl: Optional[bool] = None,
) -> None:
r"""
Apply rotary embedding to keys and queries with precomputed cos/sin values.
This is designed to be compatible with the SGL/vLLM implementation.
The result is inplace applied to the input tensors.
Parameters
----------
positions : torch.Tensor
Position indices, shape: ``(nnz)``.
query : torch.Tensor
Query tensor, shape: ``(nnz, num_q_heads * head_size)``.
key : torch.Tensor
Key tensor, shape: ``(nnz, num_k_heads * head_size)``.
cos_sin_cache : torch.Tensor
Cosine and Sine cache tensor, shape: ``(max_seq_len, rotary_dim)``.
Cosine is the first half and Sine is the second half on rotary_dim.
is_neox : bool
Whether to use Neox style RoPE, default: ``True``.
* If ``True``, the last dimension of the query/key tensor is not interleaved, i.e.,
we rotate the first half dimensions ``([..., :head_dim//2])`` and the second half
dimensions ``([..., head_dim//2:])``.
* If ``False``, the last dimension of the query/key tensor is interleaved, i.e.,
we rotate the even dimensions ``([..., ::2])`` and odd dimensions ``([..., 1::2])``.
fused_set_kv_buffer_arg : FusedSetKVBufferArg
Fuse the set-kv-buffer operation into this kernel
Note
----
The rotary dimension is determined by the cosine cache and sine cache.
"""
if cos_sin_cache.dtype != torch.float32:
raise ValueError("cos_sin_cache should be float32")
if enable_pdl is None:
# the non-fused branch does not yet support PDL, but after we switch to our impl for that branch it will
enable_pdl = is_arch_support_pdl() and (fused_set_kv_buffer_arg is not None)
if (a := fused_set_kv_buffer_arg) is not None:
assert a.k_scale is None, "k_scale is not yet supported"
assert a.v_scale is None, "v_scale is not yet supported"
assert a.cache_loc.dtype == torch.int64, f"{a.cache_loc.dtype=}"
torch.ops.sgl_kernel.apply_rope_pos_ids_cos_sin_cache.default(
_view_3d(query, head_size),
_view_3d(key, head_size),
_view_3d(query, head_size),
_view_3d(key, head_size),
cos_sin_cache,
positions.long(),
(not is_neox),
enable_pdl,
(
_view_3d(fused_set_kv_buffer_arg.value, head_size)
if fused_set_kv_buffer_arg is not None
else None
),
(
_view_3d(fused_set_kv_buffer_arg.k_buffer, head_size)
if fused_set_kv_buffer_arg is not None
else None
),
(
_view_3d(fused_set_kv_buffer_arg.v_buffer, head_size)
if fused_set_kv_buffer_arg is not None
else None
),
(
fused_set_kv_buffer_arg.cache_loc
if fused_set_kv_buffer_arg is not None
else None
),
)
def rotary_embedding(
positions: torch.Tensor,
query: torch.Tensor,

128
sgl-kernel/python/sgl_kernel/gemm.py
View File

@@ -1,4 +1,4 @@
from typing import Optional, Tuple
from typing import Optional
import torch
from sgl_kernel.utils import _get_cache_buf
@@ -140,17 +140,6 @@ sgl_per_token_group_quant_fp8 = sgl_per_token_group_quant_8bit
sgl_per_token_group_quant_int8 = sgl_per_token_group_quant_8bit
def sgl_per_tensor_quant_fp8(
input: torch.Tensor,
output_q: torch.Tensor,
output_s: torch.Tensor,
is_static: bool,
) -> None:
torch.ops.sgl_kernel.sgl_per_tensor_quant_fp8.default(
input, output_q, output_s, is_static
)
def sgl_per_token_quant_fp8(
input: torch.Tensor,
output_q: torch.Tensor,
@@ -159,54 +148,6 @@ def sgl_per_token_quant_fp8(
torch.ops.sgl_kernel.sgl_per_token_quant_fp8.default(input, output_q, output_s)
def cutlass_scaled_fp4_mm(
a: torch.Tensor,
b: torch.Tensor,
block_scale_a: torch.Tensor,
block_scale_b: torch.Tensor,
alpha: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
from sglang.jit_kernel.nvfp4 import (
cutlass_scaled_fp4_mm as jit_cutlass_scaled_fp4_mm,
)
return jit_cutlass_scaled_fp4_mm(
a,
b,
block_scale_a,
block_scale_b,
alpha,
out_dtype,
)
def scaled_fp4_quant(
input: torch.Tensor, input_global_scale: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantize input tensor to FP4 and return quantized tensor and scale.
This function quantizes the last dimension of the given tensor `input`. For
every 16 consecutive elements, a single dynamically computed scaling factor
is shared. This scaling factor is quantized using the `input_global_scale`
and is stored in a swizzled layout (see
https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x).
Args:
input: The input tensor to be quantized to FP4
input_global_scale: A scalar scaling factor for the entire tensor.
Returns:
Tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every
two values are packed into a uint8 and float8_e4m3 scaling factors
in a sizzled layout.
"""
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as jit_scaled_fp4_quant
return jit_scaled_fp4_quant(input, input_global_scale)
def qserve_w4a8_per_chn_gemm(
in_feats: torch.Tensor,
kernel: torch.Tensor,
@@ -280,73 +221,6 @@ def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape):
return output_tensor
def scaled_fp4_grouped_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
mask: torch.Tensor,
):
"""
Quantize input tensor to FP4 and return quantized tensor and scale, for
grouped gemm inputs (e.g., grouped_gemm_nt_masked for flashinfer).
Args:
input: The input tensor to be quantized to FP4, with shape (l, m, k)
l is number of groups, m is number of tokens per group, k is number of features.
input_global_scale: A scalar scaling factor for the entire tensor, with
shape (l,).
Outputs:
output: The quantized tensor in FP4, with shape (m, k // 2, l) but the physical
layout is (l, m, k // 2). `// 2` is because two fp4 values are packed into
an uint8.
output_scales: The blockscale tensor in FP8-E4M3, with shape (32, 4, rm, 4, rk, l)
but the physical layout is (l, rm, rk, 32, 4, 4).
Note:
For the shape of output_scales, `32 * 4 * rm` is a padded m to nearest multiple of 128.
`4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are
required by the NVIDIA Blackwell MMA operations.
"""
from sglang.jit_kernel.nvfp4 import (
scaled_fp4_grouped_quant as jit_scaled_fp4_grouped_quant,
)
return jit_scaled_fp4_grouped_quant(input_tensor, input_global_scale, mask)
def silu_and_mul_scaled_fp4_grouped_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
mask: torch.Tensor,
):
"""
Quantize input tensor to FP4 and return quantized tensor and scale, for
grouped gemm inputs (e.g., grouped_gemm_nt_masked for flashinfer).
Args:
input: The input tensor to be quantized to FP4, with shape (l, m, k * 2)
l is number of groups, m is number of tokens per group, k is number of features.
input_global_scale: A scalar scaling factor for the entire tensor, with
shape (l,).
mask: The mask tensor, with shape (l,)
Outputs:
output: The quantized tensor in FP4, with shape (m, k // 2, l) but the physical
layout is (l, m, k // 2). `// 2` is because two fp4 values are packed into
an uint8.
output_scales: The blockscale tensor in FP8-E4M3, with shape (32, 4, rm, 4, rk, l)
but the physical layout is (l, rm, rk, 32, 4, 4).
Note:
For the shape of output_scales, `32 * 4 * rm` is a padded m to nearest multiple of 128.
`4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are
required by the NVIDIA Blackwell MMA operations.
"""
from sglang.jit_kernel.nvfp4 import (
silu_and_mul_scaled_fp4_grouped_quant as jit_silu_and_mul_scaled_fp4_grouped_quant,
)
return jit_silu_and_mul_scaled_fp4_grouped_quant(
input_tensor,
input_global_scale,
mask,
)
# GPTQ kernels
def gptq_gemm(
a: torch.Tensor,

17
sgl-kernel/python/sgl_kernel/memory.py
View File

@@ -1,23 +1,6 @@
import torch
def set_kv_buffer_kernel(
k_cache: torch.Tensor,
v_cache: torch.Tensor,
loc: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
fallback: bool = False,
):
try:
if fallback:
raise RuntimeError("Fallback to torch implementation")
torch.ops.sgl_kernel.store_kv_cache(k_cache, v_cache, loc, k, v)
except RuntimeError: # ok, fallback to torch implementation
k_cache[loc] = k
v_cache[loc] = v
def weak_ref_tensor(tensor):
return (
torch.ops.sgl_kernel.weak_ref_tensor(tensor)

43
sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py
View File

@@ -1,7 +1,31 @@
from dataclasses import dataclass
from typing import Optional, Tuple, Union
import torch
from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace
from sglang.jit_kernel.rope import FusedSetKVBufferArg as _JitFusedSetKVBufferArg
from sglang.jit_kernel.rope import (
apply_rope_with_cos_sin_cache_inplace as _jit_apply_rope_with_cos_sin_cache_inplace,
)
@dataclass
class FusedSetKVBufferArg:
value: torch.Tensor
k_buffer: torch.Tensor
v_buffer: torch.Tensor
cache_loc: torch.Tensor
# Kept for backward compatibility with old sgl_kernel test/bench callsites.
k_scale: Optional[float] = None
v_scale: Optional[float] = None
def to_jit(self) -> _JitFusedSetKVBufferArg:
return _JitFusedSetKVBufferArg(
value=self.value,
k_buffer=self.k_buffer,
v_buffer=self.v_buffer,
cache_loc=self.cache_loc,
)
# vLLM torch native
@@ -129,14 +153,19 @@ class FlashInferRotaryEmbedding(RotaryEmbedding):
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=query,
key=key,
fused_set_kv_buffer_arg=fused_set_kv_buffer_arg,
head_size=self.head_size,
query_view = query.view(query.shape[0], -1, self.head_size)
key_view = key.view(key.shape[0], -1, self.head_size)
_jit_apply_rope_with_cos_sin_cache_inplace(
q=query_view,
k=key_view,
cos_sin_cache=self.cos_sin_cache,
positions=positions,
is_neox=self.is_neox_style,
fused_args=(
fused_set_kv_buffer_arg.to_jit()
if fused_set_kv_buffer_arg is not None
else None
),
)
return query, key