minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[core] Maintain req_pool_indices_cpu host mirror (like seq_lens_cpu) (#26425)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-26 18:59:54 -07:00
committed by GitHub
parent 6076066e38
commit 1051a8456f
4 changed files with 18 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/hisparse_coordinator.py
+1 -2
View File
@@ -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
)
python/sglang/srt/managers/schedule_batch.py
+12 -1
View File
@@ -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])