Perf opt

backends/gaudi/server/text_generation_server/layers/attention/__init__.py

@@ -11,6 +11,7 @@
     attention,
     paged_attention,
     paged_attention_mla,
+    set_block_mapping,
 )
 
 
@@ -22,6 +23,7 @@
     "get_kv_scales",
     "paged_attention",
     "paged_attention_mla",
+    "set_block_mapping",
     "SUPPORTS_WINDOWING",
     "KVCache",
     "KVCompressCache",

backends/gaudi/server/text_generation_server/layers/attention/hpu.py

@@ -8,6 +8,7 @@
 from vllm_hpu_extension.utils import ModuleFusedSDPA
 import os
 from text_generation_server.models.globals import BLOCK_SIZE
+import math
 
 SUPPORTS_WINDOWING = False
 
@@ -106,6 +107,21 @@ def attention(
     return attn_output
 
 
+def set_block_mapping(hpu_attention_meta: HPUPagedAttentionMetadata, batch_size):
+    block_mapping = torch.nn.functional.one_hot(
+        hpu_attention_meta.block_groups, num_classes=batch_size
+    )
+    dtype = hpu_attention_meta.block_usage.dtype
+    device = hpu_attention_meta.block_usage.device
+    mask = torch.arange(0, BLOCK_SIZE, device=device, dtype=torch.int32).unsqueeze(0)
+    mask = mask >= hpu_attention_meta.block_usage.unsqueeze(-1)
+    attn_bias = torch.zeros_like(mask, dtype=dtype).masked_fill_(mask, -math.inf)
+    hpu_attention_meta = hpu_attention_meta._replace(
+        attn_bias=attn_bias, block_mapping=block_mapping.to(dtype)
+    )
+    return hpu_attention_meta
+
+
 def paged_attention(
     query: torch.Tensor,
     kv_cache: KVCache,
@@ -176,4 +192,10 @@ def paged_attention_mla(
     return output.view(batch_size, head_num, -1)
 
 
-__all__ = ["SUPPORTS_WINDOWING", "attention", "paged_attention", "paged_attention_mla"]
+__all__ = [
+    "SUPPORTS_WINDOWING",
+    "attention",
+    "paged_attention",
+    "paged_attention_mla",
+    "set_block_mapping",
+]

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_cohere_modeling.py

@@ -28,6 +28,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -415,6 +416,10 @@ def forward(
         seqlen: torch.Tensor,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_tokens(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_dbrx_modeling.py

@@ -26,6 +26,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -678,6 +679,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_tokens(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_deepseek_v2_modeling.py

@@ -33,6 +33,7 @@
     Seqlen,
     attention,
     paged_attention,
+    set_block_mapping,
     HPUPagedAttentionMetadata,
 )
 from text_generation_server.layers.attention.kv_cache import KVCache, get_kv_scales
@@ -569,6 +570,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_tokens(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_deepseek_v3_modeling.py

@@ -34,6 +34,7 @@
     Seqlen,
     attention,
     paged_attention_mla,
+    set_block_mapping,
     HPUPagedAttentionMetadata,
 )
 from text_generation_server.layers.attention.kv_cache import KVCache, get_kv_scales
@@ -645,6 +646,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_tokens(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_gemma2_modeling.py

@@ -28,6 +28,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -466,6 +467,10 @@ def forward(
         adapter_data: Optional[torch.Tensor],
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
         hidden_states = inputs_embeds
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_gemma_modeling.py

@@ -28,6 +28,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -388,6 +389,10 @@ def forward(
         adapter_data: Optional[torch.Tensor],
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
         hidden_states = inputs_embeds
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_gpt2_modeling.py

@@ -27,6 +27,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -383,6 +384,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
         hidden_states = inputs_embeds
 
         residual = None

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_gptj_modeling.py

@@ -28,6 +28,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -324,6 +325,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.wte(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llama4_modeling.py

@@ -43,6 +43,7 @@
 from text_generation_server.layers.attention import (
     KVCache,
     paged_attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -548,6 +549,10 @@ def forward(
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
         attention_mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
 
         hidden_states = inputs_embeds
         bs = seqlen.input_lengths.shape[0]

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py

@@ -35,6 +35,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -549,6 +550,11 @@ def forward(
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
         cross_attention_states=None,
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
+
         hidden_states = inputs_embeds
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py

@@ -30,6 +30,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -396,6 +397,10 @@ def forward(
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
         adapter_data: Optional[torch.Tensor] = None,
     ):
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
         hidden_states = inputs_embeds
         # Get rotary cos and sin for this forward
         # Avoid to index in each layer

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py

@@ -37,6 +37,7 @@
     Seqlen,
     attention,
     paged_attention,
+    set_block_mapping,
     HPUPagedAttentionMetadata,
 )
 from text_generation_server.layers.attention.kv_cache import get_kv_scales
@@ -446,6 +447,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_tokens(input_ids)
 
         # Get rotary cos and sin for this forward
@@ -505,7 +510,6 @@ def forward(
         lm_head_indices: Optional[torch.Tensor] = None,
         adapter_data: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-
         hidden_states = self.model(
             input_ids,
             position_ids,

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_neox_modeling.py

@@ -29,6 +29,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -354,6 +355,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_in(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_phi_modeling.py

@@ -9,6 +9,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -347,6 +348,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, input_ids.shape[0]
+            )
         hidden_states = self.embed_tokens(input_ids)
 
         # Get rotary cos and sin for this forward

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_qwen2_modeling.py

@@ -8,6 +8,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -288,6 +289,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
         hidden_states = inputs_embeds
 
         cos, sin = self.layers[0].self_attn.rotary_emb.get_cos_sin(
@@ -359,7 +364,6 @@ def forward(
         lm_head_indices: Optional[torch.Tensor] = None,
         adapter_data: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-
         inputs_embeds = self.embed_tokens(input_ids)
 
         hidden_states = self.model(

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_qwen3_modeling.py

@@ -18,6 +18,7 @@
 from text_generation_server.layers.attention import (
     paged_attention,
     attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -266,7 +267,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
-
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
         hidden_states = inputs_embeds
 
         # create position embeddings to be shared across the decoder layers
@@ -334,7 +338,6 @@ def forward(
         lm_head_indices: Optional[torch.Tensor] = None,
         adapter_data: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-
         inputs_embeds = self.embed_tokens(input_ids)
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         hidden_states = self.model(