[Kernel] Migrate scattered MoE kernels to sglang.kernels (RFC #29630, Phase 2.5, 2/7) (#30786)

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

python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py

@@ -163,7 +163,7 @@ def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
 )
 @torch.inference_mode()
 def test_mxfp8_native_moe(T, H, inter, E, top_k):
-    from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
+    from sglang.kernels.ops.moe.mxfp8_moe_amd_gfx95 import (
         fused_moe_mxfp8_native,
     )
 
@@ -203,7 +203,7 @@ def test_mxfp8_native_moe(T, H, inter, E, top_k):
 @requires_gfx950
 @torch.inference_mode()
 def test_mxfp8_native_moe_ep_expert_map_filters_non_local_routes():
-    from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
+    from sglang.kernels.ops.moe.mxfp8_moe_amd_gfx95 import (
         fused_moe_mxfp8_native,
     )
 

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

@@ -133,6 +133,32 @@ for _mod, _fn in _TRITON_KERNELS:
 del _mod, _fn
 
 
+# Triton kernels migrated from srt/layers/moe (RFC #29630, Phase 2.5);
+# registered for inventory. Import them from their modules.
+_PHASE25_TRITON_KERNELS = [
+    ("ep_moe_kernels", "deepep_run_moe_deep_preprocess"),
+    ("ep_moe_kernels", "deepep_permute_triton_kernel"),
+    ("ep_moe_kernels", "deepep_post_reorder_triton_kernel"),
+    ("fused_moe_triton_kernels", "invoke_fused_moe_kernel"),
+    ("fused_moe_triton_kernels", "fused_moe_kernel"),
+    ("fused_moe_triton_kernels", "fused_moe_kernel_gptq_awq"),
+    ("mxfp8_moe_amd_gfx95", "fused_experts_mxfp8"),
+    ("rocm_moe_utils", "upscale"),
+    ("rocm_moe_utils", "upscale_mxfp4"),
+    ("router", "fused_moe_router_shim"),
+    ("deepep_waterfill_kernels", "materialize_waterfill_dispatch_fused"),
+    ("fill_padded_rows", "_fill_padded_rows"),
+]
+for _mod, _fn in _PHASE25_TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"moe.{_fn.lstrip('_')}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.moe.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
 # Packed (topk_id << 16 | bf16-weight) kernel migrated from
 # srt/layers/quantization/mxfp4_flashinfer_trtllm_moe (RFC #29630, Phase 2.5).
 register_kernel(

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

@@ -0,0 +1,322 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""DeepEP Waterfill dispatch kernels (RFC #29630, Phase 2.5).
+
+Triton kernels and the fused dispatch materializer migrated from
+``sglang.srt.layers.moe.deepep_waterfill``; the balancer policy stays in srt.
+"""
+
+from typing import NamedTuple, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+LOCAL_SHARED_MARKER = -1  # Invalid expert ID; DeepEP ignores expert_id < 0.
+_LOCAL_PREF_NUMER = 11  # local-rank preference = 11/10
+_LOCAL_PREF_DENOM = 10
+
+
+class WaterfillDispatchPlan(NamedTuple):
+    """Inputs needed by the fused DeepEP Waterfill expansion path."""
+
+    # Effective rank load consumed by the fused kernel.
+    rank_load: Tensor
+    allow_all_ranks: bool
+    target_total: int
+
+
+def _empty_expanded(topk_ids: Tensor, topk_weights: Tensor):
+    """Return empty expanded tensors for zero-token batches."""
+    topk, d = topk_ids.shape[1], topk_ids.device
+    return (
+        torch.empty(0, topk + 1, dtype=topk_ids.dtype, device=d),
+        torch.empty(0, topk + 1, dtype=topk_weights.dtype, device=d),
+    )
+
+
+@triton.jit
+def _count_routed_per_rank_kernel(
+    topk_ids_ptr,  # [num_tokens, topk]
+    counts_ptr,  # [world_size] output (atomic add)
+    num_tokens,
+    topk: tl.constexpr,
+    experts_per_rank,
+    world_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Count routed tokens per rank using block-level histogram."""
+    pid = tl.program_id(0)
+    token_idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = token_idx < num_tokens
+
+    for r in range(world_size):
+        rank_count = tl.zeros([BLOCK_SIZE], dtype=tl.int64)
+
+        for k in range(topk):
+            expert_id = tl.load(
+                topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1
+            ).to(tl.int64)
+            valid = expert_id >= 0
+            target_rank = expert_id // experts_per_rank
+            target_rank = tl.minimum(tl.maximum(target_rank, 0), world_size - 1)
+            rank_count += tl.where(
+                mask & valid & (target_rank == r),
+                tl.full([BLOCK_SIZE], 1, dtype=tl.int64),
+                tl.zeros([BLOCK_SIZE], dtype=tl.int64),
+            )
+
+        block_total = tl.sum(rank_count)
+        if block_total > 0:
+            tl.atomic_add(counts_ptr + r, block_total)
+
+
+@triton.jit
+def _waterfill_expand_kernel(
+    topk_ids_ptr,
+    topk_weights_ptr,
+    rank_load_ptr,
+    expanded_ids_ptr,
+    expanded_weights_ptr,
+    num_tokens,
+    topk: tl.constexpr,
+    old_experts_per_rank,
+    new_experts_per_rank,
+    world_size: tl.constexpr,
+    source_rank,
+    shared_weight,
+    local_marker,
+    local_pref_numer,
+    local_pref_denom,
+    precomputed_target_total,
+    ALLOW_ALL_RANKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Fused waterfill + expand. ID remap: old_id -> old_id + old_id // old_epr."""
+    pid = tl.program_id(0)
+    token_idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = token_idx < num_tokens
+
+    r_idx = tl.arange(0, world_size)
+    rank_load_vec = tl.load(rank_load_ptr + r_idx, mask=r_idx < world_size, other=0).to(
+        tl.int64
+    )
+    total_effective_k = tl.sum(rank_load_vec)
+    total_tokens_global_k = total_effective_k // topk
+    derived_target_total = (
+        total_effective_k + total_tokens_global_k + world_size - 1
+    ) // world_size
+    target_total = tl.where(
+        precomputed_target_total > 0,
+        precomputed_target_total,
+        derived_target_total,
+    )
+
+    # Step 1: Select destination rank for shared expert (waterfill sampling).
+    source_count = tl.load(rank_load_ptr + source_rank)
+    best_count = tl.where(mask, source_count, 2**30)
+    best_rank = tl.full([BLOCK_SIZE], source_rank, dtype=tl.int64)
+    has_valid = tl.zeros([BLOCK_SIZE], dtype=tl.int1)
+    src_rank_i32 = tl.full([BLOCK_SIZE], source_rank, dtype=tl.int32)
+
+    if ALLOW_ALL_RANKS:
+        candidate_mask = tl.full([BLOCK_SIZE], (1 << world_size) - 1, dtype=tl.int32)
+        for r in range(world_size):
+            target_count = tl.load(rank_load_ptr + r).to(tl.int64)
+            better = (
+                target_count * local_pref_numer < best_count * local_pref_denom
+            ) & mask
+            best_count = tl.where(better, target_count, best_count)
+            best_rank = tl.where(
+                better, tl.full([BLOCK_SIZE], r, dtype=tl.int64), best_rank
+            )
+    else:
+        candidate_mask = (tl.full([BLOCK_SIZE], 1, dtype=tl.int32) << src_rank_i32).to(
+            tl.int32
+        )
+
+    for k in range(topk):
+        expert_id = tl.load(
+            topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1
+        ).to(tl.int64)
+        valid = expert_id >= 0
+        has_valid = has_valid | valid
+
+        if not ALLOW_ALL_RANKS:
+            target_rank = expert_id // old_experts_per_rank
+            target_rank = tl.minimum(tl.maximum(target_rank, 0), world_size - 1)
+            target_rank_i32 = target_rank.to(tl.int32)
+            shift_amt = tl.where(valid, target_rank_i32, 0)
+            bit = tl.full([BLOCK_SIZE], 1, dtype=tl.int32) << shift_amt
+            candidate_mask = tl.where(
+                valid & mask, candidate_mask | bit, candidate_mask
+            )
+
+            target_count = tl.load(
+                rank_load_ptr + target_rank, mask=mask & valid, other=2**30
+            )
+
+            better = (
+                (target_count * local_pref_numer < best_count * local_pref_denom)
+                & valid
+                & mask
+            )
+            best_count = tl.where(better, target_count, best_count)
+            best_rank = tl.where(better, target_rank, best_rank)
+
+    total_w = tl.zeros([BLOCK_SIZE], dtype=tl.int32)
+    for r in range(world_size):
+        present = ((candidate_mask >> r) & 1) == 1
+        rank_load_r = tl.load(rank_load_ptr + r).to(tl.int64)
+        w = tl.where(target_total > rank_load_r, target_total - rank_load_r, 0).to(
+            tl.int32
+        )
+        w_vec = tl.full([BLOCK_SIZE], w, dtype=tl.int32)
+        w_vec = tl.where(
+            src_rank_i32 == r,
+            w_vec,
+            (w_vec * local_pref_denom) // local_pref_numer,
+        )
+        total_w += tl.where(present, w_vec, 0)
+
+    token_seed = token_idx.to(tl.uint32) ^ (
+        src_rank_i32.to(tl.uint32) * tl.full([BLOCK_SIZE], 0x9E3779B9, dtype=tl.uint32)
+    )
+    token_seed = token_seed * tl.full([BLOCK_SIZE], 1664525, dtype=tl.uint32) + tl.full(
+        [BLOCK_SIZE], 1013904223, dtype=tl.uint32
+    )
+    u = tl.where(total_w > 0, token_seed % total_w.to(tl.uint32), 0).to(tl.int32)
+
+    chosen = src_rank_i32
+    cum = tl.zeros([BLOCK_SIZE], dtype=tl.int32)
+    for r in range(world_size):
+        present = ((candidate_mask >> r) & 1) == 1
+        rank_load_r = tl.load(rank_load_ptr + r).to(tl.int64)
+        w = tl.where(target_total > rank_load_r, target_total - rank_load_r, 0).to(
+            tl.int32
+        )
+        w_vec = tl.full([BLOCK_SIZE], w, dtype=tl.int32)
+        w_vec = tl.where(
+            src_rank_i32 == r,
+            w_vec,
+            (w_vec * local_pref_denom) // local_pref_numer,
+        )
+        w_vec = tl.where(present, w_vec, 0)
+        pick = (total_w > 0) & present & (u >= cum) & (u < (cum + w_vec))
+        chosen = tl.where(pick, r, chosen)
+        cum += w_vec
+
+    best_rank = tl.where(total_w > 0, chosen.to(tl.int64), best_rank)
+
+    # Step 2: Compute shared expert ID and local mask.
+    is_local = best_rank == source_rank
+    local_shared_id = source_rank * new_experts_per_rank + old_experts_per_rank
+    remote_shared_id = best_rank * new_experts_per_rank + old_experts_per_rank
+    shared_expert_id = tl.where(
+        is_local,
+        tl.full([BLOCK_SIZE], local_shared_id, dtype=tl.int64),
+        remote_shared_id,
+    ).to(tl.int64)
+    shared_expert_id = tl.where(
+        has_valid,
+        shared_expert_id,
+        tl.full([BLOCK_SIZE], local_marker, dtype=tl.int64),
+    )
+
+    # Step 3: Copy and remap topk_ids, copy weights.
+    for k in range(topk):
+        old_id = tl.load(topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1).to(
+            tl.int64
+        )
+        valid_id = old_id >= 0
+        new_id = tl.where(valid_id, old_id + (old_id // old_experts_per_rank), old_id)
+        tl.store(expanded_ids_ptr + token_idx * (topk + 1) + k, new_id, mask=mask)
+
+    for k in range(topk):
+        val = tl.load(topk_weights_ptr + token_idx * topk + k, mask=mask, other=0.0)
+        expert_id = tl.load(
+            topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1
+        ).to(tl.int64)
+        val = tl.where(expert_id >= 0, val, 0.0)
+        tl.store(expanded_weights_ptr + token_idx * (topk + 1) + k, val, mask=mask)
+
+    # Step 4: Write shared expert column.
+    tl.store(
+        expanded_ids_ptr + token_idx * (topk + 1) + topk,
+        shared_expert_id,
+        mask=mask,
+    )
+    tl.store(
+        expanded_weights_ptr + token_idx * (topk + 1) + topk,
+        tl.where(has_valid, shared_weight, 0.0),
+        mask=mask,
+    )
+
+
+def materialize_waterfill_dispatch_fused(
+    topk_ids: Tensor,
+    topk_weights: Tensor,
+    rank_load: Tensor,
+    num_routed_experts: int,
+    world_size: int,
+    source_rank: int,
+    shared_weight: float,
+    allow_all_ranks: bool = False,
+    target_total: int = 0,
+) -> Tuple[Tensor, Tensor]:
+    """Run fused Waterfill rank selection and DeepEP TopK expansion.
+
+    The Triton kernel intentionally selects each token's shared-expert rank and
+    writes the expanded DeepEP TopK layout in one pass.
+    """
+    num_tokens = topk_ids.shape[0]
+    topk = topk_ids.shape[1]
+    old_experts_per_rank = num_routed_experts // world_size
+    new_experts_per_rank = old_experts_per_rank + 1
+    device = topk_ids.device
+
+    if num_tokens == 0:
+        return _empty_expanded(topk_ids, topk_weights)
+
+    expanded_topk_ids = torch.empty(
+        num_tokens, topk + 1, dtype=topk_ids.dtype, device=device
+    )
+    expanded_topk_weights = torch.empty(
+        num_tokens, topk + 1, dtype=topk_weights.dtype, device=device
+    )
+    BLOCK_SIZE = 256
+    grid = ((num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE,)
+    _waterfill_expand_kernel[grid](
+        topk_ids,
+        topk_weights,
+        rank_load,
+        expanded_topk_ids,
+        expanded_topk_weights,
+        num_tokens,
+        topk,
+        old_experts_per_rank,
+        new_experts_per_rank,
+        world_size,
+        source_rank,
+        shared_weight,
+        LOCAL_SHARED_MARKER,
+        _LOCAL_PREF_NUMER,
+        _LOCAL_PREF_DENOM,
+        target_total,
+        allow_all_ranks,
+        BLOCK_SIZE,
+    )
+
+    return expanded_topk_ids, expanded_topk_weights

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

@@ -0,0 +1,78 @@
+"""Fused padded-row fill for MoE top-k outputs.
+
+Migrated from ``sglang.srt.layers.moe.topk`` (RFC #29630, Phase 2.5), where two
+near-identical copies had accumulated; this keeps the later, runtime-winning
+copy (explicit raises instead of asserts).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fill_padded_rows_kernel(
+    out_ptr,
+    num_token_non_padded_ptr,
+    n_cols,
+    fill_value,
+    stride_row,
+    BLOCK_COLS: tl.constexpr,
+):
+    row = tl.program_id(0)
+    n_valid = tl.load(num_token_non_padded_ptr)
+    if row >= n_valid:
+        cols = tl.arange(0, BLOCK_COLS)
+        mask = cols < n_cols
+        ptrs = out_ptr + row * stride_row + cols
+        fill = tl.full((BLOCK_COLS,), fill_value, dtype=out_ptr.dtype.element_ty)
+        tl.store(ptrs, fill, mask=mask)
+
+
+def _can_fuse_padded_region(x: torch.Tensor) -> bool:
+    # The fused kernel uses one program per row and assumes a row-major 2D
+    # tensor (columns contiguous); fall back to eager for anything else.
+    return x.dim() == 2 and x.stride(1) == 1
+
+
+def _fill_padded_rows(
+    x: torch.Tensor,
+    num_token_non_padded: torch.Tensor,
+    fill_value,
+) -> None:
+    """Set ``x[row, :] = fill_value`` for every padded row (row index
+    ``>= num_token_non_padded``) using a single Triton launch.
+
+    Replaces the eager ``arange + (>=) + boolean index_put_`` sequence, which
+    issues several launch-latency-bound kernels per call. The grid is static
+    (one program per row) and the pad count is read from device memory inside
+    the kernel, so this is safe to capture inside a CUDA/HIP graph.
+    """
+    # Metadata-only checks (no device sync): the kernel reads a single scalar
+    # routing count from device memory, so it must be a 1-element integer tensor
+    # on the same device as ``x``. Use explicit raises (not asserts) so the
+    # checks survive ``python -O`` and invalid inputs fail loudly instead of
+    # turning into opaque Triton/memory errors.
+    if not isinstance(num_token_non_padded, torch.Tensor):
+        raise TypeError("num_token_non_padded must be a torch.Tensor")
+    if num_token_non_padded.numel() != 1:
+        raise ValueError(
+            "num_token_non_padded must be a single-element tensor, got shape "
+            f"{tuple(num_token_non_padded.shape)}"
+        )
+    if num_token_non_padded.dtype.is_floating_point:
+        raise TypeError(
+            "num_token_non_padded must be an integer tensor, got "
+            f"{num_token_non_padded.dtype}"
+        )
+    if num_token_non_padded.device != x.device:
+        raise ValueError("num_token_non_padded and x must be on the same device")
+    n_rows, n_cols = x.shape
+    _fill_padded_rows_kernel[(n_rows,)](
+        x,
+        num_token_non_padded,
+        n_cols,
+        fill_value,
+        x.stride(0),
+        BLOCK_COLS=triton.next_power_of_2(n_cols),
+    )

python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py

@@ -864,7 +864,7 @@ def _merged_experts_fused_moe_lora_add_impl(
 
         return result
 
-    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+    from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
         invoke_fused_moe_kernel,
     )
 

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

@@ -637,7 +637,7 @@ def _merged_experts_fused_moe_lora_add_impl(
 
         return result
 
-    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+    from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
         invoke_fused_moe_kernel,
     )
 

python/sglang/kernels/ops/quantization/fp8_kernel.py

@@ -381,7 +381,7 @@ def _per_token_group_quant_8bit_fuse_silu_and_mul(
 
     from deep_gemm import transform_sf_into_required_layout
 
-    from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd
+    from sglang.kernels.ops.moe.ep_moe_kernels import silu_and_mul_masked_post_quant_fwd
 
     assert column_major_scales
     assert scale_tma_aligned

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -23,7 +23,7 @@ else:
     from sgl_kernel import silu_and_mul
 
 from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
-from sglang.srt.layers.moe.ep_moe.kernels import (
+from sglang.kernels.ops.moe.ep_moe_kernels import (
     cutlass_w4_run_moe_ep_preproess,
     deepep_ll_get_cutlass_w4a8_moe_mm_data,
     deepep_permute_triton_kernel,

python/sglang/srt/layers/moe/moe_runner/aiter.py

@@ -300,7 +300,7 @@ def _pre_permute_deepep_to_aiter(
     quant_type = quant_info.quant_type
 
     if is_mori:
-        from sglang.srt.layers.moe.rocm_moe_utils import upscale, upscale_mxfp4
+        from sglang.kernels.ops.moe.rocm_moe_utils import upscale, upscale_mxfp4
 
         a1_scale = dispatch_output.hidden_states_scale
         num_local_tokens = dispatch_output.num_recv_tokens_per_expert

python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -179,10 +179,10 @@ class DeepGemmRunnerCore(MoeRunnerCore):
         running_state: dict,
     ) -> torch.Tensor:
         from sglang.jit_kernel.dsv4 import silu_and_mul_contig_post_quant
+        from sglang.kernels.ops.moe.ep_moe_kernels import tma_align_input_scale
         from sglang.kernels.ops.quantization.fp8_kernel import (
             create_per_token_group_quant_fp8_output_scale,
         )
-        from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale
 
         hidden_states = runner_input.hidden_states
         hidden_states_scale = runner_input.hidden_states_scale
@@ -565,8 +565,8 @@ class DeepGemmRunnerCore(MoeRunnerCore):
         quant_info: DeepGemmMoeQuantInfo,
         running_state: dict,
     ) -> torch.Tensor:
+        from sglang.kernels.ops.moe.ep_moe_kernels import silu_and_mul_masked_fwd
         from sglang.srt.layers import deep_gemm_wrapper
-        from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_fwd
 
         hidden_states = runner_input.hidden_states
         masked_m = runner_input.masked_m
@@ -635,7 +635,7 @@ def pre_permute_standard_to_deep_gemm(
     runner_config: MoeRunnerConfig,
     running_state: dict,
 ) -> DeepGemmRunnerInput:
-    from sglang.srt.layers.moe.ep_moe.kernels import moe_ep_deepgemm_preprocess
+    from sglang.kernels.ops.moe.ep_moe_kernels import moe_ep_deepgemm_preprocess
 
     hidden_states, topk_output = (
         dispatch_output.hidden_states,
@@ -696,7 +696,7 @@ def post_permute_deep_gemm_to_standard(
     runner_config: MoeRunnerConfig,
     running_state: dict,
 ) -> StandardCombineInput:
-    from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_deepgemm
+    from sglang.kernels.ops.moe.ep_moe_kernels import post_reorder_deepgemm
     from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
 
     hidden_states_shape = running_state["hidden_states_shape"]
@@ -782,7 +782,7 @@ def pre_permute_deepep_normal_to_deep_gemm(
     runner_config: MoeRunnerConfig,
     running_state: dict,
 ) -> DeepGemmRunnerInput:
-    from sglang.srt.layers.moe.ep_moe.kernels import ep_scatter
+    from sglang.kernels.ops.moe.ep_moe_kernels import ep_scatter
 
     (
         hidden_states,
@@ -875,7 +875,7 @@ def post_permute_deep_gemm_to_deepep_normal(
     runner_config: MoeRunnerConfig,
     running_state: dict,
 ) -> DeepEPNormalCombineInput:
-    from sglang.srt.layers.moe.ep_moe.kernels import ep_gather
+    from sglang.kernels.ops.moe.ep_moe_kernels import ep_gather
     from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPNormalCombineInput
 
     hidden_states = runner_output.hidden_states
@@ -908,10 +908,10 @@ def _varlen_deep_gemm_silu_mul_quant(
     gemm1_clamp_limit: Optional[float] = None,
     num_real_tokens: Optional[int] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
+    from sglang.kernels.ops.moe.ep_moe_kernels import silu_and_mul_masked_post_quant_fwd
     from sglang.kernels.ops.quantization.fp8_kernel import (
         sglang_per_token_group_quant_8bit,
     )
-    from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd
 
     if _MASKED_GEMM_FAST_ACT:
         assert (
@@ -951,7 +951,7 @@ def _varlen_deep_gemm_silu_mul_quant(
         and G % 4 == 0
         and D % (group_size * 4) == 0
     ):
-        from sglang.srt.layers.moe.ep_moe.kernels import (
+        from sglang.kernels.ops.moe.ep_moe_kernels import (
             silu_and_mul_masked_post_quant_packed_fwd,
         )
 

python/sglang/srt/layers/moe/moe_runner/humming.py

@@ -9,8 +9,8 @@ from weakref import WeakValueDictionary
 
 import torch
 
+from sglang.kernels.ops.moe.ep_moe_kernels import moe_permute, moe_unpermute
 from sglang.srt.environ import envs
-from sglang.srt.layers.moe.ep_moe.kernels import moe_permute, moe_unpermute
 from sglang.srt.layers.moe.fused_moe_triton.moe_fused_mul_sum import moe_fused_mul_sum
 from sglang.srt.layers.moe.moe_runner.base import (
     MoeQuantInfo,
@@ -371,7 +371,7 @@ class HummingRunnerCore(MoeRunnerCore):
 
     def apply_activation(self, inputs: torch.Tensor, outputs: torch.Tensor):
         if self.activation == "silu" and self.swiglu_limit is not None:
-            from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+            from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
                 act_and_mul_triton,
             )
 

python/sglang/srt/layers/moe/moe_runner/triton.py

@@ -84,7 +84,7 @@ class TritonRunnerCore(MoeRunnerCore):
         hooks: Optional[Any] = None,
     ) -> TritonRunnerOutput:
         if quant_info.use_mxfp8 and is_hip() and is_gfx95_supported():
-            from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
+            from sglang.kernels.ops.moe.mxfp8_moe_amd_gfx95 import (
                 fused_experts_mxfp8,
             )
 
@@ -181,7 +181,7 @@ def fused_experts_none_to_triton(
     from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
 
     if quant_info.use_mxfp8 and is_hip() and is_gfx95_supported():
-        from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
+        from sglang.kernels.ops.moe.mxfp8_moe_amd_gfx95 import (
             fused_experts_mxfp8,
         )
 

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

@@ -13,6 +13,12 @@ import torch
 import torch.nn.functional as F
 import triton.language as tl
 
+from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
+    act_and_mul_triton,
+    invoke_fused_moe_kernel,
+    moe_sum_reduce_triton,
+    support_tensor_descriptor,
+)
 from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
 from sglang.srt.environ import envs
 from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
@@ -31,12 +37,6 @@ from sglang.srt.utils import (
 from sglang.srt.utils.custom_op import register_custom_op
 
 from .fused_moe_triton_config import get_config_dtype_str, try_get_optimal_moe_config
-from .fused_moe_triton_kernels import (
-    act_and_mul_triton,
-    invoke_fused_moe_kernel,
-    moe_sum_reduce_triton,
-    support_tensor_descriptor,
-)
 from .moe_align_block_size import moe_align_block_size
 
 if TYPE_CHECKING:

python/sglang/srt/layers/moe/topk.py

@@ -31,8 +31,6 @@ from typing import (
 
 import torch
 import torch.nn.functional as F
-import triton
-import triton.language as tl
 
 from sglang.srt.runtime_context import get_parallel
 
@@ -1209,75 +1207,16 @@ def biased_grouped_topk_impl(
     return topk_weights, topk_ids
 
 
+from sglang.kernels.ops.moe.fill_padded_rows import (
+    _can_fuse_padded_region,
+    _fill_padded_rows,
+)
+
+
 def is_power_of_two(n):
     return n > 0 and math.log2(n).is_integer()
 
 
-@triton.jit
-def _fill_padded_rows_kernel(
-    out_ptr,
-    num_token_non_padded_ptr,
-    n_cols,
-    fill_value,
-    stride_row,
-    BLOCK_COLS: tl.constexpr,
-):
-    row = tl.program_id(0)
-    n_valid = tl.load(num_token_non_padded_ptr)
-    if row >= n_valid:
-        cols = tl.arange(0, BLOCK_COLS)
-        mask = cols < n_cols
-        ptrs = out_ptr + row * stride_row + cols
-        fill = tl.full((BLOCK_COLS,), fill_value, dtype=out_ptr.dtype.element_ty)
-        tl.store(ptrs, fill, mask=mask)
-
-
-def _can_fuse_padded_region(x: torch.Tensor) -> bool:
-    # The fused kernel uses one program per row and assumes a row-major 2D
-    # tensor (columns contiguous); fall back to eager for anything else.
-    return x.dim() == 2 and x.stride(1) == 1
-
-
-def _fill_padded_rows(
-    x: torch.Tensor,
-    num_token_non_padded: torch.Tensor,
-    fill_value,
-) -> None:
-    """Set ``x[row, :] = fill_value`` for every padded row (row index
-    ``>= num_token_non_padded``) using a single Triton launch.
-
-    Replaces the eager ``arange + (>=) + boolean index_put_`` sequence, which
-    issues several launch-latency-bound kernels per call. The grid is static
-    (one program per row) and the pad count is read from device memory inside
-    the kernel, so this is safe to capture inside a CUDA/HIP graph.
-    """
-    # Metadata-only checks (no device sync): the kernel reads a single scalar
-    # routing count from device memory, so it must be a 1-element integer tensor
-    # on the same device as ``x``.
-    assert isinstance(
-        num_token_non_padded, torch.Tensor
-    ), "num_token_non_padded must be a torch.Tensor"
-    assert num_token_non_padded.numel() == 1, (
-        "num_token_non_padded must be a single-element tensor, got shape "
-        f"{tuple(num_token_non_padded.shape)}"
-    )
-    assert (
-        not num_token_non_padded.dtype.is_floating_point
-    ), f"num_token_non_padded must be an integer tensor, got {num_token_non_padded.dtype}"
-    assert (
-        num_token_non_padded.device == x.device
-    ), "num_token_non_padded and x must be on the same device"
-    n_rows, n_cols = x.shape
-    _fill_padded_rows_kernel[(n_rows,)](
-        x,
-        num_token_non_padded,
-        n_cols,
-        fill_value,
-        x.stride(0),
-        BLOCK_COLS=triton.next_power_of_2(n_cols),
-    )
-
-
 def _eplb_remap_enabled() -> bool:
     # A real logical->physical mapping only exists when EPLB is enabled, the
     # initial expert placement is non-trivial, or there are redundant physical
@@ -1299,74 +1238,6 @@ def _eplb_remap_enabled() -> bool:
     )
 
 
-@triton.jit
-def _fill_padded_rows_kernel(
-    out_ptr,
-    num_token_non_padded_ptr,
-    n_cols,
-    fill_value,
-    stride_row,
-    BLOCK_COLS: tl.constexpr,
-):
-    row = tl.program_id(0)
-    n_valid = tl.load(num_token_non_padded_ptr)
-    if row >= n_valid:
-        cols = tl.arange(0, BLOCK_COLS)
-        mask = cols < n_cols
-        ptrs = out_ptr + row * stride_row + cols
-        fill = tl.full((BLOCK_COLS,), fill_value, dtype=out_ptr.dtype.element_ty)
-        tl.store(ptrs, fill, mask=mask)
-
-
-def _can_fuse_padded_region(x: torch.Tensor) -> bool:
-    # The fused kernel uses one program per row and assumes a row-major 2D
-    # tensor (columns contiguous); fall back to eager for anything else.
-    return x.dim() == 2 and x.stride(1) == 1
-
-
-def _fill_padded_rows(
-    x: torch.Tensor,
-    num_token_non_padded: torch.Tensor,
-    fill_value,
-) -> None:
-    """Set ``x[row, :] = fill_value`` for every padded row (row index
-    ``>= num_token_non_padded``) using a single Triton launch.
-
-    Replaces the eager ``arange + (>=) + boolean index_put_`` sequence, which
-    issues several launch-latency-bound kernels per call. The grid is static
-    (one program per row) and the pad count is read from device memory inside
-    the kernel, so this is safe to capture inside a CUDA/HIP graph.
-    """
-    # Metadata-only checks (no device sync): the kernel reads a single scalar
-    # routing count from device memory, so it must be a 1-element integer tensor
-    # on the same device as ``x``. Use explicit raises (not asserts) so the
-    # checks survive ``python -O`` and invalid inputs fail loudly instead of
-    # turning into opaque Triton/memory errors.
-    if not isinstance(num_token_non_padded, torch.Tensor):
-        raise TypeError("num_token_non_padded must be a torch.Tensor")
-    if num_token_non_padded.numel() != 1:
-        raise ValueError(
-            "num_token_non_padded must be a single-element tensor, got shape "
-            f"{tuple(num_token_non_padded.shape)}"
-        )
-    if num_token_non_padded.dtype.is_floating_point:
-        raise TypeError(
-            "num_token_non_padded must be an integer tensor, got "
-            f"{num_token_non_padded.dtype}"
-        )
-    if num_token_non_padded.device != x.device:
-        raise ValueError("num_token_non_padded and x must be on the same device")
-    n_rows, n_cols = x.shape
-    _fill_padded_rows_kernel[(n_rows,)](
-        x,
-        num_token_non_padded,
-        n_cols,
-        fill_value,
-        x.stride(0),
-        BLOCK_COLS=triton.next_power_of_2(n_cols),
-    )
-
-
 def _mask_topk_ids_padded_region(
     topk_ids: torch.Tensor,
     num_token_non_padded: Optional[torch.Tensor] = None,
@@ -1917,7 +1788,7 @@ def _post_process_topk_ids(
         shared_id_base = ep_rank * num_local_experts + num_local_routed
 
         # Lazy import to avoid circular-import issues
-        from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+        from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
             fused_append_remap_shared_experts_deepep,
         )
 
@@ -1938,7 +1809,7 @@ def _post_process_topk_ids(
         )
 
         # Lazy import to avoid circular-import issues
-        from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+        from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
             fused_append_shared_experts,
         )
 

python/sglang/srt/layers/moe/waterfill.py

@@ -13,313 +13,21 @@
 # ==============================================================================
 """Waterfill: shared expert as 9th routed expert, dispatched to least-loaded rank."""
 
-from typing import NamedTuple, Optional, Tuple
+from typing import Optional, Tuple
 
 import torch
-import triton
-import triton.language as tl
 from torch import Tensor
 
+from sglang.kernels.ops.moe.deepep_waterfill_kernels import (
+    LOCAL_SHARED_MARKER,
+    WaterfillDispatchPlan,
+    _count_routed_per_rank_kernel,
+    _empty_expanded,
+    materialize_waterfill_dispatch_fused,
+)
 from sglang.srt.environ import envs
 from sglang.srt.layers.moe.topk import StandardTopKOutput
 
-LOCAL_SHARED_MARKER = -1  # Invalid expert ID; DeepEP ignores expert_id < 0.
-_LOCAL_PREF_NUMER = 11  # local-rank preference = 11/10
-_LOCAL_PREF_DENOM = 10
-
-
-class WaterfillDispatchPlan(NamedTuple):
-    """Inputs needed by the fused Waterfill expansion path."""
-
-    # Effective rank load consumed by the fused kernel.
-    rank_load: Tensor
-    allow_all_ranks: bool
-    target_total: int
-
-
-def _empty_expanded(topk_ids: Tensor, topk_weights: Tensor):
-    """Return empty expanded tensors for zero-token batches."""
-    topk, d = topk_ids.shape[1], topk_ids.device
-    return (
-        torch.empty(0, topk + 1, dtype=topk_ids.dtype, device=d),
-        torch.empty(0, topk + 1, dtype=topk_weights.dtype, device=d),
-    )
-
-
-@triton.jit
-def _count_routed_per_rank_kernel(
-    topk_ids_ptr,  # [num_tokens, topk]
-    counts_ptr,  # [world_size] output (atomic add)
-    num_tokens,
-    topk: tl.constexpr,
-    experts_per_rank,
-    world_size: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """Count routed tokens per rank using block-level histogram."""
-    pid = tl.program_id(0)
-    token_idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    mask = token_idx < num_tokens
-
-    for r in range(world_size):
-        rank_count = tl.zeros([BLOCK_SIZE], dtype=tl.int64)
-
-        for k in range(topk):
-            expert_id = tl.load(
-                topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1
-            ).to(tl.int64)
-            valid = expert_id >= 0
-            target_rank = expert_id // experts_per_rank
-            target_rank = tl.minimum(tl.maximum(target_rank, 0), world_size - 1)
-            rank_count += tl.where(
-                mask & valid & (target_rank == r),
-                tl.full([BLOCK_SIZE], 1, dtype=tl.int64),
-                tl.zeros([BLOCK_SIZE], dtype=tl.int64),
-            )
-
-        block_total = tl.sum(rank_count)
-        if block_total > 0:
-            tl.atomic_add(counts_ptr + r, block_total)
-
-
-@triton.jit
-def _waterfill_expand_kernel(
-    topk_ids_ptr,
-    topk_weights_ptr,
-    rank_load_ptr,
-    expanded_ids_ptr,
-    expanded_weights_ptr,
-    num_tokens,
-    topk: tl.constexpr,
-    old_experts_per_rank,
-    new_experts_per_rank,
-    world_size: tl.constexpr,
-    source_rank,
-    shared_weight,
-    local_marker,
-    local_pref_numer,
-    local_pref_denom,
-    precomputed_target_total,
-    ALLOW_ALL_RANKS: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """Fused waterfill + expand. ID remap: old_id -> old_id + old_id // old_epr."""
-    pid = tl.program_id(0)
-    token_idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    mask = token_idx < num_tokens
-
-    r_idx = tl.arange(0, world_size)
-    rank_load_vec = tl.load(rank_load_ptr + r_idx, mask=r_idx < world_size, other=0).to(
-        tl.int64
-    )
-    total_effective_k = tl.sum(rank_load_vec)
-    total_tokens_global_k = total_effective_k // topk
-    derived_target_total = (
-        total_effective_k + total_tokens_global_k + world_size - 1
-    ) // world_size
-    target_total = tl.where(
-        precomputed_target_total > 0,
-        precomputed_target_total,
-        derived_target_total,
-    )
-
-    # Step 1: Select destination rank for shared expert (waterfill sampling).
-    source_count = tl.load(rank_load_ptr + source_rank)
-    best_count = tl.where(mask, source_count, 2**30)
-    best_rank = tl.full([BLOCK_SIZE], source_rank, dtype=tl.int64)
-    has_valid = tl.zeros([BLOCK_SIZE], dtype=tl.int1)
-    src_rank_i32 = tl.full([BLOCK_SIZE], source_rank, dtype=tl.int32)
-
-    if ALLOW_ALL_RANKS:
-        candidate_mask = tl.full([BLOCK_SIZE], (1 << world_size) - 1, dtype=tl.int32)
-        for r in range(world_size):
-            target_count = tl.load(rank_load_ptr + r).to(tl.int64)
-            better = (
-                target_count * local_pref_numer < best_count * local_pref_denom
-            ) & mask
-            best_count = tl.where(better, target_count, best_count)
-            best_rank = tl.where(
-                better, tl.full([BLOCK_SIZE], r, dtype=tl.int64), best_rank
-            )
-    else:
-        candidate_mask = (tl.full([BLOCK_SIZE], 1, dtype=tl.int32) << src_rank_i32).to(
-            tl.int32
-        )
-
-    for k in range(topk):
-        expert_id = tl.load(
-            topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1
-        ).to(tl.int64)
-        valid = expert_id >= 0
-        has_valid = has_valid | valid
-
-        if not ALLOW_ALL_RANKS:
-            target_rank = expert_id // old_experts_per_rank
-            target_rank = tl.minimum(tl.maximum(target_rank, 0), world_size - 1)
-            target_rank_i32 = target_rank.to(tl.int32)
-            shift_amt = tl.where(valid, target_rank_i32, 0)
-            bit = tl.full([BLOCK_SIZE], 1, dtype=tl.int32) << shift_amt
-            candidate_mask = tl.where(
-                valid & mask, candidate_mask | bit, candidate_mask
-            )
-
-            target_count = tl.load(
-                rank_load_ptr + target_rank, mask=mask & valid, other=2**30
-            )
-
-            better = (
-                (target_count * local_pref_numer < best_count * local_pref_denom)
-                & valid
-                & mask
-            )
-            best_count = tl.where(better, target_count, best_count)
-            best_rank = tl.where(better, target_rank, best_rank)
-
-    total_w = tl.zeros([BLOCK_SIZE], dtype=tl.int32)
-    for r in range(world_size):
-        present = ((candidate_mask >> r) & 1) == 1
-        rank_load_r = tl.load(rank_load_ptr + r).to(tl.int64)
-        w = tl.where(target_total > rank_load_r, target_total - rank_load_r, 0).to(
-            tl.int32
-        )
-        w_vec = tl.full([BLOCK_SIZE], w, dtype=tl.int32)
-        w_vec = tl.where(
-            src_rank_i32 == r,
-            w_vec,
-            (w_vec * local_pref_denom) // local_pref_numer,
-        )
-        total_w += tl.where(present, w_vec, 0)
-
-    token_seed = token_idx.to(tl.uint32) ^ (
-        src_rank_i32.to(tl.uint32) * tl.full([BLOCK_SIZE], 0x9E3779B9, dtype=tl.uint32)
-    )
-    token_seed = token_seed * tl.full([BLOCK_SIZE], 1664525, dtype=tl.uint32) + tl.full(
-        [BLOCK_SIZE], 1013904223, dtype=tl.uint32
-    )
-    u = tl.where(total_w > 0, token_seed % total_w.to(tl.uint32), 0).to(tl.int32)
-
-    chosen = src_rank_i32
-    cum = tl.zeros([BLOCK_SIZE], dtype=tl.int32)
-    for r in range(world_size):
-        present = ((candidate_mask >> r) & 1) == 1
-        rank_load_r = tl.load(rank_load_ptr + r).to(tl.int64)
-        w = tl.where(target_total > rank_load_r, target_total - rank_load_r, 0).to(
-            tl.int32
-        )
-        w_vec = tl.full([BLOCK_SIZE], w, dtype=tl.int32)
-        w_vec = tl.where(
-            src_rank_i32 == r,
-            w_vec,
-            (w_vec * local_pref_denom) // local_pref_numer,
-        )
-        w_vec = tl.where(present, w_vec, 0)
-        pick = (total_w > 0) & present & (u >= cum) & (u < (cum + w_vec))
-        chosen = tl.where(pick, r, chosen)
-        cum += w_vec
-
-    best_rank = tl.where(total_w > 0, chosen.to(tl.int64), best_rank)
-
-    # Step 2: Compute shared expert ID and local mask.
-    is_local = best_rank == source_rank
-    local_shared_id = source_rank * new_experts_per_rank + old_experts_per_rank
-    remote_shared_id = best_rank * new_experts_per_rank + old_experts_per_rank
-    shared_expert_id = tl.where(
-        is_local,
-        tl.full([BLOCK_SIZE], local_shared_id, dtype=tl.int64),
-        remote_shared_id,
-    ).to(tl.int64)
-    shared_expert_id = tl.where(
-        has_valid,
-        shared_expert_id,
-        tl.full([BLOCK_SIZE], local_marker, dtype=tl.int64),
-    )
-
-    # Step 3: Copy and remap topk_ids, copy weights.
-    for k in range(topk):
-        old_id = tl.load(topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1).to(
-            tl.int64
-        )
-        valid_id = old_id >= 0
-        new_id = tl.where(valid_id, old_id + (old_id // old_experts_per_rank), old_id)
-        tl.store(expanded_ids_ptr + token_idx * (topk + 1) + k, new_id, mask=mask)
-
-    for k in range(topk):
-        val = tl.load(topk_weights_ptr + token_idx * topk + k, mask=mask, other=0.0)
-        expert_id = tl.load(
-            topk_ids_ptr + token_idx * topk + k, mask=mask, other=-1
-        ).to(tl.int64)
-        val = tl.where(expert_id >= 0, val, 0.0)
-        tl.store(expanded_weights_ptr + token_idx * (topk + 1) + k, val, mask=mask)
-
-    # Step 4: Write shared expert column.
-    tl.store(
-        expanded_ids_ptr + token_idx * (topk + 1) + topk,
-        shared_expert_id,
-        mask=mask,
-    )
-    tl.store(
-        expanded_weights_ptr + token_idx * (topk + 1) + topk,
-        tl.where(has_valid, shared_weight, 0.0),
-        mask=mask,
-    )
-
-
-def materialize_waterfill_dispatch_fused(
-    topk_ids: Tensor,
-    topk_weights: Tensor,
-    rank_load: Tensor,
-    num_routed_experts: int,
-    world_size: int,
-    source_rank: int,
-    shared_weight: float,
-    allow_all_ranks: bool = False,
-    target_total: int = 0,
-) -> Tuple[Tensor, Tensor]:
-    """Run fused Waterfill rank selection and TopK expansion.
-
-    The Triton kernel intentionally selects each token's shared-expert rank and
-    writes the expanded TopK layout in one pass.
-    """
-    num_tokens = topk_ids.shape[0]
-    topk = topk_ids.shape[1]
-    old_experts_per_rank = num_routed_experts // world_size
-    new_experts_per_rank = old_experts_per_rank + 1
-    device = topk_ids.device
-
-    if num_tokens == 0:
-        return _empty_expanded(topk_ids, topk_weights)
-
-    expanded_topk_ids = torch.empty(
-        num_tokens, topk + 1, dtype=topk_ids.dtype, device=device
-    )
-    expanded_topk_weights = torch.empty(
-        num_tokens, topk + 1, dtype=topk_weights.dtype, device=device
-    )
-    BLOCK_SIZE = 256
-    grid = ((num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE,)
-    _waterfill_expand_kernel[grid](
-        topk_ids,
-        topk_weights,
-        rank_load,
-        expanded_topk_ids,
-        expanded_topk_weights,
-        num_tokens,
-        topk,
-        old_experts_per_rank,
-        new_experts_per_rank,
-        world_size,
-        source_rank,
-        shared_weight,
-        LOCAL_SHARED_MARKER,
-        _LOCAL_PREF_NUMER,
-        _LOCAL_PREF_DENOM,
-        target_total,
-        allow_all_ranks,
-        BLOCK_SIZE,
-    )
-
-    return expanded_topk_ids, expanded_topk_weights
-
 
 @torch.compile(dynamic=True)
 def expand_topk_with_shared_expert(

python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py

@@ -31,7 +31,7 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 if _use_aiter:
     from aiter.ops.shuffle import shuffle_weight
 
-    from sglang.srt.layers.moe.rocm_moe_utils import rocm_fused_experts_tkw1
+    from sglang.kernels.ops.moe.rocm_moe_utils import rocm_fused_experts_tkw1
 
 
 class QuarkW8A8FP8MoE(QuarkMoEScheme):

python/sglang/srt/lora/lora_moe_runner_marlin.py

@@ -25,15 +25,15 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
     from sglang.kernels.ops.activation import silu_and_mul
+    from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
+        moe_sum_reduce_triton,
+    )
     from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (
         get_scalar_type,
     )
     from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
         moe_align_block_size,
     )
-    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
-        moe_sum_reduce_triton,
-    )
     from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
 
 

python/sglang/srt/models/grok.py

@@ -22,6 +22,7 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
+from sglang.kernels.ops.moe.router import fused_moe_router_shim
 from sglang.srt.distributed import (
     tensor_model_parallel_all_reduce,
 )
@@ -40,7 +41,6 @@ from sglang.srt.layers.linear import (
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
-from sglang.srt.layers.moe.router import fused_moe_router_shim
 from sglang.srt.layers.moe.topk import TopK
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention

python/sglang/srt/models/longcat_flash.py

@@ -37,6 +37,7 @@ from typing import Iterable, List, Optional, Tuple
 import torch
 from torch import nn
 
+from sglang.kernels.ops.moe.ep_moe_kernels import zero_experts_compute_triton
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.configs import LongcatFlashConfig
 from sglang.srt.distributed import (
@@ -57,7 +58,6 @@ from sglang.srt.layers.linear import (
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
-from sglang.srt.layers.moe.ep_moe.kernels import zero_experts_compute_triton
 from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK

python/sglang/srt/models/qwen2_moe.py

@@ -392,7 +392,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             return topk_output
         shared_weights, shared_scale = shared
 
-        from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+        from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
             fused_append_shared_experts_with_weights,
         )