enhance LoRA tests and fix base model LoRA eviction in Scheduler (#16333)

python/sglang/srt/managers/scheduler.py

@@ -1954,13 +1954,10 @@ class Scheduler(
                     | set([req.lora_id])
                 )
                 if not self.tp_worker.can_run_lora_batch(new_lora_set):
-                    # If this is a LoRA request that would exceed the LoRA slot limit,
+                    # Batch would exceed the LoRA slot limit.
-                    # skip it and continue to try scheduling non-LoRA requests.
+                    # Skip this request and try scheduling it in a future iteration.
-                    # Non-LoRA requests (lora_id=None) share a single reserved slot
+                    # Note: When eviction is needed, the eviction policy prefers to
-                    # and should never cause this check to fail.
+                    # evict LoRA adapters over base model (None) - see mem_pool.py.
-                    if req.lora_id is not None:
-                        # Skip this LoRA request - it would trigger adapter eviction/loading
-                        # which is slow. We'll try to schedule it in a future iteration.
                     continue
 
             running_bs = len(self.running_batch.reqs)

python/sglang/test/lora_utils.py

@@ -25,6 +25,7 @@ class LoRAModelCase:
     decode_tolerance: float = 1e-1
     rouge_l_tolerance: float = 1.0
     max_loras_per_batch: int = 1
+    max_loaded_loras: Optional[int] = None
     skip_long_prompt: bool = False
 
     def __post_init__(self):
@@ -115,6 +116,23 @@ ALL_OTHER_MULTI_LORA_MODELS = [
     ),
 ]
 
+LORA_MODELS_QWEN3 = [
+    LoRAModelCase(
+        base="Qwen/Qwen3-4B",
+        adaptors=[
+            LoRAAdaptor(
+                name="nissenj/Qwen3-4B-lora-v2",
+                prefill_tolerance=3e-1,
+            ),
+            LoRAAdaptor(
+                name="y9760210/Qwen3-4B-lora_model",
+                prefill_tolerance=3e-1,
+            ),
+        ],
+        max_loras_per_batch=2,
+    ),
+]
+
 
 def safe_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
     """Matrix multiplication with mixed precision handling for float16"""
@@ -314,6 +332,7 @@ def run_lora_test_one_by_one(
             adaptor.name for adaptor in model_case.adaptors if adaptor.name is not None
         ],
         max_loras_per_batch=model_case.max_loras_per_batch,
+        max_loaded_loras=model_case.max_loaded_loras,
         lora_backend=backend,
         disable_cuda_graph=disable_cuda_graph,
         disable_radix_cache=disable_radix_cache,
@@ -461,6 +480,7 @@ def run_lora_test_by_batch(
             adaptor.name for adaptor in model_case.adaptors if adaptor.name is not None
         ],
         max_loras_per_batch=model_case.max_loras_per_batch,
+        max_loaded_loras=model_case.max_loaded_loras,
         lora_backend=backend,
         disable_cuda_graph=disable_cuda_graph,
         disable_radix_cache=disable_radix_cache,
@@ -651,6 +671,7 @@ def run_lora_multiple_batch_on_model_cases(
                 model_type="generation",
                 lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
                 max_loras_per_batch=len(lora_adapter_paths) + 1,
+                max_loaded_loras=model_case.max_loaded_loras,
                 sleep_on_idle=True,  # Eliminate non-determinism by forcing all requests to be processed in one batch.
                 attention_backend=attention_backend,
                 enable_deterministic_inference=enable_deterministic_inference,
@@ -702,3 +723,112 @@ def run_lora_multiple_batch_on_model_cases(
                             )
 
                     print(f"--- Batch {i+1} Comparison Passed --- ")
+
+
+def run_lora_batch_splitting_equivalence_test(
+    model_cases: List[LoRAModelCase],
+    attention_backend: str = "torch_native",
+    disable_cuda_graph: bool = True,
+    disable_radix_cache: bool = True,
+):
+    """
+    Test that SRT correctly handles batch splitting with multiple LoRA adapters.
+
+    When the number of distinct adapters (including None for base model) exceeds
+    max_loras_per_batch, SRT internally splits requests into microbatches.
+
+    This test validates:
+    1. SRT can process batches that trigger internal splitting without errors
+    2. Different adapters don't produce all identical outputs (i.e., at least one
+       output differs, indicating adapters are being applied correctly)
+
+    Args:
+        model_cases: List of LoRAModelCase configurations to test
+        attention_backend: Attention backend to use
+        disable_cuda_graph: Whether to disable CUDA graph
+        disable_radix_cache: Whether to disable radix cache
+    """
+    max_loras_per_batch = 2
+
+    def _run_test(model_case: LoRAModelCase, torch_dtype: torch.dtype):
+        lora_adapter_paths = [a.name for a in model_case.adaptors]
+        assert (
+            len(lora_adapter_paths) >= max_loras_per_batch
+        ), f"Need at least {max_loras_per_batch} adapters for this test"
+
+        max_new_tokens = 64
+        base_path = model_case.base
+
+        print(
+            f"\n========== Testing batch splitting on base '{base_path}', "
+            f"dtype={torch_dtype} =========="
+        )
+
+        prompts = [TEST_MULTIPLE_BATCH_PROMPTS[0]] * 3
+        test_cases = [
+            (
+                prompts,
+                [None, lora_adapter_paths[0], lora_adapter_paths[1]],
+            ),
+            (
+                prompts,
+                [lora_adapter_paths[0], None, lora_adapter_paths[1]],
+            ),
+            (
+                prompts,
+                [lora_adapter_paths[0], lora_adapter_paths[1], None],
+            ),
+            (
+                prompts,
+                [None, lora_adapter_paths[1], None],
+            ),
+            (
+                prompts,
+                [lora_adapter_paths[0], lora_adapter_paths[1], lora_adapter_paths[0]],
+            ),
+            (
+                prompts,
+                [None, None, None],
+            ),
+        ]
+
+        ensure_reproducibility()
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            lora_paths=lora_adapter_paths,
+            max_loras_per_batch=max_loras_per_batch,
+            max_loaded_loras=model_case.max_loaded_loras,
+            sleep_on_idle=True,
+            attention_backend=attention_backend,
+            disable_cuda_graph=disable_cuda_graph,
+            disable_radix_cache=disable_radix_cache,
+        ) as srt_runner:
+            for batch_idx, (batch_prompts, lora_paths) in enumerate(test_cases):
+                print(f"\n--- Batch {batch_idx + 1} ---")
+                print(f"  Adapters: {lora_paths}")
+
+                srt_outputs = srt_runner.batch_forward(
+                    batch_prompts,
+                    max_new_tokens=max_new_tokens,
+                    lora_paths=lora_paths,
+                )
+
+                # If different adapters are used in this batch, verify that not every
+                # output is identical (at least one should differ)
+                unique_adapters = set(lora_paths)
+                if len(unique_adapters) >= 2:
+                    all_outputs = [s.strip() for s in srt_outputs.output_strs]
+                    all_identical = all(out == all_outputs[0] for out in all_outputs)
+                    assert not all_identical, (
+                        f"Every output was identical despite using different adapters for "
+                        f"base '{base_path}', batch {batch_idx + 1}: "
+                        f"adapters={lora_paths}. Expected at least one output to differ."
+                    )
+
+                print(f"--- Batch {batch_idx + 1} passed ---")
+
+    for model_case in model_cases:
+        for torch_dtype in TORCH_DTYPES:
+            _run_test(model_case, torch_dtype)

test/manual/lora/test_lora_spec_decoding.py

@@ -17,29 +17,11 @@ import unittest
 
 from sglang.test.lora_utils import (
     CI_MULTI_LORA_MODELS,
-    LoRAAdaptor,
+    LORA_MODELS_QWEN3,
-    LoRAModelCase,
     run_lora_multiple_batch_on_model_cases,
 )
 from sglang.test.test_utils import CustomTestCase
 
-LORA_MODELS_QWEN3 = [
-    LoRAModelCase(
-        base="Qwen/Qwen3-4B",
-        adaptors=[
-            LoRAAdaptor(
-                name="nissenj/Qwen3-4B-lora-v2",
-                prefill_tolerance=3e-1,
-            ),
-            LoRAAdaptor(
-                name="y9760210/Qwen3-4B-lora_model",
-                prefill_tolerance=3e-1,
-            ),
-        ],
-        max_loras_per_batch=2,
-    ),
-]
-
 
 class TestLoRASpecDecoding(CustomTestCase):
     def test_qwen(self):

test/registered/lora/test_lora.py

@@ -1,54 +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.
-# ==============================================================================
-
-import multiprocessing as mp
-import os
-import unittest
-
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.lora_utils import (
-    ALL_OTHER_MULTI_LORA_MODELS,
-    CI_MULTI_LORA_MODELS,
-    run_lora_multiple_batch_on_model_cases,
-)
-from sglang.test.test_utils import CustomTestCase, is_in_ci
-
-register_cuda_ci(est_time=82, suite="stage-b-test-small-1-gpu")
-register_amd_ci(est_time=82, suite="stage-b-test-small-1-gpu-amd")
-
-
-class TestLoRA(CustomTestCase):
-    def test_ci_lora_models(self):
-        run_lora_multiple_batch_on_model_cases(CI_MULTI_LORA_MODELS)
-
-    def test_all_lora_models(self):
-        if is_in_ci():
-            return
-
-        filtered_models = []
-        for model_case in ALL_OTHER_MULTI_LORA_MODELS:
-            if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
-                continue
-            filtered_models.append(model_case)
-
-        run_lora_multiple_batch_on_model_cases(filtered_models)
-
-
-if __name__ == "__main__":
-    try:
-        mp.set_start_method("spawn")
-    except RuntimeError:
-        pass
-
-    unittest.main(warnings="ignore")

test/registered/lora/test_lora_qwen3.py

@@ -17,8 +17,7 @@ import unittest
 
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.lora_utils import (
-    LoRAAdaptor,
+    LORA_MODELS_QWEN3,
-    LoRAModelCase,
     run_lora_multiple_batch_on_model_cases,
 )
 
@@ -26,23 +25,6 @@ register_cuda_ci(est_time=97, suite="nightly-1-gpu", nightly=True)
 
 from sglang.test.test_utils import CustomTestCase
 
-LORA_MODELS_QWEN3 = [
-    LoRAModelCase(
-        base="Qwen/Qwen3-4B",
-        adaptors=[
-            LoRAAdaptor(
-                name="nissenj/Qwen3-4B-lora-v2",
-                prefill_tolerance=3e-1,
-            ),
-            LoRAAdaptor(
-                name="y9760210/Qwen3-4B-lora_model",
-                prefill_tolerance=3e-1,
-            ),
-        ],
-        max_loras_per_batch=2,
-    ),
-]
-
 
 class TestLoRAQwen3(CustomTestCase):
     def test_ci_lora_models(self):

test/registered/lora/test_multi_lora_backend.py

@@ -20,30 +20,20 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.lora_utils import (
     ALL_OTHER_MULTI_LORA_MODELS,
     CI_MULTI_LORA_MODELS,
+    run_lora_batch_splitting_equivalence_test,
     run_lora_multiple_batch_on_model_cases,
 )
 from sglang.test.test_utils import CustomTestCase, is_in_ci
 
-register_cuda_ci(est_time=60, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu")
-register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu-amd")
+register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd")
-
-# All prompts are used at once in a batch.
-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.
-    ### Question:
-    What do you know about llamas?
-    ### Answer:
-    """,
-]
 
 
 class TestMultiLoRABackend(CustomTestCase):
-    def test_ci_lora_models(self):
+    def test_ci_lora_models_batch_splitting(self):
+        run_lora_batch_splitting_equivalence_test(CI_MULTI_LORA_MODELS)
+
+    def test_ci_lora_models_multi_batch(self):
         run_lora_multiple_batch_on_model_cases(CI_MULTI_LORA_MODELS)
 
     def test_all_lora_models(self):