[arcusbuilds]

Description

The fused MoE auxiliary loss kernels use 32-bit flattened offsets. Reject shapes above INT_MAX in both forward launchers and the backward launcher before launching a kernel.

Fixes #3568

Type of change

• Bug fix (non-breaking change which fixes an issue)

Changes

• Check num_rows * num_cols with 64-bit arithmetic in all three launchers.
• Use 64-bit loop counters and grid-size math so shapes right at the limit don't overflow. Flattened offsets stay 32-bit.
• Tests: an oversized shape is rejected on each path, and a shape of exactly INT_MAX elements gives the right result on each path.

Verification

• pre-commit on both changed files: all hooks passed.
• Built the real fused_moe_aux_loss.cu against stub headers and ran it on an RTX 2000 Ada (8 GB). At INT_MAX x 1, before the counter change both forward paths hit an illegal memory access and backward failed to launch. After it, all three return the right result. Kernel times in ncu at base clocks are within ±1.3% of the int32 version.
• A full TE build was not available locally (fails on Windows in existing __uint128_t code), so the new PyTorch tests haven't run yet. CI covers that.

Checklist

• I have read and followed the contributing guidelines.
• The functionality is complete.
• I have added a test for the fix.
• New and existing unit tests pass locally with my changes (see Verification).

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge.

Summary

The PR guards fused MoE auxiliary loss shapes and fixes arithmetic at the largest accepted shape.

• Fused MoE loss calls reject inputs with more than INT_MAX elements.

The previous overflow finding is fixed. No new actionable issues were found.

_{Reviews (7) · Last reviewed commit: "Use 64-bit loop counters in fused MoE au..." · Reviewed by Greptile}

[denera]

Thanks for the PR!

It looks good to me overall. Just two nits:

1. Unnecessary skip condition on CUDA GPU presence.
2. fused_moe_aux_loss.cu needs #include <climits> for INT_MAX. This currently comes in transitively from another header, but as a matter of code style and header safety, we include it explicitly in other fused MoE/top-K files too. This one should adopt the same approach.

[ptrendx]

Actually, I have a different concern - if the problem is that the shapes could exceed the int32 maximum then we should actually make the kernel work with int64/size_t for the variables that could overflow.

[arcusbuilds]

Extended the kernels to use int64_t indexing instead of rejecting shapes above INT_MAX.

[denera]

Actually, I have a different concern - if the problem is that the shapes could exceed the int32 maximum then we should actually make the kernel work with int64/size_t for the variables that could overflow.

@ptrendx I took a closer look across the entire fused router stack of kernels and I would argue that #3568 is NOT actually a real bug.

The probs that feed into our aux loss kernel come from top-k and score kernels that are explicitly guarded to reject num_tokens * num_experts > INT_MAX. Consequently, there is no actual risk of the shapes overflowing unless our aux loss kernel is used together with external non-TE fused router components that do not guard the shapes like ours do.

The reported bug says:

I found this through source inspection and have not run the oversized kernel, since an overflowing offset may access memory outside the tensor.

So this bug was never actually observed in execution, and it couldn't have been in a pure-TE implementation anyway because our top-k and score kernels will never feed a shape into aux loss that would overflow.

If we want to make our fused router accept num_tokens * num_experts > INT_MAX and shift the boundary to maximum representable int64, that would require touching every fused router kernel, not just aux loss.

Otherwise, I think the right approach here is to simply add the same INT_MAX size guard into the aux loss kernel that we already have in other fused router kernels.

---

@arcusbuilds In local testing, I'm seeing a forward pass perf difference between the old int32 and the new int64 implementations. I believe this is because the compiler is unrolling the int64 loop but not unrolling the int32. Could you add #pragma unroll 4 to the forward inner loop so we don't have to leave this up to the compiler in the future regardless of whether we stick with int64 or go back to int32 per my comments above?

[arcusbuilds]

@denera Added #pragma unroll 4 to the inner row loop in both forward kernels.

[denera]

@arcusbuilds After internal discussion about this, we're going to move forward with the original recommendation of rejecting num_tokens * num_experts > INT_MAX in the aux loss just like we do in all the other fused router kernels. Sorry for the back and forth about this! Would appreciate it if you could revert the int64 change and bring back the assert (we should keep #pragma unroll 4 in the inner row loop though regardless).

Effort to support int64 max sizes in the fused router should be deferred to a separate PR because it's going to have to touch ALL of the fused router kernels, not just aux loss.

[arcusbuilds]

@denera Reverted the int64 change and brought back the INT_MAX assert in all three launchers, matching the other fused router kernels. #pragma unroll 4 is still there.

While testing the limit I found that a shape like INT_MAX x 1 passes the guard but still overflows: the int loop counters wrap on their last increment, and so does the backward grid-size math. So I made the loop counters and grid math 64-bit and kept the flattened offsets 32-bit, since the guard already bounds them. In ncu at base clocks on an RTX 2000 Ada this is within ±1.3% of the int32 version at 16384x128, 65536x256 and 262144x64. Making the offsets 64-bit as well cost 15-23% in backward, so I left them alone. I also added a test at exactly INT_MAX elements.

[denera]

LGTM! We'll merge after a clean pipeline (which is going to be slightly delayed because of some ongoing infra issues at the moment).