Repository: sgl-project/sglang
Pull request #4199: linear support deepgemm
Author: login: sleepcoo
Labels: name: high priority
Created: 2025-03-08T07:36:02Z | Merged: 2025-03-11T07:38:37Z

<!-- Thank you for your contribution! We appreciate it. The following guidelines will help improve your pull request and facilitate feedback. If anything is unclear, don't hesitate to submit your pull request and ask the maintainers for assistance. -->

## Motivation

Use deepgemm linear gemm, dependent on https://github.com/sgl-project/sglang/pull/4165

## Test case
```
#luanch server
python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code 
```
**MMLU Average accuracy: 0.878**

### performence test

 size |  origin |  deepgemm linear
-- | -- | --
Output token throughput (tok/s) |   1244.04     |  1360.05
Total token throughput (tok/s)  |   1818.26     |  1987.82

</figure>



Diff:
diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py
index 54c07f909409..4ffeab8bab21 100644
--- a/python/sglang/srt/layers/quantization/fp8_kernel.py
+++ b/python/sglang/srt/layers/quantization/fp8_kernel.py
@@ -29,10 +29,13 @@
 
 _is_cuda = torch.cuda.is_available() and torch.version.cuda
 if _is_cuda:
+    import deep_gemm
     from sgl_kernel import sgl_per_token_group_quant_fp8, sgl_per_token_quant_fp8
 
 logger = logging.getLogger(__name__)
 
+_enable_jit_deepgemm = int(os.getenv("SGL_ENABLE_JIT_DEEPGEMM", "0"))
+
 
 @triton.jit
 def _per_token_group_quant_fp8(
@@ -722,34 +725,39 @@ def grid(META):
     num_workgroups = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(
         N, config["BLOCK_SIZE_N"]
     )
-    kernel = (
-        _w8a8_block_fp8_matmul_unrolledx4
-        if (is_hip_ == True and num_workgroups <= get_device_core_count())
-        else _w8a8_block_fp8_matmul
-    )
 
-    kernel[grid](
-        A,
-        B,
-        C,
-        As,
-        Bs,
-        M,
-        N,
-        K,
-        block_n,
-        block_k,
-        A.stride(-2),
-        A.stride(-1),
-        B.stride(1),
-        B.stride(0),
-        C.stride(-2),
-        C.stride(-1),
-        As.stride(-2),
-        As.stride(-1),
-        Bs.stride(1),
-        Bs.stride(0),
-        **config,
-    )
+    # deepgemm only support bf16
+    if _is_cuda and C.dtype == torch.bfloat16 and _enable_jit_deepgemm:
+        deep_gemm.gemm_fp8_fp8_bf16_nt((A, As), (B, Bs), C)
+    else:
+        kernel = (
+            _w8a8_block_fp8_matmul_unrolledx4
+            if (is_hip_ == True and num_workgroups <= get_device_core_count())
+            else _w8a8_block_fp8_matmul
+        )
+
+        kernel[grid](
+            A,
+            B,
+            C,
+            As,
+            Bs,
+            M,
+            N,
+            K,
+            block_n,
+            block_k,
+            A.stride(-2),
+            A.stride(-1),
+            B.stride(1),
+            B.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            As.stride(-2),
+            As.stride(-1),
+            Bs.stride(1),
+            Bs.stride(0),
+            **config,
+        )
 
     return C
diff --git a/python/sglang/test/test_block_fp8.py b/python/sglang/test/test_block_fp8.py
index b3da7690ce7b..25aaf498a06e 100644
--- a/python/sglang/test/test_block_fp8.py
+++ b/python/sglang/test/test_block_fp8.py
@@ -1,4 +1,5 @@
 import itertools
+import os
 import unittest
 
 import torch
@@ -11,6 +12,8 @@
     w8a8_block_fp8_matmul,
 )
 
+_is_cuda = torch.cuda.is_available() and torch.version.cuda
+
 
 # For test
 def native_per_token_group_quant_fp8(
@@ -208,13 +211,35 @@ def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.fl
 
 
 class TestW8A8BlockFP8Matmul(unittest.TestCase):
-    OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16]
-    M = [1, 7, 83, 512, 2048]
-    N = [128, 512, 1024, 4096, 7748, 13824]
-    K = [256, 4096, 5120, 3884, 13824]
-    # BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
-    BLOCK_SIZE = [[128, 128]]
-    SEEDS = [0]
+
+    if not _is_cuda:
+        OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16]
+        M = [1, 7, 83, 512, 2048]
+        NKs = [
+            (N, K)
+            for N in [128, 512, 1024, 4096, 7748, 13824]
+            for K in [256, 4096, 5120, 3884, 13824]
+        ]
+        # BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
+        BLOCK_SIZE = [[128, 128]]
+        SEEDS = [0]
+    else:
+        # use practical shape in DeepSeek V3 for test
+        OUT_DTYPES = [torch.bfloat16]
+        M = [64, 128, 512, 1024, 4096]
+        NKs = [
+            (1536, 7168),
+            (3072, 1536),
+            (24576, 7168),
+            (4096, 512),
+            (7168, 2048),
+            (4608, 7168),
+            (512, 7168),
+            (7168, 2304),
+            (7168, 512),
+        ]
+        BLOCK_SIZE = [[128, 128]]
+        SEEDS = [0]
 
     @classmethod
     def setUpClass(cls):
@@ -222,7 +247,8 @@ def setUpClass(cls):
             raise unittest.SkipTest("CUDA is not available")
         torch.set_default_device("cuda")
 
-    def _w8a8_block_fp8_matmul(self, M, N, K, block_size, out_dtype, seed):
+    def _w8a8_block_fp8_matmul(self, M, NK, block_size, out_dtype, seed):
+        N, K = NK
         torch.manual_seed(seed)
         # NOTE(HandH1998): to avoid overflow when out_dtype = torch.half
         factor_for_scale = 1e-2
@@ -257,19 +283,17 @@ def _w8a8_block_fp8_matmul(self, M, N, K, block_size, out_dtype, seed):
     def test_w8a8_block_fp8_matmul(self):
         for params in itertools.product(
             self.M,
-            self.N,
-            self.K,
+            self.NKs,
             self.BLOCK_SIZE,
             self.OUT_DTYPES,
             self.SEEDS,
         ):
             with self.subTest(
                 M=params[0],
-                N=params[1],
-                K=params[2],
-                block_size=params[3],
-                out_dtype=params[4],
-                seed=params[5],
+                NKs=params[1],
+                block_size=params[2],
+                out_dtype=params[3],
+                seed=params[4],
             ):
                 self._w8a8_block_fp8_matmul(*params)
 
diff --git a/test/srt/test_fp8_kernel.py b/test/srt/test_fp8_kernel.py
index fe92bfd07695..dcc5d4274880 100644
--- a/test/srt/test_fp8_kernel.py
+++ b/test/srt/test_fp8_kernel.py
@@ -17,7 +17,7 @@ def setUpClass(cls):
         cls.K = 512
         cls.group_size = 128
         cls.quant_type = torch.float8_e4m3fn
-        cls.output_type = torch.float16
+        cls.output_type = torch.bfloat16
 
     @staticmethod
     def _make_A(M, K, group_size, out_dtype):


Reviews:
user: login: zhyncs
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T07:43:59Z
user: login: merrymercy
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T08:20:19Z
user: login: sleepcoo
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T08:31:52Z
user: login: merrymercy
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T10:40:04Z
user: login: zhyncs
body: - ut can't pass on H200
- torch compile not compatible `python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-R1 --tp 8 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 1`
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T21:28:43Z
user: login: zhyncs
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-11T07:38:00Z

Review comments:
diff_hunk: @@ -728,28 +729,37 @@ def grid(META):
         else _w8a8_block_fp8_matmul
     )
 
-    kernel[grid](
-        A,
-        B,
-        C,
-        As,
-        Bs,
-        M,
-        N,
-        K,
-        block_n,
-        block_k,
-        A.stride(-2),
-        A.stride(-1),
-        B.stride(1),
-        B.stride(0),
-        C.stride(-2),
-        C.stride(-1),
-        As.stride(-2),
-        As.stride(-1),
-        Bs.stride(1),
-        Bs.stride(0),
-        **config,
-    )
+    if _is_cuda:
+        deep_gemm.gemm_fp8_fp8_bf16_nt((A, As), (B, Bs), C)
+    else:
+        kernel = (
path: python/sglang/srt/layers/quantization/fp8_kernel.py
user: login: zhyncs
body: This is a duplicate. Please delete line 726 of the code.
created_at: 2025-03-10T07:43:59Z
author_association: COLLABORATOR
diff_hunk: @@ -18,6 +18,7 @@
 import os
 from typing import Any, Dict, List, Optional, Tuple
 
+import deep_gemm
path: python/sglang/srt/layers/quantization/fp8_kernel.py
user: login: merrymercy
body: is it possible to do `from sgl_kernel import deep_gemm`?
created_at: 2025-03-10T08:20:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -18,6 +18,7 @@
 import os
 from typing import Any, Dict, List, Optional, Tuple
 
+import deep_gemm
path: python/sglang/srt/layers/quantization/fp8_kernel.py
user: login: sleepcoo
body: Currently, installing deepgemm via 3rd party is quite tricky. Could we change it to import deep_gemm from sgl_kernel? In your opinion, which should be prioritized: changing the import to import deep_gemm from sgl_kernel or integrating group gemm? I can work on whichever has higher priority first.
created_at: 2025-03-10T08:31:52Z
author_association: COLLABORATOR
diff_hunk: @@ -18,6 +18,7 @@
 import os
 from typing import Any, Dict, List, Optional, Tuple
 
+import deep_gemm
path: python/sglang/srt/layers/quantization/fp8_kernel.py
user: login: merrymercy
body: Getting better speed has a higher priority
created_at: 2025-03-10T10:40:03Z
author_association: CONTRIBUTOR

PR comments:
user: login: zhyncs
created_at: 2025-03-10T07:41:33Z
body: wait for https://github.com/sgl-project/sglang/pull/4255
author_association: COLLABORATOR
user: login: merrymercy
created_at: 2025-03-10T10:40:36Z
body: I tried to release a new version of sgl kernel to unblock this. But I hit some errors https://github.com/sgl-project/sglang/commit/1a5023e05d8f286c1754a4a4af075eefa3cfc089. Will need @zhyncs 's help.
author_association: CONTRIBUTOR
user: login: zhyncs
created_at: 2025-03-10T18:33:02Z
body: v0.0.4.post1 is ready
author_association: COLLABORATOR
user: login: zhyncs
created_at: 2025-03-10T21:43:36Z
body: ref https://github.com/sgl-project/sglang/pull/3232
author_association: COLLABORATOR
user: login: sleepcoo
created_at: 2025-03-11T04:45:07Z
body: > * ut can't pass on H200
> * torch compile not compatible `python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-R1 --tp 8 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 1`

wait this https://github.com/deepseek-ai/DeepGEMM/pull/55
author_association: COLLABORATOR
user: login: zhyncs
created_at: 2025-03-11T07:36:41Z
body: There are still many issues to be fixed in this PR, such as the condition of arch older than hopper, oom of ut, etc., but these are all minor issues that can be resolved in follow-up.
author_association: COLLABORATOR