[MoE] Latent-MoE runner, fused MoE tail, SituGLU activation and single-group top-k routing

[zyongye]

Purpose

Extends the existing fused-MoE stack with the pieces needed by latent-MoE
models, split out of the much larger Kimi K3 model PR
(#50000) so the MoE work can be
reviewed independently of the model, attention, tokenizer and parser changes.

Nothing here registers a new model. Every addition is either a new kernel, a
new opt-in code path, or plumbing for a new activation/routing mode; default
execution paths are unchanged.

What's in it

Seven self-contained commits, reviewable in order:

1. [Kernel][MoE] fused single-group top-k routing — adds a dedicated
   path for the degenerate grouped-topk case (n_group == 1 && topk_group == 1),
   which is plain top-k over all experts. invokeNoAuxTc tries it first and
   falls back to the existing grouped kernels. Two kernels are dispatched from a
   compile-time (num_experts, topk) tier list: a block-per-token kernel with a
   hierarchical per-warp reduction for high expert counts, and a warp-per-token
   kernel for the large-batch case. moeTopKFuncs.cuh gains a generalized
   bitonic / odd-even-merge Sort<N> for any N in [1, 64] (was N <= 4, or
   a multiple of 4 up to 16), a reduceTopKForLane that leaves the k-th result
   in lane k instead of replicating all K results in every lane, and a
   redux.sync.max.u32 warp reduction on SM100+.

2. [Kernel] SituGLU (situ_and_mul) — the gated activation used by Kimi
   models:

   gate_out = beta * tanh(gate / beta) * sigmoid(gate)
   up_out   = linear_beta * tanh(up / linear_beta)   # if linear_beta > 0
   out      = gate_out * up_out
   

   Adds the SituAndMul CustomOp plus two CUDA kernels: situ_and_mul for the
   dense [..., 2 * d] layout and masked_situ_and_mul for the batched
   [E, T, 2 * d] MoE layout, which reads expert_num_tokens and leaves padded
   rows untouched. Both accumulate in fp32 and write straight to the output
   buffer — the pure-torch fallback allocates ~8 fp32 temporaries per call,
   which is prohibitive inside a MoE layer.

3. [MoE] wire SituGLU through fused MoE — MoEActivation.SITU plus the
   activation_situ_beta / activation_situ_linear_beta parameters, plumbed
   from FusedMoE(...) through FusedMoEConfig to the experts.

   Declared supported on Marlin and DeepGEMM FP4 only (plus the TRTLLM-Gen
   backends in commit 5). Both betas reach the kernel from FusedMoEConfig on
   every one of those paths:

   • The Marlin helpers (fused_marlin_moe, batched_fused_marlin_moe,
     _fused_marlin_moe) take the two betas as explicit arguments and forward
     them to their activation callback, alongside the existing clamp_limit /
     gemm1_alpha / gemm1_beta quant-config parameters. This keeps the
     module-level apply_moe_activation default correct instead of depending on
     the caller supplying a config-aware callback.
   • FusedMoEExpertsModular.activation accepts the betas as well, preferring
     an explicit value and falling back to self.moe_config, so it stays
     interchangeable with the module-level function as a callback.
   • Batched Marlin's callback routes SITU to masked_situ_and_mul, which
     respects expert_num_tokens instead of running over padded rows.
   • DeepGEMM FP4 calls self.activation. SILU keeps its fused gate+mul+quant
     kernels, while the general gated activations apply the activation then
     FP8-requant into the layout matching the active scale format.

   beta is required rather than defaulted — a missing value means the caller
   bypassed the config plumbing, so both the generic path and the batched
   Marlin callback assert instead of silently substituting 1.0. linear_beta
   stays optional: <= 0 signals "unset" to the kernel and passes up
   through, matching SituAndMul(linear_beta=None).

4. [MoE] DeepSeekV3 routing on the TRTLLM-Gen MXFP4 monolithic path —
   TrtLlmMxfp4ExpertsMonolithic hardcoded the routing arguments to the
   renormalize case: logits cast to BF16, and bias / group counts / routed
   scaling factor all passed as None. That silently excluded sigmoid +
   grouped-topk routing even though the TRTLLM-Gen kernel implements it and the
   caller already threads all four values into apply. Forwards them and stops
   downcasting the logits (DeepSeekV3 routing wants FP32; the renormalize path
   still receives BF16 from the caller).

5. [MoE] SituGLU for the TRTLLM-Gen MXFP4/NVFP4 backends — the TRTLLM-Gen
   SituGLU cubin computes exactly situ_and_mul, so activation_situ_beta maps
   to gatedActAlpha and activation_situ_linear_beta to gatedActBeta.
   Unlike the SwiGLU-OAI clamp/beta these act on the dequantized gate/up rather
   than the raw GEMM1 accumulator, so on the NVFP4 path they are registered
   as-is instead of being folded by g1_alphas.

   Dependency: needs a FlashInfer build exposing ActivationType.Situ and
   the is_private kernel selector ([feat] Add SITU trtllmgen MOE flashinfer-ai/flashinfer#4180). The
   currently pinned flashinfer-python==0.6.15.post1 does not have it. This
   commit is isolated so it can be dropped or held if reviewers prefer to
   land the rest first.

6. [Kernel] CuTe DSL fused latent-MoE tail — KimiK3LatentMoETailOp.
   Given the TP-partial latent routed output and the TP-partial shared-expert
   output, it produces the fully reduced hidden-size result in one fused
   sequence instead of two NCCL all-reduces plus a separate up projection.
   Three kernels back it, all over symmetric memory: CollectiveKernel
   (all-reduce + RMSNorm + reduce-scatter with a Lamport-flag early exit, so
   ranks proceed as their peers' data lands), AdaptiveUpProjectionKernel
   (multicast GEMM for the replicated latent up projection, folding the
   shared-expert add into its epilogue) and LamportCopyKernel. The op is
   process-wide cached on a contract (TP group, dtype, shapes, eps) and
   registers a CuTe DSL warmup provider so the JIT compile happens during warmup
   rather than on the first forward.

   Decode-shaped only: TP 8 or 16, BF16, SM100, up to 16 tokens per call. Gated
   behind VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION (default off).

7. [MoE] LatentMoERunner — latent MoE keeps the routed experts in a
   lower-dimensional latent space and projects back to hidden size afterwards,
   so the runner must reduce the routed output before the output transform:
   that transform contains an RMSNorm, and normalizing TP-partial values then
   summing is not the same as summing then normalizing. MoERunner gains
   _maybe_reduce_routed_output_before_transform for this, and the two
   downstream reduction helpers now take the resulting "already reduced" flag
   explicitly instead of always reading _fused_output_is_reduced.

   LatentMoERunner adds a fused path for the common case (TP > 1, shared
   expert present, un-reduced combine output, no sequence parallelism): the up
   projection is replicated, so the latent partial can be all-reduced and
   RMSNormed in one FlashInfer fused collective then up-projected locally, with
   the shared-expert all-reduce overlapped on the aux stream at decode batch
   sizes. It optionally dispatches to the commit-6 CuTe DSL tail. The base path
   stays correct at any TP size, just with two collectives instead of one.

   Also generalizes the modular-kernel output_alias fast path to CUDA. It was
   gated on ROCm + AITER, but the redundant write-back copy it avoids is
   platform-independent; ROCm behavior is unchanged.

Not duplicating an existing PR

Checked per AGENTS.md:

gh pr list --repo vllm-project/vllm --state open --search "grouped topk routing kernel"
gh pr list --repo vllm-project/vllm --state open --search "single group topk"
gh pr list --repo vllm-project/vllm --state open --search "situ activation"
gh pr list --repo vllm-project/vllm --state open --search "SituGLU"
gh pr list --repo vllm-project/vllm --state open --search "latent moe runner"
gh pr list --repo vllm-project/vllm --state open --search "moe tail fusion"

• No open PR touches SituGLU/SITU, the latent-MoE tail, or a single-group top-k
  fast path.
• Fuse grouped_topk routing for bias-free softmax models (Mistral Large 3) #46869 (group scoring max vs top2 for bias-free models) and [Kernel] Fuse softmax into grouped_topk CUDA kernel #45120 (fusing
  softmax into grouped_topk) also edit grouped_topk_kernels.cu, but for
  different features — neither adds an n_group == 1 && topk_group == 1 fast
  path. Textual conflicts are likely; happy to rebase behind whichever lands
  first.
• Opt-in two-stream + launch-elision latency optimizations for Kimi-K2.6-NVFP4 decode on B300 (TP=4) #45452 (opt-in two-stream Kimi-K2.6 decode) and [BugFix] Fix latent-MoE crash when MoE runner holds the gate (is_internal_router) #46077 (latent-MoE
  internal-router crash fix) also edit moe_runner.py, but neither overlaps
  with the reduce-before-transform change or LatentMoERunner.
• This PR is a subset of [New model] Kimi K3 #50000 (Kimi K3), split out at reviewer-friendly size.
  [New model] Kimi K3 #50000 will be rebased onto whatever lands here.

Test plan

Not yet run — this is why the PR is a draft. Filling these in before
marking it ready for review.

Intended commands:

pytest tests/kernels/moe/test_grouped_topk.py -v
pytest tests/kernels/core/test_activation.py -v -k situ
pytest tests/models/kimi_k3/test_latent_moe_tail.py -v   # needs 8x SM100

Kernel performance for the latent-MoE tail is measured with the benchmark added
in this PR:

python benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py

Test results

To be added.

Model evaluation results

To be added. Commits 3–5 and 7 can affect model output, so accuracy numbers are
required before this is ready for review.

AI assistance

AI assistance (Claude Code) was used to extract these changes from #50000 into a
reviewable commit series and to write the commit messages and this description.
The underlying implementation is from #50000. A human submitter is reviewing
every changed line and will run the tests and evals above before marking this
ready for review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zyongye.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork