Add MoE + expert-parallel Python API surface, PyTorch reference, and tests#389
Open
Anerudhan wants to merge 3 commits into
Open
Add MoE + expert-parallel Python API surface, PyTorch reference, and tests#389Anerudhan wants to merge 3 commits into
Anerudhan wants to merge 3 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>
Contributor
📝 WalkthroughWalkthroughAdds a documented MoE expert-parallel API, its Python contract implementation, a quantized PyTorch reference with forward and backward routing, and tests covering quantization, single-rank behavior, public API allocation, and distributed execution. ChangesMoE Expert-Parallel API
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant SourceRank
participant EPCollectives
participant LocalExperts
participant CombineOutput
SourceRank->>EPCollectives: Dispatch tokens and routing metadata
EPCollectives->>LocalExperts: Exchange routes by destination rank
LocalExperts->>EPCollectives: Return expert results
EPCollectives->>CombineOutput: Reduce top-k results and encode output
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Anerudhan
changed the title
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>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
test/python/fe_api/moe_ep/test_moe_ep.py (1)
624-771: 📐 Maintainability & Code Quality | 🔵 TrivialOnly a 4-rank NCCL test is present; the PR description's claimed 2-rank Gloo coverage doesn't appear anywhere in this file.
The PR objectives state tests cover "2-rank Gloo and 4-GPU NCCL expert-parallel runs," but
_distributed_worker/test_four_rank_expert_parallelonly exercisesbackend="nccl"(Line 628) with 4 ranks; there's no Gloo-backed distributed test in this file. The current cohort's own layer description also only mentions "four-rank NCCL execution," so this may simply be deferred to another commit/file not in this batch — worth confirming whether a 2-rank Gloo test is still planned here, since it's the only path that exercises the CPU-staging branch of_collective_device/_all_to_allinmoe_ep_reference.py.Want me to draft a 2-rank Gloo variant of
_distributed_worker(spawning withbackend="gloo"and no CUDA requirement) to close this gap?🤖 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 624 - 771, Add a 2-rank Gloo distributed test alongside _distributed_worker and test_four_rank_expert_parallel, using a CPU-only execution path that exercises the CPU-staging behavior in _collective_device/_all_to_all. Parameterize or adapt the worker to select the backend and device setup, preserve the existing 4-rank NCCL test, and add a Gloo test entry that spawns two ranks without requiring CUDA.test/python/fe_api/moe_ep/moe_ep_reference.py (1)
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
formatparameter shadows the builtin in three places (Ruff A002). Same root cause across all three sites: a parameter is literally namedformat. No functional impact today (the builtin isn't called in any of these bodies), but it's flagged as a Ruff error and could fail a lint gate.
test/python/fe_api/moe_ep/moe_ep_reference.py#L207-270: rename thequantize_blockwiseparameterformat(Line 209) to e.g.fmt, updating the internal_parse_format(format)call and all in-function references.test/python/fe_api/moe_ep/moe_ep_reference.py#L313-316: rename the_format_round_tripparameterformat(Line 313) tofmtand update its body accordingly.test/python/fe_api/moe_ep/test_moe_ep.py#L68-75: rename thetest_block_scaled_round_tripparameterformat(Line 75) tofmt/moe_format, and update thepytest.mark.parametrize("format,...")string (Line 69) to match.🔧 Example fix for the production function
def quantize_blockwise( tensor: torch.Tensor, - format: Union[MoeFormat, str], + fmt_arg: Union[MoeFormat, str], *, axis: int = -1, ) -> BlockScaledTensor: ... - fmt = _parse_format(format) + fmt = _parse_format(fmt_arg)🤖 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/moe_ep_reference.py` at line 1, Rename the format parameters that trigger Ruff A002 in quantize_blockwise, _format_round_trip, and test_block_scaled_round_trip to fmt or moe_format, updating each function’s internal references and the corresponding pytest.mark.parametrize argument name so all callers and bindings remain consistent.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/__init__.py`:
- Around line 281-283: Update the lazy export mapping in the cudnn package
initializer to expose BlockScaledTensor and MoeTensor from .moe_ep alongside
MoeEp and MoeFormat, preserving the existing lazy-import pattern. Add a focused
root-import test verifying both symbols can be imported from cudnn.
In `@python/cudnn/moe_ep/api.py`:
- Around line 53-84: Update BlockScaledTensor.__post_init__ to validate that
data.device and scale.device match, raising a clear ValueError when they differ;
retain the existing normalization behavior and device property.
- Around line 257-329: Validate all non-sentinel values in topk_idx against [-1,
self.num_experts) during __call__ before the early !self.generate_c return.
Reuse or factor the existing range validation from _count_local_routes so both
generate_c paths preserve identical validation without duplicating logic.
---
Nitpick comments:
In `@test/python/fe_api/moe_ep/moe_ep_reference.py`:
- Line 1: Rename the format parameters that trigger Ruff A002 in
quantize_blockwise, _format_round_trip, and test_block_scaled_round_trip to fmt
or moe_format, updating each function’s internal references and the
corresponding pytest.mark.parametrize argument name so all callers and bindings
remain consistent.
In `@test/python/fe_api/moe_ep/test_moe_ep.py`:
- Around line 624-771: Add a 2-rank Gloo distributed test alongside
_distributed_worker and test_four_rank_expert_parallel, using a CPU-only
execution path that exercises the CPU-staging behavior in
_collective_device/_all_to_all. Parameterize or adapt the worker to select the
backend and device setup, preserve the existing 4-rank NCCL test, and add a Gloo
test entry that spawns two ranks without requiring CUDA.
🪄 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: ce53cbf6-e65a-4ef6-bfa1-855baf8a4775
📒 Files selected for processing (7)
docs/fe-oss-apis/moe_ep.mddocs/fe-oss-apis/overview.mdpython/cudnn/__init__.pypython/cudnn/moe_ep/__init__.pypython/cudnn/moe_ep/api.pytest/python/fe_api/moe_ep/moe_ep_reference.pytest/python/fe_api/moe_ep/test_moe_ep.py
Comment on lines
+281
to
+283
| "moe_ep": (".moe_ep", None), | ||
| "MoeEp": (".moe_ep", "MoeEp"), | ||
| "MoeFormat": (".moe_ep", "MoeFormat"), |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Repository root: '; pwd
printf '\nFiles:\n'
git ls-files 'python/cudnn/**' | sed -n '1,120p'
printf '\nRelevant files:\n'
git ls-files 'python/cudnn/__init__.py' 'python/cudnn/moe_ep/__init__.py' 'test/python/fe_api/**' 2>/dev/null || trueRepository: NVIDIA/cudnn-frontend
Length of output: 10964
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Repository root: '; pwd
printf '\npython/cudnn/__init__.py outline:\n'
ast-grep outline python/cudnn/__init__.py --view expanded || true
printf '\npython/cudnn/moe_ep/__init__.py outline:\n'
ast-grep outline python/cudnn/moe_ep/__init__.py --view expanded || true
printf '\nfe_api tests mentioning moe_ep / BlockScaledTensor / MoeTensor:\n'
rg -n "BlockScaledTensor|MoeTensor|MoeEp|moe_ep" test/python/fe_api python/cudnn -g'*.py' || trueRepository: NVIDIA/cudnn-frontend
Length of output: 7902
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Searching for root-level imports and docs for moe_ep symbols...\n'
rg -n "from cudnn import (BlockScaledTensor|MoeTensor|MoeEp|MoeFormat)|import .*BlockScaledTensor|import .*MoeTensor|cudnn\.BlockScaledTensor|cudnn\.MoeTensor" \
python test -g'*.py' -g'*.md' || true
printf '\nTop-level lazy import mapping around moe_ep:\n'
sed -n '275,292p' python/cudnn/__init__.py
printf '\nMoeEp package init:\n'
sed -n '1,40p' python/cudnn/moe_ep/__init__.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1875
Expose BlockScaledTensor and MoeTensor from cudnn if they’re public API. python/cudnn/moe_ep/__init__.py exports both, but python/cudnn/__init__.py only wires up moe_ep, MoeEp, and MoeFormat, so from cudnn import BlockScaledTensor/MoeTensor will fail today. Add the missing lazy imports and a small root-import test if these symbols are meant to be user-facing.
🤖 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/__init__.py` around lines 281 - 283, Update the lazy export
mapping in the cudnn package initializer to expose BlockScaledTensor and
MoeTensor from .moe_ep alongside MoeEp and MoeFormat, preserving the existing
lazy-import pattern. Add a focused root-import test verifying both symbols can
be imported from cudnn.
Comment on lines
+53
to
+84
| @dataclass(frozen=True) | ||
| class BlockScaledTensor: | ||
| """Data-plus-scale result returned for MXFP8 and NVFP4 outputs.""" | ||
|
|
||
| data: torch.Tensor | ||
| scale: torch.Tensor | ||
| format: Union[MoeFormat, str] | ||
| logical_shape: Tuple[int, ...] | ||
| axis: int = -1 | ||
|
|
||
| def __post_init__(self) -> None: | ||
| fmt = _parse_format(self.format) | ||
| if fmt is MoeFormat.BF16: | ||
| raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") | ||
| logical_shape = tuple(int(dim) for dim in self.logical_shape) | ||
| axis = _normalize_axis(self.axis, len(logical_shape)) | ||
| object.__setattr__(self, "format", fmt) | ||
| object.__setattr__(self, "logical_shape", logical_shape) | ||
| object.__setattr__(self, "axis", axis) | ||
|
|
||
| @property | ||
| def shape(self) -> Tuple[int, ...]: | ||
| return self.logical_shape | ||
|
|
||
| @property | ||
| def device(self) -> torch.device: | ||
| return self.data.device | ||
|
|
||
| @property | ||
| def block_size(self) -> int: | ||
| return 32 if self.format is MoeFormat.MXFP8 else 16 | ||
|
|
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
BlockScaledTensor doesn't enforce the documented data/scale device-match contract.
moe_ep.md states "Quantized data and its scale tensor must also share a device," but __post_init__ never checks self.data.device == self.scale.device, and the .device property (lines 77-79) only reflects data.device, so MoeEp.__call__'s device-consistency loop (lines 314-322) can't catch a mismatched scale on an input BlockScaledTensor.
🛡️ Proposed fix
def __post_init__(self) -> None:
fmt = _parse_format(self.format)
if fmt is MoeFormat.BF16:
raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4")
+ if self.data.device != self.scale.device:
+ raise ValueError(f"data device {self.data.device} does not match scale device {self.scale.device}")
logical_shape = tuple(int(dim) for dim in self.logical_shape)📝 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.
Suggested change
| @dataclass(frozen=True) | |
| class BlockScaledTensor: | |
| """Data-plus-scale result returned for MXFP8 and NVFP4 outputs.""" | |
| data: torch.Tensor | |
| scale: torch.Tensor | |
| format: Union[MoeFormat, str] | |
| logical_shape: Tuple[int, ...] | |
| axis: int = -1 | |
| def __post_init__(self) -> None: | |
| fmt = _parse_format(self.format) | |
| if fmt is MoeFormat.BF16: | |
| raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") | |
| logical_shape = tuple(int(dim) for dim in self.logical_shape) | |
| axis = _normalize_axis(self.axis, len(logical_shape)) | |
| object.__setattr__(self, "format", fmt) | |
| object.__setattr__(self, "logical_shape", logical_shape) | |
| object.__setattr__(self, "axis", axis) | |
| @property | |
| def shape(self) -> Tuple[int, ...]: | |
| return self.logical_shape | |
| @property | |
| def device(self) -> torch.device: | |
| return self.data.device | |
| @property | |
| def block_size(self) -> int: | |
| return 32 if self.format is MoeFormat.MXFP8 else 16 | |
| `@dataclass`(frozen=True) | |
| class BlockScaledTensor: | |
| """Data-plus-scale result returned for MXFP8 and NVFP4 outputs.""" | |
| data: torch.Tensor | |
| scale: torch.Tensor | |
| format: Union[MoeFormat, str] | |
| logical_shape: Tuple[int, ...] | |
| axis: int = -1 | |
| def __post_init__(self) -> None: | |
| fmt = _parse_format(self.format) | |
| if fmt is MoeFormat.BF16: | |
| raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") | |
| if self.data.device != self.scale.device: | |
| raise ValueError(f"data device {self.data.device} does not match scale device {self.scale.device}") | |
| logical_shape = tuple(int(dim) for dim in self.logical_shape) | |
| axis = _normalize_axis(self.axis, len(logical_shape)) | |
| object.__setattr__(self, "format", fmt) | |
| object.__setattr__(self, "logical_shape", logical_shape) | |
| object.__setattr__(self, "axis", axis) | |
| `@property` | |
| def shape(self) -> Tuple[int, ...]: | |
| return self.logical_shape | |
| `@property` | |
| def device(self) -> torch.device: | |
| return self.data.device | |
| `@property` | |
| def block_size(self) -> int: | |
| return 32 if self.format is MoeFormat.MXFP8 else 16 |
🤖 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 53 - 84, Update
BlockScaledTensor.__post_init__ to validate that data.device and scale.device
match, raising a clear ValueError when they differ; retain the existing
normalization behavior and device property.
Comment on lines
+257
to
+329
| def _count_local_routes(self, topk_idx: torch.Tensor) -> int: | ||
| """Number of valid routes this rank's experts receive. | ||
|
|
||
| Data-dependent: single-rank counts locally, EP exchanges per-rank | ||
| route counts (the same exchange the device dispatch performs). | ||
| """ | ||
|
|
||
| flat = topk_idx.reshape(-1).to(torch.int64) | ||
| expert = flat[flat != -1] | ||
| if expert.numel() > 0 and bool(((expert < 0) | (expert >= self.num_experts)).any().item()): | ||
| raise ValueError("topk_idx contains out-of-range expert ids") | ||
| if self.ep_size == 1: | ||
| return int(expert.numel()) | ||
| destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") | ||
| send_counts = torch.bincount(destination, minlength=self.ep_size) | ||
| if send_counts.device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": | ||
| send_counts = send_counts.cpu() | ||
| recv_counts = torch.empty_like(send_counts) | ||
| dist.all_to_all_single(recv_counts, send_counts, group=self.ep_group) | ||
| return int(recv_counts.sum().item()) | ||
|
|
||
| def __call__( | ||
| self, | ||
| activation: MoeTensor, | ||
| fc1_weight: MoeTensor, | ||
| fc2_weight: MoeTensor, | ||
| topk_idx: torch.Tensor, | ||
| topk_weights: torch.Tensor, | ||
| ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: | ||
| """Allocate the result for a future backend MoE+EP launch. | ||
|
|
||
| Expected logical shapes are ``activation=(T,H)``, | ||
| ``fc1_weight=(E_local,H,2I)``, ``fc2_weight=(E_local,I,H)``, and | ||
| ``topk_idx=topk_weights=(T,K)``. | ||
|
|
||
| Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` | ||
| when the operator was constructed with ``generate_c=True``. | ||
| """ | ||
|
|
||
| activation_shape = _logical_shape(activation) | ||
| if len(activation_shape) != 2 or activation_shape[1] != self.hidden_size: | ||
| raise ValueError(f"activation logical shape must be (T, {self.hidden_size}), got {activation_shape}") | ||
| token_count = activation_shape[0] | ||
| expected_fc1 = (self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size) | ||
| expected_fc2 = (self.experts_per_rank, self.intermediate_size, self.hidden_size) | ||
| if _logical_shape(fc1_weight) != expected_fc1: | ||
| raise ValueError(f"fc1_weight logical shape must be {expected_fc1}") | ||
| if _logical_shape(fc2_weight) != expected_fc2: | ||
| raise ValueError(f"fc2_weight logical shape must be {expected_fc2}") | ||
| route_shape = (token_count, self.top_k) | ||
| if tuple(topk_idx.shape) != route_shape: | ||
| raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") | ||
| if tuple(topk_weights.shape) != route_shape: | ||
| raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") | ||
| if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: | ||
| raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") | ||
|
|
||
| device = _tensor_device(activation) | ||
| for name, tensor in ( | ||
| ("fc1_weight", fc1_weight), | ||
| ("fc2_weight", fc2_weight), | ||
| ("topk_idx", topk_idx), | ||
| ("topk_weights", topk_weights), | ||
| ): | ||
| if _tensor_device(tensor) != device: | ||
| raise ValueError(f"{name} must be on {device}, got {_tensor_device(tensor)}") | ||
|
|
||
| # TODO: dispatch routes, launch local expert FC1/SwiGLU/FC2, return | ||
| # contributions, reduce top-k, and write this allocation. | ||
| output = self._allocate_output(token_count, device) | ||
| if not self.generate_c: | ||
| return output | ||
| local_routes = self._count_local_routes(topk_idx) |
Contributor
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/api.py test/python/fe_api | sed -n '1,200p'
printf '\n## Relevant source around the reviewed lines\n'
sed -n '240,350p' python/cudnn/moe_ep/api.py
printf '\n## Documentation references for topk_idx contract\n'
rg -n "topk_idx|unused route|out-of-range|route" python/cudnn -g '*.md' -g '*.py'
printf '\n## Tests in test/python/fe_api related to moe_ep\n'
rg -n "moe_ep|MoeEp|topk_idx|generate_c|out-of-range" test/python/fe_api -g '*.py'Repository: NVIDIA/cudnn-frontend
Length of output: 42868
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## Outline of python/cudnn/moe_ep/api.py\n'
ast-grep outline python/cudnn/moe_ep/api.py --view expanded
printf '\n## Outline of test/python/fe_api/moe_ep/test_moe_ep.py if present\n'
if [ -f test/python/fe_api/moe_ep/test_moe_ep.py ]; then
ast-grep outline test/python/fe_api/moe_ep/test_moe_ep.py --view expanded
fiRepository: NVIDIA/cudnn-frontend
Length of output: 2696
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('python/cudnn/moe_ep/api.py')
text = p.read_text()
for needle in ['def _count_local_routes', 'def __call__', 'class BlockScaledTensor', 'def __post_init__']:
idx = text.find(needle)
print(f'\n## {needle} @ {idx}')
if idx != -1:
start = max(0, text.rfind('\n', 0, idx-200))
end = min(len(text), idx+1200)
print(text[start:end])
PYRepository: NVIDIA/cudnn-frontend
Length of output: 5962
Validate topk_idx before the generate_c branch
_count_local_routes() is the only place that rejects expert ids outside [-1, num_experts), and __call__() only reaches it when generate_c=True. That leaves the default path accepting malformed topk_idx values instead of raising. Move the range check ahead of the early return, or factor it into a shared helper.
🤖 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 257 - 329, Validate all non-sentinel
values in topk_idx against [-1, self.num_experts) during __call__ before the
early !self.generate_c return. Reuse or factor the existing range validation
from _count_local_routes so both generate_c paths preserve identical validation
without duplicating logic.
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>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@python/cudnn/moe_ep/api.py`:
- Around line 358-371: Update backward to mirror __call__’s input validation
before deriving device from activation: require activation, grad_output,
fc1_weight, fc2_weight, topk_idx, topk_weights, fc1_c, and route_metadata to
share that device. Validate the re-supplied tensor shapes and dtypes
consistently with __call__, including bf16 fc1_c and int32 route_metadata. Use
_count_local_routes with topk_idx to ensure route_metadata.shape[0] matches the
valid-route count before backend execution.
🪄 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: c7dce11e-fcaa-4220-bd36-637479513ff6
📒 Files selected for processing (3)
docs/fe-oss-apis/moe_ep.mdpython/cudnn/moe_ep/api.pytest/python/fe_api/moe_ep/test_moe_ep.py
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/fe-oss-apis/moe_ep.md
- test/python/fe_api/moe_ep/test_moe_ep.py
Comment on lines
+358
to
+371
| if not self.generate_c: | ||
| raise RuntimeError("backward requires the operator to be constructed with generate_c=True") | ||
| token_count = topk_idx.shape[0] | ||
| if tuple(grad_output.shape) != (token_count, self.hidden_size): | ||
| raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") | ||
| if not grad_output.is_floating_point(): | ||
| raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") | ||
| two_i = 2 * self.intermediate_size | ||
| if route_metadata.ndim != 2 or route_metadata.shape[1] != 4: | ||
| raise ValueError(f"route_metadata shape must be (local_routes, 4), got {tuple(route_metadata.shape)}") | ||
| if tuple(fc1_c.shape) != (int(route_metadata.shape[0]), two_i): | ||
| raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") | ||
|
|
||
| device = _tensor_device(activation) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
backward skips the device-consistency check __call__ enforces.
__call__ validates every input tensor is on the same device (lines 314-322), but backward derives device solely from activation and never checks grad_output, fc1_weight, fc2_weight, topk_idx, topk_weights, fc1_c, or route_metadata against it. A caller who passes a stash/tensor on the wrong device won't be caught here; it will only surface once the real backend launch consumes these tensors together.
Also, unlike __call__, backward doesn't validate the shapes of the re-supplied activation/fc1_weight/fc2_weight/topk_idx/topk_weights, nor the dtypes of fc1_c (bf16 per forward) / route_metadata (int32 per forward), nor that route_metadata.shape[0] actually matches the valid-route count implied by topk_idx (forward's _count_local_routes already computes this).
🛡️ Proposed fix for device consistency
if tuple(fc1_c.shape) != (int(route_metadata.shape[0]), two_i):
raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}")
device = _tensor_device(activation)
+ for name, tensor in (
+ ("fc1_weight", fc1_weight),
+ ("fc2_weight", fc2_weight),
+ ("topk_idx", topk_idx),
+ ("topk_weights", topk_weights),
+ ("fc1_c", fc1_c),
+ ("route_metadata", route_metadata),
+ ("grad_output", grad_output),
+ ):
+ if _tensor_device(tensor) != device:
+ raise ValueError(f"{name} must be on {device}, got {_tensor_device(tensor)}")
# TODO: re-dispatch grad_output rows, recompute SwiGLU from fc1_c,📝 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.
Suggested change
| if not self.generate_c: | |
| raise RuntimeError("backward requires the operator to be constructed with generate_c=True") | |
| token_count = topk_idx.shape[0] | |
| if tuple(grad_output.shape) != (token_count, self.hidden_size): | |
| raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") | |
| if not grad_output.is_floating_point(): | |
| raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") | |
| two_i = 2 * self.intermediate_size | |
| if route_metadata.ndim != 2 or route_metadata.shape[1] != 4: | |
| raise ValueError(f"route_metadata shape must be (local_routes, 4), got {tuple(route_metadata.shape)}") | |
| if tuple(fc1_c.shape) != (int(route_metadata.shape[0]), two_i): | |
| raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") | |
| device = _tensor_device(activation) | |
| if not self.generate_c: | |
| raise RuntimeError("backward requires the operator to be constructed with generate_c=True") | |
| token_count = topk_idx.shape[0] | |
| if tuple(grad_output.shape) != (token_count, self.hidden_size): | |
| raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") | |
| if not grad_output.is_floating_point(): | |
| raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") | |
| two_i = 2 * self.intermediate_size | |
| if route_metadata.ndim != 2 or route_metadata.shape[1] != 4: | |
| raise ValueError(f"route_metadata shape must be (local_routes, 4), got {tuple(route_metadata.shape)}") | |
| if tuple(fc1_c.shape) != (int(route_metadata.shape[0]), two_i): | |
| raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") | |
| device = _tensor_device(activation) | |
| for name, tensor in ( | |
| ("fc1_weight", fc1_weight), | |
| ("fc2_weight", fc2_weight), | |
| ("topk_idx", topk_idx), | |
| ("topk_weights", topk_weights), | |
| ("fc1_c", fc1_c), | |
| ("route_metadata", route_metadata), | |
| ("grad_output", grad_output), | |
| ): | |
| if _tensor_device(tensor) != device: | |
| raise ValueError(f"{name} must be on {device}, got {_tensor_device(tensor)}") |
🤖 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 358 - 371, Update backward to mirror
__call__’s input validation before deriving device from activation: require
activation, grad_output, fc1_weight, fc2_weight, topk_idx, topk_weights, fc1_c,
and route_metadata to share that device. Validate the re-supplied tensor shapes
and dtypes consistently with __call__, including bf16 fc1_c and int32
route_metadata. Use _count_local_routes with topk_idx to ensure
route_metadata.shape[0] matches the valid-route count before backend execution.
Anerudhan
commented
Jul 14, 2026
| top_k=top_k, | ||
| generate_c=True, | ||
| ) | ||
| api = MoeEp(**kwargs) |
Collaborator
Author
There was a problem hiding this comment.
Usage and reference
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Proposes the public contract for a fused SwiGLU MoE operator with expert parallelism (
cudnn.MoeEp), ahead of the device implementation:python/cudnn/moe_ep/—MoeEp(plain class: constructor +__call__+backward) with bf16, mxfp8, and nvfp4 public outputs. Quantized operands travel asBlockScaledTensor(data + scale bundled, logical/unswizzled layouts). Both__call__andbackwardcurrently allocate correctly-shaped results without launching a kernel.test/python/fe_api/moe_ep/moe_ep_reference.py— executable pure-PyTorch semantic reference: routing (contiguous expert sharding,-1= dropped route), variable-sizeall_to_all_singleEP dispatch, local SwiGLU experts (silu(gate) * up, optional DeepSeek-style clamp, router weight in FC1 or FC2 epilogue), Form-A per-top-k combine with quantized-combine round-trips, and block-scaled output encoding.generate_c=Truereturns(output, fc1_c, route_metadata): the expert-rank-local pre-SwiGLU BF16 accumulator and its row-aligned(local_expert, src_rank, src_token, src_slot)routing words.backward(grad_output, forward inputs..., fc1_c, route_metadata)returns FP32(grad_activation, grad_fc1_weight, grad_fc2_weight, grad_topk_weights). It is a collective (every EP rank participates):MoeEpReference.backwardre-dispatches gradient rows along the identical forward routes, recomputes SwiGLU from thefc1_cstash (no post-SwiGLU intermediate is saved), reconstructs the row permutation purely fromroute_metadata, and returns/scatters route gradients to their source(token, slot). Quantization round-trips are straight-through.docs/fe-oss-apis/moe_ep.md— design proposal: mapping to the MegaMoE kernel interface (why workspaces, peer pointer mappers, streams, and launch config are not public arguments), thefc1_c/route_metadatacontract, the backward call contract (argument/return tables), the saved-tensors-for-backward table, and the numerical conventions a fused kernel must match (straight-through quantization round-trips, bf16 stash recompute, inclusive clamp bounds).Testing
pytest test/python/fe_api/moe_ep— 25 tests:apply_topk_in_fc1modes; clamp).-1routes.fc1_c,route_metadata, and backward gradients cross-checked from deterministic seeds); skips below 4 GPUs.generate_c, and backward) that flip to hard failures once the device kernels are wired.Verified on a single-GPU SM100 machine and on 8x B200 nodes.
🤖 Generated with Claude Code