fix(v4-flash): make ExpertDistributionRecorder work for DeepSeek-V4-Flash#45
fix(v4-flash): make ExpertDistributionRecorder work for DeepSeek-V4-Flash#45yiqiliu2 wants to merge 2 commits into
Conversation
Original `_SelectExpertsSinglePassGatherer.on_select_experts` crashed on V4-Flash routing with two distinct bugs: 1. `scatter_add_` on `self._data[layer_idx, :]` raised "Index tensor must have the same number of dimensions as self tensor" depending on the stride/dtype of the incoming topk_ids tensor. 2. `is_current_stream_capturing()` in `_on_hook` lets the recorder fire during CUDA graph capture, but any host sync (`.item()`, boolean indexing, nonzero) inside the hook hits cudaErrorStreamCaptureInvalidated. Rewrite 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 with `--expert-distribution-recorder-mode stat` (cuda-graph capture succeeds) then dumping a 44 MB .pt with shape (1000, 43, 256) sum=28080 from a 3-prompt warmup. Unblocks `--init-expert-location <warmup>.pt --kt-expert-placement-strategy frequency` for V4-Flash.
…t_layer ctx The forward loop in DeepSeek-V4Model.forward iterated layers without wrapping each call in get_global_expert_distribution_recorder(). with_current_layer(i). Without this, the recorder hook's `_current_layer_idx.value` stays at None and every on_select_experts write lands in slot 0 of the gatherer's _data tensor → the dumped logical_count.pt only has data for layer 0 → torch.topk on it places all GPU experts on layer 0 (the placement-strategy 'frequency cold-start' bug we previously diagnosed). Other DeepSeek family models (deepseek_v2, qwen2_moe, qwen3_moe, exaone_moe, llada2, longcat_flash, step3p5) all wrap their forward loops with this context — V4 was an oversight on initial port. Pairs with c32f840 (cuda-graph + shape safe recorder hook); together they make `--expert-distribution-recorder-mode stat` work end-to-end on V4-Flash.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the expert distribution recording in expert_distribution.py to be compatible with V4-Flash and CUDA graph capture by using pure GPU operations and avoiding host-device synchronization. It also updates deepseek_v4.py to wrap layer calls with a recording context. Feedback focuses on adding a safety check for layer_idx is None in on_select_experts to prevent silent data corruption and questioning the exclusion of the recording context when piecewise CUDA graphs are enabled, as this leads to incorrect data being recorded for layer 0 only.
| def on_select_experts(self, layer_idx: int, topk_ids: torch.Tensor): | ||
| 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() | ||
| ) | ||
| if topk_ids is None: | ||
| return | ||
| flat_ids = topk_ids.reshape(-1).to(torch.int64) | ||
| target = self._data[layer_idx].reshape(-1) | ||
| valid = (flat_ids >= 0).to(target.dtype) | ||
| safe_ids = flat_ids.clamp(min=0) | ||
| target.scatter_add_(0, safe_ids, valid) |
There was a problem hiding this comment.
The on_select_experts method should explicitly check if layer_idx is None. If the context manager with_current_layer(i) is missing or skipped (as currently implemented for the piecewise CUDA graph mode in deepseek_v4.py), layer_idx will be None. In PyTorch, indexing a 2D tensor with None (e.g., self._data[None]) adds a leading dimension, and the subsequent reshape(-1) flattens the entire tensor. This causes scatter_add_ to incorrectly write all expert counts into the first num_experts slots of the flattened tensor, which correspond to layer 0. Adding a check for layer_idx is None prevents this silent data corruption.
| def on_select_experts(self, layer_idx: int, topk_ids: torch.Tensor): | |
| 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() | |
| ) | |
| if topk_ids is None: | |
| return | |
| flat_ids = topk_ids.reshape(-1).to(torch.int64) | |
| target = self._data[layer_idx].reshape(-1) | |
| valid = (flat_ids >= 0).to(target.dtype) | |
| safe_ids = flat_ids.clamp(min=0) | |
| target.scatter_add_(0, safe_ids, valid) | |
| def on_select_experts(self, layer_idx: int, topk_ids: torch.Tensor): | |
| if topk_ids is None or layer_idx is None: | |
| return | |
| flat_ids = topk_ids.reshape(-1).to(torch.int64) | |
| target = self._data[layer_idx].reshape(-1) | |
| valid = (flat_ids >= 0).to(target.dtype) | |
| safe_ids = flat_ids.clamp(min=0) | |
| target.scatter_add_(0, safe_ids, valid) |
| ctx = ( | ||
| nullcontext() | ||
| if get_global_server_args().enable_piecewise_cuda_graph | ||
| else get_global_expert_distribution_recorder().with_current_layer(i) | ||
| ) |
There was a problem hiding this comment.
The condition if get_global_server_args().enable_piecewise_cuda_graph disables the layer-specific recording context when piecewise CUDA graph is enabled. As noted in the PR description and the added comments, skipping this context causes all expert distribution data to be recorded into layer 0 (the "slot 0" issue). Unless there is a specific technical limitation preventing the use of this context manager during piecewise capture, it should be applied consistently to ensure accurate statistics across all execution modes. If it is indeed incompatible, it would be better to disable the recorder entirely for this mode rather than allowing it to produce incorrect data.
There was a problem hiding this comment.
Pull request overview
This PR fixes expert distribution recording for DeepSeek-V4-Flash when running with --expert-distribution-recorder-mode stat, addressing both incorrect per-layer attribution and CUDA-graph capture safety in the recorder’s gather path.
Changes:
- Wrap DeepSeek-V4’s per-layer forward calls with
ExpertDistributionRecorder.with_current_layer(i)(when not in piecewise CUDA-graph mode) so per-layer stats are attributed correctly. - Rewrite
_SelectExpertsSinglePassGatherer.on_select_expertsto use GPU-only ops (no host sync / boolean indexing) and be shape-safe for V4-Flash routing tensors.
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 per-layer recorder context (with a piecewise-cuda-graph guard) so expert stats aren’t all attributed to layer 0. |
python/sglang/srt/eplb/expert_distribution.py |
Updates the GPU gatherer’s on_select_experts implementation to be CUDA-graph-capture-safe and avoid shape issues. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from sglang.srt.environ import envs, is_large_dummy_model | ||
| from contextlib import nullcontext | ||
| from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder | ||
| from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation |
| # yiqiliu2 / 2026-05-08: wrap each layer call with the recorder's | ||
| # current_layer context. Without this, the recorder hook's | ||
| # `self._current_layer_idx.value` stays at None and every | ||
| # `on_select_experts` write lands in slot 0 of the gatherer's | ||
| # `_data` tensor, so the dumped logical_count.pt only has | ||
| # data for layer 0 and torch.topk on it places all GPU | ||
| # experts on layer 0 regardless of K. | ||
| ctx = ( | ||
| nullcontext() | ||
| if get_global_server_args().enable_piecewise_cuda_graph | ||
| else get_global_expert_distribution_recorder().with_current_layer(i) | ||
| ) |
Summary
Two coupled fixes that together make `--expert-distribution-recorder-mode stat` work end-to-end for V4-Flash.
1. `_SelectExpertsSinglePassGatherer.on_select_experts` cuda-graph-safe + shape-safe
The original used `scatter_add_` on `self._data[layer_idx, :]` and crashed with `Index tensor must have the same number of dimensions as self tensor` on V4-Flash routing tensors.
`_on_hook`'s `is_current_stream_capturing()` clause also lets the hook fire during cuda graph capture, but any host sync (`.item()`, boolean indexing, nonzero) inside the hook hits cudaErrorStreamCaptureInvalidated.
Rewrite uses pure GPU ops:
No host sync, no boolean indexing, no nonzero.
2. `deepseek_v4.py` forward loop wraps with `with_current_layer(i)` ctx
The forward loop iterated layers without the recorder context manager. Without this, every `on_select_experts` write lands in slot 0 of the gatherer's `_data` tensor → the dumped `logical_count.pt` only has data for layer 0 → `torch.topk` on it places all GPU experts on layer 0.
Other DeepSeek family models (`deepseek_v2`, `qwen2_moe`, `qwen3_moe`, `exaone_moe`, `llada2`, `longcat_flash`, `step3p5`) all wrap their forward loops with this context — V4 was an oversight on initial port.
Test plan
🤖 Generated with Claude Code