refactor(disagg): remove unreferenced dead code (#35838)
This commit is contained in:
@@ -65,7 +65,6 @@ class KVArgs:
|
|||||||
# per tensor when the single contiguous slice already matches the layout.
|
# per tensor when the single contiguous slice already matches the layout.
|
||||||
state_conv_shard_groups: List[List[Optional[List[int]]]]
|
state_conv_shard_groups: List[List[Optional[List[int]]]]
|
||||||
ib_device: str
|
ib_device: str
|
||||||
ib_traffic_class: str
|
|
||||||
gpu_id: int
|
gpu_id: int
|
||||||
kv_head_num: int
|
kv_head_num: int
|
||||||
total_kv_head_num: int
|
total_kv_head_num: int
|
||||||
|
|||||||
@@ -257,10 +257,6 @@ class StagingAllocator:
|
|||||||
offset, _, _ = self.allocations[alloc_id]
|
offset, _, _ = self.allocations[alloc_id]
|
||||||
return offset
|
return offset
|
||||||
|
|
||||||
def get_round(self, alloc_id: int) -> int:
|
|
||||||
_, _, rnd = self.allocations[alloc_id]
|
|
||||||
return rnd
|
|
||||||
|
|
||||||
def get_base_ptr(self) -> int:
|
def get_base_ptr(self) -> int:
|
||||||
return self.base_ptr
|
return self.base_ptr
|
||||||
|
|
||||||
|
|||||||
@@ -963,24 +963,6 @@ class SchedulerDisaggregationPrefillMixin:
|
|||||||
self.metrics_collector.increment_transfer_failed_reqs()
|
self.metrics_collector.increment_transfer_failed_reqs()
|
||||||
return exc
|
return exc
|
||||||
|
|
||||||
def get_transferred_rids(self: Scheduler) -> List[str]:
|
|
||||||
"""
|
|
||||||
Used by PP, get the transferred rids but **do not pop**
|
|
||||||
"""
|
|
||||||
polls = poll_and_all_reduce_attn_cp_tp_group(
|
|
||||||
[req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
|
|
||||||
self.attn_cp_cpu_group,
|
|
||||||
self.attn_tp_cpu_group,
|
|
||||||
)
|
|
||||||
|
|
||||||
transferred_rids: List[str] = []
|
|
||||||
|
|
||||||
for req, poll in zip(self.disagg_prefill_inflight_queue, polls):
|
|
||||||
if poll == KVPoll.Success or poll == KVPoll.Failed:
|
|
||||||
transferred_rids.append(req.rid)
|
|
||||||
|
|
||||||
return transferred_rids
|
|
||||||
|
|
||||||
def clear_pending_chunk_send(self: Scheduler, req: Req) -> None:
|
def clear_pending_chunk_send(self: Scheduler, req: Req) -> None:
|
||||||
"""Drop `req` from the sent-but-unconcluded chunk set.
|
"""Drop `req` from the sent-but-unconcluded chunk set.
|
||||||
|
|
||||||
|
|||||||
@@ -729,52 +729,6 @@ def _get_cp_rank_page_bounds(
|
|||||||
return local_start, local_start + n_pages
|
return local_start, local_start + n_pages
|
||||||
|
|
||||||
|
|
||||||
def page_indices_to_cp_rank_page_indices(
|
|
||||||
page_indices: np.ndarray,
|
|
||||||
total_pages: int,
|
|
||||||
cp_rank: int,
|
|
||||||
cp_size: int,
|
|
||||||
) -> np.ndarray:
|
|
||||||
"""
|
|
||||||
Filter page_indices (which are *global* page ids in the KV pool) to those
|
|
||||||
belonging to the given CP rank for this request.
|
|
||||||
|
|
||||||
For a single request, its pages occupy a contiguous global range
|
|
||||||
[first_page, first_page + total_pages). We first compute the local
|
|
||||||
split [0, total_pages) across cp_size ranks, then shift that local
|
|
||||||
range by first_page back into the global page id space and take
|
|
||||||
the intersection with page_indices.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Subset of page_indices that fall in this rank's global
|
|
||||||
[start_page, end_page) slice for the given CP rank.
|
|
||||||
"""
|
|
||||||
if cp_size <= 1:
|
|
||||||
return page_indices
|
|
||||||
|
|
||||||
if page_indices.size == 0:
|
|
||||||
return np.asarray(page_indices)
|
|
||||||
|
|
||||||
first_page = int(page_indices.min())
|
|
||||||
base = total_pages // cp_size
|
|
||||||
rem = total_pages % cp_size
|
|
||||||
|
|
||||||
if rem == 0:
|
|
||||||
local_start = cp_rank * base
|
|
||||||
local_end = local_start + base
|
|
||||||
else:
|
|
||||||
local_start = cp_rank * base + min(cp_rank, rem)
|
|
||||||
n_pages = base + (1 if cp_rank < rem else 0)
|
|
||||||
local_end = local_start + n_pages
|
|
||||||
|
|
||||||
# Map back to global page ids.
|
|
||||||
start_page = first_page + local_start
|
|
||||||
end_page = first_page + local_end
|
|
||||||
|
|
||||||
mask = (page_indices >= start_page) & (page_indices < end_page)
|
|
||||||
return np.asarray(page_indices)[mask]
|
|
||||||
|
|
||||||
|
|
||||||
def filter_kv_indices_for_cp_rank(
|
def filter_kv_indices_for_cp_rank(
|
||||||
kv_mgr: CommonKVManager,
|
kv_mgr: CommonKVManager,
|
||||||
kv_indices: np.ndarray,
|
kv_indices: np.ndarray,
|
||||||
|
|||||||
Reference in New Issue
Block a user