diff --git a/python/sglang/srt/lora/backend/ascend_backend.py b/python/sglang/srt/lora/backend/ascend_backend.py index 8141f1354..1c2464f43 100644 --- a/python/sglang/srt/lora/backend/ascend_backend.py +++ b/python/sglang/srt/lora/backend/ascend_backend.py @@ -211,6 +211,7 @@ class AscendLoRABackend(BaseLoRABackend): lora_ranks: list[int], scalings: list[float], use_cuda_graph: bool, + use_prefill_cuda_graph: bool = False, ): # Use pinned memory to avoid synchronizations during host-to-device transfer weight_indices_tensor = torch.tensor( diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 16879b547..5f719b5d9 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -19,11 +19,21 @@ class BaseLoRABackend(LoRABackendLmHeadMixing): device: the device where the backend runs. """ + # Supporting backends implement init_prefill_cuda_graph_batch_info() and + # honor use_prefill_cuda_graph in prepare_lora_batch(). + supports_prefill_cuda_graph: bool = False + def __init__(self, max_loras_per_batch: int, device: torch.device): self.max_loras_per_batch = max_loras_per_batch self.device = device self.init_lm_head_config() self._is_moe_lora = False + # Static metadata read by prefill-CUDA-graph kernels, refreshed in + # place every prefill batch. + self.prefill_cuda_graph_batch_info: LoRABatchInfo | None = None + # Request/token caps for serving a batch from the static metadata. + self.prefill_cuda_graph_max_bs: int | None = None + self.prefill_cuda_graph_max_tokens: int | None = None def run_lora_a_embedding( self, @@ -161,6 +171,13 @@ class BaseLoRABackend(LoRABackendLmHeadMixing): """ pass + def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int): + """Allocate static LoRA batch metadata for the prefill CUDA graph, + sized for the largest captured token bucket. Called before capture.""" + raise NotImplementedError( + f"LoRA backend {type(self).__name__} does not support the prefill CUDA graph." + ) + @property def is_moe_lora(self) -> bool: return self._is_moe_lora @@ -317,18 +334,12 @@ class BaseLoRABackend(LoRABackendLmHeadMixing): lora_ranks: list[int], scalings: list[float], use_cuda_graph: bool, + use_prefill_cuda_graph: bool = False, ): """Prepare the lora weights and batch info for current forward batch. - This method provides a hook for each backend to conduct its own preparation - logic for each forward batch. - - Args: - forward_batch: the ForwardBatch object for current forward pass - weight_indices: list of indices of lora weights to be applied for current batch - lora_ranks: list of lora ranks corresponding to weight_indices - scalings: list of scaling factors corresponding to weight_indices - use_cuda_graph: whether to use CUDA Graph for this batch + use_cuda_graph / use_prefill_cuda_graph select in-place updates of the + static decode / prefill CUDA graph batch info respectively. """ pass diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 3ca1a88f5..8e525cfc1 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -32,6 +32,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): """ name = "csgmv" + supports_prefill_cuda_graph = True def __init__( self, @@ -240,6 +241,34 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): req_weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32), ) + def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int): + # Worst-case chunk segments for any replay batch: ceil(N / chunk_top) + # (bounded by 16 for the small tiers) plus one per adapter group. + chunk_top = self._determine_chunk_size_for_tokens(max_num_tokens) + max_num_segments = ( + max((max_num_tokens + chunk_top - 1) // chunk_top, 16) + + self.max_loras_per_batch + ) + # Each extend request has >= 1 token, so bs is bounded by the bucket. + max_bs = max_num_tokens + with torch.device(self.device): + self.prefill_cuda_graph_batch_info = LoRABatchInfo( + bs=0, # Set per batch + use_cuda_graph=True, + seg_lens=torch.zeros(max_num_segments, dtype=torch.int32), + seg_indptr=torch.zeros(max_num_segments + 1, dtype=torch.int32), + weight_indices=torch.zeros(max_num_segments, dtype=torch.int32), + permutation=torch.zeros(max_num_tokens, dtype=torch.int32), + lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32), + scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float), + num_segments=None, # Set per batch + max_len=None, # Set per batch (chunk size) + req_seg_indptr=torch.zeros(max_bs + 1, dtype=torch.int32), + req_weight_indices=torch.zeros(max_bs, dtype=torch.int32), + ) + self.prefill_cuda_graph_max_bs = max_bs + self.prefill_cuda_graph_max_tokens = max_num_tokens + def prepare_lora_batch( self, forward_batch: ForwardBatch, @@ -247,6 +276,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): lora_ranks: list[int], scalings: list[float], use_cuda_graph: bool, + use_prefill_cuda_graph: bool = False, ): chunk_size = self._determine_chunk_size(forward_batch) @@ -276,7 +306,16 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): max_num_segments = 0 has_unused_cuda_graph_segments = False - if not use_cuda_graph: + if use_prefill_cuda_graph: + # In-place refresh of the prefill-sized static buffers; unused + # segment slots no-op via the padded seg_indptr tail. + batch_info = self.prefill_cuda_graph_batch_info + batch_info.bs = bs + batch_info.num_segments = num_segments + batch_info.max_len = chunk_size + max_num_segments = batch_info.weight_indices.shape[0] + has_unused_cuda_graph_segments = num_segments < max_num_segments + elif not use_cuda_graph: batch_info = LoRABatchInfo( bs=bs, num_segments=num_segments, @@ -409,6 +448,8 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): return dataclasses.replace( batch_info, + # lm_head LoRA runs in the eager tail, outside any captured graph. + use_cuda_graph=False, num_segments=num_segments, max_len=chunk_size, seg_indptr=seg_indptr, diff --git a/python/sglang/srt/lora/backend/torch_backend.py b/python/sglang/srt/lora/backend/torch_backend.py index e53904112..9519c17cf 100644 --- a/python/sglang/srt/lora/backend/torch_backend.py +++ b/python/sglang/srt/lora/backend/torch_backend.py @@ -197,6 +197,7 @@ class TorchNativeLoRABackend(BaseLoRABackend): lora_ranks: list[int], scalings: list[float], use_cuda_graph: bool, + use_prefill_cuda_graph: bool = False, ): # Do not use merge optimization for graph mode # Use pinned memory to avoid synchronizations during host-to-device transfer diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index e08fa2410..a412d4259 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -16,9 +16,14 @@ from sglang.srt.lora.utils import ( ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch +# Fixed segment slots (one per request) baked into the captured prefill LoRA +# kernel grids; batches with more requests fall back to eager prefill. +PREFILL_CUDA_GRAPH_LORA_SEGMENTS = 32 + class TritonLoRABackend(BaseLoRABackend): name = "triton" + supports_prefill_cuda_graph = True def __init__( self, @@ -181,6 +186,27 @@ class TritonLoRABackend(BaseLoRABackend): permutation=torch.zeros(max_tokens, dtype=torch.int32), ) + def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int): + num_slots = PREFILL_CUDA_GRAPH_LORA_SEGMENTS + mlpb = self.max_loras_per_batch + with torch.device(self.device): + # bs pinned at num_slots so the captured grids cover any replay + # batch; slots past the live batch keep seg_lens == 0 and no-op. + self.prefill_cuda_graph_batch_info = LoRABatchInfo( + bs=num_slots, + use_cuda_graph=True, + num_segments=num_slots, + seg_lens=torch.zeros(num_slots, dtype=torch.int32), + seg_indptr=torch.zeros(num_slots + 1, dtype=torch.int32), + max_len=0, + weight_indices=torch.zeros(num_slots, dtype=torch.int32), + lora_ranks=torch.zeros(mlpb, dtype=torch.int32), + scalings=torch.zeros(mlpb, dtype=torch.float), + permutation=None, + ) + self.prefill_cuda_graph_max_bs = num_slots + self.prefill_cuda_graph_max_tokens = max_num_tokens + def compute_sgemm_routing(self, use_cuda_graph: bool): """Sort tokens by adapter and build merged segments for sgemm LoRA.""" bi = self.batch_info @@ -231,6 +257,7 @@ class TritonLoRABackend(BaseLoRABackend): lora_ranks: list[int], scalings: list[float], use_cuda_graph: bool, + use_prefill_cuda_graph: bool = False, ): # Use pinned memory to avoid synchronizations during host-to-device transfer weight_indices_tensor = torch.tensor( @@ -252,6 +279,17 @@ class TritonLoRABackend(BaseLoRABackend): batch_info = self.cuda_graph_batch_info batch_info.bs = forward_batch.batch_size batch_info.num_segments = forward_batch.batch_size + elif use_prefill_cuda_graph: + batch_info = self.prefill_cuda_graph_batch_info + # bs stays pinned at the allocated slot count; slots past the + # live batch no-op via seg_lens == 0. + batch_info.num_segments = bs + batch_info.max_len = max(forward_batch.extend_seq_lens_cpu) + batch_info.seg_lens[:bs].copy_( + forward_batch.extend_seq_lens, non_blocking=True + ) + batch_info.seg_lens[bs:].zero_() + torch.cumsum(batch_info.seg_lens, dim=0, out=batch_info.seg_indptr[1:]) else: max_len = ( # Calculate max_len from the CPU copy to avoid D2H transfer. @@ -364,6 +402,9 @@ class TritonLoRABackend(BaseLoRABackend): return dataclasses.replace( batch_info, + # lm_head LoRA runs in the eager tail outside any captured prefill + # graph, on freshly allocated pruned metadata. + use_cuda_graph=False, bs=num_segments, num_segments=num_segments, max_len=max(seg_lens_cpu), diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 498cf4787..b26c08d81 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -89,6 +89,7 @@ class LoRAManager: self.pending_lora_load_events = {} self.eviction_policy = server_args.lora_eviction_policy + self.enable_dp_attention: bool = server_args.enable_dp_attention self._experts_shared_outer_override: Optional[bool] = ( server_args.experts_shared_outer_loras ) @@ -138,6 +139,54 @@ class LoRAManager: init_lora_two_stream_resources(self.device) # ===== END TO BE REFACTORED ==== + def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int): + """Allocate the static prefill-CUDA-graph LoRA metadata, sized by the + largest captured token bucket. Called before capture.""" + self.lora_backend.init_prefill_cuda_graph_batch_info( + max_num_tokens=max_num_tokens + ) + + @property + def supports_prefill_cuda_graph(self) -> bool: + """Whether LoRA kernels can be captured into the prefill CUDA graph; + excludes MoE LoRA and DP attention.""" + return ( + self.lora_backend.supports_prefill_cuda_graph + and not self.lora_backend.is_moe_lora + and not self.enable_dp_attention + ) + + @property + def prefill_cuda_graph_max_bs(self) -> Optional[int]: + """Request-count cap for prefill-graph LoRA batches; None until + init_prefill_cuda_graph_batch_info() ran.""" + return self.lora_backend.prefill_cuda_graph_max_bs + + def can_use_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool: + """Whether this batch can use the static prefill-graph LoRA metadata; + shared by prepare_lora_batch and can_run_graph so they stay consistent.""" + max_bs = self.lora_backend.prefill_cuda_graph_max_bs + max_tokens = self.lora_backend.prefill_cuda_graph_max_tokens + if max_bs is None or max_tokens is None: + return False + # DP attention: per-rank eligibility could diverge across ranks and + # desync collectives; keep LoRA prefill eager. + if self.enable_dp_attention: + return False + # Decode-CUDA-graph extend modes (TARGET_VERIFY, DLLM_EXTEND) are + # owned by the decode static batch info path. + if ( + not forward_batch.forward_mode.is_extend() + or forward_batch.forward_mode.is_cuda_graph() + ): + return False + if forward_batch.extend_num_tokens is None: + return False + return ( + forward_batch.batch_size <= max_bs + and forward_batch.extend_num_tokens <= max_tokens + ) + def init_cuda_graph_moe_buffers( self, max_bs: int, max_loras: int, compute_dtype, moe_layer ): @@ -374,6 +423,11 @@ class LoRAManager: and bs <= self.max_bs_in_cuda_graph and forward_batch.forward_mode.is_cuda_graph() ) + # Eligible extend batches refresh the static prefill batch info in + # place so captured kernels read current values at replay. + use_prefill_cuda_graph = not use_cuda_graph and self.can_use_prefill_cuda_graph( + forward_batch + ) weight_indices = [0] * len(forward_batch.lora_ids) lora_ranks = [0] * self.max_loras_per_batch @@ -394,6 +448,7 @@ class LoRAManager: lora_ranks=lora_ranks, scalings=scalings, use_cuda_graph=use_cuda_graph, + use_prefill_cuda_graph=use_prefill_cuda_graph, ) self.lora_backend.batch_info.has_active_lora = any( lora_ranks[wi] > 0 for wi in weight_indices diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index b9d70c6b9..fb026087e 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -180,6 +180,17 @@ def capture_prefill_graph( ) return eager_runner + if ( + model_runner.server_args.enable_lora + and not model_runner.lora_manager.supports_prefill_cuda_graph + ): + logger.warning( + "Disable prefill CUDA graph because the current LoRA " + "configuration does not support it (unsupported LoRA backend, " + "MoE LoRA, or DP attention)." + ) + return eager_runner + # Resolve the decoder once. Some VLM wrappers (for example Kimi-VL) # expose it as ``language_model`` rather than ``model``. try: diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 7f8f108cf..d4ecd4564 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -167,6 +167,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # --- model flags ---------------------------------------------- self.quant_config = getattr(model_runner.model, "quant_config", None) self.is_multimodal = model_runner.model_config.is_multimodal + self.enable_lora = model_runner.server_args.enable_lora # Classification/reward forwards branch on return_pooled_hidden_states; # capture must use the same flag value as replay for those models. self.capture_return_pooled_hidden_states = not model_runner.is_generation @@ -267,6 +268,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # default False; the assignment below sets the real value once the # backend type is known. self._is_full_backend = False + # Same ordering requirement: capture_prepare reads this. + self._capture_lora = False # TcPiecewise does its compile pass during backend construction. # Wrap only that path with the prefill CUDA graph failure hint. try: @@ -286,6 +289,30 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # Auto: scale request slots with the chunked prefill size. max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1) self._capture_req_slots = min(max_req, self.max_bs) + + # BCG/Full record LoRA kernels, so the metadata they read must live in + # static buffers refreshed in place per batch; unsupported LoRA + # configs were already routed to the eager runner. + self._capture_lora = self.enable_lora and isinstance( + self.backend, (BreakableCudaGraphBackend, FullCudaGraphBackend) + ) + if self._capture_lora: + model_runner.lora_manager.init_prefill_cuda_graph_batch_info( + max_num_tokens=self.max_num_tokens + ) + # Clamp Full's request slots to the LoRA segment-slot count + # rather than fail capture. + lora_max_bs = model_runner.lora_manager.prefill_cuda_graph_max_bs + if self._capture_req_slots > lora_max_bs: + logger.info( + "Clamping full prefill CUDA graph request slots from %d to %d " + "to fit the LoRA backend's static segment slots.", + self._capture_req_slots, + lora_max_bs, + ) + self._capture_req_slots = lora_max_bs + + if self._is_full_backend: self._full_cg_seq_lens_cpu = torch.zeros( (self._capture_req_slots,), dtype=torch.int64, device="cpu" ) @@ -681,6 +708,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): def can_run_graph(self, forward_batch: ForwardBatch) -> bool: if self._is_full_backend and forward_batch.batch_size > self._capture_req_slots: return False + # LoRA batches may only replay the graph when prepare_lora_batch put + # their metadata in the static buffers (same predicate); keyed off + # enable_lora, not lora_ids, which is non-None even without LoRA. + if self.enable_lora and not ( + self._capture_lora + and self.model_runner.lora_manager.can_use_prefill_cuda_graph(forward_batch) + ): + return False if forward_batch.input_embeds is not None: return False if forward_batch.replace_embeds is not None: @@ -850,7 +885,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): num_token_non_padded=self._capture_num_token_non_padded(num_tokens), num_token_non_padded_cpu=num_tokens, global_forward_mode=ForwardMode.EXTEND, - lora_ids=None, + # All-None ids are safe: kernels no-op at rank 0 and replay + # refreshes the static batch info with live values. + lora_ids=([None] * bs if self._capture_lora else None), return_pooled_hidden_states=self.capture_return_pooled_hidden_states, ) self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens) @@ -895,6 +932,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): """ num_tokens = size forward_batch, attn_backend = self.capture_prepare(num_tokens) + if forward_batch.lora_ids is not None: + # Fill the static prefill LoRA batch info the captured kernels + # will read (all-None ids: ranks stay 0, kernels no-op). + lora_manager = self.model_runner.lora_manager + assert lora_manager.can_use_prefill_cuda_graph(forward_batch), ( + f"Capture batch (req slots {self._capture_req_slots}, bucket " + f"{num_tokens}) exceeds the LoRA backend's prefill CUDA graph " + "limits; the graph would read stale LoRA metadata at replay." + ) + lora_manager.prepare_lora_batch(forward_batch) if self._is_full_backend: attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4ec37bf78..73feaf1bc 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4280,6 +4280,8 @@ class ServerArgs: "MoE A2A backend", lambda: _resolved_view(self).moe_a2a_backend != "none", ), + # Dynamo blocks LoRA under tc_piecewise (per-batch LoRABatchInfo + # rebinds break guards); breakable/full support LoRA. ("LoRA", lambda: bool(self.lora_paths) or self.enable_lora), ( "multimodal model", @@ -4355,8 +4357,6 @@ class ServerArgs: "decode context parallel (dcp_size > 1)", lambda: self.dcp_size > 1, ), - # BCG capture + LoRA adapter weights exceed host RAM headroom. - ("LoRA", lambda: bool(self.lora_paths) or bool(self.enable_lora)), # BCG bucket sizes exceed FlashInfer MoE A2A's dispatch cap. ( "MoE A2A backend", diff --git a/test/registered/cuda_graph/breakable/test_bcg_with_lora.py b/test/registered/cuda_graph/breakable/test_bcg_with_lora.py new file mode 100644 index 000000000..370d04c0c --- /dev/null +++ b/test/registered/cuda_graph/breakable/test_bcg_with_lora.py @@ -0,0 +1,305 @@ +"""LoRA under the breakable (BCG) prefill CUDA graph. + +For csgmv and triton, launches the same LoRA-enabled server with the prefill +graph breakable vs disabled and compares per-token prompt logprobs for LoRA, +base, and mixed batches. Guards against vacuous passes by scraping the logs +for graph-served prefill batches (including a window where a lone LoRA +request is the only work in flight) and asserting the adapter moves logprobs. +""" + +import os +import re +import shutil +import tempfile +import time +import unittest + +import requests +import torch + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=480, stage="base-b", runner_config="1-gpu-large") + +BASE_MODEL = "Qwen/Qwen3-4B" +LORA_ADAPTER = "nissenj/Qwen3-4B-lora-v2" +LORA_NAME = "lora0" + +# Both servers run with --enable-deterministic-inference (batch-invariant +# kernels), so logprobs must not depend on how the scheduler composes +# prefill batches — the graph-vs-eager comparison is exact by construction. +# Measured diff is 0 on H200 for both csgmv and triton; 1e-2 is headroom. +GRAPH_VS_EAGER_TOLERANCE = 1e-2 +# The adapter must move at least one prompt logprob by this much. +LORA_EFFECT_THRESHOLD = 5e-2 + +# Prompt lengths chosen to land in different captured token buckets. +PROMPTS = [ + "What is the capital of France?", + "List three benefits of regular exercise. " * 8, + "The quick brown fox jumps over the lazy dog. " * 40, +] + +MIXED_BATCH_PROMPTS = [ + "Summarize the plot of Romeo and Juliet.", + "Explain how photosynthesis works. " * 6, + "Write a haiku about the sea.", + "Describe the water cycle in detail. " * 12, +] +MIXED_BATCH_LORA_PATHS = [LORA_NAME, None, LORA_NAME, None] + +PREFILL_GRAPH_REPLAY_PATTERN = re.compile(r"Prefill batch.*cuda graph: True") + + +def _generate(text, lora_path): + resp = requests.post( + DEFAULT_URL_FOR_TEST + "/generate", + json={ + "text": text, + "sampling_params": {"max_new_tokens": 0, "temperature": 0.0}, + "return_logprob": True, + "logprob_start_len": 0, + "lora_path": lora_path, + }, + timeout=120, + ) + resp.raise_for_status() + return resp.json() + + +def _prompt_logprobs(out): + # First entry has no logprob (no preceding context). + return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:] + + +def _max_abs_diff(a, b): + ta, tb = torch.tensor(a, dtype=torch.float64), torch.tensor(b, dtype=torch.float64) + assert ta.shape == tb.shape, f"token count mismatch: {ta.shape} vs {tb.shape}" + return (ta - tb).abs().max().item() + + +class BCGLoRAServerMixin: + """Launches breakable + disabled servers for one LoRA kernel backend and + collects prompt logprobs from both. Concrete classes set lora_backend.""" + + lora_backend: str + + @classmethod + def _server_args(cls, prefill_backend): + return [ + "--enable-lora", + "--lora-paths", + f"{LORA_NAME}={LORA_ADAPTER}", + "--max-loras-per-batch", + "2", + "--lora-backend", + cls.lora_backend, + "--cuda-graph-backend-prefill", + prefill_backend, + "--cuda-graph-max-bs-prefill", + "1024", + "--cuda-graph-max-bs-decode", + "8", + "--disable-radix-cache", + "--mem-fraction-static", + "0.8", + "--random-seed", + "42", + "--enable-deterministic-inference", + ] + + @classmethod + def _collect(cls): + return { + "lora": [_prompt_logprobs(_generate(p, LORA_NAME)) for p in PROMPTS], + "base": [_prompt_logprobs(_generate(p, None)) for p in PROMPTS], + "mixed": [ + _prompt_logprobs(out) + for out in _generate(MIXED_BATCH_PROMPTS, MIXED_BATCH_LORA_PATHS) + ], + } + + @classmethod + def _log_sizes(cls): + return [os.path.getsize(path) for path in (cls.stdout_path, cls.stderr_path)] + + @classmethod + def _logs_appended_since(cls, sizes): + # Binary read: the size snapshot is a byte offset. + chunks = [] + for path, start in zip((cls.stdout_path, cls.stderr_path), sizes): + with open(path, "rb") as f: + f.seek(start) + chunks.append(f.read().decode(errors="replace")) + return "\n".join(chunks) + + @classmethod + def _probe_lora_replay_window(cls): + """Send one adapter-carrying request with nothing else in flight and + return the log text appended while it ran; any "Prefill batch" line + in the window belongs to that request.""" + # Baseline once the logs go quiet, so an earlier batch's line cannot + # drain into the window (the pipe-dump threads flush asynchronously). + sizes = cls._log_sizes() + quiet_deadline = time.monotonic() + 10 + while time.monotonic() < quiet_deadline: + time.sleep(1.0) + new_sizes = cls._log_sizes() + if new_sizes == sizes: + break + sizes = new_sizes + + _generate(PROMPTS[0], LORA_NAME) + deadline = time.monotonic() + 30 + while True: + window = cls._logs_appended_since(sizes) + if PREFILL_GRAPH_REPLAY_PATTERN.search(window) or ( + time.monotonic() > deadline + ): + return window + time.sleep(0.5) + + @classmethod + def setUpClass(cls): + cls.log_dir = tempfile.mkdtemp(prefix=f"bcg_lora_{cls.lora_backend}_") + cls.stdout_path = os.path.join(cls.log_dir, "server.out") + cls.stderr_path = os.path.join(cls.log_dir, "server.err") + cls.stdout = open(cls.stdout_path, "w") + cls.stderr = open(cls.stderr_path, "w") + process = popen_launch_server( + BASE_MODEL, + DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls._server_args("breakable"), + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + try: + cls.with_graph = cls._collect() + cls.lora_replay_window = cls._probe_lora_replay_window() + finally: + kill_process_tree(process.pid) + cls.stdout.close() + cls.stderr.close() + + process = popen_launch_server( + BASE_MODEL, + DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls._server_args("disabled"), + ) + try: + cls.without_graph = cls._collect() + finally: + kill_process_tree(process.pid) + + @classmethod + def tearDownClass(cls): + for attr in ("stdout", "stderr"): + f = getattr(cls, attr, None) + if f is not None and not f.closed: + f.close() + log_dir = getattr(cls, "log_dir", None) + if log_dir: + shutil.rmtree(log_dir, ignore_errors=True) + + def test_prefill_graph_replay_logged(self): + # "cuda graph: True" is logged only when PrefillCudaGraphRunner + # served the batch; catches an across-the-board eager fallback. + logs = "" + for path in (self.stdout_path, self.stderr_path): + with open(path) as f: + logs += f.read() + self.assertRegex( + logs, + PREFILL_GRAPH_REPLAY_PATTERN, + "No prefill batch was served by the prefill CUDA graph; the " + "breakable-vs-disabled comparison is vacuous.", + ) + + def test_lora_request_replays_prefill_graph(self): + # The probe window held a lone adapter-carrying request, so this + # fails if LoRA batches specifically fall back to eager prefill. + prefill_lines = [ + line + for line in self.lora_replay_window.splitlines() + if "Prefill batch" in line + ] + self.assertTrue( + any(PREFILL_GRAPH_REPLAY_PATTERN.search(line) for line in prefill_lines), + "LoRA probe request was not served by the prefill CUDA graph. " + f"Prefill lines in the probe window: {prefill_lines}", + ) + + def test_lora_prefill_logprobs_match_eager(self): + for i, prompt in enumerate(PROMPTS): + diff = _max_abs_diff( + self.with_graph["lora"][i], self.without_graph["lora"][i] + ) + print(f"[lora] prompt {i} ({len(prompt)} chars): max_abs_diff={diff:.5f}") + self.assertLess( + diff, + GRAPH_VS_EAGER_TOLERANCE, + f"LoRA prompt logprobs diverge between breakable prefill graph " + f"and eager prefill for prompt {i}", + ) + + def test_base_prefill_logprobs_match_eager(self): + for i, prompt in enumerate(PROMPTS): + diff = _max_abs_diff( + self.with_graph["base"][i], self.without_graph["base"][i] + ) + print(f"[base] prompt {i} ({len(prompt)} chars): max_abs_diff={diff:.5f}") + self.assertLess( + diff, + GRAPH_VS_EAGER_TOLERANCE, + f"Base-model prompt logprobs diverge between breakable prefill " + f"graph and eager prefill for prompt {i}", + ) + + def test_mixed_batch_matches_eager(self): + for i in range(len(MIXED_BATCH_PROMPTS)): + diff = _max_abs_diff( + self.with_graph["mixed"][i], self.without_graph["mixed"][i] + ) + print( + f"[mixed] req {i} (lora={MIXED_BATCH_LORA_PATHS[i]}): " + f"max_abs_diff={diff:.5f}" + ) + self.assertLess( + diff, + GRAPH_VS_EAGER_TOLERANCE, + f"Mixed LoRA/base batch logprobs diverge between breakable " + f"prefill graph and eager prefill for request {i}", + ) + + def test_lora_is_applied_under_graph(self): + # If the captured graph dropped the LoRA kernels (or read stale + # metadata), LoRA and base logprobs would coincide. + for i in range(len(PROMPTS)): + diff = _max_abs_diff(self.with_graph["lora"][i], self.with_graph["base"][i]) + print(f"[effect] prompt {i}: lora-vs-base max_abs_diff={diff:.5f}") + self.assertGreater( + diff, + LORA_EFFECT_THRESHOLD, + f"Adapter did not change prompt logprobs under the breakable " + f"prefill graph for prompt {i}; LoRA was likely not applied.", + ) + + +class TestBCGLoRACsgmv(BCGLoRAServerMixin, CustomTestCase): + lora_backend = "csgmv" + + +class TestBCGLoRATriton(BCGLoRAServerMixin, CustomTestCase): + lora_backend = "triton" + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py index d616c3b42..275da20fd 100644 --- a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py +++ b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py @@ -30,6 +30,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase): def _make_prefill_runner(self, backend): runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) runner._is_full_backend = False + runner.enable_lora = False runner.prefill_backend_name = backend runner.has_mha_companion_layers = backend == Backend.BREAKABLE runner.capture_hidden_mode = CaptureHiddenMode.NULL diff --git a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py index 2e0baefd4..5369cc7c3 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py +++ b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py @@ -19,6 +19,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase): def _make_runner(self): runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) runner._is_full_backend = False + runner.enable_lora = False runner.prefill_backend_name = Backend.TC_PIECEWISE runner.has_mha_companion_layers = False runner.capture_hidden_mode = CaptureHiddenMode.NULL diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 677e13b80..856b986d0 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1296,6 +1296,60 @@ class TestCudaGraphDisaggregationRoles(CustomTestCase): self.assertIn((Phase.DECODE, "backend"), args._cuda_graph_config_locked) +class TestPrefillCudaGraphLoRACompatibility(CustomTestCase): + """LoRA no longer auto-disables the breakable prefill CUDA graph; guards + test_bcg_with_lora.py against a rule re-disabling it (vacuous pass).""" + + def _handled_args(self, **overrides): + args = ServerArgs(model_path="dummy", **overrides) + args.model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]), + is_piecewise_cuda_graph_disabled_model=False, + is_multimodal=False, + is_multimodal_piecewise_cuda_graph_supported=False, + ) + with ( + patch("sglang.srt.utils.is_cuda", return_value=True), + patch.object(ServerArgs, "use_mla_backend", return_value=False), + ): + args._handle_cuda_graph_config() + return args + + def test_enable_lora_keeps_breakable_prefill_graph(self): + args = self._handled_args(enable_lora=True) + + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE) + + def test_lora_paths_keep_breakable_prefill_graph(self): + args = self._handled_args(lora_paths=["dummy/lora-adapter"]) + + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE) + + def test_lora_still_disables_tc_piecewise_prefill_graph(self): + # Pin the tc_piecewise LoRA rule itself, with the hardware rule + # neutralized so this runs on CPU-only CI. + args = ServerArgs(model_path="dummy", enable_lora=True) + args.model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]), + is_piecewise_cuda_graph_disabled_model=False, + is_multimodal=False, + is_multimodal_piecewise_cuda_graph_supported=False, + ) + args.cuda_graph_config = CudaGraphConfig( + prefill=PhaseConfig(backend=Backend.TC_PIECEWISE) + ) + with ( + patch("sglang.srt.server_args.is_hip", return_value=False), + patch("sglang.srt.server_args.is_npu", return_value=False), + patch("sglang.srt.server_args.is_cpu", return_value=False), + patch("sglang.srt.server_args.is_mps", return_value=False), + patch("sglang.srt.server_args.is_xpu", return_value=False), + ): + args._disable_tc_piecewise_cudagraph_if_incompatible() + + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED) + + class TestBreakableCudaGraphMultimodalAllowlist(CustomTestCase): """The BCG "multimodal model" rule exempts archs on the BCG multimodal opt-in allowlist (multimodal_breakable_cuda_graph_supported_model_archs)."""