chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
107
third_party/sglang/test/manual/lora/test_lora_cuda_graph.py
vendored
Normal file
107
third_party/sglang/test/manual/lora/test_lora_cuda_graph.py
vendored
Normal file
@@ -0,0 +1,107 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import unittest
|
||||
from typing import List
|
||||
|
||||
from sglang.test.lora_utils import (
|
||||
ALL_OTHER_LORA_MODELS,
|
||||
CI_LORA_MODELS,
|
||||
DEFAULT_PROMPTS,
|
||||
TORCH_DTYPES,
|
||||
LoRAModelCase,
|
||||
run_lora_test_by_batch,
|
||||
run_lora_test_one_by_one,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase, is_in_ci
|
||||
|
||||
TEST_CUDA_GRAPH_PADDING_PROMPTS = [
|
||||
"AI is a field of computer science focused on",
|
||||
"""
|
||||
### Instruction:
|
||||
Tell me about llamas and alpacas
|
||||
### Response:
|
||||
Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
|
||||
### Question 2:
|
||||
What do you know about llamas?
|
||||
### Answer:
|
||||
""",
|
||||
"Computer science is the study of",
|
||||
]
|
||||
|
||||
|
||||
class TestLoRACudaGraph(CustomTestCase):
|
||||
|
||||
def _run_without_cuda_graph_on_model_cases(self, model_cases: List[LoRAModelCase]):
|
||||
# Since we have already enabled CUDA graph by default in other lora tests,
|
||||
# we only need to run lora tests without CUDA graph here.
|
||||
for model_case in model_cases:
|
||||
# If skip_long_prompt is True, filter out prompts longer than 1000 characters
|
||||
prompts = (
|
||||
DEFAULT_PROMPTS
|
||||
if not model_case.skip_long_prompt
|
||||
else [p for p in DEFAULT_PROMPTS if len(p) < 1000]
|
||||
)
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
run_lora_test_one_by_one(
|
||||
prompts,
|
||||
model_case,
|
||||
torch_dtype,
|
||||
max_new_tokens=32,
|
||||
disable_cuda_graph=True,
|
||||
test_tag="without_cuda_graph",
|
||||
)
|
||||
|
||||
def _run_cuda_graph_padding_on_model_cases(self, model_cases: List[LoRAModelCase]):
|
||||
for model_case in model_cases:
|
||||
# Run a batch size of 3, which will not be captured by CUDA graph and need padding
|
||||
prompts = TEST_CUDA_GRAPH_PADDING_PROMPTS
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
run_lora_test_by_batch(
|
||||
prompts,
|
||||
model_case,
|
||||
torch_dtype,
|
||||
max_new_tokens=32,
|
||||
disable_cuda_graph=False,
|
||||
test_tag="cuda_graph_padding",
|
||||
)
|
||||
|
||||
def test_ci_lora_models(self):
|
||||
self._run_without_cuda_graph_on_model_cases(CI_LORA_MODELS)
|
||||
self._run_cuda_graph_padding_on_model_cases(CI_LORA_MODELS)
|
||||
|
||||
def test_all_lora_models(self):
|
||||
if is_in_ci():
|
||||
return
|
||||
|
||||
# Retain ONLY_RUN check here
|
||||
filtered_models = []
|
||||
for model_case in ALL_OTHER_LORA_MODELS:
|
||||
if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
|
||||
continue
|
||||
filtered_models.append(model_case)
|
||||
|
||||
self._run_without_cuda_graph_on_model_cases(filtered_models)
|
||||
self._run_cuda_graph_padding_on_model_cases(filtered_models)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
unittest.main(warnings="ignore")
|
||||
63
third_party/sglang/test/manual/lora/test_lora_llama4.py
vendored
Normal file
63
third_party/sglang/test/manual/lora/test_lora_llama4.py
vendored
Normal file
@@ -0,0 +1,63 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
SimpleNamespace(
|
||||
model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
||||
tp_size=8,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
class TestLlama4LoRA(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
def test_bringup(self):
|
||||
for model in MODELS:
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
model.model,
|
||||
self.base_url,
|
||||
timeout=3 * DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--enable-lora",
|
||||
"--max-lora-rank",
|
||||
"64",
|
||||
"--lora-target-modules",
|
||||
"all",
|
||||
"--tp-size",
|
||||
str(model.tp_size),
|
||||
"--context-length",
|
||||
"262144",
|
||||
"--attention-backend",
|
||||
"fa3",
|
||||
],
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing {model.model}: {e}")
|
||||
self.fail(f"Test failed for {model.model}: {e}")
|
||||
|
||||
finally:
|
||||
# Ensure process cleanup happens regardless of success/failure
|
||||
if process is not None and process.poll() is None:
|
||||
print(f"Cleaning up process {process.pid}")
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
218
third_party/sglang/test/manual/lora/test_lora_ops.py
vendored
Normal file
218
third_party/sglang/test/manual/lora/test_lora_ops.py
vendored
Normal file
@@ -0,0 +1,218 @@
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.lora.torch_ops.lora_ops import sgemm_lora_a_fwd, sgemm_lora_b_fwd
|
||||
from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestLoraOps(CustomTestCase):
|
||||
def test_sgemm_lora_a_fwd(self):
|
||||
batch_size = 2
|
||||
input_dim = 1024
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
|
||||
lora_ranks = random.sample(
|
||||
possible_lora_ranks,
|
||||
counts=[num_loras] * len(possible_lora_ranks),
|
||||
k=num_loras,
|
||||
)
|
||||
|
||||
max_lora_rank = max(lora_ranks)
|
||||
|
||||
possible_lora_scaling = [0.25, 0.5, 1.0, 2.0, 4.0]
|
||||
lora_scaling = random.sample(
|
||||
possible_lora_scaling,
|
||||
counts=[num_loras] * len(possible_lora_scaling),
|
||||
k=num_loras,
|
||||
)
|
||||
|
||||
inputs = torch.randn(batch_size, input_dim, dtype=dtype)
|
||||
lora_a_weights = torch.randn(num_loras, max_lora_rank, input_dim, dtype=dtype)
|
||||
lora_indices_tensor = torch.randint(
|
||||
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
|
||||
)
|
||||
seq_len_tensor = torch.ones(batch_size, dtype=torch.int32, device="cpu")
|
||||
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
|
||||
lora_scaling_tensor = torch.tensor(
|
||||
lora_scaling, dtype=torch.float16, device="cpu"
|
||||
)
|
||||
|
||||
expect_output = reference_sgmv_shrink(
|
||||
inputs,
|
||||
lora_a_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
lora_scaling_tensor,
|
||||
)
|
||||
|
||||
actual_output = sgemm_lora_a_fwd(
|
||||
inputs,
|
||||
lora_a_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
lora_scaling_tensor,
|
||||
)
|
||||
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
def test_sgemm_lora_b_fwd(self):
|
||||
batch_size = 2
|
||||
output_dim = 1024
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
|
||||
lora_ranks = random.sample(
|
||||
possible_lora_ranks,
|
||||
counts=[num_loras] * len(possible_lora_ranks),
|
||||
k=num_loras,
|
||||
)
|
||||
|
||||
max_lora_rank = max(lora_ranks)
|
||||
|
||||
inputs = torch.randn(batch_size, max_lora_rank, dtype=dtype)
|
||||
lora_b_weights = torch.randn(num_loras, output_dim, max_lora_rank, dtype=dtype)
|
||||
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
|
||||
seq_len_tensor = torch.ones(batch_size, dtype=torch.int32, device="cpu")
|
||||
lora_indices_tensor = torch.randint(
|
||||
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
|
||||
)
|
||||
slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cpu")
|
||||
|
||||
expect_output = reference_sgmv_expand(
|
||||
inputs,
|
||||
lora_b_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets,
|
||||
)
|
||||
|
||||
actual_output = sgemm_lora_b_fwd(
|
||||
inputs,
|
||||
lora_b_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets,
|
||||
)
|
||||
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
def test_sgemm_lora_a_fwd_expand(self):
|
||||
batch_size = 2
|
||||
input_dim = 1024
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
|
||||
lora_ranks = random.sample(
|
||||
possible_lora_ranks,
|
||||
counts=[num_loras] * len(possible_lora_ranks),
|
||||
k=num_loras,
|
||||
)
|
||||
|
||||
max_lora_rank = max(lora_ranks)
|
||||
|
||||
possible_lora_scaling = [0.25, 0.5, 1.0, 2.0, 4.0]
|
||||
lora_scaling = random.sample(
|
||||
possible_lora_scaling,
|
||||
counts=[num_loras] * len(possible_lora_scaling),
|
||||
k=num_loras,
|
||||
)
|
||||
|
||||
seq_len_tensor = torch.randint(
|
||||
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
|
||||
)
|
||||
|
||||
seq_len = sum(seq_len_tensor)
|
||||
|
||||
inputs = torch.randn(seq_len, input_dim, dtype=dtype)
|
||||
lora_a_weights = torch.randn(num_loras, max_lora_rank, input_dim, dtype=dtype)
|
||||
lora_indices_tensor = torch.randint(
|
||||
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
|
||||
)
|
||||
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
|
||||
lora_scaling_tensor = torch.tensor(
|
||||
lora_scaling, dtype=torch.float16, device="cpu"
|
||||
)
|
||||
|
||||
expect_output = reference_sgmv_shrink(
|
||||
inputs,
|
||||
lora_a_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
lora_scaling_tensor,
|
||||
)
|
||||
|
||||
actual_output = sgemm_lora_a_fwd(
|
||||
inputs,
|
||||
lora_a_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
lora_scaling_tensor,
|
||||
)
|
||||
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
def test_sgemm_lora_b_fwd_expand(self):
|
||||
batch_size = 2
|
||||
output_dim = 1024
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
|
||||
lora_ranks = random.sample(
|
||||
possible_lora_ranks,
|
||||
counts=[num_loras] * len(possible_lora_ranks),
|
||||
k=num_loras,
|
||||
)
|
||||
|
||||
max_lora_rank = max(lora_ranks)
|
||||
|
||||
seq_len_tensor = torch.randint(
|
||||
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
|
||||
)
|
||||
|
||||
seq_len = sum(seq_len_tensor)
|
||||
|
||||
inputs = torch.randn(seq_len, max_lora_rank, dtype=dtype)
|
||||
lora_b_weights = torch.randn(num_loras, output_dim, max_lora_rank, dtype=dtype)
|
||||
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
|
||||
lora_indices_tensor = torch.randint(
|
||||
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
|
||||
)
|
||||
slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cpu")
|
||||
|
||||
expect_output = reference_sgmv_expand(
|
||||
inputs,
|
||||
lora_b_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets,
|
||||
)
|
||||
|
||||
actual_output = sgemm_lora_b_fwd(
|
||||
inputs,
|
||||
lora_b_weights,
|
||||
lora_indices_tensor,
|
||||
seq_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets,
|
||||
)
|
||||
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
52
third_party/sglang/test/manual/lora/test_lora_spec_decoding.py
vendored
Normal file
52
third_party/sglang/test/manual/lora/test_lora_spec_decoding.py
vendored
Normal file
@@ -0,0 +1,52 @@
|
||||
# Copyright 2023-2025 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
import multiprocessing as mp
|
||||
import unittest
|
||||
|
||||
from sglang.test.lora_utils import (
|
||||
CI_MULTI_LORA_MODELS,
|
||||
LORA_MODELS_QWEN3,
|
||||
run_lora_multiple_batch_on_model_cases,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestLoRASpecDecoding(CustomTestCase):
|
||||
def test_qwen(self):
|
||||
run_lora_multiple_batch_on_model_cases(
|
||||
LORA_MODELS_QWEN3,
|
||||
attention_backend="triton",
|
||||
use_spec_decoding=True,
|
||||
disable_cuda_graph=True,
|
||||
enable_deterministic_inference=True,
|
||||
)
|
||||
|
||||
def test_llama(self):
|
||||
run_lora_multiple_batch_on_model_cases(
|
||||
CI_MULTI_LORA_MODELS,
|
||||
attention_backend="triton",
|
||||
use_spec_decoding=True,
|
||||
disable_cuda_graph=True,
|
||||
enable_deterministic_inference=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
unittest.main(warnings="ignore")
|
||||
246
third_party/sglang/test/manual/lora/test_torch_backend.py
vendored
Normal file
246
third_party/sglang/test/manual/lora/test_torch_backend.py
vendored
Normal file
@@ -0,0 +1,246 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.lora.backend.torch_backend import TorchNativeLoRABackend
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestTorchNativeLoRABackend(CustomTestCase):
|
||||
|
||||
device = "cpu"
|
||||
|
||||
# set duplicate weights to test merging during prepare_lora_batch
|
||||
weight_indices = [0, 0, 1]
|
||||
lora_ranks = [1, 1]
|
||||
scalings = [1.0, 0.5]
|
||||
seq_lens = [1, 1, 1]
|
||||
use_cuda_graph = False
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
batch_size=3,
|
||||
input_ids=torch.tensor([[1], [2], [3]], dtype=torch.int32),
|
||||
req_pool_indices=None,
|
||||
seq_lens=None,
|
||||
out_cache_loc=None,
|
||||
seq_lens_sum=3,
|
||||
extend_seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
|
||||
extend_seq_lens_cpu=seq_lens,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.backend = TorchNativeLoRABackend(max_loras_per_batch=2, device=cls.device)
|
||||
cls.backend.prepare_lora_batch(
|
||||
forward_batch=cls.forward_batch,
|
||||
weight_indices=cls.weight_indices,
|
||||
lora_ranks=cls.lora_ranks,
|
||||
scalings=cls.scalings,
|
||||
use_cuda_graph=cls.use_cuda_graph,
|
||||
)
|
||||
|
||||
def test_run_lora_a_sgemm(self):
|
||||
batch_size = 3
|
||||
input_dim = 4
|
||||
output_dim = 6
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
x = torch.randn(batch_size, input_dim, dtype=dtype)
|
||||
weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
|
||||
|
||||
weight_indices_tensor = torch.tensor(
|
||||
self.weight_indices, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
seg_len_tensor = torch.tensor(
|
||||
self.seq_lens, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
lora_ranks_tensor = torch.tensor(
|
||||
self.lora_ranks, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
scalings_tensor = torch.tensor(
|
||||
self.scalings, dtype=torch.float, device=self.device
|
||||
)
|
||||
|
||||
expect_output = reference_sgmv_shrink(
|
||||
x,
|
||||
weights,
|
||||
weight_indices_tensor,
|
||||
seg_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
scalings_tensor,
|
||||
)
|
||||
|
||||
actual_output = self.backend.run_lora_a_sgemm(x, weights)
|
||||
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
def test_run_lora_b_sgemm(self):
|
||||
batch_size = 3
|
||||
input_dim = 6
|
||||
output_dim = 4
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
x = torch.randn(batch_size, input_dim, dtype=dtype)
|
||||
weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
|
||||
_, weight_out_dim, _ = weights.shape
|
||||
|
||||
weight_indices_tensor = torch.tensor(
|
||||
self.weight_indices, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
seg_len_tensor = torch.tensor(
|
||||
self.seq_lens, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
lora_ranks_tensor = torch.tensor(
|
||||
self.lora_ranks, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
expect_output = reference_sgmv_expand(
|
||||
x,
|
||||
weights,
|
||||
weight_indices_tensor,
|
||||
seg_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets=torch.tensor(
|
||||
[0, weight_out_dim], dtype=torch.int32, device="cpu"
|
||||
),
|
||||
)
|
||||
|
||||
actual_output = self.backend.run_lora_b_sgemm(x, weights)
|
||||
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
def test_run_qkv_lora(self):
|
||||
batch_size = 3
|
||||
num_loras = 3
|
||||
input_dim = 6
|
||||
output_offset = [0, 3, 6, 9]
|
||||
output_dim = output_offset[-1]
|
||||
num_slices = len(output_offset) - 1 # 3 slices for Q, K, V
|
||||
max_lora_rank = max(self.lora_ranks)
|
||||
dtype = torch.float32
|
||||
|
||||
x = torch.randn(batch_size, input_dim, dtype=dtype)
|
||||
output_offset_cpu = torch.tensor(output_offset, dtype=torch.int32)
|
||||
qkv_lora_a = torch.randn(
|
||||
num_loras, max_lora_rank * num_slices, input_dim, dtype=dtype
|
||||
)
|
||||
qkv_lora_b = torch.randn(
|
||||
num_loras, output_dim, max_lora_rank * num_slices, dtype=dtype
|
||||
)
|
||||
|
||||
weight_indices_tensor = torch.tensor(
|
||||
self.weight_indices, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
seg_len_tensor = torch.tensor(
|
||||
self.seq_lens, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
lora_ranks_tensor = torch.tensor(
|
||||
self.lora_ranks, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
scalings_tensor = torch.tensor(
|
||||
self.scalings, dtype=torch.float, device=self.device
|
||||
)
|
||||
|
||||
expect_lora_a_output = reference_sgmv_shrink(
|
||||
x,
|
||||
qkv_lora_a,
|
||||
weight_indices_tensor,
|
||||
seg_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
scalings_tensor,
|
||||
num_slices,
|
||||
)
|
||||
|
||||
expect_output = reference_sgmv_expand(
|
||||
expect_lora_a_output,
|
||||
qkv_lora_b,
|
||||
weight_indices_tensor,
|
||||
seg_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
output_offset_cpu,
|
||||
)
|
||||
|
||||
actual_output = self.backend.run_qkv_lora(
|
||||
x, qkv_lora_a, qkv_lora_b, None, output_offset_cpu, 0
|
||||
)
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
def test_run_gate_up_lora(self):
|
||||
batch_size = 3
|
||||
input_dim = 6
|
||||
output_dim = 4
|
||||
num_loras = 3
|
||||
dtype = torch.float32
|
||||
|
||||
max_lora_rank = max(self.lora_ranks)
|
||||
|
||||
num_slices = 2
|
||||
|
||||
x = torch.randn(batch_size, input_dim, dtype=dtype)
|
||||
gate_up_lora_a = torch.randn(
|
||||
num_loras, max_lora_rank * num_slices, input_dim, dtype=dtype
|
||||
)
|
||||
gate_up_lora_b = torch.randn(
|
||||
num_loras, output_dim, max_lora_rank * num_slices, dtype=dtype
|
||||
)
|
||||
|
||||
_, weight_out_dim, _ = gate_up_lora_b.shape
|
||||
slice_size = weight_out_dim // num_slices
|
||||
output_offset = torch.tensor(
|
||||
[0, slice_size, weight_out_dim], dtype=torch.int32, device="cpu"
|
||||
)
|
||||
|
||||
weight_indices_tensor = torch.tensor(
|
||||
self.weight_indices, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
seg_len_tensor = torch.tensor(
|
||||
self.seq_lens, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
lora_ranks_tensor = torch.tensor(
|
||||
self.lora_ranks, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
scalings_tensor = torch.tensor(
|
||||
self.scalings, dtype=torch.float, device=self.device
|
||||
)
|
||||
|
||||
expect_lora_a_output = reference_sgmv_shrink(
|
||||
x,
|
||||
gate_up_lora_a,
|
||||
weight_indices_tensor,
|
||||
seg_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
scalings_tensor,
|
||||
num_slices,
|
||||
)
|
||||
|
||||
expect_output = reference_sgmv_expand(
|
||||
expect_lora_a_output,
|
||||
gate_up_lora_b,
|
||||
weight_indices_tensor,
|
||||
seg_len_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets=output_offset,
|
||||
)
|
||||
|
||||
actual_output = self.backend.run_gate_up_lora(x, gate_up_lora_a, gate_up_lora_b)
|
||||
self.assertTrue(torch.allclose(actual_output, expect_output))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user