WIP enabling llama4 models

Signed-off-by: Artur Fierka <[email protected]>

Author: afierka-intel

vllm_gaudi/attention/backends/hpu_attn.py

@@ -18,7 +18,7 @@
 
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionLayer,
-                                              AttentionMetadata, AttentionType)
+                                              AttentionMetadata, AttentionType, AttentionMetadataBuilder)
 from vllm.attention.backends.mla.common import MLACommonImpl
 from vllm.attention.backends.utils import CommonAttentionState
 from vllm_gaudi.attention.ops.hpu_paged_attn import (HPUPagedAttention,
@@ -47,6 +47,10 @@ def get_metadata_cls() -> type["AttentionMetadata"]:
     def get_state_cls() -> type["CommonAttentionState"]:
         return CommonAttentionState
 
+    @staticmethod
+    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+        return HPUAttentionMetadataBuilder
+
     @staticmethod
     def get_kv_cache_shape(
         num_blocks: int,

vllm_gaudi/attention/ops/hpu_paged_attn.py

@@ -8,6 +8,7 @@
 from typing import Optional
 
 import torch
+from vllm.attention.backends.abstract import AttentionMetadataBuilder
 from vllm_gaudi.extension import cache_ops, ops
 
 # Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.
@@ -24,6 +25,23 @@ class HPUPagedAttentionMetadata:
     alibi_blocks: Optional[torch.Tensor]
 
 
+@dataclass
+class HPUPagedAttentionMetadataBuilder(AttentionMetadataBuilder[HPUPagedAttentionMetadata]):
+
+    def __init__(self, input_builder: "ModelRunnerInputBuilderBase") -> None:
+        """Create the builder, remember some configuration and parameters."""
+        self.input_builder = input_builder
+
+    def prepare(self) -> None:
+        """Prepare for one batch."""
+        pass
+
+    def build(self, seq_lens: list[int], query_lens: list[int],
+              cuda_graph_pad_size: int, batch_size: int) -> HPUPagedAttentionMetadata:
+        """Build attention metadata with on-device tensors."""
+        return HPUPagedAttentionMetadata
+
+
 class HPUPagedAttention:
 
     @staticmethod