From fde4004429400051a555632858ddcaf950d71901 Mon Sep 17 00:00:00 2001 From: weizhoublue <45163302+weizhoublue@users.noreply.github.com> Date: Wed, 10 Jun 2026 06:03:58 +0800 Subject: [PATCH] [Fix] Reset positions tensor in CUDA graph runner when batch size differs from captured size (#24401) Signed-off-by: weizhoublue --- .../cuda_graph_buffer_registry.py | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 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 0224cea3d..f5ef566d2 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -511,15 +511,16 @@ def build_decode_registry( source: Optional[Any] = None, ) -> CudaGraphBufferRegistry: """Registry mirroring the always-on (+ mamba / mrope) FB-shared decode - buffers, with padding policies matching - ``DecodeInputBuffers.populate_from_forward_batch``: + buffers, with the per-slot padding policy that resets the padded tail on + each replay: - ``seq_lens`` / ``seq_lens_cpu`` -> FILL_SENTINEL(seq_len_fill_value) - ``req_pool_indices`` / ``out_cache_loc`` / ``mamba_track_*`` -> ZERO - - ``input_ids`` / ``positions`` / ``mrope_positions`` -> FOREACH_COPY - (head ``[:raw_n]`` is always overwritten by the copy; the old code's - full-buffer ``zero_()`` / ``fill_()`` on ``bs != raw_bs`` is therefore - equivalent to the tail-only reset the policies apply here). + - ``positions`` / ``mrope_positions`` -> ZERO: the flashinfer verify-path + plan reads the padded tail, so leaving stale out-of-range values there + triggers an illegal memory access (issue #24361). + - ``input_ids`` -> FOREACH_COPY: head ``[:raw_n]`` is overwritten by the + copy and the padded tail is not read. ``custom_mask`` / ``next_token_logits_buffer`` / ``input_embeds`` are not registered here — they are not per-replay FB copies (allocated and written @@ -545,7 +546,13 @@ def build_decode_registry( slots = [ GraphSlot("input_ids", _tokens, torch.int64, axis="tokens"), - GraphSlot("positions", _tokens, torch.int64, axis="tokens"), + GraphSlot( + "positions", + _tokens, + torch.int64, + axis="tokens", + padding_policy=PaddingPolicy.ZERO, + ), GraphSlot( "out_cache_loc", _tokens, @@ -583,6 +590,7 @@ def build_decode_registry( torch.int64, axis="tokens", slice_fn=lambda buf, n: buf[:, :n], + padding_policy=PaddingPolicy.ZERO, ), ] if enable_mamba_track: