[lucifer1004]

Description

Fixes #20540.

Kimi K3 decode with the MXFP4 TRTLLM-Gen MoE ran the router GEMM, the latent down projection, the input quantization, the routing kernel and FC1 / FC2 in order on the main stream. This PR lets the router GEMM and the routing kernel run on their own stream. Both are opt-in for now.

Commit 1: routing as its own op (mxFp4BlockScaleMoe.cpp, trtllm_gen_custom_ops.py, cpp_custom_ops.py)

• trtllm::mxe2m1_moe_routing_workspace allocates the routing outputs that FC1 / FC2 / finalize read. trtllm::mxe2m1_moe_routing runs the same routing kernel the runner runs, on the current stream.
• The MXE4M3 x MXE2M1 runner gets run_moe_routed, which takes that workspace and skips its routing, and get_tile, which reports the tile the workspace must be computed for. The runner checks that tile.
• The runner's existing per-call allocations now go through the same workspace spec, so the non-pre-routed path allocates the same tensors as before.
• In mxe4m3_mxe2m1_block_scale_moe_runner (KIMI_K3_MOE_PRE_ROUTING=1):
  • every non-tuning call records the tile and the routing inputs per (gemm1 weights, num_tokens);
  • mxe2m1_pre_routing_prepare / mxe2m1_pre_routing_launch let a caller launch the routing on another stream for a recorded shape;
  • the op joins that stream's event before reading the workspace, and uses it only if tile and routing inputs match, otherwise it routes itself;
  • different logits raise, because they would mean an unsynchronized read.
• KIMI_K3_MOE_PRE_ROUTING_MAX_TOKENS (default 8) limits pre-routing to small decode batches.

Commit 2: Kimi K3 MoE router stream (modeling_kimi_linear.py, utils.py)

• With KIMI_K3_MOE_ROUTER_STREAM=1, the model creates one more aux stream (AuxStreamType.MoeRouter).
• KimiK3MoERuntime then uses _forward_router_stream in multi-stream decode when none of the latent stream, deferred shared all-reduce or fused router+down projection applies (those are CuteDSL / MegaMoE configurations):
  • router stream: router GEMM -> router event [-> routing kernel -> routing event];
  • shared stream: wait for the router event -> shared experts with their all-reduce -> shared event;
  • main stream: latent down projection -> wait for the router event, or with pre-routing, the MoE op waits for the routing event after quantizing its input -> routed MoE -> wait for the shared event -> routed all-reduce (RMSNorm fused) and latent up projection.
• It runs the same kernels on the same inputs as the default path, and the shared all-reduce still completes before the routed one starts.
• The routed all-reduce tail is factored into _reduce_routed_latent and used by both paths.

Performance, Rubin, TP8 across 2 nodes (MNNVL), Kimi K3 with speculative decoding, MoE TP8, CUDA graphs, single request, decode step time. Measured with the same code on our integration branch:

	2K	32K	64K	96K	128K	393K
router stream vs. default	-0.10 ms	-0.08 ms			-0.47 ms
+ pre-routing vs. router stream only	-0.399 ms	-0.400 ms	-0.392 ms	-0.419 ms	-0.401 ms	-0.404 ms

Outputs are byte-identical for the router stream alone. With pre-routing, MoE outputs are bit-identical to the op routing itself, and our accuracy gates were unchanged.

Moving the routing alone, with the shared experts still behind it on one stream, was neutral to slightly negative in a whole-layer microbenchmark: the shared GEMMs then overlap FC1 / FC2. That is why the shared experts start on their own stream. Above 8 tokens per rank the same contention made pre-routing slower, hence the default token limit.

Relation to other work:

• [None][perf] Kimi K3: overlap KDA/MLA attention decode onto the aux stream #17103 forks Kimi K3 attention work onto the Attention / MoeShared aux stream and joins before the MoE. This path adds a separate router stream and keeps the shared experts on the shared stream, so the two compose.
• An earlier experiment that also folded the MoE finalize into the routed all-reduce is not part of this PR; it was slower in our measurements.

Open design question: the hand-off from the model to the op goes through module-level dicts keyed by the expert weight pointer, so the MoE / ConfigurableMoE call signatures stay unchanged. Passing the prepared routing explicitly through ConfigurableMoE.forward would be cleaner. Happy to change it if reviewers prefer that.

Test Coverage

• New tests/unittest/_torch/thop/serial/test_moe_mxe2m1_pre_routing.py (SM100 family; collected by the unittest/_torch/thop/serial entries of l0_b200.yml / l0_b300.yml):
  • pre-routed MoE output is bitwise equal to the op routing itself, for Kimi K3 routing (896 experts, top 16, DeepSeekV3 method) and Renormalize, 1 / 8 / 16 tokens, with and without a delayed router stream;
  • shapes are pre-routed only after being recorded;
  • changed routing inputs fall back to the op's own routing;
  • different logits raise;
  • workspace validation.
• tests/unittest/_torch/modeling/test_kimi_linear_modeling.py (CPU): the router stream is a separate aux stream when enabled. The existing registry test covers the default.
• GPU results at this commit (Rubin single GPU (200-SM engineering sample, SM107); C++ built from main plus the C++ changes of this PR, [#20540][perf] Kimi K3 MoE decode: router GEMM and routing kernel on a router stream #20541 and [#20545][perf] Kimi K3: fold the MoE-tail additions into the next attention-residual #20546 (they touch disjoint files)):
  • test_moe_mxe2m1_pre_routing.py: 16 passed, 0 skipped (on main: the ops do not exist);
  • test_kimi_linear_modeling.py: 3 passed (on main the new test fails);
  • test_custom_ops.py -k test_register_fake: passed (both new ops have fake implementations).
• Earlier validation of the same C++ and Python on Rubin:
  • op-level and layer-level pre-routed vs. own-routing bitwise checks (790 + 54 cases, eager and CUDA graph, with delayed router / shared streams);
  • compute-sanitizer memcheck / racecheck clean;
  • end-to-end A/B above.
• Not measured: SM100 / SM103, NVFP4 checkpoints, attention DP.
• pre-commit on the changed files: clean.

PR Checklist

• PR description clearly explains what and why.
• PR follows the TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for new code paths.
• No API changes, no new dependencies, no CODEOWNERS changes.
• Please check this after reviewing the above items as appropriate for this PR.