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

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 5, 11:00 UTC

[Attention][V0 Deprecation] Deprecate accept output buffer (#39125)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Apr 7, 2026 at 17:14 UTC 70406eb1dc191fa693258452c561855dbdf211a5

22 files changed +94 -227

tests/compile/test_config.py

+6 -4

@@ -216,12 +216,14 @@ def test_splitting_ops_dynamic():

216 compilation_config=CompilationConfig(

217 mode=CompilationMode.VLLM_COMPILE,

218 use_inductor_graph_partition=True,

219 - splitting_ops=["vllm::unified_attention"],

219 + splitting_ops=["vllm::unified_attention_with_output"],

220 )

221 )

222 # with inductor partition we use splitting_ops directly for

223 # partition rules

224 - assert config.compilation_config.splitting_ops == ["vllm::unified_attention"]

224 + assert config.compilation_config.splitting_ops == [

225 + "vllm::unified_attention_with_output"

226 + ]

227

228 # When attn_fusion pass enabled.

229 config = VllmConfig(

@@ -281,7 +283,7 @@ def test_moe_splitting_ops_deepep_ht_inductor_partition():

283 mode=CompilationMode.VLLM_COMPILE,

284 use_inductor_graph_partition=True,

285 splitting_ops=[

284 - "vllm::unified_attention",

286 + "vllm::unified_attention_with_output",

287 "vllm::moe_forward",

288 "vllm::moe_forward_shared",

289 ],

@@ -289,7 +291,7 @@ def test_moe_splitting_ops_deepep_ht_inductor_partition():

291 )

292 splitting_ops = config.compilation_config.splitting_ops

293 assert splitting_ops == [

292 - "vllm::unified_attention",

294 + "vllm::unified_attention_with_output",

295 "vllm::moe_forward",

296 "vllm::moe_forward_shared",

297 ]

vllm/config/compilation.py

+1 -3

@@ -282,7 +282,7 @@ class PassConfig:

282 """

283 enabled_fusions = [

284 f.name[len("fuse_") :]

285 - for f in fields(self)

285 + for f in fields(self) # type: ignore[arg-type]

286 if getattr(self, f.name) and f.name.startswith("fuse_")

287 ]

288

@@ -711,9 +711,7 @@ class CompilationConfig:

711 # Attention ops; used for piecewise cudagraphs

712 # Use PyTorch operator format: "namespace::name"

713 _attention_ops: ClassVar[list[str]] = [

714 - "vllm::unified_attention",

714 "vllm::unified_attention_with_output",

716 - "vllm::unified_mla_attention",

715 "vllm::unified_mla_attention_with_output",

716 "vllm::mamba_mixer2",

717 "vllm::mamba_mixer",

vllm/model_executor/layers/attention/attention.py

+53 -96

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

354 # and let torch.compile handle them.

355 self.use_direct_call = not current_platform.opaque_attention_op()

356

357 - self.use_output = self.attn_backend.accept_output_buffer

357 compilation_config = vllm_config.compilation_config

358 if prefix in compilation_config.static_forward_context:

359 raise ValueError(f"Duplicate layer name: {prefix}")

@@ -429,75 +428,62 @@ class Attention(nn.Module, AttentionLayerBase):

428 if self.impl.supports_quant_query_input:

429 query, _ = self.query_quant(query, self._q_scale)

430

432 - if self.use_output:

433 - if output_shape is None:

434 - # Handle both 2D [num_tokens, hidden] and

435 - # 3D [num_tokens, heads, head_dim] query

436 - num_tokens = query.shape[0]

437 - output_shape = torch.Size(

438 - (num_tokens, self.num_heads * self.head_size_v)

439 - )

440 - output = torch.empty(output_shape, dtype=output_dtype, device=query.device)

441 - hidden_size = output_shape[-1]

442 - # Reshape the query, key, and value tensors.

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

444 - # CPU overheads from the non-CUDA-graph regions.

445 - query = query.view(-1, self.num_heads, self.head_size)

446 - output = output.view(-1, self.num_heads, self.head_size_v)

447 - if key is not None:

448 - key = key.view(-1, self.num_kv_heads, self.head_size)

449 - if value is not None:

450 - value = value.view(-1, self.num_kv_heads, self.head_size_v)

451 - kv_cache_dummy_dep = None

452 - if self.use_direct_call:

453 - # Skip this if sharing KV cache with an earlier attention layer.

454 - if (

455 - not self.attn_backend.forward_includes_kv_cache_update

456 - and self.kv_sharing_target_layer_name is None

457 - and key is not None

458 - and value is not None

459 - ):

460 - kv_cache_dummy_dep = unified_kv_cache_update(

461 - key, value, self.layer_name

462 - )

463 - unified_attention_with_output(

464 - query,

465 - key,

466 - value,

467 - output,

468 - self.layer_name,

469 - kv_cache_dummy_dep=kv_cache_dummy_dep,

470 - )

471 - else:

472 - # Skip this if sharing KV cache with an earlier attention layer.

473 - if (

474 - not self.attn_backend.forward_includes_kv_cache_update

475 - and self.kv_sharing_target_layer_name is None

476 - and key is not None

477 - and value is not None

478 - ):

479 - kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(

480 - key, value, self.layer_name

481 - )

482 - torch.ops.vllm.unified_attention_with_output(

483 - query,

484 - key,

485 - value,

486 - output,

487 - self.layer_name,

488 - kv_cache_dummy_dep=kv_cache_dummy_dep,

431 + if output_shape is None:

432 + # Handle both 2D [num_tokens, hidden] and

433 + # 3D [num_tokens, heads, head_dim] query

434 + num_tokens = query.shape[0]

435 + output_shape = torch.Size((num_tokens, self.num_heads * self.head_size_v))

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

437 + hidden_size = output_shape[-1]

438 + # Reshape the query, key, and value tensors.

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

440 + # CPU overheads from the non-CUDA-graph regions.

441 + query = query.view(-1, self.num_heads, self.head_size)

442 + output = output.view(-1, self.num_heads, self.head_size_v)

443 + if key is not None:

444 + key = key.view(-1, self.num_kv_heads, self.head_size)

445 + if value is not None:

446 + value = value.view(-1, self.num_kv_heads, self.head_size_v)

447 + kv_cache_dummy_dep = None

448 + if self.use_direct_call:

449 + # Skip this if sharing KV cache with an earlier attention layer.

450 + if (

451 + not self.attn_backend.forward_includes_kv_cache_update

452 + and self.kv_sharing_target_layer_name is None

453 + and key is not None

454 + and value is not None

455 + ):

456 + kv_cache_dummy_dep = unified_kv_cache_update(

457 + key, value, self.layer_name

458 )

490 - return output.view(-1, hidden_size)

491 - else:

492 - assert self.attn_backend.forward_includes_kv_cache_update, (

493 - "Split KV cache update not supported when output tensor not provided."

459 + unified_attention_with_output(

460 + query,

461 + key,

462 + value,

463 + output,

464 + self.layer_name,

465 + kv_cache_dummy_dep=kv_cache_dummy_dep,

466 )

495 - if self.use_direct_call:

496 - return unified_attention(query, key, value, self.layer_name)

497 - else:

498 - return torch.ops.vllm.unified_attention(

499 - query, key, value, self.layer_name

467 + else:

468 + # Skip this if sharing KV cache with an earlier attention layer.

469 + if (

470 + not self.attn_backend.forward_includes_kv_cache_update

471 + and self.kv_sharing_target_layer_name is None

472 + and key is not None

473 + and value is not None

474 + ):

475 + kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(

476 + key, value, self.layer_name

477 )

478 + torch.ops.vllm.unified_attention_with_output(

479 + query,

480 + key,

481 + value,

482 + output,

483 + self.layer_name,

484 + kv_cache_dummy_dep=kv_cache_dummy_dep,

485 + )

486 + return output.view(-1, hidden_size)

487

488 def calc_kv_scales(self, query, key, value):

489 self._q_scale.copy_(torch.abs(query).max() / self.q_range)

@@ -633,35 +619,6 @@ def get_attention_context(

619 return attn_metadata, attn_layer, kv_cache, layer_slot_mapping

620

621

636 -@maybe_transfer_kv_layer

637 -def unified_attention(

638 - query: torch.Tensor,

639 - key: torch.Tensor,

640 - value: torch.Tensor,

641 - layer_name: str,

642 -) -> torch.Tensor:

643 - attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)

644 - output = self.impl.forward(self, query, key, value, kv_cache, attn_metadata)

645 -

646 - return output

647 -

648 -

649 -def unified_attention_fake(

650 - query: torch.Tensor,

651 - key: torch.Tensor,

652 - value: torch.Tensor,

653 - layer_name: str,

654 -) -> torch.Tensor:

655 - return torch.empty_like(query).contiguous()

656 -

657 -

658 -direct_register_custom_op(

659 - op_name="unified_attention",

660 - op_func=unified_attention,

661 - fake_impl=unified_attention_fake,

662 -)

663 -

664 -

622 def unified_kv_cache_update(

623 key: torch.Tensor,

624 value: torch.Tensor,

vllm/model_executor/layers/attention/cross_attention.py

+1 -1

@@ -133,7 +133,7 @@ def create_cross_attention_backend(

133 value: torch.Tensor,

134 kv_cache: torch.Tensor,

135 attn_metadata: AttentionMetadata,

136 - output: torch.Tensor | None = None,

136 + output: torch.Tensor,

137 output_scale: torch.Tensor | None = None,

138 output_block_scale: torch.Tensor | None = None,

139 ) -> torch.Tensor:

vllm/model_executor/layers/attention/mla_attention.py

+21 -76

@@ -494,21 +494,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):

494 self.kv_cache_dtype,

495 self._k_scale,

496 )

497 - if self.attn_backend.accept_output_buffer:

498 - output = torch.empty(output_shape, dtype=q.dtype, device=q.device)

499 - self.forward_impl(

500 - q,

501 - kv_c_normed,

502 - k_pe,

503 - self_kv_cache,

504 - attn_metadata,

505 - output=output,

506 - )

507 - return output

508 - else:

509 - return self.forward_impl(

510 - q, kv_c_normed, k_pe, self_kv_cache, attn_metadata

511 - )

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

498 + self.forward_impl(

499 + q,

500 + kv_c_normed,

501 + k_pe,

502 + self_kv_cache,

503 + attn_metadata,

504 + output=output,

505 + )

506 + return output

507 else:

508 kv_cache_dummy_dep = torch.ops.vllm.unified_mla_kv_cache_update(

509 kv_c_normed,

@@ -517,25 +512,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):

512 self.kv_cache_dtype,

513 self._k_scale,

514 )

520 - if self.attn_backend.accept_output_buffer:

521 - output = torch.empty(output_shape, dtype=q.dtype, device=q.device)

522 - torch.ops.vllm.unified_mla_attention_with_output(

523 - q,

524 - kv_c_normed,

525 - k_pe,

526 - output,

527 - self.layer_name,

528 - kv_cache_dummy_dep=kv_cache_dummy_dep,

529 - )

530 - return output

531 - else:

532 - return torch.ops.vllm.unified_mla_attention(

533 - q,

534 - kv_c_normed,

535 - k_pe,

536 - self.layer_name,

537 - kv_cache_dummy_dep=kv_cache_dummy_dep,

538 - )

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

516 + torch.ops.vllm.unified_mla_attention_with_output(

517 + q,

518 + kv_c_normed,

519 + k_pe,

520 + output,

521 + self.layer_name,

522 + kv_cache_dummy_dep=kv_cache_dummy_dep,

523 + )

524 + return output

525

526 def forward_impl(

527 self,

@@ -544,12 +530,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):

530 k_pe: torch.Tensor, # value in unified attn

531 kv_cache: torch.Tensor,

532 attn_metadata: "MLACommonMetadata",

547 - output: torch.Tensor | None = None,

533 + output: torch.Tensor,

534 output_scale: torch.Tensor | None = None,

535 output_block_scale: torch.Tensor | None = None,

536 ) -> torch.Tensor:

551 - assert output is not None, "Output tensor must be provided."

552 -

537 use_quant = output_scale is not None or output_block_scale is not None

538 if use_quant:

539 # The fusion pass has allocated output with quantized dtype

@@ -913,43 +897,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

897 out.copy_(out_new) # Copy result

898

899

916 -@maybe_transfer_kv_layer

917 -def unified_mla_attention(

918 - q: torch.Tensor,

919 - kv_c_normed: torch.Tensor,

920 - k_pe: torch.Tensor,

921 - layer_name: str,

922 - kv_cache_dummy_dep: torch.Tensor | None = None,

923 -) -> torch.Tensor:

924 - # kv_cache_dummy_dep is not used but accepting it creates a data dependency

925 - # that ensures torch.compile preserves ordering between KV cache update and

926 - # attention forward.

927 - del kv_cache_dummy_dep

928 - attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)

929 - output = layer.forward_impl(q, kv_c_normed, k_pe, kv_cache, attn_metadata)

930 -

931 - return output

932 -

933 -

934 -def unified_mla_attention_fake(

935 - q: torch.Tensor,

936 - kv_c_normed: torch.Tensor,

937 - k_pe: torch.Tensor,

938 - layer_name: str,

939 - kv_cache_dummy_dep: torch.Tensor | None = None,

940 -) -> torch.Tensor:

941 - return torch.empty_like(q).contiguous()

942 -

943 -

944 -direct_register_custom_op(

945 - op_name="unified_mla_attention",

946 - op_func=unified_mla_attention,

947 - mutates_args=[],

948 - fake_impl=unified_mla_attention_fake,

949 - dispatch_key=current_platform.dispatch_key,

950 -)

951 -

952 -

900 def unified_mla_kv_cache_update(

901 kv_c_normed: torch.Tensor,

902 k_pe: torch.Tensor,

@@ -1151,8 +1098,6 @@ CUDNN_WORKSPACE_SIZE = 12800

1098

1099

1100 class MLACommonBackend(AttentionBackend):

1154 - accept_output_buffer: bool = True

1155 -

1101 @staticmethod

1102 def get_name() -> str:

1103 return "TRITON_MLA"

vllm/model_executor/models/extract_hidden_states.py

-1

@@ -94,7 +94,6 @@ def basic_cache(

94 class CacheOnlyAttentionBackend(AttentionBackend):

95 """Attention backend that only caches KV without computing attention."""

97 - accept_output_buffer: bool = False

97 supported_dtypes: ClassVar[list[torch.dtype]] = [

98 torch.float16,

99 torch.bfloat16,

vllm/model_executor/models/whisper_causal.py

+1 -1

@@ -184,7 +184,7 @@ def create_whisper_attention_backend_with_block_pooling(

184 value: torch.Tensor,

185 kv_cache: torch.Tensor,

186 attn_metadata: AttentionMetadata,

187 - output: torch.Tensor | None = None,

187 + output: torch.Tensor,

188 output_scale: torch.Tensor | None = None,

189 output_block_scale: torch.Tensor | None = None,

190 ) -> torch.Tensor:

vllm/v1/attention/backend.py

+1 -5

@@ -53,10 +53,6 @@ class MultipleOf:

53 class AttentionBackend(ABC):

54 """Abstract class for attention backends."""

56 - # For some attention backends, we allocate an output tensor before

57 - # calling the custom op. When piecewise cudagraph is enabled, this

58 - # makes sure the output tensor is allocated inside the cudagraph.

59 - accept_output_buffer: bool = False

56 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

57 supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [

58 "auto",

@@ -779,7 +775,7 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):

775 value: torch.Tensor,

776 kv_cache: torch.Tensor,

777 attn_metadata: T,

782 - output: torch.Tensor | None = None,

778 + output: torch.Tensor,

779 output_scale: torch.Tensor | None = None,

780 output_block_scale: torch.Tensor | None = None,

781 ) -> torch.Tensor:

vllm/v1/attention/backends/cpu_attn.py

+1 -3

@@ -30,7 +30,6 @@ _CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86, CpuArchEnum.ARM, CpuArchEnum.S3

32 class CPUAttentionBackend(AttentionBackend):

33 - accept_output_buffer: bool = True

33 supported_dtypes: ClassVar[list[torch.dtype]] = [

34 torch.float16,

35 torch.bfloat16,

@@ -267,7 +266,7 @@ class CPUAttentionBackendImpl(AttentionImpl):

266 value: torch.Tensor,

267 kv_cache: torch.Tensor,

268 attn_metadata: CPUAttentionMetadata | None,

270 - output: torch.Tensor | None = None,

269 + output: torch.Tensor,

270 output_scale: torch.Tensor | None = None,

271 output_block_scale: torch.Tensor | None = None,

272 ) -> torch.Tensor:

@@ -283,7 +282,6 @@ class CPUAttentionBackendImpl(AttentionImpl):

282 Returns:

283 shape = [num_tokens, num_heads * head_size]

284 """

286 - assert output is not None, "Output tensor must be provided."

285 if output_scale is not None or output_block_scale is not None:

286 raise NotImplementedError(

287 "fused output quantization is not yet supported"

vllm/v1/attention/backends/flash_attn.py

+1 -3

@@ -62,7 +62,6 @@ logger = init_logger(__name__)

64 class FlashAttentionBackend(AttentionBackend):

65 - accept_output_buffer: bool = True

65 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

66 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

67 "auto",

@@ -664,7 +663,7 @@ class FlashAttentionImpl(AttentionImpl):

663 value: torch.Tensor,

664 kv_cache: torch.Tensor,

665 attn_metadata: FlashAttentionMetadata,

667 - output: torch.Tensor | None = None,

666 + output: torch.Tensor,

667 output_scale: torch.Tensor | None = None,

668 output_block_scale: torch.Tensor | None = None,

669 ) -> torch.Tensor:

@@ -683,7 +682,6 @@ class FlashAttentionImpl(AttentionImpl):

682 {q,k,v}_descale to be (num_sequences, num_kv_heads).

683 We use torch's .expand() to avoid duplicating values

684 """

686 - assert output is not None, "Output tensor must be provided."

685 assert self.vllm_flash_attn_version is not None, (

686 "FlashAttention version not detected."

687 )

vllm/v1/attention/backends/flash_attn_diffkv.py

+1 -2

@@ -128,7 +128,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):

128 value: torch.Tensor,

129 kv_cache: torch.Tensor,

130 attn_metadata: FlashAttentionMetadata,

131 - output: torch.Tensor | None = None,

131 + output: torch.Tensor,

132 output_scale: torch.Tensor | None = None,

133 output_block_scale: torch.Tensor | None = None,

134 ) -> torch.Tensor:

@@ -147,7 +147,6 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):

147 {q,k,v}_descale to be (num_sequences, num_kv_heads).

148 We use torch's .expand() to avoid duplicating values

149 """

150 - assert output is not None, "Output tensor must be provided."

150 assert self.vllm_flash_attn_version is not None, (

151 "FlashAttention version not detected."

152 )

vllm/v1/attention/backends/flashinfer.py

+1 -4

@@ -315,7 +315,6 @@ class BatchDCPPrefillWrapper:

315

316

317 class FlashInferBackend(AttentionBackend):

318 - accept_output_buffer: bool = True

318 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

319 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

320 "auto",

@@ -1286,7 +1285,7 @@ class FlashInferImpl(AttentionImpl):

1285 value: torch.Tensor,

1286 kv_cache: torch.Tensor,

1287 attn_metadata: FlashInferMetadata,

1289 - output: torch.Tensor | None = None,

1288 + output: torch.Tensor,

1289 output_scale: torch.Tensor | None = None,

1290 output_block_scale: torch.Tensor | None = None,

1291 ) -> torch.Tensor:

@@ -1303,8 +1302,6 @@ class FlashInferImpl(AttentionImpl):

1302 Returns:

1303 shape = [num_tokens, num_heads * head_size]

1304 """

1306 - assert output is not None, "Output tensor must be provided."

1307 -

1305 if attn_metadata is None:

1306 # Profiling run.

1307 return output.fill_(0)

vllm/v1/attention/backends/flex_attention.py

+1 -3

@@ -73,7 +73,6 @@ def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int):

75 class FlexAttentionBackend(AttentionBackend):

76 - accept_output_buffer: bool = True

76 supported_dtypes: ClassVar[list[torch.dtype]] = [

77 torch.float16,

78 torch.bfloat16,

@@ -992,7 +991,7 @@ class FlexAttentionImpl(AttentionImpl):

991 value: torch.Tensor,

992 kv_cache: torch.Tensor,

993 attn_metadata: FlexAttentionMetadata,

995 - output: torch.Tensor | None = None,

994 + output: torch.Tensor,

995 output_scale: torch.Tensor | None = None,

996 output_block_scale: torch.Tensor | None = None,

997 ) -> torch.Tensor:

@@ -1008,7 +1007,6 @@ class FlexAttentionImpl(AttentionImpl):

1007 Returns:

1008 shape = [num_tokens, num_heads * head_size]

1009 """

1011 - assert output is not None, "Output tensor must be provided."

1010 if output_scale is not None or output_block_scale is not None:

1011 raise NotImplementedError(

1012 "fused output quantization is not yet supported for FlexAttentionImpl"

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

-1

@@ -59,7 +59,6 @@ class FlashInferMLASparseBackend(AttentionBackend):

59 for models like DeepSeek-V3.2 that use index-based sparse attention.

60 """

62 - accept_output_buffer: bool = True

62 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

63 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

64 "auto",

vllm/v1/attention/backends/mla/flashmla_sparse.py

-1

@@ -78,7 +78,6 @@ structured as:

80 class FlashMLASparseBackend(AttentionBackend):

81 - accept_output_buffer: bool = True

81 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

82 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

83 "auto",

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

-1

@@ -78,7 +78,6 @@ def fetch_id_to_ragged_triton(

80 class ROCMAiterMLASparseBackend(AttentionBackend):

81 - accept_output_buffer: bool = True

81 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

82 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

83 "auto",

vllm/v1/attention/backends/mla/xpu_mla_sparse.py

-1

@@ -35,7 +35,6 @@ logger = init_logger(__name__)

37 class XPUMLASparseBackend(AttentionBackend):

38 - accept_output_buffer: bool = True

38 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

39 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

40 "auto",

vllm/v1/attention/backends/rocm_aiter_fa.py

+1 -4

@@ -744,7 +744,6 @@ class AiterFlashAttentionMetadataBuilder(

744

745

746 class AiterFlashAttentionBackend(AttentionBackend):

747 - accept_output_buffer: bool = True

747 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

748 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

749 "auto",

@@ -1037,7 +1036,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1036 value: torch.Tensor,

1037 kv_cache: torch.Tensor,

1038 attn_metadata: AiterFlashAttentionMetadata,

1040 - output: torch.Tensor | None = None,

1039 + output: torch.Tensor,

1040 output_scale: torch.Tensor | None = None,

1041 output_block_scale: torch.Tensor | None = None,

1042 ) -> torch.Tensor:

@@ -1056,8 +1055,6 @@ class AiterFlashAttentionImpl(AttentionImpl):

1055 {q,k,v}_descale to be (num_sequences, num_kv_heads).

1056 We use torch's .expand() to avoid duplicating values

1057 """

1059 - assert output is not None, "Output tensor must be provided."

1060 -

1058 if output_scale is not None or output_block_scale is not None:

1059 raise NotImplementedError(

1060 "fused output quantization is not yet supported "

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+1 -5

@@ -24,8 +24,6 @@ logger = init_logger(__name__)

26 class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):

27 - accept_output_buffer: bool = True

28 -

27 @staticmethod

28 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

29 return [MultipleOf(16)]

@@ -143,7 +141,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

141 value: torch.Tensor,

142 kv_cache: torch.Tensor,

143 attn_metadata: FlashAttentionMetadata,

146 - output: torch.Tensor | None = None,

144 + output: torch.Tensor,

145 output_scale: torch.Tensor | None = None,

146 output_block_scale: torch.Tensor | None = None,

147 ) -> torch.Tensor:

@@ -159,8 +157,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

157 Returns:

158 shape = [num_tokens, num_heads * head_size]

159 """

162 - assert output is not None, "Output tensor must be provided."

163 -

160 if output_block_scale is not None:

161 raise NotImplementedError(

162 "fused block_scale output quantization is not yet supported"

vllm/v1/attention/backends/rocm_attn.py

+1 -4

@@ -159,7 +159,6 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat

159

160

161 class RocmAttentionBackend(AttentionBackend):

162 - accept_output_buffer: bool = True

162 supported_dtypes: ClassVar[list[torch.dtype]] = [

163 torch.float16,

164 torch.bfloat16,

@@ -352,7 +351,7 @@ class RocmAttentionImpl(AttentionImpl):

351 value: torch.Tensor,

352 kv_cache: torch.Tensor,

353 attn_metadata: FlashAttentionMetadata,

355 - output: torch.Tensor | None = None,

354 + output: torch.Tensor,

355 output_scale: torch.Tensor | None = None,

356 output_block_scale: torch.Tensor | None = None,

357 ) -> torch.Tensor:

@@ -368,8 +367,6 @@ class RocmAttentionImpl(AttentionImpl):

367 Returns:

368 shape = [num_tokens, num_heads * head_size]

369 """

371 - assert output is not None, "Output tensor must be provided."

372 -

370 if output_block_scale is not None:

371 raise NotImplementedError(

372 "fused block_scale output quantization is not yet supported"

vllm/v1/attention/backends/tree_attn.py

+1 -4

@@ -30,7 +30,6 @@ logger = init_logger(__name__)

32 class TreeAttentionBackend(AttentionBackend):

33 - accept_output_buffer: bool = True

33 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

34 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

35 "auto",

@@ -368,7 +367,7 @@ class TreeAttentionImpl(AttentionImpl):

367 value: torch.Tensor,

368 kv_cache: torch.Tensor,

369 attn_metadata: TreeAttentionMetadata,

371 - output: torch.Tensor | None = None,

370 + output: torch.Tensor,

371 output_scale: torch.Tensor | None = None,

372 output_block_scale: torch.Tensor | None = None,

373 ) -> torch.Tensor:

@@ -384,8 +383,6 @@ class TreeAttentionImpl(AttentionImpl):

383 Returns:

384 shape = [num_tokens, num_heads * head_size]

385 """

387 - assert output is not None, "Output tensor must be provided."

388 -

386 if output_scale is not None or output_block_scale is not None:

387 raise NotImplementedError(

388 "fused output quantization is not yet supported for TreeAttentionImpl"

vllm/v1/attention/backends/triton_attn.py

+1 -4

@@ -262,7 +262,6 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet

262

263

264 class TritonAttentionBackend(AttentionBackend):

265 - accept_output_buffer: bool = True

265 supported_dtypes: ClassVar[list[torch.dtype]] = [

266 torch.float16,

267 torch.bfloat16,

@@ -504,7 +503,7 @@ class TritonAttentionImpl(AttentionImpl):

503 value: torch.Tensor,

504 kv_cache: torch.Tensor,

505 attn_metadata: TritonAttentionMetadata,

507 - output: torch.Tensor | None = None,

506 + output: torch.Tensor,

507 output_scale: torch.Tensor | None = None,

508 output_block_scale: torch.Tensor | None = None,

509 ) -> torch.Tensor:

@@ -520,8 +519,6 @@ class TritonAttentionImpl(AttentionImpl):

519 Returns:

520 shape = [num_tokens, num_heads * head_size]

521 """

523 - assert output is not None, "Output tensor must be provided."

524 -

522 if output_block_scale is not None:

523 raise NotImplementedError(

524 "fused block_scale output quantization is not yet supported"