minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[LoRA] Support MoE in full and breakable prefill CUDA graphs (#38578)

Browse Source
This commit is contained in:
Yanbin Jiang
2026-09-11 14:40:02 -07:00
committed by GitHub
parent 45715e7f20
commit ec30f19e4a
9 changed files with 517 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/backend/base_backend.py
+37 -39
View File
@@ -45,6 +45,8 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
# 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
# Separate scratch sized for the largest prefill token bucket.
self.prefill_moe_cg_buffers: dict | None = None
def reset_batch_state(self):
"""Idle-forward counterpart of prepare_lora_batch(): clears all
@@ -199,7 +201,9 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
"""
pass
def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int):
def init_prefill_cuda_graph_batch_info(
self, max_num_tokens: int, max_num_requests: Optional[int] = None
):
"""Allocate static LoRA batch metadata for the prefill CUDA graph,
sized for the largest captured token bucket. Called before capture."""
raise NotImplementedError(
@@ -220,24 +224,16 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
max_loras: int,
compute_dtype: torch.dtype,
moe_layer,
*,
prefill: bool = False,
):
"""Phase 1 of LoRA CUDA graph init: MoE intermediate buffers.
"""Allocate shared MoE routing buffers for decode or prefill captures.
Called once before init_memory_pool() with a representative MoE layer
to extract dimensions. All FusedMoEWithLoRA layers share the same
buffers since they execute sequentially during forward.
This is backend-agnostic because MoE LoRA always uses the same
fused Triton kernel (TritonRunnerCoreWithLoRA) regardless of which
dense LoRA backend is selected.
max_bs counts tokens. Layers reuse these buffers sequentially.
"""
base = moe_layer.base_layer
top_k = base.top_k
qinfo = moe_layer._quant_info
E, N, _ = qinfo.w13_weight.shape
hidden_dim = qinfo.w2_weight.shape[1]
device = qinfo.w13_weight.device
dtype = compute_dtype
device = moe_layer._quant_info.w13_weight.device
num_experts = base.num_experts
block_size_m = 64
@@ -247,19 +243,7 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
) * block_size_m
max_num_m_blocks = (max_num_tokens_padded + block_size_m - 1) // block_size_m
self.moe_cg_buffers = {
"intermediate_cache1": torch.empty(
(max_bs, top_k, N), device=device, dtype=dtype
),
"intermediate_cache2": torch.empty(
(max_bs * top_k, N // 2), device=device, dtype=dtype
),
"intermediate_cache3": torch.empty(
(max_bs, top_k, hidden_dim), device=device, dtype=dtype
),
"out_hidden_states": torch.empty(
(max_bs, hidden_dim), device=device, dtype=dtype
),
buffers = {
"sorted_token_ids_lora": torch.empty(
(max_loras * max_num_tokens_padded,),
device=device,
@@ -274,12 +258,6 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
(max_loras,), device=device, dtype=torch.int32
),
"adapter_enabled": torch.zeros(max_loras, dtype=torch.int32, device=device),
# int64 copy of weight_indices for index_fill_(), which requires
# LongTensor. weight_indices itself must stay int32 because the
# CUDA moe_lora_align kernel casts it to int32_t*.
"weight_indices_long": torch.zeros(
max_bs, dtype=torch.int64, device=device
),
"lora_ids": torch.arange(max_loras, dtype=torch.int32, device=device),
"cumsum_buffer": torch.zeros(
max_loras * (num_experts + 1),
@@ -291,39 +269,56 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
dtype=torch.int32,
device=device,
),
"max_num_tokens_padded": max_num_tokens_padded,
"max_num_m_blocks": max_num_m_blocks,
"token_lora_mapping": torch.full(
(max_bs,), -1, dtype=torch.int32, device=device
),
}
if prefill:
self.prefill_moe_cg_buffers = buffers
else:
self.moe_cg_buffers = buffers
def _add_moe_lora_info(
self, forward_batch: ForwardBatch, batch_info: LoRABatchInfo
) -> LoRABatchInfo:
if not self.is_moe_lora:
return batch_info
prefill = batch_info is self.prefill_cuda_graph_batch_info
if batch_info.use_cuda_graph:
adapter_enabled = self.moe_cg_buffers["adapter_enabled"]
token_lora_mapping = self.moe_cg_buffers["token_lora_mapping"]
buffers = self.prefill_moe_cg_buffers if prefill else self.moe_cg_buffers
if prefill and buffers is None:
raise RuntimeError(
"prefill MoE-LoRA CUDA graph buffers were not initialized"
)
adapter_enabled = buffers["adapter_enabled"]
token_lora_mapping = buffers["token_lora_mapping"]
else:
adapter_enabled = None
token_lora_mapping = None
num_tokens, max_len = get_batch_token_counts(forward_batch)
# Capture fixes the segment count; include every prefill request slot.
# Unused slots contain empty segments.
if (
batch_info.req_seg_indptr is not None
or batch_info.req_weight_indices is not None
):
assert batch_info.req_seg_indptr is not None
assert batch_info.req_weight_indices is not None
num_moe_segments = batch_info.bs
num_moe_segments = (
batch_info.req_weight_indices.shape[0] if prefill else batch_info.bs
)
seg_indptr = batch_info.req_seg_indptr[: num_moe_segments + 1]
req_to_lora = batch_info.req_weight_indices[:num_moe_segments]
else:
num_moe_segments = batch_info.num_segments
num_moe_segments = (
batch_info.weight_indices.shape[0]
if prefill
else batch_info.num_segments
)
seg_indptr = batch_info.seg_indptr[: num_moe_segments + 1]
req_to_lora = batch_info.weight_indices[:num_moe_segments]
@@ -422,6 +417,9 @@ def _compute_moe_lora_info(
assert num_tokens <= token_lora_mapping.shape[0], (
"num_tokens must be less than or equal to the shape of token_lora_mapping"
)
# Clear padded replay rows left by a larger batch.
if num_tokens < token_lora_mapping.shape[0]:
token_lora_mapping[num_tokens:].fill_(-1)
token_lora_mapping = token_lora_mapping[:num_tokens]
else:
token_lora_mapping = torch.empty(
python/sglang/srt/lora/backend/chunked_backend.py
+8 -3
View File
@@ -239,7 +239,9 @@ 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):
def init_prefill_cuda_graph_batch_info(
self, max_num_tokens: int, max_num_requests: Optional[int] = None
):
# 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)
@@ -247,8 +249,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
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
max_bs = max_num_tokens if max_num_requests is None else max_num_requests
with torch.device(self.device):
self.prefill_cuda_graph_batch_info = LoRABatchInfo(
bs=0, # Set per batch
@@ -370,6 +371,10 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
batch_info.permutation[: len(permutation)].copy_(permutation, non_blocking=True)
batch_info.req_seg_indptr[: bs + 1].copy_(req_seg_indptr_cpu, non_blocking=True)
batch_info.req_weight_indices[:bs].copy_(req_wi_tensor, non_blocking=True)
if use_prefill_cuda_graph:
# Captured MoE kernels read every request slot; keep the tail empty.
batch_info.req_seg_indptr[bs + 1 :].fill_(int(req_seg_indptr_cpu[-1]))
batch_info.req_weight_indices[bs:].zero_()
batch_info = self._add_moe_lora_info(forward_batch, batch_info)
python/sglang/srt/lora/backend/triton_backend.py
+60 -6
View File
@@ -18,9 +18,8 @@ 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
# Match the dense kernels' token tile.
PREFILL_CUDA_GRAPH_LORA_CHUNK_SIZE = 16
class TritonLoRABackend(BaseLoRABackend):
@@ -55,7 +54,11 @@ class TritonLoRABackend(BaseLoRABackend):
return embedding_lora_a_fwd(
input_ids=input_ids,
weights=weights,
batch_info=self.batch_info,
batch_info=(
self._sgemm_info()
if self.batch_info is self.prefill_cuda_graph_batch_info
else self.batch_info
),
vocab_size=vocab_size,
extra_embeddings=extra_embeddings,
)
@@ -200,8 +203,10 @@ 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
def init_prefill_cuda_graph_batch_info(
self, max_num_tokens: int, max_num_requests: Optional[int] = None
):
num_slots = max_num_tokens if max_num_requests is None else max_num_requests
mlpb = self.max_loras_per_batch
with torch.device(self.device):
# bs pinned at num_slots so the captured grids cover any replay
@@ -218,6 +223,22 @@ class TritonLoRABackend(BaseLoRABackend):
scalings=torch.zeros(mlpb, dtype=torch.float),
permutation=None,
)
chunk_size = PREFILL_CUDA_GRAPH_LORA_CHUNK_SIZE
# Ragged request boundaries need up to num_slots - 1 extra tiles.
num_chunks = min(
max_num_tokens,
(max_num_tokens + chunk_size - 1) // chunk_size + num_slots - 1,
65535,
)
self.prefill_cuda_graph_sgemm_batch_info = dataclasses.replace(
self.prefill_cuda_graph_batch_info,
bs=num_chunks,
num_segments=num_chunks,
max_len=chunk_size,
seg_lens=torch.zeros(num_chunks, dtype=torch.int32),
seg_indptr=torch.zeros(num_chunks + 1, dtype=torch.int32),
weight_indices=torch.zeros(num_chunks, dtype=torch.int32),
)
self.prefill_cuda_graph_max_bs = num_slots
self.prefill_cuda_graph_max_tokens = max_num_tokens
@@ -357,6 +378,39 @@ class TritonLoRABackend(BaseLoRABackend):
self.compute_sgemm_routing(use_cuda_graph)
else:
self.sgemm_batch_info = None
if use_prefill_cuda_graph:
sgemm = self.prefill_cuda_graph_sgemm_batch_info
chunk_size = PREFILL_CUDA_GRAPH_LORA_CHUNK_SIZE
num_tokens = max(1, forward_batch.extend_num_tokens)
num_chunks = min(
num_tokens,
(num_tokens + chunk_size - 1) // chunk_size
+ self.prefill_cuda_graph_max_bs
- 1,
)
# Larger grids keep the request view to fit CUDA's y/z limit.
if num_chunks <= sgemm.seg_lens.numel():
indices, lengths = merge_and_chunk_segments(
weight_indices, forward_batch.extend_seq_lens_cpu, chunk_size
)
num_segments = len(lengths)
sgemm.bs = num_chunks
sgemm.num_segments = num_segments
sgemm.weight_indices[:num_segments].copy_(
torch.tensor(
indices, dtype=torch.int32, pin_memory=True, device="cpu"
),
non_blocking=True,
)
sgemm.seg_lens[:num_segments].copy_(
torch.tensor(
lengths, dtype=torch.int32, pin_memory=True, device="cpu"
),
non_blocking=True,
)
sgemm.seg_lens[num_segments:].zero_()
torch.cumsum(sgemm.seg_lens, dim=0, out=sgemm.seg_indptr[1:])
self.sgemm_batch_info = sgemm
self.lm_head_batch_info, self.lm_head_pass_batch_infos = (
self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info)
python/sglang/srt/lora/layers.py
+5 -1
View File
@@ -1110,7 +1110,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
lora_ranks = batch_info.lora_ranks
max_lora_rank = self.down_lora_a_weights.shape[2]
cg_buffers = getattr(self.lora_backend, "moe_cg_buffers", None)
cg_buffers = (
self.lora_backend.prefill_moe_cg_buffers
if batch_info is self.lora_backend.prefill_cuda_graph_batch_info
else getattr(self.lora_backend, "moe_cg_buffers", None)
)
moe_lora_info = batch_info.moe_lora_info
assert moe_lora_info is not None
python/sglang/srt/lora/lora_manager.py
+31 -10
View File
@@ -148,23 +148,44 @@ 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."""
def init_prefill_cuda_graph_batch_info(
self, max_num_tokens: int, max_num_requests: Optional[int] = None
):
"""Allocate static LoRA metadata and MoE scratch before prefill capture."""
self.lora_backend.init_prefill_cuda_graph_batch_info(
max_num_tokens=max_num_tokens
max_num_tokens=max_num_tokens, max_num_requests=max_num_requests
)
for module in self.base_model.modules():
if isinstance(module, FusedMoEWithLoRA):
self.lora_backend.init_cuda_graph_moe_buffers(
max_bs=max_num_tokens,
max_loras=self.max_loras_per_batch,
compute_dtype=self.dtype,
moe_layer=module,
prefill=True,
)
break
@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
"""MoE LoRA supports full and breakable capture; DP attention is unsupported."""
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
if (
self.enable_dp_attention
or not self.lora_backend.supports_prefill_cuda_graph
):
return False
if self.lora_backend.is_moe_lora:
return check_cuda_graph_backend(
Phase.PREFILL, Backend.BREAKABLE
) or check_cuda_graph_backend(Phase.PREFILL, Backend.FULL)
return True
@property
def prefill_cuda_graph_max_bs(self) -> Optional[int]:
"""Request-count cap for prefill-graph LoRA batches; None until
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+1 -1
View File
@@ -366,7 +366,7 @@ def capture_prefill_graph(
logger.warning(
"Disable prefill CUDA graph because the current LoRA "
"configuration does not support it (unsupported LoRA backend, "
"MoE LoRA, or DP attention)."
"MoE LoRA without full or breakable capture, or DP attention)."
)
return result(eager_runner)
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+15 -15
View File
@@ -42,7 +42,7 @@ import dataclasses
import inspect
import logging
from collections.abc import Sequence
from contextlib import contextmanager
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
@@ -120,6 +120,7 @@ from sglang.srt.model_executor.runner_utils import (
from sglang.srt.model_executor.runner_utils.buffers import (
PrefillInputBuffers,
)
from sglang.srt.model_executor.runner_utils.capture_mode import model_capture_mode
from sglang.srt.model_executor.runner_utils.pool import (
get_or_create_global_graph_capture_stream,
)
@@ -441,20 +442,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
if self._capture_lora:
model_runner.lora_manager.init_prefill_cuda_graph_batch_info(
max_num_tokens=self.max_num_tokens
max_num_tokens=self.max_num_tokens,
max_num_requests=(
self._capture_req_slots
if self._is_full_backend
else min(self.max_num_tokens, self.max_bs)
),
)
# 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
self._full_cg_seq_lens_cpu = (
torch.zeros((self._capture_req_slots,), dtype=torch.int64, device="cpu")
if self._is_full_backend
@@ -1505,7 +1499,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self._init_forward_metadata_for_capture(forward_batch, num_tokens)
def run_once():
return self._run_forward(forward_batch, num_tokens)
# Record LoRA kernels even when capture uses base-model requests.
with (
model_capture_mode()
if self._is_full_backend and self._capture_lora
else nullcontext()
):
return self._run_forward(forward_batch, num_tokens)
# Main's monolithic BCG runner never invokes
# on_after_cuda_graph_warmup between warmup iterations — the BCG