chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-kernel/benchmark/bench_activation.py

@@ -0,0 +1,209 @@
+# Benchmarks SGLang kernels versus vLLM across
+# (kernel, dtype, batch_size, seq_len, dim) and prints speed-up.
+import argparse
+import itertools
+import os
+import re
+from typing import List, Tuple
+
+import sgl_kernel
+import torch
+import torch.nn.functional as F
+import triton
+import triton.testing
+from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
+
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm import _custom_ops as vllm_ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    vllm_ops = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+# gelu_quick is only available on HIP/ROCm platforms
+try:
+    from sgl_kernel import gelu_quick
+
+    GELU_QUICK_AVAILABLE = True
+except ImportError:
+    GELU_QUICK_AVAILABLE = False
+    gelu_quick = None
+
+if VLLM_AVAILABLE and not hasattr(vllm_ops, "silu_and_mul"):
+    vllm_ops = torch.ops._C
+
+
+def str2int_list(arg: str) -> List[int]:
+    if arg in ("", None):
+        return []
+    if re.fullmatch(r"\d+(,\d+)*", arg.strip()) is None:
+        raise argparse.ArgumentTypeError(f"Bad int list: {arg}")
+    return [int(x) for x in arg.split(",")]
+
+
+def calculate_diff(
+    kernel: str, dtype: torch.dtype, batch_size: int, seq_len: int, dim: int
+) -> bool:
+    """Compare vLLM with SGLang for one shape."""
+    device = torch.device("cuda")
+
+    if not VLLM_AVAILABLE:
+        print(
+            f"[{kernel:14s} | {str(dtype):9s} | B={batch_size:3d} | "
+            f"L={seq_len:3d} | D={dim:5d}] ⚠️  vLLM not available, skipping comparison"
+        )
+        return True
+
+    # activation-only quick GELU
+    if kernel == "gelu_quick":
+        if not GELU_QUICK_AVAILABLE:
+            print(
+                f"[{kernel:14s} | {str(dtype):9s} | B={batch_size:3d} | "
+                f"L={seq_len:3d} | D={dim:5d}] ⚠️  not available on this platform"
+            )
+            return True
+        x = torch.randn(batch_size, seq_len, dim, dtype=dtype, device=device)
+        ref_out = torch.zeros_like(x)
+        getattr(vllm_ops, kernel)(ref_out, x)
+        test_out = getattr(sgl_kernel, kernel)(x)
+    # fused activation x mul kernels
+    else:
+        x = torch.randn(batch_size, seq_len, 2 * dim, dtype=dtype, device=device)
+        ref_out = torch.zeros(batch_size, seq_len, dim, dtype=dtype, device=device)
+        getattr(vllm_ops, kernel)(ref_out, x)
+        test_out = getattr(sgl_kernel, kernel)(x)
+
+    ok = torch.allclose(ref_out, test_out, rtol=1e-3, atol=1e-5)
+    tag = "✅ match" if ok else "❌ mismatch"
+    print(
+        f"[{kernel:14s} | {str(dtype):9s} | B={batch_size:3d} | "
+        f"L={seq_len:3d} | D={dim:5d}] {tag}"
+    )
+    return ok
+
+
+# CI environment uses simplified parameters for kernels and dtypes too
+if IS_CI:
+    kernels = ["silu_and_mul"]  # Only test one kernel in CI
+    dtypes = [torch.float16]  # Only test one dtype in CI
+else:
+    kernels = ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul"]
+    if GELU_QUICK_AVAILABLE:
+        kernels.append("gelu_quick")
+    dtypes = [torch.float16, torch.bfloat16]
+
+
+def make_configs(bsizes: List[int], slens: List[int], dims_: List[int]) -> List[Tuple]:
+    return list(itertools.product(kernels, dtypes, bsizes, slens, dims_))
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    default_batch_sizes = [1]  # Single batch size for CI
+    default_seq_lens = [1]  # Single sequence length for CI
+    default_dims = [1024]  # Single dimension for CI
+else:
+    default_batch_sizes = [2**i for i in range(0, 5, 2)]  # 1,4,16
+    default_seq_lens = [2**i for i in range(0, 8, 2)]  # 1,4,16,64
+    default_dims = [2**i for i in range(10, 15)]  # 1024...16384
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["kernel", "dtype", "batch_size", "seq_len", "dim"],
+        x_vals=[],
+        line_arg="provider",
+        line_vals=["vllm", "sglang", "speedup"],
+        line_names=["vLLM", "SGL Kernel", "Speed-up (x)"],
+        styles=[("blue", "-"), ("green", "-"), ("red", "--")],
+        ylabel="µs (median)  or  × (speed-up)",
+        plot_name="activation-performance",
+        args={},
+    )
+)
+def benchmark(kernel, dtype, batch_size, seq_len, dim, provider):
+    device = torch.device("cuda")
+    in_mult = 1 if kernel == "gelu_quick" else 2
+    x = torch.randn(batch_size, seq_len, in_mult * dim, dtype=dtype, device=device)
+    y0 = torch.zeros(batch_size, seq_len, dim, dtype=dtype, device=device)
+
+    if not VLLM_AVAILABLE and provider in ["vllm", "speedup"]:
+        # Skip vLLM-related benchmarks if vLLM is not available
+        return (0, 0, 0)
+
+    if VLLM_AVAILABLE:
+        vllm_kernel = getattr(vllm_ops, kernel)
+    if kernel == "gelu_quick" and not GELU_QUICK_AVAILABLE:
+        # Skip benchmark for gelu_quick if not available
+        return (0, 0, 0)
+    sglang_kernel = getattr(sgl_kernel, kernel)
+
+    def baseline():
+        if VLLM_AVAILABLE:
+            tmp = y0.clone()
+            vllm_kernel(tmp, x)
+            return tmp
+        else:
+            return torch.zeros_like(y0)
+
+    def sglang():
+        return sglang_kernel(x)
+
+    # timing helper
+    def timed(fn):
+        for _ in range(5):
+            fn()
+        torch.cuda.synchronize()
+        ms, qmin, qmax = triton.testing.do_bench_cudagraph(
+            fn, quantiles=[0.5, 0.2, 0.8]
+        )
+        return 1000 * ms, 1000 * qmax, 1000 * qmin
+
+    if provider == "vllm":
+        return timed(baseline)
+    if provider == "sglang":
+        return timed(sglang)
+
+    # provider == "speedup"
+    t_ref, _, _ = timed(baseline)
+    t_sgl, _, _ = timed(sglang)
+    spd = t_ref / t_sgl if t_ref > 0 else 1.0
+    return (spd, spd, spd)
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser("Activation kernel benchmark")
+    p.add_argument("--batch_sizes", type=str2int_list, default=default_batch_sizes)
+    p.add_argument("--seq_lens", type=str2int_list, default=default_seq_lens)
+    p.add_argument("--dims", type=str2int_list, default=default_dims)
+    p.add_argument("--verify_only", action="store_true")
+    args = p.parse_args()
+
+    # coerce lists
+    if isinstance(args.batch_sizes, str):
+        args.batch_sizes = str2int_list(args.batch_sizes)
+    if isinstance(args.seq_lens, str):
+        args.seq_lens = str2int_list(args.seq_lens)
+    if isinstance(args.dims, str):
+        args.dims = str2int_list(args.dims)
+
+    # patch perf_report grid
+    benchmark_grid = make_configs(args.batch_sizes, args.seq_lens, args.dims)
+    if hasattr(benchmark, "benchmarks"):
+        benchmark.benchmarks.x_vals = benchmark_grid
+    else:
+        benchmark.benchmark.x_vals = benchmark_grid
+
+    if args.verify_only:
+        # Test with the first available kernel
+        test_kernel = kernels[0]
+        ok = calculate_diff(test_kernel, torch.float16, 1, 1, args.dims[0])
+        print("✅ sanity pass" if ok else "❌ mismatch")
+    else:
+        benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py

@@ -0,0 +1,689 @@
+"""
+Benchmark latency comparison between different all-reduce implementations.
+
+Compares:
+- NCCL all-reduce (may be non-deterministic)
+- Reduce-scatter + all-gather (RS+AG, deterministic but slower)
+- Deterministic 1-stage kernel (forces fixed accumulation order, deterministic)
+
+Note: The "deterministic kernel" is NOT RS+AG. It uses the 1-stage kernel where
+each GPU reads all data from all GPUs and reduces locally in a fixed order.
+
+Usage:
+    python bench_amd_deterministic_allreduce.py
+"""
+
+import multiprocessing as mp
+import os
+import socket
+import statistics
+import sys
+import time
+
+import torch
+import torch.distributed as dist
+
+# Add python directory to path to import sglang modules
+script_dir = os.path.dirname(os.path.abspath(__file__))
+python_dir = os.path.join(script_dir, "python")
+sys.path.insert(0, python_dir)
+
+# Try to import custom all-reduce if available
+from sglang.srt.environ import envs
+
+try:
+    import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as custom_ar_ops
+    from sglang.srt.distributed.device_communicators.custom_all_reduce import (
+        CustomAllreduce,
+    )
+
+    CUSTOM_AR_AVAILABLE = custom_ar_ops.IS_CUSTOM_AR_AVAILABLE
+except (ImportError, AttributeError):
+    CUSTOM_AR_AVAILABLE = False
+    CustomAllreduce = None
+
+# Note: sglang's optimized all-reduce requires full runtime initialization
+# and won't work in standalone benchmarks, so we skip it
+SGLANG_AVAILABLE = False
+
+
+def get_open_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def init_custom_ar_if_available(rank, world_size, device):
+    """Check if custom all-reduce is available and applicable."""
+    if not CUSTOM_AR_AVAILABLE or CustomAllreduce is None:
+        return False
+
+    # Custom AR works best for single-node, even number of GPUs, world_size <= 8
+    if world_size <= 8 and world_size % 2 == 0:
+        return True
+
+    return False
+
+
+def reduce_scatter_then_all_gather(tensor, rank, world_size, custom_ar=None):
+    """
+    Deterministic all-reduce using reduce-scatter + all-gather.
+    This is deterministic because it uses fixed ordering (no atomics).
+    """
+    total_size = tensor.numel()
+    if total_size % world_size != 0:
+        # Fallback to all-gather + local reduce if not divisible
+        gather_list = [torch.empty_like(tensor) for _ in range(world_size)]
+        dist.all_gather(gather_list, tensor)
+        stacked = torch.stack(gather_list, dim=0)
+        tensor.copy_(stacked.sum(dim=0))
+        return
+
+    chunk_size = total_size // world_size
+
+    # Flatten to 1D
+    tensor_flat = tensor.view(-1)
+
+    # Reduce-scatter: each rank gets its chunk of the reduced result
+    output_chunk = torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device)
+
+    # Split input into chunks for reduce-scatter
+    input_chunks = [
+        tensor_flat[i * chunk_size : (i + 1) * chunk_size].clone()
+        for i in range(world_size)
+    ]
+
+    dist.reduce_scatter(output_chunk, input_chunks)
+
+    # All-gather: broadcast each rank's chunk to all ranks
+    output_chunks = [
+        torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device)
+        for _ in range(world_size)
+    ]
+    dist.all_gather(output_chunks, output_chunk)
+
+    # Concatenate results back
+    result_flat = torch.cat(output_chunks, dim=0)
+    tensor.copy_(result_flat.view(tensor.shape))
+
+
+def worker(world_size, rank, port, results_queue):
+    envs.SGLANG_USE_1STAGE_ALLREDUCE.set("1")
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://localhost:{port}",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    # Try to initialize custom all-reduce if available
+    custom_ar = None
+    use_custom_ar = init_custom_ar_if_available(rank, world_size, device)
+    if use_custom_ar and CUSTOM_AR_AVAILABLE:
+        try:
+            # Create a gloo group for custom AR (it requires non-NCCL backend)
+            # All ranks must call new_group with the same parameters
+            from torch.distributed import new_group
+
+            dist.barrier()  # Ensure all ranks are ready
+            ar_group = new_group(backend="gloo")
+            dist.barrier()  # Ensure group creation is complete
+            custom_ar = CustomAllreduce(group=ar_group, device=device)
+            if rank == 0:
+                print("  Using custom all-reduce (deterministic)")
+        except Exception as e:
+            if rank == 0:
+                print(f"  Custom AR init failed: {e}, using NCCL fallback")
+            custom_ar = None
+            dist.barrier()  # Ensure all ranks continue even if one fails
+
+    # Test different batch sizes - similar to test_ar.py
+    batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512]
+    hidden_dim = 16384  # Fixed hidden dimension
+
+    num_trials = 10  # Same as test_ar.py
+
+    # Different seed per rank - each GPU has DIFFERENT input (like test_ar.py)
+    torch.manual_seed(42 + rank)
+
+    results = {}
+
+    for bs in batch_sizes:
+        # Create fixed input for all trials (like test_ar.py)
+        base_input = torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device)
+
+        dist.barrier()
+
+        if rank == 0:
+            print(f"\nBatch size {bs:4d}:")
+            print(f"  Testing determinism across {num_trials} trials...")
+
+        # Test all-reduce determinism
+        results_ar = []
+        latencies_ar = []
+        for trial in range(num_trials):
+            # Clone the same input for each trial
+            inp_ar = base_input.clone()
+            inp_flat_ar = inp_ar.view(-1)
+
+            # Measure latency
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            dist.all_reduce(inp_flat_ar, op=dist.ReduceOp.SUM)
+            torch.cuda.synchronize()
+            end = time.perf_counter()
+            latencies_ar.append(end - start)
+
+            # Store checksum and first values (like test_ar.py)
+            checksum = inp_flat_ar.sum().item()
+            first_vals = inp_flat_ar[:5].clone()
+            results_ar.append((checksum, first_vals))
+
+        # Test reduce-scatter + all-gather determinism
+        results_rs_ag = []
+        latencies_rs_ag = []
+        for trial in range(num_trials):
+            # Clone the same input for each trial
+            inp_rs_ag = base_input.clone()
+            inp_flat_rs_ag = inp_rs_ag.view(-1)
+
+            # Measure latency
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            reduce_scatter_then_all_gather(
+                inp_flat_rs_ag, rank, world_size, custom_ar=None
+            )
+            torch.cuda.synchronize()
+            end = time.perf_counter()
+            latencies_rs_ag.append(end - start)
+
+            # Store checksum and first values (like test_ar.py)
+            checksum = inp_flat_rs_ag.sum().item()
+            first_vals = inp_flat_rs_ag[:5].clone()
+            results_rs_ag.append((checksum, first_vals))
+
+        # Note: sglang's optimized all-reduce requires full runtime initialization
+        # and is not tested in this standalone benchmark
+        use_sglang_optimized = False
+        results_optimized_rs_ag = []
+        latencies_optimized_rs_ag = []
+
+        # Test custom all-reduce determinism (if available)
+        results_custom_ar = []
+        latencies_custom_ar = []
+        if custom_ar is not None:
+            for trial in range(num_trials):
+                # Clone the same input for each trial
+                inp_custom = base_input.clone()
+                inp_flat_custom = inp_custom.view(-1)
+
+                # Measure latency
+                torch.cuda.synchronize()
+                start = time.perf_counter()
+                reduce_scatter_then_all_gather(
+                    inp_flat_custom, rank, world_size, custom_ar=custom_ar
+                )
+                torch.cuda.synchronize()
+                end = time.perf_counter()
+                latencies_custom_ar.append(end - start)
+
+                # Store checksum and first values (like test_ar.py)
+                checksum = inp_flat_custom.sum().item()
+                first_vals = inp_flat_custom[:5].clone()
+                results_custom_ar.append((checksum, first_vals))
+
+        # Test deterministic kernel (if available)
+        results_deterministic_kernel = []
+        latencies_deterministic_kernel = []
+        deterministic_kernel_available = False
+        if custom_ar is not None:
+            # Check if input size fits in buffer
+            input_size_bytes = base_input.numel() * base_input.element_size()
+            if input_size_bytes > custom_ar.max_size:
+                if rank == 0:
+                    print(
+                        f"  Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
+                    )
+                deterministic_kernel_available = False
+            else:
+                try:
+                    deterministic_kernel_available = True
+                    for trial in range(num_trials):
+                        # Clone the same input for each trial
+                        inp_kernel = base_input.clone()
+
+                        # Measure latency
+                        torch.cuda.synchronize()
+                        start = time.perf_counter()
+                        result_kernel = custom_ar.custom_all_reduce(inp_kernel)
+                        torch.cuda.synchronize()
+                        end = time.perf_counter()
+                        latencies_deterministic_kernel.append(end - start)
+
+                        # Store checksum and first values
+                        result_flat_kernel = result_kernel.view(-1)
+                        checksum = result_flat_kernel.sum().item()
+                        first_vals = result_flat_kernel[:5].clone()
+                        results_deterministic_kernel.append((checksum, first_vals))
+                except Exception as e:
+                    if rank == 0:
+                        print(
+                            f"  Deterministic kernel test failed for batch size {bs}: {e}"
+                        )
+                    deterministic_kernel_available = False
+
+        dist.barrier()
+
+        if rank == 0:
+            # Check determinism for all-reduce
+            ar_deterministic = True
+            ar_ref_sum, ar_ref_vals = results_ar[0]
+            ar_variance = []
+            for i, (s, vals) in enumerate(results_ar[1:], 1):
+                if abs(ar_ref_sum - s) > 1e-3 or not torch.allclose(
+                    ar_ref_vals, vals, rtol=1e-3
+                ):
+                    ar_deterministic = False
+                ar_variance.append(abs(ar_ref_sum - s))
+
+            # Check determinism for reduce-scatter + all-gather
+            rs_ag_deterministic = True
+            rs_ag_ref_sum, rs_ag_ref_vals = results_rs_ag[0]
+            rs_ag_variance = []
+            for i, (s, vals) in enumerate(results_rs_ag[1:], 1):
+                if abs(rs_ag_ref_sum - s) > 1e-3 or not torch.allclose(
+                    rs_ag_ref_vals, vals, rtol=1e-3
+                ):
+                    rs_ag_deterministic = False
+                rs_ag_variance.append(abs(rs_ag_ref_sum - s))
+
+            # Check determinism for optimized RS+AG (if available)
+            optimized_rs_ag_deterministic = None
+            optimized_rs_ag_max_variance = None
+            lat_optimized_rs_ag_median = None
+            if use_sglang_optimized and results_optimized_rs_ag:
+                optimized_rs_ag_deterministic = True
+                opt_rs_ag_ref_sum, opt_rs_ag_ref_vals = results_optimized_rs_ag[0]
+                opt_rs_ag_variance = []
+                for i, (s, vals) in enumerate(results_optimized_rs_ag[1:], 1):
+                    if abs(opt_rs_ag_ref_sum - s) > 1e-3 or not torch.allclose(
+                        opt_rs_ag_ref_vals, vals, rtol=1e-3
+                    ):
+                        optimized_rs_ag_deterministic = False
+                    opt_rs_ag_variance.append(abs(opt_rs_ag_ref_sum - s))
+                optimized_rs_ag_max_variance = (
+                    max(opt_rs_ag_variance) if opt_rs_ag_variance else 0.0
+                )
+                lat_optimized_rs_ag_median = statistics.median(
+                    latencies_optimized_rs_ag
+                )
+
+            # Check determinism for custom all-reduce (if available)
+            custom_ar_deterministic = None
+            custom_ar_max_variance = None
+            lat_custom_ar_median = None
+            if custom_ar is not None and results_custom_ar:
+                custom_ar_deterministic = True
+                custom_ar_ref_sum, custom_ar_ref_vals = results_custom_ar[0]
+                custom_ar_variance = []
+                for i, (s, vals) in enumerate(results_custom_ar[1:], 1):
+                    if abs(custom_ar_ref_sum - s) > 1e-3 or not torch.allclose(
+                        custom_ar_ref_vals, vals, rtol=1e-3
+                    ):
+                        custom_ar_deterministic = False
+                    custom_ar_variance.append(abs(custom_ar_ref_sum - s))
+                custom_ar_max_variance = (
+                    max(custom_ar_variance) if custom_ar_variance else 0.0
+                )
+                lat_custom_ar_median = statistics.median(latencies_custom_ar)
+
+            # Check determinism for deterministic kernel (if available)
+            deterministic_kernel_deterministic = None
+            deterministic_kernel_max_variance = None
+            lat_deterministic_kernel_median = None
+            if deterministic_kernel_available and results_deterministic_kernel:
+                deterministic_kernel_deterministic = True
+                kernel_ref_sum, kernel_ref_vals = results_deterministic_kernel[0]
+                kernel_variance = []
+                for i, (s, vals) in enumerate(results_deterministic_kernel[1:], 1):
+                    if abs(kernel_ref_sum - s) > 1e-3 or not torch.allclose(
+                        kernel_ref_vals, vals, rtol=1e-3
+                    ):
+                        deterministic_kernel_deterministic = False
+                    kernel_variance.append(abs(kernel_ref_sum - s))
+                deterministic_kernel_max_variance = (
+                    max(kernel_variance) if kernel_variance else 0.0
+                )
+                lat_deterministic_kernel_median = statistics.median(
+                    latencies_deterministic_kernel
+                )
+
+            # Calculate latency statistics
+            lat_ar_median = statistics.median(latencies_ar)
+            lat_rs_ag_median = statistics.median(latencies_rs_ag)
+            overhead_rs_ag = ((lat_rs_ag_median - lat_ar_median) / lat_ar_median) * 100
+
+            # Calculate variance statistics
+            ar_max_variance = max(ar_variance) if ar_variance else 0.0
+            rs_ag_max_variance = max(rs_ag_variance) if rs_ag_variance else 0.0
+
+            results[bs] = {
+                "all_reduce": {
+                    "latency_median": lat_ar_median,
+                    "deterministic": ar_deterministic,
+                    "max_variance": ar_max_variance,
+                },
+                "rs_ag": {
+                    "latency_median": lat_rs_ag_median,
+                    "deterministic": rs_ag_deterministic,
+                    "max_variance": rs_ag_max_variance,
+                },
+                "custom_ar": (
+                    {
+                        "latency_median": lat_custom_ar_median,
+                        "deterministic": custom_ar_deterministic,
+                        "max_variance": custom_ar_max_variance,
+                    }
+                    if custom_ar is not None
+                    else None
+                ),
+                "deterministic_kernel": (
+                    {
+                        "latency_median": lat_deterministic_kernel_median,
+                        "deterministic": deterministic_kernel_deterministic,
+                        "max_variance": deterministic_kernel_max_variance,
+                    }
+                    if lat_deterministic_kernel_median is not None
+                    else None
+                ),
+                "optimized_rs_ag": (
+                    {
+                        "latency_median": lat_optimized_rs_ag_median,
+                        "deterministic": optimized_rs_ag_deterministic,
+                        "max_variance": optimized_rs_ag_max_variance,
+                    }
+                    if lat_optimized_rs_ag_median is not None
+                    else None
+                ),
+                "overhead_rs_ag_pct": overhead_rs_ag,
+            }
+
+            print(
+                f"    All-Reduce:     {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}"
+            )
+            print(
+                f"    RS+All-Gather:   {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}"
+            )
+            if custom_ar is not None and lat_custom_ar_median is not None:
+                overhead_custom = (
+                    (lat_custom_ar_median - lat_ar_median) / lat_ar_median
+                ) * 100
+                print(
+                    f"    Custom AR:       {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%"
+                )
+            if lat_deterministic_kernel_median is not None:
+                overhead_kernel = (
+                    (lat_deterministic_kernel_median - lat_ar_median) / lat_ar_median
+                ) * 100
+                speedup_kernel_vs_rs_ag = (
+                    (lat_rs_ag_median - lat_deterministic_kernel_median)
+                    / lat_rs_ag_median
+                ) * 100
+                print(
+                    f"    Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%"
+                )
+            if lat_optimized_rs_ag_median is not None:
+                overhead_opt = (
+                    (lat_optimized_rs_ag_median - lat_ar_median) / lat_ar_median
+                ) * 100
+                speedup_vs_rs_ag = (
+                    (lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median
+                ) * 100
+                print(
+                    f"    Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%"
+                )
+            print(f"    RS+AG Overhead: {overhead_rs_ag:+.1f}%")
+
+    if rank == 0:
+        results_queue.put(results)
+
+    dist.destroy_process_group()
+
+
+def main():
+    world_size = 8
+    available_gpus = torch.cuda.device_count()
+
+    print("=" * 80)
+    print("All-Reduce vs Reduce-Scatter + All-Gather Determinism & Latency Benchmark")
+    print("=" * 80)
+    print(f"Available GPUs: {available_gpus}")
+    print(f"Using world_size: {world_size}")
+    print(f"Hidden dimension: 16384")
+    print(f"Tensor dtype: bfloat16")
+    print(f"Trials per batch size: 10 (testing determinism)")
+    print(f"Testing batch sizes: [1, 4, 8, 16, 32, 64, 128, 256, 512]")
+    print("=" * 80)
+
+    if available_gpus < world_size:
+        print(
+            f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead"
+        )
+        world_size = available_gpus
+
+    if world_size < 2:
+        print("ERROR: Need at least 2 GPUs for this benchmark")
+        return
+
+    mp.set_start_method("spawn", force=True)
+    port = get_open_port()
+
+    results_queue = mp.Queue()
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(target=worker, args=(world_size, rank, port, results_queue))
+        p.start()
+        procs.append(p)
+
+    for p in procs:
+        p.join()
+
+    # Collect results
+    if not results_queue.empty():
+        results = results_queue.get()
+
+        print("\n" + "=" * 80)
+        print("SUMMARY")
+        print("=" * 80)
+        header = f"{'Batch':<8} {'AR (ms)':<12} {'AR Det':<8} {'RS+AG (ms)':<15} {'RS+AG Det':<10} {'RS+AG Ovh':<12}"
+        if any(r.get("custom_ar") is not None for r in results.values()):
+            header += (
+                f" {'Custom AR (ms)':<18} {'Custom AR Det':<15} {'Custom AR Ovh':<15}"
+            )
+        if any(r.get("deterministic_kernel") is not None for r in results.values()):
+            header += f" {'Det Kernel (ms)':<18} {'Det Kernel Det':<15} {'Det Kernel Ovh':<15} {'Speedup':<10}"
+        if any(r.get("optimized_rs_ag") is not None for r in results.values()):
+            header += f" {'Opt RS+AG (ms)':<18} {'Opt RS+AG Det':<15} {'Opt RS+AG Ovh':<15} {'Speedup':<10}"
+        print(header)
+        print("-" * 150)
+
+        for bs in sorted(results.keys()):
+            r = results[bs]
+            ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗"
+            rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗"
+            line = (
+                f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} "
+                f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} "
+                f"{r['overhead_rs_ag_pct']:<12.1f}"
+            )
+            if r.get("custom_ar") is not None:
+                custom_ar = r["custom_ar"]
+                custom_ar_det_str = "✓" if custom_ar["deterministic"] else "✗"
+                custom_ar_overhead = (
+                    (custom_ar["latency_median"] - r["all_reduce"]["latency_median"])
+                    / r["all_reduce"]["latency_median"]
+                ) * 100
+                line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}"
+            if r.get("deterministic_kernel") is not None:
+                det_kernel = r["deterministic_kernel"]
+                det_kernel_det_str = "✓" if det_kernel["deterministic"] else "✗"
+                det_kernel_overhead = (
+                    (det_kernel["latency_median"] - r["all_reduce"]["latency_median"])
+                    / r["all_reduce"]["latency_median"]
+                ) * 100
+                speedup_kernel = (
+                    (r["rs_ag"]["latency_median"] - det_kernel["latency_median"])
+                    / r["rs_ag"]["latency_median"]
+                ) * 100
+                line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}"
+            if r.get("optimized_rs_ag") is not None:
+                opt_rs_ag = r["optimized_rs_ag"]
+                opt_rs_ag_det_str = "✓" if opt_rs_ag["deterministic"] else "✗"
+                opt_rs_ag_overhead = (
+                    (opt_rs_ag["latency_median"] - r["all_reduce"]["latency_median"])
+                    / r["all_reduce"]["latency_median"]
+                ) * 100
+                speedup = (
+                    (r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"])
+                    / r["rs_ag"]["latency_median"]
+                ) * 100
+                line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}"
+            print(line)
+
+        print("=" * 80)
+
+        # Calculate statistics
+        overheads_rs_ag = [r["overhead_rs_ag_pct"] for r in results.values()]
+        ar_deterministic_count = sum(
+            1 for r in results.values() if r["all_reduce"]["deterministic"]
+        )
+        rs_ag_deterministic_count = sum(
+            1 for r in results.values() if r["rs_ag"]["deterministic"]
+        )
+        custom_ar_deterministic_count = sum(
+            1
+            for r in results.values()
+            if r.get("custom_ar") and r["custom_ar"]["deterministic"]
+        )
+        custom_ar_total_count = sum(
+            1 for r in results.values() if r.get("custom_ar") is not None
+        )
+
+        deterministic_kernel_deterministic_count = sum(
+            1
+            for r in results.values()
+            if r.get("deterministic_kernel")
+            and r["deterministic_kernel"]["deterministic"]
+        )
+        deterministic_kernel_total_count = sum(
+            1 for r in results.values() if r.get("deterministic_kernel") is not None
+        )
+
+        print(f"\nDeterminism Summary:")
+        print(
+            f"  All-Reduce deterministic: {ar_deterministic_count}/{len(results)} batch sizes"
+        )
+        print(
+            f"  RS+All-Gather deterministic: {rs_ag_deterministic_count}/{len(results)} batch sizes"
+        )
+        if custom_ar_total_count > 0:
+            print(
+                f"  Custom AR deterministic: {custom_ar_deterministic_count}/{custom_ar_total_count} batch sizes"
+            )
+        if deterministic_kernel_total_count > 0:
+            print(
+                f"  Deterministic Kernel deterministic: {deterministic_kernel_deterministic_count}/{deterministic_kernel_total_count} batch sizes"
+            )
+
+        print(f"\nLatency Overhead Statistics (RS+AG vs All-Reduce):")
+        avg_overhead = statistics.mean(overheads_rs_ag)
+        median_overhead = statistics.median(overheads_rs_ag)
+        min_overhead = min(overheads_rs_ag)
+        max_overhead = max(overheads_rs_ag)
+        print(f"  Average: {avg_overhead:.1f}%")
+        print(f"  Median:  {median_overhead:.1f}%")
+        print(f"  Min:     {min_overhead:.1f}%")
+        print(f"  Max:     {max_overhead:.1f}%")
+
+        if custom_ar_total_count > 0:
+            overheads_custom = []
+            for r in results.values():
+                if r.get("custom_ar") is not None:
+                    overhead = (
+                        (
+                            r["custom_ar"]["latency_median"]
+                            - r["all_reduce"]["latency_median"]
+                        )
+                        / r["all_reduce"]["latency_median"]
+                    ) * 100
+                    overheads_custom.append(overhead)
+            print(f"\nLatency Overhead Statistics (Custom AR vs All-Reduce):")
+            print(f"  Average: {statistics.mean(overheads_custom):.1f}%")
+            print(f"  Median:  {statistics.median(overheads_custom):.1f}%")
+            print(f"  Min:     {min(overheads_custom):.1f}%")
+            print(f"  Max:     {max(overheads_custom):.1f}%")
+
+        if deterministic_kernel_total_count > 0:
+            overheads_kernel = []
+            speedups_kernel = []
+            for r in results.values():
+                if r.get("deterministic_kernel") is not None:
+                    overhead = (
+                        (
+                            r["deterministic_kernel"]["latency_median"]
+                            - r["all_reduce"]["latency_median"]
+                        )
+                        / r["all_reduce"]["latency_median"]
+                    ) * 100
+                    overheads_kernel.append(overhead)
+                    speedup = (
+                        (
+                            r["rs_ag"]["latency_median"]
+                            - r["deterministic_kernel"]["latency_median"]
+                        )
+                        / r["rs_ag"]["latency_median"]
+                    ) * 100
+                    speedups_kernel.append(speedup)
+            print(
+                f"\nLatency Overhead Statistics (Deterministic Kernel vs All-Reduce):"
+            )
+            print(f"  Average: {statistics.mean(overheads_kernel):.1f}%")
+            print(f"  Median:  {statistics.median(overheads_kernel):.1f}%")
+            print(f"  Min:     {min(overheads_kernel):.1f}%")
+            print(f"  Max:     {max(overheads_kernel):.1f}%")
+            print(f"\nSpeedup Statistics (Deterministic Kernel vs RS+AG):")
+            print(f"  Average: {statistics.mean(speedups_kernel):.1f}%")
+            print(f"  Median:  {statistics.median(speedups_kernel):.1f}%")
+            print(f"  Min:     {min(speedups_kernel):.1f}%")
+            print(f"  Max:     {max(speedups_kernel):.1f}%")
+
+        # Show variance for non-deterministic cases
+        print(f"\nVariance Analysis (non-deterministic cases):")
+        for bs in sorted(results.keys()):
+            r = results[bs]
+            if not r["all_reduce"]["deterministic"]:
+                print(
+                    f"  Batch {bs}: All-Reduce max variance: {r['all_reduce']['max_variance']:.6f}"
+                )
+            if not r["rs_ag"]["deterministic"]:
+                print(
+                    f"  Batch {bs}: RS+All-Gather max variance: {r['rs_ag']['max_variance']:.6f}"
+                )
+            if r.get("custom_ar") is not None and not r["custom_ar"]["deterministic"]:
+                print(
+                    f"  Batch {bs}: Custom AR max variance: {r['custom_ar']['max_variance']:.6f}"
+                )
+            if (
+                r.get("deterministic_kernel") is not None
+                and not r["deterministic_kernel"]["deterministic"]
+            ):
+                print(
+                    f"  Batch {bs}: Deterministic Kernel max variance: {r['deterministic_kernel']['max_variance']:.6f}"
+                )
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/sgl-kernel/benchmark/bench_awq_dequant.py

@@ -0,0 +1,151 @@
+import itertools
+import os
+from typing import List, Tuple
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel import awq_dequantize
+
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm import _custom_ops as ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    ops = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+
+def vllm_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if not VLLM_AVAILABLE:
+        # Fallback to SGLang implementation
+        return sglang_awq_dequantize(qweight, scales, qzeros)
+    return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+
+
+def sglang_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    return awq_dequantize(qweight, scales, qzeros)
+
+
+def calculate_diff(qweight_row: int, qweight_col: int):
+    """Calculate difference between VLLM and SGLang implementations."""
+    device = torch.device("cuda")
+    qweight = torch.randint(
+        0,
+        torch.iinfo(torch.int32).max,
+        (qweight_row, qweight_col),
+        dtype=torch.int32,
+        device=device,
+    )
+    group_size = qweight_row
+    scales_row = qweight_row // group_size
+    scales_col = qweight_col * 8
+    scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
+    qzeros = torch.randint(
+        0,
+        torch.iinfo(torch.int32).max,
+        (scales_row, qweight_col),
+        dtype=torch.int32,
+        device=device,
+    )
+
+    if not VLLM_AVAILABLE:
+        print("⚠️ vLLM not available, skipping comparison")
+        return
+
+    vllm_out = vllm_awq_dequantize(qweight, scales, qzeros)
+    sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
+
+    output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+
+    if torch.allclose(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ):
+        print("✅ All implementations match")
+    else:
+        print("❌ Implementations differ")
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    qweight_row_range = [128]  # Single row size for CI
+    qweight_cols_range = [16]  # Single column size for CI
+else:
+    qweight_row_range = [3584, 18944, 128, 256, 512, 1024]
+    qweight_cols_range = [448, 576, 4736, 16, 32, 64, 128]
+
+configs = list(itertools.product(qweight_row_range, qweight_cols_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["qweight_row", "qweight_col"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["vllm", "sglang"] if VLLM_AVAILABLE else ["sglang"],
+        line_names=["VLLM", "SGL Kernel"] if VLLM_AVAILABLE else ["SGL Kernel"],
+        styles=[("blue", "-"), ("green", "-")] if VLLM_AVAILABLE else [("green", "-")],
+        ylabel="us",
+        plot_name="awq-dequantize-performance",
+        args={},
+    )
+)
+def benchmark(qweight_row, qweight_col, provider):
+    dtype = torch.float16
+    device = torch.device("cuda")
+    qweight = torch.randint(
+        0,
+        torch.iinfo(torch.int32).max,
+        (qweight_row, qweight_col),
+        dtype=torch.int32,
+        device=device,
+    )
+    group_size = qweight_row
+    scales_row = qweight_row // group_size
+    scales_col = qweight_col * 8
+    scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
+    qzeros = torch.randint(
+        0,
+        torch.iinfo(torch.int32).max,
+        (scales_row, qweight_col),
+        dtype=torch.int32,
+        device=device,
+    )
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "vllm":
+        if not VLLM_AVAILABLE:
+            return (0, 0, 0)
+        fn = lambda: vllm_awq_dequantize(
+            qweight.clone(), scales.clone(), qzeros.clone()
+        )
+    elif provider == "sglang":
+        fn = lambda: sglang_awq_dequantize(
+            qweight.clone(), scales.clone(), qzeros.clone()
+        )
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    # Simplify for CI environment
+    if IS_CI:
+        qweight_row, qweight_col = 128, 16  # Smaller values for CI
+    else:
+        qweight_row, qweight_col = 3584, 448
+
+    calculate_diff(qweight_row=qweight_row, qweight_col=qweight_col)
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_cutlass_mla.py

@@ -0,0 +1,175 @@
+import argparse
+import copy
+import itertools
+import os
+
+import torch
+import triton
+from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size
+
+from sglang.srt.utils import get_device_capability
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+# CI environment uses simplified parameters
+if IS_CI:
+    bs_range = [1]  # Single batch size for CI
+    qlen_range = [64]  # Single sequence length for CI
+else:
+    bs_range = [1, 8, 32, 64, 128, 256]
+    qlen_range = [1, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
+
+configs = list(itertools.product(bs_range, qlen_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len"],
+        x_vals=configs,
+        x_log=False,
+        line_arg="provider",
+        line_vals=[
+            "128 heads",
+            "64 heads",
+            "32 heads",
+            "16 heads",
+        ],
+        line_names=[
+            "128 heads",
+            "64 heads",
+            "32 heads",
+            "16 heads",
+        ],
+        styles=[("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")],
+        ylabel="GB/s",
+        plot_name="cutlass mla",
+        args={},
+    )
+)
+def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits):
+    d = 576
+    dn = 64
+    dv = 512
+
+    h_q_map = {
+        "128": 128,
+        "64": 64,
+        "32": 32,
+        "16": 16,
+    }
+    parsed_h_q = next(
+        (value for key, value in h_q_map.items() if key in provider), None
+    )
+
+    if parsed_h_q is None:
+        raise ValueError(f"Unknown head configuration in provider: {provider}")
+    h_q = parsed_h_q
+
+    seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
+    max_seq_len = seq_lens.max().item()
+    block_num = (max_seq_len + block_size - 1) // block_size
+
+    # Pad block_num so that small blocks can be packed into full 128-sized CUTLASS tiles.
+    # One 128-wide tile can hold (128 // block_size) small blocks.
+    pack_factor = 128 // block_size
+    block_num = ((block_num + pack_factor - 1) // pack_factor) * pack_factor
+
+    qn = (
+        torch.randn(h_q, batch_size, d - dn, dtype=torch.bfloat16, device="cuda")
+        * 100.0
+    )
+    qr = torch.randn(batch_size, h_q, dn, dtype=torch.bfloat16, device="cuda") * 100.0
+    block_table = torch.randint(
+        0,
+        batch_size * block_num,
+        (batch_size, block_num),
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    kv_cache = torch.randn(
+        block_table.numel(), block_size, d, dtype=torch.bfloat16, device="cuda"
+    )
+
+    workspace_size = cutlass_mla_get_workspace_size(
+        block_num * block_size, batch_size, num_kv_splits=num_kv_splits
+    )
+    workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8)
+
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+        lambda: cutlass_mla_decode(
+            qn.transpose(0, 1),
+            qr,
+            kv_cache,
+            seq_lens,
+            block_table,
+            workspace,
+            1.44,
+            num_kv_splits,
+        ),
+        quantiles=quantiles,
+    )
+
+    q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size()
+
+    gbps = (
+        lambda ms: (
+            q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()
+        )
+        * 1e-9
+        / (ms * 1e-3)
+    )
+    return gbps(ms), gbps(max_ms), gbps(min_ms)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--block-sizes",
+        nargs="+",
+        type=int,
+        default=[1, 32, 64, 128],
+        help="List of batch sizes",
+    )
+    parser.add_argument(
+        "--num-kv-splits",
+        nargs="+",
+        type=int,
+        default=[-1],
+        help="List of batch sizes",
+    )
+    args = parser.parse_args()
+
+    # Skip in CI environment or unsupported architectures
+    if IS_CI:
+        major, minor = get_device_capability()
+        if major is None or major < 10:  # Requires compute capability 10.0+
+            print("Skipping Cutlass MLA benchmark in CI environment")
+            if major is not None:
+                print(
+                    f"Cutlass MLA requires compute capability 10.0+, but found {major}.{minor}"
+                )
+            else:
+                print("Could not determine device capability")
+        else:
+            for block_size in args.block_sizes:
+                for kv_split in args.num_kv_splits:
+                    print(f"block_size={block_size}, num_kv_splits={kv_split}: ")
+                    benchmark.run(
+                        print_data=True,
+                        block_size=block_size,
+                        num_kv_splits=kv_split,
+                    )
+            print("Benchmark finished!")
+    else:
+        for block_size in args.block_sizes:
+            for kv_split in args.num_kv_splits:
+                print(f"block_size={block_size}, num_kv_splits={kv_split}: ")
+                benchmark.run(
+                    print_data=True,
+                    block_size=block_size,
+                    num_kv_splits=kv_split,
+                )
+        print("Benchmark finished!")

third_party/sglang/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py

@@ -0,0 +1,73 @@
+import argparse
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.testing
+from sgl_kernel import dsv3_fused_a_gemm
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+# CI environment uses simplified parameters
+if IS_CI:
+    num_tokens_vals = [1]  # Only test 1 value in CI
+    line_vals = ["sgl-kernel"]  # Only test sgl-kernel implementation in CI
+else:
+    num_tokens_vals = [i + 1 for i in range(16)]  # Test 1-16 in full mode
+    line_vals = ["torch", "sgl-kernel"]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens"],
+        x_vals=num_tokens_vals,
+        x_log=False,
+        line_arg="impl",
+        line_vals=line_vals,
+        line_names=(
+            ["torch (bf16)", "dsv3_fused_a_gemm"]
+            if not IS_CI
+            else ["dsv3_fused_a_gemm"]
+        ),
+        styles=[("blue", "-"), ("orange", "-")] if not IS_CI else [("orange", "-")],
+        ylabel="TFLOPs",
+        plot_name="bf16 dsv3 fused a GEMM throughput",
+        args={},
+    )
+)
+def benchmark(num_tokens, impl):
+    kHdIn = 7168
+    kHdOut = 2112
+    M, K, N = num_tokens, kHdIn, kHdOut
+
+    mat_a = torch.randn((M, K), dtype=torch.bfloat16, device="cuda").contiguous()
+    mat_b = torch.randn((N, K), dtype=torch.bfloat16, device="cuda").transpose(0, 1)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if impl == "torch":
+
+        def runner():
+            F.linear(mat_a, mat_b.T)
+
+    elif impl == "sgl-kernel":
+
+        def runner():
+            dsv3_fused_a_gemm(mat_a, mat_b)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(runner, quantiles=quantiles)
+
+    def tflops(t_ms):
+        flops = 2 * M * K * N
+        return flops / (t_ms * 1e-3) / 1e12
+
+    return tflops(ms), tflops(max_ms), tflops(min_ms)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    args = parser.parse_args()
+
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_dsv3_router_gemm.py

@@ -0,0 +1,151 @@
+import argparse
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.testing
+from sgl_kernel import dsv3_router_gemm
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+# CI environment uses simplified parameters
+if IS_CI:
+    num_tokens_vals = [1]  # Only test 1 value in CI
+    line_vals = ["sgl-kernel-256"]  # Only test one implementation in CI
+else:
+    num_tokens_vals = [i + 1 for i in range(16)]  # Test 1-16 in full mode
+    line_vals = ["torch-256", "sgl-kernel-256", "torch-384", "sgl-kernel-384"]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens"],
+        x_vals=num_tokens_vals,
+        x_log=False,
+        line_arg="impl",
+        line_vals=line_vals,
+        line_names=(
+            [
+                "torch-256",
+                "dsv3_router_gemm-256",
+                "torch-384",
+                "dsv3_router_gemm-384",
+            ]
+            if not IS_CI
+            else ["dsv3_router_gemm-256"]
+        ),
+        styles=(
+            [("blue", "-"), ("orange", "-"), ("green", "-"), ("red", "-")]
+            if not IS_CI
+            else [("orange", "-")]
+        ),
+        ylabel="TFLOPs",
+        plot_name="input-bf16-output-bf16 dsv3 router gemm throughput",
+        args={},
+    )
+)
+def benchmark_bf16_output(num_tokens, impl):
+    # M: num_tokens, K: hidden_dim, N: num_experts
+    M, K = num_tokens, 7168
+
+    if impl == "torch-256" or impl == "sgl-kernel-256":
+        N = 256
+    elif impl == "torch-384" or impl == "sgl-kernel-384":
+        N = 384
+    else:
+        raise ValueError(f"Unknown impl: {impl}")
+
+    mat_a = torch.randn((M, K), dtype=torch.bfloat16, device="cuda").contiguous()
+    mat_b = torch.randn((N, K), dtype=torch.bfloat16, device="cuda").contiguous()
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if impl == "torch-256" or impl == "torch-384":
+
+        def runner():
+            F.linear(mat_a, mat_b)
+
+    elif impl == "sgl-kernel-256" or impl == "sgl-kernel-384":
+
+        def runner():
+            dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.bfloat16)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(runner, quantiles=quantiles)
+
+    def tflops(t_ms):
+        flops = 2 * M * K * N
+        return flops / (t_ms * 1e-3) / 1e12
+
+    return tflops(ms), tflops(max_ms), tflops(min_ms)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens"],
+        x_vals=num_tokens_vals,
+        x_log=False,
+        line_arg="impl",
+        line_vals=line_vals,
+        line_names=(
+            [
+                "torch-256",
+                "dsv3_router_gemm-256",
+                "torch-384",
+                "dsv3_router_gemm-384",
+            ]
+            if not IS_CI
+            else ["dsv3_router_gemm-256"]
+        ),
+        styles=(
+            [("blue", "-"), ("orange", "-"), ("green", "-"), ("red", "-")]
+            if not IS_CI
+            else [("orange", "-")]
+        ),
+        ylabel="TFLOPs",
+        plot_name="input-bf16-output-fp32 dsv3 router gemm throughput",
+        args={},
+    )
+)
+def benchmark_float_output(num_tokens, impl):
+    # M: num_tokens, K: hidden_dim, N: num_experts
+    M, K = num_tokens, 7168
+
+    if impl == "torch-256" or impl == "sgl-kernel-256":
+        N = 256
+    elif impl == "torch-384" or impl == "sgl-kernel-384":
+        N = 384
+    else:
+        raise ValueError(f"Unknown impl: {impl}")
+
+    mat_a = torch.randn((M, K), dtype=torch.bfloat16, device="cuda").contiguous()
+    mat_b = torch.randn((N, K), dtype=torch.bfloat16, device="cuda").contiguous()
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if impl == "torch-256" or impl == "torch-384":
+
+        def runner():
+            F.linear(mat_a, mat_b).to(torch.float32)
+
+    elif impl == "sgl-kernel-256" or impl == "sgl-kernel-384":
+
+        def runner():
+            dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.float32)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(runner, quantiles=quantiles)
+
+    def tflops(t_ms):
+        flops = 2 * M * K * N
+        return flops / (t_ms * 1e-3) / 1e12
+
+    return tflops(ms), tflops(max_ms), tflops(min_ms)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    args = parser.parse_args()
+
+    benchmark_bf16_output.run(print_data=True)
+    benchmark_float_output.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py

@@ -0,0 +1,373 @@
+import argparse
+import random
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import numpy as np
+import torch
+from sgl_kernel import (
+    es_fp8_blockwise_scaled_grouped_mm,
+    fp8_blockwise_scaled_grouped_mm,
+)
+
+random.seed(28)
+
+
+def ceil_div(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+
+def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert x.dim() == 2
+    m, n = x.shape
+    pad_size = (128 - (n % 128)) % 128
+    x = torch.nn.functional.pad(x, (0, pad_size), value=0) if pad_size > 0 else x
+    x_view = x.view(m, -1, 128)
+    x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
+    fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
+    return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
+
+
+def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert x.dim() == 2
+    m, n = x.shape
+    x_padded = torch.zeros(
+        (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device
+    )
+    x_padded[:m, :n] = x
+    x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
+    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
+    x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
+    return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(
+        x_view.size(0), x_view.size(2)
+    )
+
+
+def create_unbalanced_expert_token_distribution(
+    batch_size: int, topk: int, num_experts: int
+):
+    expert_ids = np.random.randint(0, num_experts, size=(batch_size * topk,)).tolist()
+    expert_to_count = dict()
+    for expert_id in range(num_experts):
+        expert_to_count[expert_id] = 0
+    for expert_id in expert_ids:
+        expert_to_count[expert_id] += 1
+    group_ms = []
+    for expert_id in range(num_experts):
+        group_ms.append(expert_to_count[expert_id])
+    return group_ms
+
+
+def bench_es(
+    group_ms: List[int],
+    n: int,
+    k: int,
+    num_groups: int,
+    num_warmup: int,
+    num_run: int,
+) -> Tuple[float, int]:
+    device = "cuda"
+    alignment = 128
+    n_g = ceil_div(n, alignment) * alignment
+    k_g = ceil_div(k, alignment) * alignment
+    out_dtype = torch.bfloat16
+
+    expert_offsets = torch.zeros((num_groups + 1), device=device, dtype=torch.int32)
+    problem_sizes = torch.zeros((num_groups, 3), device=device, dtype=torch.int32)
+
+    a_tensors = []
+    b_tensors = []
+    a_scales_tensors = []
+    b_scales_tensors = []
+    if False:
+        print("Token Distributtion: ", group_ms[0:num_groups])
+        print("Token Count: ", sum(group_ms[0:num_groups]))
+    for g in range(num_groups):
+        m_g = group_ms[g]
+        expert_offsets[g + 1] = expert_offsets[g] + m_g
+        problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
+        if m_g != 0:
+            a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
+            a_tensors.append(a_g)
+            a_scales_tensors.append(a_scale)
+
+        b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
+        b_tensors.append(b_g)
+        b_scales_tensors.append(b_scale)
+
+    a_stack = torch.empty(
+        (expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    b_stack = torch.empty(
+        (num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+
+    _aux_idx = 0
+    for g in range(num_groups):
+        if group_ms[g] != 0:
+            a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx]
+            _aux_idx += 1
+        b_stack[g] = b_tensors[g].t()
+    b_stack = b_stack.transpose(1, 2)
+
+    a_scale_stack = torch.empty(
+        (expert_offsets[-1], k_g // 128), device=device, dtype=torch.float32
+    )
+    b_scale_stack = torch.empty(
+        (num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
+    )
+
+    _aux_idx = 0
+    for g in range(num_groups):
+        if group_ms[g] != 0:
+            a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[
+                _aux_idx
+            ]
+            _aux_idx += 1
+        b_scale_stack[g] = b_scales_tensors[g].t()
+    b_scale_stack = b_scale_stack.transpose(1, 2)
+
+    workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
+    c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
+    a_strides = torch.full(
+        (num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
+    )
+    d_strides = torch.full(
+        (num_groups,), c_out.stride(0), device=device, dtype=torch.int64
+    )
+
+    def run_cutlass():
+        es_fp8_blockwise_scaled_grouped_mm(
+            c_out,
+            a_stack,
+            b_stack,
+            a_scale_stack,
+            b_scale_stack,
+            a_strides,
+            a_strides,
+            d_strides,
+            problem_sizes,
+            expert_offsets[:-1],
+            workspace,
+        )
+
+    run_cutlass()
+    # warmup
+    for _ in range(num_warmup):
+        run_cutlass()
+    torch.cuda.synchronize()
+
+    # run
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    start_event.record()
+    for _ in range(num_run):
+        run_cutlass()
+    end_event.record()
+    end_event.synchronize()
+    torch.cuda.synchronize()
+    avg = start_event.elapsed_time(end_event) / num_run * 1000  # us
+
+    return avg, expert_offsets[-1]
+
+
+def bench_sgl(
+    group_ms: List[int],
+    n: int,
+    k: int,
+    num_groups: int,
+    num_warmup: int,
+    num_run: int,
+) -> Tuple[float, int]:
+    device = "cuda"
+    alignment = 128
+    n_g = ceil_div(n, alignment) * alignment
+    k_g = ceil_div(k, alignment) * alignment
+    out_dtype = torch.bfloat16
+
+    expert_offsets = torch.zeros((num_groups + 1), device=device, dtype=torch.int32)
+    problem_sizes = torch.zeros((num_groups, 3), device=device, dtype=torch.int32)
+    layout_sfa = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
+    layout_sfb = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
+
+    a_tensors = []
+    b_tensors = []
+    a_scales_tensors = []
+    b_scales_tensors = []
+
+    for g in range(num_groups):
+        m_g = group_ms[g]
+        expert_offsets[g + 1] = expert_offsets[g] + m_g
+        problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
+        if m_g != 0:
+            a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
+            a_tensors.append(a_g)
+            a_scales_tensors.append(a_scale)
+
+        b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
+        b_tensors.append(b_g)
+        b_scales_tensors.append(b_scale)
+
+    a_stack = torch.empty(
+        (expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    b_stack = torch.empty(
+        (num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+
+    _aux_idx = 0
+    for g in range(num_groups):
+        if group_ms[g] != 0:
+            a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx]
+            _aux_idx += 1
+        b_stack[g] = b_tensors[g].t()
+    b_stack = b_stack.transpose(1, 2)
+
+    a_scale_stack = torch.empty(
+        (expert_offsets[-1], k_g // 128), device=device, dtype=torch.float32
+    )
+    b_scale_stack = torch.empty(
+        (num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
+    )
+
+    _aux_idx = 0
+    for g in range(num_groups):
+        if group_ms[g] != 0:
+            a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[
+                _aux_idx
+            ]
+            _aux_idx += 1
+        b_scale_stack[g] = b_scales_tensors[g].t()
+    b_scale_stack = b_scale_stack.transpose(1, 2)
+
+    c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
+    a_strides = torch.full(
+        (num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
+    )
+    c_strides = torch.full(
+        (num_groups,), c_out.stride(0), device=device, dtype=torch.int64
+    )
+    workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
+    a_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    b_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    out_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    a_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    b_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+
+    def run_cutlass():
+        fp8_blockwise_scaled_grouped_mm(
+            c_out,
+            a_ptrs,
+            b_ptrs,
+            out_ptrs,
+            a_scales_ptrs,
+            b_scales_ptrs,
+            a_stack,
+            b_stack,
+            a_scale_stack,
+            b_scale_stack,
+            a_strides,
+            a_strides,
+            c_strides,
+            layout_sfa,
+            layout_sfb,
+            problem_sizes,
+            expert_offsets[:-1],
+            workspace,
+        )
+
+    # warmup
+    for _ in range(num_warmup):
+        run_cutlass()
+    torch.cuda.synchronize()
+
+    # run
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    start_event.record()
+    for _ in range(num_run):
+        run_cutlass()
+    end_event.record()
+    end_event.synchronize()
+    torch.cuda.synchronize()
+    avg = start_event.elapsed_time(end_event) / num_run * 1000  # us
+
+    return avg, expert_offsets[-1]
+
+
+benchmark_kernels = {"es": bench_es, "sgl-kernel": bench_sgl}
+
+
+@dataclass
+class ShapeArg:
+    n: int
+    k: int
+    num_groups: int
+
+
+def benchmark_one_shape(
+    shape_args: List[ShapeArg],
+    num_warmup: int,
+    num_run: int,
+):
+    for shape in shape_args:
+        for batch_size in [
+            128,
+            256,
+            384,
+            512,
+            640,
+            768,
+            896,
+            1024,
+            1280,
+            1536,
+            2048,
+            3072,
+        ]:
+            group_ms = create_unbalanced_expert_token_distribution(
+                batch_size, 8, shape.num_groups
+            )
+            print(
+                f"\nBenchmark: batch_size={batch_size}, n={shape.n}, k={shape.k}, num_groups={shape.num_groups}"
+            )
+            for kernel_name, kernel_func in benchmark_kernels.items():
+                average_time, m = kernel_func(
+                    group_ms,
+                    shape.n,
+                    shape.k,
+                    shape.num_groups,
+                    num_warmup,
+                    num_run,
+                )
+                print(f"{kernel_name}: {average_time} us")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-warmup", type=int, default=3)
+    parser.add_argument("--num-run", type=int, default=20)
+    shape_args = [
+        # DeepSeek-R1, gateup, TP = 8
+        ShapeArg(n=512, k=7168, num_groups=256),
+        # DeepSeek-R1, down, TP = 8
+        ShapeArg(n=7168, k=256, num_groups=256),
+        # DeepSeek-R1, gateup, TP = 4
+        ShapeArg(n=1024, k=7168, num_groups=256),
+        # DeepSeek-R1, down, TP = 4
+        ShapeArg(n=7168, k=512, num_groups=256),
+        # Qwen3-235B-A22B-FP8, gateup, TP = 4
+        ShapeArg(n=768, k=4096, num_groups=128),
+        # Qwen3-235B-A22B-FP8, down, TP = 4
+        ShapeArg(n=4096, k=384, num_groups=128),
+        # Qwen3-235B-A22B-FP8, gateup, TP = 2
+        ShapeArg(n=1536, k=4096, num_groups=128),
+        # Qwen3-235B-A22B-FP8, down, TP = 2
+        ShapeArg(n=4096, k=768, num_groups=128),
+    ]
+    args = parser.parse_args()
+    benchmark_one_shape(shape_args, args.num_warmup, args.num_run)
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/sgl-kernel/benchmark/bench_fp4_gemm.py

@@ -0,0 +1,389 @@
+import argparse
+import csv
+import os
+from functools import partial
+from typing import List, Tuple
+
+import torch
+import triton
+from flashinfer import mm_fp4
+from flashinfer.testing import bench_gpu_time
+
+from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
+from sglang.srt.utils import (
+    get_device_capability,
+    is_sm100_supported,
+    is_sm120_supported,
+)
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+FLOAT4_E2M1_MAX = 6.0
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+DEEPSEEK_R1_MODEL = "deepseek-ai/DeepSeek-R1-0528-FP4"
+
+# Weight shapes are in the format: ([K, N], TP_SPLIT_DIM)
+# TP split dim 0 means split K by tp size; dim 1 means split N by tp size.
+WEIGHT_SHAPES = {
+    "meta-llama/Llama-3.1-8B-Instruct": [
+        ([4096, 6144], 1),
+        ([4096, 4096], 0),
+        ([4096, 28672], 1),
+        ([14336, 4096], 0),
+    ],
+    "meta-llama/Llama-3.3-70B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 57344], 1),
+        ([28672, 8192], 0),
+    ],
+    "mistralai/Mistral-Large-Instruct-2407": [
+        ([12288, 14336], 1),
+        ([12288, 12288], 0),
+        ([12288, 57344], 1),
+        ([28672, 12288], 0),
+    ],
+    "Qwen/Qwen2.5-7B-Instruct": [
+        ([3584, 4608], 1),
+        ([3584, 3584], 0),
+        ([3584, 37888], 1),
+        ([18944, 3584], 0),
+    ],
+    "Qwen/Qwen2.5-32B-Instruct": [
+        ([5120, 7168], 1),
+        ([5120, 5120], 0),
+        ([5120, 55296], 1),
+        ([27648, 5120], 0),
+    ],
+    "Qwen/Qwen2.5-72B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 59136], 1),
+        ([29568, 8192], 0),
+    ],
+    "Qwen/Qwen3.5-27B": [
+        ([5120, 8192], 1),
+        ([6144, 5120], 0),
+        ([5120, 34816], 1),
+        ([17408, 5120], 0),
+    ],
+    "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
+        ([2048, 3072], 1),
+        ([2048, 4096], 1),
+        ([2048, 2048], 0),
+        ([2048, 576], 0),
+        ([2048, 21888], 1),
+        ([10944, 2048], 0),
+        ([2048, 2816], 1),
+        ([1408, 2048], 0),
+    ],
+}
+
+DEEPSEEK_R1_WEIGHT_SHAPES = {
+    4: [[1024, 3584], [7168, 256], [7168, 2304], [9216, 3584]],
+    8: [[512, 3584], [7168, 128], [7168, 1152], [4608, 3584]],
+}
+
+
+def get_weight_shapes(args) -> List[Tuple[int, int, str]]:
+    shapes: List[Tuple[int, int, str]] = []
+    for model in args.models:
+        if model == DEEPSEEK_R1_MODEL:
+            for tp_size in args.tp_sizes:
+                if tp_size in DEEPSEEK_R1_WEIGHT_SHAPES:
+                    selected = DEEPSEEK_R1_WEIGHT_SHAPES[tp_size]
+                else:
+                    selected = (
+                        DEEPSEEK_R1_WEIGHT_SHAPES[4] + DEEPSEEK_R1_WEIGHT_SHAPES[8]
+                    )
+                for n, packed_k in selected:
+                    shapes.append((n, packed_k, model))
+            continue
+
+        if model not in WEIGHT_SHAPES:
+            raise ValueError(f"Unsupported model: {model}")
+        for tp_size in args.tp_sizes:
+            for k_n, tp_split_dim in WEIGHT_SHAPES[model]:
+                k, n = k_n
+                if tp_split_dim == 0:
+                    k = k // tp_size
+                else:
+                    n = n // tp_size
+                packed_k = k // 2
+                shapes.append((n, packed_k, model))
+    return shapes
+
+
+if IS_CI:
+    batch_sizes = [1, 8]
+else:
+    batch_sizes = [
+        1,
+        2,
+        4,
+        8,
+        16,
+        32,
+        64,
+        128,
+        256,
+        512,
+        1024,
+        2048,
+        3072,
+        4096,
+        8192,
+        16384,
+    ]
+
+
+def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend):
+    return mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend=backend)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=batch_sizes,
+        x_log=False,
+        line_arg="provider",
+        line_vals=(
+            ["sglang_cutlass", "cutlass", "cudnn", "trtllm", "auto"]
+            if is_sm100_supported()
+            else ["sglang_cutlass", "cutlass", "cudnn", "auto"]
+        ),
+        line_names=(
+            [
+                "sglang cutlass fp4",
+                "flashinfer cutlass fp4",
+                "cudnn fp4",
+                "trtllm fp4",
+                "auto fp4 (cudnn/cutlass)",
+            ]
+            if is_sm100_supported()
+            else [
+                "sglang cutlass fp4",
+                "flashinfer cutlass fp4",
+                "cudnn fp4",
+                "auto fp4",
+            ]
+        ),
+        styles=(
+            [
+                ("red", "solid"),
+                ("orange", "solid"),
+                ("blue", "solid"),
+                ("green", "solid"),
+                ("purple", "solid"),
+            ]
+            if is_sm100_supported()
+            else [
+                ("red", "solid"),
+                ("orange", "solid"),
+                ("blue", "solid"),
+                ("purple", "solid"),
+            ]
+        ),
+        ylabel="bandwidth (GB/s)",
+        plot_name="fp4_gemm_benchmark",
+        args={},
+    )
+)
+def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
+    M = batch_size
+    packed_k = K
+    K = 2 * packed_k
+    a_dtype = torch.randn((M, K), dtype=dtype, device="cuda")
+    b_dtype = torch.randn((N, K), dtype=dtype, device="cuda")
+    a_global_scale = (
+        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
+    ).to(torch.float32)
+    b_global_scale = (
+        (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
+    ).to(torch.float32)
+    alpha = 1.0 / (a_global_scale * b_global_scale)
+    a_fp4, a_scale_interleaved = scaled_fp4_quant(a_dtype, a_global_scale)
+    b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale)
+    b_fp4_T = b_fp4.T
+    b_sf_T = b_scale_interleaved.T
+    res_fi = torch.empty((M, N), dtype=dtype, device="cuda")
+
+    if provider == "sglang_cutlass":
+        times_ms = bench_gpu_time(
+            fn=cutlass_scaled_fp4_mm,
+            input_args=(
+                a_fp4,
+                b_fp4,
+                a_scale_interleaved,
+                b_scale_interleaved,
+                alpha,
+                dtype,
+            ),
+            use_cuda_graph=True,
+        )
+    elif provider == "cutlass":
+        times_ms = bench_gpu_time(
+            fn=partial(_run_mm_fp4, backend="cutlass"),
+            input_args=(
+                a_fp4,
+                b_fp4_T,
+                a_scale_interleaved,
+                b_sf_T,
+                alpha,
+                dtype,
+                res_fi,
+            ),
+            use_cuda_graph=True,
+        )
+    elif provider == "cudnn":
+        times_ms = bench_gpu_time(
+            fn=partial(_run_mm_fp4, backend="cudnn"),
+            input_args=(
+                a_fp4,
+                b_fp4_T,
+                a_scale_interleaved,
+                b_sf_T,
+                alpha,
+                dtype,
+                res_fi,
+            ),
+            use_cuda_graph=True,
+        )
+    elif provider == "trtllm":
+        a_sf_u8 = a_scale_interleaved.to(torch.uint8)
+        b_sf_u8_T = b_sf_T.to(torch.uint8)
+        times_ms = bench_gpu_time(
+            fn=partial(_run_mm_fp4, backend="trtllm"),
+            input_args=(a_fp4, b_fp4_T, a_sf_u8, b_sf_u8_T, alpha, dtype, res_fi),
+            use_cuda_graph=True,
+        )
+    elif provider == "auto":
+        times_ms = bench_gpu_time(
+            fn=partial(_run_mm_fp4, backend="auto"),
+            input_args=(
+                a_fp4,
+                b_fp4_T,
+                a_scale_interleaved,
+                b_sf_T,
+                alpha,
+                dtype,
+                res_fi,
+            ),
+            use_cuda_graph=True,
+        )
+
+    ms = torch.tensor(times_ms).median().item()
+
+    # A: M×packed_k bytes (fp4 packed), B: N×packed_k bytes, C: M×N×element_size bytes
+    element_size = torch.finfo(dtype).bits // 8
+    total_bytes = M * packed_k + N * packed_k + M * N * element_size
+    bandwidth_gbs = total_bytes / (ms * 1e-3) / 1e9
+
+    if correctness:
+        res_cutlass = cutlass_scaled_fp4_mm(
+            a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype
+        )
+        mm_fp4(
+            a_fp4,
+            b_fp4_T,
+            a_scale_interleaved,
+            b_sf_T,
+            alpha,
+            dtype,
+            res_fi,
+            backend="cudnn",
+        )
+        assert torch.allclose(
+            res_fi, res_cutlass, atol=1e-3, rtol=1e-3
+        ), "cudnn fp4 doesn't match cutlass fp4"
+        mm_fp4(
+            a_fp4,
+            b_fp4_T,
+            a_scale_interleaved,
+            b_sf_T,
+            alpha,
+            dtype,
+            res_fi,
+            backend="trtllm",
+        )
+        assert torch.allclose(
+            res_fi, res_cutlass, atol=1e-3, rtol=1e-3
+        ), "trtllm fp4 doesn't match cutlass fp4"
+
+    if csv_file:
+        with open(csv_file, "a", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow([provider, M, N, K, ms, bandwidth_gbs])
+
+    return bandwidth_gbs
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        type=str,
+        default=[DEEPSEEK_R1_MODEL],
+        help="List of models to benchmark. Supported: Llama 8B/70B, Qwen, Mistral, DeepSeek.",
+    )
+    parser.add_argument(
+        "--tp-sizes",
+        nargs="+",
+        type=int,
+        default=[1],
+        help="List of tensor parallel sizes",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=torch.dtype,
+        default=torch.bfloat16,
+        help="Output data type",
+    )
+    parser.add_argument(
+        "--correctness",
+        action="store_true",
+        help="Check correctness",
+    )
+    parser.add_argument(
+        "--csv",
+        type=str,
+        default="results_cutlass_cudnn.csv",
+        help="CSV file to save results",
+    )
+    args = parser.parse_args()
+
+    if IS_CI:
+        args.tp_sizes = [args.tp_sizes[0]]
+
+    if args.csv:
+        with open(args.csv, "w", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow(["provider", "m", "n", "k", "time_ms", "bandwidth_gbs"])
+
+    major, minor = get_device_capability()
+    if not (is_sm100_supported() or is_sm120_supported()):
+        print("Skipping FP4 GEMM benchmark")
+        if major is not None:
+            print(f"FP4 operations require sm100+, but found sm{major}{minor}")
+        else:
+            print("Could not determine device capability")
+    else:
+        NKs = get_weight_shapes(args)
+
+        if IS_CI:
+            NKs = NKs[:2]
+
+        for N, K, model_name in NKs:
+            print(f"{model_name} N={N} packed_k={K}: ")
+            benchmark.run(
+                print_data=True,
+                N=N,
+                K=K,
+                dtype=args.dtype,
+                correctness=args.correctness,
+                csv_file=args.csv,
+            )
+        print("Benchmark finished!")

third_party/sglang/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py

@@ -0,0 +1,237 @@
+import argparse
+import copy
+import itertools
+import os
+
+import deep_gemm
+import torch
+import triton
+from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor
+from sgl_kernel import fp8_blockwise_scaled_mm
+
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    vllm_scaled_mm = None
+    VLLM_AVAILABLE = False
+
+from sglang.srt.layers.quantization.fp8_kernel import (
+    w8a8_block_fp8_matmul_triton as w8a8_block_fp8_matmul,
+)
+
+IS_CI = is_in_ci()
+
+
+def get_weight_shapes(args):
+    models_tps = list(itertools.product(args.models, args.tp_sizes))
+    # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model.
+    # cannot TP
+    total = [
+        (512 + 64, 7168),
+        ((128 + 64) * 128, 7168),
+        (128 * (128 + 128), 512),
+        (7168, 16384),
+        (7168, 18432),
+    ]
+    # N can TP
+    n_tp = [
+        (18432 * 2, 7168),
+        ((128 + 64) * 128, 7168),
+        (128 * (128 + 128), 512),
+        (24576, 1536),
+        (4096, 7168),
+    ]
+    # K can TP
+    k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)]
+    # only support Deepseek-V3
+    SUPPORT_MODEL = ["deepseek-ai/DeepSeek-V3"]
+
+    weight_shapes = []
+    for model, tp_size in models_tps:
+        assert model in SUPPORT_MODEL
+        for t in total:
+            new_t = [t[0], t[1], model]
+            weight_shapes.append(new_t)
+        for n_t in n_tp:
+            new_t = [n_t[0] // tp_size, n_t[1], model]
+            weight_shapes.append(new_t)
+        for k_t in k_tp:
+            new_t = [k_t[0], k_t[1] // tp_size, model]
+            weight_shapes.append(new_t)
+    return weight_shapes
+
+
+def cdiv(a: int, b: int) -> int:
+    """Ceiling division."""
+    return -(a // -b)
+
+
+def fp8_gemm_deepgemm(
+    x_fp8: torch.Tensor,
+    x_scale: torch.Tensor,
+    y_fp8: torch.Tensor,
+    y_scale: torch.Tensor,
+    m: int,
+    n: int,
+    k: int,
+):
+    """DeepGEMM implementation of FP8 GEMM"""
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+
+    # Run DeepGEMM kernel
+    deep_gemm.fp8_gemm_nt((x_fp8, x_scale), (y_fp8, y_scale), out)
+    return out
+
+
+def scale_shape(shape, group_shape):
+    assert len(shape) == len(group_shape)
+    return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_sizes = [1, 8]  # Simplified for CI
+else:
+    batch_sizes = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
+
+# Filter providers based on availability
+available_providers = ["sgl-kernel"]
+available_names = ["sgl-kernel"]
+available_styles = [("orange", "-")]
+
+if VLLM_AVAILABLE:
+    available_providers.insert(0, "vllm")
+    available_names.insert(0, "vllm")
+    available_styles.insert(0, ("blue", "-"))
+
+available_providers.append("triton")
+available_names.append("sglang triton")
+available_styles.append(("red", "-"))
+
+# Add deepgemm if available
+try:
+    import deep_gemm
+
+    available_providers.append("deepgemm")
+    available_names.append("deepgemm")
+    available_styles.append(("yellow", "-"))
+except ImportError:
+    pass
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=batch_sizes,
+        x_log=False,
+        line_arg="provider",
+        line_vals=available_providers,
+        line_names=available_names,
+        styles=available_styles,
+        ylabel="GB/s",
+        plot_name="fp8 blockwise scaled matmul",
+        args={},
+    )
+)
+def benchmark(batch_size, provider, N, K):
+    M = batch_size
+    fp8_info = torch.finfo(torch.float8_e4m3fn)
+    fp8_max, fp8_min = fp8_info.max, fp8_info.min
+
+    a_fp32 = (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max
+    a_fp8 = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+    b_fp32 = (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max
+    b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+    scale_a_group_shape = (1, 128)
+    scale_b_group_shape = (128, 128)
+    scale_a_shape = scale_shape(a_fp8.shape, scale_a_group_shape)
+    scale_b_shape = scale_shape(b_fp8.shape, scale_b_group_shape)
+
+    scale_a = torch.randn(scale_a_shape, device="cuda", dtype=torch.float32)
+    scale_b = torch.randn(scale_b_shape, device="cuda", dtype=torch.float32)
+
+    quantiles = [0.5, 0.2, 0.8]
+    if provider == "sgl-kernel":
+        scale_a = scale_a.t().contiguous().t()
+        b_fp8, scale_b = b_fp8.t(), scale_b.t()
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: fp8_blockwise_scaled_mm(
+                a_fp8, b_fp8, scale_a, scale_b, torch.float16
+            ),
+            quantiles=quantiles,
+        )
+    elif provider == "vllm":
+        if not VLLM_AVAILABLE:
+            return (0, 0, 0)
+        scale_a = scale_a.t().contiguous().t()
+        b_fp8, scale_b = b_fp8.t(), scale_b.t()
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, torch.float16),
+            quantiles=quantiles,
+        )
+    elif provider == "triton":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: w8a8_block_fp8_matmul(
+                a_fp8, b_fp8, scale_a, scale_b, [128, 128], torch.float16
+            ),
+            quantiles=quantiles,
+        )
+    if provider == "deepgemm":
+        scale_a_col_major = get_mn_major_tma_aligned_tensor(scale_a.clone())
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: fp8_gemm_deepgemm(
+                a_fp8, scale_a_col_major, b_fp8, scale_b, M, N, K
+            ),
+            quantiles=quantiles,
+        )
+    return ms * 1000, max_ms * 1000, min_ms * 1000  # convert to ms
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        type=str,
+        default=["deepseek-ai/DeepSeek-V3"],
+        help="List of models to benchmark",
+    )
+    parser.add_argument(
+        "--tp-sizes",
+        nargs="+",
+        type=int,
+        default=[1],
+        help="List of tensor parallel sizes",
+    )
+    args = parser.parse_args()
+
+    # Simplify for CI environment
+    if IS_CI:
+        args.models = [args.models[0]]  # Use only first model
+        args.tp_sizes = [args.tp_sizes[0]]  # Use only first TP size
+
+    NK_model_names = get_weight_shapes(args)
+
+    # Limit iterations in CI
+    if IS_CI:
+        NK_model_names = NK_model_names[:2]  # Only test first 2 shapes in CI
+
+    for N, K, model_name in NK_model_names:
+        if N % 128 != 0 or K % 128 != 0:
+            print(f"Skip {N=}, {K=} now")
+            continue
+        print(f"{model_name} N={N} K={K}: ")
+        benchmark.run(
+            print_data=True,
+            N=N,
+            K=K,
+        )
+
+    print("Benchmark finished!")

third_party/sglang/sgl-kernel/benchmark/bench_fp8_blockwise_group_gemm.py

@@ -0,0 +1,340 @@
+import argparse
+import random
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import deep_gemm
+import torch
+from sgl_kernel import fp8_blockwise_scaled_grouped_mm
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+def get_m_alignment_for_contiguous_layout():
+    return 128
+
+
+def ceil_div(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+
+def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert x.dim() == 2
+    m, n = x.shape
+    pad_size = (128 - (n % 128)) % 128
+    x = torch.nn.functional.pad(x, (0, pad_size), value=0) if pad_size > 0 else x
+    x_view = x.view(m, -1, 128)
+    x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
+    fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
+    return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
+
+
+def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert x.dim() == 2
+    m, n = x.shape
+    x_padded = torch.zeros(
+        (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device
+    )
+    x_padded[:m, :n] = x
+    x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
+    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
+    x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
+    return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(
+        x_view.size(0), x_view.size(2)
+    )
+
+
+def construct_contiguous_grouped(
+    num_groups: int, expected_m_per_group: int, k: int, n: int
+) -> Tuple[
+    int,
+    Tuple[torch.Tensor, torch.Tensor],
+    Tuple[torch.Tensor, torch.Tensor],
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+]:
+    alignment = get_m_alignment_for_contiguous_layout()
+    group_ms = [int(expected_m_per_group) for _ in range(num_groups)]
+    m = sum([ceil_div(x, alignment) * alignment for x in group_ms])
+
+    x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
+    y = torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16)
+    m_indices = torch.empty(m, device="cuda", dtype=torch.int32)
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+
+    start = 0
+    for i, group_m in enumerate(group_ms):
+        actual_end = start + group_m
+        aligned_end = start + ceil_div(group_m, alignment) * alignment
+        m_indices[start:actual_end] = i
+        m_indices[actual_end:aligned_end] = -1
+        start = aligned_end
+
+    assert m % 4 == 0, f"TMA alignment error: {m}"
+    x_fp8 = per_token_cast_to_fp8(x)
+    y_fp8 = (
+        torch.empty_like(y, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (num_groups, ceil_div(n, 128), k // 128), device="cuda", dtype=torch.float
+        ),
+    )
+    for i in range(num_groups):
+        y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i])
+
+    return m, x_fp8, y_fp8, m_indices, out
+
+
+def bench_deepgemm(
+    expected_m_per_group: int,
+    n: int,
+    k: int,
+    num_groups: int,
+    num_warmup: int,
+    num_run: int,
+) -> Tuple[float, int]:
+    # construct tensors
+    m, x_fp8, y_fp8, m_indices, out = construct_contiguous_grouped(
+        num_groups, expected_m_per_group, k, n
+    )
+
+    def run_deepgemm():
+        deep_gemm.m_grouped_fp8_gemm_nt_contiguous(x_fp8, y_fp8, out, m_indices)
+
+    # warmup
+    for _ in range(num_warmup):
+        run_deepgemm()
+    torch.cuda.synchronize()
+
+    # run
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    latencies: list[float] = []
+    start_event.record()
+    for _ in range(num_run):
+        run_deepgemm()
+    end_event.record()
+    end_event.synchronize()
+    torch.cuda.synchronize()
+    avg = start_event.elapsed_time(end_event) / num_run * 1000  # us
+
+    return avg, m
+
+
+def bench_cutlass(
+    expected_m_per_group: int,
+    n: int,
+    k: int,
+    num_groups: int,
+    num_warmup: int,
+    num_run: int,
+) -> Tuple[float, int]:
+    device = "cuda"
+    alignment = 16
+    n_g = ceil_div(n, alignment) * alignment
+    k_g = ceil_div(k, alignment) * alignment
+    out_dtype = torch.bfloat16
+
+    expert_offsets = torch.zeros((num_groups + 1), device=device, dtype=torch.int32)
+    problem_sizes = torch.zeros((num_groups, 3), device=device, dtype=torch.int32)
+    layout_sfa = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
+    layout_sfb = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
+
+    a_tensors = []
+    b_tensors = []
+    a_scales_tensors = []
+    b_scales_tensors = []
+
+    # TODO(@TianQiLin666666): Unique group_ms in all bench function
+    group_ms = [
+        alignment * ceil_div(int(expected_m_per_group), alignment)
+        for _ in range(num_groups)
+    ]
+    for g in range(num_groups):
+        m_g = group_ms[g]
+        expert_offsets[g + 1] = expert_offsets[g] + m_g
+        problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
+
+        a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
+        b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
+        a_tensors.append(a_g)
+        b_tensors.append(b_g)
+        a_scales_tensors.append(a_scale)
+        b_scales_tensors.append(b_scale)
+
+    a_stack = torch.empty(
+        (expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    b_stack = torch.empty(
+        (num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+
+    for g in range(num_groups):
+        a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g]
+        b_stack[g] = b_tensors[g].t()
+    b_stack = b_stack.transpose(1, 2)
+
+    a_scale_stack = torch.empty(
+        (expert_offsets[-1], k_g // 128), device=device, dtype=torch.float32
+    )
+    b_scale_stack = torch.empty(
+        (num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
+    )
+
+    for g in range(num_groups):
+        a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[g]
+        b_scale_stack[g] = b_scales_tensors[g].t()
+    b_scale_stack = b_scale_stack.transpose(1, 2)
+
+    c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
+    a_strides = torch.full(
+        (num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
+    )
+    c_strides = torch.full(
+        (num_groups,), c_out.stride(0), device=device, dtype=torch.int64
+    )
+    workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
+    a_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    b_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    out_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    a_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+    b_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
+
+    def run_cutlass():
+        fp8_blockwise_scaled_grouped_mm(
+            c_out,
+            a_ptrs,
+            b_ptrs,
+            out_ptrs,
+            a_scales_ptrs,
+            b_scales_ptrs,
+            a_stack,
+            b_stack,
+            a_scale_stack,
+            b_scale_stack,
+            a_strides,
+            a_strides,
+            c_strides,
+            layout_sfa,
+            layout_sfb,
+            problem_sizes,
+            expert_offsets[:-1],
+            workspace,
+        )
+
+    # warmup
+    for _ in range(num_warmup):
+        run_cutlass()
+    torch.cuda.synchronize()
+
+    # run
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    start_event.record()
+    for _ in range(num_run):
+        run_cutlass()
+    end_event.record()
+    end_event.synchronize()
+    torch.cuda.synchronize()
+    avg = start_event.elapsed_time(end_event) / num_run * 1000  # us
+
+    return avg, expert_offsets[-1]
+
+
+def bench_sglang_triton(
+    expected_m_per_group: int,
+    n: int,
+    k: int,
+    num_groups: int,
+    num_warmup: int,
+    num_run: int,
+) -> Tuple[float, int]:
+    pass
+
+
+benchmark_kernels = {
+    "deepgemm": bench_deepgemm,
+    "cutlass": bench_cutlass,
+    # "triton": bench_sglang_triton,
+}
+
+
+@dataclass
+class ShapeArg:
+    expected_m_per_group: int
+    n: int
+    k: int
+    num_groups: int
+
+
+def benchmark_one_shape(
+    shape_args: List[ShapeArg],
+    num_warmup: int,
+    num_run: int,
+):
+    for shape in shape_args:
+        print(
+            f"\nBenchmark: expected_m_per_group={shape.expected_m_per_group}, n={shape.n}, k={shape.k}, num_groups={shape.num_groups}"
+        )
+        for kernel_name, kernel_func in benchmark_kernels.items():
+            average_time, m = kernel_func(
+                shape.expected_m_per_group,
+                shape.n,
+                shape.k,
+                shape.num_groups,
+                num_warmup,
+                num_run,
+            )
+            print(f"{kernel_name}: {average_time} us")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-warmup", type=int, default=3)
+    parser.add_argument("--num-run", type=int, default=10)
+
+    # CI environment uses simplified parameters
+    if IS_CI:
+        shape_args = [
+            # Only test one simple shape in CI
+            ShapeArg(expected_m_per_group=128, n=512, k=7168, num_groups=256),
+        ]
+    else:
+        shape_args = [
+            # Prefill, DeepSeek-R1, gateup, chunk_size = 4096, TP = 8
+            ShapeArg(expected_m_per_group=128, n=512, k=7168, num_groups=256),
+            # Prefill, DeepSeek-R1, gateup, chunk_size = 8192, TP = 8
+            ShapeArg(expected_m_per_group=256, n=512, k=7168, num_groups=256),
+            # Prefill, DeepSeek-R1, gateup, chunk_size = 8192, TP = 16
+            ShapeArg(expected_m_per_group=256, n=256, k=7168, num_groups=256),
+            # Prefill, DeepSeek-R1, gateup, chunk_size = 16384, TP = 16
+            ShapeArg(expected_m_per_group=512, n=256, k=7168, num_groups=256),
+            # Decode, DeepSeek-R1, gateup, bs = 32, TP = 8
+            ShapeArg(expected_m_per_group=1, n=512, k=7168, num_groups=256),
+            # Decode, DeepSeek-R1, gateup, bs = 64, TP = 16
+            ShapeArg(expected_m_per_group=2, n=256, k=7168, num_groups=256),
+            # Prefill, DeepSeek-R1, gateup, chunk_size = 8192, EP = 8
+            ShapeArg(expected_m_per_group=256, n=4096, k=7168, num_groups=32),
+            # Prefill, DeepSeek-R1, gateup, chunk_size = 16384, EP = 16
+            ShapeArg(expected_m_per_group=512, n=4096, k=7168, num_groups=16),
+            # Decode, DeepSeek-R1, gateup, bs = 128, EP = 8
+            ShapeArg(expected_m_per_group=4, n=4096, k=7168, num_groups=32),
+            # Decode, DeepSeek-R1, gateup, bs = 256, EP = 16
+            ShapeArg(expected_m_per_group=8, n=4096, k=7168, num_groups=16),
+            # Prefill, Qwen3-235B-A22B-FP8, gateup, chunk_size = 16384, TP = 4
+            ShapeArg(expected_m_per_group=1024, n=768, k=4096, num_groups=128),
+            # Prefill, Qwen3-235B-A22B-FP8, down, chunk_size = 16384, TP = 4
+            ShapeArg(expected_m_per_group=1024, n=4096, k=384, num_groups=128),
+            # Decode, Qwen3-235B-A22B-FP8, gateup, bs = 256, TP = 4
+            ShapeArg(expected_m_per_group=16, n=768, k=4096, num_groups=128),
+            # Decode, Qwen3-235B-A22B-FP8, down, bs = 256, TP = 4
+            ShapeArg(expected_m_per_group=16, n=4096, k=384, num_groups=128),
+        ]
+    args = parser.parse_args()
+    benchmark_one_shape(shape_args, args.num_warmup, args.num_run)
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/sgl-kernel/benchmark/bench_fp8_gemm.py

@@ -0,0 +1,230 @@
+import argparse
+import copy
+import itertools
+import os
+from typing import Optional, Tuple
+
+import torch
+import triton
+from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm
+
+from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
+    from vllm._custom_ops import scaled_fp8_quant as vllm_scaled_fp8_quant
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    vllm_scaled_mm = None
+    vllm_scaled_fp8_quant = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+# Weight Shapes are in the format
+# ([K, N], TP_SPLIT_DIM)
+# Example:
+#  A shape of ([14336, 4096], 0) indicates the following GEMM shape,
+#   - TP1 : K = 14336, N = 4096
+#   - TP2 : K = 7168, N = 4096
+#  A shape of ([4096, 6144], 1) indicates the following GEMM shape,
+#   - TP1 : K = 4096, N = 6144
+#   - TP4 : K = 4096, N = 1536
+
+# TP1 shapes
+WEIGHT_SHAPES = {
+    "meta-llama/Llama-3.1-8B-Instruct": [
+        ([4096, 6144], 1),
+        ([4096, 4096], 0),
+        ([4096, 28672], 1),
+        ([14336, 4096], 0),
+    ],
+    "meta-llama/Llama-3.3-70B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 57344], 1),
+        ([28672, 8192], 0),
+    ],
+    "mistralai/Mistral-Large-Instruct-2407": [
+        ([12288, 14336], 1),
+        ([12288, 12288], 0),
+        ([12288, 57344], 1),
+        ([28672, 12288], 0),
+    ],
+    "Qwen/Qwen2.5-7B-Instruct": [
+        ([3584, 4608], 1),
+        ([3584, 3584], 0),
+        ([3584, 37888], 1),
+        ([18944, 3584], 0),
+    ],
+    "Qwen/Qwen2.5-32B-Instruct": [
+        ([5120, 7168], 1),
+        ([5120, 5120], 0),
+        ([5120, 55296], 1),
+        ([27648, 5120], 0),
+    ],
+    "Qwen/Qwen2.5-72B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 59136], 1),
+        ([29568, 8192], 0),
+    ],
+    "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
+        ([2048, 3072], 1),
+        ([2048, 4096], 1),
+        ([2048, 2048], 0),
+        ([2048, 576], 0),
+        ([2048, 21888], 1),
+        ([10944, 2048], 0),
+        ([2048, 2816], 1),
+        ([1408, 2048], 0),
+    ],
+}
+
+
+def sglang_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_type_: torch.dtype = torch.float8_e4m3fn
+    output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
+    is_static = True
+    if scale is None:
+        scale = torch.zeros(1, device=input.device, dtype=torch.float32)
+        is_static = False
+    per_tensor_quant_fp8(input, output, scale, is_static)
+
+    return output, scale
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_sizes = [1]  # Single batch size for CI
+else:
+    batch_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048]
+
+# Filter line_vals based on vLLM availability
+if VLLM_AVAILABLE:
+    line_vals = [
+        "vllm-fp8-fp16",
+        "vllm-fp8-bf16",
+        "sglang-fp8-fp16",
+        "sglang-fp8-bf16",
+    ]
+    line_names = [
+        "vllm-fp8-fp16",
+        "vllm-fp8-bf16",
+        "sglang-fp8-fp16",
+        "sglang-fp8-bf16",
+    ]
+    styles = [("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")]
+else:
+    line_vals = [
+        "sglang-fp8-fp16",
+        "sglang-fp8-bf16",
+    ]
+    line_names = [
+        "sglang-fp8-fp16",
+        "sglang-fp8-bf16",
+    ]
+    styles = [("blue", "-"), ("blue", "--")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=batch_sizes,
+        x_log=False,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="GB/s",
+        plot_name="fp8 scaled matmul",
+        args={},
+    )
+)
+def benchmark(batch_size, provider, N, K):
+    # M, N, K = batch_size, 4096, 8192
+    M = batch_size
+    a = torch.ones((M, K), device="cuda") * 5.0
+    b = torch.ones((N, K), device="cuda") * 5.0
+    # vLLM expects scalar scales, while sglang can handle per-token scales
+    scale_a_scalar = torch.randn(1, device="cuda", dtype=torch.float32)
+    scale_b_scalar = torch.randn(1, device="cuda", dtype=torch.float32)
+    scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
+    scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
+    quantiles = [0.5, 0.2, 0.8]
+
+    dtype = torch.float16 if "fp16" in provider else torch.bfloat16
+
+    if "vllm-fp8" in provider:
+        if not VLLM_AVAILABLE:
+            # Return zero if vLLM is not available
+            return (0, 0, 0)
+        a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a_scalar)
+        b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b_scalar)
+        b_fp8 = b_fp8.t()
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype),
+            quantiles=quantiles,
+        )
+    elif "sglang-fp8" in provider:
+        a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a)
+        b_fp8, scale_b_fp8 = sglang_scaled_fp8_quant(b, scale_b)
+        b_fp8 = b_fp8.t()
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: sgl_scaled_mm(
+                a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype, bias=None
+            ),
+            quantiles=quantiles,
+        )
+
+    gbps = lambda ms: (2 * M * N * K + M * N) * a.element_size() * 1e-9 / (ms * 1e-3)
+    return gbps(ms), gbps(max_ms), gbps(min_ms)
+
+
+def prepare_shapes(args):
+    KN_model_names = []
+    models_tps = list(itertools.product(args.models, args.tp_sizes))
+    for model, tp_size in models_tps:
+        assert model in WEIGHT_SHAPES
+        for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
+            KN[tp_split_dim] = KN[tp_split_dim] // tp_size
+            KN.append(model)
+            KN_model_names.append(KN)
+    return KN_model_names
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        type=str,
+        default=["meta-llama/Llama-3.1-8B-Instruct"],
+        help="List of models to benchmark",
+    )
+    parser.add_argument(
+        "--tp-sizes",
+        nargs="+",
+        type=int,
+        default=[1],
+        help="List of tensor parallel sizes",
+    )
+    args = parser.parse_args()
+
+    # Simplify for CI environment
+    if IS_CI:
+        args.models = [args.models[0]]  # Use only first model
+        args.tp_sizes = [args.tp_sizes[0]]  # Use only first TP size
+
+    KN_model_names = prepare_shapes(args)
+    for K, N, model_name in KN_model_names:
+        print(f"{model_name} N={N} K={K}: ")
+        benchmark.run(print_data=True, N=N, K=K)
+
+    print("Benchmark finished!")

third_party/sglang/sgl-kernel/benchmark/bench_int8_gemm.py

@@ -0,0 +1,183 @@
+import argparse
+import copy
+import itertools
+import os
+
+import torch
+import triton
+from sgl_kernel import int8_scaled_mm
+
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    vllm_scaled_mm = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+
+def to_int8(tensor: torch.Tensor) -> torch.Tensor:
+    return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
+
+
+WEIGHT_SHAPES = {
+    "meta-llama/Llama-3.1-8B-Instruct": [
+        ([4096, 6144], 1),
+        ([4096, 4096], 0),
+        ([4096, 28672], 1),
+        ([14336, 4096], 0),
+    ],
+    "meta-llama/Llama-3.3-70B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 57344], 1),
+        ([28672, 8192], 0),
+    ],
+    "mistralai/Mistral-Large-Instruct-2407": [
+        ([12288, 14336], 1),
+        ([12288, 12288], 0),
+        ([12288, 57344], 1),
+        ([28672, 12288], 0),
+    ],
+    "Qwen/Qwen2.5-7B-Instruct": [
+        ([3584, 4608], 1),
+        ([3584, 3584], 0),
+        ([3584, 37888], 1),
+        ([18944, 3584], 0),
+    ],
+    "Qwen/Qwen2.5-32B-Instruct": [
+        ([5120, 7168], 1),
+        ([5120, 5120], 0),
+        ([5120, 55296], 1),
+        ([27648, 5120], 0),
+    ],
+    "Qwen/Qwen2.5-72B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 59136], 1),
+        ([29568, 8192], 0),
+    ],
+    "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
+        ([2048, 3072], 1),
+        ([2048, 4096], 1),
+        ([2048, 2048], 0),
+        ([2048, 576], 0),
+        ([2048, 21888], 1),
+        ([10944, 2048], 0),
+        ([2048, 2816], 1),
+        ([1408, 2048], 0),
+    ],
+}
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_sizes = [1]  # Single batch size for CI
+else:
+    batch_sizes = [1, 16, 32, 64, 128, 256, 512, 1024, 2048]
+
+# Filter providers based on vLLM availability
+if VLLM_AVAILABLE:
+    line_vals = ["vllm", "sgl-kernel"]
+    line_names = ["vllm int8 gemm", "sgl-kernel int8 gemm"]
+    styles = [("blue", "-"), ("orange", "-")]
+else:
+    line_vals = ["sgl-kernel"]
+    line_names = ["sgl-kernel int8 gemm"]
+    styles = [("orange", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=batch_sizes,
+        x_log=False,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="GB/s",
+        plot_name="int8 scaled matmul",
+        args={},
+    )
+)
+def benchmark(batch_size, provider, N, K):
+    M = batch_size
+    a = to_int8(torch.randn((M, K), device="cuda") * 5)
+    b = to_int8(torch.randn((N, K), device="cuda").t() * 5)
+    scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
+    scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
+    bias = torch.randn((N,), device="cuda", dtype=torch.float16)
+
+    quantiles = [0.5, 0.2, 0.8]
+    if provider == "sgl-kernel":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: int8_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
+            quantiles=quantiles,
+        )
+    elif provider == "vllm":
+        if not VLLM_AVAILABLE:
+            return (0, 0, 0)
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
+            quantiles=quantiles,
+        )
+    gbps = (
+        lambda ms: (
+            (2 * M * N * K - M * N) * a.element_size()
+            + (3 * M * N) * scale_a.element_size()
+        )
+        * 1e-9
+        / (ms * 1e-3)
+    )
+    return gbps(ms), gbps(max_ms), gbps(min_ms)
+
+
+def prepare_shapes(args):
+    KN_model_names = []
+    models_tps = list(itertools.product(args.models, args.tp_sizes))
+    for model, tp_size in models_tps:
+        assert model in WEIGHT_SHAPES
+        for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
+            KN[tp_split_dim] = KN[tp_split_dim] // tp_size
+            KN.append(model)
+            KN_model_names.append(KN)
+    return KN_model_names
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        type=str,
+        default=["meta-llama/Llama-3.1-8B-Instruct"],
+        help="List of models to benchmark",
+    )
+    parser.add_argument(
+        "--tp-sizes",
+        nargs="+",
+        type=int,
+        default=[1],
+        help="List of tensor parallel sizes",
+    )
+    args = parser.parse_args()
+
+    # Skip in CI environment due to architecture compatibility issues
+    if IS_CI:
+        print(
+            "Skipping INT8 GEMM benchmark in CI environment due to architecture compatibility issues"
+        )
+        print("INT8 operations may not be supported on all GPU architectures")
+    else:
+        KN_model_names = prepare_shapes(args)
+        for K, N, model_name in KN_model_names:
+            print(f"{model_name} N={N} K={K}: ")
+            benchmark.run(print_data=True, N=N, K=K)
+
+        print("Benchmark finished!")

third_party/sglang/sgl-kernel/benchmark/bench_kimi_k2_moe_fused_gate.py

@@ -0,0 +1,114 @@
+import itertools
+import math
+import os
+
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import kimi_k2_moe_fused_gate
+
+from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+def kimi_k2_biased_topk_torch_compile(scores, bias, topk, routed_scaling_factor):
+    """Original torch.compile-based implementation"""
+    return kimi_k2_biased_topk_impl(
+        scores,
+        scores,
+        bias,
+        topk=topk,
+        renormalize=True,
+        routed_scaling_factor=routed_scaling_factor,
+    )
+
+
+def kimi_k2_biased_topk_fused_kernel(scores, bias, topk, routed_scaling_factor):
+    """Our fused CUDA kernel implementation"""
+    return kimi_k2_moe_fused_gate(
+        scores,
+        bias,
+        topk=topk,
+        renormalize=True,
+        routed_scaling_factor=routed_scaling_factor,
+    )
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    seq_length_range = [5000]  # Only test one sequence length in CI
+else:
+    seq_length_range = [
+        1,
+        8,
+        16,
+        32,
+        64,
+        128,
+        256,
+        512,
+        1024,
+        2048,
+        4096,
+        10000,
+        15000,
+        20000,
+        25000,
+        30000,
+        35000,
+        40000,
+    ]
+
+configs = [(sq,) for sq in seq_length_range]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["seq_length"],
+        x_vals=[list(_) for _ in configs],
+        line_arg="provider",
+        line_vals=["torch_compile", "fused_kernel"],
+        line_names=["Torch Compile", "Fused Kernel"],
+        styles=[("blue", "-"), ("red", "-")],
+        ylabel="us",
+        plot_name="kimi-k2-moe-fused-gate-performance",
+        args={},
+    )
+)
+def benchmark(seq_length, provider):
+    dtype = torch.float32
+    device = torch.device("cuda")
+    num_experts, topk = 384, 6  # Kimi K2 configuration
+    routed_scaling_factor = 2.872  # Kimi K2's routed scaling factor
+
+    scores = torch.randn((seq_length, num_experts), device=device, dtype=dtype)
+    bias = torch.rand(num_experts, device=device, dtype=dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "torch_compile":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: kimi_k2_biased_topk_torch_compile(
+                scores.clone(), bias.clone(), topk, routed_scaling_factor
+            ),
+            quantiles=quantiles,
+        )
+    elif provider == "fused_kernel":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: kimi_k2_biased_topk_fused_kernel(
+                scores.clone(), bias.clone(), topk, routed_scaling_factor
+            ),
+            quantiles=quantiles,
+        )
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    print("=" * 80)
+    print("Benchmarking Kimi K2 MoE Fused Gate Performance")
+    print("=" * 80)
+    print("\nPerformance vs Sequence Length (384 experts, topk=6)")
+    benchmark.run(print_data=True, save_path=".")

third_party/sglang/sgl-kernel/benchmark/bench_moe_align_block_size.py

@@ -0,0 +1,443 @@
+import argparse
+import itertools
+import os
+
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size
+
+from sglang.utils import is_in_ci
+
+try:
+    from vllm import _custom_ops as ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    ops = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+USE_RANDOM_PERM = False
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+@triton.jit
+def moe_align_block_size_stage1(
+    topk_ids_ptr,
+    tokens_cnts_ptr,
+    num_experts: tl.constexpr,
+    numel: tl.constexpr,
+    tokens_per_thread: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    start_idx = pid * tokens_per_thread
+    off_c = (pid + 1) * num_experts
+
+    for i in range(tokens_per_thread):
+        if start_idx + i < numel:
+            idx = tl.load(topk_ids_ptr + start_idx + i)
+            token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
+            tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)
+
+
+@triton.jit
+def moe_align_block_size_stage2(
+    tokens_cnts_ptr,
+    num_experts: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    last_cnt = 0
+    for i in range(1, num_experts + 1):
+        token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
+        last_cnt = last_cnt + token_cnt
+        tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)
+
+
+@triton.jit
+def moe_align_block_size_stage3(
+    total_tokens_post_pad_ptr,
+    tokens_cnts_ptr,
+    cumsum_ptr,
+    num_experts: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    last_cumsum = 0
+    off_cnt = num_experts * num_experts
+    for i in range(1, num_experts + 1):
+        token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
+        last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
+        tl.store(cumsum_ptr + i, last_cumsum)
+    tl.store(total_tokens_post_pad_ptr, last_cumsum)
+
+
+@triton.jit
+def moe_align_block_size_stage4(
+    topk_ids_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    tokens_cnts_ptr,
+    cumsum_ptr,
+    num_experts: tl.constexpr,
+    block_size: tl.constexpr,
+    numel: tl.constexpr,
+    tokens_per_thread: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    start_idx = tl.load(cumsum_ptr + pid)
+    end_idx = tl.load(cumsum_ptr + pid + 1)
+
+    for i in range(start_idx, end_idx, block_size):
+        tl.store(expert_ids_ptr + i // block_size, pid)
+
+    start_idx = pid * tokens_per_thread
+    off_t = pid * num_experts
+
+    for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread, numel)):
+        expert_id = tl.load(topk_ids_ptr + i)
+        token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
+        rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
+        tl.store(sorted_token_ids_ptr + rank_post_pad, i)
+        tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)
+
+
+def moe_align_block_size_triton(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    numel = topk_ids.numel()
+    grid = (num_experts,)
+    tokens_cnts = torch.zeros(
+        (num_experts + 1, num_experts), dtype=torch.int32, device=topk_ids.device
+    )
+    cumsum = torch.zeros((num_experts + 1,), dtype=torch.int32, device=topk_ids.device)
+    tokens_per_thread = ceil_div(numel, num_experts)
+
+    moe_align_block_size_stage1[grid](
+        topk_ids,
+        tokens_cnts,
+        num_experts,
+        numel,
+        tokens_per_thread,
+    )
+    moe_align_block_size_stage2[grid](
+        tokens_cnts,
+        num_experts,
+    )
+    moe_align_block_size_stage3[(1,)](
+        num_tokens_post_pad,
+        tokens_cnts,
+        cumsum,
+        num_experts,
+        block_size,
+    )
+    moe_align_block_size_stage4[grid](
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        tokens_cnts,
+        cumsum,
+        num_experts,
+        block_size,
+        numel,
+        tokens_per_thread,
+    )
+
+
+def calculate_diff(num_tokens, num_experts=256, block_size=128, topk=8):
+    topk_ids = torch.stack(
+        [
+            torch.randperm(num_experts, dtype=torch.int32, device="cuda")[:topk]
+            for _ in range(num_tokens)
+        ]
+    )
+
+    # SGL kernel uses dynamic padding optimization
+    max_num_tokens_padded_sgl = topk_ids.numel() + num_experts * (block_size - 1)
+    if topk_ids.numel() < num_experts + 1:
+        max_num_tokens_padded_sgl = topk_ids.numel() * block_size
+    sorted_ids_cuda = torch.empty(
+        (max_num_tokens_padded_sgl,), dtype=torch.int32, device=topk_ids.device
+    )
+    sorted_ids_cuda.fill_(topk_ids.numel())
+    max_num_m_blocks_sgl = max_num_tokens_padded_sgl // block_size
+    expert_ids_cuda = torch.zeros(
+        (max_num_m_blocks_sgl,), dtype=torch.int32, device=topk_ids.device
+    )
+    num_tokens_post_pad_cuda = torch.empty(
+        (1), dtype=torch.int32, device=topk_ids.device
+    )
+    cumsum_buffer = torch.zeros(
+        num_experts + 1, dtype=torch.int32, device=topk_ids.device
+    )
+
+    # Triton and vLLM use original padding calculation
+    max_num_tokens_padded_triton = topk_ids.numel() + num_experts * (block_size - 1)
+    max_num_m_blocks_triton = max_num_tokens_padded_triton // block_size
+    sorted_ids_triton = torch.empty(
+        (max_num_tokens_padded_triton,), dtype=torch.int32, device=topk_ids.device
+    )
+    sorted_ids_triton.fill_(topk_ids.numel())
+    expert_ids_triton = torch.zeros(
+        (max_num_m_blocks_triton,), dtype=torch.int32, device=topk_ids.device
+    )
+    num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
+
+    sorted_ids_vllm = torch.empty_like(sorted_ids_triton)
+    sorted_ids_vllm.fill_(topk_ids.numel())
+    expert_ids_vllm = torch.zeros_like(expert_ids_triton)
+    num_tokens_post_pad_vllm = torch.empty_like(num_tokens_post_pad_cuda)
+
+    # compare the performance of cuda, triton and vllm implementation
+    sgl_moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_ids_cuda,
+        expert_ids_cuda,
+        num_tokens_post_pad_cuda,
+        cumsum_buffer,
+    )
+    moe_align_block_size_triton(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_ids_triton,
+        expert_ids_triton,
+        num_tokens_post_pad_triton,
+    )
+
+    if VLLM_AVAILABLE:
+        try:
+            ops.moe_align_block_size(
+                topk_ids,
+                num_experts,
+                block_size,
+                sorted_ids_vllm,
+                expert_ids_vllm,
+                num_tokens_post_pad_vllm,
+            )
+            print(f"✅ VLLM implementation works with {num_experts} experts!")
+            vllm_works = True
+        except Exception as e:
+            print(f"❌ VLLM implementation failed with {num_experts} experts: {e}")
+            vllm_works = False
+    else:
+        print("⚠️ vLLM not available, skipping vLLM test")
+        vllm_works = False
+
+    if torch.allclose(expert_ids_cuda, expert_ids_triton) and torch.allclose(
+        num_tokens_post_pad_cuda, num_tokens_post_pad_triton
+    ):
+        print("✅ SGL and Triton implementations match")
+    else:
+        print("❌ SGL and Triton implementations do not match")
+        print("SGL expert_ids:", expert_ids_cuda)
+        print("Triton expert_ids:", expert_ids_triton)
+        print("SGL num_tokens_post_pad:", num_tokens_post_pad_cuda)
+        print("Triton num_tokens_post_pad:", num_tokens_post_pad_triton)
+
+    if (
+        vllm_works
+        and torch.allclose(expert_ids_cuda, expert_ids_vllm)
+        and torch.allclose(num_tokens_post_pad_cuda, num_tokens_post_pad_vllm)
+    ):
+        print("✅ SGL and VLLM implementations match")
+    else:
+        if not vllm_works:
+            print("⚠️ VLLM comparison skipped due to failure")
+        else:
+            print("❌ SGL and VLLM implementations do not match")
+            print("SGL expert_ids:", expert_ids_cuda)
+            print("VLLM expert_ids:", expert_ids_vllm)
+            print("SGL num_tokens_post_pad:", num_tokens_post_pad_cuda)
+            print("VLLM num_tokens_post_pad:", num_tokens_post_pad_vllm)
+
+
+# Test range
+num_tokens_range = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
+num_experts_range = [8, 32, 64, 128, 256]
+topk_range = [1, 2, 4, 8]
+
+configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
+
+
+def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
+    topk_ids = torch.zeros((num_tokens, topk), dtype=torch.int32, device="cuda")
+    for i in range(num_tokens):
+        topk_ids[i, :] = torch.randperm(num_experts, dtype=torch.int32, device="cuda")[
+            :topk
+        ]
+    return topk_ids
+
+
+def sgl_moe_align_block_size_with_empty(
+    topk_ids,
+    num_experts,
+    block_size,
+    sorted_ids,
+    expert_ids,
+    num_tokens_post_pad,
+    pad_sorted_token_ids=False,
+):
+    if not pad_sorted_token_ids:
+        sorted_ids.fill_(topk_ids.numel())
+
+    cumsum_buffer = torch.empty(
+        num_experts + 1, dtype=torch.int32, device=topk_ids.device
+    )
+
+    sgl_moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_ids.clone(),
+        expert_ids.clone(),
+        num_tokens_post_pad.clone(),
+        cumsum_buffer,
+        pad_sorted_token_ids,
+    )
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens", "num_experts", "topk"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["sgl", "sgl_fusion", "triton"],
+        line_names=["SGL", "SGL Fusion", "Triton"],
+        styles=[("blue", "-"), ("red", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="moe-align-block-size-performance",
+        args={},
+    )
+)
+def benchmark(num_tokens, num_experts, topk, provider):
+    block_size = 128
+
+    if USE_RANDOM_PERM:
+        topk_ids = get_topk_ids(num_tokens, num_experts, topk)
+    else:
+        topk_ids = torch.randint(
+            0,
+            num_experts,
+            (num_tokens, topk),
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+    # Calculate max_num_tokens_padded based on provider
+    if provider == "sgl" or provider == "sgl_fusion":
+        # Apply dynamic padding optimization for SGL kernel
+        max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+        if topk_ids.numel() < num_experts:
+            max_num_tokens_padded = topk_ids.numel() * block_size
+    else:  # triton
+        # Use original padding calculation for Triton
+        max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+
+    # Create tensors
+    sorted_ids = torch.empty(
+        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
+    )
+    max_num_m_blocks = max_num_tokens_padded // block_size
+    expert_ids = torch.empty(
+        (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
+    )
+    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
+
+    quantiles = [0.5, 0.2, 0.8]
+    if provider == "sgl":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: sgl_moe_align_block_size_with_empty(
+                topk_ids,
+                num_experts,
+                block_size,
+                sorted_ids,
+                expert_ids,
+                num_tokens_post_pad,
+            ),
+            quantiles=quantiles,
+        )
+    elif provider == "sgl_fusion":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: sgl_moe_align_block_size_with_empty(
+                topk_ids,
+                num_experts,
+                block_size,
+                sorted_ids,
+                expert_ids,
+                num_tokens_post_pad,
+                pad_sorted_token_ids=True,
+            ),
+            quantiles=quantiles,
+        )
+    elif provider == "triton":
+        sorted_ids.fill_(topk_ids.numel())
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: moe_align_block_size_triton(
+                topk_ids,
+                num_experts,
+                block_size,
+                sorted_ids.clone(),
+                expert_ids.clone(),
+                num_tokens_post_pad.clone(),
+            ),
+            quantiles=quantiles,
+        )
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default="./configs/benchmark_ops/moe_align_blocks/",
+        help="Path to save moe align benchmark results",
+    )
+    parser.add_argument(
+        "--num_experts",
+        type=int,
+        default=256,
+        choices=[8, 16, 32, 64, 128, 256],
+        help="Number of experts for benchmark",
+    )
+    parser.add_argument(
+        "--topk",
+        type=int,
+        default=8,
+        choices=[2, 4, 8],
+        help="Top-k value for benchmark",
+    )
+    parser.add_argument(
+        "--skip_full_benchmark",
+        action="store_true",
+        help="Only run the calculate_diff function, skip full benchmarking",
+    )
+    args = parser.parse_args()
+
+    # Simplify for CI environment
+    if IS_CI:
+        num_tokens = 256  # Smaller for CI
+        num_experts = 8  # Smaller for CI
+        topk = 2  # Smaller for CI
+    else:
+        num_tokens = 1024
+        num_experts = args.num_experts
+        topk = args.topk
+
+    calculate_diff(num_tokens=num_tokens, num_experts=num_experts, topk=topk)
+
+    if not args.skip_full_benchmark and not IS_CI:  # Skip full benchmark in CI
+        print(f"\n📊 Running performance benchmark for {args.num_experts} experts...")
+        benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_moe_ep_post_reorder.py

@@ -0,0 +1,85 @@
+import torch
+import triton
+
+from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_sizes = [64, 128]  # Only test 2 values in CI
+else:
+    batch_sizes = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]
+
+configs = [(bs,) for bs in batch_sizes]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=[list(_) for _ in configs],
+        line_arg="provider",
+        line_vals=["triton"],
+        line_names=["Triton Kernel"],
+        styles=[("orange", "-")],
+        ylabel="us",
+        plot_name="ep-moe-post-reorder-performance",
+        args={},
+    )
+)
+def benchmark(batch_size, provider):
+    dtype = torch.bfloat16
+    device = torch.device("cuda")
+    hidden_size, topk, start_expert_id, end_expert_id, block_size = 4096, 8, 0, 255, 512
+
+    def alloc_tensors():
+        down_output = torch.randn(
+            batch_size * topk, hidden_size, dtype=dtype, device=device
+        )
+        output = torch.zeros(batch_size, hidden_size, dtype=dtype, device=device)
+        src2dst = torch.randint(
+            0, batch_size * topk, (batch_size, topk), dtype=torch.int32, device=device
+        )
+        topk_ids = torch.randint(
+            start_expert_id,
+            end_expert_id + 1,
+            (batch_size, topk),
+            dtype=torch.int32,
+            device=device,
+        )
+        topk_weights = torch.rand(batch_size, topk, dtype=dtype, device=device)
+        return down_output, output, src2dst, topk_ids, topk_weights
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "triton":
+        d_out, out, s2d, tk_ids, tk_weights = alloc_tensors()
+
+        def run_triton():
+            post_reorder_triton_kernel[(batch_size,)](
+                d_out.view(-1),
+                out.view(-1),
+                s2d.view(-1),
+                tk_ids.view(-1),
+                tk_weights.view(-1),
+                start_expert_id,
+                end_expert_id,
+                topk,
+                hidden_size,
+                0,
+                block_size,
+            )
+
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            run_triton, quantiles=quantiles
+        )
+
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_moe_fused_gate.py

@@ -0,0 +1,86 @@
+import itertools
+import math
+import os
+
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import moe_fused_gate
+
+from sglang.srt.layers.moe.topk import biased_grouped_topk
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+def biased_grouped_topk_org(scores, bias, num_expert_group, topk_group, topk):
+    return biased_grouped_topk(
+        scores,
+        scores,
+        bias,
+        topk=topk,
+        renormalize=True,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        routed_scaling_factor=2.5,  # DeepSeek-R1 : 2.5, Kimi K2: 2.872
+    )
+
+
+def biased_grouped_topk_org_fuse_kernel(
+    scores, bias, num_expert_group, topk_group, topk
+):
+    return moe_fused_gate(scores, bias, num_expert_group, topk_group, topk)
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    seq_length_range = [5000]  # Only test one sequence length in CI
+else:
+    seq_length_range = [5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000]
+
+configs = [(sq,) for sq in seq_length_range]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["seq_length"],
+        x_vals=[list(_) for _ in configs],
+        line_arg="provider",
+        line_vals=["original", "kernel"],
+        line_names=["Original", "SGL Kernel"],
+        styles=[("blue", "-"), ("red", "-")],
+        ylabel="us",
+        plot_name="moe-fused-gate-performance",
+        args={},
+    )
+)
+def benchmark(seq_length, provider):
+    dtype = torch.float32
+    device = torch.device("cuda")
+    num_experts, num_expert_group, topk_group, topk = 256, 8, 4, 8
+
+    scores = torch.randn((seq_length, num_experts), device=device, dtype=dtype)
+    bias = torch.rand(num_experts, device=device, dtype=dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "original":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: biased_grouped_topk_org(
+                scores.clone(), bias.clone(), num_expert_group, topk_group, topk
+            ),
+            quantiles=quantiles,
+        )
+    elif provider == "kernel":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: biased_grouped_topk_org_fuse_kernel(
+                scores.clone(), bias.clone(), num_expert_group, topk_group, topk
+            ),
+            quantiles=quantiles,
+        )
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py

@@ -0,0 +1,169 @@
+import itertools
+import os
+
+import pytest
+import torch
+import triton
+from sgl_kernel import topk_sigmoid
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+def torch_topk_sigmoid_native(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    correction_bias: torch.Tensor = None,
+):
+    scores = gating_output.sigmoid()
+    if correction_bias is not None:
+        n_routed_experts = gating_output.shape[-1]
+        scores_for_choice = scores.view(
+            -1, n_routed_experts
+        ) + correction_bias.unsqueeze(0)
+        _, topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1)
+        topk_weights = scores.gather(1, topk_indices)
+    else:
+        topk_weights, topk_indices = torch.topk(scores, k=topk, dim=-1)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    return topk_weights, topk_indices
+
+
+def sglang_topk_sigmoid(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    correction_bias: torch.Tensor = None,
+):
+    num_tokens, num_experts = gating_output.shape
+
+    topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+
+    topk_sigmoid(
+        topk_weights,
+        topk_indices,
+        gating_output,
+        renormalize=renormalize,
+        correction_bias=correction_bias,
+    )
+
+    return topk_weights, topk_indices
+
+
+def get_topk_sigmoid_input(num_tokens, num_experts):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+    correction_bias = torch.randn((num_experts), dtype=torch.float32, device="cuda")
+    return gating_output, correction_bias
+
+
+def calculate_diff(num_tokens, num_experts, topk):
+    gating_output, correction_bias = get_topk_sigmoid_input(num_tokens, num_experts)
+
+    weights_torch, indices_torch = torch_topk_sigmoid_native(
+        gating_output.clone(),
+        topk,
+        True,
+        correction_bias.clone(),
+    )
+    weights_sglang, indices_sglang = sglang_topk_sigmoid(
+        gating_output.clone(),
+        topk,
+        True,
+        correction_bias.clone(),
+    )
+
+    weights_diff = torch.abs(weights_torch - weights_sglang).mean().item()
+    indices_match = torch.equal(indices_torch, indices_sglang)
+
+    if (
+        torch.allclose(weights_torch, weights_sglang, atol=1e-3, rtol=1e-3)
+        and indices_match
+    ):
+        print("✅ Torch and SGLang topk_sigmoid implementations match")
+    else:
+        print(
+            f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}"
+        )
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    num_tokens_range = [128]  # Single value for CI
+    num_experts_range = [32]  # Single value for CI
+    topk_range = [2]  # Single value for CI
+else:
+    num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768]
+    num_experts_range = [32, 64, 128, 256, 12, 512]
+    topk_range = [1, 2, 4, 8]
+
+configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
+
+
+# Filter providers based on vLLM availability
+line_vals = ["sglang", "torch"]
+line_names = ["SGLang", "Torch"]
+styles = [("blue", "-"), ("green", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens", "num_experts", "topk"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="Latency (us)",
+        plot_name="topk-sigmoid-performance",
+        args={},
+    )
+)
+def benchmark(num_tokens, num_experts, topk, provider):
+    gating_output, correction_bias = get_topk_sigmoid_input(num_tokens, num_experts)
+
+    if provider == "torch" or provider == "torch1":
+
+        def fn():
+            return torch_topk_sigmoid_native(
+                gating_output,
+                topk,
+                True,
+                correction_bias,
+            )
+
+    elif provider == "sglang" or provider == "sglang1":
+
+        def fn():
+            return sglang_topk_sigmoid(gating_output, topk, True, correction_bias)
+
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    # Simplify configs for CI environment
+    if IS_CI:
+        test_configs = [(20, 32, 2)]  # Single config for CI
+    else:
+        test_configs = [
+            (20, 256, 4),
+            (20, 256, 8),
+            (20, 12, 4),
+            (20, 12, 1),
+            (20, 512, 4),
+            (20, 512, 1),
+        ]
+
+    for num_tokens, num_experts, topk in test_configs:
+        calculate_diff(num_tokens, num_experts, topk)
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_moe_topk_softmax.py

@@ -0,0 +1,161 @@
+import itertools
+import os
+
+import pytest
+import torch
+import triton
+from sgl_kernel import topk_softmax
+
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm import _custom_ops as vllm_custom_ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    vllm_custom_ops = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+
+def vllm_topk_softmax(gating_output, topk):
+    if not VLLM_AVAILABLE:
+        # Fallback to SGLang implementation if vLLM is not available
+        return sglang_topk_softmax(gating_output, topk)
+
+    num_tokens, num_experts = gating_output.shape
+
+    topk_weights = torch.empty(
+        (num_tokens, topk), device=gating_output.device, dtype=torch.float32
+    )
+    topk_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device=gating_output.device
+    )
+    token_expert_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device=gating_output.device
+    )
+    torch.ops._moe_C.topk_softmax(
+        topk_weights, topk_indices, token_expert_indices, gating_output
+    )
+    return topk_weights, topk_indices
+
+
+def sglang_topk_softmax(gating_output, topk):
+    num_tokens, num_experts = gating_output.shape
+
+    topk_weights = torch.empty(
+        (num_tokens, topk), device=gating_output.device, dtype=torch.float32
+    )
+    topk_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device=gating_output.device
+    )
+
+    topk_softmax(
+        topk_weights=topk_weights,
+        topk_ids=topk_indices,
+        gating_output=gating_output,
+    )
+
+    return topk_weights, topk_indices
+
+
+def calculate_diff(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), device="cuda", dtype=torch.float32
+    )
+    weights_vllm, indices_vllm = vllm_topk_softmax(gating_output.clone(), topk)
+    weights_sglang, indices_sglang = sglang_topk_softmax(gating_output.clone(), topk)
+
+    weights_diff = torch.abs(weights_vllm - weights_sglang).mean().item()
+    indices_match = torch.equal(indices_vllm, indices_sglang)
+
+    if not VLLM_AVAILABLE:
+        print("⚠️ vLLM not available, skipping comparison")
+        return
+
+    if (
+        torch.allclose(weights_vllm, weights_sglang, atol=1e-3, rtol=1e-3)
+        and indices_match
+    ):
+        print("✅ VLLM and SGLang topk_softmax implementations match")
+    else:
+        print(
+            f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}"
+        )
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    num_tokens_range = [128]  # Single value for CI
+    num_experts_range = [32]  # Single value for CI
+    topk_range = [2]  # Single value for CI
+else:
+    num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768]
+    num_experts_range = [32, 64, 128, 256, 12, 512]
+    topk_range = [1, 2, 4, 8]
+
+configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
+
+
+# Filter providers based on vLLM availability
+if VLLM_AVAILABLE:
+    line_vals = ["sglang", "vllm"]
+    line_names = ["SGLang", "VLLM"]
+    styles = [("blue", "-"), ("green", "-")]
+else:
+    line_vals = ["sglang"]
+    line_names = ["SGLang"]
+    styles = [("blue", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens", "num_experts", "topk"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="Latency (us)",
+        plot_name="topk-softmax-performance",
+        args={},
+    )
+)
+def benchmark(num_tokens, num_experts, topk, provider):
+
+    gating_output = torch.randn(
+        (num_tokens, num_experts), device="cuda", dtype=torch.float32
+    )
+
+    if provider == "vllm" or provider == "vllm1":
+        if not VLLM_AVAILABLE:
+            return (0, 0, 0)
+        fn = lambda: vllm_topk_softmax(gating_output, topk)
+    elif provider == "sglang" or provider == "sglang1":
+        fn = lambda: sglang_topk_softmax(gating_output, topk)
+
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    # Simplify configs for CI environment
+    if IS_CI:
+        test_configs = [(20, 32, 2)]  # Single config for CI
+    else:
+        test_configs = [
+            (20, 256, 4),
+            (20, 256, 8),
+            (20, 12, 4),
+            (20, 12, 1),
+            (20, 512, 4),
+            (20, 512, 1),
+        ]
+
+    for num_tokens, num_experts, topk in test_configs:
+        calculate_diff(num_tokens, num_experts, topk)
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_mrope.py

@@ -0,0 +1,250 @@
+# Adapted from vLLM benchmark_mrope.py
+
+# This script benchmarks the mrope kernel (mainly for Qwen2VL and Qwen2.5VL models).
+# It generates test data, runs benchmarks, and saves results to a CSV file.
+#
+# The CSV file (named with current date/time) contains these columns:
+# model_name, tp_size, num_tokens, num_heads, num_kv_heads, head_dim, max_position,
+# rope_theta, is_neox_style, rope_scaling, dtype, torch_mean, torch_median, torch_p99,
+# torch_min, torch_max, triton_mean, triton_median, triton_p99, triton_min, triton_max,
+# speedup
+#
+# == Usage Examples ==
+#
+# Single model benchmark:
+# python3 benchmark_mrope.py --model-name Qwen/Qwen2.5-VL-7B-Instruct --tp-size 8 \
+#   --warmup-iter 10 --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024
+
+import argparse
+import time
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import AutoConfig
+
+from sglang.srt.layers.rotary_embedding import get_rope
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def get_model_config(model_name: str):
+    """Get model configuration parameters"""
+    config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+    return config
+
+
+def generate_test_data(
+    num_tokens: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    max_position_embeddings: int,
+    dtype: torch.dtype,
+    device: torch.device,
+):
+    """Generate test data for given configuration."""
+    # Create 2D positions (3, num_tokens) for multimodal case
+    positions = torch.randint(
+        0, max_position_embeddings // 4, (3, num_tokens), device=device
+    )
+
+    # Create query and key tensors
+    query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device)
+    key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
+
+    return positions, query, key
+
+
+def calculate_stats(times: list[float]) -> dict[str, float]:
+    """Calculate statistics from a list of times."""
+    times_array = np.array(times)
+    return {
+        "mean": np.mean(times_array),
+        "median": np.median(times_array),
+        "p99": np.percentile(times_array, 99),
+        "min": np.min(times_array),
+        "max": np.max(times_array),
+    }
+
+
+def benchmark_mrope(
+    model_name: str,
+    num_tokens: int,
+    head_dim: int,
+    tp_size: int,
+    num_heads: int,
+    num_kv_heads: int,
+    max_position: int = 8192,
+    rope_theta: float = 10000,
+    is_neox_style: bool = True,
+    rope_scaling: dict[str, Any] = None,
+    dtype: torch.dtype = torch.bfloat16,
+    seed: int = 0,
+    warmup_iter: int = 10,
+    benchmark_iter: int = 100,
+):
+    torch.manual_seed(seed)
+    torch.set_default_device(device)
+    # the parameters to compute the q k v size based on tp_size
+    mrope_helper_class = get_rope(
+        head_size=head_dim,
+        rotary_dim=head_dim,
+        max_position=max_position,
+        base=rope_theta,
+        is_neox_style=is_neox_style,
+        rope_scaling=rope_scaling,
+        dtype=dtype,
+    ).to(device=device)
+
+    print(80 * "=")
+    print(
+        f"Evaluating model: {model_name} "
+        f"with tp_size: {tp_size} "
+        f"and num_tokens: {num_tokens}, "
+        f"dtype: {dtype}"
+    )
+
+    # create q k v input tensors
+    # create rotary pos emb input tensors
+    positions, query, key = generate_test_data(
+        num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
+    )
+
+    # Warm up
+    for _ in range(warmup_iter):
+        mrope_helper_class.forward_native(
+            positions,
+            query.clone(),
+            key.clone(),
+        )
+
+        mrope_helper_class.forward(
+            positions,
+            query.clone(),
+            key.clone(),
+        )
+
+    torch.cuda.synchronize()
+
+    # Time reference implementation
+    torch_times = []
+    for _ in range(benchmark_iter):
+        query_clone = query.clone()
+        key_clone = key.clone()
+        torch.cuda.synchronize()
+        start_time = time.time()
+
+        mrope_helper_class.forward_native(
+            positions,
+            query_clone,
+            key_clone,
+        )
+
+        torch.cuda.synchronize()
+        torch_times.append(time.time() - start_time)
+
+    # Time triton kernel implementation
+    triton_times = []
+    for _ in range(benchmark_iter):
+        query_clone = query.clone()
+        key_clone = key.clone()
+        torch.cuda.synchronize()
+        start_time = time.time()
+        mrope_helper_class.forward(
+            positions,
+            query_clone,
+            key_clone,
+        )
+        torch.cuda.synchronize()
+        triton_times.append(time.time() - start_time)
+
+    # Calculate statistics
+    torch_stats = calculate_stats(torch_times)
+    triton_stats = calculate_stats(triton_times)
+    print(f"\nPerformance for config ({num_tokens}, {num_heads}, {num_kv_heads}):")
+
+    print(
+        f"Torch implementation: "
+        f"mean={torch_stats['mean']:.8f}s, "
+        f"median={torch_stats['median']:.8f}s, "
+        f"p99={torch_stats['p99']:.8f}s"
+    )
+
+    print(
+        f"Triton implementation: "
+        f"mean={triton_stats['mean']:.8f}s, "
+        f"median={triton_stats['median']:.8f}s, "
+        f"p99={triton_stats['p99']:.8f}s"
+    )
+
+    print(
+        f"Triton Speedup over Torch: {torch_stats['mean'] / triton_stats['mean']:.8f}x"
+    )
+
+    return torch_stats, triton_stats
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Benchmark the rotary embedding kernels."
+    )
+    parser.add_argument("--model-name", type=str, default="")
+    parser.add_argument("--tp-size", type=int, default=1)
+    parser.add_argument("--warmup-iter", type=int, default=10)
+    parser.add_argument("--benchmark-iter", type=int, default=100)
+    parser.add_argument("--dtype", type=str, choices=["bfloat16"], default="bfloat16")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--num-tokens", type=int, nargs="+", required=False)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    args = parser.parse_args()
+    print(args)
+
+    model_tp_dict = {}
+    if args.model_name == "":
+        model_tp_dict = {
+            "Qwen/Qwen2-VL-2B-Instruct": [1],
+            "Qwen/Qwen2-VL-7B-Instruct": [1],
+            "Qwen/Qwen2-VL-72B-Instruct": [2, 4, 8],
+            "Qwen/Qwen2.5-VL-3B-Instruct": [1, 2, 4, 8],
+            "Qwen/Qwen2.5-VL-7B-Instruct": [1, 2, 4, 8],
+            "Qwen/Qwen2.5-VL-72B-Instruct": [2, 4, 8],
+        }
+    else:
+        model_tp_dict[args.model_name] = [args.tp_size]
+
+    if args.num_tokens is None:
+        num_tokens_list = [2**i for i in range(0, 18)]
+    else:
+        num_tokens_list = args.num_tokens
+
+    for model_name, tp_list in model_tp_dict.items():
+        for tp_size in tp_list:
+            config = get_model_config(model_name)
+            # get the model config
+            total_num_kv_heads = config.num_key_value_heads
+            total_num_heads = config.num_attention_heads
+            num_heads = total_num_heads // tp_size
+            num_kv_heads = max(1, total_num_kv_heads // tp_size)
+            head_dim = config.hidden_size // total_num_heads
+            is_neox_style = True
+            rope_theta = config.rope_theta
+            max_position = config.max_position_embeddings
+
+            for num_tokens in num_tokens_list:
+                benchmark_mrope(
+                    model_name=model_name,
+                    num_tokens=num_tokens,
+                    head_dim=head_dim,
+                    tp_size=tp_size,
+                    num_heads=num_heads,
+                    num_kv_heads=num_kv_heads,
+                    max_position=max_position,
+                    rope_theta=rope_theta,
+                    is_neox_style=is_neox_style,
+                    rope_scaling=config.rope_scaling,
+                    dtype=getattr(torch, args.dtype),
+                    seed=args.seed,
+                    warmup_iter=args.warmup_iter,
+                    benchmark_iter=args.benchmark_iter,
+                )

third_party/sglang/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py

@@ -0,0 +1,134 @@
+import itertools
+import math
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.utils import is_in_ci
+
+# Optional imports
+try:
+    from vllm import _custom_ops as ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    ops = None
+    VLLM_AVAILABLE = False
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+IS_CI = is_in_ci()
+
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+
+def vllm_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if not VLLM_AVAILABLE:
+        # Fallback to SGLang implementation
+        return sglang_scaled_fp8_quant(input, scale)
+    return ops.scaled_fp8_quant(input, scale)
+
+
+def sglang_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_type_: torch.dtype = torch.float8_e4m3fn
+    output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
+    is_static = True
+    if scale is None:
+        scale = torch.zeros(1, device=input.device, dtype=torch.float32)
+        is_static = False
+    per_tensor_quant_fp8(input, output, scale, is_static)
+
+    return output, scale
+
+
+def calculate_diff(batch_size: int, seq_len: int):
+    """Calculate difference between VLLM and SGLang implementations."""
+    device = torch.device("cuda")
+    x = torch.rand((batch_size, seq_len), dtype=torch.float16, device=device)
+
+    if not VLLM_AVAILABLE:
+        print("⚠️ vLLM not available, skipping comparison")
+        return
+
+    vllm_out, vllm_scale = vllm_scaled_fp8_quant(x)
+    sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
+
+    scale_diff = torch.abs(vllm_scale - sglang_scale).item()
+    output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+
+    if torch.allclose(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ) and torch.allclose(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-5):
+        print("✅ All implementations match")
+    else:
+        print("❌ Implementations differ")
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_size_range = [16]  # Single batch size for CI
+    seq_len_range = [64]  # Single sequence length for CI
+else:
+    batch_size_range = [16, 32, 64, 128]
+    seq_len_range = [64, 128, 256, 512, 1024, 2048]
+
+configs = list(itertools.product(batch_size_range, seq_len_range))
+
+
+if VLLM_AVAILABLE:
+    line_vals = ["vllm", "sglang"]
+    line_names = ["VLLM", "SGL Kernel"]
+    styles = [("blue", "-"), ("green", "-")]
+else:
+    line_vals = ["sglang"]
+    line_names = ["SGL Kernel"]
+    styles = [("green", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="us",
+        plot_name="per-tensor-quant-fp8-performance",
+        args={},
+    )
+)
+def benchmark(batch_size, seq_len, provider):
+    dtype = torch.float16
+    device = torch.device("cuda")
+
+    x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "vllm":
+        fn = lambda: vllm_scaled_fp8_quant(x.clone())
+    elif provider == "sglang":
+        fn = lambda: sglang_scaled_fp8_quant(x.clone())
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    calculate_diff(batch_size=4, seq_len=4096)
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py

@@ -0,0 +1,245 @@
+import itertools
+import os
+import time
+from functools import partial
+from pathlib import Path
+
+import torch
+import triton
+from sgl_kernel.test_utils import create_per_token_group_quant_test_data
+
+from sglang.srt.layers.quantization.fp8_kernel import (
+    create_per_token_group_quant_fp8_output_scale,
+)
+from sglang.srt.layers.quantization.fp8_kernel import (
+    per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
+)
+from sglang.srt.layers.quantization.fp8_kernel import (
+    sglang_per_token_group_quant_8bit,
+)
+from sglang.srt.utils import is_hip
+from sglang.srt.utils.bench_utils import bench_kineto
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+_is_hip = is_hip()
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+
+mode_concentrated = IS_CI or (os.environ.get("SGLANG_BENCH_MODE", "") == "concentrated")
+
+if int(os.environ.get("SGLANG_NSYS_PROFILING", "0")):
+    configs = [
+        [
+            768 * 8,
+            2048,
+            128,
+            48,
+            fp8_type_,
+            dict(
+                column_major_scales=True,
+                scale_tma_aligned=True,
+                scale_ue8m0=True,
+                fuse_silu_and_mul=True,
+                # masked_layout_mode=None,
+                masked_layout_mode="balanced",
+                # masked_layout_mode="extreme",
+            ),
+        ]
+    ]
+elif mode_concentrated:
+    configs = list(
+        itertools.product(
+            [768],
+            [1536, 7168, 16384],
+            [128],
+            [None],
+            [fp8_type_],
+            [
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=False,
+                    masked_layout_mode=None,
+                ),
+            ],
+        )
+    ) + list(
+        itertools.product(
+            [768 * 8],
+            [2048],
+            [128],
+            [48],
+            [fp8_type_],
+            [
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode=None,
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode="balanced",
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode="imbalanced",
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode="extreme",
+                ),
+            ],
+        )
+    )
+else:
+    configs = list(
+        itertools.product(
+            [1, 4, 16, 64, 256, 768, 2048, 8192, 16384],
+            [1536, 7168, 16384],
+            [128],
+            [None],
+            [fp8_type_],
+            [
+                dict(
+                    column_major_scales=False,
+                    scale_tma_aligned=False,
+                    scale_ue8m0=False,
+                    fuse_silu_and_mul=False,
+                    masked_layout_mode=None,
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=False,
+                    scale_ue8m0=False,
+                    fuse_silu_and_mul=False,
+                    masked_layout_mode=None,
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=False,
+                    fuse_silu_and_mul=False,
+                    masked_layout_mode=None,
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=False,
+                    masked_layout_mode=None,
+                ),
+            ],
+        )
+    ) + list(
+        itertools.product(
+            [1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
+            [2048],
+            [128],
+            [8, 16, 32, 48],
+            [fp8_type_],
+            [
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode=None,
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode="balanced",
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode="imbalanced",
+                ),
+                dict(
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    scale_ue8m0=True,
+                    fuse_silu_and_mul=True,
+                    masked_layout_mode="extreme",
+                ),
+            ],
+        )
+    )
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=[
+            "num_tokens",
+            "hidden_dim",
+            "group_size",
+            "num_ranks",
+            "dst_dtype",
+            "flags",
+        ],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["triton", "sglang"],
+        # Triton has multi kernels and we only report the time for the core one
+        line_names=["Triton (Inaccurate)", "SGL Kernel"],
+        styles=[("blue", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="per-token-group-quant-8bit-performance",
+        args={},
+    )
+)
+def benchmark(
+    num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, provider
+):
+    print(
+        f"Testing: {num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=} {provider=}"
+    )
+
+    x, masked_m = create_per_token_group_quant_test_data(
+        num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
+    )
+
+    fn, kernel_names = {
+        "triton": (
+            triton_per_token_group_quant_8bit,
+            "_per_token_group_quant_8bit|_silu_and_mul_post_quant_kernel",
+        ),
+        "sglang": (
+            partial(sglang_per_token_group_quant_8bit, enable_v2=True),
+            "per_token_group_quant_8bit_kernel",
+        ),
+    }[provider]
+    bench_fn = lambda: fn(
+        x=x,
+        masked_m=masked_m,
+        group_size=group_size,
+        dst_dtype=dst_dtype,
+        **{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
+    )
+
+    time_s = bench_kineto(
+        bench_fn, kernel_names=kernel_names, num_tests=300 if mode_concentrated else 30
+    )
+    return time_s * 1e6
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_per_token_quant_fp8.py

@@ -0,0 +1,228 @@
+import itertools
+import os
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel import sgl_per_token_quant_fp8
+
+from sglang.utils import is_in_ci
+
+# Optional vLLM import
+try:
+    from vllm import _custom_ops as ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    ops = None
+    VLLM_AVAILABLE = False
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+IS_CI = is_in_ci()
+
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+# Get correct FP8 E4M3 maximum value
+if _is_hip:
+    FP8_E4M3_MAX = 224.0  # ROCM uses 224.0
+else:
+    # For CUDA, get the actual max value from the type
+    FP8_E4M3_MAX = float(torch.finfo(fp8_type_).max)
+
+
+def torch_per_token_quant_fp8(
+    input: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Pure PyTorch reference implementation for per-token FP8 quantization."""
+    device = input.device
+    dtype = input.dtype
+
+    # Find max absolute value per token (row) - exactly like CUDA kernel
+    max_vals = torch.abs(input).max(dim=1)[0]  # [num_tokens]
+
+    # Calculate scale per token - exactly like CUDA kernel: scale = max_value / FP8_E4M3_MAX
+    scales = max_vals / FP8_E4M3_MAX  # [num_tokens]
+
+    # No special zero handling - directly compute 1.0 / scale like CUDA kernel
+    scale_inv = 1.0 / scales  # [num_tokens]
+
+    # Quantize: input * scale_inv, then clamp to FP8 range
+    quantized_float = input * scale_inv.unsqueeze(1)  # Broadcast scale_inv
+    quantized_float = torch.clamp(quantized_float, -FP8_E4M3_MAX, FP8_E4M3_MAX)
+
+    # Convert to FP8 - use more explicit conversion
+    quantized_fp8 = quantized_float.to(fp8_type_)
+
+    return quantized_fp8, scales
+
+
+def vllm_per_token_quant_fp8(
+    input: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if not VLLM_AVAILABLE:
+        # Fallback to SGLang implementation
+        return sglang_per_token_quant_fp8(input)
+    return ops.scaled_fp8_quant(input, use_per_token_if_dynamic=True)
+
+
+def sglang_per_token_quant_fp8(
+    input: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    scale = torch.zeros(input.size(0), device=input.device, dtype=torch.float32)
+    output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
+    sgl_per_token_quant_fp8(input, output, scale)
+
+    return output, scale
+
+
+def calculate_diff(batch_size: int, seq_len: int, hidden_dim: int):
+    """Compare Torch reference, VLLM, and SGLang implementations."""
+    device = torch.device("cuda")
+    x = torch.rand(
+        (batch_size * seq_len, hidden_dim), dtype=torch.float16, device=device
+    )
+
+    # Get all three implementations
+    torch_out, torch_scale = torch_per_token_quant_fp8(x)
+    vllm_out, vllm_scale = vllm_per_token_quant_fp8(x)
+    sglang_out, sglang_scale = sglang_per_token_quant_fp8(x)
+
+    if not VLLM_AVAILABLE:
+        print("⚠️ vLLM not available, skipping vLLM comparison")
+        # Only compare Torch vs SGLang
+        torch_sglang_scale_diff = torch.abs(torch_scale - sglang_scale).mean().item()
+        torch_sglang_out_diff = (
+            torch.abs(torch_out.float() - sglang_out.float()).mean().item()
+        )
+        print(f"Scale difference (Torch vs SGLang): {torch_sglang_scale_diff:.8f}")
+        print(f"Output difference (Torch vs SGLang): {torch_sglang_out_diff:.8f}")
+        return
+
+    print(f"\n=== Comparison for hidden_dim={hidden_dim} ===")
+
+    # Compare scales
+    torch_vllm_scale_diff = torch.abs(torch_scale - vllm_scale).mean().item()
+    torch_sglang_scale_diff = torch.abs(torch_scale - sglang_scale).mean().item()
+    vllm_sglang_scale_diff = torch.abs(vllm_scale - sglang_scale).mean().item()
+
+    print(f"Scale differences:")
+    print(f"  Torch vs VLLM:   {torch_vllm_scale_diff:.8f}")
+    print(f"  Torch vs SGLang: {torch_sglang_scale_diff:.8f}")
+    print(f"  VLLM vs SGLang:  {vllm_sglang_scale_diff:.8f}")
+
+    # Compare outputs
+    torch_vllm_out_diff = torch.abs(torch_out.float() - vllm_out.float()).mean().item()
+    torch_sglang_out_diff = (
+        torch.abs(torch_out.float() - sglang_out.float()).mean().item()
+    )
+    vllm_sglang_out_diff = (
+        torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+    )
+
+    print(f"Output differences:")
+    print(f"  Torch vs VLLM:   {torch_vllm_out_diff:.8f}")
+    print(f"  Torch vs SGLang: {torch_sglang_out_diff:.8f}")
+    print(f"  VLLM vs SGLang:  {vllm_sglang_out_diff:.8f}")
+
+    # Check tolerances
+    rtol, atol = 1e-3, 1e-5
+
+    torch_vllm_match = torch.allclose(
+        torch_out.float(), vllm_out.float(), rtol=rtol, atol=atol
+    ) and torch.allclose(torch_scale, vllm_scale, rtol=rtol, atol=atol)
+    torch_sglang_match = torch.allclose(
+        torch_out.float(), sglang_out.float(), rtol=rtol, atol=atol
+    ) and torch.allclose(torch_scale, sglang_scale, rtol=rtol, atol=atol)
+
+    if hidden_dim == 1368:
+        rtol = 1e-2
+        # we found vllm sglang has diff when hidden dim is not dividable by 16
+        # and we believe SGLang is closer to Torch implementation
+
+    vllm_sglang_match = torch.allclose(
+        vllm_out.float(), sglang_out.float(), rtol=rtol, atol=atol
+    ) and torch.allclose(vllm_scale, sglang_scale, rtol=rtol, atol=atol)
+
+    print(f"Matches (rtol={rtol}, atol={atol}):")
+    print(f"  Torch vs VLLM:   {'✅' if torch_vllm_match else '❌'}")
+    print(f"  Torch vs SGLang: {'✅' if torch_sglang_match else '❌'}")
+    print(f"  VLLM vs SGLang:  {'✅' if vllm_sglang_match else '❌'}")
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_size_range = [16]  # Single batch size for CI
+    seq_len_range = [64]  # Single sequence length for CI
+    hidden_dim_range = [2048]  # Single hidden dimension for CI
+else:
+    batch_size_range = [16, 32, 64, 128]
+    seq_len_range = [64, 128, 256, 512, 1024, 2048, 4096]
+    hidden_dim_range = [1368, 2048, 4096]
+
+configs = list(itertools.product(batch_size_range, seq_len_range, hidden_dim_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len", "hidden_dim"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=(
+            ["torch", "vllm", "sglang"] if VLLM_AVAILABLE else ["torch", "sglang"]
+        ),
+        line_names=(
+            ["Torch Reference", "VLLM", "SGL Kernel"]
+            if VLLM_AVAILABLE
+            else ["Torch Reference", "SGL Kernel"]
+        ),
+        styles=(
+            [("red", "-"), ("blue", "-"), ("green", "-")]
+            if VLLM_AVAILABLE
+            else [("red", "-"), ("green", "-")]
+        ),
+        ylabel="us",
+        plot_name="per-token-dynamic-quant-fp8-performance",
+        args={},
+    )
+)
+def benchmark_quantization(batch_size, seq_len, hidden_dim, provider):
+    dtype = torch.float16
+    device = torch.device("cuda")
+
+    x = torch.randn(batch_size * seq_len, hidden_dim, device=device, dtype=dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "torch":
+        fn = lambda: torch_per_token_quant_fp8(x.clone())
+    elif provider == "vllm":
+        if not VLLM_AVAILABLE:
+            return (0, 0, 0)
+        fn = lambda: vllm_per_token_quant_fp8(x.clone())
+    elif provider == "sglang":
+        fn = lambda: sglang_per_token_quant_fp8(x.clone())
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    # Test various hidden dimensions for correctness - simplified for CI
+    if IS_CI:
+        test_dims = [2048]  # Single dimension for CI
+        batch_size, seq_len = 4, 64  # Smaller values for CI
+    else:
+        test_dims = [1368, 2048, 4096]
+        batch_size, seq_len = 4, 4096
+
+    for dim in test_dims:
+        calculate_diff(batch_size=batch_size, seq_len=seq_len, hidden_dim=dim)
+
+    print("\n" + "=" * 60)
+    print("Starting performance benchmark...")
+    benchmark_quantization.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_qserve_w4a8_gemm.py

@@ -0,0 +1,214 @@
+import argparse
+import copy
+import itertools
+import os
+
+import torch
+import triton
+from sgl_kernel import (
+    int8_scaled_mm,
+    qserve_w4a8_per_chn_gemm,
+    qserve_w4a8_per_group_gemm,
+)
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+def to_int8(tensor: torch.Tensor) -> torch.Tensor:
+    return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
+
+
+WEIGHT_SHAPES = {
+    "meta-llama/Llama-3.1-8B-Instruct": [
+        ([4096, 6144], 1),
+        ([4096, 4096], 0),
+        ([4096, 28672], 1),
+        ([14336, 4096], 0),
+    ],
+    "meta-llama/Llama-3.3-70B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 57344], 1),
+        ([28672, 8192], 0),
+    ],
+    "mistralai/Mistral-Large-Instruct-2407": [
+        ([12288, 14336], 1),
+        ([12288, 12288], 0),
+        ([12288, 57344], 1),
+        ([28672, 12288], 0),
+    ],
+    "Qwen/Qwen2.5-7B-Instruct": [
+        ([3584, 4608], 1),
+        ([3584, 3584], 0),
+        ([3584, 37888], 1),
+        ([18944, 3584], 0),
+    ],
+    "Qwen/Qwen2.5-32B-Instruct": [
+        ([5120, 7168], 1),
+        ([5120, 5120], 0),
+        ([5120, 55296], 1),
+        ([27648, 5120], 0),
+    ],
+    "Qwen/Qwen2.5-72B-Instruct": [
+        ([8192, 10240], 1),
+        ([8192, 8192], 0),
+        ([8192, 59136], 1),
+        ([29568, 8192], 0),
+    ],
+    "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
+        ([2048, 3072], 1),
+        ([2048, 4096], 1),
+        ([2048, 2048], 0),
+        ([2048, 576], 0),
+        ([2048, 21888], 1),
+        ([10944, 2048], 0),
+        ([2048, 2816], 1),
+        ([1408, 2048], 0),
+    ],
+}
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_sizes = [1, 16]  # Simplified for CI
+else:
+    batch_sizes = [1, 16, 32, 64, 128, 256, 512, 1024, 2048]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=batch_sizes,
+        x_log=False,
+        line_arg="provider",
+        line_vals=["FP16", "W8A8", "Qserve_W4A8_Per_Channel", "Qserve_W4A8_Per_Group"],
+        line_names=["FP16", "W8A8", "Qserve_W4A8_Per_Channel", "Qserve_W4A8_Per_Group"],
+        styles=[("blue", "-"), ("orange", "-"), ("green", "-"), ("red", "-")],
+        ylabel="ms",
+        plot_name="FP16_vs_W8A8_vs_Qserve_W4A8_GEMM",
+        args={},
+    )
+)
+def benchmark(batch_size, provider, N, K):
+    M = batch_size
+    # For W8A8
+    a = to_int8(torch.randn((M, K), device="cuda") * 5)
+    b = to_int8(torch.randn((N, K), device="cuda").t() * 5)
+    a_fp16 = a.to(torch.float16)
+    b_fp16 = b.to(torch.float16)
+    scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
+    scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
+
+    # For Qserve W4A8 per channel
+    a_qserve_chn = a
+    # two int4s pack into one int8
+    b_qserve_chn = to_int8(torch.randn((N, K // 2), device="cuda") * 5)
+    # b_qserve_chn = b.t().contiguous()
+    scale_a_qserve_chn = scale_a.to(torch.float16)
+    scale_b_qserve_chn = scale_b.to(torch.float16)
+    szero_b_qserve_chn = torch.randn((N,), device="cuda", dtype=torch.float16)
+    a_sum_qserve_chn = torch.randn((M,), device="cuda", dtype=torch.float16)
+
+    # For Qserve W4A8 per group
+    group_size = 128
+    assert K % group_size == 0, "K must be divisible by group_size"
+    a_qserve_group = a
+    # two int4s pack into one int8
+    b_qserve_group = to_int8(torch.randn((N, K // 2), device="cuda") * 5)
+    # b_qserve_group = b.t().contiguous()
+    scale_a_qserve_group = scale_a.to(torch.float16)
+    scale_b_qserve_group = scale_b.to(torch.float16)
+    scale_i8_b_qserve_group = to_int8(
+        torch.randn((K // group_size, N), device="cuda", dtype=torch.float16)
+    )
+    zero_i8_b_qserve_group = to_int8(
+        torch.randn((K // group_size, N), device="cuda", dtype=torch.float16)
+    )
+
+    quantiles = [0.5, 0.2, 0.8]
+    if provider == "FP16":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: torch.matmul(a_fp16, b_fp16),
+            quantiles=quantiles,
+        )
+    if provider == "W8A8":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: int8_scaled_mm(a, b, scale_a, scale_b, torch.float16),
+            quantiles=quantiles,
+        )
+    if provider == "Qserve_W4A8_Per_Channel":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: qserve_w4a8_per_chn_gemm(
+                a_qserve_chn,
+                b_qserve_chn,
+                scale_b_qserve_chn,
+                scale_a_qserve_chn,
+                szero_b_qserve_chn,
+                a_sum_qserve_chn,
+            ),
+            quantiles=quantiles,
+        )
+    if provider == "Qserve_W4A8_Per_Group":
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: qserve_w4a8_per_group_gemm(
+                a_qserve_group,
+                b_qserve_group,
+                zero_i8_b_qserve_group,
+                scale_i8_b_qserve_group,
+                scale_b_qserve_group,
+                scale_a_qserve_group,
+            ),
+            quantiles=quantiles,
+        )
+
+    return ms, max_ms, min_ms
+
+
+def prepare_shapes(args):
+    KN_model_names = []
+    models_tps = list(itertools.product(args.models, args.tp_sizes))
+    for model, tp_size in models_tps:
+        assert model in WEIGHT_SHAPES
+        for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
+            KN[tp_split_dim] = KN[tp_split_dim] // tp_size
+            KN.append(model)
+            KN_model_names.append(KN)
+    return KN_model_names
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        type=str,
+        default=["meta-llama/Llama-3.1-8B-Instruct"],
+        help="List of models to benchmark",
+    )
+    parser.add_argument(
+        "--tp-sizes",
+        nargs="+",
+        type=int,
+        default=[1],
+        help="List of tensor parallel sizes",
+    )
+    args = parser.parse_args()
+
+    # Skip in CI environment
+    if IS_CI:
+        print("Skipping QServe W4A8 GEMM benchmark in CI environment")
+        print("QServe operations may have compatibility issues in CI")
+    else:
+        KN_model_names = prepare_shapes(args)
+
+        for K, N, model_name in KN_model_names:
+            print(f"{model_name} N={N} K={K}: ")
+            benchmark.run(
+                print_data=True,
+                N=N,
+                K=K,
+            )
+
+        print("Benchmark finished!")

third_party/sglang/sgl-kernel/benchmark/bench_rmsnorm.py

@@ -0,0 +1,395 @@
+# Benchmarks SGLang RMSNorm kernels versus vLLM and FlashInfer across
+# (batch_size, seq_len, hidden_size) and prints speed-up.
+import argparse
+import itertools
+import os
+import re
+from typing import List, Optional, Tuple, Union
+
+import sgl_kernel
+import torch
+import torch.nn as nn
+import triton
+import triton.testing
+from sgl_kernel.utils import is_arch_support_pdl
+
+from sglang.utils import is_in_ci
+
+# Optional imports
+try:
+    from flashinfer.norm import fused_add_rmsnorm, rmsnorm
+
+    FLASHINFER_AVAILABLE = True
+except ImportError:
+    fused_add_rmsnorm = None
+    rmsnorm = None
+    FLASHINFER_AVAILABLE = False
+
+try:
+    from vllm import _custom_ops as vllm_ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    vllm_ops = None
+    VLLM_AVAILABLE = False
+
+IS_CI = is_in_ci()
+
+
+def str2int_list(arg: str) -> List[int]:
+    if arg in ("", None):
+        return []
+    if re.fullmatch(r"\d+(,\d+)*", arg.strip()) is None:
+        raise argparse.ArgumentTypeError(f"Bad int list: {arg}")
+    return [int(x) for x in arg.split(",")]
+
+
+class HuggingFaceRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        if residual is not None:
+            x = x + residual.to(torch.float32)
+            residual = x.to(orig_dtype)
+
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.variance_epsilon)
+        x = x.to(orig_dtype) * self.weight
+        if residual is None:
+            return x
+        else:
+            return x, residual
+
+
+def rmsnorm_naive(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+):
+    naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps)
+    naive_norm.weight = nn.Parameter(weight)
+    naive_norm = naive_norm.to(x.device)
+
+    orig_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    if residual is not None:
+        residual = residual.view(-1, residual.shape[-1])
+
+    output = naive_norm(x, residual)
+
+    if isinstance(output, tuple):
+        output = (output[0].view(orig_shape), output[1].view(orig_shape))
+    else:
+        output = output.view(orig_shape)
+    return output
+
+
+def rmsnorm_flashinfer(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+):
+    if not FLASHINFER_AVAILABLE:
+        # Fallback to naive implementation if FlashInfer is not available
+        return rmsnorm_naive(x, weight, residual, eps)
+
+    orig_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    if residual is not None:
+        residual = residual.view(-1, residual.shape[-1])
+
+    if residual is not None:
+        fused_add_rmsnorm(x, residual, weight, eps)
+        output = (x, residual)
+    else:
+        output = rmsnorm(x, weight, eps)
+
+    if isinstance(output, tuple):
+        output = (output[0].view(orig_shape), output[1].view(orig_shape))
+    else:
+        output = output.view(orig_shape)
+    return output
+
+
+def rmsnorm_vllm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+):
+    if not VLLM_AVAILABLE:
+        # Fallback to naive implementation if vLLM is not available
+        return rmsnorm_naive(x, weight, residual, eps)
+
+    orig_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    if residual is not None:
+        residual = residual.view(-1, residual.shape[-1])
+
+    if residual is not None:
+        vllm_ops.fused_add_rms_norm(x, residual, weight, eps)
+        output = (x, residual)
+    else:
+        out = torch.empty_like(x)
+        vllm_ops.rms_norm(out, x, weight, eps)
+        output = out
+
+    if isinstance(output, tuple):
+        output = (output[0].view(orig_shape), output[1].view(orig_shape))
+    else:
+        output = output.view(orig_shape)
+    return output
+
+
+def rmsnorm_sglang(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+):
+    orig_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    if residual is not None:
+        residual = residual.view(-1, residual.shape[-1])
+
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+
+    if residual is not None:
+        sgl_kernel.fused_add_rmsnorm(x, residual, weight, eps, enable_pdl=enable_pdl)
+        output = (x, residual)
+    else:
+        out = torch.empty_like(x)
+        sgl_kernel.rmsnorm(x, weight, eps, out=out, enable_pdl=enable_pdl)
+        output = out
+
+    if isinstance(output, tuple):
+        output = (output[0].view(orig_shape), output[1].view(orig_shape))
+    else:
+        output = output.view(orig_shape)
+    return output
+
+
+def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):
+    dtype = torch.bfloat16
+    x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")
+    weight = torch.ones(hidden_size, dtype=dtype, device="cuda")
+    residual = torch.randn_like(x) if use_residual else None
+
+    output_naive = rmsnorm_naive(
+        x.clone(), weight, residual.clone() if residual is not None else None
+    )
+    output_flashinfer = rmsnorm_flashinfer(
+        x.clone(), weight, residual.clone() if residual is not None else None
+    )
+    output_vllm = rmsnorm_vllm(
+        x.clone(), weight, residual.clone() if residual is not None else None
+    )
+    output_sglang = rmsnorm_sglang(
+        x.clone(), weight, residual.clone() if residual is not None else None
+    )
+
+    if use_residual:
+        output_naive = output_naive[0]
+        output_flashinfer = output_flashinfer[0]
+        output_vllm = output_vllm[0]
+        output_sglang = output_sglang[0]
+
+    print(f"Naive output={output_naive}")
+    if FLASHINFER_AVAILABLE:
+        print(f"FlashInfer output={output_flashinfer}")
+    else:
+        print("FlashInfer not available, skipped")
+    if VLLM_AVAILABLE:
+        print(f"VLLM output={output_vllm}")
+    else:
+        print("vLLM not available, skipped")
+    print(f"SGLang output={output_sglang}")
+
+    # Only compare available implementations
+    all_match = torch.allclose(output_naive, output_sglang, atol=1e-2, rtol=1e-2)
+    if FLASHINFER_AVAILABLE:
+        all_match = all_match and torch.allclose(
+            output_naive, output_flashinfer, atol=1e-2, rtol=1e-2
+        )
+    if VLLM_AVAILABLE:
+        all_match = all_match and torch.allclose(
+            output_naive, output_vllm, atol=1e-2, rtol=1e-2
+        )
+
+    if all_match:
+        print("✅ All available implementations match")
+    else:
+        print("❌ Implementations differ")
+
+
+# CI environment uses simplified parameters
+if IS_CI:
+    default_batch_sizes = [1]  # Single batch size for CI
+    default_seq_lens = [64]  # Single sequence length for CI
+    default_hidden_sizes = [4096]  # Single hidden size for CI
+else:
+    default_batch_sizes = [2**i for i in range(0, 7, 2)]  # 1, 4, 16, 64
+    default_seq_lens = [2**i for i in range(6, 11, 1)]  # 64, 128, 256, 512, 1024
+    default_hidden_sizes = [32 * 128, 48 * 128]  # 4096, 6144
+
+
+def make_configs(bsizes: List[int], slens: List[int], hsizes: List[int]) -> List[Tuple]:
+    return list(itertools.product(bsizes, slens, hsizes))
+
+
+# Filter providers based on availability
+available_providers = ["huggingface", "sglang"]
+available_names = ["HuggingFace", "SGL Kernel"]
+available_styles = [("blue", "-"), ("orange", "-")]
+
+if FLASHINFER_AVAILABLE:
+    available_providers.insert(-1, "flashinfer")
+    available_names.insert(-1, "FlashInfer")
+    available_styles.insert(-1, ("green", "-"))
+
+if VLLM_AVAILABLE:
+    available_providers.insert(-1, "vllm")
+    available_names.insert(-1, "vLLM")
+    available_styles.insert(-1, ("red", "-"))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len", "hidden_size"],
+        x_vals=[],
+        line_arg="provider",
+        line_vals=available_providers,
+        line_names=available_names,
+        styles=available_styles,
+        ylabel="µs (median)  or  × (speed-up)",
+        plot_name="rmsnorm-performance",
+        args={},
+    )
+)
+def benchmark(batch_size, seq_len, hidden_size, provider, use_residual):
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device)
+    weight = torch.ones(hidden_size, dtype=dtype, device=device)
+    residual = torch.randn_like(x) if use_residual else None
+
+    # timing helper
+    def timed(fn):
+        for _ in range(5):
+            fn()
+        torch.cuda.synchronize()
+        ms, qmin, qmax = triton.testing.do_bench_cudagraph(
+            fn, quantiles=[0.5, 0.2, 0.8]
+        )
+        return 1000 * ms, 1000 * qmax, 1000 * qmin
+
+    if provider == "huggingface":
+        return timed(
+            lambda: rmsnorm_naive(
+                x.clone(),
+                weight,
+                residual.clone() if residual is not None else None,
+            )
+        )
+    elif provider == "flashinfer":
+        if not FLASHINFER_AVAILABLE:
+            return (0, 0, 0)
+        return timed(
+            lambda: rmsnorm_flashinfer(
+                x.clone(),
+                weight,
+                residual.clone() if residual is not None else None,
+            )
+        )
+    elif provider == "vllm":
+        if not VLLM_AVAILABLE:
+            return (0, 0, 0)
+        return timed(
+            lambda: rmsnorm_vllm(
+                x.clone(),
+                weight,
+                residual.clone() if residual is not None else None,
+            )
+        )
+    elif provider == "sglang":
+        return timed(
+            lambda: rmsnorm_sglang(
+                x.clone(),
+                weight,
+                residual.clone() if residual is not None else None,
+            )
+        )
+
+    # provider == "speedup"
+    if VLLM_AVAILABLE:
+        t_ref, _, _ = timed(
+            lambda: rmsnorm_vllm(
+                x.clone(),
+                weight,
+                residual.clone() if residual is not None else None,
+            )
+        )
+    else:
+        t_ref, _, _ = timed(
+            lambda: rmsnorm_naive(
+                x.clone(),
+                weight,
+                residual.clone() if residual is not None else None,
+            )
+        )
+    t_sgl, _, _ = timed(
+        lambda: rmsnorm_sglang(
+            x.clone(),
+            weight,
+            residual.clone() if residual is not None else None,
+        )
+    )
+    spd = t_ref / t_sgl if t_ref > 0 else 1.0
+    return (spd, spd, spd)
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser("RMSNorm kernel benchmark")
+    p.add_argument("--batch_sizes", type=str2int_list, default=default_batch_sizes)
+    p.add_argument("--seq_lens", type=str2int_list, default=default_seq_lens)
+    p.add_argument("--hidden_sizes", type=str2int_list, default=default_hidden_sizes)
+    p.add_argument(
+        "--use_residual", action="store_true", help="Whether to use residual connection"
+    )
+    p.add_argument("--verify_only", action="store_true")
+    args = p.parse_args()
+
+    # coerce lists
+    if isinstance(args.batch_sizes, str):
+        args.batch_sizes = str2int_list(args.batch_sizes)
+    if isinstance(args.seq_lens, str):
+        args.seq_lens = str2int_list(args.seq_lens)
+    if isinstance(args.hidden_sizes, str):
+        args.hidden_sizes = str2int_list(args.hidden_sizes)
+
+    # patch perf_report grid
+    benchmark_grid = make_configs(args.batch_sizes, args.seq_lens, args.hidden_sizes)
+    if hasattr(benchmark, "benchmarks"):
+        benchmark.benchmarks.x_vals = benchmark_grid
+    else:
+        benchmark.benchmark.x_vals = benchmark_grid
+
+    if args.verify_only:
+        ok = calculate_diff(4, 128, args.hidden_sizes[0], args.use_residual)
+        print("✅ sanity pass" if ok else "❌ mismatch")
+    else:
+        benchmark.run(print_data=True, use_residual=args.use_residual)

third_party/sglang/sgl-kernel/benchmark/bench_rotary_embedding.py

@@ -0,0 +1,108 @@
+import itertools
+import os
+
+import torch
+import triton
+from sgl_kernel.testing.rotary_embedding import (
+    FlashInferRotaryEmbedding,
+    FusedSetKVBufferArg,
+    MHATokenToKVPool,
+    RotaryEmbedding,
+    create_inputs,
+)
+
+from sglang.srt.utils.bench_utils import bench_kineto
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+# CI environment uses simplified parameters
+if IS_CI:
+    batch_seq_configs = [(1, 1)]  # Single config for CI
+    save_kv_configs = [False]  # Single option for CI
+else:
+    batch_seq_configs = [
+        (1, 1),
+        (32, 1),
+        (128, 1),
+        (512, 1),
+        (2, 512),
+        (4, 4096),
+    ]
+    save_kv_configs = [False, True]
+
+configs = [
+    (batch_size, seq_len, save_kv_cache)
+    for batch_size, seq_len in batch_seq_configs
+    for save_kv_cache in save_kv_configs
+]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len", "save_kv_cache"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["sglang"],
+        line_names=["SGL Kernel"],
+        styles=[("green", "-")],
+        ylabel="us",
+        plot_name="bench_rotary_embedding",
+        args={},
+    )
+)
+def benchmark(batch_size, seq_len, save_kv_cache, provider):
+    device = torch.device("cuda")
+
+    num_q_heads = 32
+    num_kv_heads = 8
+    head_size = 64
+    dtype = torch.bfloat16
+
+    config = dict(
+        head_size=head_size,
+        rotary_dim=64,
+        max_position_embeddings=4096,
+        base=8000,
+        is_neox_style=True,
+        dtype=dtype,
+    )
+    rope_flashinfer = FlashInferRotaryEmbedding(**config).to(device)
+    pool_flashinfer = MHATokenToKVPool(head_num=num_kv_heads, head_dim=head_size)
+
+    inputs = create_inputs(
+        head_size=head_size,
+        batch_size=batch_size,
+        seq_len=seq_len,
+        device=device,
+        dtype=dtype,
+        num_q_heads=num_q_heads,
+        num_kv_heads=num_kv_heads,
+    )
+
+    query_flashinfer, key_flashinfer = inputs["query"].clone(), inputs["key"].clone()
+
+    bench_fn = lambda: rope_flashinfer.forward_cuda(
+        inputs["pos_ids"],
+        query_flashinfer,
+        key_flashinfer,
+        fused_set_kv_buffer_arg=(
+            FusedSetKVBufferArg(
+                value=inputs["value"],
+                k_buffer=pool_flashinfer.k_buffer[0].view(-1, num_kv_heads * head_size),
+                v_buffer=pool_flashinfer.v_buffer[0].view(-1, num_kv_heads * head_size),
+                k_scale=None,
+                v_scale=None,
+                cache_loc=inputs["out_cache_loc"],
+            )
+            if save_kv_cache
+            else None
+        ),
+    )
+
+    time_s = bench_kineto(bench_fn, kernel_names="BatchQKApplyRotaryPosIds")
+    return time_s * 1e6
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

third_party/sglang/sgl-kernel/benchmark/bench_sum_scale.py

@@ -0,0 +1,249 @@
+import os
+
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import moe_sum_reduce as moe_sum_reduce_cuda
+from triton.testing import do_bench
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+@triton.jit
+def _moe_sum_reduce_kernel(
+    input_ptr,
+    input_stride_0,
+    input_stride_1,
+    input_stride_2,
+    output_ptr,
+    output_stride_0,
+    output_stride_1,
+    token_num: int,
+    topk_num: int,
+    hidden_dim: int,
+    routed_scaling_factor: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DIM: tl.constexpr,
+    NUM_STAGE: tl.constexpr,
+):
+    input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
+    input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
+    output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64)
+
+    token_block_id = tl.program_id(0)
+    dim_block_id = tl.program_id(1)
+
+    offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM)
+
+    mask_token = offs_token < token_num
+    mask_dim = offs_dim < hidden_dim
+
+    base_ptrs = input_ptr + offs_token[:, None] * input_stride_0 + offs_dim[None, :]
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32)
+    for i in tl.range(0, topk_num, num_stages=NUM_STAGE):
+        tile = tl.load(
+            base_ptrs + i * input_stride_1,
+            mask=mask_token[:, None] & mask_dim[None, :],
+            other=0.0,
+        )
+        accumulator += tile.to(tl.float32)
+    accumulator *= routed_scaling_factor
+
+    # -------- Write back --------
+    store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :]
+    tl.store(
+        store_ptrs,
+        accumulator.to(input_ptr.dtype.element_ty),
+        mask=mask_token[:, None] & mask_dim[None, :],
+    )
+
+
+# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py
+def moe_sum_reduce_triton(
+    input: torch.Tensor, output: torch.Tensor, routed_scaling_factor: float
+):
+    assert input.is_contiguous()
+    assert output.is_contiguous()
+
+    token_num, topk_num, hidden_dim = input.shape
+    assert output.shape[0] == token_num and output.shape[1] == hidden_dim
+
+    BLOCK_M = 1
+    BLOCK_DIM = 2048
+    NUM_STAGE = 1
+    num_warps = 16
+
+    grid = (
+        triton.cdiv(token_num, BLOCK_M),
+        triton.cdiv(hidden_dim, BLOCK_DIM),
+    )
+
+    _moe_sum_reduce_kernel[grid](
+        input,
+        *input.stride(),
+        output,
+        *output.stride(),
+        token_num=token_num,
+        topk_num=topk_num,
+        hidden_dim=hidden_dim,
+        routed_scaling_factor=routed_scaling_factor,
+        BLOCK_M=BLOCK_M,
+        BLOCK_DIM=BLOCK_DIM,
+        NUM_STAGE=NUM_STAGE,
+        num_warps=num_warps,
+    )
+    return
+
+
+def compute_sum_scaled_baseline(
+    x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float
+) -> torch.Tensor:
+    torch.sum(x, dim=1, out=out)
+    out.mul_(routed_scaling_factor)
+    return out
+
+
+@torch.compile
+def compute_sum_scaled_compiled(
+    x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float
+) -> torch.Tensor:
+    torch.sum(x * routed_scaling_factor, dim=1, out=out)
+    return out
+
+
+def get_benchmark(dtype=torch.bfloat16):
+    num_tokens_range = [2**i for i in range(0, 13)]
+
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["num_tokens"],
+            x_vals=num_tokens_range,
+            line_arg="version",
+            line_vals=["baseline", "compiled", "triton", "cuda"],
+            line_names=["Original", "TorchCompile", "TritonKernel", "CudaKernel"],
+            styles=[("blue", "-"), ("green", "-"), ("red", "-"), ("yellow", "-")],
+            ylabel="us",
+            plot_name=f"sum_scaled_performance_{str(dtype).split('.')[-1]}",
+            args={},
+        )
+    )
+    def benchmark(num_tokens, version):
+        topk = 9
+        hidden_size = 4096
+        dtype = torch.bfloat16
+        scaling_factor = 0.3
+
+        x = torch.randn(num_tokens, topk, hidden_size, dtype=dtype, device="cuda")
+        out = torch.empty(num_tokens, hidden_size, dtype=dtype, device="cuda")
+
+        # Warmup
+        for _ in range(3):
+            if version == "baseline":
+                compute_sum_scaled_baseline(x, out, scaling_factor)
+            elif version == "compiled":
+                compute_sum_scaled_compiled(x, out, scaling_factor)
+            elif version == "triton":
+                moe_sum_reduce_triton(x, out, scaling_factor)
+            else:
+                moe_sum_reduce_cuda(x, out, scaling_factor)
+
+        # Benchmark
+        quantiles = [0.5, 0.2, 0.8]
+        if version == "baseline":
+            ms, min_ms, max_ms = do_bench(
+                lambda: compute_sum_scaled_baseline(x, out, scaling_factor),
+                quantiles=quantiles,
+            )
+        elif version == "compiled":
+            ms, min_ms, max_ms = do_bench(
+                lambda: compute_sum_scaled_compiled(x, out, scaling_factor),
+                quantiles=quantiles,
+            )
+        elif version == "triton":
+            ms, min_ms, max_ms = do_bench(
+                lambda: moe_sum_reduce_triton(x, out, scaling_factor),
+                quantiles=quantiles,
+            )
+        else:
+            ms, min_ms, max_ms = do_bench(
+                lambda: moe_sum_reduce_cuda(x, out, scaling_factor),
+                quantiles=quantiles,
+            )
+
+        return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+    return benchmark
+
+
+def verify_correctness(num_tokens=1024, dtype=torch.bfloat16):
+    x = torch.randn(num_tokens, 9, 4096, device="cuda", dtype=dtype)
+    scaling_factor = 0.3
+
+    out_baseline = torch.empty_like(x[:, 0])
+    compute_sum_scaled_baseline(x, out_baseline, scaling_factor)
+
+    out_compiled = torch.empty_like(out_baseline)
+    compute_sum_scaled_compiled(x, out_compiled, scaling_factor)
+
+    out_cuda = torch.empty_like(out_baseline)
+    moe_sum_reduce_cuda(x, out_cuda, scaling_factor)
+
+    triton_skipped = dtype == torch.float64
+    if not triton_skipped:
+        out_triton = torch.empty_like(out_baseline)
+        moe_sum_reduce_triton(x, out_triton, scaling_factor)
+
+    if dtype == torch.float64:
+        atol, rtol = 1e-12, 1e-12
+    elif dtype == torch.float32:
+        atol, rtol = 1e-6, 1e-6
+    else:  # bfloat16 / float16
+        atol, rtol = 1e-2, 1e-2
+
+    ok_compiled = torch.allclose(out_baseline, out_compiled, atol=atol, rtol=rtol)
+    ok_cuda = torch.allclose(out_baseline, out_cuda, atol=atol, rtol=rtol)
+    ok_triton = (
+        True
+        if triton_skipped
+        else torch.allclose(out_baseline, out_triton, atol=atol, rtol=rtol)
+    )
+
+    if ok_compiled and ok_triton and ok_cuda:
+        msg = "✅ All implementations match"
+        if triton_skipped:
+            msg += " (Triton skipped for float64)"
+        print(msg)
+    else:
+        print("❌ Implementations differ")
+        print(
+            f"Baseline vs Compiled: {(out_baseline - out_compiled).abs().max().item()}"
+        )
+        if not triton_skipped:
+            print(
+                f"Baseline vs Triton: {(out_baseline - out_triton).abs().max().item()}"
+            )
+        print(f"Baseline vs Cuda: {(out_baseline - out_cuda).abs().max().item()}")
+
+
+if __name__ == "__main__":
+    print("Running correctness verification for bfloat16...")
+    verify_correctness(dtype=torch.bfloat16)
+
+    # CI environment uses simplified parameters
+    if not IS_CI:
+        print("Running correctness verification for float64...")
+        verify_correctness(dtype=torch.float64)
+
+    print("Running correctness verification for float64...")
+    verify_correctness(dtype=torch.float64)
+
+    print("\nRunning performance benchmark for bfloat16...")
+    benchmark = get_benchmark(dtype=torch.bfloat16)
+    benchmark.run(
+        print_data=True,
+        # save_path="./configs/benchmark_ops/sum_scaled/"
+    )

third_party/sglang/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py

@@ -0,0 +1,146 @@
+import itertools
+import os
+
+import flashinfer.sampling
+import sgl_kernel
+import torch
+import triton
+import triton.testing
+
+from sglang.utils import is_in_ci
+
+IS_CI = is_in_ci()
+
+
+def torch_top_k_top_p_joint_sampling_from_probs(
+    normalized_prob, top_k, top_p, eps=1e-4
+):
+    """Reference PyTorch implementation of joint top-k top-p sampling."""
+    batch_size, vocab_size = normalized_prob.shape
+    samples = torch.empty(batch_size, dtype=torch.int64, device=normalized_prob.device)
+
+    for i in range(batch_size):
+        p_val = top_p[i].item()
+        k_val = top_k[i].item()
+
+        # top-p mask
+        sorted_prob, indices = torch.sort(normalized_prob[i], descending=False)
+        cdf = torch.cumsum(sorted_prob, dim=-1)
+        mask_top_p = torch.zeros(
+            vocab_size, dtype=torch.int32, device=normalized_prob.device
+        )
+        mask_top_p.scatter_add_(0, indices, (cdf > (1 - p_val) - eps).int())
+
+        # top-k mask
+        sorted_prob_desc, _ = torch.sort(normalized_prob[i], descending=True)
+        pivot = sorted_prob_desc[k_val - 1]
+        mask_top_k = (normalized_prob[i] >= pivot).int()
+
+        # joint mask
+        mask = torch.minimum(mask_top_p, mask_top_k).bool()
+
+        # sample from masked probs
+        masked_probs = normalized_prob[i] * mask
+        masked_probs = masked_probs / masked_probs.sum()
+        idx = torch.multinomial(masked_probs, 1)
+        samples[i] = idx
+
+    return samples
+
+
+def calculate_diff(batch_size, vocab_size, p):
+    """Compare Torch reference and SGLang kernel for correctness."""
+    torch.manual_seed(42)
+    if p == 0.1:
+        k = int(vocab_size * 0.5)
+    elif p == 0.5:
+        k = int(vocab_size * 0.1)
+    else:
+        raise ValueError("p not recognized")
+
+    device = torch.device("cuda")
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+
+    top_p_tensor = torch.full((batch_size,), p, device=device)
+    top_k_tensor = torch.full((batch_size,), k, device=device)
+
+    torch_samples = torch_top_k_top_p_joint_sampling_from_probs(
+        normalized_prob, top_k_tensor, top_p_tensor
+    )
+    sglang_samples = flashinfer.sampling.top_k_top_p_sampling_from_probs(
+        normalized_prob, top_k_tensor, top_p_tensor, filter_apply_order="joint"
+    )
+
+
+# parameter space - simplified for CI
+if IS_CI:
+    batch_size_range = [16]  # Single batch size for CI
+    vocab_size_range = [111]  # Single vocab size for CI
+    p_range = [0.1]  # Single p value for CI
+else:
+    batch_size_range = [16, 64, 128]
+    vocab_size_range = [111, 32000]
+    p_range = [0.1, 0.5]
+
+configs = list(itertools.product(batch_size_range, vocab_size_range, p_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "vocab_size", "p"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["torch", "sglang"],
+        line_names=["Torch Reference", "SGL Kernel"],
+        styles=[("red", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="top-k-top-p-joint-sampling-performance",
+        args={},
+    )
+)
+def benchmark_sampling(batch_size, vocab_size, p, provider):
+    torch.manual_seed(42)
+    if p == 0.1:
+        k = int(vocab_size * 0.5)
+    elif p == 0.5:
+        k = int(vocab_size * 0.1)
+    else:
+        raise ValueError("p not recognized")
+
+    device = torch.device("cuda")
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    top_p_tensor = torch.full((batch_size,), p, device=device)
+    top_k_tensor = torch.full((batch_size,), k, device=device)
+
+    if provider == "torch":
+        fn = lambda: torch_top_k_top_p_joint_sampling_from_probs(
+            normalized_prob.clone(), top_k_tensor, top_p_tensor
+        )
+    elif provider == "sglang":
+        fn = lambda: flashinfer.sampling.top_k_top_p_sampling_from_probs(
+            normalized_prob.clone(),
+            top_k_tensor,
+            top_p_tensor,
+            filter_apply_order="joint",
+        )
+
+    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__":
+    # Correctness check - simplified for CI
+    if IS_CI:
+        # Only test one configuration in CI
+        test_configs = [configs[0]] if configs else [(16, 111, 0.1)]
+    else:
+        test_configs = configs
+
+    for cfg in test_configs:
+        calculate_diff(*cfg)
+
+    print("\n" + "=" * 60)
+    print("Starting performance benchmark...")
+    benchmark_sampling.run(print_data=True)