From 690d4cdd942a17e45dd16b1c3346e375f38efb9e Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Fri, 29 May 2026 04:17:37 +0800 Subject: [PATCH] Revert "[CI] FA3: ascending cuda-graph capture to avoid varlen workspace IMA (#26532) (#26550)" (#26600) --- .../srt/model_executor/cuda_graph_runner.py | 24 +++---------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 702fb34da..6501b6fe1 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -530,18 +530,6 @@ 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.""" @@ -833,17 +821,11 @@ class CudaGraphRunner: self.model_runner.gpu_id, empty_cache=False, ) - # 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)) - ) + # Reverse the order to enable better memory sharing across cuda graphs. capture_range = ( - tqdm.tqdm(bs_seq) + tqdm.tqdm(list(reversed(self.capture_bs))) if get_tensor_model_parallel_rank() == 0 - else iter(bs_seq) + else reversed(self.capture_bs) ) for i, bs in enumerate(capture_range): if get_tensor_model_parallel_rank() == 0: