Add B200-Calibrated Sparse-Routing Adaptive Wave Sizing

ADAPTIVE_WAVE.md

@@ -0,0 +1,93 @@
+# MegaMoE Adaptive Wave Sizing
+
+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;
+- maintains an independent bounded cache entry for each logical receive-counter
+  tensor, so multiple MegaMoE layers can alternate on one host thread;
+- caches the sampled distribution and refreshes it every 256 launches. This
+  interval was calibrated for stationary B200 routing and amortizes the
+  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 or
+  zero-delta samples, outside the calibrated tier, or when the requested tier
+  exceeds ring capacity. Reset and zero-delta samples observe the same refresh
+  interval instead of synchronizing on every launch.
+
+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.

csrc/apis/mega.hpp

@@ -2,6 +2,8 @@
 
 #include <functional>
 #include <string>
+#include <unordered_map>
+#include <vector>
 #include <pybind11/functional.h>
 
 #include <deep_gemm/common/types.cuh>
@@ -229,6 +231,117 @@ static void fp8_fp4_mega_moe(
     // Already registered tensors
     const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer);
 
+    // Adaptive heuristics consume a per-iteration receive-count distribution.
+    // The public tensor is cumulative, so using its absolute values makes the
+    // selected JIT config drift as the counter grows. Snapshot the cumulative
+    // values and pass only a valid delta from the previous snapshot. Sampling is
+    // opt-in, restricted to the calibrated tokens/expert band, and amortized
+    // over 256 launches; the cached distribution also avoids host synchronization
+    // and JIT-config churn on the steady-state path. Each logical counter tensor
+    // has an independent bounded cache entry, so alternating MegaMoE layers do
+    // not invalidate one another. The first call, counter resets, and zero deltas
+    // fall back to the default heuristic until a new valid delta is observed.
+    const int* host_recv_stats_ptr = nullptr;
+    const float expected_tokens_per_expert =
+        static_cast<float>(num_tokens * num_topk) / num_experts_per_rank;
+    const bool is_calibrated_shape =
+        num_ranks == 8 and num_experts == 256 and num_experts_per_rank == 32 and num_topk == 8 and
+        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
+        expected_tokens_per_expert > 127.5f and expected_tokens_per_expert <= 128.5f;
+    if (use_adaptive_stats and cumulative_local_expert_recv_stats.has_value()) {
+        const auto* device_ptr = cumulative_local_expert_recv_stats->data_ptr<int>();
+        struct AdaptiveRecvStatsKey {
+            const void* tensor_impl;
+            const int* device_ptr;
+
+            bool operator==(const AdaptiveRecvStatsKey& other) const {
+                return tensor_impl == other.tensor_impl and device_ptr == other.device_ptr;
+            }
+        };
+        struct AdaptiveRecvStatsKeyHash {
+            size_t operator()(const AdaptiveRecvStatsKey& key) const {
+                const auto tensor_hash = std::hash<const void*>{}(key.tensor_impl);
+                const auto data_hash = std::hash<const int*>{}(key.device_ptr);
+                return tensor_hash ^ (data_hash + (tensor_hash << 6) + (tensor_hash >> 2));
+            }
+        };
+        struct AdaptiveRecvStatsCache {
+            // Retain the tensor identity while this entry exists. Besides making
+            // the logical identity explicit, this prevents its TensorImpl and
+            // CUDA allocation from being recycled into a false cache hit.
+            torch::Tensor counter_identity;
+            std::vector<int> previous_cumulative;
+            std::vector<int> cached_delta;
+            int calls_since_refresh = 0;
+            bool has_previous_snapshot = false;
+            bool sample_next_call = false;
+            bool has_cached_delta = false;
+        };
+        static thread_local std::unordered_map<
+            AdaptiveRecvStatsKey, AdaptiveRecvStatsCache,
+            AdaptiveRecvStatsKeyHash> caches;
+
+        // The B200 A/B calibration used a stationary routing distribution. A
+        // 256-launch interval amortizes the synchronous D2H copy while still
+        // periodically detecting workload changes. Keep the cache bounded for
+        // models that create transient counter tensors on a long-lived thread.
+        constexpr int kRefreshInterval = 256;
+        constexpr size_t kMaxCachedCounters = 64;
+        const AdaptiveRecvStatsKey cache_key = {
+            cumulative_local_expert_recv_stats->unsafeGetTensorImpl(), device_ptr};
+        auto cache_it = caches.find(cache_key);
+        if (cache_it == caches.end()) {
+            if (caches.size() >= kMaxCachedCounters)
+                caches.clear();
+            cache_it = caches.try_emplace(cache_key).first;
+            cache_it->second.counter_identity = *cumulative_local_expert_recv_stats;
+        }
+        auto& cache = cache_it->second;
+
+        // Count this launch before testing the periodic deadline so refreshes
+        // are exactly kRefreshInterval launches apart after initialization.
+        ++ cache.calls_since_refresh;
+        const bool has_matching_snapshot = cache.has_previous_snapshot and
+            cache.previous_cumulative.size() ==
+                static_cast<size_t>(cumulative_local_expert_recv_stats->numel());
+        const bool should_refresh = not has_matching_snapshot or
+            cache.sample_next_call or cache.calls_since_refresh >= kRefreshInterval;
+
+        if (should_refresh) {
+            const auto cpu = cumulative_local_expert_recv_stats->to(torch::kCPU, torch::kInt);
+            const auto* current_ptr = cpu.data_ptr<int>();
+            bool valid_delta = has_matching_snapshot;
+            int64_t delta_sum = 0;
+            cache.cached_delta.resize(cpu.numel());
+            if (valid_delta) {
+                for (int i = 0; i < cpu.numel(); ++ i) {
+                    const int delta = current_ptr[i] - cache.previous_cumulative[i];
+                    if (delta < 0) {
+                        valid_delta = false;
+                        break;
+                    }
+                    cache.cached_delta[i] = delta;
+                    delta_sum += delta;
+                }
+            }
+
+            cache.previous_cumulative.assign(current_ptr, current_ptr + cpu.numel());
+            cache.has_previous_snapshot = true;
+            // Only the initial snapshot samples again immediately so adaptivity
+            // can start on the next launch. Zero deltas and counter resets back
+            // off for the regular interval instead of synchronizing every call.
+            cache.sample_next_call = not has_matching_snapshot;
+            cache.has_cached_delta = valid_delta and delta_sum > 0;
+            cache.calls_since_refresh = 0;
+        }
+
+        if (cache.has_cached_delta)
+            host_recv_stats_ptr = cache.cached_delta.data();
+    }
+
     // Dispatch into different architectures
     if (arch_major == 10) {
         sm100_fp8_fp4_mega_moe(y,
@@ -242,7 +355,8 @@ static void fp8_fp4_mega_moe(
                                num_experts_per_rank,
                                num_tokens, num_topk,
                                hidden, intermediate_hidden,
-                               activation_clamp, fast_math);
+                               activation_clamp, fast_math,
+                               host_recv_stats_ptr);
     } else {
         DG_HOST_UNREACHABLE("Unsupported architecture");
     }