minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Bug][PD][NIXL] always send aux on is_last; only expects_state when truthy (#25699)

Browse Source 
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
1 changed files with 22 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/nixl/conn.py
+22 -16
View File
@@ -688,21 +688,24 @@ class NixlKVManager(CommonKVManager):
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]
state_xfer_handles = self.maybe_send_extra(
req.agent_name,
kv_chunk.state_indices,
dst_info.dst_state_data_ptrs,
req.dst_state_indices,
dst_info.gpu_id,
f"{req.room}_state_{self.kv_args.engine_rank}",
decode_tp_size,
decode_tp_rank=dst_info.decode_tp_rank,
dst_state_item_lens=dst_info.dst_state_item_lens,
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)
if kv_chunk.state_indices:
state_xfer_handles = self.maybe_send_extra(
req.agent_name,
kv_chunk.state_indices,
dst_info.dst_state_data_ptrs,
req.dst_state_indices,
dst_info.gpu_id,
f"{req.room}_state_{self.kv_args.engine_rank}",
decode_tp_size,
decode_tp_rank=dst_info.decode_tp_rank,
dst_state_item_lens=dst_info.dst_state_item_lens,
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
)
if kv_chunk.prefill_aux_index is None:
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
if state_indices is not None:
# Mark that we expect state data if state_indices was provided.
# 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.started_transfer = True