From 6b232c19b80f0200acb2f15c11ab72b0606fc967 Mon Sep 17 00:00:00 2001 From: alichen Date: Fri, 10 Jul 2026 04:00:49 -0700 Subject: [PATCH] Fix PP proxy buffer sizing and refresh for speculative verify Four related fixes to the pipeline-parallel proxy tensor plumbing that only manifest when PP is combined with speculative decoding (multi-token verify); plain decode (num_tokens_per_bs == 1) is bit-identical: 1. _allocate_decode_buffers / DecodeInputBuffers.create: hidden_states and residual proxy buffers were sized (max_bs, hidden) while every other token-axis buffer (input_ids, positions, and topk_indices in the same dict) uses max_num_token = max_bs * num_tokens_per_bs. Under TARGET_VERIFY the [:num_tokens] slice silently returned fewer rows, crashing rotary with mismatched query/positions during warmup. 2. DecodeCudaGraphRunner.load_batch: the pre-planned early-return path (taken when eagle_prepare_for_verify already ran load_batch) copied input_ids and positions but never refreshed the pp_proxy_tensors input buffers, so the last PP stage replayed verify graphs against stale hidden states. 3. DecodeCudaGraphRunner.execute: the PPProxyTensors output was sliced with [:self.bs] (request rows) instead of [:self.bs * self.num_tokens_per_bs] (token rows). With verify (3 tokens/request) only the first bs rows were forwarded downstream, corrupting every request after the first in a microbatch while single-request runs looked healthy. 4. cuda_graph_buffer_registry: the pp-proxy slot source indexed ppx.tensors[key] unconditionally; an entry can legitimately be absent (e.g. topk_indices when a DSA model runs a dense attention backend). Use .get() so the established None-skips-copy contract applies. Co-Authored-By: Claude Fable 5 --- .../cuda_graph_buffer_registry.py | 5 +++- .../srt/model_executor/runner/base_runner.py | 9 ++++++-- .../runner/decode_cuda_graph_runner.py | 23 ++++++++++++++++++- .../model_executor/runner_utils/buffers.py | 7 ++++-- 4 files changed, 38 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index a90b9bef9527..a4dde273af0f 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -742,7 +742,10 @@ def _fn(fb, _ctx): def _pp_source(key): def _fn(_fb, ctx): ppx = ctx.pp_proxy_tensors - return None if ppx is None else ppx.tensors[key] + # .get(): a proxy entry can be absent (e.g. topk_indices + # when a DSA model runs a dense attention backend); + # returning None skips the copy for that slot. + return None if ppx is None else ppx.tensors.get(key) return _fn diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 56420661a9d7..ee0e2266df64 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -110,12 +110,17 @@ def _allocate_decode_buffers( # mHC (e.g. DSV4) flattens residual into hidden_states (size = hc_hidden_size). is_mhc = hc_hidden_size is not None hs = hc_hidden_size if is_mhc else hidden_size + # Sized in tokens, not requests: under speculative decoding the + # verify forward carries num_tokens_per_bs tokens per request and + # _dummy_run slices these buffers to num_tokens (same as + # topk_indices below). Identical for plain decode where + # num_tokens_per_bs == 1. pp_proxy_tensors = { - "hidden_states": torch.zeros((max_bs, hs), dtype=dtype), + "hidden_states": torch.zeros((max_num_token, hs), dtype=dtype), } if not is_mhc: pp_proxy_tensors["residual"] = torch.zeros( - (max_bs, hidden_size), dtype=dtype + (max_num_token, hidden_size), dtype=dtype ) if pp_proxy_topk_size is not None: pp_proxy_tensors["topk_indices"] = torch.zeros( diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 83a1b4ac5425..5b7dabdb727a 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -912,6 +912,19 @@ def load_batch( # In speculative decoding, these two fields are still needed. self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) + if ( + pp_proxy_tensors is not None + and self.buffers.pp_proxy_tensors is not None + ): + # PP + spec verify: the pre-planned load ran without the proxy + # (eagle_prepare_for_verify has no access to it), so the + # graph's proxy input buffers must be refreshed here -- the + # captured graph reads these rows (mirrors fill_from's + # side-slot copy). + for k, v in pp_proxy_tensors.tensors.items(): + buf = self.buffers.pp_proxy_tensors.get(k) + if buf is not None: # skip markers like __msg_type__ + buf[: v.shape[0]].copy_(v) if ( self.model_runner.spec_algorithm.is_dflash() and self.model_runner.is_draft_worker @@ -1056,7 +1069,15 @@ def execute( ) else: assert isinstance(output, PPProxyTensors) - return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()}) + # Slice in token rows, not request rows: under speculative verify + # each request carries num_tokens_per_bs tokens (identical for + # plain decode, where num_tokens_per_bs == 1). + return PPProxyTensors( + { + k: v[: self.bs * self.num_tokens_per_bs] + for k, v in output.tensors.items() + } + ) def get_spec_info(self, num_tokens: int): spec_info = None diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 00d32a24945c..78b0c4ce408e 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -132,12 +132,15 @@ def create( if pp_size > 1: is_mhc = hc_hidden_size is not None hs = hc_hidden_size if is_mhc else hidden_size + # Sized in tokens, not requests: under speculative decoding + # the verify forward carries num_tokens_per_bs tokens per + # request (same as topk_indices below). pp_proxy_tensors = { - "hidden_states": torch.zeros((max_bs, hs), dtype=dtype), + "hidden_states": torch.zeros((max_num_token, hs), dtype=dtype), } if not is_mhc: pp_proxy_tensors["residual"] = torch.zeros( - (max_bs, hidden_size), dtype=dtype + (max_num_token, hidden_size), dtype=dtype ) if pp_proxy_topk_size is not None: pp_proxy_tensors["topk_indices"] = torch.zeros(