Skip to content

[None][feat] Update CuTeDSL MegaMoE kernels#16190

Open
Barry-Delaney wants to merge 10 commits into
NVIDIA:mainfrom
Barry-Delaney:user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt
Open

[None][feat] Update CuTeDSL MegaMoE kernels#16190
Barry-Delaney wants to merge 10 commits into
NVIDIA:mainfrom
Barry-Delaney:user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt

Conversation

@Barry-Delaney

@Barry-Delaney Barry-Delaney commented Jul 9, 2026 •
edited by coderabbitai Bot
Loading

Copy link
Copy Markdown
Collaborator

Dev Engineer Review

  • Updated MegaMoE CuteDSL NVFP4 path with major kernel + custom-op refactors:
    • Lazy-load + register CuteDSL MegaMoE custom op only when required Cutlass/Cute-DSL pieces exist.
    • New token-aware 8-tuple tactic model: default_megamoe_tactic(num_tokens), enumerate_megamoe_candidate_tactics(num_tokens), and stricter validate_megamoe_tactic() arity/field validation.
    • Added API knobs for in_kernel_fc2_reduce, combine_format, tactic_autotune, and optional live num_tokens; removed prior public knobs (token_back_by_dispatch, non_ubulk_fc2_store).
    • Reworked CUDA-graph-safe workspace/fence lifecycle using pointer-keyed fencing, weakref/ABA detection, and autotuner profiling scratch allocation/cleanup.
  • Refactored epilogue + quant/combine pipeline:
    • Introduced centralized quant/combine implementation (QuantImpl) and reorganized fc1/fc2 completion signaling.
    • Rebuilt fc2 routing around a redesigned Fc2OutputRouter and router helpers; removed contract-wrapped tensor plumbing in multiple paths.
    • Added/propagated combine_format through runner/kernel/op layers; updated fake registration signatures accordingly.
  • Refactored token communication for quantized token-back:
    • Added CombineFormat, TokenSrcMetadata, and GpuReleaseFlagBatchTracker.
    • Reworked token-back scheduling/publishing, NVLink barrier behavior, and metadata transport.
  • Kernel/scheduler correctness updates:
    • Standardized token validity fields in scheduler/work-tile + custom extension logic (e.g., valid_tokens_in_cta_tile, plus non-swap_ab cluster tile + counter index handling).
    • Updated low-level sync semantics in grid_sync spin-wait exit fencing and various PTX helper additions.
  • Backend integration + config:
    • Added/registered new MoE backend identifier MEGAMOE_CUTEDSL and connected it through MoE backend registries and model config/manifest enums.
    • Implemented multi-rank adaptive max_tokens_per_rank bucketing and EP/WORLD initialization tightening for MPI under mpirun.
    • Streamed weight loading introduced for NVFP4 MegaMoE CuteDSL with reload safeguards; hot weight reloading explicitly disabled.

Key checks to perform before approval:

  • Verify all API/signature changes are consistently wired across: runner ↔ custom op ↔ kernel ↔ fake registration, especially around combine_format, in_kernel_fc2_reduce, tactic_autotune, and num_tokens.
  • Confirm combine_format value handling is consistent end-to-end (token_comm/token-back + topk-reduce + epilogue) and that the new tactic tuple ordering matches the single source of truth used in autotune/enumeration/validation.
  • Ensure new workspace pointer-typing changes (cute.Pointer vs cute.Tensor) don’t break any callers and that CUDA-graph reuse guards match the intended life-cycle.

CI:

  • Main-branch pipeline 48938 failed (triggered for both commits). Failure details/regressions should be retrieved and addressed before approval.

QA Engineer Review

Test list changes

  • Modified tests/integration/test_lists/test-db/l0_b200.yml
    • Added: unittest/_torch/modules/moe/test_megamoe_streaming_load.py under the MoE backend grouping.

Test code changes (outside test lists)

  • tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
    • Added: test_initial_streaming_load_layer_atomic
    • Added: test_partial_load_rejected_before_source_materialization
  • tests/unittest/_torch/modules/moe/test_moe_backend.py
    • Added: constructor helper coverage + tests for MEGAMOE_TACTIC_AUTOTUNE behavior, default moe.tactic_autotune, and enumerate_megamoe_candidate_tactics/validate_megamoe_tactic.

Coverage / verdict

  • test_megamoe_streaming_load.py is covered by CI via the new l0_b200.yml entry.
  • test_moe_backend.py additions are not listed in tests/integration/test_lists/test-db/ (nor indicated as covered in qa/ in the provided changes).

Verdict: insufficient — CI coverage appears incomplete for the new backend tactic/autotune tests; plus CI failures (48938) require follow-up.

@Barry-Delaney Barry-Delaney self-assigned this Jul 9, 2026
@Barry-Delaney
Barry-Delaney force-pushed the user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt branch 11 times, most recently from d4cb10a to 5cded49 Compare July 17, 2026 04:39
@Barry-Delaney Barry-Delaney added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 17, 2026
@Barry-Delaney
Barry-Delaney force-pushed the user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt branch from 06e6965 to 468c64a Compare July 17, 2026 09:57
@Barry-Delaney
Barry-Delaney marked this pull request as ready for review July 17, 2026 09:57
@Barry-Delaney
Barry-Delaney requested review from a team as code owners July 17, 2026 09:57
@Barry-Delaney
Barry-Delaney requested a review from a team as a code owner July 22, 2026 05:29
@Barry-Delaney

Barry-Delaney commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@coderabbitai

coderabbitai Bot commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

MegaMoE 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.

Changes

MegaMoE CuteDSL runtime

Layer / File(s) Summary
Runtime contracts and tactic selection
tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py, tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/{contract.py,megamoe_constants.py}		 Custom-op arguments, eight-field tactics, token-aware enumeration, workspace caching, profiling scratch, fencing, lazy imports, and runtime contract evaluation are updated.
Kernel workspace and token communication
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/{kernel_fc12.py,megamoe_kernel.py,token_comm.py,sym_buffer.py,flag_batch.py}		 Workspace regions, pointer views, packed metadata, token-back modes, synchronization, release batching, and combine-format payload routing are reworked.
Epilogue pipelines and top-k reduction
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/{epilogue_refactor.py,topk_reduce.py}		 FC1/FC2 epilogues use centralized quantization and routed combine stores, while device-side BF16, MXFP8, and FP4 reductions produce token-hidden outputs.
Backend bucketing and distributed setup
tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py, tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py		 Adaptive token buckets, per-bucket providers, MPI rendezvous, profiling scratch, and direct reduced outputs are wired into backend execution.
Streamed weight lifecycle
tensorrt_llm/_torch/modules/fused_moe/quantization.py		 Streamed NVFP4 source tensors are materialized during loading, checked for complete coverage, transformed into derived weights, and shrunk afterward.
Integration and validation
tensorrt_llm/llmapi/llm_args.py, tests/unittest/_torch/modules/moe/*, tests/microbenchmarks/bench_moe/*		 The new backend is registered across configuration and benchmark paths, and tests cover tactic selection, adaptive tuning, distributed setup, and streamed loading.

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
Loading

Suggested reviewers: xxi-nv

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Description check ⚠️ Warning No PR description was provided, so the required Description, Test Coverage, and PR Checklist sections are missing. Add the template sections with a brief issue/solution summary, relevant tests, and checklist confirmation.
Docstring Coverage ⚠️ Warning Docstring coverage is 41.67% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title matches the main change and follows the repository's [None][feat] summary pattern.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 Checkov (3.3.8)
tensorrt_llm/usage/llm_args_golden_manifest.json

Traceback (most recent call last):
File "/usr/local/bin/checkov", line 2, in
from checkov.main import Checkov
ModuleNotFoundError: No module named 'checkov'

Comment @coderabbitai help to get the list of available commands.

coderabbitai[bot]
coderabbitai Bot reviewed Jul 22, 2026
View reviewed changes

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Apply 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 None for non-returning functions, avoid Any and 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 win

Add 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 None for 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

📥 Commits

Reviewing files that changed from the base of the PR and between 128d020 and ec63e67.

📒 Files selected for processing (33)
  • tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.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/sym_buffer.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py
  • tensorrt_llm/_torch/models/modeling_deepseekv4.py
  • tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tensorrt_llm/llmapi/llm_args.py
  • tensorrt_llm/usage/llm_args_golden_manifest.json
  • tests/integration/test_lists/test-db/l0_b200.yml
  • tests/microbenchmarks/bench_moe/backend.py
  • tests/microbenchmarks/bench_moe/build.py
  • tests/microbenchmarks/bench_moe/search.py
  • tests/microbenchmarks/bench_moe/timing/autotune.py
  • tests/microbenchmarks/bench_moe/utils.py
  • tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
  • tests/unittest/_torch/modules/moe/test_moe_backend.py

Comment thread tests/microbenchmarks/bench_moe/timing/autotune.py Outdated
@tensorrt-cicd

tensorrt-cicd commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #60899 [ run ] triggered by Bot. Commit: ec63e67 Link to invocation

@tensorrt-cicd

tensorrt-cicd commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #60899 [ run ] completed with state FAILURE. Commit: ec63e67
/LLM/main/L0_MergeRequest_PR pipeline #49167 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Barry-Delaney
[TRTLLM-12950][feat] MegaMoE-CuteDSL NVFP4 fused-MoE backend for Deep…
834fde2 
…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>
@Barry-Delaney
[TRTLLM-12950][feat] DeepSeek-V4 wiring and streaming weight load for…
c44b273 
… 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>
@Barry-Delaney
[TRTLLM-12950][test] reproducible MegaMoE-CuteDSL kernel microbenchmark
e0b03d0 
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>
@Barry-Delaney
[TRTLLM-12950][fix] narrow MegaMoE-CuteDSL integration scope
2e76a3e 
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
@Barry-Delaney
[TRTLLM-12950][fix] isolate MegaMoE experimental controls
7211cbb 
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
@Barry-Delaney
[TRTLLM-12950][test] narrow MegaMoE-CuteDSL test scope
b0d9194 
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
@Barry-Delaney
[TRTLLM-12950][fix] suppress MegaMoE Bandit false positives
9470830 
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
@Barry-Delaney
[TRTLLM-12950][test] remove redundant combine format tests
579dcdc 
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
@Barry-Delaney
[None][fix] Resolve DeepSeek V4 mixed-precision base layers
682a8f2 
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>
@Barry-Delaney
[TRTLLM-12950][fix] address MegaMoE review feedback
a6aa781 
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
@Barry-Delaney
Barry-Delaney force-pushed the user/jinshik/dsv4-megamoe-cutedsl-adaptive-maxt branch from 59cc034 to a6aa781 Compare July 23, 2026 02:01
@Barry-Delaney

Barry-Delaney commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61149 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

coderabbitai[bot]
coderabbitai Bot reviewed Jul 23, 2026
View reviewed changes

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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 win

Duplicate sentinel WorkTileInfo construction — extract a shared helper.

The is_swap_ab sentinel-init block in create() (Lines 662-686) is duplicated verbatim in _gen_work_from_cluster_idx() (Lines 1258-1280). Both explicitly document that the field set must track ext.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 both create() 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

📥 Commits

Reviewing files that changed from the base of the PR and between ec63e67 and a6aa781.

📒 Files selected for processing (25)
  • tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.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/sym_buffer.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py
  • tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py
  • tensorrt_llm/_torch/models/modeling_deepseekv4.py
  • tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tensorrt_llm/llmapi/llm_args.py
  • tensorrt_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

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61149 [ run ] completed with state FAILURE. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49397 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Barry-Delaney

Barry-Delaney commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61269 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61269 [ run ] completed with state FAILURE. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49504 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Barry-Delaney

Barry-Delaney commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61294 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61294 [ run ] completed with state FAILURE. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49527 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Barry-Delaney

Barry-Delaney commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61374 [ run ] triggered by Bot. Commit: a6aa781 Link to invocation

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #61374 [ run ] completed with state SUCCESS. Commit: a6aa781
/LLM/main/L0_MergeRequest_PR pipeline #49602 completed with status: 'SUCCESS'

CI Report

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@coderabbitai coderabbitai[bot] coderabbitai[bot] left review comments
@xxi-nv xxi-nv xxi-nv approved these changes
@yufeiwu-nv yufeiwu-nv yufeiwu-nv approved these changes
@arysef arysef Awaiting requested review from arysef arysef is a code owner automatically assigned from NVIDIA/trt-llm-usage-telemetry-devs
@QiJune QiJune Awaiting requested review from QiJune QiJune is a code owner automatically assigned from NVIDIA/trt-llm-noncommitted-api-review-committee
@lori-ren lori-ren Awaiting requested review from lori-ren lori-ren is a code owner automatically assigned from NVIDIA/trt-llm-runtime-devs
@bo-nv bo-nv Awaiting requested review from bo-nv bo-nv is a code owner automatically assigned from NVIDIA/trt-llm-disagg-devs
@nv-xtf nv-xtf Awaiting requested review from nv-xtf nv-xtf is a code owner automatically assigned from NVIDIA/trt-llm-disagg-devs
@sunnyqgg sunnyqgg Awaiting requested review from sunnyqgg sunnyqgg is a code owner automatically assigned from NVIDIA/trt-llm-moe-devs
@rosong11 rosong11 Awaiting requested review from rosong11 rosong11 is a code owner automatically assigned from NVIDIA/trt-llm-moe-devs
@mzweilz mzweilz Awaiting requested review from mzweilz mzweilz is a code owner automatically assigned from NVIDIA/trt-llm-infra-devs
@mlefeb01 mlefeb01 Awaiting requested review from mlefeb01 mlefeb01 is a code owner automatically assigned from NVIDIA/trt-llm-infra-devs
@lfr-0531 lfr-0531 Awaiting requested review from lfr-0531
@mikeiovine mikeiovine Awaiting requested review from mikeiovine

At least 1 approving review is required to merge this pull request.

Assignees

@Barry-Delaney Barry-Delaney

Labels

api-compatible Accepted LLM API contract change that is backwards-compatible

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@Barry-Delaney @tensorrt-cicd @xxi-nv @yufeiwu-nv