perf(moe_ep): CuTe-DSL 4.5.2 mainloop WAR — drop the 4.6.1 runtime floor - #4101
Merged
Merged
Conversation
… (4.6.1 is a perf floor) Same sweep recipe/geometry as the headline table, DSL runtime pinned to 4.5.2: compiles and runs, but generated code is 34-54% slower than 4.6.1 across every nvfp4 variant and token count (dg baseline reproduces within 1%, isolating the delta to the DSL runtime). Treat 4.6.1 as a performance floor, not just a compile-compatibility floor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… notes) Four actionable follow-ups from the 2026-07-15 vLLM 0.25.1 e2e integration (DeepSeek-V4-Flash, 4x GB200): fi_dg run-to-run nondeterminism, CUDA-graph capture (analysis updated with e2e context), MoEEpMegaLayer source-weight retention OOM at model load, and hot-path-free knob selection (offline tuning + geometry/token-bucket heuristic; in-engine knobs=auto is unusable and its tuner-harness winner lost 13% e2e). Plus four earlier analysis notes: MoEWeightPack discriminated-union refactor, multinode support, nixl_ep suppression, trtllm import. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…l load) MoEEpMegaLayer kept self._weights alive for the layer lifetime even though nothing reads it after preprocess_weights() produces the transformed tensors. In the fp4-checkpoint flow the pack is a ~3.2 GB per-layer bf16 dequant copy, so 43 MoE layers pinned ~140 GB of dead weight and OOMed a 186 GB GB200 at model load (vLLM e2e run 7, 2026-07-15). - MoEEpMegaLayer: type _weights Optional, skip storing the pack when backend.transformed_weights is supplied, release it at the end of _preprocess_weights(); memory invariant documented in the class docstring. - MoEEpSplitLayer: same pattern — the pack is only needed for init-time validation and kernel preprocessing (the kernel retains what it needs), so release it at the end of __init__. - Weakref-based release tests for all three paths; factory docstring and todo_weight_pack_retention_oom.md updated (downstream vLLM patch workaround is now a redundant no-op). Resolves flashinfer/moe_ep/todo_weight_pack_retention_oom.md. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ntract + capture guards) Implements the todo_cuda_graph.md restoration plan (single-rank scope). The per-forward sync was already gone (sync=False entry-point defaults); this adds the pieces that make capture safe and testable: - shim/comm.py: ensure_not_capturing() guard; wired after the no-op early returns in _ensure_mega_compiled, set_gate_up_clamp, apply_knobs, _release_workspace (both nvfp4 and mxfp8 frontends) and at the top of the autotune_knobs collective sweep — host-side compile/alloc/free now fails loudly instead of corrupting a capture. - shim run()/mega_moe entry points: sync=True is gated on is_current_stream_capturing (zero-break variant from the TODO). - MoEEpMegaLayer.warmup(): one full eager forward (workspace alloc, cute.compile, knobs="auto" sweep, real launch) + device sync; documents the call-on-all-ranks-before-capture contract. _ensure_workspace raises under capture with a warmup() hint. - tests: tests/moe_ep/test_mega_cuda_graph.py (GB200, MEGA_NO_DIST=1) — nvfp4 + mxfp8 capture post-warmup, 3x replay bit-exact vs eager, replay over in-place-mutated inputs, and capture-without-warmup raises; wired into run_tests.sh oracle section (excluded from unit). Plus mocked capture-guard unit tests in test_mega_layer_validation.py. Verified on GB200: unit 164 passed, graph tests 3/3 passed. Remaining (tracked in todo_cuda_graph.md): 2-rank lockstep replay + skewed staging stress tests, vLLM shared-workspace warmup adoption. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The kernel package ships DataPreprocess (src/src/inputs_process.py) — a fused CuTe-DSL kernel that quantizes bf16 activations (NVFP4 per-16 E4M3 scales / MXFP8 per-32 E8M0) and repacks routing to int64/fp32 in ONE launch — but nothing wired it: both cutedsl mega backends staged through a torch-composed path (fp32 upcast + ~a dozen elementwise/reduce launches + five buffer copies) on every forward of every layer. The vLLM e2e nsys showed these paths host-gap-bound (946k cudaLaunchKernel on fi_nvfp4 vs 100k native), so per-forward launch count is the first-order cost. - shim/quant_stage.py: fused_quant_stage() wrapping DataPreprocess with the frontends' caching pattern — one cute.compile per (topk, hidden, quant_type) with dynamic token extent, launch-args cache keyed on pointers + token count + stream, capture-guarded compile, and the capacity-tail topk_idx=-1 re-mask for torch-path parity. Offline norm_const mode only (online amax mode exists in the kernel, not wired). - nvfp4/mxfp8 backend staging: fused path is the default; FLASHINFER_MEGA_FUSED_STAGE=0 restores the torch path (bisection aid, documented in CLAUDE.md). quantize_input=False (pre-staged) unchanged. - tests/moe_ep/test_fused_quant_stage.py: fused vs torch staging must agree BIT-EXACTLY (the repo's torch quantizers are bit-matched to the kernel) across all three quant kinds, partial/full capacity, and launch-cache hit/rebuild; verified on GB200 (7 passed). CUDA-graph capture/replay (3 passed, staging kernel captured), torch-oracle section (PASS), and unit suite (171 passed) all green with the fused path active. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…hot path)
Implements the lookup half of todo_vllm_knob_heuristic.md: tuning becomes an
offline step whose winners persist, and engine-time knob resolution is a
dict lookup — no compiles, no collectives, no timing in the serving path.
- shim/knob_cache.py: JSON cache (FLASHINFER_MOE_EP_KNOB_CACHE, default
~/.cache/flashinfer/moe_ep_knob_cache.json) keyed on (device, dtype,
world_size, hidden, intermediate, num_experts, topk, combine wire) with
nearest-bucket max_tokens selection; atomic upsert writes; corrupted or
disabled cache degrades to the heuristic with a warning.
- get_symm_buffer_for_{mega,mxfp8_mega}_moe: knobs=None resolves cache-first,
then tuner.default_knobs; explicit dicts unchanged.
- autotune_knobs: on_winner callback; the nvfp4/mxfp8 wrappers record each
winner (rank 0) so any knobs="auto" run persists its result.
- knobs="auto" warns loudly at backend construction pointing at the offline
flow (it remains a collective multi-minute sweep — never in-engine).
- python -m flashinfer.moe_ep.tune: offline CLI (torchrun multi-rank or
MEGA_NO_DIST=1); deterministic candidates by default,
--allow-nondeterministic for ikr; --max-candidates for smoke runs.
- tests/moe_ep/test_knob_cache.py: round-trip/bucket/miss/upsert/disable/
corruption CPU tests + GPU test that buffer creation picks up a cached
entry. Verified on GB200: unit 181 passed, GPU test passed, CLI smoke
recorded a winner and a fresh process resolved it via lookup.
Remaining (tracked in todo_vllm_knob_heuristic.md): populate the cache for
production geometries, skewed-rank offline timing mode, vLLM matrix re-run.
AI-assisted.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Every MoEEpMegaLayer allocated its own symmetric-heap workspace and compiled kernel session; a model's MoE stack multiplies that by the layer count (43x NVSHMEM heap + 43 cute.compiles at DeepSeek-scale). The vLLM integration carried a wrapper-side _SHARED_WORKSPACE dict as a workaround — this upstreams it properly: - core/kernel/workspace_pool.py: process-level refcounted pool. MegaKernelBackend.prepare_workspace acquires by the backend's _workspace_pool_key(); base destroy() releases and only the last layer frees the heap. Sharing the buffer also shares the frontend's compiled kernel (one compile per geometry, not per layer). - nvfp4/mxfp8/deep_gemm backends: pool keys cover everything baked into the buffer — device, EP rank/world/comm group, geometry, kind, effective clamp, combine wire, config-level epilogue scalars (tensors share by object identity only: their values are baked in at creation), and canonicalized knobs. knobs="auto" sessions stay unpooled (autotune recompiles the shared frontend). - Sharing is safe because the workspace is stateless across forwards (staging overwrites inputs, the kernel tail-cleans its counters) and layers run sequentially on one stream — same argument as the validated vLLM workaround, which is now obsolete. - tests/moe_ep/test_workspace_pool.py: pool/refcount/key-semantics CPU tests, base-wiring tests with a fake backend, and a GPU test where two same-geometry nvfp4 layers share one buffer, agree bit-exactly, and survive one layer's destroy. Verified on GB200: unit 189 passed, pool GPU test passed, CUDA-graph (3) and fused-staging (7) regressions green. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…xonerated Layer-level discrimination tests for the vLLM e2e fi_dg run-to-run nondeterminism (todo_fi_dg_nondeterminism.md). Both PASS on 4x GB200: - repeated forwards stay bit-exact under gross (~100 ms) per-iteration rotating rank skew — the deep_gemm mega combine is arrival-order fixed, killing the skew/atomics hypothesis; - identical input bytes at shifted base addresses (16B/8B/128B/256B alignments) produce bit-identical staged fp8, scales, and outputs — killing the allocator-address hypothesis. With staging bytes, kernel args, is_padding handling, workspace provenance, and output-buffer freshness already equalized against native, the remaining lead is engine-side batch-formation timing (chunked-prefill/batch shapes straddling scheduler boundaries across runs); next experiment documented in the TODO. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…backend-independent Fresh smoke rerun pairs (2026-07-17) invert the 07-15 result: fi_dg 8/8 self-exact, native 3/8 divergent under identical conditions. Combined with the layer-level exoneration (skew + address probes), run-to-run nondeterminism is an engine/environment property, not an fi moe_ep defect. Reclassified in the TODO; determinism claims now require N rerun pairs per backend + the per-step batch-shape log as a control. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…on timing (schedule diff at step 5, 1-tok vs 8-tok batch); closed as fi issue AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
todo_cuda_graph.md item 5 (multirank half): collective warmup, per-rank capture (capture records without executing, so no cross-rank dependency), lockstep replays, and replay-over-mutated-inputs — all bit-exact vs eager on 2x GB200 (torchrun -np 2). AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The dg staging recipe (per_token_cast_to_fp8 use_ue8m0=True gran_k=32 use_packed_ue8m0=True) is byte-identical to the cutedsl mxfp8_e4m3 recipe — data bytes equal, and the "packed" int32 scales are the same per-32 e8m0 bytes viewed 4-per-word (verified on GB200). The deep_gemm backend now stages through the fused single-launch DataPreprocess kernel via byte views, replacing its multi-kernel torch staging (fi_dg's share of the launch soup); FLASHINFER_MEGA_FUSED_STAGE=0 keeps the torch path. The address-shift probe caught a fused-path constraint the torch path lacked: the DSL launcher requires 16-byte-aligned activation pointers. All three backends now fall back to torch staging per call via fused_quant_stage_supported() when the input is misaligned (engine tensors are allocator-aligned; the fallback covers sliced/offset views). Verified on GB200: fused-stage equivalence 8/8 (incl. dg int32-scale layout), dg oracle, 4-rank layer-vs-reference + skew + address probes all bit-exact with the fused path active, CUDA-graph tests 3/3, unit suite green. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Deleted (all resolved on this branch, full analysis preserved in git history and outcomes in tests/commits): - todo_weight_pack_retention_oom.md — fixed (b79b74d), vLLM workaround since removed from the integration patch - todo_cuda_graph.md — implemented (964d301): warmup contract + capture guards, single-rank + 2-rank lockstep replay tests green, and fi_nvfp4 runs under vLLM CUDA graphs (enforce_eager=False smoke) - todo_fi_dg_nondeterminism.md — closed: layer exonerated (skew + address probes), premise inverted on rerun pairs, mechanism = vLLM batch-formation timing (engine-level, backend-independent) - todo_vllm_knob_heuristic.md — implemented (c5a31c5): knob cache + offline tuner CLI; DSV4 geometry tuned and e2e-validated (+4.5%/+5.0%) Cross-references in remaining TODOs, TUNING.md, knob_cache.py, and test docstrings updated to point at the durable artifacts. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mask) + caller-owned ikr Steady-state nsys (graphs-mode decode, 95% GPU busy both backends) pinned the remaining fi_nvfp4 decode gap on the explicit TopkReduce combine (+32us/step vs native's in-kernel combine; the fp4 mega kernel itself is FASTER than dg at decode shapes, 124 vs 136us) and prefill on mega kernel time at chunked-prefill shapes — host overhead is exonerated. Per-call launch inventory: fi 44 vs native 42, with exactly two cuttable extras, both removed here: - compute(output=None) on the cutedsl backends returns the workspace [:n] output view instead of copying into a caller buffer (zero-copy; live-token count comes from the staging memo; layer API unchanged). - The topk_idx capacity-tail re-mask is memoized per buffer: staging only overwrites [:n], so a fill is needed only over [n, prev_n) when the batch shrinks. The torch fallback and pre-staged copy paths update the memo (note_staged_tokens) so mixed sequences stay safe; unknown buffers default to fully-live. - in_kernel_fc2_reduce is now caller-owned in the symm-buffer factories: cached/heuristic knob dicts cannot flip the config's choice (it is a correctness/determinism decision, not a perf knob); forcing ikr on mxfp8 also forces epi-warps token-back per the kernel constraint. Verified on GB200: fused-stage suite 9/9 (incl. shrink/grow + mixed memo sequences and dg layout), graph tests 4/4 (incl. zero-copy view bit-equality), nvfp4 oracle, nvfp4+mxfp8 4-rank multirank 12/12, unit suite green. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The tail-fill memoization (ea26808) was incompatible with the one-graph-per-batch-size capture pattern serving engines use: graph replays bypass host memo logic, so a small-size graph replayed after a larger one (or an eager call after replays) skipped fills over rows the larger launch left live — stale rows dispatched garbage work (observed as a 3x decode-graphs regression in vLLM). Now every captured graph bakes the full conservative [n, capacity) fill and a buffer that has ever been staged during capture permanently drops the memo; the memo only accelerates pure-eager sessions, its honest scope. New regression test drives the exact engine pattern: two graphs at different sizes + eager, interleaved (test_mega_cuda_graph.py::test_mega_layer_multi_size_graphs_and_eager_interleave). Also: python -m flashinfer.moe_ep.tune gains --live-tokens to time sweeps at decode-like live counts on a larger buffer bucket (record such winners to a separate cache file — the entry is keyed on the bucket). First use at DSV4 (24 candidates incl. ikr, live=256): the winner is NON-ikr (256x128 / fb8 / epi_warps, 290us) — ikr's small-token penalty is inherent at this geometry, confirming the e2e verdict at kernel level. Verified on GB200: fused-stage 9/9, graph tests 5/5 (incl. the new interleave reproducer), unit suite green. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… regression on capture-touched buffers) The capture-branch memo poisoning (c41c5e3) broke compute(output=None): staged_tokens() returned capacity instead of the live batch size during engine capture warmup (4096-row view for a 512-row batch). The memo now always records the actual staged count — the tail-fill decision for capture-touched buffers never consults it (they unconditionally full-fill). Interleave regression test extended with view sizing on a capture-touched buffer. Verified: fused-stage 9/9, graph tests 5/5. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ep schedule) The tuner's default routing is near-uniform, which cannot discriminate skew-sensitive knobs — production routing shows 14-28x per-layer expert load skew (cold-run FI_MOE_EP_LOAD_STATS). --skew restages the tuning routing at a target max/mean ratio (power-law popularity, bisected exponent); --sweep schedule pins tile/token-back from the cache winner (or --base-knobs) and sweeps load_balance_mode x group_hint, the axes the default candidate set fixes. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ot path Host-probe measurement (run 28): the per-layer-call host cost of the fi fast path was ~100us, of which ~70us was loop-invariant Python inside nvfp4_mega_moe()/frontend.run() — re-validation, clamp re-resolution, the 12-field inputs rebuild, and launch-cache re-keying on every call. At 43 layers x 4 ranks this is the measured generator of the inter-rank arrival skew the mega kernel then absorbs in its device barriers (run 27: WAIT 850us ~= SKEW 864us per launch, flat across layers, vs native's arrival-independent kernel + allreduce-overlapped sync). The backend now builds the shim's make_launch_thunk() once per (workspace, weights, compiled-session, stream) and steady-state calls are thunk() + a view slice — compute() host cost halves (6.3x -> 3.0x a raw torch-op launch). The STREAM is part of the cache key: thunk kwargs bind the stream at build time, and a graph capture must get its own capture-stream thunk or the kernel launch escapes the graph (caught by the multi-size interleave test: replays silently reusing stale kernel output). Knob/clamp changes null the compiled session and force a rebuild through the fully validated path (make_launch_thunk -> _prepare_launch_inputs). Verified on GB200: graph tests 5/5 (incl. both capture regressions), nvfp4 oracle, 4-rank multirank 8/8, 2-rank lockstep graph replay, unit suite. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MoEWeightPack stays a real base class whose __new__ dispatches to the frozen variant, so all construction sites, weights: MoEWeightPack annotations, and isinstance validation keep working unchanged while the mixed one-scale state raises at construction. Backends discriminate via isinstance(weights, PrequantizedMoEWeights). AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MegaMoE kernels compile 34-54% slower on nvidia-cutlass-dsl 4.5.x (TUNING.md 'CuTe-DSL runtime sensitivity'); results stay correct, so warn instead of raise. Silence with FLASHINFER_MOE_EP_SKIP_DSL_CHECK=1. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The shared validate_mega_fleet_params %128 bound was inherited from deep_gemm, whose wire format is genuinely hard-%128 (SFs packed 4-per-int32 word, static-asserted host and device side). The cutedsl kernels are tail-safe (ceil-div K/M loops, TMA OOB zero-fill, predicated epilogue stores); their true bound is TMA 16B row alignment + SF-word packing = %64. Unblocks gpt-oss-120b geometry (hidden=inter=2880). - validate_mega_fleet_params gains a per-backend alignment param; nvfp4/mxfp8 cutedsl pass 64, deep_gemm keeps the 128 default. - shim config validators + staging guards drop to %64. - fused quant-stage: support gate is now per-type (nvfp4 %64, mxfp8 %128 - its SF view needs hidden/32 % 4); unsupported shapes fall back to torch staging instead of erroring. - oracle test parametrized with (2880, 2880) to pin the K-tail and predicated-epilogue paths to the torch reference. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ng-note pointers) AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ling doc refs) pre-commit run on all branch-touched files: ruff-format reformats in validation/common.py and shim/quant_stage.py; pytest.raises narrowed to dataclasses.FrozenInstanceError (B017); explicit type-ignore[misc] on the intentional ClassVar shadowing in UnquantizedMoEWeights; two references to the deleted todo_weight_pack_union.md planning note inlined. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nch tip) Fresh 7-point default-geometry sweep (job 2419216) reproduces the 07-15 table within run noise; the 07-15 section is kept as the historical record of the layout-bug correction. Adds a pointer to the model-shapes sweep (6 real-model geometries incl. gpt-oss-120b via the %64 relaxation) in moe_ep_benchmark/model_shapes/RESULTS.md. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…bles, add vLLM e2e results, knob-cache runtime flow) - Promote the 2026-07-21 sweep to the single microbench reference table; fold the 07-15 corrected sweep into a short historical note (it reproduces at tip within noise). - Add the vLLM 0.25.1 e2e section: prefill 1.18x / decode-1k 1.07x / GSM8K 0.975 vs native, fi_dg control, and the capture-the-prefill-chunks requirement behind the decode number. - Refresh the combine-leg takeaways to 07-21 values + e2e transfer findings (ikr and quantized wires lose at DSV4 geometry). - Runtime-flow section now documents the persistent knob cache (pure lookup before the heuristic, per-role cache files) instead of the pre-c5a31c55 "no persistence" story. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…NG.md Reviewers previously had to follow an external moe_ep_benchmark pointer for the model_shapes results; the 2026-07-21 tables (deepseek_v3 / v4-flash / v4-pro, kimi_k2_6, qwen3_5_397b, gpt_oss_120b) now live inline, columns reordered to match the main sweep table. Values verified cell-by-cell against model_shapes_20260721_1109* CSVs. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MegaConfig.megakernel is duck-typed object; every kernel config declares top_k: int, so coerce + ignore in the same style as nccl_ep/handle.py. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Working notes, not PR material; they stay local-only. Also moots the review suggestion to sync their stale "Remaining" sections. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
mhoqueanik
requested review from
dhiraj113,
jimmyzho,
qiching and
qsang-nv
as code owners
July 22, 2026 20:59
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughCuTeDSL MegaMoE adds a CuTeDSL 4.5.2 NVFP4 compatibility path, fixes zero-token staging state, threads combine dtype through tuning, expands shim exports, updates descriptor utilities, and refreshes benchmark and version guidance. ChangesCuTeDSL MegaMoE updates
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Staging
participant QuantStage
participant TokenMemo
Staging->>QuantStage: stage zero-token input
QuantStage->>TokenMemo: mask routing tail with -1
QuantStage->>TokenMemo: record staged token count 0
TokenMemo-->>Staging: cleared routing state
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.py (1)
2255-2283: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPeeled-iteration body duplicates the in-loop MMA-issue logic.
The SFA/SFB s2t copy +
issue_dynamic_block_scaled_mma_tilecall here duplicates the loop body above (lines ~2229-2253) apart from thek_tile_idxand the absence of the peek/advance dance. A small local helper (closure overtCtAcc,tCtSFB_mma, etc.) takinghandle/k_tile_idxwould remove the duplication and reduce the risk of the two copies drifting if the MMA-issue call's arguments change later.♻️ Sketch of a shared helper
def _issue_mma_for_stage(handle, k_tile_idx): s2t_stage_coord = (None, None, None, None, handle.index) cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t[s2t_stage_coord], tCtSFA_compact_s2t) cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t[s2t_stage_coord], tCtSFB_compact_s2t) tile_crd = (None, None, None, handle.index) dynamic_mainloop.issue_dynamic_block_scaled_mma_tile( acc_tensor=tCtAcc, a_frag_tile=tCrA[tile_crd], b_frag_tile=tCrB[tile_crd], sfa_tensor=tCtSFA, sfb_tensor=tCtSFB_mma, k_tile_idx=k_tile_idx, valid_tokens_in_tile=work_tile_info.valid_tokens_in_cta_tile, mma_tiler_mnk=self.mma_tiler_mnk, ) handle.release()Both the main loop body and the peeled block then just call
_issue_mma_for_stage(handle, k_tile)/_issue_mma_for_stage(handle, k_tile_cnt - 1).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.py` around lines 2255 - 2283, Extract the duplicated SFA/SFB copy and MMA issuance logic from the main loop and peeled `cutedsl_452_ver_war` block into a local helper near these statements, capturing the shared tensors and configuration and accepting `handle` and `k_tile_idx`. Replace both implementations with helper calls while preserving each block’s existing wait/advance behavior and using `k_tile_cnt - 1` for the peeled iteration; keep handle release inside the shared flow.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/knob_cache.py`:
- Around line 99-100: Update `_load_entries` to return only elements that are
dictionaries when `data["entries"]` is a list, filtering out all other element
types while preserving the existing empty-list fallback for invalid or missing
entries.
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/quant_stage.py`:
- Around line 146-149: The zero-token branch in the fused staging path must
clear this buffer’s staged-row mask and cached count before returning. Update
the logic around _LAST_STAGED_N and the associated tail mask so staged_tokens()
and compute(output=None) report zero on reused pooled buffers, while preserving
normal staging for non-empty inputs.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.py`:
- Around line 2216-2227: Add an inline comment at the cutedsl_452_ver_war
conditional in the K-loop pipelining logic explaining that CuTe DSL 4.5.2 has a
code-generation regression requiring the final K tile to be peeled, and include
the relevant upstream MR reference. Keep the existing wait/advance cadence and
version-gated behavior unchanged.
In `@flashinfer/moe_ep/tune.py`:
- Around line 190-206: Thread args.combine_dtype through the NVFP4 tuning flow:
pass it to create_dummy_nvfp4_inputs and include it in the schedule sweep’s
resolve_knobs call. Ensure timing configuration, candidate pruning, and cached
winner records use the requested combine dtype instead of defaulting to bf16.
---
Nitpick comments:
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.py`:
- Around line 2255-2283: Extract the duplicated SFA/SFB copy and MMA issuance
logic from the main loop and peeled `cutedsl_452_ver_war` block into a local
helper near these statements, capturing the shared tensors and configuration and
accepting `handle` and `k_tile_idx`. Replace both implementations with helper
calls while preserving each block’s existing wait/advance behavior and using
`k_tile_cnt - 1` for the peeled iteration; keep handle release inside the shared
flow.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1e416b0f-5ece-483b-8e04-8aecd2c565c8
📒 Files selected for processing (41)
CLAUDE.mdflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/backend.pyflashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/staging.pyflashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/weights.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/backend.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/staging.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/weights.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/backend.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/staging.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/weights.pyflashinfer/moe_ep/core/kernel/base.pyflashinfer/moe_ep/core/kernel/workspace_pool.pyflashinfer/moe_ep/core/validation/common.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/__init__.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/__init__.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/autotune.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/comm.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/knob_cache.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/mxfp8.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/nvfp4.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/quant_stage.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/modes/mega_layer.pyflashinfer/moe_ep/modes/split_layer.pyflashinfer/moe_ep/tune.pyflashinfer/moe_ep/weights.pytests/moe_ep/run_tests.shtests/moe_ep/smoke_nixl_ep.pytests/moe_ep/test_fused_quant_stage.pytests/moe_ep/test_knob_cache.pytests/moe_ep/test_layer_single_gpu.pytests/moe_ep/test_mega_cuda_graph.pytests/moe_ep/test_mega_cuda_graph_multirank.pytests/moe_ep/test_mega_layer_validation.pytests/moe_ep/test_moe_ep_deep_gemm_skew_determinism.pytests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.pytests/moe_ep/test_weight_pack_union.pytests/moe_ep/test_workspace_pool.py
…updated nvidia-cutlass-dsl-internal 0.3.0+20260721052439.f58f4a5 (the 4.7.0.dev drop): nvfp4 427.0/614.2/1907.7 us @1024/2048/8192, in the fast band. Documented the internal-wheel version-scheme caveat: cutlass.__version__ reports 0.3.0, so the ==4.5.2 WAR gate is always off and the shim warning skips on internal drops — correct as long as they carry the 4.5.3+ fix. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 23, 2026
…edsl sources The vendored moe_nvfp4_swapab tree still carried the internal release-tooling markers. All but one guarded comments only, so removing them is cosmetic: mega_reference.py (SFB peer-CTA replication note, the 256xnNSF cta_tile_n=192 layout note, the SFB-slicing and TMEM-offset hack notes, the m-major FP4 TODO), moe_persistent_scheduler.py (the CLC cluster-origin semantics block), and moe_utils.py (the GeneralGroupedGemmTensormapConstructor internal example, unreferenced here -- the MoE path uses its own constructor below it). The one marker that guarded code was moe_utils.py's TensormapDescBytes, written as 128 then immediately shadowed by a marked "= 64". 64 is what every build in this tree has actually used (it is the per-slot stride and alignment in TensormapWorkspace), so it is now the single unconditional assignment -- no behavior change, but stripping the line would silently have doubled the stride. Caveat: a future re-sync from the internal kernel source will reintroduce these markers and conflict on the touched lines. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The runbook covered build/test but pointed nowhere for perf. Adds the two harnesses that live in the companion moe_ep_benchmark repo -- the single-node microbenchmark (fi_mega / model_shapes) and the vLLM 0.25.1 e2e suite -- with the clone step, the five-variant column legend, and the exact sweep commands. Numbers are the 2026-07-21/22 measurements at this branch tip on 4x GB200, nvidia-cutlass-dsl 4.5.2: microbenchmark tables for all six model geometries (dg-parity below ~512 tok/rank, crossover by 1024, up to 1.62x at 8192 and 1.86x with the nvfp4 combine wire) and the vLLM headline cells (prefill-8k 1.18x, decode-1k 1.07x, GSM8K in-band). Also records the methodology that the cells depend on: e2e_pipelined timing, CUDA-graph capture must cover the prefill chunk shapes (otherwise decode reads ~4% below native -- a config artifact), per-role offline knob caches, the two-checkpoint fairness gate, and >=3 repeats with medians given the +-35% cross-restart variance on prefill-heavy cells. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/design_docs/moe_ep_runbook.md`:
- Around line 73-85: Update the benchmarking setup instructions for the
moe_ep_benchmark checkout to switch from the default branch to the measured
4_5_2-perf-fix branch and pin the repository to the exact commit used for the
published results, preserving the existing clone workflow.
- Around line 267-285: Align the decode metric in the reproduction note with the
headline table in the “End-to-end results” section: either report output tok/s
consistently or explicitly label the reproduction’s total tok/s as a separate
result. Update the claim that the headline cells reproduce so it only applies
when the same metric is being compared.
- Around line 260-265: Update the GSM8K fairness-gate documentation near the
“Two-checkpoint fairness” guidance to define the exact allowed accuracy
difference and the pass/fail rule, including the evaluation conditions used for
both checkpoints. Ensure the documented threshold makes it possible to determine
whether scores such as 0.965 and 0.975 satisfy the gate.
- Around line 87-90: Update the introductory version statement in the runbook to
explicitly identify CuTe-DSL 4.5.2 as the measurement and supported baseline,
describe 4.6.1 only as a parity comparison reference, and state that versions
below 4.5.2 are unsupported.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.py`:
- Around line 48-50: Update the TensormapDescBytes constant used by
TensormapWorkspace.get_ptr() and size_bytes() from 64 to 128 so TMA descriptor
slots do not overlap and expert workspace is correctly allocated. Document any
independent 64-byte alignment requirement separately rather than using it as the
descriptor stride.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: d403a740-16d3-47d2-87c8-fd600adffecb
📒 Files selected for processing (4)
docs/design_docs/moe_ep_runbook.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_persistent_scheduler.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.py
💤 Files with no reviewable changes (1)
- flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_persistent_scheduler.py
Two moe_ep env-var rows (FLASHINFER_MEGA_FUSED_STAGE, FLASHINFER_MOE_EP_KNOB_CACHE) were added to the env-var reference table as part of the feature commits that introduced those knobs (75221c8, c5a31c5). CLAUDE.md is shared guidance for the whole repo, so this branch should not carry edits to it -- the knobs stay documented in the moe_ep docs alongside the code. Restores the file to its content at the merge base, so this branch's diff no longer touches CLAUDE.md. No code change. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
5 tasks
Anerudhan
approved these changes
Jul 27, 2026
# Conflicts: # flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md # flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/__init__.py # tests/moe_ep/test_layer_single_gpu.py
Collaborator
|
/bot run moe_ep |
Collaborator
- quant_stage/staging: a zero-token staging now re-masks rows the previous batch left routed and records 0 in the live-count memo (fused_quant_stage epilogue extracted into _mask_tail_and_note; the nvfp4/mxfp8 cutedsl wrapper early-outs fill + note explicitly). Regression test added. - tune.py + create_dummy_nvfp4_inputs: thread --combine-dtype into the tuning session (symm buffer config, schedule-sweep base resolution) so quantized-wire tuning times and records under the requested combine_dtype instead of the bf16 default. - knob_cache: filter non-dict elements out of a corrupted cache's entries list so lookup/record degrade instead of raising AttributeError. - kernel_fc12: document the CuTe-DSL 4.5.2 mainloop WAR (MR!27 last-k-tile peel) rationale at the point of use. - moe_utils: document why TensormapDescBytes=64 is correct (copy_tma_desc copies 8 x i64; PTX needs only 64B alignment) — value unchanged; the flagged "128 -> 64" was the vendored internal-release line that was always in effect. - runbook/TUNING.md: pin the moe_ep_benchmark revision (vllm-pr c8aefda), state 4.5.2 as the measured baseline (4.6.1 = parity reference only), define the GSM8K fairness band (fi >= native - 0.02, 200q greedy), and label the 07-22 decode revalidation as total tok/s (ratio corroboration, not the headline output-tok/s cells). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md (1)
643-643: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winEncode the DSL version requirement in the pip command.
The
# >= 4.5.2comment is not enforced, and--upgrademay install a different version than the one used for the documented validation. Usenvidia-cutlass-dsl[cu13]>=4.5.2; pin==4.5.2when reproducing the WAR-specific results.Proposed fix
- python -m pip install --upgrade "nvidia-cutlass-dsl[cu13]" # >= 4.5.2 + python -m pip install --upgrade "nvidia-cutlass-dsl[cu13]>=4.5.2"As per coding guidelines, “Keep documentation synchronized with code changes.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md` at line 643, Update the documented pip command in TUNING.md to require nvidia-cutlass-dsl[cu13]>=4.5.2 instead of leaving the version in a comment; specify ==4.5.2 for the WAR-specific reproduction instructions.Source: Coding guidelines
flashinfer/moe_ep/tune.py (1)
239-248: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winKeep MXFP8 schedule lookups on the BF16 cache key.
MXFP8 dummy sessions do not consume
combine_dtype, but this lookup now keys them by a user-supplied non-BF16 value. That can miss the actual MXFP8/BF16 cache entry and tune from an unrelated baseline. Force"bf16"whennot is_nvfp4, or reject non-BF16--combine-dtypefor MXFP8.Proposed fix
- combine_dtype=args.combine_dtype, + combine_dtype=args.combine_dtype if is_nvfp4 else "bf16",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/moe_ep/tune.py` around lines 239 - 248, Update the resolve_knobs call in the tuning flow so MXFP8 lookups (when is_nvfp4 is false) always pass "bf16" as combine_dtype, regardless of the user-supplied value; preserve args.combine_dtype for NVFP4.docs/design_docs/moe_ep_runbook.md (3)
552-552: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd language identifiers to fenced code blocks.
markdownlintreports missing fence languages at Lines 552, 639, and 653. Usetextfor these diagrams/formulas.Also applies to: 639-639, 653-653
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/design_docs/moe_ep_runbook.md` at line 552, Add the text language identifier to the fenced code blocks at the affected sections, including the diagrams and formulas around lines 552, 639, and 653, so each opening fence uses text and satisfies markdownlint.Source: Linters/SAST tools
665-667: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winState the contiguous, divisibility precondition for the owner formula.
This assumes
world_sizeis positive, fewer than or equal tonum_experts, andnum_expertsis evenly divisible byworld_size; otherwisenum_experts // world_sizecan be 0/nonzero in ways that diverge from the documentation’s contiguous per-rank expert mapping. Add a short precondition/assertion note here to keep the docs in sync with the current EP layout assumptions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/design_docs/moe_ep_runbook.md` around lines 665 - 667, Add a concise precondition note beside the owner calculation stating that world_size must be positive, no greater than num_experts, and num_experts must be evenly divisible by world_size; document that the formula relies on contiguous, equal-sized per-rank expert blocks.
668-669: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winHandle tokens with no surviving experts explicitly.
Clamping
surviving_wonly prevents division-by-zero; it still divies partial-output tokens by1e-6instead of reweighting them, andall_deadtokens with undefined output are scaled by1e6. Branch onsurviving_w == 0or drop below a chosen threshold before dividing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/design_docs/moe_ep_runbook.md` around lines 668 - 669, Update the output normalization around surviving_w so tokens with no surviving experts, or weight below the chosen threshold, are handled explicitly before division. Reweight only tokens with valid surviving experts, and drop or otherwise define the output for all_dead tokens instead of relying on clamp_min(1e-6).
♻️ Duplicate comments (1)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.py (1)
50-57: 🗄️ Data Integrity & Integration | 🟠 MajorKeep descriptor slots at 128 bytes.
TensormapDescBytesis also used byTensormapWorkspace.get_ptr()andsize_bytes()(Lines 458-474), so setting it to 64 makes adjacent descriptor slots overlap the documented 128-byteCUtensorMapobject and underallocates the workspace. NVIDIA documents tensor maps as requiring 128-byte alignment and shows 128-byte tensor-map fencing; a 64-byte copy workaround must not also become the slot stride/alignment contract. Restore a 128-byte slot stride, keeping any 64-byte copy length separate. (docs.nvidia.com)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.py` around lines 50 - 57, Restore TensormapDescBytes to 128 bytes so TensormapWorkspace.get_ptr() and size_bytes() maintain non-overlapping, correctly aligned CUtensorMap slots. Keep any 64-byte copy length used by copy_tma_desc separate from the descriptor slot stride and workspace allocation contract.Source: MCP tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/moe_ep/test_fused_quant_stage.py`:
- Around line 214-231: The test’s capability guard in
test_zero_token_stage_masks_stale_rows is too broad because _require_blackwell()
does not enforce the repository’s SM100/SM103 CUDA-version requirements. Replace
or extend that guard with flashinfer.utils.is_cvt_rs_supported(...) or the
equivalent supported Blackwell helper, preserving the test’s existing
parameterization and behavior while skipping unsupported SM100/SM103 and SM120
environments according to repository semantics.
---
Outside diff comments:
In `@docs/design_docs/moe_ep_runbook.md`:
- Line 552: Add the text language identifier to the fenced code blocks at the
affected sections, including the diagrams and formulas around lines 552, 639,
and 653, so each opening fence uses text and satisfies markdownlint.
- Around line 665-667: Add a concise precondition note beside the owner
calculation stating that world_size must be positive, no greater than
num_experts, and num_experts must be evenly divisible by world_size; document
that the formula relies on contiguous, equal-sized per-rank expert blocks.
- Around line 668-669: Update the output normalization around surviving_w so
tokens with no surviving experts, or weight below the chosen threshold, are
handled explicitly before division. Reweight only tokens with valid surviving
experts, and drop or otherwise define the output for all_dead tokens instead of
relying on clamp_min(1e-6).
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md`:
- Line 643: Update the documented pip command in TUNING.md to require
nvidia-cutlass-dsl[cu13]>=4.5.2 instead of leaving the version in a comment;
specify ==4.5.2 for the WAR-specific reproduction instructions.
In `@flashinfer/moe_ep/tune.py`:
- Around line 239-248: Update the resolve_knobs call in the tuning flow so MXFP8
lookups (when is_nvfp4 is false) always pass "bf16" as combine_dtype, regardless
of the user-supplied value; preserve args.combine_dtype for NVFP4.
---
Duplicate comments:
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.py`:
- Around line 50-57: Restore TensormapDescBytes to 128 bytes so
TensormapWorkspace.get_ptr() and size_bytes() maintain non-overlapping,
correctly aligned CUtensorMap slots. Keep any 64-byte copy length used by
copy_tma_desc separate from the descriptor slot stride and workspace allocation
contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 8ed40740-a3e9-449f-9018-bf35613bbef9
📒 Files selected for processing (11)
docs/design_docs/moe_ep_runbook.mdflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/staging.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/staging.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/knob_cache.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/nvfp4.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/quant_stage.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.pyflashinfer/moe_ep/tune.pytests/moe_ep/test_fused_quant_stage.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.py
mhoqueanik
enabled auto-merge (squash)
July 29, 2026 07:41
aleozlx
approved these changes
Jul 29, 2026
mhoqueanik
pushed a commit
to mhoqueanik/flashinfer-moe_ep
that referenced
this pull request
Jul 30, 2026
Conflict resolution policy: keep this branch's kernel_src/sm100/ layout and sm90 additions; take upstream's post-squash content fixes (zero-token staging re-mask + regression test, knob-cache non-dict entry filter, combine_dtype plumbing in tune.py, ft test target, arch-validation mock, WAR comments) rewritten onto the sm100 paths. TUNING.md 2026-07-22 revalidation paragraph takes the upstream (flashinfer-ai#4101) wording. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Anerudhan
pushed a commit
that referenced
this pull request
Aug 8, 2026
## Summary > Must be merged after #4079 and #4101 - Adds the first Hopper mega kernel to the moe_ep stack: the kernel team's SM90 pull-style FP8 CuTeDSL MegaMoE (Vincent's `hopper_megamoe`, kernel repo commit `1275b8b`), integrated as a new `sm90_pull_fp8` backend behind the existing `MegaKernelBackend` contract — fused NVSHMEM dispatch + FC1 + SwiGLU + FC2 + combine in one launch, FP8 E4M3/E5M2 with per-tensor or DeepGEMM-style blockwise scaling, native and swap-A/B layouts. - New oracle test paths: every mega kernel (SM90 **and** the three existing SM100 paths) now has a **multi-rank torch oracle** — real cross-rank EP launches judged against pure-torch global math instead of same-kernel parity. ## What's in here 1. **Two-tree `kernel_src/` restructure** — the existing Blackwell tree moves to `kernel_src/sm100/cutedsl_megamoe/`; the SM90 drop lands as `kernel_src/sm90/pull_style_cutedsl_megakernel/` (a fork of the same kernel repo, shared runtime intentionally duplicated at its own drop revision). The trees expose colliding top-level module names, so each `shim/_paths.py` now guards process exclusivity (a process runs on one arch anyway). All paths/docs/lint excludes updated; sm100 tree is byte-identical after the move. 2. **Vendored drop + shim** — verbatim `src/` plus our adapter layer (`shim/hopper_fp8.py` mirrors the sm100 `mxfp8.py` design: frozen validated config, lazy `cute.compile`, launch cache, symm-buffer allocator + compute entry per the runbook mega-kernel contract). 3. **Backend + tests** — `Sm90PullFp8MegaMoeConfig` (registered `"sm90_pull_fp8"`), weight preprocessing for both scale modes (gate/up interleave-8, K-major without repack), sm_90 arch gate, runtime requirements, and a core fix: `_init_nvshmem_after_dist` no longer imports a kernel tree (it would have broken every sm90 multirank session via the exclusivity guard). 4. **Benchmark + TUNING.md** — `benchmarks/bench_moe_ep_sm90_mega.py` reproduces the kernel team's 7-point token sweep through the FI layer; measured results and methodology documented in the sm90 tree's `TUNING.md`. 5. **Multi-rank torch oracles** — parity tests can't catch a mega kernel that is wrong but self-consistent at `world_size > 1` (comm + compute are fused, so both sides of a parity test run the same CUDA kernel). New `test_moe_ep_*_mega_multirank_torch_oracle` tests close that gap: each rank launches the fused kernel with real cross-rank NVSHMEM traffic, all-gathers the *actual* operands the kernel consumed (plain pre-swizzle weight legs; for mxfp8/sm90 also staged payloads + routing), and checks its own output slice against torch math over the global expert set, with forced cross-rank routing. Covers sm90 (`{per_tensor, blockwise} × {native, swap_ab}`) and all three SM100 paths, including the variant knobs: nvfp4 `in_kernel_fc2_reduce` + quantized combine wires (`16e2m1xbf16`, `32e4m3xe8m0`, wire modeled exactly via `combine_roundtrip_to_fp32` — newly exported through the sm100 shim boundary), mxfp8 `in_kernel_fc2_reduce`. All picked up by the existing `run_tests.sh mega` / `mega_sm90` targets. 6. **Docs** — `docs/design_docs/moe_ep_architecture.md` gains a "Torch oracles" section (methodology, independence contract, per-kernel last-passed table) and a refreshed `run_tests.sh` target list. ## Testing **SM90 (4×H100 80GB, EP4):** - `run_tests.sh unit` — 189+9 tests pass (new host-only config/registry tests included) - `run_tests.sh oracle_sm90` — kernel vs the drop's fp32 torch reference, `{per_tensor, blockwise} × {native, swap_ab}`: 5/5 pass - `run_tests.sh mega_sm90` — 4-GPU `MoEEpLayer` parity vs a direct-shim session with forced cross-rank routing: bit-exact on separate-reduce paths (incl. pre-staged fp8 inputs and repeat-forward launch-cache guards), roundoff-envelope for `in_kernel_fc2_reduce`: 5/5 × 4 ranks pass; plus the multi-rank torch oracle, 4/4 params × 4 ranks pass **SM100 regression (4×GB200, 2026-07-31):** - `run_tests.sh all` — all 8 sections pass (unit 303/303, single-GPU oracles, split multirank, split-path correctness bf16/nvfp4/ht, mega multirank, smoke): the SM90 integration and restructure leave every Blackwell path unregressed - New SM100 multi-rank oracles: 7 instances (nvfp4 ×4 variants, mxfp8 ×2, deep_gemm) pass on all 4 ranks, `rel_l2 < 0.02` per rank ## Performance 28-point sweep at the drop's DSv4-Pro geometry (384 experts, top-6, hidden 7168, inter 3072, 512–32K tokens/rank): **the FI integration adds no measurable kernel-path overhead** — within ±5% of the kernel team's own harness at most points (their numbers exclude the TopkReduce tail and report min-rank; ours don't), peak 562 TFLOPS/rank. Full tables, comparison caveats, and open items (16K-token iteration variance, fused staging port, DSL 4.6.1 A/B) in `kernel_src/sm90/pull_style_cutedsl_megakernel/TUNING.md`. ## Notes for reviewers - `kernel_src/*/src/` directories are verbatim kernel-team drops — review the shim/backend layers, not those. - Per-tensor activation dequant scales are static calibration scalars, identical across ranks by contract (documented in the config). - The torch oracles never execute the kernel under test (GEMMs are torch fp32); they deliberately share the host-side quant recipes so the compare band stays tight — that surface is covered by the preprocess-vs-plain-quant tests. See the new design-doc section for the full independence contract. - Known gaps, all documented: no tuner/knob-cache for sm90 yet; `PrequantizedMoEWeights` not wired (bf16 canonical or `preprocess_weights=False`); `reuse_dispatch_warps` token-back is perf-exercised but not yet in the correctness matrix. - A push-style SM90 tree will follow as a sibling backend under the same layout. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added Hopper (SM90) FP8 MegaMoE support, including per-tensor and blockwise scaling. - Added public configuration and weight-preprocessing APIs for SM90 FP8 workloads. - Added fused input quantization and routing-stage processing for supported Blackwell formats. - Added automatic performance tuning and cached tuning choices for improved launch efficiency. - **Performance** - Added token-sweep benchmarking tools for FP8 MegaMoE configurations. - **Documentation** - Expanded architecture, benchmarking, fault-tolerance, and kernel-development guidance. - **Tests** - Added coverage for CUDA Graph replay, distributed execution, staging correctness, and Hopper FP8 validation. <!-- 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> Co-authored-by: Md Anik <mhoqueanik@cw-dfw-cs-001-login-01.cm.cluster>
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
nvidia-cutlass-dsl==4.5.2codegen has a regression that makes the cutedsl MegaMoE nvfp4 swap-AB fc12 kernel 34–54% slower than 4.6.1 at every token count (documented inkernel_src/cutedsl_megamoe/TUNING.md, "CuTe-DSL runtime sensitivity"). This mattered because vLLM 0.25.1 pins exactly 4.5.2, forcing integrations to carry a 4.6.1 force-upgrade plus its compat chain (quack, tvm-ffi, tilelang, vendored-kernel patches).This PR ports cutedsl_megamoe MR!27 (single file:
src/moe_nvfp4_swapab/kernel_fc12.py): when the installed DSL is exactly 4.5.2, the MMA-consumer k-tile mainloop is peeled by one iteration (unconditionaltry_waitinside the loop, last tile issued after it). The gate iscutlass.const_expr-folded at trace time, so on any other version the generated kernel is byte-identical to before.Version support matrix (all measured, 4x GB200, default geometry)
cute.compile(unsupported)Support statement: >= 4.5.2 at full performance; < 4.5.2 unsupported. The regression existed only in 4.5.2 and was fixed upstream in 4.5.3, so the exact
== 4.5.2gate never affects any other version.Validation (all on pinned 4.5.2, 2026-07-22)
tests/moe_ep/run_tests.sh all— all 8 sections PASS (unit, torch-oracle, split multirank, bf16/nvfp4/ht correctness, mega multirank, smoke).Changes
kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.py: the MR!27 loop peel,const_expr-gated on== 4.5.2.kernel_src/cutedsl_megamoe/shim/__init__.py: perf-floor warning now fires only below 4.5.2, with the measured support matrix in the docstring.kernel_src/cutedsl_megamoe/TUNING.md: 2026-07-15 sensitivity section marked OBSOLETE (kept as record); reference tables re-measured on 4.5.2 and adopted; version-gap measurements and support statement added.AI-assisted (Claude Code): MR port, benchmark reruns, and docs.
Summary by CodeRabbit
Summary by CodeRabbit
Performance
Compatibility
Bug Fixes
Documentation / Tests