Skip to content

Add Flashinfer DeepGEMM SM90 for SwapAB Optimization #15514

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
Open
b8zhong wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add Flashinfer DeepGEMM SM90 for SwapAB Optimization #15514

b8zhong wants to merge 3 commits into from
+73 −0

Conversation

@b8zhong
Copy link
Collaborator

@b8zhong b8zhong commented Dec 20, 2025

Motivation

After flashinfer-ai/flashinfer#2131 in Flashinfer, we can benefit from SwapAB, where the input order is swapped to benefit when the M dimension is < 32 (e.g when BS < 32 in decoding). When it is larger, there is no benefit.

Modifications

(Requires Flashinfer nightly, and the backend currently only supports SM90)
Note that Flashinfer will compile it's own DeepGEMM. So it is separate from the DeepGEMM built in the Docker container.

Accuracy Tests

Benchmarking and Profiling

for ((N=1; N<=128; N*=2)); do
  python3 -m sglang.bench_serving \
    --backend sglang \
    --flush-cache \
    --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 1024 \
    --random-range-ratio 1.0 \
    --num-prompts $((6*N)) \
    --max-concurrency $N \
    --output-file res.jsonl
done

We can see that when the M dimension is small, there is around a 5-8% E2E benefit

image image

Checklist

@gemini-code-assist
Copy link
Contributor

gemini-code-assist bot commented Dec 20, 2025

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@Fridge003 Fridge003 mentioned this pull request Dec 21, 2025
Draft
6 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner

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

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

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

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

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@b8zhong