fix(disagg): correct DSA/SWA state-page transfer mismatch in PD disaggregation (#27004)
This commit is contained in:
@@ -98,9 +98,19 @@ def group_concurrent_contiguous(
|
||||
src_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32]
|
||||
) -> Tuple[List[npt.NDArray[np.int32]], List[npt.NDArray[np.int32]]]:
|
||||
"""Vectorised NumPy implementation."""
|
||||
if src_indices.size == 0:
|
||||
# src/dst indices are transferred pairwise, so an empty side means there is
|
||||
# nothing to transfer. Guarding both sides (not just src) avoids a cryptic
|
||||
# NumPy broadcast error from np.diff() below when only one side is empty, e.g.
|
||||
# a non-empty prefill DSA/SWA state list paired with an empty decode registration.
|
||||
if src_indices.size == 0 or dst_indices.size == 0:
|
||||
return [], []
|
||||
|
||||
if src_indices.size != dst_indices.size:
|
||||
raise ValueError(
|
||||
"group_concurrent_contiguous requires equal-length src/dst index arrays, "
|
||||
f"got {src_indices.size} and {dst_indices.size}"
|
||||
)
|
||||
|
||||
brk = np.where((np.diff(src_indices) != 1) | (np.diff(dst_indices) != 1))[0] + 1
|
||||
src_groups = np.split(src_indices, brk)
|
||||
dst_groups = np.split(dst_indices, brk)
|
||||
|
||||
@@ -947,7 +947,13 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
if last_chunk:
|
||||
self.disagg_metadata_buffers.set_buf(req)
|
||||
|
||||
seq_len = len(req.fill_ids)
|
||||
# fill_ids includes the token sampled during prefill, but decode
|
||||
# registers state pages over origin_input_ids (DecodePreallocQueue)
|
||||
# and the main pool send is clamped to end_idx above. Matching that
|
||||
# length here avoids emitting an extra state page when the sampled
|
||||
# token crosses a page boundary, which mismatched src/dst lengths in
|
||||
# group_concurrent_contiguous.
|
||||
seq_len = min(len(req.fill_ids), len(req.origin_input_ids))
|
||||
|
||||
def _mamba_payload():
|
||||
return [
|
||||
|
||||
Reference in New Issue
Block a user