[None][feat] Update CuTeDSL MegaMoE kernels#16190
Open
Barry-Delaney wants to merge 11 commits into
Open
Conversation
Barry-Delaney
force-pushed
the
user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt
branch
11 times, most recently
from
July 17, 2026 04:39
d4cb10a to
5cded49
Compare
Barry-Delaney
added
the
api-compatible
Barry-Delaney
force-pushed
the
user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt
branch
from
July 17, 2026 09:57
06e6965 to
468c64a
Compare
Barry-Delaney
marked this pull request as ready for review
July 17, 2026 09:57
Barry-Delaney
requested review from
QiJune,
arysef,
bo-nv,
lori-ren,
nv-xtf and
sunnyqgg
July 17, 2026 09:57
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughMegaMoE CuteDSL adds quantized combine formats, token-aware tactics, adaptive token buckets, CUDA-graph-safe workspace fencing, distributed token-back communication, device-side top-k reduction, and streamed NVFP4 weight loading. Backend registration, scheduler integration, model configuration, and tests are updated. ChangesMegaMoE CuteDSL runtime
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Scheduler
participant MegaMoECuteDsl
participant Sm100MegaMoEKernel
participant TokenInPullTokenBackPush
participant TopkReduce
Scheduler->>MegaMoECuteDsl: provide rank-identical token count
MegaMoECuteDsl->>MegaMoECuteDsl: select launch bucket and acquire buffers
MegaMoECuteDsl->>Sm100MegaMoEKernel: launch fused dispatch and FC1/FC2
Sm100MegaMoEKernel->>TokenInPullTokenBackPush: exchange token-back DATA and scale payloads
Sm100MegaMoEKernel->>TopkReduce: reduce combine cells when separate reduction is selected
TopkReduce-->>MegaMoECuteDsl: return token-hidden output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: 1
🧹 Nitpick comments (2)
tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py (1)
36-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winApply the required function annotations consistently across the new tests.
Both new test surfaces omit annotations required by the Python coding standard.
tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py#L36-L180: annotate helpers, nested helpers,_StreamingMoEModule.__init__, and test functions.tests/unittest/_torch/modules/moe/test_moe_backend.py#L479-L525: annotate the constructor helper and all three test functions.As per coding guidelines, “Annotate every function, use
Nonefor non-returning functions, avoidAnyand unnecessary type ignores.”🤖 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 `@tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py` around lines 36 - 180, Add explicit return and parameter annotations to every function in tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py lines 36-180, including _load_classes, _StreamingMoEModule.__init__, nested rand_u8/rand_fp8, _w13_input_scale, _w2_input_scale, _make_vanilla_weights, _fresh, _load, _streamed_numels, _assert_sources_freed, _expected_mega_fc2, _expected_mega_fc1, _expected_fc1_norm_const, and both test functions; use None for non-returning functions and concrete container/tensor types without Any. Apply the same annotations to the constructor helper and all three test functions in tests/unittest/_torch/modules/moe/test_moe_backend.py lines 479-525.Source: Coding guidelines
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py (1)
401-401: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a return type annotation to
eval_function_mapping.All return paths yield a
dict, so annotate accordingly for consistency with the rest of the module.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".Proposed change
-def eval_function_mapping(contract: Contract, **domain_coord): +def eval_function_mapping(contract: Contract, **domain_coord) -> dict:🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py` at line 401, Add a return type annotation of dict to eval_function_mapping, preserving its existing behavior and all return paths.Source: Coding guidelines
🤖 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/microbenchmarks/bench_moe/timing/autotune.py`:
- Around line 61-67: In the autotune flow around moe.forward() and
torch.cuda.synchronize(), determine the tactic_autotune flag before entering the
try block, then add conditional profiling-scratch release to the cleanup path so
it runs on both success and failure. Preserve the existing tuner-state
restoration in finally and reuse release_megamoe_profiling_scratch for enabled
backends.
---
Nitpick comments:
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py`:
- Line 401: Add a return type annotation of dict to eval_function_mapping,
preserving its existing behavior and all return paths.
In `@tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py`:
- Around line 36-180: Add explicit return and parameter annotations to every
function in tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
lines 36-180, including _load_classes, _StreamingMoEModule.__init__, nested
rand_u8/rand_fp8, _w13_input_scale, _w2_input_scale, _make_vanilla_weights,
_fresh, _load, _streamed_numels, _assert_sources_freed, _expected_mega_fc2,
_expected_mega_fc1, _expected_fc1_norm_const, and both test functions; use None
for non-returning functions and concrete container/tensor types without Any.
Apply the same annotations to the constructor helper and all three test
functions in tests/unittest/_torch/modules/moe/test_moe_backend.py lines
479-525.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2e4ad531-b0fe-43a0-b057-a5d2fce1260c
📒 Files selected for processing (33)
tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.pytensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/test_lists/test-db/l0_b200.ymltests/microbenchmarks/bench_moe/backend.pytests/microbenchmarks/bench_moe/build.pytests/microbenchmarks/bench_moe/search.pytests/microbenchmarks/bench_moe/timing/autotune.pytests/microbenchmarks/bench_moe/utils.pytests/unittest/_torch/modules/moe/test_megamoe_streaming_load.pytests/unittest/_torch/modules/moe/test_moe_backend.py
Collaborator
|
PR_Github #60899 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60899 [ run ] completed with state
|
…Seek-V4 Add the MegaMoE-CuteDSL MoE backend (moe_config.backend=MEGAMOE_CUTEDSL): a pure cute.compile JIT integration of the MegaMoE NVFP4 swap-AB kernel (vendored under cute_dsl_kernels/mega_moe_nvfp4) for DeepSeek-V4 DEP serving on Blackwell. - Vendored kernel package: dispatch/combine token communication with NVLink sense-reversing barriers, swap-AB FC1/FC2 scheduling, NVFP4 epilogue with E8M0 scale conversion, adaptive grid sizing. - Custom op wrapper: workspace/fence lifecycle keyed on symmetric-buffer identity with storage-weakref liveness (a recycled allocation address can never skip the barrier-reset fence; capture-safe eviction, latched tuning workspaces, lazy profiling scratch released after autotuner warmup), adaptive max-tokens bucket ladder. - Tactic autotuning is an explicit opt-in (tactic_autotune backend constructor parameter, default OFF): serving always uses the deterministic heuristic tactic even under global autotuner warmup; the microbenchmark's --autotune enables the sweep, which then runs through the standard AutoTuner (tuning mode only). The curated reduced space is the only tuning space. - Cross-rank combine wire format is an explicit, validated backend constructor parameter (bf16 default; 32e4m3xe8m0 / 16e2m1xbf16 for benchmarking, threaded create_moe -> backend -> op and part of the workspace/provider cache keys; quantized formats use separate-reduce and disable the incompatible in-kernel reduce). No environment variables besides MEGAMOE_MPI_INIT_TORCH_DIST (opt-out of the process-level torch.distributed bootstrap). - Backend module wires the staged post-load hooks (transform_weights / cache_derived_state), free-port MPI bootstrap for multi-rank tuning, and fail-fast guards for unsupported configs (EPLB reload, pure-DEP requirement). - moe_config.backend gains the MEGAMOE_CUTEDSL literal (golden manifest regenerated); model_engine releases the profiling scratch at the end of autotuner warmup. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
… MegaMoE-CuteDSL - DSv4 model wiring: swiglu clamp forwarded in tensor form, fused-A / kv_b_proj atomic weight groups with module-state-keyed required sets, cross-RPC partial-group stash with covering-resend audit (bucket keys pre-registered per cycle, aborts union them into a tainted set that only a covering sweep clears), and a sticky abort-cycle protocol exposed as the duck-typed abort_reload_cycle hook. - attn_sink is a module-owned parameter (graph-stable storage created at construction, in-place refresh on load, value-preserving staged hooks) so checkpoint sinks survive CUDA-graph capture, dummy init and direct-transport loads; dummy weight init skips .attn_sink and memory-tagged allocation preserves constructor values. - Streaming source load for the MegaMoE-CuteDSL quant method: 0-element placeholders materialized per load, eager per-module finalize once coverage completes (CPU staging on partial reloads keeps the GPU peak order-independent), per-expert-row streamed coverage guard, and cycle-start purge of per-cycle quant transients (_RELOAD_TRANSIENT_ATTRS) so an aborted cycle's scales can never surface in the recovery cycle's finalize. - Reload correctness guards shared by Linear/MoE: per-unit reload-coverage debt recorded before pre-reload re-registration and consumed only by keys actually rewritten, with finalize refusing while units are outstanding; a side-effect-free check_reload_capability preflight refuses EPLB/non-partial-capable quant methods before the destructive walk; the update_weights exception boundary aborts the cycle (latch cleared, duck-typed model hook) and abort_update_weights lets a coordinator broadcast the same abort to every rank. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Self-contained kernel microbench over the public create_moe/forward path: synthetic NVFP4 weights, EP/DEP via mpirun, sweeps tokens-per-rank / combine wire format (passed explicitly to the backend constructor), --autotune opting into the MegaMoE tactic sweep through the standard AutoTuner (eager and --cuda-graph timing modes), free-port bootstrap and between-point workspace reset so repeated points are independent. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Treat the non-expert weights in DeepSeek V4 NVFP4 checkpoints as FP8 block-scale weights while retaining the per-layer NVFP4 expert configuration. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Barry-Delaney
force-pushed
the
user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt
branch
from
July 23, 2026 02:01
59cc034 to
a6aa781
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #61149 [ run ] triggered by Bot. Commit: |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py (1)
662-686: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate sentinel
WorkTileInfoconstruction — extract a shared helper.The
is_swap_absentinel-init block increate()(Lines 662-686) is duplicated verbatim in_gen_work_from_cluster_idx()(Lines 1258-1280). Both explicitly document that the field set must trackext.WorkTileInfo's serialized shape exactly — a future field change only applied to one copy would silently desync MLIR serialization.♻️ Suggested consolidation
`@staticmethod` def _default_work_tile(ext, params) -> "MoEWorkTileInfo": if const_expr(params.is_swap_ab): return ext.WorkTileInfo( expert_idx=Int32(WorkTileState.DONE), tile_m_idx=Int32(0), tile_n_idx=Int32(0), cumulative_data_physical_row=Int32(0), cumulative_sf_physical_row=Int32(0), cumulative_token_block_count=Int32(0), valid_tokens_in_cta_tile=Int32(0), phase_and_peek=Int32(BlockPhase.None_), ) return ext.WorkTileInfo( expert_idx=Int32(WorkTileState.DONE), tile_m_idx=Int32(0), tile_n_idx=Int32(0), cumulative_data_physical_row=Int32(0), cumulative_sf_physical_row=Int32(0), cumulative_token_block_count=Int32(0), valid_tokens_in_cta_cluster_tile=Int32(0), phase_and_peek=Int32(BlockPhase.None_), fc1_counter_index=Int32(0), )Then call
self._default_work_tile(ext, params)from bothcreate()and_gen_work_from_cluster_idx().Also applies to: 1258-1280
🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py` around lines 662 - 686, Extract the duplicated sentinel WorkTileInfo construction into a shared _default_work_tile helper that selects the exact field set based on params.is_swap_ab. Replace both initialization blocks in create() and _gen_work_from_cluster_idx() with calls to this helper, preserving the serialized shapes and sentinel values unchanged.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py`:
- Around line 662-686: Extract the duplicated sentinel WorkTileInfo construction
into a shared _default_work_tile helper that selects the exact field set based
on params.is_swap_ab. Replace both initialization blocks in create() and
_gen_work_from_cluster_idx() with calls to this helper, preserving the
serialized shapes and sentinel values unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7a64c324-04bb-4756-b59c-604e208c2c20
📒 Files selected for processing (25)
tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.pytensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.pytensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.json
🚧 Files skipped from review as they are similar to previous changes (18)
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
- tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/init.py
- tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
- tensorrt_llm/_torch/modules/fused_moe/quantization.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py
Collaborator
|
PR_Github #61149 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61269 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61269 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61294 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61294 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61374 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61374 [ run ] completed with state |
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Shixiaowei02
approved these changes
Jul 24, 2026
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
5 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.
Dev Engineer Review
default_megamoe_tactic(num_tokens)and made candidate enumeration depend onnum_tokens.validate_megamoe_tactic()arity/field constraints.in_kernel_fc2_reduce,combine_format,tactic_autotune, and optional livenum_tokens.token_back_by_dispatch,non_ubulk_fc2_store.valid_tokens_in_cta_tile, and non-swap_abcluster-tile +fc1_counter_indexhandling).QuantImpl,combine_formatpropagation, centralized NVFP4/CombineFormat logic).CombineFormat,TokenSrcMetadata, andGpuReleaseFlagBatchTracker; updated NVLink/barrier + scheduling behavior and tail reset delegation.import_topk_reduceexport/function frommega_moe_nvfp4/__init__.py.Fc2OutputRouter) and completion signaling via batch flag tracking.megamoe_constants.py.valid_tokens_in_cta_tile).MEGAMOE_CUTEDSLand wired it through model args/manifests/registries and benchmark harnesses.max_tokens_per_rankbucketing.mpirunprocess-group bootstrapping invariants for EP/WORLD sizing.TRTLLM-12950enforces a minimum MegaMoE group hint (as called out in the provided commit summary).Key checks before approval:
combine_format,in_kernel_fc2_reduce,tactic_autotune, andnum_tokensslicing.default_megamoe_tactic(num_tokens)fallback path).MEGAMOE_CUTEDSL, validate candidate filtering behavior and EP/WORLD invariants in distributed search paths.CI status (from provided run notes):
468c64aandec63e67.a6aa781failed in multiple runs (with and without fail-fast), then a later rerun succeeded.QA Engineer Review
Test list changes
Modified (test-list / waive / QA-list related files):
tests/integration/test_lists/test-db/l0_b200.ymlunittest/_torch/modules/moe/test_megamoe_streaming_load.pytests/integration/test_lists/qa/README.mdaccuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_auto_dtypetests/integration/test_lists/test-db/README.mdtests/integration/test_lists/waives.txtunittest/disaggregated/test_kv_transfer.py SKIP (nvbugs/6403793)unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (nvbugs/6445316)unittest/executor/test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release SKIP (nvbugs/6480574)tests/integration/test_lists/test-db/l0_*.ymland README files intests/integration/test_lists/test-db/(mostly deletions/rewording per diff stats).Verdict for test-list coverage: needs follow-up (test selection/list churn is broader than the new MegaMoE test entry, and additional coverage intent is unclear from provided context).
Test code changes (outside test lists)
New/modified pytest unit tests detected from diff (files under
tests/unittest/**):tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.pytest_initial_streaming_load_layer_atomictest_partial_load_rejected_before_source_materializationtests/unittest/_torch/modules/moe/test_moe_backend.pytest_enumerate_megamoe_candidate_tactics_curated_spacetest_megamoe_cutedsl_tactic_autotune_defaults_offtest_megamoe_cutedsl_tuning_mode_forces_top_maxt_buckettests/integration/test_lists/test-db/in the provided diff):tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytest_msa_metadata_rejects_undersized_max_score_buffertest_msa_proxy_max_score_view_is_contiguous_over_stable_storetest_msa_requires_block_size_128test_resolver_selects_msa_backend_when_availabletests/unittest/_torch/executor/test_kvv2_host_tier_sizing.pytest_aggregate_across_ranks_stays_within_available_memorytest_colocated_ranks_divide_node_memory_budgettest_memlock_limit_caps_quotatest_non_positive_result_falls_back_to_device_quotatest_result_is_always_positivetest_single_rank_uses_device_quota_when_memory_is_ampletest_unknown_limits_fall_back_to_device_quotatests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.pytest_gate_bias_params_maptest_kv_scale_is_not_expanded_when_kv_heads_cover_tptest_load_weights_accepts_base_mapper_without_params_maptest_load_weights_exposes_weight_mappertest_mapper_registration_and_mx_fallbacktest_nvfp4_scale_behavior_is_preservedtest_quant_config_none_is_guardedtest_tp8_mxfp8_duplicates_kv_weight_and_scale_via_callbackstests/unittest/_torch/models/test_minimax_m3.pytest_text_norm_weights_real_loader_smoketests/unittest/_torch/sampler/test_torch_sampler.pytest_out_of_range_decay_params_rejectedtest_reject_speculative_draft_tokenstest_runtime_update_paritytest_strategy_routingtests/unittest/_torch/speculative/test_eagle3.pytest_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31tests/unittest/llmapi/apps/_test_openai_metrics.pytest_metrics_available_before_first_requesttests/unittest/llmapi/apps/_test_openai_prometheus.pytest_kv_cache_metrics_available_before_first_requesttests/unittest/tools/test_config_database_sync.pytest_profile_and_validated_commit_metadatatest_recipe_metadata_validationCoverage / verdict:
test_megamoe_streaming_load.pyis added tol0_b200test list.tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/in the provided diff.