zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-09-09 02:26:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

Introduce Comfy Compiler (CORE-389) (#15861)

Browse Source
This commit is contained in:
rattus
2026-09-05 03:18:39 +10:00
committed by GitHub
parent e3a78dab9e
commit 804eb5513a
10 changed files with 198 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfy/cli_args.py
+5
View File
@@ -181,6 +181,8 @@ parser.add_argument("--disable-dynamic-vram", action="store_true", help="Disable
parser.add_argument("--enable-dynamic-vram", action="store_true", help="Enable dynamic VRAM on systems where it's not enabled by default.")
parser.add_argument("--fast-disk", action="store_true", help="Prefer disk-backed dynamic loading and offload over unpinned RAM. Can be faster for users with fast NVME disks.")
parser.add_argument("--disable-cuda-graphs", action="store_true", help="Disable CUDA graphs.")
parser.add_argument("--disable-comfy-compiler", action="store_true", help="Disable the Comfy model compiler, including its CUDA graph subfeature.")
parser.add_argument("--assert-graph-breaks", action="store_true", help="Fail on Comfy model compiler graph breaks.")
parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")
@@ -291,6 +293,9 @@ if args.windows_standalone_build:
if args.disable_auto_launch:
args.auto_launch = False
if args.disable_comfy_compiler:
args.disable_cuda_graphs = True
if args.force_fp16:
args.fp16_unet = True
comfy/ldm/lightricks/av_model.py
+8 -2
View File
@@ -938,8 +938,11 @@ class LTXAVModel(LTXVModel):
stg_self_attn_blocks = transformer_options.get("stg_self_attn_blocks", ())
# Process transformer blocks
comfy.model_prefetch.malloc_graph_begin(self, vx.device)
for i, block in enumerate(self.transformer_blocks):
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, vx.device, block)
comfy.model_prefetch.prefetch_queue_pop(
prefetch_queue, vx.device, block, malloc_scope="block"
)
block_transformer_options = transformer_options
if i in stg_self_attn_blocks:
block_transformer_options = {**transformer_options, "stg_skip_self_attn": True}
@@ -1015,7 +1018,10 @@ class LTXAVModel(LTXVModel):
a_prompt_timestep=a_prompt_timestep,
)
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, vx.device, None)
comfy.model_prefetch.prefetch_queue_pop(
prefetch_queue, vx.device, None, malloc_scope="block"
)
comfy.model_prefetch.malloc_graph_end()
return [vx, ax]
comfy/ldm/minimax/model.py
+14 -4
View File
@@ -563,12 +563,23 @@ class MiniMaxH3Model(nn.Module):
carry = (sigma_a / sigma_v).to(audio_src.dtype)
x = [x[0], audio_src * carry]
out = comfy.patcher_extension.WrapperExecutor.new_class_executor(
compile_allocations = comfy.model_prefetch.malloc_graph_enabled(x[0].device)
if compile_allocations:
out = [torch.empty_like(x[0]), torch.empty_like(x[1])]
comfy.model_prefetch.malloc_graph_begin(self, x[0].device)
graph_out = comfy.patcher_extension.WrapperExecutor.new_class_executor(
self._forward,
self,
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
).execute(x, timestep, context, transformer_options, minimax_payload=minimax_payload,
denoise_mask=denoise_mask, audio_denoise_mask=audio_denoise_mask, **kwargs)
if compile_allocations:
out[0].copy_(graph_out[0])
out[1].copy_(graph_out[1])
del graph_out
comfy.model_prefetch.malloc_graph_end()
else:
out = graph_out
if scale != 1.0:
# d/d(sigma_v) of the carried variable
@@ -724,7 +735,7 @@ class MiniMaxH3Model(nn.Module):
blocks_replace = patches_replace.get("dit", {})
prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), device, transformer_options)
for i, block in enumerate(self.blocks):
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block, malloc_scope="block")
if ("double_block", i) in blocks_replace:
def block_wrap(args):
return {"img": block(args["img"], args["t_emb"], args["mod_segments"], args["rope_freqs"],
@@ -735,8 +746,7 @@ class MiniMaxH3Model(nn.Module):
{"original_block": block_wrap})["img"]
else:
h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)
if prefetch_queue is not None:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None)
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None, malloc_scope="block")
# target streams are single contiguous segments (audio then video, last two)
va, vb, _ = next(s for s in layout.segments if s[2] == "video")
comfy/ldm/minimax_music/ar.py
+25 -13
View File
@@ -251,8 +251,13 @@ class MiniMaxMusic3AR(nn.Module):
decode_limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES)
past = self.model.init_kv_cache(2, prompt_tokens + decode_limit + 1, device, execution_dtype)
output = self.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype)
last_hidden = output[0][:, -1]
last_hidden = output[0][:, -1].clone()
past = output[2]
del output
vbar = getattr(self, "dynamic_vbars", {}).get(device)
if vbar is not None:
comfy.model_management.reset_cast_buffers()
vbar.set_watermark(vbar.max_size)
generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar"))
decoder = self.model.audio_decoder
@@ -263,13 +268,12 @@ class MiniMaxMusic3AR(nn.Module):
"codes": torch.empty((last_hidden.shape[0], self.num_codebooks), dtype=torch.long, device=device),
"depth_hidden": torch.empty((1, last_hidden.shape[-1] * (self.num_codebooks - 1)), dtype=execution_dtype, device=device),
}
decoder._comfy_cross_step_state = depth_io
comfy.model_management._register_cross_step(decoder)
hidden_frames = []
pending_code = None
stop_token = None
pending_event = None
pending_hidden = None
pending_hidden = torch.empty(last_hidden.shape[-1] * self.num_codebooks, dtype=execution_dtype, device=device)
pending_hidden_valid = False
progress = comfy.utils.ProgressBar(decode_limit)
cuda_device = torch.device(device).type == "cuda"
vocab_mask = None
@@ -284,14 +288,16 @@ class MiniMaxMusic3AR(nn.Module):
if pending_event is not None:
pending_event.synchronize()
if int(pending_code.item()) == stop_token:
pending_hidden = None
pending_hidden_valid = False
break
if pending_hidden is not None:
hidden_frames.append(pending_hidden)
if pending_hidden_valid:
hidden_frames.append(pending_hidden.clone())
progress.update_absolute(len(hidden_frames))
if len(hidden_frames) >= decode_limit:
break
if frame_index:
comfy.model_prefetch.malloc_graph_begin(self, device)
c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask)
if pending_code is None:
pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device)
@@ -318,25 +324,31 @@ class MiniMaxMusic3AR(nn.Module):
[[decoder, self.model.audio_extra_embedding]], device, {"prefetch_dynamic_vbars": True}
)
comfy.model_prefetch.prefetch_queue_pop(
depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True, generator=generator
depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True,
generator=generator, malloc_scope="depth"
)
comfy.model_prefetch.prefetch_queue_pop(
depth_queue, device, None, malloc_scope="depth"
)
comfy.model_prefetch.prefetch_queue_pop(depth_queue, device, None)
feedback_codes = depth_io["codes"]
depth_hidden = depth_io["depth_hidden"]
frame_hidden = torch.cat((last_hidden[:1].detach(), depth_hidden), dim=-1)
if frame_index > 0:
pending_hidden = frame_hidden[0].clone()
pending_hidden.copy_(frame_hidden[0])
pending_hidden_valid = True
feedback = self._embed_audio_frame(feedback_codes, execution_dtype)
output = self.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
last_hidden = output[0][:, -1]
last_hidden.copy_(output[0][:, -1])
past = output[2]
del output, feedback, frame_hidden, depth_hidden, feedback_codes, c0_embed, c0, code_or_stop
comfy.model_prefetch.malloc_graph_end()
if pending_hidden is not None and len(hidden_frames) < decode_limit:
if pending_hidden_valid and len(hidden_frames) < decode_limit:
if pending_event is not None:
pending_event.synchronize()
if int(pending_code.item()) != stop_token:
hidden_frames.append(pending_hidden)
hidden_frames.append(pending_hidden.clone())
if not hidden_frames:
raise ValueError("MiniMax Music3 generated zero audio frames")
comfy/model_management.py
+4
View File
@@ -1373,6 +1373,7 @@ LARGEST_CASTED_WEIGHT = (None, 0)
STREAM_AIMDO_CAST_BUFFERS = {}
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
CROSS_STEP_STATE = weakref.WeakSet()
MALLOC_GRAPH_MODULES = weakref.WeakSet()
DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
@@ -1458,6 +1459,9 @@ def reset_cast_buffers():
STREAM_CAST_BUFFERS.clear()
STREAM_AIMDO_CAST_BUFFERS.clear()
for module in MALLOC_GRAPH_MODULES:
del module._comfy_malloc_graph
MALLOC_GRAPH_MODULES.clear()
soft_empty_cache()
def get_offload_stream(device):
comfy/model_prefetch.py
+79 -8
View File
@@ -1,7 +1,11 @@
import torch
import logging
import threading
import warnings
import weakref
import torch
import comfy_aimdo.malloc_graph
import comfy_aimdo.model_vbar
from comfy.cli_args import args
import comfy.memory_management
@@ -12,6 +16,41 @@ PREFETCH_QUEUES = []
GRAPH_MODULES = weakref.WeakSet()
GRAPH_WARMED_MODULES = weakref.WeakSet()
GRAPH_CAPTURE_STREAMS = {}
ACTIVE_MALLOC_GRAPHS = {}
MALLOC_GRAPH_BREAKS = 0
MALLOC_GRAPH_USED = False
def _malloc_graph_break():
global MALLOC_GRAPH_BREAKS
MALLOC_GRAPH_BREAKS += 1
logging.debug("Comfy model compiler graph break")
def malloc_graph_enabled(device):
return not args.disable_comfy_compiler and comfy.memory_management.aimdo_enabled and comfy.model_management.is_device_cuda(device)
def malloc_graph_begin(module, device):
global MALLOC_GRAPH_USED
if not malloc_graph_enabled(device):
return
graph = getattr(module, "_comfy_malloc_graph", None)
if graph is None:
graph = comfy_aimdo.malloc_graph.record(
comfy.model_management.current_stream(device), args.assert_graph_breaks
)
module._comfy_malloc_graph = graph
comfy.model_management.MALLOC_GRAPH_MODULES.add(module)
else:
graph.push()
ACTIVE_MALLOC_GRAPHS[threading.get_ident()] = graph
MALLOC_GRAPH_USED = True
def malloc_graph_end():
thread_id = threading.get_ident()
graph = ACTIVE_MALLOC_GRAPHS.get(thread_id)
if graph is not None:
if graph.pop():
_malloc_graph_break()
ACTIVE_MALLOC_GRAPHS.pop(thread_id)
def cleanup_prefetched_modules(module, comfy_modules):
for s in comfy_modules:
@@ -42,8 +81,13 @@ def _drop_graph(module):
del module._comfy_graph
def cleanup_prefetch_queues():
global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS
global PREFETCH_QUEUES
global MALLOC_GRAPH_BREAKS
global MALLOC_GRAPH_USED
graph = ACTIVE_MALLOC_GRAPHS.pop(threading.get_ident(), None)
if graph is not None:
graph.abort()
for queue in PREFETCH_QUEUES:
for entry in queue:
if entry is None or not isinstance(entry, tuple):
@@ -57,11 +101,18 @@ def cleanup_prefetch_queues():
_drop_graph(module)
GRAPH_MODULES.clear()
GRAPH_WARMED_MODULES.clear()
GRAPH_CAPTURE_STREAMS = {}
if MALLOC_GRAPH_USED:
logging.info("Comfy model compiler graph breaks: %d", MALLOC_GRAPH_BREAKS)
MALLOC_GRAPH_BREAKS = 0
MALLOC_GRAPH_USED = False
def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None):
enable_graph = enable_graph and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) and getattr(module, "_v_block", None) is not None
def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None, malloc_scope=None):
malloc_graph = ACTIVE_MALLOC_GRAPHS.get(threading.get_ident())
enable_graph = enable_graph and malloc_graph is not None and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) and getattr(module, "_v_block", None) is not None
if queue is None:
if malloc_graph is not None and malloc_scope is not None:
if malloc_graph.iterate(malloc_scope if module is not None else None):
_malloc_graph_break()
if core is not None:
core()
return
@@ -71,6 +122,13 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
capture_stream = GRAPH_CAPTURE_STREAMS.get(device)
if capture_stream is None:
capture_stream = torch.cuda.Stream(device=device)
# Keep PyTorch's persistent BLAS workspaces outside the allocation graph.
malloc_graph.pause()
with torch.cuda.stream(capture_stream):
torch.cuda.current_blas_handle()
one = torch.empty((2, 2), device=device)
torch.addmm(one[0], one, one)
malloc_graph.resume()
GRAPH_CAPTURE_STREAMS[device] = capture_stream
signature = None
@@ -82,6 +140,10 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
module._v_block_faulted = True
graph_hit = comfy_aimdo.model_vbar.vbar_signature_compare(signature, graph["signature"])
if malloc_graph is not None and malloc_scope is not None:
if malloc_graph.iterate(malloc_scope if module is not None and not graph_hit else None):
_malloc_graph_break()
consumed = queue.pop(0)
if consumed is not None:
offload_stream, prefetch_state = consumed
@@ -131,12 +193,21 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
module._v_block_faulted = True
if signature is not None:
_drop_graph(module)
malloc_graph.pause()
graph = torch.cuda.CUDAGraph()
if generator is not None:
graph.register_generator_state(generator)
malloc_graph.resume()
# Capture-time VBAR eviction is safe after prior work completes.
comfy.model_management.synchronize()
capture_stream.wait_stream(comfy.model_management.current_stream(device))
with torch.cuda.graph(graph, stream=capture_stream, capture_error_mode="thread_local"):
core()
malloc_graph.pause(sync=True)
with malloc_graph.use_stream(capture_stream):
with torch.cuda.graph(graph, stream=capture_stream, capture_error_mode="thread_local"):
malloc_graph.resume()
core()
malloc_graph.pause()
malloc_graph.resume(sync=True)
comfy.model_management.current_stream(device).wait_stream(capture_stream)
graph.replay()
module._comfy_graph = {"graph": graph, "signature": signature}
@@ -146,7 +217,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
core()
else:
capture_stream.wait_stream(comfy.model_management.current_stream(device))
with torch.cuda.stream(capture_stream):
with torch.cuda.stream(capture_stream), malloc_graph.use_stream(capture_stream):
core()
comfy.model_management.current_stream(device).wait_stream(capture_stream)
GRAPH_WARMED_MODULES.add(module)
comfy/text_encoders/gemma4.py
+29 -46
View File
@@ -228,9 +228,8 @@ class Gemma4Attention(nn.Module):
if fixed_cache is not None:
if seq_length == 1 and fixed_cache.index > 0:
# CUDA-graphable decode: write at the device-side ring/linear position
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)
fixed_cache.key.index_copy_(2, fixed_cache.position, xk)
fixed_cache.value.index_copy_(2, fixed_cache.position, xv)
output = self._decode_attention(xq, fixed_cache, attention_mask)
return self.o_proj(output), fixed_cache, None
@@ -516,61 +515,36 @@ 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 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
decode = fixed_kv and seq_len == 1
# Compiled decode needs fixed-capacity attention for a stable allocation trace;
# CUDA graph capture has additional prefetch and device requirements.
compiled_decode = decode and past_len > 0 and mask is None and self.graph_dynamic_vbar_blocks
if compiled_decode:
x = x.clone()
enable_graph = (compiled_decode
and prefetch_queue is not None
and hasattr(self.layers[0], "_v_block")
and not comfy.model_management.args.disable_cuda_graphs
and comfy.model_management.is_device_cuda(x.device))
decode_bias = None
decode_masks = None
if fixed_kv:
if decode:
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:
if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
if enable_graph:
# static buffers + per-capacity attention biases: layer graphs replay against
# stable storage, refreshed eagerly each step
if compiled_decode:
capacities = tuple(sorted({kv.key.shape[2] for kv in past_key_values if isinstance(kv, FixedKV)}))
state_key = (x.shape, x.dtype, x.device, tuple(t.shape for t in freqs_cis), capacities,
None if per_layer_inputs is None else per_layer_inputs.shape)
state = getattr(self, "_comfy_cross_step_state", None)
if state is None or state["key"] != state_key:
state = {"key": state_key,
"x": torch.empty_like(x),
"freqs_cis": [torch.empty_like(t) for t in freqs_cis],
"bias": {c: torch.empty((1, 1, 1, c), dtype=x.dtype, device=x.device) for c in capacities},
"per_layer": None if per_layer_inputs is None else torch.empty_like(per_layer_inputs),
"bias_valid": -1}
self._comfy_cross_step_state = state
comfy.model_management._register_cross_step(self)
state["x"].copy_(x)
for source, target in zip(freqs_cis, state["freqs_cis"]):
target.copy_(source)
x = state["x"]
freqs_cis = state["freqs_cis"]
if per_layer_inputs is not None:
state["per_layer"].copy_(per_layer_inputs)
per_layer_inputs = state["per_layer"]
valid = past_len + 1
for capacity, bias in state["bias"].items():
if state["bias_valid"] != past_len:
bias.fill_(min_val)
bias[..., :min(valid, capacity)] = 0
elif past_len < capacity:
bias[..., past_len:valid] = 0
state["bias_valid"] = valid
decode_bias = state["bias"]
decode_bias = {capacity: torch.full((1, 1, 1, capacity), min_val, dtype=x.dtype, device=x.device) for capacity in capacities}
for capacity, bias in decode_bias.items():
bias[..., :min(valid, capacity)] = 0
intermediate = None
all_intermediate = None
@@ -606,7 +580,7 @@ class Gemma4Transformer(nn.Module):
if shared is not None:
layer_kwargs['shared_kv'] = shared
if enable_graph:
if compiled_decode:
bias_cache = layer_kwargs.get('shared_kv', past_kv)
layer_mask = decode_bias[bias_cache.key.shape[2]]
elif decode:
@@ -619,10 +593,17 @@ class Gemma4Transformer(nn.Module):
def core():
nonlocal x
x, current_kv, shareable_kv = layer(x=x, attention_mask=layer_mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
output, current_kv, shareable_kv = layer(x=x, attention_mask=layer_mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
if compiled_decode:
x.copy_(output)
else:
x = output
result.append((current_kv, shareable_kv))
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
comfy.model_prefetch.prefetch_queue_pop(
prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph,
malloc_scope="block"
)
if result:
current_kv, shareable_kv = result[0]
@@ -641,8 +622,10 @@ class Gemma4Transformer(nn.Module):
if i == intermediate_output:
intermediate = x.clone()
if prefetch_queue is not None:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
comfy.model_prefetch.prefetch_queue_pop(
prefetch_queue, x.device, None,
malloc_scope="block"
)
if fixed_kv:
for kv in past_key_values:
@@ -707,7 +690,7 @@ 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((batch,), device=device, dtype=torch.int64),
tracker = (torch.empty((1,), 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
comfy/text_encoders/llama.py
+32 -32
View File
@@ -858,29 +858,7 @@ class Llama2_(nn.Module):
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)]
for group in freqs_cis_groups:
for tensor in group:
cross_step_state_key.append((tensor.shape, tensor.stride(), tensor.dtype, tensor.device))
cross_step_state_key = tuple(cross_step_state_key)
cross_step_state = getattr(self, "_comfy_cross_step_state", None)
if cross_step_state is None or cross_step_state["key"] != cross_step_state_key:
static_freqs_cis = []
for group in freqs_cis_groups:
static_freqs_cis.append(tuple(torch.empty_like(tensor) for tensor in group))
if not isinstance(freqs_cis, list):
static_freqs_cis = static_freqs_cis[0]
cross_step_state = {"key": cross_step_state_key, "x": torch.empty_like(x), "freqs_cis": static_freqs_cis}
self._comfy_cross_step_state = cross_step_state
comfy.model_management._register_cross_step(self)
cross_step_state["x"].copy_(x)
static_freqs_cis_groups = cross_step_state["freqs_cis"] if isinstance(freqs_cis, list) else [cross_step_state["freqs_cis"]]
for source_group, target_group in zip(freqs_cis_groups, static_freqs_cis_groups):
for source, target in zip(source_group, target_group):
target.copy_(source)
x = cross_step_state["x"]
freqs_cis = cross_step_state["freqs_cis"]
x = x.clone()
intermediate = None
all_intermediate = None
@@ -911,17 +889,24 @@ class Llama2_(nn.Module):
def core():
nonlocal x
x, current_kv = layer(
output, current_kv = layer(
x=x,
attention_mask=mask,
freqs_cis=freqs_cis,
optimized_attention=optimized_attention,
past_key_value=past_kv,
)
if enable_graph:
x.copy_(output)
else:
x = output
if next_key_values:
next_key_values[i] = current_kv
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
comfy.model_prefetch.prefetch_queue_pop(
prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph,
malloc_scope="block"
)
if fixed_kv:
next_key_values[i].advance(seq_len)
@@ -932,8 +917,10 @@ class Llama2_(nn.Module):
if i == intermediate_output:
intermediate = x.clone()
if prefetch_queue is not None:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
comfy.model_prefetch.prefetch_queue_pop(
prefetch_queue, x.device, None,
malloc_scope="block"
)
if self.norm is not None:
x = self.norm(x)
@@ -1041,9 +1028,19 @@ class BaseGenerate:
# MRoPE: prefill uses explicit 3D position_ids, decode continues from the last position
next_pos = int(position_ids[:, -1].max()) + 1 if position_ids is not None else None
compile_allocations = self.model.graph_dynamic_vbar_blocks and comfy.model_prefetch.malloc_graph_enabled(device)
decode_tokens = torch.empty((embeds.shape[0], 1), dtype=torch.long, device=device)
# Generation loop
current_input_ids = initial_input_ids
for step in tqdm(range(max_length), desc="Generating tokens"):
if step > 0:
if compile_allocations:
comfy.model_prefetch.malloc_graph_begin(self, device)
embeds = self.model.embed_tokens(decode_tokens).to(execution_dtype)
current_input_ids = decode_tokens if initial_input_ids is not None else None
position_ids = torch.tensor([[next_pos]], device=device) if next_pos is not None else None
# DeepStack visual features are injected on the prefill only; gemma4's forward lacks these kwargs.
extra = {}
if step == 0 and deepstack_embeds is not None:
@@ -1052,13 +1049,16 @@ class BaseGenerate:
x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids, position_ids=position_ids, **extra, embeds_info=(embeds_info if step == 0 else None))
logits = self.logits(x)[:, -1]
next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, initial_tokens + generated_token_ids, generator, do_sample=do_sample, presence_penalty=presence_penalty)
token_id = next_token[0].item()
decode_tokens.copy_(next_token)
del next_token, logits, x, embeds, position_ids
if step > 0 and compile_allocations:
comfy.model_prefetch.malloc_graph_end()
token_id = decode_tokens[0].item()
generated_token_ids.append(token_id)
embeds = self.model.embed_tokens(next_token).to(execution_dtype)
current_input_ids = next_token if initial_input_ids is not None else None
if next_pos is not None: # advance MRoPE position for the next (decode) step
position_ids = torch.tensor([[next_pos]], device=device)
if step > 0 and next_pos is not None:
next_pos += 1
pbar.update(1)
execution.py
+1 -1
View File
@@ -547,8 +547,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
if comfy.memory_management.aimdo_enabled:
if get_console_log_level(args.verbose) == "DEBUG":
comfy_aimdo.control.analyze()
comfy.model_management.reset_cast_buffers()
comfy.model_prefetch.cleanup_prefetch_queues()
comfy.model_management.reset_cast_buffers()
comfy_aimdo.model_vbar.vbars_reset_watermark_limits()
if has_pending_tasks:
requirements.txt
+1 -1
View File
@@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0
filelock
av>=17.0.0
comfy-kitchen==0.2.31
comfy-aimdo==0.4.15
comfy-aimdo==0.5.1
requests
simpleeval>=1.0.0
blake3