feat(moe_ep): restructure MegaMoE kernel sources into kernel_src + latest CuTeDSL MegaMoE kernels with improved performance - #3980
Merged
Conversation
… into canonical kernel_src layout
Moves CuTeDSL kernel sources from the ad-hoc
`backends/mega/kernel/cutedsl_backend_kernels/` tree into the new
`flashinfer/moe_ep/kernel_src/cutedsl_megamoe/` hierarchy so all kernel
drops from the kernel team land in one predictable location. Backend
shims (mxfp8_cutedsl, nvfp4_cutedsl) and bootstrap updated to import
from the new path.
Also drops dead code from the vendored frontend: the unused
`megamoe_frontend/api.py` re-export shim and the deprecated
`MegaMoE{Config,Frontend,Inputs}` aliases (all live callers use the
dtype-specific `MegaMoENvfp4*` / `MegaMoEMxfp8*` names).
Mega tests (nvfp4 + mxfp8 multirank, mxfp8 preprocess vs reference)
verified passing end-to-end on Blackwell via SLURM.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…nit__
Collapses the confusing two-layer split (parent nvfp4.py/mxfp8.py wrappers
over megamoe_frontend/api_nvfp4.py/api_mxfp8.py frontends) into a cleaner,
more compact layout that future kernel drops can follow:
__init__.py public API for moe_ep; talks only to shim/, re-exports the
curated symm-buffer + fused-launch surface (the sole boundary)
shim/comm.py dist bootstrap, sym heap, compile state, resolve_gate_up_clamp
(was megamoe_frontend/common.py + _util.py)
shim/nvfp4.py NVFP4 frontend + symm-buffer/launch wrappers, self-contained
shim/mxfp8.py MXFP8 frontend + symm-buffer/launch wrappers, self-contained
shim/correctness.py standalone NVFP4 smoke runner
Each dtype is now one self-contained module (frontend + wrapper merged), so
there is no more api_mxfp8.py-vs-mxfp8.py duplication. Path bootstrap moved
into shim/__init__ so the shim is import-safe. The unused ``python -m`` entry
wrappers (__main__.py) were dropped; the smoke runners remain invokable via
``python -m ...cutedsl_megamoe.shim.{nvfp4,mxfp8,correctness}``. SKILL.md
updated for the new layout.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…zone
Prepares the cutedsl_megamoe kernel_src for repeatable kernel-team drops by
making shim/ the single adaptation layer, so a new drop is a pure replace of
src/ with the only work confined to shim/.
Shim-boundary isolation (nothing outside shim/ imports src/ packages):
- New shim/kernel_helpers.py: the single re-export point for raw-kernel
helpers/constants/reference the FI backend + tests need. Light helpers are
eager; the cutlass-pulling mega_runner/mega_reference helpers are exposed
lazily (module __getattr__ here + package-level __getattr__) so importing
the package stays CPU-safe.
- backends/mega/kernel/{nvfp4,mxfp8}_cutedsl/{backend,staging,weights}.py now
import every helper/constant through the package public API instead of
reaching into common/moe_nvfp4_swapab/moe_mxfp8_glu directly; dropped the
now-redundant per-file _require_cutedsl_paths() shims.
- tests/moe_ep cutedsl tests verify only through the package boundary
(the MXFP8 torch reference is re-exported by kernel_helpers.py); added
clarifying comments to all three.
Verbatim src/ + bootstrap relocation:
- Moved path bootstrap out of src/ into shim/_paths.py (bootstrap_paths now
computes the sibling src/ dir from the shim location) and re-exported it
from shim/ and the package __init__. src/ no longer contains any injected
files (deleted src/__init__.py and src/_bootstrap_paths.py) and is now an
exact copy of the kernel-team drop's four packages.
- core/runtime/bootstrap.py reaches bootstrap_paths through the package
boundary rather than importing src._bootstrap_paths.
Docs/SKILL.md updated for the new shim layout, the verbatim-src principle, and
the kernel_helpers src-symbol audit table used when resyncing a drop.
Verified on Blackwell (4-GPU srun): mega multirank (nvfp4 + mxfp8, staged and
prestaged) and mxfp8 preprocess-vs-reference all pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…kernel API
Replaces the vendored src/ with the latest kernel-team drop (verbatim: the four
packages common/, moe_nvfp4_swapab/, moe_mxfp8_glu/, src/ copied as-is, no
injected files) and adapts the shim to the kernel's substantially changed
construct + launch + output model. All adaptation is confined to shim/; src/
is an exact copy of the drop.
Kernel API migration (mirrors the training driver
moe_ep_training/megamoe/forward_nvfp4.py / forward.py):
- Constructor: NVFP4 dropped local_rank (derived from the peer mapper now) and
takes a CombineFormat object instead of a combine_dtype string; MXFP8 keeps
local_rank but renames token_back_by_dispatch(bool) -> token_back_mode(str).
- Launch/output: the 4 combine tensors (form-A combine_output + reduced +
quantized combine staging) collapse to a single 2D (T, hidden)
output_activation -- the kernel reduces the top-k combine internally. NVFP4
passes local/shared_workspace as raw uint8 gmem pointers (make_ptr); MXFP8
keeps cute-tensor workspaces and adds fc1_c=None. Launch the full padded
buffer and slice [:n] (matches the reference driver).
- _stack_byte_reinterpretable_tensors moved to runner_common (light) -> made
an eager kernel_helpers re-export. Removed the obsolete standalone
topk-reduce path (the free fns became the TopkReduce class; unused by
moe_ep) and the dead _CompiledTopk scaffolding.
tuner.py (new): exposes the kernel tuning knobs -- correctness knobs
(in_kernel_fc2_reduce, token_back_mode, non_ubulk_fc2_store, load_balance_mode,
mma_tiler_mnk, cluster_shape_mnk) and perf knobs (group_hint, flag_batch,
epi_flag_batch), mirroring tester/solvers/inference_solver. Provides is_valid()
(mirrors filter_invalid), iter_candidates(), and a dtype-aware with_knobs();
get_symm_buffer_for_*(knobs=...) applies them before compile.
correctness.py smoke runner stays stubbed (NotImplementedError) pending a port
to the drop's MegaMoETester. Preprocess test updated: the form-A per-topk output
is gone, and the reference now uses apply_topk_in_fc1=True + plain-sum (the
kernel folds the topk weight pre-quant, so post-hoc weighting no longer matches).
SKILL.md + architecture doc updated for the new layout and the construct/launch
drop-audit surface.
Verified on Blackwell (4-GPU srun): mega multirank (nvfp4 + mxfp8, staged and
prestaged), deep_gemm control, and mxfp8 preprocess-vs-reference all pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…el call
Wires a compile-time token-count heuristic into get_symm_buffer_for_{mega,mxfp8}
_moe as the default tactic (an explicit knobs= dict still overrides it). Keyed on
the buffer's num_max_tokens (the kernel compiles once for that size):
num_tokens < 2048 -> mma_tiler (256,128,256), flag_batch 4, token_back epi_warps
num_tokens >= 2048 -> mma_tiler (256,256,256), flag_batch 8, token_back
reuse_dispatch_warps
both with cluster (2,1,1), group_hint 512, epi_flag_batch (2,4), load_balance
atomic_counter. in_kernel_fc2_reduce and combine_format stay owned by the config
/ caller (not set by the heuristic).
MXFP8's mma_tiler is kernel-fixed at (256,256) -- the NVFP4-tuned (256,128,256)
is rejected by the kernel ctor -- so tuner.default_knobs(..., include_tile=False)
drops the tile for MXFP8, which keeps its validated (256,256,128) and takes only
the non-tile perf knobs (schedule / token-back / load balance).
Verified on Blackwell (max_tokens=64 -> small profile): mega multirank (nvfp4 +
mxfp8, staged + prestaged), deep_gemm control, and mxfp8 preprocess-vs-reference
all pass. NOTE: the large profile (>=2048 tokens) and MXFP8 large-token
reuse_dispatch_warps are not covered by the 64-token suite -- verify against a
large-token workload before relying on them.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…rofile The 64-token mega suite only exercised the tuner's SMALL profile. Add num_tokens=max_tokens=2048 variants of the nvfp4 and mxfp8 mega-layer tests so the LARGE profile is actually run: - nvfp4: throughput tile mma_tiler (256,256,256) + token_back reuse_dispatch_warps - mxfp8: kernel-fixed tile (256,256) + flag_batch 8 + token_back_by_dispatch=True The mxfp8 case is the regression guard for large-token dispatch-warp token-back (mxfp8 has no non_ubulk_fc2_store knob); it compiles + runs + stays bit-exact with the direct-kernel reference, so the large profile's reuse_dispatch_warps is confirmed valid for MXFP8 -- no need to pin it back to epi_warps. Parametrized _mega_problem / _run_mega_layer with num_tokens/max_tokens (defaults unchanged at 64, so the existing small-profile tests are untouched). Verified on Blackwell (4-GPU srun): all mega sections pass, including the two new large-token tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…overhead fixes
Closes the measured gap between the FI moe_ep cutedsl backends and the
kernel-repo tester, in three parts:
Knob plumbing + online autotuning:
- Nvfp4/Mxfp8CutedslMegaMoeConfig gain `knobs`: an explicit tuner knob dict
(e.g. a kernel-repo tester sweep winner) overrides the token-count
heuristic; "auto" runs a COLLECTIVE online autotune at the first forward
(new shim/autotune.py): every rank compiles+times a curated candidate set
in lockstep (barriers around compile and timing), per-candidate medians
are all-reduced MAX (slowest rank = collective latency), and the argmin
winner is applied identically on every rank. Candidates mirror the tester
sweep restriction minus in_kernel_fc2_reduce (changes output placement /
determinism; stays config-owned). Frontends gain apply_knobs().
- tuner.default_knobs(dtype="mxfp8"): MXFP8 gets its own measured profile
(flag_batch=4 + epi_warps at all sizes; the NVFP4-large fb8+dispatch-warp
schedule measured ~5% slower for MXFP8 at 2048 tokens on 4x GB200). The
mxfp8 large-token test pins the old dispatch-warp combo via knobs= so
that regression coverage survives the heuristic change.
Launch-path overhead (was misread as kernel slowness):
- frontend.run() rebuilt all 12 cute tensor views (from_dlpack) plus a
SymBufferHost on EVERY launch; now cached on _CompiledMega keyed by input
data_ptrs + shape + stream (dropped automatically on recompile).
- make_launch_thunk() + {nvfp4,mxfp8}_mega_launch_thunk(): prebuilt bare
launch closures matching the kernel tester's perf_run timed region, for
benchmark loops and tuners.
- run() defaults reset_counters=False: workspaces are allocated zeroed and
the kernel tail-cleans its counters/flags after every launch (the
kernel-team drivers/tester never host-reset). The multirank layer tests
now run a REPEATED forward on the same session and assert both outputs
bit-exact against the direct-kernel reference (tail-cleanup contract).
Measured (4x GB200, 7168/2048/256/top-8, p50 us, kernel-parity | e2e):
@8 tok dg 215|243 nvfp4 132|269 (e2e was 465) mxfp8 381|519 (was 691)
@2048 tok dg 844|872 nvfp4 527|671 (e2e was 1121) mxfp8 1214|1325 (was 1658)
nvfp4_cutedsl is ~1.6x FASTER than deep_gemm at the kernel level at both
batch sizes and beats it e2e at 2048 tokens; the earlier 2x-slower reading
was launch-path overhead plus a geometry mismatch vs the tester problems
(nvfp4_perf.jsonl is top-6 with different inter/experts).
Verified on Blackwell (4-GPU srun): cutedsl multirank suite (nvfp4 + mxfp8 +
preprocess-vs-reference, incl. the new repeated-forward guards) 14 passed;
autotune (12 nvfp4 + 4 mxfp8 candidates x 2 token sizes) compiles and
selects winners collectively with zero failed candidates.
AI-assisted (Claude Code).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… wrappers Trims the remaining per-forward host overhead in the cutedsl e2e path: - frontend.run() fast path: when the launch cache hits (same session buffers, same resolved token count, same stream), the call goes from a cheap pointer-tuple key compare straight to the kernel launch -- input validation, _prepare_launch_inputs, and the compile-key rebuild run only on a cache miss. Safe because any config change (apply_knobs / set_gate_up_clamp) nulls self._mega, so a live entry always matches the compiled config, and the cached entry was fully validated when built. The key uses the RAW input pointers + resolved token count (slices start at row 0, so sliced views keep the same data_ptrs). - nvfp4_mega_moe / mxfp8_mega_moe gain sync (default False): the kernel and the y output copy are enqueued on the current stream and the call returns without a host sync, matching deep_gemm and normal stream semantics. sync=True blocks until y is host-readable. The online autotuner passes sync=True explicitly (it times launches with perf_counter); the multirank tests already synchronize after the wrapper (audited). Measured (4x GB200, 7168/2048/256/top-8, p50 us, e2e barrier-cold): @8 tok nvfp4 269 -> 221 (kernel floor 132; dg e2e 244) mxfp8 519 -> 473 @2048 tok nvfp4 671 -> 629 (kernel floor 534; dg e2e 866) mxfp8 1325 -> 1289 nvfp4_cutedsl now beats deep_gemm end-to-end at BOTH batch sizes; its residual e2e-vs-kernel delta (~90 us) is the output copy + the benchmark's barrier-cold per-iter launch, not per-call Python. Verified on Blackwell (4-GPU srun): cutedsl multirank suite (incl. the repeated-forward bit-exact guards) 14 passed on every rank. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…profiles The full token sweep (1..8192, e2e_pipelined) showed nvfp4 beating deep_gemm 1.3-1.7x everywhere EXCEPT a dip at 512 (1.20x) and 1024 (1.00x) -- right below the SMALL->LARGE threshold at 2048. Online autotune at those points found the cause is token_back_mode alone: every dispatch-warp candidate beat every epi_warps candidate (epi_warps +18% at 512, +35% at 1024); tile and flag_batch are second-order. default_knobs now has four NVFP4 profiles (SMALL and LARGE unchanged): <512 SMALL (256,128,256) fb4 epi_warps [winner @8] 512..1023 MID (256,128,256) fb4 reuse_dispatch_warps [winner @512] 1024..2047 MID-LARGE (256,256,256) fb4 standalone_warps [winner @1024] >=2048 LARGE (256,256,256) fb8 reuse_dispatch_warps [validated @2048] Measured with the new defaults (4x GB200, 7168/2048/256/top-8, p50 us, e2e_pipelined): 512: 287.7 -> 234.5 (1.48x vs dg); 1024: 474.1 -> 332.3 (1.43x vs dg); 2048 unchanged (537.6). The nvfp4 advantage over deep_gemm is now uniform 1.3-1.7x across the whole curve. New parametrized multirank test covers both mid profiles (512 + 1024) bit-exact against the direct-kernel reference; full cutedsl suite passed 16/16 on every rank. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…h methodology Companion doc to SKILL.md capturing what the 2026-07-14 tuning round changed and learned: the knob system (per-size profiles + provenance, online collective autotune, kernel-repo tester handoff), the launch-path work that brought the FI forward within 2-13us of the bare kernel (kwargs cache, no-reset tail-cleanup contract, async wrappers, launch thunks), and the benchmarking rules that prevent the 2x-slower misread from recurring (match geometry AND timed region; MEGA_TIMING kernel/e2e_pipelined/e2e semantics; from-idle collective start skew). Headline: nvfp4_cutedsl 1.3-1.7x faster than deep_gemm at every token count 1..8192, uniform after the mid-range profile fix. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…euse nvshmem4py's get_peer_tensor resolves the Buffer from a tracker keyed by raw address and re-views it with the tensor's exact shape; when the symmetric heap reuses an address for a SMALLER allocation (first hit when in_kernel_fc2_reduce shrinks shared_workspace from ~1 MiB to ~8 KiB), a stale larger peer entry survives and the view throws. _compute_peer_offsets only needs peer BASE ADDRESSES (the nvshmem_ptr mapping is deterministic), so read them off the peer Buffer without the shaped view. Also skip the self-peer lookup entirely: nvshmem4py resolves it to the parent tracker entry and bumps its refcount, deferring the real nvshmem free from free_tensor() to GC (the 'memory was not freed explicitly' finalize warnings). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mport)
Import both idea families from the TRT-LLM MegaMoE integration
(PR #16190; plan in moe_ep/todo_trtllm_import.md):
- in_kernel_fc2_reduce: in-flight top-k combine via cross-rank REDG
atomic-add. get_symm_buffer_for_mega_moe exposes the param;
output_activation is now ALWAYS sym-heap allocated (ikr hard-requires
it, explicit-reduce is indifferent, and the knob can then flip
per-compile without reallocation); frontend.run() and the launch thunks
enqueue the accumulate-from-zero output zero_() before every launch.
Fixes the MXFP8 twin's latent bug: it accepted the param but allocated
the REDG target rank-locally.
- combine_dtype: quantized cross-rank combine wire ('mxfp8' =
32e4m3xe8m0, 2x less combine traffic; 'nvfp4' = 16e2m1xbf16, 4x less)
on the NVFP4 backend. Default knobs auto-adjust to dispatch-warp
token-back (the only wired mode for quantized wires); tuner.is_valid
gains the matching rule. Both fields surface on
Nvfp4CutedslMegaMoeConfig.
Measured (4x GB200, 7168/2048/256/top-8): combine_dtype='nvfp4' is the
large-token winner (-16% @2048, -19% @8192 vs bf16 wire; 2.08x vs
deep_gemm); ikr's win is the multi-GB combine-staging removal, not
latency at this geometry. Details in cutedsl_megamoe/TUNING.md.
AI-assisted (Claude Code).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
nvfp4_candidates gains the in_kernel_fc2_reduce axis (24 candidates for the bf16 combine wire) plus combine_format / allow_in_kernel_fc2_reduce pruning; autotune_nvfp4_mega_moe derives the default list from the live session config (combine wire, apply_topk_in_fc1). The always-sym output (previous commit) is what makes the ikr knob flippable per-compile. mxfp8_candidates prunes ikr-incompatible dispatch-warp token-back combos instead of letting them fail at compile. An ikr winner makes the output accumulation order nondeterministic; callers needing bit-reproducibility pin in_kernel_fc2_reduce=False via explicit knobs. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
New 4-rank variants (both suites pass 15/6 on 4x GB200): - in_kernel_fc2_reduce (nvfp4 + mxfp8): tolerance verdict vs the explicit-reduce reference via _assert_ikr_close -- a row-scaled bf16 K-term accumulation band (K x 2^-8 x safety 8 x row max). A flat atol/rtol misfires here: the REDG path accumulates the K per-topk terms in bf16, so where large terms nearly cancel the achievable agreement is bounded by the ULP of the largest TERM, not the final value (the kernel repo validates ikr with a K!-ordering bitwise check for the same reason; measured band need <= 0.67x, a missing per-launch output zero overshoots ~64x so the repeated-forward regression guard stays sharp). - combine_dtype nvfp4/mxfp8: bit-exact vs a same-wire direct-shim reference (deterministic path) + rel-L2 sanity band vs the exact bf16 wire. - CPU-safe config/tuner/candidate validation tests. Also fixes the mxfp8 large-token test: pinned knobs= now reaches the layer config and the reference (it previously only reached the throwaway runtime-requirements kernel, so the pinned profile was never exercised). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… flow - TRT-LLM-import knobs section: what ikr / combine_dtype mean (combine leg only; output always bf16; exact-vs-lossy terms x fp32-vs-bf16 sum), the measured 4-variant sweep table with speedups vs deep_gemm (combine_nvfp4 up to 2.08x @8192), methodology + exact environment versions, and a reproduce runbook pointing at https://github.com/mhoqueanik/moe_ep_benchmark (MEGA_IKR / MEGA_COMBINE_DTYPE land in bench commit a15fb01). - Runtime knob-resolution flow + mermaid sequence diagrams (supersedes the unmerged docs/cutedsl-tuning-flow commit 32175498, updated for the ikr/combine changes: knobs bind once per session keyed on num_max_tokens, single-slot compile cache, no persistence; setup + steady-state forward and the knobs='auto' collective sweep). - Stale-info cleanup: headline speedup range, ikr expectation replaced with the measured verdict, e2e_pipelined staging-exclusion caveat. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Credit the NVIDIA CuTeDSL MegaMoE kernel team for the vendored src/ kernels, validation harness, and tuning groundwork the moe_ep integration builds on. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… intro Mermaid sequence diagrams broke on GitHub: a ';' inside a Note is a mermaid statement separator (parse error at 'output_activation always sym-heap'), and a bare '<' in note text can open a stray HTML tag. Replace both; also drop quotes from an else-label for parser safety. Rewrite the abrupt opening paragraph into a proper overview (what the doc covers, shared measurement setup, pointer to the reproduce recipe). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
mhoqueanik
requested review from
Anerudhan,
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
jiahanc,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww and
yzh119
as code owners
July 15, 2026 06:14
…EMENT.md Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ght layout Full per-bucket retune (online collective autotuner, 4x GB200, default geometry, kernel-mode p50 vs defaults): - NVFP4: all four non-ikr profiles CONFIRMED optimal (64/512/1024 tokens within <=1.2% run noise; @64 the tuner reproduced the SMALL profile exactly). The only wins were ikr candidates at >=2048 (-4.8% @2048, -4.2% @8192); ikr stays opt-in/autotune because its nondeterministic accumulation order would break the bit-exact multirank tests and bit-reproducibility. No NVFP4 profile change. - MXFP8: ONE real change — a new >=2048 profile with dispatch-warp token-back (fb4 + reuse_dispatch_warps): 1006.6 vs 1209.6 us @2048 (-17%), 1742.8 vs 2208.2 @4096 (-21%), 3122.2 vs 4844.0 @8192 (-35%, reaching deep_gemm parity at a 3x better accuracy point). Supersedes the 07-14 'dispatch-warp is ~5% slower for MXFP8' reading, which conflated the token-back mode with fb8. Deterministic (no ikr in the mxfp8 candidate space); acc_loss_pct unchanged at 6.358%; mega test section re-verified green with the new default. TUNING.md updated: sweep-table mxfp8 rows re-measured with the new default, profile provenance notes refreshed. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ant breakdown One table now carries dg, nvfp4 bf16/ikr/combine-wire variants, and the mxfp8 backend (new dispatch-warp default at >=2048) per token count, replacing the compact nvfp4 column and the duplicate variants table in 'Measured results'. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…d summary First table: nvfp4 combine-leg variants vs dg (unchanged columns). Second table: backend summary — dg vs best nvfp4 config (+combine_nvfp4) vs the mxfp8 backend on its new dispatch-warp default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Anerudhan
enabled auto-merge (squash)
July 15, 2026 21:35
Anerudhan
approved these changes
Jul 15, 2026
Anerudhan
disabled auto-merge
July 15, 2026 21:37
Formatting only (ruff format via pre-commit); no behavior change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Collaborator
|
[FAILED] Pipeline #58079185: 13/20 passed |
Anerudhan
enabled auto-merge (squash)
July 16, 2026 06:19
bkryu
approved these changes
Jul 16, 2026
This was referenced Jul 24, 2026
Anerudhan
pushed a commit
that referenced
this pull request
Jul 28, 2026
…, fused quant+stage launch, persistent knob cache, and prequantized weight packs (#4079) <!-- .github/pull_request_template.md --> ## 📌 Description ### Summary Follow-up to #3980. This PR makes the MegaMoE (`moe_ep`) path consumable by serving engines (vLLM-style integration): full CUDA-graph capture/replay support, a fused single-launch quant+stage hot path, engine-friendly weight lifecycle (prequantized packs, post-preprocess source release), pooled symmetric-buffer workspaces shared across layers, and a persistent knob cache so production sessions resolve tuned knobs with a pure lookup. No in-engine autotuning. ### Performance Highlights All numbers measured on a **single node, 4x GB200 (EP=4)**. The built-in knob heuristic currently supports **GB200 only** (offline-tuned knob caches expected to work on any supported device). **Microbenchmark** — DeepSeek-V3 geometry (256 experts, top-8, hidden 7168, intermediate 2048), `e2e_pipelined` p50 µs, heuristic knobs, best variant (`nvfp4 + combine_nvfp4`) vs `deep_gemm_mega` baseline: | tok/rank | deep_gemm_mega | nvfp4 + combine_nvfp4 | Speedup | |---|---:|---:|---:| | 512 | 340.0 | 321.1 | 1.06x | | 1024 | 468.0 | 363.5 | 1.29x | | 2048 | 817.2 | 529.4 | 1.54x | | 4096 | 1473.5 | 862.7 | 1.71x | | 8192 | 2993.7 | 1677.3 | 1.78x | The sweep also covers real-model MoE geometries (DeepSeek V3 / V4-Flash / V4-Pro, Kimi K2.6, Qwen3.5-397B, gpt-oss-120b — the last enabled by the %64 alignment relaxation); the pattern holds everywhere: deep_gemm-parity at small token counts, 1.6–1.9x for fp4 combine-wire at large tokens on 7168-hidden shapes. Full sweep tables (all variants), accuracy, and benchmark methodology: [`kernel_src/cutedsl_megamoe/TUNING.md`](https://github.com/mhoqueanik/flashinfer-moe_ep/blob/fi-moe_ep-framework-integration/flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md#real-model-geometry-sweep-2026-07-21). **End-to-end (vLLM)** — early e2e benchmarks with this backend integrated into vLLM 0.25.1 (DeepSeek-V4-Flash, 4x GB200 TP4/EP4, CUDA graphs capturing all recurring step shapes including prefill chunks, per-role offline knob caches) show **~18% prefill and ~7% decode throughput gain over native vLLM** (i.e.: vLLMs built-in deep_gemm mega path): | Workload | native vLLM | fi MegaMoE (nvfp4) | Speedup | |---|---:|---:|---:| | Prefill, 8k-token chunks (tok/s) | 45,701 | 53,962 | 1.18x | | Decode, 1024-seq concurrency (output tok/s) | 21,049 | 22,614 | 1.07x | | GSM8K (200q, greedy) | 0.965 | 0.975 | — | Numbers reproduce within 1% at this branch tip; details and caveats (the decode number requires graph capture to cover prefill chunk shapes) are in TUNING.md's e2e section. The vLLM integration itself will land as a separate PR. **Accuracy** — per-variant microbench rel-L2 (`acc_loss_pct` vs fp32 dense-MoE reference) is unchanged from #3980; see the accuracy table in TUNING.md. On real-model distributions the nvfp4 path is at parity-or-better (GSM8K above), consuming the single-quant nvfp4 checkpoint directly via the prequantized weight-pack path. ### Key Changes **CUDA Graph Support (cutedsl mega path)** - Warmup contract + capture guards for graph-safe forward passes - Tail-mask memo made graph-capture-aware; records the actual staged count (fixes a view-slicing regression on capture-touched buffers) - 2-rank lockstep CUDA-graph replay test for the nvfp4 mega layer **Launch-Path Optimizations** - Fused single-launch quant+stage for the cutedsl mega path; fused staging for `deep_gemm` mega with a 16B-alignment fallback (bit-identical to the multi-kernel path, toggle: `FLASHINFER_MEGA_FUSED_STAGE`) - Zero-copy output + memoized tail mask + caller-owned ikr buffers to cut per-call launch count - Cached stream-aware launch thunk in the nvfp4 backend hot path **Engine Integration: Weights & Workspaces** - `MoEWeightPack` split into `Unquantized | Prequantized` variants - Source weight pack released after preprocess (fixes OOM at model load) - Pooled symmetric-buffer workspaces shared across mega layers **Tuning Infrastructure** - Persistent knob cache (`FLASHINFER_MOE_EP_KNOB_CACHE`): offline-tuned winners resolved by pure lookup before the built-in heuristic; populated via the new offline tuner CLI `python -m flashinfer.moe_ep.tune`. Per-role cache files (prefill-tuned / decode-tuned) are the validated deployment pattern. - Skew-aware tuning (`--skew`), schedule-axis sweep (`--sweep schedule`), `--live-tokens` **Shape Coverage & Guardrails** - cutedsl mega alignment relaxed to %64 (dg keeps %128) — enables gpt-oss-120b-class geometries - Warn-once import check for the CuTe-DSL 4.6.1 performance floor **Determinism Investigation (docs)** - Reported `fi_dg` cross-run nondeterminism traced to vLLM batch-formation timing (schedule diff, 1-tok vs 8-tok batch) — engine-level and backend-independent; closed as a FlashInfer issue with probe tests retained ### Directories Affected - `flashinfer/moe_ep/backends/mega/` — CUDA graph support, fused staging, launch caching - `flashinfer/moe_ep/kernel_src/cutedsl_megamoe/` — shim knob cache, TUNING.md - `flashinfer/moe_ep/core/`, `flashinfer/moe_ep/modes/` — workspace pooling, weight-pack lifecycle - `flashinfer/moe_ep/tune.py` — offline tuner CLI - `tests/moe_ep/` — CUDA graph (single- and multi-rank), fused stage, knob cache, weight-pack, workspace-pool, determinism-probe coverage ### Next Steps (WIP) - Generalized tuning: extend the knob heuristic / offline-tuned profiles beyond GB200 - Performance support for CuTe-DSL 4.5.2 (currently 34–54% slower than 4.6.1, which is treated as the perf floor with a warn-once import check) ## 🔍 Related Issues Follow-up to #3980 (MegaMoE kernel restructure). ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). New tests: `test_mega_cuda_graph.py`, `test_mega_cuda_graph_multirank.py`, `test_fused_quant_stage.py`, `test_knob_cache.py`, `test_weight_pack_union.py`, `test_workspace_pool.py`, `test_moe_ep_deep_gemm_skew_determinism.py`. Multi-rank tests run under `tests/moe_ep/run_tests.sh`. ## Reviewer Notes - The fused quant+stage path is bit-identical to the original multi-kernel path (`tests/moe_ep/test_fused_quant_stage.py`); `FLASHINFER_MEGA_FUSED_STAGE=0` reverts it for bisection. - The knob cache keeps autotuning strictly offline: `knobs=None` sessions do a pure cache lookup, falling back to the heuristic on miss — no measurement runs inside an engine process. <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added fused input staging for supported MoE quantization workflows, with an environment-variable toggle. * Added persistent autotuning cache support and an offline tuning command. * Added workspace sharing to reduce memory usage across compatible layers. * Added optional zero-copy output handling and CUDA Graph warmup support. * Added explicit pre-quantized and unquantized weight types. * **Bug Fixes** * Improved CUDA Graph capture safety, token tracking, alignment support, and workspace cleanup. * **Documentation** * Updated configuration and tuning guidance, including cache settings and performance information. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
5 tasks
mhoqueanik
added a commit
that referenced
this pull request
Jul 30, 2026
…le test (#4221) test_deep_gemm_mega_kernel_matches_torch_reference (added in #3980) fails under plain pytest with "ValueError: ... environment variable RANK expected, but not set": its WORLD_SIZE guard only skips when the variable is set and != 1, so without torchrun it falls through to dist.init_process_group(backend="nccl"), whose default env:// rendezvous requires torchrun's RANK/MASTER_* variables. This is only reachable on CI jobs that combine plain-pytest discovery of tests/moe_ep, a cu13 image with the EP stack (deep_gemm) installed, and a capability-10 GPU — i.e. the B300 cu130 unit-test job, where it currently errors on every run. Fix: when RANK is absent, self-bootstrap a 1-rank NCCL group via an explicit tcp://127.0.0.1:<free-port> init (rank=0, world_size=1) instead of env://. Deliberately avoids os.environ mutation so RANK/MASTER_* don't leak to later tests in the same pytest process; teardown is unchanged (conftest.pytest_sessionfinish destroys the group). The torchrun path is preserved and run_tests.sh still exercises it. Verified on GB200 (single GPU): both launch modes pass with rel_l2=0.0027 vs the torch oracle; the plain-pytest mode was confirmed with RANK/WORLD_SIZE/MASTER_ADDR/MASTER_PORT/LOCAL_RANK explicitly unset. AI-assisted (Claude Code). <!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This was referenced Aug 7, 2026
mhoqueanik
added a commit
to mhoqueanik/flashinfer-moe_ep
that referenced
this pull request
Aug 13, 2026
…hinfer-ai#4449 Five of the six findings are actionable in our tree: - test_sm90_push_fp8_packaging.py: the import-boundary gate built backend_root with the pre-taxonomy flat name kernel/sm90_fp8_fp8_bf16_push_cuda, which no longer exists — Path.rglob on a missing dir yields nothing, so the gate passed vacuously. Fixed to kernel/sm90/fp8_fp8_bf16_push_cuda and added is_dir/non-empty asserts; the gate now validates 5 files / 5 marker imports (verified standalone). - _sm90_push_fp8_baseline.py: _weight_cache entries now evict via weakref callbacks when either source tensor dies, instead of pinning four dead GPU tensors per (w13, w2) pair. In-place weight mutation versioning is intentionally skipped — the oracle tests never mutate expert weights after quantization. - cutedsl_megamoe/shim/__main__.py: added, so the documented `python -m ...cutedsl_megamoe.shim` invocations in correctness.py actually resolve (delegates to correctness.main()). - cutedsl_megamoe/VENDOR.md: filled the provenance TODOs with what is honestly known — repo URL (internal GitLab, already referenced in the pending-diffs section), drop date 2026-07-13 via PR flashinfer-ai#3980, and an explicit note that the base SHA was never recorded and MUST be pinned at the next full re-sync. - test_sm90_push_fp8_orchestrator.py: RUF043 — raw-string the pytest.raises match pattern. The sixth (missing CHECK_DIM(2, out) in the vendored push_style_megamoe/src/a2a/sm90_push_a2a_ops.cu) is real but sits in a verbatim-vendored src/ tree: per kernel_src/README.md it routes to upstream (now flashinfer main's own copy of the file) rather than being patched here. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.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
4 participants
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.
Summary
Follow-up to #3852: brings the CuTeDSL MegaMoE kernels to their tuned state. Restructures kernel sources into a canonical
kernel_src/cutedsl_megamoelayout (verbatimsrc/drop zone behind an enforcedshim/boundary), drops the updated kernel, and adds new kernel knobs (in_kernel_fc2_reduce,combine_dtype), per-size default profiles, and session-aware collective online autotuning. SeeTUNING.mdfor detailed changes.Perf
Measured 2026-07-15 on one GB200 node (4x GPUs, EP=4), DeepSeek-V3-like geometry (256 experts, top-8, hidden 7168, intermediate 2048), number below checked in-line against an fp32 dense-MoE reference (see Accuracy). Steady-state full-forward p50 _s (
e2e_pipelined; bare-kernel timing tracks within a few _s), speedup vsdeep_gemm_megain parens:[ikr = in_kernel_fc2_reduce; +combine_nvfp4 / +combine_mxfp8 = the fc2 epilogue quantizes each partial IN REGISTERS to fp4 (e2m1 + bf16 scale per 16) or fp8 (e4m3 + e8m0 scale per 32) just for the wire]
combine_dtype="nvfp4"is the throughput winner across this range (1.89x vs dg and 19.9 Mtok/s at 8192); bigger wins expected multi-node where combine bytes dominate. Below ~1024 tok/rank nvfp4 and dg are at parity (e.g. 217.9 vs 211.0 _s at 8 tok/rank _ small batch is weight-load bound and fp4-vs-fp4 is a wash), with the plain bf16 wire the best small-batchconfig.
[Known issue] CuTe-DSL runtime sensitivity:
nvidia-cutlass-dsl[cu13]==4.5.2(2026-07-15)[Fix underway] Same recipe, geometry and runbook as the table above, with the only change being the CuTe-DSL runtime pinned to 4.5.2 instead of >= 4.6.1 (CSVs
moe_ep_benchmark/results/sweep_20260715_1634*..1641*_fi_mega.csv):Takeaways:
cute.compile), and every variant/point completed without error.+combine_nvfp4~+50% everywhere). The dg baseline reproduces within noise (468.8-3060.7 vs 473.1-3105.2), so the delta is the DSL runtime, not the run.Conclusion: treat 4.6.1 as a performance floor, not just a compile-compatibility floor at 4.5.2 is functional but leaves ~1.3-1.5x of nvfp4 kernel performance on the table.
Accuracy
Reported per run by the benchmark's new
acc_loss_pctpass: one un-timed forward through the exact timed configuration vs an fp32 dense-MoE ground truth over all 256 experts (global rel-L2, random data _ an upper bound; calibrated real-model weights typically lose less). Constant across token counts:+ikr+combine_mxfp8+combine_nvfp4The fp4 combine wire buys its large-token speedup for ~1.8 pt of extra quantization loss. Every compute path is additionally pinned to a single-GPU pure-torch oracle in
tests/moe_ep/(run_tests.sh oracle); the nvfp4 mega kernel is bit-exact vs the kernel-team CuTeDSL reference on identical operands.What's included
cutedsl_backend_kernelsrelocated tokernel_src/cutedsl_megamoe—src/is a verbatim kernel drop zone, backends import only through theshim/package boundary (enforced by test).in_kernel_fc2_reduceand quantized combine wire (combine_dtype= bf16 / nvfp4 / mxfp8).knobs="auto"session-aware collective online autotune (candidates include the ikr axis).TUNING.md(tuning surface, measured sweeps, benchmark methodology + pitfalls),ACKNOWLEDGEMENT.mdfor the kernel authors.Directories affected
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/— kernel sources (src/), shim boundary (shim/), docs (new home)flashinfer/moe_ep/backends/mega/— nvfp4/mxfp8 cutedsl backends: knob plumbing, autotune, launch path (old kernel dir removed)flashinfer/moe_ep/core/runtime/— bootstrap import path updatetests/moe_ep/— multirank + tuner coveragedocs/design_docs/— moe_ep architecture + runbook updatesRunbook
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md("Sweep methodology" + "Runbook" sections)docs/design_docs/moe_ep_runbook.mdSummary by CodeRabbit
New Features
"auto"autotuning on first compute, with runtime tuning-knob overrides.Documentation
Breaking Changes
Tests