[gaudi] Gemma3 sliding window support (#3280)

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

Author: sywangyi

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

@@ -2,6 +2,7 @@
 import torch
 from typing import Optional, List, Dict
 import collections
+import torch.nn.functional as F
 
 _TYPE_CACHE = {}
 
@@ -15,6 +16,12 @@ class HPUPagedAttentionMetadata:
     block_usage: Optional[torch.Tensor]
     block_groups: Optional[torch.Tensor]
     attn_bias: Optional[torch.Tensor]
+    slots_in_window_mask: Optional[torch.Tensor] = None
+    block_list_in_window: Optional[torch.Tensor] = None
+    block_mapping_in_window: Optional[torch.Tensor] = None
+    block_usage_in_window: Optional[torch.Tensor] = None
+    block_groups_in_window: Optional[torch.Tensor] = None
+    attn_bias_in_window: Optional[torch.Tensor] = None
 
 
 def subtuple(
@@ -67,6 +74,12 @@ def trim_attn_metadata(metadata: HPUPagedAttentionMetadata) -> object:
             "block_usage",
             "block_groups",
             "attn_bias",
+            "slots_in_window_mask",
+            "block_list_in_window",
+            "block_mapping_in_window",
+            "block_usage_in_window",
+            "block_groups_in_window",
+            "attn_bias_in_window",
         ],
     )
     return attention_metadata
@@ -75,6 +88,7 @@ def trim_attn_metadata(metadata: HPUPagedAttentionMetadata) -> object:
 @dataclass
 class Seqlen:
     input_lengths: torch.Tensor
+    attn_mask: Optional[torch.Tensor] = None
 
     def __init__(
         self,
@@ -86,6 +100,48 @@ def clamp(self, max):
         # Flash decoding doesn't need to clamp
         return self
 
+    def make_sliding_window_bias(
+        self,
+        seq_lens: List[int],
+        window_size: Optional[int],
+        dtype: torch.dtype,
+        padded_input_len: Optional[int],
+        padded_bs: Optional[int],
+    ) -> List[torch.Tensor]:
+        attn_biases = []
+        for seq_len in seq_lens:
+            if seq_len != 0:
+                tensor = torch.full(
+                    (1, seq_len, seq_len),
+                    dtype=dtype,
+                    fill_value=1,
+                )
+                shift = 0
+                mask = torch.tril(tensor, diagonal=shift).to(dtype)  # type: ignore
+                if window_size is not None:
+                    mask = torch.triu(mask, diagonal=shift - window_size + 1)
+                mask = F.pad(
+                    mask,
+                    (
+                        padded_input_len - seq_len,
+                        0,
+                        padded_input_len - seq_len,
+                        0,
+                        0,
+                        0,
+                    ),
+                    value=0,
+                )
+            else:
+                mask = torch.full(
+                    (1, padded_input_len, padded_input_len),
+                    dtype=dtype,
+                    fill_value=0,
+                )
+            attn_biases.append(mask)
+        attn_biases = torch.stack(attn_biases, dim=0)
+        return attn_biases.to(torch.bool)
+
 
 def _async_h2d_tensor_copy(source, device="hpu"):
     if source is None:
@@ -124,6 +180,7 @@ def trim_seqlen_metadata(metadata: Seqlen) -> object:
         "TrimmedSeqlen",
         [
             "input_lengths",
+            "attn_mask",
         ],
     )
     return attention_metadata

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

@@ -94,13 +94,13 @@ def attention(
         query,
         key,
         value,
-        attn_mask=None,
+        attn_mask=seqlen.attn_mask if window_size_left != -1 else None,
         dropout_p=0.0,
-        is_causal=causal,
+        is_causal=causal if window_size_left == -1 else False,
         scale=softmax_scale,
         softmax_mode="None",
         recompute_mode=None,
-        valid_sequence_lengths=seqlen.input_lengths,
+        valid_sequence_lengths=seqlen.input_lengths if window_size_left == -1 else None,
         padding_side="left",
     )
     attn_output = attn_output.transpose(1, 2).squeeze(0).contiguous()
@@ -119,6 +119,15 @@ def set_block_mapping(hpu_attention_meta: HPUPagedAttentionMetadata, batch_size)
     hpu_attention_meta = hpu_attention_meta._replace(
         attn_bias=attn_bias, block_mapping=block_mapping.to(dtype)
     )
+    if hpu_attention_meta.block_groups_in_window is not None:
+        block_mapping = torch.nn.functional.one_hot(
+            hpu_attention_meta.block_groups_in_window, num_classes=batch_size
+        )
+        attn_bias = torch.log(hpu_attention_meta.slots_in_window_mask.float())
+        hpu_attention_meta = hpu_attention_meta._replace(
+            attn_bias_in_window=attn_bias,
+            block_mapping_in_window=block_mapping.to(dtype),
+        )
     return hpu_attention_meta
 
 
@@ -132,17 +141,34 @@ def paged_attention(
     kv_scales: KVScales,
     softcap: Optional[float] = None,
     hpu_attention_meta: HPUPagedAttentionMetadata,
+    window_size_left: int = -1,
 ):
     batch_size, head_num, head_size = query.shape
     fp8_kv = kv_cache.dtype == torch.float8_e4m3fn
     output = ops.flat_pa(
         query=query.view(batch_size, 1, head_num * head_size),
         key_cache=kv_cache.key,
         value_cache=kv_cache.value,
-        block_list=hpu_attention_meta.block_list,
-        block_mapping=hpu_attention_meta.block_mapping,
-        block_bias=hpu_attention_meta.attn_bias,
-        block_groups=hpu_attention_meta.block_groups,
+        block_list=(
+            hpu_attention_meta.block_list
+            if window_size_left == -1
+            else hpu_attention_meta.block_list_in_window
+        ),
+        block_mapping=(
+            hpu_attention_meta.block_mapping
+            if window_size_left == -1
+            else hpu_attention_meta.block_mapping_in_window
+        ),
+        block_bias=(
+            hpu_attention_meta.attn_bias
+            if window_size_left == -1
+            else hpu_attention_meta.attn_bias_in_window
+        ),
+        block_groups=(
+            hpu_attention_meta.block_groups
+            if window_size_left == -1
+            else hpu_attention_meta.block_groups_in_window
+        ),
         block_size=BLOCK_SIZE,
         scale=softmax_scale,
         matmul_qk_op=FP8Matmul(kv_scales.key_scale) if fp8_kv else Matmul(),

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

@@ -288,6 +288,7 @@ def forward(
                 softcap=self.softcap,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.window_size,
             )
 
         return self.o_proj(

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

@@ -135,9 +135,6 @@ def __init__(
         self.causal = causal
         if is_sliding:
             self.window_size = config.sliding_window
-            # TODO: remove this hack to support local sliding window
-            config = copy.deepcopy(config)
-            config.rope_scaling = dict(rope_type="default")
             self.rotary_emb = local_rotary_emb
         else:
             self.window_size = -1
@@ -267,6 +264,7 @@ def forward(
                 softcap=self.softcap,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.window_size,
             )
 
         return self.o_proj(
@@ -425,8 +423,10 @@ def __init__(self, prefix: str, config, weights, causal: bool):
         process_group = weights.process_group
         self.tp_rank = process_group.rank()
         self.tp_world_size = process_group.size()
+        local_config = copy.deepcopy(config)
+        local_config.rope_scaling = dict(rope_type="default")
         local_rotary_emb = PositionRotaryEmbedding.static(
-            config=config,
+            config=local_config,
             dim=config.head_dim,
             base=config.rope_local_base_freq,
             device=weights.device,

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

@@ -224,6 +224,7 @@ def forward(
                 seqlen,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.max_past,
             )
 
         return self.o_proj(

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

@@ -62,7 +62,9 @@ def __init__(
     ):
         super().__init__()
         self.max_past = (
-            config.sliding_window if config.sliding_window is not None else -1
+            config.sliding_window
+            if config.use_sliding_window and config.sliding_window is not None
+            else -1
         )
         self.num_heads = config.num_attention_heads
         self.hidden_size = config.hidden_size
@@ -150,6 +152,7 @@ def forward(
                 seqlen,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.max_past,
             )
 
         return self.o_proj(attn_output.view(-1, self.num_heads * self.head_size))

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

@@ -167,6 +167,7 @@ def forward(
                 seqlen,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.max_past,
             )
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()

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

@@ -190,6 +190,7 @@ def forward(
                 seqlen,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.max_past,
             )
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()

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

@@ -280,6 +280,7 @@ def forward(
                 seqlen,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                window_size_left=self.max_past,
             )
 
         return self.o_proj(