Add B200-Calibrated Sparse-Routing Adaptive Wave Sizing by qinqinwo · Pull Request #381 · deepseek-ai/DeepGEMM

qinqinwo · 2026-07-15T11:46:44Z

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.

Size block_m from realized per-expert token counts under skewed routing instead of the uniform mean, minimizing a wall-cost objective (padded rows weighted by SM100 UMMA M-utilization efficiency). Gated by DG_MEGA_MOE_IMBALANCE_AWARE_BLOCK_M; strict fallback guarantees the chosen block_m is never coarser than the default, so the balanced case is unchanged. Adds host-side unit tests, a Zipf --skew-alpha router bias for A/B benchmarking, and an on-GPU baseline-vs-adaptive harness.

ds-review-bot · 2026-07-15T11:52:28Z

+            int calls_since_refresh = 0;
+            bool has_cached_delta = false;
+        };
+        static thread_local AdaptiveRecvStatsCache cache;


🔴 critical: Keep an adaptive cache per receive-counter tensor: When one host thread invokes multiple calibrated MegaMoE layers with distinct receive-stat tensors, this single thread_local cache is overwritten on every layer. Alternating A/B counters therefore always makes pointer_changed true, so no delta is ever valid, adaptivity never activates, and every invocation performs a synchronous device-to-host copy. Cache entries need to be maintained independently per counter tensor.

_{🤖 v6}

Fixed in 5daaaab. The cache is now a bounded thread-local map keyed by TensorImpl identity plus the CUDA data pointer. Each entry retains the tensor identity, so alternating A/B layers keep independent snapshots. The new regression alternates two counters on one host thread; B200 job 2310282 passed and showed the adaptive wave-8 config after independent snapshots.

ds-review-bot · 2026-07-15T11:52:28Z

+        const bool should_refresh = pointer_changed or not cache.has_cached_delta or
+            cache.calls_since_refresh >= kRefreshInterval;


🟡 warning: Back off sampling after zero-delta windows: If a calibrated rank receives no tokens, delta_sum == 0 leaves has_cached_delta false, and this condition forces another synchronous CUDA-to-CPU copy on every subsequent launch instead of waiting 256 launches. Highly sparse routing can therefore impose a stream synchronization on every call; failed or zero samples should still observe the refresh interval.

_{🤖 v6}

Fixed in 5daaaab. Only the initialization snapshot requests an immediate second sample. A zero or negative/reset delta clears the cached delta and then observes the normal 256-launch refresh interval, instead of synchronizing on every call. The 100%-masked zero-delta validation passed on 8xB200 in job 2310282.

ds-review-bot · 2026-07-15T11:52:28Z

@@ -92,13 +92,41 @@ static int get_num_wave_pool_tokens(
    );


🟡 warning: Block_m is shrunk under skew via a purely THEORETICAL wall-cost/MMA-efficiency model and the broad min(eff, mean_tpe) fallback. That still permits pushing block_m into the small-M region the project DELISTED in history (surviving comment here says B200 measured up to 31% regressions under skew; the sibling adaptive-wave path was adopted instead). Gate/calibrate block_m adaptivity to the narrow tokens/expert bands and active-ratio that actually WON on-GPU A/B -- or keep it opt-in off-by-default pending an on-GPU perf gate. Otherwise the adaptive path can be SLOWER than the mean baseline in production.

_{🤖 v4}

The production imbalance-aware block_m selector was removed in af06265 after the B200 regressions. The current helper returns the upstream mean tokens/expert unchanged, except for the explicit benchmark-only force override. Production adaptivity changes only the narrowly calibrated wave size.

ds-review-bot · 2026-07-15T11:52:28Z

    const auto num_sf_ring_tokens = static_cast<int>(l1_acts_sf.size(0));

    // Heuristics
+    // NOTES: `host_recv_stats` are the sampled per-local-expert token counts used


🟡 warning: The block_m selector consumes the SAME cumulative per-expert receive counter that adaptive-wave sizing uses, but LACKS the adaptive-wave guard (only use realized stats after a real dispatch; never on the first call/after a counter reset; fall back to default). Block_m is a whole-run/joint-with-wave tiling; a stale/meaningless first-iteration/garbage counter can mis-select block_m and keep switching it across iterations. The config-invariance harness misses this because it FORCES distinct block_m rather than running the auto-selector. Apply the same sampling/guard/fallback discipline.

_{🤖 v4}

This no longer applies to the current head: block_m does not consume receive statistics. The only production consumer is adaptive wave sizing, and mega.hpp passes it only a validated per-window delta after the initial snapshot/reset guard.

ds-review-bot · 2026-07-15T11:52:28Z

+        constexpr int kImbalanceFactor = 2;
+
+        // Count L1 blocks per expert assuming tokens are evenly spread across experts
+        const float num_expected_tokens_per_expert = static_cast<float>(num_tokens * num_topk) / num_experts_per_rank;


🟡 warning: Block_m is not an independent perf knob: the candidate block_m feeds the wave/pool-token sizing (alignment kLCMCandidateBlockM, pool-token clamp, num_expected_m_blocks/kImbalanceFactor). Changing block_m under skew simultaneously changes the chosen experts-per-wave in an unintended/contradictory way, switching the whole config non-monotonically. Harmonize the block_m + wave choice via a single cost objective/preference/order, and assert the config/performance choice is consistent/sensible (not just that output is bit-identical).

_{🤖 v4}

The coupled production block_m policy was removed in af06265. Production block_m is the upstream choice; the opt-in adaptive policy changes only wave 16 to wave 8 in the single measured tier. The remaining force controls are benchmark-only calibration hooks.

ds-review-bot · 2026-07-15T11:52:28Z

+    const float mean_tpe = static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
+
+    // Benchmark/calibration override. This deliberately bypasses both the
+    // imbalance gate and strict fallback so every production tier can be timed


🔵 suggestion: The FORCE_BLOCK_M override returns EARLY in the SAME helper that the imbalance path shares, after mapping block_m to a per-tier eff. The three-way precedence FORCE > IMBALANCE > mean (and FORCE bypassing BOTH the env gate and the strict fallback, plus the independent ADAPTIVE_WAVE flag) is undocumented, untested and fragile. Document/test the full env/flag precedence, e.g. by extending the existing config-invariance harness.

_{🤖 v4}

The current config-invariance harness exercises all 6 forced block_m tiers and all 6 forced wave tiers, and the source documents both controls as benchmark-only. The old production imbalance precedence discussed here was removed in af06265; forced wave still intentionally returns before the adaptive policy for calibration.

ds-review-bot · 2026-07-15T11:52:29Z

+        hidden == 7168 and intermediate_hidden == 2048;
+    const bool use_adaptive_stats =
+        get_env<int>("DG_MEGA_MOE_ADAPTIVE_WAVE", 0) != 0 and
+        is_calibrated_shape and


🟡 warning: thread_local AdaptiveRecvStatsCache keys only on device_ptr and numel. With a caching allocator, a freed tensor's device address can be reused by a different logical counter tensor with identical numel, causing a delta to be computed against a stale snapshot. The negative-delta guard catches decreases, but a monotonic-looking address collision would silently produce a wrong delta. Risk is low (opt-in + single gated shape), but consider documenting the assumption or adding a cheaper guard such as validating the tensor's storage generation/data version.

_{🤖 v3}

Addressed in 5daaaab. The cache key now includes both the logical TensorImpl identity and CUDA data pointer, and each bounded cache entry retains a shallow tensor reference. This prevents both identities and the underlying CUDA allocation from being recycled into a false hit while the entry is live.

ds-review-bot · 2026-07-15T11:52:29Z

+        constexpr int kRefreshInterval = 256;
+
+        const bool pointer_changed = cache.device_ptr != device_ptr or
+            cache.previous_cumulative.size() !=


🔵 suggestion: kRefreshInterval = 256 is a magic constant with no env override. Under a workload whose routing skew shifts within 256 iterations, the sampled distribution (and thus the wave choice) lags. Acceptable for the stated calibration point, but a brief comment explaining why 256 (vs e.g. 64) was chosen would aid future tuning.

_{🤖 v3}

Addressed in 5daaaab. The source and ADAPTIVE_WAVE.md now explain that 256 was selected for the stationary B200 calibration to amortize the synchronous D2H copy while periodically detecting distribution changes.

ds-review-bot · 2026-07-15T11:52:29Z

+    // The production default remains the strict fallback. Adaptive sizing uses
+    // the previous iteration's realized local-expert receive counts, so it is
+    // opt-in and never guesses during the first call or after a counter reset.
+    if (recv_stats == nullptr or get_env<int>("DG_MEGA_MOE_ADAPTIVE_WAVE", 0) == 0)


🔵 suggestion: get_env("DG_MEGA_MOE_ADAPTIVE_WAVE") is read here and again in mega.hpp on every launch. If get_env performs a real getenv each call this is a small redundant cost on the steady-state path; if it is cached this is fine. Worth confirming get_env is cheap.

_{🤖 v3}

Confirmed that get_env calls getenv. Fixed in 5daaaab by removing the second read from the heuristic: a non-null recv_stats pointer now means the API opt-in shape/band gate and delta validation already succeeded.

ds-review-bot · 2026-07-15T11:52:29Z

@@ -0,0 +1,88 @@
+# MegaMoE Adaptive Wave Sizing


🔵 suggestion: The PR/branch title 'imbalance aware block m' is misleading: the imbalance-aware block_m policy was actually removed, and the calibration override deliberately bypasses the imbalance gate. The delivered feature is 'opt-in adaptive wave sizing'. Consider retitling for changelog/history clarity. Also note the doc says 'EP 8, 256 experts' while the code gate checks num_experts_per_rank==32 and num_ranks==8 (=256 total); consistent but the gate never validates total num_experts directly.

_{🤖 v3}

The PR title is already updated to Add B200-Calibrated Sparse-Routing Adaptive Wave Sizing. Commit 5daaaab also makes the calibrated gate explicitly require num_experts == 256 in addition to 8 ranks and 32 experts/rank.

ds-review-bot · 2026-07-15T11:52:31Z

🤖 ds-review-bot Code Review

v6

The adaptive cache cannot support multiple MegaMoE counters on one thread and can repeatedly synchronize on zero-token ranks. These issues can disable the feature or materially regress its intended workloads.

v4

MR review result: imbalance-aware block_m (re-land). This branch was removed from the line in history after B200 A/B measured up to 31% regressions under skew, and the project instead adopted calibrated, narrow adaptive-WAVE sizing (DG_MEGA_MOE_ADAPTIVE_WAVE) which (a) never guesses on the first call/after a counter reset and (b) only enters calibrated tokens/expert bands that actually WON on-GPU A/B. The present re-land restores an OPT-IN (DG_MEGA_MOE_IMBALANCE_AWARE_BLOCK_M) block_m selector (recv_stats host snapshot is already wired in the impl/apis) that PURELY shrinks the tile-M under skew using a theoretical wall-cost/MMA-efficiency model + a broad min(eff, mean_tpe) strict fallback. It is joined at the hip with the wave/pool sizing and the FORCE_BLOCK_M override, but the joint behaviour, env/flag priority, stale-cumulative-stats guard, and perf calibration are all UNTESTED/unharmonized (the config-invariance harness only checks OUTPUT identity, which any block_m trivially satisfies).

The core reconstructed logic of the re-land (for reference, in this repo's removed commit, 7 files: heuristics/mega_moe.hpp, apis/mega.hpp, impls/..., IMBALANCE_BLOCK_M.md, scripts/bench_imbalance_block_m.sh, tests/test_imbalance_block_m.py, tests/test_mega_moe.py) is:
eff_tpe = mean for env=0/recv_stats null, OR after a FORCE_BLOCK_M override(returns first) the per-tier eff; for env=1: minimize wall_cost(b)=sum_e ceil(c_e/b)*b/sigma(b) over the candidate block_m (eta table 8/16/32/64/96/128/192), map best block_m back to its tier-upper-edge eff, then take min(eff, mean).

VERDICT: REQUEST CHANGES / conditional-on-calibration. The opt-in performance claim is unproven on real silicon and the original was delisted for regressions; merge only after a WAVE-style narrow/calibrated re-landing with on-GPU A/B wins and the guards below.

v3

Despite the branch name 'imbalance aware block m', this PR is a conservative, low-risk change. It removes an earlier imbalance-aware block_m shrink policy (which regressed up to 31% on B200 under skew) and replaces it with: (1) an opt-in adaptive wave-sizing policy (DG_MEGA_MOE_ADAPTIVE_WAVE=1) narrowly gated to a single calibrated shape/token band; (2) benchmark-only overrides (DG_MEGA_MOE_FORCE_BLOCK_M, DG_MEGA_MOE_FORCE_EXPERTS_PER_WAVE) that fail loudly on invalid values; (3) new test modes and A/B benchmarking scripts. The production default path is behaviorally unchanged: the adaptive branch only activates when the env var is set AND host recv stats are available AND the shape/band match; otherwise the refactored (but behavior-identical) upstream logic runs. The cumulative->delta conversion in mega.hpp is handled correctly, guarding against counter resets, pointer changes, and zero-sum deltas. The --validate-config-invariance test asserts bit-exact output and receive-stat equality across default, adaptive, all 6 forced block_m tiers, and all 6 forced wave tiers, which is exactly the right check for a heuristic-only change. I verified the DG_MEGA_MOE_FORCE_BLOCK_M -> fake tokens/expert mapping against the block config table and it is correct for all six tiers (16->8.0, 32->16.0, 64->32.0, 96->64.0, 128->96.0, 192->128.0), each landing in the intended branch. Documentation is honest and detailed, including rejected candidates and win rates. Overall: approve after addressing the thread_local cache aliasing note and considering a title change; no blocking correctness issues found in the production path.

Files reviewed: 7
Issues found: 🔴 1 critical | 🟡 5 warning | 🔵 5 suggestion
Inline comments posted: 10
General comments (无法定位到 diff): 1

📍 未定位到 diff 的评论

🔵 suggestion tests/test_imbalance_block_m.py:L1: The host-side unit test only COPIES the heuristic formula in pure Python (eff/edge/wall-cost/fallback + FORCE map) instead of exercising the C++ binding, so it cannot detect a behaviour/typo divergence, and it SELF-REPORTS allowed regressions/balanced-case changes that mask real problems. Tie the mirror to the actual C++ heuristic (bind/export it or cross-check), and FAIL on the strict-fallback/cost-increase/balanced-case invariants instead of merely printing them. _{🤖 v4}

qinqinwo · 2026-07-15T12:15:13Z

Review update:

Pushed 5daaaab with per-counter cache isolation, allocator-reuse protection, zero-delta/reset sampling backoff, the explicit 256-expert gate, removal of the redundant getenv, documentation, and a two-counter regression.
Rebuilt successfully and passed both calibrated multi-counter and 100%-masked zero-delta configuration-invariance tests on 8xB200 (Slurm job 2310282, exit 0).

Upstream compatibility note: #377 merged while this review was in progress. It replaces the old MegaMoE wave scheduler with a dynamic interleaved task scheduler and removes num_experts_per_wave entirely. The PR therefore now conflicts with main. A mechanical rebase would either delete this feature or reintroduce the superseded scheduler, so I have intentionally not forced such a rebase. The adaptive idea would need to be re-evaluated against the new scheduler rather than transplanted unchanged.

qinqinwo and others added 4 commits July 14, 2026 17:39

MegaMoE: add calibrated adaptive wave sizing

458dd21

MegaMoE: narrow adaptive waves to stable sparse routing

af06265

Docs: add B200 performance results in English

c343597

qinqinwo changed the title ~~Feature/imbalance aware block m~~ Add B200-Calibrated Sparse-Routing Adaptive Wave Sizing Jul 15, 2026

ds-review-bot reviewed Jul 15, 2026

View reviewed changes

Fix adaptive receive-stat cache isolation

5daaaab

		const bool should_refresh = pointer_changed or not cache.has_cached_delta or
		cache.calls_since_refresh >= kRefreshInterval;

Uh oh!

Conversation

qinqinwo commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Validation

B200 performance

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ds-review-bot commented Jul 15, 2026

🤖 ds-review-bot Code Review

v6

v4

v3

📍 未定位到 diff 的评论

Uh oh!

qinqinwo commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

qinqinwo commented Jul 15, 2026 •

edited

Loading