From b9e33d6a5be796ab0c893bf4b1b825f451a3caa8 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Wed, 22 Apr 2026 17:11:11 -0400 Subject: [PATCH] Dual MoE CUDA graph capture for lora/nolora batches (#22809) --- python/sglang/srt/layers/moe/utils.py | 25 +++++ python/sglang/srt/lora/lora_moe_runners.py | 22 +++-- .../srt/model_executor/cuda_graph_runner.py | 98 ++++++++++++++----- python/sglang/srt/server_args.py | 9 ++ 4 files changed, 126 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 61ba72177..73dde3166 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -148,6 +148,7 @@ MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None SPECULATIVE_MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None SPECULATIVE_MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None +RECORD_NOLORA_GRAPH: bool = False DEEPEP_MODE: Optional[DeepEPMode] = None IS_TBO_ENABLED: Optional[bool] = None IS_SBO_ENABLED: Optional[bool] = None @@ -162,6 +163,7 @@ def initialize_moe_config(server_args: ServerArgs): global MOE_RUNNER_BACKEND global SPECULATIVE_MOE_RUNNER_BACKEND global SPECULATIVE_MOE_A2A_BACKEND + global RECORD_NOLORA_GRAPH global DEEPEP_MODE global DEEPEP_CONFIG global IS_TBO_ENABLED @@ -172,6 +174,25 @@ def initialize_moe_config(server_args: ServerArgs): MOE_A2A_BACKEND = MoeA2ABackend(server_args.moe_a2a_backend) MOE_RUNNER_BACKEND = MoeRunnerBackend(server_args.moe_runner_backend) + # Dual CUDA graphs only validated for triton MoE backends. + _triton_ok = MOE_RUNNER_BACKEND in ( + MoeRunnerBackend.TRITON, + MoeRunnerBackend.TRITON_KERNELS, + ) + if ( + bool(server_args.record_nolora_graph) + and bool(server_args.enable_lora) + and not _triton_ok + ): + logger.warning( + f"record_nolora_graph only validated for triton MoE backend, " + f"but moe_runner_backend={server_args.moe_runner_backend}. Disabling." + ) + RECORD_NOLORA_GRAPH = ( + bool(server_args.record_nolora_graph) + and bool(server_args.enable_lora) + and _triton_ok + ) SPECULATIVE_MOE_RUNNER_BACKEND = ( MoeRunnerBackend(server_args.speculative_moe_runner_backend) if server_args.speculative_moe_runner_backend is not None @@ -227,6 +248,10 @@ def get_speculative_moe_a2a_backend() -> MoeA2ABackend: return SPECULATIVE_MOE_A2A_BACKEND +def should_record_nolora_graph() -> bool: + return RECORD_NOLORA_GRAPH + + def get_deepep_mode() -> DeepEPMode: global DEEPEP_MODE if DEEPEP_MODE is None: diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 983064154..04f8b4867 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -444,12 +444,10 @@ def _add_lora_gate_up_delta( ) if get_is_capture_mode(): - # During CUDA graph capture, always enter the LoRA path so that - # the LoRA kernels are recorded in the graph. adapter_enabled is - # all-zeros during capture, so the Triton kernel early-exits per - # program (zero overhead). During replay the tensor is updated - # in-place with the real adapter mask before graph.replay(). - has_active_lora = True + from sglang.srt.model_executor.cuda_graph_runner import get_capture_lora_variant + + # Record LoRA kernels for lora graph; skip for nolora graph. + has_active_lora = get_capture_lora_variant() != "nolora" else: num_loras = len(lora_info.lora_ranks) has_active_lora = ( @@ -549,6 +547,12 @@ def _add_lora_down_delta( if lora_info.max_lora_rank == 0: return + if get_is_capture_mode(): + from sglang.srt.model_executor.cuda_graph_runner import get_capture_lora_variant + + if get_capture_lora_variant() == "nolora": + return + M, top_k, hidden_dim = intermediate_cache.shape down_lora_a = lora_info.down_lora_a_weights @@ -629,6 +633,12 @@ def build_lora_hooks( if lora_info is None or lora_info.max_lora_rank == 0: return LoRAHooks() + if get_is_capture_mode(): + from sglang.srt.model_executor.cuda_graph_runner import get_capture_lora_variant + + if get_capture_lora_variant() == "nolora": + return LoRAHooks() + # Compute alignment / mapping (once, shared by both hooks) token_lora_mapping: torch.Tensor | None = None sorted_token_ids_reshaped: torch.Tensor | None = None diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 8a5555c42..2d2795275 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -53,7 +53,11 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer -from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend +from sglang.srt.layers.moe.utils import ( + get_deepep_mode, + get_moe_a2a_backend, + should_record_nolora_graph, +) from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -364,12 +368,25 @@ class DecodeInputBuffers(ForwardInputBuffers): # Detect whether the current forward pass is in capture mode is_capture_mode = False +# When capturing dual MoE backends, tracks which variant is being captured. +# None = not dual, "lora" = capturing lora variant, "nolora" = capturing nolora variant. +_capture_lora_variant: Optional[str] = None def get_is_capture_mode(): return is_capture_mode +def get_capture_lora_variant() -> Optional[str]: + """Return the lora variant being captured, or None if not in dual capture.""" + return _capture_lora_variant + + +def _set_capture_lora_variant(variant: Optional[str]): + global _capture_lora_variant + _capture_lora_variant = variant + + @contextmanager def model_capture_mode(): global is_capture_mode @@ -509,6 +526,19 @@ def set_global_graph_memory_pool(val): global_graph_memory_pool = val +def _default_make_graph_key(bs, stream_idx=None, variant_label=None): + """Build a graph dict key from batch size, stream index, and lora variant. + + Standalone function so it can be used by CudaGraphRunner.capture() even when + called on subclasses (e.g. EAGLEDraftCudaGraphRunner) that don't inherit from + CudaGraphRunner and thus lack the method. + """ + key = bs if stream_idx is None else f"{stream_idx}_{bs}" + if variant_label is not None: + key = f"{variant_label}_{key}" + return key + + class CudaGraphRunner: """A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile.""" @@ -555,6 +585,7 @@ class CudaGraphRunner: self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() + self.record_nolora_graph = should_record_nolora_graph() self.dllm_config = DllmConfig.from_server_args(model_runner.server_args) self.is_dllm = self.dllm_config is not None @@ -677,6 +708,20 @@ class CudaGraphRunner: def _cache_loc_dtype(self): return torch.int64 + def _make_graph_key(self, bs, stream_idx=None, variant_label=None): + """Build a graph dict key from batch size, stream index, and lora variant.""" + return _default_make_graph_key(bs, stream_idx, variant_label) + + def _resolve_lora_variant(self, forward_batch: ForwardBatch): + """Return the variant label for the given batch, or None if dual backends are off.""" + if not getattr(self, "record_nolora_graph", False): + return None + if forward_batch.lora_ids is not None and any( + uid is not None for uid in forward_batch.lora_ids + ): + return "lora" + return "nolora" + def can_run(self, forward_batch: ForwardBatch): # Disable for token embedding overrides (dynamic per-request) if forward_batch.replace_embeds is not None: @@ -692,9 +737,9 @@ class CudaGraphRunner: else: cuda_graph_bs = forward_batch.batch_size - graph_key = cuda_graph_bs - if self.enable_pdmux: - graph_key = f"{get_current_stream_idx()}_{cuda_graph_bs}" + variant_label = self._resolve_lora_variant(forward_batch) + stream_idx = get_current_stream_idx() if self.enable_pdmux else None + graph_key = self._make_graph_key(cuda_graph_bs, stream_idx, variant_label) is_bs_supported = ( graph_key in self.graphs @@ -789,6 +834,13 @@ class CudaGraphRunner: if get_tensor_model_parallel_rank() == 0 else reversed(self.capture_bs) ) + # When record_nolora_graph is set, capture each batch size twice: + # once with LoRA hooks and once without. + lora_variants = ( + [("lora", True), ("nolora", False)] + if getattr(self, "record_nolora_graph", False) + else [(None, None)] + ) for i, bs in enumerate(capture_range): if get_tensor_model_parallel_rank() == 0: avail_mem = get_available_gpu_memory( @@ -800,20 +852,21 @@ class CudaGraphRunner: f"Capturing batches ({bs=} {avail_mem=:.2f} GB)" ) - with patch_model( - self.model_runner.model, - bs in self.compile_bs, - num_tokens=bs * self.num_tokens_per_bs, - tp_group=self.model_runner.tp_group, - ) as forward: - ( - graph, - output_buffers, - ) = self.capture_one_batch_size(bs, forward, stream_idx) - # For pd_multiplexing, we need to save the graph and output buffers - key = bs if stream_idx is None else f"{stream_idx}_{bs}" - self.graphs[key] = graph - self.output_buffers[key] = output_buffers + for variant_label, variant_has_lora in lora_variants: + _set_capture_lora_variant(variant_label) + with patch_model( + self.model_runner.model, + bs in self.compile_bs, + num_tokens=bs * self.num_tokens_per_bs, + tp_group=self.model_runner.tp_group, + ) as forward: + ( + graph, + output_buffers, + ) = self.capture_one_batch_size(bs, forward, stream_idx) + key = _default_make_graph_key(bs, stream_idx, variant_label) + self.graphs[key] = graph + self.output_buffers[key] = output_buffers # Trigger CUDA graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes @@ -832,6 +885,8 @@ class CudaGraphRunner: self.stream = graph_capture_context.stream _capture_one_stream(i) + _set_capture_lora_variant(None) + if self.enable_profile_cuda_graph: self._post_process_after_profile(prof) @@ -1228,10 +1283,9 @@ class CudaGraphRunner: ) # Replay - if self.enable_pdmux: - graph_key = f"{get_current_stream_idx()}_{self.bs}" - else: - graph_key = self.bs + variant_label = self._resolve_lora_variant(forward_batch) + stream_idx = get_current_stream_idx() if self.enable_pdmux else None + graph_key = self._make_graph_key(self.bs, stream_idx, variant_label) self.graphs[graph_key].replay() output = self.output_buffers[graph_key] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index beb8398f8..4454d1926 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -538,6 +538,7 @@ class ServerArgs: "none", "deepep", "mooncake", "nixl", "mori", "ascend_fuseep", "flashinfer" ] = "none" moe_runner_backend: str = "auto" + record_nolora_graph: bool = True flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default" enable_flashinfer_allreduce_fusion: bool = False enforce_disable_flashinfer_allreduce_fusion: bool = False @@ -5407,6 +5408,14 @@ class ServerArgs: default=ServerArgs.moe_runner_backend, help="Choose the runner backend for MoE.", ) + parser.add_argument( + "--record-nolora-graph", + action=argparse.BooleanOptionalAction, + default=ServerArgs.record_nolora_graph, + help="Capture a second set of CUDA graphs without LoRA hooks. " + "Batches without active adapters replay the faster nolora graph. " + "Enabled by default.", + ) parser.add_argument( "--flashinfer-mxfp4-moe-precision", type=str,