optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1a1__py3-none-any.whl

optimum/rbln/transformers/models/distilbert/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_distilbert import RBLNDistilBertForQuestionAnsweringConfig
+from .modeling_distilbert import RBLNDistilBertForQuestionAnswering
+
+
+__all__ = ["RBLNDistilBertForQuestionAnsweringConfig", "RBLNDistilBertForQuestionAnswering"]

optimum/rbln/transformers/models/distilbert/configuration_distilbert.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_generic import RBLNModelForQuestionAnsweringConfig
+
+
+class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
+    ""

optimum/rbln/transformers/models/distilbert/modeling_distilbert.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...modeling_generic import RBLNModelForQuestionAnswering
+
+
+class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
+    rbln_model_input_names = ["input_ids", "attention_mask"]

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

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNExaoneForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Exaone models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNExaoneForCausalLM, RBLNExaoneForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNExaoneForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNExaoneForCausalLM.from_pretrained(
+        "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

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

@@ -60,7 +60,11 @@ 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, max_seq_len=max_seq_len
+            causal_lm.transformer,
+            new_layers,
+            partition_len=self.kvcache_partition_len,
+            max_seq_len=max_seq_len,
+            sliding_window_layers=self.sliding_window_layers,
         )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm

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

@@ -25,16 +25,62 @@ logger = logging.get_logger(__name__)
 
 class RBLNExaoneForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
-    The Exaone Model transformer with a language modeling head on top (linear layer with weights tied to the input
-    embeddings).
+    The Exaone Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
-    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the
-    library implements for all its model.
+    A class to convert and run pre-trained transformers based ExaoneForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers ExaoneForCausalLM model into a RBLN transformer model by:
 
-    It implements the methods to convert a pre-trained transformers Exaone model into a RBLN transformer model by:
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
 
+    **Configuration:**
+    This model uses [`RBLNExaoneForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNExaoneForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNExaoneForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNExaoneForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNExaoneForCausalLM.from_pretrained(
+            "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNExaoneForCausalLM.from_pretrained(
+            "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNExaoneForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNExaoneForCausalLMConfig
+
+        config = RBLNExaoneForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNExaoneForCausalLM.from_pretrained(
+            "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = ExaoneForCausalLMWrapper

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

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNGemmaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Gemma models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNGemmaForCausalLM, RBLNGemmaForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNGemmaForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNGemmaForCausalLM.from_pretrained(
+        "google/gemma-7b",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

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

@@ -52,7 +52,11 @@ class GemmaWrapper(DecoderOnlyWrapper):
             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, max_seq_len=max_seq_len
+            causal_lm.model,
+            new_layers,
+            partition_len=self.kvcache_partition_len,
+            max_seq_len=max_seq_len,
+            sliding_window_layers=self.sliding_window_layers,
         )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm

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

@@ -27,8 +27,57 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     A class to convert and run pre-trained transformers based GemmaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers GemmaForCausalLM model into a RBLN transformer model by:
+
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNGemmaForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNGemmaForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNGemmaForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNGemmaForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNGemmaForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNGemmaForCausalLMConfig
+
+        config = RBLNGemmaForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNGemmaForCausalLM.from_pretrained(
+            "google/gemma-7b",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = GemmaWrapper

optimum/rbln/transformers/models/gemma3/configuration_gemma3.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import rebel
 
@@ -26,7 +26,7 @@ class RBLNGemma3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
         prefill_chunk_size: Optional[int] = None,
         use_position_ids: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         # use_attention_mask and use_position_ids are always True for Gemma3
         use_attention_mask = use_attention_mask or True
@@ -53,7 +53,7 @@ class RBLNGemma3ForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/gemma3/gemma3_architecture.py

@@ -16,11 +16,9 @@ import copy
 from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
-from torch import nn
 from transformers.models.gemma3.modeling_gemma3 import Gemma3RMSNorm
 
 from ..decoderonly.decoderonly_architecture import (
-    AttentionOp,
     DecoderOnlyAttention,
     DecoderOnlyFlashAttention,
     DecoderOnlyForCausalLM,
@@ -28,7 +26,6 @@ from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyModel,
     DecoderOnlyWrapper,
     RotaryEmbedding,
-    SlidingWindowAttentionOp,
     slice_and_unsqueeze_cos_sin,
 )
 
@@ -50,13 +47,14 @@ class Gemma3ForCausalLMWrapper(DecoderOnlyWrapper):
 
     def convert_to_rbln_causal_lm(self, causal_lm: "Gemma3ForCausalLM", max_seq_len: int):
         new_layers = []
-        for layer in causal_lm.model.layers:
-            if layer.is_sliding:
+        for layer_idx, layer in enumerate(causal_lm.model.layers):
+            if layer_idx in self.sliding_window_layers:
                 new_self_attn = Gemma3Attention(
                     layer.self_attn,
                     use_attention_mask=None,  # FIXME: no use in SWA
                     use_position_ids=self.use_position_ids,
                     kvcache_block_size=self.config.sliding_window,
+                    is_sliding=True,
                 )
             else:
                 if self.attn_impl == "eager":
@@ -65,6 +63,7 @@ class Gemma3ForCausalLMWrapper(DecoderOnlyWrapper):
                         use_attention_mask=self.use_attention_mask,
                         use_position_ids=self.use_position_ids,
                         kvcache_block_size=self.kvcache_block_size,
+                        is_sliding=False,
                     )
                 elif self.attn_impl == "flash_attn":
                     new_self_attn = Gemma3FlashAttention(
@@ -85,131 +84,14 @@ class Gemma3ForCausalLMWrapper(DecoderOnlyWrapper):
             new_layers,
             partition_len=self.kvcache_partition_len,
             max_seq_len=max_seq_len,
+            sliding_window_layers=self.sliding_window_layers,
         )
-        new_causal_lm = Gemma3ForCausalLM(causal_lm, new_model)
+        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
 
-    def forward(self, *args):
-        if self.phase == "decode":
-            (
-                input_ids_or_inputs_embeds,
-                attention_mask,  # used in global layer, 2D attn_mask for padded KVcache.
-                cache_position,
-                position_ids,
-                golbal_block_tables,
-                local_block_tables,
-                *past_key_values,
-            ) = args
-            query_position = None
-
-        elif "prefill" in self.phase:
-            (
-                input_ids_or_inputs_embeds,
-                attention_mask,
-                cache_position,
-                position_ids,
-                query_position,
-                golbal_block_tables,
-                local_block_tables,
-                *past_key_values,
-            ) = args
-
-        else:
-            raise ValueError(f"Unknown phase: {self.phase}")
-
-        if input_ids_or_inputs_embeds.ndim == 2:
-            input_ids = input_ids_or_inputs_embeds
-            inputs_embeds = None
-        elif input_ids_or_inputs_embeds.ndim == 3:
-            input_ids = None
-            inputs_embeds = input_ids_or_inputs_embeds
-        else:
-            raise NotImplementedError(f"Unknown ndim of input : {input_ids_or_inputs_embeds.ndim}")
-
-        if len(past_key_values) != 2 * self.num_hidden_layers:
-            raise ValueError(
-                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.num_hidden_layers}"
-            )
-
-        # [key, value] * n_layer -> ( (key, value) ) * n_layer
-        # cache shape : batch, n_heads, 1, max_seq_len, head_dim
-        _past_key_values = []
-        for i in range(self.config.num_hidden_layers):
-            key_states = past_key_values[i * 2]
-            value_states = past_key_values[i * 2 + 1]
-            past_key_value = [key_states, value_states]
-            _past_key_values.append(past_key_value)
-        past_key_values = _past_key_values
-
-        logit = self.causal_lm(
-            input_ids=input_ids,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            position_ids=position_ids,
-            query_position=query_position,
-            past_key_values=past_key_values,
-            rotary_emb=(self.rotary_emb_global, self.rotary_emb_local),
-            global_block_tables=golbal_block_tables,
-            local_block_tables=local_block_tables,
-        )
-
-        return logit
-
-
-class Gemma3ForCausalLM(DecoderOnlyForCausalLM):
-    def forward(
-        self,
-        input_ids: torch.Tensor = None,
-        inputs_embeds: torch.Tensor = None,
-        attention_mask: torch.Tensor = None,
-        cache_position: torch.Tensor = None,
-        position_ids: torch.Tensor = None,
-        query_position: torch.Tensor = None,
-        past_key_values: Tuple[Tuple[torch.Tensor]] = None,
-        rotary_emb: nn.Module = None,
-        global_block_tables: Optional[torch.Tensor] = None,
-        local_block_tables: Optional[torch.Tensor] = None,
-    ):
-        # outputs
-        hidden_states = self.model(
-            input_ids=input_ids,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            position_ids=position_ids,
-            query_position=query_position,
-            past_key_values=past_key_values,
-            rotary_emb=rotary_emb,
-            global_block_tables=global_block_tables,
-            local_block_tables=local_block_tables,
-        )
-
-        if "prefill" in self.phase:
-            hidden_states = hidden_states[:, query_position.to(torch.int).unsqueeze(0)]
-
-        logits = self.lm_head(hidden_states)
-
-        # Apply final logit softmaxing if configured, e.g. for Gemma2
-        if getattr(self.config, "final_logit_softcapping", None) is not None:
-            logits = logits / self.config.final_logit_softcapping
-            logits = torch.tanh(logits)
-            logits = logits * self.config.final_logit_softcapping
-
-        return logits
-
 
 class Gemma3TextModel(DecoderOnlyModel):
-    def get_local_cache_positions(self, position_ids, query_position):
-        max_cache_len = self._original_mod.config.sliding_window
-        valid_input_len = 1 if query_position is None else query_position + 1
-        cache_seq_len = torch.clamp(position_ids, max=max_cache_len)[:, :1]  # past seen tokens
-        cache_offset = (
-            torch.clamp(position_ids, max=max_cache_len)[:, :1] + valid_input_len
-        )  # cache offset for next steps
-
-        return cache_seq_len, cache_offset
-
+    # Different from DecoderOnlyModel, this model has global and local rotary embeddings.
     def forward(
         self,
         input_ids: torch.Tensor = None,
@@ -254,37 +136,23 @@ class Gemma3TextModel(DecoderOnlyModel):
 
         sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
 
-        for layer in self.layers:
-            if layer.is_sliding:
-                hidden_states = layer(
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    seq_positions=sliding_cache_pos,
-                    past_key_values=past_key_values,
-                    cos=cos_local,
-                    sin=sin_local,
-                    block_tables=local_block_tables,
-                )
-            else:
-                hidden_states = layer(
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    seq_positions=seq_positions,
-                    past_key_values=past_key_values,
-                    cos=cos_global,
-                    sin=sin_global,
-                    block_tables=global_block_tables,
-                )
+        for layer_idx, layer in enumerate(self.layers):
+            is_sliding = True if layer_idx in self.sliding_window_layers else False
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                seq_positions=sliding_cache_pos if is_sliding else seq_positions,
+                past_key_values=past_key_values,
+                cos=cos_local if is_sliding else cos_global,
+                sin=sin_local if is_sliding else sin_global,
+                block_tables=local_block_tables if is_sliding else global_block_tables,
+            )
 
         hidden_states = self.get_last_layernorm()(hidden_states)
         return hidden_states
 
 
 class Gemma3DecoderLayer(DecoderOnlyLayer):
-    def __init__(self, layer, self_attn: "DecoderOnlyAttention"):
-        super().__init__(layer, self_attn)
-        self.is_sliding = self._original_mod.is_sliding
-
     def get_pre_feedforward_layernorm(self) -> Gemma3RMSNorm:
         return self._original_mod.pre_feedforward_layernorm
 
@@ -328,69 +196,10 @@ class Gemma3Attention(DecoderOnlyAttention):
         self.o_proj = self._original_mod.o_proj
         self.q_norm = self._original_mod.q_norm
         self.k_norm = self._original_mod.k_norm
-        self.is_sliding = self._original_mod.is_sliding
 
     def get_attn_scale(self):
         return self._original_mod.config.query_pre_attn_scalar**-0.5
 
-    def get_attention(self):
-        if self._original_mod.is_sliding:
-            return SlidingWindowAttentionOp(
-                self.num_heads,
-                self.head_dim,
-                self.num_key_value_heads,
-                self.use_attention_mask,
-                self.use_position_ids,
-            )
-        else:
-            return AttentionOp(
-                self.num_heads, self.head_dim, self.num_key_value_heads, self.use_attention_mask, self.use_position_ids
-            )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        seq_positions: torch.LongTensor,
-        past_key_values: Tuple[Tuple[torch.Tensor]],
-        cos: Optional[torch.Tensor] = None,
-        sin: Optional[torch.Tensor] = None,
-        block_tables: Optional[torch.Tensor] = None,
-    ):
-        batch_size, query_length, _ = hidden_states.size()
-
-        query_states, key_states, value_states = self.projection(hidden_states=hidden_states)
-
-        query_states = query_states.view(batch_size, query_length, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(batch_size, query_length, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(batch_size, query_length, self.num_key_value_heads, self.head_dim).transpose(
-            1, 2
-        )
-
-        query_states = self.q_norm(query_states)
-        key_states = self.k_norm(key_states)
-        query_states, key_states = self.apply_rotary_pos_embed(query_states, key_states, cos, sin)
-
-        batch_size = query_states.shape[0]
-        if batch_size > 1 and "prefill" in self.phase:
-            raise NotImplementedError(f"batch size should be 1 if prefill phase, but got {batch_size}.")
-
-        attn_output = self.attention(
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            past_key_state=past_key_values[self.layer_idx][0],
-            past_value_state=past_key_values[self.layer_idx][1],
-            seq_position=seq_positions,
-            scale=self.scale,
-            block_tables=block_tables,
-            block_size=self.kvcache_block_size,
-        )
-
-        attn_outputs = self.o_proj(attn_output)
-        return attn_outputs
-
 
 class Gemma3FlashAttention(DecoderOnlyFlashAttention):
     def __post_init__(self):
@@ -400,47 +209,6 @@ class Gemma3FlashAttention(DecoderOnlyFlashAttention):
         self.o_proj = self._original_mod.o_proj
         self.q_norm = self._original_mod.q_norm
         self.k_norm = self._original_mod.k_norm
-        self.is_sliding = self._original_mod.is_sliding
 
     def get_attn_scale(self):
         return self._original_mod.config.query_pre_attn_scalar**-0.5
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        seq_positions: torch.LongTensor,
-        past_key_values: Tuple[Tuple[torch.Tensor]],
-        cos: Optional[torch.Tensor] = None,
-        sin: Optional[torch.Tensor] = None,
-        block_tables: Optional[torch.Tensor] = None,
-    ):
-        batch_size, query_length, _ = hidden_states.size()
-
-        query_states, key_states, value_states = self.projection(hidden_states=hidden_states)
-
-        query_states = query_states.view(batch_size, query_length, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(batch_size, query_length, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(batch_size, query_length, self.num_key_value_heads, self.head_dim).transpose(
-            1, 2
-        )
-
-        query_states = self.q_norm(query_states)
-        key_states = self.k_norm(key_states)
-        query_states, key_states = self.apply_rotary_pos_embed(query_states, key_states, cos, sin)
-
-        attn_output = self.attention(
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            past_key_state=past_key_values[self.layer_idx][0],
-            past_value_state=past_key_values[self.layer_idx][1],
-            seq_position=seq_positions,
-            scale=self.scale,
-            block_tables=block_tables,
-            kvcache_block_size=self.kvcache_block_size,
-        )
-
-        attn_outputs = self.o_proj(attn_output)
-        return attn_outputs