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 14, 05:30 UTC
Remove redundant all gather + split (#23441)
Co-authored-by: Chenxi Yang <cxyang@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
Chenxi Yang committed
Sep 10, 2025 at 23:45 UTC
d13360183a429562ff7d62235534ef37e712409c
1 file changed
-13
vllm/model_executor/models/glm4_1v.py
-13
@@ -272,23 +272,10 @@ class Glm4vVisionAttention(nn.Module):
272
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
273
# [s, b, 3 * head * head_dim]
274
seq_len, bs, _ = qkv.shape
275
- if self.tp_size > 1:
276
- qkv = all_gather_interleave(qkv, self.qkv.hidden_size,
277
- self.tp_size)
275
276
# [s, b, 3 * head * head_dim] -> 3 * [s, b, head * head_dim]
277
q, k, v = qkv.chunk(3, dim=2)
278
282
- # 3 * [s, b, head * head_dim]
283
- if self.tp_size > 1:
284
- splitter = partial(
285
- dist_utils.split_tensor_along_last_dim,
286
- num_partitions=self.tp_size,
287
- )
288
- q = splitter(q)[self.tp_rank]
289
- k = splitter(k)[self.tp_rank]
290
- v = splitter(v)[self.tp_rank]
291
-
279
# 3 * [s, b, head * head_dim] -> 3 * [s, b, head, head_dim]
280
new_shape = (
281
seq_len,