diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py index f59c824d226e..05c3e8ac9e96 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/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, diff --git a/test/registered/unit/batch_invariant_ops/test_matmul_persistent_configs.py b/test/registered/unit/batch_invariant_ops/test_matmul_persistent_configs.py new file mode 100644 index 000000000000..b530fa35ba6e --- /dev/null +++ b/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()