Skip to content

[Feature] Apply Cublas Grouped Gemm kernel #3629

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 3 commits into from
Feb 18, 2025
Merged

[Feature] Apply Cublas Grouped Gemm kernel #3629

merged 3 commits into from
Feb 18, 2025

Conversation

Fridge003
Copy link
Collaborator

@Fridge003 Fridge003 commented Feb 17, 2025 •
edited
Loading

Motivation

#3323

Grouped Gemm kernel added in Cublas 12.5 is useful. It can be applied to MoE EP layer/Lora layer for acceleration.

Modifications

  • Add cublas_grouped_gemm in sgl-kernel library, and provides accuracy test/benchmark script.
  • Update document for this feature.

Environment:

Torch 2.5.1, Cuda 12.5, Cublas 12.5.3.2, sglang 0.4.3
Since sglang doesn't support torch 2.6 yet, to build the environment:

  1. First make sure the Cuda version is >= 12.5 with nvcc -V
  2. Then install sglang as the official document does
  3. Reinstall cublas 12.5 through pip install nvidia-cublas-cu12==12.5.3.2 so that the cublas is upgraded
  4. Compile the new sgl-kernel library.

Accuracy Test

python3 sgl-kernel/tests/test_cublas_grouped_gemm.py

Kernel Benchmark

Deepseek V2 setting

On Deepseek V2 setting with TP Size = 8 (Group Size=20), N = 3072, K = 5120:

!python3 sgl-kernel/benchmark/bench_cublas_grouped_gemm.py --models DeepSeek-V2 --tp-size 8

Result in GB per second:
截屏2025-02-17 00 56 17

Deepseek V2-Lite setting

On Deepseek V2 setting with TP Size = 2 (Group Size=32), N = 2816, K = 2048:

!python3 sgl-kernel/benchmark/bench_cublas_grouped_gemm.py --models DeepSeek-V2-Lite --tp-size 2

Result in GB per second:
截屏2025-02-17 01 00 33

Checklist

yizhang2077
yizhang2077 reviewed Feb 17, 2025
View reviewed changes
yizhang2077
yizhang2077 reviewed Feb 17, 2025
View reviewed changes
@yizhang2077
Copy link
Collaborator

Since pytorch 2.5.1 only supports cuda12.4 in official docs, and we can not change pytorch version easily, we need to update doc to guide user to reinstall pytorch if they want to use group gemm to accelerate their models.

@Fridge003 Fridge003 changed the title [Feature] Implement Cublas Grouped Gemm kernel and apply it to MoE EP [Feature] Implement Cublas Grouped Gemm kernel Feb 17, 2025
@Fridge003 Fridge003 changed the title [Feature] Implement Cublas Grouped Gemm kernel [Feature] Apply Cublas Grouped Gemm kernel Feb 17, 2025
@yizhang2077
Copy link
Collaborator

LGTM cc @zhyncs

yizhang2077
yizhang2077 reviewed Feb 18, 2025
View reviewed changes
@yizhang2077 yizhang2077 self-requested a review February 18, 2025 05:41
@zhyncs zhyncs merged commit 67fc595 into sgl-project:main Feb 18, 2025
12 of 16 checks passed
@zhyncs
Copy link
Member

zhyncs commented Feb 18, 2025

amazing work!

@Fridge003 Fridge003 deleted the deepseek branch February 18, 2025 19:12
@hebiao064 hebiao064 mentioned this pull request Mar 12, 2025
Closed
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yizhang2077 yizhang2077 yizhang2077 approved these changes

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@HandH1998 HandH1998 Awaiting requested review from HandH1998 HandH1998 is a code owner

@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner

@Ying1123 Ying1123 Awaiting requested review from Ying1123

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002

@ByronHsu ByronHsu Awaiting requested review from ByronHsu

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@Fridge003 @yizhang2077 @zhyncs