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

Implement Minimax Music 3 + Core Support for Cuda Graphs (#15570)

Browse Source
This commit is contained in:
rattus
2026-08-14 02:10:08 +10:00
committed by GitHub
parent 12666983cb
commit efd4e951a0
19 changed files with 1333 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
comfy/ops.py
View File

@@ -123,10 +123,12 @@ def materialize_meta_param(s, param_keys):
# FIXME: add n=1 cache hit fast path
def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking):
def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking, return_faulted=False):
offload_stream = None
cast_buffer = None
cast_buffer_offset = 0
if return_faulted:
fully_faulted = all(not getattr(s, param_key + "_function", []) for s in comfy_modules for param_key in ("weight", "bias"))
def ensure_offload_stream(module, required_size, check_largest):
nonlocal offload_stream
@@ -163,6 +165,8 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
for s in comfy_modules:
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
if return_faulted and (signature is None or not resident):
fully_faulted = False
prefetch = {
"signature": signature,
"resident": resident,
@@ -255,10 +259,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
prefetch["needs_cast"] = needs_cast
s._prefetch = prefetch
if return_faulted:
return offload_stream, fully_faulted
return offload_stream
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant):
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant, return_weights=True):
prefetch = getattr(s, "_prefetch", None)
@@ -298,7 +304,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
tensor = tensor.dequantize()
return tensor
if orig.dtype != dtype or len(fns) > 0:
if (return_weights and orig.dtype != dtype) or len(fns) > 0:
x = to_dequant(x, dtype)
if not resident and lowvram_fn is not None:
x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
@@ -325,7 +331,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
if prefetch["signature"] is not None:
prefetch["resident"] = True
return weight, bias
return (weight, bias) if return_weights else None
def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):