[None][feat] Update CuTeDSL MegaMoE kernels#16190
Draft
Barry-Delaney wants to merge 10 commits into
Draft
Conversation
…FP4) Integrate the MEGAMOE_CUTEDSL backend (in-kernel NVLink all-to-all fused dispatch + fc1 + fc2 + combine) for DeepSeek-V4 DEP8 serving: - adaptive max_tokens_per_rank bucketing (decode launches a small kernel instead of padding to the prefill max) - bf16 / fp8 / fp4 combine formats (fp8 combine may use the bulk fc2 store) - pure cute.compile JIT launch path (drop the AOT plain-arg path) - workspace zero-init opt (zero only the required leading bytes, once) - spin_wait dispatch-barrier loads use acquire semantics - review fixes: tactic legality, adaptive-bucket overflow, EP fence guard, reload-safe source-weight free, single-rank shared-counter init, autotune fenced-key prune (freed-workspace ABA), python -O safe EP-group guard - legacy-lint config for the vendored CuteDSL kernel sources (pinned to gitlab bangyus/cutedsl_megamoe @ ce712016) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Standalone MPI microbench driving MEGAMOE_CUTEDSL via the public create_moe -> forward path with synthetic NVFP4 weights. Times the eager forward by default; --cuda-graph gives a kernel-bound number (CUDA-graph replay; slowest EP rank, min over iters) matching CUDA-graphed serving. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…he load window The NVFP4 source expert tensors and the packed mega_* buffers used to be materialized together at init (~229 GB/rank on DSv4-Pro DEP8), which cannot fit on 180 GB-class GPUs even though the steady state (~131 GB) would. Keep the four source tensors (w3_w1_weight, w3_w1_weight_scale, w2_weight, w2_weight_scale) as 0-element placeholders outside the load window: create_weights shrinks them via the existing replace_parameter_and_save_metadata idiom, a load_weights override rematerializes them one module at a time, and the existing post-load free block re-shrinks them after packing. Peak becomes steady set plus one layer of sources. Dynamic-EPLB ranks (need_load_shared_weights) keep full sources after load exactly as before. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…ion) - Adaptive-maxT ladder cap: derive from the true per-rank chunk bound min(moe_max_num_tokens, max_num_tokens) instead of moe_max // dp_size, which under-sized the top bucket for explicit moe_config.max_num_tokens and made _select_launch_max_tokens raise on legal scheduler chunks. - Live-token trim: keep the full bucket for num_tokens == 0 (zero-token lockstep launches) so TopkReduce never computes a zero grid. - Pure-DEP invariant: reject WORLD != EP at construction (matching the op's forward-time fence guard) and pre-check before the MPI bootstrap collectives so PP/CP hybrids fail loud instead of hanging. - MPI torch.dist bootstrap: rank-0 free-port draw broadcast unconditionally over MPI (fixed 29561 collided across same-node jobs; per-rank env gating could desync the collective), loopback fallback for unresolvable hostnames. - Streaming load: keep streamed sources as placeholders in pre_reload_weights (a model-wide hot reload re-created the ~229 GB/rank init peak), rebind quant_scales right after the per-module source free (block scales otherwise accumulate until the model-wide sweep), and require full local-expert coverage at finalize when sources were rematerialized (uncovered rows are uninitialized). - Microbench: --autotune now wires AutoTuner distributed state and runs a real tuning forward (it only set an env var before, so timed runs silently used the fallback tactic), releases the tuning-only symmetric scratch before timing, and draws its rendezvous port like the backend. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…g, FC1 SF pad granularity Two review findings inside the vendored mega_moe_nvfp4 kernel (to be upstreamed to the kernel team's swapab branch as well): - token_comm token-back push_sf: promote the fc2_output_sf pool row index (and the peer combine_sf_u8 row, for symmetry with the adjacent Int64 DATA/metadata paths) to Int64 -- the Int32 offset overflows the multi-GiB SF pool at large max_tokens_per_rank x EP. - kernel_fc12 FC1 weight-SF constexpr check: pad the M side to the 128-row SF atom (as the FC2 path and the host packer already do) instead of the K-side sf_vec_size*4 granularity; behavior-identical for 128-aligned shapes like DSv4-Pro. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…ion) - Autotuner: drop the class-scope tuning-config cache -- TuningConfig embeds inputs_pre_hook, a bound method reading the per-call runner's _profiling_scratch, so a cache hit handed the autotuner a hook bound to a dead runner with a stale/cleared scratch (mixing non-symmetric profiling tensors with the live runner's peer offsets: CUDA illegal access + peer barrier hang class). The config is rebuilt per call; the shape-derivation rule is hoisted to module scope so ConstraintSpec hashing stays stable and the AutoTuner's unbounded profile lru_cache does not grow per call. - Streamed-coverage guard: verify per COMPONENT instead of counting stash entries -- an entry created by a w1-only shard used to satisfy the count while the packing loop skipped it, baking uninitialized rows; w2/w2-scale direct-copy loaders now record per-row coverage. - Dynamic EPLB: correct the free-block comment (the raw sources are kept because the inherited load-balancer plumbing registers them as migration targets, not because the kernel re-reads them) and emit a one-time warning about the ~2x expert-weight residency; derived-only migration registration (as MegaMoE-DeepGemm) is the tracked follow-up. - Full-model RLHF hot reload still accumulates sources until the model-wide finalize sweep (documented KNOWN LIMIT; per-bucket finalization is order-sensitive and needs its own PR). Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…ImplementedError fallback The kernel re-import dropped main's separate 'except NotImplementedError' import handlers (PR NVIDIA#14608) in token_comm / fc1_fc2_fuse_sched / kernel_fc12 / moe_persistent_scheduler: cutlass wheels on CUDA < 13.1 raise NotImplementedError (not ImportError) from cute.experimental, so the iket_compat no-op shim never engaged and workspace init died after the capability probe had declared the backend available. Restored as SEPARATE handlers on purpose -- the CuteDSL 4.5.0 preprocessor import-walker raises AttributeError on tuple except types, silently disabling AST preprocessing for the module. Also restores iket_compat's logger.debug notice (the re-import had reverted it to a raw print). Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…eak, EPLB residency, fence ABA, scale-family completeness) - Hot reload: eager per-module finalize on the partial path -- once a module's 4 streamed components and aux scale families are fully covered, pack + free immediately instead of accumulating every touched layer's sources until the model-wide sweep (a layer-contiguous full-model reload now peaks at ~one layer of sources instead of the ~229 GB/rank init peak). Aux-only shards for placeholder modules stash scales without rematerializing (lazy materialization keyed on the flat weights dict); an aux tail after the eager finalize refreshes scales/alphas only; interleaved-layer bucket orderings keep the sweep-time coverage guard as backstop. - Dynamic EPLB: derived-only migration registration (the MegaMoE-DeepGemm pattern). post_load_weights skips the base raw-weight registration, _eplb_migrate_raw_block_scales exempts the parent raw block-scale registration (other NVFP4 children unchanged), and the raw source set is freed unconditionally -- dynamic EPLB no longer doubles the expert-weight residency. bias+dynamic-EPLB now rejected at create time. - Autotune fence ABA: the barrier-reset fence now drops every fenced pairing sharing either workspace side before re-fencing (general-warmup real -> profiling scratch -> real transitions no longer skip into a phase-decoupled dispatch barrier); capture reaching an unfenced pairing fails loud; release_megamoe_profiling_scratch prunes scratch-paired fence keys and promotes the winner's workspace out of the tuning-evict latch (a later eviction would free a graph-captured pointer). - Aux scale families: all-or-nothing completeness per family (weight_scale_2 / w1+w3 / w2 input_scale) -- a partial family raises instead of silently mixing fresh and stale alphas or deriving the global input scale from a subset; weights-only reloads keep previous scales (fc1_norm_const rebuild skipped on an empty stash). - Microbench: --autotune now mirrors the serving general-warmup -> autotuner interleaving and asserts the fence-pairing invariant. - New unit test (1 GPU, no model weights): streamed load/reload lifecycle -- layer-atomic and split-bucket reloads byte-identical to the initial load, aux tails, partial-coverage / partial-family / aux-only-without-weights failure modes. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…n's staged-hook contract Rebasing onto main crossed the Wave-3 MoE staged-hook migration (TRTLLM-13248): nothing calls quant_method.post_load_weights anymore -- module.post_load_weights is a compat shim over transform_weights / cache_derived_state, and the base load_weights tail EAGERLY calls _finalize_shared_weights (where the raw shared-weight registration now lives). The derived-only EPLB override therefore moves from post_load_weights (dead on main) to _finalize_shared_weights: skip the base raw w3_w1/w2 weight_fns registration (the balancer would IndexError slicing the freed 0-element placeholders), drop the private CPU staging, and finalize host sharing -- idempotent across the eager call and the transform_weights safety net, mirroring the base contract. The set_initial_weight_assignments + setup_quant_scales tail of the old override is base cache_derived_state, reached via the module shim. Also: thin transform_weights / cache_derived_state guards on the backend (ConfigurableMoE delegates to them directly; base MoE dereferences quant_method without the create-if-None guard the other entry points carry), drop the dead import_topk_reduce lazy accessor (the rewritten topk_reduce exposes the TopkReduce class; the referenced functions no longer exist), and deduplicate the RcpLimit constant block. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…rdering, alpha consistency, scratch lifecycle, store alignment)
- Hot reload: partial-reload sources now stage on CPU
(_materialize_source_params gained a device param) and the finalize
packing uploads one layer transiently -- the reload GPU peak is
ordering-INDEPENDENT, covering the in-tree RLHF flow that buckets by
tensor-name suffix across all layers (previously re-accumulated every
layer's CUDA sources until the last suffix). CPU-staged stash shards
are alias-broken (an already-CPU input's .to('cpu') is a no-op, and
the stash is consumed buckets later).
- Aux scale families: bidirectional dependency -- input_scale families
without weight_scale_2 (and vice versa) are rejected; the alphas
derive from BOTH, so a one-sided refresh silently mismatches the
runtime quantization scales. The eager-finalize gate mirrors the
legal set exactly (w1/w3 input_scale may be omitted).
- Profiling scratch lifecycle: serving never called
release_megamoe_profiling_scratch on main (the call lived in a
feat-side model_engine hunk the rebase dropped) -- ported into
_run_autotuner_warmup per engine, also on the enable_autotuner=False
early return (MEGAMOE_AUTOTUNE=1 allocates the scratch regardless).
release now EVICTS tuning-latched workspaces instead of bare-clearing
the latch (default-tactic-wins left the last loser's multi-GiB
workspace resident). The spec-decode draft engine lazily re-requests
its scratch after the target's release (collective, MERGE-lockstep
safe); multi-rank tuning without a scratch fails loud instead of
peer-mapping non-symmetric tensors. Eager fences over a shared side
already baked into a captured graph fail loud (captured pairings
cannot re-fence).
- Vendored kernel: SF-scale and fp4-wire DATA store pointers now claim
their true alignment (e8m0 scales are 1 B at row base +1/+2/+3; e2m1
STG vectors are 16 B at row base +16) instead of a blanket 4/32 --
the fp8/fp4 combine formats executed under an undefined-behavior
alignment promise.
- EPLB + hot reload: the backend pre_reload_weights override now
carries the interface's fail-fast (the module walk bypasses the
interface guard; latent on main since PR NVIDIA#14608).
- Tests: suffix-ordered multi-module reload (CPU staging + bounded GPU
peak + bit-exact repack), input-scale-only and no-ws2 rejection
cases; aux-tail test ships both families.
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.