fix(disagg): correct DSA/SWA state-page transfer mismatch in PD disaggregation (#27004)

This commit is contained in:
Kevin Flansburg
2026-06-03 14:33:41 +08:00
committed by GitHub
parent dae86f51f5
commit 52f2fe456a
3 changed files with 63 additions and 2 deletions
@@ -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)
+7 -1
View File
@@ -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 [