Remove obsolete sgl-kernel legacy paths (#21528)
This commit is contained in:
Xiaoyu Zhang
@@ -1,7 +1,6 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
from sgl_kernel import downcast_fp8 as downcast_fp8_aot
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
@@ -31,9 +30,9 @@ HEAD_DIM_LIST = get_benchmark_range(
|
||||
|
||||
CONFIGS = [(sl, h, d, sl * 2) for sl in SL_LIST for h, d in HEAD_DIM_LIST]
|
||||
|
||||
LINE_VALS = ["aot", "jit"]
|
||||
LINE_NAMES = ["AOT (sgl-kernel)", "JIT (cast.cuh, 256 threads, 2D grid)"]
|
||||
STYLES = [("blue", "--"), ("orange", "-")]
|
||||
LINE_VALS = ["jit"]
|
||||
LINE_NAMES = ["JIT (cast.cuh, 256 threads, 2D grid)"]
|
||||
STYLES = [("orange", "-")]
|
||||
|
||||
|
||||
# ── Perf report ────────────────────────────────────────────────────────────────
|
||||
@@ -48,7 +47,7 @@ STYLES = [("blue", "--"), ("orange", "-")]
|
||||
line_names=LINE_NAMES,
|
||||
styles=STYLES,
|
||||
ylabel="us",
|
||||
plot_name="downcast-fp8-aot-vs-jit",
|
||||
plot_name="downcast-fp8-jit",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
@@ -61,10 +60,7 @@ def benchmark(input_sl, head, dim, out_sl, provider):
|
||||
v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
|
||||
|
||||
if provider == "aot":
|
||||
fn = lambda: downcast_fp8_aot(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
else:
|
||||
fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
@@ -84,26 +80,19 @@ def _report_bandwidth(input_sl, head, dim, dtype):
|
||||
v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
|
||||
|
||||
aot_fn = lambda: downcast_fp8_aot(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
jit_fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
|
||||
aot_ms, _, _ = triton.testing.do_bench(aot_fn, quantiles=[0.5, 0.2, 0.8])
|
||||
jit_ms, _, _ = triton.testing.do_bench(jit_fn, quantiles=[0.5, 0.2, 0.8])
|
||||
|
||||
def fmt(ms):
|
||||
return f"{ms*1000:6.2f}us {total_bytes/(ms*1e-3)/1e9:6.0f}GB/s"
|
||||
|
||||
print(
|
||||
f" sl={input_sl:5d} h={head:2d} d={dim:4d}"
|
||||
f" | aot {fmt(aot_ms)}"
|
||||
f" | jit {fmt(jit_ms)}"
|
||||
f" | speedup {aot_ms/jit_ms:.2f}x"
|
||||
)
|
||||
print(f" sl={input_sl:5d} h={head:2d} d={dim:4d}" f" | jit {fmt(jit_ms)}")
|
||||
|
||||
|
||||
def report_bandwidth():
|
||||
print(f"\n{'='*95}")
|
||||
print(" AOT (sgl-kernel) vs JIT (cast.cuh, 256 threads, 2D grid)")
|
||||
print(" JIT (cast.cuh, 256 threads, 2D grid)")
|
||||
print(f" dtype={DTYPE}, device={DEVICE}")
|
||||
print(f"{'='*95}")
|
||||
for sl in [64, 256, 1024, 2048]:
|
||||
|
||||
@@ -82,31 +82,6 @@ def torch_top_p_renorm_probs(probs, top_p, eps=1e-5):
|
||||
return renorm_probs
|
||||
|
||||
|
||||
def torch_top_k_mask_logits(logits, top_k):
|
||||
"""Vectorized PyTorch implementation of top-k logits masking."""
|
||||
batch_size, vocab_size = logits.shape
|
||||
|
||||
# Handle scalar or tensor k
|
||||
if isinstance(top_k, int):
|
||||
k_val = min(max(top_k, 1), vocab_size)
|
||||
# Get top-k indices for all batches at once
|
||||
_, topk_indices = torch.topk(logits, k_val, dim=1, largest=True)
|
||||
|
||||
# Create masked logits: start with -inf everywhere
|
||||
masked_logits = torch.full_like(logits, float("-inf"))
|
||||
# Scatter the top-k values back
|
||||
masked_logits.scatter_(1, topk_indices, logits.gather(1, topk_indices))
|
||||
else:
|
||||
# Variable k per batch - need to handle separately
|
||||
masked_logits = torch.full_like(logits, float("-inf"))
|
||||
for i in range(batch_size):
|
||||
k_val = min(max(top_k[i].item(), 1), vocab_size)
|
||||
_, topk_indices = torch.topk(logits[i], k_val, largest=True)
|
||||
masked_logits[i, topk_indices] = logits[i, topk_indices]
|
||||
|
||||
return masked_logits
|
||||
|
||||
|
||||
def calculate_diff_top_k_renorm(batch_size, vocab_size, k):
|
||||
"""Compare Torch reference and SGLang kernel for top-k renorm correctness."""
|
||||
torch.manual_seed(42)
|
||||
@@ -139,20 +114,6 @@ def calculate_diff_top_p_renorm(batch_size, vocab_size, p):
|
||||
torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
def calculate_diff_top_k_mask(batch_size, vocab_size, k):
|
||||
"""Compare Torch reference and SGLang kernel for top-k mask correctness."""
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
logits = torch.randn(batch_size, vocab_size, device=device) * 5
|
||||
top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
|
||||
|
||||
torch_output = torch_top_k_mask_logits(logits, top_k_tensor)
|
||||
sglang_output = sgl_kernel.top_k_mask_logits(logits, top_k_tensor)
|
||||
|
||||
torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
# Parameter space - simplified for CI
|
||||
if is_in_ci():
|
||||
batch_size_range = [16]
|
||||
@@ -231,38 +192,6 @@ def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "vocab_size", "k"],
|
||||
x_vals=configs_k,
|
||||
line_arg="provider",
|
||||
line_vals=["torch", "sglang"],
|
||||
line_names=["Torch Reference", "SGL Kernel"],
|
||||
styles=[("red", "-"), ("orange", "-")],
|
||||
ylabel="us",
|
||||
plot_name="top-k-mask-logits-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_top_k_mask(batch_size, vocab_size, k, provider):
|
||||
# Skip invalid configurations
|
||||
if k >= vocab_size:
|
||||
return float("nan"), float("nan"), float("nan")
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
logits = torch.randn(batch_size, vocab_size, device=device) * 5
|
||||
top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
|
||||
|
||||
if provider == "torch":
|
||||
fn = lambda: torch_top_k_mask_logits(logits.clone(), top_k_tensor)
|
||||
elif provider == "sglang":
|
||||
fn = lambda: sgl_kernel.top_k_mask_logits(logits.clone(), top_k_tensor)
|
||||
|
||||
return run_benchmark_no_cudagraph(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("Running correctness checks...")
|
||||
@@ -291,15 +220,6 @@ if __name__ == "__main__":
|
||||
batch_size, vocab_size, p = cfg
|
||||
print(f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, p={p}")
|
||||
|
||||
print("\n3. Testing top_k_mask_logits...")
|
||||
for cfg in test_configs_k:
|
||||
batch_size, vocab_size, k = cfg
|
||||
if k < vocab_size: # Skip invalid configs
|
||||
calculate_diff_top_k_mask(batch_size, vocab_size, k)
|
||||
print(
|
||||
f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}"
|
||||
)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All correctness checks passed!")
|
||||
print("=" * 60)
|
||||
@@ -314,9 +234,6 @@ if __name__ == "__main__":
|
||||
print("\n2. Benchmarking top_p_renorm_probs...")
|
||||
benchmark_top_p_renorm.run(print_data=True)
|
||||
|
||||
print("\n3. Benchmarking top_k_mask_logits...")
|
||||
benchmark_top_k_mask.run(print_data=True)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Benchmarking complete!")
|
||||
print("=" * 60)
|
||||
|
||||
Reference in New Issue
Block a user