Skip to content

enable inline asm for mxfp8+rceil+compile+cuda+sm100 - #4282

Merged
vkuzo merged 2 commits into
mainfrom
gh/vkuzo/247/head
Apr 15, 2026
Merged

enable inline asm for mxfp8+rceil+compile+cuda+sm100#4282
vkuzo merged 2 commits into
mainfrom
gh/vkuzo/247/head

Conversation

@vkuzo

@vkuzo vkuzo commented Apr 15, 2026

Copy link
Copy Markdown
Contributor

Summary:

PyTorch Core support for compile fusing inline_asm_elementwise into
surrounding ops landed in pytorch/pytorch#177922

This PR modifies the mxfp8 RCEIL cast to use the
cvt.rp.satfinite.ue8m0x2.f32 instruction when the hardware supports
it.

Cast mem bw 3.7 TB/s -> 4.1 TB/s

Test Plan:

// correctness
pytest test/prototype/mx_formats -s

// cast perf
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 python benchmarks/mx_formats/cast_bench.py --mode dim0_mxfp8_rceil
...
// before
mem_bw_gbps 3676.9529733517074
// after
mem_bw_gbps 4116.556410292409

// microbenchmark perf - see PR doc update, we see small speedups across
// the board on relu->linear benchmark

[ghstack-poisoned]
@vkuzo

vkuzo commented Apr 15, 2026

Copy link
Copy Markdown
Contributor Author

Stack from ghstack (oldest at bottom):

vkuzo added a commit that referenced this pull request Apr 15, 2026
Summary:

PyTorch Core support for compile fusing `inline_asm_elementwise` into
surrounding ops landed in pytorch/pytorch#177922

This PR modifies the mxfp8 RCEIL cast to use the
`cvt.rp.satfinite.ue8m0x2.f32` instruction when the hardware supports
it.

Cast mem bw 3.7 TB/s -> 4.1 TB/s

Test Plan:

```
// correctness
pytest test/prototype/mx_formats -s

// cast perf
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 python benchmarks/mx_formats/cast_bench.py --mode dim0_mxfp8_rceil
...
// before
mem_bw_gbps 3676.9529733517074
// after
mem_bw_gbps 4116.556410292409

// microbenchmark perf - see PR doc update, we see small speedups across
// the board on relu->linear benchmark
```
ghstack-source-id: be9a5e2
ghstack-comment-id: 4252680990
Pull-Request: #4282
@pytorch-bot

pytorch-bot Bot commented Apr 15, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

✅ No Failures

As of commit f11cf9d with merge base 00ef369 (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 Apr 15, 2026
# `torch.compiler.is_compiling()` returns False.
# Note that we need both checks, as neither of them work in both benchmark
# and e2e use cases by themselves.
(torch.compiler.is_compiling() or is_fake(descale))

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

this is not super intuitive, but it works

[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 15, 2026
Summary:

PyTorch Core support for compile fusing `inline_asm_elementwise` into
surrounding ops landed in pytorch/pytorch#177922

This PR modifies the mxfp8 RCEIL cast to use the
`cvt.rp.satfinite.ue8m0x2.f32` instruction when the hardware supports
it.

Cast mem bw 3.7 TB/s -> 4.1 TB/s

Test Plan:

```
// correctness
pytest test/prototype/mx_formats -s

// cast perf
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 python benchmarks/mx_formats/cast_bench.py --mode dim0_mxfp8_rceil
...
// before
mem_bw_gbps 3676.9529733517074
// after
mem_bw_gbps 4116.556410292409

// microbenchmark perf - see PR doc update, we see small speedups across
// the board on relu->linear benchmark
```
ghstack-source-id: e220999
ghstack-comment-id: 4252680990
Pull-Request: #4282
@vkuzo vkuzo added the module: inference quantize_ api inference flow label Apr 15, 2026
@vkuzo
vkuzo merged commit c06e481 into main Apr 15, 2026
64 of 66 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: inference quantize_ api inference flow

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants