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

[NPU] make torch_native lora backend a little bit faster (#17228)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Egor Filimonov <44640852+ssshinigami@users.noreply.github.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
VDV1985
2026-03-07 20:14:46 +03:00
committed by GitHub
co-authored by gemini-code-assist[bot] Egor Filimonov ronnie_zheng
parent 5867c3fa80
commit 45bd30e29d
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/torch_ops/lora_ops.py
+2 -3
View File
@@ -35,7 +35,7 @@ def sgemm_lora_a_fwd(
x_seq = inputs[token_offset : token_offset + seq_len, :] x_seq = inputs[token_offset : token_offset + seq_len, :]
w_seq = weights[lora_idx, : num_slices * rank, :] w_seq = weights[lora_idx, : num_slices * rank, :]
result = torch.einsum("si, oi -> so", x_seq, w_seq) result = torch.mm(x_seq, w_seq.T)
output[token_offset : token_offset + seq_len, : num_slices * rank] = ( output[token_offset : token_offset + seq_len, : num_slices * rank] = (
scaling_tensor[lora_idx] * result scaling_tensor[lora_idx] * result
) )
@@ -98,11 +98,10 @@ def sgemm_lora_b_fwd(
lora_idx, slice_start_output:slice_end_output, :rank lora_idx, slice_start_output:slice_end_output, :rank
] # (slice_dim, rank) ] # (slice_dim, rank)
result = torch.einsum("si, oi -> so", x_slice, w_slice)
output[ output[
token_offset : token_offset + seq_len, token_offset : token_offset + seq_len,
slice_start_output:slice_end_output, slice_start_output:slice_end_output,
] += result ].add_(torch.mm(x_slice, w_slice.T))
token_offset += seq_len token_offset += seq_len