xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 02:07:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>
This commit is contained in:
Woosuk Kwon 2025-09-19 06:50:56 +00:00
parent b405d78c07
commit 0d3de9e082
2 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
vllm/v1/worker/gpu/dist_utils.py
View File

@ -5,6 +5,18 @@ import torch
from vllm.distributed import tensor_model_parallel_all_gather from vllm.distributed import tensor_model_parallel_all_gather
def evenly_split(
n: int,
tp_size: int,
tp_rank: int,
) -> tuple[int, int]:
q = n // tp_size
r = n % tp_size
start = q * tp_rank + min(tp_rank, r)
end = start + q + (1 if tp_rank < r else 0)
return start, end
def pad_and_all_gather( def pad_and_all_gather(
x: torch.Tensor, x: torch.Tensor,
padded_size: int, padded_size: int,

13
vllm/v1/worker/gpu/model_runner.py
View File

@ -23,7 +23,7 @@ from vllm.v1.sample.sampler import SamplerOutput
from vllm.v1.worker.gpu.attn_utils import (get_kv_cache_spec, from vllm.v1.worker.gpu.attn_utils import (get_kv_cache_spec,
init_attn_backend, init_kv_cache) init_attn_backend, init_kv_cache)
from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.dist_utils import pad_and_all_gather from vllm.v1.worker.gpu.dist_utils import evenly_split, pad_and_all_gather
from vllm.v1.worker.gpu.input_batch import (InputBatch, InputBuffers, from vllm.v1.worker.gpu.input_batch import (InputBatch, InputBuffers,
prepare_inputs) prepare_inputs)
from vllm.v1.worker.gpu.sampler import Sampler from vllm.v1.worker.gpu.sampler import Sampler
@ -366,11 +366,11 @@ class GPUModelRunner:
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
input_batch: InputBatch, input_batch: InputBatch,
) -> SamplerOutput: ) -> SamplerOutput:
num_reqs = logits.shape[0]
sample_hidden_states = hidden_states[input_batch.logits_indices] sample_hidden_states = hidden_states[input_batch.logits_indices]
logits = self.model.compute_logits(sample_hidden_states, None) logits = self.model.compute_logits(sample_hidden_states, None)
pos = input_batch.positions[input_batch.logits_indices] pos = input_batch.positions[input_batch.logits_indices]
idx_mapping_np = input_batch.idx_mapping_np idx_mapping_np = input_batch.idx_mapping_np
num_reqs = logits.shape[0]
# When the batch size is large enough, use DP sampler. # When the batch size is large enough, use DP sampler.
tp_group = get_tp_group() tp_group = get_tp_group()
@ -378,14 +378,9 @@ class GPUModelRunner:
n = (num_reqs + tp_size - 1) // tp_size n = (num_reqs + tp_size - 1) // tp_size
use_dp_sampler = tp_size > 1 and n > 32 use_dp_sampler = tp_size > 1 and n > 32
if use_dp_sampler: if use_dp_sampler:
# Shard the inputs as evenly as possible. # NOTE(woosuk): Make sure that no rank gets zero requests.
# Make sure that no rank gets zero requests.
tp_rank = tp_group.rank tp_rank = tp_group.rank
q = num_reqs // tp_size start, end = evenly_split(num_reqs, tp_size, tp_rank)
r = num_reqs % tp_size
start = q * tp_rank + min(tp_rank, r)
end = start + q + (1 if tp_rank < r else 0)
logits = logits[start:end] logits = logits[start:end]
pos = pos[start:end] pos = pos[start:end]
idx_mapping_np = idx_mapping_np[start:end] idx_mapping_np = idx_mapping_np[start:end]