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

python/sglang/jit_kernel/benchmark/bench_renorm.py

@@ -0,0 +1,321 @@
+import itertools
+
+import sgl_kernel
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import is_in_ci
+
+
+def torch_top_k_renorm_probs(probs, top_k):
+    """Vectorized PyTorch implementation of top-k renormalization."""
+    batch_size, vocab_size = probs.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(probs, k_val, dim=1, largest=True)
+
+        # Create mask: batch_size x vocab_size
+        mask = torch.zeros_like(probs)
+        mask.scatter_(1, topk_indices, 1.0)
+
+        # Vectorized renormalization
+        masked_probs = probs * mask
+        renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + 1e-10)
+        return renorm_probs
+    else:
+        # Variable k per batch - need to handle separately
+        renorm_probs = torch.zeros_like(probs)
+        for i in range(batch_size):
+            k_val = min(max(top_k[i].item(), 1), vocab_size)
+            _, topk_indices = torch.topk(probs[i], k_val, largest=True)
+            mask = torch.zeros_like(probs[i])
+            mask[topk_indices] = 1.0
+            masked_probs = probs[i] * mask
+            renorm_probs[i] = masked_probs / (masked_probs.sum() + 1e-10)
+        return renorm_probs
+
+
+def torch_top_p_renorm_probs(probs, top_p, eps=1e-5):
+    """Vectorized PyTorch implementation of top-p renormalization."""
+    batch_size, vocab_size = probs.shape
+
+    # Handle scalar or tensor p
+    if isinstance(top_p, float):
+        p_val = top_p
+        # Vectorized implementation for uniform top_p
+        # Sort probs in descending order
+        sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=1)
+        cumsum_probs = torch.cumsum(sorted_probs, dim=1)
+
+        # Find cutoff: where cumsum exceeds top_p
+        cutoff_mask = cumsum_probs <= p_val
+        # Keep at least one token (the highest prob)
+        cutoff_mask[:, 0] = True
+
+        # Create mask in original order
+        mask = torch.zeros_like(probs)
+        mask.scatter_(1, sorted_indices, cutoff_mask.float())
+
+        # Vectorized renormalization
+        masked_probs = probs * mask
+        renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + eps)
+        return renorm_probs
+    else:
+        # Variable p per batch - need to handle separately
+        renorm_probs = torch.zeros_like(probs)
+        for i in range(batch_size):
+            p_val = top_p[i].item()
+            sorted_prob, indices = torch.sort(probs[i], descending=False)
+            cdf = torch.cumsum(sorted_prob, dim=-1)
+            mask = torch.zeros(vocab_size, dtype=torch.float32, device=probs.device)
+            mask.scatter_(0, indices, (cdf >= (1 - p_val) - eps).float())
+            masked_probs = probs[i] * mask
+            renorm_probs[i] = masked_probs / (masked_probs.sum() + eps)
+        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)
+    device = torch.device("cuda")
+
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
+    probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+
+    top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
+
+    torch_output = torch_top_k_renorm_probs(probs, top_k_tensor)
+    sglang_output = sgl_kernel.top_k_renorm_prob(probs, top_k_tensor)
+
+    torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3)
+
+
+def calculate_diff_top_p_renorm(batch_size, vocab_size, p):
+    """Compare Torch reference and SGLang kernel for top-p renorm correctness."""
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
+    probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+
+    top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32)
+
+    torch_output = torch_top_p_renorm_probs(probs, top_p_tensor)
+    sglang_output = sgl_kernel.top_p_renorm_prob(probs, top_p_tensor)
+
+    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]
+    vocab_size_range = [111]
+    k_range = [10]
+    p_range = [0.5]
+else:
+    batch_size_range = [16, 64, 128]
+    vocab_size_range = [111, 32000, 128256]
+    k_range = [10, 100, 500]
+    p_range = [0.1, 0.5, 0.9]
+
+configs_k = list(itertools.product(batch_size_range, vocab_size_range, k_range))
+configs_p = list(itertools.product(batch_size_range, vocab_size_range, p_range))
+
+
+@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 (FlashInfer)"],
+        styles=[("red", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="top-k-renorm-probs-performance",
+        args={},
+    )
+)
+def benchmark_top_k_renorm(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")
+
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
+    probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32)
+
+    if provider == "torch":
+        fn = lambda: torch_top_k_renorm_probs(probs.clone(), top_k_tensor)
+    elif provider == "sglang":
+        fn = lambda: sgl_kernel.top_k_renorm_prob(probs.clone(), top_k_tensor)
+
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "vocab_size", "p"],
+        x_vals=configs_p,
+        line_arg="provider",
+        line_vals=["torch", "sglang"],
+        line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
+        styles=[("red", "-"), ("blue", "-")],
+        ylabel="us",
+        plot_name="top-p-renorm-probs-performance",
+        args={},
+    )
+)
+def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
+    probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32)
+
+    if provider == "torch":
+        fn = lambda: torch_top_p_renorm_probs(probs.clone(), top_p_tensor)
+    elif provider == "sglang":
+        fn = lambda: sgl_kernel.top_p_renorm_prob(probs.clone(), top_p_tensor)
+
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+@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 (FlashInfer)"],
+        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)
+
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Running correctness checks...")
+    print("=" * 60)
+
+    # Correctness checks - simplified for CI
+    if is_in_ci():
+        test_configs_k = [configs_k[0]] if configs_k else [(16, 111, 10)]
+        test_configs_p = [configs_p[0]] if configs_p else [(16, 111, 0.5)]
+    else:
+        test_configs_k = configs_k[:3]  # Test first 3 configs
+        test_configs_p = configs_p[:3]
+
+    print("\n1. Testing top_k_renorm_probs...")
+    for cfg in test_configs_k:
+        batch_size, vocab_size, k = cfg
+        if k < vocab_size:  # Skip invalid configs
+            calculate_diff_top_k_renorm(batch_size, vocab_size, k)
+            print(
+                f"  ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}"
+            )
+
+    print("\n2. Testing top_p_renorm_probs...")
+    for cfg in test_configs_p:
+        calculate_diff_top_p_renorm(*cfg)
+        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)
+
+    print("\n" + "=" * 60)
+    print("Starting performance benchmarks...")
+    print("=" * 60)
+
+    print("\n1. Benchmarking top_k_renorm_probs...")
+    benchmark_top_k_renorm.run(print_data=True)
+
+    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

@@ -0,0 +1,118 @@
+# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/main/tests/test_sampling.py
+# and /sgl-workspace/sglang/sgl-kernel/tests/test_sampling.py
+
+import pytest
+import sgl_kernel
+import torch
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("k", [10, 100, 500])
+def test_top_k_renorm_probs(batch_size, vocab_size, k):
+    """Test top_k_renorm_probs kernel for correctness.
+
+    This test validates that the kernel correctly:
+    1. Identifies the top-k probabilities
+    2. Masks out non-top-k values
+    3. Renormalizes the remaining probabilities to sum to 1
+    """
+    if k > vocab_size:
+        pytest.skip("k should be less than vocab_size")
+    torch.manual_seed(42)
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    sorted_prob, _ = torch.sort(normalized_prob, descending=True)
+    pivot = sorted_prob[:, k - 1]
+    mask = (normalized_prob >= pivot.unsqueeze(-1)).int()
+    renorm_prob_ground_truth = normalized_prob.clone()
+    renorm_prob_ground_truth[mask == 0] = 0
+    renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum(
+        dim=-1, keepdim=True
+    )
+
+    renorm_prob = sgl_kernel.top_k_renorm_prob(normalized_prob, k)
+    for i in range(batch_size):
+        torch.testing.assert_close(
+            renorm_prob_ground_truth[i],
+            renorm_prob[i],
+            rtol=1e-3,
+            atol=1e-3,
+        )
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("p", [0.1, 0.5, 0.9])
+def test_top_p_renorm_probs(batch_size, vocab_size, p):
+    """Test top_p_renorm_probs kernel for correctness.
+
+    This test validates that the kernel correctly:
+    1. Computes the cumulative probability distribution
+    2. Identifies tokens in the top-p threshold
+    3. Masks out tokens outside the threshold
+    4. Renormalizes the remaining probabilities to sum to 1
+    """
+    torch.manual_seed(42)
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    sorted_prob, indices = torch.sort(normalized_prob, descending=False)
+    cdf = torch.cumsum(sorted_prob, dim=-1)
+    mask = torch.zeros(batch_size, vocab_size, dtype=torch.int32, device="cuda:0")
+    mask.scatter_add_(1, indices, (cdf >= (1 - p)).int())
+    renorm_prob_ground_truth = normalized_prob.clone()
+    renorm_prob_ground_truth[mask == 0] = 0
+    renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum(
+        dim=-1, keepdim=True
+    )
+
+    renorm_prob = sgl_kernel.top_p_renorm_prob(normalized_prob, p)
+    torch.testing.assert_close(
+        renorm_prob_ground_truth,
+        renorm_prob,
+        rtol=1e-3,
+        atol=1e-3,
+    )
+
+
+@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__":
+    pytest.main([__file__])

sgl-kernel/python/sgl_kernel/sampling.py

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