[FullCG] Support chunked cached-prefix prefill (#30825)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
paulzhang-tm
co-authored by
Claude Fable 5
Claude Fable 5
parent
1eee8fbdcc
commit
4e5a05148a
@@ -1,19 +1,22 @@
|
||||
"""Integration test for the full prefill CUDA graph backend.
|
||||
"""Integration tests for the full prefill CUDA graph backend.
|
||||
|
||||
Spins up Qwen3-8B with --cuda-graph-backend-prefill=full and checks
|
||||
mgsm_en accuracy, mirroring the breakable-CG integration test.
|
||||
The Qwen3-8B test checks end-to-end accuracy with FlashInfer. The smaller
|
||||
DeepSeek-Coder-V2-Lite test checks that an MLA radix-prefix hit selects the
|
||||
OSS FA4 cached-prefix graph variant and matches an eager cold request.
|
||||
|
||||
The attention backend is pinned to flashinfer: plain EXTEND under full
|
||||
CUDA graph requires the backend's init_forward_metadata_out_graph to
|
||||
support extend (capture-stable plan state). flashinfer and the
|
||||
FlashAttention backend (fa4; fa3 untested — needs SM90 hardware)
|
||||
implement it; flashinfer is pinned here for CI-hardware portability
|
||||
(fa4 requires Blackwell).
|
||||
implement it; the FA4 case is restricted to Blackwell.
|
||||
"""
|
||||
|
||||
import re
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import get_device_sm, kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
@@ -24,8 +27,8 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
# CI Registration — large suite to fit the integration test's server startup.
|
||||
register_cuda_ci(est_time=79, stage="base-b", runner_config="1-gpu-large")
|
||||
# OSS FA4 coverage requires Blackwell. Each test still uses only one GPU.
|
||||
register_cuda_ci(est_time=170, stage="base-b", runner_config="4-gpu-b200")
|
||||
|
||||
|
||||
class TestFullCudaGraphPrefill(CustomTestCase):
|
||||
@@ -65,5 +68,97 @@ class TestFullCudaGraphPrefill(CustomTestCase):
|
||||
self.assertGreaterEqual(score, 0.80)
|
||||
|
||||
|
||||
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
||||
class TestFullCudaGraphChunkedPrefix(unittest.TestCase):
|
||||
"""A radix-cache hit replays the OSS FA4 FullCG prefix variant."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--prefill-attention-backend=fa4",
|
||||
"--decode-attention-backend=flashinfer",
|
||||
"--disable-flashinfer-autotune",
|
||||
"--context-length=256",
|
||||
"--max-total-tokens=512",
|
||||
"--max-running-requests=1",
|
||||
"--chunked-prefill-size=256",
|
||||
"--skip-server-warmup",
|
||||
"--enable-metrics",
|
||||
"--cuda-graph-config",
|
||||
'{"decode":{"backend":"disabled"},'
|
||||
'"prefill":{"backend":"full","bs":[32],"max_bs":32,'
|
||||
'"full_prefill_max_req":1,'
|
||||
'"full_prefill_prefix_chunk_tokens":64}}',
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _generate(self, input_ids):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"input_ids": input_ids,
|
||||
"sampling_params": {"max_new_tokens": 4, "temperature": 0},
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
def _prefill_graph_count(self):
|
||||
metrics = requests.get(self.base_url + "/metrics", timeout=30).text
|
||||
match = re.search(
|
||||
r'^sglang:cuda_graph_passes_total\{[^}]*mode="prefill_cuda_graph"[^}]*\}'
|
||||
r"\s+([0-9.eE+-]+)$",
|
||||
metrics,
|
||||
re.MULTILINE,
|
||||
)
|
||||
return float(match.group(1)) if match else 0.0
|
||||
|
||||
def test_cached_prefix_replays_full_cuda_graph(self):
|
||||
prefix = list(range(1000, 1048))
|
||||
prompt = prefix + list(range(2000, 2032))
|
||||
|
||||
requests.post(self.base_url + "/flush_cache", timeout=30).raise_for_status()
|
||||
cold = self._generate(prompt)
|
||||
|
||||
requests.post(self.base_url + "/flush_cache", timeout=30).raise_for_status()
|
||||
self._generate(prefix)
|
||||
graph_count = self._prefill_graph_count()
|
||||
cached = self._generate(prompt)
|
||||
|
||||
self.assertEqual(cached["meta_info"]["cached_tokens"], len(prefix))
|
||||
self.assertEqual(cached["output_ids"], cold["output_ids"])
|
||||
self.assertEqual(self._prefill_graph_count(), graph_count + 1)
|
||||
|
||||
def test_cached_prefix_replays_two_64_token_full_cuda_graph_chunks(self):
|
||||
prefix = list(range(1000, 1128))
|
||||
prompt = prefix + list(range(2000, 2032))
|
||||
|
||||
# The 160-token cold reference fits in one scheduler prefill chunk, so
|
||||
# only the 32-token cache-hit suffix is eligible for this FullCG bucket.
|
||||
requests.post(self.base_url + "/flush_cache", timeout=30).raise_for_status()
|
||||
cold = self._generate(prompt)
|
||||
|
||||
requests.post(self.base_url + "/flush_cache", timeout=30).raise_for_status()
|
||||
self._generate(prefix)
|
||||
graph_count = self._prefill_graph_count()
|
||||
cached = self._generate(prompt)
|
||||
|
||||
self.assertEqual(cached["meta_info"]["cached_tokens"], len(prefix))
|
||||
self.assertEqual(cached["output_ids"], cold["output_ids"])
|
||||
self.assertEqual(self._prefill_graph_count(), graph_count + 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -31,6 +31,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
|
||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||
runner._is_full_backend = False
|
||||
runner.enable_lora = False
|
||||
runner._capture_chunked_prefix = False
|
||||
runner.prefill_backend_name = backend
|
||||
runner.has_mha_companion_layers = backend == Backend.BREAKABLE
|
||||
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
||||
|
||||
@@ -20,6 +20,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
|
||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||
runner._is_full_backend = False
|
||||
runner.enable_lora = False
|
||||
runner._capture_chunked_prefix = False
|
||||
runner.prefill_backend_name = Backend.TC_PIECEWISE
|
||||
runner.has_mha_companion_layers = False
|
||||
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
"""CPU coverage for chunked-prefix Full prefill CUDA-graph state."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
|
||||
from sglang.srt.model_executor.cuda_graph_config import Backend
|
||||
from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
|
||||
PrefillCudaGraphRunner,
|
||||
)
|
||||
from sglang.srt.model_executor.runner.shape_key import ShapeKey
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _FakeAttentionBackend:
|
||||
supports_full_cuda_graph_chunked_prefix = True
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def prepare_full_cuda_graph_chunked_prefix(self, forward_batch, *, in_capture):
|
||||
self.calls.append((forward_batch, in_capture))
|
||||
|
||||
|
||||
class _FakeKVIndexKernel:
|
||||
def __getitem__(self, grid):
|
||||
del grid
|
||||
|
||||
def run(
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
starts,
|
||||
seq_lens,
|
||||
cu_seq_lens,
|
||||
output,
|
||||
req_to_token_stride,
|
||||
):
|
||||
del cu_seq_lens, req_to_token_stride
|
||||
cursor = 0
|
||||
for row in range(seq_lens.numel()):
|
||||
seq_len = int(seq_lens[row])
|
||||
start = int(starts[row])
|
||||
req = int(req_pool_indices[row])
|
||||
output[cursor : cursor + seq_len].copy_(
|
||||
req_to_token[req, start : start + seq_len]
|
||||
)
|
||||
cursor += seq_len
|
||||
|
||||
return run
|
||||
|
||||
|
||||
class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
|
||||
def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
|
||||
model_runner = SimpleNamespace(
|
||||
server_args=SimpleNamespace(
|
||||
chunked_prefill_size=16,
|
||||
context_length=None,
|
||||
cuda_graph_config=SimpleNamespace(
|
||||
prefill=SimpleNamespace(
|
||||
full_prefill_prefix_chunk_tokens=None, max_bs=8
|
||||
)
|
||||
),
|
||||
),
|
||||
req_to_token_pool=SimpleNamespace(
|
||||
req_to_token=torch.empty((1, 32), dtype=torch.int32)
|
||||
),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
|
||||
(4, 16),
|
||||
)
|
||||
|
||||
model_runner.server_args.chunked_prefill_size = -1
|
||||
self.assertEqual(
|
||||
PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
|
||||
(2, 8),
|
||||
)
|
||||
model_runner.server_args.chunked_prefill_size = 16
|
||||
|
||||
model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
|
||||
24
|
||||
)
|
||||
self.assertEqual(
|
||||
PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
|
||||
(6, 24),
|
||||
)
|
||||
|
||||
model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
|
||||
256
|
||||
)
|
||||
self.assertEqual(
|
||||
PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
|
||||
(32, 128),
|
||||
)
|
||||
|
||||
# At least one token is reserved per request lane even if the requested
|
||||
# aggregate capacity is smaller than the fixed request-slot count.
|
||||
model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
|
||||
2
|
||||
)
|
||||
self.assertEqual(
|
||||
PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
|
||||
(1, 4),
|
||||
)
|
||||
|
||||
model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
|
||||
0
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "must be positive"):
|
||||
PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4)
|
||||
|
||||
def test_buffers_are_shared_across_token_buckets(self):
|
||||
backend = _FakeAttentionBackend()
|
||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||
runner._capture_req_slots = 3
|
||||
runner._prefix_chunk_len = 2
|
||||
runner._prefix_chunk_capacity = 6
|
||||
runner._prefix_max_len = 8
|
||||
runner._prefix_capture_variants = (1, 2, 4)
|
||||
runner.device = torch.device("cpu")
|
||||
runner._prefill_static_buffers = {
|
||||
"extend_prefix_lens": torch.zeros(3, dtype=torch.int64),
|
||||
"req_pool_indices": torch.tensor([2, 0, 1], dtype=torch.int64),
|
||||
}
|
||||
runner._prefix_capture_batches = {}
|
||||
runner._prefix_capture_buffers = runner._create_chunked_prefix_buffers()
|
||||
runner.model_runner = SimpleNamespace(
|
||||
attn_backend=backend,
|
||||
req_to_token_pool=SimpleNamespace(
|
||||
req_to_token=torch.arange(24, dtype=torch.int32).view(3, 8)
|
||||
),
|
||||
)
|
||||
|
||||
first = SimpleNamespace()
|
||||
second = SimpleNamespace()
|
||||
first_key = ShapeKey(size=8, variant_label="chunked_prefix:4")
|
||||
second_key = ShapeKey(size=16, variant_label="chunked_prefix:4")
|
||||
|
||||
with patch.object(
|
||||
runner_module,
|
||||
"create_chunked_prefix_cache_kv_indices",
|
||||
_FakeKVIndexKernel(),
|
||||
):
|
||||
runner._prepare_chunked_prefix_capture(first, first_key, 4)
|
||||
runner._prepare_chunked_prefix_capture(second, second_key, 4)
|
||||
|
||||
buffers = runner._prefix_capture_buffers
|
||||
self.assertIsNotNone(buffers)
|
||||
# Chunk starts are constant and prefilled at allocation.
|
||||
self.assertEqual(
|
||||
buffers.starts_cpu.tolist(),
|
||||
[[0, 0, 0], [2, 2, 2], [4, 4, 4], [6, 6, 6]],
|
||||
)
|
||||
self.assertEqual(first.extend_prefix_lens_cpu, [8, 8, 8])
|
||||
self.assertEqual(first.prefix_chunk_num_tokens, [6, 6, 6, 6])
|
||||
self.assertIs(first.prefix_chunk_starts, buffers.starts)
|
||||
self.assertIs(first.prefix_chunk_seq_lens, buffers.seq_lens)
|
||||
self.assertIs(first.prefix_chunk_cu_seq_lens, buffers.cu_seq_lens)
|
||||
self.assertIs(first.prefix_chunk_starts, second.prefix_chunk_starts)
|
||||
self.assertIs(first.prefix_chunk_seq_lens, second.prefix_chunk_seq_lens)
|
||||
self.assertIs(
|
||||
first.prefix_chunk_cu_seq_lens,
|
||||
second.prefix_chunk_cu_seq_lens,
|
||||
)
|
||||
# Per-chunk KV indices are views of one shared 2-D buffer; what
|
||||
# capture bakes into the graph is the address, so compare pointers.
|
||||
for kv_chunk_idx in (0, 3):
|
||||
self.assertEqual(
|
||||
first.prefix_chunk_kv_indices[kv_chunk_idx].data_ptr(),
|
||||
buffers.kv_indices[kv_chunk_idx].data_ptr(),
|
||||
)
|
||||
self.assertEqual(
|
||||
first.prefix_chunk_kv_indices[kv_chunk_idx].data_ptr(),
|
||||
second.prefix_chunk_kv_indices[kv_chunk_idx].data_ptr(),
|
||||
)
|
||||
|
||||
runner._prepare_chunked_prefix_replay(
|
||||
second_key,
|
||||
SimpleNamespace(batch_size=2, extend_prefix_lens_cpu=[5, 1]),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
second.prefix_chunk_seq_lens.tolist(),
|
||||
[[2, 1, 0], [2, 0, 0], [1, 0, 0], [0, 0, 0]],
|
||||
)
|
||||
self.assertEqual(
|
||||
second.prefix_chunk_kv_indices[0].tolist(),
|
||||
[16, 17, 0, 0, 0, 0],
|
||||
)
|
||||
self.assertEqual(
|
||||
second.prefix_chunk_kv_indices[1].tolist(),
|
||||
[18, 19, 0, 0, 0, 0],
|
||||
)
|
||||
self.assertEqual(
|
||||
second.prefix_chunk_kv_indices[2].tolist(),
|
||||
[20, 0, 0, 0, 0, 0],
|
||||
)
|
||||
self.assertEqual(second.prefix_chunk_kv_indices[3].tolist(), [0] * 6)
|
||||
self.assertEqual(
|
||||
backend.calls,
|
||||
[(first, True), (second, True), (second, False)],
|
||||
)
|
||||
|
||||
def test_prefix_gate_only_applies_to_chunked_prefix_variant(self):
|
||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||
runner._capture_req_slots = 4
|
||||
runner.enable_lora = False
|
||||
runner.capture_hidden_mode = None
|
||||
runner.max_num_tokens = 32
|
||||
runner.capture_num_tokens = [4]
|
||||
runner.backend = SimpleNamespace()
|
||||
runner.prefill_backend_name = Backend.FULL
|
||||
runner.has_mha_companion_layers = False
|
||||
runner._prefix_chunk_len = 2
|
||||
runner._prefix_capture_variants = (1, 2, 4)
|
||||
|
||||
forward_batch = SimpleNamespace(
|
||||
batch_size=1,
|
||||
input_ids=torch.zeros(4, dtype=torch.int64),
|
||||
input_embeds=None,
|
||||
replace_embeds=None,
|
||||
forward_mode=SimpleNamespace(is_target_verify=lambda: False),
|
||||
capture_hidden_mode=None,
|
||||
global_num_tokens_cpu=None,
|
||||
return_logprob=False,
|
||||
extend_prefix_lens_cpu=[8],
|
||||
)
|
||||
|
||||
# Prefix hits in BCG/TC-piecewise and ordinary non-MLA FullCG use the
|
||||
# normal graph topology and must retain their existing eligibility.
|
||||
runner._capture_chunked_prefix = False
|
||||
for is_full_backend in (False, True):
|
||||
with self.subTest(is_full_backend=is_full_backend):
|
||||
runner._is_full_backend = is_full_backend
|
||||
self.assertTrue(runner.can_run_graph(forward_batch))
|
||||
|
||||
# The dedicated chunked-prefix topology has a fixed captured capacity.
|
||||
runner._is_full_backend = True
|
||||
runner._capture_chunked_prefix = True
|
||||
self.assertTrue(runner.can_run_graph(forward_batch))
|
||||
self.assertEqual(
|
||||
runner._shape_key(4, forward_batch).variant_label,
|
||||
"chunked_prefix:4",
|
||||
)
|
||||
forward_batch.batch_size = 2
|
||||
# Capacity is per request, not a sum: three real chunks round up to the
|
||||
# four-chunk graph even though the aggregate prefix has eight tokens.
|
||||
forward_batch.extend_prefix_lens_cpu = [5, 3]
|
||||
self.assertTrue(runner.can_run_graph(forward_batch))
|
||||
self.assertEqual(
|
||||
runner._shape_key(4, forward_batch).variant_label,
|
||||
"chunked_prefix:4",
|
||||
)
|
||||
forward_batch.extend_prefix_lens_cpu = [9, 1]
|
||||
self.assertFalse(runner.can_run_graph(forward_batch))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user