Remove DeepGEMM for indexer GEMM in piecewise NSA path (#26494)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
2026-05-28 00:15:23 -07:00
parent 686ef50672
commit e31ea50df8
1 changed files with 1 additions and 6 deletions
@@ -162,12 +162,7 @@ if _is_cuda:
        softmax_scale: float,
        q_scale: torch.Tensor,
    ) -> torch.Tensor:
-        from sglang.srt.layers.deep_gemm_wrapper import entrypoint as deep_gemm_wrapper
-
-        out = torch.empty(
-            (x.shape[0], weight.shape[0]), dtype=torch.float32, device=x.device
-        )
-        deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, weight, out)
+        out = torch.mm(x, weight.t(), out_dtype=torch.float32)
        weights = out * n_heads_inv_sqrt
        weights = weights.unsqueeze(-1) * q_scale * softmax_scale
        return weights