fix(deepgemm): add nsa init in deepgemm init

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