From 1ce515a53def35ce90f8d40e576050469b4c96c8 Mon Sep 17 00:00:00 2001 From: Yanbin Jiang Date: Tue, 11 Aug 2026 17:37:52 -0700 Subject: [PATCH] Refocus LoRA tests on regression coverage (#34464) --- .../ops/gemm/test_chunked_sgmv_cuda_graph.py | 294 ++++++++++++++++++ .../lora/test_embedding_lora_support.py | 125 ++++++++ .../lora/test_lora_eviction_policy.py | 201 ------------ .../lora/test_lora_openai_compatible.py | 238 ++------------ .../unit/lora/test_eviction_policy.py | 83 +++++ .../unit/lora/test_lm_head_pruning.py | 87 ++++++ .../lora/test_lora_manager_tied_lm_head.py | 91 ++++++ .../unit/managers/test_io_struct.py | 33 ++ 8 files changed, 745 insertions(+), 407 deletions(-) create mode 100644 test/registered/kernels/ops/gemm/test_chunked_sgmv_cuda_graph.py create mode 100644 test/registered/lora/test_embedding_lora_support.py delete mode 100644 test/registered/lora/test_lora_eviction_policy.py create mode 100644 test/registered/unit/lora/test_eviction_policy.py create mode 100644 test/registered/unit/lora/test_lm_head_pruning.py create mode 100644 test/registered/unit/lora/test_lora_manager_tied_lm_head.py diff --git a/test/registered/kernels/ops/gemm/test_chunked_sgmv_cuda_graph.py b/test/registered/kernels/ops/gemm/test_chunked_sgmv_cuda_graph.py new file mode 100644 index 000000000..f9f8f906a --- /dev/null +++ b/test/registered/kernels/ops/gemm/test_chunked_sgmv_cuda_graph.py @@ -0,0 +1,294 @@ +# Copyright 2023-2026 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. +# ============================================================================== + +"""CUDA-graph regressions for #28371: dynamic chunked-SGMV LoRA segments.""" + +import pytest +import torch + +from sglang.kernels.ops.gemm.chunked_sgmv_expand import ( + _chunked_lora_expand_kernel, + chunked_sgmv_lora_expand_forward, +) +from sglang.kernels.ops.gemm.chunked_sgmv_shrink import ( + _chunked_lora_shrink_kernel, + chunked_sgmv_lora_shrink_forward, +) +from sglang.kernels.ops.gemm.kv_b_lora_absorbed import ( + step_a_q_fwd, + step_a_v_fwd, + step_b_q_fwd, + step_b_v_fwd, +) +from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend +from sglang.srt.lora.utils import LoRABatchInfo +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +ATOL = 1e-3 +RTOL = 1e-3 +BS = 8 +NUM_LORAS = 5 +MAX_RANK = 8 + + +def _make_batch_info(): + return LoRABatchInfo( + use_cuda_graph=True, + bs=BS, + num_segments=None, + max_len=16, + seg_lens=None, + seg_indptr=torch.zeros(BS + 1, dtype=torch.int32, device="cuda"), + weight_indices=torch.zeros(BS, dtype=torch.int32, device="cuda"), + lora_ranks=torch.zeros(NUM_LORAS, dtype=torch.int32, device="cuda"), + scalings=torch.ones(NUM_LORAS, dtype=torch.float, device="cuda"), + permutation=torch.arange(BS, dtype=torch.int32, device="cuda"), + ) + + +def _set_segment_state(batch_info, *, active): + if active: + lora_ranks = [MAX_RANK] * NUM_LORAS + weight_indices = [1, 2, 3, 4] + seg_indptr = [0, 2, 4, 6, BS] + else: + lora_ranks = [0] * NUM_LORAS + weight_indices = [0] + seg_indptr = [0, BS] + + num_segments = len(weight_indices) + batch_info.lora_ranks.copy_( + torch.tensor(lora_ranks, dtype=torch.int32, device="cuda") + ) + batch_info.weight_indices.zero_() + batch_info.weight_indices[:num_segments].copy_( + torch.tensor(weight_indices, dtype=torch.int32, device="cuda") + ) + batch_info.seg_indptr.fill_(seg_indptr[-1]) + batch_info.seg_indptr[: num_segments + 1].copy_( + torch.tensor(seg_indptr, dtype=torch.int32, device="cuda") + ) + batch_info.num_segments = num_segments + + +def _capture(call): + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + call() + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = call() + return graph, output + + +def test_shrink_replay_uses_all_current_segments(): + """A one-segment graph must match eager shrink after replaying four adapters.""" + _chunked_lora_shrink_kernel._clear_cache() + batch_info = _make_batch_info() + inputs = torch.randn(BS, 64, dtype=torch.float16, device="cuda") + weights = torch.randn(NUM_LORAS, MAX_RANK, 64, dtype=torch.float16, device="cuda") + + _set_segment_state(batch_info, active=True) + expected = chunked_sgmv_lora_shrink_forward( + inputs, weights, batch_info, num_slices=1 + ).clone() + + _set_segment_state(batch_info, active=False) + graph, captured_output = _capture( + lambda: chunked_sgmv_lora_shrink_forward( + inputs, weights, batch_info, num_slices=1 + ) + ) + + _set_segment_state(batch_info, active=True) + captured_output.zero_() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + captured_output[:, :MAX_RANK], + expected[:, :MAX_RANK], + rtol=RTOL, + atol=ATOL, + ) + + +def test_expand_replay_uses_all_current_segments(): + """A one-segment graph must match eager expand after replaying four adapters.""" + _chunked_lora_expand_kernel._clear_cache() + batch_info = _make_batch_info() + output_dim = 32 + inputs = torch.randn(BS, MAX_RANK, dtype=torch.float16, device="cuda") + weights = torch.randn( + NUM_LORAS, output_dim, MAX_RANK, dtype=torch.float16, device="cuda" + ) + slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cuda") + base_output = torch.randn(BS, output_dim, dtype=torch.float16, device="cuda") + graph_base_output = base_output.clone() + + _set_segment_state(batch_info, active=True) + expected = chunked_sgmv_lora_expand_forward( + inputs, + weights, + batch_info, + slice_offsets, + output_dim, + base_output=base_output.clone(), + ).clone() + + def run_graph(): + graph_base_output.copy_(base_output) + return chunked_sgmv_lora_expand_forward( + inputs, + weights, + batch_info, + slice_offsets, + output_dim, + base_output=graph_base_output, + ) + + _set_segment_state(batch_info, active=False) + graph, captured_output = _capture(run_graph) + + _set_segment_state(batch_info, active=True) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(captured_output, expected, rtol=RTOL, atol=ATOL) + + +def test_prepare_batch_neutralizes_static_tail_segments(): + """A smaller replay batch must not expose stale adapter metadata.""" + + class MockForwardBatch: + def __init__(self, batch_size): + self.batch_size = batch_size + self.forward_mode = ForwardMode.DECODE + + server_args = type("ServerArgs", (), {"max_lora_chunk_size": 16}) + backend = ChunkedSgmvLoRABackend( + max_loras_per_batch=NUM_LORAS, + device=torch.device("cuda"), + server_args=server_args, + ) + backend.init_cuda_graph_batch_info(max_bs_in_cuda_graph=BS, num_tokens_per_req=1) + lora_ranks = [MAX_RANK] * NUM_LORAS + scalings = [1.0] * NUM_LORAS + + backend.prepare_lora_batch( + forward_batch=MockForwardBatch(BS), + weight_indices=[0, 1, 2, 3, 4, 0, 1, 2], + lora_ranks=lora_ranks, + scalings=scalings, + use_cuda_graph=True, + ) + backend.prepare_lora_batch( + forward_batch=MockForwardBatch(2), + weight_indices=[0, 0], + lora_ranks=lora_ranks, + scalings=scalings, + use_cuda_graph=True, + ) + torch.cuda.synchronize() + + assert backend.batch_info.num_segments == 1 + torch.testing.assert_close( + backend.batch_info.weight_indices.cpu(), + torch.tensor([0] * BS, dtype=torch.int32), + ) + torch.testing.assert_close( + backend.batch_info.seg_indptr.cpu(), + torch.tensor([0, 2, 2, 2, 2, 2, 2, 2, 2], dtype=torch.int32), + ) + + +def test_absorbed_kv_b_replay_uses_all_current_segments(): + """A one-segment graph must match eager MLA Q/V updates for four adapters.""" + batch_info = _make_batch_info() + num_heads = 2 + qk_nope_head_dim = 16 + v_head_dim = 16 + kv_lora_rank = 32 + full_k_per_head = qk_nope_head_dim + v_head_dim + + q_nope = torch.randn( + BS, num_heads, qk_nope_head_dim, dtype=torch.float16, device="cuda" + ) + attn_output = torch.randn( + BS, num_heads, kv_lora_rank, dtype=torch.float16, device="cuda" + ) + a_buf = torch.randn( + NUM_LORAS, MAX_RANK, kv_lora_rank, dtype=torch.float16, device="cuda" + ) + b_buf = torch.randn( + NUM_LORAS, + num_heads * full_k_per_head, + MAX_RANK, + dtype=torch.float16, + device="cuda", + ) + base_q = torch.randn( + BS, num_heads, kv_lora_rank, dtype=torch.float16, device="cuda" + ) + base_v = torch.randn(BS, num_heads, v_head_dim, dtype=torch.float16, device="cuda") + graph_base_q = base_q.clone() + graph_base_v = base_v.clone() + + def run_kv_b(base_q_output, base_v_output): + q_lora_a = step_a_q_fwd(q_nope, b_buf, batch_info, full_k_per_head) + q_output = step_b_q_fwd(q_lora_a, a_buf, batch_info, base_q_output) + v_lora_a = step_a_v_fwd(attn_output, a_buf, batch_info) + v_output = step_b_v_fwd( + v_lora_a, + b_buf, + batch_info, + base_v_output, + qk_nope_head_dim, + v_head_dim, + ) + return q_output, v_output + + _set_segment_state(batch_info, active=True) + expected_q, expected_v = run_kv_b(base_q.clone(), base_v.clone()) + expected_q = expected_q.clone() + expected_v = expected_v.clone() + + def run_graph(): + graph_base_q.copy_(base_q) + graph_base_v.copy_(base_v) + return run_kv_b(graph_base_q, graph_base_v) + + _set_segment_state(batch_info, active=False) + graph, (captured_q, captured_v) = _capture(run_graph) + + _set_segment_state(batch_info, active=True) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(captured_q, expected_q, rtol=RTOL, atol=ATOL) + torch.testing.assert_close(captured_v, expected_v, rtol=RTOL, atol=ATOL) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/lora/test_embedding_lora_support.py b/test/registered/lora/test_embedding_lora_support.py new file mode 100644 index 000000000..4a3dde48f --- /dev/null +++ b/test/registered/lora/test_embedding_lora_support.py @@ -0,0 +1,125 @@ +# Copyright 2023-2026 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. +# ============================================================================== + +"""LoRA embedding parity against an independent Hugging Face oracle.""" + +import multiprocessing as mp +import unittest + +import numpy as np +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.runners import SRTRunner +from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase + +register_cuda_ci(est_time=150, stage="nightly", runner_config="1-gpu-large") + +MODEL_PATH = "meta-llama/Llama-2-7b-hf" +LORA_PATH = "yushengsu/sglang_lora_logprob_diff_without_tuning" +SIMILARITY_THRESHOLD = 0.9999 + + +class TestEmbeddingLoRAParity(CustomTestCase): + """Guard the end-to-end embedding request and LoRA execution path.""" + + @staticmethod + def _hf_embeddings(texts): + from peft import PeftModel + from transformers import AutoModelForCausalLM, AutoTokenizer + + base_model = AutoModelForCausalLM.from_pretrained( + MODEL_PATH, + torch_dtype=torch.float16, + trust_remote_code=True, + ).cuda() + model = PeftModel.from_pretrained(base_model, LORA_PATH) + model.eval() + + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + with torch.no_grad(): + inputs = tokenizer( + texts, padding=True, truncation=True, return_tensors="pt" + ).to("cuda") + hidden_states = model.model( + **inputs, output_hidden_states=True + ).hidden_states[-1] + last_token_indices = inputs["attention_mask"].sum(dim=1) - 1 + embeddings = hidden_states[ + torch.arange(hidden_states.shape[0], device="cuda"), + last_token_indices, + ] + embeddings = embeddings / embeddings.norm(dim=1, keepdim=True) + + result = embeddings.cpu().numpy() + del model, base_model + torch.cuda.empty_cache() + return result + + @staticmethod + def _sglang_embeddings(texts): + def extract_embeddings(response): + if not isinstance(response, list): + response = [response] + return np.asarray([item["embedding"] for item in response]) + + with SRTRunner( + MODEL_PATH, + torch_dtype=torch.float16, + model_type="embedding", + lora_paths=[LORA_PATH], + lora_backend="triton", + port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER, + trust_remote_code=True, + mem_fraction_static=0.88, + ) as runner: + base_response = runner.engine.encode(prompt=texts, lora_path=None) + lora_response = runner.engine.encode(prompt=texts, lora_path=LORA_PATH) + + return extract_embeddings(base_response), extract_embeddings(lora_response) + + def test_hf_sglang_embedding_similarity(self): + """Dropping LoRA at any embedding handoff must fail external parity.""" + texts = [ + "Hello world", + "This is a test sentence for embedding comparison", + ] + + base_embeddings, sglang_embeddings = self._sglang_embeddings(texts) + self.assertFalse( + np.allclose(base_embeddings, sglang_embeddings, rtol=1e-4, atol=1e-5), + "The requested adapter had no observable effect on embeddings", + ) + torch.cuda.empty_cache() + hf_embeddings = self._hf_embeddings(texts) + + self.assertEqual(sglang_embeddings.shape, hf_embeddings.shape) + similarities = np.sum(hf_embeddings * sglang_embeddings, axis=1) / ( + np.linalg.norm(hf_embeddings, axis=1) + * np.linalg.norm(sglang_embeddings, axis=1) + ) + np.testing.assert_array_less( + np.full_like(similarities, SIMILARITY_THRESHOLD), similarities + ) + + +if __name__ == "__main__": + try: + mp.set_start_method("spawn") + except RuntimeError: + pass + unittest.main() diff --git a/test/registered/lora/test_lora_eviction_policy.py b/test/registered/lora/test_lora_eviction_policy.py deleted file mode 100644 index 0aedcdf3b..000000000 --- a/test/registered/lora/test_lora_eviction_policy.py +++ /dev/null @@ -1,201 +0,0 @@ -# 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. -# ============================================================================== - -""" -Unit tests for LoRA eviction policies. -Tests LRU and FIFO eviction behavior. -""" - -import unittest - -from sglang.srt.lora.eviction_policy import get_eviction_policy -from sglang.test.ci.ci_register import ( - register_amd_ci, - register_cpu_ci, - register_cuda_ci, - register_xpu_ci, -) - -register_cuda_ci(est_time=10, stage="nightly", runner_config="1-gpu-large") -register_amd_ci(est_time=200, suite="nightly-amd-1-gpu", nightly=True) -register_cpu_ci(est_time=6, suite="base-c-test-cpu") -register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu") - - -class TestLoRAEvictionPolicy(unittest.TestCase): - """Unit tests for LoRA eviction policies.""" - - def _test_eviction_policy( - self, policy_name, access_sequence, candidates, expected_victim - ): - """ - Helper to test eviction policy with given access pattern. - - Args: - policy_name: Name of eviction policy ("lru" or "fifo") - access_sequence: List of adapter IDs in access order - candidates: Set of adapter IDs that can be evicted - expected_victim: Expected adapter ID to be evicted - """ - policy = get_eviction_policy(policy_name) - - # Simulate access pattern - for adapter_id in access_sequence: - policy.mark_used(adapter_id) - - # Select victim from candidates - victim = policy.select_victim(candidates) - self.assertEqual( - victim, - expected_victim, - f"{policy_name.upper()}: Expected {expected_victim}, got {victim}", - ) - - def test_lru_basic(self): - """Test LRU selects least recently used adapter.""" - self._test_eviction_policy( - "lru", - access_sequence=["lora1", "lora2", "lora3", "lora4"], - candidates={"lora1", "lora2", "lora3", "lora4"}, - expected_victim="lora1", - ) - - def test_lru_with_reuse(self): - """Test LRU updates order on reuse.""" - self._test_eviction_policy( - "lru", - access_sequence=["lora1", "lora2", "lora3", "lora4", "lora1"], - candidates={"lora1", "lora2", "lora3", "lora4"}, - expected_victim="lora2", - ) - - def test_lru_multiple_reuse(self): - """Test LRU with multiple reuses.""" - self._test_eviction_policy( - "lru", - access_sequence=["lora1", "lora2", "lora3", "lora1", "lora2"], - candidates={"lora1", "lora2", "lora3"}, - expected_victim="lora3", - ) - - def test_lru_with_subset_candidates(self): - """Test LRU with subset of candidates.""" - self._test_eviction_policy( - "lru", - access_sequence=["lora1", "lora2", "lora3", "lora4"], - candidates={"lora2", "lora3", "lora4"}, - expected_victim="lora2", - ) - - def test_lru_base_model_evicted_last(self): - """Test LRU evicts LoRA adapters before base model (None).""" - self._test_eviction_policy( - "lru", - access_sequence=["lora1", "lora2", "lora3"], - candidates={None, "lora1", "lora2", "lora3"}, - expected_victim="lora1", - ) - - def test_fifo_basic(self): - """Test FIFO selects first inserted adapter.""" - self._test_eviction_policy( - "fifo", - access_sequence=["lora1", "lora2", "lora3", "lora4"], - candidates={"lora1", "lora2", "lora3", "lora4"}, - expected_victim="lora1", - ) - - def test_fifo_ignores_reuse(self): - """Test FIFO ignores reuse.""" - self._test_eviction_policy( - "fifo", - access_sequence=[ - "lora1", - "lora2", - "lora3", - "lora4", - "lora4", - "lora3", - "lora2", - "lora1", - ], - candidates={"lora1", "lora2", "lora3", "lora4"}, - expected_victim="lora1", - ) - - def test_fifo_with_subset_candidates(self): - """Test FIFO with subset of candidates.""" - self._test_eviction_policy( - "fifo", - access_sequence=["lora1", "lora2", "lora3", "lora4"], - candidates={"lora2", "lora3", "lora4"}, - expected_victim="lora2", - ) - - def test_fifo_base_model_evicted_last(self): - """Test FIFO evicts LoRA adapters before base model (None).""" - self._test_eviction_policy( - "fifo", - access_sequence=["lora1", "lora2", "lora3"], - candidates={None, "lora1", "lora2", "lora3"}, - expected_victim="lora1", - ) - - def test_policy_remove(self): - """Test that remove() correctly updates internal state.""" - lru = get_eviction_policy("lru") - lru.mark_used("lora1") - lru.mark_used("lora2") - lru.mark_used("lora3") - - # Remove lora1, so lora2 becomes LRU - lru.remove("lora1") - victim = lru.select_victim({"lora1", "lora2", "lora3"}) - self.assertEqual(victim, "lora2") - - def test_eviction_policy_factory(self): - """Test eviction policy factory function.""" - # Test valid policies - lru = get_eviction_policy("lru") - fifo = get_eviction_policy("fifo") - - self.assertIsNotNone(lru) - self.assertIsNotNone(fifo) - - # Test invalid policy - with self.assertRaises(ValueError): - get_eviction_policy("invalid_policy") - - def test_lru_vs_fifo_behavior(self): - """Test that LRU and FIFO behave differently.""" - access_sequence = ["lora1", "lora2", "lora3", "lora1"] - candidates = {"lora1", "lora2", "lora3"} - - lru = get_eviction_policy("lru") - for adapter_id in access_sequence: - lru.mark_used(adapter_id) - lru_victim = lru.select_victim(candidates) - - fifo = get_eviction_policy("fifo") - for adapter_id in access_sequence: - fifo.mark_used(adapter_id) - fifo_victim = fifo.select_victim(candidates) - - self.assertNotEqual(lru_victim, fifo_victim) - self.assertEqual(lru_victim, "lora2") - self.assertEqual(fifo_victim, "lora1") - - -if __name__ == "__main__": - unittest.main(verbosity=2) diff --git a/test/registered/lora/test_lora_openai_compatible.py b/test/registered/lora/test_lora_openai_compatible.py index 79d1250a6..24b163d33 100644 --- a/test/registered/lora/test_lora_openai_compatible.py +++ b/test/registered/lora/test_lora_openai_compatible.py @@ -1,24 +1,11 @@ -""" -End-to-end tests for OpenAI-compatible LoRA adapter usage. - -Tests the model:adapter syntax and backward compatibility with explicit lora_path. - -Usage: - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_model_adapter_syntax - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_explicit_lora_path - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_priority_model_over_explicit - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_base_model_no_adapter - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_completions_api_with_adapter - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_streaming_with_adapter - python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRADisabledError.test_lora_disabled_error -""" +"""End-to-end negative-branch contracts for OpenAI-compatible LoRA routing.""" import unittest import openai from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( DEFAULT_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -28,34 +15,27 @@ from sglang.test.test_utils import ( ) register_cuda_ci(est_time=180, stage="nightly", runner_config="1-gpu-large") -register_amd_ci(est_time=150, suite="nightly-amd-1-gpu", nightly=True) -def get_real_lora_adapter() -> str: - """Use a real LoRA adapter from Hugging Face.""" - return "codelion/Llama-3.2-1B-Instruct-tool-calling-lora" - - -def setup_class(cls, enable_lora=True): - """Setup test class with LoRA-enabled server.""" +def setup_class(cls, *, enable_lora): + """Start the shared server for one routing contract.""" cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST - # Use real LoRA adapter - cls.lora_adapter_path = get_real_lora_adapter() - other_args = [ "--max-running-requests", "10", - "--disable-radix-cache", # Disable cache for cleaner tests + "--disable-radix-cache", ] if enable_lora: other_args.extend( [ "--enable-lora", - "--lora-paths", - f"tool_calling={cls.lora_adapter_path}", + "--max-lora-rank", + "8", + "--lora-target-modules", + "q_proj", ] ) @@ -69,7 +49,7 @@ def setup_class(cls, enable_lora=True): class TestLoRAOpenAICompatible(CustomTestCase): - """Test OpenAI-compatible LoRA adapter usage.""" + """Verify that ``model:adapter`` reaches the LoRA registry.""" @classmethod def setUpClass(cls): @@ -77,205 +57,51 @@ class TestLoRAOpenAICompatible(CustomTestCase): @classmethod def tearDownClass(cls): - kill_process_tree(cls.process.pid) + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) - def test_model_adapter_syntax(self): - """Test the new model:adapter syntax works correctly.""" - response = self.client.chat.completions.create( - # ← New OpenAI-compatible syntax - model=f"{self.model}:tool_calling", - messages=[{"role": "user", "content": "What tools do you have available?"}], - max_tokens=50, - temperature=0, - ) + def test_unknown_model_adapter_is_rejected(self): + """An adapter suffix must not be ignored and routed to the base model.""" + with self.assertRaises(openai.APIError) as context: + self.client.chat.completions.create( + model=f"{self.model}:nonexistent", + messages=[{"role": "user", "content": "Hello!"}], + max_tokens=1, + ) - self.assertIsNotNone(response.choices[0].message.content) - self.assertGreater(len(response.choices[0].message.content), 0) - print(f"Model adapter syntax response: {response.choices[0].message.content}") - - def test_explicit_lora_path(self): - """Test backward compatibility with explicit lora_path via extra_body.""" - response = self.client.chat.completions.create( - model=self.model, - messages=[{"role": "user", "content": "What tools do you have available?"}], - # ← Legacy explicit method - extra_body={"lora_path": "tool_calling"}, - max_tokens=50, - temperature=0, - ) - - self.assertIsNotNone(response.choices[0].message.content) - self.assertGreater(len(response.choices[0].message.content), 0) - print(f"Explicit lora_path response: {response.choices[0].message.content}") - - def test_priority_model_over_explicit(self): - """Test that model:adapter syntax takes precedence over explicit lora_path.""" - # This test verifies the priority logic in _resolve_lora_path - response = self.client.chat.completions.create( - # ← Model specifies tool_calling adapter - model=f"{self.model}:tool_calling", - messages=[{"role": "user", "content": "What tools do you have available?"}], - # ← Both specify same adapter - extra_body={"lora_path": "tool_calling"}, - max_tokens=50, - temperature=0, - ) - - # Should use tool_calling adapter (model parameter takes precedence) - self.assertIsNotNone(response.choices[0].message.content) - self.assertGreater(len(response.choices[0].message.content), 0) - print(f"Priority test response: {response.choices[0].message.content}") - - def test_base_model_no_adapter(self): - """Test using base model without any adapter.""" - response = self.client.chat.completions.create( - model=self.model, # ← No adapter specified - messages=[{"role": "user", "content": "Hello, how are you?"}], - max_tokens=30, - temperature=0, - ) - - self.assertIsNotNone(response.choices[0].message.content) - self.assertGreater(len(response.choices[0].message.content), 0) - print(f"Base model response: {response.choices[0].message.content}") - - def test_completions_api_with_adapter(self): - """Test completions API with LoRA adapter.""" - response = self.client.completions.create( - model=f"{self.model}:tool_calling", # ← Using model:adapter syntax - prompt="What tools do you have available?", - max_tokens=50, - temperature=0, - ) - - self.assertIsNotNone(response.choices[0].text) - self.assertGreater(len(response.choices[0].text), 0) - print(f"Completions API response: {response.choices[0].text}") - - def test_streaming_with_adapter(self): - """Test streaming with LoRA adapter.""" - stream = self.client.chat.completions.create( - model=f"{self.model}:tool_calling", - messages=[{"role": "user", "content": "What tools do you have available?"}], - max_tokens=50, - temperature=0, - stream=True, - ) - - collected_content = "" - for chunk in stream: - if chunk.choices[0].delta.content: - collected_content += chunk.choices[0].delta.content - - self.assertGreater(len(collected_content), 0) - print(f"Streaming response: {collected_content}") - - def test_multiple_adapters(self): - """Test using different adapters in sequence.""" - # Test tool_calling adapter - tool_response = self.client.chat.completions.create( - model=f"{self.model}:tool_calling", - messages=[{"role": "user", "content": "What tools do you have available?"}], - max_tokens=30, - temperature=0, - ) - - # Test base model without adapter - base_response = self.client.chat.completions.create( - model=self.model, - messages=[{"role": "user", "content": "Hello, how are you?"}], - max_tokens=30, - temperature=0, - ) - - self.assertIsNotNone(tool_response.choices[0].message.content) - self.assertIsNotNone(base_response.choices[0].message.content) - print( - f"Tool calling adapter response: {tool_response.choices[0].message.content}" - ) - print(f"Base model response: {base_response.choices[0].message.content}") + error_message = str(context.exception) + self.assertIn("never been loaded", error_message) + self.assertIn("nonexistent", error_message) class TestLoRADisabledError(CustomTestCase): - """Test error handling when LoRA is disabled.""" + """Verify the disabled-LoRA request contract.""" @classmethod def setUpClass(cls): - setup_class(cls, enable_lora=False) # ← LoRA disabled + setup_class(cls, enable_lora=False) @classmethod def tearDownClass(cls): - kill_process_tree(cls.process.pid) + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) def test_lora_disabled_error(self): - """Test that using LoRA adapter when LoRA is disabled raises appropriate error.""" + """A requested adapter must fail clearly when LoRA is disabled.""" with self.assertRaises(openai.APIError) as context: self.client.chat.completions.create( - model=f"{self.model}:tool_calling", # ← Trying to use adapter + model=f"{self.model}:tool_calling", messages=[ {"role": "user", "content": "What tools do you have available?"} ], - max_tokens=50, + max_tokens=1, ) - # Verify the error message contains helpful guidance error_message = str(context.exception) self.assertIn("LoRA", error_message) self.assertIn("not enabled", error_message) - print(f"Expected error message: {error_message}") - - -class TestLoRAEdgeCases(CustomTestCase): - """Test edge cases for LoRA adapter usage.""" - - @classmethod - def setUpClass(cls): - setup_class(cls, enable_lora=True) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_model_with_colon_no_adapter(self): - """Test model parameter ending with colon (empty adapter).""" - response = self.client.chat.completions.create( - model=f"{self.model}:", # ← Model ends with colon - messages=[{"role": "user", "content": "Hello!"}], - max_tokens=30, - temperature=0, - ) - - # Should work as base model (no adapter) - self.assertIsNotNone(response.choices[0].message.content) - print(f"Model with colon response: {response.choices[0].message.content}") - - def test_explicit_lora_path_none(self): - """Test explicit lora_path set to None.""" - response = self.client.chat.completions.create( - model=self.model, - messages=[{"role": "user", "content": "Hello!"}], - extra_body={"lora_path": None}, # ← Explicitly None - max_tokens=30, - temperature=0, - ) - - # Should work as base model - self.assertIsNotNone(response.choices[0].message.content) - print( - f"Explicit None lora_path response: {response.choices[0].message.content}" - ) - - def test_invalid_adapter_name(self): - """Test using non-existent adapter name.""" - with self.assertRaises(openai.APIError) as context: - self.client.chat.completions.create( - model=f"{self.model}:nonexistent", # ← Non-existent adapter - messages=[{"role": "user", "content": "Hello!"}], - max_tokens=30, - ) - - error_message = str(context.exception) - print(f"Invalid adapter error: {error_message}") + self.assertIn("tool_calling", error_message) + self.assertIn("--enable-lora", error_message) if __name__ == "__main__": diff --git a/test/registered/unit/lora/test_eviction_policy.py b/test/registered/unit/lora/test_eviction_policy.py new file mode 100644 index 000000000..ef64a54a7 --- /dev/null +++ b/test/registered/unit/lora/test_eviction_policy.py @@ -0,0 +1,83 @@ +# 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. +# ============================================================================== + +"""Focused contract tests for LoRA adapter eviction policies.""" + +import unittest + +from sglang.srt.lora.eviction_policy import get_eviction_policy +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestLoRAEvictionPolicy(CustomTestCase): + """Protect the behavioral differences that callers rely on.""" + + @staticmethod + def _make_policy(policy_name, access_sequence): + policy = get_eviction_policy(policy_name) + for adapter_id in access_sequence: + policy.mark_used(adapter_id) + return policy + + def test_reuse_changes_lru_but_not_fifo(self): + """A repeated access refreshes LRU recency, not FIFO insertion order.""" + access_sequence = ["lora1", "lora2", "lora3", "lora1"] + candidates = {"lora1", "lora2", "lora3"} + + lru = self._make_policy("lru", access_sequence) + fifo = self._make_policy("fifo", access_sequence) + + self.assertEqual(lru.select_victim(candidates), "lora2") + self.assertEqual(fifo.select_victim(candidates), "lora1") + + def test_selection_skips_older_noncandidates(self): + """The oldest ineligible adapter must not displace an eligible one.""" + for policy_name in ("lru", "fifo"): + with self.subTest(policy=policy_name): + policy = self._make_policy(policy_name, ["lora1", "lora2", "lora3"]) + self.assertEqual(policy.select_victim({"lora2", "lora3"}), "lora2") + + def test_base_model_is_evicted_only_as_last_resort(self): + """Regression for #14795: keep the base slot while an adapter can move.""" + for policy_name in ("lru", "fifo"): + with self.subTest(policy=policy_name): + policy = self._make_policy(policy_name, ["lora1", "lora2"]) + self.assertEqual( + policy.select_victim({None, "lora1", "lora2"}), "lora1" + ) + self.assertIsNone(policy.select_victim({None})) + + def test_remove_excludes_adapter_from_future_selection(self): + """Unloaded adapters must not remain eligible through stale policy state.""" + for policy_name in ("lru", "fifo"): + with self.subTest(policy=policy_name): + policy = self._make_policy(policy_name, ["lora1", "lora2", "lora3"]) + policy.remove("lora1") + self.assertEqual( + policy.select_victim({"lora1", "lora2", "lora3"}), "lora2" + ) + + def test_unknown_policy_is_rejected(self): + """A configuration typo must not silently select a fallback policy.""" + with self.assertRaisesRegex( + ValueError, "Unknown eviction policy: invalid_policy" + ): + get_eviction_policy("invalid_policy") + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/lora/test_lm_head_pruning.py b/test/registered/unit/lora/test_lm_head_pruning.py new file mode 100644 index 000000000..1827df561 --- /dev/null +++ b/test/registered/unit/lora/test_lm_head_pruning.py @@ -0,0 +1,87 @@ +# Copyright 2023-2026 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. +# ============================================================================== + +"""Keep LoRA lm_head segmentation aligned with logits-state pruning.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor +from sglang.srt.lora.utils import get_lm_head_pruned_lens +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestLMHeadPruning(CustomTestCase): + def test_lora_segments_match_pruned_rows_per_request(self): + """Equal total rows must not hide routing one request through another LoRA.""" + cases = ( + { + "name": "without_logprobs", + "extend_seq_lens": [4, 5, 6], + "return_logprob": False, + "logprob_start_lens": None, + }, + { + "name": "with_logprobs", + "extend_seq_lens": [4, 5, 6], + "return_logprob": True, + "logprob_start_lens": [0, 5, 3], + }, + ) + + for case in cases: + with self.subTest(case=case["name"]): + extend_seq_lens = case["extend_seq_lens"] + sequence_ids = torch.repeat_interleave( + torch.arange(len(extend_seq_lens)), + torch.tensor(extend_seq_lens), + ) + hidden_states = sequence_ids[:, None].expand(-1, 4).float() + metadata = LogitsMetadata( + forward_mode=ForwardMode.EXTEND, + extend_return_logprob=case["return_logprob"], + extend_seq_lens=torch.tensor(extend_seq_lens), + extend_seq_lens_cpu=extend_seq_lens, + extend_logprob_start_lens_cpu=case["logprob_start_lens"], + ) + pruned_states = LogitsProcessor._get_pruned_states( + None, hidden_states, None, None, metadata + )[0] + actual_lens = torch.bincount( + pruned_states[:, 0].long(), minlength=len(extend_seq_lens) + ).tolist() + + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + batch_size=len(extend_seq_lens), + return_logprob=case["return_logprob"], + extend_seq_lens_cpu=extend_seq_lens, + extend_logprob_start_lens_cpu=case["logprob_start_lens"], + ) + self.assertEqual( + get_lm_head_pruned_lens(forward_batch), + actual_lens, + "LoRA lm_head segment lengths must match the exact per-request " + "rows forwarded by logits pruning", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/lora/test_lora_manager_tied_lm_head.py b/test/registered/unit/lora/test_lora_manager_tied_lm_head.py new file mode 100644 index 000000000..94fdfec2c --- /dev/null +++ b/test/registered/unit/lora/test_lora_manager_tied_lm_head.py @@ -0,0 +1,91 @@ +# Copyright 2023-2026 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. +# ============================================================================== + +"""Regression for #18634: LoRA wrapping of an object-shared lm_head.""" + +import sys +import types +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.lora import lora_manager as lora_manager_module +from sglang.srt.lora.lora_manager import LoRAManager +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _TiedEmbedding(torch.nn.Module): + def __init__(self): + super().__init__() + self.org_vocab_size = 8 + self.embedding_dim = 4 + self.weight = torch.nn.Parameter(torch.randn(8, 4)) + + +class _ParallelLMHead(torch.nn.Module): + def __init__( + self, + num_embeddings, + embedding_dim, + params_dtype, + org_num_embeddings, + ): + super().__init__() + self.weight = torch.nn.Parameter( + torch.empty(num_embeddings, embedding_dim, dtype=params_dtype) + ) + + +class TestTiedLMHeadLoRA(CustomTestCase): + def test_tied_head_gets_independent_wrapper_with_shared_base_weight(self): + """An lm_head-only adapter must survive tied input/output embeddings.""" + model = torch.nn.Module() + tied_embedding = _TiedEmbedding() + model.embed_tokens = tied_embedding + model.lm_head = tied_embedding + + manager = LoRAManager.__new__(LoRAManager) + manager.base_model = model + manager.base_hf_config = SimpleNamespace(num_hidden_layers=0) + manager.target_modules = {"lm_head"} + wrapped_lm_head = object() + + inkling_module = types.ModuleType("sglang.srt.models.inkling_common.dense_mlp") + inkling_module.InklingBatchDenseMLP = type("InklingBatchDenseMLP", (), {}) + + with ( + patch.object(lora_manager_module, "ParallelLMHead", _ParallelLMHead), + patch.object( + manager, "set_lora_module", return_value=wrapped_lm_head + ) as set_lora_module, + patch.dict( + sys.modules, + {"sglang.srt.models.inkling_common.dense_mlp": inkling_module}, + ), + ): + manager.init_lora_modules() + + self.assertIsNot(model.lm_head, model.embed_tokens) + self.assertIs(model.lm_head.weight, model.embed_tokens.weight) + self.assertIs(manager.lm_head_module, wrapped_lm_head) + set_lora_module.assert_called_once_with("lm_head", model.lm_head) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index 46427c5f8..755cbf5ae 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -740,6 +740,39 @@ class TestEmbeddingReqInputGetItem(CustomTestCase): [cross_encoder_req[0].priority, cross_encoder_req[1].priority], [3, 3] ) + def test_lora_identity_survives_batch_split(self): + """Each embedding subrequest must retain its adapter path and resolved ID.""" + cases = ( + (["Hello", "World"], False), + ( + [["query 1", "document 1"], ["query 2", "document 2"]], + True, + ), + ) + for text, is_cross_encoder_request in cases: + with self.subTest(cross_encoder=is_cross_encoder_request): + req = EmbeddingReqInput( + text=text, + is_cross_encoder_request=is_cross_encoder_request, + lora_path="adapter", + lora_id=["id-0", "id-1"], + ) + req.normalize_batch_and_arguments() + + self.assertEqual(req.lora_path, ["adapter", "adapter"]) + self.assertEqual( + [(req[i].lora_path, req[i].lora_id) for i in range(2)], + [("adapter", "id-0"), ("adapter", "id-1")], + ) + + def test_lora_path_count_must_match_embedding_batch(self): + """A partial adapter list must not silently route remaining items to base.""" + req = EmbeddingReqInput( + text=["first", "second"], lora_path=["only-one-adapter"] + ) + with self.assertRaisesRegex(ValueError, "must match batch size"): + req.normalize_batch_and_arguments() + if __name__ == "__main__": unittest.main()