Improve small size performance in cutedsl fp4 #2517
Improve small size performance in cutedsl fp4 #2517vincentzed wants to merge 4 commits intoflashinfer-ai:mainfrom
Conversation
Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the
✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @vincentzed, 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 optimizes the performance of CuteDSL for FP4 Mixture-of-Experts (MoE) operations, particularly for small token counts. The primary improvement involves streamlining the output buffer initialization process by removing a redundant zeroing operation and allowing this initialization to occur asynchronously, overlapping with other computations. This change reduces overhead, leading to better latency for low token/s scenarios, as evidenced by the provided benchmarks showing CuteDSL now outperforming TRTLLM at 8 tokens. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
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 performance optimization for small batch sizes in the CuteDSL FP4 MoE implementation. The change replaces a custom kernel for zeroing the output buffer with a standard PyTorch zero_() call and avoids a redundant zeroing operation inside the GEMM kernel. This is a sensible optimization that, according to the description, yields significant performance improvements. The implementation looks correct and the logic is sound. My only suggestion is to add documentation for the new parameter introduced to control the zeroing behavior, which will improve code clarity and maintainability.
| cluster_shape_mn: Tuple[int, int] = (2, 1), | ||
| raster_along_m: bool = False, | ||
| sm_count: Optional[int] = None, | ||
| zero_output: bool = True, |
There was a problem hiding this comment.
The new zero_output parameter should be documented in the function's docstring to explain its purpose. Additionally, the docstring for the out parameter should be updated to clarify that its zero-initialization is now conditional on this new parameter.
For example:
out: Optional output tensor, shape (seq_len, n). Created if None.
This tensor is used for atomic accumulation. If `zero_output` is True (the default),
it will be zero-initialized by this function. Otherwise, the caller is responsible
for ensuring it is zero-initialized before calling.
...
sm_count: Number of SMs to use. Default: max available.
zero_output: If True, zero-initialize the `out` tensor before the kernel launch.
Set to False if the output tensor is already zeroed. Default: True.
|
Closing as covered by #2811 |
1 participant
📌 Description
As a hunch, cuteDSL should be able to match trt performance for most low tok/s. This PR is meant to more often choose CuteDSL for fp4 moe. So here we compare ncu measure, cutlass, trt, cutedsl, and find an easy way to more often select cutedsl on toks = small case.
void at::native::vectorized_elementwise_kernel<4, ...FillFunctor<c10::BFloat16>...>is moved to the overlap with gemm (this happen before gemm2 finalize), unless output = None.It is 4.3 us in ncu.
At 8 toks, CuteDSL wins over TRT, and the Deltas are as follows:
Before: CuteDSL/Trt: 1,2,4,8: 0.88x,0.92x,0.93x,0.99x
After:
Which is consistent with remove the fill (~4us)- which only matters at small size.
I can upload ncu as well, if it is useful.
🔍 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