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