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 9, 12:30 UTC
[Perf] Optimize dcp allocate tensor (#33102)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 27, 2026 at 17:24 UTC
3a6d5cbefd97a3dee07ba1756d8b5a9052801403
1 file changed
+3
-7
vllm/v1/attention/ops/common.py
+3
-7
@@ -195,14 +195,10 @@ def _cp_lse_common(
195
if ctx is None:
196
ctx = CPTritonContext()
197
198
- lses = torch.empty(
199
- (cp_group.world_size,) + cp_attn_lse.shape,
200
- dtype=cp_attn_lse.dtype,
201
- device=cp_attn_lse.device,
202
- )
203
-
198
cp_attn_lse = cp_attn_lse.contiguous()
205
- lses = cp_group.all_gather(cp_attn_lse, dim=0).view_as(lses)
199
+ lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape(
200
+ (cp_group.world_size,) + cp_attn_lse.shape
201
+ )
202
out, lse = correct_attn_out(
203
cp_attn_out,
204
lses,