minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
test/registered/e2e/lora/test_lora_moe_prefill_cuda_graph.py
+189
View File
@@ -0,0 +1,189 @@
# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""MoE LoRA prefill graphs must replay adapters and match eager prefill."""
import os
import unittest
import torch
import sglang as sgl
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.lora_utils import (
MOE_BASE_MODEL_PATH,
MOE_LORA_PATH,
MOE_LORA_TEST_PROMPTS,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=600, stage="extra-a", runner_config="1-gpu-large")
# Missing adapters shift logprobs by 7-17.
LOGPROB_THRESHOLD = 1.0
MAX_NEW_TOKENS = 8
# Keep padded buckets within the runner's 2x token-count limit.
PREFILL_GRAPH_BATCH_SIZES = [512, 1024, 2048, 4096]
class TestMoELoRAPrefillCudaGraph(CustomTestCase):
def test_prefill_graph_matches_eager(self):
from prometheus_client import REGISTRY
prompts = (MOE_LORA_TEST_PROMPTS * 4)[:64]
lora_paths = [None if i % 3 == 1 else "moe_lora" for i in range(len(prompts))]
results = {}
for lora_backend, backend in (
("triton", "disabled"),
("triton", "breakable"),
("triton", "full"),
("csgmv", "breakable"),
("csgmv", "full"),
):
label = f"{lora_backend}/{backend}"
prefill_graph = backend != "disabled"
if prefill_graph:
# Isolate replay counts between engines.
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
kwargs = dict(
model_path=MOE_BASE_MODEL_PATH,
enable_lora=True,
lora_paths={"moe_lora": MOE_LORA_PATH},
max_loras_per_batch=2,
lora_backend=lora_backend,
attention_backend="flashinfer",
trust_remote_code=True,
enable_tokenizer_batch_encode=True,
enable_metrics=prefill_graph,
disable_radix_cache=True,
mem_fraction_static=0.8,
max_running_requests=len(prompts),
chunked_prefill_size=4096,
cuda_graph_max_bs_decode=4,
cuda_graph_backend_prefill=backend,
)
if prefill_graph:
kwargs["cuda_graph_bs_prefill"] = PREFILL_GRAPH_BATCH_SIZES
if backend == "full":
kwargs["cuda_graph_config"] = {
"prefill": {"full_prefill_max_req": len(prompts)}
}
collectors_before = set(REGISTRY._collector_to_names)
engine = None
try:
engine = sgl.Engine(**kwargs)
# Compare prefill outputs before decoding.
prompt_out = engine.generate(
prompts,
sampling_params={"max_new_tokens": 0, "temperature": 0.0},
return_logprob=True,
logprob_start_len=0,
lora_path=lora_paths,
)
prompt_logprobs = [
torch.tensor(
[lp for lp, _, _ in o["meta_info"]["input_token_logprobs"][1:]]
)
for o in prompt_out
]
gen_out = engine.generate(
prompts,
sampling_params={
"max_new_tokens": MAX_NEW_TOKENS,
"temperature": 0.0,
},
lora_path=lora_paths,
return_logprob=True,
logprob_start_len=-1,
top_logprobs_num=5,
)
if prefill_graph:
from prometheus_client import CollectorRegistry, multiprocess
# Wait for scheduler metric reporting after the response.
engine.get_server_info()
registry = CollectorRegistry()
multiprocess.MultiProcessCollector(registry)
samples = [
sample
for metric in registry.collect()
for sample in metric.samples
if sample.name == "sglang:cuda_graph_passes_total"
]
passes = {
mode: sum(
sample.value
for sample in samples
if sample.labels.get("mode") == mode
)
for mode in ("prefill_cuda_graph", "prefill_none")
}
self.assertEqual(
passes,
{"prefill_cuda_graph": 2, "prefill_none": 0},
f"{label}: expected two 64-request graph prefills",
)
results[label] = {
"prompt_logprobs": prompt_logprobs,
"next_token_scores": [
(
o["meta_info"]["output_token_logprobs"][0][1],
{
token: lp
for lp, token, _ in o["meta_info"][
"output_top_logprobs"
][0]
},
)
for o in gen_out
],
}
finally:
if engine is not None:
engine.shutdown()
for collector in set(REGISTRY._collector_to_names) - collectors_before:
REGISTRY.unregister(collector)
torch.cuda.empty_cache()
eager = results.pop("triton/disabled")
for backend, graph in results.items():
for i, prompt in enumerate(prompts):
e_lp, g_lp = eager["prompt_logprobs"][i], graph["prompt_logprobs"][i]
self.assertEqual(e_lp.numel(), g_lp.numel(), f"prompt {i}: token count")
max_diff = (e_lp - g_lp).abs().max().item()
self.assertLess(
max_diff,
LOGPROB_THRESHOLD,
f"{backend}, prompt {i} ({prompt[:40]!r}): logprobs drift "
f"{max_diff:.2e} from eager",
)
# Compare first-token scores; argmax can flip near ties.
e_token, e_scores = eager["next_token_scores"][i]
g_token, g_scores = graph["next_token_scores"][i]
for token in {e_token, g_token}:
self.assertIn(
token, e_scores, f"{backend}, prompt {i}: eager top-5"
)
self.assertIn(
token, g_scores, f"{backend}, prompt {i}: graph top-5"
)
self.assertLess(
abs(e_scores[token] - g_scores[token]),
LOGPROB_THRESHOLD,
f"{backend}, prompt {i}: first-token logprob drift",
)
if __name__ == "__main__":
unittest.main()
test/registered/lora/test_moe_lora_info.py
+171 -1
View File
@@ -1,15 +1,23 @@
import sys
from types import SimpleNamespace
import pytest
import torch
from sglang.srt.lora.backend.base_backend import _compute_moe_lora_info
from sglang.srt.lora.backend.base_backend import (
BaseLoRABackend,
_compute_moe_lora_info,
)
from sglang.srt.lora.backend.triton_backend import TritonLoRABackend
from sglang.srt.lora.utils import LoRABatchInfo
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd")
@@ -74,6 +82,67 @@ def test_compute_moe_lora_info_expands_segments(use_preallocated_buffers: bool):
assert actual_mapping.data_ptr() == token_lora_mapping.data_ptr()
def test_moe_graph_metadata_uses_matching_static_buffers():
"""Capture fixes the align kernel's request count; include empty tail slots."""
num_slots, max_loras = 8, 4
backend = BaseLoRABackend.__new__(BaseLoRABackend)
backend._is_moe_lora = True
backend.prefill_cuda_graph_batch_info = None
with torch.device(DEVICE):
backend.moe_cg_buffers = {
"adapter_enabled": torch.zeros(max_loras, dtype=torch.int32),
"token_lora_mapping": torch.full((8,), 7, dtype=torch.int32),
}
backend.prefill_moe_cg_buffers = {
"adapter_enabled": torch.zeros(max_loras, dtype=torch.int32),
"token_lora_mapping": torch.full((64,), 7, dtype=torch.int32),
}
for prefill, buffers in (
(False, backend.moe_cg_buffers),
(True, backend.prefill_moe_cg_buffers),
):
with torch.device(DEVICE):
info = LoRABatchInfo(
bs=num_slots,
use_cuda_graph=True,
num_segments=2,
seg_lens=torch.tensor(
[5, 3] + [0] * (num_slots - 2), dtype=torch.int32
),
seg_indptr=torch.zeros(num_slots + 1, dtype=torch.int32),
max_len=5,
weight_indices=torch.tensor(
[2, 1] + [0] * (num_slots - 2), dtype=torch.int32
),
lora_ranks=torch.tensor([0, 16, 8, 0], dtype=torch.int32),
scalings=torch.zeros(max_loras, dtype=torch.float),
permutation=None,
)
if prefill:
backend.prefill_cuda_graph_batch_info = info
torch.cumsum(info.seg_lens, dim=0, out=info.seg_indptr[1:])
forward_batch = SimpleNamespace(
forward_mode=ForwardMode.EXTEND,
batch_size=2,
extend_num_tokens=8,
extend_seq_lens_cpu=[5, 3],
)
moe = backend._add_moe_lora_info(forward_batch, info).moe_lora_info
torch.get_device_module(DEVICE).synchronize()
assert moe.adapter_enabled.data_ptr() == buffers["adapter_enabled"].data_ptr()
assert (
moe.token_lora_mapping.data_ptr()
== buffers["token_lora_mapping"].data_ptr()
)
if prefill:
assert moe.seg_indptr.shape[0] == num_slots + 1
assert moe.req_to_lora.shape[0] == num_slots
assert torch.all(moe.seg_indptr[2:] == 8)
assert torch.all(buffers["token_lora_mapping"][8:] == -1)
def test_compute_moe_lora_info_rejects_undercovered_launch():
device = DEVICE
seg_indptr = torch.tensor([0, 300], dtype=torch.int32, device=device)
@@ -92,5 +161,106 @@ def test_compute_moe_lora_info_rejects_undercovered_launch():
)
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.version.hip is not None,
reason="requires CUDA graph capture",
)
class TestDenseLoRAPrefillGraph(CustomTestCase):
def test_replay_preserves_ragged_adapters(self):
"""Ragged replays retain adapters without a token bucket per request."""
device, dtype = torch.device("cuda"), torch.float16
capacity, num_requests, rank, width = 1024, 64, 32, 64
ranks, scalings = [0, 16, 32], [0.0, 0.5, 1.0]
backend = TritonLoRABackend(max_loras_per_batch=3, device=device)
backend.init_prefill_cuda_graph_batch_info(
capacity, max_num_requests=num_requests
)
generator = torch.Generator().manual_seed(0)
cpu_a, cpu_b, cpu_embedding = [
torch.randint(-4, 5, shape, generator=generator).float() / 16
for shape in ((3, rank, width), (3, width, rank), (3, rank, width))
]
a_weights, b_weights, embedding_weights = [
weight.to(device=device, dtype=dtype)
for weight in (cpu_a, cpu_b, cpu_embedding)
]
x = torch.empty((capacity, width), device=device, dtype=dtype)
input_ids = torch.empty(capacity, device=device, dtype=torch.int64)
output = torch.full_like(x, 0.25)
ragged = [1, 3, 7, 15, 16, 17, 23, 31] * 8
ragged[-1] += capacity - sum(ragged)
cases = (
([capacity], [0]),
(ragged, [(i + 1) % 3 for i in range(num_requests)]),
([1, 17], [2, 0]),
)
for phase, (lengths, adapters) in enumerate(cases):
cpu_x = torch.randint(-4, 5, x.shape, generator=generator).float() / 16
cpu_ids = (torch.arange(capacity) + phase) % width
x.copy_(cpu_x)
input_ids.copy_(cpu_ids)
backend.prepare_lora_batch(
SimpleNamespace(
forward_mode=ForwardMode.EXTEND,
batch_size=len(lengths),
extend_num_tokens=sum(lengths),
extend_seq_lens_cpu=lengths,
extend_seq_lens=torch.tensor(
lengths, device=device, dtype=torch.int32
),
return_logprob=False,
),
weight_indices=adapters,
lora_ranks=ranks,
scalings=scalings,
use_cuda_graph=False,
use_prefill_cuda_graph=True,
)
if phase == 0:
info = backend._sgemm_info()
# Allow one partial 16-token tile per request, not a bucket per slot.
assert info.bs * info.max_len <= capacity + 16 * num_requests
graph, stream = torch.cuda.CUDAGraph(), torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
for capture in (False, True):
with (
torch.cuda.graph(graph, stream=stream)
if capture
else torch.cuda.stream(stream)
):
a_output = backend.run_lora_a_sgemm(x, a_weights)
backend.run_lora_b_sgemm(
a_output, b_weights, base_output=output
)
embedding_output = backend.run_lora_a_embedding(
input_ids, embedding_weights, vocab_size=width
)
torch.cuda.synchronize()
continue
output.fill_(0.25)
graph.replay()
torch.cuda.synchronize()
expected = torch.full((capacity, width), 0.25, dtype=dtype)
expected_embedding = torch.zeros((capacity, rank), dtype=dtype)
start = 0
for length, adapter in zip(lengths, adapters):
rows, r = slice(start, start + length), ranks[adapter]
if r:
expected_a = (cpu_x[rows] @ cpu_a[adapter, :r].T).to(dtype)
delta = (
expected_a.float() @ cpu_b[adapter, :, :r].T * scalings[adapter]
).to(dtype)
expected[rows] += delta
expected_embedding[rows, :r] = cpu_embedding[adapter, :r][
:, cpu_ids[rows]
].T.to(dtype)
start += length
torch.testing.assert_close(output.cpu(), expected, atol=1e-3, rtol=1e-3)
torch.testing.assert_close(
embedding_output.cpu(), expected_embedding, atol=0, rtol=0
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))