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

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
+        max_bs counts tokens. Layers reuse these buffers sequentially.
-        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.
         """
         base = moe_layer.base_layer
         top_k = base.top_k
-        qinfo = moe_layer._quant_info
+        device = moe_layer._quant_info.w13_weight.device
-        E, N, _ = qinfo.w13_weight.shape
-        hidden_dim = qinfo.w2_weight.shape[1]
-        device = qinfo.w13_weight.device
-        dtype = compute_dtype
         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 = {
+        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
-            ),
             "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"]
+            buffers = self.prefill_moe_cg_buffers if prefill else self.moe_cg_buffers
-            token_lora_mapping = self.moe_cg_buffers["token_lora_mapping"]
+            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

@@ -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 if max_num_requests is None else max_num_requests
-        max_bs = max_num_tokens
         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

@@ -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
+# Match the dense kernels' token tile.
-# kernel grids; batches with more requests fall back to eager prefill.
+PREFILL_CUDA_GRAPH_LORA_CHUNK_SIZE = 16
-PREFILL_CUDA_GRAPH_LORA_SEGMENTS = 32
 
 
 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):
+    def init_prefill_cuda_graph_batch_info(
-        num_slots = PREFILL_CUDA_GRAPH_LORA_SEGMENTS
+        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

@@ -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

@@ -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):
+    def init_prefill_cuda_graph_batch_info(
-        """Allocate the static prefill-CUDA-graph LoRA metadata, sized by the
+        self, max_num_tokens: int, max_num_requests: Optional[int] = None
-        largest captured token bucket. Called before capture."""
+    ):
+        """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;
+        """MoE LoRA supports full and breakable capture; DP attention is unsupported."""
-        excludes MoE LoRA and DP attention."""
+        from sglang.srt.model_executor.cuda_graph_config import (
-        return (
+            Backend,
-            self.lora_backend.supports_prefill_cuda_graph
+            Phase,
-            and not self.lora_backend.is_moe_lora
+            check_cuda_graph_backend,
-            and not self.enable_dp_attention
         )
 
+        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

@@ -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

@@ -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

@@ -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

@@ -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__]))