minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Use Torch torch.mm for Deepseek V3.2 Indexer GEMM (#23856)

Browse Source 
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-05-11 00:41:05 -07:00
committed by GitHub
co-authored by b8zhong
parent 1d80a1a9fe
commit 1df9edcd01
1 changed files with 2 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+2 -9
View File
@@ -266,15 +266,8 @@ class Indexer(MultiPlatformOp):
# avoiding an expensive FP8-to-bf16 dequantization. # avoiding an expensive FP8-to-bf16 dequantization.
if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3:
x = x[2] x = x[2]
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: if _is_cuda:
weight = self.weights_proj.weight return torch.mm(x, self.weights_proj.weight.t(), out_dtype=torch.float32)
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
weights, _ = self.weights_proj(x) weights, _ = self.weights_proj(x)
if _is_hip: if _is_hip: