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

[LoRA] Torch native backend: rework implementation and updated tests (#15187)

Browse Source
This commit is contained in:
Vladimir Serov
2025-12-30 06:48:51 +03:00
committed by GitHub
parent 0d003e34b0
commit 9e263c2162
14 changed files with 958 additions and 1000 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/sglang/srt/lora/backend/chunked_backend.py
View File

@@ -5,7 +5,7 @@ from sglang.srt.lora.triton_ops import (
chunked_sgmv_lora_expand_forward,
chunked_sgmv_lora_shrink_forward,
)
from sglang.srt.lora.utils import LoRABatchInfo
from sglang.srt.lora.utils import LoRABatchInfo, generate_sequence_lengths
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.server_args import ServerArgs
@@ -283,24 +283,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
"""
with torch.device("cpu"):
seq_weight_indices = torch.tensor(seq_weight_indices, dtype=torch.int32)
if forward_batch.forward_mode.is_decode():
seg_lens_cpu = torch.ones(forward_batch.batch_size, dtype=torch.int32)
elif forward_batch.forward_mode.is_target_verify():
seg_lens_cpu = torch.full(
size=(forward_batch.batch_size,),
fill_value=forward_batch.spec_info.draft_token_num,
dtype=torch.int32,
)
elif forward_batch.forward_mode.is_extend():
seg_lens_cpu = torch.tensor(
forward_batch.extend_seq_lens_cpu,
dtype=torch.int32,
)
else:
raise ValueError(
f"Unsupported forward mode: {forward_batch.forward_mode}"
)
seg_lens_cpu = generate_sequence_lengths(forward_batch)
row_weight_indices = torch.repeat_interleave(
seq_weight_indices, seg_lens_cpu

276
python/sglang/srt/lora/backend/torch_backend.py
View File

@@ -1,11 +1,29 @@
from dataclasses import dataclass
from typing import Optional
import torch
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
from sglang.srt.lora.torch_ops import sgmv_expand, sgmv_expand_slice, sgmv_shrink
from sglang.srt.lora.utils import LoRABatchInfo
from sglang.srt.lora.torch_ops import sgemm_lora_a_fwd, sgemm_lora_b_fwd
from sglang.srt.lora.utils import LoRABatchInfo, generate_sequence_lengths
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@dataclass
class TorchNativeLoRABatchInfo(LoRABatchInfo):
# ranks of each lora adapter, in shape (lora_num,) placed on cpu device
lora_ranks_cpu: Optional[torch.Tensor] = None
# Indice pointers of each segment in shape (num_segments + 1, ) placed on cpu device
seg_indptr_cpu: Optional[torch.Tensor] = None
# Lengths of each segments in shape (num_segments,) placed on cpu device
seg_lens_cpu: Optional[torch.Tensor] = None
# The index of lora adapter used by each segment, in shape (num_segments,) placed on cpu device
weight_indices_cpu: Optional[torch.Tensor] = None
class TorchNativeLoRABackend(BaseLoRABackend):
name = "torch_native"
@@ -20,27 +38,15 @@ class TorchNativeLoRABackend(BaseLoRABackend):
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
output_tensor = sgemm_lora_a_fwd(
inputs=x,
weights=weights,
weight_indices=self.batch_info.weight_indices_cpu,
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
scaling_tensor=self.batch_info.scalings,
num_slices=1,
)
sgmv_shrink(
x,
weights,
output_tensor,
self.batch_info.seg_lens,
self.batch_info.weight_indices,
1.0,
)
scaling = torch.repeat_interleave(
self.batch_info.scalings[self.batch_info.weight_indices],
self.batch_info.seg_lens,
output_size=total_seq_len,
).unsqueeze(-1)
output_tensor = output_tensor * scaling
return output_tensor
@@ -52,23 +58,18 @@ class TorchNativeLoRABackend(BaseLoRABackend):
*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
sgmv_expand(
x,
weights,
output_tensor,
self.batch_info.seg_lens,
self.batch_info.weight_indices,
True,
output_offset = torch.tensor(
[0, weight_out_dim], dtype=torch.int32, device="cpu"
)
output_tensor = sgemm_lora_b_fwd(
inputs=x,
weights=weights,
weight_indices=self.batch_info.weight_indices_cpu,
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
slice_offsets=output_offset,
base_output=base_output,
)
return output_tensor
@@ -86,52 +87,25 @@ class TorchNativeLoRABackend(BaseLoRABackend):
**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
lora_a_output = sgemm_lora_a_fwd(
inputs=x,
weights=qkv_lora_a,
weight_indices=self.batch_info.weight_indices_cpu,
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
scaling_tensor=self.batch_info.scalings,
num_slices=num_slices,
)
sgmv_shrink(
x,
qkv_lora_a,
lora_a_output,
self.batch_info.seg_lens,
self.batch_info.weight_indices,
1.0,
)
scaling = torch.repeat_interleave(
self.batch_info.scalings[self.batch_info.weight_indices],
self.batch_info.seg_lens,
output_size=total_seq_len,
).unsqueeze(-1)
lora_a_output = 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
sgmv_expand_slice(
lora_a_output[:, (max_rank * slice_id) : (max_rank * (slice_id + 1))],
qkv_lora_b[:, slice_offset:slice_offset_next],
output_tensor,
self.batch_info.seg_lens,
self.batch_info.weight_indices,
slice_offset,
slice_size,
True,
)
output_tensor = sgemm_lora_b_fwd(
inputs=lora_a_output,
weights=qkv_lora_b,
weight_indices=self.batch_info.weight_indices_cpu,
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
slice_offsets=output_offset_cpu,
base_output=base_output,
)
return output_tensor
@@ -144,54 +118,32 @@ class TorchNativeLoRABackend(BaseLoRABackend):
*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
output_offset = torch.tensor(
[0, slice_size, weight_out_dim], dtype=torch.int32, device="cpu"
)
sgmv_shrink(
x,
gate_up_lora_a,
lora_a_output,
self.batch_info.seg_lens,
self.batch_info.weight_indices,
1.0,
)
scaling = torch.repeat_interleave(
self.batch_info.scalings[self.batch_info.weight_indices],
self.batch_info.seg_lens,
output_size=total_seq_len,
).unsqueeze(-1)
lora_a_output = lora_a_output * scaling
slice_offset = 0
for slice_id in range(num_slices):
sgmv_expand_slice(
lora_a_output[:, (max_rank * slice_id) : (max_rank * (slice_id + 1))],
gate_up_lora_b[:, slice_offset : slice_offset + slice_size],
output_tensor,
self.batch_info.seg_lens,
self.batch_info.weight_indices,
slice_offset,
slice_size,
True,
)
slice_offset += slice_size
lora_a_output = sgemm_lora_a_fwd(
inputs=x,
weights=gate_up_lora_a,
weight_indices=self.batch_info.weight_indices_cpu,
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
scaling_tensor=self.batch_info.scalings,
num_slices=num_slices,
)
output_tensor = sgemm_lora_b_fwd(
inputs=lora_a_output,
weights=gate_up_lora_b,
weight_indices=self.batch_info.weight_indices_cpu,
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
slice_offsets=output_offset,
base_output=base_output,
)
return output_tensor
@@ -201,19 +153,19 @@ class TorchNativeLoRABackend(BaseLoRABackend):
num_tokens_per_bs: int,
):
with torch.device("cuda"):
self.cuda_graph_batch_info = LoRABatchInfo(
bs=max_bs_in_cuda_graph,
self.cuda_graph_batch_info = TorchNativeLoRABatchInfo(
use_cuda_graph=True,
num_segments=None,
bs=max_bs_in_cuda_graph,
num_segments=self.max_loras_per_batch,
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,
seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
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,
max_len=num_tokens_per_bs,
)
# Initialize seg_indptr for CUDA graph as they remain constant
@@ -232,10 +184,36 @@ class TorchNativeLoRABackend(BaseLoRABackend):
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"
original_seq_lens_cpu = generate_sequence_lengths(forward_batch, device="cpu")
original_weight_indices_tensor = torch.tensor(
weight_indices, dtype=torch.int32, device="cpu"
)
unique_weight_indices_tensor, inverse_weight_indices_tensor = (
torch.unique_consecutive(
original_weight_indices_tensor, return_inverse=True
)
)
seg_lens_cpu = (
torch.zeros_like(
unique_weight_indices_tensor, dtype=torch.int32, device="cpu"
)
.scatter_add_(
0,
inverse_weight_indices_tensor,
original_seq_lens_cpu,
)
.pin_memory()
)
seg_indptr_cpu = torch.zeros(
(len(seg_lens_cpu) + 1,), dtype=torch.int32, pin_memory=True
)
seg_indptr_cpu[1:] = torch.cumsum(seg_lens_cpu, dim=0)
# Use pinned memory to avoid synchronizations during host-to-device transfer
weight_indices_tensor = unique_weight_indices_tensor.pin_memory()
lora_ranks_tensor = torch.tensor(
lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
)
@@ -253,27 +231,17 @@ class TorchNativeLoRABackend(BaseLoRABackend):
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)
max_len = max(seg_lens_cpu)
batch_info = LoRABatchInfo(
batch_info = TorchNativeLoRABatchInfo(
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,
seg_lens=torch.empty((bs,), dtype=torch.int32, device=self.device),
seg_indptr=torch.empty(
(bs + 1,), dtype=torch.int32, device=self.device
),
weight_indices=torch.empty(
(bs,), dtype=torch.int32, device=self.device
),
@@ -294,4 +262,14 @@ class TorchNativeLoRABackend(BaseLoRABackend):
scalings_tensor, non_blocking=True
)
batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True)
batch_info.seg_indptr[: len(seg_indptr_cpu)].copy_(
seg_indptr_cpu, non_blocking=True
)
batch_info.seg_lens[: len(seg_lens_cpu)].copy_(seg_lens_cpu, non_blocking=True)
batch_info.lora_ranks_cpu = lora_ranks_tensor
batch_info.seg_indptr_cpu = seg_indptr_cpu
batch_info.seg_lens_cpu = seg_lens_cpu
batch_info.weight_indices_cpu = weight_indices_tensor
self.batch_info = batch_info

7
python/sglang/srt/lora/torch_ops/__init__.py
View File

@@ -1,7 +1,6 @@
from .lora_ops import sgmv_expand, sgmv_expand_slice, sgmv_shrink
from .lora_ops import sgemm_lora_a_fwd, sgemm_lora_b_fwd
__all__ = [
"sgmv_expand",
"sgmv_expand_slice",
"sgmv_shrink",
"sgemm_lora_a_fwd",
"sgemm_lora_b_fwd",
]

186
python/sglang/srt/lora/torch_ops/lora_ops.py
View File

@@ -1,125 +1,109 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Optional
import torch
def sgmv_expand(
def sgemm_lora_a_fwd(
inputs: torch.Tensor,
lora_b_weights: torch.Tensor,
output_tensor: torch.Tensor,
seq_len_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
add_inputs: bool = False,
weights: torch.Tensor,
weight_indices: torch.Tensor,
seg_len_tensor: torch.Tensor,
lora_ranks: torch.Tensor,
scaling_tensor: torch.Tensor,
num_slices: int = 1,
):
total_seq_len, _ = inputs.shape
exploded_indices = torch.repeat_interleave(
lora_indices_tensor, seq_len_tensor, output_size=total_seq_len
total_seq_len, input_dim = inputs.shape
if weights.numel() == 0:
return torch.zeros(total_seq_len, 0, dtype=inputs.dtype, device=inputs.device)
num_loras, weight_out_dim, _ = weights.shape
max_rank = weight_out_dim // num_slices
output = torch.zeros(
total_seq_len, num_slices * max_rank, dtype=inputs.dtype, device=inputs.device
)
bgmv_expand(inputs, lora_b_weights, output_tensor, exploded_indices, add_inputs)
token_offset = 0
for lora_idx, seq_len, rank in zip(
weight_indices, seg_len_tensor, lora_ranks[weight_indices]
):
if seq_len == 0:
continue
if rank > 0:
x_seq = inputs[token_offset : token_offset + seq_len, :]
w_seq = weights[lora_idx, : num_slices * rank, :]
result = torch.einsum("si, oi -> so", x_seq, w_seq)
output[token_offset : token_offset + seq_len, : num_slices * rank] = (
scaling_tensor[lora_idx] * result
)
token_offset += seq_len
return output
def bgmv_expand(
def sgemm_lora_b_fwd(
inputs: torch.Tensor,
lora_b_weights: torch.Tensor,
output_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
add_inputs: bool = True,
weights: torch.Tensor,
weight_indices: torch.Tensor,
seg_len_tensor: torch.Tensor,
lora_ranks: torch.Tensor,
slice_offsets: torch.Tensor,
base_output: Optional[torch.Tensor] = None,
):
selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
if len(selected_loras.shape) == 4:
selected_loras = selected_loras.squeeze(dim=1)
inputs = inputs.to(dtype=output_tensor.dtype)
outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
total_seq_len, _ = inputs.shape
num_loras, weight_out_dim, _ = weights.shape
total_output_dim = slice_offsets[-1].item() if len(slice_offsets) > 0 else 0
limit = output_tensor.shape[0]
if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
limit = 1
if weights.numel() == 0:
return torch.zeros(
total_seq_len, total_output_dim, dtype=inputs.dtype, device=inputs.device
)
# LoRA adapter and model may add different amounts of padding to output
common_len = min(outputs.shape[1], output_tensor.shape[1])
num_slices = len(slice_offsets) - 1
if add_inputs:
output_tensor[:, :common_len] += outputs[:limit, :common_len]
if base_output is not None:
output = base_output
else:
output_tensor[:, :common_len] = outputs[:limit, :common_len]
output = torch.zeros(
total_seq_len, total_output_dim, dtype=inputs.dtype, device=inputs.device
)
token_offset = 0
for lora_idx, seq_len, rank in zip(
weight_indices, seg_len_tensor, lora_ranks[weight_indices]
):
if seq_len == 0:
continue
def sgmv_shrink(
inputs: torch.Tensor,
lora_a_weights: torch.Tensor,
output_tensor: torch.Tensor,
seq_len_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
scaling: float,
):
total_seq_len, _ = inputs.shape
exploded_indices = torch.repeat_interleave(
lora_indices_tensor, seq_len_tensor, output_size=total_seq_len
)
if rank == 0:
token_offset += seq_len
continue
bgmv_shrink(inputs, lora_a_weights, output_tensor, exploded_indices, scaling)
for slice_idx in range(num_slices):
slice_start_input = slice_idx * rank
slice_end_input = (slice_idx + 1) * rank
slice_start_output = slice_offsets[slice_idx]
slice_end_output = slice_offsets[slice_idx + 1]
def bgmv_shrink(
inputs: torch.Tensor,
lora_a_weights: torch.Tensor,
output_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
scaling: float = 1.0,
):
selected_loras = lora_a_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
if len(selected_loras.shape) == 4:
selected_loras = selected_loras.squeeze(dim=1)
inputs = inputs.to(dtype=output_tensor.dtype)
outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
x_slice = inputs[
token_offset : token_offset + seq_len :,
slice_start_input:slice_end_input,
] # (seq_len, rank)
w_slice = weights[
lora_idx, slice_start_output:slice_end_output, :rank
] # (slice_dim, rank)
output_tensor[:, : outputs.shape[1]] = scaling * outputs[:]
result = torch.einsum("si, oi -> so", x_slice, w_slice)
output[
token_offset : token_offset + seq_len,
slice_start_output:slice_end_output,
] += result
token_offset += seq_len
def sgmv_expand_slice(
inputs: torch.Tensor,
lora_b_weights: torch.Tensor,
output_tensor: torch.Tensor,
seq_len_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
slice_offset: int,
slice_size: int,
add_inputs: bool = False,
):
total_seq_len, _ = inputs.shape
exploded_indices = torch.repeat_interleave(
lora_indices_tensor, seq_len_tensor, output_size=total_seq_len
)
bgmv_expand_slice(
inputs,
lora_b_weights,
output_tensor,
exploded_indices,
slice_offset,
slice_size,
add_inputs,
)
def bgmv_expand_slice(
inputs: torch.Tensor,
lora_b_weights: torch.Tensor,
output_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
slice_offset: int,
slice_size: int,
add_inputs: bool = True,
):
selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
inputs = inputs.to(dtype=output_tensor.dtype)
if len(selected_loras.shape) == 4:
selected_loras = selected_loras.squeeze(dim=1)
outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
if add_inputs:
output_tensor[:, slice_offset : slice_offset + slice_size] += outputs[:]
else:
output_tensor[:, slice_offset : slice_offset + slice_size] = outputs[:]
return output

29
python/sglang/srt/lora/utils.py
View File

@@ -4,6 +4,7 @@ from typing import Iterable, Optional, Set, Tuple
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils.hf_transformers_utils import AutoConfig
@@ -151,3 +152,31 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj"]
def generate_sequence_lengths(
forward_batch: ForwardBatch, device: Optional[torch.device] = None
) -> torch.Tensor:
device = torch.get_default_device() if device is None else device
with torch.device(device):
if forward_batch.forward_mode.is_decode():
seg_lens = torch.ones(forward_batch.batch_size, dtype=torch.int32)
elif forward_batch.forward_mode.is_target_verify():
seg_lens = torch.full(
size=(forward_batch.batch_size,),
fill_value=forward_batch.spec_info.draft_token_num,
dtype=torch.int32,
)
elif forward_batch.forward_mode.is_extend():
seg_lens = (
forward_batch.extend_seq_lens
if forward_batch.extend_seq_lens.device == device
else torch.tensor(
forward_batch.extend_seq_lens_cpu,
dtype=torch.int32,
)
)
else:
raise ValueError(f"Unsupported forward mode: {forward_batch.forward_mode}")
return seg_lens

147
python/sglang/test/lora_utils.py
View File

@@ -1,6 +1,6 @@
import dataclasses
import random
from typing import List
from typing import List, Optional
import torch
@@ -116,6 +116,151 @@ ALL_OTHER_MULTI_LORA_MODELS = [
]
def safe_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""Matrix multiplication with mixed precision handling for float16"""
result = torch.matmul(a.float(), b.float())
return result.to(a.dtype)
def reference_sgmv_shrink(
x: torch.Tensor,
weights: torch.Tensor,
weight_indices: torch.Tensor,
seq_lengths: torch.Tensor,
lora_ranks: torch.Tensor,
lora_scalings: torch.Tensor,
num_slices: int = 1,
) -> torch.Tensor:
"""
Simple sequence-level reference implementation of SGMV shrink operation.
Args:
x: (total_seq_len, input_dim) - Input activations
weights: (num_loras, num_slices * max_rank, input_dim) - LoRA A weights
weight_indices: LoRA idx for each sequence
seq_lengths: Length of each sequence
lora_ranks: LoRA rank for each LoRA adapters
lora_scalings: LoRA scaling for each LoRA adapters
num_slices: Number of slices (3 for QKV, 2 for gate_up, 1 for others)
Returns:
output: (total_seq_len, num_slices * max_rank) - Intermediate activations
"""
if weights.numel() == 0:
total_seq_len = x.shape[0]
return torch.zeros(total_seq_len, 0, dtype=x.dtype, device=x.device)
total_seq_len, _ = x.shape
_, weight_out_dim, _ = weights.shape
max_rank = weight_out_dim // num_slices
output = torch.zeros(
total_seq_len, num_slices * max_rank, dtype=x.dtype, device=x.device
)
token_offset = 0
for lora_idx, seq_len, rank, scaling in zip(
weight_indices,
seq_lengths,
lora_ranks[weight_indices],
lora_scalings[weight_indices],
):
if seq_len == 0:
continue
if rank > 0:
x_seq = x[token_offset : token_offset + seq_len, :]
w_seq = weights[lora_idx, : num_slices * rank, :]
result = safe_matmul(x_seq, w_seq.t())
output[token_offset : token_offset + seq_len, : num_slices * rank] = (
scaling * result
)
token_offset += seq_len
return output
def reference_sgmv_expand(
x: torch.Tensor,
weights: torch.Tensor,
weight_indices: torch.Tensor,
seq_lengths: torch.Tensor,
lora_ranks: torch.Tensor,
slice_offsets: torch.Tensor,
base_output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Simple sequence-level reference implementation of SGMV expand operation.
Args:
x: (total_seq_len, num_slices * max_rank) - Intermediate activations
weights: (num_loras, output_dim, max_rank) - LoRA B weights
weight_indices: LoRA idx for each sequence
seq_lengths: Length of each sequence
lora_ranks: LoRA rank for each LoRA adapters
slice_offsets: Tensor defining slice boundaries
base_output: Optional base output to accumulate into
Returns:
output: (total_seq_len, total_output_dim) - Final output
"""
if weights.numel() == 0:
total_seq_len = x.shape[0]
total_output_dim = slice_offsets[-1].item() if len(slice_offsets) > 0 else 0
return torch.zeros(
total_seq_len, total_output_dim, dtype=x.dtype, device=x.device
)
total_seq_len, _ = x.shape
num_slices = len(slice_offsets) - 1
if base_output is not None:
output = base_output.clone()
else:
total_output_dim = slice_offsets[-1].item()
output = torch.zeros(
total_seq_len, total_output_dim, dtype=x.dtype, device=x.device
)
token_offset = 0
for lora_idx, seq_len, rank in zip(
weight_indices, seq_lengths, lora_ranks[weight_indices]
):
if seq_len == 0:
continue
if rank > 0:
# Extract sequence intermediate activations
x_seq = x[
token_offset : token_offset + seq_len, : num_slices * rank
] # (seq_len, num_slices * rank)
for slice_idx in range(num_slices):
slice_start_input = slice_idx * rank
slice_end_input = (slice_idx + 1) * rank
slice_start_output = slice_offsets[slice_idx].item()
slice_end_output = slice_offsets[slice_idx + 1].item()
x_slice = x_seq[:, slice_start_input:slice_end_input] # (seq_len, rank)
w_slice = weights[
lora_idx, slice_start_output:slice_end_output, :rank
] # (slice_dim, rank)
result = safe_matmul(x_slice, w_slice.t()) # (seq_len, slice_dim)
output[
token_offset : token_offset + seq_len,
slice_start_output:slice_end_output,
] += result
token_offset += seq_len
return output
def run_lora_test_one_by_one(
prompts: List[str],
model_case: LoRAModelCase,