Update attn_mask_utils.py

attn_mask_utils.py

@@ -30,7 +30,6 @@ def find_prefix_seq_length_by_pe(
     return seq_len
 
 
-
 def update_causal_mask_with_pad_non_visible_2d(
     input_ids: torch.Tensor,
     attn_mask_2d: torch.Tensor,
@@ -96,6 +95,58 @@ def update_causal_mask_with_pad_non_visible_2d(
     return attn_mask_2d
 
 
+def update_causal_mask_with_pad_non_visible_2d_for_ssd_cache(
+    input_ids: torch.Tensor,
+    attn_mask_2d: torch.Tensor,
+    block_size: int = 4,
+    use_cache: bool = True,
+    causal_attn: bool = False
+) -> torch.Tensor:
+    """
+    Updates a 2D attention mask for Self-Speculative Decoding generate
+
+    Details is avaliabe in Appendix B Figure 5. 
+
+    Args:
+        input_ids: Input token IDs (unused in current implementation)
+        attn_mask_2d: 2D attention mask matrix of shape [seq_len, seq_len] where:
+            - 0.0 indicates allowed attention
+            - -inf indicates masked attention
+        block_size: Size of the diffusion window
+        use_cache: Whether key-value cache is being used
+        causal_attn: If True, maintains strict causal masking throughout
+    
+    Returns:
+        Modified attention mask with updated visibility patterns
+    """
+
+    q_len, kv_len = attn_mask_2d.shape
+
+    if q_len == kv_len:
+        # prefill
+        return update_causal_mask_for_one_gen_window_2d(
+            input_ids = input_ids,
+            attn_mask_2d = attn_mask_2d,
+            block_size = block_size,
+            use_cache = use_cache,
+            causal_attn = causal_attn
+        )
+
+    start_ix = q_len - block_size
+    start_jx = kv_len - block_size
+    for ix in range(block_size-1, -1, -1):
+        attn_mask_2d[start_ix:start_ix+block_size, start_jx:start_jx+block_size] = 0.0
+        attn_mask_2d[start_ix+block_size:, start_jx-ix:start_jx+block_size] = -float('inf')
+        
+        start_ix = start_ix - ix - block_size
+        start_jx = start_jx - ix - block_size 
+
+    attn_mask_2d[start_ix+block_size:, start_jx+block_size-1] = -float('inf')
+
+    return attn_mask_2d
+
+
+
 def update_causal_mask_for_one_gen_window_2d(
     input_ids: torch.Tensor,
     attn_mask_2d: torch.Tensor,