vllm-output
vllm/vllm/v1/core/sched/scheduler.py · ef9339e · L2349–2458 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
2349 stopped_preempted_reqs: set[Request] = set()
2350 # Stopped resumable requests to re-enqueue, in output order.
2351 continuing_reqs: list[Request] = []
2352 for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
2353 assert num_tokens_scheduled > 0
2354 request = self.requests.get(req_id)
2355 output_is_stale = False
2356 if request is not None:
2357 request.num_in_flight_tokens -= num_tokens_scheduled
2358 # Drain any stale share (see _preempt_request) in lockstep.
2359 if request.num_stale_output_tokens > 0:
2360 output_is_stale = True
2361 request.num_stale_output_tokens -= num_tokens_scheduled
2362 assert request.num_stale_output_tokens >= 0
2363 if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
2364 # skip failed or rescheduled requests from KV load failure
2365 continue
2366 if request is None or request.is_finished():
2367 # The request is already finished. This can happen if the
2368 # request is aborted while the model is executing it (e.g.,
2369 # in pipeline parallelism or in async scheduling).
2370 # NOTE(Kuntai): When delay_free_blocks=True (for async KV
2371 # cache transfer in KV connector), the aborted request will not
2372 # be set to None (in order to finish async KV transfer).
2373 # In this case, we use is_finished() to check.
2374 continue
2375
2376 # Drop-mode stale output (same-step resume) is discarded entirely.
2377 if output_is_stale and request.drop_stale_output:
2378 continue
2379
2380 req_index = model_runner_output.req_id_to_index[req_id]
2381 generated_token_ids = (
2382 sampled_token_ids[req_index] if sampled_token_ids else []
2383 )
2384
2385 scheduled_spec_token_ids = (
2386 scheduler_output.scheduled_spec_decode_tokens.get(req_id)
2387 )
2388 if scheduled_spec_token_ids and (
2389 generated_token_ids or self.num_sampled_tokens_per_step == 0
2390 ):
2391 num_draft_tokens = len(scheduled_spec_token_ids)
2392 num_sampled = self.num_sampled_tokens_per_step
2393 num_accepted = max(len(generated_token_ids) - num_sampled, 0)
2394 num_rejected = num_draft_tokens - num_accepted
2395 # Rejections roll back num_computed_tokens (and, under async
2396 # scheduling, num_output_placeholders, which covers the spec
2397 # tokens). A stale rejection count predates the preemption
2398 # rollback and must not apply.
2399 if not output_is_stale:
2400 if request.num_computed_tokens > 0:
2401 request.num_computed_tokens -= num_rejected
2402 if request.num_output_placeholders > 0:
2403 request.num_output_placeholders -= num_rejected
2404 spec_decoding_stats = self.make_spec_decoding_stats(
2405 spec_decoding_stats,
2406 num_draft_tokens=num_draft_tokens,
2407 num_accepted_tokens=num_accepted,
2408 num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
2409 request_id=req_id,
2410 )
2411 if request.spec_decode_metrics is not None:
2412 # Exclude grammar-invalidated drafts from the proposed
2413 # count, mirroring make_spec_decoding_stats; the accepted
2414 # bucket (j) is unaffected.
2415 adj_draft_tokens = num_draft_tokens
2416 if scheduler_output.num_invalid_spec_tokens:
2417 adj_draft_tokens -= (
2418 scheduler_output.num_invalid_spec_tokens.get(req_id, 0)
2419 )
2420 request.spec_decode_metrics.observe(
2421 num_draft_tokens=adj_draft_tokens,
2422 num_accepted=num_accepted,
2423 detailed=self.spec_decode_metrics_level == "detailed",
2424 )
2425
2426 # Free encoder inputs only after the step has actually executed.
2427 if request.has_encoder_inputs:
2428 self._free_encoder_inputs(request)
2429
2430 stopped = False
2431 new_logprobs = None
2432 new_sampling_mask = None
2433 new_token_ids = generated_token_ids
2434 pooler_output = pooler_outputs[req_index] if pooler_outputs else None
2435 kv_transfer_params = None
2436 ec_transfer_params = None
2437 prefill_stats = None
2438 watermarking = None
2439 status_before_stop = request.status
2440
2441 # Check for stop and update request status.
2442 if new_token_ids:
2443 has_output_placeholders = request.num_output_placeholders > 0
2444 new_token_ids, stopped = self._update_request_with_output(
2445 request, new_token_ids, is_stale=output_is_stale
2446 )
2447 if (
2448 not output_is_stale
2449 and status_before_stop == RequestStatus.RUNNING
2450 and (has_output_placeholders or self.num_spec_tokens > 0)
2451 ):
2452 self.kv_cache_manager.cache_blocks(
2453 request,
2454 request.num_computed_tokens - request.num_in_flight_tokens,
2455 )
2456 elif request.pooling_params and pooler_output is not None:
2457 # Pooling stops as soon as there is output.
2458 request.status = RequestStatus.FINISHED_STOPPED