optimum-rbln 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

optimum/rbln/transformers/models/phi/phi_architecture.py

@@ -21,302 +21,101 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-import math
-from typing import Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
-import torch.nn as nn
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-)
+from transformers import PhiForCausalLM
 
-from ...cache_utils import RebelDynamicCache
-from ..decoderonly import (
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyForCausalLM,
+    DecoderOnlyLayer,
+    DecoderOnlyModel,
     DecoderOnlyWrapper,
-    apply_rotary_pos_emb,
-    slice_and_unsqueeze_cos_sin,
+    apply_rotary_pos_emb_partial,
 )
 
 
-class PhiWrapper(DecoderOnlyWrapper):
-    def get_forward_dict(self):
-        forward_dict = {}
-        forward_dict.update(
-            {
-                "wrapper": PhiModel.forward,
-                "model": PhiDecoderLayer.forward,
-                "decoder_layer": PhiAttention.forward,
-            }
-        )
-        return forward_dict
-
-
-class PhiAttention:
-    def _attn(self, query_state, key_state, value_state, attn_mask, past_key_value, batch_idx=0, is_prefill=False):
-        # 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)
-
-        query_state = query_state.view(
-            1,
-            self.num_key_value_heads,
-            self.num_heads // self.num_key_value_heads,
-            -1,
-            self.head_dim,
-        )
-
-        key_state, value_state = past_key_value.update(key_state, value_state, self.layer_idx, batch_idx, is_prefill)
-
-        # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
-        attn_weights = torch.matmul(
-            query_state.to(torch.float32),
-            key_state.to(torch.float32).transpose(3, 4),
-        ) / math.sqrt(self.head_dim)
-        attn_weights = attn_weights + attn_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_state.dtype)
-        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_state)
-
-        # reshape for removing repeat_kv
-        attn_output = attn_output.view(1, self.num_heads, -1, self.head_dim)
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(1, -1, self.num_heads * self.head_dim)
+if TYPE_CHECKING:
+    from transformers import PhiForCausalLM
 
-        return attn_output, key_state, value_state
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        past_key_value: Optional[RebelDynamicCache] = None,
-        batch_index: Optional[int] = None,
-        output_attentions: bool = False,
-        cos: Optional[torch.Tensor] = None,
-        sin: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        bsz, q_len, _ = hidden_states.size()
 
+class PhiWrapper(DecoderOnlyWrapper):
+    def convert_to_rbln_causal_lm(self, causal_lm: "PhiForCausalLM"):
+        new_layers = []
+        for layer in causal_lm.model.layers:
+            if self.attn_impl == "eager":
+                new_self_attn = PhiAttention(layer.self_attn)
+            elif self.attn_impl == "flash_attn":
+                raise NotImplementedError(f"flash attn for {self.__class__} is not implemented yet.")
+            else:
+                raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")
+            new_layer = PhiLayer(layer, new_self_attn)
+            new_layers.append(new_layer)
+        new_model = PhiModel(causal_lm.model, new_layers)
+        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
+        return new_causal_lm
+
+
+class PhiAttention(DecoderOnlyAttention):
+    def __post_init__(self):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.o_proj = self._original_mod.dense
+        self.qk_layernorm = self._original_mod.qk_layernorm
+        self.rotary_ndims = self._original_mod.rotary_ndims
+
+    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
 
         if self.qk_layernorm:
-            query_states = self.q_layernorm(query_states)
-            key_states = self.k_layernorm(key_states)
-
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-        # Partial rotary embedding
-        query_rot, query_pass = (
-            query_states[..., : self.rotary_ndims],
-            query_states[..., self.rotary_ndims :],
-        )
-        key_rot, key_pass = (
-            key_states[..., : self.rotary_ndims],
-            key_states[..., self.rotary_ndims :],
-        )
+            query_states = self._original_mod.q_layernorm(query_states)
+            key_states = self._original_mod.k_layernorm(key_states)
 
-        # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
-        query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin)
+        return query_states, key_states, value_states
 
-        # [batch_size, seq_length, num_heads, head_dim]
-        query_states = torch.cat((query_rot, query_pass), dim=-1)
-        key_states = torch.cat((key_rot, key_pass), dim=-1)
+    def apply_rotary_pos_embed(self, query_states, key_states, cos, sin):
+        return apply_rotary_pos_emb_partial(query_states, key_states, cos, sin, ndim=self.rotary_ndims)
 
-        # Decoder (bsz > 1)
-        if bsz > 1:
-            iterate_results = {"key_states": [], "value_states": [], "attn_output": []}
-            for b in range(bsz):
-                attn_output, key_state, value_state = PhiAttention._attn(
-                    self,
-                    query_states[b].unsqueeze(0),
-                    key_states[b].unsqueeze(0),
-                    value_states[b].unsqueeze(0),
-                    attention_mask[b].unsqueeze(0),
-                    past_key_value,
-                    batch_idx=b,
-                    is_prefill=False,
-                )
-                iterate_results["key_states"].append(key_state)
-                iterate_results["value_states"].append(value_state)
-                iterate_results["attn_output"].append(attn_output)
 
-            key_states = torch.cat(iterate_results["key_states"], dim=0)
-            value_states = torch.cat(iterate_results["value_states"], dim=0)
-            attn_output = torch.cat(iterate_results["attn_output"], dim=0)
-        # Prefill & Decoder (bsz == 1)
-        else:
-            attn_output, key_states, value_states = PhiAttention._attn(
-                self,
-                query_states,
-                key_states,
-                value_states,
-                attention_mask,
-                past_key_value,
-                batch_idx=batch_index,
-                is_prefill=True,
-            )
+class PhiLayer(DecoderOnlyLayer):
+    def get_post_attention_layernorm(self):
+        raise NotImplementedError
 
-        attn_output = self.dense(attn_output)
-
-        if not output_attentions:
-            attn_weights = None
-
-        return attn_output, attn_weights, key_states, value_states
-
-
-class PhiDecoderLayer:
     def forward(
         self,
         hidden_states: torch.Tensor,
-        layer_idx: int,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[RebelDynamicCache] = None,
-        output_attentions: Optional[bool] = None,
-        use_cache: Optional[bool] = None,
-        batch_ids: Optional[torch.LongTensor] = None,
+        attention_mask: torch.Tensor,
+        seq_positions: torch.LongTensor,
+        batch_position: torch.Tensor,
+        past_key_values: Tuple[Tuple[torch.Tensor]],
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
-        forward_dict: Optional[Dict[str, classmethod]] = None,
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`):
-                input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-            position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
-                Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
-                `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-        """
-
+    ):
         residual = hidden_states
 
-        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.get_pre_attention_layernorm()(hidden_states)
 
-        # Self Attention
-        attn_outputs, self_attn_weights, key_states, value_states = forward_dict["decoder_layer"](
-            self.self_attn,
+        attn_outputs, present_key_values = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            batch_index=batch_ids,
-            use_cache=use_cache,
+            seq_positions=seq_positions,
+            batch_position=batch_position,
+            past_key_values=past_key_values,
             cos=cos,
             sin=sin,
-            **kwargs,
         )
-        past_key_value.assign(key_states, value_states, layer_idx)
 
-        attn_outputs = self.resid_dropout(attn_outputs)
+        feed_forward_hidden_states = self._original_mod.mlp(hidden_states)
 
-        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
         hidden_states = attn_outputs + feed_forward_hidden_states + residual
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        if use_cache:
-            outputs += (past_key_value,)
-
-        return outputs
-
-
-class PhiModel:
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[RebelDynamicCache] = None,
-        batch_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = True,
-        output_attentions: Optional[bool] = False,
-        output_hidden_states: Optional[bool] = False,
-        cache_pos_for_partitions: Optional[torch.Tensor] = None,
-        kvcache_partition_size: Optional[torch.Tensor] = None,
-        forward_dict: Optional[Dict[str, classmethod]] = None,
-        rotary_pos_emb=None,
-    ) -> BaseModelOutputWithPast:
-        # retrieve input_ids and inputs_embeds
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
-            )
-
-        # embed positions
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        hidden_states = inputs_embeds
-        attention_mask = (1 - attention_mask) * torch.finfo(torch.float16).min
-
-        # get cos,sin vector
-        cos, sin = rotary_pos_emb(inputs_embeds, attention_mask.shape[-1])
-        cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-
-        for layer_idx, decoder_layer in enumerate(self.layers):
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-            layer_outputs = forward_dict["model"](
-                decoder_layer,
-                hidden_states,
-                layer_idx,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_values,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                batch_ids=batch_ids,
-                cos=cos,
-                sin=sin,
-                cache_pos_for_partitions=cache_pos_for_partitions,
-                kvcache_partition_size=kvcache_partition_size,
-                forward_dict=forward_dict,
-            )
 
-            hidden_states = layer_outputs[0]
+        return hidden_states, present_key_values
 
-            updated_cache = layer_outputs[2 if output_attentions else 1]
 
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.final_layernorm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        # convert RebelDynamicCache to legacy Tuple[Tuple[torch.Tensor]]
-        next_cache = updated_cache.to_legacy_cache()
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
+class PhiModel(DecoderOnlyModel):
+    def get_last_layernorm(self):
+        return self._original_mod.final_layernorm

optimum/rbln/transformers/models/qwen2/qwen2_architecture.py

@@ -21,7 +21,6 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-
 from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper
 
 

optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -26,13 +26,14 @@ import logging
 from abc import ABC
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
-import rebel  # noqa: F401
-import torch  # noqa: F401
+import rebel
+import torch
+from rebel.compile_context import CompileContext
 from transformers import AutoModelForSeq2SeqLM, GenerationConfig, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
-from ....modeling_base import RBLNModel
-from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ....modeling import RBLNModel
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
@@ -66,7 +67,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
 class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     """
     This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence-to-sequence language modeling head) when created with the from_pretrained() class method.
-    This model inherits from [`RBLNBaseModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
     A class to convert and run pre-trained transformers based Seq2SeqLM models on RBLN devices.
     It implements the methods to convert a pre-trained transformers Seq2SeqLM model into a RBLN transformer model by:
@@ -88,49 +89,42 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNConfig):
         wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
-        wrapped_model.encoder.encoder_max_length = rbln_config.model_cfg["enc_max_seq_len"]
-        wrapped_model.encoder.decoder_max_length = rbln_config.model_cfg["dec_max_seq_len"]
+        enc_compile_config = rbln_config.compile_cfgs[0]
+        dec_compile_config = rbln_config.compile_cfgs[1]
 
-        wrapped_model.decoder.encoder_max_length = rbln_config.model_cfg["enc_max_seq_len"]
-        wrapped_model.decoder.decoder_max_length = rbln_config.model_cfg["dec_max_seq_len"]
+        context = CompileContext(use_weight_sharing=False)
 
-        enc_rbln_compile_config = rbln_config.compile_cfgs[0]
-        dec_rbln_compile_config = rbln_config.compile_cfgs[1]
+        enc_example_inputs = enc_compile_config.get_dummy_inputs(fill=0)
 
-        enc_example_inputs = enc_rbln_compile_config.get_dummy_inputs(fill=0)
-        dec_example_inputs = dec_rbln_compile_config.get_dummy_inputs(fill=0)
+        # Mark encoder's static tensors (cross kv states)
+        static_tensors = {}
+        for (name, _, _), tensor in zip(enc_compile_config.input_info, enc_example_inputs):
+            if "key_value_states" in name:
+                static_tensors[name] = tensor
+                context.mark_static_address(tensor)
 
-        enc_example_inputs[3].fill_(0)
-        dec_example_inputs[4].fill_(-1)
+        dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
 
-        enc_scripted_model = torch.jit.trace(wrapped_model.encoder, enc_example_inputs, check_trace=False)
-        dec_scripted_model = torch.jit.trace(wrapped_model.decoder, dec_example_inputs, check_trace=False)
+        # Mark decoder's static tensors (self kv states)
+        for (name, _, _), tensor in zip(dec_compile_config.input_info, dec_example_inputs):
+            if "key_value_states" in name:
+                context.mark_static_address(tensor)
 
-        enc_ir = rebel.torchscript_to_ir(
-            enc_scripted_model,
-            input_names=[v[0] for v in enc_rbln_compile_config.input_info],
-            name=enc_rbln_compile_config.mod_name,
+        compiled_encoder = super().compile(
+            wrapped_model.encoder,
+            enc_compile_config,
+            example_inputs=enc_example_inputs,
+            compile_context=context,
         )
-        dec_ir = rebel.torchscript_to_ir(
-            dec_scripted_model,
-            input_names=[v[0] for v in dec_rbln_compile_config.input_info],
-            name=dec_rbln_compile_config.mod_name,
-        )
-
-        connections = [
-            (enc_ir.outputs[0], enc_ir.inputs[2], dec_ir.inputs[6]),
-            (dec_ir.outputs[1], dec_ir.inputs[5]),
-        ]
 
-        compiled_model = rebel.compile(
-            enc_ir,
-            dec_ir,
-            connections=connections,
-            fusion=enc_rbln_compile_config.fusion,
-            npu=enc_rbln_compile_config.npu,
-            tensor_parallel_size=enc_rbln_compile_config.tensor_parallel_size,
+        compiled_decoder = super().compile(
+            wrapped_model.decoder,
+            dec_compile_config,
+            example_inputs=dec_example_inputs,
+            compile_context=context,
         )
-        return compiled_model
+
+        return {"encoder": compiled_encoder, "decoder": compiled_decoder}
 
     @classmethod
     def _get_rbln_config(
@@ -204,7 +198,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 ],
                 "float32",
             ),
-            ("batch_idx", [], "int32"),
+            ("batch_position", [], "int16"),
         ]
 
         dec_input_info = [
@@ -216,17 +210,16 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 [rbln_batch_size, 1],
                 "int32",
             ),
-            ("batch_position", [], "int32"),
         ]
         dec_input_info.extend(
             [
                 (
-                    "self_key_value_states",
+                    "cross_key_value_states",
                     [
                         n_layer * 2,
                         rbln_batch_size,
                         n_head,
-                        rbln_dec_max_seq_len,
+                        rbln_enc_max_seq_len,
                         d_kv,
                     ],
                     "float32",
@@ -236,24 +229,24 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         dec_input_info.extend(
             [
                 (
-                    "cross_key_value_states",
+                    f"self_key_value_states_{i}",
                     [
-                        n_layer * 2,
                         rbln_batch_size,
                         n_head,
-                        rbln_enc_max_seq_len,
+                        rbln_dec_max_seq_len,
                         d_kv,
                     ],
                     "float32",
                 )
+                for i in range(n_layer * 2)
             ]
         )
-        enc_rbln_compile_config = RBLNCompileConfig(mod_name="encoder", input_info=enc_input_info)
-        dec_rbln_compile_config = RBLNCompileConfig(mod_name="decoder", input_info=dec_input_info)
+        enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
+        dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
-            compile_cfgs=[enc_rbln_compile_config, dec_rbln_compile_config],
+            compile_cfgs=[enc_compile_config, dec_compile_config],
             rbln_kwargs=rbln_kwargs,
         )
 
@@ -270,12 +263,21 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     @classmethod
     def _create_runtimes(
-        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_device_map: Dict[str, int],
+        activate_profiler: Optional[bool] = None,
     ) -> List[rebel.Runtime]:
-        device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
+        if any(model_name not in rbln_device_map for model_name in ["encoder", "decoder"]):
+            cls._raise_missing_compiled_file_error(["encoder", "decoder"])
+
         return [
-            compiled_models[0].create_runtime("encoder", tensor_type="pt", device=device_val),
-            compiled_models[0].create_runtime("decoder", tensor_type="pt", device=device_val),
+            compiled_models[0].create_runtime(
+                tensor_type="pt", device=rbln_device_map["encoder"], activate_profiler=activate_profiler
+            ),
+            compiled_models[1].create_runtime(
+                tensor_type="pt", device=rbln_device_map["decoder"], activate_profiler=activate_profiler
+            ),
         ]
 
     def can_generate(self):
@@ -340,57 +342,11 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             attention_mask=dec_attention_mask,
             encoder_attention_mask=attention_mask,
             cache_position=cache_position,
-            batch_position=torch.tensor(0, dtype=torch.int32),
         )
-        lm_logits = decoder_output.logits[0]
+        lm_logits = decoder_output.logits
 
         return Seq2SeqLMOutput(logits=lm_logits)
 
-    def vllm_forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        cache_position: Union[List[torch.Tensor], torch.Tensor] = None,  # vllm keyword argument
-        batch_idx: Optional[torch.LongTensor] = None,
-        enc_lengths: List[int] = None,  # vllm return current attention_mask length
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor]:
-        # When using vllm, need the output of the encoder (ex. vocab_size + 100) and use that value act as start_token_id in decoder (ex. vocab_size + 99)
-        # encoder
-        if batch_idx is not None:
-            enc_attention_mask = torch.zeros(1, self.rbln_config.model_cfg["enc_max_seq_len"], dtype=torch.float32)
-            enc_attention_mask[0][: enc_lengths[batch_idx] + 1] = 1
-            padding_need = self.rbln_config.model_cfg["enc_max_seq_len"] - input_ids.shape[-1]
-            input_ids = torch.nn.functional.pad(input_ids, (0, padding_need))
-            _ = self.encoder(input_ids, enc_attention_mask, batch_idx=batch_idx.to(torch.int32))
-            logits = torch.zeros(1, 1, self.config.vocab_size + 100)
-            logits[0][0][-1] = 1
-        # decoder
-        else:
-            input_ids[input_ids == (self.config.vocab_size + 99)] = self.config.decoder_start_token_id
-            cache_position[cache_position != 0] = cache_position[cache_position != 0] - 2
-
-            enc_attention_mask = torch.zeros(
-                self.rbln_config.model_cfg["batch_size"],
-                self.rbln_config.model_cfg["enc_max_seq_len"],
-                dtype=torch.float32,
-            )
-            dec_attention_mask = torch.zeros(
-                self.rbln_config.model_cfg["batch_size"],
-                self.rbln_config.model_cfg["dec_max_seq_len"],
-                dtype=torch.float32,
-            )
-            for batch_idx in range(self.rbln_config.model_cfg["batch_size"]):
-                enc_attention_mask[batch_idx, : enc_lengths[batch_idx] + 1] = 1
-
-            logits = self._forward_decoder(
-                attention_mask=enc_attention_mask,
-                decoder_input_ids=input_ids,
-                decoder_attention_mask=dec_attention_mask,
-                cache_position=cache_position,
-            ).logits
-
-        return Seq2SeqLMOutput(logits=logits)
-
     def _prepare_encoder_decoder_kwargs_for_generation(
         self,
         inputs_tensor: torch.Tensor,
@@ -426,15 +382,14 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         )
 
         # 3. make sure that encoder returns `ModelOutput`
-        model_input_name = model_input_name if model_input_name is not None else self.main_input_name
         encoder_kwargs["return_dict"] = True
         encoder_kwargs["output_hidden_states"] = False
         encoder_kwargs["output_attentions"] = False
 
         for b in range(batch_size):
-            batch_idx = torch.tensor(b, dtype=torch.int32)
+            batch_position = torch.tensor(b, dtype=torch.int16)
             encoder_kwargs["input_ids"] = inputs_tensor[b].unsqueeze(0)
             encoder_kwargs["attention_mask"] = model_kwargs["attention_mask"][b].unsqueeze(0).to(torch.float32)
-            model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, batch_idx=batch_idx)
+            model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, batch_position=batch_position)
 
         return model_kwargs