/test/registered/jit/benchmark/bench_renorm.py

import itertools

import sgl_kernel
import torch
import triton
import triton.testing

from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci

register_cuda_ci(
    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)


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


# 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"],
        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)

    return run_benchmark_no_cudagraph(fn)


@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"],
        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)

    return run_benchmark_no_cudagraph(fn)


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("\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("\n" + "=" * 60)
    print("Benchmarking complete!")
    print("=" * 60)