File "/public_sglang_ci/runner-l1e-gpu-4567/_work/sglang/sglang/python/sglang/srt/layers/moe/moe_runner/triton.py", line 40, in <module>
    from sgl_kernel import gelu_and_mul, moe_sum_reduce, silu_and_mul
ImportError: cannot import name 'moe_sum_reduce' from 'sgl_kernel' (/usr/local/lib/python3.10/dist-packages/sgl_kernel/__init__.py)
ERROR

[2025-09-29 02:36:08] INFO:     127.0.0.1:35010 - "POST /v1/chat/completions HTTP/1.1" 200 OK
[2025-09-29 02:36:08] INFO:     127.0.0.1:35022 - "POST /v1/chat/completions HTTP/1.1" 200 OK
[2025-09-29 02:36:08] INFO:     127.0.0.1:35044 - "POST /v1/chat/completions HTTP/1.1" 200 OK
[2025-09-29 02:36:08] INFO:     127.0.0.1:35068 - "POST /v1/chat/completions HTTP/1.1" 200 OK
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 64/64 [00:30<00:00,  2.12it/s]
Writing report to /tmp/mmlu_deepseek-ai_DeepSeek-Coder-V2-Lite-Instruct.html
{'other': np.float64(0.0), 'other:std': np.float64(0.0), 'score:std': np.float64(0.0), 'stem': np.float64(0.0), 'stem:std': np.float64(0.0), 'humanities': np.float64(0.0), 'humanities:std': np.float64(0.0), 'social_sciences': np.float64(0.0), 'social_sciences:std': np.float64(0.0), 'score': np.float64(0.0)}
Writing results to /tmp/mmlu_deepseek-ai_DeepSeek-Coder-V2-Lite-Instruct.json
Total latency: 30.266 s
Score: 0.000
E
======================================================================
ERROR: test_mgsm_en (__main__.TestHierarchicalMLA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/utils.py", line 2259, in retry
    return fn()
  File "/usr/local/lib/python3.10/dist-packages/sglang/test/test_utils.py", line 1437, in <lambda>
    lambda: super(CustomTestCase, self)._callTestMethod(method),
AssertionError: np.float64(0.0) not greater than 0.8

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/sglang/test/test_utils.py", line 1436, in _callTestMethod
    retry(
  File "/usr/local/lib/python3.10/dist-packages/sglang/srt/utils.py", line 2262, in retry
    raise Exception(f"retry() exceed maximum number of retries.")
Exception: retry() exceed maximum number of retries.

[2025-10-07 22:32:23] Decode batch. #running-req: 2, #token: 3489, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:32:36] Decode batch. #running-req: 2, #token: 3569, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:32:49] Decode batch. #running-req: 2, #token: 3649, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:33:01] Decode batch. #running-req: 2, #token: 3729, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:33:14] Decode batch. #running-req: 2, #token: 3809, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:33:27] Decode batch. #running-req: 2, #token: 3889, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:33:40] Decode batch. #running-req: 2, #token: 3969, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:33:53] Decode batch. #running-req: 2, #token: 4049, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:34:06] Decode batch. #running-req: 2, #token: 4129, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:34:19] Decode batch. #running-req: 2, #token: 4209, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:34:31] Decode batch. #running-req: 2, #token: 4289, token usage: 0.00, cuda graph: True, gen throughput (token/s): 6.23, #queue-req: 0,
[2025-10-07 22:34:37] INFO:     127.0.0.1:57532 - "POST /v1/chat/completions HTTP/1.1" 200 OK
[2025-10-07 22:34:37] INFO:     127.0.0.1:57650 - "POST /v1/chat/completions HTTP/1.1" 200 OK
100%|██████████| 250/250 [11:01<00:00,  2.65s/it]
.
----------------------------------------------------------------------
Ran 1 test in 743.456s

OK
Writing report to /tmp/mgsm_en_neuralmagic_DeepSeek-Coder-V2-Lite-Instruct-FP8.html
{'en': np.float64(0.864), 'en:std': np.float64(0.3427885645700568), 'group_latin': np.float64(0.864), 'group_latin:std': np.float64(0.3427885645700568), 'score:std': np.float64(0.3427885645700568), 'score': np.float64(0.864)}
Writing results to /tmp/mgsm_en_neuralmagic_DeepSeek-Coder-V2-Lite-Instruct-FP8.json
Total latency: 661.457 s
Score: 0.864

import os

import torch
import triton
import triton.language as tl
from sgl_kernel import moe_sum_reduce as moe_sum_reduce_cuda
from triton.testing import do_bench

@triton.jit
def _moe_sum_reduce_kernel(
    input_ptr,
    input_stride_0,
    input_stride_1,
    input_stride_2,
    output_ptr,
    output_stride_0,
    output_stride_1,
    token_num: int,
    topk_num: int,
    hidden_dim: int,
    routed_scaling_factor: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_DIM: tl.constexpr,
    NUM_STAGE: tl.constexpr,
):
    input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
    input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
    output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64)

    token_block_id = tl.program_id(0)
    dim_block_id = tl.program_id(1)

    offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM)

    mask_token = offs_token < token_num
    mask_dim = offs_dim < hidden_dim

    base_ptrs = input_ptr + offs_token[:, None] * input_stride_0 + offs_dim[None, :]

    accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32)
    for i in tl.range(0, topk_num, num_stages=NUM_STAGE):
        tile = tl.load(
            base_ptrs + i * input_stride_1,
            mask=mask_token[:, None] & mask_dim[None, :],
            other=0.0,
        )
        accumulator += tile.to(tl.float32)
    accumulator *= routed_scaling_factor

    # -------- Write back --------
    store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :]
    tl.store(
        store_ptrs,
        accumulator.to(input_ptr.dtype.element_ty),
        mask=mask_token[:, None] & mask_dim[None, :],
    )


# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py
def moe_sum_reduce_triton(
    input: torch.Tensor, output: torch.Tensor, routed_scaling_factor: float
):
    assert input.is_contiguous()
    assert output.is_contiguous()

    token_num, topk_num, hidden_dim = input.shape
    assert output.shape[0] == token_num and output.shape[1] == hidden_dim

    BLOCK_M = 1
    BLOCK_DIM = 2048
    NUM_STAGE = 1
    num_warps = 16

    grid = (
        triton.cdiv(token_num, BLOCK_M),
        triton.cdiv(hidden_dim, BLOCK_DIM),
    )

    _moe_sum_reduce_kernel[grid](
        input,
        *input.stride(),
        output,
        *output.stride(),
        token_num=token_num,
        topk_num=topk_num,
        hidden_dim=hidden_dim,
        routed_scaling_factor=routed_scaling_factor,
        BLOCK_M=BLOCK_M,
        BLOCK_DIM=BLOCK_DIM,
        NUM_STAGE=NUM_STAGE,
        num_warps=num_warps,
    )
    return


def compute_sum_scaled_baseline(
    x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float
) -> torch.Tensor:
    torch.sum(x, dim=1, out=out)
    out.mul_(routed_scaling_factor)
    return out


@torch.compile
def compute_sum_scaled_compiled(
    x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float
) -> torch.Tensor:
    torch.sum(x * routed_scaling_factor, dim=1, out=out)
    return out



def verify_correctness(num_tokens=1024, dtype=torch.bfloat16):
    x = torch.randn(num_tokens, 9, 4096, device="cuda", dtype=dtype)
    scaling_factor = 0.3

    out_baseline = torch.empty_like(x[:, 0])
    compute_sum_scaled_baseline(x, out_baseline, scaling_factor)

    out_compiled = torch.empty_like(out_baseline)
    compute_sum_scaled_compiled(x, out_compiled, scaling_factor)

    out_cuda = torch.empty_like(out_baseline)
    moe_sum_reduce_cuda(x, out_cuda, scaling_factor)

    triton_skipped = False
    # triton_skipped = dtype == torch.float64
    # if not triton_skipped:
    out_triton = torch.empty_like(out_baseline)
    moe_sum_reduce_triton(x, out_triton, scaling_factor)

    if dtype == torch.float64:
        atol, rtol = 1e-12, 1e-12
    elif dtype == torch.float32:
        atol, rtol = 1e-6, 1e-6
    else:  # bfloat16 / float16
        atol, rtol = 1e-2, 1e-2

    ok_compiled = torch.allclose(out_baseline, out_compiled, atol=atol, rtol=rtol)
    ok_cuda = torch.allclose(out_baseline, out_cuda, atol=atol, rtol=rtol)
    ok_triton = torch.allclose(out_baseline, out_triton, atol=atol, rtol=rtol)

    if ok_compiled and ok_triton and ok_cuda:
        msg = "✅ All implementations match"
        if triton_skipped:
            msg += " (Triton skipped for float64)"
        print(msg)
    else:
        print(ok_compiled)
        print(ok_triton)
        print(ok_cuda)
        print("❌ Implementations differ")
        print(
            f"Baseline vs Compiled: {(out_baseline - out_compiled).abs().max().item()}"
        )
        print(
            f"Baseline vs Triton: {(out_baseline - out_triton).abs().max().item()}"
        )
        print(f"Baseline vs Cuda: {(out_baseline - out_cuda).abs().max().item()}")


if __name__ == "__main__":

    print("Running correctness verification for float64...")
    verify_correctness(dtype=torch.float64)

Running correctness verification for float64...
True
False
True
❌ Implementations differ
Baseline vs Compiled: 1.3322676295501878e-15
Baseline vs Triton: 8.264526956125451e-07
Baseline vs Cuda: 1.7763568394002505e-15