[LoRA] Support LoRA under the breakable/full prefill CUDA graph (#30988)
This commit is contained in:
@@ -211,6 +211,7 @@ class AscendLoRABackend(BaseLoRABackend):
|
|||||||
lora_ranks: list[int],
|
lora_ranks: list[int],
|
||||||
scalings: list[float],
|
scalings: list[float],
|
||||||
use_cuda_graph: bool,
|
use_cuda_graph: bool,
|
||||||
|
use_prefill_cuda_graph: bool = False,
|
||||||
):
|
):
|
||||||
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
||||||
weight_indices_tensor = torch.tensor(
|
weight_indices_tensor = torch.tensor(
|
||||||
|
|||||||
@@ -19,11 +19,21 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
|
|||||||
device: the device where the backend runs.
|
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):
|
def __init__(self, max_loras_per_batch: int, device: torch.device):
|
||||||
self.max_loras_per_batch = max_loras_per_batch
|
self.max_loras_per_batch = max_loras_per_batch
|
||||||
self.device = device
|
self.device = device
|
||||||
self.init_lm_head_config()
|
self.init_lm_head_config()
|
||||||
self._is_moe_lora = False
|
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(
|
def run_lora_a_embedding(
|
||||||
self,
|
self,
|
||||||
@@ -161,6 +171,13 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
|
|||||||
"""
|
"""
|
||||||
pass
|
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
|
@property
|
||||||
def is_moe_lora(self) -> bool:
|
def is_moe_lora(self) -> bool:
|
||||||
return self._is_moe_lora
|
return self._is_moe_lora
|
||||||
@@ -317,18 +334,12 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
|
|||||||
lora_ranks: list[int],
|
lora_ranks: list[int],
|
||||||
scalings: list[float],
|
scalings: list[float],
|
||||||
use_cuda_graph: bool,
|
use_cuda_graph: bool,
|
||||||
|
use_prefill_cuda_graph: bool = False,
|
||||||
):
|
):
|
||||||
"""Prepare the lora weights and batch info for current forward batch.
|
"""Prepare the lora weights and batch info for current forward batch.
|
||||||
|
|
||||||
This method provides a hook for each backend to conduct its own preparation
|
use_cuda_graph / use_prefill_cuda_graph select in-place updates of the
|
||||||
logic for each forward batch.
|
static decode / prefill CUDA graph batch info respectively.
|
||||||
|
|
||||||
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
|
|
||||||
"""
|
"""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
name = "csgmv"
|
name = "csgmv"
|
||||||
|
supports_prefill_cuda_graph = True
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -240,6 +241,34 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
|
|||||||
req_weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
|
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(
|
def prepare_lora_batch(
|
||||||
self,
|
self,
|
||||||
forward_batch: ForwardBatch,
|
forward_batch: ForwardBatch,
|
||||||
@@ -247,6 +276,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
|
|||||||
lora_ranks: list[int],
|
lora_ranks: list[int],
|
||||||
scalings: list[float],
|
scalings: list[float],
|
||||||
use_cuda_graph: bool,
|
use_cuda_graph: bool,
|
||||||
|
use_prefill_cuda_graph: bool = False,
|
||||||
):
|
):
|
||||||
chunk_size = self._determine_chunk_size(forward_batch)
|
chunk_size = self._determine_chunk_size(forward_batch)
|
||||||
|
|
||||||
@@ -276,7 +306,16 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
|
|||||||
max_num_segments = 0
|
max_num_segments = 0
|
||||||
has_unused_cuda_graph_segments = False
|
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(
|
batch_info = LoRABatchInfo(
|
||||||
bs=bs,
|
bs=bs,
|
||||||
num_segments=num_segments,
|
num_segments=num_segments,
|
||||||
@@ -409,6 +448,8 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
|
|||||||
|
|
||||||
return dataclasses.replace(
|
return dataclasses.replace(
|
||||||
batch_info,
|
batch_info,
|
||||||
|
# lm_head LoRA runs in the eager tail, outside any captured graph.
|
||||||
|
use_cuda_graph=False,
|
||||||
num_segments=num_segments,
|
num_segments=num_segments,
|
||||||
max_len=chunk_size,
|
max_len=chunk_size,
|
||||||
seg_indptr=seg_indptr,
|
seg_indptr=seg_indptr,
|
||||||
|
|||||||
@@ -197,6 +197,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
|
|||||||
lora_ranks: list[int],
|
lora_ranks: list[int],
|
||||||
scalings: list[float],
|
scalings: list[float],
|
||||||
use_cuda_graph: bool,
|
use_cuda_graph: bool,
|
||||||
|
use_prefill_cuda_graph: bool = False,
|
||||||
):
|
):
|
||||||
# Do not use merge optimization for graph mode
|
# Do not use merge optimization for graph mode
|
||||||
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
||||||
|
|||||||
@@ -16,9 +16,14 @@ from sglang.srt.lora.utils import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
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):
|
class TritonLoRABackend(BaseLoRABackend):
|
||||||
name = "triton"
|
name = "triton"
|
||||||
|
supports_prefill_cuda_graph = True
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -181,6 +186,27 @@ class TritonLoRABackend(BaseLoRABackend):
|
|||||||
permutation=torch.zeros(max_tokens, dtype=torch.int32),
|
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):
|
def compute_sgemm_routing(self, use_cuda_graph: bool):
|
||||||
"""Sort tokens by adapter and build merged segments for sgemm LoRA."""
|
"""Sort tokens by adapter and build merged segments for sgemm LoRA."""
|
||||||
bi = self.batch_info
|
bi = self.batch_info
|
||||||
@@ -231,6 +257,7 @@ class TritonLoRABackend(BaseLoRABackend):
|
|||||||
lora_ranks: list[int],
|
lora_ranks: list[int],
|
||||||
scalings: list[float],
|
scalings: list[float],
|
||||||
use_cuda_graph: bool,
|
use_cuda_graph: bool,
|
||||||
|
use_prefill_cuda_graph: bool = False,
|
||||||
):
|
):
|
||||||
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
# Use pinned memory to avoid synchronizations during host-to-device transfer
|
||||||
weight_indices_tensor = torch.tensor(
|
weight_indices_tensor = torch.tensor(
|
||||||
@@ -252,6 +279,17 @@ class TritonLoRABackend(BaseLoRABackend):
|
|||||||
batch_info = self.cuda_graph_batch_info
|
batch_info = self.cuda_graph_batch_info
|
||||||
batch_info.bs = forward_batch.batch_size
|
batch_info.bs = forward_batch.batch_size
|
||||||
batch_info.num_segments = 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:
|
else:
|
||||||
max_len = (
|
max_len = (
|
||||||
# Calculate max_len from the CPU copy to avoid D2H transfer.
|
# Calculate max_len from the CPU copy to avoid D2H transfer.
|
||||||
@@ -364,6 +402,9 @@ class TritonLoRABackend(BaseLoRABackend):
|
|||||||
|
|
||||||
return dataclasses.replace(
|
return dataclasses.replace(
|
||||||
batch_info,
|
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,
|
bs=num_segments,
|
||||||
num_segments=num_segments,
|
num_segments=num_segments,
|
||||||
max_len=max(seg_lens_cpu),
|
max_len=max(seg_lens_cpu),
|
||||||
|
|||||||
@@ -89,6 +89,7 @@ class LoRAManager:
|
|||||||
self.pending_lora_load_events = {}
|
self.pending_lora_load_events = {}
|
||||||
|
|
||||||
self.eviction_policy = server_args.lora_eviction_policy
|
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] = (
|
self._experts_shared_outer_override: Optional[bool] = (
|
||||||
server_args.experts_shared_outer_loras
|
server_args.experts_shared_outer_loras
|
||||||
)
|
)
|
||||||
@@ -138,6 +139,54 @@ class LoRAManager:
|
|||||||
init_lora_two_stream_resources(self.device)
|
init_lora_two_stream_resources(self.device)
|
||||||
# ===== END TO BE REFACTORED ====
|
# ===== 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(
|
def init_cuda_graph_moe_buffers(
|
||||||
self, max_bs: int, max_loras: int, compute_dtype, moe_layer
|
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 bs <= self.max_bs_in_cuda_graph
|
||||||
and forward_batch.forward_mode.is_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)
|
weight_indices = [0] * len(forward_batch.lora_ids)
|
||||||
lora_ranks = [0] * self.max_loras_per_batch
|
lora_ranks = [0] * self.max_loras_per_batch
|
||||||
@@ -394,6 +448,7 @@ class LoRAManager:
|
|||||||
lora_ranks=lora_ranks,
|
lora_ranks=lora_ranks,
|
||||||
scalings=scalings,
|
scalings=scalings,
|
||||||
use_cuda_graph=use_cuda_graph,
|
use_cuda_graph=use_cuda_graph,
|
||||||
|
use_prefill_cuda_graph=use_prefill_cuda_graph,
|
||||||
)
|
)
|
||||||
self.lora_backend.batch_info.has_active_lora = any(
|
self.lora_backend.batch_info.has_active_lora = any(
|
||||||
lora_ranks[wi] > 0 for wi in weight_indices
|
lora_ranks[wi] > 0 for wi in weight_indices
|
||||||
|
|||||||
@@ -180,6 +180,17 @@ def capture_prefill_graph(
|
|||||||
)
|
)
|
||||||
return eager_runner
|
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)
|
# Resolve the decoder once. Some VLM wrappers (for example Kimi-VL)
|
||||||
# expose it as ``language_model`` rather than ``model``.
|
# expose it as ``language_model`` rather than ``model``.
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -167,6 +167,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
# --- model flags ----------------------------------------------
|
# --- model flags ----------------------------------------------
|
||||||
self.quant_config = getattr(model_runner.model, "quant_config", None)
|
self.quant_config = getattr(model_runner.model, "quant_config", None)
|
||||||
self.is_multimodal = model_runner.model_config.is_multimodal
|
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;
|
# Classification/reward forwards branch on return_pooled_hidden_states;
|
||||||
# capture must use the same flag value as replay for those models.
|
# capture must use the same flag value as replay for those models.
|
||||||
self.capture_return_pooled_hidden_states = not model_runner.is_generation
|
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
|
# default False; the assignment below sets the real value once the
|
||||||
# backend type is known.
|
# backend type is known.
|
||||||
self._is_full_backend = False
|
self._is_full_backend = False
|
||||||
|
# Same ordering requirement: capture_prepare reads this.
|
||||||
|
self._capture_lora = False
|
||||||
# TcPiecewise does its compile pass during backend construction.
|
# TcPiecewise does its compile pass during backend construction.
|
||||||
# Wrap only that path with the prefill CUDA graph failure hint.
|
# Wrap only that path with the prefill CUDA graph failure hint.
|
||||||
try:
|
try:
|
||||||
@@ -286,6 +289,30 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
# Auto: scale request slots with the chunked prefill size.
|
# Auto: scale request slots with the chunked prefill size.
|
||||||
max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1)
|
max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1)
|
||||||
self._capture_req_slots = min(max_req, self.max_bs)
|
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._full_cg_seq_lens_cpu = torch.zeros(
|
||||||
(self._capture_req_slots,), dtype=torch.int64, device="cpu"
|
(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:
|
def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
|
||||||
if self._is_full_backend and forward_batch.batch_size > self._capture_req_slots:
|
if self._is_full_backend and forward_batch.batch_size > self._capture_req_slots:
|
||||||
return False
|
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:
|
if forward_batch.input_embeds is not None:
|
||||||
return False
|
return False
|
||||||
if forward_batch.replace_embeds is not None:
|
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=self._capture_num_token_non_padded(num_tokens),
|
||||||
num_token_non_padded_cpu=num_tokens,
|
num_token_non_padded_cpu=num_tokens,
|
||||||
global_forward_mode=ForwardMode.EXTEND,
|
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,
|
return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
|
||||||
)
|
)
|
||||||
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
|
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
|
||||||
@@ -895,6 +932,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
"""
|
"""
|
||||||
num_tokens = size
|
num_tokens = size
|
||||||
forward_batch, attn_backend = self.capture_prepare(num_tokens)
|
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:
|
if self._is_full_backend:
|
||||||
attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
|
attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -4280,6 +4280,8 @@ class ServerArgs:
|
|||||||
"MoE A2A backend",
|
"MoE A2A backend",
|
||||||
lambda: _resolved_view(self).moe_a2a_backend != "none",
|
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),
|
("LoRA", lambda: bool(self.lora_paths) or self.enable_lora),
|
||||||
(
|
(
|
||||||
"multimodal model",
|
"multimodal model",
|
||||||
@@ -4355,8 +4357,6 @@ class ServerArgs:
|
|||||||
"decode context parallel (dcp_size > 1)",
|
"decode context parallel (dcp_size > 1)",
|
||||||
lambda: self.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.
|
# BCG bucket sizes exceed FlashInfer MoE A2A's dispatch cap.
|
||||||
(
|
(
|
||||||
"MoE A2A backend",
|
"MoE A2A backend",
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -30,6 +30,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
|
|||||||
def _make_prefill_runner(self, backend):
|
def _make_prefill_runner(self, backend):
|
||||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||||
runner._is_full_backend = False
|
runner._is_full_backend = False
|
||||||
|
runner.enable_lora = False
|
||||||
runner.prefill_backend_name = backend
|
runner.prefill_backend_name = backend
|
||||||
runner.has_mha_companion_layers = backend == Backend.BREAKABLE
|
runner.has_mha_companion_layers = backend == Backend.BREAKABLE
|
||||||
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
|
|||||||
def _make_runner(self):
|
def _make_runner(self):
|
||||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||||
runner._is_full_backend = False
|
runner._is_full_backend = False
|
||||||
|
runner.enable_lora = False
|
||||||
runner.prefill_backend_name = Backend.TC_PIECEWISE
|
runner.prefill_backend_name = Backend.TC_PIECEWISE
|
||||||
runner.has_mha_companion_layers = False
|
runner.has_mha_companion_layers = False
|
||||||
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
||||||
|
|||||||
@@ -1296,6 +1296,60 @@ class TestCudaGraphDisaggregationRoles(CustomTestCase):
|
|||||||
self.assertIn((Phase.DECODE, "backend"), args._cuda_graph_config_locked)
|
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):
|
class TestBreakableCudaGraphMultimodalAllowlist(CustomTestCase):
|
||||||
"""The BCG "multimodal model" rule exempts archs on the BCG multimodal
|
"""The BCG "multimodal model" rule exempts archs on the BCG multimodal
|
||||||
opt-in allowlist (multimodal_breakable_cuda_graph_supported_model_archs)."""
|
opt-in allowlist (multimodal_breakable_cuda_graph_supported_model_archs)."""
|
||||||
|
|||||||
Reference in New Issue
Block a user