[PR #381] Add B200-Calibrated Sparse-Routing Adaptive Wave Sizing #404

New issue

Open

opened 2026-10-26 21:17:20 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:20 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/381
Author: @qinqinwo
Created: 7/15/2026
Status: 🔄 Open

Base: main ← Head: feature/imbalance-aware-block-m

📝 Commits (5)

3f25a3a MegaMoE: imbalance-aware block_m selection (opt-in, strict fallback)
458dd21 MegaMoE: add calibrated adaptive wave sizing
af06265 MegaMoE: narrow adaptive waves to stable sparse routing
c343597 Docs: add B200 performance results in English
5daaaab Fix adaptive receive-stat cache isolation

📊 Changes

7 files changed (+714 additions, -49 deletions)

View changed files

➕ ADAPTIVE_WAVE.md (+93 -0)
📝 csrc/apis/mega.hpp (+115 -1)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+127 -42)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+6 -2)
➕ scripts/bench_adaptive_wave_ab.sh (+113 -0)
➕ scripts/bench_mega_moe_wave_size.sh (+68 -0)
📝 tests/test_mega_moe.py (+192 -4)

📄 Description

MegaMoE processes a fixed number of local experts in each L1→L2 scheduler wave.
The upstream heuristic derives that number from the mean tokens per expert. That
is a good general fallback, but it cannot see the realized routing distribution.

Set DG_MEGA_MOE_ADAPTIVE_WAVE=1 to enable the opt-in B200 FP8×FP4 policy. The policy:

reads a delta of cumulative_local_expert_recv_stats from the preceding window;
caches the sampled distribution and refreshes it every 256 launches, avoiding a
synchronous device-to-host copy on the steady-state path;
only changes the calibrated shape (EP 8, 256 experts, top-k 8, hidden 7168,
intermediate 2048) and 127.5 < expected tokens/expert <= 128.5 band;
skips receive-stat sampling entirely outside that calibrated shape and band;
uses 8 experts/wave when the active-expert ratio is at or below 0.92;
balanced and moderate-skew routing retain the upstream size;
falls back to the upstream wave size on the first call, after counter resets,
outside the calibrated tier, or when the requested tier exceeds ring capacity.

The deliberately narrow gate is based on same-process, order-balanced 8×B200
measurements. A broader lower bound was rejected after the 96 tokens/expert
point regressed 1.48% under high skew. A balanced-routing wave-12 candidate was
also removed after an independent 6x30 repeat regressed 1.28% with 0/6 wins.
Other broad candidate policies looked promising in process-per-config sweeps,
but did not survive interleaved A/B validation.

Validation

Correctness and configuration invariance:

python3 tests/test_mega_moe.py \
  --num-processes 8 --num-tokens 512 --num-max-tokens-per-rank 512 \
  --num-experts 256 --num-topk 8 --hidden 7168 --intermediate-hidden 2048 \
  --skew-alpha 1.5 --validate-config-invariance

Robust baseline/adaptive A/B (four measurements per side by default,
alternating order, taking the slowest of all eight ranks in each repetition,
then reporting the median):

bash scripts/bench_adaptive_wave_ab.sh

B200 performance

The final-source validation used 8×B200, six alternating baseline/adaptive
measurements per case, and 30 profiled kernel calls per measurement. Distributed
latency is reduced as median_repeat(max_rank(latency)): each repetition first
takes the slowest of all eight ranks, then the six slowest-rank samples are
reduced by their median.

The common-token sweep used 256 experts, top-k 8, EP 8, hidden 7168, and
intermediate hidden 2048. Values below are observed baseline/adaptive deltas:

tokens/rank	expected TPE	alpha 0.0	alpha 1.0	alpha 1.5
64	16	+0.933%	-1.223%	+0.752%
128	32	-1.015%	-0.077%	-0.361%
256	64	-0.246%	-0.280%	-0.168%
384	96	-0.763%	+0.676%	-1.009%
512	128	-0.789%	+0.041%	+1.268%
1024	256	+0.059%	-0.702%	+0.043%
2048	512	+0.087%	+0.041%	+0.114%

Only the 512-token, alpha-1.5 cell changes the production kernel configuration
(upstream wave 16 to adaptive wave 8). Its exact-delivery-source result was
352.642 → 348.225 µs, a 1.268% speedup with 5/6 wins. The same wave-8 branch
was positive across six independent runs and multiple B200 nodes: 0.634%,
1.365%, 2.231%, 1.580%, 2.072%, and 1.268% (median 1.473%).

Every other cell retains the upstream block and wave configuration. Their
approximately -1.23% to +0.93% variation is the measured same-configuration
noise band and is not attributed to the policy. Two broader candidates were
explicitly rejected: 384-token high-skew wave 8 regressed 1.481% with 0/6 wins,
and 512-token balanced wave 12 regressed 1.276% with 0/6 wins in an independent
repeat.

Forced-wave calibration is available separately:

bash scripts/bench_mega_moe_wave_size.sh

DG_MEGA_MOE_FORCE_EXPERTS_PER_WAVE and DG_MEGA_MOE_FORCE_BLOCK_M are
benchmark-only overrides. They fail loudly on invalid values and are not used by
the production adaptive policy.

co-author with agent.

_{🔄 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/deepseek-ai/DeepGEMM/pull/381 **Author:** [@qinqinwo](https://github.com/qinqinwo) **Created:** 7/15/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feature/imbalance-aware-block-m` --- ### 📝 Commits (5) - [`3f25a3a`](https://github.com/deepseek-ai/DeepGEMM/commit/3f25a3a427f7c3feae977c6849beafeab506ccee) MegaMoE: imbalance-aware block_m selection (opt-in, strict fallback) - [`458dd21`](https://github.com/deepseek-ai/DeepGEMM/commit/458dd21b376ff9bfa13d069bf7e1ef84e04e0eb9) MegaMoE: add calibrated adaptive wave sizing - [`af06265`](https://github.com/deepseek-ai/DeepGEMM/commit/af06265d39c22c340370ef78485f94f7c9eac25f) MegaMoE: narrow adaptive waves to stable sparse routing - [`c343597`](https://github.com/deepseek-ai/DeepGEMM/commit/c343597a793e6bc1133dee41dfc3ef23ee80b5c2) Docs: add B200 performance results in English - [`5daaaab`](https://github.com/deepseek-ai/DeepGEMM/commit/5daaaab0318d02e0511ad069d43109c0c93ac2ba) Fix adaptive receive-stat cache isolation ### 📊 Changes **7 files changed** (+714 additions, -49 deletions) <details> <summary>View changed files</summary> ➕ `ADAPTIVE_WAVE.md` (+93 -0) 📝 `csrc/apis/mega.hpp` (+115 -1) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+127 -42) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+6 -2) ➕ `scripts/bench_adaptive_wave_ab.sh` (+113 -0) ➕ `scripts/bench_mega_moe_wave_size.sh` (+68 -0) 📝 `tests/test_mega_moe.py` (+192 -4) </details> ### 📄 Description MegaMoE processes a fixed number of local experts in each L1→L2 scheduler wave. The upstream heuristic derives that number from the mean tokens per expert. That is a good general fallback, but it cannot see the realized routing distribution. Set `DG_MEGA_MOE_ADAPTIVE_WAVE=1` to enable the opt-in B200 FP8×FP4 policy. The policy: - reads a delta of `cumulative_local_expert_recv_stats` from the preceding window; - caches the sampled distribution and refreshes it every 256 launches, avoiding a synchronous device-to-host copy on the steady-state path; - only changes the calibrated shape (EP 8, 256 experts, top-k 8, hidden 7168, intermediate 2048) and `127.5 < expected tokens/expert <= 128.5` band; - skips receive-stat sampling entirely outside that calibrated shape and band; - uses 8 experts/wave when the active-expert ratio is at or below 0.92; balanced and moderate-skew routing retain the upstream size; - falls back to the upstream wave size on the first call, after counter resets, outside the calibrated tier, or when the requested tier exceeds ring capacity. The deliberately narrow gate is based on same-process, order-balanced 8×B200 measurements. A broader lower bound was rejected after the 96 tokens/expert point regressed 1.48% under high skew. A balanced-routing wave-12 candidate was also removed after an independent 6x30 repeat regressed 1.28% with 0/6 wins. Other broad candidate policies looked promising in process-per-config sweeps, but did not survive interleaved A/B validation. ## Validation Correctness and configuration invariance: ```bash python3 tests/test_mega_moe.py \ --num-processes 8 --num-tokens 512 --num-max-tokens-per-rank 512 \ --num-experts 256 --num-topk 8 --hidden 7168 --intermediate-hidden 2048 \ --skew-alpha 1.5 --validate-config-invariance ``` Robust baseline/adaptive A/B (four measurements per side by default, alternating order, taking the slowest of all eight ranks in each repetition, then reporting the median): ```bash bash scripts/bench_adaptive_wave_ab.sh ``` ## B200 performance The final-source validation used 8×B200, six alternating baseline/adaptive measurements per case, and 30 profiled kernel calls per measurement. Distributed latency is reduced as `median_repeat(max_rank(latency))`: each repetition first takes the slowest of all eight ranks, then the six slowest-rank samples are reduced by their median. The common-token sweep used 256 experts, top-k 8, EP 8, hidden 7168, and intermediate hidden 2048. Values below are observed baseline/adaptive deltas: | tokens/rank | expected TPE | alpha 0.0 | alpha 1.0 | alpha 1.5 | |---:|---:|---:|---:|---:| | 64 | 16 | +0.933% | -1.223% | +0.752% | | 128 | 32 | -1.015% | -0.077% | -0.361% | | 256 | 64 | -0.246% | -0.280% | -0.168% | | 384 | 96 | -0.763% | +0.676% | -1.009% | | 512 | 128 | -0.789% | +0.041% | **+1.268%** | | 1024 | 256 | +0.059% | -0.702% | +0.043% | | 2048 | 512 | +0.087% | +0.041% | +0.114% | Only the 512-token, alpha-1.5 cell changes the production kernel configuration (upstream wave 16 to adaptive wave 8). Its exact-delivery-source result was 352.642 → 348.225 µs, a **1.268% speedup with 5/6 wins**. The same wave-8 branch was positive across six independent runs and multiple B200 nodes: 0.634%, 1.365%, 2.231%, 1.580%, 2.072%, and 1.268% (median **1.473%**). Every other cell retains the upstream block and wave configuration. Their approximately -1.23% to +0.93% variation is the measured same-configuration noise band and is not attributed to the policy. Two broader candidates were explicitly rejected: 384-token high-skew wave 8 regressed 1.481% with 0/6 wins, and 512-token balanced wave 12 regressed 1.276% with 0/6 wins in an independent repeat. Forced-wave calibration is available separately: ```bash bash scripts/bench_mega_moe_wave_size.sh ``` `DG_MEGA_MOE_FORCE_EXPERTS_PER_WAVE` and `DG_MEGA_MOE_FORCE_BLOCK_M` are benchmark-only overrides. They fail loudly on invalid values and are not used by the production adaptive policy. co-author with agent. --- <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-26 21:17:20 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:36 +00:00

[PR #404] feat: add opt-in MegaMoE EVICT_FIRST weight loads #422