fix(v4-flash): make ExpertDistributionRecorder cuda-graph-safe + per-layer for DeepSeek-V4-Flash

[yiqiliu2]

Summary

Two coupled bug fixes in --expert-distribution-recorder-mode stat on DeepSeek-V4-Flash. Without these, V4-Flash users hit:

• RuntimeError: Index tensor must have the same number of dimensions as self tensor and
• CUDA error: operation not permitted when stream is capturing during cuda graph capture, and
• Even after the above are fixed, every layer's routing data lands in slot 0 of the recorder's data tensor, so dumped logical_count.pt only has data for layer 0 → downstream frequency-based GPU expert placement puts all 172 experts on layer 0 → useless.

Fix 1: _SelectExpertsSinglePassGatherer.on_select_experts (cuda-graph + shape safe)

_on_hook lets the recorder fire during cuda graph capture (the is_current_stream_capturing() clause). The original implementation:

def on_select_experts(self, layer_idx, topk_ids):
    topk_ids = topk_ids.flatten()
    mask = topk_ids != -1
    self._data[layer_idx, :].scatter_add_(
        dim=0, index=topk_ids.masked_fill(~mask, 0).long(), src=mask.int()
    )

This crashes on V4-Flash routing tensors with the dimension-mismatch error. An intermediate rewrite using bool(valid_mask.any()) + flat_ids[valid_mask] synced device→host (.item() / nonzero), forbidden during cuda graph capture.

The new version uses pure GPU ops:

• clamp(-1 → 0) so all indices are in range
• src = (id >= 0).to(dtype) so invalid entries contribute 0 (no-op writes)

No host sync, no boolean indexing, no nonzero. Verified by booting V4-Flash with --expert-distribution-recorder-mode stat; cuda graph capture succeeds and a 44 MB .pt dumps cleanly.

Fix 2: deepseek_v4.py forward loop wraps each layer with with_current_layer(i)

The other DeepSeek/Qwen models in this repo all wrap their for i in range(...) layer loop in
get_global_expert_distribution_recorder().with_current_layer(i) (gated by enable_piecewise_cuda_graph). deepseek_v4.py was missing this. Without it, self._current_layer_idx.value stays at None throughout the forward pass and every on_select_experts write lands in slot 0 of the gatherer's _data tensor.

Symptom: dumped logical_count.pt has counts only on layer 0; remaining 42 layers report 0. torch.topk on this in kt_ep_wrapper.generate_gpu_experts_masks then places all GPU expert quota on layer 0 regardless of K.

Test plan

• Boot V4-Flash with --expert-distribution-recorder-mode stat. Before: cuda graph capture crashes with the recorder error above. After: capture succeeds in 7.79s.
• Send a 3-prompt 32-token warmup, dump record. Before: file empty / single-layer. After: 44 MB file with non-zero counts on 40/43 layers (the 3 dense layers correctly have 0).
• Use the dumped .pt as --init-expert-location with --kt-num-gpu-experts 4 --kt-expert-placement-strategy frequency. Before: all 172 experts placed on layer 0. After: distributed 2-9 per MoE layer (sum 172).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request updates the expert distribution recording logic to ensure compatibility with V4-Flash and CUDA graph capture by replacing host-device synchronization with pure GPU operations. It also integrates the recorder into the DeepseekV4 model's forward pass using a layer-specific context manager. Feedback includes a recommendation to explicitly check for a null layer index to prevent data corruption and a suggestion to optimize the layer loop by hoisting global lookups and re-evaluating the piecewise CUDA graph gate.

[gemini-code-assist]

It is recommended to explicitly check if layer_idx is None here. If the context manager with_current_layer(i) is not active (e.g., due to the nullcontext gate in deepseek_v4.py), layer_idx will be None. In PyTorch, self._data[None] creates a new dimension, and the subsequent reshape(-1) flattens the entire tensor. This causes scatter_add_ to record data into the first num_experts slots of the flattened tensor, which effectively attributes all routing data to layer 0. Adding or layer_idx is None to the check ensures that no data is recorded when the layer context is missing, preventing the "slot 0" bug.

Suggested change

	if topk_ids is None:
	return
	if topk_ids is None or layer_idx is None:
	return

[gemini-code-assist]

The nullcontext() gate for enable_piecewise_cuda_graph prevents the expert distribution recorder from tracking the current layer index when piecewise CUDA graphs are enabled. This will cause the recorder to either skip these layers (if a None check is added to on_select_experts) or incorrectly attribute all data to layer 0. If expert distribution recording is intended to be supported with piecewise CUDA graphs, this gate should be removed. Additionally, consider hoisting the calls to get_global_server_args() and get_global_expert_distribution_recorder() outside the loop to avoid redundant lookups in every iteration.

[Copilot]

Pull request overview

This PR targets DeepSeek-V4-Flash expert-distribution recording (--expert-distribution-recorder-mode stat) to avoid CUDA-graph capture issues and to attribute routing statistics to the correct layer.

Changes:

• Updates _SelectExpertsSinglePassGatherer.on_select_experts to use GPU-only ops for scatter-add accumulation.
• Wraps DeepSeek-V4’s per-layer forward calls with the expert distribution recorder’s with_current_layer(i) context (conditionally).

Reviewed changes

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

File	Description
python/sglang/srt/models/deepseek_v4.py	Adds a per-layer recorder context around each decoder layer call to attribute routing stats to the intended layer.
python/sglang/srt/eplb/expert_distribution.py	Rewrites the stat gatherer’s on_select_experts accumulation to avoid host-sync patterns and handle V4-Flash routing tensors.

---

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