From 91a45ea37e043979731c6725c383fa73759a2aab Mon Sep 17 00:00:00 2001 From: Anupa Sajikumar Date: Tue, 8 Sep 2026 10:16:41 +0530 Subject: [PATCH] [Intel][XPU][LoRA] Enable LoRA on Intel XPU (#30345) --- .github/workflows/pr-test-xpu.yml | 4 +- python/sglang/srt/arg_groups/overrides.py | 2 + .../srt/layers/rotary_embedding/base.py | 16 +- .../srt/layers/rotary_embedding/mrope.py | 11 +- .../sglang/srt/lora/backend/base_backend.py | 10 +- .../srt/lora/backend/chunked_backend.py | 2 +- .../sglang/srt/lora/backend/torch_backend.py | 2 +- .../sglang/srt/lora/backend/triton_backend.py | 2 +- python/sglang/srt/lora/lora_moe_runners.py | 2 +- python/sglang/srt/lora/lora_overlap_loader.py | 21 +- python/sglang/test/lora_utils.py | 52 +- test/manual/lora/test_lora_xpu_basics.py | 1038 +++++++++++++++++ test/manual/lora/test_torch_backend.py | 14 +- .../lora/test_fused_moe_lora_kernel.py | 56 +- .../test_lora_moe_vllm_sgl_logprob_diff.py | 27 +- .../lora/test_lora_overlap_loading.py | 9 +- .../lora/test_lora_qwen3_8b_logprob_diff.py | 23 +- test/registered/lora/test_moe_lora_info.py | 16 +- .../lora/test_virtual_experts_kernels.py | 33 +- 19 files changed, 1262 insertions(+), 78 deletions(-) create mode 100644 test/manual/lora/test_lora_xpu_basics.py diff --git a/.github/workflows/pr-test-xpu.yml b/.github/workflows/pr-test-xpu.yml index ff284c96a..2ab203de7 100644 --- a/.github/workflows/pr-test-xpu.yml +++ b/.github/workflows/pr-test-xpu.yml @@ -110,7 +110,7 @@ jobs: timeout-minutes: 60 run: | docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install --upgrade pip - docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9" + docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9" accelerate docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip uninstall -y flashinfer-python sgl-kernel sglang docker exec ci_sglang_xpu cp /sglang-checkout/python/pyproject_xpu.toml /sglang-checkout/python/pyproject.toml # Fetch tags so setuptools_scm resolves a real version instead of @@ -202,7 +202,7 @@ jobs: timeout-minutes: 60 run: | docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install --upgrade pip - docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9" + docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install pytest expecttest ray huggingface_hub tabulate "lmcache>=0.3.9" accelerate docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip uninstall -y flashinfer-python sgl-kernel sglang docker exec ci_sglang_xpu cp /sglang-checkout/python/pyproject_xpu.toml /sglang-checkout/python/pyproject.toml # Fetch tags so setuptools_scm resolves a real version instead of diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index c02e823af..bbfaf34ea 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -564,6 +564,8 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset( def supports_mamba_cache_extra_buffer(view: Any, model_arch: str) -> bool: """Whether ``model_arch`` supports the extra_buffer strategy on the configured linear-attention backend (pure read).""" + if get_platform().is_xpu: + return False if model_arch in _MAMBA_EXTRA_BUFFER_ARCHS: return view.linear_attn_backend == "triton" return False diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index 778ab7e9b..a8bda5656 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -72,7 +72,15 @@ if _is_hip: ) if _is_xpu: - from sgl_kernel import fused_qk_rope_with_cos_sin_cache_inplace + try: + from sgl_kernel import fused_qk_rope_with_cos_sin_cache_inplace + except ImportError: + fused_qk_rope_with_cos_sin_cache_inplace = None + logger.warning( + "sgl_kernel.fused_qk_rope_with_cos_sin_cache_inplace is unavailable; " + "XPU rotary embedding will use the generic rotary_embedding kernel. " + "Upgrade sgl_kernel to enable the fused XPU kernel." + ) class RotaryEmbedding(BaseFusedOp): @@ -454,7 +462,11 @@ class RotaryEmbedding(BaseFusedOp): positions = torch.add(positions, offsets) if offsets is not None else positions # Fused_qk_rope only supports aligned head_size - if self.head_size in [128, 256, 512]: + if fused_qk_rope_with_cos_sin_cache_inplace is not None and self.head_size in [ + 128, + 256, + 512, + ]: num_tokens = positions.size(0) q_rope = query.view(num_tokens, -1, self.head_size) k_rope = key.view(num_tokens, -1, self.head_size) diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index f60c24f4d..777edc7c5 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -39,7 +39,10 @@ if _is_npu: import torch_npu if _is_xpu: - from sgl_kernel import multimodal_rotary_embedding + try: + from sgl_kernel import multimodal_rotary_embedding + except ImportError: + multimodal_rotary_embedding = None from sglang.kernels.ops.attention.mrope import apply_interleaved_rope_triton @@ -309,7 +312,11 @@ class MRotaryEmbedding(RotaryEmbedding): ) -> Tuple[torch.Tensor, torch.Tensor]: assert positions.ndim in (1, 2) self._match_cos_sin_cache_dtype(query) - if positions.ndim == 2 and self.mrope_section: + if ( + multimodal_rotary_embedding is not None + and positions.ndim == 2 + and self.mrope_section + ): multimodal_rotary_embedding( query, key, diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 29723d929..e5ddb3adb 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -440,10 +440,8 @@ def _compute_moe_lora_info( adapter_enabled.zero_() has_segments = weight_indices.numel() != 0 - use_cuda_kernel = ( - num_tokens != 0 and has_segments and seg_indptr.device.type == "cuda" - ) - if use_cuda_kernel: + needs_launch = num_tokens != 0 and has_segments + if needs_launch: block_size = 256 tiles_per_segment = triton.cdiv(max_len, block_size) grid_size = tiles_per_segment * weight_indices.numel() @@ -451,6 +449,10 @@ def _compute_moe_lora_info( f"MoE LoRA token-mapping launch under-covers tokens: " f"{grid_size=} {block_size=} {num_tokens=}" ) + + # Triton kernel on CUDA only; every other device (e.g. XPU) falls through to + # the native torch path below, which yields the same mapping. + if needs_launch and seg_indptr.device.type == "cuda": _compute_moe_lora_info_kernel[(grid_size,)]( seg_indptr, lora_ranks, diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 0f0f74c03..eb88d8273 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -223,7 +223,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): (num_tokens_per_req + MIN_CHUNK_SIZE - 1) // MIN_CHUNK_SIZE ) * max_bs_in_cuda_graph max_num_tokens = max_bs_in_cuda_graph * num_tokens_per_req - with torch.device("cuda"): + with torch.device(self.device): self.cuda_graph_batch_info = LoRABatchInfo( bs=max_bs_in_cuda_graph, use_cuda_graph=True, diff --git a/python/sglang/srt/lora/backend/torch_backend.py b/python/sglang/srt/lora/backend/torch_backend.py index 4e3ada965..c88e06403 100644 --- a/python/sglang/srt/lora/backend/torch_backend.py +++ b/python/sglang/srt/lora/backend/torch_backend.py @@ -166,7 +166,7 @@ class TorchNativeLoRABackend(BaseLoRABackend): max_bs_in_cuda_graph: int, num_tokens_per_req: int, ): - with torch.device("cuda"): + with torch.device(self.device): self.cuda_graph_batch_info = TorchNativeLoRABatchInfo( use_cuda_graph=True, bs=max_bs_in_cuda_graph, diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index 12656ecbe..aa578c399 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -163,7 +163,7 @@ class TritonLoRABackend(BaseLoRABackend): ): max_tokens = max_bs_in_cuda_graph * num_tokens_per_req mlpb = self.max_loras_per_batch - with torch.device("cuda"): + with torch.device(self.device): self.cuda_graph_batch_info = LoRABatchInfo( bs=max_bs_in_cuda_graph, use_cuda_graph=True, diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 7ffa971c7..8b1d3e5e4 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -227,7 +227,7 @@ def _compute_lora_alignment( device = topk_ids.device - use_naive = ( + use_naive = _is_xpu or ( cg is None and M * topk_ids.shape[1] * _SPARSITY_FACTOR <= lora_info.num_experts * max_loras diff --git a/python/sglang/srt/lora/lora_overlap_loader.py b/python/sglang/srt/lora/lora_overlap_loader.py index 9e9d725e1..c05a4e997 100644 --- a/python/sglang/srt/lora/lora_overlap_loader.py +++ b/python/sglang/srt/lora/lora_overlap_loader.py @@ -44,16 +44,19 @@ class LoRAOverlapLoader: lora_pipeline_load_status = self._check_overlap_load_status(lora_id) if lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADING: return False - elif lora_pipeline_load_status == LoRAOverlapLoadStatus.NOT_LOADED: - res = self._try_start_overlap_load(lora_id, running_loras) - if res: - logger.debug(f"Loading LoRA adapter {lora_id} asynchronously") - - return False - else: - assert lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADED + elif lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADED: return True + assert lora_pipeline_load_status == LoRAOverlapLoadStatus.NOT_LOADED + if not self._try_start_overlap_load(lora_id, running_loras): + return False + logger.debug(f"Loading LoRA adapter {lora_id} asynchronously") + + # Report an already-finished copy as LOADED, or a sibling's load could + # evict it before it is ever scheduled. No-op while still in flight. + self._drain_completed_overlap_loads() + return self._check_overlap_load_status(lora_id) == LoRAOverlapLoadStatus.LOADED + def _check_overlap_load_status( self, lora_id: Optional[str] ) -> LoRAOverlapLoadStatus: @@ -74,7 +77,7 @@ class LoRAOverlapLoader: if event.query() ] for lora_id, event in completed_loads: - torch.cuda.current_stream().wait_event(event) + self.device_module.current_stream().wait_event(event) del self.lora_to_overlap_load_event[lora_id] def _try_start_overlap_load( diff --git a/python/sglang/test/lora_utils.py b/python/sglang/test/lora_utils.py index f5c65b18e..4d5ca6cca 100644 --- a/python/sglang/test/lora_utils.py +++ b/python/sglang/test/lora_utils.py @@ -4,9 +4,36 @@ from typing import List, Optional import torch +from sglang.srt.utils import is_xpu from sglang.test.runners import HFRunner, SRTRunner from sglang.test.test_utils import calculate_rouge_l +_IS_XPU = is_xpu() + + +def _assert_lora_output_match( + srt_str: str, hf_str: str, rouge_tol: float, context: str +): + """Compare SRT vs HF greedy output strings. + + Everywhere except XPU we keep the historical strict exact-match (SGLang and HF + kernels agree numerically enough for greedy argmax to pick identical tokens). + On XPU, small kernel-level fp differences can make greedy decoding diverge + after a shared prefix even when the LoRA math is correct, so we fall back to + the same ROUGE-L tolerance the per-adaptor comparison path uses. + """ + srt_str = srt_str.strip(" ") + hf_str = hf_str.strip(" ") + if not _IS_XPU: + assert srt_str == hf_str, (srt_str, hf_str) + return + rouge_score = calculate_rouge_l([srt_str], [hf_str])[0] + if rouge_score < rouge_tol: + raise AssertionError( + f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} for {context}. " + f"SRT: {srt_str!r} HF: {hf_str!r}" + ) + @dataclasses.dataclass class LoRAAdaptor: @@ -624,17 +651,22 @@ def run_lora_test_by_batch( print("HF output:", hf_output_str) print("SRT no lora output:", srt_no_lora_outputs.output_strs[i].strip()) print("HF no lora output:", hf_no_lora_outputs.output_strs[i].strip()) - assert srt_outputs.output_strs[i].strip(" ") == hf_outputs.output_strs[i].strip( - " " - ), ( - srt_outputs.output_strs[i].strip(" "), - hf_outputs.output_strs[i].strip(" "), + rouge_tol = ( + adaptors[i].rouge_l_tolerance + if adaptors[i].rouge_l_tolerance is not None + else model_case.rouge_l_tolerance ) - assert srt_no_lora_outputs.output_strs[i].strip( - " " - ) == hf_no_lora_outputs.output_strs[i].strip(" "), ( - srt_no_lora_outputs.output_strs[i].strip(" "), - hf_no_lora_outputs.output_strs[i].strip(" "), + _assert_lora_output_match( + srt_outputs.output_strs[i], + hf_outputs.output_strs[i], + rouge_tol, + f"base '{base_path}', adaptor '{adaptor_names[i]}', backend '{backend}' (LoRA)", + ) + _assert_lora_output_match( + srt_no_lora_outputs.output_strs[i], + hf_no_lora_outputs.output_strs[i], + rouge_tol, + f"base '{base_path}', backend '{backend}' (no-LoRA baseline)", ) diff --git a/test/manual/lora/test_lora_xpu_basics.py b/test/manual/lora/test_lora_xpu_basics.py new file mode 100644 index 000000000..02d67d577 --- /dev/null +++ b/test/manual/lora/test_lora_xpu_basics.py @@ -0,0 +1,1038 @@ +# 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. +# ============================================================================== + +""" +Core LoRA serving features: accuracy against HF+PEFT, multi-adapter batching, +the drainer, runtime load/unload, pinning, pool eviction, the radix cache, and +embedding models. + +Adapters are synthetic with known-distinct weights, so most cases assert +nearest-reference routing rather than exact output. + +The cases are in two groups: each feature in isolation, then feature +*combinations* with graph capture OFF (the "Feature combinations" section at the +bottom). Device-graph capture and tensor parallelism -- including the graph-ON +combinations -- live in test_lora_comb_xpu.py. +""" + +import atexit +import multiprocessing as mp +import os +import shutil +import tempfile +import unittest +from typing import List, Optional + +# XPU cannot be re-initialized in a forked child and SRTRunner spawns one, so the +# start method is forced here (pytest never runs the __main__ block below). +try: + mp.set_start_method("spawn") +except RuntimeError: + pass + +import torch + +from sglang.test.lora_utils import BACKENDS +from sglang.test.runners import HFRunner, SRTRunner +from sglang.test.test_utils import ( + DEFAULT_PORT_FOR_SRT_TEST_RUNNER, + CustomTestCase, + calculate_rouge_l, + empty_gpu_cache, +) + +DENSE_BACKENDS = [*BACKENDS, "torch_native"] + +MEM_FRACTION_STATIC = 0.80 +MAX_NEW_TOKENS = 16 +LOGPROB_THRESHOLD = 2e-1 + +TEST_PROMPTS = [ + "AI is a field of computer science focused on", + "The capital of France is", +] + +# > 0 enables the scheduler's LoRADrainer, which frees slots for starving +# adapters once max_loras_per_batch fills. Must not change outputs. +LORA_DRAIN_WAIT_THRESHOLD = 0.1 + + +# ============================================================================== +# Synthetic LoRA adapter fixtures +# +# Shared by the tests that need adapters with known-distinct weights rather than +# a published checkpoint: routing assertions have to tell adapters apart, which +# real adapters trained on the same task do not reliably allow. +# ============================================================================== + +SYNTH_BASE_MODEL = "Qwen/Qwen2.5-0.5B" + +# q_proj/v_proj fuse into the column-parallel qkv_proj, which is what makes the +# TP cases a real test of sharding (LoRA A replicated, LoRA B sliced). +SYNTH_TARGET_MODULES = ["q_proj", "v_proj"] + +SYNTH_LORA_RANK = 8 +SYNTH_MAX_LORA_RANK = 16 # engines started without adapters must declare a max + +# Calibrated: 0.02 shifts logprobs (enough to compare against HF) but leaves the +# argmax unchanged, so multi-adapter routing tests need 0.04 for distinct tokens. +SINGLE_LORA_B_STD = 0.02 +MULTI_LORA_B_STD = 0.04 + +_temp_dirs = [] + + +def lora_temp_dir(prefix: str) -> str: + """Temp dir removed at interpreter exit. Fixtures are cached across test + classes, so cleanup cannot live in tearDownClass.""" + path = tempfile.mkdtemp(prefix=prefix) + _temp_dirs.append(path) + return path + + +@atexit.register +def _cleanup_lora_temp_dirs(): + for path in _temp_dirs: + if os.path.exists(path): + shutil.rmtree(path, ignore_errors=True) + + +def create_attention_lora_adapter( + base_model_name: str, + output_dir: str, + seed: int = 0, + lora_b_std: float = SINGLE_LORA_B_STD, +): + """ + Create a randomly-initialized (untrained) LoRA adapter on the attention + projections. PEFT zeroes lora_B, which would make the adapter a no-op, so + lora_B gets non-zero weights and the adapter has a verifiable effect. + + Both lora_A and lora_B are seeded from one generator: PEFT draws lora_A from + the *global* RNG, so seeding only lora_B leaves B@A varying run-to-run, which + can make an adapter too weak to move any greedy token (reading as identical + to base). Seeding both removes that flakiness at the source. + """ + from peft import LoraConfig, get_peft_model + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained( + base_model_name, torch_dtype=torch.float16, device_map="cpu" + ) + peft_model = get_peft_model( + model, + LoraConfig( + r=SYNTH_LORA_RANK, + lora_alpha=2 * SYNTH_LORA_RANK, + target_modules=SYNTH_TARGET_MODULES, + lora_dropout=0, + bias="none", + task_type="CAUSAL_LM", + ), + ) + + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for name, param in peft_model.named_parameters(): + if "lora_A" in name: + param.normal_( + mean=0.0, std=1.0 / param.shape[1] ** 0.5, generator=generator + ) + elif "lora_B" in name: + param.normal_(mean=0.0, std=lora_b_std, generator=generator) + + peft_model.save_pretrained(output_dir) + + # A silently empty adapter would make every routing assertion vacuous. + from safetensors import safe_open + + keys = safe_open( + os.path.join(output_dir, "adapter_model.safetensors"), framework="pt" + ).keys() + assert any(t in k for k in keys for t in SYNTH_TARGET_MODULES), ( + f"Expected {SYNTH_TARGET_MODULES} LoRA weights in adapter, got: {sorted(keys)}" + ) + + del peft_model, model + empty_gpu_cache() + + +# Adapters are referred to by identity, not list index, so one cached adapter (or +# isolated-reference run) serves classes with different slot layouts. +ADAPTER_A = "A" +ADAPTER_B = "B" +ADAPTER_C = "C" +BASE_SLOT = None # "no adapter" -- the base model +ADAPTER_SEEDS = {ADAPTER_A: 1, ADAPTER_B: 2, ADAPTER_C: 3} + +_adapter_dirs = {} + + +def synth_adapter_path( + identity: str, + base_model: str = SYNTH_BASE_MODEL, + lora_b_std: float = MULTI_LORA_B_STD, +) -> str: + """Shared adapter for ``identity``, built on first use and cached.""" + key = (identity, base_model, lora_b_std) + if key not in _adapter_dirs: + path = lora_temp_dir(f"sglang_test_lora_{identity}_{lora_b_std}_") + create_attention_lora_adapter( + base_model, path, seed=ADAPTER_SEEDS[identity], lora_b_std=lora_b_std + ) + _adapter_dirs[key] = path + return _adapter_dirs[key] + + +def synth_slot_paths(slots, base_model: str = SYNTH_BASE_MODEL, **kwargs): + """Adapter path per slot, ``None`` for the base-model slot.""" + return [ + None if s is BASE_SLOT else synth_adapter_path(s, base_model, **kwargs) + for s in slots + ] + + +def assert_nearest_reference_routing( + test_case, + *, + batched: List[str], + reference: dict, + slots: List[Optional[str]], + context: str = "", +): + """ + Assert each batched output is closer (ROUGE-L) to its own adapter's isolated + output than to any other adapter's. + + Exact reproduction is not required: batching changes kernel tiling, so tokens + can differ slightly from a bs=1 run even when routing is correct. Nearest- + reference is the strongest assertion that stays robust to that -- a swapped or + dropped adapter moves the argmax to another slot, which this catches. + + The base slot is excluded from the reference set: base output is a prefix-like + attractor that every adapter's output stays near, so including it would make + the argmax uninformative. + """ + test_case.assertEqual( + len(batched), + len(slots), + f"{context} expected {len(slots)} outputs, got {len(batched)}", + ) + + adapter_slots = [i for i, s in enumerate(slots) if s is not BASE_SLOT] + for i in adapter_slots: + sims = { + j: calculate_rouge_l([batched[i]], [reference[slots[j]]])[0] + for j in adapter_slots + } + best = max(adapter_slots, key=lambda j: sims[j]) + print( + f"{context} slot {i} (adapter={slots[i]}) self={sims[i]:.4f} " + f"argmax=slot{best}" + ) + test_case.assertEqual( + best, + i, + f"{context} slot {i} (adapter={slots[i]}): batched output is closer to " + f"another adapter's isolated output than to its own " + f"(sims={ {j: round(sims[j], 3) for j in adapter_slots} }); " + f"wrong adapter routed?", + ) + + # Distinct adapters on an identical prompt must not collapse. + test_case.assertFalse( + all(s == batched[0] for s in batched), + f"{context} all outputs identical despite distinct adapters {slots}; " + f"routing had no effect.", + ) + + +def lora_runner( + *, + backend: str = "triton", + disable_radix: bool = True, + model_type: str = "generation", + **extra, +) -> SRTRunner: + """SRTRunner with the kwargs shared by every case here. Graph capture is off + throughout, for the combination cases as well as the isolated ones, so that a + failure is never ambiguous between the feature and graph capture; the graph-ON + interactions are covered in test_lora_comb_xpu.py.""" + return SRTRunner( + SYNTH_BASE_MODEL, + torch_dtype=torch.float16, + model_type=model_type, + lora_backend=backend, + lora_target_modules=SYNTH_TARGET_MODULES, + disable_cuda_graph=True, + disable_radix_cache=disable_radix, + mem_fraction_static=MEM_FRACTION_STATIC, + port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER, + **extra, + ) + + +def generate(runner: SRTRunner, lora_paths, batched: bool = False): + """Serve one request per entry of ``lora_paths`` on a shared prompt, returning + the stripped outputs. ``batched`` forces all of them into one physical batch + (sleep_on_idle on the runner keeps the scheduler from dribbling them out).""" + forward = runner.batch_forward if batched else runner.forward + outputs = forward( + [TEST_PROMPTS[0]] * len(lora_paths), + max_new_tokens=MAX_NEW_TOKENS, + lora_paths=list(lora_paths), + ) + return [s.strip() for s in outputs.output_strs] + + +_references = {} + + +def isolated_reference(slots): + """ + Per-adapter baseline for the routing check, keyed by adapter identity so one + run serves classes with different slot layouts. forward() runs each request as + its own batch, with enough slots that nothing splits or drains. + + The backend is pinned to triton deliberately: the reference is the fixed + fingerprint every backend is matched against, so a mis-routing backend cannot + hide by shifting its own reference alongside its output. + """ + key = tuple(slots) + if key in _references: + return _references[key] + + paths = synth_slot_paths(slots) + adapters = [p for p in paths if p is not None] + with lora_runner( + lora_paths=adapters, max_loras_per_batch=len(adapters) + 1 + ) as runner: + outputs = generate(runner, paths) + empty_gpu_cache() + + _references[key] = {s: outputs[i] for i, s in enumerate(slots)} + return _references[key] + + +# ============================================================================== +# Accuracy vs HuggingFace + PEFT +# ============================================================================== + + +class TestLoRAAccuracy(CustomTestCase): + """ + A basic attention LoRA must match HF+PEFT on every dense backend. The HF + reference is backend-independent, so it is computed once. This is the only + *absolute* numerical check in the file; everything else asserts relative + properties (routing, round-trip, determinism). + """ + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.lora_path = synth_adapter_path(ADAPTER_A, lora_b_std=SINGLE_LORA_B_STD) + with HFRunner( + SYNTH_BASE_MODEL, torch_dtype=torch.float16, model_type="generation" + ) as hf_runner: + cls.hf_outputs = hf_runner.forward( + TEST_PROMPTS, + max_new_tokens=MAX_NEW_TOKENS, + lora_paths=[cls.lora_path] * len(TEST_PROMPTS), + ) + empty_gpu_cache() + + def _check_backend(self, backend: str): + with lora_runner( + backend=backend, lora_paths=[self.lora_path], max_loras_per_batch=1 + ) as runner: + srt_outputs = runner.forward( + TEST_PROMPTS, + max_new_tokens=MAX_NEW_TOKENS, + lora_paths=[self.lora_path] * len(TEST_PROMPTS), + ) + empty_gpu_cache() + + for phase, srt, hf in ( + ( + "prefill", + srt_outputs.top_input_logprobs, + self.hf_outputs.top_input_logprobs, + ), + ( + "decode", + srt_outputs.top_output_logprobs, + self.hf_outputs.top_output_logprobs, + ), + ): + for i in range(len(TEST_PROMPTS)): + max_diff = torch.max( + torch.abs(torch.tensor(srt[i]) - torch.tensor(hf[i])) + ).item() + print( + f"[{backend}] prompt {i} {phase} logprob max_diff " + f"(SGLang vs HF): {max_diff:.6e}" + ) + self.assertLess( + max_diff, + LOGPROB_THRESHOLD, + f"[{backend}] prompt {i}: {phase} logprob diff {max_diff:.6e} " + f"exceeds threshold {LOGPROB_THRESHOLD:.0e}", + ) + + def test_lora_matches_hf(self): + for backend in DENSE_BACKENDS: + with self.subTest(backend=backend): + self._check_backend(backend) + + +# ============================================================================== +# Multi-LoRA routing: one physical batch, every slot on its own adapter +# ============================================================================== + + +class _RoutingTestCase(CustomTestCase): + """ + Base for the multi-adapter routing tests. Subclasses declare the slot layout + and any extra server args; the batched run and its reference share the rest, + and the assertion lives here so the two halves cannot drift apart. + + max_loras_per_batch below the slot count is how the drainer case forces a + running adapter out for a starving one. + """ + + SLOTS = [ADAPTER_A, ADAPTER_B, BASE_SLOT] + MAX_LORAS_PER_BATCH = 3 + BACKENDS = BACKENDS + RUNNER_ARGS = {} + DISABLE_RADIX = True + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.paths = synth_slot_paths(cls.SLOTS) + cls.reference = isolated_reference(cls.SLOTS) + + def _check_backend(self, backend: str): + adapters = [p for p in self.paths if p is not None] + with lora_runner( + backend=backend, + disable_radix=self.DISABLE_RADIX, + lora_paths=adapters, + max_loras_per_batch=self.MAX_LORAS_PER_BATCH, + sleep_on_idle=True, + **self.RUNNER_ARGS, + ) as runner: + batched = generate(runner, self.paths, batched=True) + empty_gpu_cache() + + assert_nearest_reference_routing( + self, + batched=batched, + reference=self.reference, + slots=self.SLOTS, + context=f"[{backend}]", + ) + + def _check_routing(self): + for backend in self.BACKENDS: + with self.subTest(backend=backend): + self._check_backend(backend) + + +class TestMultiLoRABatching(_RoutingTestCase): + """Multi-LoRA batching (S-LoRA/Punica): two distinct adapters plus base in one + physical batch, each request routed to its own adapter. Three slots for three + requests, so nothing splits or drains.""" + + BACKENDS = DENSE_BACKENDS + + def test_routing(self): + self._check_routing() + + +class TestLoRADrainWaitThreshold(_RoutingTestCase): + """Drainer (lora_drain_wait_threshold > 0): three distinct adapters through + two slots forces the scheduler to drain a running adapter for a starving one. + Draining is scheduling fairness, not correctness -- it changes *when* a request + runs, never which adapter it uses -- so routing must be unaffected.""" + + SLOTS = [ADAPTER_A, ADAPTER_B, ADAPTER_C] # no base: 3 through 2 slots + MAX_LORAS_PER_BATCH = 2 + RUNNER_ARGS = {"lora_drain_wait_threshold": LORA_DRAIN_WAIT_THRESHOLD} + + def test_routing(self): + self._check_routing() + + +class TestLoRAOverlapLoading(_RoutingTestCase): + """Overlap loading end-to-end: weights are staged on a side device stream + concurrently with compute. It changes only how/when weights are staged, never + the numerics, so the routing check applies unchanged. + + max_loaded_loras is required with overlap loading, within + [max_loras_per_batch, 2x].""" + + RUNNER_ARGS = {"enable_lora_overlap_loading": True, "max_loaded_loras": 6} + + def test_routing(self): + self._check_routing() + + +class TestLoRARadixCacheRouting(_RoutingTestCase): + """Radix prefix cache ENABLED (every other case here disables it). Radix keys + prefixes by adapter via extra_key, so the same prompt under two adapters lands + in disjoint namespaces. All slots share one prompt -- the worst case for + cross-contamination; if radix ever shared KV across adapters, a slot would look + closest to another adapter's output.""" + + DISABLE_RADIX = False + + def test_routing(self): + self._check_routing() + + +# ============================================================================== +# Radix cache hit +# ============================================================================== + + +class TestLoRARadixCacheHit(CustomTestCase): + """ + Serving the same (adapter, prompt) twice in one engine hits the first + request's inserted prefix; with greedy decoding a correct hit reuses the exact + KV, so equality is exact and a mismatch means corrupted or wrong-adapter KV. + Serving a different adapter on the same prompt probes the other direction: the + cache must not collapse distinct adapters onto one entry. + """ + + def test_radix_cache_hit_reproduces_and_keys_on_adapter(self): + paths = synth_slot_paths([ADAPTER_A, ADAPTER_B]) + with lora_runner( + disable_radix=False, lora_paths=paths, max_loras_per_batch=2 + ) as runner: + (a_first,) = generate(runner, paths[:1]) + (a_second,) = generate(runner, paths[:1]) # cache hit + (b,) = generate(runner, paths[1:]) + empty_gpu_cache() + + print(f"[radix-hit] A first={a_first!r} A second={a_second!r} B={b!r}") + self.assertEqual( + a_second, + a_first, + "radix cache hit did not reproduce adapter A's output (the prefix " + "cache may have corrupted or dropped the adapter KV)", + ) + self.assertNotEqual( + b, + a_first, + "adapter B produced adapter A's cached output on the same prompt " + "(the radix cache is not keying on the LoRA adapter)", + ) + + +# ============================================================================== +# Runtime adapter mutation +# ============================================================================== + + +class TestDynamicLoRAUpdate(CustomTestCase): + """ + Dynamic load/unload (/load_lora_adapter, /unload_lora_adapter) on an engine + that starts with NO adapters: load two at runtime, serve them by name, unload + one and reload it. load/unload return a LoRAUpdateOutput and never raise; only + a forward referencing an unknown name raises, and the probe uses a NEVER + loaded name since an unloaded one would be transparently reloaded. + """ + + def _check_load_unload_reload(self, backend: str): + with lora_runner( + backend=backend, + # No initial adapters -> enable_lora + max_lora_rank + + # lora_target_modules are all required. + lora_paths=None, + enable_lora=True, + max_lora_rank=SYNTH_MAX_LORA_RANK, + max_loras_per_batch=2, + max_loaded_loras=2, + ) as runner: + load_a = runner.load_lora_adapter( + lora_name="adapter_a", lora_path=synth_adapter_path(ADAPTER_A) + ) + load_b = runner.load_lora_adapter( + lora_name="adapter_b", lora_path=synth_adapter_path(ADAPTER_B) + ) + self.assertTrue(load_a.success, f"load_a: {load_a.error_message}") + self.assertTrue(load_b.success, f"load_b: {load_b.error_message}") + self.assertEqual( + set(load_b.loaded_adapters or []), + {"adapter_a", "adapter_b"}, + f"unexpected loaded set: {load_b.loaded_adapters}", + ) + + base, out_a, out_b = generate(runner, [None, "adapter_a", "adapter_b"]) + print(f"[{backend}] base={base!r} a={out_a!r} b={out_b!r}") + self.assertNotEqual(out_a, base, "adapter_a matched base") + self.assertNotEqual(out_b, base, "adapter_b matched base") + self.assertNotEqual(out_a, out_b, "both adapters identical") + + with self.assertRaises( + ValueError, msg="a never-loaded name must raise on forward" + ): + generate(runner, ["adapter_never_loaded"]) + + unload = runner.unload_lora_adapter(lora_name="adapter_a") + self.assertTrue(unload.success, f"unload: {unload.error_message}") + self.assertNotIn( + "adapter_a", + set(unload.loaded_adapters or []), + f"adapter_a still loaded: {unload.loaded_adapters}", + ) + (b_after_unload,) = generate(runner, ["adapter_b"]) + self.assertEqual( + b_after_unload, out_b, "adapter_b changed after unloading adapter_a" + ) + + # Reload explicitly, so the round-trip tested is the load path rather + # than the forward-time auto-reload fallback. + reload_a = runner.load_lora_adapter( + lora_name="adapter_a", lora_path=synth_adapter_path(ADAPTER_A) + ) + self.assertTrue(reload_a.success, f"reload: {reload_a.error_message}") + (a_reloaded,) = generate(runner, ["adapter_a"]) + self.assertEqual( + a_reloaded, out_a, "adapter_a changed after unload + reload" + ) + empty_gpu_cache() + + def test_dynamic_load_unload_reload(self): + self._check_load_unload_reload("triton") + + def test_dynamic_load_unload_reload_torch_native(self): + self._check_load_unload_reload("torch_native") + + +class TestPinnedLoRAAdapters(CustomTestCase): + """ + Pinned adapters (load_lora_adapter(pinned=True)) are excluded from eviction + victim selection, so churning the single free slot must never evict one. Also + probes the pin-all-slots guard: with 2 slots at most 1 may be pinned, and the + rejection is success=False, not an exception. + """ + + def test_pinned_adapter_loads_and_survives_eviction(self): + with lora_runner( + lora_paths=None, + enable_lora=True, + max_lora_rank=SYNTH_MAX_LORA_RANK, + max_loras_per_batch=2, + # Roomy so eviction is driven by the per-batch pool, not this cap. + max_loaded_loras=4, + ) as runner: + pinned = runner.load_lora_adapter( + lora_name="pinned_a", + lora_path=synth_adapter_path(ADAPTER_A), + pinned=True, + ) + unpinned = runner.load_lora_adapter( + lora_name="unpinned_b", lora_path=synth_adapter_path(ADAPTER_B) + ) + self.assertTrue(pinned.success, f"pinned load: {pinned.error_message}") + self.assertTrue( + unpinned.success, f"unpinned load: {unpinned.error_message}" + ) + + over_pin = runner.load_lora_adapter( + lora_name="pinned_c", + lora_path=synth_adapter_path(ADAPTER_C), + pinned=True, + ) + self.assertFalse( + over_pin.success, + "pinning a second adapter should be rejected (would pin all slots)", + ) + self.assertIn( + "pin all slots", + over_pin.error_message or "", + f"unexpected over-pinning error: {over_pin.error_message!r}", + ) + + (baseline,) = generate(runner, ["pinned_a"]) + for _ in range(3): + # pinned_c was rejected by the guard, so churn with unpinned_b. + generate(runner, ["unpinned_b"]) + (after_churn,) = generate(runner, ["pinned_a"]) + empty_gpu_cache() + + print(f"[pinning] baseline={baseline!r} after churn={after_churn!r}") + self.assertEqual( + after_churn, + baseline, + "pinned adapter output changed after serving other adapters (it may " + "have been evicted despite being pinned)", + ) + + +class TestLoRAEvictionEndToEnd(CustomTestCase): + """ + End-to-end pool eviction with one batch slot: serving a second adapter forces + the first out, and serving it again reloads it (mem_pool.py + select_victim/evict/reload). Asserts correctness under eviction rather than + which adapter was evicted -- output is policy-invariant, so victim identity is + test_lora_eviction_policy.py's job. + """ + + def _check_eviction(self, *, policy: str, backend: str = "triton"): + label = f"{policy},{backend}" + paths = synth_slot_paths([ADAPTER_A, ADAPTER_B]) + with lora_runner( + backend=backend, + lora_paths=paths, + max_loras_per_batch=1, + lora_eviction_policy=policy, + ) as runner: + (a_first,) = generate(runner, paths[:1]) + (b_first,) = generate(runner, paths[1:]) # evicts A + (a_again,) = generate(runner, paths[:1]) # must reload A + (b_again,) = generate(runner, paths[1:]) # must reload B + empty_gpu_cache() + + print(f"[evict:{label}] A={a_first!r} B={b_first!r}") + self.assertNotEqual( + a_first, b_first, f"[{label}] adapters A and B produced identical output" + ) + self.assertEqual( + a_again, a_first, f"[{label}] adapter A changed after eviction + reload" + ) + self.assertEqual( + b_again, b_first, f"[{label}] adapter B changed after eviction + reload" + ) + + def test_eviction_lru(self): + self._check_eviction(policy="lru") + + def test_eviction_fifo(self): + self._check_eviction(policy="fifo") + + def test_eviction_lru_torch_native(self): + self._check_eviction(policy="lru", backend="torch_native") + + def test_eviction_fifo_torch_native(self): + self._check_eviction(policy="fifo", backend="torch_native") + + +# ============================================================================== +# Embedding LoRA +# ============================================================================== + + +class TestEmbeddingLoRA(CustomTestCase): + """ + LoRA on an embedding model (last-token pooling, L2-normalized): LoRA on + q_proj/v_proj perturbs attention, which propagates to the pooled embedding. + Checked HF-free via cosine similarity (unit vectors, so cosine == dot): + unit-norm, the adapter moves the embedding off base, the same adapter + reproduces, and two adapters stay distinct -- the embedding analog of the + routing check used for generation. A is encoded twice to separate "the adapter + moved the embedding" from "the embedding is nondeterministic". + + SRTRunner.forward drops lora_path on the embedding path, so the adapter is + applied via engine.encode(prompt=..., lora_path=...) directly. + """ + + EFFECT_MAX_COS = 0.999 # base-vs-adapter must be below this + REPEAT_MIN_COS = 0.999 # the same adapter twice must be above this + + @staticmethod + def _cos(a, b) -> float: + ta = torch.tensor(a, dtype=torch.float32) + tb = torch.tensor(b, dtype=torch.float32) + return float(torch.dot(ta, tb) / (ta.norm() * tb.norm() + 1e-12)) + + def _check_embedding_backend(self, backend: str): + paths = synth_slot_paths([ADAPTER_A, ADAPTER_B]) + with lora_runner( + backend=backend, + model_type="embedding", + lora_paths=paths, + max_loras_per_batch=3, + ) as runner: + + def encode(lora_path): + resp = runner.engine.encode(prompt=TEST_PROMPTS[0], lora_path=lora_path) + return (resp[0] if isinstance(resp, list) else resp)["embedding"] + + vectors = { + "base": encode(None), + "a": encode(paths[0]), + "a2": encode(paths[0]), + "b": encode(paths[1]), + } + empty_gpu_cache() + + for name in ("base", "a", "b"): + norm = torch.tensor(vectors[name], dtype=torch.float32).norm().item() + print(f"[{backend}] {name} ||v||={norm:.4f}") + self.assertAlmostEqual( + norm, 1.0, delta=1e-2, msg=f"{name} embedding is not unit-norm" + ) + + cos_base_a = self._cos(vectors["base"], vectors["a"]) + cos_base_b = self._cos(vectors["base"], vectors["b"]) + cos_a_a = self._cos(vectors["a"], vectors["a2"]) + cos_a_b = self._cos(vectors["a"], vectors["b"]) + print( + f"[{backend}] cos(base,A)={cos_base_a:.5f} cos(base,B)={cos_base_b:.5f} " + f"cos(A,A)={cos_a_a:.5f} cos(A,B)={cos_a_b:.5f}" + ) + + self.assertLess( + cos_base_a, + self.EFFECT_MAX_COS, + f"adapter A did not change the embedding (cos(base,A)={cos_base_a:.5f})", + ) + self.assertLess( + cos_base_b, + self.EFFECT_MAX_COS, + f"adapter B did not change the embedding (cos(base,B)={cos_base_b:.5f})", + ) + self.assertGreater( + cos_a_a, + self.REPEAT_MIN_COS, + f"adapter A embedding not reproducible (cos(A,A)={cos_a_a:.5f})", + ) + self.assertLess( + cos_a_b, + cos_a_a, + f"adapter A and B embeddings not distinct (cos(A,B)={cos_a_b:.5f} " + f"not < cos(A,A)={cos_a_a:.5f})", + ) + + def test_embedding_lora_applies_and_routes(self): + self._check_embedding_backend("triton") + + def test_embedding_lora_applies_and_routes_torch_native(self): + self._check_embedding_backend("torch_native") + + +# ============================================================================== +# Feature combinations (graph OFF) +# +# Everything above isolates ONE serving feature. test_lora_comb_xpu.py covers +# combinations, but every case there holds disable_cuda_graph=False, so a failure +# has two candidate causes: the feature interaction, or graph capture. These cases +# close that gap by combining the same features with graph capture OFF, matching +# the rest of this file -- so a failure here localizes the bug to the interaction +# itself, and a case that fails there but passes here implicates graph capture. +# +# Only combinations whose parts are individually covered above appear here; each +# pairs features that contend for the SAME state, which is where interactions +# actually break: +# +# B1. drainer x radix cache -- both reorder/reuse work across adapters +# B2. overlap loading x drainer -- graph-OFF twin of C5, to separate the +# interaction from graph capture +# B3. overlap loading x eviction -- a slot is staged asynchronously while the +# victim selector is reclaiming slots +# B4. pinning x drainer -- the drainer picks adapters to displace and +# must respect the pin +# B5. dynamic load x radix cache -- an adapter loaded at runtime must not +# inherit a previous adapter's cached prefix +# ============================================================================== + + +class TestDrainerWithRadixCache(_RoutingTestCase): + """B1. Drainer x radix cache. Both features reuse work across adapters: the + drainer reorders *which* adapter runs when, and radix reuses KV keyed by + adapter. Combined, a prefix inserted before a drain must not be handed to the + adapter that displaced it. Three adapters through two slots forces draining + while every slot shares one prompt (maximal prefix contention).""" + + SLOTS = [ADAPTER_A, ADAPTER_B, ADAPTER_C] + MAX_LORAS_PER_BATCH = 2 + RUNNER_ARGS = {"lora_drain_wait_threshold": LORA_DRAIN_WAIT_THRESHOLD} + DISABLE_RADIX = False + + def test_routing(self): + self._check_routing() + + +class TestOverlapLoadingWithDrainer(_RoutingTestCase): + """B2. Overlap loading x drainer. Overlap loading stages weights on a side + stream while the drainer concurrently displaces adapters, so a staged copy can + land in a slot the drainer has already reclaimed. This is the graph-OFF twin + of C5 in test_lora_comb_xpu.py -- if that fails and this passes, the fault is + graph capture rather than the interaction.""" + + SLOTS = [ADAPTER_A, ADAPTER_B, ADAPTER_C] + MAX_LORAS_PER_BATCH = 2 + RUNNER_ARGS = { + "enable_lora_overlap_loading": True, + "max_loaded_loras": 4, + "lora_drain_wait_threshold": LORA_DRAIN_WAIT_THRESHOLD, + } + + def test_routing(self): + self._check_routing() + + +class TestOverlapLoadingWithEviction(CustomTestCase): + """B3. Overlap loading x pool eviction. With a single batch slot, every switch + evicts -- and with overlap loading the replacement is staged asynchronously, so + a reload races the victim's teardown. Asserts the same reload-fidelity property + as TestLoRAEvictionEndToEnd (each adapter reproduces its own first output), which + is what a torn or half-staged copy would break.""" + + def _check_eviction_with_overlap(self, backend: str): + paths = synth_slot_paths([ADAPTER_A, ADAPTER_B]) + with lora_runner( + backend=backend, + lora_paths=paths, + max_loras_per_batch=1, + enable_lora_overlap_loading=True, + max_loaded_loras=2, + ) as runner: + (a_first,) = generate(runner, paths[:1]) + (b_first,) = generate(runner, paths[1:]) # evicts A while staging B + (a_again,) = generate(runner, paths[:1]) # restage A over B + (b_again,) = generate(runner, paths[1:]) # restage B over A + empty_gpu_cache() + + print(f"[overlap+evict:{backend}] A={a_first!r} B={b_first!r}") + self.assertNotEqual( + a_first, + b_first, + f"[{backend}] adapters A and B produced identical output", + ) + self.assertEqual( + a_again, + a_first, + f"[{backend}] adapter A changed after eviction + overlapped reload", + ) + self.assertEqual( + b_again, + b_first, + f"[{backend}] adapter B changed after eviction + overlapped reload", + ) + + def test_overlap_loading_with_eviction(self): + self._check_eviction_with_overlap("triton") + + def test_overlap_loading_with_eviction_csgmv(self): + self._check_eviction_with_overlap("csgmv") + + +class TestPinnedAdapterWithDrainer(CustomTestCase): + """B4. Pinning x drainer. Pinning excludes an adapter from eviction victim + selection; the drainer independently chooses adapters to displace when slots + are contended. Together the drainer must not displace the pinned adapter. + TestPinnedLoRAAdapters churns a single unpinned adapter with the drainer off -- + here the drainer is active and two adapters contend for the one free slot.""" + + def test_pinned_adapter_survives_drainer(self): + with lora_runner( + lora_paths=None, + enable_lora=True, + max_lora_rank=SYNTH_MAX_LORA_RANK, + max_loras_per_batch=2, + max_loaded_loras=4, + lora_drain_wait_threshold=LORA_DRAIN_WAIT_THRESHOLD, + ) as runner: + pinned = runner.load_lora_adapter( + lora_name="pinned_a", + lora_path=synth_adapter_path(ADAPTER_A), + pinned=True, + ) + self.assertTrue(pinned.success, f"pinned load: {pinned.error_message}") + for name, identity in (("churn_b", ADAPTER_B), ("churn_c", ADAPTER_C)): + loaded = runner.load_lora_adapter( + lora_name=name, lora_path=synth_adapter_path(identity) + ) + self.assertTrue(loaded.success, f"{name}: {loaded.error_message}") + + (baseline,) = generate(runner, ["pinned_a"]) + # Two unpinned adapters contending for the single free slot is what + # engages the drainer; the pinned slot must stay untouched throughout. + for _ in range(3): + generate(runner, ["churn_b", "churn_c"], batched=True) + (after_churn,) = generate(runner, ["pinned_a"]) + empty_gpu_cache() + + print(f"[pin+drain] baseline={baseline!r} after churn={after_churn!r}") + self.assertEqual( + after_churn, + baseline, + "pinned adapter output changed after the drainer displaced adapters " + "(it may have been drained despite being pinned)", + ) + + +class TestDynamicLoadWithRadixCache(CustomTestCase): + """B5. Dynamic load/unload x radix cache. An adapter loaded at runtime reuses a + pool slot a previous adapter occupied, while radix holds prefixes keyed by + adapter. A freshly loaded adapter must not inherit the cached prefix of the one + it replaced -- the failure mode is silent (plausible text from stale KV), so it + is asserted as "B does not reproduce A's cached output on the same prompt". + TestLoRARadixCacheHit covers the same property with adapters supplied at + startup; the runtime-load path reaches it through different bookkeeping.""" + + def test_runtime_loaded_adapter_does_not_inherit_cached_prefix(self): + with lora_runner( + disable_radix=False, + lora_paths=None, + enable_lora=True, + max_lora_rank=SYNTH_MAX_LORA_RANK, + max_loras_per_batch=1, # one slot: B must reuse A's + max_loaded_loras=2, + ) as runner: + load_a = runner.load_lora_adapter( + lora_name="adapter_a", lora_path=synth_adapter_path(ADAPTER_A) + ) + self.assertTrue(load_a.success, f"load_a: {load_a.error_message}") + (a_first,) = generate(runner, ["adapter_a"]) + (a_cached,) = generate(runner, ["adapter_a"]) # radix hit + + # Load B *after* A's prefix is in the cache, so B is the first request + # to occupy the slot with a populated cache behind it. + load_b = runner.load_lora_adapter( + lora_name="adapter_b", lora_path=synth_adapter_path(ADAPTER_B) + ) + self.assertTrue(load_b.success, f"load_b: {load_b.error_message}") + (b_first,) = generate(runner, ["adapter_b"]) + + unload = runner.unload_lora_adapter(lora_name="adapter_b") + self.assertTrue(unload.success, f"unload: {unload.error_message}") + (a_after,) = generate(runner, ["adapter_a"]) + empty_gpu_cache() + + print(f"[dynload+radix] A={a_first!r} A(cached)={a_cached!r} B={b_first!r}") + self.assertEqual( + a_cached, + a_first, + "radix cache hit did not reproduce adapter A's output", + ) + self.assertNotEqual( + b_first, + a_first, + "a runtime-loaded adapter reproduced the previous adapter's cached " + "output on the same prompt (stale prefix reused across adapters)", + ) + self.assertEqual( + a_after, + a_first, + "adapter A changed after a runtime-loaded adapter reused its slot", + ) + + +if __name__ == "__main__": + unittest.main(warnings="ignore") diff --git a/test/manual/lora/test_torch_backend.py b/test/manual/lora/test_torch_backend.py index b88245a69..2cba8316e 100644 --- a/test/manual/lora/test_torch_backend.py +++ b/test/manual/lora/test_torch_backend.py @@ -103,18 +103,20 @@ class TestTorchNativeLoRABackend(CustomTestCase): self.lora_ranks, dtype=torch.int32, device=self.device ) + output_offset = torch.tensor( + [0, weight_out_dim], dtype=torch.int32, device="cpu" + ) + 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" - ), + slice_offsets=output_offset, ) - actual_output = self.backend.run_lora_b_sgemm(x, weights) + actual_output = self.backend.run_lora_b_sgemm(x, weights, output_offset) self.assertTrue(torch.allclose(actual_output, expect_output)) @@ -237,7 +239,9 @@ class TestTorchNativeLoRABackend(CustomTestCase): slice_offsets=output_offset, ) - actual_output = self.backend.run_gate_up_lora(x, gate_up_lora_a, gate_up_lora_b) + actual_output = self.backend.run_gate_up_lora( + x, gate_up_lora_a, gate_up_lora_b, output_offset + ) self.assertTrue(torch.allclose(actual_output, expect_output)) diff --git a/test/registered/lora/test_fused_moe_lora_kernel.py b/test/registered/lora/test_fused_moe_lora_kernel.py index efcc63b7a..8b63796ce 100644 --- a/test/registered/lora/test_fused_moe_lora_kernel.py +++ b/test/registered/lora/test_fused_moe_lora_kernel.py @@ -11,7 +11,8 @@ from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora # IMPORT PREBUILT KERNEL # ============================================================================== from sglang.kernels.ops.moe.moe_lora_align import moe_lora_align_block_size -from sglang.srt.utils import set_random_seed +from sglang.srt.lora.lora_moe_runners import _naive_moe_lora_align_block_size +from sglang.srt.utils import get_device, is_xpu, set_random_seed from sglang.test.ci.ci_register import register_cuda_ci # ============================================================================== @@ -160,23 +161,40 @@ def use_fused_moe_lora_kernel( adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=device) lora_ids = torch.arange(max_loras, dtype=torch.int32, device=device) - # call kernel - moe_lora_align_block_size( - topk_ids, - seg_indptr, - req_to_lora, - num_experts, - block_size, - max_loras, - max_num_tokens_padded, - max_num_m_blocks, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - adapter_enabled, - lora_ids, - None, # maybe_expert_map - ) + # call kernel — the fused CUDA align kernel exists only for CUDA; on XPU + # use the pure-torch native alignment. + if is_xpu(): + sorted_token_ids, expert_ids, num_tokens_post_padded = ( + _naive_moe_lora_align_block_size( + topk_ids, + seg_indptr, + req_to_lora, + num_experts, + block_size, + max_loras, + max_num_tokens_padded, + max_num_m_blocks, + adapter_enabled, + device, + ) + ) + else: + moe_lora_align_block_size( + topk_ids, + seg_indptr, + req_to_lora, + num_experts, + block_size, + max_loras, + max_num_tokens_padded, + max_num_m_blocks, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + adapter_enabled, + lora_ids, + None, # maybe_expert_map + ) config = { "BLOCK_SIZE_M": 16, @@ -264,7 +282,7 @@ def use_torch( DTYPES = [torch.float32, torch.float16, torch.bfloat16] -DEVICES = [f"cuda:{0}"] +DEVICES = [get_device(0)] SEED = [42] diff --git a/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py b/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py index 2986f63d8..7217d57f7 100644 --- a/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py +++ b/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py @@ -284,7 +284,29 @@ REFERENCE_STATS = { class TestMoELoraRegression(unittest.TestCase): - def test_sglang_moe_parity_strict(self): + def test_sglang_moe_parity_flashinfer(self): + # flashinfer is CUDA-only. + from sglang.srt.utils import is_cuda + + if not is_cuda(): + self.skipTest("flashinfer backend requires CUDA") + self._run_parity_strict(attention_backend="flashinfer") + + def test_sglang_moe_parity_intel_xpu(self): + from sglang.srt.utils import is_xpu + + if not is_xpu(): + self.skipTest("intel_xpu backend requires XPU") + self._run_parity_strict(attention_backend="intel_xpu") + + def test_sglang_moe_parity_triton(self): + from sglang.srt.utils import is_cuda, is_xpu + + if is_cuda() or is_xpu(): + self.skipTest("triton backend is the fallback for non-CUDA/non-XPU") + self._run_parity_strict(attention_backend="triton") + + def _run_parity_strict(self, *, attention_backend, **runner_kwargs): with SRTRunner( model_path=MOE_BASE_MODEL_PATH, @@ -295,8 +317,9 @@ class TestMoELoraRegression(unittest.TestCase): tp_size=1, trust_remote_code=True, disable_radix_cache=True, - attention_backend="flashinfer", + attention_backend=attention_backend, mem_fraction_static=0.80, + **runner_kwargs, ) as srt_runner: srt_outputs = srt_runner.forward( MOE_LORA_TEST_PROMPTS, diff --git a/test/registered/lora/test_lora_overlap_loading.py b/test/registered/lora/test_lora_overlap_loading.py index aa3f2c51e..e88e5c6b7 100644 --- a/test/registered/lora/test_lora_overlap_loading.py +++ b/test/registered/lora/test_lora_overlap_loading.py @@ -23,6 +23,7 @@ from torch.cuda import Stream as CudaStream from sglang.srt.lora.lora_manager import LoRAManager from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader, LoRAOverlapLoadStatus +from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.lora_utils import ( CI_MULTI_LORA_MODELS, @@ -33,6 +34,8 @@ from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=237, stage="base-b", runner_config="1-gpu-large") register_amd_ci(est_time=75, suite="stage-b-test-1-gpu-small-amd") +DEVICE = get_device(0) + class TestLoRAOverlapLoading(CustomTestCase): def test_ci_lora_models_batch_splitting(self): @@ -56,6 +59,8 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase): self.mock_stream = MagicMock(spec=CudaStream) self.mock_stream_context = MagicMock() self.mock_event = MagicMock(spec=CudaEvent) + # Default to an in-flight copy; cases needing completion flip query(). + self.mock_event.query.return_value = False self.mock_device_module.Stream.return_value = self.mock_stream self.mock_device_module.stream.return_value = self.mock_stream_context @@ -64,7 +69,7 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase): self.mock_torch.cuda.current_stream.return_value = MagicMock(spec=CudaStream) self.mock_lora_manager = MagicMock(spec=LoRAManager) - self.mock_lora_manager.device = "cuda:0" + self.mock_lora_manager.device = DEVICE self.mock_lora_manager.memory_pool = MagicMock() self.mock_lora_manager.memory_pool.uid_to_buffer_id = {} self.mock_lora_manager.validate_lora_batch.return_value = True @@ -166,7 +171,7 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase): def test_pending_load_is_synchronized_before_unload(self): manager = LoRAManager.__new__(LoRAManager) - manager.device = torch.device("cuda:0") + manager.device = torch.device(DEVICE) manager.pending_lora_load_events = {} manager.memory_pool = MagicMock() manager.configs = {"lora_A": object()} diff --git a/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py index 9e5809ae5..078c92823 100644 --- a/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py @@ -50,6 +50,25 @@ DECODE_ATTENTION_BACKEND = "fa4" KL_THRESHOLD = 5e-3 +def attention_backend_kwargs(): + """Engine attention-backend kwargs for the current platform. + + fa4 and flashinfer are CUDA-only: fa4 dispatches into the CUTLASS CUTE DSL + kernel, which cannot import off CUDA. On XPU the equivalent fused path is + the intel_xpu backend, so select it there instead of forcing a backend the + device has no kernels for. + """ + from sglang.srt.utils import is_xpu + + if is_xpu(): + return {"attention_backend": "intel_xpu"} + return { + "attention_backend": "flashinfer", + "prefill_attention_backend": PREFILL_ATTENTION_BACKEND, + "decode_attention_backend": DECODE_ATTENTION_BACKEND, + } + + def kl_v2(a, b): a = torch.tensor(a) if not torch.is_tensor(a) else a b = torch.tensor(b) if not torch.is_tensor(b) else b @@ -137,9 +156,7 @@ class TestLoRAQwen3_8BLogprobDiff(CustomTestCase): max_lora_rank=MAX_LORA_RANK, lora_paths={"my_lora": adapter_path}, lora_backend=LORA_BACKEND, - attention_backend="flashinfer", - prefill_attention_backend=PREFILL_ATTENTION_BACKEND, - decode_attention_backend=DECODE_ATTENTION_BACKEND, + **attention_backend_kwargs(), ) try: diff --git a/test/registered/lora/test_moe_lora_info.py b/test/registered/lora/test_moe_lora_info.py index fbd3a49d7..478c6b8f4 100644 --- a/test/registered/lora/test_moe_lora_info.py +++ b/test/registered/lora/test_moe_lora_info.py @@ -4,10 +4,18 @@ import pytest import torch from sglang.srt.lora.backend.base_backend import _compute_moe_lora_info -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import ( + register_amd_ci, + register_cuda_ci, + register_xpu_ci, +) register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd") +register_xpu_ci(est_time=20, suite="stage-a-test-1-gpu-xpu") + +DEVICE = get_device() def _expected_adapter_enabled( @@ -25,7 +33,7 @@ def _expected_adapter_enabled( @pytest.mark.parametrize("use_preallocated_buffers", [False, True]) def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool): - device = "cuda" + device = DEVICE seg_lens = torch.tensor([5, 1, 7, 3, 9, 2], dtype=torch.int32, device=device) seg_indptr = torch.zeros((seg_lens.numel() + 1,), dtype=torch.int32, device=device) seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) @@ -54,7 +62,7 @@ def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool): token_lora_mapping, max_len=int(seg_lens.max().item()), ) - torch.cuda.synchronize() + torch.get_device_module(device).synchronize() expected_mapping = torch.repeat_interleave(weight_indices, seg_lens) expected_enabled = _expected_adapter_enabled(lora_ranks, weight_indices) @@ -67,7 +75,7 @@ def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool): def test_compute_moe_lora_info_rejects_undercovered_launch(): - device = "cuda" + device = DEVICE seg_indptr = torch.tensor([0, 300], dtype=torch.int32, device=device) weight_indices = torch.tensor([0], dtype=torch.int32, device=device) lora_ranks = torch.tensor([16], dtype=torch.int32, device=device) diff --git a/test/registered/lora/test_virtual_experts_kernels.py b/test/registered/lora/test_virtual_experts_kernels.py index 4db23391b..e6c8d6596 100644 --- a/test/registered/lora/test_virtual_experts_kernels.py +++ b/test/registered/lora/test_virtual_experts_kernels.py @@ -15,9 +15,10 @@ Covers two regression bugs that surface only with `--lora-use-virtual-experts` wrap, or past-end) and don't get assigned to a real expert in the consumer-block table. -Both kernels run on CUDA. The fallback is gated on `virtual_num_experts >= 1024` -in production, but we exercise it directly here at smaller sizes for cheaper -iteration; one test sticks to the >1024 regime to mirror the production trigger. +Both kernels run on CUDA or XPU; only the CUDA-JIT align variant is CUDA-only. +The fallback is gated on `virtual_num_experts >= 1024` in production, but we +exercise it directly here at smaller sizes for cheaper iteration; one test +sticks to the >1024 regime to mirror the production trigger. Usage: python -m pytest test/registered/lora/test_virtual_experts_kernels.py -v @@ -27,10 +28,19 @@ import unittest import torch -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import get_device, is_cuda, is_xpu +from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=14, stage="base-b", runner_config="1-gpu-small") +register_xpu_ci(est_time=20, suite="stage-a-test-1-gpu-xpu") + + +def _require_accelerator_device(): + if not (is_cuda() or is_xpu()): + raise unittest.SkipTest("CUDA or XPU required") + return get_device(0) + from sglang.kernels.ops.moe.virtual_experts import ( _align_block_size_jit, @@ -45,9 +55,7 @@ class TestFusedVirtualTopkIdsPreservesSentinels(CustomTestCase): @classmethod def setUpClass(cls): - if not torch.cuda.is_available(): - raise unittest.SkipTest("CUDA required") - cls.device = "cuda:0" + cls.device = _require_accelerator_device() def test_negative_sentinels_preserved(self): # Mix of valid topk_ids in [0, num_experts), -1 sentinels (typical @@ -139,11 +147,9 @@ class _AlignBlockSizeSentinelBucketBase(CustomTestCase): @classmethod def setUpClass(cls): - if not torch.cuda.is_available(): - raise unittest.SkipTest("CUDA required") if cls is _AlignBlockSizeSentinelBucketBase: raise unittest.SkipTest("Base class") - cls.device = "cuda:0" + cls.device = _require_accelerator_device() def _align(self, topk_ids, block_size, num_experts): raise NotImplementedError @@ -264,6 +270,13 @@ class TestAlignBlockSizeTorchSentinelBucket(_AlignBlockSizeSentinelBucketBase): return _align_block_size_torch(topk_ids, block_size, num_experts) +@unittest.skipIf( + is_xpu(), + "_align_block_size_jit builds a CUDA JIT kernel via tvm_ffi.load_inline " + "(requires a CUDA/nvcc install); it cannot run on XPU. The torch variant " + "(TestAlignBlockSizeTorchSentinelBucket) covers the same alignment logic " + "on that platform.", +) class TestAlignBlockSizeJitSentinelBucket(_AlignBlockSizeSentinelBucketBase): """Test the CUDA JIT kernel path (with fused_sanitize_expert_ids, as in production)."""