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 10, 10:15 UTC
[Refactor] Remove numpy split in async scheduling (#32034)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 9, 2026 at 14:09 UTC
28ae32a5d3717d4d5fc73c2338cca909775910c0
1 file changed
+6
-12
vllm/v1/engine/async_llm.py
+6
-12
@@ -9,7 +9,6 @@ from collections.abc import AsyncGenerator, Iterable, Mapping
9
from copy import copy
10
from typing import Any, cast
11
12
-import numpy as np
12
import torch
13
14
import vllm.envs as envs
@@ -32,7 +31,6 @@ from vllm.transformers_utils.config import maybe_register_config_serialize_by_va
31
from vllm.usage.usage_lib import UsageContext
32
from vllm.utils.async_utils import cancel_task_threadsafe
33
from vllm.utils.collection_utils import as_list
35
-from vllm.utils.math_utils import cdiv
34
from vllm.v1.engine import EngineCoreRequest
35
from vllm.v1.engine.core_client import EngineCoreClient
36
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
@@ -495,6 +493,7 @@ class AsyncLLM(EngineClient):
493
log_stats = self.log_stats
494
logger_manager = self.logger_manager
495
input_processor = self.input_processor
496
+ chunk_size = envs.VLLM_V1_OUTPUT_PROC_CHUNK_SIZE
497
498
async def output_handler():
499
try:
@@ -510,15 +509,10 @@ class AsyncLLM(EngineClient):
509
# Split outputs into chunks of at most
510
# VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the
511
# event loop for too long.
513
- if num_outputs <= envs.VLLM_V1_OUTPUT_PROC_CHUNK_SIZE:
514
- slices = (outputs.outputs,)
515
- else:
516
- slices = np.array_split(
517
- outputs.outputs,
518
- cdiv(num_outputs, envs.VLLM_V1_OUTPUT_PROC_CHUNK_SIZE),
519
- )
520
-
521
- for i, outputs_slice in enumerate(slices):
512
+ engine_core_outputs = outputs.outputs
513
+ for start in range(0, num_outputs, chunk_size):
514
+ end = start + chunk_size
515
+ outputs_slice = engine_core_outputs[start:end]
516
# 2) Process EngineCoreOutputs.
517
processed_outputs = output_processor.process_outputs(
518
outputs_slice, outputs.timestamp, iteration_stats
@@ -527,7 +521,7 @@ class AsyncLLM(EngineClient):
521
assert not processed_outputs.request_outputs
522
523
# Allow other asyncio tasks to run between chunks
530
- if i + 1 < len(slices):
524
+ if end < num_outputs:
525
await asyncio.sleep(0)
526
527
# 3) Abort any reqs that finished due to stop strings.