Add output_scale to new attn backends

Signed-off-by: Luka Govedič <lgovedic@redhat.com>

Author: ProExpertProg

vllm/attention/backends/dual_chunk_flash_attn.py

@@ -370,6 +370,8 @@ def forward(  # type: ignore
         value: torch.Tensor,
         kv_cache: torch.Tensor,
         attn_metadata: DualChunkFlashAttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+        output_scale: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Forward pass with DualChunkFlashAttention.
         Args:
@@ -383,6 +385,13 @@ def forward(  # type: ignore
         Returns:
             shape = [num_tokens, num_heads * head_size]
         """
+        assert output is None, "Output tensor not supported for DualChunk"
+
+        if output_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported"
+                " for FlashAttentionImpl")
+
         (
             query,
             query_succ,

vllm/v1/attention/backends/flashinfer.py

@@ -547,6 +547,7 @@ def forward(
         kv_cache: torch.Tensor,
         attn_metadata: FlashInferMetadata,
         output: Optional[torch.Tensor] = None,
+        output_scale: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Forward pass with FlashInfer.
 
@@ -561,6 +562,11 @@ def forward(
         """
         assert output is not None, "Output tensor must be provided."
 
+        if output_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported"
+                " for FlashInferImpl")
+
         if attn_metadata is None:
             # Profiling run.
             return output

vllm/v1/attention/backends/flex_attention.py

@@ -414,6 +414,7 @@ def forward(
         kv_cache: torch.Tensor,
         attn_metadata: FlexAttentionMetadata,
         output: Optional[torch.Tensor] = None,
+        output_scale: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Forward pass with FLexAttention.
 
@@ -427,6 +428,12 @@ def forward(
             shape = [num_tokens, num_heads * head_size]
         """
         assert output is not None, "Output tensor must be provided."
+
+        if output_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported"
+                " for FlexAttentionImpl")
+
         enable_gqa = self.num_kv_heads != self.num_heads
 
         if attn_metadata is None: