[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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user