diff --git a/comfy/text_encoders/anima_cache.py b/comfy/text_encoders/anima_cache.py index 8392acd17..4a194bb56 100644 --- a/comfy/text_encoders/anima_cache.py +++ b/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 = [] - next_key_values = past_key_values - for index in range(suffix_embeds.shape[1]): - output = transformer( - None, - None, - embeds=suffix_embeds[:, index:index + 1], - num_tokens=[1], - intermediate_output=intermediate_output, - final_layer_norm_intermediate=final_layer_norm_intermediate, - dtype=dtype, - embeds_info=embeds_info, - past_key_values=next_key_values, - ) - suffix_outputs.append(output[0]) - next_key_values = output[2] - suffix_hidden = torch.cat(suffix_outputs, dim=1) + output = transformer( + None, + None, + embeds=suffix_embeds, + num_tokens=[suffix_embeds.shape[1]], + intermediate_output=intermediate_output, + final_layer_norm_intermediate=final_layer_norm_intermediate, + dtype=dtype, + embeds_info=embeds_info, + past_key_values=past_key_values, + ) + suffix_hidden = output[0] + next_key_values = output[2] else: output = transformer( None, diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index 3f98fb0a5..01b65a1cf 100644 --- a/comfy/text_encoders/llama.py +++ b/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: diff --git a/tests-unit/comfy_test/test_anima_cache.py b/tests-unit/comfy_test/test_anima_cache.py index 7a3a036b4..e0f8a438d 100644 --- a/tests-unit/comfy_test/test_anima_cache.py +++ b/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.num_tokens == [(3,), (1,), (1,)] + assert transformer.calls == [(3, True), (2, True)] + 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):