Motivation

TritonRunnerCore.run (in moe_runner/triton.py) and fused_experts_impl (in fused_moe_triton/fused_moe.py) had grown to ~95% identical logic — same two-kernel + activation + combine pipeline, same platform dispatch ladders, same activation variants — but with subtle divergences (runner missed filter_expert, TMA, enable_fused_moe_sum_all_reduce, non-gated activations, the sgl-kernel moe_sum_reduce path, and the HIP small-token torch.compile branch). Every future kernel change had to be duplicated, with a constant risk of drift — and the LoRA path, which is the sole live consumer of the runner, was the one getting the stale copy.

This PR reconciles both paths and merges them onto a single shared implementation while preserving every current behavior (LoRA hooks, fused-func fast path, end-to-end model accuracy).

Modifications

Three commits, each independently reviewable:

1. refactor(moe): sync triton runner pipeline with fused_experts_impl (9463977)

   • Drop the 64K-token chunk loop in fused_experts_impl (single-shot over all tokens). max_block_m no longer needed.
   • Bring the runner path up to parity with fused_experts_impl: filter_expert, down_moe_use_tma (TMA), enable_fused_moe_sum_all_reduce, non-gated silu/gelu/relu2, sgl_kernel.moe_sum_reduce on CUDA, HIP small-token torch.compile branch, PyTorch fallbacks when vllm_ops is missing.
   • Unify platform flags/imports (get_bool_env_var, get_moe_padding_size, _has_vllm_ops) across the two files.
   • pre_permute_standard_to_triton now calls try_get_optimal_moe_config(..., return_down_config=True) and stashes down_config/down_moe_use_tma in running_state.
   • intermediate_cache1/2/3 allocated at their original logical sites (cache2 right before activation, cache3 right before the second kernel); each del'd once fully consumed so the caching allocator can reuse memory.
   • LoRA hooks (after_gate_up/after_down) and the or hooks widening of _use_intermediate are preserved as the only runner-only additions.

2. refactor(moe): extract shared _fused_moe_kernel_sequence helper (3e4b8c5)

   • Extract two helpers into fused_moe.py:
     • _prepare_fused_moe_run: resolves padded_size, config_dtype, optimal config (with down_config + TMA flag), and moe_align_block_size. Used by fused_experts_impl and pre_permute_standard_to_triton.
     • _fused_moe_kernel_sequence: runs the kernel/activation/kernel/combine sequence over already-aligned inputs. Takes an optional LoRA hooks object; the second-kernel write target and the CUDA topk==1, routed==1.0 shortcut both use a unified _use_intermediate = not no_combine and (topk != 1 or hooks) guard, which is a no-op for hooks=None and preserves the runner's widening when hooks are present.
   • Callers collapse to thin adapters:
     • fused_experts_impl: assertions → _prepare_fused_moe_run → _fused_moe_kernel_sequence(hooks=None).
     • TritonRunnerCore.run: derive filter_expert → _fused_moe_kernel_sequence(hooks=hooks) on pre-aligned runner inputs.
     • pre_permute_standard_to_triton: _prepare_fused_moe_run → stash config/down_config/down_moe_use_tma in running_state.
   • moe_runner/triton.py loses all its platform-dispatch plumbing; the entire CUDA/HIP/XPU/vllm-fallback ladder lives in one place.

3. refactor(moe): move triton util modules under moe_runner/triton_utils (bae9029)

   • Relocate the pure-Triton MoE utilities out of layers/moe/fused_moe_triton/ into layers/moe/moe_runner/triton_utils/, co-locating them with the runner that drives them. FusedMoE / FusedMoeWeightScaleSupported / fused_marlin_moe / layer.py / triton_kernels_moe.py stay (higher-level or unrelated).
   • Files moved (git-rename preserved): fused_moe.py, fused_moe_triton_config.py, fused_moe_triton_kernels.py, moe_align_block_size.py.
   • fused_moe_triton/__init__.py now re-exports the public surface (fused_experts, moe_align_block_size, try_get_optimal_moe_config, get_config_file_name, override_config, get_config) from the new location, so existing callers using the package-level API are unchanged. 24 direct submodule importers (tests, benchmarks, models, quantization, lora, topk, runner/triton.py, 3rdparty tuning) are updated to the new path.
   • _config / override_config / get_config (the context-manager-based config override) move into the new triton_utils/__init__.py; the fused_moe_triton_config.py late-import of get_config is repointed accordingly.

Net effect: moe_runner/triton.py went from ~500 LOC of duplicated pipeline to ~240 LOC of pure adapter code.

Accuracy Tests

Unit parity (bit-identical vs fused_experts), re-run after each commit:

[OK] M=32 topk=2 act=silu                diff=0.00e+00
[OK] M=32 topk=1 act=silu                diff=0.00e+00
[OK] M=32 topk=4 act=silu                diff=0.00e+00
[OK] M=64 topk=2 act=silu routed=2.5     diff=0.00e+00
[OK] M=8  topk=2 act=silu                diff=0.00e+00
[OK] M=32 topk=2 act=gelu                diff=0.00e+00
[OK] M=32 topk=2 act=silu no_combine     diff=0.00e+00
[OK] M=32 topk=2 act=silu inplace        diff=0.00e+00
[OK] M=128 topk=4 act=silu               diff=0.00e+00
[OK] M=128 topk=4 act=silu routed=0.5    diff=0.00e+00
[OK] M=70000 topk=2 act=silu             diff=0.00e+00

M=70000 exercises what used to be the chunked path (>64K tokens).

GLM-4.5-Air-FP8 GSM8K with SGLANG_CI_DISABLE_MOE_FUSED_FUNC=1 (forces the pre_permute → run → post_permute path that LoRA depends on, 100 examples, TP=2):

test/registered/moe/test_glm4_moe_models.py threshold is 0.80; prior baseline ~0.85.

Speed Tests and Profiling

No measurable perf impact expected — the runner path was missing features (TMA, enable_fused_moe_sum_all_reduce, sgl-kernel moe_sum_reduce, small-token torch.compile branch) that are now enabled; the fused-impl path is unchanged except the dormant chunk loop is gone. No dedicated benchmarks run.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. (existing test/registered/moe/test_glm4_moe_models.py covers the live path; parity check ran locally.)
• Update documentation according to Write documentations. (no user-facing API change)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.