attn_mask for sliding_window_decode

Author: rebel-thkim

src/optimum/rbln/ops/sliding_window_attn.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 
+from typing import Optional
+
 import torch
 from torch import Tensor
-from typing import Optional
 
 
 @torch.library.custom_op(
@@ -95,6 +96,7 @@ def paged_sliding_window_attn_decode(
     scale: Tensor,
     block_table: Tensor,
     block_size: int,
+    attn_mask: Tensor,
     s_aux: Optional[Tensor] = None,
 ) -> Tensor:
     return torch.empty_like(q)
@@ -112,6 +114,7 @@ def paged_sliding_window_attn_decode_fake(
     scale: Tensor,
     block_table: Tensor,
     block_size: int,
+    attn_mask: Tensor,
     s_aux: Optional[Tensor] = None,
 ) -> Tensor:
     return torch.empty_like(q)

src/optimum/rbln/transformers/models/colpali/modeling_colpali.py

@@ -14,8 +14,7 @@
 
 import bisect
 from pathlib import Path
-from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from transformers import PretrainedConfig, PreTrainedModel

src/optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -372,7 +372,11 @@ def get_local_cache_positions(self, position_ids, query_position):
             torch.clamp(position_ids, max=max_cache_len)[:, :1] + valid_input_len
         )  # cache offset for next steps
 
-        return cache_seq_len, cache_offset
+        # Causal mask for sliding window attention
+        attn_mask = torch.arange(max_cache_len)[None, :] - cache_seq_len
+        attn_mask = torch.where(attn_mask > 0, 0.0, 1.0)[:, None, :, None]
+
+        return cache_seq_len, cache_offset, attn_mask
 
     def get_last_layernorm(self) -> nn.LayerNorm:
         return self._original_mod.norm
@@ -458,7 +462,7 @@ def forward(
 
         # Get local cache positions for sliding window layers
         if len(self.sliding_window_layers) > 0:
-            sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
+            sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position, hidden_states)
 
         for layer_idx, layer in enumerate(self.layers):
             is_sliding = True if layer_idx in self.sliding_window_layers else False
@@ -1128,6 +1132,9 @@ def forward(
         if self.phase == "prefill" or self.phase == "image_prefill":
             op_args["is_bidirectional"] = self.phase == "image_prefill"  # FIXME, Hard-coded for Gemma3.
 
+        if self.phase == "decode":
+            op_args["attn_mask"] = attn_mask
+
         if s_aux is not None:
             op_args["s_aux"] = s_aux
 

src/optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -97,11 +97,9 @@ def redirect(func):
 
     def can_generate(self):
         return True
-    
-    
+
     @classmethod
     def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
-        
         with no_init_weights():
             model_cls_name = model.model.language_model.__class__.__name__
             causal_model_cls_name = model_cls_name.replace("TextModel", "ForCausalLM")

src/optimum/rbln/transformers/models/gpt_oss/gpt_oss_architecture.py

@@ -24,7 +24,6 @@
     DecoderOnlyAttention,
     DecoderOnlyLayer,
     DecoderOnlyWrapper,
-    DecoderOnlyAttention,
 )
 
 
@@ -120,7 +119,9 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
         hidden_states = hidden_states.repeat(num_experts, 1)
         hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
 
-        gate_up = torch.bmm(hidden_states, self.gate_up_proj.to(hidden_states.dtype)) + self.gate_up_proj_bias[..., None, :].to(hidden_states.dtype)
+        gate_up = torch.bmm(hidden_states, self.gate_up_proj.to(hidden_states.dtype)) + self.gate_up_proj_bias[
+            ..., None, :
+        ].to(hidden_states.dtype)
         gate, up = gate_up[..., ::2], gate_up[..., 1::2]
         gate = gate.clamp(min=None, max=self.limit)
         up = up.clamp(min=-self.limit, max=self.limit)

src/optimum/rbln/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Union, TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, Union
 
 from transformers import PretrainedConfig
 
 from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
 from .gpt_oss_architecture import RBLNGptOssWrapper
 
+
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
-    from transformers import PreTrainedModel
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
 
 class RBLNGptOssForCausalLM(RBLNDecoderOnlyModelForCausalLM):

src/optimum/rbln/transformers/models/grounding_dino/modeling_grounding_dino.py

@@ -204,8 +204,7 @@ def save_torch_artifacts(
         save_dict["bbox_embed"] = model.bbox_embed.state_dict()
 
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
-        
-    
+
     @classmethod
     def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
         model.encoder = model.model.encoder

src/optimum/rbln/transformers/models/siglip/modeling_siglip.py

@@ -66,7 +66,9 @@ class RBLNSiglipVisionModel(RBLNModel):
     _tp_support = False
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig
+    ) -> torch.nn.Module:
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,