Remove obsolete sgl-kernel legacy paths (#21528)

python/sglang/jit_kernel/benchmark/bench_cast.py

@@ -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]:

python/sglang/jit_kernel/benchmark/bench_renorm.py

@@ -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)

python/sglang/jit_kernel/tests/test_renorm.py

@@ -82,44 +82,5 @@ def test_top_p_renorm_probs(batch_size, vocab_size, p):
     )
 
 
-@pytest.mark.parametrize("batch_size", [1, 99, 989])
-@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
-@pytest.mark.parametrize("k", [10, 100, 500])
-@pytest.mark.parametrize("neginf_input", [False, True])
-def test_top_k_mask_logits(batch_size, vocab_size, k, neginf_input):
-    """Test top_k_mask_logits kernel for correctness.
-
-    This test validates that the kernel correctly:
-    1. Identifies the top-k logits
-    2. Masks non-top-k values to -inf
-    3. Preserves the top-k values
-    4. Handles negative infinity inputs gracefully
-
-    The test verifies correctness by comparing softmax(top_k_mask_logits(logits))
-    with top_k_renorm_prob(probs), which should be equivalent.
-    """
-    if k > vocab_size:
-        pytest.skip("k should be less than vocab_size")
-    torch.manual_seed(42)
-    logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5
-    if neginf_input:
-        # Randomly assign some logits to -inf to test edge cases
-        num_neginf = torch.randint(1, vocab_size * batch_size, (1,)).item()
-        idxs = torch.randperm(batch_size * vocab_size, device="cuda:0")[:num_neginf]
-        logits[idxs // vocab_size, idxs % vocab_size] = -float("inf")
-
-    probs = torch.softmax(logits, dim=-1)
-    masked_logits = sgl_kernel.top_k_mask_logits(logits, k)
-    renormed_probs = torch.softmax(masked_logits, dim=-1)
-    renormed_probs_ref = sgl_kernel.top_k_renorm_prob(probs, k)
-
-    torch.testing.assert_close(
-        renormed_probs,
-        renormed_probs_ref,
-        rtol=1e-3,
-        atol=1e-3,
-    )
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))