optimum-rbln 0.8.2a4__py3-none-any.whl → 0.8.2a5__py3-none-any.whl

optimum/rbln/transformers/models/qwen3/modeling_qwen3.py

@@ -12,32 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING
 
-import rebel
-import torch
-from rebel.compile_context import CompileContext
-from transformers import PretrainedConfig, PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.modeling_utils import no_init_weights
+from transformers import PretrainedConfig
 
-from ....configuration_utils import RBLNCompileConfig
-from ....modeling import RBLNModel
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
-from ..decoderonly.modeling_decoderonly import set_default_values, validate_attention_method
-from .configuration_qwen3 import RBLNQwen3ModelConfig
-from .qwen3_architecture import Qwen3ModelWrapper, Qwen3Wrapper
+from ...models.decoderonly import (
+    RBLNDecoderOnlyModel,
+    RBLNDecoderOnlyModelForCausalLM,
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
+from .qwen3_architecture import Qwen3Wrapper
 
 
 logger = logging.get_logger(__name__)
 
 if TYPE_CHECKING:
     from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
         PretrainedConfig,
     )
 
@@ -63,315 +54,6 @@ class RBLNQwen3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         return super().forward(*args, **kwargs)
 
 
-class RBLNQwen3Model(RBLNModel):
-    _decoder_wrapper_cls = Qwen3ModelWrapper
+class RBLNQwen3Model(RBLNDecoderOnlyModel):
+    _decoder_wrapper_cls = Qwen3Wrapper
     _use_rotary_emb = True
-
-    def __post_init__(self, **kwargs):
-        artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
-        self.embed_tokens = self._create_embedding_layer()
-        self.embed_tokens.load_state_dict(artifacts["embed_tokens"])
-        self.block_tables = torch.arange(
-            self.rbln_config.max_seq_len / self.rbln_config.kvcache_block_size, dtype=torch.int16
-        )
-        self.causal_mask = 1 - torch.triu(
-            torch.ones(1, 1, self.rbln_config.prefill_chunk_size, self.rbln_config.prefill_chunk_size), diagonal=1
-        )
-
-    @classmethod
-    def save_torch_artifacts(
-        cls,
-        model: PreTrainedModel,
-        save_dir_path: Path,
-        subfolder: str,
-        rbln_config: RBLNQwen3ModelConfig,
-    ):
-        save_dict = {}
-        save_dict["embed_tokens"] = model.get_input_embeddings().state_dict()
-        torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
-
-    def _create_embedding_layer(self):
-        with no_init_weights():
-            embed_tokens = torch.nn.Embedding(
-                self.config.vocab_size,
-                self.config.hidden_size,
-                self.config.pad_token_id,
-            )
-        return embed_tokens
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    @classmethod
-    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: "RBLNQwen3ModelConfig"):
-        wrapper_cfg = {
-            "max_seq_len": rbln_config.max_seq_len,
-            "attn_impl": rbln_config.attn_impl,
-            "kvcache_partition_len": rbln_config.kvcache_partition_len,
-            "kvcache_block_size": rbln_config.kvcache_block_size,
-            "use_rotary_emb": cls._use_rotary_emb,
-            "use_attention_mask": rbln_config.use_attention_mask,
-            "cache_impl": rbln_config.cache_impl,
-            "sliding_window": rbln_config.sliding_window,
-            "sliding_window_layers": rbln_config.sliding_window_layers,
-        }
-        return cls._decoder_wrapper_cls(model, **wrapper_cfg).eval()
-
-    @classmethod
-    @torch.inference_mode()
-    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNQwen3ModelConfig):
-        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
-
-        rbln_compile_configs = rbln_config.compile_cfgs
-        prefill_compile_config = rbln_compile_configs[0]
-
-        context = CompileContext(use_weight_sharing=False)
-
-        meta_tensor_names = [name for name, _, _ in prefill_compile_config.input_info if "past_key_values" in name]
-        prefill_example_inputs = prefill_compile_config.get_dummy_inputs(fill=0, meta_tensor_names=meta_tensor_names)
-
-        static_tensors = {}
-        for (name, _, _), tensor in zip(prefill_compile_config.input_info, prefill_example_inputs):
-            if "past_key_values" in name:
-                static_tensors[name] = tensor
-                context.mark_static_address(tensor)
-
-        def compile_model(wrapped_model, compile_config, example_inputs, compile_context):
-            try:
-                original_linear = torch.nn.functional.linear
-                torch.nn.functional.linear = torch.ops.rbln_custom_ops.linear
-                compiled_model = RBLNModel.compile(
-                    wrapped_model,
-                    compile_config,
-                    example_inputs=example_inputs,
-                    compile_context=compile_context,
-                    create_runtimes=rbln_config.create_runtimes,
-                    device=rbln_config.device,
-                )
-                return compiled_model
-            finally:
-                torch.nn.functional.linear = original_linear
-
-        wrapped_model.phase = "prefill"
-        compiled_prefill = compile_model(wrapped_model, prefill_compile_config, prefill_example_inputs, context)
-
-        compiled_models = {"prefill": compiled_prefill}
-        return compiled_models
-
-    @classmethod
-    def get_input_info(
-        cls,
-        batch_size: int,
-        query_length: int,
-        rbln_config: RBLNQwen3ModelConfig,
-        model_config: PretrainedConfig,
-    ):
-        input_info = RBLNDecoderOnlyModelForCausalLM.get_input_info(
-            batch_size,
-            query_length,
-            rbln_config=rbln_config,
-            model_config=model_config,
-        )
-
-        if rbln_config.sliding_window is None:
-            # remove query position
-            input_info.pop(3)
-
-        return input_info
-
-    @classmethod
-    def _update_rbln_config(
-        cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
-        model: Optional["PreTrainedModel"] = None,
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[RBLNQwen3ModelConfig] = None,
-    ) -> RBLNQwen3ModelConfig:
-        if rbln_config.max_seq_len is None:
-            rbln_config.max_seq_len = getattr(model_config, "max_position_embeddings", None) or getattr(
-                model_config, "n_positions", None
-            )
-        if rbln_config.max_seq_len is None:
-            raise ValueError("`max_seq_len` should be specified.")
-
-        rbln_config.attn_impl, rbln_config.kvcache_partition_len, rbln_config.kvcache_block_size = set_default_values(
-            attn_impl=rbln_config.attn_impl,
-            kvcache_partition_len=rbln_config.kvcache_partition_len,
-            kvcache_block_size=rbln_config.kvcache_block_size,
-            max_seq_len=rbln_config.max_seq_len,
-        )
-
-        validate_attention_method(
-            attn_impl=rbln_config.attn_impl,
-            kvcache_partition_len=rbln_config.kvcache_partition_len,
-            kvcache_block_size=rbln_config.kvcache_block_size,
-            max_seq_len=rbln_config.max_seq_len,
-        )
-
-        # only compile prefill cb -> always batch_size 1
-        required_num_blocks = rbln_config.max_seq_len // rbln_config.kvcache_block_size
-        max_num_blocks = required_num_blocks
-
-        if rbln_config.attn_impl == "flash_attn":
-            estimated_max_num_blocks = RBLNDecoderOnlyModelForCausalLM.get_maximum_num_blocks(
-                config=model_config,
-                tensor_parallel_size=rbln_config.tensor_parallel_size or 1,
-                kvcache_block_size=rbln_config.kvcache_block_size,
-                nbits_per_param=16 if not rbln_config.quantization else 4,
-                n_model_params=sum(p.numel() for p in model.parameters()),
-                num_runtimes=1 + len(rbln_config.decoder_batch_sizes),
-            )
-
-            max_num_blocks = min(max_num_blocks, estimated_max_num_blocks)
-
-            flash_min_blocks = rbln_config.max_seq_len // rbln_config.kvcache_block_size + 1
-            if max_num_blocks < flash_min_blocks:
-                max_num_blocks = flash_min_blocks
-
-        if rbln_config.kvcache_num_blocks is None:
-            rbln_config.kvcache_num_blocks = max_num_blocks
-
-        prefill_input_info = cls.get_input_info(
-            batch_size=1,
-            query_length=rbln_config.prefill_chunk_size,
-            rbln_config=rbln_config,
-            model_config=model_config,
-        )
-
-        prefill_compile_config = RBLNCompileConfig(compiled_model_name="prefill", input_info=prefill_input_info)
-        rbln_config.set_compile_cfgs([prefill_compile_config])
-
-        return rbln_config
-
-    @classmethod
-    def _create_runtimes(
-        cls,
-        compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_config: RBLNQwen3ModelConfig,
-    ) -> List[rebel.Runtime]:
-        expected_model_names = ["prefill"]
-        if any(model_name not in rbln_config.device_map for model_name in expected_model_names):
-            cls._raise_missing_compiled_file_error(expected_model_names)
-
-        return [
-            rebel.Runtime(
-                compiled_models[0],
-                tensor_type="pt",
-                device=rbln_config.device_map["prefill"],
-                activate_profiler=rbln_config.activate_profiler,
-            ),
-        ]
-
-    def _preprocess_chunked_prefill(
-        self,
-        inputs: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_embed: Optional[torch.Tensor] = None,
-    ):
-        # valid sequence length of inputs_embeds
-        query_length = inputs.shape[1] if attention_mask is None else torch.sum(attention_mask.view(-1)).item()
-
-        # extract valid inputs
-        inputs = inputs[:, attention_mask.bool()] if attention_mask is not None else inputs
-        if position_embed is not None:
-            position_embed = (
-                position_embed[:, :, :, attention_mask.bool(), :] if attention_mask is not None else position_embed
-            )
-
-        if self.rbln_config.use_attention_mask:
-            chunked_attention_mask = (
-                torch.zeros(
-                    1, 1, self.rbln_config.prefill_chunk_size, self.rbln_config.max_seq_len, dtype=torch.float32
-                )
-                if self.rbln_config.use_attention_mask
-                else None
-            )
-        else:
-            chunked_attention_mask = None
-
-        # padding for chunked prefill
-        padding_size = (
-            self.rbln_config.prefill_chunk_size - (query_length % self.rbln_config.prefill_chunk_size)
-        ) % self.rbln_config.prefill_chunk_size
-        padded_len = query_length + padding_size
-
-        inputs = torch.nn.functional.pad(inputs, (0, padding_size))
-        position_embed = (
-            None if position_embed is None else torch.nn.functional.pad(position_embed, (0, 0, 0, padding_size))
-        )
-        cache_position = torch.arange(padded_len, dtype=torch.int32).unsqueeze(0)
-
-        return inputs, chunked_attention_mask, position_embed, cache_position, query_length
-
-    def _chunked_prefill_forward(
-        self,
-        inputs_embeds: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_embed: Optional[torch.Tensor] = None,
-    ):
-        padded_input, chunked_attention_mask, padded_position_embed, cache_position, query_length = (
-            self._preprocess_chunked_prefill(inputs_embeds, attention_mask, position_embed)
-        )
-
-        # chunked prefill
-        last_hidden_states = []
-        for step in range(0, query_length, self.rbln_config.prefill_chunk_size):
-            # Extract the current chunk of inputs and cache positions
-            input_chunk = padded_input[:, step : step + self.rbln_config.prefill_chunk_size]
-            cache_pos_chunk = cache_position[:, step : step + self.rbln_config.prefill_chunk_size]
-
-            model_args = {
-                "input_ids": input_chunk,
-                "cache_position": cache_pos_chunk,
-                "block_tables": self.block_tables,
-            }
-
-            if chunked_attention_mask is not None:
-                if step >= self.rbln_config.prefill_chunk_size:
-                    chunked_attention_mask[:, :, :, step - self.rbln_config.prefill_chunk_size : step] = 1
-                chunked_attention_mask[:, :, :, step : step + self.rbln_config.prefill_chunk_size] = self.causal_mask
-                model_args["attention_mask"] = chunked_attention_mask
-
-            last_hidden_states_chunk = self.model[0](**model_args)
-            last_hidden_states.append(last_hidden_states_chunk)
-
-        last_hidden_states = torch.concat(last_hidden_states, dim=-2)[:, :query_length]
-
-        return self._postprocess_chunked_prefill(last_hidden_states, attention_mask)
-
-    def _postprocess_chunked_prefill(
-        self, last_hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
-    ):
-        # index copy for attention mask
-        if attention_mask is not None:
-            new_last_hidden_states = torch.full(
-                (1, attention_mask.shape[-1], last_hidden_states.shape[-1]),
-                fill_value=1e-10,
-                dtype=last_hidden_states.dtype,
-            )
-            mask_indices = torch.nonzero(attention_mask, as_tuple=True)[0]
-            new_last_hidden_states.index_copy_(dim=-2, index=mask_indices, source=last_hidden_states)
-        else:
-            new_last_hidden_states = last_hidden_states
-        return new_last_hidden_states
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.LongTensor] = None,
-        position_embed: Optional[torch.Tensor] = None,
-        **kwargs,
-    ):
-        inputs = inputs_embeds if inputs_embeds is not None else input_ids
-        batch_size = inputs.shape[0]
-        all_last_hidden_states = []
-        for b_idx in range(batch_size):
-            last_hidden_states = self._chunked_prefill_forward(
-                inputs[b_idx : b_idx + 1],
-                attention_mask[b_idx] if attention_mask is not None else None,
-                position_embed[b_idx : b_idx + 1] if position_embed is not None else None,
-            )
-            all_last_hidden_states.append(last_hidden_states)
-
-        return BaseModelOutputWithPast(last_hidden_state=torch.concat(all_last_hidden_states, dim=0))

optimum/rbln/transformers/models/qwen3/qwen3_architecture.py

@@ -12,15 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
-import torch.nn as nn
-from transformers import PreTrainedModel
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
-    DecoderOnlyLayer,
     DecoderOnlyWrapper,
-    RotaryEmbedding,
 )
 
 
@@ -37,239 +32,3 @@ class Qwen3Attention(DecoderOnlyAttention):
         self.o_proj = self._original_mod.o_proj
         self.q_norm = self._original_mod.q_norm
         self.k_norm = self._original_mod.k_norm
-
-
-class Qwen3ModelWrapper(nn.Module):
-    def __init__(
-        self,
-        model,
-        attn_impl=None,
-        use_inputs_embeds=None,
-        use_attention_mask=None,
-        use_rotary_emb=None,
-        cache_impl=None,
-        kvcache_partition_len=None,
-        max_seq_len=None,
-        kvcache_block_size=None,
-        sliding_window=None,
-        sliding_window_layers=None,
-    ):
-        super().__init__()
-        self.config = model.config
-
-        if use_rotary_emb:
-            rotary_embs = self.get_rotary_emb(max_seq_len=max_seq_len)
-            if isinstance(rotary_embs, tuple):
-                self.rotary_emb_global, self.rotary_emb_local = rotary_embs
-            else:
-                self.rotary_emb = rotary_embs
-        else:
-            self.rotary_emb = None
-
-        self._original_mod = model
-        self.use_inputs_embeds = use_inputs_embeds
-        self.attn_impl = attn_impl
-        self.cache_impl = cache_impl
-        self.use_attention_mask = use_attention_mask
-        self.kvcache_partition_len = kvcache_partition_len
-        self.kvcache_block_size = kvcache_block_size
-        self.max_seq_len = max_seq_len
-        self.sliding_window = sliding_window
-        self.sliding_window_layers = sliding_window_layers
-        self.model = self.convert_to_rbln_model(model)
-
-    def get_rotary_emb(self, max_seq_len):
-        return RotaryEmbedding(config=self.config, max_seq_len_cached=max_seq_len)
-
-    def convert_to_rbln_model(self, base_model: PreTrainedModel):
-        for layer_idx, layer in enumerate(base_model.layers):
-            is_sliding = layer_idx in self.sliding_window_layers
-            new_self_attn = Qwen3Attention(
-                layer.self_attn,
-                self.use_attention_mask if not is_sliding else True,
-                use_position_ids=None,
-                kvcache_block_size=self.sliding_window
-                if layer_idx in self.sliding_window_layers
-                else self.kvcache_block_size,
-                is_sliding=is_sliding,
-                attn_impl=self.attn_impl if not is_sliding else "eager",
-                kvcache_partition_len=self.kvcache_partition_len,
-            )
-            base_model.layers[layer_idx] = DecoderOnlyLayer(layer, new_self_attn)
-
-        return base_model
-
-    @property
-    def hidden_multiplier(self):
-        return 1
-
-    def get_last_layernorm(self) -> nn.LayerNorm:
-        return self._original_mod.norm
-
-    def get_embedding(self) -> nn.Embedding:
-        return self._original_mod.embed_tokens
-
-    def get_pos_embedding(self) -> nn.Embedding:
-        raise NotImplementedError(
-            "The 'get_pos_embedding' method is not implemented. Please define this method in a subclass."
-        )
-
-    def convert_sequence_positions_for_flash_attn(self, seq_positions, max_seq_len):
-        if self.attn_impl not in ["flash_attn"]:
-            raise NotImplementedError(f"Unknown attn_impl ({self.attn_impl}).")
-        partition_len = self.kvcache_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_local_cache_positions(self, position_ids, query_position):
-        max_cache_len = self.model.config.sliding_window
-        valid_input_len = 1 if query_position is None else query_position + 1
-        cache_seq_len = torch.clamp(position_ids, max=max_cache_len)[:, :1]  # past seen tokens
-        cache_offset = (
-            torch.clamp(position_ids, max=max_cache_len)[:, :1] + valid_input_len
-        )  # cache offset for next steps
-
-        return cache_seq_len, cache_offset
-
-    def prepare_forward_args(self, *args):
-        args = list(args)
-        input_ids = None if self.use_inputs_embeds else args.pop(0)
-        inputs_embeds = args.pop(0) if self.use_inputs_embeds else None
-        cache_position = args.pop(0)
-        global_block_tables = args.pop(0) if self.cache_impl in ["hybrid", "static"] else None
-        local_block_tables = args.pop(0) if self.cache_impl in ["hybrid", "sliding_window"] else None
-        query_position = args.pop(0) if self.sliding_window else None
-        attention_mask = args.pop(0) if self.use_attention_mask else None
-        position_ids = None
-        past_key_values = args
-
-        if len(past_key_values) != 2 * self.config.num_hidden_layers:
-            raise ValueError(
-                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.config.num_hidden_layers}"
-            )
-
-        # [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
-
-        if hasattr(self, "rotary_emb_global") and hasattr(self, "rotary_emb_local"):
-            rotary_emb = (self.rotary_emb_global, self.rotary_emb_local)
-        else:
-            rotary_emb = self.rotary_emb
-
-        return (
-            input_ids,
-            inputs_embeds,
-            cache_position,
-            global_block_tables,
-            local_block_tables,
-            attention_mask,
-            position_ids,
-            query_position,
-            past_key_values,
-            rotary_emb,
-        )
-
-    def forward(self, *args):
-        (
-            input_ids,
-            inputs_embeds,
-            cache_position,
-            global_block_tables,
-            local_block_tables,
-            attention_mask,
-            position_ids,
-            query_position,
-            past_key_values,
-            rotary_emb,
-        ) = self.prepare_forward_args(*args)
-
-        # 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"
-            )
-
-        # embed positions
-        if inputs_embeds is None:
-            inputs_embeds = self.get_embedding()(input_ids)
-
-        hidden_states = inputs_embeds * self.hidden_multiplier
-
-        # get cos,sin vector if needed
-        position_ids = position_ids if position_ids is not None else cache_position
-        if rotary_emb is not None:
-            if isinstance(rotary_emb, torch.Tensor):
-                cos = rotary_emb[0]
-                sin = rotary_emb[1]
-            else:
-                cos, sin = rotary_emb(hidden_states, self.max_seq_len)  # dtype carrier, max_seq_len
-                cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
-        else:
-            batch_size = inputs_embeds.shape[0]
-            if position_ids.shape[0] > 1:
-                position_embeds = []
-                for b_idx in range(batch_size):
-                    position_embed = self.get_pos_embedding()(position_ids[b_idx])
-                    position_embeds.append(position_embed)
-
-                position_embeds = torch.cat(position_embeds, dim=0).unsqueeze(1)
-            else:
-                position_embeds = self.get_pos_embedding()(position_ids)
-            hidden_states = hidden_states + position_embeds
-            cos, sin = None, None
-
-        # Get sequence positions for flash attention
-        if self.attn_impl == "flash_attn":
-            seq_positions = cache_position[:, 0]
-            seq_positions = self.convert_sequence_positions_for_flash_attn(
-                seq_positions=seq_positions, max_seq_len=self.max_seq_len
-            )
-        else:
-            seq_positions = cache_position[:, :1]
-
-        # Get local cache positions for sliding window layers
-        if len(self.sliding_window_layers) > 0:
-            sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
-
-        for layer_idx, layer in enumerate(self.model.layers):
-            is_sliding = True if layer_idx in self.sliding_window_layers else False
-            hidden_states = layer(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                seq_positions=sliding_cache_pos if is_sliding else seq_positions,
-                past_key_values=past_key_values,
-                cos=cos,
-                sin=sin,
-                block_tables=local_block_tables if is_sliding else global_block_tables,
-            )
-
-        hidden_states = self.get_last_layernorm()(hidden_states)
-        return hidden_states
-
-
-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(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[cache_position].unsqueeze(unsqueeze_dim)
-        sin = sin[cache_position].unsqueeze(unsqueeze_dim)
-
-    return cos, sin

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

@@ -14,8 +14,6 @@
 
 from typing import Any, Dict, Optional
 
-import rebel
-
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
 
@@ -39,7 +37,6 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
             enc_max_seq_len (Optional[int]): Maximum sequence length for the encoder.
             dec_max_seq_len (Optional[int]): Maximum sequence length for the decoder.
             use_attention_mask (Optional[bool]): Whether to use attention masks during inference.
-                This is automatically set to True for RBLN-CA02 devices.
             pad_token_id (Optional[int]): The ID of the padding token in the vocabulary.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
@@ -55,12 +52,5 @@ class RBLNModelForSeq2SeqLMConfig(RBLNModelConfig):
         self.dec_max_seq_len = dec_max_seq_len
 
         self.use_attention_mask = use_attention_mask
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            if self.use_attention_mask is False:
-                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
-            self.use_attention_mask = True
-        else:
-            self.use_attention_mask = self.use_attention_mask or False
 
         self.pad_token_id = pad_token_id

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

@@ -14,8 +14,6 @@
 
 from typing import Any, Dict
 
-import rebel
-
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
 
@@ -45,7 +43,6 @@ class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
             batch_size (int, optional): The batch size for inference. Defaults to 1.
             token_timestamps (bool, optional): Whether to output token timestamps during generation. Defaults to False.
             use_attention_mask (bool, optional): Whether to use attention masks during inference. This is automatically
-                set to True for RBLN-CA02 devices.
             enc_max_seq_len (int, optional): Maximum sequence length for the encoder.
             dec_max_seq_len (int, optional): Maximum sequence length for the decoder.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
@@ -64,10 +61,4 @@ class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
         self.dec_max_seq_len = dec_max_seq_len
 
         self.use_attention_mask = use_attention_mask
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            if self.use_attention_mask is False:
-                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
-            self.use_attention_mask = True
-        else:
-            self.use_attention_mask = self.use_attention_mask or False
+        self.use_attention_mask = self.use_attention_mask or False

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

@@ -73,6 +73,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         decoder_input_ids: torch.Tensor = None,
         decoder_attention_mask: torch.Tensor = None,
         cache_position: torch.Tensor = None,
+        block_tables: torch.Tensor = None,
     ):
         inputs_bsz = decoder_input_ids.shape[0]
         padded_bsz = self.batch_size - inputs_bsz
@@ -89,11 +90,14 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
                     )
                 decoder_attention_mask[b_idx, : decoding_step + 1] = 1
 
+        if block_tables is None:
+            block_tables = self.default_block_tables
+
         outputs = super().forward(
             decoder_input_ids,
             decoder_attention_mask if self.use_attention_mask else None,
             cache_position,
-            block_tables=self.default_block_tables,
+            block_tables=block_tables,
         )
 
         if isinstance(outputs, torch.Tensor):