fix: piecewise_cuda_graph get correct qo_indptr (#21452)

Co-authored-by: Avery Huang <averyh@nvidia.com>

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -71,6 +71,7 @@ class ForwardContext:
         self.quant_config = None
         self.moe_layers = None
         self.moe_fusions = None
+        self.num_tokens: Optional[int] = None
 
     def set_forward_batch(self, forward_batch: ForwardBatch):
         self.forward_batch = forward_batch
@@ -104,6 +105,7 @@ def set_forward_context(
     quant_config: Any,
     moe_layers: List[Any],
     moe_fusions: List[Any],
+    num_tokens: Optional[int] = None,
 ):
     global _forward_context
     _forward_context = ForwardContext()
@@ -112,6 +114,7 @@ def set_forward_context(
     _forward_context.set_quant_config(quant_config)
     _forward_context.set_moe_layers(moe_layers)
     _forward_context.set_moe_fusions(moe_fusions)
+    _forward_context.num_tokens = num_tokens
     try:
         yield
     finally:

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

@@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
 import torch
 
 from sglang.kernel_api_logging import debug_kernel_api
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
+from sglang.srt.compilation.piecewise_context_manager import (
+    get_forward_context,
+    is_in_piecewise_cuda_graph,
+)
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -147,6 +150,7 @@ class FlashInferAttnBackend(AttentionBackend):
         self.max_context_len = model_runner.model_config.context_len
         self.skip_prefill = skip_prefill
         self.is_multimodal = model_runner.model_config.is_multimodal
+        self.page_size = model_runner.page_size
 
         assert not (
             model_runner.sliding_window_size is not None
@@ -1192,6 +1196,7 @@ class FlashInferIndicesUpdaterPrefill:
         self.q_data_type = model_runner.dtype
         self.sliding_window_size = model_runner.sliding_window_size
         self.attn_backend = attn_backend
+        self.page_size = attn_backend.page_size
 
         # Buffers and wrappers
         self.kv_indptr = attn_backend.kv_indptr
@@ -1381,8 +1386,13 @@ class FlashInferIndicesUpdaterPrefill:
             # Normal extend
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
+            # Reserve extra space in kv_indices for a potential piecewise CUDA graph
+            # dummy request (see below). Worst case: static_num_tokens extra pages.
+            fwd_ctx = get_forward_context()
+            pcg_num_tokens = fwd_ctx.num_tokens if fwd_ctx is not None else None
+            extra_kv = pcg_num_tokens if pcg_num_tokens is not None else 0
             kv_indices = torch.empty(
-                paged_kernel_lens_sum + 256,
+                paged_kernel_lens_sum + extra_kv + 256,
                 dtype=torch.int32,
                 device=req_pool_indices.device,
             )
@@ -1397,6 +1407,40 @@ class FlashInferIndicesUpdaterPrefill:
             )
             qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
             qo_indptr = qo_indptr[: bs + 1]
+
+            # Piecewise CUDA graph padding: input_ids are padded to static_num_tokens,
+            # so q.shape[0] == static_num_tokens but qo_indptr[-1] == actual tokens.
+            # Append a dummy request for the padding tokens so that
+            # qo_indptr[-1] == static_num_tokens, satisfying flashinfer's shape check
+            # without corrupting the causal masks of real requests.
+            # The dummy request's KV indices all point to slot 0 (a scratch location);
+            # its attention output is discarded via the [:raw_num_tokens] slice in replay.
+            bs_eff = bs
+            # extend_num_tokens is a Python int (== sum of seq_lens - prefix_lens),
+            # and paged_kernel_lens_sum is also a Python int (== kv_indptr[-1]),
+            # so this block requires no CPU-GPU synchronisation.
+            actual_qo_tokens = (
+                fwd_ctx.forward_batch.extend_num_tokens if fwd_ctx is not None else None
+            )
+            if (
+                pcg_num_tokens is not None
+                and actual_qo_tokens is not None
+                and pcg_num_tokens > actual_qo_tokens
+            ):
+                pad_tokens = pcg_num_tokens - actual_qo_tokens
+                num_dummy_pages = (pad_tokens + self.page_size - 1) // self.page_size
+                kv_start = (
+                    paged_kernel_lens_sum  # equals kv_indptr[-1], no .item() needed
+                )
+                kv_indices[kv_start : kv_start + num_dummy_pages] = 0
+                qo_indptr = torch.cat(
+                    [qo_indptr, qo_indptr.new_tensor([pcg_num_tokens])]
+                )
+                kv_indptr = torch.cat(
+                    [kv_indptr, kv_indptr.new_tensor([kv_start + num_dummy_pages])]
+                )
+                bs_eff = bs + 1
+
             custom_mask = None
         else:
             assert isinstance(spec_info, SpecInput)
@@ -1408,6 +1452,7 @@ class FlashInferIndicesUpdaterPrefill:
                     self.req_to_token,
                 )
             )
+            bs_eff = bs
 
         # extend part
         if use_ragged:
@@ -1449,7 +1494,7 @@ class FlashInferIndicesUpdaterPrefill:
             qo_indptr,
             kv_indptr,
             kv_indices,
-            self.kv_last_page_len[:bs],
+            self.kv_last_page_len[:bs_eff],
             self.num_qo_heads,
             self.num_kv_heads,
             self.head_dim,

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -754,9 +754,13 @@ class PiecewiseCudaGraphRunner:
         forward_batch: ForwardBatch,
         **kwargs,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
+        num_tokens = len(forward_batch.input_ids)
+        index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
+        static_num_tokens = self.capture_num_tokens[index]
         with enable_piecewise_cuda_graph():
-            # Due to the dispatch kernel for MLA model, we init the metadata with original forward_batch
-            self.model_runner.attn_backend.init_forward_metadata(forward_batch)
+            # Prepare static buffers first so set_forward_context can carry num_tokens
+            # into call_begin_forward (via ForwardContext.num_tokens), eliminating the
+            # need for a separate global and allowing pre-calculation of dummy-page count.
             static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
             # Replay
             with set_forward_context(
@@ -765,7 +769,10 @@ class PiecewiseCudaGraphRunner:
                 self.quant_config,
                 self.moe_layers,
                 self.moe_fusions,
+                num_tokens=static_num_tokens,
             ):
+                # Due to the dispatch kernel for MLA model, we init the metadata with original forward_batch
+                self.model_runner.attn_backend.init_forward_metadata(forward_batch)
                 output = self.model_runner.model.forward(
                     static_forward_batch.input_ids,
                     static_forward_batch.positions,