Skip to content

Add blackwell dual gemm fusion support for MLP - #41887

Merged
BBuf merged 18 commits into
sgl-project:mainfrom
DevashishLal-CB:dev/dlal/dual-gemm
Oct 9, 2026
Merged

BBuf merged 18 commits into
sgl-project:mainfrom
DevashishLal-CB:dev/dlal/dual-gemm

Conversation

@DevashishLal-CB

@DevashishLal-CB DevashishLal-CB commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Adds a Blackwell CuTe DSL dual-GEMM kernel that fuses:

gate = input @ gate_weight
up = input @ up_weight
output = SiLU(gate) * up

Tested on B200

  • Llama 3 8B FP8, BS1, 5.2% Offline Gains
  • Qwen 2 7B FP8, BS1, 7.8% Offline Gains

Qwen 2 7B FP8

Metric BS1 BS2 BS4 BS6 BS8 BS10 BS12 BS14 BS16 Mean
Benchmark duration +7.56% +9.08% +7.76% +7.80% +7.75% +8.33% +8.49% +8.56% +8.59% +8.21%
Input token throughput +8.15% +9.99% +8.42% +8.49% +8.42% +9.09% +9.29% +9.37% +9.40% +8.96%
Output token throughput +8.15% +9.99% +8.42% +8.49% +8.42% +9.09% +9.29% +9.37% +9.40% +8.96%
Total token throughput +8.15% +9.99% +8.42% +8.49% +8.42% +9.09% +9.29% +9.37% +9.40% +8.96%
Mean E2E latency +7.54% +9.09% +7.77% +7.83% +7.77% +8.34% +8.51% +8.58% +8.61% +8.22%
Median E2E latency +7.52% +9.07% +7.78% +7.25% +7.17% +7.74% +7.91% +7.96% +8.06% +7.83%
Mean TTFT +23.55% +9.22% +32.92% +22.68% +25.48% +22.56% +23.14% +16.34% +22.23% +22.01%
Median TTFT +24.56% -1.21% +5.82% +3.47% +2.05% +3.21% +6.06% -2.51% +3.30% +4.97%
Mean TPOT +7.54% +9.16% +7.60% +7.52% +7.84% +8.19% +8.13% +8.68% +8.62% +8.14%
Median TPOT +7.54% +9.16% +7.25% +7.17% +7.49% +7.86% +7.80% +8.36% +8.30% +7.88%
Mean ITL +7.54% +9.16% +7.60% +7.52% +7.84% +8.19% +8.13% +8.68% +8.62% +8.14%
Median ITL +7.17% +8.81% +7.28% +7.20% +7.14% +7.89% +7.83% +8.39% +8.33% +7.78%

Llama 3 8B FP8 Online

Metric BS1 BS2 BS4 BS6 BS8 BS10 BS12 BS14 BS16 Mean
Benchmark duration +4.00% +3.95% +4.95% +3.59% +3.14% +4.01% +4.07% +3.88% +3.72% +3.92%
Input token throughput +4.16% +4.12% +5.19% +3.75% +3.24% +4.19% +4.23% +4.03% +3.85% +4.08%
Output token throughput +4.16% +4.12% +5.19% +3.75% +3.24% +4.19% +4.23% +4.03% +3.85% +4.08%
Total token throughput +4.16% +4.12% +5.19% +3.75% +3.24% +4.19% +4.23% +4.03% +3.85% +4.08%
Mean E2E latency +3.99% +3.96% +4.94% +3.62% +3.16% +4.03% +4.06% +3.87% +3.71% +3.93%
Median E2E latency +3.99% +3.94% +4.94% +3.11% +2.57% +3.49% +3.51% +3.36% +3.20% +3.57%
Mean TTFT +20.97% +6.26% +26.55% +18.49% +10.19% +16.00% +17.00% +14.27% +11.20% +15.66%
Median TTFT +22.58% +0.16% +18.08% +3.84% -3.39% +0.00% +11.11% -1.01% -0.35% +5.67%
Mean TPOT +3.92% +3.85% +4.92% +3.68% +2.90% +3.56% +3.85% +3.70% +3.64% +3.78%
Median TPOT +3.92% +3.85% +4.56% +3.32% +2.55% +3.21% +3.51% +3.38% +3.32% +3.51%
Mean ITL +3.92% +3.85% +4.92% +3.68% +2.90% +3.56% +3.85% +3.70% +3.64% +3.78%
Median ITL +3.92% +3.86% +4.56% +2.96% +2.91% +3.57% +3.52% +3.39% +3.33% +3.56%

Modifications

  • Supports BF16, FP16, and FP8 inputs.
  • FP8 supports:
    • Static per-tensor
    • Static per-token
    • Dynamic per-tensor
    • Dynamic per-token
  • Supports decode batches from 1–16 tokens.
  • Uses TMA loads, tcgen05 MMA, warp specialization, one/two-CTA tactics, and resident-grid dynamic reductions.
  • Adds B200-tuned tactic selection based on batch size, hidden size, intermediate size, quantization mode, and dtype.
  • Gates the fused path to SM10x, aligned shapes, and intermediate sizes up to 18,944.
  • Supports prequantized activations and avoids redundant down-projection quantization.
  • Adds a reusable non-nn.Module DualGemm integration layer.
  • Integrates the fused path into Llama and Qwen2 MLPs.
  • Adds kernel correctness tests, layer integration tests, and benchmarks against the existing SGLang operator sequence.

Accuracy Tests

Started the server with max running requests 16

Llama 3 8B FP8

# Fusion
Total latency: 37.918 s
Score: 0.466
Output throughput: 4273.517 token/s

# No Fusion
Total latency: 42.018 s
Score: 0.466
Output throughput: 4122.628 token/s

Qwen 2 7B FP8

# Fusion
Total latency: 38.209 s
Score: 0.840
Output throughput: 5071.280 token/s

# No Fusion
Total latency: 41.280 s
Score: 0.828
Output throughput: 4677.125 token/s

Speed Tests and Profiling

FP8 Kernel Test

====================================================================================================================================
    num_tokens  hidden_size  intermediate_size        quantization |     cutedsl(us)     sglang(us) |   cutedsl(GB/s)   sglang(GB/s)
------------------------------------------------------------------------------------------------------------------------------------
0            1          128                128   static_per_tensor |          3.4211         9.7478 |            9.27           3.25
1            1          128                128    static_per_token |          3.4406         9.7293 |            9.22           3.26
2            1          128                128  dynamic_per_tensor |          8.6221        10.5062 |            3.68           3.02
3            1          128                128   dynamic_per_token |          8.6630         8.0291 |            3.66           3.95
4            1         4096              14336   static_per_tensor |         20.8625        32.6123 |         5248.61        3357.60
5            1         4096              14336    static_per_token |         20.8802        32.5925 |         5244.17        3359.64
6            1         4096              14336  dynamic_per_tensor |         25.8636        30.3979 |         4233.71        3602.19
7            1         4096              14336   dynamic_per_token |         25.7954        29.8712 |         4244.91        3665.70
8            2         4096              14336   static_per_tensor |         21.1627        32.5827 |         5174.97        3361.17
9            2         4096              14336    static_per_token |         21.1532        32.2999 |         5177.28        3390.61
10           2         4096              14336  dynamic_per_tensor |         25.8344        30.4567 |         4239.17        3595.80
11           2         4096              14336   dynamic_per_token |         25.7859        29.0530 |         4247.13        3769.54
12           4         4096              14336   static_per_tensor |         21.5043        33.4510 |         5094.35        3274.96
13           4         4096              14336    static_per_token |         21.4839        33.3824 |         5099.19        3281.68
14           4         4096              14336  dynamic_per_tensor |         25.8051        30.9833 |         4245.30        3535.79
15           4         4096              14336   dynamic_per_token |         25.8907        30.3192 |         4231.26        3613.23
16           8         4096              14336   static_per_tensor |         22.0699        34.8645 |         4966.90        3144.15
17           8         4096              14336    static_per_token |         22.0312        33.9870 |         4975.63        3225.32
18           8         4096              14336  dynamic_per_tensor |         26.1269        31.3344 |         4195.64        3498.36
19           8         4096              14336   dynamic_per_token |         26.3808        30.9638 |         4155.26        3540.24
20          16         4096              14336   static_per_tensor |         22.2845        34.1528 |         4925.24        3213.69
21          16         4096              14336    static_per_token |         22.2845        33.9288 |         4925.25        3234.91
22          16         4096              14336  dynamic_per_tensor |         26.8486        32.2511 |         4087.98        3403.18
23          16         4096              14336   dynamic_per_token |         27.7163        31.0632 |         3960.01        3533.34
24           1         3584              18944   static_per_tensor |         23.2813        36.4122 |         5439.01        3477.59
25           1         3584              18944    static_per_token |         23.2612        36.2315 |         5443.71        3494.94
26           1         3584              18944  dynamic_per_tensor |         28.5619        34.6253 |         4433.43        3657.07
27           1         3584              18944   dynamic_per_token |         28.4816        35.3183 |         4445.93        3585.31
28           2         3584              18944   static_per_tensor |         23.4620        38.0182 |         5398.01        3331.25
29           2         3584              18944    static_per_token |         23.4632        36.3322 |         5397.72        3485.83
30           2         3584              18944  dynamic_per_tensor |         28.3511        34.5653 |         4467.13        3664.02
31           2         3584              18944   dynamic_per_token |         28.2811        33.7016 |         4478.18        3757.92
32           4         3584              18944   static_per_tensor |         23.8450        37.8780 |         5313.06        3344.69
33           4         3584              18944    static_per_token |         23.8234        37.4563 |         5317.88        3382.34
34           4         3584              18944  dynamic_per_tensor |         28.7235        35.0155 |         4410.68        3618.11
35           4         3584              18944   dynamic_per_token |         28.7824        34.3451 |         4401.64        3688.73
36           8         3584              18944   static_per_tensor |         24.4257        38.9923 |         5190.18        3251.25
37           8         3584              18944    static_per_token |         24.4154        37.7575 |         5192.38        3357.58
38           8         3584              18944  dynamic_per_tensor |         28.8125        35.0365 |         4399.95        3618.33
39           8         3584              18944   dynamic_per_token |         29.0839        35.3082 |         4358.90        3590.49
40          16         3584              18944   static_per_tensor |         24.9076        38.8414 |         5096.50        3268.20
41          16         3584              18944    static_per_token |         24.9076        37.6370 |         5096.51        3372.79
42          16         3584              18944  dynamic_per_tensor |         30.0668        35.5489 |         4221.99        3570.91
43          16         3584              18944   dynamic_per_token |         30.9810        34.4144 |         4097.41        3688.62
====================================================================================================================================

FP16/BF16 Kernel Test

================================================================================================================================
    num_tokens  hidden_size  intermediate_size           dtype |     cutedsl(us)     sglang(us) |   cutedsl(GB/s)   sglang(GB/s)
--------------------------------------------------------------------------------------------------------------------------------
0            1          128                128  torch.bfloat16 |          2.9286         4.5059 |           21.00          13.65
1            1          128                128   torch.float16 |          3.2973         3.7891 |           18.66          16.23
2            1         4096              14336  torch.bfloat16 |         39.4240        42.9674 |         5549.52        5091.87
3            1         4096              14336   torch.float16 |         39.3731        43.1107 |         5556.69        5074.94
4            2         4096              14336  torch.bfloat16 |         39.3731        43.5200 |         5557.56        5028.00
5            2         4096              14336   torch.float16 |         39.4346        44.3398 |         5548.91        4935.04
6            4         4096              14336  torch.bfloat16 |         39.6902        44.9936 |         5514.89        4864.85
7            4         4096              14336   torch.float16 |         39.5165        44.4416 |         5539.14        4925.28
8            8         4096              14336  torch.bfloat16 |         39.7520        45.9776 |         5509.78        4763.73
9            8         4096              14336   torch.float16 |         39.6189        45.7523 |         5528.29        4787.18
10          16         4096              14336  torch.bfloat16 |         39.9466        43.4179 |         5489.82        5050.89
11          16         4096              14336   torch.float16 |         40.1514        43.8893 |         5461.82        4996.65
12           1         3584              18944  torch.bfloat16 |         42.4352        51.5693 |         5961.36        4905.47
13           1         3584              18944   torch.float16 |         42.5782        47.8925 |         5941.34        5282.07
14           2         3584              18944  torch.bfloat16 |         42.5997        48.0259 |         5939.33        5268.27
15           2         3584              18944   torch.float16 |         42.7830        48.8656 |         5913.88        5177.74
16           4         3584              18944  torch.bfloat16 |         42.8650        50.0118 |         5904.53        5060.75
17           4         3584              18944   torch.float16 |         42.9856        49.1728 |         5887.96        5147.10
18           8         3584              18944  torch.bfloat16 |         43.0083        49.9098 |         5888.75        5074.47
19           8         3584              18944   torch.float16 |         43.1322        48.8042 |         5871.85        5189.42
20          16         3584              18944  torch.bfloat16 |         43.3971        48.6608 |         5843.73        5211.61
21          16         3584              18944   torch.float16 |         43.6547        49.3978 |         5809.25        5133.86
================================================================================================================================

E2E Llama 8B FP8

# RedHatAI/Meta-Llama-3.1-8B-FP8 BS1 ISL 1024 OSL 1024 Default
python3 -m sglang.benchmark.one_batch --model RedHatAI/Meta-Llama-3.1-8B-FP8 --batch-size 1 --input-len 1024 --output-len 1024
Benchmark ...
Prefill. latency: 0.01137 s, throughput:  90049.06 token/s
Decode 0. Batch size: 1, latency: 0.00313 s, throughput:    319.79 token/s
Decode 1. Batch size: 1, latency: 0.00292 s, throughput:    341.95 token/s
Decode 2. Batch size: 1, latency: 0.00290 s, throughput:    345.12 token/s
Decode 3. Batch size: 1, latency: 0.00292 s, throughput:    342.77 token/s
Decode 4. Batch size: 1, latency: 0.00291 s, throughput:    343.89 token/s
Decode.  median latency: 0.00289 s, median throughput:    345.47 token/s
Total. latency:  2.977 s, throughput:    688.05 token/s

# RedHatAI/Meta-Llama-3.1-8B-FP8 BS1 ISL 1024 OSL 1024 Fusion Enabled
python3 -m sglang.benchmark.one_batch --model RedHatAI/Meta-Llama-3.1-8B-FP8 --batch-size 1 --input-len 1024 --output-len 1024
Benchmark ...
Prefill. latency: 0.01096 s, throughput:  93458.73 token/s
Decode 0. Batch size: 1, latency: 0.00285 s, throughput:    351.21 token/s
Decode 1. Batch size: 1, latency: 0.00277 s, throughput:    361.26 token/s
Decode 2. Batch size: 1, latency: 0.00281 s, throughput:    355.48 token/s
Decode 3. Batch size: 1, latency: 0.00281 s, throughput:    356.06 token/s
Decode 4. Batch size: 1, latency: 0.00282 s, throughput:    354.40 token/s
Decode.  median latency: 0.00275 s, median throughput:    363.47 token/s
Total. latency:  2.828 s, throughput:    724.11 token/s

E2E Qwen 7B FP8

# RedHatAI/Qwen2-7B-Instruct-FP8 BS1 ISL 1024 OSL 1024 Default
python3 -m sglang.benchmark.one_batch --model RedHatAI/Qwen2-7B-Instruct-FP8 --batch-size 1 --input-len 1024 --output-len 1024
Benchmark ...
Prefill. latency: 0.01151 s, throughput:  88946.00 token/s
Decode 0. Batch size: 1, latency: 0.00293 s, throughput:    340.95 token/s
Decode 1. Batch size: 1, latency: 0.00286 s, throughput:    349.96 token/s
Decode 2. Batch size: 1, latency: 0.00290 s, throughput:    345.31 token/s
Decode 3. Batch size: 1, latency: 0.00287 s, throughput:    348.82 token/s
Decode 4. Batch size: 1, latency: 0.00284 s, throughput:    352.13 token/s
Decode.  median latency: 0.00284 s, median throughput:    352.05 token/s
Total. latency:  2.921 s, throughput:    701.22 token/s

# RedHatAI/Qwen2-7B-Instruct-FP8 BS1 ISL 1024 OSL 1024 Fusion Enabled
python3 -m sglang.benchmark.one_batch --model RedHatAI/Qwen2-7B-Instruct-FP8 --batch-size 1 --input-len 1024 --output-len 1024
Benchmark ...
Prefill. latency: 0.01154 s, throughput:  88716.99 token/s
Decode 0. Batch size: 1, latency: 0.00276 s, throughput:    361.74 token/s
Decode 1. Batch size: 1, latency: 0.00266 s, throughput:    376.01 token/s
Decode 2. Batch size: 1, latency: 0.00266 s, throughput:    375.28 token/s
Decode 3. Batch size: 1, latency: 0.00265 s, throughput:    377.30 token/s
Decode 4. Batch size: 1, latency: 0.00266 s, throughput:    376.19 token/s
Decode.  median latency: 0.00263 s, median throughput:    380.03 token/s
Total. latency:  2.707 s, throughput:    756.43 token/s

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ⏳ Run #37860676131
Latest PR Test (Extra): ❌ Run #37860675592
Latest PR Test (AMD ROCm 10): ❌ Run #37860676222

support both static and dynamic quant

Signed-off-by: Devashish Lal <devcode@meta.com>
Signed-off-by: Devashish Lal <devcode@meta.com>
Signed-off-by: Devashish Lal <devcode@meta.com>
support upto BS 16
support all for quant permutations

Signed-off-by: Devashish Lal <devcode@meta.com>
Signed-off-by: Devashish Lal <devcode@meta.com>
@DevashishLal-CB

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 30, 2026
Signed-off-by: Devashish Lal <devcode@meta.com>
Signed-off-by: Devashish Lal <devcode@meta.com>
removed test only envar

Signed-off-by: Devashish Lal <devcode@meta.com>
@DevashishLal-CB

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci

Comment thread python/sglang/srt/layers/dual_gemm.py Outdated
if not self.mode.is_quantized:
from sglang.kernels.ops.gemm import dual_gemm_swiglu

return dual_gemm_swiglu(x, self.gate_up_proj.weight)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Could we fall back when gate/up LoRA is enabled? This keeps a reference to the original layer, so replacing gate_up_proj with a LoRA wrapper later does not affect it. Small batches will use the base weight directly and skip the adapter.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

instead of saving the reference during init we can pass it to the call

MAX_DUAL_GEMM_DECODE_TOKENS,
hidden_size,
local_intermediate_size,
):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Could we include the dynamic-quantization residency check here and fall back when it fails? For Qwen2's 18944 intermediate size on a 132-SM device, BS1 passes this gate but needs 296 CTAs against a capacity of 264. The kernel then raises instead of using the existing path.

@DevashishLal-CB DevashishLal-CB Oct 8, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

yeah, make sense also the MAX_DUAL_GEMM_INTERMEDIATE_SIZE is only needed for dynamic quant we can enable this for any size of static quant

removed max intermediate size limit for unquant and static quant
updated residency check to ensure correct sm usage

Signed-off-by: Devashish Lal <devcode@meta.com>
@DevashishLal-CB
DevashishLal-CB force-pushed the dev/dlal/dual-gemm branch 2 times, most recently from eed6d5e to 7051359 Compare October 8, 2026 17:03
# Conflicts:
#	python/sglang/srt/layers/quantization/fp8_utils.py
@DevashishLal-CB

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

1 similar comment
@DevashishLal-CB

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

@BBuf
BBuf merged commit 4debae7 into sgl-project:main Oct 9, 2026
152 of 177 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

jit-kernel run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants