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

Remove obsolete sgl-kernel legacy paths (#21528)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-04-01 09:00:20 +08:00
committed by GitHub
parent a8759dd9af
commit cdd7d6a227
17 changed files with 14 additions and 1708 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/tests/test_hadamard.py
-86
View File
@@ -1,86 +0,0 @@
import math
import sys
import pytest
import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from scipy.linalg import hadamard
try:
from sgl_kernel import hadamard_transform
except Exception:
pytest.skip(
"sgl-kernel hadamard interface was removed (migrated to jit_kernel)",
allow_module_level=True,
)
def hadamard_transform_ref(x, scale=1.0):
"""
x: (..., dim)
out: (..., dim)
"""
if hadamard is None:
raise ImportError("Please install scipy")
x_shape = x.shape
dim = x.shape[-1]
x = x.reshape(-1, dim)
log_dim = math.ceil(math.log2(dim))
dim_padded = 2**log_dim
if dim != dim_padded:
x = F.pad(x, (0, dim_padded - dim))
out = F.linear(
x,
torch.tensor(hadamard(dim_padded, dtype=float), dtype=x.dtype, device=x.device),
)
out = out * scale
return out[..., :dim].reshape(*x_shape)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
"dim",
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 137, 1024, 2048, 4096, 8192, 16384, 32768],
)
def test_fast_hadamard_transform(dim, dtype):
device = "cuda"
if dtype == torch.float32:
rtol, atol = 3e-4, 3e-3
elif dtype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
else: # float16
rtol, atol = 3e-3, 5e-3
torch.random.manual_seed(0)
batch_size = 15
x = torch.randn(batch_size, dim, device=device, dtype=dtype)
x_ref = x.detach().clone().to(torch.float32)
x_pt = x.detach().clone()
scale = 1 / math.sqrt(dim)
out = hadamard_transform(x, scale=scale)
out_ref = hadamard_transform_ref(x_ref, scale=scale)
out_pt = hadamard_transform_ref(x_pt, scale=scale)
torch.testing.assert_close(
out_pt.float(),
out_ref,
rtol=rtol,
atol=atol,
msg="Reference implementations mismatch",
)
torch.testing.assert_close(
out.float(),
out_ref,
rtol=rtol,
atol=atol,
msg="fast_hadamard_transform output mismatch",
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
sgl-kernel/tests/test_merge_state.py
-143
View File
@@ -1,143 +0,0 @@
# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/55576c626421b5ee7e7ebe74afd26465c8ae863f/flashinfer/triton/kernels/cascade.py
import sys
from typing import List
import pytest
import torch
import triton
import triton.language as tl
from sgl_kernel import merge_state
def check_input(x: torch.Tensor):
assert x.is_cuda, f"{str(x)} must be a CUDA Tensor"
assert x.is_contiguous(), f"{str(x)} must be contiguous"
def check_dim(d, x: torch.Tensor):
assert x.dim() == d, f"{str(x)} must be a {d}D tensor"
def check_shape(a: torch.Tensor, b: torch.Tensor):
assert a.dim() == b.dim(), "tensors should have same dim"
for i in range(a.dim()):
assert a.size(i) == b.size(
i
), f"tensors shape mismatch, {a.size()} and {b.size()}"
def check_device(tensors: List[torch.Tensor]):
device = tensors[0].device
for t in tensors:
assert (
t.device == device
), f"All tensors should be on the same device, but got {device} and {t.device}"
@triton.jit
def state_merge(o, m, d, other_o, other_m, other_d):
m_max = tl.maximum(m, other_m)
d = d * tl.exp2(m - m_max) + other_d * tl.exp2(other_m - m_max)
o = o * tl.exp2(m - m_max) + other_o * tl.exp2(other_m - m_max)
return o, m_max, d
@triton.jit
def state_normalize(o, m, d):
o = o / d
return o, m, d
@triton.jit
def state_get_lse(o, m, d):
return m + tl.log2(d)
@triton.jit
def merge_state_kernel(
v_a_ptr,
s_a_ptr,
v_b_ptr,
s_b_ptr,
v_merged_ptr,
s_merged_ptr,
num_heads,
head_dim,
bdx: tl.constexpr,
bdy: tl.constexpr,
):
pos = tl.program_id(axis=0)
for tx in tl.range(bdx):
for head_idx in tl.range(bdy):
s_a_val = tl.load(s_a_ptr + pos * num_heads + head_idx)
s_b_val = tl.load(s_b_ptr + pos * num_heads + head_idx)
offsets = (pos * num_heads + head_idx) * head_dim + tx
v_a = tl.load(v_a_ptr + offsets)
v_b = tl.load(v_b_ptr + offsets)
v_merged, s_max, d = state_merge(
o=v_a, m=s_a_val, d=1, other_o=v_b, other_m=s_b_val, other_d=1
)
v_merged, s_max, d = state_normalize(v_merged, s_max, d)
v_merged_offset = (pos * num_heads + head_idx) * head_dim + tx
tl.store(v_merged_ptr + v_merged_offset, v_merged)
if s_merged_ptr:
tl.store(
s_merged_ptr + pos * num_heads + head_idx,
tl.log2(d) + s_max,
)
def merge_state_triton(
v_a: torch.Tensor, s_a: torch.Tensor, v_b: torch.Tensor, s_b: torch.Tensor
):
check_input(v_a)
check_input(s_a)
check_input(v_b)
check_input(s_b)
check_device([v_a, s_a, v_b, s_b])
check_dim(3, v_a)
check_dim(2, s_a)
check_dim(3, v_b)
check_dim(2, s_b)
check_shape(v_a, v_b)
check_shape(s_a, s_b)
assert v_a.size(0) == s_a.size(0)
assert v_a.size(1) == s_b.size(1)
s_a = s_a.to(torch.float32)
s_b = s_b.to(torch.float32)
seq_len = v_a.size(0)
num_heads = v_a.size(1)
head_dim = v_a.size(2)
v_merged = torch.empty_like(v_a).to(s_a.device)
s_merged = torch.empty((seq_len, num_heads)).to(s_a.device)
bdx = head_dim
bdy = num_heads
merge_state_kernel[lambda meta: (seq_len,)](
v_a, s_a, v_b, s_b, v_merged, s_merged, num_heads, head_dim, bdx=bdx, bdy=bdy
)
return v_merged, s_merged
@pytest.mark.parametrize("seq_len", [2048])
@pytest.mark.parametrize("num_heads", [32])
@pytest.mark.parametrize("head_dim", [128])
def test_merge_state(seq_len, num_heads, head_dim):
va = torch.randn(seq_len, num_heads, head_dim).half().to("cuda:0")
sa = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda:0")
vb = torch.randn(seq_len, num_heads, head_dim).half().to("cuda:0")
sb = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda:0")
v_merged, s_merged = merge_state_triton(va, sa, vb, sb)
v_merged_std, s_merged_std = merge_state(va, sa, vb, sb)
assert torch.allclose(v_merged, v_merged_std, atol=1e-2)
assert torch.allclose(s_merged, s_merged_std, atol=1e-2)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
sgl-kernel/tests/test_merge_state_v2.py
+7 -27
View File
@@ -5,7 +5,7 @@ import pytest
import torch
import triton
import triton.language as tl
from sgl_kernel import merge_state, merge_state_v2
from sgl_kernel import merge_state_v2
@triton.jit
@@ -146,11 +146,9 @@ def generate_markdown_table():
global all_case_info
table_header = (
"| tokens | heads | headsize | dtype "
"| device | torch | triton | v1 | v2 | speedup(vs triton) | speedup(vs v1)|"
)
table_separator = (
"| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |"
"| device | torch | triton | v2 | speedup(vs triton) |"
)
table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"
def shortly_dtype(dtype: torch.dtype) -> str:
return str(dtype).removeprefix("torch.")
@@ -169,21 +167,17 @@ def generate_markdown_table():
device,
time_torch,
time_triton,
time_v1,
time_v2,
) = info
dtype = shortly_dtype(dtype)
device = shortly_device(device)
improved_triton = time_triton / time_v2
improved_v1 = time_v1 / time_v2
print(
f"| {num_tokens} | {num_heads} | {head_size} "
f"| {dtype} | {device} | {time_torch:.4f}ms "
f"| {time_triton:.4f}ms "
f"| {time_v1:.4f}ms "
f"| {time_v2:.4f}ms "
f"| {improved_triton:.4f}x "
f"| {improved_v1:.4f}x |"
f"| {improved_triton:.4f}x |"
)
@@ -259,11 +253,6 @@ def test_merge_attn_states(
prefix_lse_ = prefix_lse
suffix_lse_ = suffix_lse
if fn_type == "cuda_v1":
# merge_state v1 kernel not support float32
if output_dtype not in (torch.half, torch.bfloat16):
return 0, output_fn, output_lse_fn
total_time = 0
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
@@ -316,29 +305,21 @@ def test_merge_attn_states(
fn_type="triton",
)
# 2. Run the merge_state V1 kernel
output_v1 = output.clone()
output_lse_v1 = output_lse.clone()
time_v1, output_v1, output_lse_v1 = perf_kernel_fn(
output_v1, output_lse_v1, merge_state, fn_type="cuda_v1"
)
# 3. Run the merge_state V2 kernel
# 2. Run the merge_state V2 kernel
output_v2 = output.clone()
output_lse_v2 = output_lse.clone()
time_v2, output_v2, output_lse_v2 = perf_kernel_fn(
output_v2, output_lse_v2, merge_state_v2, fn_type="cuda_v2"
)
# 4. Performance compare
# 3. Performance compare
improved = time_triton / time_v2
print(f" Torch time: {time_torch:.6f}ms")
print(f" Triton time: {time_triton:.6f}ms")
print(f"CUDA v1 time: {time_v1:.6f}ms")
print(f"CUDA v2 time: {time_v2:.6f}ms, Performance: {improved:.5f}x")
print("-" * 100)
# 5. Correctness compare
# 4. Correctness compare
# Liger Kernel: Efficient Triton Kernels for LLM Training
# https://arxiv.org/pdf/2410.10989, 3.3 Correctness
# use rtol = 1e-2 for bfloat16.
@@ -387,7 +368,6 @@ def test_merge_attn_states(
device,
time_torch,
time_triton,
time_v1,
time_v2,
)
)