Skip to content

Make _emulated_mxfp8_scaled_grouped_mm_2d_2d torch.compile compatible#3906

Merged
danielvegamyhre merged 2 commits intopytorch:mainfrom
graphcore-research:_emulated_mxfp8_scaled_grouped_mm_2d2d_compile_compat
Feb 19, 2026
Merged

Make _emulated_mxfp8_scaled_grouped_mm_2d_2d torch.compile compatible#3906
danielvegamyhre merged 2 commits intopytorch:mainfrom
graphcore-research:_emulated_mxfp8_scaled_grouped_mm_2d2d_compile_compat

Conversation

@lyprince
Copy link
Copy Markdown
Contributor

  • Replace loop-based dequantization in _emulated_mxfp8_scaled_grouped_mm_2d_2d with vectorized ops, removing the @torch.compiler.disable decorator
  • The previous implementation used offs.tolist() and a Python loop with data-dependent control flow, which prevented torch.compile from tracing through the emulated MXFP8 wgrad path
  • Remove the pytest.skip for EMULATED + use_compile=True in test_mxfp8_grouped_gemm_with_dq_fwd_bwd

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Feb 17, 2026

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3906

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 1fd57f5 with merge base 396d78b (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 17, 2026
@danielvegamyhre danielvegamyhre self-requested a review February 17, 2026 17:12
@danielvegamyhre danielvegamyhre added mx module: training quantize_ api training flow moe labels Feb 17, 2026
@danielvegamyhre
Copy link
Copy Markdown
Contributor

this is much better, thanks!

@lyprince
Copy link
Copy Markdown
Contributor Author

Looks like CI is blocked by PR label check?

@danielvegamyhre danielvegamyhre added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Feb 18, 2026
@danielvegamyhre danielvegamyhre merged commit 01b37b2 into pytorch:main Feb 19, 2026
22 of 23 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: training quantize_ api training flow moe mx topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants