diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 3e5789b77..40c204468 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2529,6 +2529,14 @@ class ModelRunner(ModelRunnerKVCacheMixin): k, 1, # num_groups=1 for standard GEMM ) + if module.weight.dim() == 2 and module.weight.dtype == torch.bfloat16: + n, k = module.weight.shape + _maybe_compile_deep_gemm_one_type_all( + DeepGemmKernelType.GEMM_NT_BF16BF16F32, + n, + k, + 1, + ) # deep_gemm.preload_kernels() def init_threads_binding(self):