Skip to content

fix(metal): revert inactive NAX MoE simdgroup skipping - #4404

Closed
jonathan308 wants to merge 1 commit into
ml-explore:mainfrom
jonathan308:fix/nax-moe-simdgroup-stability
Closed

jonathan308 wants to merge 1 commit into
ml-explore:mainfrom
jonathan308:fix/nax-moe-simdgroup-stability

Conversation

@jonathan308

Copy link
Copy Markdown
Contributor

Summary

Revert #4352's inactive-simdgroup NAX MoE optimization after a physical mixed-Mac tensor-parallel A/B isolated it as the trigger for repeatable lost JACCL completions.

Evidence

The detailed physical reproducer and measurements are linked in the issue.

Why a revert

The issue is hardware/workload-specific and current CI cannot exercise a heterogeneous M3 Ultra + M5 Max JACCL pair. Restoring the pre-#4352 code is the smallest proven-safe change. A follow-up can reintroduce the optimization with a safe uniform execution scheme or an explicit capability gate.

This PR is a draft pending maintainer guidance and any additional isolated-kernel reproducer they want.

Refs #4403.

@nastya236

Copy link
Copy Markdown
Collaborator

Thanks for flagging this. I will look into it soon.

@zcbenz

zcbenz commented Aug 25, 2026

Copy link
Copy Markdown
Member

/cc @RohanGautam

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Sep 1, 2026
@zcbenz

zcbenz commented Sep 5, 2026

Copy link
Copy Markdown
Member

Since there is an issue tracking the problem and we are not going to revert the change, I'm closing this PR.

@zcbenz zcbenz closed this Sep 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants