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 13, 09:30 UTC
[Refactor] Remove redundant TP gather/split in split_qkv in QwenVL (#28271)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Canlin Guo committed
Nov 12, 2025 at 23:56 UTC
bc5bd45c7d1abbac4a63d97d383212c108e55308
2 files changed
+1
-42
vllm/model_executor/models/qwen2_5_vl.py
-30
@@ -291,25 +291,6 @@ class Qwen2_5_VisionMLP(nn.Module):
291
return x_down
292
293
294
-def all_gather_interleave(local_tensor, hidden_size: int, tp_size: int):
295
- """All-gather the input tensor interleavely across model parallel group."""
296
- import torch.distributed as dist
297
-
298
- gathered_tensors = [torch.zeros_like(local_tensor) for _ in range(tp_size)]
299
- dist.all_gather(
300
- gathered_tensors, local_tensor, group=parallel_state.get_tp_group().device_group
301
- )
302
-
303
- gathered_tensors_split = [
304
- torch.split(tensor, hidden_size // tp_size, -1) for tensor in gathered_tensors
305
- ]
306
- ordered_tensors = [
307
- tensor for pair in zip(*gathered_tensors_split) for tensor in pair
308
- ]
309
- result_tensor = torch.cat(ordered_tensors, dim=-1)
310
- return result_tensor
311
-
312
-
294
class Qwen2_5_VisionAttention(nn.Module):
295
def __init__(
296
self,
@@ -383,21 +364,10 @@ class Qwen2_5_VisionAttention(nn.Module):
364
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
365
# [s, b, 3 * head * head_dim]
366
seq_len, bs, _ = qkv.shape
386
- if self.tp_size > 1:
387
- qkv = all_gather_interleave(qkv, self.qkv.hidden_size, self.tp_size)
367
368
# [s, b, 3 * head * head_dim] -> 3 * [s, b, head * head_dim]
369
q, k, v = qkv.chunk(3, dim=2)
370
392
- # 3 * [s, b, head * head_dim]
393
- if self.tp_size > 1:
394
- splitter = partial(
395
- dist_utils.split_tensor_along_last_dim, num_partitions=self.tp_size
396
- )
397
- q = splitter(q)[self.tp_rank]
398
- k = splitter(k)[self.tp_rank]
399
- v = splitter(v)[self.tp_rank]
400
-
371
# 3 * [s, b, head * head_dim] -> 3 * [s, b, head, head_dim]
372
new_shape = (
373
seq_len,
vllm/model_executor/models/qwen2_vl.py
+1
-12
@@ -50,7 +50,7 @@ from vllm.attention.layer import (
50
)
51
from vllm.config import VllmConfig
52
from vllm.config.multimodal import BaseDummyOptions
53
-from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
53
+from vllm.distributed import parallel_state
54
from vllm.distributed import utils as dist_utils
55
from vllm.logger import init_logger
56
from vllm.model_executor.layers.activation import QuickGELU
@@ -396,21 +396,10 @@ class Qwen2VisionAttention(nn.Module):
396
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
397
# [s, b, 3 * head * head_dim]
398
seq_len, bs, _ = qkv.shape
399
- if self.tp_size > 1:
400
- qkv = tensor_model_parallel_all_gather(qkv)
399
400
# [s, b, 3 * head * head_dim] -> 3 * [s, b, head * head_dim]
401
q, k, v = qkv.chunk(3, dim=2)
402
405
- # 3 * [s, b, head * head_dim]
406
- if self.tp_size > 1:
407
- splitter = partial(
408
- dist_utils.split_tensor_along_last_dim, num_partitions=self.tp_size
409
- )
410
- q = splitter(q)[self.tp_rank]
411
- k = splitter(k)[self.tp_rank]
412
- v = splitter(v)[self.tp_rank]
413
-
403
# 3 * [s, b, head * head_dim] -> 3 * [s, b, head, head_dim]
404
new_shape = (
405
seq_len,