optimum-rbln 0.7.2rc1__py3-none-any.whl → 0.7.3__py3-none-any.whl

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

@@ -32,11 +32,13 @@ if TYPE_CHECKING:
 
 
 class PhiWrapper(DecoderOnlyWrapper):
-    def convert_to_rbln_causal_lm(self, causal_lm: "PhiForCausalLM"):
+    def convert_to_rbln_causal_lm(self, causal_lm: "PhiForCausalLM", max_seq_len: int):
         new_layers = []
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
-                new_self_attn = PhiAttention(layer.self_attn)
+                new_self_attn = PhiAttention(
+                    layer.self_attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                )
             elif self.attn_impl == "flash_attn":
                 raise NotImplementedError(f"flash attn for {self.__class__} is not implemented yet.")
             else:
@@ -81,30 +83,30 @@ class PhiLayer(DecoderOnlyLayer):
         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,
+        block_tables: Optional[torch.Tensor] = None,
     ):
         residual = hidden_states
 
         hidden_states = self.get_pre_attention_layernorm()(hidden_states)
 
-        attn_outputs, present_key_values = self.self_attn(
+        attn_output = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
             seq_positions=seq_positions,
-            batch_position=batch_position,
             past_key_values=past_key_values,
             cos=cos,
             sin=sin,
+            block_tables=block_tables,
         )
 
         feed_forward_hidden_states = self._original_mod.mlp(hidden_states)
 
-        hidden_states = attn_outputs + feed_forward_hidden_states + residual
+        hidden_states = attn_output + feed_forward_hidden_states + residual
 
-        return hidden_states, present_key_values
+        return hidden_states
 
 
 class PhiModel(DecoderOnlyModel):

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

@@ -50,11 +50,18 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         runtime: rebel.Runtime,
         batch_size: int,
         dec_max_seq_len: int,
+        support_paged_causal_attn: Optional[bool] = None,
+        use_attention_mask: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(runtime, **kwargs)
         self.batch_size = batch_size
         self.dec_max_seq_len = dec_max_seq_len
+        self.use_attention_mask = use_attention_mask
+        if support_paged_causal_attn:
+            self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
+        else:
+            self.default_block_tables = None
 
     def forward(
         self,
@@ -62,6 +69,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         attention_mask: Optional[torch.FloatTensor] = None,
         decoder_attention_mask: Optional[torch.BoolTensor] = None,
         cache_position: Optional[torch.Tensor] = None,
+        block_tables: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> Tuple[torch.FloatTensor]:
         batch_size = decoder_input_ids.shape[0]
@@ -73,19 +81,24 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         if batch_size != cache_position.shape[0]:
             raise RuntimeError(f"Cache position size mismatch: got {cache_position.shape[0]}, expected {batch_size}.")
 
-        for b_idx in range(self.batch_size):
-            decoding_step = cache_position[b_idx].item()
-            if not (0 <= decoding_step < self.dec_max_seq_len):
-                raise ValueError(
-                    f"Decoding step {decoding_step} out of bounds for attention mask with shape {self.dec_attn_mask.shape}."
-                )
-            decoder_attention_mask[b_idx, : decoding_step + 1] = 1
+        if self.use_attention_mask:
+            for b_idx in range(self.batch_size):
+                decoding_step = cache_position[b_idx].item()
+                if not (0 <= decoding_step < self.dec_max_seq_len):
+                    raise ValueError(
+                        f"Decoding step {decoding_step} out of bounds for attention mask with shape {self.dec_attn_mask.shape}."
+                    )
+                decoder_attention_mask[b_idx, : decoding_step + 1] = 1
+
+        if block_tables is None:
+            block_tables = self.default_block_tables
 
         lm_logits = super().forward(
             decoder_input_ids,
-            decoder_attention_mask,
+            decoder_attention_mask if self.use_attention_mask else None,
             attention_mask,
             cache_position,
+            block_tables=block_tables,
         )
 
         return Seq2SeqLMOutput(logits=lm_logits)
@@ -106,16 +119,24 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
+    support_paged_causal_attn = None
 
     def __post_init__(self, **kwargs):
         batch_size = self.rbln_config.model_cfg["batch_size"]
         dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
+        self.use_attention_mask = self.rbln_config.model_cfg.get("use_attention_mask", None)
+
         self.encoder = RBLNRuntimeEncoder(
             runtime=self.model[0],
             main_input_name="input_ids",
         )
         self.decoder = RBLNRuntimeDecoder(
-            runtime=self.model[1], main_input_name="input_ids", batch_size=batch_size, dec_max_seq_len=dec_max_seq_len
+            runtime=self.model[1],
+            main_input_name="input_ids",
+            batch_size=batch_size,
+            dec_max_seq_len=dec_max_seq_len,
+            support_paged_causal_attn=self.support_paged_causal_attn,
+            use_attention_mask=self.use_attention_mask,
         )
 
     @classmethod
@@ -172,6 +193,16 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         rbln_batch_size = rbln_kwargs.get("batch_size", None)
         rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
 
+        if cls.support_paged_causal_attn:
+            rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
+            if rbln_use_attention_mask is None:
+                rbln_use_attention_mask = False
+                rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
+                if rbln_npu == "RBLN-CA02":
+                    rbln_use_attention_mask = True
+        else:
+            rbln_use_attention_mask = True
+
         n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
         n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
         d_kv = (
@@ -232,12 +263,11 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 ],
                 "float32",
             ),
-            ("batch_position", [], "int16"),
+            ("block_tables", [1], "int16"),
         ]
 
         dec_input_info = [
             ("input_ids", [rbln_batch_size, 1], "int64"),
-            ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"),
             ("encoder_attention_mask", [rbln_batch_size, rbln_enc_max_seq_len], "float32"),
             (
                 "cache_position",
@@ -275,6 +305,12 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 for i in range(n_layer * 2)
             ]
         )
+
+        if cls.support_paged_causal_attn:
+            dec_input_info.insert(3, ("block_tables", [rbln_batch_size, 1], "int16"))
+        if rbln_use_attention_mask:
+            dec_input_info.insert(1, ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"))
+
         enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
         dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
@@ -290,6 +326,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 "dec_max_seq_len": rbln_dec_max_seq_len,
                 "batch_size": rbln_batch_size,
                 "pad_token_id": rbln_pad_token_id,
+                "use_attention_mask": rbln_use_attention_mask,
             }
         )
 
@@ -400,9 +437,9 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         encoder_kwargs["output_attentions"] = False
 
         for b in range(batch_size):
-            batch_position = torch.tensor(b, dtype=torch.int16)
+            block_tables = torch.tensor([b], dtype=torch.int16)
             encoder_kwargs["input_ids"] = inputs_tensor[b].unsqueeze(0)
             encoder_kwargs["attention_mask"] = model_kwargs["attention_mask"][b].unsqueeze(0).to(torch.float32)
-            model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, batch_position=batch_position)
+            model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, block_tables=block_tables)
 
         return model_kwargs

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

@@ -12,13 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch import nn
 from transformers.utils import logging
 
-from ....ops import register_rbln_custom_attention, register_rbln_custom_cache_update
+from ....ops import (
+    register_rbln_custom_cache_update,
+    register_rbln_custom_paged_attention,
+    register_rbln_custom_paged_causal_attention,
+)
 
 
 logger = logging.get_logger(__name__)
@@ -87,7 +91,7 @@ class Seq2SeqEncoderWrapper(nn.Module):
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
         cross_key_values: torch.Tensor,
-        batch_position: torch.Tensor,
+        b_idx: torch.Tensor,
     ) -> Tuple[torch.Tensor]:
         # 1. get encoder last_hidden_states
         encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
@@ -110,11 +114,9 @@ class Seq2SeqEncoderWrapper(nn.Module):
 
         # 3. update the cross_attention's past_key_value direct to the device-dram for optimization.
         batch_axis = torch.tensor(1, dtype=torch.int16)
-        cross_key_values = torch.ops.rbln_custom_ops.rbln_cache_update(
-            cross_key_values, cross_kv, batch_position, batch_axis
-        )
+        enc_out = torch.ops.rbln_custom_ops.rbln_cache_update(cross_key_values, cross_kv, b_idx[0], batch_axis)
 
-        return cross_key_values
+        return enc_out
 
 
 class Seq2SeqDecoderWrapper(nn.Module):
@@ -131,9 +133,10 @@ class Seq2SeqDecoderWrapper(nn.Module):
         **kwargs: Additional arguments for decoder configuration.
     """
 
-    def __init__(self, model: nn.Module, **kwargs):
+    def __init__(self, model: nn.Module, use_attention_mask: bool = True, **kwargs):
         super().__init__()
         self.config = model.config
+        self.use_attention_mask = use_attention_mask
         self.__post_init__(model, **kwargs)
 
     def __post_init__(self, model: nn.Module, **kwargs):
@@ -143,7 +146,11 @@ class Seq2SeqDecoderWrapper(nn.Module):
         It is inspired by the BART architecture, but it is designed to be flexible and can be overridden
         by subclasses to modify or add custom attributes as necessary.
         """
-        register_rbln_custom_attention()
+        if self.use_attention_mask:
+            register_rbln_custom_paged_attention()
+        else:
+            register_rbln_custom_paged_causal_attention()
+
         self.num_layers = self.config.decoder_layers
         self.conditional_generation = self.convert_to_rbln_conditional_generation(model)
 
@@ -160,13 +167,23 @@ class Seq2SeqDecoderWrapper(nn.Module):
 
     def forward(
         self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        cache_position: torch.Tensor,
-        cross_kv_cache: torch.Tensor,
-        *self_kv_cache: torch.Tensor,
+        *args,
     ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
+        if self.use_attention_mask:
+            (
+                input_ids,
+                attention_mask,
+                encoder_attention_mask,
+                cache_position,
+                block_tables,
+                cross_kv_cache,
+                *self_kv_cache,
+            ) = args
+
+        else:
+            attention_mask = None
+            (input_ids, encoder_attention_mask, cache_position, block_tables, cross_kv_cache, *self_kv_cache) = args
+
         self_past_key_values = ()
         cross_past_key_values = ()
         for i in range(0, self.num_layers * 2, 2):
@@ -174,18 +191,17 @@ class Seq2SeqDecoderWrapper(nn.Module):
             cross_past_key_values = cross_past_key_values + ((cross_kv_cache[i], cross_kv_cache[i + 1]),)
 
         # decode
-        lm_logits, self_present_key_values = self.conditional_generation(
+        lm_logits = self.conditional_generation(
             input_ids=input_ids,
             attention_mask=attention_mask,
             encoder_attention_mask=encoder_attention_mask,
             self_past_key_values=self_past_key_values,
             cross_past_key_values=cross_past_key_values,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
 
-        outputs = (lm_logits,) + self_present_key_values
-
-        return outputs
+        return lm_logits
 
 
 class Seq2SeqForConditionalGeneration(nn.Module):
@@ -228,14 +244,16 @@ class Seq2SeqForConditionalGeneration(nn.Module):
         self_past_key_values,
         cross_past_key_values,
         cache_position,
+        block_tables: Optional[torch.Tensor] = None,
     ):
-        hidden_states, self_present_key_values = self.decoder(
+        hidden_states = self.decoder(
             input_ids=input_ids,
             attention_mask=attention_mask,
             encoder_attention_mask=encoder_attention_mask,
             self_past_key_values=self_past_key_values,
             cross_past_key_values=cross_past_key_values,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
 
         if self.has_rescaling and self.config.tie_word_embeddings:
@@ -243,7 +261,7 @@ class Seq2SeqForConditionalGeneration(nn.Module):
 
         lm_logits = self.lm_head(hidden_states)
 
-        return lm_logits, self_present_key_values
+        return lm_logits
 
 
 class Seq2SeqDecoder(torch.nn.Module):
@@ -292,6 +310,7 @@ class Seq2SeqDecoder(torch.nn.Module):
         self_past_key_values: torch.Tensor,
         cross_past_key_values: torch.Tensor,
         cache_position: torch.Tensor,
+        block_tables: Optional[torch.Tensor] = None,
     ):
         # embedding
         hidden_states = self.get_embedding()(input_ids)
@@ -303,24 +322,23 @@ class Seq2SeqDecoder(torch.nn.Module):
             hidden_states = self.apply_position_embedding(hidden_states, cache_position)
 
         # iterate decoder_layer
-        self_present_key_values = ()
         for decoder_layer, self_past_key_value, cross_past_key_value in zip(
             self.layers, self_past_key_values, cross_past_key_values
         ):
-            hidden_states, self_present_key_value = decoder_layer(
+            hidden_states = decoder_layer(
                 hidden_states,
                 attention_mask=attention_mask,
                 encoder_attention_mask=encoder_attention_mask,
                 self_past_key_value=self_past_key_value,
                 cross_past_key_value=cross_past_key_value,
                 cache_position=cache_position,
+                block_tables=block_tables,
             )
-            self_present_key_values += self_present_key_value
 
         if self.final_layer_norm is not None:
             hidden_states = self.final_layer_norm(hidden_states)
 
-        return hidden_states, self_present_key_values
+        return hidden_states
 
 
 class Seq2SeqDecoderLayer(torch.nn.Module):
@@ -373,17 +391,19 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
         self_past_key_value: Tuple[torch.Tensor],
         cross_past_key_value: Tuple[torch.Tensor],
         cache_position: torch.Tensor,
+        block_tables: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
         dummy_encoder_hidden_states = torch.zeros(1, encoder_attention_mask.shape[-1])
 
         # Self Attention Block
         residual = hidden_states
         hidden_states = self.pre_self_attn_layer_norm(hidden_states)
-        hidden_states, self_attn_past_key_value = self.self_attn(
+        hidden_states = self.self_attn(
             hidden_states=hidden_states,
             past_key_value=self_past_key_value,
             attention_mask=attention_mask,
             cache_position=cache_position,
+            block_tables=block_tables,
         )
         hidden_states = residual + hidden_states
         hidden_states = self.post_self_attn_layer_norm(hidden_states)
@@ -403,14 +423,14 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
         # Feed-Forward Block
         hidden_states = self.ff_layer(hidden_states)
 
-        return hidden_states, self_attn_past_key_value
+        return hidden_states
 
 
 class Seq2SeqSelfAttention(nn.Module):
-    def __init__(self, attn):
+    def __init__(self, attn, **kwargs):
         super().__init__()
         self._original_mod = attn
-        self.__post_init__()
+        self.__post_init__(**kwargs)
 
     def __post_init__(self, **kwargs):
         """
@@ -442,6 +462,7 @@ class Seq2SeqSelfAttention(nn.Module):
         past_key_value: Tuple[torch.Tensor],
         attention_mask: torch.Tensor,
         cache_position: torch.Tensor,
+        block_tables: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
         bsz, tgt_len, _ = hidden_states.size()
 
@@ -450,23 +471,26 @@ class Seq2SeqSelfAttention(nn.Module):
         key_states = self._shape(key_states, -1, bsz)
         value_states = self._shape(value_states, -1, bsz)
 
-        attn_output, key_states, value_states = self.attn_decode(
+        block_size = past_key_value[0].shape[-2]
+        args = [
             query_states,
             key_states,
             value_states,
-            attention_mask.unsqueeze(
-                2
-            ),  # Unsqueeze group axis since CustomKernel expects it for group query attention
             past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim),
             past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim),
             cache_position,
             torch.tensor(1.0, dtype=torch.float32),  # scale
-        )
+            block_tables,
+            block_size,
+        ]
+        if attention_mask is not None:
+            args.insert(3, attention_mask.unsqueeze(2))
+
+        attn_output = self.attn_decode(*args)
 
         attn_output = attn_output.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)
         attn_output = attn_output.reshape(bsz, -1, self.num_heads * self.head_dim)
 
         attn_output = self.out_proj(attn_output)
-        present_key_value = (key_states, value_states)
 
-        return attn_output, present_key_value
+        return attn_output

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -120,7 +120,7 @@ class RBLNT5EncoderModel(RBLNModel):
         if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
             raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
 
-        signature_params = inspect.signature(cls.hf_class.forward).parameters.keys()
+        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
 
         if rbln_model_input_names is None:
             for tokenizer in preprocessors:
@@ -189,6 +189,8 @@ class RBLNT5EncoderModel(RBLNModel):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    support_causal_paged_attn = False
+
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = rbln_config.model_cfg["enc_max_seq_len"]

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -18,7 +18,7 @@ import torch
 from torch import nn
 from transformers.utils import logging
 
-from ....ops import register_rbln_custom_attention_add_softmax
+from ....ops import register_rbln_custom_add_softmax_attention
 from ..seq2seq.seq2seq_architecture import (
     Seq2SeqDecoder,
     Seq2SeqDecoderLayer,
@@ -55,7 +55,7 @@ class T5EncoderWrapper(Seq2SeqEncoderWrapper):
 
 class T5DecoderWrapper(Seq2SeqDecoderWrapper):
     def __post_init__(self, model, dec_max_seq_len: int = None):
-        register_rbln_custom_attention_add_softmax()
+        register_rbln_custom_add_softmax_attention()
         self.num_layers = self.config.num_layers
         self.conditional_generation = self.convert_to_rbln_conditional_generation(model, dec_max_seq_len)
 
@@ -71,6 +71,34 @@ class T5DecoderWrapper(Seq2SeqDecoderWrapper):
 
         return new_model
 
+    def forward(
+        self,
+        input_ids,
+        attention_mask,
+        encoder_attention_mask,
+        cache_position,
+        cross_kv_cache,
+        *self_kv_cache,
+    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
+        self_past_key_values = ()
+        cross_past_key_values = ()
+
+        for i in range(0, self.num_layers * 2, 2):
+            self_past_key_values = self_past_key_values + ((self_kv_cache[i], self_kv_cache[i + 1]),)
+            cross_past_key_values = cross_past_key_values + ((cross_kv_cache[i], cross_kv_cache[i + 1]),)
+
+        # decode
+        lm_logits = self.conditional_generation(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            encoder_attention_mask=encoder_attention_mask,
+            self_past_key_values=self_past_key_values,
+            cross_past_key_values=cross_past_key_values,
+            cache_position=cache_position,
+        )
+
+        return lm_logits
+
 
 class T5ForConditionalGeneration(Seq2SeqForConditionalGeneration):
     has_rescaling = True
@@ -134,7 +162,7 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
         self.out_proj = self._original_mod.o
         self.num_heads = self._original_mod.n_heads
         self.head_dim = self._original_mod.key_value_proj_dim
-        self.attn_decode = torch.ops.rbln_custom_ops.attn_decode_add_softmax
+        self.attn_decode = torch.ops.rbln_custom_ops.add_softmax_attn_decode
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_states = self.q_proj(hidden_states)
@@ -142,6 +170,40 @@ class T5LayerSelfAttention(Seq2SeqSelfAttention):
         value_states = self.v_proj(hidden_states)
         return query_states, key_states, value_states
 
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        past_key_value: Tuple[torch.Tensor],
+        attention_mask: torch.Tensor,
+        cache_position: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
+        bsz, tgt_len, _ = hidden_states.size()
+
+        query_states, key_states, value_states = self.projection(hidden_states=hidden_states)
+        query_states = self._shape(query_states, tgt_len, bsz)
+        key_states = self._shape(key_states, -1, bsz)
+        value_states = self._shape(value_states, -1, bsz)
+
+        attn_output = self.attn_decode(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask.unsqueeze(
+                2
+            ),  # Unsqueeze group axis since CustomKernel expects it for group query attention
+            past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim),
+            past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim),
+            cache_position,
+            torch.tensor(1.0, dtype=torch.float32),  # scale
+        )
+
+        attn_output = attn_output.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, -1, self.num_heads * self.head_dim)
+
+        attn_output = self.out_proj(attn_output)
+        return attn_output
+
 
 class T5CrossAttention(nn.Module):
     def __init__(self, attn):