Skip to content

[moe training] add bench script for fp8 rowwise kernels and update autotune configs #2697

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.

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

Open
wants to merge 1 commit into
base: danielvegamyhre/stack/30
Choose a base branch
from

Conversation

danielvegamyhre
Copy link
Contributor

@danielvegamyhre danielvegamyhre commented Aug 5, 2025

Stacked PRs:


[moe training] add bench script for fp8 rowwise kernels and update autotune configs

Performance vs torch.compile

It's faster for llama4 shape (16, 5120, 4*5120), but slower for skinny shapes.

input_shape            torch_time_us    triton_time_us
-------------------  ---------------  ----------------
(8, (4096, 1024))              92.32           113.248
(16, (20480, 5120))          5746.62          4075.38

There is more we can do, for example writing row major outputs was roughly 2x faster, but we need the outputs in col-major. I can probably look with NCU and figure out what's going on but for now this is a start.

Copy link

pytorch-bot bot commented Aug 5, 2025

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

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

@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from 59e34b1 to 362cfb2 Compare August 5, 2025 23:12
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/30 branch from af159db to f6688be Compare August 5, 2025 23:12
@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 Aug 5, 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 Aug 5, 2025
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 5, 2025 23:44
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from 362cfb2 to 14acee2 Compare August 5, 2025 23:44
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 5, 2025 23:44
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 5, 2025 23:56
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from 14acee2 to e3c97ec Compare August 5, 2025 23:56
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 5, 2025 23:57
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 6, 2025 00:00
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from e3c97ec to 5118bcf Compare August 6, 2025 00:00
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 6, 2025 00:01
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 6, 2025 00:13
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from 5118bcf to 0305192 Compare August 6, 2025 00:13
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 6, 2025 00:13
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 6, 2025 00:56
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from 0305192 to c621ce5 Compare August 6, 2025 00:56
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 6, 2025 00:56
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 6, 2025 01:19
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from c621ce5 to ab43973 Compare August 6, 2025 01:20
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 6, 2025 01:20
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 6, 2025 01:36
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from ab43973 to 1fb9ee1 Compare August 6, 2025 01:36
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 6, 2025 01:36
print(tabulate(rows, headers=headers))


def benchmark_cuda_function_in_microseconds(f, *args):
Copy link
Contributor

Choose a reason for hiding this comment

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

we have so many of these, maybe reuse in a separate PR?

@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/30 branch from 6704fd3 to c789281 Compare August 6, 2025 17:24
danielvegamyhre added a commit that referenced this pull request Aug 6, 2025
…totune configs

stack-info: PR: #2697, branch: danielvegamyhre/stack/31
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from 1fb9ee1 to f124235 Compare August 6, 2025 17:24
…totune configs

stack-info: PR: #2697, branch: danielvegamyhre/stack/31
@danielvegamyhre danielvegamyhre changed the base branch from danielvegamyhre/stack/30 to main August 6, 2025 17:28
@danielvegamyhre danielvegamyhre force-pushed the danielvegamyhre/stack/31 branch from f124235 to 73b26f1 Compare August 6, 2025 17:28
@danielvegamyhre danielvegamyhre changed the base branch from main to danielvegamyhre/stack/30 August 6, 2025 17:28
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. topic: not user facing Use this tag if you don't want this PR to show up in release notes
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants