Piecewise Cuda Graph set default (#16331)

This commit is contained in:
Yuwei An
2026-03-02 23:18:07 +08:00
committed by GitHub
parent 468e3dc56b
commit c64274c746
34 changed files with 505 additions and 143 deletions
@@ -137,7 +137,7 @@ class TestVLMPiecewiseCudaGraph(CustomTestCase):
"--trust-remote-code",
"--piecewise-cuda-graph-max-tokens",
"8192",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--tp=8",
"--piecewise-cuda-graph-compiler=eager",
"--disable-radix-cache",
@@ -140,7 +140,7 @@ class TestVLMViTCudaGraph(CustomTestCase):
other_args=[
"--mm-attention-backend",
"fa3",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-max-tokens",
"8192",
"--chunked-prefill-size",
@@ -38,7 +38,7 @@ class TestDisaggregationPiecewiseCudaGraph(PDDisaggregationServerBase):
"prefill",
"--tp",
"1",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
+8 -1
View File
@@ -111,7 +111,14 @@ class TestMLADeepseekV3Fa3Fp8Kvcache(CustomTestCase):
]
if is_cuda():
other_args.extend(
["--attention-backend", "fa3", "--cuda-graph-max-bs", "2"]
[
"--attention-backend",
"fa3",
"--mem-fraction-static",
"0.8",
"--cuda-graph-max-bs",
"2",
]
)
cls.process = popen_launch_server(
cls.model,
@@ -66,7 +66,7 @@ CI_MODELS = [
# the complete set of models to test sglang's generation model
ALL_MODELS = [
*CI_MODELS,
ModelCase("Qwen/Qwen2-1.5B"),
ModelCase("Qwen/Qwen2-1.5B", decode_tolerance=7e-2),
ModelCase("Qwen/Qwen2.5-14B-Instruct"),
ModelCase("HuggingFaceTB/SmolLM-135M-Instruct", skip_long_prompt=True),
ModelCase("allenai/OLMo-1B-0724-hf", decode_tolerance=8e-2, skip_long_prompt=True),
@@ -41,7 +41,6 @@ class TestKimiLinearPiecewiseCudaGraph(CustomTestCase):
"--tp",
"2",
"--trust-remote-code",
"--enable-piecewise-cuda-graph",
],
)
@@ -40,9 +40,6 @@ class TestQwen3NextPiecewiseCudaGraph(CustomTestCase):
other_args=[
"--tp",
"4",
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
],
)
@@ -27,7 +27,6 @@ class TestPiecewiseCudaGraphTP(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--tp",
@@ -27,7 +27,6 @@ class TestPiecewiseCudaGraphQwen3MoE(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
],
@@ -65,7 +64,7 @@ class TestPiecewiseCudaGraphGPTQ(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-piecewise-cuda-graph"],
other_args=[],
)
@classmethod
@@ -101,7 +100,7 @@ class TestPiecewiseCudaGraphAWQ(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-piecewise-cuda-graph"],
other_args=[],
)
@classmethod
@@ -33,7 +33,7 @@ class TestPiecewiseCudaGraphCorrectness(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-piecewise-cuda-graph"],
other_args=[],
)
@classmethod
@@ -57,8 +57,7 @@ class TestPiecewiseCudaGraphBenchmark(CustomTestCase):
def test_latency(self):
prefill_latency, _, _ = run_bench_one_batch(
DEFAULT_MODEL_NAME_FOR_TEST,
other_args=["--enable-piecewise-cuda-graph"],
DEFAULT_MODEL_NAME_FOR_TEST, other_args=[]
)
self.assertLess(prefill_latency, 0.015)
@@ -76,7 +75,6 @@ class TestPiecewiseCudaGraphLlama31FP4(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--quantization",
"modelopt_fp4",
"--mem-fraction-static",
@@ -112,7 +110,6 @@ class TestPiecewiseCudaGraphDeepSeek(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--piecewise-cuda-graph-max-tokens",
@@ -152,7 +149,6 @@ class TestPiecewiseCudaGraphFP8(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--quantization",
"modelopt_fp8",
"--kv-cache-dtype",
@@ -191,7 +187,6 @@ class TestPiecewiseCudaGraphQwen25VL(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--disable-radix-cache",
@@ -232,7 +227,6 @@ class TestPiecewiseCudaGraphInternVL25(CustomTestCase):
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
"--disable-radix-cache",
@@ -273,7 +267,6 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
model_path=model_path,
enable_multimodal=True,
is_embedding=True,
enable_piecewise_cuda_graph=True,
piecewise_cuda_graph_compiler="eager",
)
out = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0]["embedding"]
@@ -284,7 +277,7 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
model_path=model_path,
enable_multimodal=True,
is_embedding=True,
enable_piecewise_cuda_graph=False,
disable_piecewise_cuda_graph=True,
)
out_without_pcg = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0][
"embedding"
@@ -100,7 +100,6 @@ class TestDeepseekV3FP4PiecewiseCudaGraph(CustomTestCase):
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--enable-piecewise-cuda-graph",
"--kv-cache-dtype",
"fp8_e4m3",
"--model-loader-extra-config",
@@ -31,11 +31,15 @@ class TestSWARadixCacheKL(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
# Use a lower mem-fraction-static to avoid OOM during input logprobs
# gathering. With PCG enabled, more memory is reserved for CUDA graph
# captures, so the static fraction should be lower.
other_args=[
"--tp-size",
"1",
"--mem-fraction-static",
"0.75",
"0.70",
"--disable-piecewise-cuda-graph",
],
)
@@ -36,7 +36,7 @@ class TestPiecewiseGraphPrefillCorrectness(CustomTestCase):
"ascend",
"--cuda-graph-bs",
128,
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-tokens",
*TOKENS_TO_CAPTURE,
],
@@ -77,7 +77,7 @@ class TestPiecewiseGraphPrefillBenchmark(CustomTestCase):
0.8,
"--attention-backend",
"ascend",
"--enable-piecewise-cuda-graph",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-tokens",
]
+ TOKENS_TO_CAPTURE,