Batch cached Anima Qwen suffixes

comfy/text_encoders/anima_cache.py

@@ -90,23 +90,19 @@ def forward(transformer, cache_owner, tokens, attention_mask, embeds, num_tokens
 
     suffix_embeds = embeds[:, common:]
     if common > 0 and suffix_embeds.shape[1] > 1:
-        suffix_outputs = []
+        output = transformer(
-        next_key_values = past_key_values
+            None,
-        for index in range(suffix_embeds.shape[1]):
+            None,
-            output = transformer(
+            embeds=suffix_embeds,
-                None,
+            num_tokens=[suffix_embeds.shape[1]],
-                None,
+            intermediate_output=intermediate_output,
-                embeds=suffix_embeds[:, index:index + 1],
+            final_layer_norm_intermediate=final_layer_norm_intermediate,
-                num_tokens=[1],
+            dtype=dtype,
-                intermediate_output=intermediate_output,
+            embeds_info=embeds_info,
-                final_layer_norm_intermediate=final_layer_norm_intermediate,
+            past_key_values=past_key_values,
-                dtype=dtype,
+        )
-                embeds_info=embeds_info,
+        suffix_hidden = output[0]
-                past_key_values=next_key_values,
+        next_key_values = output[2]
-            )
-            suffix_outputs.append(output[0])
-            next_key_values = output[2]
-        suffix_hidden = torch.cat(suffix_outputs, dim=1)
     else:
         output = transformer(
             None,

comfy/text_encoders/llama.py

@@ -737,7 +737,7 @@ class Llama2_(nn.Module):
             mask = mask.masked_fill(mask.to(torch.bool), torch.finfo(x.dtype).min / 4)
 
         if seq_len > 1:
-            causal_mask = torch.empty(past_len + seq_len, past_len + seq_len, dtype=x.dtype, device=x.device).fill_(torch.finfo(x.dtype).min / 4).triu_(1)
+            causal_mask = torch.empty(seq_len, past_len + seq_len, dtype=x.dtype, device=x.device).fill_(torch.finfo(x.dtype).min / 4).triu_(past_len + 1)
             if mask is not None:
                 mask += causal_mask
             else:

tests-unit/comfy_test/test_anima_cache.py

@@ -3,7 +3,7 @@ import asyncio
 import pytest
 import torch
 
-from comfy.text_encoders import anima_cache
+from comfy.text_encoders import anima_cache, llama
 
 
 def tokens(*ids):
@@ -18,8 +18,6 @@ class TinyCausalTransformer:
     def __call__(self, _, attention_mask, embeds, num_tokens, intermediate_output, final_layer_norm_intermediate, dtype, embeds_info, past_key_values=None):
         self.calls.append((embeds.shape[1], past_key_values is not None))
         self.num_tokens.append(tuple(num_tokens))
-        if past_key_values and embeds.shape[1] > 1:
-            raise RuntimeError("cached multi-token suffix would use an invalid causal mask")
         prefix = 0
         if past_key_values:
             prefix = past_key_values[0][2]
@@ -124,8 +122,48 @@ def test_diverging_suffix_matches_full_causal_forward(prefix_cache):
     output, _ = cached_forward(transformer, [1, 2, 4, 5], owner=owner)
 
     assert torch.equal(output, torch.tensor([[[1.0], [3.0], [7.0], [12.0]]]))
-    assert transformer.calls == [(3, True), (1, True), (1, True)]
+    assert transformer.calls == [(3, True), (2, True)]
-    assert transformer.num_tokens == [(3,), (1,), (1,)]
+    assert transformer.num_tokens == [(3,), (2,)]
+    cached = prefix_cache[owner]
+    assert cached[2] == (1, 2, 4, 5)
+    assert torch.equal(cached[3], output)
+    key, value, length = cached[4][0]
+    assert length == 4
+    assert torch.equal(key.flatten(), torch.tensor([1.0, 2.0, 4.0, 5.0]))
+    assert torch.equal(value.flatten(), torch.tensor([2.0, 4.0, 8.0, 10.0]))
+
+
+def test_llama_cached_multi_token_causal_mask_uses_absolute_positions(monkeypatch):
+    masks = []
+
+    class Layer:
+        def __call__(self, x, attention_mask, freqs_cis, optimized_attention, past_key_value):
+            masks.append(attention_mask.clone())
+            return x, None
+
+    class Model:
+        layers = (Layer(),)
+        norm = None
+
+        def get_past_len(self, past_key_values):
+            return past_key_values[0][2]
+
+        def compute_freqs_cis(self, position_ids, device):
+            return None
+
+    monkeypatch.setattr(llama, "optimized_attention_for_device", lambda *args, **kwargs: None)
+    x = torch.zeros((1, 2, 1))
+    past_key_values = [(torch.empty((1, 1, 3, 1)), torch.empty((1, 1, 3, 1)), 3)]
+
+    llama.Llama2_.forward(Model(), None, embeds=x, past_key_values=past_key_values)
+
+    blocked = torch.finfo(x.dtype).min / 4
+    expected = torch.tensor([
+        [0.0, 0.0, 0.0, 0.0, blocked],
+        [0.0, 0.0, 0.0, 0.0, 0.0],
+    ])
+    assert masks[0].shape == (2, 5)
+    assert torch.equal(masks[0], expected)
 
 
 def test_cache_is_isolated_by_owner_and_transformer_identity(prefix_cache):