chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-kernel/tests/conftest.py

@@ -0,0 +1,14 @@
+import pytest
+import torch
+
+
+# This fixture ensures the torch defaults don't get left in modified states between
+# tests (e.g., when a test fails before restoring the original value), which
+# can cause subsequent tests to fail.
+@pytest.fixture(autouse=True)
+def reset_torch_defaults():
+    orig_default_device = torch.get_default_device()
+    orig_default_dtype = torch.get_default_dtype()
+    yield
+    torch.set_default_dtype(orig_default_dtype)
+    torch.set_default_device(orig_default_device)

third_party/sglang/sgl-kernel/tests/spatial/test_greenctx_stream.py

@@ -0,0 +1,27 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import create_greenctx_stream_by_value, get_sm_available
+
+
+def test_green_ctx():
+    A = torch.randn(5120, 5120).cuda()
+    B = torch.randn(5120, 5120).cuda()
+    C = torch.matmul(A, B)
+    sm_counts = get_sm_available(0)
+    stream_group = create_greenctx_stream_by_value(sm_counts // 2, sm_counts // 2, 0)
+    with torch.cuda.stream(stream_group[0]):
+        for _ in range(100):
+            result_0 = torch.matmul(A, B)
+    with torch.cuda.stream(stream_group[1]):
+        for _ in range(100):
+            result_1 = torch.matmul(A, B)
+    torch.cuda.synchronize()
+    assert torch.allclose(result_0, C)
+    assert torch.allclose(result_1, C)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/speculative/test_eagle_utils.py

@@ -0,0 +1,89 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import verify_tree_greedy
+
+
+def test_verify_tree_greedy():
+    candidates = torch.tensor(
+        [
+            [0, 1, 2, 3, 4, 5],
+            [7, 8, 9, 10, 11, 12],
+        ],
+        dtype=torch.int64,
+        device="cuda",
+    )
+    retrive_index = torch.tensor(
+        [
+            [0, 1, 2, 3, 4, 5],
+            [6, 7, 8, 9, 10, 11],
+        ],
+        dtype=torch.int64,
+        device="cuda",
+    )
+    retrive_next_token = torch.tensor(
+        [
+            [1, 2, -1, 4, 5, -1],
+            [4, 2, 3, -1, 5, -1],
+        ],
+        dtype=torch.int64,
+        device="cuda",
+    )
+    retrive_next_sibling = torch.tensor(
+        [
+            [-1, 3, -1, -1, -1, -1],
+            [-1, -1, -1, -1, 1, -1],
+        ],
+        dtype=torch.int64,
+        device="cuda",
+    )
+
+    target_logits = torch.full((2, 6, 20), 1, dtype=torch.float32, device="cuda")
+    target_logits[0, 0, 3] = 10
+    target_logits[0, 3, 4] = 10
+    target_logits[0, 4, 5] = 10
+    target_logits[1, 0, 11] = 10
+    target_logits[1, 4, 12] = 10
+    for i in range(target_logits.shape[0]):
+        for j in range(target_logits.shape[1]):
+            if torch.max(target_logits[i][j]) < 10:
+                target_logits[i][j][18] = 10
+
+    target_predict = torch.argmax(target_logits, dim=-1)
+    predict_shape = (12,)
+
+    bs = candidates.shape[0]
+    num_spec_step = 4
+
+    predicts = torch.full(
+        predict_shape, -1, dtype=torch.int32, device="cuda"
+    )  # mutable
+    accept_index = torch.full(
+        (bs, num_spec_step), -1, dtype=torch.int32, device="cuda"
+    )  # mutable
+    accept_token_num = torch.full((bs,), 0, dtype=torch.int32, device="cuda")  # mutable
+
+    verify_tree_greedy(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates,
+        retrive_index=retrive_index,
+        retrive_next_token=retrive_next_token,
+        retrive_next_sibling=retrive_next_sibling,
+        target_predict=target_predict,
+    )
+
+    # Check the expected output.
+    assert predicts.tolist() == [3, -1, -1, 4, 5, 18, 11, -1, -1, -1, 12, 18]
+    assert accept_index.tolist() == [
+        [0, 3, 4, 5],
+        [6, 10, 11, -1],
+    ]
+    assert accept_token_num.tolist() == [3, 2]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/speculative/test_ngram_utils.py

@@ -0,0 +1,78 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import reconstruct_indices_from_tree_mask
+
+
+def test_reconstruct_indices_from_tree_mask():
+    bs = 1
+    num_branch_token = 4
+    seq_lens = torch.tensor([12], device="cuda", dtype=torch.int64)
+
+    retrive_index = torch.full(
+        (bs, num_branch_token), -1, device="cuda", dtype=torch.int64
+    )
+    retrive_next_token = torch.full(
+        (bs, num_branch_token), -1, device="cuda", dtype=torch.int64
+    )
+    retrive_next_sibling = torch.full(
+        (bs, num_branch_token), -1, device="cuda", dtype=torch.int64
+    )
+    positions = torch.empty((bs * num_branch_token), device="cuda", dtype=torch.int64)
+
+    tree_mask = torch.tensor(
+        [
+            1,
+            0,
+            0,
+            0,
+            1,
+            1,
+            0,
+            0,
+            1,
+            0,
+            1,
+            0,
+            1,
+            0,
+            1,
+            1,
+        ],
+        device="cuda",
+        dtype=torch.int32,
+    ).to(torch.bool)
+
+    reconstruct_indices_from_tree_mask(
+        tree_mask,
+        seq_lens,
+        positions,  # mutable
+        retrive_index,  # mutable
+        retrive_next_token,  # mutable
+        retrive_next_sibling,  # mutable
+        bs,
+        num_branch_token,
+    )
+    # print(f"debug: \n\n{tree_mask=}, {retrive_index=}, {retrive_next_token=}, {retrive_next_sibling=}, {positions=}\n\n")
+    assert retrive_index.tolist() == [
+        [0, 1, 2, 3],
+    ], f"{retrive_index=}"
+    assert retrive_next_token.tolist() == [
+        [1, -1, 3, -1],
+    ], f"{retrive_next_token=}"
+    assert retrive_next_sibling.tolist() == [
+        [-1, 2, -1, -1],
+    ], f"{retrive_next_sibling=}"
+    assert positions.tolist() == [
+        12,
+        13,
+        13,
+        14,
+    ], f"{positions=}"
+
+
+if __name__ == "__main__":
+    test_reconstruct_indices_from_tree_mask()
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/speculative/test_speculative_sampling.py

@@ -0,0 +1,131 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import tree_speculative_sampling_target_only
+
+test_cases = [
+    (
+        1,
+        1,
+        [3, -1, -1, 4, 5, 18, 11, -1, -1, -1, 12, 18],
+        [[0, 3, 4, 5], [6, 10, 11, -1]],
+        [3, 2],
+    ),
+    (
+        0,  # threshold_single
+        0,  # threshold_acc
+        [1, 2, 18, -1, -1, -1, 11, -1, -1, -1, 12, 18],
+        [[0, 1, 2, -1], [6, 10, 11, -1]],
+        [2, 2],
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "threshold_single, threshold_acc, expected_predicts, expected_accept_index, expected_accept_token_num",
+    test_cases,
+)
+def test_tree_speculative_sampling_target_only(
+    threshold_single,
+    threshold_acc,
+    expected_predicts,
+    expected_accept_index,
+    expected_accept_token_num,
+):
+    """
+    Tests the tree_speculative_sampling_target_only function using Pytest parameterization.
+    """
+    device = "cuda"
+
+    candidates = torch.tensor(
+        [
+            [0, 1, 2, 3, 4, 5],
+            [7, 8, 9, 10, 11, 12],
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+    retrive_index = torch.tensor(
+        [
+            [0, 1, 2, 3, 4, 5],
+            [6, 7, 8, 9, 10, 11],
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+    retrive_next_token = torch.tensor(
+        [
+            [1, 2, -1, 4, 5, -1],
+            [4, 2, 3, -1, 5, -1],
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+    retrive_next_sibling = torch.tensor(
+        [
+            [-1, 3, -1, -1, -1, -1],
+            [-1, -1, -1, -1, 1, -1],
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+
+    target_logits = torch.full((2, 6, 20), 1, dtype=torch.float32, device=device)
+    target_logits[0, 0, 3] = 10
+    target_logits[0, 3, 4] = 10
+    target_logits[0, 4, 5] = 10
+    target_logits[1, 0, 11] = 10
+    target_logits[1, 4, 12] = 10
+
+    for i in range(target_logits.shape[0]):
+        for j in range(target_logits.shape[1]):
+            if torch.max(target_logits[i, j]) < 10:
+                target_logits[i, j, 18] = 10
+
+    temperatures = torch.tensor([0.01, 0.01], dtype=torch.float32, device=device)
+    bs, num_draft_tokens = candidates.shape
+    num_spec_step = len(expected_accept_index[0])
+    predict_shape = (len(expected_predicts),)
+
+    predicts = torch.full(predict_shape, -1, dtype=torch.int32, device=device)
+    accept_index = torch.full((bs, num_spec_step), -1, dtype=torch.int32, device=device)
+    accept_token_num = torch.full((bs,), 0, dtype=torch.int32, device=device)
+
+    expanded_temperature = temperatures.unsqueeze(1).unsqueeze(1)
+    target_probs = F.softmax(target_logits / expanded_temperature, dim=-1)
+    draft_probs = torch.full_like(target_probs, 0, dtype=torch.float32, device=device)
+    coins = torch.rand(bs, num_draft_tokens, device=device, dtype=torch.float32)
+    coins_for_final_sampling = torch.rand(bs, device=device).to(torch.float32)
+
+    tree_speculative_sampling_target_only(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates,
+        retrive_index=retrive_index,
+        retrive_next_token=retrive_next_token,
+        retrive_next_sibling=retrive_next_sibling,
+        uniform_samples=coins,
+        uniform_samples_for_final_sampling=coins_for_final_sampling,
+        target_probs=target_probs,
+        draft_probs=draft_probs,
+        threshold_single=threshold_single,
+        threshold_acc=threshold_acc,
+        deterministic=True,
+    )
+
+    assert (
+        predicts.tolist() == expected_predicts
+    ), f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})"
+    assert (
+        accept_index.tolist() == expected_accept_index
+    ), f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})"
+    assert (
+        accept_token_num.tolist() == expected_accept_token_num
+    ), f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_activation.py

@@ -0,0 +1,41 @@
+# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/4e8eb1879f9c3ba6d75511e5893183bf8f289a62/tests/test_activation.py
+
+import sys
+
+import pytest
+import sgl_kernel
+import torch
+
+
+@pytest.mark.parametrize("dim", [128, 256, 512, 2048, 4096, 11008, 16384])
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16])
+@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 16, 32, 64, 128, 512])
+def test_fused_silu_mul(dim, batch_size, seq_len):
+    x = torch.randn(batch_size, seq_len, 2 * dim).to(0).to(torch.float16)
+    y_ref = x[..., dim:] * torch.nn.functional.silu(x[..., :dim])
+    y = sgl_kernel.silu_and_mul(x)
+    torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("dim", [128, 256, 512, 2048, 4096, 11008, 16384])
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16])
+@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 16, 32, 64, 128, 512])
+def test_fused_gelu_tanh_mul(dim, batch_size, seq_len):
+    x = torch.randn(batch_size, seq_len, 2 * dim).to(0).to(torch.float16)
+    y_ref = x[..., dim:] * torch.nn.functional.gelu(x[..., :dim], approximate="tanh")
+    y = sgl_kernel.gelu_tanh_and_mul(x)
+    torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("dim", [128, 256, 512, 2048, 4096, 11008, 16384])
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16])
+@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 16, 32, 64, 128, 512])
+def test_fused_gelu_mul(dim, batch_size, seq_len):
+    x = torch.randn(batch_size, seq_len, 2 * dim).to(0).to(torch.float16)
+    y_ref = x[..., dim:] * torch.nn.functional.gelu(x[..., :dim], approximate="none")
+    y = sgl_kernel.gelu_and_mul(x)
+    torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py

@@ -0,0 +1,251 @@
+"""
+Test deterministic custom all-reduce kernel behavior with batch size invariance.
+
+This test uses the 1-stage all-reduce kernel which is inherently deterministic
+due to fixed accumulation ordering (each GPU reads all data from all GPUs and
+reduces locally in a fixed order - no atomics, no race conditions).
+
+Note: This is NOT a reduce-scatter + all-gather (RS+AG) approach.
+
+This test compares:
+1. Deterministic kernel (same batch size)
+2. Deterministic kernel (different batch size)
+
+Usage:
+    pytest test_amd_deterministic_custom_allreduce.py
+"""
+
+import multiprocessing as mp
+import socket
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from sglang.srt.environ import envs
+
+
+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 worker(world_size, rank, port):
+    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 import and use deterministic kernel
+    try:
+        from torch.distributed import new_group
+
+        from sglang.srt.distributed.device_communicators.custom_all_reduce import (
+            CustomAllreduce,
+        )
+
+        # Create gloo group for custom AR
+        dist.barrier()
+        ar_group = new_group(backend="gloo")
+        dist.barrier()
+
+        custom_ar = CustomAllreduce(group=ar_group, device=device)
+
+        if custom_ar is None or custom_ar.disabled:
+            if rank == 0:
+                print("✗ Custom AR not available or disabled")
+            dist.destroy_process_group()
+            return
+    except Exception as e:
+        if rank == 0:
+            print(f"✗ Failed to initialize deterministic kernel: {e}")
+            import traceback
+
+            traceback.print_exc()
+        dist.destroy_process_group()
+        return
+
+    num_trials = 10
+
+    # Matrix sizes similar to real model layers
+    # Format: (batch_size, hidden_dim) - typical tensor shape for all-reduce
+    BS = 50  # max batch_size (1..BS)
+    hidden_dim = 16384  # hidden dimension / intermediate dimension
+
+    # Different seed per rank - each GPU has DIFFERENT input
+    torch.manual_seed(42 + rank)
+
+    # Create fixed inputs for all trials
+    # Single request: (hidden_dim,)
+    base_input = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+    base_input_rand = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+
+    # Check if inputs fit in buffer
+    # Buffer size is max_size bytes, input size is numel * element_size bytes
+    input_size_bytes = base_input.numel() * base_input.element_size()
+    if input_size_bytes > custom_ar.max_size and rank == 0:
+        print(
+            f"Warning: Input size ({input_size_bytes/(1024*1024):.1f} MB) exceeds buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
+        )
+        print("  Using unregistered mode (will copy to buffer)")
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 1: Deterministic kernel (same batch size) - should be DETERMINISTIC
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 1: Deterministic kernel (same batch size)")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = []
+    for trial in range(num_trials):
+        # Clone the same input
+        inp = base_input.clone()
+
+        result = custom_ar.custom_all_reduce(inp)
+        torch.cuda.synchronize()
+
+        # Store checksum
+        checksum = result.view(-1).sum().item()
+        first_vals = result.view(-1)[:5].clone()
+        results_allreduce_only.append((checksum, first_vals))
+
+        if rank == 0:
+            print(
+                f"  Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+            )
+
+    # Check determinism
+    if rank == 0:
+        ref_sum, ref_vals = results_allreduce_only[0]
+        all_match = True
+        for i, (s, vals) in enumerate(results_allreduce_only[1:], 1):
+            if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3):
+                all_match = False
+                print(f"  Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
+
+        if all_match:
+            print("  ✓ DETERMINISTIC KERNEL (fixed BS): DETERMINISTIC (as expected)")
+        else:
+            print(
+                "  ✗ DETERMINISTIC KERNEL (fixed BS): NON-DETERMINISTIC (unexpected!)"
+            )
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 2: Deterministic kernel (different batch size) - should be DETERMINISTIC
+    # [a], [a, x], [a, x, x], ...
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 2: Deterministic kernel (different batch size)")
+        print("Batches: [a], [a,x], [a,x,x], ...")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = {trial: [] for trial in range(num_trials)}
+    for trial in range(num_trials):
+        for bs in range(1, BS + 1):
+            # Construct batch: (batch_size, hidden_dim)
+            # First element is base_input, rest are base_input_rand
+            batch = torch.stack([base_input] + [base_input_rand] * (bs - 1), dim=0)
+            # Shape: (bs, hidden_dim)
+
+            # Flatten for all-reduce: (bs * hidden_dim,)
+            batch_flat = batch.view(-1)
+
+            result_flat = custom_ar.custom_all_reduce(batch_flat)
+            torch.cuda.synchronize()
+
+            # Reshape back to (bs, hidden_dim)
+            batch_out = result_flat.view(bs, hidden_dim)
+
+            # Only compare output corresponding to first request
+            out_first_req = batch_out[0].clone()
+            checksum = out_first_req.sum().item()
+            first_vals = out_first_req[:5].clone()
+            results_allreduce_only[trial].append((bs, checksum, first_vals))
+
+            if rank == 0:
+                print(
+                    f"  Batch size {bs:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+                )
+
+    # Check determinism
+    if rank == 0:
+        for trial in range(num_trials):
+            results = results_allreduce_only[trial]
+
+            _, ref_sum, ref_vals = results[0]
+            all_match = True
+            for _, s, vals in results[1:]:
+                if abs(ref_sum - s) > 1e-3 or not torch.allclose(
+                    ref_vals, vals, rtol=1e-3
+                ):
+                    all_match = False
+
+        if all_match:
+            print("  ✓ DETERMINISTIC KERNEL (variant BS): DETERMINISTIC")
+        else:
+            print("  ✗ DETERMINISTIC KERNEL (variant BS): NON-DETERMINISTIC")
+
+    dist.barrier()
+
+    dist.destroy_process_group()
+
+
+def main():
+    world_size = 8
+    available_gpus = torch.cuda.device_count()
+
+    print("=" * 70)
+    print("Deterministic Kernel All-Reduce Determinism Test")
+    print("=" * 70)
+    print(f"Available GPUs: {available_gpus}")
+    print(f"Using world_size: {world_size}")
+
+    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 test")
+        return
+
+    mp.set_start_method("spawn", force=True)
+    port = get_open_port()
+
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(target=worker, args=(world_size, rank, port))
+        p.start()
+        procs.append(p)
+
+    for p in procs:
+        p.join()
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+    reason="Requires at least 2 CUDA GPUs",
+)
+def test_deterministic_custom_allreduce():
+    """Test that deterministic custom all-reduce produces consistent results."""
+    main()
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py

@@ -0,0 +1,208 @@
+"""
+Test to confirm non-determinism of default NCCL all-reduce with batch size invariance.
+
+This test uses the default torch.distributed.all_reduce (NCCL) which can be
+NON-DETERMINISTIC due to tree-based reduction algorithms that don't guarantee
+fixed accumulation order for bfloat16/float16.
+
+This test compares:
+1. Default all-reduce (same batch size) - should be DETERMINISTIC
+2. Default all-reduce (different batch size) - typically NON-DETERMINISTIC for bfloat16
+
+Usage:
+    pytest test_amd_nccl_allreduce_determinism.py
+"""
+
+import multiprocessing as mp
+import socket
+
+import pytest
+import torch
+import torch.distributed as dist
+
+
+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 worker(world_size, rank, port):
+    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,
+    )
+
+    num_trials = 10
+
+    # Matrix sizes similar to real model layers
+    # Format: (batch_size, hidden_dim) - typical tensor shape for all-reduce
+    BS = 50  # max batch_size (1..BS)
+    hidden_dim = 16384  # hidden dimension / intermediate dimension
+
+    # Different seed per rank - each GPU has DIFFERENT input
+    torch.manual_seed(42 + rank)
+
+    # Create fixed inputs for all trials
+    # Single request: (hidden_dim,)
+    base_input = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+    base_input_rand = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 1: Default all-reduce (same batch size) - should be DETERMINISTIC
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 1: Default NCCL all_reduce (same batch size)")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = []
+    for trial in range(num_trials):
+        # Clone the same input
+        inp = base_input.clone()
+
+        # Use default NCCL all-reduce
+        dist.all_reduce(inp)
+        torch.cuda.synchronize()
+
+        # Store checksum
+        checksum = inp.view(-1).sum().item()
+        first_vals = inp.view(-1)[:5].clone()
+        results_allreduce_only.append((checksum, first_vals))
+
+        if rank == 0:
+            print(
+                f"  Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+            )
+
+    # Check determinism
+    if rank == 0:
+        ref_sum, ref_vals = results_allreduce_only[0]
+        all_match = True
+        for i, (s, vals) in enumerate(results_allreduce_only[1:], 1):
+            if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3):
+                all_match = False
+                print(f"  Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
+
+        if all_match:
+            print("  ✓ DEFAULT ALL_REDUCE (fixed BS): DETERMINISTIC (as expected)")
+        else:
+            print("  ✗ DEFAULT ALL_REDUCE (fixed BS): NON-DETERMINISTIC (unexpected!)")
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 2: Default all-reduce (different batch size) - typically NON-DETERMINISTIC
+    # [a], [a, x], [a, x, x], ...
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 2: Default NCCL all_reduce (different batch size)")
+        print("Batches: [a], [a,x], [a,x,x], ...")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = {trial: [] for trial in range(num_trials)}
+    for trial in range(num_trials):
+        for bs in range(1, BS + 1):
+            # Construct batch: (batch_size, hidden_dim)
+            # First element is base_input, rest are base_input_rand
+            batch = torch.stack([base_input] + [base_input_rand] * (bs - 1), dim=0)
+            # Shape: (bs, hidden_dim)
+
+            # Flatten for all-reduce: (bs * hidden_dim,)
+            batch_flat = batch.view(-1)
+
+            # Use default NCCL all-reduce
+            dist.all_reduce(batch_flat)
+            torch.cuda.synchronize()
+
+            # Reshape back to (bs, hidden_dim)
+            batch_out = batch_flat.view(bs, hidden_dim)
+
+            # Only compare output corresponding to first request
+            out_first_req = batch_out[0].clone()
+            checksum = out_first_req.sum().item()
+            first_vals = out_first_req[:5].clone()
+            results_allreduce_only[trial].append((bs, checksum, first_vals))
+
+            if rank == 0:
+                print(
+                    f"  Batch size {bs:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+                )
+
+    # Check determinism
+    if rank == 0:
+        for trial in range(num_trials):
+            results = results_allreduce_only[trial]
+
+            _, ref_sum, ref_vals = results[0]
+            all_match = True
+            for _, s, vals in results[1:]:
+                if abs(ref_sum - s) > 1e-3 or not torch.allclose(
+                    ref_vals, vals, rtol=1e-3
+                ):
+                    all_match = False
+
+        if all_match:
+            print("  ✓ DEFAULT ALL_REDUCE (variant BS): DETERMINISTIC")
+        else:
+            print("  ✗ DEFAULT ALL_REDUCE (variant BS): NON-DETERMINISTIC")
+
+    dist.barrier()
+
+    dist.destroy_process_group()
+
+
+def main():
+    world_size = 8
+    available_gpus = torch.cuda.device_count()
+
+    print("=" * 70)
+    print("Default NCCL All-Reduce Determinism Test")
+    print("=" * 70)
+    print(f"Available GPUs: {available_gpus}")
+    print(f"Using world_size: {world_size}")
+
+    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 test")
+        return
+
+    mp.set_start_method("spawn", force=True)
+    port = get_open_port()
+
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(target=worker, args=(world_size, rank, port))
+        p.start()
+        procs.append(p)
+
+    for p in procs:
+        p.join()
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+    reason="Requires at least 2 CUDA GPUs",
+)
+def test_nccl_allreduce_determinism():
+    """Test NCCL all-reduce determinism behavior with varying batch sizes."""
+    main()
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/sgl-kernel/tests/test_apply_token_bitmask_inplace.py

@@ -0,0 +1,25 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import apply_token_bitmask_inplace_cuda
+
+
+def test_apply_token_bitmask_inplace_kernel():
+    neginf = float("-inf")
+    bool_mask = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.bool)
+    logits = torch.tensor(
+        [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], dtype=torch.float32
+    )
+    expected = torch.where(bool_mask, logits, neginf)
+
+    logits_gpu = logits.to("cuda")
+    bitmask = torch.tensor([0b1010101010], dtype=torch.int32).to("cuda")
+    apply_token_bitmask_inplace_cuda(logits_gpu, bitmask)
+    torch.cuda.synchronize()
+    torch.testing.assert_close(logits_gpu, expected.to("cuda"))
+
+
+if __name__ == "__main__":
+    test_apply_token_bitmask_inplace_kernel()
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_awq_dequant.py

@@ -0,0 +1,116 @@
+import itertools
+import sys
+from typing import Optional, Tuple
+
+import pytest
+import torch
+from sgl_kernel import awq_dequantize
+
+
+def reverse_awq_order(t: torch.Tensor):
+    bits = 4
+    AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
+    reverse_order_tensor = torch.arange(
+        t.shape[-1],
+        dtype=torch.int32,
+        device=t.device,
+    )
+    reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
+    reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
+    reverse_order_tensor = reverse_order_tensor.view(-1)
+
+    t = t[:, reverse_order_tensor] & 0xF
+    return t
+
+
+# qweights - [R     , C // 8], int32
+# scales   - [R // G, C     ], float16
+# zeros    - [R // G, C // 8], int32
+def awq_dequantize_torch(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int
+) -> torch.Tensor:
+
+    if group_size == -1:
+        group_size = qweight.shape[0]
+
+    bits = 4
+    shifts = torch.arange(0, 32, bits, device=qzeros.device)
+
+    iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(
+        torch.int8
+    )
+
+    iweights = iweights.view(iweights.shape[0], -1)
+
+    zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(
+        torch.int8
+    )
+    zeros = zeros.view(qzeros.shape[0], -1)
+    zeros = reverse_awq_order(zeros)
+
+    iweights = reverse_awq_order(iweights)
+
+    iweights = torch.bitwise_and(iweights, (2**bits) - 1)
+    zeros = torch.bitwise_and(zeros, (2**bits) - 1)
+
+    scales = scales.repeat_interleave(group_size, dim=0)
+    zeros = zeros.repeat_interleave(group_size, dim=0)
+    return (iweights - zeros) * scales
+
+
+def sglang_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> torch.Tensor:
+    return awq_dequantize(qweight, scales, qzeros)
+
+
+@pytest.mark.parametrize(
+    "qweight_row,qweight_col,is_bf16_act",
+    list(
+        itertools.product(
+            [3584, 18944, 128, 256, 512, 1024, 1536],
+            [448, 576, 4736, 16, 32, 64, 128, 72],
+            [True, False],
+        )
+    ),
+)
+def test_awq_dequant_compare_implementations(
+    qweight_row: int, qweight_col: int, is_bf16_act: bool
+):
+    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
+
+    if is_bf16_act:
+        scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
+    else:
+        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,
+    )
+
+    # Run both implementations
+    torch_out = awq_dequantize_torch(qweight, scales, qzeros, group_size)
+    sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
+
+    # Compare results
+    torch.testing.assert_close(
+        torch_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_bmm_fp8.py

@@ -0,0 +1,45 @@
+# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/4e8eb1879f9c3ba6d75511e5893183bf8f289a62/tests/test_bmm_fp8.py
+
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import bmm_fp8
+
+
+def to_float8(x, dtype=torch.float8_e4m3fn):
+    finfo = torch.finfo(dtype)
+    min_val, max_val = x.aminmax()
+    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
+    scale = finfo.max / amax
+    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
+    return x_scl_sat.to(dtype), scale.float().reciprocal()
+
+
+@pytest.mark.parametrize("input_dtype", [torch.float8_e4m3fn, torch.float8_e5m2])
+@pytest.mark.parametrize("mat2_dtype", [torch.float8_e4m3fn, torch.float8_e5m2])
+@pytest.mark.parametrize("res_dtype", [torch.bfloat16, torch.float16])
+def test_bmm_fp8(input_dtype, mat2_dtype, res_dtype):
+    if input_dtype == torch.float8_e5m2 and mat2_dtype == torch.float8_e5m2:
+        pytest.skip("Invalid combination: both input and mat2 are e5m2")
+
+    input = torch.randn([16, 48, 64], device="cuda", dtype=torch.bfloat16)
+    input_fp8, input_inv_s = to_float8(input, dtype=input_dtype)
+
+    # mat2 row  major -> column major
+    mat2 = torch.randn([16, 80, 64], device="cuda", dtype=torch.bfloat16).transpose(
+        -2, -1
+    )
+    mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=mat2_dtype)
+
+    res = torch.empty([16, 48, 80], device="cuda", dtype=res_dtype)
+    bmm_fp8(input_fp8, mat2_fp8, input_inv_s, mat2_inv_s, res_dtype, res)
+
+    reference = torch.bmm(input, mat2)
+    cos_sim = F.cosine_similarity(reference.reshape(-1), res.reshape(-1), dim=0)
+    assert cos_sim > 0.99
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_causal_conv1d.py

@@ -0,0 +1,490 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/main/tests/kernels/mamba/test_causal_conv1d.py
+import sys
+from typing import Optional
+
+import torch
+from sgl_kernel import causal_conv1d_fwd
+from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel
+
+PAD_SLOT_ID = -1
+
+
+def causal_conv1d_fn(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    query_start_loc: Optional[torch.Tensor] = None,
+    cache_indices: Optional[torch.Tensor] = None,
+    has_initial_state: Optional[torch.Tensor] = None,
+    conv_states: Optional[torch.Tensor] = None,
+    activation: Optional[str] = "silu",
+    pad_slot_id: int = PAD_SLOT_ID,
+):
+    """
+    x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen
+        sequences are concatenated from left to right for varlen
+    weight: (dim, width)
+    bias: (dim,)
+    query_start_loc: (batch + 1) int32
+        The cumulative sequence lengths of the sequences in
+        the batch, used to index into sequence. prepended by 0.
+        for example: query_start_loc = torch.Tensor([0,10,16,17]),
+        x.shape=(dim,17)
+    cache_indices: (batch)  int32
+        indicates the corresponding state index,
+        like so: conv_state = conv_states[cache_indices[batch_id]]
+    has_initial_state: (batch) bool
+        indicates whether should the kernel take the current state as initial
+        state for the calculations
+    conv_states: (...,dim,width - 1) itype
+        updated inplace if provided
+    activation: either None or "silu" or "swish"
+    pad_slot_id: int
+            if cache_indices is passed, lets the kernel identify padded
+            entries that will not be processed,
+            for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
+            in this case, the kernel will not process entries at
+            indices 0 and 3
+    out: (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    if x.stride(-1) != 1:
+        x = x.contiguous()
+    bias = bias.contiguous() if bias is not None else None
+
+    causal_conv1d_fwd(
+        x,
+        weight,
+        bias,
+        conv_states,
+        query_start_loc,
+        cache_indices,
+        has_initial_state,
+        activation in ["silu", "swish"],
+        pad_slot_id,
+    )
+    return x
+
+
+def causal_conv1d_update(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    activation: Optional[str] = None,
+    cache_seqlens: Optional[torch.Tensor] = None,
+    conv_state_indices: Optional[torch.Tensor] = None,
+    pad_slot_id: int = PAD_SLOT_ID,
+):
+    """
+    x: (batch, dim) or (batch, dim, seqlen)
+    conv_state: (batch, dim, state_len), where state_len >= width - 1
+    weight: (dim, width)
+    bias: (dim,)
+    cache_seqlens: (batch,), dtype int32.
+        If not None, the conv_state is treated as a circular buffer.
+        The conv_state will be updated by copying x to the conv_state
+        starting at the index
+        @cache_seqlens % state_len.
+    conv_state_indices: (batch,), dtype int32
+        If not None, the conv_state is a larger tensor along the batch dim,
+        and we are selecting the batch coords specified by conv_state_indices.
+        Useful for a continuous batching scenario.
+    pad_slot_id: int
+            if cache_indices is passed, lets the kernel identify padded
+            entries that will not be processed,
+            for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
+            in this case, the kernel will not process entries at
+            indices 0 and 3
+    out: (batch, dim) or (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError(
+            f"activation must be None, silu, or swish, actual: {activation}"
+        )
+    activation_val = activation in ["silu", "swish"]
+    unsqueeze = x.dim() == 2
+    if unsqueeze:
+        x = x.unsqueeze(-1)
+    causal_conv1d_update_kernel(
+        x,
+        conv_state,
+        weight,
+        bias,
+        activation_val,
+        cache_seqlens,
+        conv_state_indices,
+        pad_slot_id,
+    )
+    if unsqueeze:
+        x = x.squeeze(-1)
+    return x
+
+
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+def causal_conv1d_ref(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    initial_states: Optional[torch.Tensor] = None,
+    return_final_states: bool = False,
+    final_states_out: Optional[torch.Tensor] = None,
+    activation: Optional[str] = "silu",
+):
+    """
+    x: (batch, dim, seqlen)
+    weight: (dim, width)
+    bias: (dim,)
+    initial_states: (batch, dim, width - 1)
+    final_states_out: (batch, dim, width - 1)
+
+    out: (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    dtype_in = x.dtype
+    x = x.to(weight.dtype)
+    seqlen = x.shape[-1]
+    dim, width = weight.shape
+    if initial_states is None:
+        out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim)
+    else:
+        x = torch.cat([initial_states, x], dim=-1)
+        out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim)
+    out = out[..., :seqlen]
+    if return_final_states:
+        final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to(
+            dtype_in
+        )  # (batch, dim, width - 1)
+        if final_states_out is not None:
+            final_states_out.copy_(final_states)
+        else:
+            final_states_out = final_states
+    out = (out if activation is None else F.silu(out)).to(dtype=dtype_in)
+    return (out, None) if not return_final_states else (out, final_states_out)
+
+
+def causal_conv1d_update_ref(
+    x, conv_state, weight, bias=None, activation=None, cache_seqlens=None
+):
+    """
+    x: (batch, dim) or (batch, dim, seqlen)
+    conv_state: (batch, dim, state_len), where state_len >= width - 1
+    weight: (dim, width)
+    bias: (dim,)
+    cache_seqlens: (batch,), dtype int32.
+        If not None, the conv_state is treated as a circular buffer.
+        The conv_state will be updated by copying x to the
+        conv_state starting at the index
+        @cache_seqlens % state_len before performing the convolution.
+
+    out: (batch, dim) or (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    dtype_in = x.dtype
+    unsqueeze = x.dim() == 2
+    if unsqueeze:
+        x = x.unsqueeze(-1)
+    batch, dim, seqlen = x.shape
+    width = weight.shape[1]
+    state_len = conv_state.shape[-1]
+    assert conv_state.shape == (batch, dim, state_len)
+    assert weight.shape == (dim, width)
+    if cache_seqlens is None:
+        x_new = torch.cat([conv_state, x], dim=-1).to(
+            weight.dtype
+        )  # (batch, dim, state_len + seqlen)
+        conv_state.copy_(x_new[:, :, -state_len:])
+    else:
+        width_idx = torch.arange(
+            -(width - 1), 0, dtype=torch.long, device=x.device
+        ).unsqueeze(0) + cache_seqlens.unsqueeze(1)
+        width_idx = (
+            torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1)
+        )
+        x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype)
+        copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(
+            0
+        ) + cache_seqlens.unsqueeze(1)
+        copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1)
+        conv_state.scatter_(2, copy_idx, x)
+    out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[
+        :, :, -seqlen:
+    ]
+    if unsqueeze:
+        out = out.squeeze(-1)
+    return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
+
+
+@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
+@pytest.mark.parametrize("silu_activation", [True])
+@pytest.mark.parametrize("has_bias", [True])
+@pytest.mark.parametrize("has_initial_state", [True, False])
+@pytest.mark.parametrize("width", [4])
+@pytest.mark.parametrize(
+    "seqlen", [1, 8, 16, 32, 64, 128, 256, 512, 784, 1024, 1025, 2048, 4096]
+)
+@pytest.mark.parametrize("dim", [64])
+@pytest.mark.parametrize("batch", [1])
+def test_causal_conv1d(
+    batch, dim, seqlen, width, has_bias, silu_activation, has_initial_state, itype
+):
+    device = "cuda"
+    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
+    if itype == torch.bfloat16:
+        rtol, atol = 1e-2, 5e-2
+    x = torch.randn(batch, dim, seqlen, device=device, dtype=itype).contiguous()
+
+    weight = torch.randn(dim, width, device=device, dtype=itype)
+    bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
+    if has_initial_state:
+        initial_states = torch.randn(batch, dim, width - 1, device=device, dtype=itype)
+        has_initial_state_tensor = torch.ones(batch, dtype=torch.bool, device=x.device)
+    else:
+        initial_states = None
+        has_initial_state_tensor = None
+    x_ref = x.clone()
+    weight_ref = weight.clone()
+    bias_ref = bias.clone() if bias is not None else None
+    initial_states_ref = initial_states.clone() if initial_states is not None else None
+    activation = None if not silu_activation else "silu"
+    out = causal_conv1d_fn(
+        x,
+        weight,
+        bias,
+        activation=activation,
+        conv_states=initial_states,
+        has_initial_state=has_initial_state_tensor,
+    )
+    out_ref, final_states_ref = causal_conv1d_ref(
+        x_ref,
+        weight_ref,
+        bias_ref,
+        initial_states=initial_states_ref,
+        return_final_states=True,
+        activation=activation,
+    )
+    if has_initial_state:
+        assert initial_states is not None and final_states_ref is not None
+        assert torch.allclose(initial_states, final_states_ref, rtol=rtol, atol=atol)
+    assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("itype", [torch.bfloat16])
+@pytest.mark.parametrize("silu_activation", [False, True])
+@pytest.mark.parametrize("has_bias", [False, True])
+@pytest.mark.parametrize("seqlen", [1])
+@pytest.mark.parametrize("width", [4])
+@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
+def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, itype):
+    device = "cuda"
+    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
+    if itype == torch.bfloat16:
+        rtol, atol = 1e-2, 5e-2
+
+    batch = 2
+    x = torch.randn(batch, dim, seqlen, device=device, dtype=itype)
+    x_ref = x.clone()
+    conv_state = torch.randn(batch, dim, width - 1, device=device, dtype=itype)
+
+    weight = torch.randn(dim, width, device=device, dtype=itype)
+    bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
+    conv_state_ref = conv_state.detach().clone()
+    activation = None if not silu_activation else "silu"
+    out = causal_conv1d_update(x, conv_state, weight, bias, activation=activation)
+    out_ref = causal_conv1d_update_ref(
+        x_ref, conv_state_ref, weight, bias, activation=activation
+    )
+
+    assert torch.equal(conv_state, conv_state_ref)
+    assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("silu_activation", [False, True])
+@pytest.mark.parametrize("has_bias", [False, True])
+@pytest.mark.parametrize("seqlen", [1, 4, 5])
+@pytest.mark.parametrize("width", [2, 3, 4])
+@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
+# tests correctness in case subset of the sequences are padded
+@pytest.mark.parametrize("with_padding", [True, False])
+def test_causal_conv1d_update_with_batch_gather(
+    with_padding, dim, width, seqlen, has_bias, silu_activation, itype
+):
+    device = "cuda"
+    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
+    if itype == torch.bfloat16:
+        rtol, atol = 1e-2, 5e-2
+
+    batch_size = 3
+    padding = 5 if with_padding else 0
+    padded_batch_size = batch_size + padding
+    total_entries = 10 * batch_size
+
+    x = torch.randn(padded_batch_size, dim, 1, device=device, dtype=itype)
+    x_ref = x.clone()
+
+    conv_state_indices = torch.randperm(total_entries)[:batch_size].to(
+        dtype=torch.int32, device=device
+    )
+    unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device)
+    unused_states_bool[conv_state_indices] = False
+    padded_state_indices = torch.concat(
+        [
+            conv_state_indices,
+            torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
+        ],
+        dim=0,
+    )
+    conv_state = torch.randn(total_entries, dim, width - 1, device=device, dtype=itype)
+    conv_state_for_padding_test = conv_state.clone()
+
+    weight = torch.randn(dim, width, device=device, dtype=itype)
+    bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
+    conv_state_ref = conv_state[conv_state_indices, :].detach().clone()
+    activation = None if not silu_activation else "silu"
+    out = causal_conv1d_update(
+        x,
+        conv_state,
+        weight,
+        bias,
+        activation=activation,
+        conv_state_indices=padded_state_indices,
+        pad_slot_id=PAD_SLOT_ID,
+    )
+    out_ref = causal_conv1d_update_ref(
+        x_ref[:batch_size], conv_state_ref, weight, bias, activation=activation
+    )
+
+    assert torch.equal(conv_state[conv_state_indices, :], conv_state_ref)
+    assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol)
+    assert torch.equal(
+        conv_state[unused_states_bool], conv_state_for_padding_test[unused_states_bool]
+    )
+
+
+@pytest.mark.parametrize("itype", [torch.bfloat16])
+@pytest.mark.parametrize("silu_activation", [True])
+@pytest.mark.parametrize("has_bias", [True])
+@pytest.mark.parametrize("width", [4])
+@pytest.mark.parametrize(
+    "seqlen", [8, 16, 32, 64, 128, 256, 512, 784, 1024, 2048, 2049, 4096]
+)
+@pytest.mark.parametrize("dim", [64, 4096])
+# tests correctness in case subset of the sequences are padded
+@pytest.mark.parametrize("with_padding", [True, False])
+def test_causal_conv1d_varlen(
+    with_padding, dim, seqlen, width, has_bias, silu_activation, itype
+):
+    device = "cuda"
+    torch.cuda.empty_cache()
+    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
+    if itype == torch.bfloat16:
+        rtol, atol = 1e-2, 5e-2
+
+    seqlens = []
+    batch_size = 4
+    if seqlen < 10:
+        batch_size = 1
+    padding = 3 if with_padding else 0
+    padded_batch_size = batch_size + padding
+    nsplits = padded_batch_size - 1
+
+    eos_pos = torch.randperm(seqlen - 1)[:nsplits].sort().values
+    seqlens.append(
+        torch.diff(
+            torch.cat([torch.tensor([-1]), eos_pos, torch.tensor([seqlen - 1])])
+        ).tolist()
+    )
+    assert sum(seqlens[-1]) == seqlen
+    assert all(s > 0 for s in seqlens[-1])
+
+    total_entries = batch_size * 10
+    cumsum = torch.cumsum(torch.tensor(seqlens[0]), dim=0).to(torch.int32)
+    cumsum = torch.concat([torch.tensor([0], dtype=torch.int32), cumsum], dim=0)
+    x = torch.randn(1, 4096 + dim + 64, seqlen, device=device, dtype=itype)[
+        :, 4096 : 4096 + dim, :
+    ]
+    weight = torch.randn(dim, width, device=device, dtype=itype)
+    bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
+    x_ref = x.clone()
+    weight_ref = weight.clone()
+    bias_ref = bias.clone() if bias is not None else None
+    activation = None if not silu_activation else "silu"
+    final_states = torch.randn(
+        total_entries, dim, width - 1, device=x.device, dtype=x.dtype
+    )
+    final_states_ref = final_states.clone()
+    has_initial_states = torch.randint(
+        0, 2, (cumsum.shape[0] - 1,), dtype=torch.bool, device=x.device
+    )
+    state_indices = torch.randperm(total_entries, dtype=torch.int32, device=x.device)[
+        :batch_size
+    ]
+    padded_state_indices = torch.concat(
+        [
+            state_indices,
+            torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
+        ],
+        dim=-1,
+    )
+
+    out = causal_conv1d_fn(
+        x.squeeze(0),
+        weight,
+        bias,
+        cumsum.cuda(),
+        padded_state_indices,
+        has_initial_states,
+        final_states,
+        activation,
+        PAD_SLOT_ID,
+    )
+    out_ref = []
+    out_ref_b = []
+
+    splits = [torch.split(var, seqlens[0], dim=-1) for var in (x_ref)]
+    for i in range(len(seqlens[0])):
+        x_s = [v[i].unsqueeze(0) for v in splits][0]
+        if padded_state_indices[i] == PAD_SLOT_ID:
+            continue
+        out_ref_b.append(
+            causal_conv1d_ref(
+                x_s,
+                weight_ref,
+                bias_ref,
+                activation=activation,
+                return_final_states=True,
+                final_states_out=final_states_ref[padded_state_indices[i]].unsqueeze(0),
+                initial_states=(
+                    final_states_ref[padded_state_indices[i]].unsqueeze(0)
+                    if has_initial_states[i]
+                    else None
+                ),
+            )
+        )
+    out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=2))
+    out_ref_tensor = torch.cat(out_ref, dim=0)
+
+    unpadded_out = out[:, : out_ref_tensor.shape[-1]]
+    assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol)
+    assert torch.allclose(
+        final_states[state_indices],
+        final_states_ref[state_indices],
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_copy.py

@@ -0,0 +1,18 @@
+import sys
+
+import pytest
+import sgl_kernel
+import torch
+from sgl_kernel.elementwise import copy_to_gpu_no_ce
+
+
+@pytest.mark.parametrize("size", [64, 72])
+def test_copy_to_gpu_no_ce(size):
+    tensor_cpu = torch.randint(0, 1000000, (size,), dtype=torch.int32, device="cpu")
+    tensor_gpu = torch.empty_like(tensor_cpu, device="cuda")
+    copy_to_gpu_no_ce(tensor_cpu, tensor_gpu)
+    assert torch.all(tensor_cpu.cuda() == tensor_gpu)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_custom_allreduce.py

@@ -0,0 +1,185 @@
+import ctypes
+import multiprocessing as mp
+import random
+import socket
+import unittest
+from typing import Any, List, Optional
+
+import sgl_kernel.allreduce as custom_ops
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
+
+
+def _run_correctness_worker(world_size, rank, distributed_init_port, test_sizes):
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+    distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+    dist.init_process_group(
+        backend="nccl",
+        init_method=distributed_init_method,
+        rank=rank,
+        world_size=world_size,
+    )
+    group = dist.group.WORLD
+
+    try:
+        device = torch.device(f"cuda:{rank}")
+        max_size = 8192 * 1024
+        meta_ptrs = TestCustomAllReduce.create_shared_buffer(
+            custom_ops.meta_size() + max_size, group=group
+        )
+
+        rank_data = torch.empty(8 * 1024 * 1024, dtype=torch.uint8, device=device)
+        buffer_ptrs = TestCustomAllReduce.create_shared_buffer(max_size, group=group)
+
+        custom_ptr = custom_ops.init_custom_ar(meta_ptrs, rank_data, rank, True)
+        custom_ops.register_buffer(custom_ptr, buffer_ptrs)
+
+        test_loop = 10
+        for sz in test_sizes:
+            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+                for _ in range(test_loop):
+                    inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device)
+                    inp1_ref = inp1.clone()
+                    out1 = torch.empty_like(inp1)
+
+                    custom_ops.all_reduce(
+                        custom_ptr, inp1, out1, buffer_ptrs[rank], max_size
+                    )
+
+                    dist.all_reduce(inp1_ref, group=group)
+
+                    torch.testing.assert_close(out1, inp1_ref)
+
+    finally:
+        dist.barrier(group=group)
+        if custom_ptr is not None:
+            custom_ops.dispose(custom_ptr)
+        if buffer_ptrs:
+            TestCustomAllReduce.free_shared_buffer(buffer_ptrs, group)
+        if meta_ptrs:
+            TestCustomAllReduce.free_shared_buffer(meta_ptrs, group)
+
+        dist.destroy_process_group(group=group)
+
+
+def get_open_port() -> int:
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("127.0.0.1", 0))
+            return s.getsockname()[1]
+    except OSError:
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+            s.bind(("::1", 0))
+            return s.getsockname()[1]
+
+
+def multi_process_parallel(
+    world_size: int, test_target: Any, target_args: tuple = ()
+) -> None:
+    mp.set_start_method("spawn", force=True)
+
+    procs = []
+    distributed_init_port = get_open_port()
+    for i in range(world_size):
+        proc_args = (world_size, i, distributed_init_port) + target_args
+        proc = mp.Process(target=test_target, args=proc_args, name=f"Worker-{i}")
+        proc.start()
+        procs.append(proc)
+
+    for i in range(world_size):
+        procs[i].join()
+        assert (
+            procs[i].exitcode == 0
+        ), f"Process {i} failed with exit code {procs[i].exitcode}"
+
+
+class TestCustomAllReduce(unittest.TestCase):
+    test_sizes = [
+        512,
+        2560,
+        4096,
+        5120,
+        7680,
+        32768,
+        262144,
+        524288,
+        1048576,
+        2097152,
+    ]
+    world_sizes = [2, 4, 8]
+
+    @staticmethod
+    def create_shared_buffer(
+        size_in_bytes: int, group: Optional[ProcessGroup] = None
+    ) -> List[int]:
+        lib = CudaRTLibrary()
+        pointer = lib.cudaMalloc(size_in_bytes)
+        handle = lib.cudaIpcGetMemHandle(pointer)
+        if group is None:
+            group = dist.group.WORLD
+        world_size = dist.get_world_size(group=group)
+        rank = dist.get_rank(group=group)
+
+        handle_bytes = ctypes.string_at(ctypes.addressof(handle), ctypes.sizeof(handle))
+        input_tensor = torch.ByteTensor(list(handle_bytes)).to(f"cuda:{rank}")
+        gathered_tensors = [torch.empty_like(input_tensor) for _ in range(world_size)]
+        dist.all_gather(gathered_tensors, input_tensor, group=group)
+
+        handles = []
+        handle_type = type(handle)
+        for tensor in gathered_tensors:
+            bytes_list = tensor.cpu().tolist()
+            bytes_data = bytes(bytes_list)
+            handle_obj = handle_type()
+            ctypes.memmove(ctypes.addressof(handle_obj), bytes_data, len(bytes_data))
+            handles.append(handle_obj)
+
+        pointers: List[int] = []
+        for i, h in enumerate(handles):
+            if i == rank:
+                pointers.append(pointer.value)
+            else:
+                try:
+                    opened_ptr = lib.cudaIpcOpenMemHandle(h)
+                    pointers.append(opened_ptr.value)
+                except Exception as e:
+                    print(f"Rank {rank}: Failed to open IPC handle from rank {i}: {e}")
+                    raise
+
+        dist.barrier(group=group)
+        return pointers
+
+    @staticmethod
+    def free_shared_buffer(
+        pointers: List[int], group: Optional[ProcessGroup] = None
+    ) -> None:
+        if group is None:
+            group = dist.group.WORLD
+        rank = dist.get_rank(group=group)
+        lib = CudaRTLibrary()
+        if pointers and len(pointers) > rank and pointers[rank] is not None:
+            lib.cudaFree(ctypes.c_void_p(pointers[rank]))
+        dist.barrier(group=group)
+
+    def test_correctness(self):
+        for world_size in self.world_sizes:
+            available_gpus = torch.cuda.device_count()
+            if world_size > available_gpus:
+                print(
+                    f"Skipping world_size={world_size}, requires {world_size} GPUs, found {available_gpus}"
+                )
+                continue
+
+            print(f"Running test for world_size={world_size}")
+            multi_process_parallel(
+                world_size, _run_correctness_worker, target_args=(self.test_sizes,)
+            )
+            print(f"custom allreduce tp = {world_size}: OK")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/sgl-kernel/tests/test_cutlass_mla.py

@@ -0,0 +1,106 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size
+from torch import Tensor
+
+# Disable tests on SM103 until the accuracy issues are fixed.
+if torch.cuda.get_device_capability() != (10, 0):
+    pytest.skip(
+        reason="Cutlass MLA Requires compute capability of 10.",
+        allow_module_level=True,
+    )
+
+
+def ref_mla(
+    out: Tensor,  # (bs, num_heads, v_head_dim)
+    query: Tensor,  # (bs, num_heads, head_dim)
+    kv_cache: Tensor,  # (num_blocks, block_size, head_dim)
+    scale: float,
+    block_tables: Tensor,  # (bs, max_num_blocks)
+    seq_lens: Tensor,  # (bs,)
+):
+    bs, num_heads, v_head_dim = out.shape
+    head_dim = query.shape[2]
+
+    for i in range(bs):
+        # gather and flatten KV-cache
+        kv = kv_cache[block_tables[i]]  # (max_num_blocks, block_size, head_dim)
+        kv = kv.view(1, -1, head_dim)[:, : seq_lens[i]]  # (1, seq_len, head_dim)
+        v = kv[:, :, :v_head_dim]
+
+        q = query[i].view(num_heads, 1, head_dim)
+        o = F.scaled_dot_product_attention(q, kv, v, scale=scale, enable_gqa=True)
+        out[i] = o.view(num_heads, v_head_dim)
+
+    return out
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("mean_seq_len", [128, 1024, 4096])
+@pytest.mark.parametrize("bs", [1, 2, 4])
+@pytest.mark.parametrize("varlen", [False, True])
+@pytest.mark.parametrize("block_size", [1, 16, 64, 128])
+@pytest.mark.parametrize("num_heads", [16, 32, 64, 128])
+@pytest.mark.parametrize("num_kv_splits", [-1, 1])
+def test_cutlass_mla_decode(
+    dtype: torch.dtype,
+    mean_seq_len: int,
+    bs: int,
+    varlen: bool,
+    block_size: int,
+    num_heads: int,
+    num_kv_splits: int,
+):
+    torch.set_default_dtype(dtype)
+    torch.set_default_device("cuda")
+    torch.manual_seed(42)
+
+    d = 576
+    h_q = num_heads
+    dv = 512
+
+    q_nope_dim = 128
+    q_pe_dim = 64
+    scale = (q_nope_dim + q_pe_dim) ** (-0.5)
+    if varlen:
+        seq_lens = torch.empty(bs).normal_(mean_seq_len, mean_seq_len / 2)
+        seq_lens = seq_lens.clip(2).to(torch.int32)
+    else:
+        seq_lens = torch.full((bs,), mean_seq_len, dtype=torch.int32)
+    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
+
+    # Lager q values to detect split kv error
+    q = torch.randn(bs, h_q, d) * 100.0
+    block_table = torch.randint(0, bs * block_num, (bs, block_num), dtype=torch.int32)
+
+    kv_cache = torch.randn(block_table.numel(), block_size, d)
+
+    workspace_size = cutlass_mla_get_workspace_size(
+        block_num * block_size, bs, num_kv_splits=num_kv_splits
+    )
+    workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8)
+
+    q_nope = torch.empty((h_q, bs, dv)).transpose(0, 1)
+    q_nope.copy_(q[:, :, :dv])
+    q_pe = q[:, :, dv:].clone()
+
+    out_ref = q.new_zeros(bs, h_q, dv)
+    ref_mla(out_ref, q, kv_cache, scale, block_table, seq_lens)
+    out = cutlass_mla_decode(
+        q_nope, q_pe, kv_cache, seq_lens, block_table, workspace, scale, num_kv_splits
+    )
+
+    torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py

@@ -0,0 +1,287 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import cutlass_w4a8_moe_mm
+from utils import is_hopper
+
+from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+
+
+def pack_int4_values_to_int8(int4_values_interleaved: torch.Tensor) -> torch.Tensor:
+    if int4_values_interleaved.shape[-1] % 2 != 0:
+        raise ValueError(
+            "the last dim size of int4_values_interleaved tensor must be even."
+        )
+
+    input_tensor_int8 = int4_values_interleaved.to(torch.int8)
+
+    low_nibbles = input_tensor_int8[..., 0::2]
+    high_nibbles = input_tensor_int8[..., 1::2]
+
+    packed_tensor = (high_nibbles << 4) | (low_nibbles & 0x0F)
+
+    return packed_tensor.to(torch.int8)
+
+
+def pack_interleave(num_experts, ref_weight, ref_scale):
+    n, k = ref_weight.shape[1], ref_weight.shape[2]
+
+    weight = pack_int4_values_to_int8(ref_weight.cpu()).cuda()
+    w_q = weight.view((num_experts, n, k // 2)).view(torch.int8)
+    w_q = w_q.contiguous()
+
+    alignment = 4 if k % 512 == 0 else 1
+    scale_interleaved = ref_scale.reshape(
+        ref_scale.shape[0],
+        ref_scale.shape[1],
+        (ref_scale.shape[2] // alignment),
+        alignment,
+    )  # [E, N, K/4, 4]
+    scale_interleaved = scale_interleaved.permute(0, 2, 1, 3)  # [E, K/4, N, 4]
+    scale_interleaved = scale_interleaved.reshape(
+        ref_scale.shape[0],
+        ref_scale.shape[2] // alignment,
+        ref_scale.shape[1] * alignment,
+    )  # [E, K/4, N*4]
+    w_scale = scale_interleaved.contiguous()
+
+    return w_q, w_scale
+
+
+@pytest.mark.skipif(
+    not is_hopper(),
+    reason="cutlass_w4a8_moe_mm is only supported on sm90",
+)
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16])
+def test_int4_fp8_grouped_gemm_single_expert(batch_size):
+    # Test parameters
+    num_experts = 1
+    m = batch_size  # batch size
+    k = 512  # input dimension
+    n = 1024  # output dimension
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    device = "cuda"
+    debug = False
+
+    print(f"\nTesting with batch_size={batch_size}")
+
+    # Create input tensors with ones
+    if debug:
+        a = torch.ones(m, k, dtype=torch.bfloat16, device=device)
+        ref_w = torch.ones(num_experts, n, k, dtype=torch.int8, device=device)
+        ref_w_scale = torch.ones(num_experts, n, k // 128, dtype=dtype, device=device)
+    else:
+        a = torch.randn(m, k, dtype=dtype, device=device)
+        ref_w = torch.randint(
+            -8, 8, (num_experts, n, k), dtype=torch.int8, device=device
+        )
+        affine_coeff = 0.005
+        ref_w_scale = (
+            torch.randn(num_experts, n, k // 128, dtype=dtype, device=device)
+            * affine_coeff
+        )
+
+    w, w_scale = pack_interleave(num_experts, ref_w, ref_w_scale)
+
+    # Create expert offsets and problem sizes
+    expert_offsets = torch.tensor([0, m], dtype=torch.int32, device=device)
+    problem_sizes = torch.tensor([[n, m, k]], dtype=torch.int32, device=device)
+
+    a_strides = torch.full((num_experts, 3), k, device=device, dtype=torch.int64)
+    c_strides = torch.full((num_experts, 3), n, device=device, dtype=torch.int64)
+    b_strides = a_strides
+    s_strides = c_strides
+
+    # Quantize input
+    a_q, a_scale = _per_tensor_quant_fp8(a)
+
+    # Create output tensor
+    c = torch.empty((m, n), dtype=torch.bfloat16, device=device)
+    cutlass_w4a8_moe_mm(
+        c,
+        a_q,
+        w,
+        a_scale,
+        w_scale,
+        expert_offsets[:-1],
+        problem_sizes,
+        a_strides,
+        b_strides,
+        c_strides,
+        s_strides,
+        128,
+        8,
+    )
+    c = c.to(dtype)
+
+    # Reference implementation
+    experts_selection_result = torch.full((m,), 0)
+    c_ref = ref_grouped_gemm(
+        c, a_q, a_scale, ref_w, ref_w_scale, num_experts, experts_selection_result
+    )
+
+    # Compare results
+    try:
+        torch.testing.assert_close(c, c_ref, rtol=1e-2, atol=0.1)
+    except AssertionError as e:
+        # torch.set_printoptions(threshold=10_000)
+        print(f"  FAILURE: tensors are NOT close.")
+        print(f"    Ref tensor: {c_ref.flatten()}")
+        print(f"    Cutlass tensor: {c.flatten()}")
+        print(
+            f"    Max absolute difference: {torch.max(torch.abs(c.to(c_ref.dtype) - c_ref))}"
+        )
+        print(
+            f"    Mean absolute difference: {torch.mean(torch.abs(c.to(c_ref.dtype) - c_ref))}"
+        )
+        print(f"    AssertionError: {e}")
+        raise
+
+
+def _per_tensor_quant_fp8(
+    x: torch.Tensor,
+    dtype: torch.dtype = torch.float8_e4m3fn,
+):
+    assert x.is_contiguous(), "`x` is not contiguous"
+
+    x_q = torch.empty_like(x, device=x.device, dtype=dtype)
+    x_s = torch.empty(
+        1,
+        device=x.device,
+        dtype=torch.float32,
+    )
+    per_tensor_quant_fp8(x, x_q, x_s, is_static=False)
+    return x_q, x_s
+
+
+@pytest.mark.skipif(
+    not is_hopper(),
+    reason="cutlass_w4a8_moe_mm is only supported on sm90",
+)
+@pytest.mark.parametrize("batch_size", [2, 4, 8, 16, 32])
+@pytest.mark.parametrize("k", [256, 512, 1024, 2048, 4096, 7168])
+@pytest.mark.parametrize("n", [256, 512, 1024, 2048, 7168])
+@pytest.mark.parametrize("num_experts", [2, 4, 6, 8])
+def test_int4_fp8_grouped_gemm_multi_experts(batch_size, k, n, num_experts):
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    device = "cuda"
+    debug = False
+
+    print(
+        f"\nTesting with batch_size={batch_size}, k={k}, n={n}, num_experts={num_experts}"
+    )
+
+    if debug:
+        a = torch.ones(batch_size, k, dtype=torch.bfloat16, device=device)
+        ref_w = torch.ones(num_experts, n, k, dtype=torch.int8, device=device)
+        ref_w_scale = torch.ones(num_experts, n, k // 128, dtype=dtype, device=device)
+    else:
+        a = torch.randn(batch_size, k, dtype=dtype, device=device)
+        ref_w = torch.randint(
+            -8, 8, (num_experts, n, k), dtype=torch.int8, device=device
+        )
+        affine_coeff = 0.005
+        ref_w_scale = (
+            torch.randn(num_experts, n, k // 128, dtype=dtype, device=device)
+            * affine_coeff
+        )
+
+    w, w_scale = pack_interleave(num_experts, ref_w, ref_w_scale)
+
+    # random select experts
+    experts_selection_result = torch.randint(
+        0, num_experts, (batch_size,), device=device
+    )
+    permutation = torch.argsort(experts_selection_result)
+    expert_token_counts = torch.bincount(
+        experts_selection_result, minlength=num_experts
+    )
+
+    # Create problem sizes and offsets for active experts
+    problem_sizes = []
+    for i in range(num_experts):
+        problem_sizes.append([n, expert_token_counts[i].item(), k])
+    problem_sizes = torch.tensor(problem_sizes, dtype=torch.int32, device=device)
+
+    expert_offsets = []
+    offset = 0
+    for i in range(num_experts):
+        expert_offsets.append(offset)
+        offset += problem_sizes[i][1].item()
+    expert_offsets = torch.tensor(expert_offsets, dtype=torch.int32, device=device)
+
+    # Permute input and quantize
+    a_q, a_scale = _per_tensor_quant_fp8(a)
+    a_q_perm = a_q[permutation]
+
+    # Create stride tensors
+    a_strides = torch.full((num_experts, 3), k, device=device, dtype=torch.int64)
+    c_strides = torch.full((num_experts, 3), n, device=device, dtype=torch.int64)
+    b_strides = a_strides
+    s_strides = c_strides
+
+    c_perm = torch.empty((batch_size, n), dtype=torch.bfloat16, device=device)
+    cutlass_w4a8_moe_mm(
+        c_perm,
+        a_q_perm,
+        w,
+        a_scale,
+        w_scale,
+        expert_offsets,
+        problem_sizes,
+        a_strides,
+        b_strides,
+        c_strides,
+        s_strides,
+        128,
+        8,
+    )
+
+    # Un-permute the result
+    c = torch.empty_like(c_perm)
+    c[permutation] = c_perm
+    c = c.to(dtype)
+
+    c_ref = ref_grouped_gemm(
+        c, a_q, a_scale, ref_w, ref_w_scale, num_experts, experts_selection_result
+    )
+
+    # Compare results
+    try:
+        torch.testing.assert_close(c, c_ref, rtol=1e-2, atol=0.1)
+    except AssertionError as e:
+        print(f"  FAILURE: tensors are NOT close.")
+        print(
+            f"    Max absolute difference: {torch.max(torch.abs(c.to(c_ref.dtype) - c_ref))}"
+        )
+        print(
+            f"    Mean absolute difference: {torch.mean(torch.abs(c.to(c_ref.dtype) - c_ref))}"
+        )
+        print(f"    AssertionError: {e}")
+        raise
+
+
+def ref_grouped_gemm(
+    c, a_q, a_scale, w, w_scale, num_experts, experts_selection_result
+):
+    dtype = torch.bfloat16
+    c_ref = torch.zeros_like(c)
+    for i in range(num_experts):
+        token_idx = torch.where(experts_selection_result == i)[0]
+        if len(token_idx) == 0:
+            continue
+        a = a_q[token_idx]
+
+        ref_w_scale_repeat = w_scale[i].repeat_interleave(128, dim=1).to(torch.float32)
+        ref_w = w[i].to(torch.float32) * ref_w_scale_repeat
+        c = torch.matmul(a.to(torch.float32), ref_w.t()) * a_scale
+        c_ref[token_idx] = c.to(dtype)
+
+    return c_ref
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_dsv3_fused_a_gemm.py

@@ -0,0 +1,34 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import dsv3_fused_a_gemm
+
+
+@pytest.mark.parametrize("num_tokens", [i + 1 for i in range(16)])
+def test_dsv3_fused_a_gemm(num_tokens):
+    kHdIn = 7168
+    kHdOut = 2112
+
+    mat_a = torch.randn(
+        (num_tokens, kHdIn), dtype=torch.bfloat16, device="cuda"
+    ).contiguous()
+    mat_b = torch.randn((kHdOut, kHdIn), dtype=torch.bfloat16, device="cuda").transpose(
+        0, 1
+    )
+    output = torch.empty(
+        (num_tokens, kHdOut), dtype=torch.bfloat16, device="cuda"
+    ).contiguous()
+
+    ref = F.linear(mat_a, mat_b.T)
+
+    output = dsv3_fused_a_gemm(mat_a, mat_b)
+
+    assert torch.allclose(
+        output, ref, rtol=1e-2, atol=1e-3
+    ), "Fused GEMM output mismatch with torch.nn.functional.linear reference"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_dsv3_router_gemm.py

@@ -0,0 +1,37 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+from sgl_kernel import dsv3_router_gemm
+
+
+@pytest.mark.parametrize("num_tokens", [i + 1 for i in range(16)])
+@pytest.mark.parametrize("num_experts", [256, 384])
+def test_dsv3_router_gemm(num_tokens, num_experts):
+    hidden_dim = 7168
+
+    mat_a = torch.randn(
+        (num_tokens, hidden_dim), dtype=torch.bfloat16, device="cuda"
+    ).contiguous()
+    mat_b = torch.randn(
+        (num_experts, hidden_dim), dtype=torch.bfloat16, device="cuda"
+    ).contiguous()
+
+    bf16_ref = F.linear(mat_a, mat_b)
+    float_ref = bf16_ref.to(torch.float32)
+
+    bf16_output = dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.bfloat16)
+    float_output = dsv3_router_gemm(mat_a, mat_b, out_dtype=torch.float32)
+
+    assert torch.allclose(
+        bf16_output, bf16_ref, rtol=1e-2, atol=1e-3
+    ), "Router GEMM output in bf16 dtype mismatch with torch.nn.functional.linear reference"
+
+    assert torch.allclose(
+        float_output, float_ref, rtol=1e-2, atol=1e-3
+    ), "Router GEMM output in float32 dtype mismatch with torch.nn.functional.linear reference"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_es_fp8_blockwise_moe.py

@@ -0,0 +1,206 @@
+import random
+import sys
+from typing import Tuple
+
+import pytest
+import torch
+from sgl_kernel import es_fp8_blockwise_scaled_grouped_mm
+
+
+def cdiv(a: int, b: int) -> int:
+    return -(a // -b)
+
+
+def scale_shape(shape, group_shape):
+    return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))
+
+
+def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(
+        dtype=torch.float8_e4m3fn
+    )
+
+
+# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+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 baseline_scaled_mm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale_a: torch.Tensor,
+    scale_b: torch.Tensor,
+    out_dtype: type[torch.dtype],
+) -> torch.Tensor:
+
+    def group_broadcast(t, shape):
+        for i, s in enumerate(shape):
+            if t.shape[i] != s and t.shape[i] != 1:
+                assert s % t.shape[i] == 0
+                t = (
+                    t.unsqueeze(i + 1)
+                    .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
+                    .flatten(i, i + 1)
+                )
+        return t
+
+    scale_a = group_broadcast(scale_a, a.shape)
+    scale_b = group_broadcast(scale_b, b.shape)
+
+    return torch.mm(
+        (scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32))
+    ).to(out_dtype)
+
+
+def is_sm100_supported(device=None) -> bool:
+    return (torch.cuda.get_device_capability(device)[0] == 10) and (
+        torch.version.cuda >= "12.8"
+    )
+
+
+def is_sm90_supported(device=None) -> bool:
+    return (torch.cuda.get_device_capability(device)[0] == 9) and (
+        torch.version.cuda >= "12.3"
+    )
+
+
+@pytest.mark.skipif(
+    not is_sm90_supported(),
+    reason="es_fp8_blockwise_scaled_grouped_mm at sgl-kernel is only supported on sm90",
+)
+@pytest.mark.parametrize("num_experts", [8, 16, 32, 64, 128])
+@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
+def test_fp8_blockwise_scaled_grouped_mm(num_experts, out_dtype):
+    device = "cuda"
+    alignment = 128
+    n_g = random.randint(1, 64) * alignment
+    k_g = random.randint(1, 64) * alignment
+
+    expert_offsets = torch.zeros((num_experts + 1), device=device, dtype=torch.int32)
+    problem_sizes = torch.zeros((num_experts, 3), device=device, dtype=torch.int32)
+    a_tensors = []
+    b_tensors = []
+    a_scales_tensors = []
+    b_scales_tensors = []
+    baseline_tensors = []
+
+    for g in range(num_experts):
+        m_g = random.randint(1, 256)
+        expert_offsets[g + 1] = expert_offsets[g] + m_g
+        problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
+
+        a = torch.randn((m_g, k_g), device=device, dtype=out_dtype)  # (M, K):(K, 1)
+        b = torch.randn((n_g, k_g), device=device, dtype=out_dtype).t()  # (K, N):(1, K)
+
+        a_g, a_scale = per_token_cast_to_fp8(
+            a
+        )  # ag -- (M, K):(K, 1), a_scale() -- (M, k):(k, 1)
+        b_g, b_scale = per_block_cast_to_fp8(
+            b
+        )  # bg -- (K, N):(N, 1), b_scale() -- (k, n):(n, 1)
+        a_tensors.append(a_g)
+        b_tensors.append(b_g)
+        a_scales_tensors.append(a_scale)
+        b_scales_tensors.append(b_scale)
+
+        baseline = torch.mm(a, b)
+        baseline_tensors.append(baseline)
+    a_stack = torch.empty(
+        (expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    b_stack = torch.empty(
+        (num_experts, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    a_scale_stack = torch.empty(
+        (expert_offsets[-1], (k_g // 128)), device=device, dtype=torch.float32
+    )
+    b_scale_stack = torch.empty(
+        (num_experts, n_g // 128, k_g // 128), device=device, dtype=torch.float32
+    )
+
+    for g in range(num_experts):
+        # Matrix A is Row-Major.
+        a_stack[expert_offsets[g] : expert_offsets[g + 1], :] = a_tensors[
+            g
+        ]  # a_stack[expert_offsets[g] : expert_offsets[g + 1], :] -- (M, K):(K, 1)
+        b_stack[g] = b_tensors[g].t()  # b_stack[g] -- (N, K):(K, 1)
+
+        # We need K-Major scale factor
+        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[g] -- (k, n):(n, 1), we need transpose & contiguous later
+    b_stack = b_stack.transpose(1, 2)  # Transpose Matrix B to Column-Major.
+    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_experts,), a_stack.stride(0), device=device, dtype=torch.int64
+    )
+    d_strides = torch.full(
+        (num_experts,), c_out.stride(0), device=device, dtype=torch.int64
+    )
+
+    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,
+    )
+
+    for g in range(num_experts):
+        baseline = baseline_tensors[g]
+        actual = c_out[expert_offsets[g] : expert_offsets[g + 1]]
+        diff = calc_diff(actual, baseline)
+        assert diff < 0.001
+        print(
+            f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, out_dtype={out_dtype}, diff={diff:.5f}: OK"
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_es_mxfp8_blockscaled_moe.py

@@ -0,0 +1,156 @@
+import random
+import sys
+
+import pytest
+import torch
+from sgl_kernel import (
+    es_sm100_mxfp8_blockscaled_grouped_mm,
+    es_sm100_mxfp8_blockscaled_grouped_quant,
+)
+
+random.seed(42)
+torch.manual_seed(42)
+torch.cuda.manual_seed(42)
+torch.cuda.manual_seed_all(42)
+
+
+def align(val: int, alignment: int = 128) -> int:
+    return int((val + alignment - 1) // alignment * alignment)
+
+
+# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def is_sm100_supported(device=None) -> bool:
+    return (torch.cuda.get_device_capability(device)[0] == 10) and (
+        torch.version.cuda >= "12.8"
+    )
+
+
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason="test_es_sm100_mxfp8_blockscaled_grouped_mm at sgl-kernel is only supported on sm100",
+)
+@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
+@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
+def test_es_sm100_mxfp8_blockscaled_grouped_mm(num_experts, out_dtype):
+    device = "cuda"
+    alignment = 128
+    n_g = random.randint(1, 64) * alignment
+    k_g = random.randint(1, 64) * alignment
+
+    expert_offset = 0
+    expert_offsets = []
+    aux_expert_offset = 0
+    aux_expert_offsets = []
+    a_blockscale_offset = 0
+    a_blockscale_offsets = []
+    b_blockscale_offset = 0
+    b_blockscale_offsets = []
+    problem_sizes = []
+    aux_problem_sizes = []
+    a_list = []
+    b_list = []
+    ref_d_list = []
+
+    for g in range(num_experts):
+        m_g = random.randint(1, 512)
+        expert_offsets.append(expert_offset)
+        expert_offset += m_g
+        aux_expert_offsets.append(aux_expert_offset)
+        aux_expert_offset += n_g
+        a_blockscale_offsets.append(a_blockscale_offset)
+        a_blockscale_offset += align(m_g, 128)
+        b_blockscale_offsets.append(b_blockscale_offset)
+        b_blockscale_offset += n_g  # n_g already align to 128
+        problem_sizes.append([m_g, n_g, k_g])
+        aux_problem_sizes.append([n_g, m_g, k_g])
+
+        a = torch.normal(
+            0.0, std=1.0, size=(m_g, k_g), device=device, dtype=out_dtype
+        )  # (M, K):(K, 1)
+        b = torch.normal(
+            0.0, std=1.0, size=(n_g, k_g), device=device, dtype=out_dtype
+        )  # (N, K):(K, 1)
+
+        a_list.append(a)
+        b_list.append(b)
+        ref_d = a @ b.T
+        ref_d_list.append(ref_d)
+    a = torch.concat(a_list, dim=0)
+    b = torch.concat(b_list, dim=0)
+
+    _problem_sizes = torch.tensor(problem_sizes).to(device=device, dtype=torch.int32)
+    _aux_problem_sizes = torch.tensor(aux_problem_sizes).to(
+        device=device, dtype=torch.int32
+    )
+    _expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32)
+    _aux_expert_offsets = torch.tensor(aux_expert_offsets).to(
+        device=device, dtype=torch.int32
+    )
+    _a_blockscale_offsets = torch.tensor(a_blockscale_offsets).to(
+        device=device, dtype=torch.int32
+    )
+    _b_blockscale_offsets = torch.tensor(b_blockscale_offsets).to(
+        device=device, dtype=torch.int32
+    )
+
+    a_quant = torch.zeros_like(a, dtype=torch.float8_e4m3fn, device=device)
+    a_scale_factor = torch.zeros(
+        (a_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device
+    )
+
+    b_quant = torch.zeros_like(b, dtype=torch.float8_e4m3fn, device=device)
+    b_scale_factor = torch.zeros(
+        (num_experts, n_g, k_g // 32), dtype=torch.uint8, device=device
+    )
+
+    es_sm100_mxfp8_blockscaled_grouped_quant(
+        a,
+        _problem_sizes,
+        _expert_offsets,
+        _a_blockscale_offsets,
+        a_quant,
+        a_scale_factor,
+    )
+
+    es_sm100_mxfp8_blockscaled_grouped_quant(
+        b,
+        _aux_problem_sizes,
+        _aux_expert_offsets,
+        _b_blockscale_offsets,
+        b_quant,
+        b_scale_factor,
+    )
+    b_quant = b_quant.view(num_experts, n_g, k_g).transpose(1, 2)
+    b_scale_factor = b_scale_factor.view(num_experts, n_g, k_g // 32).transpose(1, 2)
+
+    d = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
+    es_sm100_mxfp8_blockscaled_grouped_mm(
+        d,
+        a_quant,
+        b_quant,
+        a_scale_factor,
+        b_scale_factor,
+        _problem_sizes,
+        _expert_offsets,
+        _a_blockscale_offsets,
+    )
+
+    for g in range(num_experts):
+        baseline = ref_d_list[g]
+        actual = d[expert_offsets[g] : (expert_offsets[g] + problem_sizes[g][0])]
+        diff = calc_diff(actual, baseline)
+        assert diff < 0.001
+        print(
+            f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, out_dtype={out_dtype}, diff={diff:.5f}: OK"
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_flash_attn_sparse.py

@@ -0,0 +1,493 @@
+import math
+import sys
+from typing import List, Optional, Tuple
+
+import pytest
+import torch
+from einops import rearrange, repeat
+from sgl_kernel.sparse_flash_attn import (
+    convert_vertical_slash_indexes,
+    convert_vertical_slash_indexes_mergehead,
+    sparse_attn_func,
+)
+from test_flash_attention import construct_local_mask, is_fa3_supported
+
+
+def ref_attn(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    attn_bias=None,
+    dropout_p=0.0,
+    dropout_mask=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite window size
+    softcap=0.0,
+    upcast=True,
+    reorder_ops=False,
+    key_leftpad=None,
+):
+    """
+    Arguments:
+        q: (batch_size, seqlen_q, nheads, head_dim)
+        k: (batch_size, seqlen_k, nheads_k, head_dim)
+        v: (batch_size, seqlen_k, nheads_k, head_dim)
+        query_padding_mask: (batch_size, seqlen_q)
+        key_padding_mask: (batch_size, seqlen_k)
+        attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k)
+        dropout_p: float
+        dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k)
+        causal: whether to apply causal masking
+        window_size: (int, int), left and right window size
+        upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast
+            output back to fp16/bf16.
+        reorder_ops: whether to change the order of operations (scaling k instead of scaling q, etc.)
+            without changing the math. This is to estimate the numerical error from operation
+            reordering.
+    Output:
+        output: (batch_size, seqlen_q, nheads, head_dim)
+        lse: (batch_size, nheads, seqlen_q)
+    """
+    if causal:
+        window_size = (window_size[0], 0)
+    dtype_og = q.dtype
+    if upcast:
+        q, k, v = q.float(), k.float(), v.float()
+    seqlen_q, seqlen_k = q.shape[1], k.shape[1]
+    k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2])
+    v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2])
+    d = q.shape[-1]
+    if not reorder_ops:
+        scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k)
+    else:
+        scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d))
+
+    lse_ref = scores.logsumexp(dim=-1)
+
+    if softcap > 0:
+        scores = scores / softcap
+        scores = scores.tanh()
+        scores = scores * softcap
+    if key_padding_mask is not None:
+        scores.masked_fill_(
+            rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")
+        )
+    if window_size[0] >= 0 or window_size[1] >= 0:
+        local_mask = construct_local_mask(
+            seqlen_q,
+            seqlen_k,
+            window_size,
+            query_padding_mask,
+            key_padding_mask,
+            q.device,
+            key_leftpad=key_leftpad,
+        )
+        scores.masked_fill_(local_mask, float("-inf"))
+    if attn_bias is not None:
+        scores = scores + attn_bias
+    attention = torch.softmax(scores, dim=-1).to(v.dtype)
+    # Some rows might be completely masked out so we fill them with zero instead of NaN
+    if window_size[0] >= 0 or window_size[1] >= 0:
+        attention = attention.masked_fill(
+            torch.all(local_mask, dim=-1, keepdim=True), 0.0
+        )
+    # We want to mask here so that the attention matrix doesn't have any NaNs
+    # Otherwise we'll get NaN in dV
+    if query_padding_mask is not None:
+        attention = attention.masked_fill(
+            rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0
+        )
+    dropout_scaling = 1.0 / (1 - dropout_p)
+    # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling
+    # output = torch.einsum('bhts,bshd->bthd', attention_drop , v)
+    if dropout_mask is not None:
+        attention_drop = attention.masked_fill(~dropout_mask, 0.0)
+    else:
+        attention_drop = attention
+    output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling)
+    if query_padding_mask is not None:
+        output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0)
+
+    return output.to(dtype=dtype_og), lse_ref
+
+
+def ref_paged_attn(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    query_lens: List[int],
+    kv_lens: List[int],
+    block_tables: torch.Tensor,
+    scale: float,
+    sliding_window: Optional[int] = None,
+    soft_cap: Optional[float] = None,
+) -> torch.Tensor:
+    num_seqs = len(query_lens)
+    block_tables = block_tables.cpu().numpy()
+    _, block_size, num_kv_heads, head_size = key_cache.shape
+
+    outputs: List[torch.Tensor] = []
+    start_idx = 0
+    for i in range(num_seqs):
+        query_len = query_lens[i]
+        kv_len = kv_lens[i]
+        # clone to avoid clobbering the query tensor
+        q = query[start_idx : start_idx + query_len].clone()
+        q *= scale
+
+        num_kv_blocks = (kv_len + block_size - 1) // block_size
+        block_indices = block_tables[i, :num_kv_blocks]
+
+        k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
+        k = k[:kv_len]
+        v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
+        v = v[:kv_len]
+
+        if q.shape[1] != k.shape[1]:
+            k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
+            v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
+        attn = torch.einsum("qhd,khd->hqk", q, k).float()
+        empty_mask = torch.ones(query_len, kv_len)
+        mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
+        if sliding_window is not None:
+            sliding_window_mask = (
+                torch.triu(
+                    empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
+                )
+                .bool()
+                .logical_not()
+            )
+            mask |= sliding_window_mask
+        if soft_cap is not None:
+            attn = soft_cap * torch.tanh(attn / soft_cap)
+        attn.masked_fill_(mask, float("-inf"))
+        attn = torch.softmax(attn, dim=-1).to(v.dtype)
+        out = torch.einsum("hqk,khd->qhd", attn, v)
+
+        outputs.append(out)
+        start_idx += query_len
+
+    return torch.cat(outputs, dim=0)
+
+
+@pytest.mark.skipif(
+    not is_fa3_supported(),
+    reason="flash_attn at sgl-kernel is only supported on sm90 or sm80",
+)
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize(
+    "seq_lens",
+    [
+        (1, 1),
+        (1, 1024),
+        (1, 2048),
+        (1023, 2049),
+        (1023, 1023),
+        (32, 32),
+        (65, 65),
+        (129, 129),
+    ],
+)
+@pytest.mark.parametrize("num_heads", [1, 2, 4])
+@pytest.mark.parametrize("head_size", [128])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("NNZ_S", [0, 1, 2, 3, 7, 15, 32])
+@torch.inference_mode()
+def test_sparse_attention(
+    batch_size,
+    seq_lens,
+    num_heads,
+    head_size,
+    dtype,
+    NNZ_S,
+) -> None:
+    torch.set_default_device("cuda")
+    torch.cuda.manual_seed_all(0)
+    block_size_M = 64
+    block_size_N = 64
+    seqlen_q, seqlen_k = seq_lens
+    q = torch.randn(
+        batch_size, seqlen_q, num_heads, head_size, dtype=dtype, requires_grad=False
+    )
+    k = torch.randn(
+        batch_size, seqlen_k, num_heads, head_size, dtype=dtype, requires_grad=False
+    )
+    v = torch.randn(
+        batch_size, seqlen_k, num_heads, head_size, dtype=dtype, requires_grad=False
+    )
+    NUM_ROWS = (seqlen_q + block_size_M - 1) // block_size_M
+    if NNZ_S * block_size_N > seqlen_k:
+        return
+    NNZ_V = seqlen_k - NNZ_S * block_size_N
+    block_count = torch.tensor(
+        [NNZ_S] * batch_size * NUM_ROWS * num_heads, dtype=torch.int32
+    ).reshape(batch_size, num_heads, NUM_ROWS)
+    column_count = torch.tensor(
+        [NNZ_V] * batch_size * NUM_ROWS * num_heads, dtype=torch.int32
+    ).reshape(batch_size, num_heads, NUM_ROWS)
+    block_offset = torch.tensor(
+        [[i * block_size_N for i in range(NNZ_S)]] * batch_size * NUM_ROWS * num_heads,
+        dtype=torch.int32,
+    ).reshape(batch_size, num_heads, NUM_ROWS, NNZ_S)
+    column_index = torch.tensor(
+        [[NNZ_S * block_size_N + i for i in range(NNZ_V)]]
+        * batch_size
+        * NUM_ROWS
+        * num_heads,
+        dtype=torch.int32,
+    ).reshape(batch_size, num_heads, NUM_ROWS, NNZ_V)
+    out, lse = sparse_attn_func(
+        q,
+        k,
+        v,
+        block_count,
+        block_offset,
+        column_count,
+        column_index,
+        return_softmax_lse=True,
+    )
+
+    ref_out, ref_lse = ref_attn(q, k, v)
+
+    torch.testing.assert_close(
+        out, ref_out, atol=2e-2, rtol=1e-2
+    ), f"{torch.max(torch.abs(out - ref_out))}"
+    torch.testing.assert_close(
+        lse, ref_lse, atol=2e-2, rtol=1e-2
+    ), f"{torch.max(torch.abs(lse - ref_lse))}"
+
+
+# sparse attention utils
+# origin
+@pytest.mark.skipif(
+    not is_fa3_supported(),
+    reason="flash_attn at sgl-kernel is only supported on sm90 or sm80",
+)
+@pytest.mark.parametrize("causal", [True, False])
+def test_convert_vertical_slash_indexes(causal):
+    # Prepare small, hand-checkable inputs
+    q_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda")  # [BATCH]
+    kv_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda")
+    vertical_indexes = torch.tensor(
+        [[[1, 3]]], dtype=torch.int32, device="cuda"
+    )  # [BATCH, N_HEADS, NNZ_V]
+    slash_indexes = torch.tensor(
+        [[[2]]], dtype=torch.int32, device="cuda"
+    )  # [BATCH, N_HEADS, NNZ_S]
+    context_size = 4
+    block_size_M = 2
+    block_size_N = 2
+
+    # Call your CUDA kernel wrapper
+    block_count, block_offset, column_count, column_index = (
+        convert_vertical_slash_indexes(
+            q_seqlens,
+            kv_seqlens,
+            vertical_indexes,
+            slash_indexes,
+            context_size,
+            block_size_M,
+            block_size_N,
+            causal=causal,
+        )
+    )
+
+    # Manually create expected outputs for this input
+    # There are 2 rows (blocks): row0 (tokens 0-1), row1 (tokens 2-3)
+    # Fill these expected tensors based on your CUDA kernel's logic
+    # For demonstration, we assume:
+    # - block_count: how many slash indices fall into each block
+    # - block_offset: the value of those indices
+    # - column_count: number of valid vertical indices per block
+    # - column_index: the actual vertical indices
+
+    expected_column_index = torch.tensor(
+        [[[[0, 0], [0, 0]]]], dtype=torch.int32, device="cuda"
+    )
+
+    # If causal=False, update these tensors according to expected behavior
+    if not causal:
+        # Update these values if your kernel produces different output in non-causal mode
+        expected_column_index = torch.tensor(
+            [[[[1, 0], [1, 3]]]], dtype=torch.int32, device="cuda"
+        )
+
+    # Assert that outputs match expectations
+    assert torch.equal(column_index, expected_column_index)
+
+
+# mergehead
+@pytest.mark.skipif(
+    not is_fa3_supported(),
+    reason="flash_attn at sgl-kernel is only supported on sm90 or sm80",
+)
+@pytest.mark.parametrize("causal", [True, False])
+def test_convert_vertical_slash_indexes_mergehead(causal):
+    # Prepare small, hand-checkable inputs for mergehead version
+    q_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda")
+    kv_seqlens = torch.tensor([4], dtype=torch.int32, device="cuda")
+    vertical_indexes = torch.tensor(
+        [
+            [
+                [1, 3],  # head 0
+                [2, 0],  # head 1
+            ]
+        ],
+        dtype=torch.int32,
+        device="cuda",
+    )  # [BATCH, N_HEADS, NNZ_V]
+    slash_indexes = torch.tensor(
+        [
+            [
+                [2, 0],  # head 0
+                [1, 3],  # head 1
+            ]
+        ],
+        dtype=torch.int32,
+        device="cuda",
+    )  # [BATCH, N_HEADS, NNZ_S]
+    vertical_indices_count = torch.tensor([2, 1], dtype=torch.int32, device="cuda")
+    slash_indices_count = torch.tensor([1, 2], dtype=torch.int32, device="cuda")
+    context_size = 4
+    block_size_M = 2
+    block_size_N = 2
+
+    # Call your CUDA kernel wrapper
+    block_count, block_offset, column_count, column_index = (
+        convert_vertical_slash_indexes_mergehead(
+            q_seqlens,
+            kv_seqlens,
+            vertical_indexes,
+            slash_indexes,
+            vertical_indices_count,
+            slash_indices_count,
+            context_size,
+            block_size_M,
+            block_size_N,
+            causal=causal,
+        )
+    )
+
+    # Manually create expected outputs for this input
+    # For demonstration, assume:
+    # - batch=1, head=2, num_rows=2, nnz_v=2, nnz_s=2
+    # Fill these expected tensors according to your kernel's behavior
+
+    expected_column_index = torch.tensor(
+        [[[[1, 0], [1, 3]], [[-1079459945, -1077788999], [-1080050043, -1104625879]]]],
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    if not causal:
+        # If non-causal mode output is different, update these values
+        expected_column_index = torch.tensor(
+            [[[[1, 0], [1, 3]], [[2, -1077788999], [2, -1104625879]]]],
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+    # Assert that outputs match expectations
+    assert torch.equal(column_index, expected_column_index)
+
+
+# skip cause use fa2 for test
+# @pytest.mark.parametrize("seq_lens", [[(1024, 1328)],
+#                                     [(1024, 1328), (1, 2048)],
+#                                     [(1025, 1328), (2, 2048)],
+#                                     [(1025, 2049), (2, 1281)],
+#                                     ])
+# @pytest.mark.parametrize("head_size", [128])
+# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+# @torch.inference_mode()
+# def test_sparse_attention_varlen(
+#         seq_lens,
+#         head_size,
+#         dtype,
+# ) -> None:
+#     torch.set_default_device("cuda")
+#     torch.cuda.manual_seed_all(0)
+#     block_size_M = 64
+#     block_size_N = 64
+#     num_seqs = len(seq_lens)
+#     query_lens = [x[0] for x in seq_lens]
+#     kv_lens = [x[1] for x in seq_lens]
+#     num_heads = 1
+#     query = torch.randn(sum(query_lens),
+#                         num_heads,
+#                         head_size,
+#                         dtype=dtype)
+#     key = torch.randn(sum(kv_lens),
+#                     num_heads,
+#                     head_size,
+#                     dtype=dtype)
+#     value = torch.randn_like(key)
+#     cu_query_lens = torch.tensor([0] + query_lens,
+#                                 dtype=torch.int32).cumsum(dim=0,
+#                                                         dtype=torch.int32)
+#     cu_kv_lens = torch.tensor([0] + kv_lens,
+#                                 dtype=torch.int32).cumsum(dim=0,
+#                                                         dtype=torch.int32)
+#     max_query_len = max(query_lens)
+#     max_kv_len = max(kv_lens)
+
+#     NUM_ROWS = (max_query_len + block_size_M - 1) // block_size_M
+#     NNZ_S = 20
+#     NNZ_V = 2048
+#     batch_size = len(query_lens)
+
+#     block_counts = []
+#     column_counts = []
+#     block_offsets = []
+#     column_indices = []
+#     for b in range(batch_size):
+#         block_counts.append(torch.tensor([NNZ_S] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS))
+#         columns = kv_lens[b] - NNZ_S * block_size_N
+#         column_counts.append(torch.tensor([columns] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS))
+#         block_offsets.append(torch.tensor([[i * block_size_N for i in range(NNZ_S)]] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS, NNZ_S))
+#         column_indices.append(torch.tensor([[NNZ_S * block_size_N + i for i in range(NNZ_V)]] * NUM_ROWS * num_heads, dtype=torch.int32).reshape(num_heads, NUM_ROWS, NNZ_V))
+#     block_count = torch.concat(block_counts).reshape(batch_size, num_heads, NUM_ROWS)
+#     column_count = torch.concat(column_counts).reshape(batch_size, num_heads, NUM_ROWS)
+#     block_offset = torch.concat(block_offsets).reshape(batch_size, num_heads, NUM_ROWS, NNZ_S)
+#     column_index = torch.concat(column_indices).reshape(batch_size, num_heads, NUM_ROWS, NNZ_V)
+#     out, lse = sparse_attn_varlen_func(
+#         query,
+#         key,
+#         value,
+#         block_count,
+#         block_offset,
+#         column_count,
+#         column_index,
+#         cu_seqlens_q=cu_query_lens,
+#         cu_seqlens_k=cu_kv_lens,
+#         max_seqlen_q=max_query_len,
+#         max_seqlen_k=max_kv_len,
+#         return_softmax_lse=True,
+#     )
+
+#     max_num_blocks_per_seq = (max_kv_len + 2048 - 1) // 2048
+#     block_tables = torch.randint(0,
+#                                  2048,
+#                                  (len(query_lens), max_num_blocks_per_seq),
+#                                  dtype=torch.int32)
+#     scale = head_size**-0.5
+
+#     ref_out, ref_lse, _ = ref_paged_attn(
+#         query,
+#         key,
+#         value,
+#         query_lens=query_lens,
+#         kv_lens=kv_lens,
+#         block_tables=block_tables,
+#         scale=scale
+#     )
+
+#     torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2), \
+#         f"{torch.max(torch.abs(out - ref_out))}"
+#     torch.testing.assert_close(lse, ref_lse, atol=2e-2, rtol=1e-2), \
+#         f"{torch.max(torch.abs(lse - ref_lse))}"
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_flashmla.py

@@ -0,0 +1,663 @@
+import math
+import random
+import sys
+from typing import Optional, Tuple
+
+import pytest
+import torch
+import triton
+from sgl_kernel.flash_mla import (
+    flash_mla_sparse_fwd,
+    flash_mla_with_kvcache,
+    get_mla_metadata,
+)
+
+# ================ prefill usage ================ #
+S_Q_PREFILL = [1, 62]
+KV_TOPK_PREFILL = [
+    # Regular shapes
+    (128, 128),
+    (256, 256),
+    (512, 512),
+    # Irregular shapes
+    (592, 128),
+    (1840, 256),
+    (1592, 384),
+    (1521, 512),
+    # Irregular shapes with OOB TopK
+    (95, 128),
+    (153, 256),
+    (114, 384),
+]
+
+# ================= decode usage ================= #
+B_DECODE = [1, 2, 6, 64]
+S_Q_DECODE = [1, 2, 4]
+S_K_DECODE = [20, 140, 4096]
+IS_VARLEN = [False, True]
+CAUSAL_TOPK = [(True, None), (False, None), (False, 128), (False, 2048)]
+DTYPE = [torch.float16, torch.bfloat16]
+
+
+def is_sm90_supported(device=None) -> bool:
+    return (torch.cuda.get_device_capability(device)[0] == 9) and (
+        torch.version.cuda >= "12.3"
+    )
+
+
+def quantize_k_cache(
+    input_k_cache: torch.Tensor,  # (num_blocks, block_size, h_k, d)
+    dv: int,
+    tile_size: int = 128,
+) -> torch.Tensor:
+    """
+    Quantize the k-cache
+    Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size()
+    For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md
+    """
+    assert dv % tile_size == 0
+    num_tiles = dv // tile_size
+    num_blocks, block_size, h_k, d = input_k_cache.shape
+    assert h_k == 1
+    input_k_cache = input_k_cache.squeeze(2)  # [num_blocks, block_size, d]
+    input_elem_size = input_k_cache.element_size()
+
+    result = torch.empty(
+        (num_blocks, block_size, dv + num_tiles * 4 + input_elem_size * (d - dv)),
+        dtype=torch.float8_e4m3fn,
+        device=input_k_cache.device,
+    )
+    result_k_nope_part = result[..., :dv]
+    result_k_scale_factor = result[..., dv : dv + num_tiles * 4].view(torch.float32)
+    result_k_rope_part = result[..., dv + num_tiles * 4 :].view(input_k_cache.dtype)
+    result_k_rope_part[:] = input_k_cache[..., dv:]
+
+    for tile_idx in range(0, num_tiles):
+        cur_scale_factors_inv = (
+            torch.abs(
+                input_k_cache[..., tile_idx * tile_size : (tile_idx + 1) * tile_size]
+            )
+            .max(dim=-1)
+            .values
+            / 448.0
+        )  # [num_blocks, block_size]
+        result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv
+
+        cur_scale_factors_inv.unsqueeze_(-1)  # [num_blocks, block_size, 1]
+        cur_quantized_nope = (
+            input_k_cache[
+                ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
+            ].float()
+            / cur_scale_factors_inv.float()
+        ).to(torch.float8_e4m3fn)
+        result_k_nope_part[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
+            cur_quantized_nope
+        )
+
+    result = result.view(num_blocks, block_size, 1, -1)
+    return result
+
+
+def dequantize_k_cache(
+    quant_k_cache: torch.Tensor,  # (num_blocks, block_size, 1, bytes_per_token)
+    dv: int = 512,
+    tile_size: int = 128,
+    d: int = 576,
+) -> torch.Tensor:
+    """
+    De-quantize the k-cache
+    """
+    assert dv % tile_size == 0
+    num_tiles = dv // tile_size
+    num_blocks, block_size, h_k, _ = quant_k_cache.shape
+    assert h_k == 1
+    result = torch.empty(
+        (num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
+    )
+
+    quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
+
+    input_nope = quant_k_cache[..., :dv]
+    input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
+    input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
+    result[..., dv:] = input_rope
+
+    for tile_idx in range(0, num_tiles):
+        cur_nope = input_nope[
+            ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
+        ].to(torch.float32)
+        cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
+        result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
+            cur_nope * cur_scales
+        )
+
+    result = result.view(num_blocks, block_size, 1, d)
+    return result
+
+
+def cdiv(x: int, y: int):
+    return (x + y - 1) // y
+
+
+def get_window_size(causal, window):
+    if window > 0:
+        window_size = (window - 1, 0) if causal else (window - 1, window - 1)
+    else:
+        window_size = (-1, -1)
+    return window_size
+
+
+def get_attn_bias(s_q, s_k, causal, window):
+    attn_bias = torch.zeros(s_q, s_k, dtype=torch.float32, device="cuda")
+    if causal:
+        temp_mask = torch.ones(s_q, s_k, dtype=torch.bool, device="cuda").tril(
+            diagonal=s_k - s_q
+        )
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+    if window > 0:
+        temp_mask = torch.ones(s_q, s_k, dtype=torch.bool, device="cuda").tril(
+            diagonal=s_k - s_q - window
+        )
+        attn_bias.masked_fill_(temp_mask, float("-inf"))
+        temp_mask = torch.ones(s_q, s_k, dtype=torch.bool, device="cuda").tril(
+            diagonal=s_k - s_q + window - 1
+        )
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+    return attn_bias
+
+
+def sdpa(query, key, value, attn_bias, softmax_scale=None):
+    query = query.float().transpose(-3, -2)
+    key = key.float().transpose(-3, -2)
+    value = value.float().transpose(-3, -2)
+    key = key.repeat_interleave(h // h_k, dim=-3)
+    value = value.repeat_interleave(h // h_k, dim=-3)
+    if softmax_scale is None:
+        softmax_scale = query.shape[-1] ** (-0.5)
+    attn_weight = (query @ key.transpose(-2, -1)) * softmax_scale
+    attn_weight += attn_bias
+    lse = attn_weight.logsumexp(dim=-1)
+    attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
+    return attn_weight.to(query.dtype) @ value, lse
+
+
+def sdpa_checkpoint(*args, **kwargs):
+    return checkpoint(sdpa, *args, use_reentrant=False, **kwargs)
+
+
+def reference_torch_prefill(
+    s_q, s_kv, topk, indices, q, kv, sm_scale: float
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor:
+        return torch.logsumexp(a * math.log(2), dim=dim) * math.log2(math.e)
+
+    indices = indices[0, :, 0, :]  # [s_q, topk]
+    invalid_indices_mask = (indices < 0) | (indices >= s_kv)
+    qs = q[0, :, :, :].float()  # [s_q, h_q, d_qk]
+    kvs = kv[0, :, 0, :].float()  # [s_kv, d_qk]
+
+    kvs = torch.index_select(
+        kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten()
+    ).view(
+        s_q, topk, 576
+    )  # [s_q, topk, d_qk]
+    attn_score = qs @ kvs.transpose(1, 2)  # [s_q, h_q, topk]
+    attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf"))
+    attn_score *= sm_scale * math.log2(math.e)
+    max_logits = torch.max(attn_score, dim=-1)[0]  # [s_q, h_q]
+    lse = log2sumexp2(attn_score, dim=-1)  # [s_q, h_q]
+    attn_score = torch.exp2(attn_score - lse.unsqueeze(-1))  # [s_q, h_q, topk]
+    result = attn_score @ kvs[:, :, :512]
+    return (max_logits, lse, result)
+
+
+def reference_torch_decode(
+    cache_seqlens: torch.Tensor,  # [batch_size]
+    block_table: torch.Tensor,  # [batch_size, ?]
+    q: torch.Tensor,  # [batch_size, s_q, h_q, d]
+    blocked_k: torch.Tensor,  # [?, block_size, h_kv, d]
+    dv: int,
+    is_causal: bool,
+    indices: Optional[torch.Tensor] = None,  # [batch_size, s_q, topk]
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    A reference implementation in PyTorch
+    """
+
+    def get_topk_attn_mask(s_q: int, s_k: int, indices: torch.Tensor):
+        mask = torch.zeros(s_q, s_k, dtype=torch.bool, device="cuda")
+        for i in range(s_q):
+            cur_indices = indices[i]
+            valid_indices = cur_indices[cur_indices != -1]
+            mask[i, valid_indices] = True
+        return mask
+
+    def scaled_dot_product_attention(
+        batch_idx: int,
+        query: torch.Tensor,  # [h_q, s_q, d]
+        kv: torch.Tensor,  # [h_kv, s_k, d]
+        dv: int,
+        is_causal,
+        indices: Optional[torch.Tensor],  # [s_q, topk]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        h_q = query.size(0)
+        h_kv = kv.size(0)
+        s_q = query.shape[-2]
+        s_k = kv.shape[-2]
+        query = query.float()
+        kv = kv.float()
+        if h_kv != 1:
+            kv = kv.repeat_interleave(h_q // h_kv, dim=0)
+        kv[kv != kv] = 0.0
+        attn_weight = query @ kv.transpose(-2, -1)  # [h_q, s_q, s_k]
+        if (is_causal and query.size(1) > 1) or indices is not None:
+            mask = torch.ones(s_q, s_k, dtype=torch.bool, device="cuda")
+            if is_causal:
+                assert indices is None
+                mask = mask.tril(diagonal=s_k - s_q)
+            if indices is not None:
+                mask &= get_topk_attn_mask(s_q, s_k, indices)
+            attn_bias = torch.zeros(s_q, s_k, dtype=torch.float, device="cuda")
+            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
+            attn_weight += attn_bias.to(q.dtype)
+        attn_weight /= math.sqrt(query.size(-1))
+        lse = attn_weight.logsumexp(dim=-1)  # [h_q, s_q]
+        attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
+        output = attn_weight @ kv[..., :dv]  # [h_q, s_q, dv]
+        # Correct for q tokens which has no attendable k
+        lonely_q_mask = lse == float("-inf")
+        output[lonely_q_mask.unsqueeze(-1).broadcast_to(h_q, s_q, dv)] = 0.0
+        lse[lonely_q_mask] = float("+inf")
+
+        return output, lse
+
+    b, s_q, h_q, d = q.size()
+    block_size = blocked_k.size(1)
+    h_kv = blocked_k.size(2)
+    cache_seqlens_cpu = cache_seqlens.cpu()
+    out_ref = torch.empty(b, s_q, h_q, dv, dtype=torch.float32, device="cuda")
+    lse_ref = torch.empty(b, h_q, s_q, dtype=torch.float32, device="cuda")
+    for i in range(b):
+        cur_len = cache_seqlens_cpu[i].item()
+        cur_num_blocks = cdiv(cur_len, block_size)
+        cur_block_indices = block_table[i][0:cur_num_blocks]
+        cur_kv = blocked_k[cur_block_indices].view(-1, h_kv, d)[:cur_len, ...]
+        cur_out, cur_lse = scaled_dot_product_attention(
+            i,
+            q[i].transpose(0, 1),
+            cur_kv.transpose(0, 1),
+            dv,
+            is_causal,
+            indices[i] if indices is not None else None,
+        )
+        out_ref[i] = cur_out.transpose(0, 1)
+        lse_ref[i] = cur_lse
+    out_ref = out_ref.to(torch.bfloat16)
+    return out_ref, lse_ref
+
+
+@pytest.mark.parametrize("s_q", S_Q_PREFILL)
+@pytest.mark.parametrize("kv_topk", KV_TOPK_PREFILL)
+@torch.inference_mode()
+def test_flashmla_prefill(
+    s_q: int,
+    kv_topk: Tuple[int, int],
+):
+
+    torch.cuda.empty_cache()
+
+    q = torch.randn((1, s_q, 128, 576), dtype=torch.bfloat16, device="cuda") / 10
+    kv = torch.randn((1, kv_topk[0], 1, 576), dtype=torch.bfloat16, device="cuda") / 10
+
+    q.clamp_(-10, 10)
+    kv.clamp_(-10, 10)
+
+    indices = torch.full(
+        (1, s_q, 1, kv_topk[1]), kv_topk[0], dtype=torch.int32, device="cuda"
+    )
+    for s in range(s_q):
+        # NOTE We use the following method to generate indices so that most indices lies within [s_kv-20000, s_kv), which is more realistic for sparse attention
+        near_mask = (
+            torch.randint(0, 32, (min(kv_topk[1], kv_topk[0]),), device="cuda") < 31
+        )
+        cur_indices = torch.randperm(kv_topk[0], device="cuda")[: kv_topk[1]]
+        cur_indices[near_mask] = torch.randint(
+            max(0, kv_topk[0] - 20000),
+            kv_topk[0] - 1,
+            (near_mask.sum().item(),),
+            device="cuda",
+        )
+        if len(cur_indices) < kv_topk[1]:
+            cur_indices = torch.cat(
+                [
+                    cur_indices,
+                    torch.full(
+                        (kv_topk[1] - len(cur_indices),), 2147480000, device="cuda"
+                    ),
+                ]
+            )
+        cur_indices = cur_indices[torch.randperm(kv_topk[1], device="cuda")]
+        indices[0, s, 0] = cur_indices
+    indices = indices.to(q.device)
+
+    sm_scale = 1 / math.sqrt(576)
+    torch.cuda.synchronize()
+
+    ans_out, ans_max_logits, ans_lse = flash_mla_sparse_fwd(
+        q.squeeze(0), kv.squeeze(0), indices.squeeze(0), sm_scale=sm_scale
+    )
+
+    ans_out, ans_max_logits, ans_lse = (
+        ans_out.float(),
+        ans_max_logits.float(),
+        ans_lse.float(),
+    )
+
+    torch.cuda.synchronize()
+    ref_max_logits, ref_lse, ref_out = reference_torch_prefill(
+        s_q, kv_topk[0], kv_topk[1], indices, q, kv, sm_scale
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(ans_out, ref_out, atol=8e-4, rtol=2.01 / 128)
+    torch.testing.assert_close(
+        ans_max_logits,
+        ref_max_logits,
+        atol=1e-6,
+        rtol=2.01 / 65536,
+    )
+    torch.testing.assert_close(ans_lse, ref_lse, atol=1e-6, rtol=2.01 / 65536)
+
+
+@pytest.mark.skipif(not is_sm90_supported(), reason="SM90 required for FP8 support")
+@pytest.mark.parametrize("b", B_DECODE)
+@pytest.mark.parametrize("s_q", S_Q_DECODE)
+@pytest.mark.parametrize("s_k", S_K_DECODE)
+@pytest.mark.parametrize("is_varlen", IS_VARLEN)
+@pytest.mark.parametrize("causal_topk", CAUSAL_TOPK)
+@pytest.mark.parametrize("dtype", DTYPE)
+@torch.inference_mode()
+def test_flash_mla_decode(
+    b: int,
+    s_q: int,
+    s_k: int,
+    is_varlen: bool,
+    causal_topk: Tuple[bool, Optional[int]],
+    dtype: torch.dtype,
+):
+    d = 576
+    dv = 512
+    block_size = 64
+    h_q = 128
+    h_kv = 1
+    is_causal = causal_topk[0]
+    topk = causal_topk[1]
+
+    # Generating test data
+    torch.cuda.synchronize()
+
+    cache_seqlens_cpu = torch.full((b,), s_k, dtype=torch.int32, device="cpu")
+    if is_varlen:
+        for i in range(b):
+            cache_seqlens_cpu[i] = max(random.normalvariate(s_k, s_k / 2), s_q)
+
+    max_seqlen = cache_seqlens_cpu.max().item()
+    max_seqlen_pad = cdiv(max_seqlen, 256) * 256
+    cache_seqlens = cache_seqlens_cpu.cuda()
+
+    q = torch.randn(b, s_q, 128, d, dtype=torch.bfloat16, device="cuda")
+    q.clamp_(min=-1.0, max=1.0)
+
+    block_table = torch.arange(
+        b * max_seqlen_pad // block_size, dtype=torch.int32, device="cuda"
+    ).view(b, max_seqlen_pad // block_size)
+    block_table = block_table.view(-1)[torch.randperm(block_table.numel())].view(b, -1)
+    blocked_k = (
+        torch.randn(
+            block_table.numel(),
+            block_size,
+            h_kv,
+            d,
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+        / 10
+    )
+    blocked_k.clamp_(min=-1.0, max=1.0)
+
+    if topk is None:
+        for i in range(b):
+            cur_len = cache_seqlens_cpu[i].item()
+            cur_num_blocks = cdiv(cur_len, block_size)
+            blocked_k[block_table[i][cur_num_blocks:]] = float("nan")
+            if cur_len % block_size != 0:
+                blocked_k[block_table[i][cur_num_blocks - 1]][
+                    cur_len % block_size :
+                ] = float("nan")
+            block_table[i][cur_num_blocks:] = 2147480000
+        abs_indices = None
+        indices_in_kvcache = None
+    else:
+        block_table_cpu = block_table.cpu()
+        abs_indices = torch.empty(b, s_q, topk, dtype=torch.int32, device="cpu")
+        indices_in_kvcache = torch.empty(b, s_q, topk, dtype=torch.int32, device="cpu")
+        for i in range(b):
+            # Generate indices
+            for j in range(s_q):
+                cur_abs_indices = torch.randperm(
+                    int(cache_seqlens_cpu[i].item()), device="cpu"
+                )[:topk]
+                cur_blocked_indices = block_table_cpu[
+                    i, cur_abs_indices // block_size
+                ] * block_size + (cur_abs_indices % block_size)
+                if len(cur_abs_indices) < topk:
+                    pad_len = topk - len(cur_abs_indices)
+                    cur_abs_indices = torch.cat(
+                        [cur_abs_indices, torch.full((pad_len,), -1, device="cpu")]
+                    )
+                    cur_blocked_indices = torch.cat(
+                        [cur_blocked_indices, torch.full((pad_len,), -1, device="cpu")]
+                    )
+
+                # Mask KV
+                perm = torch.randperm(topk, device="cpu")
+                cur_abs_indices = cur_abs_indices[perm]
+                cur_blocked_indices = cur_blocked_indices[perm]
+
+                abs_indices[i, j, :] = cur_abs_indices
+                indices_in_kvcache[i, j, :] = cur_blocked_indices
+
+        # Mask nonused KV as NaN
+        all_indices = indices_in_kvcache.flatten().tolist()
+        all_indices = list(set(all_indices))
+        if -1 in all_indices:
+            all_indices.remove(-1)
+        all_indices = torch.tensor(all_indices, dtype=torch.int32, device="cpu")
+
+        blocked_k = blocked_k.view(-1, h_kv, d)
+        nonused_indices_mask = torch.ones(
+            blocked_k.size(0) * blocked_k.size(1), dtype=torch.bool, device="cpu"
+        )
+        nonused_indices_mask[all_indices] = False
+        blocked_k[nonused_indices_mask, :, :] = float("nan")
+        blocked_k = blocked_k.view(-1, block_size, h_kv, d)
+
+        abs_indices = abs_indices.to(q.device)
+        indices_in_kvcache = indices_in_kvcache.to(q.device)
+
+    is_fp8 = topk is not None
+    if is_fp8:
+        # The quantization error may be too large to be distinguished from wrong kernels
+        # So we quantize and de-quantize kv-cache here to mitigate quantization error
+        blocked_k_quantized = quantize_k_cache(blocked_k, dv, 128)
+        blocked_k_dequantized = dequantize_k_cache(blocked_k_quantized)
+        blocked_k = blocked_k_dequantized
+
+    # Get schedule metadata
+    torch.cuda.synchronize()
+    tile_scheduler_metadata, num_splits = get_mla_metadata(
+        cache_seqlens, s_q * h_q // h_kv, h_kv, h_q, is_fp8, topk
+    )
+    torch.cuda.synchronize()
+
+    out_ans, lse_ans = flash_mla_with_kvcache(
+        q,
+        blocked_k if not is_fp8 else blocked_k_quantized,  # type: ignore
+        block_table,
+        cache_seqlens,
+        dv,
+        tile_scheduler_metadata,
+        num_splits,
+        causal=is_causal,
+        is_fp8_kvcache=is_fp8,
+        indices=indices_in_kvcache,
+    )
+
+    out_ref, lse_ref = reference_torch_decode(
+        cache_seqlens, block_table, q, blocked_k, dv, is_causal, abs_indices
+    )
+    torch.testing.assert_close(out_ans, out_ref, atol=8e-4, rtol=2.01 / 128)
+    torch.testing.assert_close(lse_ans, lse_ref, atol=1e-6, rtol=8.01 / 65536)
+
+
+@pytest.mark.skipif(not is_sm90_supported(), reason="SM90 required for FP8 support")
+@pytest.mark.parametrize("b", [128])
+@pytest.mark.parametrize("s_q", [1, 2])
+@pytest.mark.parametrize("mean_sk", [4096, 8192, 16384])
+@pytest.mark.parametrize("h_q", [16, 32, 64, 128])
+@pytest.mark.parametrize("h_kv", [1])
+@pytest.mark.parametrize("d", [576])
+@pytest.mark.parametrize("dv", [512])
+@pytest.mark.parametrize("block_size", [64])
+@pytest.mark.parametrize("causal", [True])
+@pytest.mark.parametrize("varlen", [False, True])
+@pytest.mark.parametrize("torch_dtype", [torch.float8_e4m3fn])
+@torch.inference_mode()
+def test_flash_mla_fp8(
+    b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype
+):
+    device = torch.device("cuda:0")
+    init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype
+    torch.set_default_dtype(init_dtype)
+    torch.set_default_device(device)
+    torch.cuda.set_device(device)
+    torch.manual_seed(0)
+    random.seed(0)
+
+    use_fp8 = torch_dtype == torch.float8_e4m3fn
+    cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
+    if varlen:
+        for i in range(b):
+            cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2), s_q)
+    total_seqlens = cache_seqlens.sum().item()
+    max_seqlen = cache_seqlens.max().item()
+    max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256
+
+    q = torch.randn(b, s_q, h_q, d)
+    block_table = torch.arange(
+        b * max_seqlen_pad // block_size, dtype=torch.int32
+    ).view(b, max_seqlen_pad // block_size)
+    blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d)
+    for i in range(b):
+        blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item() :] = (
+            float("nan")
+        )
+    blocked_v = blocked_k[..., :dv]
+
+    tile_scheduler_metadata, num_splits = get_mla_metadata(
+        cache_seqlens, s_q * h_q // h_kv, h_kv, is_fp8_kvcache=use_fp8
+    )
+
+    init_dtype = q.dtype
+    if use_fp8:
+        fp8_dtype = torch.float8_e4m3fn
+        descale_q = torch.ones((1), dtype=torch.float32)
+        descale_k = torch.ones((1), dtype=torch.float32)
+
+        q = q.to(fp8_dtype)
+        blocked_k = blocked_k.to(fp8_dtype)
+        blocked_v = blocked_v.to(fp8_dtype)
+    else:
+        descale_q = None
+        descale_k = None
+
+    def flash_mla():
+        return flash_mla_with_kvcache(
+            q,
+            blocked_k,
+            block_table,
+            cache_seqlens,
+            dv,
+            tile_scheduler_metadata,
+            num_splits,
+            causal=causal,
+            descale_q=descale_q,
+            descale_k=descale_k,
+        )
+
+    def scaled_dot_product_attention(query, key, value, is_causal=False):
+        query = query.float()
+        key = key.float()
+        value = value.float()
+        key = key.repeat_interleave(h_q // h_kv, dim=0)
+        value = value.repeat_interleave(h_q // h_kv, dim=0)
+        attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))
+        if is_causal:
+            s_q = query.shape[-2]
+            s_k = key.shape[-2]
+            attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype)
+            temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q)
+            attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+            attn_bias.to(query.dtype)
+            attn_weight += attn_bias
+        lse = attn_weight.logsumexp(dim=-1)
+        attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
+        return attn_weight @ value, lse
+
+    def ref_mla():
+        q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q
+        blocked_k_ = (
+            (blocked_k.to(torch.float) * descale_k).to(init_dtype)
+            if use_fp8
+            else blocked_k
+        )
+        blocked_v_ = (
+            (blocked_v.to(torch.float) * descale_k).to(init_dtype)
+            if use_fp8
+            else blocked_v
+        )
+        out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)
+        lse = torch.empty(b, h_q, s_q, dtype=torch.float32)
+        for i in range(b):
+            begin = i * max_seqlen_pad
+            end = begin + cache_seqlens[i]
+            out_i, lse_i = scaled_dot_product_attention(
+                q_[i].transpose(0, 1),
+                blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1),
+                blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1),
+                is_causal=causal,
+            )
+            out[i] = out_i.transpose(0, 1)
+            lse[i] = lse_i
+        return out, lse
+
+    def cal_diff(
+        x: torch.Tensor, y: torch.Tensor, name: str, use_fp8: bool = False
+    ) -> None:
+        x, y = x.double(), y.double()
+        cos_diff = 1 - 2 * (x * y).sum().item() / max(
+            (x * x + y * y).sum().item(), 1e-12
+        )
+        if use_fp8:
+            assert cos_diff < 1e-4
+        else:
+            assert cos_diff < 1e-5
+
+    out_flash, lse_flash = flash_mla()
+    out_torch, lse_torch = ref_mla()
+    cal_diff(out_flash, out_torch, "out", use_fp8)
+    cal_diff(lse_flash, lse_torch, "lse")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_fp8_blockwise_gemm.py

@@ -0,0 +1,94 @@
+import os
+import random
+import sys
+from typing import Optional, Type
+
+import pytest
+import torch
+from sgl_kernel import fp8_blockwise_scaled_mm
+
+
+def cdiv(a: int, b: int) -> int:
+    return -(a // -b)
+
+
+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)))
+
+
+def baseline_scaled_mm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale_a: torch.Tensor,
+    scale_b: torch.Tensor,
+    out_dtype: Type[torch.dtype],
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # We treat N-dimensional group scaling as extended numpy-style broadcasting
+    # in numpy simply stretches dimensions with an extent of 1 to match the
+    # the target shape by repeating the data along that dimension (broadcasting)
+    # , we extend these semantics to say if the extent of a dimension in the
+    # source shape is not 1 and does not match the target shape we repeat each
+    # element along that dimension src_shape[dim] // target_shape[dim] times
+    # example if we have:
+    #       a = [[1, 2], and target_shape = (2, 4)
+    #            [3, 4]]
+    # then we would expand a to:
+    #       a = [[1, 1, 2, 2],
+    #            [3, 3, 4, 4]]
+    # NOTE this function this function does not explicitly broadcast dimensions
+    # with an extent of 1, since this can be done implicitly by pytorch
+    def group_broadcast(t, shape):
+        for i, s in enumerate(shape):
+            if t.shape[i] != s and t.shape[i] != 1:
+                assert s % t.shape[i] == 0
+                t = (
+                    t.unsqueeze(i + 1)
+                    .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
+                    .flatten(i, i + 1)
+                )
+        return t
+
+    scale_a = group_broadcast(scale_a, a.shape)
+    scale_b = group_broadcast(scale_b, b.shape)
+    output = torch.mm(
+        (scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32))
+    ).to(out_dtype)
+    if bias is not None:
+        output = output + bias
+    return output
+
+
+def _test_accuracy_once(M, N, K, out_dtype, device):
+    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=device) - 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=device) - 0.5) * 2 * fp8_max
+    b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn).t()
+    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=device, dtype=torch.float32) * 0.001
+    scale_b = torch.randn(scale_b_shape, device=device, dtype=torch.float32) * 0.001
+    scale_a = scale_a.t().contiguous().t()
+    scale_b = scale_b.t().contiguous().t()
+    o = baseline_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype)
+    o1 = fp8_blockwise_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype)
+    rtol = 0.02
+    atol = 1
+    torch.testing.assert_close(o, o1, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("M", [1, 3, 5, 127, 128, 512, 1024, 4096])
+@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 14080])
+@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 14080, 16384])
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
+def test_accuracy(M, N, K, out_dtype):
+    _test_accuracy_once(M, N, K, out_dtype, "cuda")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_fp8_blockwise_moe.py

@@ -0,0 +1,222 @@
+import random
+import sys
+from typing import Tuple
+
+import pytest
+import torch
+from sgl_kernel import fp8_blockwise_scaled_grouped_mm
+
+
+def cdiv(a: int, b: int) -> int:
+    return -(a // -b)
+
+
+def scale_shape(shape, group_shape):
+    return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))
+
+
+def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(
+        dtype=torch.float8_e4m3fn
+    )
+
+
+# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+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 baseline_scaled_mm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale_a: torch.Tensor,
+    scale_b: torch.Tensor,
+    out_dtype: type[torch.dtype],
+) -> torch.Tensor:
+
+    def group_broadcast(t, shape):
+        for i, s in enumerate(shape):
+            if t.shape[i] != s and t.shape[i] != 1:
+                assert s % t.shape[i] == 0
+                t = (
+                    t.unsqueeze(i + 1)
+                    .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
+                    .flatten(i, i + 1)
+                )
+        return t
+
+    scale_a = group_broadcast(scale_a, a.shape)
+    scale_b = group_broadcast(scale_b, b.shape)
+
+    return torch.mm(
+        (scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32))
+    ).to(out_dtype)
+
+
+def is_blackwell_supported(device=None) -> bool:
+    return (torch.cuda.get_device_capability(device)[0] in [10, 12]) and (
+        torch.version.cuda >= "12.8"
+    )
+
+
+def is_sm90_supported(device=None) -> bool:
+    return (torch.cuda.get_device_capability(device)[0] == 9) and (
+        torch.version.cuda >= "12.3"
+    )
+
+
+@pytest.mark.skipif(
+    not (is_blackwell_supported() or is_sm90_supported()),
+    reason="fp8_blockwise_scaled_grouped_mm at sgl-kernel is only supported on sm100 or sm90",
+)
+@pytest.mark.parametrize("num_experts", [8, 16, 32, 64, 128])
+@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
+def test_fp8_blockwise_scaled_grouped_mm(num_experts, out_dtype):
+    device = "cuda"
+    alignment = 128
+    n_g = random.randint(1, 64) * 128
+    k_g = random.randint(1, 64) * 128
+
+    expert_offsets = torch.zeros((num_experts + 1), device=device, dtype=torch.int32)
+    problem_sizes = torch.zeros((num_experts, 3), device=device, dtype=torch.int32)
+    layout_sfa = torch.zeros((num_experts, 5), device=device, dtype=torch.int32)
+    layout_sfb = torch.zeros((num_experts, 5), device=device, dtype=torch.int32)
+
+    a_tensors = []
+    b_tensors = []
+    a_scales_tensors = []
+    b_scales_tensors = []
+    baseline_tensors = []
+
+    for g in range(num_experts):
+        m_g = random.randint(1, 256)
+        expert_offsets[g + 1] = expert_offsets[g] + m_g
+        problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
+
+        a = torch.randn((m_g, k_g), device=device, dtype=out_dtype)  # (M, K):(K, 1)
+        b = torch.randn((n_g, k_g), device=device, dtype=out_dtype).t()  # (K, N):(1, K)
+
+        a_g, a_scale = per_token_cast_to_fp8(
+            a
+        )  # ag -- (M, K):(K, 1), a_scale() -- (M, k):(k, 1)
+        b_g, b_scale = per_block_cast_to_fp8(
+            b
+        )  # bg -- (K, N):(N, 1), b_scale() -- (k, n):(n, 1)
+        a_tensors.append(a_g)
+        b_tensors.append(b_g)
+        a_scales_tensors.append(a_scale)
+        b_scales_tensors.append(b_scale)
+
+        baseline = torch.mm(a, b)
+        baseline_tensors.append(baseline)
+    a_stack = torch.empty(
+        (expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    b_stack = torch.empty(
+        (num_experts, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
+    )
+    a_scale_stack = torch.empty(
+        (expert_offsets[-1], (k_g // 128)), device=device, dtype=torch.float32
+    )
+    b_scale_stack = torch.empty(
+        (num_experts, n_g // 128, k_g // 128), device=device, dtype=torch.float32
+    )
+
+    for g in range(num_experts):
+        # Matrix A is Row-Major.
+        a_stack[expert_offsets[g] : expert_offsets[g + 1], :] = a_tensors[
+            g
+        ]  # a_stack[expert_offsets[g] : expert_offsets[g + 1], :] -- (M, K):(K, 1)
+        b_stack[g] = b_tensors[g].t()  # b_stack[g] -- (N, K):(K, 1)
+
+        # We need K-Major scale factor
+        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[g] -- (k, n):(n, 1), we need transpose & contiguous later
+    b_stack = b_stack.transpose(1, 2)  # Transpose Matrix B to Column-Major.
+    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_experts,), a_stack.stride(0), device=device, dtype=torch.int64
+    )
+    c_strides = torch.full(
+        (num_experts,), 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_experts,), device=device, dtype=torch.int64)
+    b_ptrs = torch.empty((num_experts,), device=device, dtype=torch.int64)
+    out_ptrs = torch.empty((num_experts,), device=device, dtype=torch.int64)
+    a_scales_ptrs = torch.empty((num_experts,), device=device, dtype=torch.int64)
+    b_scales_ptrs = torch.empty((num_experts,), device=device, dtype=torch.int64)
+
+    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,
+    )
+
+    for g in range(num_experts):
+        baseline = baseline_tensors[g]
+        actual = c_out[expert_offsets[g] : expert_offsets[g + 1]]
+        diff = calc_diff(actual, baseline)
+        assert diff < 0.001
+        print(
+            f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, out_dtype={out_dtype}, diff={diff:.5f}: OK"
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_fp8_gemm.py

@@ -0,0 +1,51 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import fp8_scaled_mm
+
+
+def torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias):
+    o = torch.matmul(a.to(torch.float32), b.to(torch.float32))
+    o = o.to(torch.float32)
+    temp1 = o * scale_a.view(-1, 1)
+    temp2 = temp1 * scale_b.view(1, -1)
+    final = temp2.to(out_dtype)
+    if bias is not None:
+        final = final + bias.view(1, -1)
+    return final
+
+
+def _test_accuracy_once(M, N, K, with_bias, out_dtype, device):
+    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=device) - 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=device) - 0.5) * 2 * fp8_max
+    b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+    scale_a = torch.randn((M,), device=device, dtype=torch.float32) * 0.001
+    scale_b = torch.randn((N,), device=device, dtype=torch.float32) * 0.001
+    if with_bias:
+        bias = torch.randn((N,), device=device, dtype=out_dtype)
+    else:
+        bias = None
+    b_fp8 = b_fp8.t()
+    o = torch_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias)
+    o1 = fp8_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias)
+    rtol = 0.02
+    atol = 1
+    torch.testing.assert_close(o, o1, rtol=rtol, atol=atol)
+    print(f"M={M}, N={N}, K={K}, with_bias={with_bias}, out_dtype={out_dtype}: OK")
+
+
+@pytest.mark.parametrize("M", [1, 128, 512, 1024, 4096])
+@pytest.mark.parametrize("N", [16, 128, 512, 1024, 4096])
+@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("with_bias", [True, False])
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
+def test_accuracy(M, N, K, with_bias, out_dtype):
+    _test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_fused_qk_norm_rope.py

@@ -0,0 +1,235 @@
+import pytest
+import torch
+from sgl_kernel import fused_qk_norm_rope as sgl_fused_qk_norm_rope
+
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.rotary_embedding import get_rope
+from sglang.srt.server_args import (
+    ServerArgs,
+    get_global_server_args,
+    set_global_server_args_for_scheduler,
+)
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_npu,
+    is_xpu,
+)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_cpu = is_cpu()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_npu = is_npu()
+_is_xpu = is_xpu()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+@torch.inference_mode()
+def torch_ref_rms_norm_rope(
+    qkv,
+    num_heads_q,
+    num_heads_k,
+    num_heads_v,
+    head_dim,
+    eps,
+    q_weight,
+    k_weight,
+    base,
+    is_neox,
+    position_ids,
+    partial_rotary_factor,
+):
+    """
+    PyTorch reference implementation of RMSNorm+RoPE for verification.
+
+    Uses SGLang's own RMSNorm and RotaryEmbedding modules to ensure consistency
+    with the expected behavior of the fused kernel.
+
+    Args:
+        qkv: Combined QKV tensor of shape [num_tokens, hidden_size]
+        num_heads_q: Number of query heads
+        num_heads_k: Number of key heads
+        num_heads_v: Number of value heads (unused for normalization/RoPE but needed for tensor splitting)
+        head_dim: Dimension of each head
+        eps: Epsilon value for RMS normalization
+        q_weight: RMSNorm weights for query [head_dim]
+        k_weight: RMSNorm weights for key [head_dim]
+        base: Base value for RoPE calculations
+        is_neox: Whether to use NeoX style RoPE
+        position_ids: Position IDs for RoPE of shape [num_tokens]
+        partial_rotary_factor: Partial rotary factor
+
+    Returns:
+        Combined tensor with Q and K parts normalized and RoPE applied
+    """
+    # Get input shape information
+    num_tokens = qkv.shape[0]
+    hidden_size = qkv.shape[1]
+
+    # Calculate dimensions for Q, K, V segments
+    q_size = num_heads_q * head_dim
+    k_size = num_heads_k * head_dim
+    v_size = num_heads_v * head_dim
+
+    # Verify dimensions match
+    assert (
+        hidden_size == q_size + k_size + v_size
+    ), f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}"
+
+    # Split the tensor into Q, K, V parts
+    q = qkv[:, :q_size]
+    k = qkv[:, q_size : q_size + k_size]
+    v = qkv[:, q_size + k_size :]
+
+    rotary_emb = get_rope(
+        head_dim,
+        rotary_dim=head_dim,
+        max_position=8192,
+        base=10000,
+        is_neox_style=is_neox,
+        rope_scaling=None,
+        dual_chunk_attention_config=None,
+        partial_rotary_factor=partial_rotary_factor,
+    )
+    rotary_emb = rotary_emb.to(qkv.device)
+
+    # Create and apply RMSNorm modules with custom weights
+    q_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype)
+    q_norm.weight.data.copy_(q_weight)
+    k_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype)
+    k_norm.weight.data.copy_(k_weight)
+
+    q_by_head = q.reshape(-1, head_dim)
+    q_by_head = q_norm(q_by_head)
+    k_by_head = k.reshape(-1, head_dim)
+    k_by_head = k_norm(k_by_head)
+    q = q_by_head.view(q.shape)
+    k = k_by_head.view(k.shape)
+
+    [q_rope, k_rope] = rotary_emb(
+        position_ids,
+        q,
+        k,
+        fused_set_kv_buffer_arg=None,
+    )
+
+    # Combine Q, K, V back together
+    result = torch.cat([q_rope, k_rope, v], dim=1)
+
+    return result
+
+
+head_dims = [64, 128]
+# (Q heads, K heads, V heads)
+num_heads_groups = [
+    (16, 8, 8),  # Qwen3-0.6B, Qwen3-1.7B
+    (32, 8, 8),  # Qwen3-4B, Qwen3-8B, Qwen3-30B-A3B
+    (40, 8, 8),  # Qwen3-14B
+    (64, 8, 8),  # Qwen3-32B, Qwen3-235B-A22B
+    (12, 1, 1),  # GLM4.6 TP8
+]
+num_tokens_list = [1, 3, 8, 32, 256]
+is_neox_list = [False, True]
+dtypes = [torch.bfloat16]
+partial_rotary_factor_list = [1.0, 0.5]
+
+
+@pytest.mark.skipif(not _is_cuda, reason="Skipping CUDA/ROCm only tests.")
+@pytest.mark.parametrize("head_dim", head_dims)
+@pytest.mark.parametrize("num_heads_group", num_heads_groups)
+@pytest.mark.parametrize("num_tokens", num_tokens_list)
+@pytest.mark.parametrize("is_neox", is_neox_list)
+@pytest.mark.parametrize("dtype", dtypes)
+@pytest.mark.parametrize("partial_rotary_factor", partial_rotary_factor_list)
+def test_fused_qk_norm_rope(
+    head_dim, num_heads_group, num_tokens, is_neox, dtype, partial_rotary_factor
+):
+    """
+    Test the fused QK RMSNorm + RoPE operation with various configurations.
+
+    This test verifies that the fused kernel correctly applies:
+    1. RMSNorm to both query (Q) and key (K) portions of the QKV tensor
+    2. Rotary Position Embeddings (RoPE) to the normalized Q and K
+    3. Leaves the value (V) portion unchanged
+
+    Args:
+        head_dim: Dimension of each attention head
+        num_heads_group: Tuple of (num_heads_q, num_heads_k, num_heads_v)
+        num_tokens: Number of tokens to process
+        dtype: Data type (float16 or bfloat16)
+    """
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+    device = "cuda"
+    torch_dtype = dtype
+
+    # Unpack head counts
+    num_heads_q, num_heads_k, num_heads_v = num_heads_group
+
+    # Calculate total hidden dimension
+    hidden_size = (num_heads_q + num_heads_k + num_heads_v) * head_dim
+
+    # Generate random inputs directly as 2D [num_tokens, hidden_size]
+    torch.random.manual_seed(0)
+    qkv = torch.randn(num_tokens, hidden_size, dtype=torch_dtype, device=device)
+    qkv_copy = qkv.clone()
+
+    # Generate position IDs with +100 offset to test decoding scenarios
+    position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) + 100
+
+    # Generate random weights for RMSNorm
+    q_weight = torch.randn(head_dim, dtype=torch_dtype, device=device) * 5.0
+    k_weight = torch.randn(head_dim, dtype=torch_dtype, device=device) * 5.0
+
+    # Set RMSNorm and RoPE parameters
+    eps = 1e-5
+    base = 10000.0
+
+    factor, low, high, attention_factor = 1.0, 0, 0, 1.0
+    # Run the custom fusedQKNormRope operation
+    sgl_fused_qk_norm_rope(
+        qkv,
+        num_heads_q,
+        num_heads_k,
+        num_heads_v,
+        head_dim,
+        eps,
+        q_weight,
+        k_weight,
+        base,
+        is_neox,
+        position_ids,
+        factor,
+        low,
+        high,
+        attention_factor,
+        int(head_dim * partial_rotary_factor),
+    )
+    output = qkv  # This op is inplace
+
+    # Compute reference output using TensorRT-LLM modules
+    ref_output = torch_ref_rms_norm_rope(
+        qkv_copy,
+        num_heads_q,
+        num_heads_k,
+        num_heads_v,
+        head_dim,
+        eps,
+        q_weight,
+        k_weight,
+        base,
+        is_neox,
+        position_ids,
+        partial_rotary_factor,
+    )
+
+    # Compare outputs from custom kernel vs reference implementation
+    torch.testing.assert_close(
+        output,
+        ref_output,
+        rtol=5e-2,
+        atol=1e-1,
+    )

third_party/sglang/sgl-kernel/tests/test_gguf.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from gguf import GGMLQuantizationType, GGUFReader, ReaderTensor, dequantize
+from huggingface_hub import snapshot_download
+from sgl_kernel import (
+    ggml_dequantize,
+    ggml_moe_a8,
+    ggml_moe_a8_vec,
+    ggml_moe_get_block_size,
+    ggml_mul_mat_a8,
+    ggml_mul_mat_vec_a8,
+)
+
+GGUF_SAMPLE = snapshot_download("Isotr0py/test-gguf-sample")
+GGUF_SAMPLE_MOE = snapshot_download("SzymonOzog/test-gguf-moe-sample")
+
+
+def get_gguf_sample_tensors(
+    hidden_size: int, quant_type: GGMLQuantizationType
+) -> list[ReaderTensor]:
+    sample_dir = GGUF_SAMPLE
+    filename = f"Quant_{quant_type.name}_{hidden_size}.gguf"
+    sample_file = Path(sample_dir) / filename
+    return GGUFReader(sample_file).tensors
+
+
+def get_gguf_MoE_tensors(
+    hidden_size: int, quant_type: GGMLQuantizationType
+) -> list[ReaderTensor]:
+    sample_dir = GGUF_SAMPLE_MOE
+    filename = f"Quant_{quant_type.name}_{hidden_size}.gguf"
+    sample_file = Path(sample_dir) / filename
+    return GGUFReader(sample_file).tensors
+
+
+DTYPES = [torch.bfloat16]  # [torch.half, torch.bfloat16, torch.float32]
+# Hidden_size for testing, must match the sample file in HF repo,
+# we have `hidden_size = 256, 1024` for test in HF repo currently.
+HIDDEN_SIZES = [256, 1024]
+NUM_TOKENS = [7, 2050]  # Arbitrary values for testing
+SEEDS = [0]
+QUANT_TYPES = [
+    # i-matrix
+    GGMLQuantizationType.IQ1_M,
+    GGMLQuantizationType.IQ1_S,
+    GGMLQuantizationType.IQ2_S,
+    GGMLQuantizationType.IQ2_XS,
+    GGMLQuantizationType.IQ3_S,
+    GGMLQuantizationType.IQ3_XXS,
+    GGMLQuantizationType.IQ4_NL,
+    GGMLQuantizationType.IQ4_XS,
+    # k-quants
+    GGMLQuantizationType.Q2_K,
+    GGMLQuantizationType.Q3_K,
+    GGMLQuantizationType.Q4_K,
+    GGMLQuantizationType.Q5_K,
+    GGMLQuantizationType.Q6_K,
+    # standard quantization
+    GGMLQuantizationType.Q4_0,
+    GGMLQuantizationType.Q5_0,
+    GGMLQuantizationType.Q8_0,
+]
+
+
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("quant_type", QUANT_TYPES)
+@torch.inference_mode()
+def test_dequantize(
+    hidden_size: int, dtype: torch.dtype, quant_type: GGMLQuantizationType
+):
+    tensors = get_gguf_sample_tensors(hidden_size, quant_type)
+    for tensor in tensors:
+        shape_str = tensor.name.split("_")[-1]
+        shape = map(int, shape_str.split("x"))
+
+        ref_output = torch.tensor(
+            dequantize(tensor.data, quant_type), device="cuda"
+        ).to(dtype)
+        output = ggml_dequantize(
+            torch.tensor(tensor.data, device="cuda"), quant_type, *list(shape), dtype
+        )
+
+        torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=4e-2)
+
+
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("quant_type", QUANT_TYPES)
+@torch.inference_mode()
+def test_mmvq(hidden_size: int, dtype: torch.dtype, quant_type: GGMLQuantizationType):
+
+    tensors = get_gguf_sample_tensors(hidden_size, quant_type)
+    x = torch.rand((1, hidden_size), dtype=dtype, device="cuda")
+    for tensor in tensors:
+        weight = torch.tensor(dequantize(tensor.data, quant_type), device="cuda").to(
+            dtype
+        )
+        ref_output = x @ weight.T
+
+        qweight = torch.tensor(tensor.data, device="cuda")
+        output = ggml_mul_mat_vec_a8(qweight, x, quant_type, qweight.shape[0]).to(dtype)
+
+        # NOTE(FlamingoPg): There can be occasional errors, Loosen the granularity of gguf bf16 verification.
+        atols = {torch.half: 1, torch.bfloat16: 1.5, torch.float: 1}
+        rtols = {torch.half: 1e-1, torch.bfloat16: 3e1, torch.float: 1e-1}
+
+        torch.testing.assert_close(
+            output, ref_output, atol=atols[dtype], rtol=rtols[dtype]
+        )
+
+
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize(
+    "quant_type",
+    [
+        # k-quants
+        GGMLQuantizationType.Q2_K,
+        GGMLQuantizationType.Q3_K,
+        GGMLQuantizationType.Q4_K,
+        GGMLQuantizationType.Q5_K,
+        GGMLQuantizationType.Q6_K,
+        # standard quants
+        GGMLQuantizationType.Q4_0,
+        GGMLQuantizationType.Q5_0,
+        GGMLQuantizationType.Q8_0,
+    ],
+)
+@torch.inference_mode()
+def test_mmq(
+    num_tokens: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    quant_type: GGMLQuantizationType,
+):
+
+    tensors = get_gguf_sample_tensors(hidden_size, quant_type)
+    x = torch.rand((num_tokens, hidden_size), dtype=dtype, device="cuda")
+    for tensor in tensors:
+        weight = torch.tensor(dequantize(tensor.data, quant_type), device="cuda").to(
+            dtype
+        )
+        ref_output = x @ weight.T
+
+        qweight = torch.tensor(tensor.data, device="cuda")
+        output = ggml_mul_mat_a8(qweight, x, quant_type, qweight.shape[0])
+        atols = {torch.half: 1, torch.bfloat16: 1.5, torch.float: 1.2}
+        # test matrix has inputs centered around 0 and lower precision from
+        # bfloat16 tends to accumulate and can greatly inflate rtol
+        # since outputs are also very close to 0
+        rtols = {torch.half: 1e-1, torch.bfloat16: 1e4, torch.float: 2e1}
+        torch.testing.assert_close(
+            output, ref_output, atol=atols[dtype], rtol=rtols[dtype]
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_gptq_kernel.py

@@ -0,0 +1,133 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import gptq_gemm
+
+from sglang.srt.layers.quantization.utils import pack_cols, pack_rows
+
+
+def torch_dequantize(q_weight, q_zeros, scales, g_idx, use_shuffle, bit, K, N):
+    assert bit == 4, "Reference dequantization only supports 4-bit"
+    group_size = K // scales.shape[0]
+    pack_factor = 32 // bit
+
+    # unpack q_weight: (K//pack_factor, N) -> (K, N)
+    unpacked_q_weight = torch.empty(
+        q_weight.shape[0] * pack_factor,
+        q_weight.shape[1],
+        dtype=torch.uint8,
+        device=q_weight.device,
+    )
+    for i in range(pack_factor):
+        unpacked_q_weight[i::pack_factor, :] = (q_weight >> (i * 4)) & 0x0F
+
+    # unpack q_zeros: (num_groups, N//pack_factor) -> (num_groups, N)
+    unpacked_q_zeros = torch.empty(
+        q_zeros.shape[0],
+        q_zeros.shape[1] * pack_factor,
+        dtype=torch.uint8,
+        device=q_zeros.device,
+    )
+    for i in range(pack_factor):
+        unpacked_q_zeros[:, i::pack_factor] = (q_zeros >> (i * 4)) & 0x0F
+
+    unpacked_q_zeros += 1
+    unpacked_q_zeros = unpacked_q_zeros.to(scales.dtype)
+
+    scale_zeros = unpacked_q_zeros * scales  # (num_groups, N)
+
+    current_g_idx = torch.tensor(
+        [i // group_size for i in range(K)], dtype=torch.int32, device=q_weight.device
+    )
+
+    scale_mat = scales[current_g_idx]  # (K, N)
+    scale_zeros_mat = scale_zeros[current_g_idx]  # (K, N)
+
+    # dequant: weight * scale - scale_zeros
+    dequantized_b = unpacked_q_weight.to(scales.dtype) * scale_mat - scale_zeros_mat
+
+    return dequantized_b.reshape(K, N)
+
+
+def torch_gptq_gemm(
+    a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit
+):
+    K, N = a.shape[1], b_q_weight.shape[1]
+
+    b_dequant = torch_dequantize(
+        b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit, K, N
+    )
+    c = torch.matmul(a, b_dequant)
+    return c
+
+
+def _test_gptq_gemm_once(M, N, K, bit, group_size, use_shuffle, dtype, device="cuda"):
+
+    b_fp = torch.randn(K, N, dtype=dtype, device=device)
+
+    assert K % group_size == 0, "K must be divisible by group_size"
+    num_groups = K // group_size
+
+    if use_shuffle:
+        return
+    else:
+        g_idx = torch.tensor(
+            [i // group_size for i in range(K)], dtype=torch.int32, device=device
+        )
+        b_shuffled = b_fp[g_idx]
+
+    b_grouped = b_shuffled.reshape(num_groups, group_size, N)
+
+    b_max = torch.max(b_grouped, dim=1, keepdim=True)[0]
+    b_min = torch.min(b_grouped, dim=1, keepdim=True)[0]
+
+    scales = (b_max - b_min) / (2**bit - 1)
+    scales = scales.clamp(min=1e-6)
+
+    zeros_float = (-b_min / scales).round()
+
+    q_b = (
+        (b_grouped / scales + zeros_float).round().clamp(0, 2**bit - 1).to(torch.uint8)
+    )
+
+    q_zeros_unpacked = zeros_float.to(torch.uint8) - 1
+
+    b_q_weight = pack_rows(q_b.reshape(K, N), bit, K, N)
+
+    q_zeros_unpacked = q_zeros_unpacked.reshape(num_groups, N)
+    b_gptq_qzeros = pack_cols(q_zeros_unpacked, bit, num_groups, N)
+    b_gptq_scales = scales.squeeze(1)
+
+    a = torch.randn(M, K, dtype=dtype, device=device)
+
+    c_ref = torch_gptq_gemm(
+        a, b_q_weight, b_gptq_qzeros, b_gptq_scales, g_idx, use_shuffle, bit
+    )
+    c_out = gptq_gemm(
+        a, b_q_weight, b_gptq_qzeros, b_gptq_scales, g_idx, use_shuffle, bit
+    )
+
+    rtol = 4e-2
+    atol = 4e-2
+    torch.testing.assert_close(c_ref, c_out, rtol=rtol, atol=atol)
+    print(
+        f"✅ Test passed: M={M}, N={N}, K={K}, bit={bit}, group_size={group_size}, use_shuffle={use_shuffle}, dtype={dtype}"
+    )
+
+
+@pytest.mark.parametrize("M", [1, 8, 128])
+@pytest.mark.parametrize("N", [2048, 4096])
+@pytest.mark.parametrize("K", [2048, 4096])
+@pytest.mark.parametrize("bit", [4])
+@pytest.mark.parametrize("group_size", [128])
+@pytest.mark.parametrize("use_shuffle", [False])
+@pytest.mark.parametrize("dtype", [torch.float16])
+def test_gptq_gemm(M, N, K, bit, group_size, use_shuffle, dtype):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    _test_gptq_gemm_once(M, N, K, bit, group_size, use_shuffle, dtype, "cuda")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

third_party/sglang/sgl-kernel/tests/test_int8_gemm.py

@@ -0,0 +1,50 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import int8_scaled_mm
+from utils import is_sm10x
+
+
+def to_int8(tensor: torch.Tensor) -> torch.Tensor:
+    return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
+
+
+def torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias):
+    o = torch.matmul(a.to(torch.float32), b.to(torch.float32))
+    if bias is not None:
+        o = o.to(torch.float32) * scale_a.view(-1, 1) * scale_b.view(1, -1) + bias
+    else:
+        o = o.to(torch.float32) * scale_a.view(-1, 1) * scale_b.view(1, -1)
+    return o.to(out_dtype)
+
+
+def _test_accuracy_once(M, N, K, with_bias, out_dtype, device):
+    a = to_int8(torch.randn((M, K), device=device) * 5)
+    b = to_int8(torch.randn((N, K), device=device).t() * 5)
+    scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
+    scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
+    if with_bias:
+        bias = torch.randn((N,), device="cuda", dtype=out_dtype) * 10
+    else:
+        bias = None
+    o = int8_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
+    o1 = torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
+    torch.testing.assert_close(o, o1)
+
+
+@pytest.mark.skipif(
+    is_sm10x(),
+    reason="int8_scaled_mm is only supported on sm90 and lower",
+)
+@pytest.mark.parametrize("M", [1, 16, 32, 64, 128, 512, 1024, 4096, 8192])
+@pytest.mark.parametrize("N", [16, 128, 512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("with_bias", [True, False])
+@pytest.mark.parametrize("out_dtype", [torch.float16, torch.bfloat16])
+def test_accuracy(M, N, K, with_bias, out_dtype):
+    _test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_kimi_k2_moe_fused_gate.py

@@ -0,0 +1,126 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import kimi_k2_moe_fused_gate
+
+from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl
+
+
+@pytest.mark.parametrize(
+    "seq_length",
+    list(range(1, 10))
+    + [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536],
+)
+@pytest.mark.parametrize("topk", [6])  # Kimi K2 uses topk=6
+@pytest.mark.parametrize("dtype", [torch.float32])
+@pytest.mark.parametrize("apply_routed_scaling_factor_on_output", [False, True])
+def test_kimi_k2_moe_fused_gate(
+    seq_length, topk, dtype, apply_routed_scaling_factor_on_output
+):
+    num_experts = 384  # Kimi K2: only support 384 experts
+    renormalize = True
+    routed_scaling_factor = 2.872  # Kimi K2's routed scaling factor
+
+    torch.manual_seed(seq_length)
+    tensor = torch.rand((seq_length, num_experts), dtype=dtype, device="cuda")
+    scores = tensor.clone()
+    bias = torch.rand(num_experts, dtype=dtype, device="cuda")
+
+    # Test our fused kernel
+    output, indices = kimi_k2_moe_fused_gate(
+        tensor,
+        bias,
+        topk=topk,
+        renormalize=renormalize,
+        routed_scaling_factor=routed_scaling_factor,
+        apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+    )
+
+    # Reference implementation
+    ref_output, ref_indices = kimi_k2_biased_topk_impl(
+        scores,
+        scores,
+        bias,
+        topk=topk,
+        renormalize=renormalize,
+        routed_scaling_factor=routed_scaling_factor,
+        apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+    )
+
+    # Check weights match (after sorting)
+    # Weights are the most important - they determine the actual MoE output
+    output_check = torch.allclose(
+        ref_output.sort()[0].to(torch.float32),
+        output.sort()[0].to(torch.float32),
+        rtol=1e-02,
+        atol=1e-03,
+    )
+
+    assert output_check, (
+        f"Output mismatch at seq_length {seq_length}, dtype {dtype}, "
+        f"num_experts {num_experts}, topk {topk}, "
+        f"apply_routed_scaling_factor_on_output {apply_routed_scaling_factor_on_output}"
+    )
+
+
+@pytest.mark.parametrize("seq_length", [1024, 4096])
+@pytest.mark.parametrize("num_experts", [384])
+@pytest.mark.parametrize("topk", [6])
+def test_kimi_k2_specific_case(seq_length, num_experts, topk):
+    """Test specifically for Kimi K2 configuration: 384 experts, topk=6"""
+    dtype = torch.float32
+    renormalize = True
+    routed_scaling_factor = 2.872
+
+    torch.manual_seed(42)
+    tensor = torch.rand((seq_length, num_experts), dtype=dtype, device="cuda")
+    scores = tensor.clone()
+    bias = torch.rand(num_experts, dtype=dtype, device="cuda")
+
+    output, indices = kimi_k2_moe_fused_gate(
+        tensor,
+        bias,
+        topk=topk,
+        renormalize=renormalize,
+        routed_scaling_factor=routed_scaling_factor,
+        apply_routed_scaling_factor_on_output=False,
+    )
+
+    ref_output, ref_indices = kimi_k2_biased_topk_impl(
+        scores,
+        scores,
+        bias,
+        topk=topk,
+        renormalize=renormalize,
+        routed_scaling_factor=routed_scaling_factor,
+        apply_routed_scaling_factor_on_output=False,
+    )
+
+    # Verify output shapes
+    assert output.shape == (seq_length, topk)
+    assert indices.shape == (seq_length, topk)
+    assert output.dtype == torch.float32
+    assert indices.dtype == torch.int32
+
+    # Verify weights are normalized (sum to 1 per token if renormalize=True)
+    if renormalize:
+        weight_sums = output.sum(dim=-1)
+        assert torch.allclose(
+            weight_sums, torch.ones_like(weight_sums), rtol=1e-3, atol=1e-4
+        )
+
+    # Check weights match (after sorting)
+    # Weights are the most important - they determine the actual MoE output
+    output_check = torch.allclose(
+        ref_output.sort()[0].to(torch.float32),
+        output.sort()[0].to(torch.float32),
+        rtol=1e-02,
+        atol=1e-03,
+    )
+
+    assert output_check, f"Output mismatch for Kimi K2 specific case"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_kvcacheio.py

@@ -0,0 +1,708 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel.kvcacheio import (
+    transfer_kv_all_layer,
+    transfer_kv_all_layer_direct_lf_pf,
+    transfer_kv_all_layer_lf_ph,
+    transfer_kv_all_layer_mla,
+    transfer_kv_direct,
+    transfer_kv_per_layer,
+    transfer_kv_per_layer_direct_pf_lf,
+    transfer_kv_per_layer_mla,
+)
+
+from sglang.srt.utils import is_hip
+
+
+def ref_copy_with_indices(src_pool, dst_pool, src_indices, dst_indices):
+    dst_pool[dst_indices] = src_pool[src_indices].to(dst_pool.device)
+
+
+def ref_copy_with_indices_pf_direct(
+    src_pool, dst_pool, src_indices, dst_indices, page_size, layer_id, lf_to_pf=False
+):
+    if lf_to_pf:
+        for i in range(0, len(src_indices), page_size):
+            dst_pool[dst_indices[i] // page_size][layer_id] = src_pool[layer_id][
+                src_indices[i : i + page_size]
+            ].to(dst_pool.device)
+    else:
+        for i in range(0, len(src_indices), page_size):
+            dst_pool[layer_id][dst_indices[i : i + page_size]] = src_pool[
+                src_indices[i] // page_size
+            ][layer_id].to(dst_pool.device)
+
+
+def ref_copy_with_indices_page_head(
+    src_pool,
+    dst_pool,
+    src_indices,
+    dst_indices,
+    page_size,
+    layer_id,
+    head_num,
+    lf_to_ph=False,
+):
+    if lf_to_ph:
+        for head_id in range(head_num):
+            for i in range(0, len(src_indices)):
+                dst_pool[dst_indices[i] // page_size][head_id][
+                    dst_indices[i] % page_size
+                ][layer_id] = src_pool[layer_id][src_indices[i]][head_id].to(
+                    dst_pool.device
+                )
+    else:
+        for head_id in range(head_num):
+            for i in range(0, len(src_indices)):
+                dst_pool[layer_id][dst_indices[i]][head_id] = src_pool[
+                    src_indices[i] // page_size
+                ][head_id][src_indices[i] % page_size][layer_id].to(dst_pool.device)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_items_to_transfer", [1, 128, 1024])
+@pytest.mark.parametrize("page_size", [1, 16, 64])
+@pytest.mark.parametrize("item_size", [256])
+@pytest.mark.parametrize("total_items_in_pool", [10240])
+@pytest.mark.parametrize("is_mla", [False, True])
+@pytest.mark.parametrize("all_layers", [False, True])
+def test_transfer_kv(
+    dtype: torch.dtype,
+    num_items_to_transfer: int,
+    item_size: int,
+    page_size: int,
+    total_items_in_pool: int,
+    is_mla: bool,
+    all_layers: bool,
+):
+    """
+    Tests the per-layer transfer functions, treating tensors as memory pools.
+    """
+
+    original_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    device = "cuda"
+    torch.cuda.manual_seed(42)
+
+    num_layers = 4  # A small number of layers for pool creation
+
+    total_pages_in_pool = total_items_in_pool // page_size
+    num_pages_to_transfer = num_items_to_transfer // page_size
+    if num_pages_to_transfer == 0:
+        torch.set_default_dtype(original_dtype)
+        return
+    page_indices = torch.randperm(total_pages_in_pool, dtype=torch.int64)
+    src_indices_host = torch.cat(
+        [
+            torch.arange(p * page_size, (p + 1) * page_size)
+            for p in page_indices[:num_pages_to_transfer]
+        ]
+    )
+    src_indices_device = src_indices_host.to(device)
+    dst_indices_host = torch.cat(
+        [
+            torch.arange(p * page_size, (p + 1) * page_size)
+            for p in page_indices[num_pages_to_transfer : 2 * num_pages_to_transfer]
+        ]
+    )
+    dst_indices_device = dst_indices_host.to(device)
+
+    # Prepare memory pools based on whether it's an MLA case.
+    if is_mla:
+        src_pool_host = torch.randn(
+            num_layers, total_items_in_pool, item_size
+        ).pin_memory()
+        dst_pool_ref = torch.zeros_like(src_pool_host).to(device)
+        dst_pool_kernel = torch.zeros_like(dst_pool_ref)
+        dst_pool_direct = torch.zeros_like(dst_pool_ref)
+    else:
+        src_k_pool = torch.randn(
+            num_layers, total_items_in_pool, item_size
+        ).pin_memory()
+        src_v_pool = torch.randn(
+            num_layers, total_items_in_pool, item_size
+        ).pin_memory()
+        dst_k_pool_ref = torch.zeros_like(src_k_pool).to(device)
+        dst_v_pool_ref = torch.zeros_like(src_v_pool).to(device)
+        dst_k_pool_kernel = torch.zeros_like(dst_k_pool_ref)
+        dst_v_pool_kernel = torch.zeros_like(dst_v_pool_ref)
+        dst_k_pool_direct = torch.zeros_like(dst_k_pool_ref)
+        dst_v_pool_direct = torch.zeros_like(dst_v_pool_ref)
+
+    torch.cuda.synchronize()
+
+    # We will test the per-layer function on the first layer (index 0) of the pool.
+    layer_idx_to_test = 0
+
+    if is_mla:
+        if not all_layers:
+            ref_copy_with_indices(
+                src_pool_host[layer_idx_to_test],
+                dst_pool_ref[layer_idx_to_test],
+                src_indices_host,
+                dst_indices_device,
+            )
+            transfer_kv_per_layer_mla(
+                src_pool_host[layer_idx_to_test],
+                dst_pool_kernel[layer_idx_to_test],
+                src_indices_device,
+                dst_indices_device,
+                item_size=item_size * dtype.itemsize,
+            )
+            transfer_kv_direct(
+                [src_pool_host[layer_idx_to_test]],
+                [dst_pool_direct[layer_idx_to_test]],
+                src_indices_host,
+                dst_indices_device,
+                page_size=page_size,
+            )
+        else:
+            for layer_id in range(num_layers):
+                ref_copy_with_indices(
+                    src_pool_host[layer_id],
+                    dst_pool_ref[layer_id],
+                    src_indices_host,
+                    dst_indices_device,
+                )
+            src_layers_device = torch.tensor(
+                [src_pool_host[layer_id].data_ptr() for layer_id in range(num_layers)],
+                dtype=torch.uint64,
+                device=device,
+            )
+            dst_layers_device = torch.tensor(
+                [
+                    dst_pool_kernel[layer_id].data_ptr()
+                    for layer_id in range(num_layers)
+                ],
+                dtype=torch.uint64,
+                device=device,
+            )
+            transfer_kv_all_layer_mla(
+                src_layers_device,
+                dst_layers_device,
+                src_indices_device,
+                dst_indices_device,
+                item_size=item_size * dtype.itemsize,
+                num_layers=num_layers,
+            )
+            transfer_kv_direct(
+                [src_pool_host[layer_id] for layer_id in range(num_layers)],
+                [dst_pool_direct[layer_id] for layer_id in range(num_layers)],
+                src_indices_host,
+                dst_indices_device,
+                page_size=page_size,
+            )
+        torch.cuda.synchronize()
+        torch.testing.assert_close(dst_pool_kernel, dst_pool_ref)
+        torch.testing.assert_close(dst_pool_direct, dst_pool_ref)
+    else:
+        if not all_layers:
+            ref_copy_with_indices(
+                src_k_pool[layer_idx_to_test],
+                dst_k_pool_ref[layer_idx_to_test],
+                src_indices_host,
+                dst_indices_device,
+            )
+            ref_copy_with_indices(
+                src_v_pool[layer_idx_to_test],
+                dst_v_pool_ref[layer_idx_to_test],
+                src_indices_host,
+                dst_indices_device,
+            )
+            transfer_kv_per_layer(
+                src_k_pool[layer_idx_to_test],
+                dst_k_pool_kernel[layer_idx_to_test],
+                src_v_pool[layer_idx_to_test],
+                dst_v_pool_kernel[layer_idx_to_test],
+                src_indices_device,
+                dst_indices_device,
+                item_size=item_size * dtype.itemsize,
+            )
+            transfer_kv_direct(
+                [src_k_pool[layer_idx_to_test], src_v_pool[layer_idx_to_test]],
+                [
+                    dst_k_pool_direct[layer_idx_to_test],
+                    dst_v_pool_direct[layer_idx_to_test],
+                ],
+                src_indices_host,
+                dst_indices_device,
+                page_size=page_size,
+            )
+        else:
+            for layer_id in range(num_layers):
+                ref_copy_with_indices(
+                    src_k_pool[layer_id],
+                    dst_k_pool_ref[layer_id],
+                    src_indices_host,
+                    dst_indices_device,
+                )
+                ref_copy_with_indices(
+                    src_v_pool[layer_id],
+                    dst_v_pool_ref[layer_id],
+                    src_indices_host,
+                    dst_indices_device,
+                )
+
+            src_k_layers_device = torch.tensor(
+                [src_k_pool[layer_id].data_ptr() for layer_id in range(num_layers)],
+                dtype=torch.uint64,
+                device=device,
+            )
+            src_v_layers_device = torch.tensor(
+                [src_v_pool[layer_id].data_ptr() for layer_id in range(num_layers)],
+                dtype=torch.uint64,
+                device=device,
+            )
+            dst_k_layers_device = torch.tensor(
+                [
+                    dst_k_pool_kernel[layer_id].data_ptr()
+                    for layer_id in range(num_layers)
+                ],
+                dtype=torch.uint64,
+                device=device,
+            )
+            dst_v_layers_device = torch.tensor(
+                [
+                    dst_v_pool_kernel[layer_id].data_ptr()
+                    for layer_id in range(num_layers)
+                ],
+                dtype=torch.uint64,
+                device=device,
+            )
+            transfer_kv_all_layer(
+                src_k_layers_device,
+                dst_k_layers_device,
+                src_v_layers_device,
+                dst_v_layers_device,
+                src_indices_device,
+                dst_indices_device,
+                item_size=item_size * dtype.itemsize,
+                num_layers=num_layers,
+            )
+            transfer_kv_direct(
+                [src_k_pool[layer_id] for layer_id in range(num_layers)]
+                + [src_v_pool[layer_id] for layer_id in range(num_layers)],
+                [dst_k_pool_direct[layer_id] for layer_id in range(num_layers)]
+                + [dst_v_pool_direct[layer_id] for layer_id in range(num_layers)],
+                src_indices_host,
+                dst_indices_device,
+                page_size=page_size,
+            )
+        torch.cuda.synchronize()
+        torch.testing.assert_close(dst_k_pool_kernel, dst_k_pool_ref)
+        torch.testing.assert_close(dst_v_pool_kernel, dst_v_pool_ref)
+        torch.testing.assert_close(dst_k_pool_direct, dst_k_pool_ref)
+        torch.testing.assert_close(dst_v_pool_direct, dst_v_pool_ref)
+
+    torch.set_default_dtype(original_dtype)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_items_to_transfer", [128, 1024, 8192])
+@pytest.mark.parametrize("page_size", [16, 64, 128])
+@pytest.mark.parametrize("item_size", [256])
+@pytest.mark.parametrize("total_items_in_pool", [20480])
+@pytest.mark.parametrize("is_mla", [False, True])
+@pytest.mark.parametrize("lf_to_pf", [False, True])
+def test_transfer_kv_pf_direct(
+    dtype: torch.dtype,
+    num_items_to_transfer: int,
+    item_size: int,
+    page_size: int,
+    total_items_in_pool: int,
+    is_mla: bool,
+    lf_to_pf: bool,
+):
+    original_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    device = "cuda"
+    torch.cuda.manual_seed(42)
+    test_stream = torch.cuda.Stream()
+
+    num_layers = 4
+
+    total_pages_in_pool = total_items_in_pool // page_size
+    num_pages_to_transfer = num_items_to_transfer // page_size
+    if num_pages_to_transfer == 0:
+        torch.set_default_dtype(original_dtype)
+        return
+    page_indices = torch.randperm(total_pages_in_pool, dtype=torch.int64)
+    src_indices_host = torch.cat(
+        [
+            torch.arange(p * page_size, (p + 1) * page_size)
+            for p in page_indices[:num_pages_to_transfer]
+        ]
+    )
+    src_indices_device = src_indices_host.to(device)
+    dst_indices_host = torch.cat(
+        [
+            torch.arange(p * page_size, (p + 1) * page_size)
+            for p in page_indices[num_pages_to_transfer : 2 * num_pages_to_transfer]
+        ]
+    )
+    dst_indices_device = dst_indices_host.to(device)
+
+    # We will test the per-layer function on the first layer (index 0) of the pool.
+    layer_idx_to_test = 0
+
+    if lf_to_pf:
+        if is_mla:
+            src_pool = torch.randn(num_layers, total_items_in_pool, item_size).to(
+                device
+            )
+            src_pool_ptrs = [src_pool[i] for i in range(num_layers)]
+            dst_pool_ref = torch.zeros(
+                total_pages_in_pool, num_layers, page_size, item_size
+            ).pin_memory()
+            dst_pool_direct = torch.zeros_like(dst_pool_ref)
+            torch.cuda.synchronize()
+
+            with torch.cuda.stream(test_stream):
+                transfer_kv_all_layer_direct_lf_pf(
+                    src_pool_ptrs,
+                    [dst_pool_direct],
+                    src_indices_host,
+                    dst_indices_host,
+                    page_size,
+                )
+            test_stream.synchronize()
+
+            for i in range(num_layers):
+                ref_copy_with_indices_pf_direct(
+                    src_pool,
+                    dst_pool_ref,
+                    src_indices_device,
+                    dst_indices_host,
+                    page_size,
+                    i,
+                    lf_to_pf=True,
+                )
+            torch.cuda.synchronize()
+            torch.testing.assert_close(dst_pool_direct, dst_pool_ref)
+
+        else:
+            src_k_pool = torch.randn(num_layers, total_items_in_pool, item_size).to(
+                device
+            )
+            src_k_pool_ptrs = [src_k_pool[i] for i in range(num_layers)]
+            src_v_pool = torch.randn(num_layers, total_items_in_pool, item_size).to(
+                device
+            )
+            src_v_pool_ptrs = [src_v_pool[i] for i in range(num_layers)]
+            dst_k_pool_ref = torch.zeros(
+                total_pages_in_pool, num_layers, page_size, item_size
+            ).pin_memory()
+            dst_v_pool_ref = torch.zeros_like(dst_k_pool_ref)
+            dst_k_pool_direct = torch.zeros_like(dst_k_pool_ref)
+            dst_v_pool_direct = torch.zeros_like(dst_v_pool_ref)
+            torch.cuda.synchronize()
+
+            with torch.cuda.stream(test_stream):
+                transfer_kv_all_layer_direct_lf_pf(
+                    src_k_pool_ptrs + src_v_pool_ptrs,
+                    [dst_k_pool_direct, dst_v_pool_direct],
+                    src_indices_host,
+                    dst_indices_host,
+                    page_size,
+                )
+            test_stream.synchronize()
+
+            for i in range(num_layers):
+                ref_copy_with_indices_pf_direct(
+                    src_k_pool,
+                    dst_k_pool_ref,
+                    src_indices_device,
+                    dst_indices_host,
+                    page_size,
+                    i,
+                    lf_to_pf=True,
+                )
+                ref_copy_with_indices_pf_direct(
+                    src_v_pool,
+                    dst_v_pool_ref,
+                    src_indices_device,
+                    dst_indices_host,
+                    page_size,
+                    i,
+                    lf_to_pf=True,
+                )
+            torch.cuda.synchronize()
+            torch.testing.assert_close(dst_k_pool_direct, dst_k_pool_ref)
+            torch.testing.assert_close(dst_v_pool_direct, dst_v_pool_ref)
+    else:
+        if is_mla:
+            src_pool = torch.randn(
+                total_pages_in_pool, num_layers, page_size, item_size
+            ).pin_memory()
+
+            dst_pool_ref = torch.zeros(num_layers, total_items_in_pool, item_size).to(
+                device
+            )
+            dst_pool_direct = torch.zeros_like(dst_pool_ref)
+            dst_pool_direct_ptrs = [dst_pool_direct[i] for i in range(num_layers)]
+            torch.cuda.synchronize()
+
+            with torch.cuda.stream(test_stream):
+                transfer_kv_per_layer_direct_pf_lf(
+                    [src_pool],
+                    [dst_pool_direct_ptrs[layer_idx_to_test]],
+                    src_indices_host,
+                    dst_indices_host,
+                    layer_idx_to_test,
+                    page_size,
+                )
+            test_stream.synchronize()
+
+            ref_copy_with_indices_pf_direct(
+                src_pool,
+                dst_pool_ref,
+                src_indices_host,
+                dst_indices_device,
+                page_size,
+                layer_idx_to_test,
+                lf_to_pf=False,
+            )
+            torch.cuda.synchronize()
+            torch.testing.assert_close(dst_pool_direct, dst_pool_ref)
+        else:
+            src_k_pool = torch.randn(
+                total_pages_in_pool, num_layers, page_size, item_size
+            ).pin_memory()
+            src_v_pool = torch.randn(
+                total_pages_in_pool, num_layers, page_size, item_size
+            ).pin_memory()
+
+            dst_k_pool_ref = torch.zeros(num_layers, total_items_in_pool, item_size).to(
+                device
+            )
+            dst_k_pool_direct = torch.zeros_like(dst_k_pool_ref)
+            dst_k_pool_direct_ptrs = [dst_k_pool_direct[i] for i in range(num_layers)]
+
+            dst_v_pool_ref = torch.zeros_like(dst_k_pool_ref)
+            dst_v_pool_direct = torch.zeros_like(dst_v_pool_ref)
+            dst_v_pool_direct_ptrs = [dst_v_pool_direct[i] for i in range(num_layers)]
+            torch.cuda.synchronize()
+
+            with torch.cuda.stream(test_stream):
+                transfer_kv_per_layer_direct_pf_lf(
+                    [src_k_pool, src_v_pool],
+                    [
+                        dst_k_pool_direct_ptrs[layer_idx_to_test],
+                        dst_v_pool_direct_ptrs[layer_idx_to_test],
+                    ],
+                    src_indices_host,
+                    dst_indices_host,
+                    layer_idx_to_test,
+                    page_size,
+                )
+            test_stream.synchronize()
+
+            ref_copy_with_indices_pf_direct(
+                src_k_pool,
+                dst_k_pool_ref,
+                src_indices_host,
+                dst_indices_device,
+                page_size,
+                layer_idx_to_test,
+                lf_to_pf=False,
+            )
+            ref_copy_with_indices_pf_direct(
+                src_v_pool,
+                dst_v_pool_ref,
+                src_indices_host,
+                dst_indices_device,
+                page_size,
+                layer_idx_to_test,
+                lf_to_pf=False,
+            )
+
+            torch.cuda.synchronize()
+            torch.testing.assert_close(dst_k_pool_direct, dst_k_pool_ref)
+            torch.testing.assert_close(dst_v_pool_direct, dst_v_pool_ref)
+    torch.set_default_dtype(original_dtype)
+
+
+@pytest.mark.skipif(is_hip(), reason="HIP is not supported for this test")
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_items_to_transfer", [256, 1024])
+@pytest.mark.parametrize("page_size", [16, 64, 128])
+@pytest.mark.parametrize("item_size", [1024])
+@pytest.mark.parametrize("head_num", [8, 16])
+@pytest.mark.parametrize("total_items_in_pool", [4096])
+@pytest.mark.parametrize("lf_to_ph", [False, True])
+def test_transfer_kv_page_head(
+    dtype: torch.dtype,
+    num_items_to_transfer: int,
+    page_size: int,
+    item_size: int,
+    head_num: int,
+    total_items_in_pool: int,
+    lf_to_ph: bool,
+):
+    original_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    device = "cuda"
+    torch.cuda.manual_seed(42)
+
+    num_layers = 4
+
+    total_pages_in_pool = total_items_in_pool // page_size
+    num_pages_to_transfer = num_items_to_transfer // page_size
+    if num_pages_to_transfer == 0:
+        torch.set_default_dtype(original_dtype)
+        return
+
+    assert item_size % head_num == 0
+    head_dim = item_size // head_num
+
+    page_indices = torch.randperm(total_pages_in_pool, dtype=torch.int64)
+    src_indices_host = torch.cat(
+        [
+            torch.arange(p * page_size, (p + 1) * page_size)
+            for p in page_indices[:num_pages_to_transfer]
+        ]
+    )
+    src_indices_device = src_indices_host.to(device)
+    dst_indices_host = torch.cat(
+        [
+            torch.arange(p * page_size, (p + 1) * page_size)
+            for p in page_indices[num_pages_to_transfer : 2 * num_pages_to_transfer]
+        ]
+    )
+    dst_indices_device = dst_indices_host.to(device)
+
+    # We will test the per-layer function on the first layer (index 0) of the pool.
+    layer_idx_to_test = 0
+
+    if lf_to_ph:
+        src_k_pool = torch.randn(
+            num_layers, total_items_in_pool, head_num, head_dim
+        ).to(device)
+        src_v_pool = torch.randn(
+            num_layers, total_items_in_pool, head_num, head_dim
+        ).to(device)
+        src_k_pool_ptrs = [src_k_pool[i] for i in range(num_layers)]
+        src_k_pool_ptrs = torch.tensor(
+            [x.data_ptr() for x in src_k_pool_ptrs],
+            dtype=torch.uint64,
+            device=device,
+        )
+        src_v_pool_ptrs = [src_v_pool[i] for i in range(num_layers)]
+        src_v_pool_ptrs = torch.tensor(
+            [x.data_ptr() for x in src_v_pool_ptrs],
+            dtype=torch.uint64,
+            device=device,
+        )
+
+        dst_k_pool_ref = torch.zeros(
+            total_pages_in_pool, head_num, page_size, num_layers, head_dim
+        ).pin_memory()
+        dst_v_pool_ref = torch.zeros_like(dst_k_pool_ref).pin_memory()
+
+        dst_k_pool_kernel = torch.zeros_like(dst_k_pool_ref).pin_memory()
+        dst_v_pool_kernel = torch.zeros_like(dst_v_pool_ref).pin_memory()
+        torch.cuda.synchronize()
+
+        transfer_kv_all_layer_lf_ph(
+            src_k_pool_ptrs,
+            dst_k_pool_kernel,
+            src_v_pool_ptrs,
+            dst_v_pool_kernel,
+            src_indices_device,
+            dst_indices_device,
+            item_size * dtype.itemsize,
+            item_size * num_layers * dtype.itemsize,
+            num_layers,
+            page_size,
+            head_num,
+        )
+        torch.cuda.synchronize()
+
+        for i in range(num_layers):
+            ref_copy_with_indices_page_head(
+                src_k_pool,
+                dst_k_pool_ref,
+                src_indices_device,
+                dst_indices_host,
+                page_size,
+                i,
+                head_num,
+                lf_to_ph=True,
+            )
+            ref_copy_with_indices_page_head(
+                src_v_pool,
+                dst_v_pool_ref,
+                src_indices_device,
+                dst_indices_host,
+                page_size,
+                i,
+                head_num,
+                lf_to_ph=True,
+            )
+        torch.cuda.synchronize()
+        torch.testing.assert_close(dst_k_pool_kernel, dst_k_pool_ref)
+        torch.testing.assert_close(dst_v_pool_kernel, dst_v_pool_ref)
+    else:
+        from sgl_kernel.kvcacheio import transfer_kv_per_layer_ph_lf
+
+        src_k_pool = torch.randn(
+            total_pages_in_pool, head_num, page_size, num_layers, head_dim
+        ).pin_memory()
+        src_v_pool = torch.randn(
+            total_pages_in_pool, head_num, page_size, num_layers, head_dim
+        ).pin_memory()
+
+        dst_k_pool_ref = torch.zeros(
+            num_layers, total_items_in_pool, head_num, head_dim
+        ).to(device)
+        dst_v_pool_ref = torch.zeros_like(dst_k_pool_ref)
+        dst_k_pool_kernel = torch.zeros_like(dst_k_pool_ref)
+        dst_v_pool_kernel = torch.zeros_like(dst_v_pool_ref)
+        dst_k_pool_kernel_ptrs = [dst_k_pool_kernel[i] for i in range(num_layers)]
+        dst_v_pool_kernel_ptrs = [dst_v_pool_kernel[i] for i in range(num_layers)]
+        torch.cuda.synchronize()
+
+        transfer_kv_per_layer_ph_lf(
+            src_k_pool,
+            dst_k_pool_kernel_ptrs[layer_idx_to_test],
+            src_v_pool,
+            dst_v_pool_kernel_ptrs[layer_idx_to_test],
+            src_indices_device,
+            dst_indices_device,
+            layer_idx_to_test,
+            item_size * dtype.itemsize,
+            item_size * num_layers * dtype.itemsize,
+            page_size,
+            head_num,
+        )
+
+        ref_copy_with_indices_page_head(
+            src_k_pool,
+            dst_k_pool_ref,
+            src_indices_host,
+            dst_indices_device,
+            page_size,
+            layer_idx_to_test,
+            head_num,
+            lf_to_ph=False,
+        )
+        ref_copy_with_indices_page_head(
+            src_v_pool,
+            dst_v_pool_ref,
+            src_indices_host,
+            dst_indices_device,
+            page_size,
+            layer_idx_to_test,
+            head_num,
+            lf_to_ph=False,
+        )
+        torch.cuda.synchronize()
+        torch.testing.assert_close(dst_k_pool_kernel, dst_k_pool_ref)
+        torch.testing.assert_close(dst_v_pool_kernel, dst_v_pool_ref)
+    torch.set_default_dtype(original_dtype)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_merge_state_v2.py

@@ -0,0 +1,381 @@
+import sys
+from typing import Optional
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import merge_state_v2
+
+
+@triton.jit
+def merge_state_kernel(
+    output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_merged
+    output_lse,  # [NUM_TOKENS, NUM_HEADS] s_merged
+    prefix_output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_a
+    prefix_lse,  # [NUM_TOKENS, NUM_HEADS] s_a
+    suffix_output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_b
+    suffix_lse,  # [NUM_TOKENS, NUM_HEADS] s_b
+    HEAD_SIZE: tl.constexpr,
+    PADDED_HEAD_SIZE: tl.constexpr,
+    OUTPUT_LSE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    num_tokens = tl.num_programs(0)
+    head_idx = tl.program_id(1)
+    num_heads = tl.num_programs(1)
+
+    p_lse = tl.load(prefix_lse + token_idx * num_heads + head_idx)
+    s_lse = tl.load(suffix_lse + token_idx * num_heads + head_idx)
+    p_lse = float("-inf") if p_lse == float("inf") else p_lse
+    s_lse = float("-inf") if s_lse == float("inf") else s_lse
+
+    max_lse = tl.maximum(p_lse, s_lse)
+    p_lse = p_lse - max_lse
+    s_lse = s_lse - max_lse
+    out_se = tl.exp(p_lse) + tl.exp(s_lse)
+
+    if OUTPUT_LSE:
+        out_lse = tl.log(out_se) + max_lse
+        tl.store(output_lse + token_idx * num_heads + head_idx, out_lse)
+
+    head_arange = tl.arange(0, PADDED_HEAD_SIZE)
+    head_mask = head_arange < HEAD_SIZE
+    p_out = tl.load(
+        prefix_output
+        + token_idx * num_heads * HEAD_SIZE
+        + head_idx * HEAD_SIZE
+        + head_arange,
+        mask=head_mask,
+    )
+    s_out = tl.load(
+        suffix_output
+        + token_idx * num_heads * HEAD_SIZE
+        + head_idx * HEAD_SIZE
+        + head_arange,
+        mask=head_mask,
+    )
+
+    p_scale = tl.exp(p_lse) / out_se
+    s_scale = tl.exp(s_lse) / out_se
+    out = p_out * p_scale + s_out * s_scale
+    tl.store(
+        output + token_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE + head_arange,
+        out,
+        mask=head_mask,
+    )
+
+
+def merge_state_triton(
+    prefix_output: torch.Tensor,
+    prefix_lse: torch.Tensor,
+    suffix_output: torch.Tensor,
+    suffix_lse: torch.Tensor,
+    output: Optional[torch.Tensor] = None,
+    output_lse: Optional[torch.Tensor] = None,
+) -> None:
+    num_tokens = output.shape[0]
+    num_query_heads = output.shape[1]
+    head_size = output.shape[2]
+    padded_head_size = triton.next_power_of_2(head_size)
+    # Avoid creating new tensors if they are already provided
+    if output is None:
+        output = torch.empty_like(prefix_output)
+    if output_lse is None:
+        output_lse = torch.empty_like(prefix_lse)
+
+    merge_state_kernel[(num_tokens, num_query_heads)](
+        output,
+        output_lse,
+        prefix_output,
+        prefix_lse,
+        suffix_output,
+        suffix_lse,
+        head_size,
+        padded_head_size,
+        output_lse is not None,
+    )
+    return output, output_lse
+
+
+# Naive PyTorch Implements of Merge Attn States
+def merge_state_torch(
+    prefix_output: torch.Tensor,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
+    prefix_lse: torch.Tensor,  # [NUM_TOKENS, NUM_HEADS]
+    suffix_output: torch.Tensor,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
+    suffix_lse: torch.Tensor,  # [NUM_TOKENS, NUM_HEADS]
+    output: Optional[torch.Tensor] = None,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
+    output_lse: Optional[torch.Tensor] = None,  # [NUM_TOKENS, NUM_HEADS]
+):
+    # Avoid creating new tensors if they are already provided
+    if output is None:
+        output = torch.empty_like(prefix_output)
+    if output_lse is None:
+        output_lse = torch.empty_like(prefix_lse)
+    p_lse = prefix_lse
+    s_lse = suffix_lse
+    # inf -> -inf
+    p_lse[p_lse == torch.inf] = -torch.inf
+    s_lse[s_lse == torch.inf] = -torch.inf
+    # max_lse [NUM_HEADS, NUM_TOKENS]
+    max_lse = torch.maximum(p_lse, s_lse)
+    p_lse = p_lse - max_lse
+    s_lse = s_lse - max_lse
+    p_lse_exp = torch.exp(p_lse)
+    s_lse_exp = torch.exp(s_lse)
+    out_se = p_lse_exp + s_lse_exp
+    if output_lse is not None:
+        output_lse = torch.log(out_se) + max_lse
+    p_scale = p_lse_exp / out_se
+    s_scale = s_lse_exp / out_se
+    p_scale = p_scale.unsqueeze(2)  # [NUM_TOKENS, NUM_HEADS, 1]
+    s_scale = s_scale.unsqueeze(2)  # [NUM_TOKENS, NUM_HEADS, 1]
+    output = prefix_output * p_scale + suffix_output * s_scale
+    return output, output_lse
+
+
+NUM_BATCH_TOKENS = [256, 512, 613, 1024, 1536]
+NUM_QUERY_HEADS = [8, 16, 32]
+HEAD_SIZES = [32, 48, 64, 128, 256]
+DTYPES = [torch.half, torch.bfloat16]
+
+all_case_info: list[tuple] = []
+
+
+def generate_markdown_table():
+    global all_case_info
+    table_header = (
+        "| tokens | heads | headsize | dtype "
+        "| device | torch | triton | v2 | speedup(vs triton) |"
+    )
+    table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"
+
+    def shortly_dtype(dtype: torch.dtype) -> str:
+        return str(dtype).removeprefix("torch.")
+
+    def shortly_device(device: str) -> str:
+        return device.removeprefix("NVIDIA").strip()
+
+    print(table_header)
+    print(table_separator)
+    for info in all_case_info:
+        (
+            num_tokens,
+            num_heads,
+            head_size,
+            dtype,
+            device,
+            time_torch,
+            time_triton,
+            time_v2,
+        ) = info
+        dtype = shortly_dtype(dtype)
+        device = shortly_device(device)
+        improved_triton = time_triton / time_v2
+        print(
+            f"| {num_tokens} | {num_heads} | {head_size} "
+            f"| {dtype} | {device} | {time_torch:.4f}ms "
+            f"| {time_triton:.4f}ms "
+            f"| {time_v2:.4f}ms "
+            f"| {improved_triton:.4f}x |"
+        )
+
+
+@pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS)
+@pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("output_dtype", DTYPES)
+@torch.inference_mode()
+def test_merge_attn_states(
+    num_tokens: int, num_query_heads: int, head_size: int, output_dtype: torch.dtype
+):
+    if not torch.cuda.is_available():
+        pytest.skip(
+            "Currently only support compare triton merge_attn_states "
+            "with custom cuda merge_attn_states kernel"
+        )
+
+    NUM_TOKENS = num_tokens
+    NUM_HEADS = num_query_heads
+    HEAD_SIZE = head_size
+
+    print(
+        f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "
+        f"HEAD_SIZE:{HEAD_SIZE}, DTYPE: {output_dtype}, "
+        f"Device: {torch.cuda.get_device_name()}"
+    )
+
+    # prefix_lse and suffix_lse contain inf and normal values
+    prefix_lse = torch.randn(NUM_TOKENS, NUM_HEADS, dtype=torch.float32, device="cuda")
+    suffix_lse = torch.randn(NUM_TOKENS, NUM_HEADS, dtype=torch.float32, device="cuda")
+
+    # Generate boolean masks
+    mask_prefix = torch.rand(NUM_TOKENS, NUM_HEADS) < 0.1
+    mask_suffix = torch.rand(NUM_TOKENS, NUM_HEADS) < 0.1
+    # Ensure that the same position is not True at the same time
+    combined_mask = torch.logical_and(mask_prefix, mask_suffix)
+    mask_prefix = torch.logical_and(mask_prefix, ~combined_mask)
+    mask_suffix = torch.logical_and(mask_suffix, ~combined_mask)
+
+    prefix_lse[mask_prefix] = float("inf")
+    suffix_lse[mask_suffix] = float("inf")
+
+    # Other input tensors (need to be initialized but
+    # no actual calculation needed)
+    output = torch.zeros(
+        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
+    )
+    output_lse = torch.zeros(
+        (NUM_TOKENS, NUM_HEADS), dtype=torch.float32, device="cuda"
+    )
+    prefix_output = torch.randn(
+        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
+    )
+    suffix_output = torch.randn(
+        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
+    )
+
+    warmup_times = 2
+    repeat_times = 20
+
+    def perf_kernel_fn(
+        output_fn: torch.Tensor,
+        output_lse_fn: torch.Tensor,
+        kernel_fn: callable,
+        fn_type: str = "torch",
+    ):
+        # Avoid inplace inf -> -inf, we have to use prefix_lse
+        # and suffix_lse for other kernel.
+        if fn_type == "torch":
+            prefix_lse_ = prefix_lse.clone()
+            suffix_lse_ = suffix_lse.clone()
+        else:
+            prefix_lse_ = prefix_lse
+            suffix_lse_ = suffix_lse
+
+        total_time = 0
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        try:
+            for _ in range(warmup_times):
+                output_fn, output_lse_fn = kernel_fn(
+                    prefix_output,
+                    prefix_lse_,
+                    suffix_output,
+                    suffix_lse_,
+                    output_fn,
+                    output_lse_fn,
+                )
+            torch.cuda.synchronize()
+
+            for _ in range(repeat_times):
+                start.record()
+                output_fn, output_lse_fn = kernel_fn(
+                    prefix_output,
+                    prefix_lse_,
+                    suffix_output,
+                    suffix_lse_,
+                    output_fn,
+                    output_lse_fn,
+                )
+                end.record()
+                torch.cuda.synchronize()
+                total_time += start.elapsed_time(end)
+
+            avg_time = total_time / repeat_times
+            return avg_time, output_fn, output_lse_fn
+        except Exception as e:
+            return 0, output_fn, output_lse_fn
+
+    # 0. Run the Torch kernel
+    output_torch = output.clone()
+    output_lse_torch = output_lse.clone()
+    time_torch, output_torch, output_lse_torch = perf_kernel_fn(
+        output_torch, output_lse_torch, merge_state_torch, fn_type="torch"
+    )
+
+    # 1. Run the Triton kernel
+    output_ref_triton = output.clone()
+    output_lse_ref_triton = output_lse.clone()
+    time_triton, output_ref_triton, output_lse_ref_triton = perf_kernel_fn(
+        output_ref_triton,
+        output_lse_ref_triton,
+        merge_state_triton,
+        fn_type="triton",
+    )
+
+    # 2. Run the merge_state V2 kernel
+    output_v2 = output.clone()
+    output_lse_v2 = output_lse.clone()
+    time_v2, output_v2, output_lse_v2 = perf_kernel_fn(
+        output_v2, output_lse_v2, merge_state_v2, fn_type="cuda_v2"
+    )
+
+    # 3. Performance compare
+    improved = time_triton / time_v2
+    print(f"  Torch time: {time_torch:.6f}ms")
+    print(f" Triton time: {time_triton:.6f}ms")
+    print(f"CUDA v2 time: {time_v2:.6f}ms, Performance: {improved:.5f}x")
+    print("-" * 100)
+
+    # 4. Correctness compare
+    # Liger Kernel: Efficient Triton Kernels for LLM Training
+    # https://arxiv.org/pdf/2410.10989, 3.3 Correctness
+    # use rtol = 1e-2 for bfloat16.
+    rtol = 1e-2 if output_dtype == torch.bfloat16 else 1e-3
+
+    def diff(a: torch.Tensor, b: torch.Tensor):
+        max_diff = torch.max(torch.abs(a.float() - b.float()))
+        return max_diff
+
+    # Use Triton output as reference because we want to replace
+    # the Triton kernel with custom CUDA kernel for merge attn
+    # states operation.
+    output_ref = output_ref_triton
+    output_lse_ref = output_lse_ref_triton
+    torch.testing.assert_close(
+        output_v2.float(), output_ref.float(), atol=1e-3, rtol=rtol
+    )
+    print("Output all match, max abs diff:")
+    print(f"(Triton  vs Torch) : {diff(output_torch, output_ref)}")
+    print(f"(CUDA v2 vs Torch) : {diff(output_torch, output_v2)}")
+    print(f"(CUDA v2 vs Triton): {diff(output_ref, output_v2)}")
+    print("-" * 100)
+
+    torch.testing.assert_close(
+        output_lse_v2.float(), output_lse_ref.float(), atol=1e-3, rtol=rtol
+    )
+    print("Output LSE all match, max abs diff:")
+    print(f"(Triton  vs Torch) : {diff(output_lse_torch, output_lse_ref)}")
+    print(f"(CUDA v2 vs Torch) : {diff(output_lse_torch, output_lse_v2)}")
+    print(f"(CUDA v2 vs Triton): {diff(output_lse_ref, output_lse_v2)}")
+    print("-" * 100)
+
+    print(
+        "All output values test passed! All inf values "
+        "are correctly replaced with -inf."
+    )
+    print("-" * 100)
+
+    device = torch.cuda.get_device_name()
+    all_case_info.append(
+        (
+            NUM_TOKENS,
+            NUM_HEADS,
+            HEAD_SIZE,
+            output_dtype,
+            device,
+            time_torch,
+            time_triton,
+            time_v2,
+        )
+    )
+    if len(all_case_info) == (
+        len(NUM_BATCH_TOKENS) * len(HEAD_SIZES) * len(NUM_QUERY_HEADS) * len(DTYPES)
+    ):
+        generate_markdown_table()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_moe_align.py

@@ -0,0 +1,275 @@
+import itertools
+import sys
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import moe_align_block_size, moe_sum
+
+
+def is_hip() -> bool:
+    return torch.version.hip is not None
+
+
+_is_hip = is_hip()
+
+
+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,
+    )
+
+
+@pytest.mark.parametrize(
+    "block_size,num_tokens,topk,num_experts,pad_sorted_token_ids",
+    list(
+        itertools.product(
+            [32, 64, 128, 256],  # block_size
+            [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],  # num_tokens
+            [1, 2, 4, 8, 16, 32, 64],  # topk
+            [64, 160, 256, 257, 260, 264],  #  num_experts
+            [True, False],  # pad_sorted_token_ids
+        )
+    ),
+)
+def test_moe_align_block_size_compare_implementations(
+    block_size, num_tokens, topk, num_experts, pad_sorted_token_ids
+):
+
+    topk_ids = torch.argsort(torch.rand(num_tokens, num_experts, device="cuda"), dim=1)[
+        :, :topk
+    ]
+
+    max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1)
+    if topk_ids.numel() < num_experts + 1:
+        max_num_tokens_padded = topk_ids.numel() * block_size
+
+    sorted_ids_cuda = torch.empty(
+        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
+    )
+    if not pad_sorted_token_ids:
+        sorted_ids_cuda.fill_(topk_ids.numel())
+    max_num_m_blocks = max_num_tokens_padded // block_size
+    expert_ids_cuda = torch.zeros(
+        (max_num_m_blocks,), 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.empty(
+        num_experts + 2, dtype=torch.int32, device=topk_ids.device
+    )
+
+    sorted_ids_triton = torch.empty_like(sorted_ids_cuda)
+    sorted_ids_triton.fill_(topk_ids.numel())
+    expert_ids_triton = torch.zeros_like(expert_ids_cuda)
+    num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
+
+    moe_align_block_size(
+        topk_ids,
+        num_experts + 1,
+        block_size,
+        sorted_ids_cuda,
+        expert_ids_cuda,
+        num_tokens_post_pad_cuda,
+        cumsum_buffer,
+        pad_sorted_token_ids,
+    )
+
+    moe_align_block_size_triton(
+        topk_ids,
+        num_experts + 1,
+        block_size,
+        sorted_ids_triton,
+        expert_ids_triton,
+        num_tokens_post_pad_triton,
+    )
+
+    assert torch.allclose(expert_ids_cuda, expert_ids_triton, atol=0, rtol=0), (
+        f"Expert IDs mismatch for block_size={block_size}, "
+        f"num_tokens={num_tokens}, topk={topk}\n"
+        f"CUDA expert_ids: {expert_ids_cuda}\n"
+        f"Triton expert_ids: {expert_ids_triton}"
+    )
+
+    assert torch.allclose(
+        num_tokens_post_pad_cuda, num_tokens_post_pad_triton, atol=0, rtol=0
+    ), (
+        f"Num tokens post pad mismatch for block_size={block_size}, "
+        f"num_tokens={num_tokens}, topk={topk}\n"
+        f"CUDA num_tokens_post_pad: {num_tokens_post_pad_cuda}\n"
+        f"Triton num_tokens_post_pad: {num_tokens_post_pad_triton}"
+    )
+
+    # Select an expert to check
+    expert_idx = expert_ids_cuda.max().item()
+
+    # Get the first and last block id where expert_ids_cuda == expert_idx
+    matching_indices = torch.where(expert_ids_cuda == expert_idx)[0]
+    block_sorted_start = matching_indices[0].item() * block_size
+    block_sorted_end = min(
+        (matching_indices[-1].item() + 1) * block_size, num_tokens_post_pad_cuda.item()
+    )
+
+    selected_sorted_ids_cuda = sorted_ids_cuda[
+        block_sorted_start:block_sorted_end
+    ].sort()[0]
+    selected_sorted_ids_triton = sorted_ids_triton[
+        block_sorted_start:block_sorted_end
+    ].sort()[0]
+
+    assert torch.allclose(
+        selected_sorted_ids_cuda,
+        selected_sorted_ids_triton,
+        atol=0,
+        rtol=0,
+    ), (
+        f"Sorted IDs mismatch for block_size={block_size}, "
+        f"num_tokens={num_tokens}, topk={topk}\n"
+        f"CUDA sorted_ids: {selected_sorted_ids_cuda}\n"
+        f"Triton sorted_ids: {selected_sorted_ids_triton}"
+    )
+
+
+@pytest.mark.parametrize("m", [1, 33, 64, 222])
+@pytest.mark.parametrize("topk", [2, 6])
+@pytest.mark.parametrize("k", [128, 511, 1024])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.skipif(_is_hip, reason="Skip for AMD GPU")
+def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
+    input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
+    actual = torch.empty((m, k), device="cuda", dtype=dtype)
+
+    expected = input.sum(dim=1)
+    moe_sum(input, actual)
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_moe_fused_gate.py

@@ -0,0 +1,219 @@
+import sys
+from typing import Optional
+
+import pytest
+import torch
+from sgl_kernel import moe_fused_gate
+
+
+def biased_grouped_topk_impl(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    correction_bias: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: Optional[int] = None,
+    topk_group: Optional[int] = None,
+    num_fused_shared_experts: int = 0,
+    routed_scaling_factor: Optional[float] = None,
+    apply_routed_scaling_factor_on_output: Optional[bool] = False,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+
+    scores = gating_output.sigmoid()
+    num_token = scores.shape[0]
+    num_experts = scores.shape[1]
+    scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
+    group_scores = (
+        scores_for_choice.view(num_token, num_expert_group, -1)
+        .topk(2, dim=-1)[0]
+        .sum(dim=-1)
+    )  # [n, n_group]
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
+        1
+    ]  # [n, top_k_group]
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
+        .reshape(num_token, -1)
+    )  # [n, e]
+    tmp_scores = scores_for_choice.masked_fill(
+        ~score_mask.bool(), float("-inf")
+    )  # [n, e]
+
+    topk_excluding_shared = topk - num_fused_shared_experts
+    _, routed_topk_ids = torch.topk(
+        tmp_scores,
+        k=topk_excluding_shared,
+        dim=-1,
+        sorted=False,
+    )
+    routed_topk_weights = scores.gather(1, routed_topk_ids)
+
+    if num_fused_shared_experts > 0:
+        topk_ids = torch.empty(
+            (num_token, topk),
+            dtype=routed_topk_ids.dtype,
+            device=routed_topk_ids.device,
+        )
+        topk_weights = torch.empty(
+            (num_token, topk),
+            dtype=routed_topk_weights.dtype,
+            device=routed_topk_weights.device,
+        )
+        topk_ids[:, :topk_excluding_shared] = routed_topk_ids
+        topk_weights[:, :topk_excluding_shared] = routed_topk_weights
+
+        scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
+        routed_sum = routed_topk_weights.sum(dim=-1, keepdim=True)
+
+        for i in range(num_fused_shared_experts):
+            topk_ids[:, topk_excluding_shared + i] = num_experts + i
+            topk_weights[:, topk_excluding_shared + i] = routed_sum[:, 0] / scale
+    else:
+        topk_ids = routed_topk_ids
+        topk_weights = routed_topk_weights
+
+    if renormalize:
+        if num_fused_shared_experts > 0:
+            topk_weights_sum = topk_weights[:, :topk_excluding_shared].sum(
+                dim=-1, keepdim=True
+            )
+        else:
+            topk_weights_sum = topk_weights.sum(dim=-1, keepdim=True)
+        topk_weights = topk_weights / topk_weights_sum
+        if apply_routed_scaling_factor_on_output:
+            scale = (
+                1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
+            )
+            topk_weights *= scale
+
+    topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+    return topk_weights, topk_ids
+
+
+def biased_grouped_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    correction_bias: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: Optional[int] = None,
+    topk_group: Optional[int] = None,
+    num_fused_shared_experts: int = 0,
+    routed_scaling_factor: Optional[float] = None,
+    num_token_non_padded: Optional[torch.Tensor] = None,
+    apply_routed_scaling_factor_on_output: Optional[bool] = False,
+):
+    return biased_grouped_topk_impl(
+        hidden_states,
+        gating_output,
+        correction_bias,
+        topk,
+        renormalize,
+        num_expert_group,
+        topk_group,
+        num_fused_shared_experts=num_fused_shared_experts,
+        routed_scaling_factor=routed_scaling_factor,
+        apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+    )
+
+
+@pytest.mark.parametrize(
+    "seq_length",
+    list(range(1, 10))
+    + [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536],
+)
+@pytest.mark.parametrize(
+    "params",
+    [
+        (128, 4, 2, 4),
+        (256, 8, 4, 8),  # deepseek v3
+        (512, 16, 8, 16),
+    ],
+)
+@pytest.mark.parametrize("num_fused_shared_experts", [0, 1, 2])
+@pytest.mark.parametrize("apply_routed_scaling_factor_on_output", [False, True])
+def test_moe_fused_gate_combined(
+    seq_length, params, num_fused_shared_experts, apply_routed_scaling_factor_on_output
+):
+    num_experts, num_expert_group, topk_group, topk = params
+    dtype = torch.float32
+
+    torch.manual_seed(seq_length)
+    tensor = torch.rand((seq_length, num_experts), dtype=dtype, device="cuda")
+    scores = tensor.clone()
+    bias = torch.rand(num_experts, dtype=dtype, device="cuda")
+    topk = topk + num_fused_shared_experts
+
+    output, indices = moe_fused_gate(
+        tensor,
+        bias,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        topk=topk,
+        num_fused_shared_experts=num_fused_shared_experts,
+        routed_scaling_factor=2.5,
+        apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+    )
+    ref_output, ref_indices = biased_grouped_topk(
+        scores,
+        scores,
+        bias,
+        topk=topk,
+        renormalize=True,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        num_fused_shared_experts=num_fused_shared_experts,
+        routed_scaling_factor=2.5,
+        apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+    )
+
+    # When num_fused_shared_experts > 0, ignore the comparison of the last topk dimension
+    if num_fused_shared_experts > 0:
+        original_indices = indices.clone()
+        original_ref_indices = ref_indices.clone()
+
+        indices = indices[:, :-1]
+        ref_indices = ref_indices[:, :-1]
+
+        valid_min = num_experts
+        valid_max = num_experts + num_fused_shared_experts
+        shared_indices = original_indices[:, -1]
+        shared_ref_indices = original_ref_indices[:, -1]
+        if shared_indices is not None:
+            assert torch.all(
+                (shared_indices >= valid_min) & (shared_indices < valid_max)
+            ), f"Shared expert indices out of range: found values outside [{valid_min}, {valid_max})"
+        if shared_ref_indices is not None:
+            assert torch.all(
+                (shared_ref_indices >= valid_min) & (shared_ref_indices < valid_max)
+            ), f"Shared expert reference indices out of range: found values outside [{valid_min}, {valid_max})"
+
+    idx_check = torch.allclose(
+        ref_indices.sort()[0].to(torch.int32),
+        indices.sort()[0].to(torch.int32),
+        rtol=1e-04,
+        atol=1e-05,
+    )
+    output_check = torch.allclose(
+        ref_output.sort()[0].to(torch.float32),
+        output.sort()[0].to(torch.float32),
+        rtol=1e-02,
+        atol=1e-03,
+    )
+
+    assert idx_check, (
+        f"Indices mismatch at seq_length {seq_length}, dtype {dtype}, "
+        f"params {params}, num_fused_shared_experts {num_fused_shared_experts}"
+    )
+    assert output_check, (
+        f"Output mismatch at seq_length {seq_length}, dtype {dtype}, "
+        f"params {params}, num_fused_shared_experts {num_fused_shared_experts}"
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_moe_topk_sigmoid.py

@@ -0,0 +1,184 @@
+import itertools
+import sys
+
+import pytest
+import torch
+from sgl_kernel import topk_sigmoid
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 64, 128, 256],  # num_experts
+            [1, 2, 4],  # topk
+        )
+    ),
+)
+def test_topk_sigmoid(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+
+    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,
+    )
+
+    # Native torch implementation
+    sigmoid_output = torch.sigmoid(gating_output)
+    topk_weights_ref, topk_indices_ref = torch.topk(sigmoid_output, topk, dim=-1)
+
+    # Verify the top-k weights and indices match the torch native ones
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
+
+    assert torch.allclose(
+        topk_indices_ref.int(), topk_indices, atol=0, rtol=0
+    ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk, dtype",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 64, 128, 256],  # num_experts
+            [1, 2, 4],  # topk
+            [torch.float16, torch.bfloat16, torch.float32],  # dtype
+        )
+    ),
+)
+def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype):
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+
+    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,
+    )
+
+    topk_weights_ref = torch.empty(
+        (num_tokens, topk), dtype=torch.float32, device="cuda"
+    )
+    topk_indices_ref = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+
+    topk_sigmoid(
+        topk_weights_ref,
+        topk_indices_ref,
+        gating_output.float(),
+    )
+
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}"
+
+    assert torch.allclose(
+        topk_indices_ref.int(), topk_indices, atol=0, rtol=0
+    ), f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}"
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 64, 128, 256],  # num_experts
+            [1, 2, 4],  # topk
+        )
+    ),
+)
+def test_topk_sigmoid_renormalize(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.bfloat16, device="cuda"
+    )
+
+    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=True,
+    )
+
+    topk_weights_ref = torch.empty(
+        (num_tokens, topk), dtype=torch.float32, device="cuda"
+    )
+    topk_indices_ref = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    token_expert_indices_ref = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device="cuda"
+    )
+
+    topk_sigmoid(
+        topk_weights_ref,
+        topk_indices_ref,
+        gating_output,
+    )
+    topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
+
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}"
+
+    assert torch.allclose(
+        topk_indices_ref.int(), topk_indices, atol=0, rtol=0
+    ), f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}"
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 48, 64, 128, 256],  # num_experts
+            [1, 2, 4],  # topk
+        )
+    ),
+)
+def test_topk_sigmoid_renormalize_correction_bias(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+    correction_bias = torch.randn((num_experts), dtype=torch.float32, device="cuda")
+
+    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=True,
+        correction_bias=correction_bias,
+    )
+
+    # Native torch implementation
+    sigmoid_output = torch.sigmoid(gating_output)
+    sigmoid_scores = sigmoid_output.view(-1, num_experts) + correction_bias.unsqueeze(0)
+    _, topk_indices_ref = torch.topk(sigmoid_scores, k=topk, dim=-1)
+    topk_weights_ref = sigmoid_output.gather(1, topk_indices_ref)
+    topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
+
+    # Verify the top-k weights and indices match the torch native ones
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
+
+    assert torch.allclose(
+        topk_indices_ref.int(), topk_indices, atol=0, rtol=0
+    ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_moe_topk_softmax.py

@@ -0,0 +1,184 @@
+import itertools
+import sys
+
+import pytest
+import torch
+from sgl_kernel import topk_softmax
+
+
+def compare_topk_values(gating_output, topk_indices_ref, topk_indices):
+    values_ref = torch.gather(gating_output, 1, topk_indices_ref)
+    values = torch.gather(gating_output, 1, topk_indices)
+    return torch.equal(values_ref, values)
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [512],  # num_experts
+            [1, 2, 3, 4, 5, 8],  # topk
+        )
+    ),
+)
+def test_topkfast_softmax(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+
+    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_softmax(
+        topk_weights,
+        topk_indices,
+        gating_output,
+    )
+
+    # Native torch implementation
+    softmax_output = torch.softmax(gating_output, dim=-1)
+    topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1)
+
+    # Verify the top-k weights and indices match the torch native ones
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
+
+    assert compare_topk_values(
+        gating_output, topk_indices_ref.int(), topk_indices
+    ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 64, 128, 256],  # num_experts
+            [1, 2, 4],  # topk
+        )
+    ),
+)
+def test_topk_softmax(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+
+    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_softmax(
+        topk_weights,
+        topk_indices,
+        gating_output,
+    )
+
+    # Native torch implementation
+    softmax_output = torch.softmax(gating_output, dim=-1)
+    topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1)
+
+    # Verify the top-k weights and indices match the torch native ones
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
+
+    assert compare_topk_values(
+        gating_output, topk_indices_ref.int(), topk_indices
+    ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk, dtype",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 64, 128, 256, 512],  # num_experts
+            [1, 2, 4],  # topk
+            [torch.float16, torch.bfloat16, torch.float32],  # dtype
+        )
+    ),
+)
+def test_topk_softmax_dtype_regression(num_tokens, num_experts, topk, dtype):
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+
+    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_softmax(
+        topk_weights,
+        topk_indices,
+        gating_output,
+    )
+
+    topk_weights_ref = torch.empty(
+        (num_tokens, topk), dtype=torch.float32, device="cuda"
+    )
+    topk_indices_ref = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+
+    topk_softmax(
+        topk_weights_ref,
+        topk_indices_ref,
+        gating_output.float(),
+    )
+
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}"
+
+    assert compare_topk_values(
+        gating_output, topk_indices_ref.int(), topk_indices
+    ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(
+        itertools.product(
+            [1, 16, 128, 512, 1024, 2048],  # num_tokens
+            [4, 8, 16, 32, 64, 128, 256, 512],  # num_experts
+            [1, 2, 4],  # topk
+        )
+    ),
+)
+def test_topk_softmax_renormalize(num_tokens, num_experts, topk):
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.bfloat16, device="cuda"
+    )
+
+    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_softmax(
+        topk_weights,
+        topk_indices,
+        gating_output,
+        renormalize=True,
+    )
+
+    topk_weights_ref = torch.empty(
+        (num_tokens, topk), dtype=torch.float32, device="cuda"
+    )
+    topk_indices_ref = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    token_expert_indices_ref = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device="cuda"
+    )
+
+    topk_softmax(
+        topk_weights_ref,
+        topk_indices_ref,
+        gating_output,
+    )
+    topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
+
+    assert torch.allclose(
+        topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
+    ), f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}"
+
+    assert compare_topk_values(
+        gating_output, topk_indices_ref.int(), topk_indices
+    ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_mscclpp.py

@@ -0,0 +1,146 @@
+import multiprocessing as mp
+import os
+import socket
+import unittest
+from enum import IntEnum
+from typing import Any
+
+import sgl_kernel.allreduce as custom_ops
+import torch
+import torch.distributed as dist
+
+
+class MscclContextSelection(IntEnum):
+    MSCCL1SHOT1NODELL = 1
+    MSCCL1SHOT2NODELL = 2
+
+
+def _run_correctness_worker(world_size, rank, distributed_init_port, test_sizes):
+    device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
+    torch.cuda.set_device(device)
+    distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+    dist.init_process_group(
+        backend="nccl",
+        init_method=distributed_init_method,
+        rank=rank,
+        world_size=world_size,
+    )
+    group = dist.group.WORLD
+    cpu_group = torch.distributed.new_group(list(range(world_size)), backend="gloo")
+    if rank == 0:
+        unique_id = [custom_ops.mscclpp_generate_unique_id()]
+    else:
+        unique_id = [None]
+    dist.broadcast_object_list(
+        unique_id, src=0, device=torch.device("cpu"), group=cpu_group
+    )
+    unique_id = unique_id[0]
+    rank_to_node, rank_to_ib = list(range(world_size)), list(range(world_size))
+    for r in range(world_size):
+        rank_to_node[r] = r // 8
+        rank_to_ib[r] = rank % 8
+    MAX_BYTES = 2**20
+    scratch = torch.empty(
+        MAX_BYTES * 8, dtype=torch.bfloat16, device=torch.cuda.current_device()
+    )
+    put_buffer = torch.empty(
+        MAX_BYTES, dtype=torch.bfloat16, device=torch.cuda.current_device()
+    )
+    print(f"[{rank}] start mscclpp_context init")
+    nranks_per_node = torch.cuda.device_count()
+    selection = int(MscclContextSelection.MSCCL1SHOT1NODELL)
+    mscclpp_context = custom_ops.mscclpp_init_context(
+        unique_id,
+        rank,
+        world_size,
+        scratch,
+        put_buffer,
+        nranks_per_node,
+        rank_to_node,
+        rank_to_ib,
+        selection,
+    )
+    try:
+        test_loop = 10
+        for sz in test_sizes:
+            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+                if sz * dtype.itemsize > MAX_BYTES:
+                    continue
+                if rank == 0:
+                    print(f"mscclpp allreduce test sz {sz}, dtype {dtype}")
+                for _ in range(test_loop):
+                    inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device)
+                    inp1_ref = inp1.clone()
+                    out1 = torch.empty_like(inp1)
+                    custom_ops.mscclpp_allreduce(
+                        mscclpp_context, inp1, out1, nthreads=512, nblocks=21
+                    )
+                    dist.all_reduce(inp1_ref, group=group)
+                    torch.testing.assert_close(out1, inp1_ref)
+    finally:
+        dist.barrier(group=group)
+        dist.destroy_process_group(group=group)
+
+
+def get_open_port() -> int:
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("127.0.0.1", 0))
+            return s.getsockname()[1]
+    except OSError:
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+            s.bind(("::1", 0))
+            return s.getsockname()[1]
+
+
+def multi_process_parallel(
+    world_size: int, test_target: Any, target_args: tuple = ()
+) -> None:
+    mp.set_start_method("spawn", force=True)
+
+    procs = []
+    distributed_init_port = get_open_port()
+    for i in range(world_size):
+        proc_args = (world_size, i, distributed_init_port) + target_args
+        proc = mp.Process(target=test_target, args=proc_args, name=f"Worker-{i}")
+        proc.start()
+        procs.append(proc)
+
+    for i in range(world_size):
+        procs[i].join()
+        assert (
+            procs[i].exitcode == 0
+        ), f"Process {i} failed with exit code {procs[i].exitcode}"
+
+
+class TestMSCCLAllReduce(unittest.TestCase):
+    test_sizes = [
+        512,
+        2560,
+        4096,
+        5120,
+        7680,
+        32768,
+        262144,
+        524288,
+    ]
+    world_sizes = [8]
+
+    def test_correctness(self):
+        for world_size in self.world_sizes:
+            available_gpus = torch.cuda.device_count()
+            if world_size > available_gpus:
+                print(
+                    f"Skipping world_size={world_size}, found {available_gpus} and now ray is not supported here"
+                )
+                continue
+
+            print(f"Running test for world_size={world_size}")
+            multi_process_parallel(
+                world_size, _run_correctness_worker, target_args=(self.test_sizes,)
+            )
+            print(f"custom allreduce tp = {world_size}: OK")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/sgl-kernel/tests/test_norm.py

@@ -0,0 +1,144 @@
+# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/4e8eb1879f9c3ba6d75511e5893183bf8f289a62/tests/test_norm.py
+
+import sys
+
+import pytest
+import sgl_kernel
+import torch
+from sgl_kernel.utils import is_arch_support_pdl
+
+
+def llama_rms_norm(x, w, eps=1e-6):
+    orig_dtype = x.dtype
+    x = x.float()
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + eps)
+    x = x * w.float()
+    x = x.to(orig_dtype)
+    return x
+
+
+def gemma_rms_norm(x, w, eps=1e-6):
+    orig_dtype = x.dtype
+    x = x.float()
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + eps)
+    x = x * (1.0 + w.float())
+    x = x.to(orig_dtype)
+    return x
+
+
+def gemma_fused_add_rms_norm(x, residual, w, eps=1e-6):
+    orig_dtype = x.dtype
+    x = x + residual
+    residual = x
+    x = x.float()
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + eps)
+    x = x * (1.0 + w.float())
+    x = x.to(orig_dtype)
+    return x, residual
+
+
+def fused_add_rms_norm(x, residual, weight, eps):
+    orig_dtype = x.dtype
+    x = x.to(torch.float32)
+    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 + eps)
+    x = (x * weight.float()).to(orig_dtype)
+    return x, residual
+
+
+@pytest.mark.parametrize("batch_size", [1, 19, 99, 989])
+@pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384])
+@pytest.mark.parametrize("dtype", [torch.float16])
+@pytest.mark.parametrize("specify_out", [True, False])
+def test_norm(batch_size, hidden_size, dtype, specify_out):
+    x = torch.randn(batch_size, hidden_size).to(0).to(dtype)
+    w = torch.randn(hidden_size).to(0).to(dtype)
+
+    y_ref = llama_rms_norm(x, w)
+    enable_pdl = is_arch_support_pdl()
+    if specify_out:
+        y = torch.empty_like(x)
+        sgl_kernel.rmsnorm(x, w, out=y, enable_pdl=enable_pdl)
+    else:
+        y = sgl_kernel.rmsnorm(x, w, enable_pdl=enable_pdl)
+
+    torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("batch_size", [1, 19, 99, 989])
+@pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+def test_fused_add_rmsnorm(batch_size, hidden_size, dtype):
+    eps = 1e-6
+
+    x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
+    residual = torch.randn_like(x)
+    weight = torch.randn(hidden_size, dtype=dtype, device="cuda")
+
+    x_native, residual_native = fused_add_rms_norm(
+        x.clone(), residual.clone(), weight, eps
+    )
+
+    x_fused = x.clone()
+    residual_fused = residual.clone()
+    enable_pdl = is_arch_support_pdl()
+    sgl_kernel.fused_add_rmsnorm(
+        x_fused, residual_fused, weight, eps, enable_pdl=enable_pdl
+    )
+
+    torch.testing.assert_close(x_fused, x_native, rtol=1e-3, atol=1e-3)
+    torch.testing.assert_close(residual_fused, residual_native, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("batch_size", [1, 19, 99, 989])
+@pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384])
+@pytest.mark.parametrize("dtype", [torch.float16])
+@pytest.mark.parametrize("specify_out", [True, False])
+def test_gemma_norm(batch_size, hidden_size, dtype, specify_out):
+    x = torch.randn(batch_size, hidden_size).to(0).to(dtype)
+    w = torch.randn(hidden_size).to(0).to(dtype)
+
+    y_ref = gemma_rms_norm(x, w)
+    enable_pdl = is_arch_support_pdl()
+    if specify_out:
+        y = torch.empty_like(x)
+        sgl_kernel.gemma_rmsnorm(x, w, out=y, enable_pdl=enable_pdl)
+    else:
+        y = sgl_kernel.gemma_rmsnorm(x, w, enable_pdl=enable_pdl)
+
+    torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("batch_size", [1, 19, 99, 989])
+@pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384])
+@pytest.mark.parametrize("dtype", [torch.float16])
+def test_gemma_fused_add_rmsnorm(batch_size, hidden_size, dtype):
+    eps = 1e-6
+
+    x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
+    residual = torch.randn_like(x)
+    weight = torch.randn(hidden_size, dtype=dtype, device="cuda")
+
+    x_native, residual_native = gemma_fused_add_rms_norm(
+        x.clone(), residual.clone(), weight, eps
+    )
+
+    x_fused = x.clone()
+    residual_fused = residual.clone()
+    enable_pdl = is_arch_support_pdl()
+    sgl_kernel.gemma_fused_add_rmsnorm(
+        x_fused, residual_fused, weight, eps, enable_pdl=enable_pdl
+    )
+
+    torch.testing.assert_close(x_fused, x_native, rtol=1e-3, atol=1e-3)
+    torch.testing.assert_close(residual_fused, residual_native, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_per_token_group_quant_8bit.py

@@ -0,0 +1,186 @@
+import itertools
+import os
+import sys
+import time
+from pathlib import Path
+
+import pytest
+import torch
+from sgl_kernel.test_utils import (
+    assert_all_close_or_tiny_diff,
+    create_per_token_group_quant_test_data,
+)
+
+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 get_bool_env_var, is_hip
+
+_is_hip = is_hip()
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+configs = list(
+    itertools.product(
+        [1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192],  # num_tokens
+        [128, 256, 384, 512, 1024, 1536, 1664, 2048, 4096, 7168, 16384],  # hidden_dim
+        [16, 32, 64, 128],  # group_size
+        [None],  # num_ranks
+        [fp8_type_, torch.int8],  # dtype
+        [
+            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, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
+        # TODO support more
+        [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",
+            ),
+        ],
+    )
+)
+
+
+@pytest.mark.parametrize(
+    "num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags", configs
+)
+def test_per_token_group_quant_with_column_major(
+    num_tokens,
+    hidden_dim,
+    group_size,
+    num_ranks,
+    dst_dtype,
+    flags,
+):
+    print(
+        f"{num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=}"
+    )
+
+    arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
+    if flags["scale_ue8m0"] and (arch_major <= 9):
+        pytest.skip("Only Blackwell need ue8m0 fusion")
+        return
+
+    if (flags["scale_ue8m0"] and (group_size != 128)) or (
+        (dst_dtype == torch.int8) and flags["column_major_scales"]
+    ):
+        pytest.skip()
+        return
+
+    x, masked_m = create_per_token_group_quant_test_data(
+        num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
+    )
+
+    # print("hack data!!!")
+    # x = torch.full_like(x, fill_value=100)
+
+    execute_kwargs = dict(
+        x=x,
+        masked_m=masked_m,
+        group_size=group_size,
+        eps=1e-10,
+        dst_dtype=dst_dtype,
+        **{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
+    )
+
+    def _postprocess(x_q, x_s):
+        if masked_m is not None:
+            print(f"Mask tokens after {masked_m} to be zero")
+            for i in range(len(masked_m)):
+                x_q[i, masked_m[i] :, :] = 0
+                x_s[i, masked_m[i] :, :] = 0
+        return x_q, x_s
+
+    x_q_triton, x_s_triton = _postprocess(
+        *triton_per_token_group_quant_8bit(**execute_kwargs)
+    )
+    x_q_sglang, x_s_sglang = _postprocess(
+        *sglang_per_token_group_quant_8bit(**execute_kwargs, enable_v2=True)
+    )
+
+    try:
+        assert_all_close_or_tiny_diff(x_q_triton, x_q_sglang)
+        torch.testing.assert_close(
+            x_s_triton.contiguous(),
+            x_s_sglang.contiguous(),
+            rtol=1e-3,
+            atol=1e-5,
+            msg=lambda message: message + f" {x_s_triton=} {x_s_sglang=}",
+        )
+    except AssertionError:
+        print(
+            f"{x.shape=} {x_q_triton.shape=} {x_s_triton.shape=} {x_q_sglang.shape=} {x_s_sglang.shape=}"
+        )
+        print(f"{x=}")
+        print(f"{masked_m=}")
+        print(f"{x_q_triton=}")
+        print(f"{x_s_triton=}")
+        print(f"{x_q_sglang=}")
+        print(f"{x_s_sglang=}")
+
+        raise
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_per_token_quant_fp8.py

@@ -0,0 +1,58 @@
+import itertools
+import sys
+from typing import Optional, Tuple
+
+import pytest
+import torch
+from sgl_kernel import sgl_per_token_quant_fp8
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+
+def torch_per_token_quant_fp8(tensor, inv_scale):
+    # The reference implementation that fully aligns to
+    # the kernel being tested.
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    inv_scale = inv_scale.view(-1, 1)
+    scale = inv_scale.reciprocal()
+    qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    return qweight
+
+
+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)
+    scale = scale.reshape(-1, 1)
+
+    return output, scale
+
+
+@pytest.mark.parametrize(
+    "num_tokens,hidden_dim",
+    list(itertools.product([128, 256, 512], [512, 1076, 1368, 2048, 4096])),
+)
+def test_per_token_quant_compare_implementations(
+    num_tokens: int,
+    hidden_dim: int,
+):
+    device = torch.device("cuda")
+    x = torch.rand((num_tokens, hidden_dim), dtype=torch.float16, device=device)
+
+    sglang_out, sglang_scale = sglang_per_token_quant_fp8(x)
+    torch_out = torch_per_token_quant_fp8(x, sglang_scale)
+
+    torch.testing.assert_close(
+        sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_qserve_w4a8_per_chn_gemm.py

@@ -0,0 +1,120 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import qserve_w4a8_per_chn_gemm
+
+
+# Adapted from https://github.com/mit-han-lab/omniserve/blob/main/omniserve/modeling/layers/quantized_linear/w4a8_linear.py
+def convert_to_qserve_format(qweight, scale, zero):
+    assert qweight.min() >= 0 and qweight.max() <= 15, "Quantized weight out of range"
+    in_features = qweight.shape[1]
+    out_features = qweight.shape[0]
+    assert in_features % 32 == 0, "Input features must be divisible by 32"
+    assert out_features % 32 == 0, "Output features must be divisible by 32"
+
+    # ---- Repack the weight ---- #
+    # pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4)
+    qweight_unpack_reorder = (
+        qweight.reshape(
+            out_features // 32,
+            2,
+            2,
+            8,
+            in_features // 32,
+            2,
+            4,
+            4,
+        )
+        .permute(0, 4, 3, 6, 1, 5, 2, 7)
+        .contiguous()
+    )
+    qweight_unpack_reorder = (
+        qweight_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, 4)
+        .contiguous()
+        .to(torch.int8)
+    )
+    # B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous()
+    # [16, 0, 17, 1, ...]
+    qweight_unpack_repacked = (
+        qweight_unpack_reorder[..., 1] << 4
+    ) + qweight_unpack_reorder[..., 0]
+    qweight_unpack_repacked = qweight_unpack_repacked.reshape(
+        out_features // 32, in_features // 32, 32, 16
+    )
+    qweight_unpack_repacked = qweight_unpack_repacked.reshape(
+        out_features, in_features // 2
+    ).contiguous()
+
+    # ---- Pack the scales ---- #
+    scale = scale.reshape(out_features).to(torch.float16).contiguous()
+    szero = zero.reshape(out_features).to(torch.float16).contiguous() * scale
+
+    return qweight_unpack_repacked, scale, szero
+
+
+# INT4 Quantization
+def asym_quantize_tensor(tensor):
+    tensor_min = tensor.min(dim=-1, keepdim=True)[0]
+    tensor_max = tensor.max(dim=-1, keepdim=True)[0]
+    q_min = 0
+    q_max = 15
+    tensor_scale = (tensor_max - tensor_min) / (q_max - q_min)
+    tensor_zero = q_min - torch.round(tensor_min / tensor_scale)
+    tensor_q = torch.clamp(
+        torch.round(tensor / tensor_scale) + tensor_zero, q_min, q_max
+    ).to(torch.int8)
+    return tensor_q, tensor_scale.to(torch.float16), tensor_zero.to(torch.int8)
+
+
+# INT8 Quantization
+def sym_quantize_tensor(tensor):
+    tensor_scale = tensor.abs().max(dim=-1, keepdim=True)[0] / 127
+    tensor_q = torch.clamp(torch.round(tensor / tensor_scale), -128, 127).to(torch.int8)
+    return tensor_q, tensor_scale.to(torch.float16)
+
+
+def torch_w4a8_per_chn_gemm(a, b, a_scale, b_scale, b_zero, out_dtype):
+    print(a.shape)
+    print(b.shape)
+    print(b_zero.shape)
+    o = torch.matmul(
+        a.to(torch.float16), (b.to(torch.float16) - b_zero.to(torch.float16)).t()
+    )
+    o = o * a_scale.view(-1, 1) * b_scale.view(1, -1)
+    return o.to(out_dtype)
+
+
+def _test_accuracy_once(M, N, K, out_dtype, device):
+    # to avoid overflow, multiply 0.01
+    a = torch.randn((M, K), device=device, dtype=torch.float32) * 0.01
+    b = torch.randn((N, K), device=device, dtype=torch.float32) * 0.01
+
+    # symmetric quantize a
+    a_q, a_scale = sym_quantize_tensor(a)
+    # asymmetric quantize b
+    b_q, b_scale, b_zero = asym_quantize_tensor(b)
+    # convert to qserve format
+    b_q_format, b_scale_format, b_szero_format = convert_to_qserve_format(
+        b_q, b_scale, b_zero
+    )
+
+    # cal sum of every row of a
+    a_sum = a.sum(dim=-1, keepdim=True).to(torch.float16)
+    out = qserve_w4a8_per_chn_gemm(
+        a_q, b_q_format, b_scale_format, a_scale, b_szero_format, a_sum
+    )
+    ref_out = torch_w4a8_per_chn_gemm(a_q, b_q, a_scale, b_scale, b_zero, out_dtype)
+    torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-2)
+
+
+@pytest.mark.parametrize("M", [1, 16, 32, 64, 128, 512, 1024, 4096, 8192])
+@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("out_dtype", [torch.float16])
+def test_accuracy(M, N, K, out_dtype):
+    _test_accuracy_once(M, N, K, out_dtype, "cuda")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_qserve_w4a8_per_group_gemm.py

@@ -0,0 +1,185 @@
+import sys
+
+import pytest
+import torch
+from sgl_kernel import qserve_w4a8_per_group_gemm
+
+
+# Adapted from https://github.com/mit-han-lab/omniserve/blob/main/omniserve/modeling/layers/quantized_linear/w4a8_linear.py
+def convert_to_qserve_format(qweight, chn_scale, scale_i8, zero_i8, group_size):
+    assert qweight.min() >= 0 and qweight.max() <= 15, "Quantized weight out of range"
+    in_features = qweight.shape[1]
+    out_features = qweight.shape[0]
+    assert in_features % 32 == 0, "Input features must be divisible by 32"
+    assert out_features % 32 == 0, "Output features must be divisible by 32"
+    assert group_size == 128, "Group size must be 128"
+    assert (
+        in_features % group_size == 0
+    ), "Input features must be divisible by group_size"
+
+    # ---- Repack the weight ---- #
+    # pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4)
+    qweight_unpack_reorder = (
+        qweight.reshape(
+            out_features // 32,
+            2,
+            2,
+            8,
+            in_features // 32,
+            2,
+            4,
+            4,
+        )
+        .permute(0, 4, 3, 6, 1, 5, 2, 7)
+        .contiguous()
+    )
+    qweight_unpack_reorder = (
+        qweight_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, 4)
+        .contiguous()
+        .to(torch.int8)
+    )
+    # B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous()
+    # [16, 0, 17, 1, ...]
+    qweigth_unpack_repacked = (
+        qweight_unpack_reorder[..., 1] << 4
+    ) + qweight_unpack_reorder[..., 0]
+    qweigth_unpack_repacked = qweigth_unpack_repacked.reshape(
+        out_features // 32, in_features // 32, 32, 16
+    )
+    qweigth_unpack_repacked = qweigth_unpack_repacked.reshape(
+        out_features, in_features // 2
+    )
+
+    # ---- Pack the scales ---- #
+    chn_scale = chn_scale.reshape(out_features)
+
+    scale_i8 = (
+        scale_i8.reshape(out_features, in_features // group_size)
+        .transpose(0, 1)
+        .contiguous()
+    )
+    scale_i8 = scale_i8.reshape(in_features // group_size, out_features // 32, 32)
+    scale_i8 = (
+        scale_i8.reshape(in_features // group_size, out_features // 32, 4, 8)
+        .transpose(-2, -1)
+        .contiguous()
+    )
+    scale_i8 = scale_i8.reshape(in_features // group_size, out_features).contiguous()
+
+    # ---- Pack the zeros ---- #
+    zero_i8 = -zero_i8
+    # zero_i8 = zero_i8.int()  # convert to 2-complement
+
+    zero_i8 = (
+        zero_i8.reshape(out_features, in_features // group_size)
+        .transpose(0, 1)
+        .contiguous()
+    )
+    zero_i8 = zero_i8.reshape(in_features // group_size, out_features // 32, 32)
+    # for the last dimension, organize as 0, 8, 16, 24, 1, 9, 17, 25, ... following the requirement of tensor core gemm
+    zero_i8 = (
+        zero_i8.reshape(in_features // group_size, out_features // 32, 4, 8)
+        .transpose(-2, -1)
+        .contiguous()
+    )
+    zero_i8 = (
+        zero_i8.reshape(in_features // group_size, out_features).contiguous() * scale_i8
+    )
+
+    return qweigth_unpack_repacked, chn_scale, scale_i8, zero_i8
+
+
+# Progressive Group INT4 Quantization
+def progressive_group_quantize_tensor(tensor, group_size):
+    assert group_size == 128, "Group size must be 128"
+    assert (
+        tensor.shape[-1] % group_size == 0
+    ), "Input features must be divisible by group_size"
+    # Channel scale
+    # NOTE(HandH1998): use protective quantization range
+    chn_scale = tensor.abs().max(dim=-1, keepdim=True)[0] / 119
+    tensor_i8 = torch.clamp(torch.round(tensor / chn_scale), -119, 119)
+
+    # Group scale
+    tensor_i8 = tensor_i8.reshape(-1, group_size)
+    tensor_i8_min = tensor_i8.min(dim=-1, keepdim=True)[0]
+    tensor_i8_max = tensor_i8.max(dim=-1, keepdim=True)[0]
+    q_min = 0
+    q_max = 15
+    scale_i8 = torch.round((tensor_i8_max - tensor_i8_min) / (q_max - q_min))
+    zero_i8 = q_min - torch.round(tensor_i8_min / scale_i8)
+    tensor_q = (
+        torch.clamp(torch.round(tensor_i8 / scale_i8) + zero_i8, q_min, q_max)
+        .reshape(tensor.shape[0], -1)
+        .to(torch.int8)
+    )
+    return (
+        tensor_q,
+        chn_scale.to(torch.float16),
+        scale_i8.reshape(tensor.shape[0], -1).to(torch.int8),
+        zero_i8.reshape(tensor.shape[0], -1).to(torch.int8),
+    )
+
+
+# INT8 Quantization
+def sym_quantize_tensor(tensor):
+    tensor_scale = tensor.abs().max(dim=-1, keepdim=True)[0] / 127
+    tensor_q = torch.clamp(torch.round(tensor / tensor_scale), -128, 127).to(torch.int8)
+    return tensor_q, tensor_scale.to(torch.float16)
+
+
+def torch_w4a8_per_group_gemm(
+    a, b, a_scale, b_chn_scale, b_scale_i8, b_zero_i8, group_size, out_dtype
+):
+    assert group_size == 128, "Group size must be 128"
+    b_dq = (
+        b.reshape(-1, group_size).to(torch.float32)
+        - b_zero_i8.reshape(-1, 1).to(torch.float32)
+    ) * b_scale_i8.reshape(-1, 1).to(torch.float32)
+    b_dq = b_dq.reshape(b.shape[0], b.shape[1])
+    o = torch.matmul(a.to(torch.float32), b_dq.t())
+    o = o * a_scale.view(-1, 1) * b_chn_scale.view(1, -1)
+    return o.to(out_dtype)
+
+
+def _test_accuracy_once(M, N, K, group_size, out_dtype, device):
+    # to avoid overflow, multiply 0.01
+    a = torch.randn((M, K), device=device, dtype=torch.float32) * 0.01
+    b = torch.randn((N, K), device=device, dtype=torch.float32) * 0.01
+
+    # symmetric quantize a
+    a_q, a_scale = sym_quantize_tensor(a)
+    # asymmetric quantize b
+    b_q, b_chn_scale, b_scale_i8, b_zero_i8 = progressive_group_quantize_tensor(
+        b, group_size
+    )
+    # convert to qserve format
+    b_q_format, b_chn_scale_format, b_scale_i8_format, b_zero_i8_format = (
+        convert_to_qserve_format(b_q, b_chn_scale, b_scale_i8, b_zero_i8, group_size)
+    )
+
+    out = qserve_w4a8_per_group_gemm(
+        a_q,
+        b_q_format,
+        b_zero_i8_format,
+        b_scale_i8_format,
+        b_chn_scale_format,
+        a_scale,
+    )
+    ref_out = torch_w4a8_per_group_gemm(
+        a_q, b_q, a_scale, b_chn_scale, b_scale_i8, b_zero_i8, group_size, out_dtype
+    )
+    torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-5)
+
+
+@pytest.mark.parametrize("M", [1, 16, 32, 64, 128, 512, 1024, 4096, 8192])
+@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
+@pytest.mark.parametrize("group_size", [128])
+@pytest.mark.parametrize("out_dtype", [torch.float16])
+def test_accuracy(M, N, K, group_size, out_dtype):
+    _test_accuracy_once(M, N, K, group_size, out_dtype, "cuda")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_sampling.py

@@ -0,0 +1,188 @@
+# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/93e1a2634e22355b0856246b032b285ad1d1da6b/tests/test_sampling.py
+
+import sys
+
+import flashinfer.sampling
+import pytest
+import sgl_kernel
+import torch
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("k", [100])
+@pytest.mark.parametrize("p", [0.1, 0.5])
+def test_top_k_top_p_sampling_from_probs_logits_top_k_first_alignment(
+    batch_size, vocab_size, k, p
+):
+    torch.manual_seed(42)
+    logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5
+    generator_logits = torch.Generator("cuda:0")
+    generator_probs = generator_logits.clone_state()
+    samples = flashinfer.sampling.top_k_top_p_sampling_from_logits(
+        logits, k, p, filter_apply_order="top_k_first", generator=generator_logits
+    )
+    samples_ref = flashinfer.sampling.top_k_top_p_sampling_from_probs(
+        torch.softmax(logits, dim=-1),
+        k,
+        p,
+        filter_apply_order="top_k_first",
+        generator=generator_probs,
+    )
+    assert torch.all(samples == samples_ref)
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("k", [100])
+@pytest.mark.parametrize("p", [0.1, 0.5])
+def test_top_k_top_p_sampling_from_probs_logits_joint_alignment(
+    batch_size, vocab_size, k, p
+):
+    torch.manual_seed(42)
+    logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5
+    generator_logits = torch.Generator("cuda:0")
+    generator_probs = generator_logits.clone_state()
+    samples = flashinfer.sampling.top_k_top_p_sampling_from_logits(
+        logits, k, p, filter_apply_order="joint", generator=generator_logits
+    )
+    samples_ref = flashinfer.sampling.top_k_top_p_sampling_from_probs(
+        torch.softmax(logits, dim=-1),
+        k,
+        p,
+        filter_apply_order="joint",
+        generator=generator_probs,
+    )
+    assert torch.all(samples == samples_ref)
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("p", [0.1, 0.5])
+def test_top_k_top_p_joint_sampling_from_probs(batch_size, vocab_size, p):
+    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")
+    eps = 1e-4
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    # top-p mask
+    sorted_prob, indices = torch.sort(normalized_prob, descending=False)
+    cdf = torch.cumsum(sorted_prob, dim=-1)
+    mask_top_p = torch.zeros(batch_size, vocab_size, dtype=torch.int32, device="cuda:0")
+    mask_top_p.scatter_add_(1, indices, (cdf > (1 - p) - eps).int())
+    # top-k mask
+    sorted_prob, _ = torch.sort(normalized_prob, descending=True)
+    pivot = sorted_prob[:, k - 1]
+    mask_top_k = (normalized_prob >= pivot.unsqueeze(-1)).int()
+    # overall mask
+    mask = torch.minimum(mask_top_p, mask_top_k)
+    top_p_tensor = torch.full((batch_size,), p, device="cuda:0")
+    top_k_tensor = torch.full((batch_size,), k, device="cuda:0")
+
+    num_trails = 1000
+    for _ in range(num_trails):
+        samples = flashinfer.sampling.top_k_top_p_sampling_from_probs(
+            normalized_prob,
+            top_k_tensor,
+            top_p_tensor,
+            filter_apply_order="joint",
+        )
+        assert torch.all(samples < vocab_size) and torch.all(samples >= 0)
+        assert torch.all(mask[torch.arange(batch_size), samples] == 1), normalized_prob[
+            torch.arange(batch_size), samples
+        ]
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("p", [0.1, 0.5, 0.9])
+def test_top_p_renorm_probs(batch_size, vocab_size, p):
+    torch.manual_seed(42)
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    sorted_prob, indices = torch.sort(normalized_prob, descending=False)
+    cdf = torch.cumsum(sorted_prob, dim=-1)
+    mask = torch.zeros(batch_size, vocab_size, dtype=torch.int32, device="cuda:0")
+    mask.scatter_add_(1, indices, (cdf >= (1 - p)).int())
+    renorm_prob_ground_truth = normalized_prob.clone()
+    renorm_prob_ground_truth[mask == 0] = 0
+    renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum(
+        dim=-1, keepdim=True
+    )
+
+    renorm_prob = sgl_kernel.top_p_renorm_prob(normalized_prob, p)
+    torch.testing.assert_close(
+        renorm_prob_ground_truth,
+        renorm_prob,
+        rtol=1e-3,
+        atol=1e-3,
+    )
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("k", [10, 100, 500])
+def test_top_k_renorm_probs(batch_size, vocab_size, k):
+    if k > vocab_size:
+        pytest.skip("k should be less than vocab_size")
+    torch.manual_seed(42)
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    sorted_prob, _ = torch.sort(normalized_prob, descending=True)
+    pivot = sorted_prob[:, k - 1]
+    mask = (normalized_prob >= pivot.unsqueeze(-1)).int()
+    renorm_prob_ground_truth = normalized_prob.clone()
+    renorm_prob_ground_truth[mask == 0] = 0
+    renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum(
+        dim=-1, keepdim=True
+    )
+
+    renorm_prob = sgl_kernel.top_k_renorm_prob(normalized_prob, k)
+    for i in range(batch_size):
+        torch.testing.assert_close(
+            renorm_prob_ground_truth[i],
+            renorm_prob[i],
+            rtol=1e-3,
+            atol=1e-3,
+        )
+
+
+@pytest.mark.parametrize("batch_size", [1, 99, 989])
+@pytest.mark.parametrize("vocab_size", [111, 32000, 128256])
+@pytest.mark.parametrize("p", [0.05, 0.1, 0.2, 0.7, 1])
+def test_min_p_sampling(batch_size, vocab_size, p):
+    torch.manual_seed(42)
+    pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0")
+    normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
+    sorted_prob, indices = torch.sort(normalized_prob, descending=False)
+    # scale min-p
+    top_probs = sorted_prob[:, -1].unsqueeze(-1)
+    scaled_p = p * top_probs
+    # min-p mask
+    mask = torch.zeros(batch_size, vocab_size, dtype=torch.int32, device="cuda:0")
+    mask.scatter_add_(1, indices, (sorted_prob >= scaled_p).int())
+    min_p_tensor = torch.full((batch_size,), p, device="cuda:0")
+
+    num_trails = 1000
+    for _ in range(num_trails):
+        samples = flashinfer.sampling.min_p_sampling_from_probs(
+            normalized_prob,
+            min_p_tensor,
+        )
+
+        assert torch.all(mask[torch.arange(batch_size), samples] == 1), samples[
+            torch.nonzero(mask[torch.arange(batch_size), samples] == 0)
+        ]
+
+        assert torch.all(mask[torch.arange(batch_size), samples] == 1), samples[
+            torch.nonzero(mask[torch.arange(batch_size), samples] == 0)
+        ]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_topk.py

@@ -0,0 +1,253 @@
+import sys
+from typing import Any, Optional
+
+import pytest
+import torch
+from sgl_kernel import (
+    fast_topk_transform_fused,
+    fast_topk_transform_ragged_fused,
+    fast_topk_v2,
+)
+
+
+def _ref_torch_impl(
+    score: torch.Tensor,
+    seq_len: int,
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert score.dim() == 2
+    if row_starts is None:
+        return torch.topk(score[:, :seq_len], topk, dim=-1, sorted=False).indices
+    else:
+        ks = row_starts.cpu().tolist()
+        ke = (row_starts + seq_len).tolist()
+        scores = []
+        for i, (start, end) in enumerate(zip(ks, ke)):
+            scores.append(score[i, start:end].unsqueeze(0))
+        score = torch.cat(scores, dim=0)
+        return torch.topk(score, topk, dim=-1, sorted=False).indices
+
+
+def _ref_torch_transform_decode_impl(
+    score: torch.Tensor,
+    seq_len: int,
+    src_page_table: torch.Tensor,
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    batch_size, _ = score.shape
+    assert score.shape[0] == src_page_table.shape[0]
+    assert seq_len >= topk
+    indices = _ref_torch_impl(score, seq_len, topk, row_starts=row_starts)
+    topk_indices = torch.empty(
+        (batch_size, topk), dtype=torch.int32, device=score.device
+    )
+    for i in range(batch_size):
+        topk_indices[i] = src_page_table[i, indices[i]]
+    return topk_indices
+
+
+def _ref_torch_transform_ragged_impl(
+    score: torch.Tensor,
+    seq_len: int,
+    topk_indices_offset: torch.Tensor,
+    topk: int,
+    row_starts: torch.Tensor,
+) -> torch.Tensor:
+    assert score.shape[0] == topk_indices_offset.shape[0]
+    assert seq_len >= topk
+    indices = _ref_torch_impl(score, seq_len, topk, row_starts=row_starts)
+
+    mask = indices != -1
+    topk_indices_offset = topk_indices_offset.unsqueeze(1)
+    return torch.where(mask, indices + topk_indices_offset, indices)
+
+
+MAX_SEQ_LEN = 131072
+
+
+def assert_equal(
+    score: torch.Tensor,
+    indices_ref: torch.Tensor,
+    indices_our: torch.Tensor,
+    bs: int,
+    k: int,
+    seq_len: int,
+    topk_indices_offset: Optional[torch.Tensor] = None,
+    max_permit_error: int = 0,
+):
+    indices_our_cpu = indices_our.cpu().tolist()
+    indices_ref_cpu = indices_ref.cpu().tolist()
+
+    wrong_values = 0
+    for i in range(bs):
+        indices_ref_set_i = set(indices_ref_cpu[i])
+        indices_our_set_i = set(indices_our_cpu[i])
+        more = indices_our_set_i - indices_ref_set_i
+        less = indices_ref_set_i - indices_our_set_i
+        offset = topk_indices_offset[i].item() if topk_indices_offset is not None else 0
+        if len(more) > 0 or len(less) > 0:
+            # check whether more values are the same with less values
+            # if so, either one is acceptable, since their values are the same
+            more_values = sorted(score[i, idx - offset].item() for idx in more)
+            less_values = sorted(score[i, idx - offset].item() for idx in less)
+            if more_values != less_values:
+                wrong_values += len(more)
+                print(
+                    f"{bs=}, {k=}, {seq_len=}, {i=}, {more=}, {less=} failed, with {more_values=}, {less_values=}"
+                )
+        assert wrong_values <= max_permit_error, f"{wrong_values=}, {max_permit_error=}"
+
+
+@pytest.mark.parametrize("bs", [1, 132, 256, 4096])
+@pytest.mark.parametrize("k", [2048])  # we only support 2048 now
+@pytest.mark.parametrize("seq_len", [2048, 4096, 16384, 65536])
+@pytest.mark.parametrize("has_row_starts", [True, False])
+@torch.inference_mode()
+def test_topk_kernel(bs: int, k: int, seq_len: int, has_row_starts: bool) -> None:
+    torch.manual_seed(42)
+
+    stream = torch.cuda.Stream()
+    torch.cuda.set_stream(stream)
+    score = torch.randn(bs, MAX_SEQ_LEN, dtype=torch.float32, device="cuda")
+    lengths = torch.full((bs,), seq_len, dtype=torch.int32, device="cuda")
+
+    if has_row_starts:
+        row_starts = torch.randint(0, 2048, (bs,), dtype=torch.int32, device="cuda")
+    else:
+        row_starts = None
+
+    indices_ref = _ref_torch_impl(score, seq_len, k, row_starts=row_starts)
+    indices_our = fast_topk_v2(score, lengths, k, row_starts=row_starts)
+
+    # sort and compare
+    indices_ref = torch.sort(indices_ref, dim=-1).values
+    indices_our = torch.sort(indices_our, dim=-1).values
+
+    # Tests can pass with max_permit_error=3, set to 5 for safety
+    assert_equal(score, indices_ref, indices_our, bs, k, seq_len, max_permit_error=5)
+
+
+@pytest.mark.parametrize("bs", [1, 132, 256, 4096])
+@pytest.mark.parametrize("k", [2048])  # we only support 2048 now
+@pytest.mark.parametrize("seq_len", [2048, 4096, 16384, 65536])
+@pytest.mark.parametrize("mode", ["extend", "decode", "target_verify"])
+@torch.inference_mode()
+def test_topk_transform_kernel(bs: int, k: int, seq_len: int, mode: str) -> None:
+    torch.manual_seed(42)
+
+    stream = torch.cuda.Stream()
+    torch.cuda.set_stream(stream)
+
+    # NOTE: for decode, cumulative seqlens_q is just 0..=bs
+    # NOTE: since page table is arange, they equal topk indices
+    if mode == "decode":
+        step = 1
+    else:
+        step = 4 if bs % 4 == 0 else 1
+    num_tokens = bs
+    bs = bs // step
+
+    if mode == "extend":
+        row_starts = torch.randint(0, 2048, (bs,), dtype=torch.int32, device="cuda")
+    else:
+        row_starts = None
+
+    score = torch.randn(bs, MAX_SEQ_LEN, dtype=torch.float32, device="cuda")
+    lengths = torch.full((bs,), seq_len, dtype=torch.int32, device="cuda")
+    cu_seqlens_q = torch.arange(
+        0, num_tokens + 1, step=step, dtype=torch.int32, device="cuda"
+    )
+    src_page_table = torch.arange(0, seq_len, dtype=torch.int32, device="cuda")
+    src_page_table = src_page_table.unsqueeze(0).expand(bs, -1)
+
+    dst_page_table_ref = _ref_torch_transform_decode_impl(
+        score=score,
+        seq_len=seq_len,
+        src_page_table=src_page_table,
+        topk=k,
+        row_starts=row_starts,
+    )
+    dst_page_table_our = fast_topk_transform_fused(
+        score=score,
+        lengths=lengths,
+        page_table_size_1=src_page_table,
+        cu_seqlens_q=cu_seqlens_q,
+        topk=k,
+        row_starts=row_starts,
+    )
+
+    # sort and compare
+    dst_page_table_our = torch.sort(dst_page_table_our, dim=-1).values
+    dst_page_table_ref = torch.sort(dst_page_table_ref, dim=-1).values
+
+    assert_equal(
+        score,
+        dst_page_table_ref,
+        dst_page_table_our,
+        bs,
+        k,
+        seq_len,
+        max_permit_error=5,
+    )
+
+
+@pytest.mark.parametrize("bs", [1, 132, 256, 4096])
+@pytest.mark.parametrize("k", [2048])  # we only support 2048 now
+@pytest.mark.parametrize("seq_len", [2048, 4096, 16384, 65536])
+@pytest.mark.parametrize("has_row_starts", [True, False])
+@torch.inference_mode()
+def test_topk_transform_ragged_kernel(
+    bs: int, k: int, seq_len: int, has_row_starts: bool
+) -> None:
+    # Used in prefill only
+    torch.manual_seed(42)
+
+    stream = torch.cuda.Stream()
+    torch.cuda.set_stream(stream)
+    # bs: # of q tokens
+    score = torch.randn(bs, MAX_SEQ_LEN, dtype=torch.float32, device="cuda")
+    # kv_len
+    if has_row_starts:
+        row_starts = torch.randint(0, 2048, (bs,), dtype=torch.int32, device="cuda")
+    else:
+        row_starts = None
+    lengths = torch.full((bs,), seq_len, dtype=torch.int32, device="cuda")
+    topk_indices_offset = torch.randint(
+        0, 1024, (bs,), dtype=torch.int32, device="cuda"
+    )
+
+    dst_page_table_ref = _ref_torch_transform_ragged_impl(
+        score=score,
+        seq_len=seq_len,
+        topk_indices_offset=topk_indices_offset,
+        topk=k,
+        row_starts=row_starts,
+    )
+    dst_page_table_our = fast_topk_transform_ragged_fused(
+        score=score,
+        lengths=lengths,
+        topk_indices_offset=topk_indices_offset,
+        topk=k,
+        row_starts=row_starts,
+    )
+
+    # sort and compare
+    dst_page_table_our = torch.sort(dst_page_table_our, dim=-1).values
+    dst_page_table_ref = torch.sort(dst_page_table_ref, dim=-1).values
+
+    assert_equal(
+        score,
+        dst_page_table_ref,
+        dst_page_table_our,
+        bs,
+        k,
+        seq_len,
+        topk_indices_offset,
+        max_permit_error=5,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/test_torch_defaults_reset.py

@@ -0,0 +1,18 @@
+import sys
+
+import pytest
+import torch
+
+
+def test_change_torch_defaults():
+    torch.set_default_device("cpu:0")
+    torch.set_default_dtype(torch.float16)
+
+
+def test_check_torch_defaults():
+    assert torch.get_default_device() == torch.device("cpu")
+    assert torch.get_default_dtype() == torch.float32
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/sgl-kernel/tests/utils.py

@@ -0,0 +1,9 @@
+import torch
+
+
+def is_sm10x():
+    return torch.cuda.get_device_capability() >= (10, 0)
+
+
+def is_hopper():
+    return torch.cuda.get_device_capability() == (9, 0)