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

Batch CP HiCache backup submits across requests

Browse Source 
CP HiCache write reservations must stay per radix node, but the transfer descriptor does not need to be per request. This changes the layer-end hook to group pending write states for the same source and layer, so bs>1 prefill emits one target D2H descriptor and one draft D2H descriptor per layer while preserving per-node metadata, rollback, and ack semantics.\n\nConstraint: CP shared-KV HiCache metadata, host slots, and radix acknowledgements remain per request/node.\nConstraint: TAI direct transfer kernels already accept flattened page descriptors, so no tai-kernel change is required.\nRejected: Merge HiCache reservations or radix nodes | would complicate rollback and split handling.\nRejected: Add collective synchronization for grouped backup | grouping is local descriptor construction and must not add rank-level sync.\nConfidence: high\nScope-risk: moderate\nDirective: Keep target and draft source notifications separate; final ack must wait for both when draft HiCache is attached.\nTested: local py_compile for cache_controller.py and test_hicache_controller_cp.py\nTested: local git diff --check\nTested: remote pytest test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 85 passed, 3 warnings\nNot-tested: full ETE bs>1 CP HiCache replay with admission gate removed\nNot-tested: Nsight/throughput validation of reduced D2H submit count
This commit is contained in:
laoyao0822
2026-06-03 10:23:30 +08:00
parent 19dcd6c4dc
commit b3913046b6
3 changed files with 655 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

296
python/sglang/srt/managers/cache_controller.py
View File

@@ -452,18 +452,12 @@ class HiCacheController:
return
if source not in ("target", "draft"):
raise ValueError(f"Unknown CP HiCache layer backup source: {source}")
reservations = [
state.reservation for state in list(self.pending_layer_writes.values())
]
for reservation in reservations:
if reservation.node_id not in self.pending_layer_writes:
continue
self.submit_write_cp_layer(
reservation,
layer_id,
submit_target=(source == "target"),
submit_draft=(source == "draft"),
)
self._submit_write_cp_layer_states(
list(self.pending_layer_writes.values()),
layer_id,
submit_target=(source == "target"),
submit_draft=(source == "draft"),
)
def _start_storage_threads(self):
"""Start storage prefetch/backup threads and their queues.
@@ -1148,6 +1142,193 @@ class HiCacheController:
state.start_event.record()
return state
def _layer_write_state_done(self, state: HiCacheLayerWriteState) -> bool:
target_done = len(state.completed_target_layers) >= self.layer_num
draft_done = (
state.draft_host_indices is None
or len(state.completed_draft_layers)
>= self.draft_mem_pool_device.layer_num
)
return target_done and draft_done
@staticmethod
def _record_tensor_on_stream(tensor: Optional[torch.Tensor], stream) -> None:
if tensor is not None and tensor.is_cuda:
tensor.record_stream(stream)
@staticmethod
def _concat_layer_write_tensors(tensors: List[torch.Tensor]) -> torch.Tensor:
if len(tensors) == 1:
return tensors[0]
return torch.cat(tensors)
def _append_layer_write_ack(self, state: HiCacheLayerWriteState) -> None:
if state.ack_appended:
return
state.ack_appended = True
reservation = state.reservation
self.pending_layer_writes.pop(reservation.node_id, None)
self.ack_write_queue.append(
HiCacheAck(state.start_event, state.finish_event, [reservation.node_id])
)
logger.info(
"[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s",
reservation.node_id,
reservation.metadata.logical_len,
state.total_layers,
reservation.draft_host_indices is not None,
)
def _submit_write_cp_layer_states(
self,
states: List[HiCacheLayerWriteState],
layer_id: int,
*,
submit_target: bool = True,
submit_draft: bool = True,
) -> None:
"""Submit one layer for a group of reserved CP host backups.
Radix/HiCache metadata remains per node. Only the transfer descriptor is
grouped so bs>1 batches issue one target D2H and one draft D2H per layer
instead of one D2H per request per layer.
"""
if not states:
return
unique_states: List[HiCacheLayerWriteState] = []
seen_state_ids = set()
for state in states:
state_id = id(state)
if state_id in seen_state_ids:
continue
seen_state_ids.add(state_id)
unique_states.append(state)
for state in unique_states:
if layer_id < 0 or layer_id >= state.total_layers:
raise ValueError(
f"layer_id={layer_id} is outside CP HiCache backup layer range "
f"[0, {state.total_layers}) for node_id={state.reservation.node_id}"
)
target_states = [
state
for state in unique_states
if submit_target
and layer_id < self.layer_num
and layer_id not in state.completed_target_layers
]
draft_states = [
state
for state in unique_states
if submit_draft
and state.draft_host_indices is not None
and layer_id < self.draft_mem_pool_device.layer_num
and layer_id not in state.completed_draft_layers
]
if not target_states and not draft_states:
return
active_states: List[HiCacheLayerWriteState] = []
seen_active_ids = set()
for state in target_states + draft_states:
state_id = id(state)
if state_id in seen_active_ids:
continue
seen_active_ids.add(state_id)
active_states.append(state)
layer_event = device_module.Event()
layer_event.record()
for state in active_states:
state.layer_events.append(layer_event)
target_transfer_states = [
state
for state in target_states
if state.physical_device_indices is not None
and len(state.physical_device_indices) > 0
]
draft_transfer_states = [
state
for state in draft_states
if state.physical_device_indices is not None
and len(state.physical_device_indices) > 0
]
grouped_tensors: List[torch.Tensor] = []
final_states: List[HiCacheLayerWriteState] = []
with device_module.stream(self.write_stream):
for state in active_states:
state.start_event.wait(self.write_stream)
layer_event.wait(self.write_stream)
if target_transfer_states:
target_host_indices = self._concat_layer_write_tensors(
[state.host_indices for state in target_transfer_states]
)
target_device_indices = self._concat_layer_write_tensors(
[
state.physical_device_indices
for state in target_transfer_states
]
)
self.mem_pool_host.backup_from_device_per_layer(
self.mem_pool_device,
target_host_indices,
target_device_indices,
layer_id,
self.io_backend,
)
grouped_tensors.extend([target_host_indices, target_device_indices])
if draft_transfer_states:
draft_host_indices = self._concat_layer_write_tensors(
[state.draft_host_indices for state in draft_transfer_states]
)
draft_device_indices = self._concat_layer_write_tensors(
[
state.physical_device_indices
for state in draft_transfer_states
]
)
self.draft_mem_pool_host.backup_from_device_per_layer(
self.draft_mem_pool_device,
draft_host_indices,
draft_device_indices,
layer_id,
self.io_backend,
)
grouped_tensors.extend([draft_host_indices, draft_device_indices])
for tensor in grouped_tensors:
self._record_tensor_on_stream(tensor, self.write_stream)
for state in target_states:
state.completed_target_layers.add(layer_id)
for state in draft_states:
state.completed_draft_layers.add(layer_id)
final_states = [
state
for state in active_states
if self._layer_write_state_done(state) and not state.ack_appended
]
for state in final_states:
state.finish_event.record()
self._record_tensor_on_stream(state.host_indices, self.write_stream)
self._record_tensor_on_stream(
state.physical_device_indices, self.write_stream
)
self._record_tensor_on_stream(
state.draft_host_indices, self.write_stream
)
for state in final_states:
self._append_layer_write_ack(state)
def submit_write_cp_layer(
self,
reservation: HiCacheWriteReservation,
@@ -1163,92 +1344,11 @@ class HiCacheController:
"""
state = self._get_or_create_layer_write_state(reservation)
if layer_id < 0 or layer_id >= state.total_layers:
raise ValueError(
f"layer_id={layer_id} is outside CP HiCache backup layer range "
f"[0, {state.total_layers}) for node_id={reservation.node_id}"
)
needs_target = (
submit_target
and layer_id < self.layer_num
and layer_id not in state.completed_target_layers
)
needs_draft = (
submit_draft
and state.draft_host_indices is not None
and layer_id < self.draft_mem_pool_device.layer_num
and layer_id not in state.completed_draft_layers
)
if not needs_target and not needs_draft:
return
layer_event = device_module.Event()
layer_event.record()
state.layer_events.append(layer_event)
if len(reservation.physical_device_indices) > 0:
with device_module.stream(self.write_stream):
state.start_event.wait(self.write_stream)
layer_event.wait(self.write_stream)
if needs_target:
self.mem_pool_host.backup_from_device_per_layer(
self.mem_pool_device,
state.host_indices,
state.physical_device_indices,
layer_id,
self.io_backend,
)
if needs_draft:
self.draft_mem_pool_host.backup_from_device_per_layer(
self.draft_mem_pool_device,
state.draft_host_indices,
state.physical_device_indices,
layer_id,
self.io_backend,
)
if needs_target:
state.completed_target_layers.add(layer_id)
if needs_draft:
state.completed_draft_layers.add(layer_id)
target_done = len(state.completed_target_layers) >= self.layer_num
draft_done = (
state.draft_host_indices is None
or len(state.completed_draft_layers)
>= self.draft_mem_pool_device.layer_num
)
if not target_done or not draft_done:
return
if state.ack_appended:
return
with device_module.stream(self.write_stream):
state.finish_event.record()
if state.host_indices is not None and state.host_indices.is_cuda:
state.host_indices.record_stream(self.write_stream)
if (
state.physical_device_indices is not None
and state.physical_device_indices.is_cuda
):
state.physical_device_indices.record_stream(self.write_stream)
if (
state.draft_host_indices is not None
and state.draft_host_indices.is_cuda
):
state.draft_host_indices.record_stream(self.write_stream)
state.ack_appended = True
self.pending_layer_writes.pop(reservation.node_id, None)
self.ack_write_queue.append(
HiCacheAck(state.start_event, state.finish_event, [reservation.node_id])
)
logger.info(
"[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s",
reservation.node_id,
reservation.metadata.logical_len,
state.total_layers,
reservation.draft_host_indices is not None,
self._submit_write_cp_layer_states(
[state],
layer_id,
submit_target=submit_target,
submit_draft=submit_draft,
)
def submit_write_cp_per_layer(