Add MoE + expert-parallel (MoeEp) Python API with MegaMoE CuTe DSL backend - #448
Add MoE + expert-parallel (MoeEp) Python API with MegaMoE CuTe DSL backend#448mhoqueanik wants to merge 10 commits into
Conversation
…tests MoeEp exposes the constructor/forward contract for the fused SwiGLU MoE with EP dispatch (bf16, mxfp8, and nvfp4 public outputs); forward allocates the output representation until the device kernel lands. The pure-PyTorch MoeEpReference implements the full semantics — routing, variable-size all-to-all dispatch, local experts, Form-A top-k combine, and block-scaled quantization — and is validated by test/python/fe_api/moe_ep against naive per-token references, including 2-rank (gloo) and 4-GPU (NCCL) expert-parallel runs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The doc predated the training-integration work: it still listed "no backward API" and "no auxiliary FC1 capture" as first-version boundaries and described the reference as inference-forward only. Refresh the status line, return-value contract, reference-scope notes, and validation matrix to match MoeEpReference.backward, the generate_c=True (output, fc1_c, route_metadata) return, and the current test suite. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The proposed-API section only showed the constructor and __call__; add the backward signature, its argument/return tables, and the collective requirement. Give MoeEp a matching allocation-only backward stub (validating the generate_c stash shapes) with a passing allocation test and a strict-xfail gate against MoeEpReference.backward. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MoeEp.__call__/backward now delegate to the megamoe training package (SM100 MegaMoE MXFP8 mega-kernel + bwd_impl="mega" backward) when the configuration is supported, falling back to the original allocate-only stubs otherwise (CUDNN_MOE_EP_BACKEND=auto|megamoe|none, CUDNN_MEGAMOE_ROOT locates the package). The backend adapts the public contract to the kernel: weight layout transpose + gate/up half swap, logical fc1_c/route_metadata extraction from the kernel pools (32-block de-interleave, contract row order), reference-bit-exact host quantization for mxfp8/nvfp4 output_format, and fp32 gradient conversion from mega_backward. megamoe_backend_parity.py validates MoeEp-on-megamoe against MoeEpReference at kernel shapes (single-rank and torchrun EP). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdds the public ChangesMoE Expert Parallel API
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (3)
test/python/fe_api/moe_ep/test_moe_ep.py (2)
847-855: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider promoting this out of L0.
The module-level
pytestmark = pytest.mark.L0(line 19) applies here too, but a 4-rank NCCL spawn is not a fast test. Overriding it with a higher level keeps L0 quick.As per coding guidelines, "keep
L0tests fast and place large parameter sweeps at higher levels."♻️ Proposed tweak
+@pytest.mark.L2 `@pytest.mark.skipif`( not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < 4, reason="requires NCCL and at least 4 GPUs", ) def test_four_rank_expert_parallel(tmp_path):🤖 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 `@test/python/fe_api/moe_ep/test_moe_ep.py` around lines 847 - 855, Override the module-level L0 marker on test_four_rank_expert_parallel with the appropriate higher-level pytest marker, while preserving its existing skip conditions and four-rank NCCL behavior.Source: Coding guidelines
216-216: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePin the accumulator's device.
torch.zeros(...)defaults to CPU, so this shared helper silently breaks the moment it's called with CUDA inputs.device=a.devicecosts nothing.♻️ Proposed tweak
- acc = torch.zeros(a.shape[1], dtype=torch.float32) + acc = torch.zeros(a.shape[1], dtype=torch.float32, device=a.device)🤖 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 `@test/python/fe_api/moe_ep/test_moe_ep.py` at line 216, Update the accumulator initialization in the shared helper to create the tensor on the same device as input tensor a by reusing a.device, while preserving its shape and float32 dtype.python/cudnn/moe_ep/_megamoe.py (1)
214-271: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated quantizer with only its MXFP8 branch pinned. The backend re-implements the reference
quantize_blockwise, and the only bit-exactness check covers MXFP8, leaving the NVFP4 nibble path free to drift from the reference.
python/cudnn/moe_ep/_megamoe.py#L214-L271: import/share the reference quantizer instead of keeping a hand copy of_quantize_output/_nearest_e2m1_codes(or move the single implementation into a non-test module both sides can use).test/python/fe_api/moe_ep/megamoe_backend_parity.py#L139-L147: add an NVFP4 case alongside the MXFP8torch.equalcheck so both branches are verified.🤖 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 `@python/cudnn/moe_ep/_megamoe.py` around lines 214 - 271, Replace the duplicated _quantize_output and _nearest_e2m1_codes implementation in python/cudnn/moe_ep/_megamoe.py:214-271 with the shared reference quantizer, or move that single implementation into a non-test module used by both callers. In test/python/fe_api/moe_ep/megamoe_backend_parity.py:139-147, add an NVFP4 torch.equal parity assertion alongside the existing MXFP8 check; no other site changes are required.
🤖 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/fe-oss-apis/moe_ep.md`:
- Around line 3-8: The status paragraph in the MoeEp API documentation is
outdated because MoeEp.__call__ now uses the MegaMoE device backend when
configured. Update the forward and backward descriptions to distinguish the
backend path from the fallback implementation, document CUDNN_MOE_EP_BACKEND and
CUDNN_MEGAMOE_ROOT as selectors, and retain the uninitialized-storage/no-kernel
behavior only for the fallback path.
In `@python/cudnn/moe_ep/_megamoe.py`:
- Around line 336-345: Rename the ambiguous I local in _to_kernel_weights and
the corresponding locals at the referenced locations to inter or intermediate,
updating every use consistently while preserving the existing tensor slicing and
shape transformations.
- Around line 400-412: The backend selection and _extract_stash collective path
must use a consistent decision across all EP ranks. When ep_group is set, make
maybe_create’s availability result collective by all-reducing the per-rank
success flag, and ensure every rank either selects the same backend or falls
back together before executing the two collectives in _extract_stash. Preserve
the existing behavior for non-distributed execution.
- Around line 454-462: Update the backward path around the visible fc1_c and
route_metadata parameters to honor the module docstring: extract the current
pool stash and cheaply validate both arguments, checking fc1_c shape and
route_metadata equality before computing gradients. Reject mismatches rather
than silently ignoring stale forward arguments; keep the pool stash as the
source of truth after validation.
- Around line 132-140: Replace tempfile.mktemp in the process-group
initialization with tempfile.mkdtemp to create a private 0700 temporary
directory, and pass that directory path to dist.FileStore while preserving the
existing prefix and initialization behavior.
- Around line 209-211: Update _weight_signature to use a tensor-object identity
or weak-reference-based invalidation signal rather than only data_ptr() and
_version, including the underlying data tensor for BlockScaledTensor. Ensure
replacing a weight with a fresh tensor cannot match the prior signature, so
_sync_weights refreshes kernel weights.
In `@python/cudnn/moe_ep/api.py`:
- Around line 139-180: The MoE EP API lacks the required frontend-only
scaffolding. In python/cudnn/moe_ep/api.py lines 139-180, make MoeEp subclass
APIBase and implement check_support() returning bool while setting
self._is_supported, compile(), execute(..., current_stream=None), and
moe_ep_wrapper() that allocates outputs and returns a TupleDict with documented
key order; alternatively document the proposal-stage deviation in
docs/fe-oss-apis/moe_ep.md. In python/cudnn/moe_ep/__init__.py lines 4-6, export
moe_ep_wrapper alongside MoeEp through __all__.
- Around line 386-398: Update MoeEp.backward to resolve or validate _backend
before dispatching to _backend.backward, rather than treating an unresolved
backend as a valid path. Reuse the same backend-resolution logic as __call__,
ensuring backward either invokes the available MegaMoE backend or raises clearly
when it cannot be resolved, never returning uninitialized gradients.
- Around line 63-98: Strengthen BlockScaledTensor.__post_init__ to validate data
and scale shapes against the normalized logical_shape/axis, require the
supported scale dtype, and ensure data and scale are on the same device. Update
dequantize’s NVFP4 unpacking path to reinterpret self.data as torch.uint8 before
applying nibble bit operations, preserving support for both uint8 and
torch.float4_e2m1fn_x2 payloads.
- Around line 331-337: Update the backend initialization flow in the enclosing
forward method so `_MegamoeBackend` is not permanently sized from the first
call’s token_count when max_tokens_per_rank is unset. Cache or recreate backends
keyed by both device and token_count, or otherwise initialize them using an
explicit upper bound, while preserving reuse for compatible calls and ensuring
later larger batches do not exceed _max_tokens.
In `@test/python/fe_api/moe_ep/megamoe_backend_parity.py`:
- Around line 115-121: Update the parity failure aggregation near the existing
fc1_c diagnostics so fc1_c shape mismatches and relative-error violations
contribute to failures alongside the other forward and metadata checks. Reuse
the established tolerance and failure-reporting pattern, ensuring an incorrect
fc1_c stash causes the test to report failure rather than only printing inf.
In `@test/python/fe_api/moe_ep/test_moe_ep.py`:
- Around line 498-520: Add a module- or test-level CUDA availability gate for
the GPU-dependent Moe EP tests, including the strict xfail variants, using the
project’s supported capability checks so CPU-only runners skip them instead of
passing spuriously. Apply the gate around the relevant tests in test_moe_ep.py
while preserving their existing assertions and xfail behavior on supported CUDA
environments.
---
Nitpick comments:
In `@python/cudnn/moe_ep/_megamoe.py`:
- Around line 214-271: Replace the duplicated _quantize_output and
_nearest_e2m1_codes implementation in python/cudnn/moe_ep/_megamoe.py:214-271
with the shared reference quantizer, or move that single implementation into a
non-test module used by both callers. In
test/python/fe_api/moe_ep/megamoe_backend_parity.py:139-147, add an NVFP4
torch.equal parity assertion alongside the existing MXFP8 check; no other site
changes are required.
In `@test/python/fe_api/moe_ep/test_moe_ep.py`:
- Around line 847-855: Override the module-level L0 marker on
test_four_rank_expert_parallel with the appropriate higher-level pytest marker,
while preserving its existing skip conditions and four-rank NCCL behavior.
- Line 216: Update the accumulator initialization in the shared helper to create
the tensor on the same device as input tensor a by reusing a.device, while
preserving its shape and float32 dtype.
🪄 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: 68f67f5b-0bc8-4458-96bf-e5755031f123
📒 Files selected for processing (9)
docs/fe-oss-apis/moe_ep.mddocs/fe-oss-apis/overview.mdpython/cudnn/__init__.pypython/cudnn/moe_ep/__init__.pypython/cudnn/moe_ep/_megamoe.pypython/cudnn/moe_ep/api.pytest/python/fe_api/moe_ep/megamoe_backend_parity.pytest/python/fe_api/moe_ep/moe_ep_reference.pytest/python/fe_api/moe_ep/test_moe_ep.py
| Status: public API stub, implementation proposal, and executable PyTorch | ||
| reference for both forward and backward. The API currently allocates | ||
| uninitialized output storage without launching a device kernel. Its numerical | ||
| comparison is therefore marked as a strict expected failure under | ||
| `test/python/fe_api/moe_ep`; remove that marker when backend execution is | ||
| connected. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Status paragraph is stale relative to the MegaMoE backend in this PR.
MoeEp.__call__ now dispatches to the MegaMoE device backend when available (api.py lines 157-164, selected via CUDNN_MOE_EP_BACKEND / CUDNN_MEGAMOE_ROOT), so "allocates uninitialized output storage without launching a device kernel" is only the fallback path. Same for the backward note at lines 156-160. Consider documenting the backend-selection env vars and the fallback semantics here.
🤖 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/fe-oss-apis/moe_ep.md` around lines 3 - 8, The status paragraph in the
MoeEp API documentation is outdated because MoeEp.__call__ now uses the MegaMoE
device backend when configured. Update the forward and backward descriptions to
distinguish the backend path from the fallback implementation, document
CUDNN_MOE_EP_BACKEND and CUDNN_MEGAMOE_ROOT as selectors, and retain the
uninitialized-storage/no-kernel behavior only for the fallback path.
Source: Path instructions
| import tempfile | ||
|
|
||
| store_path = tempfile.mktemp(prefix="cudnn_moe_ep_pg_") | ||
| dist.init_process_group( | ||
| "nccl", | ||
| store=dist.FileStore(store_path, 1), | ||
| world_size=1, | ||
| rank=0, | ||
| ) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Replace tempfile.mktemp with a private temp directory.
mktemp is deprecated and racy in shared /tmp (another process can pre-create/symlink the path before FileStore opens it). A mkdtemp directory keeps the same FileStore semantics with 0700 ownership.
🔒 Proposed fix
- import tempfile
+ import tempfile
- store_path = tempfile.mktemp(prefix="cudnn_moe_ep_pg_")
+ store_path = os.path.join(tempfile.mkdtemp(prefix="cudnn_moe_ep_pg_"), "store")
dist.init_process_group(📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import tempfile | |
| store_path = tempfile.mktemp(prefix="cudnn_moe_ep_pg_") | |
| dist.init_process_group( | |
| "nccl", | |
| store=dist.FileStore(store_path, 1), | |
| world_size=1, | |
| rank=0, | |
| ) | |
| import tempfile | |
| store_path = os.path.join(tempfile.mkdtemp(prefix="cudnn_moe_ep_pg_"), "store") | |
| dist.init_process_group( | |
| "nccl", | |
| store=dist.FileStore(store_path, 1), | |
| world_size=1, | |
| rank=0, | |
| ) |
🧰 Tools
🪛 ast-grep (0.45.0)
[warning] 133-133: The function mktemp is deprecated. When using this function, it is possible for an attacker to modify the created file before the filename is returned. Use NamedTemporaryFile() instead and pass it the delete=False parameter.
Context: tempfile.mktemp(prefix="cudnn_moe_ep_pg_")
Note: [CWE-377]: Insecure Temporary File [OWASP A01:2021]: Broken Access Control
(avoid-mktemp-python)
🪛 Ruff (0.16.0)
[error] 134-134: Use of insecure and deprecated function (mktemp)
(S306)
🤖 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 `@python/cudnn/moe_ep/_megamoe.py` around lines 132 - 140, Replace
tempfile.mktemp in the process-group initialization with tempfile.mkdtemp to
create a private 0700 temporary directory, and pass that directory path to
dist.FileStore while preserving the existing prefix and initialization behavior.
Source: Linters/SAST tools
| def _weight_signature(tensor: MoeTensor) -> Tuple: | ||
| data = tensor.data if isinstance(tensor, BlockScaledTensor) else tensor | ||
| return (data.data_ptr(), data._version) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== File map ==\n'
git ls-files 'python/cudnn/moe_ep/_megamoe.py' 'python/cudnn/moe_ep/*.py' 'test/python/fe_api/*' | sed 's#^`#-` #'
printf '\n== Relevant lines in python/cudnn/moe_ep/_megamoe.py ==\n'
nl -ba python/cudnn/moe_ep/_megamoe.py | sed -n '180,280p'
printf '\n== Search for weight signature / sync logic ==\n'
rg -n "_weight_signature|_sync_weights|data_ptr\(|_version|weight cache|cache key" python/cudnn/moe_ep -S
printf '\n== Related tensor replacement / weakref usage ==\n'
rg -n "weakref|id\(|data_ptr\(|_version" python/cudnn/moe_ep -S
printf '\n== Test coverage nearby ==\n'
rg -n "_weight_signature|_sync_weights|megamoe|moe_ep" test/python -SRepository: NVIDIA/cudnn-frontend
Length of output: 3673
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== python/cudnn/moe_ep/_megamoe.py (selected range) ==\n'
sed -n '180,280p' python/cudnn/moe_ep/_megamoe.py | cat -n
printf '\n== Full outline of python/cudnn/moe_ep/_megamoe.py ==\n'
python3 - <<'PY'
from pathlib import Path
p = Path("python/cudnn/moe_ep/_megamoe.py")
for i, line in enumerate(p.read_text().splitlines(), 1):
if "_weight_signature" in line or "_sync_weights" in line or "cache" in line or "weight" in line:
print(f"{i}: {line}")
PY
printf '\n== moe_ep tests and references ==\n'
sed -n '1,260p' test/python/fe_api/moe_ep/test_moe_ep.py | cat -n
printf '\n== backend parity test ==\n'
sed -n '1,260p' test/python/fe_api/moe_ep/megamoe_backend_parity.py | cat -n
printf '\n== Search for tensor replacement / identity assumptions in moe_ep ==\n'
rg -n "reassign|replace|new_tensor|weakref|data_ptr\(|_version|signature|sync" python/cudnn/moe_ep test/python/fe_api/moe_ep -SRepository: NVIDIA/cudnn-frontend
Length of output: 24711
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== _MegamoeBackend weight sync implementation ==\n'
sed -n '286,380p' python/cudnn/moe_ep/_megamoe.py | cat -n
printf '\n== Forward/backward call sites around sync ==\n'
sed -n '380,490p' python/cudnn/moe_ep/_megamoe.py | cat -n
printf '\n== Search for layer.refresh_weights and weight mutation semantics ==\n'
rg -n "refresh_weights|_weight_sig|_layer =|build.*weights|setattr\(.*weight|weight =" python/cudnn/moe_ep -S
printf '\n== Public API entrypoints for MoeEp ==\n'
sed -n '1,220p' python/cudnn/moe_ep/__init__.py | cat -n
sed -n '1,240p' python/cudnn/moe_ep/api.py | cat -nRepository: NVIDIA/cudnn-frontend
Length of output: 21744
🏁 Script executed:
python3 - <<'PY'
import gc
import torch
print("torch:", torch.__version__)
x = torch.empty(8)
print("initial version:", x._version, "ptr:", x.data_ptr())
x.add_(1)
print("after in-place add_ version:", x._version, "ptr:", x.data_ptr())
# Try to observe pointer reuse for fresh tensors of the same size.
seen = {}
reuse = []
for i in range(20000):
t = torch.empty(1024)
ptr = t.data_ptr()
if ptr in seen:
reuse.append((i, seen[ptr], ptr, t._version))
break
seen[ptr] = i
del t
if i % 100 == 0:
gc.collect()
print("pointer reuse observed:", bool(reuse))
if reuse:
i, j, ptr, ver = reuse[0]
print("first reuse at iter", i, "previous iter", j, "ptr", ptr, "version", ver)
# Show that a brand-new tensor typically starts at version 0.
y = torch.empty(8)
print("fresh tensor version:", y._version)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 277
🏁 Script executed:
python3 - <<'PY'
import gc
import torch
print("torch:", torch.__version__)
x = torch.empty(8)
print("initial version:", x._version, "ptr:", x.data_ptr())
x.add_(1)
print("after in-place add_ version:", x._version, "ptr:", x.data_ptr())
seen = {}
reuse = []
for i in range(20000):
t = torch.empty(1024)
ptr = t.data_ptr()
if ptr in seen:
reuse.append((i, seen[ptr], ptr, t._version))
break
seen[ptr] = i
del t
if i % 100 == 0:
gc.collect()
print("pointer reuse observed:", bool(reuse))
if reuse:
i, j, ptr, ver = reuse[0]
print("first reuse at iter", i, "previous iter", j, "ptr", ptr, "version", ver)
y = torch.empty(8)
print("fresh tensor version:", y._version)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 277
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== BlockScaledTensor / MoeTensor usage ==\n'
rg -n "BlockScaledTensor|MoeTensor|dequantize\(" python/cudnn/moe_ep test/python/fe_api -S
printf '\n== weight-related docs and examples ==\n'
sed -n '1,220p' test/python/fe_api/moe_ep/moe_ep_reference.py | cat -n
sed -n '1,240p' python/cudnn/moe_ep/api.py | cat -n
printf '\n== any tests exercising repeated calls with changing weights ==\n'
rg -n "forward\(.*fc1_weight|backward\(.*fc1_weight|same tensor|reassign|clone\(\)|copy_\(" test/python/fe_api/moe_ep -SRepository: NVIDIA/cudnn-frontend
Length of output: 28377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '380,460p' test/python/fe_api/moe_ep/test_moe_ep.py | cat -nRepository: NVIDIA/cudnn-frontend
Length of output: 3112
🌐 Web query:
PyTorch tensor _version in-place operations version counter and caching allocator pointer reuse data_ptr reused same address
💡 Result:
In PyTorch, the relationship between tensor versioning and memory management is critical for autograd safety, but they operate as distinct mechanisms [1][2]. Tensor Versioning and Autograd Every tensor maintains a version counter, accessible via ._version, which is tracked by the autograd engine [1][3]. When an in-place operation is performed on a tensor, its version counter is incremented [1][3]. This mechanism allows the autograd engine to detect stale data [1]. If a SavedVariable (used during the backward pass) holds a reference to a tensor, the engine compares the saved version with the current version of the tensor; a mismatch indicates that an illegal in-place modification occurred, triggering an error [1][3]. If you perform custom in-place operations that bypass standard PyTorch functions (such as modifying memory directly via a custom kernel using the tensor's data_ptr), PyTorch cannot automatically detect the change [4][5]. To maintain autograd integrity in these cases, you must manually increment the version counter using torch.autograd.graph.increment_version(tensor) [4][5]. Note that torch.compile may sometimes produce different version counter behaviors compared to eager mode, as the compiler may fuse operations in ways that do not map 1:1 to eager-mode counter increments [6]. Caching Allocator and Memory Reuse The PyTorch Caching Allocator (e.g., CUDACachingAllocator) manages memory independently of the autograd version counter [7][8]. It caches and reuses memory blocks to avoid the overhead of frequent cudaMalloc/cudaFree calls [9][8]. When a tensor is deleted or its storage is freed, the memory address (the data_ptr) is returned to the allocator's pool and may be assigned to a new tensor request [7][9]. It is entirely normal for a newly allocated tensor to reuse the same memory address (data_ptr) as a previously freed tensor [9][8]. The autograd system does not rely on the memory address to determine if a tensor has been modified; it relies strictly on the version counter and associated metadata [1][2]. Therefore, the reuse of a data_ptr does not cause autograd to incorrectly identify a new tensor as being a modified version of an old one [1]. Autograd effectively tracks the identity of the tensor object and its specific version state, ensuring that memory reuse by the allocator remains transparent to the functional correctness of the autograd graph [1][3].
Citations:
- 1: https://github.com/pytorch/pytorch/blob/17776db2ee8cb9248233c22f33810345859352b8/torch/csrc/autograd/variable.h
- 2: https://discuss.pytorch.org/t/solved-torch-tensor-version-why-1-for-class-attributes-e-g-self-x/41114
- 3: https://stackoverflow.com/questions/68326500/getting-the-autograd-counter-of-a-tensor-in-pytorch
- 4: https://docs.pytorch.org/docs/stable/generated/torch.autograd.graph.increment_version.md
- 5: https://docs.pytorch.org/docs/2.11/generated/torch.autograd.graph.increment_version.html
- 6: Inconsistent Tensor._version Behaviour with torch.compile() pytorch/pytorch#135441
- 7: https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp
- 8: https://docs.pytorch.org/devlogs/eager/2026-06-01-cuda-caching-allocator/
- 9: https://github.com/pytorch/pytorch/blob/ec673ecd/torch/cuda/memory.py
_weight_signature can skip required weight refreshes.
(data_ptr, _version) is not a stable cache key across tensor replacement: a fresh weight tensor can come back at the same address with version 0, match the previous signature, and leave _sync_weights using stale kernel weights. Use tensor object identity/weakrefs or another stronger invalidation signal here.
🤖 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 `@python/cudnn/moe_ep/_megamoe.py` around lines 209 - 211, Update
_weight_signature to use a tensor-object identity or weak-reference-based
invalidation signal rather than only data_ptr() and _version, including the
underlying data tensor for BlockScaledTensor. Ensure replacing a weight with a
fresh tensor cannot match the prior signature, so _sync_weights refreshes kernel
weights.
| def _to_kernel_weights(self, fc1_weight: MoeTensor, fc2_weight: MoeTensor): | ||
| """(E,H,2I) gate-first / (E,I,H) -> w13 (E,2I,H) up-first / w2 (E,H,I).""" | ||
|
|
||
| I = self._op.intermediate_size | ||
| fc1 = _dequant_to_bf16(fc1_weight) | ||
| gate = fc1[..., :I].transpose(1, 2) | ||
| up = fc1[..., I:].transpose(1, 2) | ||
| w13 = torch.cat([up, gate], dim=1).contiguous() | ||
| w2 = _dequant_to_bf16(fc2_weight).transpose(1, 2).contiguous() | ||
| return w13, w2 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous I locals.
Ruff flags I as an ambiguous variable name (E741) here and at lines 397 and 477. inter / intermediate reads better and clears the lint.
🧰 Tools
🪛 Ruff (0.16.0)
[error] 339-339: Ambiguous variable name: I
(E741)
🤖 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 `@python/cudnn/moe_ep/_megamoe.py` around lines 336 - 345, Rename the ambiguous
I local in _to_kernel_weights and the corresponding locals at the referenced
locations to inter or intermediate, updating every use consistently while
preserving the existing tensor slicing and shape transformations.
Source: Linters/SAST tools
| if world > 1: | ||
| t_all = torch.full((world,), T, dtype=torch.int64, device=device) | ||
| dist.all_gather_into_tensor( | ||
| t_all, | ||
| torch.tensor([T], dtype=torch.int64, device=device), | ||
| group=op.ep_group, | ||
| ) | ||
| if not bool((t_all == T).all().item()): | ||
| raise RuntimeError("megamoe backend requires the same token count on every " f"EP rank, got {t_all.tolist()}") | ||
| ids_all = torch.empty((world, T, K), dtype=topk_idx.dtype, device=device) | ||
| dist.all_gather_into_tensor(ids_all, topk_idx.contiguous(), group=op.ep_group) | ||
| else: | ||
| ids_all = topk_idx.view(1, T, K) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Verify the collective ordering constraint is safe.
_extract_stash runs two collectives, but only when op.generate_c is set (via forward). Any rank that takes the non-generate_c path, or that fell back to the allocate-only backend because maybe_create returned None, will not participate — an asymmetric fallback deadlocks the group instead of raising. Since maybe_create swallows all failures per-rank under the default auto policy, that asymmetry is reachable. Consider making the backend decision itself collective (all-reduce the availability flag) when ep_group is not None.
🤖 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 `@python/cudnn/moe_ep/_megamoe.py` around lines 400 - 412, The backend
selection and _extract_stash collective path must use a consistent decision
across all EP ranks. When ep_group is set, make maybe_create’s availability
result collective by all-reducing the per-rank success flag, and ensure every
rank either selects the same backend or falls back together before executing the
two collectives in _extract_stash. Preserve the existing behavior for
non-distributed execution.
| class MoeEp: | ||
| """Fused SwiGLU MoE operator with contiguous expert parallel sharding. | ||
|
|
||
| Global expert ``e`` belongs to group-relative EP rank | ||
| ``e // experts_per_rank``. The constructor captures static configuration; | ||
| calling the instance accepts runtime tensors for this rank. | ||
|
|
||
| With ``generate_c=True`` (training integration), ``__call__`` additionally | ||
| returns ``fc1_c`` and ``route_metadata``. ``fc1_c`` is the raw pre-SwiGLU | ||
| FC1 accumulator for every route this rank's experts processed, BF16, shape | ||
| ``(local_routes, 2 * intermediate)``. Rows are grouped by local expert | ||
| (ascending) and ordered within each expert by source rank, then the source | ||
| rank's token-major route order. The rows are captured before the gate/up | ||
| clamp and carry no router weight. ``route_metadata`` is Int32 | ||
| ``(local_routes, 4)`` with columns | ||
| ``(local_expert, src_rank, src_token, src_slot)``, row-aligned with | ||
| ``fc1_c``, identifying each route for the backward gradient re-dispatch. | ||
|
|
||
| Backend execution: when the MegaMoE device backend is available and the | ||
| configuration is supported (see ``moe_ep._megamoe``), ``__call__`` and | ||
| ``backward`` launch the fused SM100 kernels and return real results. | ||
| Otherwise they fall back to returning newly allocated, uninitialized | ||
| output storage with the correct public representation (the original | ||
| API-stub behavior). ``CUDNN_MOE_EP_BACKEND=megamoe|auto|none`` selects | ||
| the policy; ``CUDNN_MEGAMOE_ROOT`` locates the megamoe package. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| *, | ||
| num_experts: int, | ||
| hidden_size: int, | ||
| intermediate_size: int, | ||
| top_k: int, | ||
| ep_group: Optional[dist.ProcessGroup] = None, | ||
| max_tokens_per_rank: Optional[int] = None, | ||
| output_format: Union[MoeFormat, str] = MoeFormat.BF16, | ||
| combine_format: Union[MoeFormat, str] = MoeFormat.BF16, | ||
| apply_topk_in_fc1: bool = True, | ||
| gate_up_clamp: Optional[float] = None, | ||
| generate_c: bool = False, | ||
| ) -> None: |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
MoE EP package does not follow the frontend-only OSS API contract. The root cause is a single missing piece of scaffolding: no APIBase subclass and no <operation>_wrapper(), which then leaves the family __init__.py with nothing to export beyond the class.
python/cudnn/moe_ep/api.py#L139-L180: makeMoeEpsubclassAPIBasewithcheck_support() -> bool(settingself._is_supported),compile(), andexecute(..., current_stream=None), and add amoe_ep_wrapper()that allocates outputs and returns aTupleDictwith a documented key order — or document the proposal-stage deviation indocs/fe-oss-apis/moe_ep.md.python/cudnn/moe_ep/__init__.py#L4-L6: once the wrapper exists, export it alongsideMoeEpin__all__.
As per coding guidelines: "Every frontend-only OSS kernel API must subclass APIBase, implement check_support() -> bool, compile(), and execute(..., current_stream=None)" and "Kernel family __init__.py files must export the API class and wrapper through __all__."
📍 Affects 2 files
python/cudnn/moe_ep/api.py#L139-L180(this comment)python/cudnn/moe_ep/__init__.py#L4-L6
🤖 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 `@python/cudnn/moe_ep/api.py` around lines 139 - 180, The MoE EP API lacks the
required frontend-only scaffolding. In python/cudnn/moe_ep/api.py lines 139-180,
make MoeEp subclass APIBase and implement check_support() returning bool while
setting self._is_supported, compile(), execute(..., current_stream=None), and
moe_ep_wrapper() that allocates outputs and returns a TupleDict with documented
key order; alternatively document the proposal-stage deviation in
docs/fe-oss-apis/moe_ep.md. In python/cudnn/moe_ep/__init__.py lines 4-6, export
moe_ep_wrapper alongside MoeEp through __all__.
Source: Coding guidelines
| if not self._backend_resolved: | ||
| from . import _megamoe | ||
|
|
||
| self._backend = _megamoe.maybe_create(self, device, token_count) | ||
| self._backend_resolved = True | ||
| if self._backend is not None: | ||
| return self._backend.forward(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -t f '_megamoe.py' python/cudnn | xargs -r ast-grep outline --items all
rg -nP -C8 'def maybe_create' python/cudnnRepository: NVIDIA/cudnn-frontend
Length of output: 2533
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## _megamoe.py around maybe_create and backend methods\n'
sed -n '170,360p' python/cudnn/moe_ep/_megamoe.py
printf '\n## moe_ep/api.py around backend resolution\n'
sed -n '250,360p' python/cudnn/moe_ep/api.pyRepository: NVIDIA/cudnn-frontend
Length of output: 12743
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('python/cudnn/moe_ep/_megamoe.py')
lines = p.read_text().splitlines()
for i in range(178, 361):
if i <= len(lines):
if i in range(179, 204) or i in range(274, 361):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 4980
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C6 'self\._device|def forward|_max_tokens|token_count' python/cudnn/moe_ep/_megamoe.pyRepository: NVIDIA/cudnn-frontend
Length of output: 3616
Backend size is frozen on the first call. maybe_create() builds _MegamoeBackend with that call’s token_count when max_tokens_per_rank is unset, and forward() fails once a later batch exceeds _max_tokens. Cache per (device, token_count) or size the backend from an explicit upper bound.
🤖 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 `@python/cudnn/moe_ep/api.py` around lines 331 - 337, Update the backend
initialization flow in the enclosing forward method so `_MegamoeBackend` is not
permanently sized from the first call’s token_count when max_tokens_per_rank is
unset. Cache or recreate backends keyed by both device and token_count, or
otherwise initialize them using an explicit upper bound, while preserving reuse
for compatible calls and ensuring later larger batches do not exceed
_max_tokens.
|
|
||
| device = _tensor_device(activation) | ||
| if self._backend is not None: | ||
| return self._backend.backward( | ||
| grad_output, | ||
| activation, | ||
| fc1_weight, | ||
| fc2_weight, | ||
| topk_idx, | ||
| topk_weights, | ||
| fc1_c, | ||
| route_metadata, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
backward never resolves the backend, so it silently returns garbage if called before any forward.
_backend is only populated inside __call__ (line 331-335). A caller that constructs MoeEp(generate_c=True), runs forward through a different instance (or restores a stash across instances), and then calls backward gets uninitialized torch.empty gradients even though the MegaMoE backend is available — a silent wrong-results path rather than an error. Resolve the backend here too, or raise when it is unresolved.
🐛 Resolve the backend in `backward` as well
device = _tensor_device(activation)
+ if not self._backend_resolved:
+ from . import _megamoe
+
+ self._backend = _megamoe.maybe_create(self, device, token_count)
+ self._backend_resolved = True
if self._backend is not None:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| device = _tensor_device(activation) | |
| if self._backend is not None: | |
| return self._backend.backward( | |
| grad_output, | |
| activation, | |
| fc1_weight, | |
| fc2_weight, | |
| topk_idx, | |
| topk_weights, | |
| fc1_c, | |
| route_metadata, | |
| ) | |
| device = _tensor_device(activation) | |
| if not self._backend_resolved: | |
| from . import _megamoe | |
| self._backend = _megamoe.maybe_create(self, device, token_count) | |
| self._backend_resolved = True | |
| if self._backend is not None: | |
| return self._backend.backward( | |
| grad_output, | |
| activation, | |
| fc1_weight, | |
| fc2_weight, | |
| topk_idx, | |
| topk_weights, | |
| fc1_c, | |
| route_metadata, | |
| ) |
🤖 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 `@python/cudnn/moe_ep/api.py` around lines 386 - 398, Update MoeEp.backward to
resolve or validate _backend before dispatching to _backend.backward, rather
than treating an unresolved backend as a valid path. Reuse the same
backend-resolution logic as __call__, ensuring backward either invokes the
available MegaMoE backend or raises clearly when it cannot be resolved, never
returning uninitialized gradients.
| fc1_err = rel_err(fc1_c, ref_fc1_c) if fc1_c.shape == ref_fc1_c.shape else float("inf") | ||
| print( | ||
| f"[rank {rank}] fwd rel_err={fwd_err:.3e} " | ||
| f"fc1_c shape={tuple(fc1_c.shape)} vs ref {tuple(ref_fc1_c.shape)} " | ||
| f"rel_err={fc1_err:.3e} metadata equal={meta_ok}", | ||
| flush=True, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
fc1_c parity never affects the exit status.
fc1_err is computed and printed but excluded from failures at line 150, and a shape mismatch only yields inf in the log. Since fc1_c is the stash the whole backward path depends on (including the de-interleave order in _megamoe._extract_stash), a wrong stash would still print PASS.
💚 Proposed fix
- failures = [k for k, v in {**errs, "forward": fwd_err}.items() if v > tol]
+ failures = [k for k, v in {**errs, "forward": fwd_err, "fc1_c": fc1_err}.items() if v > tol]🤖 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 `@test/python/fe_api/moe_ep/megamoe_backend_parity.py` around lines 115 - 121,
Update the parity failure aggregation near the existing fc1_c diagnostics so
fc1_c shape mismatches and relative-error violations contribute to failures
alongside the other forward and metadata checks. Reuse the established tolerance
and failure-reporting pattern, ensuring an incorrect fc1_c stash causes the test
to report failure rather than only printing inf.
| @pytest.mark.parametrize("input_format", [MoeFormat.MXFP8, MoeFormat.NVFP4]) | ||
| def test_moe_ep_api_accepts_block_scaled_inputs(input_format): | ||
| """The API takes data+scale bundles where the kernel takes separate sf args.""" | ||
|
|
||
| from cudnn import MoeEp | ||
|
|
||
| torch.manual_seed(29) | ||
| device = torch.device("cuda") | ||
| experts, tokens, hidden, intermediate = 2, 4, 32, 16 | ||
| activation, fc1_weight, fc2_weight = _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device) | ||
|
|
||
| api = MoeEp(num_experts=experts, hidden_size=hidden, intermediate_size=intermediate, top_k=1) | ||
| output = api( | ||
| activation, | ||
| fc1_weight, | ||
| fc2_weight, | ||
| torch.tensor([[0], [1], [0], [1]], dtype=torch.int64, device=device), | ||
| torch.ones(tokens, 1, device=device), | ||
| ) | ||
|
|
||
| assert isinstance(output, torch.Tensor) | ||
| assert output.shape == (tokens, hidden) | ||
| assert output.dtype == torch.bfloat16 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Missing CUDA gating on the GPU tests.
These tests hard-code torch.device("cuda") with no availability skip, so they error on CPU-only runners. Worse, the xfail(strict=True) variants (lines 405-409, 556-559, 635-638, 676-679) would swallow a "no CUDA" RuntimeError as the expected failure and report green without exercising anything. Add a module- or test-level skip.
As per coding guidelines, "Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks, cudnn.backend_version(), and torch.cuda.get_device_capability()."
💚 Suggested gate
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+
+
`@pytest.mark.parametrize`("input_format", [MoeFormat.MXFP8, MoeFormat.NVFP4])
+@requires_cuda
def test_moe_ep_api_accepts_block_scaled_inputs(input_format):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @pytest.mark.parametrize("input_format", [MoeFormat.MXFP8, MoeFormat.NVFP4]) | |
| def test_moe_ep_api_accepts_block_scaled_inputs(input_format): | |
| """The API takes data+scale bundles where the kernel takes separate sf args.""" | |
| from cudnn import MoeEp | |
| torch.manual_seed(29) | |
| device = torch.device("cuda") | |
| experts, tokens, hidden, intermediate = 2, 4, 32, 16 | |
| activation, fc1_weight, fc2_weight = _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device) | |
| api = MoeEp(num_experts=experts, hidden_size=hidden, intermediate_size=intermediate, top_k=1) | |
| output = api( | |
| activation, | |
| fc1_weight, | |
| fc2_weight, | |
| torch.tensor([[0], [1], [0], [1]], dtype=torch.int64, device=device), | |
| torch.ones(tokens, 1, device=device), | |
| ) | |
| assert isinstance(output, torch.Tensor) | |
| assert output.shape == (tokens, hidden) | |
| assert output.dtype == torch.bfloat16 | |
| requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") | |
| `@pytest.mark.parametrize`("input_format", [MoeFormat.MXFP8, MoeFormat.NVFP4]) | |
| `@requires_cuda` | |
| def test_moe_ep_api_accepts_block_scaled_inputs(input_format): | |
| """The API takes data+scale bundles where the kernel takes separate sf args.""" | |
| from cudnn import MoeEp | |
| torch.manual_seed(29) | |
| device = torch.device("cuda") | |
| experts, tokens, hidden, intermediate = 2, 4, 32, 16 | |
| activation, fc1_weight, fc2_weight = _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device) | |
| api = MoeEp(num_experts=experts, hidden_size=hidden, intermediate_size=intermediate, top_k=1) | |
| output = api( | |
| activation, | |
| fc1_weight, | |
| fc2_weight, | |
| torch.tensor([[0], [1], [0], [1]], dtype=torch.int64, device=device), | |
| torch.ones(tokens, 1, device=device), | |
| ) | |
| assert isinstance(output, torch.Tensor) | |
| assert output.shape == (tokens, hidden) | |
| assert output.dtype == torch.bfloat16 |
🤖 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 `@test/python/fe_api/moe_ep/test_moe_ep.py` around lines 498 - 520, Add a
module- or test-level CUDA availability gate for the GPU-dependent Moe EP tests,
including the strict xfail variants, using the project’s supported capability
checks so CPU-only runners skip them instead of passing spuriously. Apply the
gate around the relevant tests in test_moe_ep.py while preserving their existing
assertions and xfail behavior on supported CUDA environments.
Source: Coding guidelines
Copy the megamoe training package, its pt reference package, and the four cutedsl_megamoe kernel subpackages (common, moe_mxfp8_glu, moe_nvfp4_swapab, src) verbatim into python/cudnn/moe_ep/_megamoe_backend, preserving the sibling layout megamoe/repo_path.py expects. The bundle carries both megakernels: the CuTe DSL NVFP4 forward (forward_nvfp4.py + moe_nvfp4_swapab) and the FP8 mega backward (bwd_kernel). Provenance (source repos and commit hashes) is recorded in the bundle README. _megamoe.py now defaults to the bundled tree via bundled_root(); CUDNN_MEGAMOE_ROOT becomes an optional override for kernel development. The parity driver picks up the same default in multi-rank mode, the vendored path is excluded from black so it stays verbatim, and run_pr_tests.sh runs the full validation (pytest suite, vendored-import check, single-rank kernel parity) in the container with no external checkout. Validated on GB200 with CUDNN_MEGAMOE_ROOT unset: 18 passed / 7 xfailed, parity fwd rel_err 6.5e-2, bwd rel_err <= 6.6e-2, metadata equal, PASS. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
CUDNN_MEGAMOE_ROOT is gone: _megamoe.py resolves the megamoe package only from the in-tree _megamoe_backend bundle (bundled_root() now raises a clear BackendUnavailable if the bundle is missing), and the parity driver, run_pr_tests.sh, and all docstrings/README lose the override mentions. CUDNN_MOE_EP_BACKEND remains the only backend env knob. Also state in _megamoe.py, the bundle README, and the moe_ep design doc that the backward (bprop) implementation in megamoe/bwd_kernel comes from the Flashinfer team, not the FastKernel team. Revalidated on GB200: 18 passed / 7 xfailed, parity fwd rel_err 6.5e-2, bwd rel_err <= 6.6e-2, metadata equal, PASS. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Covers the one-shot run_pr_tests.sh wrapper and the step-by-step equivalents (pytest suite with the wheel graft, single-rank and 4-rank kernel-vs-oracle parity), the expected baselines and how to read them (the 7 strict xfails, MXFP8-vs-FP32 tolerance band, hard invariants), and the operational gotchas (backend policy knob, supported envelope, one-backward-per-forward, frozen process mode, compile cost). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
New RUNBOOK §2a: build the environment from nvcr.io/nvidia/pytorch:26.05-py3 by installing the bundle's own requirements file, with the load-bearing pins called out (cutlass-dsl 4.5.2, nvshmem4py-cu13, torch>=2.10), a smoke-check snippet, and the sqsh-baking variant for SLURM/pyxis. Vendor cutedsl_megamoe/ci/requirements.txt (verbatim) into the bundle so the environment recipe ships in-tree like the kernels do. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Drop the internal container-image path and cluster account/partition from the known-good-container note and the one-shot srun example; validation now shows a plain in-environment invocation plus a generic SLURM/pyxis job using the image baked per section 2a. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
🟠 Major comments (24)
test/python/fe_api/moe_ep/run_pr_tests.sh-8-14 (1)
8-14: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winValidation mode is not enforced across the wrapper and runbook.
The pytest phase can inherit
auto, even though the runbook documents an uninitialized-memory fallback. This can produce a false validation result.
test/python/fe_api/moe_ep/run_pr_tests.sh#L8-L14: setCUDNN_MOE_EP_BACKEND=megamoeon the pytest command.test/python/fe_api/moe_ep/RUNBOOK.md#L85-L86: make the one-shot command enforce strict backend mode or state that the wrapper does so.test/python/fe_api/moe_ep/RUNBOOK.md#L91-L95: apply the same requirement to the SLURM command.test/python/fe_api/moe_ep/RUNBOOK.md#L113-L114: set strict backend mode on the direct pytest command.🤖 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 `@test/python/fe_api/moe_ep/run_pr_tests.sh` around lines 8 - 14, Enforce strict Megablocks validation mode consistently: update run_pr_tests.sh lines 8-14 to set CUDNN_MOE_EP_BACKEND=megamoe on the pytest command; update RUNBOOK.md lines 85-86, 91-95, and 113-114 so the one-shot, SLURM, and direct pytest commands respectively set or explicitly rely on that strict backend setting.test/python/fe_api/moe_ep/run_pr_tests.sh-5-6 (1)
5-6: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMake setup fail closed. Add
set -euo pipefail, assignCLONEbefore exporting it, and quote$CLONEand$SITEincp,grep, the append redirection, and bothcdcommands.export CLONE=$(...)can return success with an empty value when the innercdfails. Without fail-fast handling, setup failures can leave the wrapper in the wrong directory and produce misleading results.🤖 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 `@test/python/fe_api/moe_ep/run_pr_tests.sh` around lines 5 - 6, Update the setup script near the existing `set -x` and `CLONE` initialization to enable `set -euo pipefail`, assign `CLONE` before exporting it, and quote `$CLONE` and `$SITE` in the specified `cp`, `grep`, append redirection, and both `cd` commands. Preserve the existing setup flow while ensuring directory and command failures stop execution.Source: Linters/SAST tools
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/inputs_process.py-543-548 (1)
543-548: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReplace the generic E8M0 cast at
inputs_process.py:548.
stg_e8m0_from_f32states that.to(Float8E8M0FNU)does not lower correctly. The preceding PTX round trip does not fix the subsequent generic cast. Use a PTX-backed E8M0 conversion and add compiled-kernel coverage for this path.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/inputs_process.py` around lines 543 - 548, Replace the generic Float32.to(cutlass.Float8E8M0FNU) conversion assigned to sf_e8m0 with the PTX-backed stg_e8m0_from_f32 conversion used by the backend. Add compiled-kernel coverage for this E8M0 conversion path, verifying the generated kernel lowers and executes correctly.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/kernel_fc12.py-529-534 (1)
529-534: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
_compute_stagescan return a non-positivenum_ab_stage.
num_ab_stageis an unchecked floor division. Whenmisc_budget + c_bytes_totalapproachessmem_capacity // occupancy, the result becomes0or negative. The value then flows intoPipelineTmaUmma.create(num_stages=...)and into thesA/sBSMEM layouts, so the failure surfaces far from its cause. The MegaMoE subclass adds a dispatch SMEM region tomisc_budget, which makes this reachable for largehidden.Raise a clear error here instead.
🛡️ Proposed guard
num_ab_stage = ( smem_capacity // occupancy - fixed_overhead ) // ab_bytes_per_stage + if num_ab_stage < 1: + raise ValueError( + f"SMEM budget leaves no AB stage: capacity/occupancy=" + f"{smem_capacity // occupancy}, fixed_overhead={fixed_overhead} " + f"(misc={misc_budget}, c={c_bytes_total}), " + f"ab_bytes_per_stage={ab_bytes_per_stage}." + ) return num_acc_stage, num_ab_stage, num_sched_stages🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/kernel_fc12.py` around lines 529 - 534, Update _compute_stages to validate the calculated num_ab_stage before returning it. If the value is zero or negative, raise a clear error describing that available shared memory cannot support a positive number of AB stages, instead of allowing it to reach PipelineTmaUmma.create or the sA/sB layouts.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/moe_persistent_scheduler.py-1207-1212 (1)
1207-1212: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDerive the
paramsslice length from its serialized values.
MoEStaticSchedulerParams.__extract_mlir_values__serializes onlyInt32fields. Withstatic_expert_shape, all three fields are Python integers, so the serialized length is zero. The hardcoded three-value slice then causes__new_from_mlir_values__to fail its length assertion and shifts every subsequent field. Uselen(extract_mlir_values(self.params)), as in the dynamic scheduler.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/moe_persistent_scheduler.py` around lines 1207 - 1212, Update the parameter deserialization logic around new_from_mlir_values to derive the params slice length from len(extract_mlir_values(self.params)) instead of assuming three values. Increment idx by that computed length so static_expert_shape produces an empty params slice and subsequent fields remain correctly aligned.python/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_turboquant_numerics.py-178-189 (1)
178-189: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRun the distributed teardown even when an assertion fails.
The asserts at lines 179-183 raise out of
main. That skipsfinalize_dist_and_nvshmem()at line 188. Undertorchrun --nproc_per_node=4, one failing rank exits without finalizing NVSHMEM or the process group, so the remaining ranks block at their next collective and the job hangs instead of failing fast.megamoe/tests/test_bwd_v0_parity.pyavoids this by collecting failures into a list and always reaching teardown.Put the measurement body in
try/finally, or collect failures and assert after teardown.🛡️ Proposed teardown fix
- if rank == 0: - print("TURBOQUANT NUMERICS PASS") - if not _NO_DIST: - finalize_dist_and_nvshmem() - return 0 + if rank == 0: + print("TURBOQUANT NUMERICS PASS") + return 0Then wrap the call site so teardown always runs:
if __name__ == "__main__": try: rc = main() finally: if not _NO_DIST: from src.bootstrap import finalize_dist_and_nvshmem finalize_dist_and_nvshmem() sys.exit(rc)🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_turboquant_numerics.py` around lines 178 - 189, Ensure distributed teardown always executes when assertions in main fail: move finalize_dist_and_nvshmem out of the normal-success path and wrap the measurement/assertion flow in try/finally, or collect failures and defer assertion until after teardown. Update main and its __main__ call site so _NO_DIST remains respected and failures still propagate after finalize_dist_and_nvshmem runs.python/cudnn/moe_ep/_megamoe_backend/megamoe/turboquant.py-85-88 (1)
85-88: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
load_weightsforwards an fp32w13and a bf16w2.
rotate_hiddencalls.float()at line 58 and never casts back, so line 88 passes an fp32w13tosuper().load_weights()whilew2keeps its original dtype. The sibling test inmegamoe/tests/test_turboquant_numerics.pyperforms the same fold and casts back explicitly withrotate_hidden(w13g.float(), q).bfloat16(). Restore the input dtype so the mixin matches the plain forward path and the test.Also note the fp32
w13doubles the transient host/device memory for the weight fold on large expert counts.🔧 Proposed fix
def load_weights(self, w13: torch.Tensor, w2: torch.Tensor) -> None: # Fold Q into fc1's K(hidden) axis; w2 is untouched (its K is the # intermediate dim, whose quant lives inside the kernel). - super().load_weights(rotate_hidden(w13, self.q_fp32), w2) + w13_rot = rotate_hidden(w13, self.q_fp32).to(w13.dtype) + super().load_weights(w13_rot, w2)🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/turboquant.py` around lines 85 - 88, Update MegamoeBackend.load_weights to preserve w13’s original dtype after applying rotate_hidden: perform the rotation using the existing q_fp32 value, then cast the folded result back to w13.dtype before passing it to super().load_weights. Leave w2 unchanged and preserve the existing weight-folding behavior.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_fc12_common.py-113-116 (1)
113-116: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
ProblemDescvalidates MXFP8 shapes with NVFP4 constants.This descriptor is kind-parameterized, and Lines 117-128 correctly select the interleave and SF vector size from
self.kind. Two nearby checks do not:
- Line 113 requires
hidden % Nvfp4BlockSize == 0for every kind.Nvfp4BlockSizeis 16, but MXFP8 uses a 32-element block (Mxfp8BlockSize, already imported at Line 52 and used at Line 123). Anmxfp8_e4m3problem withhidden = 16orhidden = 48passes this guard while violating the MXFP8 block contract.- Lines 198-222 pass
Nvfp4DataDtypetocheck_tma_leading_dim_alignforactivation,fc1_weight,fc2_weight, andfc1_outputregardless of kind. MXFP8 stores one byte per element, not two, so the derived byte count is wrong for MXFP8. The direction is conservative, so no unaligned MXFP8 shape is accepted; the visible effect is a misleading rejection and an NVFP4-specific "2 fp4 packed per byte" error for an MXFP8 problem with an oddhidden.Derive both from
self.kind, the same way Line 123 already does.🐛 Proposed fix
- if self.hidden <= 0 or self.hidden % Nvfp4BlockSize != 0: + _block = Nvfp4BlockSize if self.kind == "nvfp4" else Mxfp8BlockSize + if self.hidden <= 0 or self.hidden % _block != 0: raise ValueError( - f"hidden ({self.hidden}) must be a positive multiple of {Nvfp4BlockSize}." + f"hidden ({self.hidden}) must be a positive multiple of " + f"{_block} (kind={self.kind!r})." )from moe_nvfp4_swapab.runner_common import ( check_tma_leading_dim_align as _check_tma_leading_dim_align, ) + _tma_data_dtype = kind_data_dtype(self.kind) _check_tma_leading_dim_align( "activation", {"k_major": self.hidden}[self.fc1_activation_layout], - Nvfp4DataDtype, + _tma_data_dtype, ) _check_tma_leading_dim_align( "fc1_weight", {"k_major": self.hidden}[self.fc1_weight_layout], - Nvfp4DataDtype, + _tma_data_dtype, ) _check_tma_leading_dim_align( "fc2_weight", {"k_major": self.intermediate // 2, "n_major": self.hidden}[ self.fc2_weight_layout ], - Nvfp4DataDtype, + _tma_data_dtype, ) _check_tma_leading_dim_align( "fc1_output (kernel-internal, fixed k_major)", self.intermediate // 2, - Nvfp4DataDtype, + _tma_data_dtype, )Also applies to: 198-227
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_fc12_common.py` around lines 113 - 116, Update ProblemDesc validation to derive the hidden-dimension block size from self.kind, using Mxfp8BlockSize for MXFP8 and Nvfp4BlockSize otherwise. In the check_tma_leading_dim_align calls for activation, fc1_weight, fc2_weight, and fc1_output, likewise select the kind-appropriate data dtype so MXFP8 uses its one-byte-per-element representation while NVFP4 retains Nvfp4DataDtype.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/megamoe_kernel_mxfp8.py-234-274 (1)
234-274: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
dedup_dispatchandcombine_pre_reduceare narrowed onselfbut forwarded raw toTokenInPullTokenBackPush.Line 234 narrows the flag:
self.dedup_dispatch = dedup_dispatch and num_topk > 1. Line 267 narrows the second one against it. The workspace layout and__call__then follow the narrowed values:
_build_shared_region_specsallocatesdup_linkonly whenself.dedup_dispatch(Line 665)._build_local_region_specsallocatesreduce_listonly whenself.combine_pre_reduce(Line 600).__call__setsdup_link = Noneandreduce_list = Noneon the same narrowed flags (Lines 1003-1010).Lines 350 and 352 instead forward the raw parameters. With
dedup_dispatch=Trueandnum_topk == 1, the token-comm helper is configured with dedup enabled while nodup_linkregion exists andtoken_comm_args.dup_linkisNone, so the dedup fan-out dereferences a null tensor inside the kernel.combine_pre_reducehas the same shape againstreduce_list.The validation at Lines 270-274 also reads the raw
dedup_dispatch, so a caller that requestscombine_pre_reduce=Truewithnum_topk == 1passes validation and gets the feature silently disabled. Line 306 already uses the narrowedself.combine_pre_reduce, which shows the intent.Forward the narrowed flags and validate against them.
🐛 Proposed fix
self.combine_pre_reduce = ( combine_pre_reduce and self.dedup_dispatch ) - if combine_pre_reduce and not (dedup_dispatch and combine_in_flight_reduce): + if combine_pre_reduce and not ( + self.dedup_dispatch and combine_in_flight_reduce + ): raise ValueError( "combine_pre_reduce requires dedup_dispatch=True and " - "combine_in_flight_reduce=True." + "combine_in_flight_reduce=True (dedup_dispatch also requires " + "num_topk > 1)." )token_back_schedule_mode=self.token_back_schedule_mode, - dedup_dispatch=dedup_dispatch, + dedup_dispatch=self.dedup_dispatch, token_back_reduce_topk=combine_in_flight_reduce, - combine_pre_reduce=combine_pre_reduce, + combine_pre_reduce=self.combine_pre_reduce, )Also applies to: 344-353
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/megamoe_kernel_mxfp8.py` around lines 234 - 274, Use the normalized instance flags consistently in the token-communication setup: update the `TokenInPullTokenBackPush` construction to receive `self.dedup_dispatch` and `self.combine_pre_reduce` rather than the raw parameters. Change the `combine_pre_reduce` validation to check the narrowed `self.combine_pre_reduce`, preserving the existing `num_topk > 1` gating and ensuring helper configuration matches allocated workspace regions.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_fc12.py-220-232 (1)
220-232: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
fc2_output_dtype=cutlass.BFloat16is hardcoded while the buffer followsproblem.fc2_output_dtype.
--fc2_output_dtype fp16is a supported CLI value.parse_output_dtypeaccepts it, andProblemDesc.__post_init__inrunner_fc12_common.pypermitstorch.float16._alloc_fc2_outputat Line 190 then allocates the buffer withproblem.fc2_output_dtype, but Line 226 always instantiates the kernel withcutlass.BFloat16.With
--fc2_output_dtype fp16the kernel writes bf16 bit patterns into a buffer the host reads as fp16, sovalidate()compares reinterpreted bytes and reports meaningless diffs. Either derive the kernel dtype fromproblem.fc2_output_dtype, or reject fp16 for this kind so the failure is explicit.🐛 Option A: derive the kernel dtype from the problem descriptor
def _instantiate_kernel(self, common_kwargs: dict): import cutlass from moe_nvfp4_swapab.kernel_fc12 import Sm100SwapABSwigluFp4Fc12Kernel + _cutlass_out_dtype = { + torch.bfloat16: cutlass.BFloat16, + torch.float16: cutlass.Float16, + }[self.problem.fc2_output_dtype] return Sm100SwapABSwigluFp4Fc12Kernel( **common_kwargs, - fc2_output_dtype=cutlass.BFloat16, + fc2_output_dtype=_cutlass_out_dtype,🛡️ Option B: reject the unsupported dtype up front
def _instantiate_kernel(self, common_kwargs: dict): import cutlass from moe_nvfp4_swapab.kernel_fc12 import Sm100SwapABSwigluFp4Fc12Kernel + if self.problem.fc2_output_dtype is not torch.bfloat16: + raise ValueError( + "the NVFP4 fused fc12 kernel emits bf16 only; got " + f"fc2_output_dtype={self.problem.fc2_output_dtype}" + ) return Sm100SwapABSwigluFp4Fc12Kernel(🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_fc12.py` around lines 220 - 232, Update _instantiate_kernel so fc2_output_dtype matches self.problem.fc2_output_dtype, including the supported torch.float16/fp16 configuration, using the appropriate CUTLASS dtype mapping; keep the kernel output dtype consistent with the buffer allocated by _alloc_fc2_output.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/common/moe_utils.py-178-181 (1)
178-181: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRaise on an unsupported fp8 type instead of returning
None.The
elsebranch prints a device-side message and executes a barereturn, so the function returnsNone. The caller at line 269 then evaluatescute.arch.rcp_approx(qpvscale_up)onNone. That produces the opaque trace error "None to integer conversion is not supported", whichmoe_mxfp8_glu/run_functional_tests.shlines 118-123 record as the reason the e5m2 test cases were removed from the functional suite.
fp8_typeis a compile-time value here, so an unsupported type is a host-side configuration error. Raise it at trace time with the offending type in the message. The diagnostic then points at the real cause instead of at an unrelated arithmetic site.
cvt_f32x4_to_f8x4_pack_i32has the sameprintfand barereturnpattern at lines 228-231. Apply the same change there.🐛 Proposed fix for lines 178-181
else: - with cute.arch.elect_one(): - cute.printf("error: unsupported fp8 element type") - return + raise TypeError(f"cvt_f32_to_f8_to_f32: unsupported fp8 element type {fp8_type}")Apply the same change in
cvt_f32x4_to_f8x4_pack_i32:else: - with cute.arch.elect_one(): - cute.printf("error: unsupported fp8 element type") - return + raise TypeError( + f"cvt_f32x4_to_f8x4_pack_i32: unsupported fp8 element type {fp8_type}" + )🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/common/moe_utils.py` around lines 178 - 181, Replace the unsupported-type else branches in the fp8 scale helper and cvt_f32x4_to_f8x4_pack_i32 with host-side exceptions that include the offending fp8_type; remove the device-side printf and bare returns so invalid compile-time configurations fail during tracing.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/flag_batch.py-75-81 (1)
75-81: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReplace Python
orwith a traced logical OR. DynamicBoolean.__bool__raisesPHASE_DYNAMIC_TO_STATIC_BOOL, so line 75 cannot trace. Combine the predicates with|. Keepcutlass.const_expr(not no_fire)because callers passno_fireascutlass.Constexpr.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/flag_batch.py` around lines 75 - 81, Update the condition in the batch-processing flow around `_make(...).fire()` to combine `cumulated == flush_threshold` and `next_phase != self.phase` with the traced logical-OR operator `|` instead of Python `or`. Preserve `cutlass.const_expr(not no_fire)` for the `no_fire` check, since it is passed as a `cutlass.Constexpr`.python/cudnn/moe_ep/_megamoe_backend/pt/config.py-40-55 (1)
40-55: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winValidate
EpConfigagainstprocess_group. When a group is supplied and distributed is initialized, compare its world size and rank withep_sizeandep_rank. Anep_sizemismatch givesall_to_all_singlesplit lists with the wrong length and can fail or hang. Anep_rankmismatch assigns the wrong local expert range. Keep thetorch.distributedimport local and add coverage for both mismatches.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/pt/config.py` around lines 40 - 55, Update EpConfig.__post_init__ to validate an optional process_group when torch.distributed is initialized: compare its world size with ep_size and its rank with ep_rank, raising clear ValueErrors on either mismatch before distributed communication. Keep the torch.distributed import local, and add coverage for both world-size and rank mismatch cases.Source: Coding guidelines
python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/weights_bwd.py-65-65 (1)
65-65: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winPreserve the input device instead of calling
.cuda().Both lines call
.cuda(), which targets the current CUDA device. Under expert parallelism the caller can holdw13andw2on a device other than the current one, so the returned operands land on the wrong device. The consumer inmegamoe/bwd_kernel/backward.py(lines 128-140) then performs an in-placecopy_into existing buffers, which masks the mismatch behind an implicit cross-device copy. Use the source tensor's device.🔧 Proposed change
- g1_f32 = w2.float().cuda().transpose(1, 2).contiguous() + g1_f32 = w2.to(device=w2.device, dtype=torch.float32).transpose(1, 2).contiguous()- g2_f32 = interleave_gate_up(w13).float().cuda().transpose(1, 2).contiguous() + g2_f32 = ( + interleave_gate_up(w13) + .to(device=w13.device, dtype=torch.float32) + .transpose(1, 2) + .contiguous() + )Also applies to: 73-73
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/weights_bwd.py` at line 65, Update the tensor conversions in the weight backward path, including the lines assigning g1_f32 and the corresponding w13 operand, to preserve each source tensor’s device instead of calling .cuda(). Keep the existing float conversion, transpose, and contiguous operations unchanged.python/cudnn/moe_ep/_megamoe_backend/megamoe/repo_path.py-17-29 (1)
17-29: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy liftRemove the global import-path shim and
MEGAMOE_REPOoverride.
- The lazy backend path reaches
megamoe.repo_path;import cudnndoes not. Keep this backend import lazy.- The shim exposes
src,common, andptthrough process-globalsys.pathentries. Use package-qualified imports or an isolated loader to prevent module shadowing.- Remove
MEGAMOE_REPOif the bundled kernels are authoritative. Remove the internal GitLab URL from the error message.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/repo_path.py` around lines 17 - 29, Update the repo_path module’s REPO_ROOT setup to use the bundled cutedsl_megamoe location without honoring MEGAMOE_REPO, and remove the sys.path insertion loop. Keep the backend import lazy so importing cudnn does not load megamoe.repo_path, and replace shim-dependent imports with package-qualified imports or an isolated loader. Simplify the missing-repository ImportError to remove the internal GitLab URL and override guidance.Source: Coding guidelines
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/mega_runner.py-2730-2750 (1)
2730-2750: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReturn a non-zero exit code when the kernel launch is skipped.
maincatchesNotImplementedErrorat line 2733 and leavesreturn_codeat 0. Line 2749 then callsos._exit(0). A run that never launched the kernel and never validated therefore reports success. Any CI job orrun_mega_tests.shwrapper that checks the exit code treats that as a pass.Set a non-zero
return_codein the handler, or re-raise.🔧 Proposed fix
except NotImplementedError as exc: # Expected until the MegaMoE kernel side is wired; the host # orchestration above is the part being smoke-tested for now. + return_code = 1 if rank == 0: print(f"[mega_runner] kernel launch skipped: {exc}")🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/mega_runner.py` around lines 2730 - 2750, Update the NotImplementedError handler in main to set return_code to a non-zero value before teardown, or re-raise the exception, so skipped kernel launches cause os._exit to report failure while preserving normal successful execution.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py-100-111 (1)
100-111: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
valid_ab_tupleis undefined in bothab_dtypevalidators. Both kernels bindvalid_abfromVALID_AB_DTYPE_SF_SIZEand then format an undefinedvalid_ab_tuplein the error message. An unsupportedab_dtypetherefore raisesNameErrorwhile building the message instead of the intendedValueError. The shared root cause is one wrong identifier duplicated by copy-paste.
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py#L100-L111: changevalid_ab_tupletovalid_abat line 106.python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_fc12.py#L100-L111: changevalid_ab_tupletovalid_abat line 106.🐛 Proposed fix (identical in both files)
raise ValueError( f"ab_dtype={ab_dtype.__name__} is not valid for " f"sf_vec_size={sf_vec_size}. " - f"Expected one of: {[t.__name__ for t in valid_ab_tuple]}." + f"Expected one of: {[t.__name__ for t in valid_ab]}." )🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py` around lines 100 - 111, Replace the undefined valid_ab_tuple identifier with the existing valid_ab variable in both ab_dtype validators: python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py lines 100-111 and python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_fc12.py lines 100-111. Preserve the intended ValueError message listing the valid dtypes.Source: Linters/SAST tools
python/cudnn/moe_ep/_megamoe_backend/pt/tests/parity_ep_vs_reference_fp4.py-108-113 (1)
108-113: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSet the CUDA device before
init_process_group.Line 108 initializes the NCCL process group, and line 112 selects the device afterwards. NCCL binds the communicator to the current device during initialization. If every rank still points at device 0 at that moment, the initialization can hang or the communicator can bind to the wrong device.
Read
LOCAL_RANKand calltorch.cuda.set_devicefirst, then initialize the process group.🔧 Proposed fix for the initialization order
- dist.init_process_group("nccl") - rank = dist.get_rank() - world = dist.get_world_size() - local_rank = int(os.environ.get("LOCAL_RANK", rank % max(torch.cuda.device_count(), 1))) - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size()🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/pt/tests/parity_ep_vs_reference_fp4.py` around lines 108 - 113, In the distributed setup before rank/world queries, read LOCAL_RANK and call torch.cuda.set_device first, then invoke dist.init_process_group("nccl") so communicator initialization uses the correct GPU. Preserve the existing local_rank fallback and device construction after initialization.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/dynamic_mainloop.py-147-153 (1)
147-153: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMask the extracted SF IDs before inserting them into
idesc.
Int32right shifts are signed. When bit 31 of an SF TMEM address is set, sign extension writes outside thea_sf_id_orb_sf_id_fields and can corrupt other descriptor fields, includingk_size_. Extract the top two bits, mask withInt32(0x3), then shift them into their target fields.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/dynamic_mainloop.py` around lines 147 - 153, Update the SF ID packing in the descriptor construction flow around sfa_top, sfb_top, and idesc so each signed right-shifted address value is masked with Int32(0x3) before shifting into the a_sf_id_ or b_sf_id_ field. Preserve the existing top-two-bit extraction and ensure no sign-extension bits can modify other descriptor fields such as k_size_.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/topk_reduce.py-322-327 (1)
322-327: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDerive the output store width from
out_dtype.All three kernels hardcode
num_bits_per_copyfor the final store: 128 in_reduce_bf16, 256 in_reduce_mxfp8and_reduce_fp4. Each value is only correct for a 16-bitout_dtype._mark_alignmentalready computes the byte width ashidden_per_thread * out_dtype.width // 8, so the alignment adapts while the copy width does not. With an fp32reduced_outputthe mxfp8 kernel needs 512 bits per thread and stores 256, which drops half of every thread's output. Compute the width fromhidden_per_thread * out_dtype.widthin all three kernels.🛡️ Proposed width derivation
out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) + store_bits = hidden_per_thread * out_dtype.width cute.copy( - cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=256), - out, self._mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=store_bits), + out, self._mark_alignment(dst, store_bits // 8), )Apply the same change in
_reduce_bf16and_reduce_fp4.Also applies to: 408-413, 500-505
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/topk_reduce.py` around lines 322 - 327, Update the final copy operations in _reduce_bf16, _reduce_mxfp8, and _reduce_fp4 to derive num_bits_per_copy from hidden_per_thread multiplied by out_dtype.width instead of hardcoded constants. Keep the existing _mark_alignment calculation and ensure the copy width matches the full per-thread output for every supported out_dtype.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py-205-208 (1)
205-208: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
allow_overlap_accand the tile-shape condition are overwritten in both epilogues. Each constructor computes_overlapping_accumfromallow_overlap_accand the_cta_tile_n == EpiWarpCount * EpilogueTileN * 2condition, then immediately assignsTrue. The parameter and the shape guard have no effect, while the TMEM column budget and the256 - _num_sf_tmem_colsphase offsets assume the overlap layout.
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py#L205-L208: delete the unconditionalself._overlapping_accum = Trueand reject an unsupportedcta_tile_nin__init__, or removeallow_overlap_accand document that overlap is mandatory.python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/epilogue_bwd.py#L225-L228: apply the same decision here, next to the existinggenerate_canduse_stg_fc1validation block.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py` around lines 205 - 208, The overlap flag is unconditionally overriding the validated configuration in both epilogue constructors. In python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py#L205-L208, remove the unconditional assignment and validate or reject unsupported cta_tile_n values while preserving allow_overlap_acc; apply the same decision in python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/epilogue_bwd.py#L225-L228 near the generate_c and use_stg_fc1 validation block.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py-243-250 (1)
243-250: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winQuantized combine without
static_expert_shapecollapses the fc2 scale stride in both epilogues. Whencombine_format.is_quantizedis true andstatic_expert_shapeisNone, both constructors set_fc2_sf_block_pad = 0and_hidden_fc2 = 0._stg_sf_fc2then computespool_token_global * 0, so every pool token writes its E8M0 scales to the same row, and thetoken_back_by_dispatchdata store multiplies the pool row stride by_hidden_fc2 = 0. The combine output dequantizes with the wrong scale for every token after the first, with no error raised. Both constructors already reject unsupported configurations, so add the check there.
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py#L243-L250: raise in__init__whencombine_format.is_quantizedandstatic_expert_shape is None, instead of falling through to the zero-padding branch.python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/epilogue_bwd.py#L263-L270: add the same check beside the existinggenerate_canduse_stg_fc1validation.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py` around lines 243 - 250, Reject quantized combine configurations without static_expert_shape in both constructors: update epilogue_mxfp8.py lines 243-250 and epilogue_bwd.py lines 263-270 to raise during __init__ when combine_format.is_quantized and static_expert_shape is None, alongside the existing validation, instead of assigning zero _fc2_sf_block_pad and _hidden_fc2 values.python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/epilogue_bwd.py-148-152 (1)
148-152: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy liftExtract the shared epilogue members instead of copying the forward file.
This module duplicates a large part of
cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.py, including the class nameGluMxfp8Epilogue.Fc2OutputDest,tma_store_fc1_output,_store_fc1_c_subtile,_subtile_fc2_tmem_tensor,_advance_fc2_tmem_tensor,_acc_pipeline_consumer_release,_run_fc2_subtile,_write_sf_fc2_buffer,_stg_sf_fc2, and every codegen-time property are byte-identical to their forward counterparts. The__init__body diverges only in the BWD validation block and_subtile_cnt.Two identical class names in the same package also make tracebacks and imports ambiguous. Move the shared members into a common base class or module, keep the BWD subclass limited to
_subtile_local_tmem_tensor_single,_run_fc1_task_tile,_run_fc1_subtile,_stg_sf_fc1, and the added validation, and give the subclass a distinct name. The other findings on this pair of files already show the cost: each one must be fixed in both copies.Also applies to: 382-428, 866-879, 1074-1082
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/epilogue_bwd.py` around lines 148 - 152, Extract the byte-identical members and codegen-time properties from GluMxfp8Epilogue into a shared base class or module, then make the BWD-specific epilogue inherit from it. Rename the BWD subclass to a distinct name and keep it limited to _subtile_local_tmem_tensor_single, _run_fc1_task_tile, _run_fc1_subtile, _stg_sf_fc1, the BWD validation, and its differing _subtile_cnt; preserve the shared forward behavior through inheritance.python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/mega_runner.py-885-921 (1)
885-921: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReport a failure exit code and always tear down.
Two problems live in this block.
return_codeis initialized to 0 and never reassigned, so aNotImplementedErrorfromtester.run()prints the skip message and the process still exits 0. A CI job that gates on the exit status treats an unrun kernel as a pass. Any other exception propagates out ofmain()before the nvshmem free loop andfinalize_dist_and_nvshmem()at line 920 run, which leaves the symmetric heap allocated and the other ranks blocked in their next collective.🛡️ Proposed exit-code and teardown fix
return_code = 0 try: tester.run() except NotImplementedError as exc: if rank == 0: print(f"[mega_runner_mxfp8] kernel launch skipped: {exc}") - - if not _NO_DIST: - tester._compiled_kernel = None - tester._kernel = None - gc.collect() - torch.cuda.synchronize() - try: - import nvshmem.core + return_code = 2 + finally: + if not _NO_DIST: + tester._compiled_kernel = None + tester._kernel = None + gc.collect() + torch.cuda.synchronize() + try: + import nvshmem.core + ... + except ImportError: + pass + gc.collect() + finalize_dist_and_nvshmem() return return_codeKeep the existing free loop body inside the
finallyblock. If a skipped launch must stay non-fatal, set a distinct non-zero code or print an explicit marker the harness can assert on.🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/mega_runner.py` around lines 885 - 921, Update the runner cleanup around tester.run() so every execution path performs the existing nvshmem free loop and finalize_dist_and_nvshmem() in a finally block, including unexpected exceptions. Set return_code to a distinct non-zero value when tester.run() raises NotImplementedError while preserving the skip message, and return that code so skipped launches cannot report success.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5e0cc53d-9bf2-4186-893c-db7236035046
📒 Files selected for processing (110)
.pre-commit-config.yamldocs/fe-oss-apis/moe_ep.mdpython/cudnn/moe_ep/_megamoe.pypython/cudnn/moe_ep/_megamoe_backend/README.mdpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/README.mdpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/ci/requirements.txtpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/common/__init__.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/common/host_utils.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/common/megamoe_constants.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/common/moe_utils.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/__init__.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/epilogue_mxfp8.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/mega_reference_mxfp8.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/mega_runner.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/megamoe_kernel_mxfp8.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/run_functional_tests.shpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/run_mega_tests.shpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/runner_common.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_mxfp8_glu/runner_fc12.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/__init__.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/benchmark_p2p.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/contract.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/custom_ext.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/dynamic_mainloop.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/epilogue.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/epilogue_refactor.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/fc1_fc2_fuse_sched.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/kernel_fc12.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/mega_reference.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/mega_runner.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/megamoe_kernel.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/moe_persistent_scheduler.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/moe_utils.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/run_functional_tests.shpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/run_mega_tests.shpython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_common.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_fc12.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/runner_fc12_common.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/simulate_fc1_fc2_sched.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/moe_nvfp4_swapab/topk_reduce.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/__init__.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/bootstrap.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/cleanup_kernel.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/config.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/dispatch_kernel.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/flag_batch.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/grid_sync.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/iket_compat.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/inputs_process.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/ptx_helpers.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/reference.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/sf_swizzle.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/sym_buffer.pypython/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/token_comm.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/BWD_DESIGN.mdpython/cudnn/moe_ep/_megamoe_backend/megamoe/README.mdpython/cudnn/moe_ep/_megamoe_backend/megamoe/__init__.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/__init__.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/backward.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/epilogue_bwd.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_fc12.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/moe_utils_bwd.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/repro_mega.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/test_bwd_fc12.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/weights_bwd.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_v0.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/forward.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/forward_nvfp4.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/fp8_bwd.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/pools.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/quant_kernels.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/repo_path.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/__init__.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/probe_bwd_contracts.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/smoke_generate_c.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_bwd_v0_parity.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_forward_parity.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_hybrid_training.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_hybrid_training_dist.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/tests/test_turboquant_numerics.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/training.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/turboquant.pypython/cudnn/moe_ep/_megamoe_backend/megamoe/weights.pypython/cudnn/moe_ep/_megamoe_backend/pt/__init__.pypython/cudnn/moe_ep/_megamoe_backend/pt/comm/__init__.pypython/cudnn/moe_ep/_megamoe_backend/pt/comm/base.pypython/cudnn/moe_ep/_megamoe_backend/pt/comm/torch_dist.pypython/cudnn/moe_ep/_megamoe_backend/pt/config.pypython/cudnn/moe_ep/_megamoe_backend/pt/dispatch_combine.pypython/cudnn/moe_ep/_megamoe_backend/pt/experts.pypython/cudnn/moe_ep/_megamoe_backend/pt/experts_fp4.pypython/cudnn/moe_ep/_megamoe_backend/pt/layer.pypython/cudnn/moe_ep/_megamoe_backend/pt/layer_fp4.pypython/cudnn/moe_ep/_megamoe_backend/pt/quant.pypython/cudnn/moe_ep/_megamoe_backend/pt/reference.pypython/cudnn/moe_ep/_megamoe_backend/pt/reference_fp4.pypython/cudnn/moe_ep/_megamoe_backend/pt/routing.pypython/cudnn/moe_ep/_megamoe_backend/pt/tests/parity_ep_vs_reference.pypython/cudnn/moe_ep/_megamoe_backend/pt/tests/parity_ep_vs_reference_fp4.pypython/cudnn/moe_ep/_megamoe_backend/pt/tests/run_tests.shpython/cudnn/moe_ep/_megamoe_backend/pt/tests/test_fp4_qat_numerics.pypython/cudnn/moe_ep/_megamoe_backend/pt/tests/test_quant_vs_kernel.pypython/cudnn/moe_ep/_megamoe_backend/pt/tests/test_reference.pypython/cudnn/moe_ep/_megamoe_backend/pt/tests/test_scaled_grouped_mm_mxfp8.pypython/cudnn/moe_ep/api.pytest/python/fe_api/moe_ep/RUNBOOK.mdtest/python/fe_api/moe_ep/megamoe_backend_parity.pytest/python/fe_api/moe_ep/run_pr_tests.sh
🚧 Files skipped from review as they are similar to previous changes (4)
- test/python/fe_api/moe_ep/megamoe_backend_parity.py
- python/cudnn/moe_ep/_megamoe.py
- python/cudnn/moe_ep/api.py
- docs/fe-oss-apis/moe_ep.md
| dout_q, dout_sf = mxfp8_rowquant(dout.to(torch.bfloat16)) | ||
| bwd.my_dout.view(torch.uint8).copy_(dout_q.view(torch.uint8)) | ||
| bwd.my_dout_sf.view(torch.uint8)[:, : H // 32].copy_( | ||
| dout_sf.contiguous().view(torch.uint8) | ||
| ) | ||
| bwd.recv_count_sum.copy_( | ||
| torch.tensor(counts, device=device, dtype=torch.int64) | ||
| ) | ||
| bwd.act_pool.view(torch.uint8).zero_() | ||
| bwd.act_sf_swz.view(torch.uint8).zero_() | ||
| bwd.fc1_ready_counter.zero_() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
The my_dout staging copy fails when T < max_tokens_per_rank.
bwd.my_dout has shape (self.T, H) with self.T = fwd.cfg.max_tokens_per_rank (Line 77). dout_q has T rows. Line 332 copies into the full buffer, so copy_ raises a shape error for every step with fewer tokens than the configured capacity. The SF copy on Line 333 already slices columns but also not rows.
Slice both destinations to [:T]. Zero the padding tail as well, otherwise the peer-pull can read stale rows from the previous step.
🐛 Proposed fix
dout_q, dout_sf = mxfp8_rowquant(dout.to(torch.bfloat16))
- bwd.my_dout.view(torch.uint8).copy_(dout_q.view(torch.uint8))
- bwd.my_dout_sf.view(torch.uint8)[:, : H // 32].copy_(
+ my_dout_u8 = bwd.my_dout.view(torch.uint8)
+ my_dout_sf_u8 = bwd.my_dout_sf.view(torch.uint8)
+ if T < bwd.T:
+ my_dout_u8[T:].zero_()
+ my_dout_sf_u8[T:].zero_()
+ my_dout_u8[:T].copy_(dout_q.view(torch.uint8))
+ my_dout_sf_u8[:T, : H // 32].copy_(
dout_sf.contiguous().view(torch.uint8)
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| dout_q, dout_sf = mxfp8_rowquant(dout.to(torch.bfloat16)) | |
| bwd.my_dout.view(torch.uint8).copy_(dout_q.view(torch.uint8)) | |
| bwd.my_dout_sf.view(torch.uint8)[:, : H // 32].copy_( | |
| dout_sf.contiguous().view(torch.uint8) | |
| ) | |
| bwd.recv_count_sum.copy_( | |
| torch.tensor(counts, device=device, dtype=torch.int64) | |
| ) | |
| bwd.act_pool.view(torch.uint8).zero_() | |
| bwd.act_sf_swz.view(torch.uint8).zero_() | |
| bwd.fc1_ready_counter.zero_() | |
| dout_q, dout_sf = mxfp8_rowquant(dout.to(torch.bfloat16)) | |
| my_dout_u8 = bwd.my_dout.view(torch.uint8) | |
| my_dout_sf_u8 = bwd.my_dout_sf.view(torch.uint8) | |
| if T < bwd.T: | |
| my_dout_u8[T:].zero_() | |
| my_dout_sf_u8[T:].zero_() | |
| my_dout_u8[:T].copy_(dout_q.view(torch.uint8)) | |
| my_dout_sf_u8[:T, : H // 32].copy_( | |
| dout_sf.contiguous().view(torch.uint8) | |
| ) | |
| bwd.recv_count_sum.copy_( | |
| torch.tensor(counts, device=device, dtype=torch.int64) | |
| ) | |
| bwd.act_pool.view(torch.uint8).zero_() | |
| bwd.act_sf_swz.view(torch.uint8).zero_() | |
| bwd.fc1_ready_counter.zero_() |
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/backward.py` around
lines 331 - 341, Update the staging copies in the backward path around
mxfp8_rowquant so both bwd.my_dout and bwd.my_dout_sf destinations are sliced to
[:T] before copying dout_q and dout_sf. After copying, explicitly zero the
remaining padding rows in both buffers from T onward so peer-pull cannot read
stale data.
| with cute.arch.elect_one(): | ||
| pull_buffer_warp_ptr = pull_buffer_ptr + ( | ||
| warp_idx * Int32(self.hidden_bytes) | ||
| ) | ||
| tma_src_addr = ( | ||
| inp_tok_local_base | ||
| + cur_peer_offset | ||
| + Int64(src_token * Int32(self.hidden_bytes)) | ||
| ) | ||
| tma_load_1d_raw( | ||
| pull_buffer_warp_ptr, | ||
| tma_src_addr, | ||
| pull_mbar_ptr + warp_idx, | ||
| Int32(self.hidden_bytes), | ||
| ) | ||
| cute.arch.sync_warp() | ||
|
|
||
| sf_passes: cutlass.Constexpr[int] = ( | ||
| self.sf_uint32_per_token + 31 | ||
| ) // 32 | ||
| sf_vals = [] | ||
| for _ in cutlass.range_constexpr(0, sf_passes, 1): | ||
| sf_vals.append(Int32(0)) | ||
| for i in cutlass.range_constexpr(0, sf_passes, 1): | ||
| j = Int32(i * self.warp_threads) + lane_idx | ||
| if j < Int32(self.sf_uint32_per_token): | ||
| sf_addr = ( | ||
| inp_sf_local_base | ||
| + cur_peer_offset | ||
| + Int64( | ||
| (src_token * Int32(self.sf_uint32_per_token) + j) | ||
| * Int32(4) | ||
| ) | ||
| ) | ||
| sf_vals[i] = ldg_b32_raw(sf_addr) | ||
|
|
||
| for i in cutlass.range_constexpr(0, sf_passes, 1): | ||
| j = Int32(i * self.warp_threads) + lane_idx | ||
| if j < Int32(self.sf_uint32_per_token): | ||
| sf_int32_pos = sf_atom_int32_offset( | ||
| sf_token_in_pool_axis, | ||
| j, | ||
| num_k_atoms=self.sf_uint32_per_token, | ||
| ) | ||
| fc1_input_sf_buffer[sf_int32_pos] = sf_vals[i] | ||
| cute.arch.sync_warp() | ||
|
|
||
| with cute.arch.elect_one(): | ||
| cute.arch.mbarrier_arrive_and_expect_tx( | ||
| pull_mbar_ptr + warp_idx, Int32(self.hidden_bytes) | ||
| ) | ||
| cute.arch.mbarrier_wait(pull_mbar_ptr + warp_idx, phase_bit) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Extract the forward comm's per-slot barrier/TMA ordering to confirm the expected sequence.
set -uo pipefail
F=python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/token_comm.py
ast-grep outline "$F" --items all
echo "== barrier / TMA call order in the forward pull path"
rg -n 'mbarrier_init|mbarrier_arrive_and_expect_tx|mbarrier_wait|tma_load_1d|tma_store_1d|cp_async_bulk' "$F"
echo "== same order in the reviewed backward gather"
rg -n 'mbarrier_init|mbarrier_arrive_and_expect_tx|mbarrier_wait|tma_load_1d|tma_store_1d|cp_async_bulk' \
python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.pyRepository: NVIDIA/cudnn-frontend
Length of output: 4617
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "== forward dispatch pull sequence =="
sed -n '1115,1270p' python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/token_comm.py
echo "== backward gather sequence =="
sed -n '180,275p' python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.py
echo "== documented forward sequence =="
rg -n -C 8 'arrive_and_expect_tx|tma_load_1d|per-slot|pull' \
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/README.mdRepository: NVIDIA/cudnn-frontend
Length of output: 16379
🌐 Web query:
CUDA PTX mbarrier.arrive.expect_tx ordering before cp.async.bulk TMA transaction barrier expected transaction count
💡 Result:
The mbarrier.arrive.expect_tx PTX instruction is a fused operation that performs two distinct functions to manage asynchronous memory operations, specifically those using the Tensor Memory Accelerator (TMA) [1][2][3]. 1. Transaction Expectation: It increases the barrier's "transaction debt" by a specified byte count (tx-count). This informs the barrier that it must track an additional number of asynchronous transactions (typically bytes) before the current barrier phase can be considered complete [2][4][5]. 2. Thread Arrival: It performs a standard thread arrival on the barrier, decrementing the pending-arrival count [2][4][3]. The critical ordering requirement is that mbarrier.arrive.expect_tx must be executed before the associated TMA cp.async.bulk operation is issued [1][6]. This ensures that the barrier object is properly initialized with the expected transaction count so that the TMA engine can correctly report completion as data arrives [1][4][5]. Regarding completion semantics: - The barrier phase completes only when both criteria are met: the pending-arrival count reaches zero, AND the pending transaction byte count reaches zero [7][1][2][4]. - The cp.async.bulk instructions used with TMA are typically configured with the mbarrier::complete_tx::bytes modifier [8][1][9]. As the TMA engine transfers data, it implicitly uses this modifier to perform complete_tx operations on the specified barrier, decrementing the transaction byte count [2][9][10]. - Consumers (e.g., MMA instructions) wait on the barrier phase to flip using mbarrier.try_wait or similar instructions [1][9][10]. Once the phase flips, it guarantees that all threads have arrived and all expected TMA bytes have been committed to shared memory [1][4][5]. In summary, the mbarrier.arrive.expect_tx instruction sets the "debt" that the TMA operation later "pays off" via complete_tx notifications; issuing the expectation before the transfer is essential to prevent race conditions where the hardware reports completion before the barrier is configured to track it [1][2][6].
Citations:
- 1: https://mlc.ai/modern-gpu-programming-for-mlsys/chapter_tma/index.html
- 2: https://cudacourseh100.github.io/pages/lesson-3.html
- 3: [MLIR][NVVM] Update mbarrier.arrive.expect_tx Op llvm/llvm-project#169922
- 4: https://mlc.ai/modern-gpu-programming-for-mlsys/chapter_async_barriers/index.html
- 5: https://accelerated-computing.academy/fall25/resources/tma-interface/
- 6: https://nvlabs.github.io/cuda-oxide/advanced/tensor-memory-accelerator.html
- 7: https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/async-barriers.html
- 8: https://nvidia.github.io/cccl/unstable/libcudacxx/ptx/instructions/cp_async_bulk.html
- 9: https://research.colfax-intl.com/tutorial-hopper-tma/
- 10: https://gh.evko.io/nvopen-tools/tileiras/dialects/nvvm/mbarrier-ops.html
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "== num_global_warps and token_idx definitions in backward kernel =="
rg -n -C 8 'num_global_warps|sm_count|token_idx|num_sms' \
python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.py
echo "== TokenCommArgs sm_count declarations and extraction =="
rg -n -C 10 'sm_count|__extract_mlir_values__|__extract_mlir_attributes__|_CONST_FIELDS' \
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/token_comm.py
echo "== backward constructor and launch call sites =="
rg -n -C 10 'BwdGatherComm|kernel_bwd_mega|num_sms|sm_count' \
python/cudnn/moe_ep/_megamoe_backend/megamoe \
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/srcRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "== TMA and barrier helper implementations =="
rg -n -C 20 'def tma_load_1d_raw|def tma_store_1d|mbarrier_arrive_and_expect_tx|mbarrier_wait' \
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src \
python/cudnn/moe_ep/_megamoe_backend/megamoe
echo "== focused backward class signatures and constructor values =="
sed -n '1,180p' python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.py
sed -n '280,380p' python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.py
echo "== all local transaction-order patterns =="
python3 - <<'PY'
from pathlib import Path
import re
roots = [
Path("python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src"),
Path("python/cudnn/moe_ep/_megamoe_backend/megamoe"),
]
for root in roots:
for path in root.rglob("*.py"):
text = path.read_text(errors="replace")
calls = [(m.start(), m.group(0).split("(")[0]) for m in re.finditer(
r"\b(?:tma_load_1d_raw|mbarrier_arrive_and_expect_tx|mbarrier_wait)\s*\(",
text,
)]
for (pos1, name1), (pos2, name2) in zip(calls, calls[1:]):
if name1.endswith("tma_load_1d_raw") and name2.endswith("mbarrier_arrive_and_expect_tx"):
print(f"{path}: load -> expect")
if name1.endswith("mbarrier_arrive_and_expect_tx") and name2.endswith("tma_load_1d_raw"):
print(f"{path}: expect -> load")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 48116
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "== TokenCommArgs construction sites =="
rg -n -C 12 'TokenCommArgs\s*\(' \
python/cudnn/moe_ep/_megamoe_backend
echo "== sm_count assignment and kernel argument declarations =="
rg -n -C 6 'sm_count\s*=|sm_count:' \
python/cudnn/moe_ep/_megamoe_backend \
-g '*.py'
echo "== base TokenInPullTokenBackPush constructor =="
rg -n -C 12 'class TokenInPullTokenBackPush|def __init__' \
python/cudnn/moe_ep/_megamoe_backend/cutedsl_megamoe/src/token_comm.py \
python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.pyRepository: NVIDIA/cudnn-frontend
Length of output: 49669
Register the TMA transaction count before issuing the load.
tma_load_1d_raw uses complete_tx::bytes, so mbarrier_arrive_and_expect_tx must run first. Otherwise, mbarrier_wait does not reliably guarantee that pull_buffer contains the transferred row, which can corrupt the FC1 A operand or hang. Apply the same ordering fix in the forward implementation.
🤖 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 `@python/cudnn/moe_ep/_megamoe_backend/megamoe/bwd_kernel/kernel_bwd_mega.py`
around lines 201 - 252, In the backward pull path, update the warp elected block
around tma_load_1d_raw so mbarrier_arrive_and_expect_tx registers hidden_bytes
before issuing the TMA load, while keeping mbarrier_wait after the load and
preserving the existing phase handling. Apply the identical ordering change in
the corresponding forward implementation.
2 participants
Wire the MegaMoE backend into
cudnn.MoeEpThis PR completes the
cudnn.MoeEpAPI introduced in #389 by replacing the allocate-only implementation with a real device backend.MoeEp.__call__andbackwardnow execute fused SM100 (GB200) MegaMoE megakernels:The kernel sources are now bundled inside cuDNN, so a fresh clone requires no external checkout or path environment variables.
The bundle includes:
Backend behavior
CUDNN_MOE_EP_BACKEND=auto|megamoe|noneis the only backend knob.auto(default) executes the bundled backend when supported and otherwise falls back to the allocate-only implementation from Add MoE + expert-parallel Python API surface, PyTorch reference, and tests #389 with a one-time warning.megamoeraises on backend failures.nonealways uses the stub implementation.Supported configurations:
apply_topk_in_fc1=Truegenerate_c/backward path)combine_format="bf16"for trainingInternal compute uses MXFP8, so outputs match the FP32 reference within ~1e-2–7e-2 relative tolerance.
route_metadataand quantized outputs remain bit-exact.Validation
6.5e-2; all four backward gradients ≤6.6e-2(gate0.10); metadata equal; output quantizer bit-exacttest/python/fe_api/moe_ep/run_pr_tests.shtest/python/fe_api/moe_ep/RUNBOOK.mdFiles touched
Integration
python/cudnn/moe_ep/_megamoe.pypython/cudnn/moe_ep/api.pyMoeEp.__call__/backwarddelegate to backend while preserving fallbacktest/python/fe_api/moe_ep/megamoe_backend_parity.pytest/python/fe_api/moe_ep/run_pr_tests.shtest/python/fe_api/moe_ep/RUNBOOK.mddocs/fe-oss-apis/moe_ep.md.pre-commit-config.yamlVendored backend
103 files under
python/cudnn/moe_ep/_megamoe_backend/containing the MegaMoE Python package, CuTe DSL kernels (MXFP8 + NVFP4), FP8 mega backward, PyTorch reference package, NVSHMEM infrastructure, and provenance/requirements.Summary by CodeRabbit
New Features
Documentation
Tests