@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a86b4c58e8

Fork 0

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 attn output view kernel (#26680)

Signed-off-by: Boyuan Feng <boyuan@meta.com> Signed-off-by: Boyuan Feng <fby.1994@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Boyuan Feng committed Oct 14, 2025 at 15:53 UTC a86b4c58e8f72f4903d873d25510f53f7577366f

10 files changed +12 -12

vllm/attention/layer.py

+3 -3

@@ -346,7 +346,7 @@ class Attention(nn.Module, AttentionLayerBase):

346

347 if self.use_output:

348 output_shape = output_shape if output_shape is not None else query.shape

349 - output = torch.zeros(output_shape, dtype=output_dtype, device=query.device)

349 + output = torch.empty(output_shape, dtype=output_dtype, device=query.device)

350 hidden_size = output_shape[-1]

351 # Reshape the query, key, and value tensors.

352 # NOTE(woosuk): We do this outside the custom op to minimize the

@@ -705,7 +705,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

705 self.calc_kv_scales(q, kv_c_normed, k_pe)

706

707 if self.attn_backend.accept_output_buffer:

708 - output = torch.zeros(output_shape, dtype=q.dtype, device=q.device)

708 + output = torch.empty(output_shape, dtype=q.dtype, device=q.device)

709 self.impl.forward(

710 self,

711 q,

@@ -722,7 +722,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

722 )

723 else:

724 if self.attn_backend.accept_output_buffer:

725 - output = torch.zeros(output_shape, dtype=q.dtype, device=q.device)

725 + output = torch.empty(output_shape, dtype=q.dtype, device=q.device)

726 torch.ops.vllm.unified_mla_attention_with_output(

727 q,

728 kv_c_normed,

vllm/v1/attention/backends/flash_attn.py

+1 -1

@@ -530,7 +530,7 @@ class FlashAttentionImpl(AttentionImpl):

530

531 if attn_metadata is None:

532 # Profiling run.

533 - return output

533 + return output.fill_(0)

534

535 attn_type = self.attn_type

536

vllm/v1/attention/backends/flashinfer.py

+1 -1

@@ -857,7 +857,7 @@ class FlashInferImpl(AttentionImpl):

857

858 if attn_metadata is None:

859 # Profiling run.

860 - return output

860 + return output.fill_(0)

861

862 if self.bmm1_scale is None:

863 self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale

vllm/v1/attention/backends/flex_attention.py

+1 -1

@@ -767,7 +767,7 @@ class FlexAttentionImpl(AttentionImpl):

767

768 if attn_metadata is None:

769 # Profiling run.

770 - return output

770 + return output.fill_(0)

771 # query = self.view_as_4d(query).permute(0, 2, 1, 3)

772 # return torch.empty_like(query)

773

vllm/v1/attention/backends/rocm_aiter_fa.py

+1 -1

@@ -485,7 +485,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

485

486 if attn_metadata is None:

487 # Profiling run.

488 - return output

488 + return output.fill_(0)

489

490 # IMPORTANT!

491 # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+1 -1

@@ -130,7 +130,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

130

131 if attn_metadata is None:

132 # Profiling run.

133 - return output

133 + return output.fill_(0)

134

135 assert attn_metadata.use_cascade is False

136

vllm/v1/attention/backends/rocm_attn.py

+1 -1

@@ -299,7 +299,7 @@ class RocmAttentionImpl(AttentionImpl):

299

300 if attn_metadata is None:

301 # Profiling run.

302 - return output

302 + return output.fill_(0)

303

304 assert attn_metadata.use_cascade is False

305

vllm/v1/attention/backends/tree_attn.py

+1 -1

@@ -379,7 +379,7 @@ class TreeAttentionImpl(AttentionImpl):

379

380 if attn_metadata is None:

381 # Profiling run.

382 - return output

382 + return output.fill_(0)

383

384 # Cache the input KVs.

385 key_cache, value_cache = kv_cache.unbind(0)

vllm/v1/attention/backends/triton_attn.py

+1 -1

@@ -298,7 +298,7 @@ class TritonAttentionImpl(AttentionImpl):

298

299 if attn_metadata is None:

300 # Profiling run.

301 - return output

301 + return output.fill_(0)

302

303 assert attn_metadata.use_cascade is False

304

vllm/v1/attention/backends/xformers.py

+1 -1

@@ -354,7 +354,7 @@ class XFormersAttentionImpl(AttentionImpl):

354

355 if attn_metadata is None:

356 # Profiling run.

357 - return output

357 + return output.fill_(0)

358

359 # Cache the input KVs.

360 key_cache, value_cache = kv_cache.unbind(0)