optimum-rbln 0.2.1a1__py3-none-any.whl → 0.2.1a2__py3-none-any.whl

optimum/rbln/__init__.py

@@ -71,6 +71,7 @@ _import_structure = {
         "RBLNRobertaForSequenceClassification",
         "RBLNRobertaForMaskedLM",
         "RBLNViTForImageClassification",
+        "RBLNBertForMaskedLM",
     ],
     "diffusers": [
         "RBLNStableDiffusionPipeline",
@@ -141,6 +142,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
         RBLNBartModel,
+        RBLNBertForMaskedLM,
         RBLNBertForQuestionAnswering,
         RBLNBertModel,
         RBLNCLIPTextModel,

optimum/rbln/__version__.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1a1'
+__version__ = version = '0.2.1a2'
 __version_tuple__ = version_tuple = (0, 2, 1)

optimum/rbln/ops/attn.py

@@ -152,16 +152,16 @@ def register_rbln_custom_attention_add_softmax():
         """
         return (
             q,
-            torch.empty(1, *kcache.shape[1:], device=kcache.device),
-            torch.empty(1, *vcache.shape[1:], device=vcache.device),
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
         )
 
     @register_fake("rbln_custom_ops::attn_decode_add_softmax")
     def attn_decode_add_softmax_abstract(q, k, v, m, kcache, vcache, seq, partition):
         return (
             q,
-            torch.empty(1, *kcache.shape[1:], device=kcache.device),
-            torch.empty(1, *vcache.shape[1:], device=vcache.device),
+            torch.empty(*kcache.shape, device=kcache.device),
+            torch.empty(*vcache.shape, device=vcache.device),
         )
 
     torch.library.define(

optimum/rbln/transformers/__init__.py

@@ -35,6 +35,7 @@ _import_structure = {
         "RBLNBartForConditionalGeneration",
         "RBLNBartModel",
         "RBLNBertModel",
+        "RBLNBertForMaskedLM",
         "RBLNBertForQuestionAnswering",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
@@ -92,6 +93,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
         RBLNBartModel,
+        RBLNBertForMaskedLM,
         RBLNBertForQuestionAnswering,
         RBLNBertModel,
         RBLNCLIPTextModel,

optimum/rbln/transformers/models/__init__.py

@@ -33,7 +33,7 @@ _import_structure = {
         "RBLNAutoModelForVision2Seq",
     ],
     "bart": ["RBLNBartForConditionalGeneration", "RBLNBartModel"],
-    "bert": ["RBLNBertModel", "RBLNBertForQuestionAnswering"],
+    "bert": ["RBLNBertModel", "RBLNBertForQuestionAnswering", "RBLNBertForMaskedLM"],
     "clip": ["RBLNCLIPTextModel", "RBLNCLIPTextModelWithProjection", "RBLNCLIPVisionModel"],
     "dpt": ["RBLNDPTForDepthEstimation"],
     "exaone": ["RBLNExaoneForCausalLM"],
@@ -67,7 +67,7 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
     )
     from .bart import RBLNBartForConditionalGeneration, RBLNBartModel
-    from .bert import RBLNBertForQuestionAnswering, RBLNBertModel
+    from .bert import RBLNBertForMaskedLM, RBLNBertForQuestionAnswering, RBLNBertModel
     from .clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection, RBLNCLIPVisionModel
     from .dpt import RBLNDPTForDepthEstimation
     from .exaone import RBLNExaoneForCausalLM

optimum/rbln/transformers/models/bert/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .modeling_bert import RBLNBertForQuestionAnswering, RBLNBertModel
+from .modeling_bert import RBLNBertForMaskedLM, RBLNBertForQuestionAnswering, RBLNBertModel

optimum/rbln/transformers/models/bert/modeling_bert.py

@@ -20,7 +20,7 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ...modeling_generic import RBLNModelForQuestionAnswering
+from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForQuestionAnswering
 
 
 logger = logging.getLogger(__name__)
@@ -100,5 +100,9 @@ class RBLNBertModel(RBLNModel):
         return rbln_config
 
 
+class RBLNBertForMaskedLM(RBLNModelForMaskedLM):
+    rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
+
+
 class RBLNBertForQuestionAnswering(RBLNModelForQuestionAnswering):
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]

optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -544,15 +544,19 @@ class DecoderOnlyAttention(nn.Module):
         super().__init__()
         self._original_mod = self_attn
         self.layer_idx = self_attn.layer_idx
-        self.num_heads = self._original_mod.num_heads
+        self.num_heads = getattr(self._original_mod, "num_heads", None) or getattr(
+            self._original_mod.config, "num_attention_heads"
+        )
         self.head_dim = self._original_mod.head_dim
         self._phase = "prefill"
         self.scale = torch.tensor(self.get_attn_scale())
 
         if hasattr(self._original_mod, "num_key_value_heads"):
             self.num_key_value_heads = self._original_mod.num_key_value_heads
+        elif hasattr(self._original_mod, "config") and hasattr(self._original_mod.config, "num_key_value_heads"):
+            self.num_key_value_heads = self._original_mod.config.num_key_value_heads
         else:
-            self.num_key_value_heads = self._original_mod.num_heads
+            self.num_key_value_heads = self.num_heads
 
         self.attention = self.get_attention()
         self.__post_init__()

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

@@ -420,7 +420,7 @@ class Seq2SeqSelfAttention(nn.Module):
         pass
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int) -> torch.Tensor:
-        return tensor.view(bsz, 1, seq_len, 1, self.num_heads, self.head_dim).transpose(2, 4)
+        return tensor.view(bsz, seq_len, 1, self.num_heads, self.head_dim).transpose(1, 3)
 
     def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Projects input hidden states into query, key, and value representations.
@@ -450,38 +450,21 @@ class Seq2SeqSelfAttention(nn.Module):
         key_states = self._shape(key_states, -1, bsz)
         value_states = self._shape(value_states, -1, bsz)
 
-        all_key_states = []
-        all_value_states = []
-        all_attn_output = []
-        for b_idx in range(bsz):
-            query_state = query_states[b_idx]
-            key_state = key_states[b_idx]
-            value_state = value_states[b_idx]
-            attn_mask = attention_mask[b_idx].unsqueeze(0).unsqueeze(2)
-            past_key_state = past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim)
-            past_value_state = past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim)
-
-            attn_output, key_state, value_state = self.attn_decode(
-                query_state,
-                key_state,
-                value_state,
-                attn_mask,
-                past_key_state,
-                past_value_state,
-                cache_position[b_idx][0],
-                torch.tensor(1.0, dtype=torch.float32),  # scale
-            )
-
-            attn_output = attn_output.view(1, self.num_heads, -1, self.head_dim).transpose(1, 2)
-            attn_output = attn_output.reshape(1, -1, self.num_heads * self.head_dim)
-
-            all_key_states.append(key_state.squeeze(2))
-            all_value_states.append(value_state.squeeze(2))
-            all_attn_output.append(attn_output)
+        attn_output, key_states, value_states = self.attn_decode(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask.unsqueeze(
+                2
+            ),  # Unsqueeze group axis since CustomKernel expects it for group query attention
+            past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim),
+            past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim),
+            cache_position.squeeze(1),
+            torch.tensor(1.0, dtype=torch.float32),  # scale
+        )
 
-        key_states = torch.cat(all_key_states, dim=0)
-        value_states = torch.cat(all_value_states, dim=0)
-        attn_output = torch.cat(all_attn_output, dim=0)
+        attn_output = attn_output.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, -1, self.num_heads * self.head_dim)
 
         attn_output = self.out_proj(attn_output)
         present_key_value = (key_states, value_states)

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

@@ -147,6 +147,11 @@ class T5CrossAttention(nn.Module):
     def __init__(self, attn):
         super().__init__()
         self.attn = attn
+        self.q = attn.q
+        self.o = attn.o
+        self.n_heads = attn.n_heads
+        self.key_value_proj_dim = attn.key_value_proj_dim
+        self.inner_dim = attn.inner_dim
 
     def forward(
         self,
@@ -155,9 +160,27 @@ class T5CrossAttention(nn.Module):
         attention_mask: torch.Tensor = None,
         key_value_states: torch.Tensor = None,
     ):
-        return self.attn(
-            hidden_states=hidden_states,
-            past_key_value=past_key_value,
-            position_bias=attention_mask,
-            key_value_states=key_value_states,
-        )
+        batch_size = hidden_states.shape[0]
+
+        query_states = self.q(hidden_states)
+        query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
+
+        # reuse k,v, cross_attentions
+        key_states = past_key_value[0]
+        value_states = past_key_value[1]
+
+        # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
+        scores = torch.matmul(query_states, key_states.transpose(3, 2))
+        scores += attention_mask
+
+        # (batch_size, n_heads, seq_length, key_length)
+        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, -1, self.inner_dim)
+        attn_output = self.o(attn_output)
+
+        outputs = (attn_output, past_key_value)
+
+        return outputs

{optimum_rbln-0.2.1a1.dist-info → optimum_rbln-0.2.1a2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.2.1a1
+Version: 0.2.1a2
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai
@@ -28,7 +28,7 @@ Requires-Dist: packaging>=24.1
 Requires-Dist: torch<=2.5.1
 Requires-Dist: torchaudio<=2.5.1
 Requires-Dist: torchvision<=0.20.1
-Requires-Dist: transformers==4.45.2
+Requires-Dist: transformers==4.48.3
 Description-Content-Type: text/markdown
 
 

{optimum_rbln-0.2.1a1.dist-info → optimum_rbln-0.2.1a2.dist-info}/RECORD

@@ -1,5 +1,5 @@
-optimum/rbln/__init__.py,sha256=zTum8bjtVKdT9dq56vyQZy1eG2UUhEbiaGyvtokVDok,6036
-optimum/rbln/__version__.py,sha256=Hh8jxA0Z8Gd3AqkgqwtySq_B5l73iIwQzKQ1Y-52wbw,413
+optimum/rbln/__init__.py,sha256=sLCjJu_MLZEKDOwHIlJP4u4GzGZx-1kqHTYGw5B4xDg,6096
+optimum/rbln/__version__.py,sha256=RUdGjwBqCynJjA53NyAfXHitjfqAoWaGNCXiO_9cNsw,413
 optimum/rbln/modeling.py,sha256=OQGLkzlE3vD3O-ZeE1Z0jK-QCqWy1V46pSCOdmehFTI,8267
 optimum/rbln/modeling_base.py,sha256=sU5Tr3SmhQZPsbKz5xo-FqMU1gC4Xd4m9xZVIx2NY7I,20359
 optimum/rbln/modeling_config.py,sha256=7104bxmrvKW4Q6XTruQayiIGl8GHDFmPkJ3cknMIInE,11335
@@ -34,26 +34,26 @@ optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_x
 optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py,sha256=3aB1Rw-OgKytQOHwOaShbEvq_XVHPOGvsGm8pstEmKU,930
 optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py,sha256=MzVP1wscaO1sUIiBIPJqG6zuGyez9VUbA42-JSIm-mk,930
 optimum/rbln/ops/__init__.py,sha256=MbAHevg59fXQOFgrnsFFZ97s8-YrgvYCYML_sgKEEfM,816
-optimum/rbln/ops/attn.py,sha256=rB0xYhUxpb1o1JISMNJnPffr7qF8DwnuecOSMWPVGoA,9728
+optimum/rbln/ops/attn.py,sha256=QYvSMg4ps_PenHwpLVhuYRoOAFvHIo19nY0ZEdj4nTE,9700
 optimum/rbln/ops/flash_attn.py,sha256=Zn5nkouY3kk6MBivQpPjgGh4oepjpi8F3tnTrmrNfpg,2304
 optimum/rbln/ops/kv_cache_update.py,sha256=9W4WCO1Dtfy0u5i978JJRa7uLbqrfR2lHuoPynb07fw,3143
-optimum/rbln/transformers/__init__.py,sha256=8NHC8MpmWc8fteSBtFNJ729LON6FRHI2J7zifCAAXZ4,4107
+optimum/rbln/transformers/__init__.py,sha256=SdOjpa4Iufo6aOJPvjQwD_vz28dmmNV9AtF5Cz9ajLM,4167
 optimum/rbln/transformers/modeling_alias.py,sha256=yx7FnZQWAnrWzivaO5hI7T6i-fyLzt2tMIXG2oDNbPo,1657
 optimum/rbln/transformers/modeling_generic.py,sha256=SD7XjpjnCn-ejNAUWgkaaHV6Fv2Y6K-hbXEXXb9W_H4,18177
 optimum/rbln/transformers/modeling_rope_utils.py,sha256=3zwkhYUyTZhxCJUSmwCc88iiY1TppRWEY9ShwUqNB2k,14293
-optimum/rbln/transformers/models/__init__.py,sha256=GBCGLGdo_HMEkGhZSevOgt-M9KMETaKUqRmgEf3WpQE,3639
+optimum/rbln/transformers/models/__init__.py,sha256=wucrA1ybpDfNcrySwdVeK5PZEYl-3ONXJvGpHGTvteo,3683
 optimum/rbln/transformers/models/auto/__init__.py,sha256=GvGbb3ZpMv-h6euXeZ42jSizoOfrL2O1uvpAnfKxYEo,1034
 optimum/rbln/transformers/models/auto/auto_factory.py,sha256=IK9jFrJ3EEzYQa9_aKpcp2TO68M5YGkA-HcfBVpA2QU,7027
 optimum/rbln/transformers/models/auto/modeling_auto.py,sha256=Un9qoqdy3dO8JBza_bTJF_6_fRVNM9QisihSgTRFI-o,3933
 optimum/rbln/transformers/models/bart/__init__.py,sha256=32HPe0_GIO0hp9U464Iv6Jd7M-1nop9g8hA1UZMHhyw,674
 optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=dTkgMpNkyh4vT_mZU5tQ5bvH_lRZfRjaJ1gIHvJkmgs,5479
 optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=ADRbE-5N3xJ60AzzjJ4BZs_THmB71qs4XTr9iFqsEqE,5667
-optimum/rbln/transformers/models/bert/__init__.py,sha256=_pEfofte9Ss8y2dgCbwl3FCHr6X2LNSm9VMMMS1vVh0,670
-optimum/rbln/transformers/models/bert/modeling_bert.py,sha256=ikAcgsx5zIeMbmyUWusWTuX35q_gm8ZehEll-xOIJFc,4497
+optimum/rbln/transformers/models/bert/__init__.py,sha256=YVV7k_laU6yJBawZrgjIWjRmIF-Y4oQQHqyf8lsraQs,691
+optimum/rbln/transformers/models/bert/modeling_bert.py,sha256=flzCLWqCaCnQLfWCVSRmQke_CEAXzcG0DOoUM8EAqkg,4649
 optimum/rbln/transformers/models/clip/__init__.py,sha256=ssJqlEt318ti2QaEakGh_tO3Ap1VSPCVF-ymUuvjAJs,698
 optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=LGytQ33WGV2qqXnj_1dMiPN63ytL1JlNQlc3aXaG1bA,5705
 optimum/rbln/transformers/models/decoderonly/__init__.py,sha256=pDogsdpJKKB5rqnVFrRjwfhUvOSV-jZ3oARMsqSvOOQ,665
-optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=73sWpSDgm6_ztiGPsGEljJo-a7AisTAqFb7s-UaShVE,36292
+optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=BjQHwoPZfM-KUQzxm4AU-PdmoMgLxnCG6kfSpGjUvrk,36578
 optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=mAgRRMGVHvTUjJBDlmUOjNhSNjprKSD7tLeFknrx0Rw,25810
 optimum/rbln/transformers/models/dpt/__init__.py,sha256=gP1tkR3XMNlHq1GT87ugIVvb2o_1eAUg1JaniXjy1Lw,651
 optimum/rbln/transformers/models/dpt/modeling_dpt.py,sha256=HS_f7bL2CvbWP_7NGMzPYb0GdHRE2xHF0e1DkzlRdRE,3411
@@ -85,10 +85,10 @@ optimum/rbln/transformers/models/qwen2/modeling_qwen2.py,sha256=9-aFDvjMzPNUyGOz
 optimum/rbln/transformers/models/qwen2/qwen2_architecture.py,sha256=XlNAMYAcDLohnSAhIFGKOPuCB5XLgzYs5ABWdeQSaZs,720
 optimum/rbln/transformers/models/seq2seq/__init__.py,sha256=EmEMV4rOYqKyruX85d0fR73-b8N6BSD6CPcbpYdBuVk,651
 optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=4GHlLf6xm7a8YegYVX-zhIqk4ouwHCzQuj8Z-jXSFJw,15407
-optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=JTcLA-xnSs6KD8xUA6RY0QMWoXChnHvmhyKa_4039UY,19053
+optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=15yoF-wyhcLcK-Z2MOUmyPlkOMNTVOJ013uBepqtpxA,18387
 optimum/rbln/transformers/models/t5/__init__.py,sha256=1skR1RmnG62WTAP3-F5P1x-V_ReFhMyirH3u56vWwvc,675
 optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=MFs-3yYviV1QqSpsTB2GarTEs9wGH5AYofksLQLMBXg,8043
-optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=_K_IROovNH7zZIj2E1datCLIWbRpAg181p03GDTTtXE,6209
+optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=kkjErS42mW2jv5O_xL7BaKobvvqy7BGmYOowKyHakvI,7189
 optimum/rbln/transformers/models/wav2vec2/__init__.py,sha256=YpgA0K-vyg9veh0eL_jxauosbRpb_kpGKHvvQLBspKM,649
 optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py,sha256=Ws0tw2j9Mp8BREP6nI-Ann_U0rhkqofaQFCKoepDYRA,3837
 optimum/rbln/transformers/models/whisper/__init__.py,sha256=ktnNe5ri3ycCWZ_W_voFB9y9-vgGgxS1X9s8LBRZmWc,665
@@ -108,7 +108,7 @@ optimum/rbln/utils/model_utils.py,sha256=DfD_Z2qvZHqcddXqnzTM1AN8khanj3-DXK2lJvV
 optimum/rbln/utils/runtime_utils.py,sha256=5-DYniyP59nx-mrrbi7AqA77L85b4Cm5oLpaxidSyss,3699
 optimum/rbln/utils/save_utils.py,sha256=W5ON-90xLcz1suFlZwOhmYB5Mf9XSTu00xrfTfyL88U,3608
 optimum/rbln/utils/submodule.py,sha256=oZoGrItB8WqY4i-K9WJPlLlcLohc1YGB9OHB8_XZw3A,4071
-optimum_rbln-0.2.1a1.dist-info/METADATA,sha256=QqxzPcaWJ57Dtaws6oRFO_tCjZPjZkWQldxpWHXpg9w,5300
-optimum_rbln-0.2.1a1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.2.1a1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.2.1a1.dist-info/RECORD,,
+optimum_rbln-0.2.1a2.dist-info/METADATA,sha256=hwr_UclFUgkNbtwCanMy_5lF56FPX8WG0GGz3B0cJCc,5300
+optimum_rbln-0.2.1a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.2.1a2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.2.1a2.dist-info/RECORD,,