From 99dad105fd85a5c6fc46e22c88e702f2511c63be Mon Sep 17 00:00:00 2001 From: Glen Liu <62917497+glenliu21@users.noreply.github.com> Date: Sun, 1 Feb 2026 17:52:08 -0500 Subject: [PATCH] [TestFix] rewrite LoRA overlap loading tests (#18047) --- python/sglang/test/lora_utils.py | 4 + .../lora/test_lora_overlap_loading.py | 97 +++---------------- 2 files changed, 18 insertions(+), 83 deletions(-) diff --git a/python/sglang/test/lora_utils.py b/python/sglang/test/lora_utils.py index 1f8d64f73..34de202be 100644 --- a/python/sglang/test/lora_utils.py +++ b/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, diff --git a/test/registered/lora/test_lora_overlap_loading.py b/test/registered/lora/test_lora_overlap_loading.py index 3b3b6d72d..c13507c2b 100644 --- a/test/registered/lora/test_lora_overlap_loading.py +++ b/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, - TORCH_DTYPES, - LoRAModelCase, - ensure_reproducibility, + run_lora_batch_splitting_equivalence_test, + run_lora_multiple_batch_on_model_cases, ) -from sglang.test.runners import SRTRunner -from sglang.test.test_utils import CustomTestCase, calculate_rouge_l +from sglang.test.test_utils import CustomTestCase -register_cuda_ci( - est_time=300, - 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", -) +register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu") +register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd") -class TestLoRAPipelineLoading(CustomTestCase): - - def _run_mixed_batch_test( - self, - 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, +class TestLoRAOverlapLoading(CustomTestCase): + def test_ci_lora_models_batch_splitting(self): + run_lora_batch_splitting_equivalence_test( + CI_MULTI_LORA_MODELS, enable_lora_overlap_loading=True ) - results_no_overlap_loading = [] - with SRTRunner( - base_path, enable_lora_overlap_loading=False, **common_args - ) 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) + def test_ci_lora_models_multi_batch(self): + run_lora_multiple_batch_on_model_cases( + CI_MULTI_LORA_MODELS, enable_lora_overlap_loading=True + ) if __name__ == "__main__":