From a371eacd86096fbd2495462f033d196d6f800a5d Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Sun, 5 Apr 2026 00:13:52 +0800
Subject: [PATCH] fix(deepgemm): add nsa init in deepgemm init

---
 python/sglang/srt/model_executor/model_runner.py | 8 ++++++++
 1 file changed, 8 insertions(+)

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):