A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
[Bugfix][Disaggregated] patch the inflight batching on the decode node in SimpleConnector to avoid hangs in SimpleBuffer (nccl based) (#13987)
Signed-off-by: Mathis Felardos <mathis@mistral.ai>
Mathis Felardos committed
Feb 28, 2025 at 08:53 UTC
b9e41734c583a5d2593dfa1ee17d609deca917e0
1 file changed
+15
-2
vllm/distributed/kv_transfer/kv_connector/simple_connector.py
+15
-2
@@ -214,6 +214,7 @@ class SimpleConnector(KVConnectorBase):
214
215
input_tokens_tensor = model_input.input_tokens
216
seq_lens = model_input.attn_metadata.seq_lens
217
+ num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
218
slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
219
220
hidden_or_intermediate_states_for_one_req = []
@@ -225,9 +226,21 @@ class SimpleConnector(KVConnectorBase):
226
# enumerate different requests
227
# FIXME(Kuntai): This impl assumes that all requests are prefill.
228
for idx, slen in enumerate(seq_lens):
228
-
229
start_pos = sum(seq_lens[:idx])
230
end_pos = start_pos + slen
231
+
232
+ if start_pos >= num_prefill_tokens:
233
+ # This can happen during inflight batching. See:
234
+ # vllm/worker/model_runner.py::_prepare_model_input_tensors:
235
+ # - input_tokens[:num_prefill_tokens] contains prefill tokens.
236
+ # - input_tokens[num_prefill_tokens:] contains decode tokens.
237
+ logger.warning("You should set --enable_chunked_prefill=False "
238
+ "and --max_num_batched_tokens "
239
+ "should be equal to max_seq_len_to_capture")
240
+ bypass_model_exec = False
241
+ assert start_pos == num_prefill_tokens
242
+ break
243
+
244
current_tokens = input_tokens_tensor[start_pos:end_pos]
245
num_tokens = slen
246
@@ -288,7 +301,7 @@ class SimpleConnector(KVConnectorBase):
301
# Here we will fall back to normal model forwarding
302
# But optionally you can adjust model_input so that you only do
303
# prefilling on those tokens that are missing KV caches.
291
- logger.debug(
304
+ logger.warning(
305
"[rank%d]: Failed to receive all KVs and hidden "
306
"states, redo model forwarding.", torch.distributed.get_rank())
307
hidden_or_intermediate_states = None