[pd]: (Bug Fix) Incorrect out_cache_loc slicing in prepare_for_prebuilt (#24230)

Co-authored-by: Shangming Cai <csmthu@gmail.com>
2026-05-03 18:35:16 +08:00
parent 2bfc5d3bb1
commit 44ca2d01fc
2 changed files with 37 additions and 2 deletions
@@ -42,9 +42,10 @@ class ScheduleBatchDisaggregationDecodeMixin:
        offset = 0
        for i, req in enumerate(reqs):
            req_pool_indices.append(req.req_pool_idx)
+            pre_len = len(req.prefix_indices)

            chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][
-                : req.extend_input_len
+                pre_len : pre_len + req.extend_input_len
            ]
            assert (
                offset + req.extend_input_len <= total_size
@@ -52,7 +53,6 @@ class ScheduleBatchDisaggregationDecodeMixin:
            out_cache_loc[offset : offset + req.extend_input_len] = chunk
            offset += req.extend_input_len

-            pre_len = len(req.prefix_indices)
            seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1)
            seq_lens.append(seq_len)
            if len(req.output_ids) == 0: