[Perf] DSv4.1: Make the breakable prefill CUDA graph usable above 16K contexts, and add a serving-shape startup warm-up

python/sglang/srt/entrypoints/deepseek_v4_warmup.py

@@ -0,0 +1,214 @@
+"""DeepSeek-V4 serving-shape warm-up (``--warmups deepseek_v4_serving_shapes``).
+
+Graph capture and the server's own warm-up run only short decode-shaped batches, so
+prefill sizes, prefix hits, chunked prompts, multi-request extends, eager decode above
+the decode graphs, mixed batches and prompts past the prefill graph's context bucket
+compiled their Triton and JIT kernel variants inside the first requests that reached
+them (about 8 s for the first long prompt on a cold cache). These internal requests take
+the real serving path at each shape class before the server listens, alternating greedy
+and the server's default sampling; the cache is flushed afterwards so no prefix or KV
+state survives.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from typing import TYPE_CHECKING, List, Optional, Sequence, Union
+
+import numpy as np
+
+from sglang.srt.managers.io_struct import GenerateReqInput
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.tokenizer_manager import TokenizerManager
+
+logger = logging.getLogger(__name__)
+
+_RID_PREFIX = "deepseek-v4-warmup-"
+
+# (concurrent, prompt lengths, max_new_tokens). "k+n" resends source prompt k plus n
+# new tokens, a prefix hit with an n-token extend ("k+0" extends one token). Triton
+# specializes integer arguments as 1, divisible by 16, or other, and the mHC prefill
+# kernel switches at 4,096 rows; 20,011 tokens take several chunks with a prefix.
+_PLAN = (
+    (False, (1, 13, 16, 211, 4096, 4817, 20011), 8),
+    (False, ("5+0", "5+7", "5+160", "6+4321", "6+0"), 8),
+    (True, (300, 5000), 16),
+    (True, (64, 2000, 9000), 16),
+    (True, (17, 32, 1100, 4100), 16),
+    (True, ("2+5", "3+40", "4+4099"), 16),
+)
+# A batched request reaches the scheduler as one message: one n-request extend, then
+# decode at batch n. Kernels keyed on next_power_of_2(bs) or bs * verify width see
+# every size, with a total token count divisible by 16 and not.
+_GROUP_SIZES = range(1, 17)
+# One eager verify step above the decode graphs, plus the second step at graph sizes.
+_GROUP_NEW_TOKENS = 4
+# Decodes that a prompt prefill joins (mixed batches): 1, then 2 decode rows. 256 tokens
+# outlast both 8K chunks of the prompt by more than 10x.
+_MIXED = ((1, 12000), (2, 9500))
+_MIXED_DECODE_TOKENS = 256
+# Past the prefill graph's context bucket prefill runs eager, where the mHC kernels keyed
+# on the row count compile on first use: an aligned 8K chunk, an unaligned tail, and
+# three extends on that prompt.
+_ABOVE_GRAPH_TAIL = (8192, 1001)
+_ABOVE_GRAPH_EXTENDS = (211, 2064, 5000)
+
+Length = Union[int, str]
+
+
+async def warm_serving_shapes(
+    disaggregation_mode: str, tokenizer_manager: TokenizerManager
+) -> None:
+    from sglang.srt.runtime_context import get_schedule
+
+    if disaggregation_mode != "null":
+        logger.warning("deepseek_v4_serving_shapes is for aggregated serving; skipped.")
+        return
+    await run_serving_shapes(
+        tokenizer_manager,
+        mixed_chunk=get_schedule().enable_mixed_chunk,
+        prefill_graph_context=_prefill_graph_context(),
+    )
+
+
+async def run_serving_shapes(
+    tokenizer_manager: TokenizerManager,
+    *,
+    mixed_chunk: bool,
+    prefill_graph_context: Optional[int],
+) -> None:
+    start = time.perf_counter()
+    warm = _Warmup(tokenizer_manager)
+    for concurrent, lengths, max_new_tokens in _PLAN:
+        prompts = warm.prompts(lengths)
+        if concurrent:
+            await asyncio.gather(*(warm.run(p, max_new_tokens) for p in prompts))
+        else:
+            for prompt in prompts:
+                await warm.run(prompt, max_new_tokens)
+    for size in _GROUP_SIZES:
+        for divisible in (True, False):
+            await warm.run_group(_group_lengths(size, divisible), greedy=divisible)
+    if mixed_chunk:
+        for num_decodes, prefill_len in _MIXED:
+            await warm.run_mixed(num_decodes, prefill_len)
+    if prefill_graph_context is not None:
+        await warm.run_above_graph(prefill_graph_context)
+    flushed = await _flush(tokenizer_manager)
+    logger.info(
+        "DeepSeek-V4 serving-shape warm-up: %d requests in %.1f s; cache flushed: %s.",
+        warm.count,
+        time.perf_counter() - start,
+        flushed,
+    )
+    if not flushed:
+        raise RuntimeError("deepseek_v4_serving_shapes could not flush the cache.")
+
+
+def _prefill_graph_context() -> Optional[int]:
+    """Longest context the breakable prefill graph replays, or None when it is off."""
+    from sglang.srt.model_executor.cuda_graph_config import Backend
+    from sglang.srt.runtime_context import get_exec
+
+    prefill = get_exec().graph.cuda_graph_config.prefill
+    if prefill.backend != Backend.BREAKABLE:
+        return None
+    bounds = [b for b in (prefill.max_context_size, prefill.max_seq_len) if b]
+    return min(bounds) if bounds else None
+
+
+def _group_lengths(size: int, divisible_by_16: bool) -> List[int]:
+    lengths = [64 + 37 * i for i in range(size)]
+    total = sum(lengths)
+    if divisible_by_16:
+        lengths[-1] += -total % 16
+    elif total % 16 == 0:
+        lengths[-1] += 1
+    return lengths
+
+
+async def _flush(tokenizer_manager: TokenizerManager) -> bool:
+    # The flush is refused while the last requests are still being released.
+    for _ in range(50):
+        if (await tokenizer_manager.flush_cache()).success:
+            return True
+        await asyncio.sleep(0.1)
+    return False
+
+
+async def _drain(gen) -> None:
+    async for _ in gen:
+        pass
+
+
+class _Warmup:
+    def __init__(self, tokenizer_manager: TokenizerManager):
+        self.tokenizer_manager = tokenizer_manager
+        self.rng = np.random.default_rng(0)
+        self.sources: List[List[int]] = []
+        self.count = 0
+
+    def prompts(self, lengths: Sequence[Length]) -> List[List[int]]:
+        out = []
+        for length in lengths:
+            if isinstance(length, str):
+                base, extra = (int(x) for x in length.split("+"))
+                out.append(self.sources[base] + self._tokens(extra))
+            else:
+                out.append(self._tokens(length))
+                self.sources.append(out[-1])
+        return out
+
+    def _tokens(self, n: int) -> List[int]:
+        return self.rng.integers(1000, 60000, size=n).tolist()
+
+    def _request(self, prompt, max_new_tokens: int, *, stream: bool = False):
+        self.count += 1
+        params = {"max_new_tokens": max_new_tokens, "ignore_eos": True}
+        if self.count % 2 == 0:
+            params["temperature"] = 0.0
+        return GenerateReqInput(
+            input_ids=prompt,
+            sampling_params=params,
+            rid=f"{_RID_PREFIX}{self.count}",
+            stream=stream,
+        )
+
+    async def run(self, prompt: List[int], max_new_tokens: int) -> None:
+        req = self._request(prompt, max_new_tokens)
+        await _drain(self.tokenizer_manager.generate_request(req, None))
+
+    async def run_group(self, lengths: List[int], *, greedy: bool) -> None:
+        prompts = self.prompts(lengths)
+        self.count += len(prompts)
+        params = {"max_new_tokens": _GROUP_NEW_TOKENS, "ignore_eos": True}
+        if greedy:
+            params["temperature"] = 0.0
+        req = GenerateReqInput(
+            input_ids=prompts,
+            sampling_params=params,
+            rid=[f"{_RID_PREFIX}{self.count}-{i}" for i in range(len(prompts))],
+        )
+        await _drain(self.tokenizer_manager.generate_request(req, None))
+
+    async def run_mixed(self, num_decodes: int, prefill_len: int) -> None:
+        streams = []
+        for prompt in self.prompts((500,) * num_decodes):
+            req = self._request(prompt, _MIXED_DECODE_TOKENS, stream=True)
+            gen = self.tokenizer_manager.generate_request(req, None)
+            await gen.__anext__()  # decoding from here on
+            streams.append(asyncio.ensure_future(_drain(gen)))
+        await self.run(self.prompts((prefill_len,))[0], 8)
+        await asyncio.gather(*streams)
+
+    async def run_above_graph(self, graph_context: int) -> None:
+        length = graph_context + sum(_ABOVE_GRAPH_TAIL)
+        if length + max(_ABOVE_GRAPH_EXTENDS) > self.tokenizer_manager.context_len:
+            return
+        base = len(self.sources)
+        await self.run(self.prompts((length,))[0], 2)
+        for prompt in self.prompts([f"{base}+{n}" for n in _ABOVE_GRAPH_EXTENDS]):
+            await self.run(prompt, 2)

python/sglang/srt/entrypoints/warmup.py

@@ -159,3 +159,13 @@ async def prefill_shapes(disaggregation_mode: str, tokenizer_manager: TokenizerM
             generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST
 
         await tokenizer_manager.generate_request(generate_req_input, None).__anext__()
+
+
+@warmup("deepseek_v4_serving_shapes")
+async def deepseek_v4_serving_shapes(
+    disaggregation_mode: str, tokenizer_manager: TokenizerManager
+):
+    """Compile DeepSeek-V4 serving-time kernel variants before the server listens."""
+    from sglang.srt.entrypoints.deepseek_v4_warmup import warm_serving_shapes
+
+    await warm_serving_shapes(disaggregation_mode, tokenizer_manager)

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -959,6 +959,18 @@ def _prefill_graph_max_seq_len() -> Optional[int]:
     return get_exec().graph.cuda_graph_config.prefill.max_seq_len
 
 
+def _low_ratio_graph_fits_candidates(
+    *, low_ratios: Tuple[int, ...], candidate_window: int, max_seq_len: Optional[int]
+) -> bool:
+    """The captured ratio-1/2 indexer takes the full top-k, which equals the candidate
+    top-k only while every compressed position of the graph's context fits the window."""
+    if not low_ratios or candidate_window <= 0:
+        return True
+    if max_seq_len is None:
+        return False
+    return max_seq_len // min(low_ratios) <= candidate_window
+
+
 @dataclass
 class DSV4Metadata:
     core_attn_metadata: DSV4AttnMetadata
@@ -1172,12 +1184,20 @@ def __init__(
         self.has_c128: bool = 128 in self.present_ratios
         cfg = model_runner.model_config.hf_text_config
         self.is_dsv41: bool = getattr(cfg, "model_type", None) == "deepseek_v41"
+        candidate_topk_blocks = getattr(cfg, "candidate_topk_blocks", 0)
+        candidate_block_size = getattr(cfg, "candidate_block_size", 0)
         self.prefill_candidates, self.decode_candidates = make_candidate_indexer(
             token_to_kv_pool=self.token_to_kv_pool,
             req_to_token=self.req_to_token,
             page_size=self.page_size,
-            candidate_topk_blocks=getattr(cfg, "candidate_topk_blocks", 0),
-            candidate_block_size=getattr(cfg, "candidate_block_size", 0),
+            candidate_topk_blocks=candidate_topk_blocks,
+            candidate_block_size=candidate_block_size,
+        )
+        # Wider prefill graphs keep ratio-1/2 sources in their eager break.
+        self._low_ratio_graph_fits_candidates = _low_ratio_graph_fits_candidates(
+            low_ratios=self.low_ratios,
+            candidate_window=candidate_topk_blocks * candidate_block_size,
+            max_seq_len=_prefill_graph_max_seq_len(),
         )
         self.full_topk_indexer = make_full_topk_indexer(
             token_to_kv_pool=self.token_to_kv_pool, req_to_token=self.req_to_token
@@ -1562,6 +1582,7 @@ def low_ratio_prefill_graph(self) -> bool:
             and has_dense_fp4_indexer()
             and is_sm100_or_newer()
             and get_parallel().attn_cp_size == 1
+            and self._low_ratio_graph_fits_candidates
         )
 
     def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
@@ -3034,6 +3055,14 @@ def forward_low_ratio_sources(
             # Bucket-sized under the prefill graph; an eager break sees the
             # live rows only and falls through.
             req, pos = hoisted_req, hoisted_pos
+        elif (
+            hoisted_req is not None
+            and meta.late_layer_tail is not None
+            and hoisted_req.shape[0] > positions.shape[0]
+        ):
+            # Late layers in the prefill graph: the eager break sees the tail block's
+            # real rows, which lead it; the batch's extend lens would cover every row.
+            req, pos = hoisted_req[: positions.shape[0]], positions
         else:
             req = token_req_indices(forward_batch, num_tokens=positions.shape[0])
             pos = positions
@@ -3219,6 +3248,8 @@ def _low_ratio_compress_fused(
 
         pool = self.token_to_kv_pool
         core = self.forward_metadata.core_metadata
+        # The prefill graph's eager break passes live rows against bucket-sized metadata.
+        num_tokens = pos.shape[0]
         compressor = layer.compressor
         layer_id = layer.layer_id
         # Contiguous complex64 freqs_cis gives a real/imag-interleaved view without copying.
@@ -3239,14 +3270,14 @@ def _low_ratio_compress_fused(
                 compressor.wkv(x),
                 compressor.norm.weight.data,
                 pos,
-                core.raw_out_loc,
+                core.raw_out_loc[:num_tokens],
                 compressor.norm.eps,
                 freqs_cis,
                 kv_cache,
                 page_size=page_size,
                 layout=kv_layout,
             )
-            out_loc = core.c1_out_loc
+            out_loc = core.c1_out_loc[:num_tokens]
         else:
             # CompressStatePool stores each request's pending pairs in a position ring.
             # KVAndScore rows use | kv | score |, addressed as req * ring_size + pos % ring_size.
@@ -3257,7 +3288,7 @@ def _low_ratio_compress_fused(
                 compressor.norm.weight.data,
                 pos,
                 req,
-                core.raw_out_loc,
+                core.raw_out_loc[:num_tokens],
                 compressor.norm.eps,
                 freqs_cis,
                 kv_cache,
@@ -3278,7 +3309,7 @@ def _low_ratio_compress_fused(
                     ring_size=state.ring_size,
                     layout=kv_layout,
                 )
-            out_loc = core.c2_out_loc
+            out_loc = core.c2_out_loc[:num_tokens]
 
         indexer = layer.indexer
         if indexer is not None and indexer.owns_k:

python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py

@@ -41,6 +41,7 @@
 
 from sglang.kernels.ops.attention.dsv4.dequant_k_cache import DIM_NOPE, DIM_ROPE
 from sglang.srt.utils import ceil_align
+from sglang.srt.utils.common import async_h2d
 
 # FlashMLA sparse prefill asserts ``params.topk % B_TOPK == 0``. B_TOPK is 64
 # for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both.
@@ -563,7 +564,9 @@ def ensure_compressed(
             seq_lens_cpu = self.seq_lens.tolist()
         lengths_cpu = [int(length) // compress_ratio for length in seq_lens_cpu]
         total_compressed = sum(lengths_cpu)
-        lengths = torch.tensor(lengths_cpu, dtype=torch.int32, device=device)
+        # Pinned upload: inside the breakable prefill graph the host runs ahead,
+        # and a pageable copy would wait for every queued graph segment.
+        lengths = async_h2d(lengths_cpu, dtype=torch.int32, device=device)
         ends = lengths.cumsum(0, dtype=torch.int32)
         compressed_base = ends - lengths
         swa_base = total_compressed + self.swa_offsets[:-1]

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -802,6 +802,31 @@ def test_sparse_prefill_compressed_gather_packs_request_lengths(self):
                 )
                 self.assertIs(cache.ensure_compressed(ratio, table, 64), gather)
 
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
+    def test_sparse_prefill_gather_does_not_wait_for_queued_gpu_work(self):
+        """The gather runs in an eager break of the prefill graph, where the host is
+        ahead of the GPU; uploading its lengths must not drain the queued replay."""
+
+        def cache_on_cuda():
+            cache = self._make_sparse_prefill_cache(257)
+            for name in ("seq_lens", "query_start_loc", "query_pos", "swa_offsets"):
+                setattr(cache, name, getattr(cache, name).cuda())
+            cache.seq_lens_cpu = [257, 257]
+            return cache
+
+        table = torch.zeros((2, 64), dtype=torch.int32, device="cuda")
+        # The first call loads its kernels, which synchronizes on its own.
+        cache_on_cuda().ensure_compressed(1, table, c_page_size=64)
+        cache = cache_on_cuda()
+        torch.cuda.synchronize()
+        replay_done = torch.cuda.Event()
+        # About 0.5 s of queued GPU work stands in for the replayed graph segments.
+        torch.cuda._sleep(1_000_000_000)
+        replay_done.record()
+        cache.ensure_compressed(1, table, c_page_size=64)
+        self.assertFalse(replay_done.query(), "the gather waited for queued GPU work")
+        torch.cuda.synchronize()
+
     def test_sparse_prefill_c128_uses_live_extent(self):
         from sglang.srt.layers.attention.dsv4 import sparse_prefill_utils