optimum-rbln 0.7.2rc2__py3-none-any.whl → 0.7.3a1__py3-none-any.whl

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

@@ -19,7 +19,12 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig, PreTrainedModel
 
-from ....ops import register_rbln_custom_attention, register_rbln_custom_flash_attention
+from ....ops import (
+    register_rbln_custom_causal_masked_attention,
+    register_rbln_custom_flash_causal_masked_attention,
+    register_rbln_custom_flash_masked_attention,
+    register_rbln_custom_masked_attention,
+)
 from ....utils import logging
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
 
@@ -128,6 +133,7 @@ class DecoderOnlyWrapper(nn.Module):
         max_seq_len: int,
         use_rotary_emb: bool,
         attn_impl: str,
+        use_attention_mask: bool,
         kvcache_partition_len: Optional[int] = None,
     ):
         super().__init__()
@@ -139,12 +145,19 @@ class DecoderOnlyWrapper(nn.Module):
             self.rotary_emb = None
 
         self.attn_impl = attn_impl
+        self.use_attention_mask = use_attention_mask
         if self.attn_impl == "flash_attn":
             self.kvcache_partition_len = kvcache_partition_len or DEFAULT_FLASH_ATTN_PARTITION_LENGTH
-            register_rbln_custom_flash_attention()
+            if self.use_attention_mask:
+                register_rbln_custom_flash_masked_attention()
+            else:
+                register_rbln_custom_flash_causal_masked_attention()
         elif self.attn_impl == "eager":
             self.kvcache_partition_len = None
-            register_rbln_custom_attention()
+            if self.use_attention_mask:
+                register_rbln_custom_masked_attention()
+            else:
+                register_rbln_custom_causal_masked_attention()
         else:
             raise ValueError(f"Unknown attn_impl : {self.attn_impl}")
 
@@ -154,7 +167,7 @@ class DecoderOnlyWrapper(nn.Module):
                 f" or equal to max_seq_len({max_seq_len})!"
             )
 
-        self.causal_lm = self.convert_to_rbln_causal_lm(causal_lm)
+        self.causal_lm = self.convert_to_rbln_causal_lm(causal_lm, max_seq_len)
 
         self.num_hidden_layers = getattr(self.config, "num_hidden_layers", None) or getattr(self.config, "n_layer")
         self._phase = "prefill"
@@ -162,21 +175,25 @@ class DecoderOnlyWrapper(nn.Module):
     def get_rotary_emb(self, max_seq_len):
         return RotaryEmbedding(config=self.config, max_seq_len_cached=max_seq_len)
 
-    def convert_to_rbln_causal_lm(self, causal_lm: PreTrainedModel):
+    def convert_to_rbln_causal_lm(self, causal_lm: PreTrainedModel, max_seq_len: int):
         new_layers = []
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
-                new_self_attn = DecoderOnlyAttention(layer.self_attn)
+                new_self_attn = DecoderOnlyAttention(layer.self_attn, self.use_attention_mask)
             elif self.attn_impl == "flash_attn":
                 new_self_attn = DecoderOnlyFlashAttention(
-                    layer.self_attn, kvcache_partition_len=self.kvcache_partition_len
+                    layer.self_attn,
+                    kvcache_partition_len=self.kvcache_partition_len,
+                    use_attention_mask=self.use_attention_mask,
                 )
             else:
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")
 
             new_layer = DecoderOnlyLayer(layer, new_self_attn)
             new_layers.append(new_layer)
-        new_model = DecoderOnlyModel(causal_lm.model, new_layers, partition_len=self.kvcache_partition_len)
+        new_model = DecoderOnlyModel(
+            causal_lm.model, new_layers, partition_len=self.kvcache_partition_len, max_seq_len=max_seq_len
+        )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
 
@@ -191,23 +208,42 @@ class DecoderOnlyWrapper(nn.Module):
 
     def forward(self, *args):
         if self.phase == "decode":
-            (
-                input_ids_or_inputs_embeds,
-                attention_mask,
-                cache_position,
-                *past_key_values,
-            ) = args
+            if self.use_attention_mask:
+                (
+                    input_ids_or_inputs_embeds,
+                    cache_position,
+                    attention_mask,
+                    *past_key_values,
+                ) = args
+            else:
+                (
+                    input_ids_or_inputs_embeds,
+                    cache_position,
+                    *past_key_values,
+                ) = args
+                attention_mask = None
             batch_position = torch.tensor(0, dtype=torch.int16)
             query_position = None
         elif self.phase == "prefill":
-            (
-                input_ids_or_inputs_embeds,
-                attention_mask,
-                cache_position,
-                batch_position,
-                query_position,
-                *past_key_values,
-            ) = args
+            if self.use_attention_mask:
+                (
+                    input_ids_or_inputs_embeds,
+                    cache_position,
+                    attention_mask,
+                    batch_position,
+                    query_position,
+                    *past_key_values,
+                ) = args
+            else:
+                (
+                    input_ids_or_inputs_embeds,
+                    cache_position,
+                    batch_position,
+                    query_position,
+                    *past_key_values,
+                ) = args
+                attention_mask = None
+
         else:
             raise ValueError(f"Unknown phase: {self.phase}")
 
@@ -338,12 +374,13 @@ class DecoderOnlyModel(nn.Module):
         _phase: Current processing phase ("prefill" or "decode")
     """
 
-    def __init__(self, model, layers: List["DecoderOnlyLayer"], partition_len=None):
+    def __init__(self, model, layers: List["DecoderOnlyLayer"], partition_len=None, max_seq_len=None):
         super().__init__()
         self._original_mod = model
         self.layers = nn.ModuleList(layers)
         self._phase = "prefill"
         self.partition_len = partition_len
+        self.max_seq_len = max_seq_len
 
     @property
     def phase(self):
@@ -410,7 +447,7 @@ class DecoderOnlyModel(nn.Module):
 
         # get cos,sin vector if needed
         if rotary_emb is not None:
-            cos, sin = rotary_emb(hidden_states, attention_mask.shape[-1])  # dtype carrier, max_seq_len
+            cos, sin = rotary_emb(hidden_states, self.max_seq_len)  # dtype carrier, max_seq_len
             cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, cache_position)
         else:
             batch_size = inputs_embeds.shape[0]
@@ -542,7 +579,7 @@ class DecoderOnlyAttention(nn.Module):
         self_attn: Original attention module from the base model
     """
 
-    def __init__(self, self_attn):
+    def __init__(self, self_attn, use_attention_mask):
         super().__init__()
         self._original_mod = self_attn
         self.layer_idx = self_attn.layer_idx
@@ -560,6 +597,7 @@ class DecoderOnlyAttention(nn.Module):
         else:
             self.num_key_value_heads = self.num_heads
 
+        self.use_attention_mask = use_attention_mask
         self.attention = self.get_attention()
         self.__post_init__()
 
@@ -573,7 +611,7 @@ class DecoderOnlyAttention(nn.Module):
         self.attention.phase = phase
 
     def get_attention(self):
-        return AttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads)
+        return AttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads, self.use_attention_mask)
 
     def __post_init__(self):
         self.q_proj = self._original_mod.q_proj
@@ -648,12 +686,13 @@ class DecoderOnlyAttention(nn.Module):
 
 
 class AttentionOp(nn.Module):
-    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int):
+    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int, use_attention_mask: bool):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.num_key_value_heads = num_key_value_heads
         self.phase = "prefill"
+        self.use_attention_mask = use_attention_mask
 
     def forward(
         self,
@@ -686,7 +725,8 @@ class AttentionOp(nn.Module):
         # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
         key_state = key_state.unsqueeze(2)  # 1, 32, 1, 128, 128
         value_state = value_state.unsqueeze(2)
-        attn_mask = attn_mask.unsqueeze(2)
+        if self.use_attention_mask:
+            attn_mask = attn_mask.unsqueeze(2)
 
         if self.phase == "decode":
             batch_size = key_state.shape[0]
@@ -702,29 +742,52 @@ class AttentionOp(nn.Module):
         )
 
         if self.phase == "decode":
-            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.attn_decode(
-                query_state,
-                key_state,
-                value_state,
-                attn_mask,
-                past_key_state.unsqueeze(2),
-                past_value_state.unsqueeze(2),
-                seq_position,
-                scale,
-            )
+            if self.use_attention_mask:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.masked_attn_decode(
+                    query_state,
+                    key_state,
+                    value_state,
+                    attn_mask,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    seq_position,
+                    scale,
+                )
+            else:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.causal_masked_attn_decode(
+                    query_state,
+                    key_state,
+                    value_state,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    seq_position,
+                    scale,
+                )
 
         else:
-            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.attn_prefill(
-                query_state,
-                key_state,
-                value_state,
-                attn_mask,
-                past_key_state.unsqueeze(2),
-                past_value_state.unsqueeze(2),
-                batch_position,
-                seq_position,
-                scale,
-            )
+            if self.use_attention_mask:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.masked_attn_prefill(
+                    query_state,
+                    key_state,
+                    value_state,
+                    attn_mask,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    batch_position,
+                    seq_position,
+                    scale,
+                )
+            else:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.causal_masked_attn_prefill(
+                    query_state,
+                    key_state,
+                    value_state,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    batch_position,
+                    seq_position,
+                    scale,
+                )
 
         attn_output = attn_output.view(batch_size, self.num_heads, -1, self.head_dim)
         attn_output = attn_output.transpose(1, 2).contiguous()
@@ -826,12 +889,19 @@ class RotaryEmbedding(nn.Module):
 
 
 class DecoderOnlyFlashAttention(DecoderOnlyAttention):
-    def __init__(self, self_attn, kvcache_partition_len):
+    def __init__(self, self_attn, kvcache_partition_len, use_attention_mask):
         self.kvcache_partition_size = kvcache_partition_len
-        super().__init__(self_attn=self_attn)
+        # self.use_attention_mask = use_attention_mask
+        super().__init__(self_attn=self_attn, use_attention_mask=use_attention_mask)
 
     def get_attention(self):
-        return FlashAttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads, self.kvcache_partition_size)
+        return FlashAttentionOp(
+            self.num_heads,
+            self.head_dim,
+            self.num_key_value_heads,
+            self.kvcache_partition_size,
+            self.use_attention_mask,
+        )
 
     def forward(
         self,
@@ -878,8 +948,20 @@ class DecoderOnlyFlashAttention(DecoderOnlyAttention):
 
 
 class FlashAttentionOp(AttentionOp):
-    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int, kvcache_partition_len: int):
-        super().__init__(num_heads=num_heads, head_dim=head_dim, num_key_value_heads=num_key_value_heads)
+    def __init__(
+        self,
+        num_heads: int,
+        head_dim: int,
+        num_key_value_heads: int,
+        kvcache_partition_len: int,
+        use_attention_mask: bool,
+    ):
+        super().__init__(
+            num_heads=num_heads,
+            head_dim=head_dim,
+            num_key_value_heads=num_key_value_heads,
+            use_attention_mask=use_attention_mask,
+        )
         self.kvcache_partition_size = kvcache_partition_len
 
     def forward(
@@ -897,7 +979,8 @@ class FlashAttentionOp(AttentionOp):
         # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
         key_state = key_state.unsqueeze(2)
         value_state = value_state.unsqueeze(2)
-        attn_mask = attn_mask.unsqueeze(2)
+        if self.use_attention_mask:
+            attn_mask = attn_mask.unsqueeze(2)
 
         if self.phase == "decode":
             batch_size = key_state.shape[0]
@@ -913,30 +996,55 @@ class FlashAttentionOp(AttentionOp):
         )
 
         if self.phase == "decode":
-            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_attn_decode(
-                query_state,
-                key_state,
-                value_state,
-                attn_mask,
-                past_key_state.unsqueeze(2),
-                past_value_state.unsqueeze(2),
-                seq_position,
-                scale,
-                self.kvcache_partition_size,
-            )
+            if self.use_attention_mask:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_masked_attn_decode(
+                    query_state,
+                    key_state,
+                    value_state,
+                    attn_mask,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    seq_position,
+                    scale,
+                    self.kvcache_partition_size,
+                )
+            else:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_causal_masked_attn_decode(
+                    query_state,
+                    key_state,
+                    value_state,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    seq_position,
+                    scale,
+                    self.kvcache_partition_size,
+                )
         else:
-            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_attn_prefill(
-                query_state,
-                key_state,
-                value_state,
-                attn_mask,
-                past_key_state.unsqueeze(2),
-                past_value_state.unsqueeze(2),
-                batch_position,
-                seq_position,
-                scale,
-                self.kvcache_partition_size,
-            )
+            if self.use_attention_mask:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_masked_attn_prefill(
+                    query_state,
+                    key_state,
+                    value_state,
+                    attn_mask,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    batch_position,
+                    seq_position,
+                    scale,
+                    self.kvcache_partition_size,
+                )
+            else:
+                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_causal_masked_attn_prefill(
+                    query_state,
+                    key_state,
+                    value_state,
+                    past_key_state.unsqueeze(2),
+                    past_value_state.unsqueeze(2),
+                    batch_position,
+                    seq_position,
+                    scale,
+                    self.kvcache_partition_size,
+                )
 
         # reshape for removing repeat_kv
         attn_output = attn_output.view(batch_size, self.num_heads, -1, self.head_dim)

optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -50,12 +50,15 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         phase: str,
         batch_size: int,
         dec_attn_mask: torch.Tensor,
+        use_attention_mask: bool,
         **kwargs: Any,
     ) -> None:
         super().__init__(runtime, **kwargs)
         self.phase = phase
         self.batch_size = batch_size
 
+        self.use_attention_mask = use_attention_mask
+
         # shared tensor between prefill and decode phase
         self.dec_attn_mask = dec_attn_mask
 
@@ -110,7 +113,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         if batch_size != cache_position.shape[0]:
             raise RuntimeError(f"Cache position size mismatch: got {cache_position.shape[0]}, expected {batch_size}.")
 
-        if attention_mask is None:
+        if self.use_attention_mask and attention_mask is None:
             for b_idx in range(batch_size):
                 decoding_step = cache_position[b_idx].item()
                 if not (0 <= decoding_step < self.dec_attn_mask.shape[-1]):
@@ -119,10 +122,12 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                     )
                 self.dec_attn_mask[b_idx, :, :, decoding_step] = 1
 
+            attention_mask = self.dec_attn_mask
+
         logits = super().forward(
             inputs,
-            self.dec_attn_mask if attention_mask is None else attention_mask,
             cache_position,
+            attention_mask if self.use_attention_mask else None,
         )
 
         return logits
@@ -156,7 +161,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             )
 
         # Initialize attention mask for chunked processing
-        chunked_attention_mask = torch.zeros(1, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.float32)
+        if self.use_attention_mask:
+            chunked_attention_mask = torch.zeros(1, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.float32)
 
         # Buffer for storing output logits
         out_buffers = [
@@ -195,28 +201,41 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             input_chunk = inputs[:, step : step + self.prefill_chunk_size]
             cache_pos_chunk = cache_position[:, step : step + self.prefill_chunk_size]
 
-            # Update attention mask to ensure proper causal behavior
-            if step >= self.prefill_chunk_size:
-                chunked_attention_mask[:, :, :, step - self.prefill_chunk_size : step] = 1
-            chunked_attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
+            if self.use_attention_mask:
+                # Update attention mask to ensure proper causal behavior
+                if step >= self.prefill_chunk_size:
+                    chunked_attention_mask[:, :, :, step - self.prefill_chunk_size : step] = 1
+                chunked_attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
 
             # Define batch position and query position
             batch_position = torch.tensor(batch_idx, dtype=torch.int16)
             query_position = torch.tensor((query_length - 1) % self.prefill_chunk_size, dtype=torch.int16)
 
+            if self.use_attention_mask:
+                args = (
+                    input_chunk,
+                    cache_pos_chunk,
+                    chunked_attention_mask,
+                    batch_position,
+                    query_position,
+                )
+            else:
+                args = (
+                    input_chunk,
+                    cache_pos_chunk,
+                    batch_position,
+                    query_position,
+                )
             # Forward pass for the current chunk
             logits = super().forward(
-                input_chunk,
-                chunked_attention_mask,
-                cache_pos_chunk,
-                batch_position,
-                query_position,
+                *args,
                 out=out_buffers,
             )
 
-        # Update decoder attention mask with processed KV-cache length from prefill phase
-        self.dec_attn_mask[batch_idx].fill_(0)
-        self.dec_attn_mask[batch_idx, :, :, :query_length] = 1
+        if self.use_attention_mask:
+            # Update decoder attention mask with processed KV-cache length from prefill phase
+            self.dec_attn_mask[batch_idx].fill_(0)
+            self.dec_attn_mask[batch_idx, :, :, :query_length] = 1
 
         return logits
 
@@ -256,6 +275,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         self.batch_size = self.rbln_config.model_cfg["batch_size"]
         self.max_seq_len = self.rbln_config.model_cfg["max_seq_len"]
         self.prefill_chunk_size = self.rbln_config.model_cfg["prefill_chunk_size"]
+        self.use_attention_mask = self.rbln_config.model_cfg["use_attention_mask"]
 
         main_input_name = self.main_input_name
         if self.rbln_config.model_cfg["use_inputs_embeds"]:
@@ -282,6 +302,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             vocab_size=self.config.vocab_size,
             max_seq_len=self.max_seq_len,
             prefill_chunk_size=self.prefill_chunk_size,
+            use_attention_mask=self.use_attention_mask,
         )
         self.decoder = RBLNRuntimeModel(
             runtime=self.model[1],
@@ -290,6 +311,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             phase="decode",
             batch_size=self.batch_size,
             dec_attn_mask=dec_attn_mask,
+            use_attention_mask=self.use_attention_mask,
         )
 
     @classmethod
@@ -363,7 +385,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
 
-        val = getattr(self.hf_class, __name, None) or getattr(PreTrainedModel, __name)
+        val = getattr(self.get_hf_class(), __name, None) or getattr(PreTrainedModel, __name)
         if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
             return redirect(val)
         return val
@@ -388,6 +410,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         wrapper_cfg["attn_impl"] = rbln_config.model_cfg.get("attn_impl")
         wrapper_cfg["kvcache_partition_len"] = rbln_config.model_cfg.get("kvcache_partition_len")
         wrapper_cfg["use_rotary_emb"] = cls._use_rotary_emb
+        wrapper_cfg["use_attention_mask"] = rbln_config.model_cfg.get("use_attention_mask")
 
         return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()
 
@@ -448,11 +471,18 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
         rbln_batch_size = rbln_kwargs.get("batch_size", None)
         rbln_use_inputs_embeds = rbln_kwargs.get("use_inputs_embeds", None)
+        rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
         rbln_attn_impl = rbln_kwargs.get("attn_impl", None)
         rbln_kvcache_partition_len = rbln_kwargs.get("kvcache_partition_len", None)
         rbln_quantization = QuantizationManager.validate_quantization_config(rbln_kwargs.get("quantization", None))
         rbln_prefill_chunk_size = rbln_kwargs.get("prefill_chunk_size", None)
 
+        if rbln_use_attention_mask is None:
+            rbln_use_attention_mask = False
+            rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
+            if rbln_npu == "RBLN-CA02":
+                rbln_use_attention_mask = True
+
         if rbln_prefill_chunk_size is None:
             rbln_prefill_chunk_size = 128
         elif rbln_prefill_chunk_size % 64 != 0 or rbln_prefill_chunk_size == 0:
@@ -495,13 +525,20 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
             input_info = [
                 main_input,
-                ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "float32"),
                 (
                     "cache_position",
                     [batch_size, query_length],
                     "int32",
                 ),
             ]
+
+            if rbln_use_attention_mask:
+                input_info.extend(
+                    [
+                        ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "float32"),
+                    ]
+                )
+
             if query_length > 1:
                 input_info.extend(
                     [
@@ -555,6 +592,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                 "max_seq_len": rbln_max_seq_len,
                 "batch_size": rbln_batch_size,
                 "prefill_chunk_size": rbln_prefill_chunk_size,
+                "use_attention_mask": rbln_use_attention_mask,
                 "use_inputs_embeds": rbln_use_inputs_embeds,
                 "kvcache_partition_len": rbln_kvcache_partition_len,
                 "attn_impl": rbln_attn_impl,

optimum/rbln/transformers/models/exaone/exaone_architecture.py

@@ -36,11 +36,11 @@ logger = logging.get_logger(__name__)
 class ExaoneForCausalLMWrapper(DecoderOnlyWrapper):
     """A wrapper class for the Exaone model with a language modeling head."""
 
-    def convert_to_rbln_causal_lm(self, causal_lm: "ExaoneForCausalLM"):
+    def convert_to_rbln_causal_lm(self, causal_lm: "ExaoneForCausalLM", max_seq_len: int):
         new_layers = []
         for layer in causal_lm.transformer.h:
             if self.attn_impl == "eager":
-                new_self_attn = ExaoneAttention(layer.attn.attention)
+                new_self_attn = ExaoneAttention(layer.attn.attention, self.use_attention_mask)
             elif self.attn_impl == "flash_attn":
                 new_self_attn = ExaoneFlashAttention(
                     layer.attn.attention, kvcache_partition_len=self.kvcache_partition_len
@@ -50,7 +50,9 @@ class ExaoneForCausalLMWrapper(DecoderOnlyWrapper):
 
             new_layer = ExaoneLayer(layer, new_self_attn)
             new_layers.append(new_layer)
-        new_model = ExaoneModel(causal_lm.transformer, new_layers, partition_len=self.kvcache_partition_len)
+        new_model = ExaoneModel(
+            causal_lm.transformer, new_layers, partition_len=self.kvcache_partition_len, max_seq_len=max_seq_len
+        )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
 

optimum/rbln/transformers/models/gemma/gemma_architecture.py

@@ -29,11 +29,11 @@ if TYPE_CHECKING:
 
 
 class GemmaWrapper(DecoderOnlyWrapper):
-    def convert_to_rbln_causal_lm(self, causal_lm: "GemmaForCausalLM"):
+    def convert_to_rbln_causal_lm(self, causal_lm: "GemmaForCausalLM", max_seq_len: int):
         new_layers = []
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
-                new_self_attn = DecoderOnlyAttention(layer.self_attn)
+                new_self_attn = DecoderOnlyAttention(layer.self_attn, self.use_attention_mask)
             elif self.attn_impl == "flash_attn":
                 new_self_attn = DecoderOnlyFlashAttention(
                     layer.self_attn, kvcache_partition_len=self.kvcache_partition_len
@@ -42,7 +42,9 @@ class GemmaWrapper(DecoderOnlyWrapper):
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")
             new_layer = DecoderOnlyLayer(layer, new_self_attn)
             new_layers.append(new_layer)
-        new_model = GemmaModel(causal_lm.model, new_layers, partition_len=self.kvcache_partition_len)
+        new_model = GemmaModel(
+            causal_lm.model, new_layers, partition_len=self.kvcache_partition_len, max_seq_len=max_seq_len
+        )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm