optimum-rbln 0.8.2a0__tar.gz → 0.8.2a1__tar.gz

{optimum_rbln-0.8.2a0 → optimum_rbln-0.8.2a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.2a0
+Version: 0.8.2a1
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.8.2a0 → optimum_rbln-0.8.2a1}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.2a0'
-__version_tuple__ = version_tuple = (0, 8, 2, 'a0')
+__version__ = version = '0.8.2a1'
+__version_tuple__ = version_tuple = (0, 8, 2, 'a1')

{optimum_rbln-0.8.2a0 → optimum_rbln-0.8.2a1}/src/optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -231,15 +231,16 @@ class DecoderOnlyWrapper(nn.Module):
     def convert_to_rbln_causal_lm(self, causal_lm: PreTrainedModel, max_seq_len: int):
         new_layers = []
         for layer_idx, layer in enumerate(self.get_decoder_layers(causal_lm)):
+            is_sliding = layer_idx in self.sliding_window_layers
             new_self_attn = self.get_rbln_attn_class()(
                 self.get_attn_layer(layer),
-                self.use_attention_mask,
+                self.use_attention_mask if not is_sliding else True,
                 self.use_position_ids,
                 kvcache_block_size=self.sliding_window
                 if layer_idx in self.sliding_window_layers
                 else self.kvcache_block_size,
-                is_sliding=layer_idx in self.sliding_window_layers,
-                attn_impl=self.attn_impl,
+                is_sliding=is_sliding,
+                attn_impl=self.attn_impl if not is_sliding else "eager",
                 kvcache_partition_len=self.kvcache_partition_len,
             )
             new_layer = self.get_rbln_layer_class()(layer, new_self_attn)
@@ -720,10 +721,6 @@ class DecoderOnlyAttention(nn.Module):
         self.use_position_ids = use_position_ids
         self.is_sliding = is_sliding
         self.attn_impl = attn_impl
-
-        if self.is_sliding and self.attn_impl != "eager":
-            raise NotImplementedError("Sliding window attention is only supported with eager attention.")
-
         self.kvcache_partition_len = kvcache_partition_len
 
         setattr(self, self.get_attention_name(), self.create_attention_op())
@@ -873,7 +870,8 @@ class AttentionOp(nn.Module):
 
     def get_attn_op_name(self):
         phase = "decode" if self.phase == "decode" else "prefill"
-        if self.use_attention_mask:
+
+        if self.use_attention_mask and not self.use_position_ids:
             attn_op_name = "paged_attn_"
         else:
             attn_op_name = "paged_causal_attn_"
@@ -1074,7 +1072,7 @@ class FlashAttentionOp(AttentionOp):
 
     def get_attn_op_name(self):
         phase = "decode" if self.phase == "decode" else "prefill"
-        if self.use_attention_mask:
+        if self.use_attention_mask and not self.use_position_ids:
             attn_op_name = "paged_flash_attn_"
         else:
             attn_op_name = "paged_flash_causal_attn_"
@@ -1128,7 +1126,7 @@ class FlashAttentionOp(AttentionOp):
             "partition": self.kvcache_partition_size,
         }
 
-        if self.use_attention_mask != self.use_position_ids:
+        if self.use_attention_mask:
             op_args["mask"] = attn_mask
 
         if self.phase == "prefill" or self.phase == "image_prefill":
@@ -1151,8 +1149,8 @@ class FlashAttentionOp(AttentionOp):
 class SlidingWindowAttentionOp(AttentionOp):
     def get_attn_op_name(self):
         phase = "decode" if self.phase == "decode" else "prefill"
-        if self.use_attention_mask:
-            raise NotImplementedError("Attention mask is not supported for sliding window attention.")
+        if not self.use_attention_mask:
+            raise NotImplementedError("Attention mask is needed for sliding window attention.")
 
         attn_op_name = "paged_sliding_window_attn_" + phase
         return attn_op_name