[Bug][PD][NIXL] always send aux on is_last; only expects_state when truthy (#25699)
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
co-authored by
Shangming Cai
parent
c2a212bfe2
commit
87c3c96bc8
@@ -688,8 +688,9 @@ class NixlKVManager(CommonKVManager):
|
|||||||
|
|
||||||
handles.append(kv_xfer_handle)
|
handles.append(kv_xfer_handle)
|
||||||
|
|
||||||
if kv_chunk.is_last and kv_chunk.state_indices:
|
if kv_chunk.is_last:
|
||||||
dst_info = self.decode_kv_args_table[req.agent_name]
|
dst_info = self.decode_kv_args_table[req.agent_name]
|
||||||
|
if kv_chunk.state_indices:
|
||||||
state_xfer_handles = self.maybe_send_extra(
|
state_xfer_handles = self.maybe_send_extra(
|
||||||
req.agent_name,
|
req.agent_name,
|
||||||
kv_chunk.state_indices,
|
kv_chunk.state_indices,
|
||||||
@@ -702,7 +703,9 @@ class NixlKVManager(CommonKVManager):
|
|||||||
dst_state_item_lens=dst_info.dst_state_item_lens,
|
dst_state_item_lens=dst_info.dst_state_item_lens,
|
||||||
dst_state_dim_per_tensor=dst_info.dst_state_dim_per_tensor,
|
dst_state_dim_per_tensor=dst_info.dst_state_dim_per_tensor,
|
||||||
)
|
)
|
||||||
handles.extend(h for h in state_xfer_handles if h is not None)
|
handles.extend(
|
||||||
|
h for h in state_xfer_handles if h is not None
|
||||||
|
)
|
||||||
|
|
||||||
if kv_chunk.prefill_aux_index is None:
|
if kv_chunk.prefill_aux_index is None:
|
||||||
raise RuntimeError("Missing aux index for last chunk")
|
raise RuntimeError("Missing aux index for last chunk")
|
||||||
@@ -1974,8 +1977,11 @@ class NixlKVReceiver(CommonKVReceiver):
|
|||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
# Mark that we expect state data if state_indices was provided
|
# Mark that we expect state data if state_indices was provided.
|
||||||
if state_indices is not None:
|
# Match the prefill-side truthy check: an empty list means the
|
||||||
|
# model has no state types (e.g. dense LLaMA/Qwen), and prefill
|
||||||
|
# won't send state notifs, so we must not expect them.
|
||||||
|
if state_indices:
|
||||||
self.kv_mgr.transfer_statuses[self.bootstrap_room].expects_state = True
|
self.kv_mgr.transfer_statuses[self.bootstrap_room].expects_state = True
|
||||||
|
|
||||||
self.started_transfer = True
|
self.started_transfer = True
|
||||||
|
|||||||
Reference in New Issue
Block a user