diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 6f52307b2..ecb54d9c4 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -77,6 +77,7 @@ class ScheduleBatchDisaggregationDecodeMixin: self.req_pool_indices = torch.tensor( req_pool_indices, dtype=torch.int64, device=self.device ) + self.req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64) self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=self.device) self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) self.orig_seq_lens = torch.tensor( diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 7d5a2dd1c..e03005546 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -444,9 +444,8 @@ class HiSparseCoordinator: out_cache_loc: torch.Tensor, req_pool_indices: torch.Tensor, seq_lens_cpu: torch.Tensor, + req_pool_indices_cpu: torch.Tensor, ) -> None: - req_pool_indices_cpu = req_pool_indices.cpu() - self._eager_backup_previous_token( seq_lens, req_pool_indices, seq_lens_cpu, req_pool_indices_cpu ) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 644f29560..ccb7c00a7 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1504,6 +1504,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): req_pool_indices: torch.Tensor = None # shape: [b], int64 seq_lens: torch.Tensor = None # shape: [b], int64 seq_lens_cpu: torch.Tensor = None # shape: [b], int64 + # CPU mirror of req_pool_indices; schedule-path only, stale in spec draft window + req_pool_indices_cpu: torch.Tensor = None # shape: [b], int64 # The output locations of the KV cache out_cache_loc: torch.Tensor = None # shape: [b], int64 @@ -1847,7 +1849,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.extend_num_tokens = extend_num_tokens # Allocate memory - out_cache_loc, req_pool_indices_tensor, _ = alloc_for_extend(self) + out_cache_loc, req_pool_indices_tensor, req_pool_indices_cpu = alloc_for_extend( + self + ) # Set fields input_embeds = [] @@ -2006,6 +2010,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.input_ids = input_ids_tensor self.req_pool_indices = req_pool_indices_tensor + self.req_pool_indices_cpu = req_pool_indices_cpu self.orig_seq_lens = orig_seq_lens_tensor self.out_cache_loc = out_cache_loc self.input_embeds = ( @@ -2379,6 +2384,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.orig_seq_lens = torch.empty(0, dtype=torch.int32, device=self.device) self.out_cache_loc = torch.empty(0, dtype=torch.int64, device=self.device) self.req_pool_indices = torch.empty(0, dtype=torch.int64, device=self.device) + self.req_pool_indices_cpu = torch.empty(0, dtype=torch.int64) self.seq_lens_sum = 0 self.extend_num_tokens = 0 self.sampling_info = SamplingBatchInfo.from_schedule_batch( @@ -2471,6 +2477,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.out_cache_loc, self.req_pool_indices, self.seq_lens_cpu, + self.req_pool_indices_cpu, ) if get_global_server_args().enable_mamba_extra_buffer(): @@ -2530,6 +2537,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): if self.multimodal_inputs is not None: self.multimodal_inputs = [self.multimodal_inputs[i] for i in keep_indices] self.req_pool_indices = self.req_pool_indices[keep_indices_device] + self.req_pool_indices_cpu = self.req_pool_indices_cpu[keep_indices] self.seq_lens = self.seq_lens[keep_indices_device] self.seq_lens_cpu = self.seq_lens_cpu[keep_indices] self.orig_seq_lens = self.orig_seq_lens[keep_indices_device] @@ -2582,6 +2590,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.req_pool_indices = torch.cat( [self.req_pool_indices, other.req_pool_indices] ) + self.req_pool_indices_cpu = torch.cat( + [self.req_pool_indices_cpu, other.req_pool_indices_cpu] + ) self.seq_lens = torch.cat([self.seq_lens, other.seq_lens]) self.seq_lens_cpu = torch.cat([self.seq_lens_cpu, other.seq_lens_cpu]) self.orig_seq_lens = torch.cat([self.orig_seq_lens, other.orig_seq_lens]) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 94d31c6c3..643bb4f94 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -428,14 +428,14 @@ def alloc_req_slots( def alloc_for_extend( batch: ScheduleBatch, -) -> tuple[torch.Tensor, torch.Tensor, list[int]]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Allocate KV cache for extend batch and write to req_to_token_pool. Returns: out_cache_loc: allocated cache locations - req_pool_indices_device: request pool indices at a device tensor - req_pool_indices: request pool indices as list + req_pool_indices_device: request pool indices as a device tensor + req_pool_indices_cpu: request pool indices as a CPU tensor (host mirror) """ # free out-of-window swa tokens batch.maybe_evict_swa() @@ -489,7 +489,7 @@ def alloc_for_extend( batch.req_to_token_pool, ) - return out_cache_loc, req_pool_indices_device, req_pool_indices + return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu def alloc_paged_token_slots_decode(