diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index 20a08a997..a49e9ad47 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 257cadf15..4fe8aec31 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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, diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 3fd7c3adc..b1935d21c 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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, diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_abort.py index ce53425d9..9643ce686 100644 --- a/test/registered/scheduler/test_abort.py +++ b/test/registered/scheduler/test_abort.py @@ -317,6 +317,8 @@ class TestAbortWithWaitingTimeout(WaitingTimeoutMixin, CustomTestCase): timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=[ "--max-running-requests=1", + # Disable PCG to avoid padding in flashinfer backend. Ref: https://github.com/sgl-project/sglang/pull/21452 + "--disable-piecewise-cuda-graph", ], ) diff --git a/test/registered/scheduler/test_priority_scheduling.py b/test/registered/scheduler/test_priority_scheduling.py index 339db350a..72321b694 100644 --- a/test/registered/scheduler/test_priority_scheduling.py +++ b/test/registered/scheduler/test_priority_scheduling.py @@ -41,6 +41,8 @@ class TestPriorityScheduling(CustomTestCase): "--max-queued-requests", # Enforce max queued request number is 3 "3", "--enable-priority-scheduling", # Enable priority scheduling + # Disable PCG to avoid padding in flashinfer backend. Ref: https://github.com/sgl-project/sglang/pull/21452 + "--disable-piecewise-cuda-graph", ), return_stdout_stderr=(cls.stdout, cls.stderr), ) @@ -247,6 +249,7 @@ class TestPrioritySchedulingMultipleRunningRequests(CustomTestCase): "--max-queued-requests", # Enforce max queued request number is 3 "3", "--enable-priority-scheduling", # Enable priority scheduling + "--disable-piecewise-cuda-graph", ), return_stdout_stderr=(cls.stdout, cls.stderr), )