[CI] Bound the CUDA graph capture range in test launches and lift the spec fixture's admission cap (#33776)

This commit is contained in:
Liangsheng Yin
2026-08-05 20:00:30 -07:00
committed by GitHub
parent f33f6a522f
commit 9bd1461757
9 changed files with 26 additions and 6 deletions
@@ -58,8 +58,12 @@ class SpecEagleServerBase(CustomTestCase):
attention_backend = "flashinfer" attention_backend = "flashinfer"
# Primary axis: False -> overlap scheduler; True -> synchronous (non-overlap). # Primary axis: False -> overlap scheduler; True -> synchronous (non-overlap).
disable_overlap = False disable_overlap = False
mem_fraction_static = 0.85 # Leaves ~3.3GB on a 32GB card for the verify logits and activations at a
max_running_requests = 8 # cap of 64; higher OOMs, lower starves the KV pool into capping the batch.
mem_fraction_static = 0.80
# The eval kits drive 128 client threads, so a small cap just serializes them.
# Capture follows: capture_bs is clipped to req_to_token_pool.size (cap + 1).
max_running_requests = 64
chunked_prefill_size = 128 chunked_prefill_size = 128
# bf16 rather than fp16: fp16 activations can overflow (-> Inf -> NaN) on # bf16 rather than fp16: fp16 activations can overflow (-> Inf -> NaN) on
# degenerate draft branches in verify and trip the CI NaN asserts. # degenerate draft branches in verify and trip the CI NaN asserts.
+9
View File
@@ -952,6 +952,15 @@ def popen_launch_server(
other_args = list(other_args) other_args = list(other_args)
other_args += ["--device", str(device)] other_args += ["--device", str(device)]
# Prefill dominates capture time: the bucket list runs to chunked_prefill_size
# (8192 on H100-class GPUs) and its largest buckets cost seconds each, while
# 97% of CI prefill batches are under 1024 tokens -- a server that serves one
# test file captures the rest and never replays it. Decode is left alone: its
# capture cost is per-phase, not per-bucket. Pass the flag to opt out.
prefill_flag = "--cuda-graph-max-bs-prefill"
if not any(str(arg).startswith(prefill_flag) for arg in other_args):
other_args = list(other_args) + [prefill_flag, "1024"]
# CI-specific: Validate cache and enable offline mode if complete # CI-specific: Validate cache and enable offline mode if complete
if env is None: if env is None:
env = os.environ.copy() env = os.environ.copy()
@@ -40,6 +40,8 @@ class _Core(EagleLlama2Base):
attention_backend = "intel_amx" attention_backend = "intel_amx"
disable_overlap = True disable_overlap = True
mem_fraction_static = 0.3 mem_fraction_static = 0.3
# CPU decode is compute-bound; a wider batch buys nothing here.
max_running_requests = 8
gsm8k_num_examples = 64 gsm8k_num_examples = 64
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
@@ -20,6 +20,8 @@ class TestEagle3ParityCPU(SpecParityKit, Eagle3Base):
attention_backend = "intel_amx" attention_backend = "intel_amx"
disable_overlap = True disable_overlap = True
mem_fraction_static = 0.3 mem_fraction_static = 0.3
# CPU decode is compute-bound; a wider batch buys nothing here.
max_running_requests = 8
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
@@ -32,6 +32,8 @@ class _Core(EagleLlama2Base):
attention_backend = "intel_amx" attention_backend = "intel_amx"
disable_overlap = True disable_overlap = True
mem_fraction_static = 0.3 mem_fraction_static = 0.3
# CPU decode is compute-bound; a wider batch buys nothing here.
max_running_requests = 8
gsm8k_num_examples = 64 gsm8k_num_examples = 64
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
@@ -30,7 +30,6 @@ class TestEagle3Perf(Eagle3Base, SpecPerfKit):
class TestEagleLlama2Retract(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): class TestEagleLlama2Retract(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
"""Retract under a small KV budget; must not leak.""" """Retract under a small KV budget; must not leak."""
max_running_requests = 64
extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract
env_overrides = ( env_overrides = (
(envs.SGLANG_TEST_RETRACT, True), (envs.SGLANG_TEST_RETRACT, True),
@@ -46,7 +45,6 @@ class TestEagle3Topk16V2Retract(Eagle3Base, SpecAccuracyKit, SpecFeatureKit):
spec_tokens = 64 spec_tokens = 64
disable_overlap = False disable_overlap = False
cuda_graph_max_bs_decode = 5 cuda_graph_max_bs_decode = 5
max_running_requests = 64
gsm8k_accept_len_thres = 2.4 gsm8k_accept_len_thres = 2.4
extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract
env_overrides = ( env_overrides = (
@@ -40,6 +40,9 @@ class TestEagle3Topk16(
spec_tokens = 64 spec_tokens = 64
disable_overlap = False disable_overlap = False
enable_return_hidden_states = True enable_return_hidden_states = True
# Verify materializes bs * spec_tokens fp32 logit rows: 262MB here, but
# 2.1GB at the fixture's 64 -- OOM on a 32GB card.
max_running_requests = 8
cuda_graph_max_bs_decode = 5 cuda_graph_max_bs_decode = 5
acc_length_thres = 3.1 acc_length_thres = 3.1
batch_accept_len_thres = 1.75 batch_accept_len_thres = 1.75
@@ -31,6 +31,8 @@ class TestEagle3Page4Topk8(Eagle3Base, SpecAccuracyKit, SpecLogprobKit, SpecFeat
# Preset accept-length values are topk=1 numbers -- loose for a topk=8 # Preset accept-length values are topk=1 numbers -- loose for a topk=8
# tree; tighten once CI reports the actuals. # tree; tighten once CI reports the actuals.
gsm8k_accept_len_thres = 2.0 gsm8k_accept_len_thres = 2.0
# Same fp32 logits budget as the topk=16 classes (bs * spec_tokens * vocab).
max_running_requests = 16
cuda_graph_max_bs_decode = 5 cuda_graph_max_bs_decode = 5
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
@@ -32,8 +32,6 @@ class TestEagle3Triton(
"""Overlap scheduler on triton (kits listed in bases).""" """Overlap scheduler on triton (kits listed in bases)."""
attention_backend = "triton" attention_backend = "triton"
max_running_requests = 64
cuda_graph_max_bs_decode = 64
gsm8k_num_examples = 200 gsm8k_num_examples = 200
gsm8k_check_accept_len = False gsm8k_check_accept_len = False
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)