kvcache-ai/sglang

Fork 0

mirror of https://github.com/kvcache-ai/sglang.git synced 2026-11-05 01:07:39 +00:00

[PR #45] fix(v4-flash): make ExpertDistributionRecorder work for DeepSeek-V4-Flash #45

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: pr/expert-distribution-recorder-v4flash-fix

📝 Commits (2)

5c3410e fix(v4-flash): cuda-graph-safe + shape-safe ExpertDistribution recorder
1e6496b fix: wrap deepseek_v4 forward layers 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 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:

`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.

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

Boot sglang server with `--expert-distribution-recorder-mode stat` for V4-Flash — capture succeeds, no crash
Hit `/start_expert_distribution_record` + a few warmup completions + `/dump_expert_distribution_record` — produces a 44 MB .pt with non-zero counts on 40/43 layers (was 1/43 before fix #2)
Frequency-based GPU expert placement using the dumped .pt distributes 172 experts as 2..9 per layer (was all 172 on layer 0)

🤖 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/45 **Author:** [@yiqiliu2](https://github.com/yiqiliu2) **Created:** 5/8/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `pr/expert-distribution-recorder-v4flash-fix` --- ### 📝 Commits (2) - [`5c3410e`](https://github.com/kvcache-ai/sglang/commit/5c3410e0fce45c719e1be7e48aa36c29151b287d) fix(v4-flash): cuda-graph-safe + shape-safe ExpertDistribution recorder - [`1e6496b`](https://github.com/kvcache-ai/sglang/commit/1e6496be454b84d0f670b73267653d316f57f096) [fix](v4-flash): wrap deepseek_v4 forward layers 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 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: - \`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. ### 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 - [x] Boot sglang server with \`--expert-distribution-recorder-mode stat\` for V4-Flash — capture succeeds, no crash - [x] Hit \`/start_expert_distribution_record\` + a few warmup completions + \`/dump_expert_distribution_record\` — produces a 44 MB .pt with non-zero counts on 40/43 layers (was 1/43 before fix #2) - [x] Frequency-based GPU expert placement using the dumped .pt distributes 172 experts as 2..9 per layer (was all 172 on layer 0) 🤖 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:43 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/sglang#45

No description provided.

Rows
Columns