From 8b2d29176f7f1b9b72f9dff11a0848027d136910 Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Wed, 19 Aug 2026 04:58:57 +1000 Subject: [PATCH] Prs/ace graphs (#15698) --- comfy/text_encoders/ace15.py | 110 ++++++++++++++++++++++------------ comfy/text_encoders/gemma4.py | 16 ++--- comfy/text_encoders/llama.py | 37 ++++++++---- 3 files changed, 106 insertions(+), 57 deletions(-) diff --git a/comfy/text_encoders/ace15.py b/comfy/text_encoders/ace15.py index 853f021ae..3ad519314 100644 --- a/comfy/text_encoders/ace15.py +++ b/comfy/text_encoders/ace15.py @@ -4,9 +4,29 @@ from comfy import sd1_clip import torch import math import yaml +import comfy.ops import comfy.utils +def _audio_logits(model, x, audio_start, audio_end, eos_token=None): + input = x[:, -1:] + module = model.embed_tokens + + offload_stream = None + if module.comfy_cast_weights: + weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True) + else: + weight = module.weight.to(x) + + logits = torch.nn.functional.linear(input, weight[audio_start:audio_end], None)[:, -1] + eos_logits = None + if eos_token is not None: + eos_logits = torch.nn.functional.linear(input, weight[eos_token:eos_token + 1], None)[:, -1] + + comfy.ops.uncast_bias_weight(module, weight, None, offload_stream) + return logits, eos_logits + + def sample_manual_loop_no_classes( model, ids=None, @@ -34,48 +54,43 @@ def sample_manual_loop_no_classes( execution_dtype = torch.float32 embeds, attention_mask, num_tokens, embeds_info = model.process_tokens(ids, device) + embeds = embeds.to(execution_dtype) embeds_batch = embeds.shape[0] - output_audio_codes = [] - past_key_values = [] + output_audio_codes = torch.empty((max_new_tokens,), device=device, dtype=torch.long) + generated_tokens = 0 generator = torch.Generator(device=device) generator.manual_seed(seed) - model_config = model.transformer.model.config - past_kv_shape = [embeds_batch, model_config.num_key_value_heads, embeds.shape[1] + min_tokens, model_config.head_dim] - - for x in range(model_config.num_hidden_layers): - past_key_values.append((torch.empty(past_kv_shape, device=device, dtype=execution_dtype), torch.empty(past_kv_shape, device=device, dtype=execution_dtype), 0)) + past_key_values = model.transformer.model.init_kv_cache(embeds_batch, embeds.shape[1] + max_new_tokens, device, execution_dtype) + fixed_kv = isinstance(past_key_values[0], comfy.text_encoders.llama.FixedKV) progress_bar = comfy.utils.ProgressBar(max_new_tokens) + sampling_logits = None for step in comfy.utils.model_trange(max_new_tokens, desc="LM sampling"): - outputs = model.transformer(None, attention_mask, embeds=embeds.to(execution_dtype), num_tokens=num_tokens, intermediate_output=None, dtype=execution_dtype, embeds_info=embeds_info, past_key_values=past_key_values) - next_token_logits = model.transformer.logits(outputs[0])[:, -1] + outputs = model.transformer(None, attention_mask, embeds=embeds, num_tokens=num_tokens, intermediate_output=None, dtype=execution_dtype, embeds_info=embeds_info, past_key_values=past_key_values) past_key_values = outputs[2] - if cfg_scale != 1.0: - cond_logits = next_token_logits[0:1] - uncond_logits = next_token_logits[1:2] - cfg_logits = uncond_logits + cfg_scale * (cond_logits - uncond_logits) - else: - cfg_logits = next_token_logits[0:1] - use_eos_score = eos_token_id is not None and eos_token_id < audio_start_id and min_tokens < step - if use_eos_score: - eos_score = cfg_logits[:, eos_token_id].clone() + audio_logits, eos_logits = _audio_logits(model.transformer.model, outputs[0], audio_start_id, audio_end_id, eos_token_id if use_eos_score else None) + if cfg_scale != 1.0: + cfg_logits = audio_logits[1:2] + cfg_scale * (audio_logits[0:1] - audio_logits[1:2]) + if use_eos_score: + cond_eos = eos_logits[0:1, 0] + uncond_eos = eos_logits[1:2, 0] + eos_score = uncond_eos + cfg_scale * (cond_eos - uncond_eos) + else: + cfg_logits = audio_logits[0:1] + if use_eos_score: + eos_score = eos_logits[0:1, 0] remove_logit_value = torch.finfo(cfg_logits.dtype).min - # Only generate audio tokens - cfg_logits[:, :audio_start_id] = remove_logit_value - cfg_logits[:, audio_end_id:] = remove_logit_value - if use_eos_score: - cfg_logits[:, eos_token_id] = eos_score + cfg_logits = torch.cat((eos_score.unsqueeze(1), cfg_logits), dim=1) if top_k is not None and top_k > 0: - top_k_vals, _ = torch.topk(cfg_logits, top_k) - min_val = top_k_vals[..., -1, None] - cfg_logits[cfg_logits < min_val] = remove_logit_value + top_k_values = torch.topk(cfg_logits, min(top_k, cfg_logits.shape[-1])).values + cfg_logits[cfg_logits < top_k_values[..., -1, None]] = remove_logit_value if min_p is not None and min_p > 0: probs = torch.softmax(cfg_logits, dim=-1) @@ -89,28 +104,40 @@ def sample_manual_loop_no_classes( sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 - indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + indices_to_remove = torch.zeros_like(cfg_logits, dtype=torch.bool) + indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove) cfg_logits[indices_to_remove] = remove_logit_value if temperature > 0: cfg_logits = cfg_logits / temperature - next_token = torch.multinomial(torch.softmax(cfg_logits, dim=-1), num_samples=1, generator=generator).squeeze(1) + if sampling_logits is None: + sampling_logits = cfg_logits.new_empty((cfg_logits.shape[0], model.transformer.model.vocab_size)) + sampling_logits.fill_(remove_logit_value) + if use_eos_score: + sampling_logits[:, eos_token_id] = cfg_logits[:, 0] + cfg_logits = cfg_logits[:, 1:] + sampling_logits[:, audio_start_id:audio_end_id] = cfg_logits + next_token = torch.multinomial(torch.softmax(sampling_logits, dim=-1), num_samples=1, generator=generator).squeeze(1) else: next_token = torch.argmax(cfg_logits, dim=-1) + if use_eos_score: + next_token = torch.where(next_token == 0, eos_token_id, next_token + audio_start_id - 1) + else: + next_token += audio_start_id - token = next_token.item() - - if token == eos_token_id: + if eos_token_id is not None and next_token.item() == eos_token_id: break - embed, _, _, _ = model.process_tokens([[token]], device) - embeds = embed.repeat(embeds_batch, 1, 1) - attention_mask = torch.cat([attention_mask, torch.ones((embeds_batch, 1), device=device, dtype=attention_mask.dtype)], dim=1) + input_ids = next_token.repeat(embeds_batch).unsqueeze(1) + embeds = model.transformer.get_input_embeddings()(input_ids, out_dtype=execution_dtype) + if not fixed_kv: + attention_mask = torch.cat([attention_mask, torch.ones((embeds_batch, 1), device=device, dtype=attention_mask.dtype)], dim=1) - output_audio_codes.append(token - audio_start_id) + output_audio_codes[generated_tokens] = next_token[0] - audio_start_id + generated_tokens += 1 progress_bar.update_absolute(step) - return output_audio_codes + return output_audio_codes[:generated_tokens].tolist() def generate_audio_codes(model, positive, negative, min_tokens=1, max_tokens=1024, seed=0, cfg_scale=2.0, temperature=0.85, top_p=0.9, top_k=0, min_p=0.000): @@ -286,7 +313,10 @@ class ACE15TEModel(torch.nn.Module): self.qwen3_06b = Qwen3_06BModel(device=device, dtype=dtype, model_options=model_options) if model is not None: setattr(self, self.lm_model, model(device=device, dtype=dtype_llama, model_options=model_options)) - + ar_model = getattr(self, self.lm_model) + ar_model.transformer.model.fixed_kv = True + ar_model.transformer.model.prefetch_dynamic_vbars = True + ar_model.transformer.model.graph_dynamic_vbar_blocks = True self.dtypes = set([dtype, dtype_llama]) def encode_token_weights(self, token_weight_pairs): @@ -319,6 +349,12 @@ class ACE15TEModel(torch.nn.Module): if lm_model is not None: lm_model.reset_clip_options() + def get_dynamic_vram__units(self): + if self.lm_model is None: + return ([], []) + model = getattr(self, self.lm_model) + return model.transformer.model.get_dynamic_vram__units() + def load_sd(self, sd): if "model.layers.0.post_attention_layernorm.weight" in sd: shape = sd["model.layers.0.post_attention_layernorm.weight"].shape diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index da116ab27..0d8f0fcc7 100644 --- a/comfy/text_encoders/gemma4.py +++ b/comfy/text_encoders/gemma4.py @@ -226,10 +226,11 @@ class Gemma4Attention(nn.Module): present_key_value = None fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None if fixed_cache is not None: - if seq_length == 1: + if seq_length == 1 and fixed_cache.index > 0: # CUDA-graphable decode: write at the device-side ring/linear position - fixed_cache.key.index_copy_(2, fixed_cache.position, xk) - fixed_cache.value.index_copy_(2, fixed_cache.position, xv) + position = fixed_cache.position.view(batch_size, 1, 1, 1).expand_as(xk) + fixed_cache.key.scatter_(2, position, xk) + fixed_cache.value.scatter_(2, position, xv) output = self._decode_attention(xq, fixed_cache, attention_mask) return self.o_proj(output), fixed_cache, None @@ -515,7 +516,7 @@ class Gemma4Transformer(nn.Module): fixed_kv = (past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV)) - decode = fixed_kv and seq_len == 1 + decode = fixed_kv and past_len > 0 and seq_len == 1 # mirror the conditions under which prefetch_queue_pop can actually capture, so # eager fallbacks keep the sliced decode path instead of the full-capacity one enable_graph = (decode and mask is None and self.graph_dynamic_vbar_blocks @@ -525,12 +526,13 @@ class Gemma4Transformer(nn.Module): and comfy.model_management.is_device_cuda(x.device)) decode_bias = None decode_masks = None - if decode: + if fixed_kv: prepared = set() for kv in past_key_values: if isinstance(kv, FixedKV) and id(kv.position) not in prepared: kv.prepare(seq_len) prepared.add(id(kv.position)) + if decode: if mask is not None: decode_masks = {} for kv in past_key_values: @@ -705,8 +707,8 @@ class Gemma4Base(BaseLlama, BaseGenerate, torch.nn.Module): cache_cls = RingKV if sliding else FixedKV tracker = trackers.get((cache_cls, length)) if tracker is None: - tracker = (torch.empty((1,), device=device, dtype=torch.int64), - torch.empty((batch,), device=device, dtype=torch.int32)) + tracker = (torch.empty((batch,), device=device, dtype=torch.int64), + torch.zeros((batch,), device=device, dtype=torch.int32)) trackers[(cache_cls, length)] = tracker # zero-init: decode attends full capacity with masked tails, 0*0 stays finite key = torch.zeros((batch, kv_heads, length, head_dim), device=device, dtype=execution_dtype) diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index f182e5147..49ba5dfa9 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -26,8 +26,8 @@ class FixedKV: seqlen: torch.Tensor def prepare(self, num_tokens): - self.position.fill_(self.index) - self.seqlen.fill_(self.index + num_tokens) + self.position.copy_(self.seqlen) + self.seqlen.add_(num_tokens) def advance(self, num_tokens): self.index += num_tokens @@ -571,17 +571,25 @@ class Attention(nn.Module): xq = xq.transpose(1, 2) xk = xk.transpose(1, 2) xv = xv.transpose(1, 2) - if seq_length == 1: + if seq_length == 1 and fixed_cache.index > 0: # CUDA-graphable decode path. - fixed_cache.key.index_copy_(1, fixed_cache.position, xk) - fixed_cache.value.index_copy_(1, fixed_cache.position, xv) + position = fixed_cache.position.view(batch_size, 1, 1, 1).expand_as(xk) + fixed_cache.key.scatter_(1, position, xk) + fixed_cache.value.scatter_(1, position, xv) output = comfy_kitchen.flash_attention_decode(xq, fixed_cache.key, fixed_cache.value, fixed_cache.seqlen) return self.o_proj(output.view(batch_size, seq_length, self.inner_size)), fixed_cache - fixed_cache.key[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xk) - fixed_cache.value[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xv) - xk = fixed_cache.key[:, :fixed_cache.index + seq_length] - xv = fixed_cache.value[:, :fixed_cache.index + seq_length] + if attention_mask is None or attention_mask.ndim < 4: + fixed_cache.key[:, :seq_length].copy_(xk) + fixed_cache.value[:, :seq_length].copy_(xv) + else: + valid = attention_mask[:, 0, -1, -seq_length:] == 0 + indices = torch.arange(seq_length, device=xk.device).expand(batch_size, -1) + indices = indices.masked_fill(~valid, seq_length).sort(dim=1).values.clamp_max_(seq_length - 1) + indices = indices.view(batch_size, seq_length, 1, 1).expand_as(xk) + fixed_cache.key[:, :seq_length].copy_(xk.gather(1, indices)) + fixed_cache.value[:, :seq_length].copy_(xv.gather(1, indices)) + fixed_cache.seqlen.copy_(valid.sum(dim=1)) xq = xq.transpose(1, 2) xk = xk.transpose(1, 2) @@ -796,8 +804,8 @@ class Llama2_(nn.Module): if fixed_kv: key = torch.empty((batch, capacity, self.config.num_key_value_heads, self.config.head_dim), device=device, dtype=dtype) value = torch.empty_like(key) - position = torch.empty((1,), device=device, dtype=torch.int64) - seqlen = torch.empty((batch,), device=device, dtype=torch.int32) + position = torch.empty((batch,), device=device, dtype=torch.int64) + seqlen = torch.zeros((batch,), device=device, dtype=torch.int32) caches.append(FixedKV(key, value, 0, position, seqlen)) else: key = torch.empty((batch, self.config.num_key_value_heads, capacity, self.config.head_dim), device=device, dtype=dtype) @@ -824,6 +832,10 @@ class Llama2_(nn.Module): past_len = 0 if past_key_values is not None and len(past_key_values) > 0: past_len = self.get_past_len(past_key_values) + fixed_kv = past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV) + fixed_kv_decode = fixed_kv and past_len > 0 and seq_len == 1 + if fixed_kv_decode: + attention_mask = None if position_ids is None: position_ids = torch.arange(past_len, past_len + seq_len, device=x.device).unsqueeze(0) @@ -844,8 +856,7 @@ class Llama2_(nn.Module): optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True) - fixed_kv = past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV) - enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv and seq_len == 1 and mask is None + enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv_decode if enable_graph: freqs_cis_groups = freqs_cis if isinstance(freqs_cis, list) else [freqs_cis] cross_step_state_key = [(x.shape, x.stride(), x.dtype, x.device)]