optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/transformers/models/qwen2_vl/qwen2_vl_architecture.py

@@ -0,0 +1,165 @@
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from ..decoderonly.decoderonly_architecture import (
+    DecoderOnlyWrapper,
+    apply_rotary_pos_emb,
+)
+
+
+class Qwen2VisionTransformerWrapper(nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self._original_mod = model
+        self.merger = model.merger
+        self.blocks = self.wrap_vision_blocks(model.blocks)
+
+    def wrap_vision_blocks(self, blocks: torch.nn.ModuleList):
+        wrapped_blocks = []
+        for i, block in enumerate(blocks):
+            wrapped_blocks.append(Qwen2VLVisionBlock(block))
+        return nn.ModuleList(wrapped_blocks)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        full_attn_masks: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ):
+        full_attn_masks = (1 - full_attn_masks) * torch.finfo(torch.float32).min
+
+        for block in self.blocks:
+            hidden_states = block(hidden_states, full_attn_masks, [cos, sin])
+
+        return self.merger(hidden_states)
+
+
+class Qwen2VLVisionBlock(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self._origin_model = model
+        self.norm1 = model.norm1
+        self.norm2 = model.norm2
+
+        self.attn = VisionAttention(model.attn)
+        self.mlp = model.mlp
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            attn_masks,
+            position_embeddings,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+
+class VisionAttention(nn.Module):
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+        self._origin_model = model
+        self.num_heads = model.num_heads
+        self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
+        self.qkv = model.qkv
+        self.proj = model.proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_masks: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        hidden_states = hidden_states.unsqueeze(0)
+        q, k, v = (
+            self.qkv(hidden_states).reshape(1, seq_length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+        )
+
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)
+        attn_weights = attn_weights + attn_masks
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(1, seq_length, -1)
+        attn_output = self.proj(attn_output).squeeze(0)
+
+        return attn_output
+
+
+class Qwen2VL_LanguageModelWrapper(DecoderOnlyWrapper):
+    def prepare_forward_args(self, *args):
+        args = list(args)
+        input_ids = None if self.rbln_config.use_inputs_embeds else args.pop(0)
+        inputs_embeds = args.pop(0) if self.rbln_config.use_inputs_embeds else None
+        cache_position = args.pop(0)
+        global_block_tables = args.pop(0)
+        local_block_tables = None
+        position_embeds = args.pop(0)
+        query_position = args.pop(0) if self.phase == "prefill" else None
+        position_ids = None
+        attention_mask = args.pop(0) if self.rbln_config.use_attention_mask else None
+        lora_int_id = args.pop(0) if self.rbln_config.lora_config else None
+        past_key_values = args
+
+        if len(past_key_values) != 2 * self.num_hidden_layers:
+            raise ValueError(
+                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.num_hidden_layers}"
+            )
+
+        # [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
+
+        return (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            global_block_tables,
+            local_block_tables,
+            query_position,
+            attention_mask,
+            position_ids,
+            lora_int_id,
+            past_key_values,
+            position_embeds,
+        )
+
+    def convert_to_rbln_class(self, model: PreTrainedModel, max_seq_len: int):
+        new_layers = []
+
+        for layer_idx, layer in enumerate(model.model.language_model.layers):
+            is_sliding = layer_idx in self.rbln_config.sliding_window_layers
+            new_self_attn = self.get_rbln_attn_class()(
+                self.get_attn_layer(layer), self.rbln_config, is_sliding=is_sliding
+            )
+            new_layer = self.get_rbln_layer_class()(layer, new_self_attn)
+            new_layers.append(new_layer)
+
+        new_model = self.get_rbln_model_class()(
+            model.model.language_model,
+            new_layers,
+            self.rbln_config,
+            use_learned_pos_emb=self.__class__._use_learned_pos_emb,
+        )
+
+        new_model = self.get_rbln_causal_lm_class()(model.model, new_model)
+        return new_model

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -43,7 +43,7 @@ class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     """
 
 
-class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelConfig):
     """
     Configuration class for RBLN Qwen3 models.
 

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

@@ -12,37 +12,76 @@
 # 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,
-    )
+    from transformers import PretrainedConfig
 
 
 class RBLNQwen3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Qwen3 Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    A class to convert and run pre-trained transformers based Qwen3ForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen3ForCausalLM model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+    **Configuration:**
+    This model uses [`RBLNQwen3ForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen3ForCausalLMConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen3ForCausalLMConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen3ForCausalLM
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen3ForCausalLM.from_pretrained(
+            "Qwen/Qwen3-4B",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 40_960,
+            "tensor_parallel_size": 4,
+            "kvcache_partition_len": 8192,
+        }
+        model = RBLNQwen3ForCausalLM.from_pretrained(
+            "Qwen/Qwen3-4B",
+            export=True,
+            rbln_config=rbln_config
+        )
+        # Using a RBLNQwen3ForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNQwen3ForCausalLMConfig
+        config = RBLNQwen3ForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=40_960,
+            tensor_parallel_size=4,
+            kvcache_partition_len=8192,
+        )
+        model = RBLNQwen3ForCausalLM.from_pretrained(
+            "Qwen/Qwen3-4B",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
     _decoder_wrapper_cls = Qwen3Wrapper
 
     @classmethod
@@ -63,315 +102,32 @@ class RBLNQwen3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
         return super().forward(*args, **kwargs)
 
 
-class RBLNQwen3Model(RBLNModel):
-    _decoder_wrapper_cls = Qwen3ModelWrapper
-    _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)
+class RBLNQwen3Model(RBLNDecoderOnlyModel):
+    """
+    The bare Qwen3 Model outputting raw hidden-states without any specific head on top.
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    A class to convert and run pre-trained transformers based Qwen3Model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers Qwen3Model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+    **Configuration:**
+    This model uses [`RBLNQwen3ModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNQwen3ModelConfig`] or a dictionary conforming to its structure.
+    See the [`RBLNQwen3ModelConfig`] class for all available configuration options.
+    Examples:
+        ```python
+        from optimum.rbln import RBLNQwen3Model
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNQwen3Model.from_pretrained(
+            "Qwen/Qwen3-Embedding-4B",
+            export=True,
+            rbln_batch_size=1,
+            rbln_max_seq_len=40_960,
+            rbln_tensor_parallel_size=4,
+            rbln_kvcache_partition_len=8192,
         )
+    """
 
-        # 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))
+    _decoder_wrapper_cls = Qwen3Wrapper
+    _use_rotary_emb = True