diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index e5ddb3adb..99c2bc690 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -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( diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index eb88d8273..c22549897 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -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) diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index aa578c399..fad45cfaf 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -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) diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index c5777e6df..1e0e604bb 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -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 diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 564a0aa1d..682116f42 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -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 diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 9baacfee7..d59182854 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -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) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 49d24ab38..ab330617b 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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 diff --git a/test/registered/e2e/lora/test_lora_moe_prefill_cuda_graph.py b/test/registered/e2e/lora/test_lora_moe_prefill_cuda_graph.py new file mode 100644 index 000000000..f43e3a581 --- /dev/null +++ b/test/registered/e2e/lora/test_lora_moe_prefill_cuda_graph.py @@ -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() diff --git a/test/registered/lora/test_moe_lora_info.py b/test/registered/lora/test_moe_lora_info.py index 478c6b8f4..484253791 100644 --- a/test/registered/lora/test_moe_lora_info.py +++ b/test/registered/lora/test_moe_lora_info.py @@ -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__]))