[Perf] [AMD] batch-invariant: Tune the bf16 matmul for gfx950

python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py

@@ -10,7 +10,7 @@
 import triton.language as tl
 
 from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
-from sglang.srt.utils import is_npu
+from sglang.srt.utils import is_gfx95_supported, is_npu
 from sglang.srt.utils.common import (
     calc_diff,
     get_bool_env_var,
@@ -19,6 +19,7 @@
 )
 
 _is_npu = is_npu()
+_is_gfx95 = is_gfx95_supported()
 if _is_npu:
     import torch_npu
 
@@ -170,6 +171,49 @@ def matmul_kernel_persistent(
         tl.store(c_ptrs, c, mask=c_mask)
 
 
+def _matmul_persistent_configs() -> dict:
+    configs = {
+        torch.bfloat16: {
+            "BLOCK_SIZE_M": 128,
+            "BLOCK_SIZE_N": 128,
+            "BLOCK_SIZE_K": 64,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 3,
+            "num_warps": 8,
+        },
+        torch.float16: {
+            "BLOCK_SIZE_M": 128,
+            "BLOCK_SIZE_N": 256,
+            "BLOCK_SIZE_K": 64,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 3,
+            "num_warps": 8,
+        },
+        torch.float32: {
+            "BLOCK_SIZE_M": 128,
+            "BLOCK_SIZE_N": 128,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 3,
+            "num_warps": 8,
+        },
+    }
+    if _is_gfx95:
+        # Measured on MI350X: faster than the default on average over M = 1-4096
+        # on dense and MoE linear shapes up to N ~ 8k, neutral on wider N; needs
+        # gfx950's 160 KB of LDS.
+        configs[torch.bfloat16] = {
+            "BLOCK_SIZE_M": 128,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 128,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 3,
+            "num_warps": 8,
+            "matrix_instr_nonkdim": 16,
+        }
+    return configs
+
+
 def _matmul_persistent_triton(
     a: torch.Tensor,
     b: torch.Tensor,
@@ -198,32 +242,7 @@ def grid(META):
             ),
         )
 
-    configs = {
-        torch.bfloat16: {
-            "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": 128,
-            "BLOCK_SIZE_K": 64,
-            "GROUP_SIZE_M": 8,
-            "num_stages": 3,
-            "num_warps": 8,
-        },
-        torch.float16: {
-            "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": 256,
-            "BLOCK_SIZE_K": 64,
-            "GROUP_SIZE_M": 8,
-            "num_stages": 3,
-            "num_warps": 8,
-        },
-        torch.float32: {
-            "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": 128,
-            "BLOCK_SIZE_K": 32,
-            "GROUP_SIZE_M": 8,
-            "num_stages": 3,
-            "num_warps": 8,
-        },
-    }
+    configs = _matmul_persistent_configs()
     # print(a.device, b.device, c.device)
     matmul_kernel_persistent[grid](
         a,

test/registered/unit/batch_invariant_ops/test_matmul_persistent_configs.py

@@ -0,0 +1,58 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.batch_invariant_ops import batch_invariant_ops
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+# gfx950 LDS per CU.
+GFX950_LDS_BYTES = 160 * 1024
+
+
+class TestMatmulPersistentConfigs(CustomTestCase):
+    def test_gfx950_bf16_preset(self):
+        """The gfx950 bf16 preset is a measured choice; a change must be deliberate
+        and still fit gfx950's LDS."""
+        with patch.object(batch_invariant_ops, "_is_gfx95", True):
+            cfg = batch_invariant_ops._matmul_persistent_configs()[torch.bfloat16]
+        self.assertEqual(
+            cfg,
+            {
+                "BLOCK_SIZE_M": 128,
+                "BLOCK_SIZE_N": 64,
+                "BLOCK_SIZE_K": 128,
+                "GROUP_SIZE_M": 8,
+                "num_stages": 3,
+                "num_warps": 8,
+                "matrix_instr_nonkdim": 16,
+            },
+        )
+        tile_bytes = (
+            cfg["BLOCK_SIZE_M"] * cfg["BLOCK_SIZE_K"]
+            + cfg["BLOCK_SIZE_K"] * cfg["BLOCK_SIZE_N"]
+        ) * torch.bfloat16.itemsize
+        self.assertLessEqual(cfg["num_stages"] * tile_bytes, GFX950_LDS_BYTES)
+
+    def test_other_gpus_keep_the_default(self):
+        """The preset needs more LDS than gfx942 has, so it must stay gfx950-only."""
+        with patch.object(batch_invariant_ops, "_is_gfx95", False):
+            configs = batch_invariant_ops._matmul_persistent_configs()
+        self.assertEqual(
+            configs[torch.bfloat16],
+            {
+                "BLOCK_SIZE_M": 128,
+                "BLOCK_SIZE_N": 128,
+                "BLOCK_SIZE_K": 64,
+                "GROUP_SIZE_M": 8,
+                "num_stages": 3,
+                "num_warps": 8,
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()