From 1df9edcd015dba4ce30864a3202c74a7ffe6a356 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 11 May 2026 03:41:05 -0400 Subject: [PATCH] Use Torch `torch.mm` for Deepseek V3.2 Indexer GEMM (#23856) Co-authored-by: b8zhong --- python/sglang/srt/layers/attention/nsa/nsa_indexer.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 28854f2f6..854c76ecc 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -266,15 +266,8 @@ class Indexer(MultiPlatformOp): # avoiding an expensive FP8-to-bf16 dequantization. if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: x = x[2] - if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: - weight = self.weights_proj.weight - out = torch.empty( - (x.shape[0], weight.shape[0]), - dtype=torch.float32, - device=x.device, - ) - deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, weight, out) - return out + if _is_cuda: + return torch.mm(x, self.weights_proj.weight.t(), out_dtype=torch.float32) weights, _ = self.weights_proj(x) if _is_hip: