optimum-rbln 0.8.1rc1__py3-none-any.whl → 0.8.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of optimum-rbln might be problematic. Click here for more details.

Files changed (119) hide show
  1. optimum/rbln/__init__.py +58 -9
  2. optimum/rbln/__version__.py +2 -2
  3. optimum/rbln/configuration_utils.py +24 -5
  4. optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py +2 -2
  5. optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl_cosmos.py +2 -2
  6. optimum/rbln/diffusers/configurations/models/configuration_controlnet.py +2 -2
  7. optimum/rbln/diffusers/configurations/models/configuration_prior_transformer.py +2 -2
  8. optimum/rbln/diffusers/configurations/models/configuration_transformer_cosmos.py +5 -2
  9. optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py +2 -2
  10. optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py +2 -2
  11. optimum/rbln/diffusers/configurations/models/configuration_vq_model.py +2 -2
  12. optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py +3 -3
  13. optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py +2 -2
  14. optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py +4 -4
  15. optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py +2 -2
  16. optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py +2 -2
  17. optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py +2 -2
  18. optimum/rbln/diffusers/modeling_diffusers.py +4 -5
  19. optimum/rbln/diffusers/models/__init__.py +3 -13
  20. optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py +1 -0
  21. optimum/rbln/diffusers/models/autoencoders/autoencoder_kl_cosmos.py +1 -0
  22. optimum/rbln/diffusers/models/autoencoders/vq_model.py +1 -0
  23. optimum/rbln/diffusers/models/transformers/transformer_cosmos.py +1 -1
  24. optimum/rbln/diffusers/pipelines/__init__.py +1 -5
  25. optimum/rbln/diffusers/pipelines/cosmos/configuration_cosmos_guardrail.py +12 -4
  26. optimum/rbln/diffusers/pipelines/cosmos/cosmos_guardrail.py +4 -28
  27. optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py +1 -1
  28. optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py +1 -1
  29. optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +1 -6
  30. optimum/rbln/modeling.py +4 -5
  31. optimum/rbln/modeling_base.py +18 -14
  32. optimum/rbln/ops/kv_cache_update.py +5 -0
  33. optimum/rbln/ops/linear.py +7 -0
  34. optimum/rbln/transformers/__init__.py +60 -0
  35. optimum/rbln/transformers/configuration_generic.py +4 -4
  36. optimum/rbln/transformers/modeling_attention_utils.py +252 -0
  37. optimum/rbln/transformers/modeling_generic.py +1 -4
  38. optimum/rbln/transformers/models/__init__.py +45 -30
  39. optimum/rbln/transformers/models/bart/bart_architecture.py +2 -7
  40. optimum/rbln/transformers/models/blip_2/configuration_blip_2.py +2 -2
  41. optimum/rbln/transformers/models/blip_2/modeling_blip_2.py +1 -5
  42. optimum/rbln/transformers/models/clip/configuration_clip.py +14 -3
  43. optimum/rbln/transformers/models/clip/modeling_clip.py +123 -28
  44. optimum/rbln/transformers/models/colpali/colpali_architecture.py +1 -4
  45. optimum/rbln/transformers/models/colpali/configuration_colpali.py +2 -2
  46. optimum/rbln/transformers/models/colpali/modeling_colpali.py +2 -10
  47. optimum/rbln/transformers/models/decoderonly/__init__.py +2 -2
  48. optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py +214 -45
  49. optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +323 -454
  50. optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +579 -362
  51. optimum/rbln/transformers/models/exaone/exaone_architecture.py +17 -42
  52. optimum/rbln/transformers/models/gemma/__init__.py +2 -2
  53. optimum/rbln/transformers/models/gemma/configuration_gemma.py +9 -1
  54. optimum/rbln/transformers/models/gemma/gemma_architecture.py +3 -44
  55. optimum/rbln/transformers/models/gemma/modeling_gemma.py +22 -1
  56. optimum/rbln/transformers/models/gemma3/configuration_gemma3.py +21 -9
  57. optimum/rbln/transformers/models/gemma3/gemma3_architecture.py +9 -63
  58. optimum/rbln/transformers/models/gemma3/modeling_gemma3.py +200 -292
  59. optimum/rbln/transformers/models/gpt2/__init__.py +2 -2
  60. optimum/rbln/transformers/models/gpt2/configuration_gpt2.py +31 -3
  61. optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +19 -24
  62. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +18 -1
  63. optimum/rbln/transformers/models/idefics3/configuration_idefics3.py +2 -2
  64. optimum/rbln/transformers/models/idefics3/modeling_idefics3.py +3 -9
  65. optimum/rbln/transformers/models/llama/__init__.py +2 -2
  66. optimum/rbln/transformers/models/llama/configuration_llama.py +9 -1
  67. optimum/rbln/transformers/models/llama/modeling_llama.py +22 -1
  68. optimum/rbln/transformers/models/llava/__init__.py +16 -0
  69. optimum/rbln/transformers/models/llava/configuration_llava.py +54 -0
  70. optimum/rbln/transformers/models/llava/modeling_llava.py +419 -0
  71. optimum/rbln/transformers/models/llava_next/configuration_llava_next.py +20 -3
  72. optimum/rbln/transformers/models/llava_next/modeling_llava_next.py +6 -16
  73. optimum/rbln/transformers/models/midm/midm_architecture.py +14 -22
  74. optimum/rbln/transformers/models/mistral/__init__.py +2 -2
  75. optimum/rbln/transformers/models/mistral/configuration_mistral.py +9 -1
  76. optimum/rbln/transformers/models/mistral/mistral_architecture.py +1 -1
  77. optimum/rbln/transformers/models/mistral/modeling_mistral.py +26 -3
  78. optimum/rbln/transformers/models/opt/__init__.py +2 -2
  79. optimum/rbln/transformers/models/opt/configuration_opt.py +8 -1
  80. optimum/rbln/transformers/models/opt/modeling_opt.py +41 -1
  81. optimum/rbln/transformers/models/opt/opt_architecture.py +16 -25
  82. optimum/rbln/transformers/models/pegasus/__init__.py +17 -0
  83. optimum/rbln/transformers/models/pegasus/configuration_pegasus.py +34 -0
  84. optimum/rbln/transformers/models/pegasus/modeling_pegasus.py +69 -0
  85. optimum/rbln/transformers/models/pegasus/pegasus_architecture.py +161 -0
  86. optimum/rbln/transformers/models/phi/__init__.py +2 -2
  87. optimum/rbln/transformers/models/phi/configuration_phi.py +9 -1
  88. optimum/rbln/transformers/models/phi/modeling_phi.py +10 -1
  89. optimum/rbln/transformers/models/phi/phi_architecture.py +16 -22
  90. optimum/rbln/transformers/models/pixtral/__init__.py +16 -0
  91. optimum/rbln/transformers/models/pixtral/configuration_pixtral.py +43 -0
  92. optimum/rbln/transformers/models/pixtral/modeling_pixtral.py +315 -0
  93. optimum/rbln/transformers/models/pixtral/pixtral_architecture.py +73 -0
  94. optimum/rbln/transformers/models/qwen2/__init__.py +2 -2
  95. optimum/rbln/transformers/models/qwen2/configuration_qwen2.py +9 -1
  96. optimum/rbln/transformers/models/qwen2/modeling_qwen2.py +27 -1
  97. optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py +3 -3
  98. optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +5 -15
  99. optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py +1 -4
  100. optimum/rbln/transformers/models/qwen3/__init__.py +16 -0
  101. optimum/rbln/transformers/models/qwen3/configuration_qwen3.py +71 -0
  102. optimum/rbln/transformers/models/qwen3/modeling_qwen3.py +133 -0
  103. optimum/rbln/transformers/models/qwen3/qwen3_architecture.py +31 -0
  104. optimum/rbln/transformers/models/seq2seq/configuration_seq2seq.py +2 -12
  105. optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py +3 -1
  106. optimum/rbln/transformers/models/siglip/__init__.py +2 -6
  107. optimum/rbln/transformers/models/siglip/modeling_siglip.py +2 -2
  108. optimum/rbln/transformers/models/time_series_transformer/configuration_time_series_transformer.py +2 -2
  109. optimum/rbln/transformers/models/time_series_transformer/modeling_time_series_transformer.py +3 -5
  110. optimum/rbln/transformers/models/whisper/configuration_whisper.py +3 -12
  111. optimum/rbln/transformers/models/whisper/modeling_whisper.py +8 -2
  112. optimum/rbln/transformers/models/xlm_roberta/__init__.py +2 -8
  113. optimum/rbln/utils/depreacate_utils.py +16 -0
  114. optimum/rbln/utils/hub.py +8 -47
  115. optimum/rbln/utils/runtime_utils.py +31 -5
  116. {optimum_rbln-0.8.1rc1.dist-info → optimum_rbln-0.8.2.dist-info}/METADATA +1 -1
  117. {optimum_rbln-0.8.1rc1.dist-info → optimum_rbln-0.8.2.dist-info}/RECORD +119 -102
  118. {optimum_rbln-0.8.1rc1.dist-info → optimum_rbln-0.8.2.dist-info}/WHEEL +0 -0
  119. {optimum_rbln-0.8.1rc1.dist-info → optimum_rbln-0.8.2.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,315 @@
1
+ # Copyright 2025 Rebellions Inc. All rights reserved.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from pathlib import Path
16
+ from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
17
+
18
+ import rebel
19
+ import torch
20
+ import torch.nn as nn
21
+ from transformers import PixtralVisionConfig, PixtralVisionModel
22
+ from transformers.modeling_outputs import BaseModelOutput
23
+ from transformers.modeling_utils import no_init_weights
24
+ from transformers.models.pixtral.modeling_pixtral import PixtralRMSNorm, PixtralRotaryEmbedding
25
+
26
+ from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
27
+ from ....modeling import RBLNModel
28
+ from ....utils.logging import get_logger
29
+ from ....utils.runtime_utils import RBLNPytorchRuntime
30
+ from .configuration_pixtral import RBLNPixtralVisionModelConfig
31
+ from .pixtral_architecture import PixtralAttention
32
+
33
+
34
+ logger = get_logger(__name__)
35
+
36
+ if TYPE_CHECKING:
37
+ from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
38
+
39
+ from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
40
+
41
+
42
+ class RBLNRuntimePixtralVisionModel(RBLNPytorchRuntime):
43
+ mandatory_members = ["main_input_name"]
44
+
45
+ def __init__(
46
+ self,
47
+ runtime: rebel.Runtime,
48
+ config: PixtralVisionConfig,
49
+ rbln_config: RBLNPixtralVisionModelConfig,
50
+ **kwargs: Any,
51
+ ) -> None:
52
+ super().__init__(runtime, **kwargs)
53
+ self.patch_positional_embedding = PixtralRotaryEmbedding(config)
54
+ self.patch_size = config.patch_size
55
+ self.image_size = config.image_size
56
+ self.hidden_size = config.hidden_size
57
+ self.max_image_size = rbln_config.max_image_size
58
+
59
+ def forward(
60
+ self,
61
+ pixel_values: torch.Tensor,
62
+ image_sizes: torch.Tensor,
63
+ output_hidden_states: Optional[bool] = None,
64
+ return_dict: Optional[bool] = None,
65
+ **kwargs,
66
+ ):
67
+ if pixel_values.shape[2] > self.max_image_size[0] or pixel_values.shape[3] > self.max_image_size[1]:
68
+ raise ValueError("The height() and width of pixel_values can't be larger than max_image_size.")
69
+
70
+ if pixel_values.shape[2] != self.max_image_size[0] or pixel_values.shape[3] != self.max_image_size[1]:
71
+ padded_pixel_values = [
72
+ torch.nn.functional.pad(
73
+ image,
74
+ pad=(
75
+ 0,
76
+ self.max_image_size[1] - pixel_values.shape[3],
77
+ 0,
78
+ self.max_image_size[0] - pixel_values.shape[2],
79
+ ),
80
+ )
81
+ for image in pixel_values
82
+ ]
83
+ pixel_values = torch.stack(padded_pixel_values)
84
+
85
+ batch_size, _, H_max, W_max = pixel_values.shape
86
+ H_max_p = H_max // self.patch_size
87
+ W_max_p = W_max // self.patch_size
88
+
89
+ final_hidden_states = None
90
+
91
+ last_hidden_state_list = []
92
+ if output_hidden_states:
93
+ batch_hidden_states_list = []
94
+
95
+ for i in range(batch_size):
96
+ h_patched_original = image_sizes[i, 0] // self.patch_size
97
+ w_patched_original = image_sizes[i, 1] // self.patch_size
98
+
99
+ single_pixel_values = pixel_values[i : i + 1]
100
+ patch_embed = self.patch_conv(single_pixel_values)
101
+ patch_embed_seq = patch_embed[:, :, :h_patched_original, :w_patched_original].flatten(2).transpose(1, 2)
102
+ patch_embed_seq = self.ln_pre(patch_embed_seq)
103
+ patch_embed_seq = nn.functional.pad(
104
+ patch_embed_seq, (0, 0, 0, H_max_p * W_max_p - patch_embed_seq.shape[1]), "constant", value=0
105
+ )
106
+
107
+ max_w_from_config = self.image_size // self.patch_size
108
+ mesh = torch.meshgrid(torch.arange(h_patched_original), torch.arange(w_patched_original), indexing="ij")
109
+ h_grid, v_grid = torch.stack(mesh, dim=-1).reshape(-1, 2).chunk(2, -1)
110
+ ids = h_grid * max_w_from_config + v_grid
111
+ position_ids = ids[:, 0]
112
+
113
+ position_embeddings = self.patch_positional_embedding(patch_embed_seq, position_ids)
114
+ cos = nn.functional.pad(
115
+ position_embeddings[0],
116
+ (0, 0, 0, H_max_p * W_max_p - position_embeddings[0].shape[0]),
117
+ "constant",
118
+ value=0,
119
+ )
120
+ sin = nn.functional.pad(
121
+ position_embeddings[1],
122
+ (0, 0, 0, H_max_p * W_max_p - position_embeddings[1].shape[0]),
123
+ "constant",
124
+ value=0,
125
+ )
126
+
127
+ attention_mask = torch.full(
128
+ (1, patch_embed_seq.shape[-2]), fill_value=torch.finfo(patch_embed_seq.dtype).min
129
+ )
130
+ attention_mask[:, : h_patched_original * w_patched_original] = 0
131
+
132
+ transformer_output = super().forward(patch_embed_seq, attention_mask, cos, sin)
133
+
134
+ last_hidden_state_list.append(transformer_output[0][:, : h_patched_original * w_patched_original, :])
135
+ hidden_states = transformer_output[1:]
136
+
137
+ if output_hidden_states:
138
+ batch_hidden_states_list.append(
139
+ [hidden_state[:, : h_patched_original * w_patched_original, :] for hidden_state in hidden_states]
140
+ )
141
+
142
+ final_last_hidden_state = torch.cat(last_hidden_state_list, dim=1)
143
+
144
+ if output_hidden_states:
145
+ hidden_states = [
146
+ torch.cat(
147
+ [batch_hidden_states[layer_idx] for batch_hidden_states in batch_hidden_states_list],
148
+ dim=1,
149
+ )
150
+ for layer_idx in range(len(batch_hidden_states_list[0]))
151
+ ]
152
+
153
+ final_hidden_states = tuple(hidden_states)
154
+
155
+ if not return_dict:
156
+ return tuple(v for v in (final_last_hidden_state, final_hidden_states) if v is not None)
157
+
158
+ # TODO: output_attentions
159
+ return BaseModelOutput(
160
+ last_hidden_state=final_last_hidden_state,
161
+ hidden_states=final_hidden_states,
162
+ )
163
+
164
+
165
+ class _PixtralVisionModel(torch.nn.Module):
166
+ def __init__(self, model: PixtralVisionModel, output_hidden_states: bool):
167
+ super().__init__()
168
+ self.transformer = self.convert_to_rbln_pixtral_vision_model(model)
169
+ self.output_hidden_states = output_hidden_states
170
+
171
+ def convert_to_rbln_pixtral_vision_model(self, model: nn.Module):
172
+ for layer in model.transformer.layers:
173
+ layer.attention = PixtralAttention(layer.attention)
174
+ return model.transformer
175
+
176
+ def forward(self, patch_embeds, attention_mask, position_embeddings_1, position_embeddings_2):
177
+ output = self.transformer(
178
+ inputs_embeds=patch_embeds,
179
+ attention_mask=attention_mask,
180
+ position_embeddings=(position_embeddings_1, position_embeddings_2),
181
+ output_hidden_states=self.output_hidden_states,
182
+ return_dict=False,
183
+ )
184
+ return output
185
+
186
+
187
+ class RBLNPixtralVisionModel(RBLNModel):
188
+ """
189
+ RBLN optimized Pixtral vision encoder model.
190
+
191
+ This class provides hardware-accelerated inference for Pixtral vision encoders
192
+ on RBLN devices, supporting image encoding for multimodal tasks.
193
+ """
194
+
195
+ def __post_init__(self, **kwargs):
196
+ artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
197
+ with no_init_weights():
198
+ self.patch_conv = nn.Conv2d(
199
+ in_channels=self.config.num_channels,
200
+ out_channels=self.config.hidden_size,
201
+ kernel_size=self.config.patch_size,
202
+ stride=self.config.patch_size,
203
+ bias=False,
204
+ )
205
+ self.ln_pre = PixtralRMSNorm(self.config.hidden_size, eps=1e-5)
206
+ self.patch_conv.load_state_dict(artifacts["patch_conv"])
207
+ self.ln_pre.load_state_dict(artifacts["ln_pre"])
208
+ self.model = RBLNRuntimePixtralVisionModel(
209
+ self.model[0],
210
+ main_input_name="pixel_values",
211
+ config=self.config,
212
+ rbln_config=self.rbln_config,
213
+ patch_conv=self.patch_conv,
214
+ ln_pre=self.ln_pre,
215
+ )
216
+
217
+ @classmethod
218
+ def save_torch_artifacts(
219
+ cls,
220
+ model: "PreTrainedModel",
221
+ save_dir_path: Path,
222
+ subfolder: str,
223
+ rbln_config: RBLNModelConfig,
224
+ ):
225
+ save_dict = {}
226
+ save_dict["patch_conv"] = model.get_input_embeddings().state_dict()
227
+ save_dict["ln_pre"] = model.ln_pre.state_dict()
228
+ torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
229
+
230
+ @classmethod
231
+ def wrap_model_if_needed(
232
+ cls, model: torch.nn.Module, rbln_config: RBLNPixtralVisionModelConfig
233
+ ) -> torch.nn.Module:
234
+ wrapper_cfg = {
235
+ "output_hidden_states": rbln_config.output_hidden_states,
236
+ }
237
+ return _PixtralVisionModel(model, **wrapper_cfg).eval()
238
+
239
+ @classmethod
240
+ def update_rbln_config_using_pipe(
241
+ cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
242
+ ) -> "RBLNDiffusionMixinConfig":
243
+ return rbln_config
244
+
245
+ @classmethod
246
+ def _update_rbln_config(
247
+ cls,
248
+ preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
249
+ model: Optional["PreTrainedModel"] = None,
250
+ model_config: "PixtralVisionConfig" = None,
251
+ rbln_config: Optional[RBLNPixtralVisionModelConfig] = None,
252
+ ) -> RBLNPixtralVisionModelConfig:
253
+ if rbln_config.max_image_size is None:
254
+ rbln_config.max_image_size = (model_config.image_size, model_config.image_size)
255
+
256
+ if rbln_config.output_hidden_states is None:
257
+ rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
258
+
259
+ num_total_patches = (rbln_config.max_image_size[0] // model_config.patch_size) * (
260
+ rbln_config.max_image_size[1] // model_config.patch_size
261
+ )
262
+
263
+ rbln_compile_config = RBLNCompileConfig(
264
+ input_info=[
265
+ (
266
+ "patch_embeds",
267
+ [1, num_total_patches, model_config.hidden_size],
268
+ "float32",
269
+ ),
270
+ ("attention_mask", [1, num_total_patches], "float32"),
271
+ (
272
+ "position_embeddings_1",
273
+ [
274
+ num_total_patches,
275
+ model_config.head_dim,
276
+ ],
277
+ "float32",
278
+ ),
279
+ (
280
+ "position_embeddings_2",
281
+ [
282
+ num_total_patches,
283
+ model_config.head_dim,
284
+ ],
285
+ "float32",
286
+ ),
287
+ ]
288
+ )
289
+
290
+ rbln_config.set_compile_cfgs([rbln_compile_config])
291
+ return rbln_config
292
+
293
+ def forward(
294
+ self,
295
+ pixel_values: Optional[torch.FloatTensor] = None,
296
+ image_sizes: Optional[torch.FloatTensor] = None,
297
+ output_hidden_states: Optional[bool] = None,
298
+ return_dict: bool = True,
299
+ **kwargs,
300
+ ) -> Union[Tuple, BaseModelOutput]:
301
+ output_hidden_states = (
302
+ output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
303
+ )
304
+
305
+ if output_hidden_states != self.rbln_config.output_hidden_states:
306
+ raise ValueError(
307
+ f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
308
+ f"Please compile again with the correct argument."
309
+ )
310
+
311
+ output = self.model(
312
+ pixel_values, image_sizes, output_hidden_states=output_hidden_states, return_dict=return_dict
313
+ )
314
+
315
+ return output
@@ -0,0 +1,73 @@
1
+ # Copyright 2025 Rebellions Inc. All rights reserved.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from typing import Optional, Tuple
16
+
17
+ import torch
18
+ import torch.nn as nn
19
+
20
+ from ..decoderonly.decoderonly_architecture import apply_rotary_pos_emb
21
+
22
+
23
+ class PixtralAttention(nn.Module):
24
+ def __init__(self, self_attention):
25
+ super().__init__()
26
+ self.original_model = self_attention
27
+ self.num_heads = getattr(self.original_model, "num_heads", None) or getattr(
28
+ self.original_model.config, "num_attention_heads"
29
+ )
30
+ self.head_dim = self.original_model.head_dim
31
+ self.scaling = self.head_dim**-0.5
32
+
33
+ self.__post_init__()
34
+
35
+ def __post_init__(self):
36
+ self.q_proj = self.original_model.q_proj
37
+ self.k_proj = self.original_model.k_proj
38
+ self.v_proj = self.original_model.v_proj
39
+ self.o_proj = self.original_model.o_proj
40
+
41
+ def forward(
42
+ self,
43
+ hidden_states: torch.Tensor,
44
+ attention_mask: Optional[torch.Tensor] = None,
45
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
46
+ output_attentions: Optional[bool] = False,
47
+ ):
48
+ batch_size, patches, _ = hidden_states.size()
49
+
50
+ query_states = self.q_proj(hidden_states)
51
+ key_states = self.k_proj(hidden_states)
52
+ value_states = self.v_proj(hidden_states)
53
+
54
+ # TODO: return output attention
55
+ query_states = query_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
56
+ key_states = key_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
57
+ value_states = value_states.view(batch_size, patches, 1, self.num_heads, self.head_dim).transpose(1, 3)
58
+
59
+ cos, sin = position_embeddings
60
+ cos = cos[None, None, None, :, :]
61
+ sin = sin[None, None, None, :, :]
62
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
63
+
64
+ attn_weights = torch.matmul(query_states, key_states.transpose(3, 4)) * self.scaling
65
+ attn_weights = attn_weights + attention_mask
66
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
67
+ attn_output = torch.matmul(attn_weights, value_states)
68
+ attn_output = attn_output.transpose(1, 3)
69
+
70
+ attn_output = attn_output.reshape(batch_size, patches, -1)
71
+ attn_output = self.o_proj(attn_output)
72
+
73
+ return attn_output, _
@@ -12,5 +12,5 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig
16
- from .modeling_qwen2 import RBLNQwen2ForCausalLM
15
+ from .configuration_qwen2 import RBLNQwen2ForCausalLMConfig, RBLNQwen2ModelConfig
16
+ from .modeling_qwen2 import RBLNQwen2ForCausalLM, RBLNQwen2Model
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
15
+ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
16
16
 
17
17
 
18
18
  class RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
40
40
  )
41
41
  ```
42
42
  """
43
+
44
+
45
+ class RBLNQwen2ModelConfig(RBLNDecoderOnlyModelConfig):
46
+ """
47
+ Configuration class for RBLN Qwen2 models.
48
+
49
+ This class is an alias of RBLNDecoderOnlyModelConfig.
50
+ """
@@ -15,7 +15,11 @@
15
15
  from transformers import PretrainedConfig
16
16
 
17
17
  from ....utils import logging
18
- from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
18
+ from ...models.decoderonly import (
19
+ RBLNDecoderOnlyModel,
20
+ RBLNDecoderOnlyModelForCausalLM,
21
+ RBLNDecoderOnlyModelForCausalLMConfig,
22
+ )
19
23
  from .qwen2_architecture import QWEN2Wrapper
20
24
 
21
25
 
@@ -95,3 +99,25 @@ class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
95
99
  rbln_config.sliding_window = model_config.sliding_window
96
100
  rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
97
101
  return rbln_config
102
+
103
+
104
+ class RBLNQwen2Model(RBLNDecoderOnlyModel):
105
+ """
106
+ The Qwen2 Model transformer without a language modeling head.
107
+ This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
108
+ """
109
+
110
+ _decoder_wrapper_cls = QWEN2Wrapper
111
+
112
+ @classmethod
113
+ def _update_sliding_window_config(
114
+ cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
115
+ ):
116
+ # https://github.com/huggingface/transformers/issues/35896
117
+ # There seems to be a bug in transformers(v4.52.4). Therefore, similar to when attn_implementation is eager,
118
+ # we set all layers to use sliding window in this version. This should be updated once the bug is fixed.
119
+
120
+ rbln_config.cache_impl = "sliding_window"
121
+ rbln_config.sliding_window = model_config.sliding_window
122
+ rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
123
+ return rbln_config
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, List, Optional, Union
15
+ from typing import Any, List, Optional, Union
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
@@ -33,7 +33,7 @@ class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausal
33
33
  self,
34
34
  visual: Optional[RBLNModelConfig] = None,
35
35
  use_inputs_embeds: bool = True,
36
- **kwargs: Dict[str, Any],
36
+ **kwargs: Any,
37
37
  ):
38
38
  super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
39
39
  if not self.use_inputs_embeds:
@@ -53,7 +53,7 @@ class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
53
53
  mechanisms for processing images and videos.
54
54
  """
55
55
 
56
- def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
56
+ def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Any):
57
57
  """
58
58
  Args:
59
59
  max_seq_lens (Optional[Union[int, List[int]]]): Maximum sequence lengths for Vision
@@ -17,12 +17,7 @@ from pathlib import Path
17
17
  from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
18
18
 
19
19
  import torch
20
- from transformers import (
21
- AutoModelForVision2Seq,
22
- PretrainedConfig,
23
- PreTrainedModel,
24
- Qwen2_5_VLForConditionalGeneration,
25
- )
20
+ from transformers import AutoModelForVision2Seq, PretrainedConfig, PreTrainedModel, Qwen2_5_VLForConditionalGeneration
26
21
  from transformers.modeling_utils import no_init_weights
27
22
  from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
28
23
  Qwen2_5_VisionPatchEmbed,
@@ -34,7 +29,7 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
34
29
  from ....configuration_utils import RBLNCompileConfig
35
30
  from ....modeling import RBLNModel
36
31
  from ....utils.logging import get_logger
37
- from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyOutput
32
+ from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyForCausalLMOutput, RBLNDecoderOnlyModelForCausalLM
38
33
  from .configuration_qwen2_5_vl import (
39
34
  RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
40
35
  RBLNQwen2_5_VLForConditionalGenerationConfig,
@@ -45,12 +40,7 @@ from .qwen2_5_vl_architecture import Qwen2_5_VisionTransformerWrapper, Qwen2_5_V
45
40
  logger = get_logger(__name__)
46
41
 
47
42
  if TYPE_CHECKING:
48
- from transformers import (
49
- AutoFeatureExtractor,
50
- AutoProcessor,
51
- AutoTokenizer,
52
- PretrainedConfig,
53
- )
43
+ from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
54
44
 
55
45
 
56
46
  class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
@@ -595,7 +585,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
595
585
  generate_idx: Optional[torch.Tensor] = None,
596
586
  return_dict: Optional[bool] = None,
597
587
  **kwargs,
598
- ) -> RBLNDecoderOnlyOutput:
588
+ ) -> RBLNDecoderOnlyForCausalLMOutput:
599
589
  # Prefill
600
590
  if cache_position is None:
601
591
  inputs_embeds, position_embed, rope_deltas = self._preprocess_prefill(
@@ -637,7 +627,7 @@ class RBLNQwen2_5_VLForConditionalGeneration(RBLNDecoderOnlyModelForCausalLM):
637
627
  if not return_dict:
638
628
  return logits, generate_idx
639
629
  else:
640
- return RBLNDecoderOnlyOutput(
630
+ return RBLNDecoderOnlyForCausalLMOutput(
641
631
  logits=logits,
642
632
  generate_idx=generate_idx,
643
633
  )
@@ -4,10 +4,7 @@ from typing import Tuple
4
4
  import torch
5
5
  import torch.nn as nn
6
6
 
7
- from ..decoderonly.decoderonly_architecture import (
8
- DecoderOnlyWrapper,
9
- apply_rotary_pos_emb,
10
- )
7
+ from ..decoderonly.decoderonly_architecture import DecoderOnlyWrapper, apply_rotary_pos_emb
11
8
 
12
9
 
13
10
  class Qwen2_5_VisionTransformerWrapper(nn.Module):
@@ -0,0 +1,16 @@
1
+ # Copyright 2025 Rebellions Inc. All rights reserved.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from .configuration_qwen3 import RBLNQwen3ForCausalLMConfig, RBLNQwen3ModelConfig
16
+ from .modeling_qwen3 import RBLNQwen3ForCausalLM, RBLNQwen3Model
@@ -0,0 +1,71 @@
1
+ # Copyright 2025 Rebellions Inc. All rights reserved.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
16
+
17
+
18
+ class RBLNQwen3ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
19
+ """
20
+ Configuration class for RBLN Qwen3 models.
21
+
22
+ This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
23
+
24
+ Example usage:
25
+ ```python
26
+ from optimum.rbln import RBLNQwen3ForCausalLM, RBLNQwen3ForCausalLMConfig
27
+
28
+ # Create a configuration object
29
+ config = RBLNQwen3ForCausalLMConfig(
30
+ batch_size=1,
31
+ max_seq_len=40960,
32
+ tensor_parallel_size=4,
33
+ kvcache_partition_len=16384
34
+ )
35
+
36
+ # Use the configuration with from_pretrained
37
+ model = RBLNQwen3ForCausalLM.from_pretrained(
38
+ "Qwen/Qwen3-4B",
39
+ export=True,
40
+ rbln_config=config
41
+ )
42
+ ```
43
+ """
44
+
45
+
46
+ class RBLNQwen3ModelConfig(RBLNDecoderOnlyModelConfig):
47
+ """
48
+ Configuration class for RBLN Qwen3 models.
49
+
50
+ This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
51
+
52
+ Example usage:
53
+ ```python
54
+ from optimum.rbln import RBLNQwen3Model, RBLNQwen3ModelConfig
55
+
56
+ # Create a configuration object
57
+ config = RBLNQwen3ModelConfig(
58
+ batch_size=1,
59
+ max_seq_len=40960,
60
+ tensor_parallel_size=4,
61
+ kvcache_partition_len=16384
62
+ )
63
+
64
+ # Use the configuration with from_pretrained
65
+ model = RBLNQwen3Model.from_pretrained(
66
+ "Qwen/Qwen3-Embedding-4B",
67
+ export=True,
68
+ rbln_config=config
69
+ )
70
+ ```
71
+ """