perf: further optimize GLM-5.2 serving

[CUHKSZzxy]

Depends on #4827.

Summary

• Align MoE router GEMM dtype and dispatch: GLM-5.2 keeps FP32 router logits through BF16 cuBLAS, while DeepSeek variants retain model-dtype routing.
• Avoid redundant sparse MLA index recompilation.
• Add opt-in FlashInfer and symmetric-memory all-reduce backends for Ray, with capability checks and NCCL fallback.
• Enable programmatic dependent launch between FP8 quantization and DeepGEMM linear kernels on supported Hopper systems.

Performance

GLM-5.2 FP8 TP8 on 8×H200 GPUs with FP8 KV cache, MTP5, concurrency 16, and five-turn SWE-Smith:

Metric	NCCL	Symmetric memory	Change
Completion throughput	846.79 tok/s	906.60 tok/s	+7.06%
Steady completion throughput	1,291.26 tok/s	1,406.08 tok/s	+8.89%
TPOT	14.88 ms	13.69 ms	-8.00%
Success	80/80	80/80	unchanged

Isolated CUDA-graph measurements of the GLM attention projection chain show PDL reducing GPU time by 9.0% at one token and 12.7% at 64 tokens. A full-model serving timeline confirms that the PDL-enabled quantization-to-DeepGEMM chain is active.

Validation

• Pre-commit checks passed.
• Targeted router GEMM, FP8 GEMM, MLA, and communicator tests passed.
• Both full-model communication variants completed 80/80 SWE-Smith requests.

Assistance

Assisted with Codex + GPT-5.6-Sol xHigh, reviewed manually

[Copilot]

Pull request overview

This PR extends the GLM-5.2 CUDA serving optimization stack by adding a small-token FP32 router GEMM fast-path, reducing redundant NSA/MLA index specialization work, and introducing opt-in CUDA all-reduce backends (FlashInfer + PyTorch symmetric-memory) with Ray-specific capability gating and fallback behavior. It also enables Triton/DeepGEMM programmatic dependent launch (PDL) coordination on supported Hopper systems.

Changes:

• Add a Triton FP32 router GEMM kernel + model-side dispatch to accelerate MoE routing for measured shapes.
• Add opt-in CUDA communicator support (FlashInfer fusion and symmetric-memory all-reduce), plumbed through DistContext and Ray executor initialization.
• Reduce redundant MLA/NSA decode index recompilations by caching specializations across query modes; enable PDL-triggered dependent launches for FP8 quantization + DeepGEMM.

Reviewed changes

Copilot reviewed 25 out of 25 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/pytorch/test_cuda_communicator.py	Adds unit tests for communicator dispatch/close + RMSNorm fusion behavior.
tests/pytorch/kernel/test_mla_attention.py	Updates NSA index updater tests to use the new cached update_decode API and disables dynamic compile where needed.
tests/pytorch/kernel/test_fp32_router_gemm.py	Adds CUDA tests validating FP32 router GEMM correctness, fallback, and CUDA graph behavior.
requirements/runtime_cuda.txt	Adds flashinfer-python runtime dependency for opt-in FlashInfer communicator support.
lmdeploy/pytorch/third_party/deep_gemm/init.py	Enables DeepGEMM PDL when supported; raises on missing DeepGEMM import.
lmdeploy/pytorch/spec_decode/base.py	Propagates communicator_builder when building draft dist contexts.
lmdeploy/pytorch/nn/norm.py	Allows RMSNorm to consume pending all-reduce (optionally fused with residual+norm).
lmdeploy/pytorch/models/deepseek_v32.py	Defers attention/MLP all-reduce into downstream RMSNorm when supported.
lmdeploy/pytorch/models/deepseek_v2.py	Routes MoE gate logits through new FP32 router GEMM dispatch; adds all_reduce control for MoE/MLP.
lmdeploy/pytorch/kernels/cuda/utils.py	Adds supports_pdl() capability helper.
lmdeploy/pytorch/kernels/cuda/fp32_router_gemm.py	Introduces Triton small-token FP32 router GEMM kernel and support checks.
lmdeploy/pytorch/kernels/cuda/blocked_gemm_fp8.py	Adds optional PDL-dependent launch path to FP8 quantization kernels.
lmdeploy/pytorch/envs.py	Adds env toggles for FlashInfer and symmetric-memory all-reduce selection.
lmdeploy/pytorch/engine/executor/ray_executor.py	Adds Ray device binding logic and symmetric-memory visibility gating.
lmdeploy/pytorch/engine/executor/base_worker.py	Passes backend communicator builder into DistContext.
lmdeploy/pytorch/distributed.py	Adds communicator plumbing to DistGroup/DistContext and TP communicator construction.
lmdeploy/pytorch/backends/dlinfer/op_backend.py	Asserts CUDA communicator env toggles are disabled for DLInfer.
lmdeploy/pytorch/backends/cuda/symm_mem_allreduce.py	Adds symmetric-memory all-reduce backend implementation.
lmdeploy/pytorch/backends/cuda/op_backend.py	Uses CUDA communicator when eligible, otherwise falls back to base communicator.
lmdeploy/pytorch/backends/cuda/flashinfer_allreduce.py	Adds FlashInfer all-reduce + fused residual+RMSNorm backend implementation.
lmdeploy/pytorch/backends/cuda/communicator.py	Adds CUDA communicator dispatcher and config gating helpers.
lmdeploy/pytorch/backends/communicator.py	Introduces DeviceCommunicator abstraction and default implementation.
lmdeploy/pytorch/backends/base.py	Adds backend hook for building communicators.
lmdeploy/pytorch/backends/cuda/blockedf8_modules.py	Enables DeepGEMM PDL usage to coordinate FP8 quant + GEMM on Hopper.
lmdeploy/pytorch/backends/cuda/attention/mla.py	Refactors NSA decode index updater to cache specializations by query/layout.

---

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

[CUHKSZzxy]

Using cpu_group here is intentional. LMDeploy builds the CPU and GPU TP groups from the same ordered rank set. FlashInfer uses the process group for workspace rendezvous/control-plane setup, while the collective itself runs through its CUDA kernels; vLLM likewise passes get_tp_group().cpu_group when creating the FlashInfer workspace. The NCCL fallback still uses device_group. I also updated the unit test to assert that both optimized backends receive cpu_group, so this contract is covered.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 26 changed files in this pull request and generated 1 comment.

[grimoire]

LGTM