[Inkling] silu_and_mul: replace helion kernels with plain Triton (#33903)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/pyproject.toml

@@ -33,7 +33,6 @@ dependencies = [
   "flash-attn-4>=4.0.0b18",
   "flashinfer_python[cu13]==0.6.15.post1", # keep it aligned with jit-cache version in Dockerfile
   "gguf",
-  "helion==1.4",
   "humming-kernels[cu13]==0.1.10",
   "interegular",
   "IPython",

python/pyproject_other.toml

@@ -32,7 +32,6 @@ runtime_base = [
   "einops",
   "fastapi",
   "gguf",
-  "helion==1.4",
   "interegular",
   "IPython",
   "llguidance>=1.7.6,<2.0.0",

python/sglang/kernels/ops/moe/inkling_moe.py

@@ -1,142 +1,195 @@
-from functools import partial
-
-import helion
-import helion.language as hl
 import torch
 import triton
 import triton.language as tl
 
 from sglang.kernels.jit.utils import is_arch_support_pdl
-from sglang.srt.layers.moe.moe_runner.triton_utils.helion_utils import (
-    get_model_depths,
-    helion_aot_autotune,
-)
 from sglang.srt.utils.common import is_sm121
 
 DEFAULT_BLOCK_SIZE = 4096
 BLOCK_SIZE_M = 128
 
 
-def silu_and_mul_key(
+@triton.jit(do_not_specialize=["M"])
-    gateup_output: torch.Tensor,
+def silu_and_mul_interleaved_kernel(
-    topk_weights: torch.Tensor | None,
+    gateup_out_ptr,  # type: ignore  # [M, 2N], rows are [g0, u0, g1, u1, ...]
-    out_dtype: object | None = None,
+    topk_weights_ptr,  # type: ignore  # [M] (unused when HAS_TOPK_WEIGHTS=False)
+    down_inp_ptr,  # type: ignore  # [M, N]
+    M,  # type: ignore
+    stride_gm,  # type: ignore
+    N: tl.constexpr,
+    HAS_TOPK_WEIGHTS: tl.constexpr,
+    INT64_INDEX: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
 ):
-    # Keep this stable across inputs for Helion AOT autotune.
+    # Flat grid, N-tile fastest: consecutive CTAs read consecutive spans of
-    del out_dtype
+    # the same rows, which measures faster than row-fastest rasterization.
-    return gateup_output.shape[1], (gateup_output.dtype,), (topk_weights is not None,)
+    pid = tl.program_id(0)
+    if INT64_INDEX:
+        pid = pid.to(tl.int64)
+    NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_N)
+    pid_m = pid // NUM_BLOCKS_N
+    pid_n = pid % NUM_BLOCKS_N
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+    offs_2n = pid_n * (2 * BLOCK_N) + tl.arange(0, 2 * BLOCK_N)
+    mask_2n = offs_2n < 2 * N
+    if BLOCK_M == 1:
+        # The grid covers M exactly; eliding the row mask measurably helps
+        # at bandwidth-bound shapes.
+        mask_mn = mask_n[None, :]
+        mask_m2n = mask_2n[None, :]
+    else:
+        mask_m = offs_m < M
+        mask_mn = mask_m[:, None] & mask_n[None, :]
+        mask_m2n = mask_m[:, None] & mask_2n[None, :]
+
+    # Contiguous load of the [g, u] pairs, deinterleaved in registers via
+    # tl.split; element-strided gather loads halve the effective bandwidth.
+    gateup = tl.load(
+        gateup_out_ptr + offs_m[:, None] * stride_gm + offs_2n[None, :],
+        mask=mask_m2n,
+        other=0.0,
+    )
+    gate, up = tl.split(tl.reshape(gateup, (BLOCK_M, BLOCK_N, 2)))
+    gate = gate.to(tl.float32)
+    up = up.to(tl.float32)
+
+    down_inp = gate * tl.sigmoid(gate) * up
+    if HAS_TOPK_WEIGHTS:
+        if BLOCK_M == 1:
+            weight_scale = tl.load(
+                topk_weights_ptr + offs_m, eviction_policy="evict_last"
+            ).to(tl.float32)
+        else:
+            weight_scale = tl.load(
+                topk_weights_ptr + offs_m,
+                mask=offs_m < M,
+                eviction_policy="evict_last",
+            ).to(tl.float32)
+        down_inp = down_inp * weight_scale[:, None]
+
+    offs_mn = offs_m[:, None] * N + offs_n[None, :]
+    tl.store(
+        down_inp_ptr + offs_mn,
+        down_inp.to(down_inp_ptr.dtype.element_ty),
+        mask=mask_mn,
+    )
 
 
-def silu_and_mul_inputs(sizes: list[int]):
+@triton.jit(do_not_specialize=["M"])
-    # Used only for Helion autotune input generation.
+def silu_and_mul_non_interleaved_kernel(
-    inputs = []
+    gateup_out_ptr,  # type: ignore  # [M, 2N], rows are [g0, ..., gN-1, u0, ..., uN-1]
-    numel = 2**30
+    topk_weights_ptr,  # type: ignore  # [M] (unused when HAS_TOPK_WEIGHTS=False)
-    with torch.device("cuda"):
+    down_inp_ptr,  # type: ignore  # [M, N]
-        for size in sizes:
+    M,  # type: ignore
-            x = torch.randn(numel // size, 2 * size, dtype=torch.bfloat16)
+    stride_gm,  # type: ignore
-            inputs.append((x, None, None))
+    N: tl.constexpr,
-            inputs.append((x, torch.randn(numel // size, dtype=torch.bfloat16), None))
+    HAS_TOPK_WEIGHTS: tl.constexpr,
-    return inputs
+    INT64_INDEX: tl.constexpr,
-
+    BLOCK_M: tl.constexpr,
-
+    BLOCK_N: tl.constexpr,
-@helion_aot_autotune(
-    "silu_and_mul_interleaved",
-    kernel_key=silu_and_mul_key,
-    primary_inputs=partial(silu_and_mul_inputs, sizes=[512, 2048, 48 * 96, 6144, 8192]),
-    secondary_inputs=partial(
-        silu_and_mul_inputs,
-        sizes=[512]
-        + [i * 96 for i in get_model_depths()]
-        + list(range(1024, 8192 + 1, 1024)),
-    ),
-)
-@helion.kernel(static_shapes=False)
-def _silu_and_mul_helion_interleaved_kernel(
-    gateup_output,
-    topk_weights: torch.Tensor | None = None,
-    out_dtype: hl.constexpr | None = None,
 ):
-    """
+    # Flat grid, N-tile fastest: consecutive CTAs read consecutive spans of
-    Interleaved version of silu_and_mul using Helion kernel.
+    # the same rows, which measures faster than row-fastest rasterization.
-    Input format: [gate[0], up[0], gate[1], up[1], ...]
+    pid = tl.program_id(0)
-    This matches the interleaved w13 weight format.
+    if INT64_INDEX:
-    """
+        pid = pid.to(tl.int64)
-    batch_size, hidden_size = gateup_output.shape
+    NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_N)
-    hidden_size = hl.specialize(hidden_size)
+    pid_m = pid // NUM_BLOCKS_N
-    assert hidden_size % 2 == 0, f"{hidden_size=}"
+    pid_n = pid % NUM_BLOCKS_N
 
-    half_hidden_size = hidden_size // 2
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    down_input = gateup_output.new_empty(
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-        batch_size, half_hidden_size, dtype=out_dtype or gateup_output.dtype
+    mask_n = offs_n < N
+    if BLOCK_M == 1:
+        # The grid covers M exactly; eliding the row mask measurably helps
+        # at bandwidth-bound shapes.
+        mask_mn = mask_n[None, :]
+    else:
+        mask_m = offs_m < M
+        mask_mn = mask_m[:, None] & mask_n[None, :]
+
+    # Cache hints match the tuned helion config (numerics-neutral): the gate
+    # stream stays in L2 while the up stream of the same rows loads.
+    offs_gate = offs_m[:, None] * stride_gm + offs_n[None, :]
+    gate = tl.load(
+        gateup_out_ptr + offs_gate,
+        mask=mask_mn,
+        other=0.0,
+        eviction_policy="evict_last",
+    ).to(tl.float32)
+    up = tl.load(
+        gateup_out_ptr + offs_gate + N,
+        mask=mask_mn,
+        other=0.0,
+        eviction_policy="evict_first",
+    ).to(tl.float32)
+
+    down_inp = gate * tl.sigmoid(gate) * up
+    if HAS_TOPK_WEIGHTS:
+        if BLOCK_M == 1:
+            weight_scale = tl.load(
+                topk_weights_ptr + offs_m, eviction_policy="evict_last"
+            ).to(tl.float32)
+        else:
+            weight_scale = tl.load(
+                topk_weights_ptr + offs_m,
+                mask=offs_m < M,
+                eviction_policy="evict_last",
+            ).to(tl.float32)
+        down_inp = down_inp * weight_scale[:, None]
+
+    offs_mn = offs_m[:, None] * N + offs_n[None, :]
+    tl.store(
+        down_inp_ptr + offs_mn,
+        down_inp.to(down_inp_ptr.dtype.element_ty),
+        mask=mask_mn,
     )
-    for batch_tile, hidden_tile in hl.tile([batch_size, half_hidden_size]):
-        gate_output = gateup_output[batch_tile, 2 * hidden_tile.index].to(torch.float32)
-        up_output = gateup_output[batch_tile, 2 * hidden_tile.index + 1].to(
-            torch.float32
-        )
-        silu_mul_output = gate_output * torch.sigmoid(gate_output) * up_output
-        if topk_weights is not None:
-            weight_scale = topk_weights[batch_tile, None].to(torch.float32)
-            silu_mul_output = silu_mul_output * weight_scale
-        down_input[batch_tile, hidden_tile] = silu_mul_output
-    return down_input
 
 
-@helion_aot_autotune(
+def _select_silu_and_mul_config(
-    "silu_and_mul",
+    M: int, N: int, interleaved: bool
-    kernel_key=silu_and_mul_key,
+) -> tuple[int, int, int]:
-    primary_inputs=partial(silu_and_mul_inputs, sizes=[512, 2048, 48 * 96, 6144, 8192]),
+    """(BLOCK_M, BLOCK_N, num_warps) heuristic from a GB300 (sm_100) sweep.
-    secondary_inputs=partial(
+
-        silu_and_mul_inputs,
+    One row per CTA keeps the grid large enough to fill the GPU at decode row
-        sizes=[512]
+    counts. BLOCK_N=512 avoids a quarter-masked tail block for N=1536-style
-        + [i * 96 for i in get_model_depths()]
+    widths; other widths prefer the widest contiguous load even when the tail
-        + list(range(1024, 8192 + 1, 1024)),
+    block is masked. The non-interleaved kernel reads two separate streams per
-    ),
+    row and prefers (2, 512) tiles at prefill row counts.
-)
-@helion.kernel(static_shapes=False)
-def _silu_and_mul_helion_non_interleaved_kernel(
-    gateup_output,
-    topk_weights: torch.Tensor | None = None,
-    out_dtype: hl.constexpr | None = None,
-):
     """
-    Non-interleaved version of silu_and_mul using Helion kernel.
+    if not interleaved and M > 4096 and N % 512 == 0:
-    Input format: [gate[0], gate[1], ..., gate[N-1], up[0], up[1], ..., up[N-1]]
+        return 2, 512, 1
-    """
+    BLOCK_M = 1
-    batch_size, hidden_size = gateup_output.shape
+    if N % 512 == 0 and N % 1024 != 0:
-    hidden_size = hl.specialize(hidden_size)
+        BLOCK_N = 512
-    assert hidden_size % 2 == 0, f"{hidden_size=}"
+    else:
-
+        BLOCK_N = min(1024, triton.next_power_of_2(N))
-    half_hidden_size = hidden_size // 2
+    if M <= 4096:
-    down_input = gateup_output.new_empty(
+        num_warps = 4
-        batch_size, half_hidden_size, dtype=out_dtype or gateup_output.dtype
+    else:
-    )
+        num_warps = 2 if interleaved else 1
-    for batch_tile, hidden_tile in hl.tile([batch_size, half_hidden_size]):
+    return BLOCK_M, BLOCK_N, num_warps
-        gate_output = gateup_output[batch_tile, hidden_tile.index].to(torch.float32)
-        up_output = gateup_output[batch_tile, hidden_tile.index + half_hidden_size].to(
-            torch.float32
-        )
-        silu_mul_output = gate_output * torch.sigmoid(gate_output) * up_output
-        if topk_weights is not None:
-            weight_scale = topk_weights[batch_tile, None].to(torch.float32)
-            silu_mul_output = silu_mul_output * weight_scale
-        down_input[batch_tile, hidden_tile] = silu_mul_output
-    return down_input
 
 
-def silu_and_mul_helion(
+def silu_and_mul(
     gateup_output: torch.Tensor,
     topk_weights: torch.Tensor | None = None,
     out_dtype: torch.dtype | None = None,
     use_interleaved: bool = True,
 ) -> torch.Tensor:
-    """
+    """silu(gate) * up over routed rows, optionally scaled by per-row weights.
-    Unified silu_and_mul function using Helion kernel.
+
-    Supports both interleaved and non-interleaved input formats.
+    Bitwise-identical port of the former Helion kernels (fp32 math,
+    gate * sigmoid(gate) * up, weight scale last, one rounding cast at the
+    store).
 
     Args:
         gateup_output: Input tensor of shape (batch_size, hidden_size)
-        topk_weights: Optional topk weights tensor
+        topk_weights: Optional per-row weights tensor of shape (batch_size,)
         out_dtype: Optional output dtype
         use_interleaved: If True, expects interleaved format [gate[0], up[0], gate[1], up[1], ...]
                         If False, expects non-interleaved format [gate[0], ..., gate[N-1], up[0], ..., up[N-1]]
@@ -144,20 +197,51 @@ def silu_and_mul_helion(
     Returns:
         Output tensor of shape (batch_size, hidden_size // 2)
     """
-    if use_interleaved:
+    assert gateup_output.ndim == 2, f"{gateup_output.shape=}"
-        return _silu_and_mul_helion_interleaved_kernel(
+    assert gateup_output.stride(1) == 1, f"{gateup_output.stride()=}"
-            gateup_output, topk_weights, out_dtype
+    M, hidden_size = gateup_output.shape
+    assert hidden_size % 2 == 0, f"{hidden_size=}"
+    N = hidden_size // 2
+    if topk_weights is not None:
+        assert topk_weights.is_contiguous(), f"{topk_weights.stride()=}"
+        assert topk_weights.shape == (M,), f"{topk_weights.shape=} {M=}"
+
+    down_input = torch.empty(
+        (M, N), device=gateup_output.device, dtype=out_dtype or gateup_output.dtype
     )
-    else:
+    if M == 0:
-        return _silu_and_mul_helion_non_interleaved_kernel(
+        return down_input
-            gateup_output, topk_weights, out_dtype
+
+    BLOCK_M, BLOCK_N, num_warps = _select_silu_and_mul_config(M, N, use_interleaved)
+    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
+    kernel = (
+        silu_and_mul_interleaved_kernel
+        if use_interleaved
+        else silu_and_mul_non_interleaved_kernel
     )
+    kernel[grid](
+        gateup_out_ptr=gateup_output,
+        topk_weights_ptr=topk_weights,
+        down_inp_ptr=down_input,
+        M=M,
+        stride_gm=gateup_output.stride(0),
+        N=N,
+        HAS_TOPK_WEIGHTS=topk_weights is not None,
+        INT64_INDEX=gateup_output.nbytes >= 2**31 or down_input.nbytes >= 2**31,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+    )
+    return down_input
 
 
 # ---------------------------------------------------------------------------
-# Triton silu_and_mul
+# Persistent-grid Triton silu_and_mul
-# Used by InklingBatchDenseMLP._swiglu because the helion kernel above produces
+# Used by InklingBatchDenseMLP._swiglu. It predates the plain-Triton port of
-# NaN for small shared-expert batches in EP+DP configs.
+# silu_and_mul above: the helion kernels that silu_and_mul replaced were
+# reported to produce NaN for small shared-expert batches in EP+DP configs.
+# Unlike silu_and_mul it can read the row count from a device tensor (M_ptr)
+# and supports PDL.
 # ---------------------------------------------------------------------------
 
 

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_100.json

@@ -1,39 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
-    "1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
-    "4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 0,
-    "(1536, (2,), (False,))": 1,
-    "(1536, (2,), (True,))": 1,
-    "(2048, (2,), (False,))": 4,
-    "(2048, (2,), (True,))": 1,
-    "(3072, (2,), (False,))": 0,
-    "(3072, (2,), (True,))": 0,
-    "(4096, (2,), (False,))": 4,
-    "(4096, (2,), (True,))": 1,
-    "(4608, (2,), (False,))": 0,
-    "(4608, (2,), (True,))": 0,
-    "(6144, (2,), (False,))": 0,
-    "(6144, (2,), (True,))": 1,
-    "(7680, (2,), (False,))": 1,
-    "(7680, (2,), (True,))": 1,
-    "(8192, (2,), (False,))": 4,
-    "(8192, (2,), (True,))": 1,
-    "(9216, (2,), (False,))": 0,
-    "(9216, (2,), (True,))": 0,
-    "(10240, (2,), (False,))": 0,
-    "(10240, (2,), (True,))": 1,
-    "(12288, (2,), (False,))": 0,
-    "(12288, (2,), (True,))": 1,
-    "(14336, (2,), (False,))": 0,
-    "(14336, (2,), (True,))": 1,
-    "(16384, (2,), (False,))": 4,
-    "(16384, (2,), (True,))": 1
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_121.json

@@ -1,39 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
-    "1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
-    "4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 0,
-    "(1536, (2,), (False,))": 1,
-    "(1536, (2,), (True,))": 1,
-    "(2048, (2,), (False,))": 4,
-    "(2048, (2,), (True,))": 1,
-    "(3072, (2,), (False,))": 0,
-    "(3072, (2,), (True,))": 0,
-    "(4096, (2,), (False,))": 4,
-    "(4096, (2,), (True,))": 1,
-    "(4608, (2,), (False,))": 0,
-    "(4608, (2,), (True,))": 0,
-    "(6144, (2,), (False,))": 0,
-    "(6144, (2,), (True,))": 1,
-    "(7680, (2,), (False,))": 1,
-    "(7680, (2,), (True,))": 1,
-    "(8192, (2,), (False,))": 4,
-    "(8192, (2,), (True,))": 1,
-    "(9216, (2,), (False,))": 0,
-    "(9216, (2,), (True,))": 0,
-    "(10240, (2,), (False,))": 0,
-    "(10240, (2,), (True,))": 1,
-    "(12288, (2,), (False,))": 0,
-    "(12288, (2,), (True,))": 1,
-    "(14336, (2,), (False,))": 0,
-    "(14336, (2,), (True,))": 1,
-    "(16384, (2,), (False,))": 4,
-    "(16384, (2,), (True,))": 1
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_90.json

@@ -1,40 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 1,
-    "(1536, (2,), (False,))": 2,
-    "(1536, (2,), (True,))": 2,
-    "(2048, (2,), (False,))": 3,
-    "(2048, (2,), (True,))": 3,
-    "(3072, (2,), (False,))": 1,
-    "(3072, (2,), (True,))": 1,
-    "(4096, (2,), (False,))": 3,
-    "(4096, (2,), (True,))": 3,
-    "(4608, (2,), (False,))": 1,
-    "(4608, (2,), (True,))": 1,
-    "(6144, (2,), (False,))": 3,
-    "(6144, (2,), (True,))": 3,
-    "(7680, (2,), (False,))": 2,
-    "(7680, (2,), (True,))": 2,
-    "(8192, (2,), (False,))": 3,
-    "(8192, (2,), (True,))": 3,
-    "(9216, (2,), (False,))": 1,
-    "(9216, (2,), (True,))": 1,
-    "(10240, (2,), (False,))": 3,
-    "(10240, (2,), (True,))": 3,
-    "(12288, (2,), (False,))": 3,
-    "(12288, (2,), (True,))": 3,
-    "(14336, (2,), (False,))": 3,
-    "(14336, (2,), (True,))": 3,
-    "(16384, (2,), (False,))": 3,
-    "(16384, (2,), (True,))": 3
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_95.json

@@ -1,40 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'pointer'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 1,
-    "(1536, (2,), (False,))": 2,
-    "(1536, (2,), (True,))": 2,
-    "(2048, (2,), (False,))": 3,
-    "(2048, (2,), (True,))": 3,
-    "(3072, (2,), (False,))": 1,
-    "(3072, (2,), (True,))": 1,
-    "(4096, (2,), (False,))": 3,
-    "(4096, (2,), (True,))": 3,
-    "(4608, (2,), (False,))": 1,
-    "(4608, (2,), (True,))": 1,
-    "(6144, (2,), (False,))": 3,
-    "(6144, (2,), (True,))": 3,
-    "(7680, (2,), (False,))": 2,
-    "(7680, (2,), (True,))": 2,
-    "(8192, (2,), (False,))": 3,
-    "(8192, (2,), (True,))": 3,
-    "(9216, (2,), (False,))": 1,
-    "(9216, (2,), (True,))": 1,
-    "(10240, (2,), (False,))": 3,
-    "(10240, (2,), (True,))": 3,
-    "(12288, (2,), (False,))": 3,
-    "(12288, (2,), (True,))": 3,
-    "(14336, (2,), (False,))": 3,
-    "(14336, (2,), (True,))": 3,
-    "(16384, (2,), (False,))": 3,
-    "(16384, (2,), (True,))": 3
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_sm_100.json

@@ -1,39 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
-    "1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
-    "4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 0,
-    "(1536, (2,), (False,))": 1,
-    "(1536, (2,), (True,))": 1,
-    "(2048, (2,), (False,))": 4,
-    "(2048, (2,), (True,))": 1,
-    "(3072, (2,), (False,))": 0,
-    "(3072, (2,), (True,))": 0,
-    "(4096, (2,), (False,))": 4,
-    "(4096, (2,), (True,))": 1,
-    "(4608, (2,), (False,))": 0,
-    "(4608, (2,), (True,))": 0,
-    "(6144, (2,), (False,))": 0,
-    "(6144, (2,), (True,))": 1,
-    "(7680, (2,), (False,))": 1,
-    "(7680, (2,), (True,))": 1,
-    "(8192, (2,), (False,))": 4,
-    "(8192, (2,), (True,))": 1,
-    "(9216, (2,), (False,))": 0,
-    "(9216, (2,), (True,))": 0,
-    "(10240, (2,), (False,))": 0,
-    "(10240, (2,), (True,))": 1,
-    "(12288, (2,), (False,))": 0,
-    "(12288, (2,), (True,))": 1,
-    "(14336, (2,), (False,))": 0,
-    "(14336, (2,), (True,))": 1,
-    "(16384, (2,), (False,))": 4,
-    "(16384, (2,), (True,))": 1
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_sm_90.json

@@ -1,40 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 1,
-    "(1536, (2,), (False,))": 2,
-    "(1536, (2,), (True,))": 2,
-    "(2048, (2,), (False,))": 3,
-    "(2048, (2,), (True,))": 3,
-    "(3072, (2,), (False,))": 1,
-    "(3072, (2,), (True,))": 1,
-    "(4096, (2,), (False,))": 3,
-    "(4096, (2,), (True,))": 3,
-    "(4608, (2,), (False,))": 1,
-    "(4608, (2,), (True,))": 1,
-    "(6144, (2,), (False,))": 3,
-    "(6144, (2,), (True,))": 3,
-    "(7680, (2,), (False,))": 2,
-    "(7680, (2,), (True,))": 2,
-    "(8192, (2,), (False,))": 3,
-    "(8192, (2,), (True,))": 3,
-    "(9216, (2,), (False,))": 1,
-    "(9216, (2,), (True,))": 1,
-    "(10240, (2,), (False,))": 3,
-    "(10240, (2,), (True,))": 3,
-    "(12288, (2,), (False,))": 3,
-    "(12288, (2,), (True,))": 3,
-    "(14336, (2,), (False,))": 3,
-    "(14336, (2,), (True,))": 3,
-    "(16384, (2,), (False,))": 3,
-    "(16384, (2,), (True,))": 3
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_sm_95.json

@@ -1,40 +0,0 @@
-{
-  "useful_configs": {
-    "0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'pointer'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
-    "3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
-  },
-  "hash_configs": {
-    "(512, (2,), (False,))": 0,
-    "(512, (2,), (True,))": 1,
-    "(1024, (2,), (False,))": 0,
-    "(1024, (2,), (True,))": 1,
-    "(1536, (2,), (False,))": 2,
-    "(1536, (2,), (True,))": 2,
-    "(2048, (2,), (False,))": 3,
-    "(2048, (2,), (True,))": 3,
-    "(3072, (2,), (False,))": 1,
-    "(3072, (2,), (True,))": 1,
-    "(4096, (2,), (False,))": 3,
-    "(4096, (2,), (True,))": 3,
-    "(4608, (2,), (False,))": 1,
-    "(4608, (2,), (True,))": 1,
-    "(6144, (2,), (False,))": 3,
-    "(6144, (2,), (True,))": 3,
-    "(7680, (2,), (False,))": 2,
-    "(7680, (2,), (True,))": 2,
-    "(8192, (2,), (False,))": 3,
-    "(8192, (2,), (True,))": 3,
-    "(9216, (2,), (False,))": 1,
-    "(9216, (2,), (True,))": 1,
-    "(10240, (2,), (False,))": 3,
-    "(10240, (2,), (True,))": 3,
-    "(12288, (2,), (False,))": 3,
-    "(12288, (2,), (True,))": 3,
-    "(14336, (2,), (False,))": 3,
-    "(14336, (2,), (True,))": 3,
-    "(16384, (2,), (False,))": 3,
-    "(16384, (2,), (True,))": 3
-  }
-}

python/sglang/srt/layers/moe/moe_runner/triton_utils/helion_utils.py

@@ -1,296 +0,0 @@
-import copy
-import dataclasses
-import functools
-import json
-import logging
-import os
-from collections import defaultdict
-from collections.abc import Iterable
-from enum import Enum
-from pathlib import Path
-from typing import Any, Callable
-
-import helion  # noqa: F401
-import torch
-from triton.testing import do_bench
-
-logger = logging.getLogger(__name__)
-AutotuneInputFn = Callable[[], Iterable[tuple[Any, ...]]]
-
-
-def get_model_depths() -> list[int]:
-    return [8, 16, 24, 32, 40, 48, 64]
-
-
-def get_cuda_device_capability() -> tuple[int, int]:
-    if torch.cuda.is_available():
-        dev = torch.cuda.current_device()
-        cc_major, cc_minor = torch.cuda.get_device_capability(dev)
-        return cc_major, cc_minor
-    return (0, 0)
-
-
-class AOTAutotuneMode(Enum):
-    NONE = "none"
-    CREATE = "create"
-    RETUNE = "retune"
-
-    @classmethod
-    def from_str(cls, mode: str) -> "AOTAutotuneMode":
-        return cls[mode.upper()]
-
-
-def load_autotune_data_from_json(
-    path: Path,
-) -> tuple[dict[int, Any], dict[tuple[Any, Any, Any], int]]:
-    with open(path, "r") as f:
-        json_obj = json.load(f)
-        useful_configs = {
-            int(k): eval(v) for k, v in json_obj["useful_configs"].items()
-        }
-        hash_configs = {tuple(eval(k)): v for k, v in json_obj["hash_configs"].items()}
-    return useful_configs, hash_configs
-
-
-def bind_and_compile_kernel(
-    kernel: helion.Kernel, args: Any, config: helion.Config
-) -> Callable:
-    has_int64 = any([i.numel() >= 2**31 for i in args if isinstance(i, torch.Tensor)])
-    if has_int64:
-        new_settings = dataclasses.replace(kernel.settings, index_dtype=torch.int64)
-        kernel = helion.Kernel(
-            kernel.fn,
-            configs=kernel.configs,
-            settings=new_settings,
-            key=kernel._key_fn,
-        )
-    config = copy.deepcopy(config)
-    return kernel.bind(args).compile_config(config, allow_print=False)
-
-
-def helion_aot_autotune(
-    config_path: str,
-    kernel_key: Callable[..., tuple[Any, Any, Any]],
-    primary_inputs: AutotuneInputFn,
-    secondary_inputs: AutotuneInputFn | None = None,
-    int64_threshold: Callable[..., bool] | None = None,
-    warn_on_hash_miss: bool = False,
-):
-    """
-    A decorator that automatically tunes and dispatches a Helion kernel based off of the kernel_key and provided inputs.
-
-    The general flow is this:
-    1. We first run helion_kernel.autotune on all primary_inputs. This will give us a list of configs, one per primary_input.
-    2. We benchmark every config on every primary_input and secondary_input.
-    3. For every primary input/secondary input, we keep the config that is the fastest (NB: We aim to do some deduplication by reusing configs if they're within some threshold of the fastest config).
-    4. We'll save the configs and the dispatch choices to a json file.
-    5. We create a dispatch function that will lookup to see whether each kernel is one we've tuned for in step 2. If so, we'll use that config. Otherwise, we'll use some heuristic to (deterministically) find a reasonable config for the kernel.
-
-    There are 3 modes (set by env variable HELION_AOT_AUTOTUNE):
-    - none: No autotuning will be done. We will skip to step 5. If the config json file doesn't exist, we'll raise an error.
-    - retune: We only benchmark the existing useful configs on the primary/secondary inputs. We'll skip to step 2. This finishes much faster than create (albeit won't fully retune for each shape), but requires create to have been run first. For example, for rmsnorm, retune takes maybe one minute for 10 shapes, but create might take 30 minutes.
-    - create: The kernel will be fully autotuned, starting from step 1.
-
-    You can also minimize the kernels to be autotuned by setting the env variable HELION_AOT_AUTOTUNE_KERNEL to the name of the kernel. For example, if you want to only autotune rmsnorm, you can set HELION_AOT_AUTOTUNE_KERNEL="rms_norm_fwd".
-
-    kernel_key: Callable[..., tuple[Any, Any, Any]]
-    returns:
-        (numeric_key, hash_key, exact_key)
-    The semantics are that if all 3 match a saved key, we'll use that config.
-    Otherwise, we'll use some heuristic to find a reasonable config for the
-    kernel. The heuristic is:
-    - We require exact_key to match the saved config.
-    - We prioritize hash_key that match the saved config.
-    - We then prioritize the highest numeric_key that's <= the current numeric_key.
-    """
-    # Threshold for how much faster a config has to be for some shape to be considered "useful"
-    threshold = 1.01
-    # How many ms to run the kernel when retuning
-    retune_rep_ms = 1000
-
-    def inner_autotune(kernel: helion.Kernel):
-        helion_dir = Path(__file__).parent / "configs"
-
-        cc_major, cc_minor = get_cuda_device_capability()
-        if cc_minor == 3:
-            cc_minor = 0
-        gpu_arch = f"sm_{cc_major}{cc_minor}"
-
-        path = helion_dir / Path(f"{config_path}_{gpu_arch}.json")
-
-        @functools.wraps(kernel_key)
-        def wrapped_kernel_key(*inps: Any) -> tuple[Any, Any, Any]:
-            """
-            A wrapper that handles dtype specially (since dtype is not
-            serializable to json).
-            """
-
-            numeric_key, hash_key, exact_key = kernel_key(*inps)
-            assert isinstance(hash_key, tuple)
-            assert isinstance(exact_key, tuple)
-            hash_key = tuple(
-                i.itemsize if isinstance(i, torch.dtype) else i for i in hash_key
-            )
-            if exact_key is not None:
-                exact_key = tuple(
-                    i.itemsize if isinstance(i, torch.dtype) else i for i in exact_key
-                )
-            return numeric_key, hash_key, exact_key
-
-        autotune_mode_str = os.environ.get("HELION_AOT_AUTOTUNE", "none")
-        autotune_mode = AOTAutotuneMode.from_str(autotune_mode_str)
-        autotune_kernel = os.environ.get("HELION_AOT_AUTOTUNE_KERNEL", "all")
-        if autotune_kernel != "all" and autotune_kernel != config_path:
-            autotune_mode = AOTAutotuneMode.NONE
-
-        @functools.cache
-        def get_configs_from_autotuning(autotune_mode: AOTAutotuneMode):
-            if autotune_mode == AOTAutotuneMode.NONE:
-                if not path.exists():
-                    raise RuntimeError(
-                        f"Helion kernel not tuned yet. Run with HELION_AOT_AUTOTUNE=create HELION_AOT_AUTOTUNE_KERNEL={config_path} to generate the config at {path}"
-                    )
-                return load_autotune_data_from_json(path)
-
-            inputs = sorted(
-                list(primary_inputs()), key=lambda x: wrapped_kernel_key(*x)[0]
-            )
-            if autotune_mode == AOTAutotuneMode.CREATE:
-                useful_configs = []
-                for idx, input in enumerate(inputs):
-                    print(
-                        f"Autotuning for {config_path} with key: ",
-                        wrapped_kernel_key(*input),
-                    )
-                    config = kernel.autotune(input)
-                    useful_configs.append(repr(config))
-            elif autotune_mode == AOTAutotuneMode.RETUNE:
-                with open(path, "r") as f:
-                    json_obj = json.load(f)
-                    useful_configs = list(json_obj["useful_configs"].values())
-            else:
-                raise RuntimeError(f"Unexpected autotune mode: {autotune_mode}")
-
-            logger.info("Candidate useful configs: ")
-            for idx, config in enumerate(useful_configs):
-                logger.info(f"{idx}:, Config: {config}")
-
-            input_timings = []
-            if secondary_inputs is not None:
-                inputs += list(secondary_inputs())
-                inputs = sorted(inputs, key=lambda x: kernel_key(*x)[0])
-
-            for input in inputs:
-                cur_input_key = wrapped_kernel_key(*input)
-                timings = []
-                for idx, config in enumerate(useful_configs):
-                    try:
-                        cur_kernel = bind_and_compile_kernel(
-                            kernel, input, eval(config)
-                        )
-                        timings.append(
-                            do_bench(lambda: cur_kernel(*input), rep=retune_rep_ms)
-                        )  # noqa: B023
-                    except Exception as e:
-                        logger.info(f"Error compiling config {config}: {e}")
-                        timings.append(float("inf"))
-                input_timings.append((cur_input_key, timings))
-            for idx, (key, timing) in enumerate(input_timings):
-                logger.info(
-                    f"Key {key} timings: {' '.join([f'{i:.5f}' for i in timing])}"
-                )
-
-            hash_configs_timings: dict[tuple[Any, Any], tuple[float, int | None]] = (
-                defaultdict(lambda: (float("inf"), None))
-            )
-            for cur_kernel_key, input_timings in input_timings:
-                for config_idx, (config, timing) in enumerate(
-                    zip(useful_configs, input_timings)
-                ):
-                    if timing < hash_configs_timings[cur_kernel_key][0] * threshold:
-                        hash_configs_timings[cur_kernel_key] = (timing, config_idx)
-
-            kept_configs = {}
-            hash_configs = {k: v[1] for k, v in hash_configs_timings.items()}
-            for key, config_idx in hash_configs.items():
-                assert config_idx is not None
-                kept_configs[config_idx] = useful_configs[config_idx]
-
-            json_obj = {
-                "useful_configs": kept_configs,
-                "hash_configs": {repr(k): v for k, v in hash_configs.items()},
-            }
-
-            with open(path, "w") as f:
-                json.dump(json_obj, f, indent=2)
-                f.write("\n")
-            return load_autotune_data_from_json(path)
-
-        cached_kernels = {}
-
-        def default_int64_threshold(*args: Any) -> bool:
-            return any(
-                [i.numel() >= 2**31 for i in args if isinstance(i, torch.Tensor)]
-            )
-
-        used_int64_threshold = (
-            int64_threshold if int64_threshold is not None else default_int64_threshold
-        )
-
-        def wrapped_func(*args: Any):
-            nonlocal cached_kernels
-            cur_kernel_key = wrapped_kernel_key(*args)
-            has_int64 = used_int64_threshold(*args)
-            size1 = tuple(
-                tuple(shape == 1 for shape in arg.shape)
-                for arg in args
-                if isinstance(arg, torch.Tensor)
-            )
-            dtypes = tuple(arg.dtype for arg in args if isinstance(arg, torch.Tensor))
-            scalar_args = tuple(a for a in args if not isinstance(a, torch.Tensor))
-            key = (cur_kernel_key, has_int64, size1, dtypes, scalar_args)
-            if key in cached_kernels:
-                out = cached_kernels[key](*args)
-                return out
-            if has_int64:
-                kernel.settings = dataclasses.replace(
-                    kernel.settings, index_dtype=torch.int64
-                )
-            useful_configs, hash_configs = get_configs_from_autotuning(autotune_mode)
-            key_to_config = {k: useful_configs[v] for k, v in hash_configs.items()}
-            if key not in cached_kernels and cur_kernel_key in hash_configs:
-                cached_kernels[key] = bind_and_compile_kernel(
-                    kernel, args, key_to_config[cur_kernel_key]
-                )
-            else:
-                if warn_on_hash_miss:
-                    logger.warning(
-                        f"No config found for key {cur_kernel_key} for kernel={config_path}. Finding best match. This is *not* a correctness issue, but means that the performance of this kernel could potentially be improved."
-                    )
-                used_config = None
-
-                def config_key_sort(
-                    config_key: tuple[Any, Any, Any],
-                ) -> tuple[Any, ...]:
-                    return (
-                        config_key[2] == cur_kernel_key[2],
-                        config_key[1] == cur_kernel_key[1],
-                        config_key[0] <= cur_kernel_key[0],
-                        config_key[0],
-                    )
-
-                sorted_keys = sorted(
-                    hash_configs.keys(), key=config_key_sort, reverse=True
-                )
-                used_config = key_to_config[sorted_keys[0]]
-                if len(used_config) == 3:
-                    assert (
-                        used_config[2] == cur_kernel_key[2]
-                    ), "Exact key not found in configs"
-                cached_kernels[key] = bind_and_compile_kernel(kernel, args, used_config)
-            out = cached_kernels[key](*args)
-            return out
-
-        return wrapped_func
-
-    return inner_autotune

python/sglang/srt/models/inkling_common/moe.py

@@ -25,7 +25,7 @@ from sglang.kernels.ops.moe.inkling_moe import (
     post_reorder,
     pre_reorder,
     select_grouped_gemm_block_m,
-    silu_and_mul_helion,
+    silu_and_mul,
 )
 from sglang.kernels.ops.moe.sigmoid_gate_topk_renorm import (
     sigmoid_gate_topk_renorm,
@@ -572,7 +572,7 @@ def activation(
                 *gateup_output.shape[:-1], gateup_output.shape[-1] // 2, dtype=out_dtype
             )
 
-        return silu_and_mul_helion(
+        return silu_and_mul(
             gateup_output, topk_weights, out_dtype, use_interleaved=use_interleaved
         )
     raise ValueError(f"Unsupported activation: {activation_type}")

test/registered/kernels/ops/moe/test_inkling_silu_and_mul.py

@@ -0,0 +1,159 @@
+"""Numerics tests for the plain-Triton silu_and_mul (former helion kernels).
+
+The kernels compute silu(gate) * up (* weight) in fp32 with one rounding cast
+at the store, in the exact operation order of the helion kernels they
+replaced, so bf16 outputs sit within 1 bf16 ulp of a same-order torch fp32
+reference (tl.sigmoid and torch.sigmoid can differ in the last fp32 ulp).
+
+The interleaved and non-interleaved kernels share the operation order, so the
+two layouts must produce bitwise-identical outputs for the same gate/up data.
+
+The small-batch tests pin the shared-expert shape family where the deleted
+helion kernels were reported to produce NaN in EP+DP configs (see the
+InklingBatchDenseMLP._swiglu comment); the port must stay finite and correct
+there.
+"""
+
+import pytest
+import torch
+
+from sglang.kernels.ops.moe.inkling_moe import silu_and_mul
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+
+def _reference(
+    gateup: torch.Tensor,
+    weights: torch.Tensor | None,
+    out_dtype: torch.dtype | None,
+    interleaved: bool,
+) -> torch.Tensor:
+    xf = gateup.float()
+    if interleaved:
+        gate, up = xf[:, 0::2], xf[:, 1::2]
+    else:
+        n = gateup.shape[1] // 2
+        gate, up = xf[:, :n], xf[:, n:]
+    out = gate * torch.sigmoid(gate) * up
+    if weights is not None:
+        out = out * weights.float()[:, None]
+    return out.to(out_dtype or gateup.dtype)
+
+
+def _ulp_diff_bf16(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    assert a.dtype == b.dtype == torch.bfloat16
+    mask = (1 << 15) - 1
+    ai = a.contiguous().view(torch.int16).to(torch.int64)
+    bi = b.contiguous().view(torch.int16).to(torch.int64)
+    ai = torch.where(ai < 0, -(ai & mask), ai)
+    bi = torch.where(bi < 0, -(bi & mask), bi)
+    return (ai - bi).abs()
+
+
+def _check(m: int, n: int, interleaved: bool, use_weights: bool, out_dtype=None):
+    torch.manual_seed(m * 7919 + n)
+    x = torch.randn(m, 2 * n, dtype=torch.bfloat16, device="cuda")
+    w = torch.randn(m, dtype=torch.bfloat16, device="cuda") if use_weights else None
+    out = silu_and_mul(x, w, out_dtype, use_interleaved=interleaved)
+    assert out.shape == (m, n)
+    assert out.isfinite().all(), "kernel produced non-finite values"
+    ref = _reference(x, w, out_dtype, interleaved)
+    if out.dtype == torch.bfloat16:
+        assert int(_ulp_diff_bf16(out, ref).max()) <= 1
+    else:
+        torch.testing.assert_close(out, ref, rtol=1e-6, atol=1e-6)
+    # Deterministic: an exact repeat must be bitwise identical.
+    again = silu_and_mul(x, w, out_dtype, use_interleaved=interleaved)
+    assert torch.equal(out, again)
+    return out
+
+
+# d42 serves N=2048 (tp1) / 1024 (tp2); d66 serves N=3072/1536/768.
+# Rows: decode <= 256 reqs x topk 6; prefill up to 8192 tokens x topk 6.
+@requires_cuda
+@pytest.mark.parametrize(
+    ("m", "n"),
+    [(1, 1024), (6, 768), (384, 1024), (1536, 2048), (4096, 3072), (49152, 1024)],
+)
+@pytest.mark.parametrize("interleaved", [True, False])
+@pytest.mark.parametrize("use_weights", [False, True])
+def test_matches_reference(m: int, n: int, interleaved: bool, use_weights: bool):
+    _check(m, n, interleaved, use_weights)
+
+
+@requires_cuda
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32])
+def test_out_dtype(out_dtype: torch.dtype):
+    out = _check(512, 1024, True, True, out_dtype=out_dtype)
+    assert out.dtype == out_dtype
+
+
+@requires_cuda
+@pytest.mark.parametrize(("m", "n"), [(7, 384), (33, 960), (5, 512)])
+def test_odd_widths(m: int, n: int):
+    _check(m, n, True, False)
+    _check(m, n, False, True)
+
+
+@requires_cuda
+def test_layouts_bitwise_consistent():
+    """Same gate/up data through both layouts must match bitwise."""
+    torch.manual_seed(0)
+    m, n = 1234, 1536
+    gate = torch.randn(m, n, dtype=torch.bfloat16, device="cuda")
+    up = torch.randn(m, n, dtype=torch.bfloat16, device="cuda")
+    w = torch.randn(m, dtype=torch.bfloat16, device="cuda")
+    x_il = torch.stack([gate, up], dim=2).reshape(m, 2 * n).contiguous()
+    x_nil = torch.cat([gate, up], dim=1).contiguous()
+    out_il = silu_and_mul(x_il, w, None, use_interleaved=True)
+    out_nil = silu_and_mul(x_nil, w, None, use_interleaved=False)
+    assert torch.equal(out_il, out_nil)
+
+
+@requires_cuda
+@pytest.mark.parametrize("m", [1, 2, 4, 8, 16])
+@pytest.mark.parametrize("n", [768, 1024, 1536])
+def test_small_shared_expert_batches(m: int, n: int):
+    """Regression for the helion NaN report: tiny gamma-weighted batches.
+
+    m = n_shared_experts * tokens_on_dp_rank (2 shared experts, 1-8 tokens);
+    n = the per-rank shared-expert width (d66 tp4 = 768, d42 tp2 = 1024).
+    """
+    out = _check(m, n, True, True)
+    assert out.isfinite().all()
+
+
+@requires_cuda
+def test_zero_rows():
+    x = torch.empty(0, 2048, dtype=torch.bfloat16, device="cuda")
+    out = silu_and_mul(x, None, None, use_interleaved=True)
+    assert out.shape == (0, 1024)
+
+
+@requires_cuda
+def test_int64_offsets():
+    """Element offsets beyond 2**31 must address correctly (INT64_INDEX path)."""
+    if torch.cuda.get_device_properties(0).total_memory < 16 * 2**30:
+        pytest.skip("needs >= 16 GB GPU memory")
+    torch.manual_seed(0)
+    n = 1024
+    m = 2**20 + 8  # numel = m * 2n = 2**31 + 16384 > 2**31
+    x = torch.randn(m, 2 * n, dtype=torch.bfloat16, device="cuda")
+    assert x.numel() > 2**31
+    for interleaved in (True, False):
+        out = silu_and_mul(x, None, None, use_interleaved=interleaved)
+        # The tail rows sit at element offsets >= 2**31.
+        tail_ref = _reference(x[-4:], None, None, interleaved)
+        assert int(_ulp_diff_bf16(out[-4:], tail_ref).max()) <= 1
+        assert out[-4:].isfinite().all()
+        del out
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))