add blockwise gemm cute dsl #1922
Conversation
Summary of ChangesHello @Amir-19, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a comprehensive test suite for the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new test file for the blockwise GEMM cute DSL. The tests are well-structured, but there are several areas for improvement in terms of code clarity, style, and maintainability. I've identified some unused variables, a magic number that should be a constant, a debug print statement, and a violation of PEP 8 naming conventions. Addressing these points will make the new test code cleaner and easier to maintain.
| l, m, n, cd_major == "m", get_cutlass_dtype(c_dtype), device=device | ||
| ) | ||
| sfa_torch_cpu = cutlass_torch.matrix( | ||
| l, m, math.ceil(k / 128), True, get_cutlass_dtype(scale_dtype), device=device |
There was a problem hiding this comment.
The magic number 128 is used in multiple places and appears to represent scale granularity. It should be defined as a constant at the module level, e.g., SCALE_GRANULARITY = 128, and used throughout the file to improve readability and maintainability. This constant should be used on lines 32, 36, 37, 205, 207, 210, 214, and 217.
| l, m, math.ceil(k / 128), True, get_cutlass_dtype(scale_dtype), device=device | |
| l, m, math.ceil(k / SCALE_GRANULARITY), True, get_cutlass_dtype(scale_dtype), device=device |
| a_tensor, a_torch = cutlass_torch.cute_tensor_like( | ||
| a_torch_cpu, | ||
| get_cutlass_dtype(ab_dtype), | ||
| is_dynamic_layout=True, | ||
| assumed_align=16, | ||
| ) | ||
| b_tensor, b_torch = cutlass_torch.cute_tensor_like( | ||
| b_torch_cpu, | ||
| get_cutlass_dtype(ab_dtype), | ||
| is_dynamic_layout=True, | ||
| assumed_align=16, | ||
| ) | ||
| c_tensor, c_torch = cutlass_torch.cute_tensor_like( | ||
| c_torch_cpu, | ||
| get_cutlass_dtype(c_dtype), | ||
| is_dynamic_layout=True, | ||
| assumed_align=16, | ||
| ) | ||
| sfa_tensor, sfa_torch = cutlass_torch.cute_tensor_like( | ||
| sfa_torch_cpu, | ||
| get_cutlass_dtype(scale_dtype), | ||
| is_dynamic_layout=True, | ||
| assumed_align=16, | ||
| ) | ||
| sfb_tensor, sfb_torch = cutlass_torch.cute_tensor_like( | ||
| sfb_torch_cpu, | ||
| get_cutlass_dtype(scale_dtype), | ||
| is_dynamic_layout=True, | ||
| assumed_align=16, | ||
| ) | ||
|
|
||
| return ( | ||
| a_tensor, | ||
| a_torch, | ||
| b_tensor, | ||
| b_torch, | ||
| c_tensor, | ||
| c_torch, | ||
| sfa_tensor, | ||
| sfa_torch, | ||
| sfb_tensor, | ||
| sfb_torch, | ||
| a_torch_cpu, | ||
| b_torch_cpu, | ||
| c_torch_cpu, | ||
| sfa_torch_cpu, | ||
| sfb_torch_cpu, | ||
| ) |
There was a problem hiding this comment.
The create_tensors function creates and returns several tensors that are not used in the test (a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, and c_torch_cpu). This can be simplified by not creating or returning these unused variables, which improves code clarity and maintainability. The cute_tensor_like calls can be updated to assign the unused cute tensor to _, and the unused c_torch_cpu can be removed from the return tuple.
_, a_torch = cutlass_torch.cute_tensor_like(
a_torch_cpu,
get_cutlass_dtype(ab_dtype),
is_dynamic_layout=True,
assumed_align=16,
)
_, b_torch = cutlass_torch.cute_tensor_like(
b_torch_cpu,
get_cutlass_dtype(ab_dtype),
is_dynamic_layout=True,
assumed_align=16,
)
_, c_torch = cutlass_torch.cute_tensor_like(
c_torch_cpu,
get_cutlass_dtype(c_dtype),
is_dynamic_layout=True,
assumed_align=16,
)
_, sfa_torch = cutlass_torch.cute_tensor_like(
sfa_torch_cpu,
get_cutlass_dtype(scale_dtype),
is_dynamic_layout=True,
assumed_align=16,
)
_, sfb_torch = cutlass_torch.cute_tensor_like(
sfb_torch_cpu,
get_cutlass_dtype(scale_dtype),
is_dynamic_layout=True,
assumed_align=16,
)
return (
a_torch,
b_torch,
c_torch,
sfa_torch,
sfb_torch,
a_torch_cpu,
b_torch_cpu,
sfa_torch_cpu,
sfb_torch_cpu,
)|
|
||
| sm_count = get_num_sm(device) | ||
|
|
||
| print(f"device: {device}") |
There was a problem hiding this comment.
This print statement appears to be for debugging and should be removed from the final test code to avoid polluting test output.
| ( | ||
| a_tensor, | ||
| a_torch, | ||
| b_tensor, | ||
| b_torch, | ||
| c_tensor, | ||
| c_torch, | ||
| sfa_tensor, | ||
| sfa_torch, | ||
| sfb_tensor, | ||
| sfb_torch, | ||
| a_torch_cpu, | ||
| b_torch_cpu, | ||
| c_torch_cpu, | ||
| sfa_torch_cpu, | ||
| sfb_torch_cpu, | ||
| ) = create_tensors( | ||
| l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype, sf_dtype, device | ||
| ) |
There was a problem hiding this comment.
The destructuring of the return value from create_tensors should be updated to match the simplified function that no longer returns unused tensors.
(
a_torch,
b_torch,
c_torch,
sfa_torch,
sfb_torch,
a_torch_cpu,
b_torch_cpu,
sfa_torch_cpu,
sfb_torch_cpu,
) = create_tensors(
l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype, sf_dtype, device
)
Amir-19
force-pushed
the
blockwise_gemm
branch
from
f203299 to
b21ff9d
Compare
1 participant
📌 Description
🔍 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes