[CI] FA3: ascending cuda-graph capture to avoid varlen workspace IMA (#26532) (#26550)

This commit is contained in:
Liangsheng Yin
2026-05-28 00:50:35 -07:00
committed by GitHub
parent 8f21b3e2ef
commit 8dca6291c7
@@ -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: