Update modeling_sliding_llama.py

modeling_sliding_llama.py

@@ -54,15 +54,22 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "LlamaConfig"
 
-def attn_causal(b, h, q_idx, kv_idx):
+def attn_causal(b, h, q_idx, kv_idx, q_len):
     causal_mask = q_idx >= kv_idx
-    return causal_mask
+    if q_len > 1:
+        return causal_mask
+    else:
+        return q_idx >= -1
 
-def stream_attn_causal(b, h, q_idx, kv_idx, attn_sink_size, sliding_window):
-    causal_mask = q_idx >= kv_idx
-    window_mask = q_idx - kv_idx <= sliding_window 
+def stream_attn_causal(b, h, q_idx, kv_idx, attn_sink_size, sliding_window, q_len, kv_len):
+    causal_mask = q_idx >= kv_idx 
     sink_mask = kv_idx < attn_sink_size
-    return causal_mask & (window_mask | sink_mask)
+    if q_len > 1:
+        window_mask = q_idx - kv_idx <= sliding_window
+        return causal_mask & (window_mask | sink_mask)
+    else:
+        window_mask = kv_len - kv_idx <= sliding_window
+        return window_mask | sink_mask
 
 
 class LlamaRMSNorm(nn.Module):
@@ -343,17 +350,19 @@ class LlamaStreamingFlexAttention(LlamaAttention):
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
+
         sliding_window_size = self.config.sliding_windows[self.layer_idx]
         if sliding_window_size > 0:
             block_mask = create_block_mask(
-                lambda b, h, q_idx, kv_idx: stream_attn_causal(b, h, q_idx, kv_idx, 4, sliding_window_size),
-                B=None, H=None, Q_LEN=query_states.shape[-2], KV_LEN=key_states.shape[-2]
+                lambda b, h, q_idx, kv_idx: stream_attn_causal(b, h, q_idx, kv_idx, 4, sliding_window_size, q_len, key_states.shape[-2]),
+                B=None, H=None, Q_LEN=q_len, KV_LEN=key_states.shape[-2]
             )
         else:
             block_mask = create_block_mask(
-                lambda b, h, q_idx, kv_idx: attn_causal(b, h, q_idx, kv_idx),
-                B=None, H=None, Q_LEN=query_states.shape[-2], KV_LEN=key_states.shape[-2]
+                lambda b, h, q_idx, kv_idx: attn_causal(b, h, q_idx, kv_idx, q_len),
+                B=None, H=None, Q_LEN=q_len, KV_LEN=key_states.shape[-2]
             )
+
         attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask)
 
         if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
@@ -363,9 +372,9 @@ class LlamaStreamingFlexAttention(LlamaAttention):
             )
 
         attn_output = attn_output.transpose(1, 2).contiguous()
-
+        
         attn_output = attn_output.reshape(bsz, q_len, -1)
-
+        
         if self.config.pretraining_tp > 1:
             attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
             o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)