[None][feat] Update CuTeDSL MegaMoE kernels - #16190
Merged
Barry-Delaney merged 10 commits intoJul 29, 2026
Merged
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
Collaborator
|
PR_Github #61875 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61839 [ run ] completed with state |
Collaborator
|
PR_Github #61875 [ 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>
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 27, 2026 17:17
bea62a0 to
6283748
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61953 [ run ] triggered by Bot. Commit: |
![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: 7
🧹 Nitpick comments (8)
tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py (5)
1233-1244: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNit: docstring grammar and
_frozen_naming."Froze at the point calling
_freeze()" reads as a typo — e.g. "Attributes become read-only once_freeze()is called." Also_frozen_(trailing underscore) is unidiomatic;_frozenmatches the rest of the file.🤖 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/epilogue_refactor.py` around lines 1233 - 1244, Update _ImmutableAfterInit’s docstring to clearly state that attributes become read-only when _freeze() is called, and rename the internal _frozen_ flag to _frozen consistently in __setattr__ and _freeze.
49-74: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value
has_side_effects=Trueon a pure conversion blocks CSE/hoisting.
cvt_f32_to_ue8m0_to_f32is a value-in/value-out conversion with no memory or state effects, yet it is marked as having side effects; the siblingmax_abshelpers inQuantImplcorrectly passFalse. It is invoked once per scale block in the mxfp8 loops, so marking it pure lets the compiler dedupe/schedule it.♻️ Proposed change
asm_tmpl, "=r,f", - has_side_effects=True, + has_side_effects=False, is_align_stack=False,🤖 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/epilogue_refactor.py` around lines 49 - 74, Update the llvm.inline_asm call in cvt_f32_to_ue8m0_to_f32 to mark the conversion as side-effect-free by passing False for has_side_effects, matching the pure max_abs helper usage in QuantImpl.
2708-2725: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNit:
packed_i32is recast toFloat32, notInt32.The name says i32 while the recast (and the
reg_tensorfed to the transposes) iscutlass.Float32used purely as a 32-bit bit-carrier for bf16x2. In bit-manipulation code that mismatch is easy to trip over —packed_b32orpacked_bf16x2would read truer.🤖 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/epilogue_refactor.py` around lines 2708 - 2725, Rename packed_i32 to a name that reflects its Float32 bit-carrier role, such as packed_b32 or packed_bf16x2, throughout the transpose setup and autovec_copy calls. Preserve the existing cutlass.Float32 recast and tensor behavior.
97-99: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
lane_idxwith its real type instead ofAny.The comment already documents it as
Int32;Optional[cutlass.Int32]makes the field self-describing and the across-lane precondition in__post_init__easier to read.As per coding guidelines: "avoid
Anyand 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 `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py` around lines 97 - 99, Update the lane_idx annotation in the relevant dataclass to use Optional[cutlass.Int32] instead of Optional[Any], preserving its optional default and existing across-lane behavior. Remove any now-unused Any import if applicable.Source: Coding guidelines
1828-1833: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueConsider asserting the single-block assumption behind
sfc_regs[0].
sfc_regshascute.size(weighted) // sf_vec_sizeentries; taking[0]is only correct while that is exactly 1 (true today: 16 regs / 16). A one-lineconst_exprcheck documents and enforces it.🤖 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/epilogue_refactor.py` around lines 1828 - 1833, In the quantization flow around `quant(weighted, norm_const=norm_const)`, add a `const_expr` assertion that `cute.size(weighted) // sf_vec_size` equals 1 before accessing `sfc_regs[0]`. Preserve the existing single-block behavior while enforcing the assumption explicitly.tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py (2)
350-366: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInitialize
combine_formatin__init__instead of lazily in_setup_attributes.
_setup_attributes()runs per__call__, so thehasattrguard is doing double duty as both "subclass injected it" and "already defaulted". Settingself.combine_format = getattr(self, "combine_format", None) or CombineFormat.parse("bf16")in the constructor (after the subclass hook point) keeps the externally visible member owned by__init__and drops the trace-timehasattr.As per coding guidelines: "initialize externally visible class members in the constructor".
🤖 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/kernel_fc12.py` around lines 350 - 366, Move the default initialization of combine_format from _setup_attributes into the constructor, after the MegaMoE subclass injection hook, using the existing value when present and parsing bf16 only when absent or unset. Then remove the hasattr-based fallback from _setup_attributes while continuing to pass self.combine_format to SwapABSwigluFp4Epilogue.Source: Coding guidelines
799-812: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueValidation reads well; consider
cute.round_upfor the local helper.
cute.round_upis already used elsewhere in this class (_setup_attributes), so the localround_upis a small duplicate. Also, the two nestedcutlass.const_exprguards could collapse into oneand.Also applies to: 846-869
🤖 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/kernel_fc12.py` around lines 799 - 812, Replace the local round_up helper in the enclosing kernel method with the existing cute.round_up utility, matching its use in _setup_attributes. Simplify the nested cutlass.const_expr guards in the affected sections around the tensor layout and related validation into a single guard using and, while preserving the current compile-time conditions and behavior.tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py (1)
774-782: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the bootstrap exception.
except Exceptionhere also swallows programming errors from the import/attribute path and downgrades them to a debug log, leavingtorch.distributedsilently uninitialized so the failure resurfaces later as an opaque rendezvous error. Catch the concrete failures instead.As per coding guidelines: "Catch specific exceptions rather than using broad or bare exception handling such as
except:".♻️ Proposed change
- except Exception as e: # not under MPI either -> leave uninitialized + except (ImportError, RuntimeError, AttributeError) as e: + # not under MPI either -> leave uninitialized logger.debug( f"[MegaMoECuteDsl] MPI rank query failed ({e!r}); " "skipping torch.distributed bootstrap." ) return🤖 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/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py` around lines 774 - 782, Update the MPI query exception handling around mpi_world_size() and mpi_rank() to catch only the concrete exceptions those calls use to indicate that MPI is unavailable, rather than broad Exception. Preserve the debug log and early return for that expected non-MPI case, while allowing programming or import/attribute errors to propagate.Sources: Coding guidelines, Linters/SAST 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 `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py`:
- Around line 2312-2323: Extract the shared STG alignment calculation into a
cached router property named _data_stg_align_bytes, using the existing
data_output element width formula without the redundant nested min. Update both
prefetch and get_data_dst to pass this property as byte_align and assumed_align
respectively, removing their duplicated inline calculations.
- Around line 1045-1048: Update epi_flag_batch handling in the epilogue
initialization and Sm100SwapABSwigluFp4Fc12Kernel.name() so the value used for
kernel behavior matches the cache-key value. Reject values outside [1, 32]
before storing or naming them, or normalize them before name() reads them; do
not silently retain raw out-of-range values that generate distinct keys for
identical kernels.
- Around line 2288-2294: Remove the unreachable self.metadata is None branch
from the destination-pointer logic after prefetch, and de-indent the remaining
metadata-dependent body accordingly. Keep destination address computation
centralized through get_data_dst and preserve the existing metadata path
behavior.
- Around line 491-502: Update _amax_lane’s sf_vec_size == 16 path so all lanes
execute both warp reductions: compute separate masked maxima for the lower and
upper 16-lane halves, then select the result based on first_half. Remove the
divergent if (!first_half) reduction and ensure each block’s amax only includes
its own 16 lanes.
- Around line 2942-2948: Update the copy_bytes calculation in the epilogue copy
path so 4-bit combine formats round the byte count up when copy_elems is odd,
preserving the trailing half-byte element; keep the existing calculation
unchanged for byte-aligned formats and retain the valid_hidden_this_cta_tile
bounds behavior.
- Around line 1976-1997: Add an explicit constructor precondition in
SwapABFc2Epilogue.__init__ (or SwapABSwigluFp4Epilogue.__init__) requiring
token_comm_args when combine_format.is_quantized is true. Reject the invalid
combination before tracing so accesses to token_comm_args.fc2_output_sf and
fc2_done_counter remain safe.
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py`:
- Around line 912-941: Rename the later `intermediate_downproj_padded` binding
used for the fc2 weight-SF view to a distinct descriptive symbol, while
preserving the earlier value used by `fc1_output_sf_gemm_for_fc2_load`. Update
all references in the `tile_atom_to_shape_SF` arguments and the `ValueError`
message to use the new symbol.
---
Nitpick comments:
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py`:
- Around line 1233-1244: Update _ImmutableAfterInit’s docstring to clearly state
that attributes become read-only when _freeze() is called, and rename the
internal _frozen_ flag to _frozen consistently in __setattr__ and _freeze.
- Around line 49-74: Update the llvm.inline_asm call in cvt_f32_to_ue8m0_to_f32
to mark the conversion as side-effect-free by passing False for
has_side_effects, matching the pure max_abs helper usage in QuantImpl.
- Around line 2708-2725: Rename packed_i32 to a name that reflects its Float32
bit-carrier role, such as packed_b32 or packed_bf16x2, throughout the transpose
setup and autovec_copy calls. Preserve the existing cutlass.Float32 recast and
tensor behavior.
- Around line 97-99: Update the lane_idx annotation in the relevant dataclass to
use Optional[cutlass.Int32] instead of Optional[Any], preserving its optional
default and existing across-lane behavior. Remove any now-unused Any import if
applicable.
- Around line 1828-1833: In the quantization flow around `quant(weighted,
norm_const=norm_const)`, add a `const_expr` assertion that `cute.size(weighted)
// sf_vec_size` equals 1 before accessing `sfc_regs[0]`. Preserve the existing
single-block behavior while enforcing the assumption explicitly.
In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py`:
- Around line 350-366: Move the default initialization of combine_format from
_setup_attributes into the constructor, after the MegaMoE subclass injection
hook, using the existing value when present and parsing bf16 only when absent or
unset. Then remove the hasattr-based fallback from _setup_attributes while
continuing to pass self.combine_format to SwapABSwigluFp4Epilogue.
- Around line 799-812: Replace the local round_up helper in the enclosing kernel
method with the existing cute.round_up utility, matching its use in
_setup_attributes. Simplify the nested cutlass.const_expr guards in the affected
sections around the tensor layout and related validation into a single guard
using and, while preserving the current compile-time conditions and behavior.
In `@tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py`:
- Around line 774-782: Update the MPI query exception handling around
mpi_world_size() and mpi_rank() to catch only the concrete exceptions those
calls use to indicate that MPI is unavailable, rather than broad Exception.
Preserve the debug log and early return for that expected non-MPI case, while
allowing programming or import/attribute errors to propagate.
🪄 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: c4ea1b52-e5a8-4e6a-80d6-e1511658bf39
📒 Files selected for processing (32)
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/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
🚧 Files skipped from review as they are similar to previous changes (26)
- tests/integration/test_lists/test-db/l0_b200.yml
- tests/microbenchmarks/bench_moe/search.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
- tests/microbenchmarks/bench_moe/build.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py
- tensorrt_llm/_torch/modules/fused_moe/create_moe.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
- tests/unittest/_torch/modules/moe/test_moe_backend.py
- tests/microbenchmarks/bench_moe/utils.py
- tests/microbenchmarks/bench_moe/backend.py
- tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
- tests/microbenchmarks/bench_moe/timing/autotune.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.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/fc1_fc2_fuse_sched.py
- tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.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/modules/fused_moe/quantization.py
Collaborator
|
PR_Github #61953 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62094 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62094 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62249 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62249 [ run ] completed with state |
juney-nvidia
approved these changes
Jul 29, 2026
Barry-Delaney
enabled auto-merge (squash)
July 29, 2026 02:43
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
7 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
num_tokens), with newdefault_/enumerate_/validate_megamoe_tactic(num_tokens)and stricter 8-field validation (includingflag_batch/epi_flag_batch).tactic_autotune+ explicitAutoTuner.choose_one) while default execution remains deterministic (tactic=-1), including updated runner/op caching and workspace behavior for CUDA-graph capture safety (fenced-reset tracking with capture guards; reduced/targeted workspace zeroing; safer workspace caching keyed to tactic).in_kernel_fc2_reduceandcombine_format(including symmetric memory contract updates such ascombine_outputshape and symm providercombine_k=1).cute.Pointer/byte-base view) and revised region sizing/zero-prefix logic.CombineFormat+ packedTokenSrcMetadata, reworks token-back scheduling with batching viaGpuReleaseFlagBatchTracker, new combine/scratch buffers (combine_sf/fc2_output_sf), and updated NVLink barrier + tail-reset ordering.valid_tokens_in_cta_tile(swap-AB vs non-swap paths) and adjusts related counters/sentinels.output_activation(internally collapsing the combine/topk dimension) and aligns multi-rank launch token bucketing viaset_adaptive_launch_tokens.MEGAMOE_CUTEDSLbackend support across LLM API config, microbenchmark backend registry, and fused-comm scheduling candidate logic.Key review focus areas:
validate_megamoe_tacticenforces the intended field order everywhere.combine_format→ symm provider sizing/contracts → kernel epilogue/quant/dequant paths).QA Engineer Review
Tests list changes (test-db/)
tests/integration/test_lists/test-db/l0_b200.ymlunittest/_torch/modules/moe/test_megamoe_streaming_load.pyVerdict: needs follow-up (CBTS/CI coverage data for the newly added backend/unit tests is not provided here; however the streaming-load test is explicitly added to test-db.)
Test code changes (outside test-list)
Modified/added:
tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.pytest_initial_streaming_load_layer_atomictest_partial_load_rejected_before_source_materializationtests/integration/test_lists/test-db/l0_b200.yml(added entry above).Modified/added:
tests/unittest/_torch/modules/moe/test_moe_backend.pymoe.tactic_autotunedefaulting when env var unsetenumerate_megamoe_candidate_tactics/default_megamoe_tactic/validate_megamoe_tactic(including invalid tactic error)test-dbchange.