[TestFix] rewrite LoRA overlap loading tests (#18047)

python/sglang/test/lora_utils.py

@@ -640,6 +640,7 @@ def run_lora_multiple_batch_on_model_cases(
     disable_cuda_graph: bool = True,
     enable_deterministic_inference: bool = False,
     disable_radix_cache: bool = True,
+    enable_lora_overlap_loading: Optional[bool] = None,
 ):
     for model_case in model_cases:
         for torch_dtype in TORCH_DTYPES:
@@ -673,6 +674,7 @@ def run_lora_multiple_batch_on_model_cases(
                 torch_dtype=torch_dtype,
                 model_type="generation",
                 lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
+                enable_lora_overlap_loading=enable_lora_overlap_loading,
                 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.
@@ -733,6 +735,7 @@ def run_lora_batch_splitting_equivalence_test(
     attention_backend: str = "torch_native",
     disable_cuda_graph: bool = True,
     disable_radix_cache: bool = True,
+    enable_lora_overlap_loading: Optional[bool] = None,
 ):
     """
     Test that SRT correctly handles batch splitting with multiple LoRA adapters.
@@ -801,6 +804,7 @@ def run_lora_batch_splitting_equivalence_test(
             torch_dtype=torch_dtype,
             model_type="generation",
             lora_paths=lora_adapter_paths,
+            enable_lora_overlap_loading=enable_lora_overlap_loading,
             max_loras_per_batch=max_loras_per_batch,
             max_loaded_loras=model_case.max_loaded_loras,
             sleep_on_idle=True,

test/registered/lora/test_lora_overlap_loading.py

@@ -18,97 +18,28 @@ End-to-end tests for the --enable-lora-overlap-loading server argument.
 import multiprocessing as mp
 import unittest
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.lora_utils import (
     CI_MULTI_LORA_MODELS,
-    TEST_MULTIPLE_BATCH_PROMPTS,
+    run_lora_batch_splitting_equivalence_test,
-    TORCH_DTYPES,
+    run_lora_multiple_batch_on_model_cases,
-    LoRAModelCase,
-    ensure_reproducibility,
 )
-from sglang.test.runners import SRTRunner
+from sglang.test.test_utils import CustomTestCase
-from sglang.test.test_utils import CustomTestCase, calculate_rouge_l
 
-register_cuda_ci(
+register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu")
-    est_time=300,
+register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd")
-    suite="stage-b-test-small-1-gpu",
-    disabled="Flaky test - outputs differ between overlap/no-overlap loading modes. See https://github.com/sgl-project/sglang/actions/runs/21320657015/job/61370002606",
-)
 
 
-class TestLoRAPipelineLoading(CustomTestCase):
+class TestLoRAOverlapLoading(CustomTestCase):
-
+    def test_ci_lora_models_batch_splitting(self):
-    def _run_mixed_batch_test(
+        run_lora_batch_splitting_equivalence_test(
-        self,
+            CI_MULTI_LORA_MODELS, enable_lora_overlap_loading=True
-        model_case: LoRAModelCase,
-        torch_dtype,
-    ):
-        base_path = model_case.base
-        adaptor_paths = [a.name for a in model_case.adaptors]
-        print(
-            f"\n========== Testing mixed batch LoRA overlap loading on base '{base_path}' "
-            f"with dtype={torch_dtype} ==========\n"
-        )
-        ensure_reproducibility()
-        max_new_tokens = 32
-
-        prompts = TEST_MULTIPLE_BATCH_PROMPTS[:3]
-        configs = [
-            [None, adaptor_paths[0], adaptor_paths[1]],
-            [adaptor_paths[0], None, adaptor_paths[1]],
-            [adaptor_paths[0], adaptor_paths[1], None],
-            [adaptor_paths[1], adaptor_paths[0], adaptor_paths[1]],
-        ]
-        common_args = dict(
-            torch_dtype=torch_dtype,
-            model_type="generation",
-            tp_size=model_case.tp_size,
-            lora_paths=adaptor_paths,
-            max_loras_per_batch=model_case.max_loras_per_batch,
-            max_loaded_loras=model_case.max_loaded_loras,
-            disable_cuda_graph=True,
-            disable_radix_cache=True,
-            mem_fraction_static=0.65,
-            sleep_on_idle=True,
         )
 
-        results_no_overlap_loading = []
+    def test_ci_lora_models_multi_batch(self):
-        with SRTRunner(
+        run_lora_multiple_batch_on_model_cases(
-            base_path, enable_lora_overlap_loading=False, **common_args
+            CI_MULTI_LORA_MODELS, enable_lora_overlap_loading=True
-        ) as runner:
+        )
-            for lora_paths in configs:
-                results_no_overlap_loading.append(
-                    runner.batch_forward(
-                        prompts, max_new_tokens=max_new_tokens, lora_paths=lora_paths
-                    ).output_strs
-                )
-
-        results_overlap_loading = []
-        with SRTRunner(
-            base_path, enable_lora_overlap_loading=True, **common_args
-        ) as runner:
-            for lora_paths in configs:
-                results_overlap_loading.append(
-                    runner.batch_forward(
-                        prompts, max_new_tokens=max_new_tokens, lora_paths=lora_paths
-                    ).output_strs
-                )
-
-        for i, (res_no_overlap_loading, res_overlap_loading) in enumerate(
-            zip(results_no_overlap_loading, results_overlap_loading)
-        ):
-            scores = calculate_rouge_l(res_overlap_loading, res_no_overlap_loading)
-            for j, score in enumerate(scores):
-                assert score >= model_case.rouge_l_tolerance, (
-                    f"Batch {i} prompt {j} mismatch: {score}\n"
-                    f"Overlap loading: {res_overlap_loading[j]}\n"
-                    f"No overlap loading: {res_no_overlap_loading[j]}"
-                )
-
-    def test_mixed_batch(self):
-        for model_case in CI_MULTI_LORA_MODELS:
-            for dtype in TORCH_DTYPES:
-                self._run_mixed_batch_test(model_case, dtype)
 
 
 if __name__ == "__main__":