Skip to content

mx: small speedup with dim0 cast#1980

Merged
vkuzo merged 72 commits into
mainfrom
gh/vkuzo/86/head
Apr 1, 2025
Merged

mx: small speedup with dim0 cast#1980
vkuzo merged 72 commits into
mainfrom
gh/vkuzo/86/head

Conversation

@vkuzo

@vkuzo vkuzo commented Mar 28, 2025

Copy link
Copy Markdown
Contributor

Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415

Reviewers:

Subscribers:

Tasks:

Tags:

vkuzo added 30 commits March 21, 2025 06:59
@vkuzo
Update
af6ae2f 
[ghstack-poisoned]
@vkuzo
Update
45120de 
[ghstack-poisoned]
@vkuzo
Update
5527e72 
[ghstack-poisoned]
@vkuzo
Update
478b9e1 
[ghstack-poisoned]
@vkuzo
Update
571775d 
[ghstack-poisoned]
@vkuzo
Update
fd30558 
[ghstack-poisoned]
@vkuzo
Update
b0cd056 
[ghstack-poisoned]
@vkuzo
Update
26b49fd 
[ghstack-poisoned]
@vkuzo
Update
ba10a02 
[ghstack-poisoned]
@vkuzo
Update
483cdfd 
[ghstack-poisoned]
@vkuzo
Update
32005c9 
[ghstack-poisoned]
@vkuzo
Update
e341c2e 
[ghstack-poisoned]
@vkuzo
Update
7ecd79f 
[ghstack-poisoned]
@vkuzo
Update
ca3c4cf 
[ghstack-poisoned]
@vkuzo
Update
0de11cf 
[ghstack-poisoned]
@vkuzo
Update
912e4dc 
[ghstack-poisoned]
@vkuzo
Update
fb5662a 
[ghstack-poisoned]
@vkuzo
Update
f245d64 
[ghstack-poisoned]
@vkuzo
Update
9e5b8f8 
[ghstack-poisoned]
@vkuzo
Update
e5bdecb 
[ghstack-poisoned]
@vkuzo
Update
4c2ad8c 
[ghstack-poisoned]
@vkuzo
Update
c1ceef1 
[ghstack-poisoned]
@vkuzo
Update
65bfff0 
[ghstack-poisoned]
@vkuzo
Update
0ff3a93 
[ghstack-poisoned]
@vkuzo
Update
71a5548 
[ghstack-poisoned]
@vkuzo
Update
0576d0d 
[ghstack-poisoned]
@vkuzo
Update
f98453f 
[ghstack-poisoned]
@vkuzo
Update
81dc214 
[ghstack-poisoned]
@vkuzo
Update
5d60f24 
[ghstack-poisoned]
@vkuzo
Update
a313055 
[ghstack-poisoned]
vkuzo added 4 commits March 28, 2025 13:00
@vkuzo
Update
84c77d7 
[ghstack-poisoned]
@vkuzo
Update
749564b 
[ghstack-poisoned]
@vkuzo
Update
f63479e 
[ghstack-poisoned]
@vkuzo
Update
8b0f250 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 28, 2025
@vkuzo
mx: small speedup with dim0 cast
886e3e5 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
vkuzo added 5 commits March 28, 2025 13:02
@vkuzo
Update
c603f09 
[ghstack-poisoned]
@vkuzo
Update
42fb0e9 
[ghstack-poisoned]
@vkuzo
Update
a16f576 
[ghstack-poisoned]
@vkuzo
Update
b890654 
[ghstack-poisoned]
@vkuzo
Update
35179ab 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 28, 2025
@vkuzo
mx: small speedup with dim0 cast
0c6b205 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
vkuzo added 4 commits March 28, 2025 13:03
@vkuzo
Update
83e1e2e 
[ghstack-poisoned]
@vkuzo
Update
c41cc19 
[ghstack-poisoned]
@vkuzo
Update
8ba3018 
[ghstack-poisoned]
@vkuzo
Update
f437e00 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 28, 2025
@vkuzo
mx: small speedup with dim0 cast
ccdcb9a 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
vkuzo added 3 commits March 28, 2025 13:03
@vkuzo
Update
a62e00b 
[ghstack-poisoned]
@vkuzo
Update
195d904 
[ghstack-poisoned]
@vkuzo
Update
bdb0996 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Mar 28, 2025
@vkuzo
mx: small speedup with dim0 cast
6a4c0b3 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
vkuzo added a commit that referenced this pull request Apr 1, 2025
@vkuzo
mx: small speedup with dim0 cast
06ac6dd 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
vkuzo added 2 commits April 1, 2025 09:40
@vkuzo
Update
8a5050e 
[ghstack-poisoned]
@vkuzo
Update
635ca65 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 1, 2025
@vkuzo
mx: small speedup with dim0 cast
661643b 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
@vkuzo
Update
0d5b763 
[ghstack-poisoned]
@vkuzo vkuzo changed the base branch from gh/vkuzo/85/head to main April 1, 2025 16:41
vkuzo added a commit that referenced this pull request Apr 1, 2025
@vkuzo
mx: small speedup with dim0 cast
2feaed7 
Summary:

Removes the unnecessary cast to bfloat16 in the MX dim0 casting code.
This is a 2.6% speedup on 16k by 16k shape:
https://www.internalfb.com/phabricator/paste/view/P1769373804

Note: this PR also includes a couple of cleanups around e8m0 dtype and
NaN handling, I found them while coding this PR. Leaving them together
instead of
separate PR since they are all safe.

Test Plan:

```bash
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 152.90741052631583
mem_bw_gbps 5321.488168553876
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$
(pytorch) [vasiliy@devgpu023.atn1 ~/local/ao (20250321_mx_dim1_triton_kernel)]$ python benchmarks/mx_formats/cast_bench.py --mode dim0_mx --M 16384 --K 16384
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.8.0a0+git25309a1
triton version: 3.3.0
mode: dim0_mx
time_us 149.03950980392162
mem_bw_gbps 5459.5924065404415
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 47fb1df
ghstack-comment-id: 2762318741
Pull Request resolved: #1980
@vkuzo vkuzo merged commit aafc1ba into main Apr 1, 2025
liangel-02 pushed a commit that referenced this pull request Aug 25, 2025
@vkuzo @liangel-02
mx: small speedup with dim0 cast (#1980)
56dda90 
* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer

@HDCharles HDCharles HDCharles approved these changes

Reviewers whose approvals may not affect merge requirements

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: performance Use this tag if this PR improves the performance of a feature

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@vkuzo @HDCharles @facebook-github-bot