optimum-rbln 0.7.5a0__py3-none-any.whl → 0.7.5rc0__py3-none-any.whl

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

@@ -146,7 +146,10 @@ class DecoderOnlyWrapper(nn.Module):
         max_seq_len: int,
         use_rotary_emb: bool,
         attn_impl: str,
+        use_inputs_embeds: bool,
         use_attention_mask: bool,
+        use_position_ids: bool,
+        use_learned_pos_emb: Optional[bool] = None,
         kvcache_partition_len: Optional[int] = None,
         kvcache_block_size: Optional[int] = None,
     ):
@@ -154,13 +157,21 @@ class DecoderOnlyWrapper(nn.Module):
         self.config = causal_lm.config
 
         if use_rotary_emb:
-            self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
+            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.attn_impl = attn_impl
         self.kvcache_block_size = kvcache_block_size
         self.use_attention_mask = use_attention_mask
+        self.use_position_ids = use_position_ids
+        self.use_inputs_embeds = use_inputs_embeds
+        self.use_learned_pos_emb = use_learned_pos_emb
+
         if self.attn_impl == "flash_attn":
             self.kvcache_partition_len = kvcache_partition_len or DEFAULT_FLASH_ATTN_PARTITION_LENGTH
         elif self.attn_impl == "eager":
@@ -188,7 +199,10 @@ class DecoderOnlyWrapper(nn.Module):
         for layer in causal_lm.model.layers:
             if self.attn_impl == "eager":
                 new_self_attn = DecoderOnlyAttention(
-                    layer.self_attn, self.use_attention_mask, kvcache_block_size=self.kvcache_block_size
+                    layer.self_attn,
+                    self.use_attention_mask,
+                    self.use_position_ids,
+                    kvcache_block_size=self.kvcache_block_size,
                 )
             elif self.attn_impl == "flash_attn":
                 new_self_attn = DecoderOnlyFlashAttention(
@@ -196,6 +210,7 @@ class DecoderOnlyWrapper(nn.Module):
                     kvcache_partition_len=self.kvcache_partition_len,
                     kvcache_block_size=self.kvcache_block_size,
                     use_attention_mask=self.use_attention_mask,
+                    use_position_ids=self.use_position_ids,
                 )
             else:
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")
@@ -209,6 +224,7 @@ class DecoderOnlyWrapper(nn.Module):
             partition_len=self.kvcache_partition_len,
             max_seq_len=max_seq_len,
             kvcache_block_size=self.kvcache_block_size,
+            use_learned_pos_emb=self.use_learned_pos_emb,
         )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
@@ -222,24 +238,16 @@ class DecoderOnlyWrapper(nn.Module):
         self._phase = phase
         self.causal_lm.phase = phase
 
-    def forward_common(
-        self,
-        input_ids_or_inputs_embeds: torch.Tensor,
-        cache_position: torch.Tensor,
-        attention_mask: torch.Tensor,
-        query_position: torch.Tensor,
-        block_tables: torch.Tensor,
-        rotary_emb: Union[nn.Module, torch.Tensor],
-        *past_key_values: List[torch.Tensor],
-    ):
-        if input_ids_or_inputs_embeds.ndim == 2:
-            input_ids = input_ids_or_inputs_embeds
-            inputs_embeds = None
-        elif input_ids_or_inputs_embeds.ndim == 3:
-            input_ids = None
-            inputs_embeds = input_ids_or_inputs_embeds
-        else:
-            raise NotImplementedError(f"Unknown ndim of input : {input_ids_or_inputs_embeds.ndim}")
+    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)
+        block_tables = args.pop(0)
+        query_position = args.pop(0) if self.phase == "prefill" else None
+        attention_mask = args.pop(0) if self.use_attention_mask else None
+        position_ids = args.pop(0) if self.use_position_ids else None
+        past_key_values = args
 
         if len(past_key_values) != 2 * self.num_hidden_layers:
             raise ValueError(
@@ -256,11 +264,37 @@ class DecoderOnlyWrapper(nn.Module):
             _past_key_values.append(past_key_value)
         past_key_values = _past_key_values
 
+        return (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            block_tables,
+            query_position,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            self.rotary_emb,
+        )
+
+    def forward(self, *args):
+        (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            block_tables,
+            query_position,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            rotary_emb,
+        ) = self.prepare_forward_args(*args)
+
         logit = self.causal_lm(
             input_ids=input_ids,
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             cache_position=cache_position,
+            position_ids=position_ids,
             query_position=query_position,
             past_key_values=past_key_values,
             rotary_emb=rotary_emb,
@@ -269,58 +303,6 @@ class DecoderOnlyWrapper(nn.Module):
 
         return logit
 
-    def forward(self, *args):
-        if self.phase == "decode":
-            if self.use_attention_mask:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    attention_mask,
-                    block_tables,
-                    *past_key_values,
-                ) = args
-            else:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    block_tables,
-                    *past_key_values,
-                ) = args
-                attention_mask = None
-            query_position = None
-        elif self.phase == "prefill":
-            if self.use_attention_mask:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    attention_mask,
-                    query_position,
-                    block_tables,
-                    *past_key_values,
-                ) = args
-            else:
-                (
-                    input_ids_or_inputs_embeds,
-                    cache_position,
-                    query_position,
-                    block_tables,
-                    *past_key_values,
-                ) = args
-                attention_mask = None
-
-        else:
-            raise ValueError(f"Unknown phase: {self.phase}")
-
-        return self.forward_common(
-            input_ids_or_inputs_embeds,
-            cache_position,
-            attention_mask,
-            query_position,
-            block_tables,
-            self.rotary_emb,
-            *past_key_values,
-        )
-
 
 class DecoderOnlyForCausalLM(nn.Module):
     """A specialized wrapper for Causal Language Models optimized for RBLN compilation.
@@ -367,6 +349,7 @@ class DecoderOnlyForCausalLM(nn.Module):
         inputs_embeds: torch.Tensor = None,
         attention_mask: torch.Tensor = None,
         cache_position: torch.Tensor = None,
+        position_ids: torch.Tensor = None,
         query_position: torch.Tensor = None,
         past_key_values: Tuple[Tuple[torch.Tensor]] = None,
         rotary_emb: nn.Module = None,
@@ -378,6 +361,7 @@ class DecoderOnlyForCausalLM(nn.Module):
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             cache_position=cache_position,
+            position_ids=position_ids,
             past_key_values=past_key_values,
             rotary_emb=rotary_emb,
             block_tables=block_tables,
@@ -387,6 +371,13 @@ class DecoderOnlyForCausalLM(nn.Module):
             hidden_states = hidden_states[:, query_position.to(torch.int).unsqueeze(0)]
 
         logits = self.lm_head(hidden_states)
+
+        # Apply final logit softmaxing if configured, e.g. for Gemma2
+        if getattr(self.config, "final_logit_softcapping", None) is not None:
+            logits = logits / self.config.final_logit_softcapping
+            logits = torch.tanh(logits)
+            logits = logits * self.config.final_logit_softcapping
+
         return logits
 
 
@@ -404,7 +395,13 @@ class DecoderOnlyModel(nn.Module):
     """
 
     def __init__(
-        self, model, layers: List["DecoderOnlyLayer"], partition_len=None, max_seq_len=None, kvcache_block_size=None
+        self,
+        model,
+        layers: List["DecoderOnlyLayer"],
+        partition_len=None,
+        max_seq_len=None,
+        kvcache_block_size=None,
+        use_learned_pos_emb=None,
     ):
         super().__init__()
         self._original_mod = model
@@ -413,6 +410,7 @@ class DecoderOnlyModel(nn.Module):
         self.partition_len = partition_len
         self.kvcache_block_size = kvcache_block_size
         self.max_seq_len = max_seq_len
+        self.use_learned_pos_emb = use_learned_pos_emb
 
     @property
     def phase(self):
@@ -457,11 +455,12 @@ class DecoderOnlyModel(nn.Module):
     def forward(
         self,
         input_ids: torch.Tensor = None,
-        inputs_embeds: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
         attention_mask: torch.Tensor = None,
         cache_position: torch.Tensor = None,
+        position_ids: torch.Tensor = None,
         past_key_values: Tuple[Tuple[torch.Tensor]] = None,
-        rotary_emb: nn.Module = None,
+        rotary_emb: Optional[Union[nn.Module, torch.Tensor]] = None,
         block_tables: Optional[torch.Tensor] = None,
     ):
         # retrieve input_ids and inputs_embeds
@@ -477,24 +476,38 @@ class DecoderOnlyModel(nn.Module):
         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, cache_position)
+                cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
+
+        elif self.use_learned_pos_emb:
+            batch_size = inputs_embeds.shape[0]
+            hidden_all = []
+            for i in range(batch_size):
+                positions_idx = position_ids[i]
+                position_weight = self.get_pos_embedding().weight[2:]
+                position = position_weight[positions_idx]
+                batch_hidden = position + inputs_embeds[i]
+                hidden_all.append(batch_hidden)
+            hidden_states = torch.stack(hidden_all, dim=0)
+            cos, sin = None, None
+
         else:
             batch_size = inputs_embeds.shape[0]
-            if cache_position.shape[0] > 1:
+            if position_ids.shape[0] > 1:
                 position_embeds = []
                 for b_idx in range(batch_size):
-                    position_embed = self.get_pos_embedding()(cache_position[b_idx])
+                    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()(cache_position)
+                position_embeds = self.get_pos_embedding()(position_ids)
             hidden_states = hidden_states + position_embeds
             cos, sin = None, None
 
@@ -612,7 +625,7 @@ class DecoderOnlyAttention(nn.Module):
         self_attn: Original attention module from the base model
     """
 
-    def __init__(self, self_attn, use_attention_mask, kvcache_block_size):
+    def __init__(self, self_attn, use_attention_mask, use_position_ids, kvcache_block_size):
         super().__init__()
         self._original_mod = self_attn
         self.layer_idx = self_attn.layer_idx
@@ -631,6 +644,7 @@ class DecoderOnlyAttention(nn.Module):
             self.num_key_value_heads = self.num_heads
 
         self.use_attention_mask = use_attention_mask
+        self.use_position_ids = use_position_ids
         self.attention = self.get_attention()
         self.kvcache_block_size = kvcache_block_size
         self.__post_init__()
@@ -645,7 +659,9 @@ class DecoderOnlyAttention(nn.Module):
         self.attention.phase = phase
 
     def get_attention(self):
-        return AttentionOp(self.num_heads, self.head_dim, self.num_key_value_heads, self.use_attention_mask)
+        return AttentionOp(
+            self.num_heads, self.head_dim, self.num_key_value_heads, self.use_attention_mask, self.use_position_ids
+        )
 
     def __post_init__(self):
         self.q_proj = self._original_mod.q_proj
@@ -718,13 +734,16 @@ class DecoderOnlyAttention(nn.Module):
 
 
 class AttentionOp(nn.Module):
-    def __init__(self, num_heads: int, head_dim: int, num_key_value_heads: int, use_attention_mask: bool):
+    def __init__(
+        self, num_heads: int, head_dim: int, num_key_value_heads: int, use_attention_mask: bool, use_position_ids: bool
+    ):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.num_key_value_heads = num_key_value_heads
         self.phase = "prefill"
         self.use_attention_mask = use_attention_mask
+        self.use_position_ids = use_position_ids
 
     def forward(
         self,
@@ -757,7 +776,8 @@ class AttentionOp(nn.Module):
         # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
         key_state = key_state.unsqueeze(2)  # 1, 32, 1, 128, 128
         value_state = value_state.unsqueeze(2)
-        if self.use_attention_mask:
+
+        if self.use_attention_mask and not self.use_position_ids:
             attn_mask = attn_mask.unsqueeze(2)
 
         if self.phase == "decode":
@@ -774,7 +794,7 @@ class AttentionOp(nn.Module):
         )
 
         if self.phase == "decode":
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_attn_decode(
                     q=query_state,
                     k=key_state,
@@ -798,10 +818,11 @@ class AttentionOp(nn.Module):
                     scale=scale,
                     block_table=block_tables,
                     block_size=block_size,
+                    mask=attn_mask if self.use_position_ids else None,
                 )
 
         else:
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_attn_prefill(
                     q=query_state,
                     k=key_state,
@@ -825,6 +846,8 @@ class AttentionOp(nn.Module):
                     scale=scale,
                     block_table=block_tables,
                     block_size=block_size,
+                    is_bidirectional=True if self.phase == "image_prefill" else False,  # FIXME, Hard-coded for Gemma3.
+                    mask=attn_mask if self.use_position_ids else None,
                 )
 
         attn_output = attn_output.view(batch_size, self.num_heads, -1, self.head_dim)
@@ -926,10 +949,13 @@ class RotaryEmbedding(nn.Module):
 
 
 class DecoderOnlyFlashAttention(DecoderOnlyAttention):
-    def __init__(self, self_attn, kvcache_partition_len, kvcache_block_size, use_attention_mask):
+    def __init__(self, self_attn, kvcache_partition_len, kvcache_block_size, use_attention_mask, use_position_ids):
         self.kvcache_partition_size = kvcache_partition_len
         super().__init__(
-            self_attn=self_attn, use_attention_mask=use_attention_mask, kvcache_block_size=kvcache_block_size
+            self_attn=self_attn,
+            use_attention_mask=use_attention_mask,
+            use_position_ids=use_position_ids,
+            kvcache_block_size=kvcache_block_size,
         )
 
     def get_attention(self):
@@ -939,6 +965,7 @@ class DecoderOnlyFlashAttention(DecoderOnlyAttention):
             self.num_key_value_heads,
             self.kvcache_partition_size,
             self.use_attention_mask,
+            self.use_position_ids,
         )
 
     def forward(
@@ -990,12 +1017,14 @@ class FlashAttentionOp(AttentionOp):
         num_key_value_heads: int,
         kvcache_partition_len: int,
         use_attention_mask: bool,
+        use_position_ids: bool,
     ):
         super().__init__(
             num_heads=num_heads,
             head_dim=head_dim,
             num_key_value_heads=num_key_value_heads,
             use_attention_mask=use_attention_mask,
+            use_position_ids=use_position_ids,
         )
         self.kvcache_partition_size = kvcache_partition_len
 
@@ -1015,7 +1044,7 @@ class FlashAttentionOp(AttentionOp):
         # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
         key_state = key_state.unsqueeze(2)
         value_state = value_state.unsqueeze(2)
-        if self.use_attention_mask:
+        if self.use_attention_mask and not self.use_position_ids:
             attn_mask = attn_mask.unsqueeze(2)
 
         if self.phase == "decode":
@@ -1032,7 +1061,7 @@ class FlashAttentionOp(AttentionOp):
         )
 
         if self.phase == "decode":
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_flash_attn_decode(
                     q=query_state,
                     k=key_state,
@@ -1058,9 +1087,10 @@ class FlashAttentionOp(AttentionOp):
                     block_table=block_tables,
                     block_size=kvcache_block_size,
                     partition=self.kvcache_partition_size,
+                    mask=attn_mask if self.use_position_ids else None,
                 )
         else:
-            if self.use_attention_mask:
+            if self.use_attention_mask and not self.use_position_ids:
                 attn_output = torch.ops.rbln_custom_ops.paged_flash_attn_prefill(
                     q=query_state,
                     k=key_state,
@@ -1086,6 +1116,8 @@ class FlashAttentionOp(AttentionOp):
                     block_table=block_tables,
                     block_size=kvcache_block_size,
                     partition=self.kvcache_partition_size,
+                    is_bidirectional=True if self.phase == "image_prefill" else False,
+                    mask=attn_mask if self.use_position_ids else None,
                 )
 
         # reshape for removing repeat_kv
@@ -1094,3 +1126,70 @@ class FlashAttentionOp(AttentionOp):
         attn_output = attn_output.reshape(batch_size, -1, self.num_heads * self.head_dim)
 
         return attn_output
+
+
+class SlidingWindowAttentionOp(AttentionOp):
+    def forward(
+        self,
+        query_state: torch.Tensor,
+        key_state: torch.Tensor,
+        value_state: torch.Tensor,
+        attn_mask: torch.Tensor,
+        past_key_state: torch.Tensor,
+        past_value_state: torch.Tensor,
+        seq_position: Tuple[torch.Tensor],
+        scale: torch.Tensor,
+        block_tables: torch.Tensor,
+        block_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
+        key_state = key_state.unsqueeze(2)
+        value_state = value_state.unsqueeze(2)
+
+        if self.phase == "decode":
+            batch_size = key_state.shape[0]
+        else:
+            batch_size = 1
+
+        query_state = query_state.view(
+            batch_size,
+            self.num_key_value_heads,
+            self.num_heads // self.num_key_value_heads,
+            -1,  # seq len
+            self.head_dim,
+        )
+
+        if self.phase == "decode":
+            attn_output = torch.ops.rbln_custom_ops.paged_sliding_window_attn_decode(
+                q=query_state,
+                k=key_state,
+                v=value_state,
+                kcache=past_key_state.unsqueeze(2),
+                vcache=past_value_state.unsqueeze(2),
+                cache_seq_len=seq_position[0],
+                cache_offset=seq_position[1],
+                scale=scale,
+                block_table=block_tables,
+                block_size=block_size,
+            )
+        else:
+            attn_output = torch.ops.rbln_custom_ops.paged_sliding_window_attn_prefill(
+                q=query_state,
+                k=key_state,
+                v=value_state,
+                kcache=past_key_state.unsqueeze(2),
+                vcache=past_value_state.unsqueeze(2),
+                cache_seq_len=seq_position[0],
+                cache_offset=seq_position[1],
+                scale=scale,
+                block_table=block_tables,
+                block_size=block_size,
+                is_bidirectional=True if self.phase == "image_prefill" else False,
+            )
+
+        # reshape for removing repeat_kv
+        attn_output = attn_output.view(batch_size, self.num_heads, -1, self.head_dim)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, -1, self.num_heads * self.head_dim)
+
+        return attn_output