[HiSparse-pd] Add device-buffer budget and fix logical pool admission in decode side (#22453)
This commit is contained in:
@@ -673,6 +673,21 @@ class DecodePreallocQueue:
|
|||||||
failed_reqs.append(decode_req)
|
failed_reqs.append(decode_req)
|
||||||
indices_to_remove.add(i)
|
indices_to_remove.add(i)
|
||||||
|
|
||||||
|
# HiSparse physical constraint: max requests by device buffer capacity.
|
||||||
|
# Each admitted req needs padded_buffer_size from hisparse device pool.
|
||||||
|
# waiting_queue reqs already have device buffers (allocated in admit_request_direct),
|
||||||
|
# only transfer_queue reqs are pending device buffer allocation.
|
||||||
|
hisparse_req_budget = float("inf")
|
||||||
|
if self.scheduler.enable_hisparse:
|
||||||
|
hisparse_avail = (
|
||||||
|
self.token_to_kv_pool_allocator.hisparse_attn_allocator.available_size()
|
||||||
|
)
|
||||||
|
hisparse_req_budget = max(
|
||||||
|
0,
|
||||||
|
hisparse_avail // self.scheduler.hisparse_coordinator.padded_buffer_size
|
||||||
|
- len(self.transfer_queue.queue),
|
||||||
|
)
|
||||||
|
|
||||||
# Then, preallocate the remaining requests if possible
|
# Then, preallocate the remaining requests if possible
|
||||||
for i, decode_req in enumerate(self.queue):
|
for i, decode_req in enumerate(self.queue):
|
||||||
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
|
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
|
||||||
@@ -690,6 +705,9 @@ class DecodePreallocQueue:
|
|||||||
if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0:
|
if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0:
|
||||||
break
|
break
|
||||||
|
|
||||||
|
if hisparse_req_budget <= 0:
|
||||||
|
break
|
||||||
|
|
||||||
# Memory estimation: don't add if the projected memory cannot be met
|
# Memory estimation: don't add if the projected memory cannot be met
|
||||||
# TODO: add new_token ratio
|
# TODO: add new_token ratio
|
||||||
origin_input_len = len(decode_req.req.origin_input_ids)
|
origin_input_len = len(decode_req.req.origin_input_ids)
|
||||||
@@ -714,6 +732,7 @@ class DecodePreallocQueue:
|
|||||||
break
|
break
|
||||||
|
|
||||||
allocatable_tokens -= required_tokens_for_request
|
allocatable_tokens -= required_tokens_for_request
|
||||||
|
hisparse_req_budget -= 1
|
||||||
dst_kv_indices = self._pre_alloc(decode_req.req)
|
dst_kv_indices = self._pre_alloc(decode_req.req)
|
||||||
|
|
||||||
origin_input_len = len(decode_req.req.origin_input_ids)
|
origin_input_len = len(decode_req.req.origin_input_ids)
|
||||||
@@ -819,7 +838,14 @@ class DecodePreallocQueue:
|
|||||||
and len(self.scheduler.running_batch.reqs) > 0
|
and len(self.scheduler.running_batch.reqs) > 0
|
||||||
else 0
|
else 0
|
||||||
)
|
)
|
||||||
available_size = self.token_to_kv_pool_allocator.available_size()
|
if self.scheduler.enable_hisparse:
|
||||||
|
# HiSparse pre-alloc only allocates logical indices (alloc_logical_only),
|
||||||
|
# so the logical pool is the binding constraint for admission control.
|
||||||
|
available_size = (
|
||||||
|
self.token_to_kv_pool_allocator.logical_attn_allocator.available_size()
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
available_size = self.token_to_kv_pool_allocator.available_size()
|
||||||
allocatable_tokens = available_size - max(
|
allocatable_tokens = available_size - max(
|
||||||
# preserve some space for future decode
|
# preserve some space for future decode
|
||||||
self.num_reserved_decode_tokens
|
self.num_reserved_decode_tokens
|
||||||
|
|||||||
Reference in New Issue
Block a user