Skip to content

add differentiable mxfp8 grouped gemm with dynamic quant (forward pass) #2627

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Jul 30, 2025
Merged

add differentiable mxfp8 grouped gemm with dynamic quant (forward pass) #2627

merged 1 commit into from
Jul 30, 2025

Conversation

danielvegamyhre
Copy link
Contributor

@danielvegamyhre danielvegamyhre commented Jul 29, 2025 •
edited
Loading

Stacked PRs:

add differentiable mxfp8 grouped gemm with dynamic quant (forward pass)

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jul 29, 2025 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit 5376f65 with merge base 9834869 (image):
💚 Looks good so far! There are no failures yet. 💚

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

danielvegamyhre added a commit that referenced this pull request Jul 29, 2025
@danielvegamyhre
add differentiable mxfp8 grouped gemm with dynamic quant (forward pass)
b3962a0 
stack-info: PR: #2627, branch: danielvegamyhre/stack/23
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/23 branch from 453201b to b3962a0 Compare July 29, 2025 19:45
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/22 branch from e29fb79 to fa77af5 Compare July 29, 2025 19:45
@danielvegamyhre danielvegamyhre mentioned this pull request Jul 29, 2025
Merged
@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 Jul 29, 2025
@danielvegamyhre danielvegamyhre added the topic: not user facing Use this tag if you don't want this PR to show up in release notes label Jul 29, 2025
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/22 to main July 29, 2025 19:50
danielvegamyhre added a commit that referenced this pull request Jul 29, 2025
@danielvegamyhre
add differentiable mxfp8 grouped gemm with dynamic quant (forward pass)
9792e76 
stack-info: PR: #2627, branch: danielvegamyhre/stack/23
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/23 branch from b3962a0 to 9792e76 Compare July 29, 2025 19:50
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/22 July 29, 2025 19:50
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/22 to main July 29, 2025 21:36
danielvegamyhre added a commit that referenced this pull request Jul 29, 2025
@danielvegamyhre
add differentiable mxfp8 grouped gemm with dynamic quant (forward pass)
fd92301 
stack-info: PR: #2627, branch: danielvegamyhre/stack/23
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/23 branch from 9792e76 to fd92301 Compare July 29, 2025 21:36
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/22 July 29, 2025 21:37
vkuzo
vkuzo reviewed Jul 30, 2025
View reviewed changes
torchao/prototype/moe_training/scaled_grouped_mm.py
offs: Optional[torch.Tensor] = None,
block_size: int = 32,
out_dtype: Optional[torch.dtype] = torch.bfloat16,
) -> torch.Tensor:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add an emulated flag and assert that it's True until we have a real kernel, to make the intent crystal clear?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

vkuzo
vkuzo reviewed Jul 30, 2025
View reviewed changes
test/prototype/moe_training/test_scaled_grouped_mm.py Outdated
@@ -249,3 +249,25 @@ def test_emulate_mxfp8_grouped_gemm(M, K, N, num_experts):
sqnr = compute_error(ref_out, out)
min_sqnr = 27.0
assert sqnr >= min_sqnr, f"sqnr {sqnr} is too low, must be >= {min_sqnr}"


@pytest.mark.parametrize("M", (1024, 4096))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: might be good to test one case where MKN are the same, and one where they are all different. If you want to do that and keep # of tests manageable, it would probably be iterating on MKN in one go instead of iterating on each individually.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated to parameterize M,N,K together and test "all same" and "all different" cases

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ergh my stack-pr got in a weird state, changes didn't go through somehow... let me try again

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok it's updated now.

@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/22 to main July 30, 2025 15:31
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/22 July 30, 2025 15:31
@danielvegamyhre danielvegamyhre mentioned this pull request Jul 30, 2025
Merged
danielvegamyhre added a commit that referenced this pull request Jul 30, 2025
@danielvegamyhre
add differentiable mxfp8 grouped gemm with dynamic quant (forward pass)
18991a4 
stack-info: PR: #2627, branch: danielvegamyhre/stack/23
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/23 branch from fd92301 to 18991a4 Compare July 30, 2025 15:31
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/22 to main July 30, 2025 15:31
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/23 branch from 18991a4 to fd04d1a Compare July 30, 2025 15:38
@danielvegamyhre
add differentiable mxfp8 grouped gemm with dynamic quant (forward pass)
5376f65 
stack-info: PR: #2627, branch: danielvegamyhre/stack/23
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/23 branch from fd04d1a to 5376f65 Compare July 30, 2025 18:00
@danielvegamyhre danielvegamyhre merged commit a8ca216 into main Jul 30, 2025
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@vkuzo vkuzo vkuzo approved these changes

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: not user facing Use this tag if you don't want this PR to show up in release notes
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@danielvegamyhre @vkuzo