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

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

@@ -22,623 +22,740 @@
 # 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 ....ops import register_rbln_custom_attention, register_rbln_custom_flash_attention
 from ....utils import logging
-from ...cache_utils import RebelDynamicCache
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
 
 
 logger = logging.get_logger(__name__)
-"""
-##############################################################################
-# RBLN custom operation (python interface)
-# torch.compile custom operation
-# torch.library.define - kernel declaration
-# torch.library.impl - kernel implementation
-# torch.library.impl_abstract - symbolic trace
-##############################################################################
-"""
-
-# RBLN custom op(flash attention decode)
-torch.library.define(
-    "rbln_custom_ops::flash_attn_decode",
-    "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d) -> Tensor[]",
-)
-
-
-@torch.library.impl("rbln_custom_ops::flash_attn_decode", "cpu")
-def flash_attn_decode_cpu(q, k, v, mask, kcache, vcache, seq, partition):
-    """
-    WORKAROUND:
-    Partition is declared as an argument to the function, even though it is
-    not actually used in the CPU implementation, this allows the rbln compiler
-    to perform flash attention operations with partition as an argument.
-    """
-    assert kcache.dim() == k.dim()
-    assert vcache.dim() == v.dim()
-    assert k.size(-2) == v.size(-2)
-    assert partition.dim() == 1
-    b = 0
-    if seq.dim() == 1:
-        s = seq[0]
-    elif seq.dim() == 0:
-        s = seq
-    else:
-        assert False
-    e = s + k.size(-2)
-    updated_k = kcache[b].unsqueeze(0).slice_scatter(k, dim=-2, start=s, end=e)
-    updated_v = vcache[b].unsqueeze(0).slice_scatter(v, dim=-2, start=s, end=e)
-    attn_weight = torch.matmul(q, updated_k.transpose(3, 4)) / math.sqrt(128)
-    attn_weight = attn_weight + mask
-    attn_weight = nn.functional.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype)
-    attn_output = torch.matmul(attn_weight, updated_v)
-    return attn_output, torch.empty_like(kcache), torch.empty_like(vcache)
 
+DEFAULT_FLASH_ATTN_PARTITION_LENGTH = 16_384
+DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH = 32_768
+MIN_FLASH_ATTN_MAX_SEQ_LEN = 8_192
+MIN_FLASH_ATTN_PARTITION_LENGTH = 4_096
+MAX_FLASH_ATTN_PARTITION_LENGTH = 32_768
 
-@torch.library.impl_abstract("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)
 
+def validate_attention_method(
+    rbln_attn_impl: str, rbln_kvcache_partition_len: int, rbln_max_seq_len: int
+) -> Tuple[str, int]:
+    if rbln_kvcache_partition_len is not None:
+        if rbln_attn_impl == "eager":
+            raise ValueError(
+                f"`rbln_kvcache_partition_len` is set to {rbln_kvcache_partition_len}, but KV cache partitioning"
+                " is not supported with 'eager' attention. Please set `rbln_kvcache_partition_len` to None, "
+                "or switch `rbln_attn_impl` to 'flash_attn' to use KV cache partitioning."
+            )
+        elif rbln_attn_impl is None:
+            rbln_attn_impl = "flash_attn"
+            logger.warning(
+                "A non-null `rbln_kvcache_partition_len` was provided, but `rbln_attn_impl` was not explicitly set. "
+                "Since KV cache partitioning is only supported with flash attention, "
+                "`rbln_attn_impl` has been automatically switched to 'flash_attn'."
+            )
 
-# RBLN custom op(flash attention prefill)
-torch.library.define(
-    "rbln_custom_ops::flash_attn_prefill",
-    "(Tensor x, Tensor y, Tensor z, Tensor w, Tensor a, Tensor b, Tensor c, Tensor d, Tensor e) -> Tensor[]",
-)
+    rbln_attn_impl = "eager" if rbln_attn_impl is None else rbln_attn_impl
+    if rbln_attn_impl not in ["eager", "flash_attn"]:
+        raise ValueError(f"Unknown `rbln_attn_impl` : {rbln_attn_impl}. (Available : 'eager', 'flash_attn`)")
+
+    if rbln_kvcache_partition_len is None and rbln_attn_impl == "flash_attn":
+        rbln_kvcache_partition_len = DEFAULT_FLASH_ATTN_PARTITION_LENGTH
+
+    ## Checking Constraints...
+    # Constraint of eager attention:
+    # - `max_seq_len` <= 32k
+
+    # Constraints of flash attention:
+    # 1. `max_seq_len` should be multiple of `partition_len`.
+    # 2. 4k <= `partition_len` <= 32k.
+    # 3. `max_seq_len` should be larger then 8k.
+    if rbln_attn_impl == "eager" and rbln_max_seq_len > DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH:
+        raise ValueError(
+            f"`rbln_max_seq_len` is set to {rbln_max_seq_len}, "
+            f"which exceeds the limit of {DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH} for 'eager' attention. "
+            f"Please reduce the `rbln_max_seq_len` to {DEFAULT_MAX_EAGER_ATTN_SEQUENCE_LENGTH} or lower,"
+            " or consider switching `rbln_attn_impl` to 'flash_attn' for larger sequence lengths."
+        )
 
+    if rbln_attn_impl == "flash_attn":
+        if rbln_max_seq_len // rbln_kvcache_partition_len < 2 or rbln_max_seq_len % rbln_kvcache_partition_len != 0:
+            raise ValueError(
+                f"`rbln_max_seq_len` ({rbln_max_seq_len}) must be a multiple of `rbln_kvcache_partition_len` ({rbln_kvcache_partition_len}) "
+                f"when using 'flash_attn'. Please adjust either value to meet this requirement."
+            )
+        elif not (MIN_FLASH_ATTN_PARTITION_LENGTH <= rbln_kvcache_partition_len <= MAX_FLASH_ATTN_PARTITION_LENGTH):
+            raise ValueError(
+                f"`rbln_kvcache_partition_len` ({rbln_kvcache_partition_len}) is out of the supported range for 'flash_attn' "
+                f"({MIN_FLASH_ATTN_PARTITION_LENGTH} <= `rbln_kvcache_partition_len` <= {MAX_FLASH_ATTN_PARTITION_LENGTH}). "
+                f"Please provide a valid value within this range."
+            )
+        elif rbln_max_seq_len < MIN_FLASH_ATTN_MAX_SEQ_LEN:
+            raise ValueError(
+                f"`rbln_max_seq_len` ({rbln_max_seq_len}) is too small for 'flash_attn'. The minimum "
+                f"supported value is {MIN_FLASH_ATTN_MAX_SEQ_LEN}. Please increase `rbln_max_seq_len` to meet "
+                "this requirement, or consider switching `rbln_attn_impl` to 'eager' for shorter lengths."
+            )
 
-@torch.library.impl("rbln_custom_ops::flash_attn_prefill", "cpu")
-def flash_attn_prefill_cpu(q, k, v, mask, kcache, vcache, batch, seq, partition):
-    """
-    WORKAROUND:
-    Partition is declared as an argument to the function, even though it is
-    not actually used in the CPU implementation, this allows the rbln compiler
-    to perform flash attention operations with partition as an argument.
+    return rbln_attn_impl, rbln_kvcache_partition_len
+
+
+class DecoderOnlyWrapper(nn.Module):
+    """A wrapper class for decoder-only language models that handles RBLN-specific optimizations and requirements.
+
+    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
+
+    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)
+
+    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
+        attn_impl (str): The attention implementation to use.
+            - "eager": Uses the standard attention.
+            - "flash_attn": Uses flash attention. When set,
+              the key/value cache is partitioned into chunks of length
+              `kvcache_partition_len`.
+        kvcache_partition_len (Optional[int]): Length of KV cache partitions for flash attention.
+            This is only relevant if `attn_impl` is set to "flash_attn`
     """
-    assert kcache.dim() == k.dim()
-    assert vcache.dim() == v.dim()
-    assert k.size(-2) == v.size(-2)
-    assert partition.dim() == 1
-    if batch.dim() == 1:
-        b = batch[0]
-    elif batch.dim() == 0:
-        b = batch
-    else:
-        assert False
-    if seq.dim() == 1:
-        s = seq[0]
-    elif seq.dim() == 0:
-        s = seq
-    else:
-        assert False
-    e = s + k.size(-2)
-    updated_k = kcache[b].unsqueeze(0).slice_scatter(k, dim=-2, start=s, end=e)
-    updated_v = vcache[b].unsqueeze(0).slice_scatter(v, dim=-2, start=s, end=e)
-    attn_weight = torch.matmul(q, updated_k.transpose(3, 4)) / math.sqrt(128)
-    attn_weight = attn_weight + mask
-    attn_weight = nn.functional.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype)
-    attn_output = torch.matmul(attn_weight, updated_v)
-    return attn_output, torch.empty_like(kcache), torch.empty_like(vcache)
 
+    def __init__(
+        self,
+        causal_lm: PreTrainedModel,
+        max_seq_len: int,
+        use_rotary_emb: bool,
+        attn_impl: str,
+        kvcache_partition_len: Optional[int] = None,
+    ):
+        super().__init__()
+        self.config = causal_lm.config
 
-@torch.library.impl_abstract("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)
-
+        if use_rotary_emb:
+            self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
+        else:
+            self.rotary_emb = None
+
+        self.attn_impl = attn_impl
+        if self.attn_impl == "flash_attn":
+            self.kvcache_partition_len = kvcache_partition_len or DEFAULT_FLASH_ATTN_PARTITION_LENGTH
+            register_rbln_custom_flash_attention()
+        elif self.attn_impl == "eager":
+            self.kvcache_partition_len = None
+            register_rbln_custom_attention()
+        else:
+            raise ValueError(f"Unknown attn_impl : {self.attn_impl}")
 
-# RBLN custom op(cache update)
-torch.library.define("rbln_custom_ops::rbln_cache_update", "(Tensor x, Tensor y, Tensor z, Tensor w) -> Tensor")
+        if kvcache_partition_len and kvcache_partition_len > max_seq_len:
+            raise ValueError(
+                f"kvcache_partition_len({kvcache_partition_len}) should be lower"
+                f" or equal to max_seq_len({max_seq_len})!"
+            )
 
+        self.causal_lm = self.convert_to_rbln_causal_lm(causal_lm)
 
-@torch.library.impl("rbln_custom_ops::rbln_cache_update", "cpu")
-def rbln_cache_update_cpu(cache, value, batch, seq):
-    updated_cache = cache[batch].slice_scatter(value, dim=-2, start=batch[0], end=batch[0] + seq[0])
-    return updated_cache
+        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)
 
-@torch.library.impl_abstract("rbln_custom_ops::rbln_cache_update")
-def rbln_cache_update_abstract(cache, value, batch, seq):
-    return torch.empty_like(cache)
+    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, partition_len=self.kvcache_partition_len)
+        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
+        return new_causal_lm
+
+    @property
+    def phase(self) -> str:
+        return self._phase
+
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.causal_lm.phase = phase
+
+    def forward(self, *args):
+        if self.phase == "decode":
+            (
+                input_ids_or_inputs_embeds,
+                attention_mask,
+                cache_position,
+                *past_key_values,
+            ) = args
+            batch_position = torch.tensor(0, dtype=torch.int16)
+            query_position = None
+        elif self.phase == "prefill":
+            (
+                input_ids_or_inputs_embeds,
+                attention_mask,
+                cache_position,
+                batch_position,
+                query_position,
+                *past_key_values,
+            ) = args
+        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}")
 
-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)
+        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}"
+            )
 
-        query_state = query_state.view(
-            1,
-            self.num_key_value_heads,
-            self.num_heads // self.num_key_value_heads,
-            -1,
-            self.head_dim,
+        # [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_state, value_state = past_key_value.update(
-            key_state, value_state, self.layer_idx, batch_idx, read_first_step=is_prefill
-        )
+        # ((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
 
-        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)
+        return logit, present_key_values
 
-        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
+class DecoderOnlyForCausalLM(nn.Module):
+    """A specialized wrapper for Causal Language Models optimized for RBLN compilation.
 
-    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)
+    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:
-            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,
-                )
+    The class serves as an intermediate layer between DecoderOnlyWrapper and the core model,
+    focusing on maintaining correct model behavior while enabling RBLN-specific optimizations.
 
-                iterate_results["key_states"].append(key_state)
-                iterate_results["value_states"].append(value_state)
-                iterate_results["attn_output"].append(attn_output)
+    Args:
+        causal_lm (PreTrainedModel): Original Huggingface causal language model
+        model (DecoderOnlyModel): RBLN-optimized model instance
 
-            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 = DecoderOnlyAttention._attn(
-                self,
-                query_states,
-                key_states,
-                value_states,
-                attention_mask,
-                past_key_value,
-                batch_idx=batch_index,
-                is_prefill=True,
-            )
-
-        attn_output = self.o_proj(attn_output)
+    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")
+    """
 
-        if not output_attentions:
-            attn_weight = None
+    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"
 
-        return attn_output, attn_weight, key_states, value_states
+    @property
+    def phase(self):
+        return self._phase
 
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.model.phase = phase
 
-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)
-
-        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,
-                )
-
-                # 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,
-                )
+        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(1, q_len, self.num_heads * self.head_dim)
+        if self.phase == "prefill":
+            hidden_states = hidden_states[:, query_position.to(torch.int).unsqueeze(0)]
 
-                all_key_states.append(key_state)
-                all_value_states.append(value_state)
-                all_attn_output.append(attn_output)
+        logits = self._original_mod.lm_head(hidden_states)
+        output = (logits, present_key_values)
+        return output
 
-            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)
 
-        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,
-            )
+class DecoderOnlyModel(nn.Module):
+    """A modified decoder-only model implementation optimized for RBLN compilation.
 
-            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,
-            )
+    Args:
+        model: Original Huggingface model to adapt
+        layers (List[DecoderOnlyLayer]): Modified transformer layers optimized for RBLN
 
-            # 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)
+    Attributes:
+        _original_mod: Reference to original Huggingface model
+        layers: ModuleList of RBLN-optimized transformer layers
+        _phase: Current processing phase ("prefill" or "decode")
+    """
 
-        attn_output = self.o_proj(attn_output)
+    def __init__(self, model, layers: List["DecoderOnlyLayer"], partition_len=None):
+        super().__init__()
+        self._original_mod = model
+        self.layers = nn.ModuleList(layers)
+        self._phase = "prefill"
+        self.partition_len = partition_len
+
+    @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
+
+    @property
+    def attn_impl(self) -> str:
+        return "eager" if self.partition_len is None else "flash_attn"
+
+    @property
+    def hidden_multiplier(self):
+        return 1
+
+    def convert_sequence_positions_for_flash_attn(self, seq_positions, max_seq_len):
+        if self.attn_impl != "flash_attn":
+            raise NotImplementedError(f"Unknown attn_impl ({self.attn_impl}).")
+
+        partition_len = self.partition_len
+        num_partition = max_seq_len // partition_len
+
+        cs = seq_positions.repeat(num_partition, 1).transpose(0, 1)
+        pidx = torch.arange(num_partition)
+        cache_pos_for_partitions = torch.clamp(cs - pidx * partition_len, 0, partition_len)
+        return cache_pos_for_partitions
+
+    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."
+        )
 
-        if not output_attentions:
-            attn_weight = None
+    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,
+        past_key_values: Tuple[Tuple[torch.Tensor]] = None,
+        rotary_emb: nn.Module = None,
+    ):
+        # 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"
+            )
 
-        return attn_output, attn_weight, key_states, value_states
+        # embed positions
+        if inputs_embeds is None:
+            inputs_embeds = self.get_embedding()(input_ids)
 
+        hidden_states = inputs_embeds * self.hidden_multiplier
 
-DECODERONLY_ATTENTION_CLASSES = {
-    "eager": DecoderOnlyAttention,
-    "flash_attn_rbln": DecoderOnlyFlashAttention,
-    # "sdpa": DecoderOnlySdpaAttention,
-}
+        # 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:
+            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,)
+        seq_positions = cache_position[:, 0]
+        if self.attn_impl == "flash_attn":
+            max_seq_len = past_key_values[0][0].shape[-2]
+            seq_positions = self.convert_sequence_positions_for_flash_attn(
+                seq_positions=seq_positions, max_seq_len=max_seq_len
+            )
 
+        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,
+                seq_positions=seq_positions,
+                batch_position=batch_position,
+                past_key_values=present_key_values,
+                cos=cos,
+                sin=sin,
+            )
 
-class DecoderOnlyWrapper(torch.nn.Module):
-    def __init__(self, model, max_seq_len, kvcache_partition_len=None):
-        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)
-
-        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"
+        hidden_states = self.get_last_layernorm()(hidden_states)
+        return hidden_states, present_key_values
 
-    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
 
-    def forward(
-        self,
-        input_ids_or_inputs_embeds,
-        attention_mask,
-        cache_position,
-        batch_position,
-        query_idx,
-        *past_key_values,
-    ):
-        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}")
+class DecoderOnlyLayer(nn.Module):
+    """A single transformer layer adapted for RBLN compilation with static shapes.
 
-        # 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,
-        )
+    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
 
-        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)
+    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
 
-            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
-        else:
-            raise NotImplementedError(f"Unknown attn_implementation: {self.attn_implementation}")
+    Args:
+        layer: Original transformer layer module to wrap
+        self_attn (DecoderOnlyAttention): Modified attention module optimized for RBLN
 
-        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,
-        )
+    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")
+    """
 
-        hidden_states = outputs[0]
-        if seq_len != 1:
-            hidden_states = hidden_states[:, query_idx.to(torch.int).unsqueeze(0)]
+    def __init__(self, layer, self_attn: "DecoderOnlyAttention"):
+        super().__init__()
+        self._original_mod = layer
+        self.self_attn = self_attn
+        self._phase = "prefill"
 
-        logits = self.lm_head(hidden_states)
+    @property
+    def phase(self):
+        return self._phase
 
-        output = (logits,) + outputs[1:]
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.self_attn.phase = phase
 
-        return output, batch_position + query_idx
+    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,
+        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,
-        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.get_pre_attention_layernorm()(hidden_states)
 
-        hidden_states = self.input_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,
+            seq_positions=seq_positions,
+            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
+
+
+class DecoderOnlyAttention(nn.Module):
+    """Attention implementation for decoder-only models optimized for RBLN compilation.
+
+    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
+    """
+
+    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.scale = torch.tensor(self.get_attn_scale())
+
+        if hasattr(self._original_mod, "num_key_value_heads"):
+            self.num_key_value_heads = self._original_mod.num_key_value_heads
+        else:
+            self.num_key_value_heads = self._original_mod.num_heads
+
+        self.attention = self.get_attention()
+        self.__post_init__()
+
+    @property
+    def phase(self):
+        return self._phase
+
+    @phase.setter
+    def phase(self, phase: str):
+        self._phase = phase
+        self.attention.phase = phase
 
-        if output_attentions:
-            outputs += (self_attn_weight,)
+    def get_attention(self):
+        return AttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads)
 
-        if use_cache:
-            outputs += (past_key_value,)
+    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
 
-        return outputs
+    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 get_attn_scale(self):
+        return 1 / math.sqrt(self.head_dim)
 
-class DecoderOnlyModel:
     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.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"
-            )
+        hidden_states: torch.Tensor,
+        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,
+    ):
+        batch_size, query_length, _ = hidden_states.size()
 
-        # 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_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}.")
+
+        # TODO(jongho): flash attn legacy. (clone)
+        _seq_positions = seq_positions.clone().unsqueeze(1)
+
+        _key_states = []
+        _value_states = []
+        _attn_outputs = []
+        for b in range(batch_size):
+            seq_position = _seq_positions[b][0]
+            attn_output, key_state, value_state = self.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,
+                past_key_state=past_key_values[self.layer_idx][0],
+                past_value_state=past_key_values[self.layer_idx][1],
+                batch_position=b if self.phase == "decode" else batch_position,
+                seq_position=seq_position,
+                scale=self.scale,
             )
+            _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)
 
-            hidden_states = layer_outputs[0]
+        attn_outputs = self.o_proj(attn_outputs)
+        past_key_values[self.layer_idx] = key_states, value_states
+        return attn_outputs, past_key_values
 
-            updated_cache = layer_outputs[2 if output_attentions else 1]
 
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
+class AttentionOp(nn.Module):
+    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.phase = "prefill"
 
-        hidden_states = self.norm(hidden_states)
+    def forward(
+        self,
+        query_state: torch.Tensor,
+        key_state: torch.Tensor,
+        value_state: torch.Tensor,
+        attn_mask: torch.Tensor,
+        batch_position: torch.Tensor,
+        past_key_state: torch.Tensor,
+        past_value_state: torch.Tensor,
+        seq_position: torch.Tensor,
+        scale: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """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 ∈ {0, 1}
+            batch_position: Batch index for cache lookup
+            past_key_state: Previous key cache states
+            past_value_state: Previous value cache states
+            seq_position: Current position in sequence
+            scale: Scale applied to attn weights
+
+        Returns:
+            Tuple of (attention_output, key_state, value_state)
+        """
+        # 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)
 
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
+        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,
+        )
 
-        # convert RebelDynamicCache to legacy Tuple[Tuple[torch.Tensor]]
-        next_cache = updated_cache.to_legacy_cache()
+        if self.phase == "decode":
+            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.attn_decode(
+                query_state,
+                key_state,
+                value_state,
+                attn_mask,
+                past_key_state.unsqueeze(2),
+                past_value_state.unsqueeze(2),
+                seq_position,
+                scale,
+            )
 
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
+        else:
+            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.attn_prefill(
+                query_state,
+                key_state,
+                value_state,
+                attn_mask,
+                past_key_state.unsqueeze(2),
+                past_value_state.unsqueeze(2),
+                batch_position,
+                seq_position,
+                scale,
+            )
 
+        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.squeeze(2), value_state.squeeze(2)
 
-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 +775,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 +811,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 +833,127 @@ 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):
+        self.kvcache_partition_size = kvcache_partition_len
+        super().__init__(self_attn=self_attn)
+
+    def get_attention(self):
+        return FlashAttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads, self.kvcache_partition_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        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,
+    ):
+        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
+
+        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)
+
+        _key_states = []
+        _value_states = []
+        _attn_outputs = []
+        for b in range(batch_size):
+            seq_position = seq_positions[b][0]  # FIXME: Remove take-take pattern matching
+            attn_output, key_state, value_state = self.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,
+                past_key_state=past_key_values[self.layer_idx][0],
+                past_value_state=past_key_values[self.layer_idx][1],
+                batch_position=b if self.phase == "decode" else batch_position,
+                seq_position=seq_position,
+                scale=self.scale,
+            )
+            _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
+
+
+class FlashAttentionOp(AttentionOp):
+    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int, kvcache_partition_len: int):
+        super().__init__(num_heads=num_heads, head_dim=head_dim, num_key_value_heads=num_key_value_heads)
+        self.kvcache_partition_size = kvcache_partition_len
+
+    def forward(
+        self,
+        query_state,
+        key_state,
+        value_state,
+        attn_mask,
+        batch_position,
+        past_key_state,
+        past_value_state,
+        seq_position,
+        scale,
+    ):
+        # 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,  # seq len
+            self.head_dim,
+        )
+
+        if self.phase == "decode":
+            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_attn_decode(
+                query_state,
+                key_state,
+                value_state,
+                attn_mask,
+                past_key_state.unsqueeze(2),
+                past_value_state.unsqueeze(2),
+                seq_position,
+                scale,
+                self.kvcache_partition_size,
+            )
+        else:
+            attn_output, key_state, value_state = torch.ops.rbln_custom_ops.flash_attn_prefill(
+                query_state,
+                key_state,
+                value_state,
+                attn_mask,
+                past_key_state.unsqueeze(2),
+                past_value_state.unsqueeze(2),
+                batch_position,
+                seq_position,
+                scale,
+                self.kvcache_partition_size,
+            )
+
+        # 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