fix(deepgemm): add nsa init in deepgemm init

2026-04-05 00:13:52 +08:00
parent d9b6e90b35
commit a371eacd86
1 changed files with 8 additions and 0 deletions
--- 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):