From 8dca6291c7a4111692d4f173299c8a960ae12e80 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 28 May 2026 00:50:35 -0700 Subject: [PATCH] [CI] FA3: ascending cuda-graph capture to avoid varlen workspace IMA (#26532) (#26550) --- .../srt/model_executor/cuda_graph_runner.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 6501b6fe1..702fb34da 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -530,6 +530,18 @@ def set_global_graph_memory_pool(val): global_graph_memory_pool = val +def _ci_use_ascending_capture_order(server_args) -> bool: + """Whether CI + FA3 forces ascending cuda-graph capture (FA3 varlen IMA workaround, #26532).""" + if not envs.SGLANG_IS_IN_CI.get(): + return False + prefill_backend, decode_backend = server_args.get_attention_backends() + return "fa3" in ( + prefill_backend, + decode_backend, + server_args.speculative_draft_attention_backend, + ) + + class CudaGraphRunner: """A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile.""" @@ -821,11 +833,17 @@ class CudaGraphRunner: self.model_runner.gpu_id, empty_cache=False, ) - # Reverse the order to enable better memory sharing across cuda graphs. + # Reverse for memory sharing; CI+FA3 uses ascending to dodge the + # FA3 varlen workspace-slot IMA (#26532). + bs_seq = ( + list(self.capture_bs) + if _ci_use_ascending_capture_order(self.model_runner.server_args) + else list(reversed(self.capture_bs)) + ) capture_range = ( - tqdm.tqdm(list(reversed(self.capture_bs))) + tqdm.tqdm(bs_seq) if get_tensor_model_parallel_rank() == 0 - else reversed(self.capture_bs) + else iter(bs_seq) ) for i, bs in enumerate(capture_range): if get_tensor_model_parallel_rank() == 0: