expert distributions

[CUHKSZzxy]

Recorder adapted from dlBLAS, endpoint refers to SGLang.

Usage

1. Start the server with the env variables

LMDEPLOY_DUMP_EXPERT_DISTRIBUTION=1 \
LMDEPLOY_EXPERT_DUMP_DIR="./path_to_save_dump" \
LMDEPLOY_EXPERT_DUMP_VISUALIZE=1 \
lmdeploy serve api_server Qwen/Qwen3-30B-A3B --backend pytorch --tp 2 --log-level INFO

• LMDEPLOY_DUMP_EXPERT_DISTRIBUTION to enable the expert distribution feature.
• LMDEPLOY_EXPERT_DUMP_DIR specifies the path to save dump files.
• LMDEPLOY_EXPERT_DUMP_VISUALIZE turns on visualization, will produce heatmap figs under dump dir.

2. Start / dump / stop

curl -X POST http://localhost:23333/start_expert_distribution_record
curl -X POST http://localhost:23333/dump_expert_distribution_record
curl -X POST http://localhost:23333/stop_expert_distribution_record

Result

Data: sharegpt, prompts 1000, random in 1024, random out 1024.

• Model: Qwen3-30B-A3B, TP=2.

• Model: Qwen3.5-35B-A3B, TP=2.

[Copilot]

Pull request overview

Adds an MoE expert-dispatch distribution recorder to the PyTorch backend, enabling periodic dumps of routed expert counts to JSON for debugging/analysis (intended for eager-mode runs).

Changes:

• Introduce ExpertsDistributionRecorder utility (real vs no-op based on env flags) that aggregates and dumps expert token counts.
• Hook the recorder into Qwen3 MoE and DeepSeek V2 MoE forward passes to record topk_ids.
• Add new environment variables for enabling/disabling and configuring dump output.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File	Description
lmdeploy/pytorch/models/utils/expert_distribution_recorder.py	New recorder implementation that tracks expert token counts and periodically dumps JSON.
lmdeploy/pytorch/models/qwen3_moe.py	Records MoE router topk_ids per forward for Qwen3 MoE blocks.
lmdeploy/pytorch/models/deepseek_v2.py	Records MoE router topk_ids per forward for DeepSeek V2 MoE blocks; stores layer_idx.
lmdeploy/pytorch/envs.py	Adds env configuration knobs for expert distribution dumping.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.