optimum-rbln 0.7.3a1__py3-none-any.whl → 0.7.3a3__py3-none-any.whl

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_cache_update, register_rbln_custom_masked_attention
+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)
@@ -111,7 +115,7 @@ 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
+            cross_key_values, cross_kv, b_idx[0], batch_axis
         )
 
         return cross_key_values
@@ -131,9 +135,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 +148,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_masked_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 +169,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):
@@ -181,6 +200,7 @@ class Seq2SeqDecoderWrapper(nn.Module):
             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
@@ -228,6 +248,7 @@ 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(
             input_ids=input_ids,
@@ -236,6 +257,7 @@ class Seq2SeqForConditionalGeneration(nn.Module):
             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:
@@ -292,6 +314,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)
@@ -314,6 +337,7 @@ class Seq2SeqDecoder(torch.nn.Module):
                 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
 
@@ -373,6 +397,7 @@ 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])
 
@@ -384,6 +409,7 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
             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)
@@ -407,10 +433,10 @@ class Seq2SeqDecoderLayer(torch.nn.Module):
 
 
 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 +468,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,18 +477,22 @@ 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, key_states, value_states = 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)

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

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
+import rebel
 import torch
 from transformers import (
     AutoModelForTextEncoding,
@@ -22,7 +23,7 @@ from transformers import (
     T5EncoderModel,
     T5ForConditionalGeneration,
 )
-from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
 from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
 from ....modeling import RBLNModel
@@ -57,6 +58,63 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         )
 
 
+class RBLNRuntimeEncoder(RBLNPytorchRuntime):
+    mandatory_members = ["main_input_name"]
+
+    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+        _ = super().forward(*args, **kwargs)
+        return BaseModelOutput(last_hidden_state=torch.tensor([1.0]))
+
+
+class RBLNRuntimeDecoder(RBLNPytorchRuntime):
+    mandatory_members = ["main_input_name"]
+
+    def __init__(
+        self,
+        runtime: rebel.Runtime,
+        batch_size: int,
+        dec_max_seq_len: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.batch_size = batch_size
+        self.dec_max_seq_len = dec_max_seq_len
+
+    def forward(
+        self,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor]:
+        batch_size = decoder_input_ids.shape[0]
+        if batch_size != self.batch_size:
+            raise RuntimeError(
+                f"Batch size mismatch: got {batch_size}, expected {self.batch_size} (compiled batch size)."
+            )
+
+        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
+
+        lm_logits = super().forward(
+            decoder_input_ids,
+            decoder_attention_mask,
+            attention_mask,
+            cache_position,
+        )
+
+        return Seq2SeqLMOutput(logits=lm_logits)
+
+
 class T5EncoderWrapper(torch.nn.Module):
     def __init__(self, model: "T5EncoderModel") -> None:
         super().__init__()
@@ -189,6 +247,21 @@ class RBLNT5EncoderModel(RBLNModel):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    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.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,
+        )
+
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = rbln_config.model_cfg["enc_max_seq_len"]
@@ -206,3 +279,139 @@ class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
             return redirect(val)
 
         return val
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model_config: "PretrainedConfig",
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_enc_max_seq_len = rbln_kwargs.get("enc_max_seq_len", None)
+        rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
+
+        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 = (
+            model_config.d_kv
+            if hasattr(model_config, "d_kv")
+            else model_config.d_model // model_config.encoder_attention_heads
+        )
+
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        rbln_pad_token_id = getattr(model_config, "pad_token_id", None)
+        if rbln_pad_token_id is None:
+            rbln_pad_token_id = getattr(model_config, "bos_token_id", None)
+            if rbln_pad_token_id is None:
+                rbln_pad_token_id = getattr(model_config, "eos_token_id", None)
+                if rbln_pad_token_id is None:
+                    rbln_pad_token_id = -1
+
+        if rbln_enc_max_seq_len is None:
+            rbln_enc_max_seq_len = max_position_embeddings
+            if rbln_enc_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_enc_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_enc_max_seq_len is None:
+                    raise ValueError("`rbln_enc_max_seq_len` should be specified!")
+        if max_position_embeddings is not None and rbln_enc_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_dec_max_seq_len is None:
+            rbln_dec_max_seq_len = max_position_embeddings
+            if rbln_dec_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_dec_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_dec_max_seq_len is None:
+                    raise ValueError("`rbln_dec_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_dec_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_dec_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        # model input info
+        enc_input_info = [
+            ("input_ids", [1, rbln_enc_max_seq_len], "int64"),
+            ("attention_mask", [1, rbln_enc_max_seq_len], "float32"),
+            (
+                "cross_key_value_states",
+                [
+                    n_layer * 2,
+                    rbln_batch_size,
+                    n_head,
+                    rbln_enc_max_seq_len,
+                    d_kv,
+                ],
+                "float32",
+            ),
+            ("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",
+                [rbln_batch_size, 1],
+                "int32",
+            ),
+        ]
+        dec_input_info.extend(
+            [
+                (
+                    "cross_key_value_states",
+                    [
+                        n_layer * 2,
+                        rbln_batch_size,
+                        n_head,
+                        rbln_enc_max_seq_len,
+                        d_kv,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+        dec_input_info.extend(
+            [
+                (
+                    f"self_key_value_states_{i}",
+                    [
+                        rbln_batch_size,
+                        n_head,
+                        rbln_dec_max_seq_len,
+                        d_kv,
+                    ],
+                    "float32",
+                )
+                for i in range(n_layer * 2)
+            ]
+        )
+
+        enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
+        dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[enc_compile_config, dec_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "enc_max_seq_len": rbln_enc_max_seq_len,
+                "dec_max_seq_len": rbln_dec_max_seq_len,
+                "batch_size": rbln_batch_size,
+                "pad_token_id": rbln_pad_token_id,
+            }
+        )
+
+        return rbln_config

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,36 @@ 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_present_key_values = 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,
+        )
+
+        outputs = (lm_logits,) + self_present_key_values
+
+        return outputs
+
 
 class T5ForConditionalGeneration(Seq2SeqForConditionalGeneration):
     has_rescaling = True
@@ -134,7 +164,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 +172,42 @@ 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, key_states, value_states = 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)
+        present_key_value = (key_states, value_states)
+
+        return attn_output, present_key_value
+
 
 class T5CrossAttention(nn.Module):
     def __init__(self, attn):

optimum/rbln/transformers/models/whisper/whisper_architecture.py

@@ -25,7 +25,7 @@ from transformers.modeling_outputs import (
 )
 from transformers.utils import logging
 
-from ....ops import register_rbln_custom_cache_update
+from ....ops import register_rbln_custom_add_softmax_attention, register_rbln_custom_cache_update
 
 
 logger = logging.get_logger(__name__)
@@ -34,6 +34,7 @@ logger = logging.get_logger(__name__)
 class WhisperWrapper:
     def __init__(self, model, rbln_token_timestamps):
         register_rbln_custom_cache_update()
+        register_rbln_custom_add_softmax_attention()
         self.encoder = WhisperEncoderWrapper(model)
         self.decoder = WhisperDecoderWrapper(model, output_attentions=rbln_token_timestamps)
 
@@ -213,7 +214,7 @@ class WhisperDecoderLayer(nn.Module):
         # Self Attention Block
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(hidden_states)
-        hidden_states, _, self_present_key_value = self.self_attn(
+        hidden_states, self_present_key_value = self.self_attn(
             hidden_states=hidden_states,
             past_key_value=self_past_key_value,
             attention_mask=attention_mask,
@@ -224,7 +225,7 @@ class WhisperDecoderLayer(nn.Module):
         # Cross-Attention Block
         residual = hidden_states
         hidden_states = self.encoder_attn_layer_norm(hidden_states)
-        hidden_states, cross_attn_weights, cross_present_key_value = self.encoder_attn(
+        hidden_states, cross_attn_weights = self.encoder_attn(
             hidden_states=hidden_states,
             past_key_value=cross_past_key_value,
         )
@@ -258,19 +259,8 @@ class WhisperAttention(nn.Module):
 
 
 class WhisperSelfAttention(WhisperAttention):
-    def rbln_cache_update(
-        self,
-        past_key_value: torch.Tensor,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        cache_position: torch.Tensor,
-    ):
-        s_idx = torch.tensor(cache_position, dtype=torch.int16)
-        axis = torch.tensor(2, dtype=torch.int16)
-
-        key_states = torch.ops.rbln_custom_ops.rbln_cache_update(past_key_value[0], key_states, s_idx, axis)
-        value_states = torch.ops.rbln_custom_ops.rbln_cache_update(past_key_value[1], value_states, s_idx, axis)
-        return key_states, value_states
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int) -> torch.Tensor:
+        return tensor.view(bsz, seq_len, 1, self.num_heads, self.head_dim).transpose(1, 3)
 
     def forward(
         self,
@@ -285,22 +275,27 @@ class WhisperSelfAttention(WhisperAttention):
 
         key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
         value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-        key_states, value_states = self.rbln_cache_update(past_key_value, key_states, value_states, cache_position)
 
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3))
-        attn_weights = attn_weights + attention_mask
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+        attn_output, key_states, value_states = torch.ops.rbln_custom_ops.add_softmax_attn_decode(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask.unsqueeze(2),
+            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.expand(bsz, 1),
+            torch.tensor(1.0, dtype=torch.float32),  # scale
+        )
 
-        attn_output = torch.matmul(attn_weights, value_states)
         attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
         attn_output = attn_output.transpose(1, 2)
         attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
         attn_output = self.out_proj(attn_output)
 
-        return attn_output, attn_weights, (key_states, value_states)
+        return attn_output, (key_states, value_states)
 
 
-class WhisperCrossAttention(WhisperSelfAttention):
+class WhisperCrossAttention(WhisperAttention):
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -322,4 +317,4 @@ class WhisperCrossAttention(WhisperSelfAttention):
         attn_output = attn_output.reshape(batch_size, query_len, self.embed_dim)
         attn_output = self.out_proj(attn_output)
 
-        return attn_output, attn_weights, (key_states, value_states)
+        return attn_output, attn_weights

{optimum_rbln-0.7.3a1.dist-info → optimum_rbln-0.7.3a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.3a1
+Version: 0.7.3a3
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai