[Fix] Reset positions tensor in CUDA graph runner when batch size differs from captured size (#24401)

Signed-off-by: weizhoublue <weizhou.lan@daocloud.io>
This commit is contained in:
weizhoublue
2026-06-09 15:03:58 -07:00
committed by GitHub
parent 2fef951fe8
commit fde4004429
@@ -511,15 +511,16 @@ def build_decode_registry(
source: Optional[Any] = None, source: Optional[Any] = None,
) -> CudaGraphBufferRegistry: ) -> CudaGraphBufferRegistry:
"""Registry mirroring the always-on (+ mamba / mrope) FB-shared decode """Registry mirroring the always-on (+ mamba / mrope) FB-shared decode
buffers, with padding policies matching buffers, with the per-slot padding policy that resets the padded tail on
``DecodeInputBuffers.populate_from_forward_batch``: each replay:
- ``seq_lens`` / ``seq_lens_cpu`` -> FILL_SENTINEL(seq_len_fill_value) - ``seq_lens`` / ``seq_lens_cpu`` -> FILL_SENTINEL(seq_len_fill_value)
- ``req_pool_indices`` / ``out_cache_loc`` / ``mamba_track_*`` -> ZERO - ``req_pool_indices`` / ``out_cache_loc`` / ``mamba_track_*`` -> ZERO
- ``input_ids`` / ``positions`` / ``mrope_positions`` -> FOREACH_COPY - ``positions`` / ``mrope_positions`` -> ZERO: the flashinfer verify-path
(head ``[:raw_n]`` is always overwritten by the copy; the old code's plan reads the padded tail, so leaving stale out-of-range values there
full-buffer ``zero_()`` / ``fill_()`` on ``bs != raw_bs`` is therefore triggers an illegal memory access (issue #24361).
equivalent to the tail-only reset the policies apply here). - ``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 ``custom_mask`` / ``next_token_logits_buffer`` / ``input_embeds`` are not
registered here — they are not per-replay FB copies (allocated and written registered here — they are not per-replay FB copies (allocated and written
@@ -545,7 +546,13 @@ def build_decode_registry(
slots = [ slots = [
GraphSlot("input_ids", _tokens, torch.int64, axis="tokens"), 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( GraphSlot(
"out_cache_loc", "out_cache_loc",
_tokens, _tokens,
@@ -583,6 +590,7 @@ def build_decode_registry(
torch.int64, torch.int64,
axis="tokens", axis="tokens",
slice_fn=lambda buf, n: buf[:, :n], slice_fn=lambda buf, n: buf[:, :n],
padding_policy=PaddingPolicy.ZERO,
), ),
] ]
if enable_mamba_track: if enable_mamba_track: