[PR #46] fix(v4-flash): make ExpertDistributionRecorder cuda-graph-safe + per-layer for DeepSeek-V4-Flash #47

New issue

Open

opened 2026-10-29 21:06:45 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:06:45 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/46
Author: @yiqiliu2
Created: 5/8/2026
Status: 🔄 Open

Base: main ← Head: upstream-pr/fix-v4-flash-recorder

📝 Commits (2)

103bc87 fix(v4-flash): cuda-graph-safe + shape-safe ExpertDistribution recorder
d6e0e89 fix(v4-flash): wrap deepseek_v4 forward loop with recorder current_layer ctx

📊 Changes

2 files changed (+40 additions, -12 deletions)

View changed files

📝 python/sglang/srt/eplb/expert_distribution.py (+19 -6)
📝 python/sglang/srt/models/deepseek_v4.py (+21 -6)

📄 Description

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

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/46 **Author:** [@yiqiliu2](https://github.com/yiqiliu2) **Created:** 5/8/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `upstream-pr/fix-v4-flash-recorder` --- ### 📝 Commits (2) - [`103bc87`](https://github.com/kvcache-ai/sglang/commit/103bc879d5f2ca4a35eb9a0c6fb845737323e882) fix(v4-flash): cuda-graph-safe + shape-safe ExpertDistribution recorder - [`d6e0e89`](https://github.com/kvcache-ai/sglang/commit/d6e0e892ea1d304457609016de18f2e8701ba400) fix(v4-flash): wrap deepseek_v4 forward loop with recorder current_layer ctx ### 📊 Changes **2 files changed** (+40 additions, -12 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/eplb/expert_distribution.py` (+19 -6) 📝 `python/sglang/srt/models/deepseek_v4.py` (+21 -6) </details> ### 📄 Description ## 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: ```python 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 - [x] 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. - [x] 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). - [x] 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](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:06:45 +00:00

bsakaguchi referenced this issue

2026-10-29 21:06:45 +00:00

[PR #47] [MERGED] Refactor: isolate DeepSeek V4 Flash behind plugin registries #48

bsakaguchi referenced this issue

2026-10-29 21:06:47 +00:00

[PR #52] [CLOSED] fix(parser): align reasoning & tool-call parsers with upstream sglang #52