Add blackwell dual gemm fusion support for MLP - #41887
Conversation
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>
3a7f875 to
f74bb89
Compare
Signed-off-by: Devashish Lal <devcode@meta.com>
Signed-off-by: Devashish Lal <devcode@meta.com>
|
/tag-and-rerun-ci |
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>
|
/tag-and-rerun-ci |
| 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) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
instead of saving the reference during init we can pass it to the call
| MAX_DUAL_GEMM_DECODE_TOKENS, | ||
| hidden_size, | ||
| local_intermediate_size, | ||
| ): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
Signed-off-by: Devashish Lal <devcode@meta.com>
e876110 to
eed6d5e
Compare
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>
eed6d5e to
7051359
Compare
# Conflicts: # python/sglang/srt/layers/quantization/fp8_utils.py
|
/rerun-failed-ci |
1 similar comment
|
/rerun-failed-ci |
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
Qwen 2 7B FP8
Llama 3 8B FP8 Online
Modifications
Accuracy Tests
Started the server with max running requests 16
Llama 3 8B FP8
Qwen 2 7B FP8
Speed Tests and Profiling
FP8 Kernel Test
FP16/BF16 Kernel Test
E2E Llama 8B FP8
E2E Qwen 7B FP8
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ⏳ Run #37860676131
Latest PR Test (Extra): ❌ Run #37860675592
Latest PR Test (AMD ROCm 10): ❌ Run #37860676222