From 80ac11eda3ac50748bff2bdbd0642369c1324e50 Mon Sep 17 00:00:00 2001 From: cctry Date: Wed, 1 Jul 2026 22:12:56 -0700 Subject: [PATCH] Fix capture-mode detection during breakable CUDA graph capture (#29866) --- .../sglang/srt/model_executor/runner_utils/capture_mode.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py index f14c55ce0..293ec1b25 100644 --- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py +++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py @@ -24,6 +24,10 @@ from typing import Optional import torch +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) + # Detect whether the current forward pass is in capture mode. is_capture_mode = False @@ -33,7 +37,7 @@ _capture_lora_variant: Optional[str] = None def get_is_capture_mode() -> bool: - return is_capture_mode + return is_capture_mode or is_in_breakable_cuda_graph() def compile_in_capture_mode(func):