Remove DeepGEMM for indexer GEMM in piecewise NSA path (#26494)
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
@@ -162,12 +162,7 @@ if _is_cuda:
|
|||||||
softmax_scale: float,
|
softmax_scale: float,
|
||||||
q_scale: torch.Tensor,
|
q_scale: torch.Tensor,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
from sglang.srt.layers.deep_gemm_wrapper import entrypoint as deep_gemm_wrapper
|
out = torch.mm(x, 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)
|
|
||||||
weights = out * n_heads_inv_sqrt
|
weights = out * n_heads_inv_sqrt
|
||||||
weights = weights.unsqueeze(-1) * q_scale * softmax_scale
|
weights = weights.unsqueeze(-1) * q_scale * softmax_scale
|
||||||
return weights
|
return weights
|
||||||
|
|||||||
Reference in New Issue
Block a user