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

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

@@ -22,20 +22,23 @@
 # from Rebellions Inc.
 
 import math
-from typing import Dict, Optional, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-)
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_4
 
 from ....utils import logging
-from ...cache_utils import RebelDynamicCache
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
 
 
+if is_torch_greater_or_equal_than_2_4:
+    register_fake = torch.library.register_fake
+else:
+    register_fake = torch.library.impl_abstract
+
+
 logger = logging.get_logger(__name__)
 """
 ##############################################################################
@@ -83,7 +86,7 @@ def flash_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, partition):
     return attn_output, torch.empty_like(kcache), torch.empty_like(vcache)
 
 
-@torch.library.impl_abstract("rbln_custom_ops::flash_attn_decode")
+@register_fake("rbln_custom_ops::flash_attn_decode")
 def flash_attn_decode_abstract(q, k, v, m, kcache, vcache, seq, partition):
     return torch.empty_like(q), torch.empty_like(kcache), torch.empty_like(vcache)
 
@@ -129,7 +132,7 @@ def flash_attn_prefill_cpu(q, k, v, mask, kcache, vcache, batch, seq, partition)
     return attn_output, torch.empty_like(kcache), torch.empty_like(vcache)
 
 
-@torch.library.impl_abstract("rbln_custom_ops::flash_attn_prefill")
+@register_fake("rbln_custom_ops::flash_attn_prefill")
 def flash_attn_prefill_abstract(q, k, v, m, kcache, vcache, batch, seq, partition):
     return torch.empty_like(q), torch.empty_like(kcache), torch.empty_like(vcache)
 
@@ -144,501 +147,599 @@ def rbln_cache_update_cpu(cache, value, batch, seq):
     return updated_cache
 
 
-@torch.library.impl_abstract("rbln_custom_ops::rbln_cache_update")
+@register_fake("rbln_custom_ops::rbln_cache_update")
 def rbln_cache_update_abstract(cache, value, batch, seq):
     return torch.empty_like(cache)
 
 
-class DecoderOnlyAttention:
-    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)
+class DecoderOnlyWrapper(nn.Module):
+    """A wrapper class for decoder-only language models that handles RBLN-specific optimizations and requirements.
 
-        query_state = query_state.view(
-            1,
-            self.num_key_value_heads,
-            self.num_heads // self.num_key_value_heads,
-            -1,
-            self.head_dim,
-        )
+    This wrapper is designed to:
+    1. Convert Huggingface decoder models for RBLN compilation with static shapes
+    2. Handle input/model mapping and additional information supply (e.g., positional embeddings)
+    3. Manage different attention implementations (standard and flash attention)
+    4. Support both prefill and decode phases
 
-        key_state, value_state = past_key_value.update(
-            key_state, value_state, self.layer_idx, batch_idx, read_first_step=is_prefill
-        )
+    Notes:
+    - Wrapper must only receive positional arguments in forward() due to torch.jit.trace dependency
+    - Wrapper should not contain neural network graph operations (including memory view handling)
 
-        attn_weight = torch.matmul(query_state, key_state.transpose(3, 4)) / math.sqrt(self.head_dim)
-        attn_weight += attn_mask
-        attn_weight = nn.functional.softmax(attn_weight, dim=-1, dtype=torch.float32).to(query_state.dtype)
-        attn_output = torch.matmul(attn_weight, value_state)
+    Args:
+        causal_lm (PreTrainedModel): The Huggingface causal language model to wrap
+        max_seq_len (int): Maximum sequence length for position embeddings and cache sizes
+        use_rotary_emb (bool): Whether to use rotary position embeddings
+        kvcache_partition_len (Optional[int]): Length of KV cache partitions for flash attention.
+            If provided, uses flash attention; if None, uses standard attention
+    """
 
-        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)
+    def __init__(self, causal_lm: PreTrainedModel, max_seq_len, use_rotary_emb: bool, kvcache_partition_len=None):
+        super().__init__()
+        self.config = causal_lm.config
 
-        return attn_output, key_state, value_state
+        if use_rotary_emb:
+            self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
+        else:
+            self.rotary_emb = None
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[RebelDynamicCache] = None,
-        batch_index: Optional[torch.Tensor] = 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()
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
+        if kvcache_partition_len is not None:
+            # WORKAROUND : for passing partition length as a value to the rbln compiler.
+            # What is actually used is the shape of this tensor.
+            self.attn_impl = "flash_attn"
+            logger.info(f"Using flash-attention. (partition length : {kvcache_partition_len})")
+        else:
+            self.attn_impl = "eager"
+        self.kvcache_partition_len = kvcache_partition_len
 
-        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)
-
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        # 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 = DecoderOnlyAttention._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,
+        self.causal_lm = self.convert_to_rbln_causal_lm(causal_lm)
+
+        self.num_hidden_layers = getattr(self.config, "num_hidden_layers", None) or getattr(self.config, "n_layer")
+        self._phase = "prefill"
+
+    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):
+        new_layers = []
+        for layer in causal_lm.model.layers:
+            if self.attn_impl == "eager":
+                new_self_attn = DecoderOnlyAttention(layer.self_attn)
+            elif self.attn_impl == "flash_attn":
+                new_self_attn = DecoderOnlyFlashAttention(
+                    layer.self_attn, kvcache_partition_len=self.kvcache_partition_len
                 )
+            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)
+        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
+        return new_causal_lm
 
-                iterate_results["key_states"].append(key_state)
-                iterate_results["value_states"].append(value_state)
-                iterate_results["attn_output"].append(attn_output)
+    @property
+    def phase(self):
+        return self._phase
 
-            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)
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.causal_lm.phase = phase
+
+    def forward(
+        self,
+        input_ids_or_inputs_embeds,
+        attention_mask,
+        cache_position,
+        batch_position,
+        query_position,
+        *past_key_values,
+    ):
+        if input_ids_or_inputs_embeds.ndim == 2:
+            # It is input_ids
+            input_ids = input_ids_or_inputs_embeds
+            inputs_embeds = None
+        elif input_ids_or_inputs_embeds.ndim == 3:
+            # It is inputs_embeds
+            input_ids = None
+            inputs_embeds = input_ids_or_inputs_embeds
         else:
-            attn_output, key_states, value_states = DecoderOnlyAttention._attn(
-                self,
-                query_states,
-                key_states,
-                value_states,
-                attention_mask,
-                past_key_value,
-                batch_idx=batch_index,
-                is_prefill=True,
+            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)} != {self.num_hidden_layers}"
             )
 
-        attn_output = self.o_proj(attn_output)
+        seq_len = input_ids_or_inputs_embeds.shape[1]
+        if seq_len == 1:
+            self.phase = "decode"
+        else:
+            self.phase = "prefill"
+
+        # [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, present_key_values = self.causal_lm(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            batch_position=batch_position,
+            query_position=query_position,
+            past_key_values=past_key_values,
+            rotary_emb=self.rotary_emb,
+        )
+
+        # ((key, value)) * n_layer -> [key, value] * n_layer
+        _present_key_values = ()
+        for i in range(self.num_hidden_layers):
+            key_states = present_key_values[i][0]
+            value_states = present_key_values[i][1]
+            _present_key_values = _present_key_values + (key_states, value_states)
+        present_key_values = _present_key_values
 
-        if not output_attentions:
-            attn_weight = None
+        # batch_position + query_position is dummy output node to keep the number of outputs
+        return logit, present_key_values, batch_position + query_position
 
-        return attn_output, attn_weight, key_states, value_states
 
+class DecoderOnlyForCausalLM(nn.Module):
+    """A specialized wrapper for Causal Language Models optimized for RBLN compilation.
 
-class DecoderOnlyFlashAttention:
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[RebelDynamicCache] = None,
-        batch_index: Optional[torch.Tensor] = None,
-        output_attentions: bool = False,
-        cos: Optional[torch.Tensor] = None,
-        sin: Optional[torch.Tensor] = None,
-        cache_pos_for_partitions: Optional[torch.Tensor] = None,
-        kvcache_partition_size: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        bsz, q_len, _ = hidden_states.size()
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
+    This class adapts Huggingface's CausalLM (or similar models) for RBLN deployment by:
+    1. Managing model phases (prefill/decode) throughout the computation graph
+    2. Handling output shape alignments for static compilation
+    3. Coordinating between the original model and RBLN-optimized components
 
-        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)
-
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        # Decoder (bsz > 1)
-        if bsz > 1:
-            all_key_states = []
-            all_value_states = []
-            all_attn_output = []
-
-            for b in range(bsz):
-                query_state = query_states[b].unsqueeze(0)
-                attn_mask = attention_mask[b].unsqueeze(0)
-                key_state = key_states[b].unsqueeze(0)
-                value_state = value_states[b].unsqueeze(0)
-
-                # 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,
-                    q_len,
-                    self.head_dim,
-                )
+    The class serves as an intermediate layer between DecoderOnlyWrapper and the core model,
+    focusing on maintaining correct model behavior while enabling RBLN-specific optimizations.
 
-                # RBLN custom flash attention(decode), dummy batch index
-                sidx = cache_pos_for_partitions[b][0]
-                attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_attn_decode(
-                    query_state,
-                    key_state,
-                    value_state,
-                    attn_mask,
-                    past_key_value.key_cache[self.layer_idx].unsqueeze(2),
-                    past_key_value.value_cache[self.layer_idx].unsqueeze(2),
-                    sidx,
-                    kvcache_partition_size,
-                )
+    Args:
+        causal_lm (PreTrainedModel): Original Huggingface causal language model
+        model (DecoderOnlyModel): RBLN-optimized model instance
 
-                # reshape for removing repeat_kv
-                attn_output = attn_output.view(1, self.num_heads, q_len, self.head_dim)
-                attn_output = attn_output.transpose(1, 2).contiguous()
-                attn_output = attn_output.reshape(1, q_len, self.num_heads * self.head_dim)
+    Attributes:
+        config: Configuration from the original causal language model
+        _original_mod: Reference to the original model for components like lm_head
+        model: RBLN-optimized decoder model instance
+        _phase: Current processing phase ("prefill" or "decode")
+    """
 
-                all_key_states.append(key_state)
-                all_value_states.append(value_state)
-                all_attn_output.append(attn_output)
+    def __init__(self, causal_lm: PreTrainedModel, model):
+        super().__init__()
+        self.config = causal_lm.config
+        self._original_mod = causal_lm
+        self.model = model
+        self._phase = "prefill"
 
-            key_states = torch.cat(all_key_states, dim=0)
-            value_states = torch.cat(all_value_states, dim=0)
-            attn_output = torch.cat(all_attn_output, dim=0)
+    @property
+    def phase(self):
+        return self._phase
 
-        else:
-            # reshape for removing repeat_kv
-            key_states = key_states.unsqueeze(2)
-            value_states = value_states.unsqueeze(2)
-            attention_mask = attention_mask.unsqueeze(2)
-            query_states = query_states.view(
-                1,
-                self.num_key_value_heads,
-                self.num_heads // self.num_key_value_heads,
-                q_len,
-                self.head_dim,
-            )
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.model.phase = phase
 
-            assert batch_index.dim() == 0
-            assert not output_attentions
-            bidx = batch_index
-            sidx = cache_pos_for_partitions[0][0]
-            attn_output, key_states, value_states = torch.ops.rbln_custom_ops.flash_attn_prefill(
-                query_states,
-                key_states,
-                value_states,
-                attention_mask,
-                past_key_value.key_cache[self.layer_idx].unsqueeze(2),
-                past_key_value.value_cache[self.layer_idx].unsqueeze(2),
-                bidx,
-                sidx,
-                kvcache_partition_size,
-            )
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        inputs_embeds: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        cache_position: torch.Tensor = None,
+        batch_position: torch.Tensor = None,
+        query_position: torch.Tensor = None,
+        past_key_values: Tuple[Tuple[torch.Tensor]] = None,
+        rotary_emb: nn.Module = None,
+    ):
+        # outputs
+        hidden_states, present_key_values = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            batch_position=batch_position,
+            past_key_values=past_key_values,
+            rotary_emb=rotary_emb,
+        )
 
-            # reshape for removing repeat_kv
-            attn_output = attn_output.view(1, self.num_heads, q_len, self.head_dim)
-            attn_output = attn_output.transpose(1, 2).contiguous()
-            attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim)
+        if self.phase == "prefill":
+            hidden_states = hidden_states[:, query_position.to(torch.int).unsqueeze(0)]
 
-        attn_output = self.o_proj(attn_output)
+        logits = self._original_mod.lm_head(hidden_states)
+        output = (logits, present_key_values)
+        return output
 
-        if not output_attentions:
-            attn_weight = None
 
-        return attn_output, attn_weight, key_states, value_states
+class DecoderOnlyModel(nn.Module):
+    """A modified decoder-only model implementation optimized for RBLN compilation.
 
+    Args:
+        model: Original Huggingface model to adapt
+        layers (List[DecoderOnlyLayer]): Modified transformer layers optimized for RBLN
 
-DECODERONLY_ATTENTION_CLASSES = {
-    "eager": DecoderOnlyAttention,
-    "flash_attn_rbln": DecoderOnlyFlashAttention,
-    # "sdpa": DecoderOnlySdpaAttention,
-}
+    Attributes:
+        _original_mod: Reference to original Huggingface model
+        layers: ModuleList of RBLN-optimized transformer layers
+        _phase: Current processing phase ("prefill" or "decode")
+    """
 
+    mask_fmin = torch.finfo(torch.float16).min
 
-class DecoderOnlyWrapper(torch.nn.Module):
-    def __init__(self, model, max_seq_len, kvcache_partition_len=None):
+    def __init__(self, model, layers: List["DecoderOnlyLayer"]):
         super().__init__()
-        self.config = model.config
-        self.model = model.model
-        self.lm_head = model.lm_head
-        self.max_seq_len = max_seq_len
-        self.rotary_emb = RotaryEmbedding(config=self.config, max_seq_len_cached=max_seq_len)
+        self._original_mod = model
+        self.layers = nn.ModuleList(layers)
+        self._phase = "prefill"
 
-        if kvcache_partition_len is not None:
-            # WORKAROUND : for passing partition length as a value to the rbln compiler.
-            # What is actually used is the shape of this tensor.
-            self.kvcache_partition_size = torch.zeros(kvcache_partition_len, dtype=torch.int32)
-            self.attn_implementation = "flash_attn_rbln"
-            logger.info(f"Using rbln-flash-attention. (partition length : {kvcache_partition_len})")
-        else:
-            self.kvcache_partition_size = None
-            self.attn_implementation = "eager"
+    @property
+    def phase(self):
+        return self._phase
+
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        for layer in self.layers:
+            layer.phase = phase
 
-    def get_forward_dict(self):
-        forward_dict = {
-            "wrapper": DecoderOnlyModel.forward,
-            "model": DecoderOnlyDecoderLayer.forward,
-            "decoder_layer": DECODERONLY_ATTENTION_CLASSES[self.attn_implementation].forward,
-        }
-        return forward_dict
+    @property
+    def hidden_multiplier(self):
+        return 1
+
+    def get_last_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.norm
+
+    def get_embedding(self) -> nn.Embedding:
+        return self._original_mod.embed_tokens
+
+    def get_pos_embedding(self) -> nn.Embedding:
+        raise NotImplementedError(
+            "The 'get_pos_embedding' method is not implemented. Please define this method in a subclass."
+        )
 
     def forward(
         self,
-        input_ids_or_inputs_embeds,
-        attention_mask,
-        cache_position,
-        batch_position,
-        query_idx,
-        *past_key_values,
+        input_ids: torch.Tensor = None,
+        inputs_embeds: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        cache_position: torch.Tensor = None,
+        batch_position: torch.Tensor = None,
+        past_key_values: Tuple[Tuple[torch.Tensor]] = None,
+        rotary_emb: nn.Module = None,
     ):
-        if input_ids_or_inputs_embeds.ndim == 2:
-            # input_ids
-            input_ids = input_ids_or_inputs_embeds
-            inputs_embeds = None
-        elif input_ids_or_inputs_embeds.ndim == 3:
-            # inputs_embeds
-            input_ids = None
-            inputs_embeds = input_ids_or_inputs_embeds
-        else:
-            raise NotImplementedError(f"Unknown ndim of input : {input_ids_or_inputs_embeds.ndim}")
+        # 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"
+            )
 
-        # Formatting list of past_kv to DynamicCache class.
-        past_key_values = RebelDynamicCache.from_input_format(
-            cache_position,
-            self.config.num_hidden_layers,
-            *past_key_values,
-        )
+        # embed positions
+        if inputs_embeds is None:
+            inputs_embeds = self.get_embedding()(input_ids)
 
-        batch_size = input_ids_or_inputs_embeds.size()[0]
-        seq_len = input_ids_or_inputs_embeds.size()[1]
-
-        if self.attn_implementation == "eager":
-            cache_pos_for_partitions = None
-        elif self.attn_implementation == "flash_attn_rbln":
-            p_len = self.kvcache_partition_size.size()[0]
-            num_partition = self.max_seq_len // p_len
-            if self.max_seq_len % p_len > 0:
-                raise ValueError(
-                    f"The partition length({p_len}) must be exactly divisible by the max_seq_len({self.max_seq_len})."
-                )
-            cache_pos_for_partitions = torch.zeros((batch_size, num_partition), dtype=torch.int32)
+        hidden_states = inputs_embeds * self.hidden_multiplier
+        attention_mask = (1 - attention_mask) * self.mask_fmin
 
-            if batch_size > 1:  # decode
-                for b_idx in range(batch_size):
-                    decoding_step = cache_position[b_idx]
-                    cache_pos = decoding_step
-                    for p_idx in range(num_partition):
-                        input_0 = torch.tensor(cache_pos - p_len * p_idx, dtype=torch.int32)
-                        input_1 = torch.tensor(p_len, dtype=torch.int32)
-                        min = torch.minimum(input_0, input_1)
-                        cache_pos_for_partition = torch.maximum(min, torch.tensor(0, dtype=torch.int32))
-                        cache_pos_for_partitions[b_idx][p_idx] = cache_pos_for_partition
-            else:  # prefill
-                cache_pos = cache_position[0][0]
-                for p_idx in range(num_partition):
-                    input_0 = torch.tensor(cache_pos - p_len * p_idx, dtype=torch.int32)
-                    input_1 = torch.tensor(p_len, dtype=torch.int32)
-                    min = torch.minimum(input_0, input_1)
-                    cache_pos_for_partition = torch.maximum(min, torch.tensor(0, dtype=torch.int32))
-                    cache_pos_for_partitions[0][p_idx] = cache_pos_for_partition
+        # 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 = slice_and_unsqueeze_cos_sin(cos, sin, cache_position)
         else:
-            raise NotImplementedError(f"Unknown attn_implementation: {self.attn_implementation}")
+            batch_size = inputs_embeds.shape[0]
+            if cache_position.shape[0] > 1:
+                position_embeds = []
+                for b_idx in range(batch_size):
+                    position_embed = self.get_pos_embedding()(cache_position[b_idx])
+                    position_embeds.append(position_embed)
+
+                position_embeds = torch.cat(position_embeds, dim=0).unsqueeze(1)
+            else:
+                position_embeds = self.get_pos_embedding()(cache_position)
+            hidden_states = hidden_states + position_embeds
+            cos, sin = None, None
+
+        # (batch, seq_len) -> (batch,)
+        current_steps = cache_position[:, 0]
+
+        present_key_values = past_key_values
+        for layer in self.layers:
+            hidden_states, present_key_values = layer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                current_steps=current_steps,
+                batch_position=batch_position,
+                past_key_values=present_key_values,
+                cos=cos,
+                sin=sin,
+            )
 
-        forward_dict = self.get_forward_dict()
-        outputs = forward_dict["wrapper"](
-            self.model,
-            input_ids=input_ids,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            position_ids=cache_position,
-            past_key_values=past_key_values,
-            batch_ids=batch_position,
-            rotary_pos_emb=self.rotary_emb,
-            cache_pos_for_partitions=cache_pos_for_partitions,
-            kvcache_partition_size=self.kvcache_partition_size,
-            forward_dict=forward_dict,
-        )
+        hidden_states = self.get_last_layernorm()(hidden_states)
+        return hidden_states, present_key_values
+
+
+class DecoderOnlyLayer(nn.Module):
+    """A single transformer layer adapted for RBLN compilation with static shapes.
+
+    This layer implements a modified transformer block that includes:
+    1. Self-attention mechanism (either standard or flash attention)
+    2. Feed-forward network (FFN)
+    3. Layer normalization
+    4. Residual connections
 
-        hidden_states = outputs[0]
-        if seq_len != 1:
-            hidden_states = hidden_states[:, query_idx.to(torch.int).unsqueeze(0)]
+    The layer is specifically designed to:
+    - Support compilation to RBLN custom ops
+    - Maintain static tensor shapes throughout computations
+    - Handle both prefill and decode phases efficiently
+    - Manage attention state transitions properly
 
-        logits = self.lm_head(hidden_states)
+    Args:
+        layer: Original transformer layer module to wrap
+        self_attn (DecoderOnlyAttention): Modified attention module optimized for RBLN
 
-        output = (logits,) + outputs[1:]
+    Attributes:
+        _original_mod: Reference to original layer for accessing components
+        self_attn: Modified attention mechanism mapped to RBLN ops at compile time
+        phase: Current operation phase ("prefill" or "decode")
+    """
+
+    def __init__(self, layer, self_attn: "DecoderOnlyAttention"):
+        super().__init__()
+        self._original_mod = layer
+        self.self_attn = self_attn
+        self._phase = "prefill"
+
+    @property
+    def phase(self):
+        return self._phase
 
-        return output, batch_position + query_idx
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.self_attn.phase = phase
 
+    def get_pre_attention_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.input_layernorm
+
+    def get_post_attention_layernorm(self) -> nn.LayerNorm:
+        return self._original_mod.post_attention_layernorm
 
-class DecoderOnlyDecoderLayer:
     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.Tensor] = None,
+        attention_mask: torch.Tensor,
+        current_steps: torch.LongTensor,
+        batch_position: torch.Tensor,
+        past_key_values: Tuple[Tuple[torch.Tensor]],
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
-        cache_pos_for_partitions: Optional[torch.Tensor] = None,
-        kvcache_partition_size: Optional[torch.Tensor] = None,
-        forward_dict: Optional[Dict[str, classmethod]] = None,
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+    ):
         residual = hidden_states
 
-        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.get_pre_attention_layernorm()(hidden_states)
 
-        hidden_states, self_attn_weight, k, v = forward_dict["decoder_layer"](
-            self.self_attn,
+        hidden_states, 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,
+            current_steps=current_steps,
+            batch_position=batch_position,
+            past_key_values=past_key_values,
             cos=cos,
             sin=sin,
-            cache_pos_for_partitions=cache_pos_for_partitions,
-            kvcache_partition_size=kvcache_partition_size,
-            **kwargs,
         )
-        past_key_value.assign(k, v, layer_idx)
-
         hidden_states = residual + hidden_states
 
         # Fully Connected
         residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.get_post_attention_layernorm()(hidden_states)
+        hidden_states = self._original_mod.mlp(hidden_states)
         hidden_states = residual + hidden_states
 
-        outputs = (hidden_states,)
+        return hidden_states, present_key_values
 
-        if output_attentions:
-            outputs += (self_attn_weight,)
 
-        if use_cache:
-            outputs += (past_key_value,)
+class DecoderOnlyAttention(nn.Module):
+    """Attention implementation for decoder-only models optimized for RBLN compilation.
 
-        return outputs
+    This class implements a modified version of the standard attention mechanism that:
+    1. Supports static shape requirements for RBLN compilation
+    2. Handles explicit batch and position management
 
+    Args:
+        self_attn: Original attention module from the base model
+    """
 
-class DecoderOnlyModel:
-    def forward(
+    def __init__(self, self_attn):
+        super().__init__()
+        self._original_mod = self_attn
+        self.layer_idx = self_attn.layer_idx
+        self.num_heads = self._original_mod.num_heads
+        self.head_dim = self._original_mod.head_dim
+        self.phase = "prefill"
+        self.__post_init__()
+
+    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.o_proj
+        self.num_key_value_heads = self._original_mod.num_key_value_heads
+
+    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Projects input hidden states into query, key, and value representations.
+
+        Args:
+            hidden_states: Input tensor of shape [batch_size, seq_len, hidden_dim]
+
+        Returns:
+            Tuple of (query_states, key_states, value_states)
+        """
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        return query_states, key_states, value_states
+
+    def apply_rotary_pos_embed(self, query_states, key_states, cos, sin):
+        return apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+    def rbln_attention(
         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.Tensor] = 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"
-            )
+        query_state,
+        key_state,
+        value_state,
+        attn_mask,
+        batch_idx,
+        past_key_state,
+        past_value_state,
+        current_step,
+        # below are designed for Midm, GPT which requires to support scaling for attention weights
+        # TODO(jongho): Merge and manage scales generally
+        layer_idx=None,
+        scale_attn_weights: bool = None,
+        scale_attn_by_inverse_layer_idx: bool = None,
+        scale_qk_by_inverse_layer_idx: bool = None,
+    ):
+        """Compute attention with static shapes and explicit cache management.
+
+        Args:
+            query_state: Query tensor [1, num_heads, 1, head_dim]
+            key_state: Key tensor [1, num_heads, seq_len, head_dim]
+            value_state: Value tensor [1, num_heads, seq_len, head_dim]
+            attn_mask: Attention mask tensor
+            batch_idx: Batch index for cache lookup
+            past_key_state: Previous key cache states
+            past_value_state: Previous value cache states
+            current_step: Current position in sequence
+
+        Returns:
+            Tuple of (attention_output, key_state, value_state)
+        """
+        # Implementation details.
+        # 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)
 
-        # 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,
-            )
+        query_state = query_state.view(
+            1,
+            self.num_key_value_heads,
+            self.num_heads // self.num_key_value_heads,
+            -1,  # seq len
+            self.head_dim,
+        )  #
+
+        kend = current_step + key_state.shape[-2]
+        vend = current_step + value_state.shape[-2]
+
+        key_state = (
+            past_key_state[batch_idx]
+            .unsqueeze(0)
+            .unsqueeze(2)
+            .slice_scatter(key_state, dim=-2, start=current_step, end=kend)
+        )
+        value_state = (
+            past_value_state[batch_idx]
+            .unsqueeze(0)
+            .unsqueeze(2)
+            .slice_scatter(value_state, dim=-2, start=current_step, end=vend)
+        )
+
+        attn_weight = torch.matmul(query_state, key_state.transpose(3, 4))
+        attn_weight = attn_weight / math.sqrt(self.head_dim)
+
+        if layer_idx is not None and (scale_attn_by_inverse_layer_idx or scale_qk_by_inverse_layer_idx):
+            attn_weight = attn_weight / float(layer_idx + 1)
+
+        attn_weight += attn_mask
 
-            hidden_states = layer_outputs[0]
+        if layer_idx is not None and scale_qk_by_inverse_layer_idx:
+            attn_weight = attn_weight * float(layer_idx + 1)
 
-            updated_cache = layer_outputs[2 if output_attentions else 1]
+        attn_weight = nn.functional.softmax(attn_weight, dim=-1)
 
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
+        attn_output = torch.matmul(attn_weight, value_state)
 
-        hidden_states = self.norm(hidden_states)
+        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)
 
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
+        return attn_output, key_state, value_state
 
-        # convert RebelDynamicCache to legacy Tuple[Tuple[torch.Tensor]]
-        next_cache = updated_cache.to_legacy_cache()
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        current_steps: torch.LongTensor,
+        batch_position: torch.Tensor,
+        past_key_values: Tuple[Tuple[torch.Tensor]],
+        cos: Optional[torch.Tensor] = None,  # (batch, 1, prefill_size, head_dim)
+        sin: Optional[torch.Tensor] = None,
+    ):
+        batch_size, query_length, _ = hidden_states.size()
 
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
+        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
         )
+        # b, num_head, query, head_dim
+
+        if cos is not None and sin is not None:
+            query_states, key_states = self.apply_rotary_pos_embed(query_states, key_states, cos, sin)
+
+        if batch_size > 1 and self.phase == "prefill":
+            raise NotImplementedError(f"batch size should be 1 if prefill phase, but got {batch_size}.")
+
+        _key_states = []
+        _value_states = []
+        _attn_outputs = []
+        for b in range(batch_size):
+            current_step = current_steps[b]
+            attn_output, key_state, value_state = self.rbln_attention(
+                query_states[b].unsqueeze(0),
+                key_states[b].unsqueeze(0),
+                value_states[b].unsqueeze(0),
+                attention_mask[b].unsqueeze(0)
+                if self.phase == "decode"
+                else attention_mask,  # TODO(jongho): fix when msoftmax is supported
+                past_key_state=past_key_values[self.layer_idx][0],
+                past_value_state=past_key_values[self.layer_idx][1],
+                batch_idx=b if self.phase == "decode" else batch_position,
+                current_step=current_step,
+            )
+            _key_states.append(key_state)
+            _value_states.append(value_state)
+            _attn_outputs.append(attn_output)
+        key_states = torch.cat(_key_states, dim=0)
+        value_states = torch.cat(_value_states, dim=0)
+        attn_outputs = torch.cat(_attn_outputs, dim=0)
 
+        attn_outputs = self.o_proj(attn_outputs)
+        past_key_values[self.layer_idx] = key_states, value_states
+        return attn_outputs, past_key_values
 
-def slice_and_unsqueeze_cos_sin(cos, sin, position_ids, unsqueeze_dim=1):
-    """Slice cos[position_ids], sin[position_ids] vector for the query."""
-    if position_ids.shape[0] > 1:
+
+def slice_and_unsqueeze_cos_sin(cos, sin, cache_position, unsqueeze_dim=1):
+    """Slice cos[cache_position], sin[cache_position] vector for the query."""
+    if cache_position.shape[0] > 1:
         cos_all = []
         sin_all = []
-        for i in range(position_ids.shape[0]):
-            cos_all.append(cos[position_ids[i : i + 1]].unsqueeze(unsqueeze_dim))
-            sin_all.append(sin[position_ids[i : i + 1]].unsqueeze(unsqueeze_dim))
+        for i in range(cache_position.shape[0]):
+            cos_all.append(cos[cache_position[i : i + 1]].unsqueeze(unsqueeze_dim))
+            sin_all.append(sin[cache_position[i : i + 1]].unsqueeze(unsqueeze_dim))
         cos = torch.cat(cos_all, dim=0)
         sin = torch.cat(sin_all, dim=0)
     else:
-        cos = cos[position_ids].unsqueeze(unsqueeze_dim)
-        sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+        cos = cos[cache_position].unsqueeze(unsqueeze_dim)
+        sin = sin[cache_position].unsqueeze(unsqueeze_dim)
 
     return cos, sin
 
@@ -658,6 +759,26 @@ def apply_rotary_pos_emb(q, k, cos, sin):
     return q_embed, k_embed
 
 
+def apply_rotary_pos_emb_partial(query_states, key_states, cos, sin, ndim) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Partial rotary embedding
+    query_rot, query_pass = (
+        query_states[..., :ndim],
+        query_states[..., ndim:],
+    )
+    key_rot, key_pass = (
+        key_states[..., :ndim],
+        key_states[..., ndim:],
+    )
+
+    # [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)
+
+    # [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)
+    return query_states, key_states
+
+
 class RotaryEmbedding(nn.Module):
     """RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
 
@@ -674,14 +795,14 @@ class RotaryEmbedding(nn.Module):
             rope_type = "default"
 
         inv_freq, attention_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, max_seq_len_cached)
-        position_ids = torch.arange(0, max_seq_len_cached, dtype=torch.float32)
-        position_ids_expanded = position_ids[:, None]
+        cache_position = torch.arange(0, max_seq_len_cached, dtype=torch.float32)
+        cache_position_expanded = cache_position[:, None]
 
         if rope_type == "dynamic":
-            freqs = position_ids_expanded.float() * inv_freq.float()
+            freqs = cache_position_expanded.float() * inv_freq.float()
         else:
             inv_freq_expanded = inv_freq[None, :]
-            freqs = position_ids_expanded.float() @ inv_freq_expanded.float()
+            freqs = cache_position_expanded.float() @ inv_freq_expanded.float()
 
         emb = torch.cat((freqs, freqs), dim=-1)
 
@@ -696,3 +817,142 @@ class RotaryEmbedding(nn.Module):
             self._cos_cached[:seq_len].to(dtype=x.dtype),
             self._sin_cached[:seq_len].to(dtype=x.dtype),
         )
+
+
+class DecoderOnlyFlashAttention(DecoderOnlyAttention):
+    def __init__(self, self_attn, kvcache_partition_len):
+        super().__init__(self_attn=self_attn)
+        self.kvcache_partition_size = torch.zeros(kvcache_partition_len, dtype=torch.int32)
+
+    def get_cache_pos_for_partitions(self, current_steps, batch_size, max_seq_len):
+        partition_len = self.kvcache_partition_size.size()[0]
+        num_partition = max_seq_len // partition_len
+        cache_pos_for_partitions = torch.zeros((batch_size, num_partition), dtype=torch.int32)
+        if self.phase == "decode":
+            for b_idx in range(batch_size):
+                cache_pos = current_steps[b_idx]
+                for p_idx in range(num_partition):
+                    cache_pos_for_partitions[b_idx][p_idx] = torch.clamp(
+                        cache_pos - partition_len * p_idx, 0, partition_len
+                    )
+        else:  # prefill
+            cache_pos = current_steps[0]
+            for p_idx in range(num_partition):
+                cache_pos_for_partitions[0][p_idx] = torch.clamp(cache_pos - partition_len * p_idx, 0, partition_len)
+
+        return cache_pos_for_partitions
+
+    def rbln_flash_attention(
+        self,
+        query_state,
+        key_state,
+        value_state,
+        attn_mask,
+        batch_idx,
+        past_key_state,
+        past_value_state,
+        cache_pos_for_partitions,
+    ):
+        # 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)
+
+        query_state = query_state.view(
+            1,
+            self.num_key_value_heads,
+            self.num_heads // self.num_key_value_heads,
+            -1,  # seq len
+            self.head_dim,
+        )
+
+        # RBLN custom flash attention(decode), dummy batch index
+        if self.phase == "decode":
+            sidx = cache_pos_for_partitions[batch_idx][0]
+            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),
+                sidx,
+                self.kvcache_partition_size,
+            )
+        else:
+            sidx = cache_pos_for_partitions[0][0]
+            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_idx,
+                sidx,
+                self.kvcache_partition_size,
+            )
+
+        # 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)
+
+        return attn_output, key_state, value_state
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        current_steps: torch.LongTensor,
+        batch_position: torch.Tensor,
+        past_key_values: Tuple[Tuple[torch.Tensor]],
+        cos: Optional[torch.Tensor] = None,
+        sin: 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
+        )
+        # b, num_head, query, head_dim
+
+        max_seq_len = past_key_values[self.layer_idx][0].shape[-2]
+
+        if cos is not None and sin is not None:
+            query_states, key_states = self.apply_rotary_pos_embed(query_states, key_states, cos, sin)
+
+        cache_pos_for_partitions = self.get_cache_pos_for_partitions(
+            current_steps, batch_size=batch_size, max_seq_len=max_seq_len
+        )  # batch_size, num_partitions
+
+        _key_states = []
+        _value_states = []
+        _attn_outputs = []
+        for b in range(batch_size):
+            attn_output, key_state, value_state = self.rbln_flash_attention(
+                query_states[b].unsqueeze(0),
+                key_states[b].unsqueeze(0),
+                value_states[b].unsqueeze(0),
+                attention_mask[b].unsqueeze(0)
+                if self.phase == "decode"
+                else attention_mask,  # TODO(jongho): fix when msoftmax is supported
+                past_key_state=past_key_values[self.layer_idx][0],
+                past_value_state=past_key_values[self.layer_idx][1],
+                batch_idx=b if self.phase == "decode" else batch_position,
+                cache_pos_for_partitions=cache_pos_for_partitions,
+            )
+            _key_states.append(key_state)
+            _value_states.append(value_state)
+            _attn_outputs.append(attn_output)
+        key_states = torch.cat(_key_states, dim=0)
+        value_states = torch.cat(_value_states, dim=0)
+        attn_outputs = torch.cat(_attn_outputs, dim=0)
+
+        attn_outputs = self.o_proj(attn_outputs)
+        past_key_values[self.layer_idx] = key_states, value_states
+        return attn_outputs, past_key_values