APC - Remove prompt attn with context and use existing implementation

requirements-hpu.txt

@@ -8,4 +8,4 @@ pandas
 tabulate
 setuptools>=61
 setuptools-scm>=8
-vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@1d157d0
+vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@adobrzyn/move_attn_with_context

vllm/attention/backends/hpu_attn.py

@@ -252,16 +252,16 @@ def forward(
                     valid_seq_lengths=attn_metadata.seq_lens_tensor,
                     **self.common_attention_args())
             else:
-                # TODO: enable FusedSDPA
-                out = HPUPagedAttention.forward_prefix(
-                    query=query.view(query_shape),
-                    key=key.view(kv_shape),
-                    value=value.view(kv_shape),
-                    key_cache=key_cache,
-                    value_cache=value_cache,
-                    block_list=attn_metadata.block_list,
-                    attn_bias=attn_metadata.attn_bias,
-                    **self.common_attention_args())
+                out = ops.prompt_attention(impl=self.prefill_impl,
+                                           query=query.view(query_shape),
+                                           key=key.view(kv_shape),
+                                           value=value.view(kv_shape),
+                                           key_cache=key_cache,
+                                           value_cache=value_cache,
+                                           block_list=attn_metadata.block_list,
+                                           is_causal=True,
+                                           attn_bias=attn_metadata.attn_bias,
+                                           **self.common_attention_args())
             output = out.reshape(batch_size, seq_len, hidden_size)
         else:
             # Decoding run.

vllm/attention/ops/hpu_paged_attn.py

@@ -63,10 +63,6 @@ def write_to_paged_cache(key: torch.Tensor, value: torch.Tensor,
     def forward_decode(**kwargs) -> torch.Tensor:
         return ops.flat_pa(**kwargs)
 
-    @staticmethod
-    def forward_prefix(**kwargs) -> torch.Tensor:
-        return ops.prompt_attention_with_context(**kwargs)
-
     @staticmethod
     def swap_blocks(
         src_kv_cache: Tuple[torch.Tensor, torch.Tensor],