perf(pytorch): optimize compact blocked FP8 MoE and route preparation

[grimoire]

Motivation

MoE execution has two avoidable overheads:

• route sorting and compact-block preparation require multiple CUDA graph nodes;
• the original blocked FP8 kernel launches a padded grid for every expert, wasting CTAs when routed tokens are sparse or unevenly distributed.

Changes

• Add a shared Triton route-preparation implementation:
  • use a single CTA on Hopper when both routes and experts are at most 2,048;
  • use a fused scalable parallel path for larger workloads;
  • preserve the existing _get_sorted_idx and _get_sorted_idx_blocks APIs.
• Extend compact blocked FP8 MoE dispatch using token count, route density, expert count, and projection dimensions.
• Use transposed MMA for measured small-M, long-K gate/up tiles:
  • block_m <= 32;
  • input features >= 2,048.
• Keep the original MMA orientation for down projections and block_m >= 64.
• Group CUDA MoE kernels under lmdeploy/pytorch/kernels/cuda/moe and update imports.
• Preserve parallel and original-kernel fallbacks for unsupported shapes and local-expert preparation.

No public operator API is changed.

Performance

Measured on one NVIDIA H200 inside CUDA graphs.

Route preparation with random top-k-8 routing:

Routes / experts	Contract	Parallel	Single CTA	Speedup
512 / 256	sorted indices	6.44 us	2.72 us	2.36x
512 / 256	compact blocks	7.04 us	3.23 us	2.18x
2,048 / 512	sorted indices	7.65 us	5.81 us	1.32x
2,048 / 512	compact blocks	8.17 us	6.64 us	1.23x

Representative exact-checkpoint blocked FP8 results:

Workload	Before	After	Improvement
Qwen3.5-35B TP2, 128 tokens	143.2 us	135.4 us	5.4%
Qwen3.5-35B TP2, 512 tokens	173.8 us	151.7 us	12.7%
Qwen3.5-122B TP4, 512 tokens	238.7 us	208.3 us	12.7%
GLM-5.2 TP8, 64-token fused MoE	290.8 us	280.5 us	3.5%
Qwen3.5-397B TP4, 512-token fused MoE	464.1 us	459.6 us	1.0%

Checkpoint weights and FP8 scales are real; activations and random/hot routing inputs are synthetic.

[Copilot]

Pull request overview

This PR optimizes PyTorch CUDA MoE execution by reducing route-preparation overhead (including a Hopper single-CTA path) and improving compact blocked FP8 MoE dispatch decisions, while reorganizing MoE CUDA kernels under lmdeploy/pytorch/kernels/cuda/moe and updating downstream imports/tests accordingly.

Changes:

• Adds a shared Triton route-preparation implementation (single-CTA on Hopper for small workloads, parallel scalable path otherwise) while preserving _get_sorted_idx* APIs.
• Extends compact blocked FP8 MoE dispatch heuristics and adds an optional transposed MMA path for measured small-M / long-K gate-up tiles.
• Moves/organizes MoE kernels under lmdeploy/pytorch/kernels/cuda/moe, updates backend wiring, and expands tests to cover new routing and config-selection behavior.

Reviewed changes

Copilot reviewed 21 out of 23 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/pytorch/kernel/test_v4_fp4_grouped_gemm.py	Updates imports to new cuda.moe.* module layout for v4 FP4 grouped GEMM/SwiGLU quant tests.
tests/pytorch/kernel/test_moe_route.py	Updates noaux-TC routing import to the new MoE routing module.
tests/pytorch/kernel/test_fused_moe.py	Updates imports and adds extensive coverage for new sorted-index/compact-block route preparation policies and implementations.
tests/pytorch/kernel/test_fused_moe_static_fp8.py	Updates imports for static FP8 MoE path to new cuda.moe.* modules.
tests/pytorch/kernel/test_fuse_moe_v4_fp4.py	Updates v4 FP4 fused MoE imports to new module locations.
tests/pytorch/kernel/test_fuse_moe_blocked_fp8.py	Adds coverage for compact blocked FP8 gate config (incl. transposed MMA) and updates imports to new cuda.moe.* modules.
lmdeploy/pytorch/kernels/cuda/moe/w8a8.py	Fixes relative imports after move; keeps w8a8 MoE implementation under cuda/moe.
lmdeploy/pytorch/kernels/cuda/moe/v4_swiglu_quant.py	Updates imports to reference parent CUDA utilities after move under cuda/moe.
lmdeploy/pytorch/kernels/cuda/moe/v4_fp4.py	Updates imports to reference parent CUDA utilities after move under cuda/moe.
lmdeploy/pytorch/kernels/cuda/moe/v4_fp4_grouped_gemm.py	Updates imports to reference relocated v4 FP4 fused launcher and parent utilities.
lmdeploy/pytorch/kernels/cuda/moe/route_single_group.py	Introduces a Triton fused single-group top-k router specialized for the supported configuration.
lmdeploy/pytorch/kernels/cuda/moe/route_noaux_tc.py	Introduces Triton noaux-TC routing kernel with autotuning keyed by routing shape parameters.
lmdeploy/pytorch/kernels/cuda/moe/fused_moe.py	Implements optimized route preparation (single-CTA + parallel paths), reorganizes kernels, and updates compact route metadata building.
lmdeploy/pytorch/kernels/cuda/moe/ep.py	Updates import paths after module move under cuda/moe.
lmdeploy/pytorch/kernels/cuda/moe/ep_fp8.py	Updates import paths and EP gather import after module move under cuda/moe.
lmdeploy/pytorch/kernels/cuda/moe/blocked_fp8.py	Extends compact blocked FP8 selection logic and adds transposed MMA option for gate/up compact kernel.
lmdeploy/pytorch/kernels/cuda/moe/init.py	Adds package marker for the new cuda.moe module tree.
lmdeploy/pytorch/kernels/cuda/init.py	Re-exports fused_moe/fused_moe_w8a8 from new cuda.moe.* locations to preserve import surface.
lmdeploy/pytorch/backends/cuda/moe/v4_fp4.py	Updates backend to import v4 FP4 fused MoE from cuda.moe.*.
lmdeploy/pytorch/backends/cuda/moe/static_fp8.py	Updates backend to import static FP8 MoE from cuda.moe.w8a8.
lmdeploy/pytorch/backends/cuda/moe/default.py	Updates EP and renormalize imports to cuda.moe.*.
lmdeploy/pytorch/backends/cuda/moe/blocked_fp8.py	Updates backend imports to new compact blocked FP8 / EP-FP8 / renormalize module locations.
lmdeploy/pytorch/backends/cuda/moe_router.py	Switches router implementation to new cuda.moe.route_* routing kernels and retains gating logic for the specialized fused path.

Suppressed comments (1)

lmdeploy/pytorch/kernels/cuda/moe/fused_moe.py:666

• _get_sorted_idx_blocks_parallel launches Triton kernels that assume expert_offset/local_num_experts describe a valid slice of [0, num_experts), and that block_m > 0. If a caller passes invalid values, the kernels can read/write out of bounds (e.g., in _route_prefix_kernel when loading Counts + expert_offset + local_expert_offsets). Add explicit argument validation before launching kernels to fail fast with a clear error.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.