optimum-rbln 0.7.3a4__tar.gz → 0.7.3a5__tar.gz

optimum_rbln-0.7.3a5/.github/version.yaml

@@ -0,0 +1 @@
+rebel_compiler_version: 0.7.3.dev196+ge9a294be

{optimum_rbln-0.7.3a4 → optimum_rbln-0.7.3a5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.3a4
+Version: 0.7.3a5
 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

{optimum_rbln-0.7.3a4 → optimum_rbln-0.7.3a5}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.3a4'
-__version_tuple__ = version_tuple = (0, 7, 3, 'a4')
+__version__ = version = '0.7.3a5'
+__version_tuple__ = version_tuple = (0, 7, 3, 'a5')

{optimum_rbln-0.7.3a4 → optimum_rbln-0.7.3a5}/src/optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -685,27 +685,28 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
             else:
                 rbln_kvcache_block_size = rbln_kvcache_partition_len
 
-        max_num_blocks, nbytes_per_block = cls.get_maximum_num_blocks(
-            config=model_config,
-            tensor_parallel_size=rbln_kwargs.get("tensor_parallel_size", 1),
-            kvcache_block_size=rbln_kvcache_block_size,
-            nbits_per_param=16 if rbln_quantization is None else 4,  # TODO(jongho): FIX Ad-hoc
-            n_model_params=rbln_kwargs["n_model_params"],
-        )
-        model_num_blocks = (rbln_max_seq_len // rbln_kvcache_block_size) * rbln_batch_size
-        rbln_kvcache_num_blocks = min(model_num_blocks, max_num_blocks)
+        rbln_kvcache_num_blocks = (rbln_max_seq_len // rbln_kvcache_block_size) * rbln_batch_size
+        if rbln_attn_impl == "flash_attn":
+            max_num_blocks, _ = cls.get_maximum_num_blocks(
+                config=model_config,
+                tensor_parallel_size=rbln_kwargs.get("tensor_parallel_size", 1),
+                kvcache_block_size=rbln_kvcache_block_size,
+                nbits_per_param=16 if rbln_quantization is None else 4,  # TODO(jongho): FIX Ad-hoc
+                n_model_params=rbln_kwargs["n_model_params"],
+            )
+            rbln_kvcache_num_blocks = min(rbln_kvcache_num_blocks, max_num_blocks)
 
-        required_blocks = rbln_max_seq_len // rbln_kvcache_block_size + 1
-        if rbln_kvcache_num_blocks < required_blocks:
-            rbln_kvcache_num_blocks = required_blocks
+            required_blocks = rbln_max_seq_len // rbln_kvcache_block_size + 1
+            if rbln_kvcache_num_blocks < required_blocks:
+                rbln_kvcache_num_blocks = required_blocks
 
-        logger.info(f"[KVCache] Compiling with num_blocks: {rbln_kvcache_num_blocks}")
+            logger.info(f"[KVCache] Compiling with num_blocks: {rbln_kvcache_num_blocks}")
 
-        if rbln_kvcache_num_blocks < rbln_batch_size:
-            raise RuntimeError(
-                f"Batch size ({rbln_batch_size}) exceeds available KV cache blocks ({rbln_kvcache_num_blocks}). "
-                "Ensure the number of blocks is at least equal to the batch size."
-            )
+            if rbln_kvcache_num_blocks < rbln_batch_size:
+                raise RuntimeError(
+                    f"Batch size ({rbln_batch_size}) exceeds available KV cache blocks ({rbln_kvcache_num_blocks}). "
+                    "Ensure the number of blocks is at least equal to the batch size."
+                )
 
         num_attention_heads = getattr(model_config, "n_head", None) or getattr(model_config, "num_attention_heads")
         num_key_value_heads = getattr(model_config, "num_key_value_heads", None) or num_attention_heads
@@ -805,9 +806,6 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                 "kvcache_block_size": rbln_kvcache_block_size,
                 "attn_impl": rbln_attn_impl,
                 "kvcache_num_blocks": rbln_kvcache_num_blocks,
-                "model_num_blocks": model_num_blocks,
-                "max_num_blocks": max_num_blocks,
-                "nbytes_per_block": nbytes_per_block,
             }
         )
 

{optimum_rbln-0.7.3a4 → optimum_rbln-0.7.3a5}/src/optimum/rbln/transformers/models/phi/phi_architecture.py

@@ -92,7 +92,7 @@ class PhiLayer(DecoderOnlyLayer):
 
         hidden_states = self.get_pre_attention_layernorm()(hidden_states)
 
-        attn_outputs, present_key_values = self.self_attn(
+        attn_output = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
             seq_positions=seq_positions,
@@ -104,9 +104,9 @@ class PhiLayer(DecoderOnlyLayer):
 
         feed_forward_hidden_states = self._original_mod.mlp(hidden_states)
 
-        hidden_states = attn_outputs + feed_forward_hidden_states + residual
+        hidden_states = attn_output + feed_forward_hidden_states + residual
 
-        return hidden_states, present_key_values
+        return hidden_states
 
 
 class PhiModel(DecoderOnlyModel):

optimum_rbln-0.7.3a4/.github/version.yaml

@@ -1 +0,0 @@
-rebel_compiler_version: 0.7.3.dev162+g81024cdb