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: