From 804eb5513a9dec3c0e624044ba72a2c026d92491 Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:18:39 +1000 Subject: [PATCH] Introduce Comfy Compiler (CORE-389) (#15861) --- comfy/cli_args.py | 5 ++ comfy/ldm/lightricks/av_model.py | 10 +++- comfy/ldm/minimax/model.py | 18 +++++-- comfy/ldm/minimax_music/ar.py | 38 +++++++++----- comfy/model_management.py | 4 ++ comfy/model_prefetch.py | 87 +++++++++++++++++++++++++++++--- comfy/text_encoders/gemma4.py | 75 +++++++++++---------------- comfy/text_encoders/llama.py | 64 +++++++++++------------ execution.py | 2 +- requirements.txt | 2 +- 10 files changed, 198 insertions(+), 107 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index d02b92d0a..fdb6e329d 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -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 diff --git a/comfy/ldm/lightricks/av_model.py b/comfy/ldm/lightricks/av_model.py index c60148e2a..d253b0144 100644 --- a/comfy/ldm/lightricks/av_model.py +++ b/comfy/ldm/lightricks/av_model.py @@ -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] diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py index b83863c66..10170ce38 100644 --- a/comfy/ldm/minimax/model.py +++ b/comfy/ldm/minimax/model.py @@ -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") diff --git a/comfy/ldm/minimax_music/ar.py b/comfy/ldm/minimax_music/ar.py index 2a8935318..78a4c7c86 100644 --- a/comfy/ldm/minimax_music/ar.py +++ b/comfy/ldm/minimax_music/ar.py @@ -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") diff --git a/comfy/model_management.py b/comfy/model_management.py index d56e69e48..def0af364 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -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): diff --git a/comfy/model_prefetch.py b/comfy/model_prefetch.py index bdde5137a..22a91cd6a 100644 --- a/comfy/model_prefetch.py +++ b/comfy/model_prefetch.py @@ -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) diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index 0d8f0fcc7..96a547b60 100644 --- a/comfy/text_encoders/gemma4.py +++ b/comfy/text_encoders/gemma4.py @@ -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 diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index 49ba5dfa9..a61c5adc7 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -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) diff --git a/execution.py b/execution.py index 62f3d7c52..48c6f641b 100644 --- a/execution.py +++ b/execution.py @@ -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: diff --git a/requirements.txt b/requirements.txt index 1599c011a..7f08247dc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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