[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:
ishandhanani
2026-05-19 10:04:04 +08:00
committed by GitHub
co-authored by Shangming Cai
parent c2a212bfe2
commit 87c3c96bc8
+10 -4
View File
@@ -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