mxtensor: add pre-swizzle support #3200
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Stack from ghstack (oldest at bottom): |
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3200
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
vkuzo
added a commit
that referenced
this pull request
Oct 17, 2025
Summary: Adds the ability to pre-swizzle scales for `MXTensor`, and turns it on for the inference workflow. For activations, this is no-change for now but if we write a fused kernel we'll hook into the pre-swizzled path. For weights, this is a performance win in this PR as now we swizzle ahead of time. Rough magnitude of the weight pre-swizzling win: on M, K, N == 4096, 4096, 4096, the inference fwd speedup on mxfp8 increases from 1.24x to 1.30x Test Plan: ```bash // correctness CUDA_VISIBLE_DEVICES=5 pytest test/prototype/mx_formats/ -s // performance CUDA_VISIBLE_DEVICES=5 python benchmarks/float8/float8_inference_roofline.py ~/local/tmp/20251017_test.csv --recipe_name mxfp8_cublas --shape_gen_name pow2_extended // before: https://www.internalfb.com/phabricator/paste/view/P1996942931 // after: https://www.internalfb.com/phabricator/paste/view/P1996941798 ``` Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 46b8d23 ghstack-comment-id: 3415966576 Pull-Request: #3200
meta-cla
bot
added
the
CLA Signed
This was referenced Oct 17, 2025
vkuzo
commented
Oct 17, 2025
| torch.float8_e5m2, | ||
| torch.uint8, | ||
| ), "unsupported" | ||
| if elem_dtype in ( |
There was a problem hiding this comment.
this code doesn't really have a strong purpose, removing instead of making it handle swizzling
vkuzo
commented
Oct 17, 2025
| return mx_tensor | ||
|
|
||
|
|
||
| def _swizzle_aware_slice( |
There was a problem hiding this comment.
extracted this out, the only things that change are the various shape calculations (fp8 vs fp4 data, 32 vs 16 block size)
vkuzo
added
the
topic: improvement
drisspg
reviewed
Oct 17, 2025
| Output: sliced qdata and scale, does the right thing for unswizzled and swizzled scales | ||
| """ | ||
|
|
||
| M, K = x.shape[0], x.shape[1] |
There was a problem hiding this comment.
Nit I should probably have used a (m/n) generic term like rows, columns
drisspg
approved these changes
Oct 17, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
Adds the ability to pre-swizzle scales for
MXTensor,and turns it on for the inference workflow.
For activations, this is no-change for now but if we write a fused
kernel we'll hook into the pre-swizzled path.
For weights, this is a performance win already as now we swizzle ahead of
time.
Rough magnitude of the weight pre-swizzling win:
on M, K, N == 4096, 4096, 4096, the inference fwd speedup on mxfp8
increases from 1.24x to 1.30x
Name of
_is_swizzled_scalesis not final, but IMO we should finalize it in a future PR together with NVFP4Tensor. For now I'm staying consistent with NVFP4Tensor.Test Plan:
// correctness CUDA_VISIBLE_DEVICES=5 pytest test/prototype/mx_formats/ -s // performance CUDA_VISIBLE_DEVICES=5 python benchmarks/float8/float8_inference_roofline.py ~/local/tmp/20251017_test.csv --recipe_name mxfp8_cublas --shape_gen_name pow2_extended // before: https://www.internalfb.com/phabricator/paste/view/P1996942931 // after: https://www.internalfb.com/phabricator/paste/view/P1996941798Reviewers:
Subscribers:
Tasks:
Tags: