[LoRA] Support MoE in full and breakable prefill CUDA graphs (#38578)
This commit is contained in:
Yanbin Jiang
@@ -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()
|
||||
@@ -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__]))
|
||||
|
||||
Reference in New Issue
Block a user