Confine req-pool-idx assignment to the pool allocator (#25726)

2026-05-19 09:22:51 +08:00
parent 6bcbf4de35
commit 5067da3eb6
1 changed files with 1 additions and 4 deletions
@@ -1774,9 +1774,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
        self.extend_num_tokens = extend_num_tokens

        # Allocate memory
-        out_cache_loc, req_pool_indices_tensor, req_pool_indices = alloc_for_extend(
-            self
-        )
+        out_cache_loc, req_pool_indices_tensor, _ = alloc_for_extend(self)

        # Set fields
        input_embeds = []
@@ -1792,7 +1790,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
        mamba_track_seqlens_cpu = []

        for i, (req, seq_len, pre_len) in enumerate(zip(reqs, seq_lens, prefix_lens)):
-            req.req_pool_idx = req_pool_indices[i]
            assert seq_len - pre_len == req.extend_input_len

            req.extend_batch_idx += 1