Skip to content

feat: cutlass fp8 gemm bringup for SM120 & SM121 #1610

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 10 commits into from
Sep 4, 2025
Merged

feat: cutlass fp8 gemm bringup for SM120 & SM121 #1610

merged 10 commits into from
Sep 4, 2025

Conversation

yongwww
Copy link
Collaborator

@yongwww yongwww commented Aug 30, 2025 •
edited
Loading

📌 Description

It depends on #1608, mainly the cutlass fp8 gemm support for sm120/121, will rebase after #1608 lands.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit (or used your preferred method).
  • I have installed the hooks with pre-commit install.
  • I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

  • Tests have been added or updated as needed.
  • All tests are passing (unittest, etc.).

Reviewer Notes

nvmbreughe
nvmbreughe reviewed Sep 2, 2025
View reviewed changes
flashinfer/gemm.py
self,
inputs: List[torch.Tensor],
tactic: int = -1,
do_preparation: bool = False,
Copy link
Contributor

Choose a reason for hiding this comment

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

seems like this parameter is unused?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Good catch. They’re part of TunableRunner, keeping them for consistency with the others.

@yongwww yongwww force-pushed the sm120_cutlass_fp8_gemm branch from 296583e to c172dfd Compare September 3, 2025 00:17
@yongwww yongwww marked this pull request as ready for review September 3, 2025 00:29
@aleozlx
Copy link
Collaborator

aleozlx commented Sep 3, 2025

looks good no further comments from me

yzh119
yzh119 reviewed Sep 3, 2025
View reviewed changes
csrc/gemm_groupwise_sm120.cu Outdated
constexpr int SCALE_GRANULARITY_M = 1; /* Always 1 for SM120 */ \
constexpr int SCALE_GRANULARITY_K = 128; /* Always 128 for SM120 per CUTLASS requirement */ \
if (scale_granularity_m != 1) { \
TORCH_CHECK(false, "SM120 only supports scale_granularity_m=1"); \
Copy link
Collaborator

@yzh119 yzh119 Sep 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Didn't see this constraint in https://github.com/NVIDIA/cutlass/blob/b2dd65dc864e09688245b316ac46c4a6cd07e15c/include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp, what's the error message if you set it to 128?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

will run into static assertion failed with "Scale Granularity M must evenly divide the tile shape M."

Copy link
Collaborator

@yzh119 yzh119 Sep 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

This error should only happend when ScaleGranularityM == 128 and TileShapeM = 64:

https://github.com/NVIDIA/cutlass/blob/b2dd65dc864e09688245b316ac46c4a6cd07e15c/include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp#L134

which is the case when MmaSM == 1 here: https://github.com/flashinfer-ai/flashinfer/pull/1610/files#diff-0977093a8d2429e66dab4cc40f31563717098cb5aca4354a814e4208f58f068bR78, which you disabled in https://github.com/flashinfer-ai/flashinfer/pull/1610/files#diff-68929275a79ec730031c1d5bec894f35ba6e932a08841fdae63087a6937c0f4fR70

Copy link
Collaborator Author

@yongwww yongwww Sep 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

right. I used a standalone test (not in this pr) to trigger that error message. will go with the #1610 (comment)

@yzh119
Copy link
Collaborator

yzh119 commented Sep 4, 2025 •
edited
Loading

There might be some misunderstanding of MmaSM here: we use it in sm100 gemm because sm100 supports tcgen05 and 2-cta mode (where 2 ctas cooperatively perform a mma computation).

However, sm120 do not have tcgen05 and 2-cta mma, so MmaSM doesn't make sense here, it should always be 1 in sm120.

The condition used separate to the two cases in https://github.com/flashinfer-ai/flashinfer/pull/1610/files#diff-0977093a8d2429e66dab4cc40f31563717098cb5aca4354a814e4208f58f068bR78 should not be MmaSM == 1, but Cooperative/PingPong schedule: https://github.com/NVIDIA/cutlass/blob/b2dd65dc864e09688245b316ac46c4a6cd07e15c/examples/87_blackwell_geforce_gemm_blockwise/87b_blackwell_geforce_fp8_bf16_gemm_groupwise.cu#L120-L123.

Please consider the following changes:

  1. remove the confusing MmaSM argument, there is no concept of 2SM in sm120.
  2. if we want to support both PingPong and Cooperative GEMM, please refer to https://github.com/NVIDIA/cutlass/blob/b2dd65dc864e09688245b316ac46c4a6cd07e15c/examples/87_blackwell_geforce_gemm_blockwise/87b_blackwell_geforce_fp8_bf16_gemm_groupwise.cu#L166-L171

@yongwww yongwww force-pushed the sm120_cutlass_fp8_gemm branch from 98fdd75 to 3a5cd77 Compare September 4, 2025 02:19
@yongwww
Copy link
Collaborator Author

yongwww commented Sep 4, 2025

Thanks, @yzh119, @nvmbreughe , @aleozlx for the helpful and insightful comments! I’ve incorporated them. Please take a look. For the PingPong gemm, I left it as a todo for now; the current default in the cutlass examples is cooperative gemm.

yzh119
yzh119 reviewed Sep 4, 2025
View reviewed changes
csrc/gemm_groupwise_sm120.cu Outdated
constexpr int SCALE_GRANULARITY_K = 128; /* equal tile K dimension*/ \
if (scale_granularity_m != 1) { \
TORCH_CHECK(false, \
"SM120 only supports scale_granularity_m=1 to ensure compatibility with all " \
Copy link
Collaborator

@yzh119 yzh119 Sep 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Is this still the case after your changes? If not, let's add 128 back.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Good catch! A divisor of 128 should be valid. I added scale_granularity_m=1 and scale_granularity_m=128, the change was: https://github.com/flashinfer-ai/flashinfer/pull/1610/files#diff-68929275a79ec730031c1d5bec894f35ba6e932a08841fdae63087a6937c0f4fR44-R52

@yongwww yongwww force-pushed the sm120_cutlass_fp8_gemm branch from b0ed858 to 5f40472 Compare September 4, 2025 16:17
@yongwww yongwww force-pushed the sm120_cutlass_fp8_gemm branch from 3c22b87 to 305be25 Compare September 4, 2025 17:04
@yongwww yongwww merged commit 90abf04 into flashinfer-ai:main Sep 4, 2025
2 checks passed
@yongwww yongwww deleted the sm120_cutlass_fp8_gemm branch September 4, 2025 20:47
@voipmonitor voipmonitor mentioned this pull request Sep 7, 2025
Merged
4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yzh119 yzh119 yzh119 approved these changes

+1 more reviewer

@nvmbreughe nvmbreughe nvmbreughe left review comments

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@yongwww @aleozlx @yzh119 @nvmbreughe