304 lines
10 KiB
Python
304 lines
10 KiB
Python
import torch
|
|
|
|
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
|
|
from sglang.srt.lora.utils import LoRABatchInfo
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
|
from sglang.srt.utils import is_npu
|
|
|
|
if is_npu():
|
|
import sgl_kernel_npu # noqa: F401
|
|
import torch_npu # noqa: F401
|
|
|
|
|
|
class AscendLoRABackend(BaseLoRABackend):
|
|
name = "ascend"
|
|
|
|
def __init__(
|
|
self,
|
|
max_loras_per_batch: int,
|
|
device: torch.device,
|
|
**kwargs,
|
|
):
|
|
super().__init__(max_loras_per_batch, device)
|
|
|
|
def run_lora_a_sgemm(
|
|
self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
|
|
) -> torch.Tensor:
|
|
|
|
total_seq_len, _ = x.shape
|
|
_, weight_out_dim, _ = weights.shape
|
|
|
|
output_tensor = torch.zeros(
|
|
(total_seq_len, weight_out_dim), dtype=x.dtype, device=x.device
|
|
)
|
|
torch.ops.npu.sgmv_shrink(
|
|
x,
|
|
weights,
|
|
self.batch_info.weight_indices,
|
|
self.batch_info.seg_lens,
|
|
output_tensor,
|
|
1.0,
|
|
)
|
|
scaling = (
|
|
self.batch_info.scalings.gather(0, self.batch_info.weight_indices)
|
|
.repeat_interleave(self.batch_info.seg_lens, output_size=total_seq_len)
|
|
.unsqueeze(-1)
|
|
)
|
|
output_tensor *= scaling
|
|
|
|
return output_tensor
|
|
|
|
def run_lora_b_sgemm(
|
|
self,
|
|
x: torch.Tensor,
|
|
weights: torch.Tensor,
|
|
base_output: torch.Tensor = None,
|
|
*args,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
total_seq_len, _ = x.shape
|
|
_, weight_out_dim, _ = weights.shape
|
|
|
|
if base_output is None:
|
|
output_tensor = torch.zeros(
|
|
(total_seq_len, weight_out_dim), device=x.device, dtype=x.dtype
|
|
)
|
|
else:
|
|
output_tensor = base_output
|
|
|
|
torch.ops.npu.sgmv_expand(
|
|
x,
|
|
weights,
|
|
self.batch_info.weight_indices,
|
|
self.batch_info.seg_lens,
|
|
output_tensor,
|
|
0,
|
|
weight_out_dim,
|
|
)
|
|
|
|
return output_tensor
|
|
|
|
def run_qkv_lora(
|
|
self,
|
|
x: torch.Tensor,
|
|
qkv_lora_a: torch.Tensor,
|
|
qkv_lora_b: torch.Tensor,
|
|
output_offset: torch.Tensor,
|
|
output_offset_cpu: torch.Tensor,
|
|
max_qkv_out_dim: int,
|
|
base_output: torch.Tensor = None,
|
|
*args,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
num_slices = 3
|
|
assert isinstance(qkv_lora_b, torch.Tensor)
|
|
|
|
total_seq_len, _ = x.shape
|
|
_, weight_intermediate_dim, _ = qkv_lora_a.shape
|
|
_, weight_out_dim, _ = qkv_lora_b.shape
|
|
max_rank = weight_intermediate_dim // num_slices
|
|
|
|
if base_output is None:
|
|
output_tensor = torch.zeros(
|
|
(total_seq_len, weight_out_dim), device=x.device, dtype=x.dtype
|
|
)
|
|
else:
|
|
output_tensor = base_output
|
|
|
|
lora_a_output = torch.zeros(
|
|
total_seq_len, weight_intermediate_dim, dtype=x.dtype, device=x.device
|
|
)
|
|
torch.ops.npu.sgmv_shrink(
|
|
x,
|
|
qkv_lora_a,
|
|
self.batch_info.weight_indices,
|
|
self.batch_info.seg_lens,
|
|
lora_a_output,
|
|
1.0,
|
|
)
|
|
|
|
scaling = (
|
|
self.batch_info.scalings.gather(0, self.batch_info.weight_indices)
|
|
.repeat_interleave(self.batch_info.seg_lens, output_size=total_seq_len)
|
|
.unsqueeze(-1)
|
|
)
|
|
lora_a_output *= scaling
|
|
|
|
for slice_id in range(num_slices):
|
|
slice_offset = output_offset_cpu[slice_id]
|
|
slice_offset_next = output_offset_cpu[slice_id + 1]
|
|
slice_size = slice_offset_next - slice_offset
|
|
torch.ops.npu.sgmv_expand(
|
|
lora_a_output[:, (max_rank * slice_id) : (max_rank * (slice_id + 1))],
|
|
qkv_lora_b[:, slice_offset:slice_offset_next],
|
|
self.batch_info.weight_indices,
|
|
self.batch_info.seg_lens,
|
|
output_tensor,
|
|
slice_offset,
|
|
slice_size,
|
|
)
|
|
|
|
return output_tensor
|
|
|
|
def run_gate_up_lora(
|
|
self,
|
|
x: torch.Tensor,
|
|
gate_up_lora_a: torch.Tensor,
|
|
gate_up_lora_b: torch.Tensor,
|
|
base_output: torch.Tensor = None,
|
|
*args,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
|
|
num_slices = 2
|
|
assert isinstance(gate_up_lora_b, torch.Tensor)
|
|
|
|
total_seq_len, _ = x.shape
|
|
_, weight_intermediate_dim, _ = gate_up_lora_a.shape
|
|
_, weight_out_dim, _ = gate_up_lora_b.shape
|
|
slice_size = weight_out_dim // num_slices
|
|
max_rank = weight_intermediate_dim // num_slices
|
|
|
|
if base_output is None:
|
|
output_tensor = torch.zeros(
|
|
(total_seq_len, weight_out_dim), device=x.device, dtype=x.dtype
|
|
)
|
|
else:
|
|
output_tensor = base_output
|
|
|
|
lora_a_output = torch.zeros(
|
|
total_seq_len, weight_intermediate_dim, dtype=x.dtype, device=x.device
|
|
)
|
|
|
|
torch.ops.npu.sgmv_shrink(
|
|
x,
|
|
gate_up_lora_a,
|
|
self.batch_info.weight_indices,
|
|
self.batch_info.seg_lens,
|
|
lora_a_output,
|
|
1.0,
|
|
)
|
|
|
|
scaling = (
|
|
self.batch_info.scalings.gather(0, self.batch_info.weight_indices)
|
|
.repeat_interleave(self.batch_info.seg_lens, output_size=total_seq_len)
|
|
.unsqueeze(-1)
|
|
)
|
|
lora_a_output *= scaling
|
|
|
|
slice_offset = 0
|
|
for slice_id in range(num_slices):
|
|
torch.ops.npu.sgmv_expand(
|
|
lora_a_output[:, (max_rank * slice_id) : (max_rank * (slice_id + 1))],
|
|
gate_up_lora_b[:, slice_offset : slice_offset + slice_size],
|
|
self.batch_info.weight_indices,
|
|
self.batch_info.seg_lens,
|
|
output_tensor,
|
|
slice_offset,
|
|
slice_size,
|
|
)
|
|
slice_offset += slice_size
|
|
|
|
return output_tensor
|
|
|
|
def init_cuda_graph_batch_info(
|
|
self,
|
|
max_bs_in_cuda_graph: int,
|
|
num_tokens_per_bs: int,
|
|
):
|
|
with torch.device("npu"):
|
|
self.npu_graph_batch_info = LoRABatchInfo(
|
|
bs=max_bs_in_cuda_graph,
|
|
use_cuda_graph=True,
|
|
num_segments=None,
|
|
seg_lens=torch.full(
|
|
(max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
|
|
),
|
|
seg_indptr=torch.empty(max_bs_in_cuda_graph + 1, dtype=torch.int32),
|
|
max_len=num_tokens_per_bs,
|
|
weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
|
|
lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
|
|
scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
|
|
permutation=None,
|
|
)
|
|
|
|
# Initialize seg_indptr for NPU graph as they remain constant
|
|
# across batches.
|
|
torch.cumsum(
|
|
self.npu_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
|
|
dim=0,
|
|
out=self.npu_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
|
|
)
|
|
|
|
def prepare_lora_batch(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
weight_indices: list[int],
|
|
lora_ranks: list[int],
|
|
scalings: list[float],
|
|
use_cuda_graph: bool,
|
|
):
|
|
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
|
weight_indices_tensor = torch.tensor(
|
|
weight_indices, dtype=torch.int32, pin_memory=True, device="cpu"
|
|
)
|
|
lora_ranks_tensor = torch.tensor(
|
|
lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
|
|
)
|
|
scalings_tensor = torch.tensor(
|
|
scalings, dtype=torch.float, pin_memory=True, device="cpu"
|
|
)
|
|
|
|
bs = forward_batch.batch_size
|
|
|
|
if use_cuda_graph:
|
|
assert (
|
|
self.npu_graph_batch_info is not None
|
|
), "NPU Graph batch info is not initialized."
|
|
batch_info = self.npu_graph_batch_info
|
|
batch_info.bs = forward_batch.batch_size
|
|
batch_info.num_segments = forward_batch.batch_size
|
|
else:
|
|
max_len = (
|
|
# Calculate max_len from the CPU copy to avoid D2H transfer.
|
|
max(forward_batch.extend_seq_lens_cpu)
|
|
if forward_batch.forward_mode.is_extend()
|
|
else 1
|
|
)
|
|
seg_lens = (
|
|
forward_batch.extend_seq_lens
|
|
if forward_batch.forward_mode.is_extend()
|
|
else torch.ones(bs, dtype=torch.int32, device=self.device)
|
|
)
|
|
seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
|
|
seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
|
|
|
|
batch_info = LoRABatchInfo(
|
|
bs=forward_batch.batch_size,
|
|
num_segments=forward_batch.batch_size,
|
|
max_len=max_len,
|
|
use_cuda_graph=False,
|
|
seg_lens=seg_lens,
|
|
seg_indptr=seg_indptr,
|
|
weight_indices=torch.empty(
|
|
(bs,), dtype=torch.int32, device=self.device
|
|
),
|
|
lora_ranks=torch.empty(
|
|
(self.max_loras_per_batch,), dtype=torch.int32, device=self.device
|
|
),
|
|
scalings=torch.empty(
|
|
(self.max_loras_per_batch,), dtype=torch.float, device=self.device
|
|
),
|
|
permutation=None,
|
|
)
|
|
|
|
# Copy to device asynchronously
|
|
batch_info.lora_ranks[: self.max_loras_per_batch].copy_(
|
|
lora_ranks_tensor, non_blocking=True
|
|
)
|
|
batch_info.scalings[: self.max_loras_per_batch].copy_(
|
|
scalings_tensor, non_blocking=True
|
|
)
|
|
batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True)
|
|
self.batch_info = batch_info
|