[hao-aaron]

Purpose

Part of #64418 (load plan), item L2. MegaMoE NVFP4 experts now follow the loader contract (§5):

• each checkpoint tensor loads on its own, in any order and in any split across load_weights calls;
• the loader only views and copy_s;
• combining tensors happens in post-load processing, which reruns on every reload.

Before this PR, DeepGemmMegaMoEExperts (MiniMax-M3 NVFP4 on deep_gemm_mega_moe) did three things a load-plan trace cannot replay:

• it computed the L2 activation scale as a running max over a CPU copy of every expert's w2_input_scale;
• it paired w1 and w3 weight_scale_2 through loader state and wrote both into one slot;
• it ran block scales through repeat_interleave + clone even when nothing is expanded.

Changes

In deep_gemm_mega_moe.py:

• w13_weight_scale_2 is now (2, local experts), rows w1 and w3. Each shard writes its own slot. Row 0 is a contiguous (e,) view, and the kernel reads it directly as l1_alpha. EPLB gets the two rows as separate expert tensors.
• w2_input_scale is now (logical experts,). It is written for every expert on every rank, as the CPU copy was, so all ranks agree on the L2 activation scale. It is not an EPLB expert tensor.
• New process_weights_after_loading. It checks that w1 and w3 match (same error as before, now raised after loading), then sets l2_activation_scale = max(w2_input_scale) or 1.0. Direct, layerwise and IPC weight-cache loads all rerun it. A side effect on the IPC path: base left l2_activation_scale at 1.0 there, and it is now computed.
• The _l2_alpha buffer is marked weight_loader_numel = 0. No loader writes it, so layerwise reload doesn't wait for it.
• _pack_scale expands or copies only when it has to. NVFP4 1x16 and MX 1x32 block scales become a pure view chain (byte→int32 reinterpretation plus interleave/UTCCP views). The DeepSeek-V4/Kimi-K3 MXFP4 scales get this too. Bytes are unchanged.
• Dummy loads: w13_weight_scale_2 gets 1e-3, as ModelOpt NVFP4 MoE does, so w1 and w3 match. w2_input_scale gets 1.0, which keeps the dummy L2 scale at 1.0 as before.

Block scales stay in the kernel layout. That layout is a permutation of the checkpoint bytes, so it is already a view (§4). Packing them in PWAL instead would cost a second copy of the scales or a non-idempotent in-place transform.

Not in this PR:

• Passing the L2 activation scale to the kernels as a tensor. That is [RFC]: Direct Weight Reload for RL Training #54477 F16 and needs [SM100] NVFP4 MegaMoE: take the L2 activation scale as a device tensor DeepGEMM#49. Until then, a reload updates l2_activation_scale, but graphs captured earlier keep the old value.
• MiniMax-M3 compressed-tensors conversions (1/x, fixed 1/16). They are per tensor, but they are not views, so those names stay fallbacks.
• Other MegaMoE layers. Step-5's loader already conforms, but its PWAL runs once (an F15-style fix). Muse Spark is F15.

This does not duplicate other work: #60096 (merged) added the kernel-layout loading this builds on, and no open PR touches these scales.

Test plan

• Cold-start equivalence, using a one-off script that is not in this PR. A tiny EP layer with one replica loads random NVFP4/MXFP4 bytes (with and without a fused MXFP8 shared expert) one shard per call, in shuffled order. Every tensor the kernels read is compared against base.

• New CPU tests in tests/kernels/moe/test_deep_gemm_mega_moe.py:

  • direct_reload_matches_cold_start (MXFP4 and NVFP4)
  • layerwise_reload_does_not_buffer_until_finalize
  • routed_loaders_only_view_and_copy (traced with FakeRDTTensor)
  • eplb_views_move_whole_experts
  • loads_and_processes_meta_tensors

  Existing tests were updated to run post-load processing.

• GB200: tests/kernels/moe/test_deep_gemm_mega_moe.py and tests/models/test_minimax_m3_mega_moe.py. These check the kernel layout against DeepGEMM's transform, run MegaMoE against a dequantized reference for every format, and check the MiniMax-M3.1 block against its QAT reference.

Test result

• Cold start: every kernel-read tensor is byte-equal to base in all configs. This needs w1 and w3 weight_scale_2 to be bitwise equal. Within the 1e-5 tolerance the check allows, base used whichever was loaded last; this PR always uses w1.
• CPU (CUDA_VISIBLE_DEVICES=""): 27 passed; the GPU-only kernel-layout tests are covered by the GB200 run.
• GB200: tests/kernels/moe/test_deep_gemm_mega_moe.py + tests/models/test_minimax_m3_mega_moe.py, 74 passed (prepare-kernel bitwise tests deselected; that code is unchanged).

A/B

Setup: A = main 4a0dfc7b2f, B = this PR on the same commit. Both are Python overlays on the 4a0dfc7b2f nightly wheel. Each row ran on one node, with the arms back to back.
GSM8K: 5-shot, greedy, 1319 questions. Throughput: vllm bench serve, random 1024 in / 256 out, concurrency 64, about 2.5 min.
Every arm's server log has a probe line showing that DeepGemmMegaMoEExperts was built, with the expected format (nvfp4 or mxfp4) on every rank. The probe is the same in both arms.

Model (hardware, TP)	GSM8K A → B	Other	Throughput tok/s A → B	KV cache A → B
MiniMax-M3.1-Flash, MXFP8 + NVFP4 MegaMoE, cold start (GB200, TP4 EP)	0.9613 → 0.9636		3721 → 3610	2.515M → 2.515M
MiniMax-M3.1-Flash, after scrambled sync + real sync, VLLM_RELOAD_MODE=direct (GB200, TP2 EP)	0.9666 → 0.9583	free / reserved after sync, GiB: 3.31 / 121.84 → 3.31 / 121.84; real-weights sync: 72.9 s → 101.9 s	2661 → 2630	721k → 721k
MiniMax-M3.1-Flash, after scrambled sync + real sync, layerwise, --gpu-memory-utilization 0.85 (GB200, TP2 EP)	0.9629 → 0.9613	free / reserved after sync, GiB: 26.20 / 121.84 → 26.20 / 121.84; peak nvidia-smi during sync: 169.6 GiB → 169.6 GiB	2635 → 2605	401k → 401k
MiniMax-M3.1-Flash, scrambled sync, layerwise, default --gpu-memory-utilization 0.9 (GB200, TP2 EP)	both arms: CUDA OOM during the scrambled sync	free before the sync: 4.4 GiB per rank (both arms); the OOM is a 1 GiB packed-transfer chunk with 124.71 GiB allocated, identical in A and B	n/a	721k → 721k
DeepSeek-V4-Flash, MXFP4 MegaMoE, cold start (GB200, TP2 EP)	0.9462 → 0.9462; swapped: 0.9492 → 0.9484		2497 → 2839; swapped: 3195 → 2858	6.157M → 6.156M; swapped: 6.156M → 6.157M

Notes:

• How to read the noise.

  • GSM8K differs by at most 0.8 pt and has no sign.
  • Throughput is within 3% on MiniMax. The largest gap is TP4 at -3.0%, which sits right at the swap threshold, so no swap was run.
  • On DeepSeek-V4 the arm that ran second was 12–28% faster in both orders. The swap removes the gap: the A mean is 2846 tok/s and the B mean is 2848 tok/s.
  • KV cache tokens match to within 0.01%.

• Sync rows. Both arms use the same protocol:

  1. cold start;
  2. one packed NCCL update of the whole checkpoint with every tensor's elements rolled by one (same names, shapes and dtypes);
  3. one update of the real checkpoint;
  4. then the measurements.

  Device digests of every registered and off-registry tensor confirm, in both arms and both modes:

  • the scrambled update changed 3160 of 4104 registered tensors (both ranks counted);
  • the real update restored all 4104 to their cold-start bytes;
  • no tensor moved.

  The real-weights sync time in the direct row (72.9 s → 101.9 s) is a single sample. The scrambled step took 91 s → 95 s, and the layerwise row has 73 s → 72 s, so it looks like trainer-side checkpoint read noise.

• Layerwise at the default --gpu-memory-utilization 0.9, TP2. Both arms OOM at the same point, with the same allocation:

  • PyTorch had 124.71 GiB allocated when a 1 GiB packed-transfer chunk was requested;
  • the cold start had left only about 4.4 GiB free per rank;
  • the base OOM'd in job 43864 and this PR in job 43877.

  So this PR neither causes nor fixes it. With 0.85 (the row above) both arms complete. Direct mode needs no extra memory and completes at 0.9.

• Layerwise buffering is unchanged by this PR here. In both arms, layerwise reload logs a 4080 MB buffer for one DeepGemmMegaMoEExperts + GateLinear pair. That buffer comes from the checkpoint's tensor order across layers, not from _l2_alpha. The other buffer warnings (12 per rank) are also identical, as is the peak memory during the sync.

• Serve args.

  • MiniMax: --enable-expert-parallel --max-model-len 4096 --kernel-config '{"moe_backend":"deep_gemm_mega_moe","linear_backend_per_quant":{"mxfp8":"flashinfer_trtllm"}}'. GSM8K uses chat completions with reasoning_effort=low and max_tokens 2048. Plain completions degenerate on this chat model (0.08).
  • DeepSeek-V4: --enable-expert-parallel --moe-backend deep_gemm_mega_moe --kv-cache-dtype fp8 --block-size 256 --attention_config.indexer_kv_dtype=mxfp4 --tokenizer-mode deepseek_v4, the CI MegaMoE config without MTP, plain completions.

• Main bugs found along the way. Both are in A too and unrelated to this PR.

  • MiniMax-M3.1 MegaMoE with TP > 1 crashes at the first forward with the auto-selected FlashInfer MXFP8 linear. The fused norm + MXFP8 quant ([Perf][MiniMax-M3] (NVIDIA-only) Fused norm+quant MXFP8 #64852) hands a QuantizedActivation to the sequence-parallel all-gather ([Perf][MiniMax-M3] Sequence parallelism and GEMM-RS for TP + MegaMoE #64659): 'QuantizedActivation' object has no attribute 'contiguous' in sp_all_gather. The TRT-LLM MXFP8 linear avoids it.
  • DeepSeek-V4-Flash MegaMoE at TP4 with --max-model-len 4096 --max-num-batched-tokens 16384 --max-num-seqs 128 crashes during CUDA graph capture: max() iterable argument is empty in sparse_attn_indexer. TP2 with the CI args works.

• Skipped.

  • Mapika/MiniMax-M3-NVFP4: the cache entry is empty, and its ModelOpt NVFP4 config has no moe_quant_algo, so --moe-backend deep_gemm_mega_moe rejects it. MiniMax-M3.1-Flash (MXFP8 + NVFP4 experts, the format MegaMoE targets) was used instead.
  • Kimi-K3 MXFP4 MegaMoE: not run, as DeepSeek-V4-Flash already covers the MXFP4 _pack_scale view path.
  • TP4 sync rows: the NCCL trainer needs a fifth GPU, and the nodes have four.

AI assistance (Claude Code) was used; the submitter reviews every changed line.

🤖 Generated with Claude Code