optimum-rbln 0.1.9__py3-none-any.whl → 0.1.12__py3-none-any.whl

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

@@ -30,7 +30,6 @@ import torch
 from transformers import (
     AutoModelForSpeechSeq2Seq,
     AutoProcessor,
-    GenerationMixin,
     PretrainedConfig,
     WhisperForConditionalGeneration,
     WhisperModel,
@@ -38,8 +37,9 @@ from transformers import (
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
 from ....modeling_base import RBLNModel
-from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
+from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
 from ....utils.runtime_utils import RBLNPytorchRuntime
+from .generation_whisper import RBLNWhisperGenerationMixin
 from .whisper_architecture import (
     _WhisperDecoderWrapper,
     _WhisperEncoderWrapper,
@@ -59,20 +59,47 @@ if TYPE_CHECKING:
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        _ = super().forward(input_features=kwargs["input_features"])
-        return BaseModelOutput(last_hidden_state=torch.tensor([1.0]))
+    def forward(self, input_features: torch.Tensor = None):
+        # backward compatibility transformers==4.40.2
+        # https://github.com/huggingface/transformers/blob/4fdf58afb72b0754da30037fc800b6044e7d9c99/src/transformers/pipelines/automatic_speech_recognition.py#L494
+
+        n_pad_to_batch = self.batch_size - input_features.shape[0]
+        if n_pad_to_batch > 0:
+            input_features = torch.nn.functional.pad(input_features, (0, 0, 0, 0, 0, n_pad_to_batch))
+
+        _ = super().forward(input_features=input_features)
+
+        # dummy output for generation
+        return BaseModelOutput(last_hidden_state=torch.tensor([[-1.0]]))
 
 
 class RBLNRuntimeDecoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        outputs = super().forward(*args, **kwargs)
-        return Seq2SeqLMOutput(logits=outputs)
+    def forward(
+        self,
+        decoder_input_ids: torch.Tensor = None,
+        decoder_attention_mask: torch.Tensor = None,
+        cache_position: torch.Tensor = None,
+    ):
+        inputs_bsz = decoder_input_ids.shape[0]
+        padded_bsz = self.batch_size - inputs_bsz
+        if padded_bsz > 0:
+            decoder_input_ids = torch.nn.functional.pad(decoder_input_ids, (0, 0, 0, padded_bsz))
+
+        outputs = super().forward(
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            cache_position=cache_position,
+        )
+
+        if isinstance(outputs, torch.Tensor):
+            return Seq2SeqLMOutput(logits=outputs[:inputs_bsz], cross_attentions=None)
+        else:
+            return Seq2SeqLMOutput(logits=outputs[0][:inputs_bsz], cross_attentions=outputs[1][:, :inputs_bsz])
 
 
-class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
+class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin):
     """
     The Whisper Model with a language modeling head. Can be used for automatic speech recognition.
     This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
@@ -83,20 +110,30 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForSpeechSeq2Seq
     main_input_name = "input_ids"
 
     def __post_init__(self, **kwargs):
-        self.batch_size = self.rbln_config[DEFAULT_COMPILED_MODEL_NAME][0].batch_size
-        self.enc_max_seq_len = self.rbln_config.meta["input_max_length"]
-        self.dec_max_seq_len = self.rbln_config.meta["rbln_dec_max_seq_len"]
+        super().__post_init__(**kwargs)
 
-        self.encoder = RBLNRuntimeEncoder(runtime=self.model[0], main_input_name="input_features")
-        self.decoder = RBLNRuntimeDecoder(runtime=self.model[1], main_input_name="input_ids")
-        self.forced_decoder_ids = self.config.forced_decoder_ids
+        self.batch_size = self.rbln_config.model_cfg["batch_size"]
+        self.dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
+        self.rbln_token_timestamps = self.rbln_config.model_cfg["token_timestamps"]
+
+        self.encoder = RBLNRuntimeEncoder(
+            runtime=self.model[0], main_input_name="input_features", batch_size=self.batch_size
+        )
+        self.decoder = RBLNRuntimeDecoder(
+            runtime=self.model[1], main_input_name="input_ids", batch_size=self.batch_size
+        )
 
-        # used in GenerationMixin.generate()
+        # skip encoder &  first decoder when language detected
+        self.is_language_detected = False
+        self.language_cross = None
+
+        # Used in GenerationMixin.generate()
+        # transformers/models/whisper/generation_whisper.py, line 505, in generate
+        #     input_stride = self.model.encoder.conv1.stride[0] * self.model.encoder.conv2.stride[0]
         self.model = WhisperModel(self.config)
         self.pad_token_id = self.config.pad_token_id
 
@@ -127,63 +164,32 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
         # TODO(jongho): implement
         raise NotImplementedError
 
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        decoder_attention_mask=None,
-        input_features=None,  # Must be explicit
-        **kwargs,
-    ):
-        max_seq_len = self.dec_max_seq_len
-        cur_seq_len = input_ids.shape[-1]
-        input_ids = input_ids[:, cur_seq_len - 1 : cur_seq_len].contiguous()
-        decoder_attention_mask = torch.zeros(self.batch_size, max_seq_len, dtype=torch.int64)
-        decoder_attention_mask[:, :cur_seq_len] = 1
-        cache_position = torch.tensor(cur_seq_len - 1, dtype=torch.int32)
-
-        return {
-            "decoder_input_ids": input_ids,
-            "decoder_attention_mask": decoder_attention_mask,
-            "cache_position": cache_position,
-        }
-
-    @classmethod
-    def update_kwargs(cls, kwargs):
-        kwargs.update(
-            {
-                "torchscript": True,
-                "return_dict": False,
-                "use_cache": True,
-            }
-        )
-        return kwargs
-
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model, rbln_config: RBLNConfig):
+        rbln_token_timestamps = rbln_config.model_cfg["token_timestamps"]
         wrapped_encoder = _WhisperEncoderWrapper(model).eval()
-        wrapped_decoder = _WhisperDecoderWrapper(model).eval()
+        wrapped_decoder = _WhisperDecoderWrapper(model, output_attentions=rbln_token_timestamps).eval()
 
-        enc_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
-        dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
+        enc_rbln_compile_config = rbln_config.compile_cfgs[0]
+        dec_rbln_compile_config = rbln_config.compile_cfgs[1]
 
-        enc_example_inputs = enc_rbln_runtime_config.get_dummy_inputs(fill=1)
-        dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=1)
+        enc_example_inputs = enc_rbln_compile_config.get_dummy_inputs(fill=1)
+        dec_example_inputs = dec_rbln_compile_config.get_dummy_inputs(fill=1)
 
-        enc_scripted_model = torch.jit.trace(wrapped_encoder, enc_example_inputs[0], check_trace=False)
+        enc_scripted_model = torch.jit.trace(wrapped_encoder, enc_example_inputs, check_trace=False)
         dec_scripted_model = torch.jit.trace(wrapped_decoder, dec_example_inputs, check_trace=False)
 
         enc_ir = rebel.torchscript_to_ir(
             enc_scripted_model,
-            input_names=[v[0] for v in enc_rbln_runtime_config.input_info],
-            name=enc_rbln_runtime_config.rbln_mod_name,
+            input_names=[v[0] for v in enc_rbln_compile_config.input_info],
+            name=enc_rbln_compile_config.mod_name,
         )
         dec_ir = rebel.torchscript_to_ir(
             dec_scripted_model,
-            input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
-            name=dec_rbln_runtime_config.rbln_mod_name,
+            input_names=[v[0] for v in dec_rbln_compile_config.input_info],
+            name=dec_rbln_compile_config.mod_name,
         )
-        dec_ir.batch_size = dec_rbln_runtime_config.batch_size
 
         # Caching encoder/decoder I/O
         connections = [
@@ -194,9 +200,9 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
             enc_ir,
             dec_ir,
             connections=connections,
-            fusion=enc_rbln_runtime_config.fusion,
-            npu=enc_rbln_runtime_config.npu,
-            tensor_parallel_size=enc_rbln_runtime_config.tensor_parallel_size,
+            fusion=enc_rbln_compile_config.fusion,
+            npu=enc_rbln_compile_config.npu,
+            tensor_parallel_size=enc_rbln_compile_config.tensor_parallel_size,
         )
         return compiled_model
 
@@ -205,42 +211,26 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor"],
         model_config: "PretrainedConfig",
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {}
-
-        input_max_length = 3000
-        rbln_enc_num_mel_bins = getattr(model_config, "num_mel_bins", None)
-        if rbln_enc_num_mel_bins is None:
-            for feature_extractor in preprocessors:
-                if hasattr(feature_extractor, "feature_size"):
-                    rbln_enc_num_mel_bins = feature_extractor.feature_size
-                    break
-            raise ValueError("`rbln_enc_num_mel_bins` should be specified!")
-
-        rbln_enc_max_seq_len = getattr(model_config, "max_source_positions", None)
-        if rbln_enc_max_seq_len is None:
-            raise ValueError("`rbln_enc_max_seq_len` should be specified!")
-
-        rbln_dec_max_seq_len = getattr(model_config, "max_length", None)
-        if rbln_dec_max_seq_len is None:
-            raise ValueError("`rbln_dec_max_seq_len` should be specified!")
-
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_token_timestamps = rbln_kwargs.get("token_timestamps", False)
         rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-        decoder_batch_size = rbln_batch_size
 
-        meta["rbln_dec_max_seq_len"] = rbln_dec_max_seq_len
-        meta["rbln_enc_max_seq_len"] = rbln_enc_max_seq_len
-        meta["num_mel_bins"] = rbln_enc_num_mel_bins
-        meta["input_max_length"] = input_max_length
-        meta["decoder_batch_size"] = decoder_batch_size
-        meta["forced_decoder_ids"] = model_config.forced_decoder_ids
+        expected_seq_len = model_config.max_source_positions * 2
+        num_mel_bins = model_config.num_mel_bins
+        enc_max_seq_len = model_config.max_source_positions
+
+        # 'whisper-large-v3-turbo' doesn't have 'max_length', but PretrainedConfig have default value for the key 'max_length'
+        rbln_dec_max_seq_len = getattr(model_config, "max_target_positions", None)
+        if rbln_dec_max_seq_len is None:
+            rbln_dec_max_seq_len = model_config.max_length
 
         # model input info
-        enc_input_info = [("input_features", [rbln_batch_size, rbln_enc_num_mel_bins, input_max_length], "float32")]
+        enc_input_info = [("input_features", [rbln_batch_size, num_mel_bins, expected_seq_len], "float32")]
         dec_input_info = [
-            ("decoder_input_ids", [decoder_batch_size, 1], "int64"),
-            ("decoder_attention_mask", [decoder_batch_size, rbln_dec_max_seq_len], "int64"),
+            ("decoder_input_ids", [rbln_batch_size, 1], "int64"),
+            ("decoder_attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "int64"),
             ("cache_position", [], "int32"),
         ]
         dec_input_info.extend(
@@ -249,7 +239,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
                     "self_key_value_states",
                     [
                         model_config.decoder_layers * 2,
-                        decoder_batch_size,
+                        rbln_batch_size,
                         model_config.decoder_attention_heads,
                         rbln_dec_max_seq_len,
                         model_config.d_model // model_config.encoder_attention_heads,
@@ -266,7 +256,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
                         model_config.decoder_layers * 2,
                         rbln_batch_size,
                         model_config.decoder_attention_heads,
-                        rbln_enc_max_seq_len,
+                        enc_max_seq_len,
                         model_config.d_model // model_config.encoder_attention_heads,
                     ],
                     "float32",
@@ -274,15 +264,21 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
             ]
         )
 
-        enc_rbln_runtime_config = RBLNRuntimeConfig(rbln_mod_name="encoder", input_info=enc_input_info)
-        dec_rbln_runtime_config = RBLNRuntimeConfig(rbln_mod_name="decoder", input_info=dec_input_info)
+        enc_rbln_compile_config = RBLNCompileConfig(mod_name="encoder", input_info=enc_input_info)
+        dec_rbln_compile_config = RBLNCompileConfig(mod_name="decoder", input_info=dec_input_info)
 
-        enc_rbln_runtime_config.batch_size = rbln_batch_size
-        dec_rbln_runtime_config.batch_size = decoder_batch_size
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[enc_rbln_compile_config, dec_rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
 
-        rbln_config = RBLNConfig.from_rbln_runtime_configs(
-            [enc_rbln_runtime_config, dec_rbln_runtime_config],
-            _rbln_meta=meta,
+        rbln_config.model_cfg.update(
+            {
+                "batch_size": rbln_batch_size,
+                "dec_max_seq_len": rbln_dec_max_seq_len,
+                "token_timestamps": rbln_token_timestamps,
+            }
         )
 
         return rbln_config
@@ -297,18 +293,83 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, GenerationMixin):
             compiled_models[0].create_runtime("decoder", tensor_type="pt", device=device_val),
         ]
 
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        cache_position: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # need for support transformers>=4.45.0
+        **kwargs,
+    ):
+        """
+        whisper don't use attention_mask,
+            attention_mask (`torch.Tensor`)`, *optional*):
+                Whisper does not support masking of the `input_features`, this argument is preserved for compatibility,
+                but it is not used. By default the silence in the input log mel spectrogram are ignored.
+        """
+        return {
+            "input_ids": input_ids,
+            "cache_position": cache_position,
+        }
+
+    # https://github.com/huggingface/transformers/blob/174890280b340b89c5bfa092f6b4fb0e2dc2d7fc/src/transformers/generation/utils.py#L512
+    def _prepare_encoder_decoder_kwargs_for_generation(
+        self, inputs_tensor: torch.Tensor, model_kwargs, *args, **kwargs
+    ) -> Dict[str, Any]:
+        if not self.is_language_detected:
+            model_kwargs["encoder_outputs"] = self.encoder(input_features=inputs_tensor)
+            self.decoder_attention_mask = torch.zeros(self.batch_size, self.dec_max_seq_len, dtype=torch.int64)
+        else:
+            model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=torch.tensor([[-1.0]]))
+
+        return model_kwargs
+
     def forward(
         self,
-        decoder_input_ids: Optional[torch.LongTensor] = None,
-        decoder_attention_mask: Optional[torch.LongTensor] = None,
+        input_ids: Optional[torch.LongTensor] = None,
         cache_position: Optional[torch.Tensor] = None,
+        input_features: Optional[torch.Tensor] = None,
+        decoder_input_ids: Optional[torch.Tensor] = None,
+        encoder_outputs: Optional[Seq2SeqLMOutput] = None,
         **kwargs,
     ) -> Seq2SeqLMOutput:
-        decoder_output = self.decoder(
-            decoder_input_ids=decoder_input_ids,
-            decoder_attention_mask=decoder_attention_mask,
-            cache_position=cache_position,
-        )
-        lm_logits = decoder_output.logits
-
-        return Seq2SeqLMOutput(logits=lm_logits)
+        # default decoder pass
+        if input_features is None and encoder_outputs is None:
+            cross_attentions = []
+            for step in cache_position:
+                # skip step 0 if language_detection has been processed
+                if step == 0 and self.is_language_detected:
+                    cross_attentions.append(self.language_cross)
+                    self.is_language_detected = False
+                else:
+                    self.decoder_attention_mask[:, step] = 1
+                    decoder_output = self.decoder(
+                        decoder_input_ids=input_ids[:, step : step + 1].contiguous(),
+                        decoder_attention_mask=self.decoder_attention_mask,
+                        cache_position=step.to(torch.int32),
+                    )
+                    cross_attentions.append(decoder_output.cross_attentions)
+                    lm_logits = decoder_output.logits
+
+            if self.rbln_token_timestamps:
+                cross_attentions = torch.cat(cross_attentions, dim=-2)
+            else:
+                cross_attentions = None
+
+            return Seq2SeqLMOutput(logits=lm_logits, cross_attentions=cross_attentions)
+
+        # detect language pass
+        # https://github.com/huggingface/transformers/blob/174890280b340b89c5bfa092f6b4fb0e2dc2d7fc/src/transformers/models/whisper/generation_whisper.py#L1442
+        else:
+            if encoder_outputs is None:
+                self.encoder(input_features=input_features.contiguous())
+            self.decoder_attention_mask = torch.zeros(self.batch_size, self.dec_max_seq_len, dtype=torch.int64)
+            self.is_language_detected = True
+            self.decoder_attention_mask[:, 0] = 1
+            decoder_output = self.decoder(
+                decoder_input_ids=decoder_input_ids.contiguous(),
+                decoder_attention_mask=self.decoder_attention_mask,
+                cache_position=torch.zeros([], dtype=torch.int32),
+            )
+            lm_logits = decoder_output.logits
+            self.language_cross = decoder_output.cross_attentions
+            return Seq2SeqLMOutput(logits=lm_logits)

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

@@ -55,7 +55,6 @@ class _WhisperAttention(WhisperAttention):
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         attention_mask: Optional[torch.Tensor] = None,
         cache_position: Optional[torch.Tensor] = None,
-        **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, tgt_len, _ = hidden_states.size()
         is_cross_attention = key_value_states is not None
@@ -99,6 +98,7 @@ class _WhisperAttention(WhisperAttention):
         if attention_mask is not None:
             attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
         attn_weights = nn.functional.softmax(attn_weights, dim=-1)
 
         attn_output = torch.bmm(attn_weights, value_states)
@@ -109,7 +109,9 @@ class _WhisperAttention(WhisperAttention):
         attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
         attn_output = self.out_proj(attn_output)
 
-        return attn_output, None, present_key_value
+        attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+
+        return attn_output, attn_weights, present_key_value
 
 
 class _WhisperSdpaAttention(WhisperSdpaAttention):
@@ -186,6 +188,7 @@ class _WhisperDecoderLayer(WhisperDecoderLayer):
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         cache_position: Optional[torch.Tensor] = None,
         attn_impl: str = "eager",
+        output_attentions: bool = False,
     ) -> torch.Tensor:
         # Self Attention Block
         residual = hidden_states
@@ -205,14 +208,22 @@ class _WhisperDecoderLayer(WhisperDecoderLayer):
         residual = hidden_states
         hidden_states = self.encoder_attn_layer_norm(hidden_states)
         cross_attn_past_key_value = past_key_value[2:] if past_key_value is not None else None
-
-        hidden_states, _, cross_attn_present_key_value = ATTN_FORWARD_MAP[attn_impl](
-            self.encoder_attn,
-            hidden_states=hidden_states,
-            key_value_states=encoder_hidden_states,
-            past_key_value=cross_attn_past_key_value,
-            cache_position=cache_position,
-        )
+        if output_attentions:
+            hidden_states, cross_attn_weights, cross_attn_present_key_value = _WhisperAttention.forward(
+                self.encoder_attn,
+                hidden_states=hidden_states,
+                key_value_states=encoder_hidden_states,
+                past_key_value=cross_attn_past_key_value,
+                cache_position=cache_position,
+            )
+        else:
+            hidden_states, cross_attn_weights, cross_attn_present_key_value = ATTN_FORWARD_MAP[attn_impl](
+                self.encoder_attn,
+                hidden_states=hidden_states,
+                key_value_states=encoder_hidden_states,
+                past_key_value=cross_attn_past_key_value,
+                cache_position=cache_position,
+            )
         hidden_states = residual + hidden_states
         present_key_value = present_key_value + cross_attn_present_key_value
 
@@ -223,7 +234,7 @@ class _WhisperDecoderLayer(WhisperDecoderLayer):
         hidden_states = self.fc2(hidden_states)
         hidden_states = residual + hidden_states
 
-        return hidden_states, present_key_value
+        return hidden_states, present_key_value, cross_attn_weights
 
 
 class _WhisperPositionalEmbedding(WhisperPositionalEmbedding):
@@ -243,6 +254,7 @@ class _WhisperDecoder(WhisperDecoder):
         past_key_values: Optional[torch.Tensor] = None,
         cache_position: Optional[torch.Tensor] = None,
         attn_impl: str = "eager",
+        output_attentions: bool = False,
         **kwargs,
     ):
         input_shape = input_ids.size()
@@ -266,6 +278,7 @@ class _WhisperDecoder(WhisperDecoder):
             )
 
         next_decoder_cache = ()
+        all_cross_attentions = () if output_attentions else None
         # iterate decoder_layer
         for idx, decoder_layer in enumerate(self.layers):
             past_key_value = past_key_values[idx] if past_key_values is not None else None
@@ -277,10 +290,13 @@ class _WhisperDecoder(WhisperDecoder):
                 past_key_value=past_key_value,
                 cache_position=cache_position,
                 attn_impl=attn_impl,
+                output_attentions=output_attentions,
             )
             hidden_states = layer_outputs[0]
 
             next_decoder_cache += (layer_outputs[1],)
+            if output_attentions:
+                all_cross_attentions += (layer_outputs[2],)
 
         # layer_norm
         hidden_states = self.layer_norm(hidden_states)
@@ -288,17 +304,19 @@ class _WhisperDecoder(WhisperDecoder):
         return BaseModelOutputWithPastAndCrossAttentions(
             last_hidden_state=hidden_states,
             past_key_values=next_decoder_cache,
+            cross_attentions=all_cross_attentions,
         )
 
 
 class _WhisperDecoderWrapper(torch.nn.Module):
-    def __init__(self, model):
+    def __init__(self, model, output_attentions: bool = False):
         super().__init__()
         self.proj_out = model.proj_out
         self.config = model.config
         self.decoder = model.get_decoder()
         self.num_layers = self.config.decoder_layers
         self.attn_impl = self.config._attn_implementation
+        self.output_attentions = output_attentions
 
     def forward(
         self,
@@ -329,6 +347,7 @@ class _WhisperDecoderWrapper(torch.nn.Module):
             past_key_values=kv_cache,
             encoder_hidden_states=torch.tensor([1]),
             attn_impl=self.attn_impl,
+            output_attentions=self.output_attentions,
         )
         sequence_output = decoder_outputs[0]
         lm_logits = self.proj_out(sequence_output)
@@ -341,7 +360,12 @@ class _WhisperDecoderWrapper(torch.nn.Module):
             self_kv_cache.append(past_key_values[i][1])
         self_kv_cache = torch.stack(self_kv_cache, dim=0)
 
-        return lm_logits, self_kv_cache
+        if self.output_attentions:
+            # deocder's cross attention is used for token_timestamps
+            cross_attention = torch.stack(decoder_outputs[2], dim=0)
+            return lm_logits, self_kv_cache, cross_attention
+        else:
+            return lm_logits, self_kv_cache
 
 
 class _WhisperEncoderWrapper(torch.nn.Module):
@@ -363,6 +387,7 @@ class _WhisperEncoderWrapper(torch.nn.Module):
         input_features: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]:
         encoder_outputs = self.encoder(input_features=input_features)
+
         last_hidden_states = encoder_outputs[0]
 
         encoder_batch_size = input_features.shape[0]
@@ -388,13 +413,15 @@ class _WhisperEncoderWrapper(torch.nn.Module):
             encoder_hidden_states=last_hidden_states,
             past_key_values=dummy_past_key_value,
             attn_impl=self.attn_impl,
+            output_attentions=False,
         )
 
         first_past_kv = decoder_outputs[1]
 
-        encoder_kv = []
+        cross_kv = []
         for layer_out in first_past_kv:  # for layer
-            encoder_kv.append(torch.stack(layer_out[2:], dim=0))
-        encoder_kv = torch.stack(encoder_kv, dim=0)
+            cross_kv.append(layer_out[2])
+            cross_kv.append(layer_out[3])
+        cross_kv = torch.stack(cross_kv, dim=0)
 
-        return encoder_kv
+        return cross_kv

optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py

@@ -22,13 +22,13 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import torch
-from transformers import AutoModel, PretrainedConfig, PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
+from transformers import PretrainedConfig, PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
 
 from ....modeling_base import RBLNModel
-from ....modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 logger = logging.getLogger(__name__)
@@ -38,7 +38,6 @@ if TYPE_CHECKING:
 
 
 class RBLNXLMRobertaModel(RBLNModel):
-    auto_model_class = AutoModel  # feature extraction
     original_model_class = XLMRobertaModel
     original_config_class = XLMRobertaConfig
 
@@ -53,8 +52,7 @@ class RBLNXLMRobertaModel(RBLNModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
-        rbln_config_kwargs: Optional[Dict[str, Any]] = None,
-        rbln_constructor_kwargs: Optional[Dict[str, Any]] = None,
+        rbln_kwargs: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> "PreTrainedModel":
         model: "PreTrainedModel" = super().get_pytorch_model(
@@ -66,8 +64,7 @@ class RBLNXLMRobertaModel(RBLNModel):
             subfolder=subfolder,
             local_files_only=local_files_only,
             trust_remote_code=trust_remote_code,
-            rbln_config_kwargs=rbln_config_kwargs,
-            rbln_constructor_kwargs=rbln_constructor_kwargs,
+            rbln_kwargs=rbln_kwargs,
             library_name="transformers",
         )
 
@@ -78,10 +75,12 @@ class RBLNXLMRobertaModel(RBLNModel):
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_model_input_names: Optional[List[str]] = None,
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs={},
     ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
         max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
             model_config, "max_position_embeddings", None
         )
@@ -111,12 +110,15 @@ class RBLNXLMRobertaModel(RBLNModel):
             for model_input_name in rbln_model_input_names
         ]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        rbln_runtime_config.batch_size = rbln_batch_size
-
-        meta = {"rbln_max_seq_len": rbln_max_seq_len}
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
 
     def forward(
         self,