Skip to content

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

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

fix(metal): revert inactive NAX MoE simdgroup skipping#4404
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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants