diff --git a/QUANTIZATION.md b/QUANTIZATION.md index 300822029..a3fd13d14 100644 --- a/QUANTIZATION.md +++ b/QUANTIZATION.md @@ -152,6 +152,27 @@ Example: To create compatible checkpoints, use any quantization tool provided the output follows the checkpoint format described above and uses a layout defined in `QUANT_ALGOS`. +### Diffusion attention preferences + +A diffusion attention module can have a `.comfy_attention.config` entry whose +uint8 tensor contains UTF-8 JSON: + +```json +{"attention": "comfy_kitchen_int8"} +``` + +Use the module that performs attention, such as `transformer_blocks.0.attn` for +Qwen Image 2.1 or `blocks.0.attn` for MiniMax H3. + +Only `comfy_kitchen_int8` is supported. Invalid targets and other method names +are ignored with a warning during loading, leaving normal attention selection. +Kitchen INT8 support is checked when each preference is loaded +for the primary device; unsupported devices keep normal attention selection. +Explicit attention overrides retain priority. +The `ComfyAttention` child module loads and saves its own metadata through normal +state-dict loading and saving. Preferences do not enable weight +quantization. Text encoder and VAE loaders do not apply these preferences. + ### Weight Quantization Weight quantization is straightforward - compute the scaling factor directly from the weight tensor using the absolute maximum method described earlier. Each layer's weights are quantized independently and stored with their corresponding `weight_scale` parameter. diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py index 5441f4fff..c4b7a68aa 100644 --- a/comfy/ldm/minimax/model.py +++ b/comfy/ldm/minimax/model.py @@ -25,7 +25,7 @@ import comfy.model_prefetch import comfy.ops import comfy.patcher_extension import comfy.quant_ops -from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention +from comfy.ldm.modules.attention import ComfyAttention, AttentionTensorContainer, optimized_attention FRAME_PER_TOKEN = (1, 4, 4, 4, 4) FRAME_RESCALE = 5.0 / 3.0 @@ -158,6 +158,7 @@ def rope_rotation_table(angles, dtype): class Attention(nn.Module): def __init__(self, hidden, heads, head_dim, eps, gate_compress=False, dtype=None, device=None, operations=None): super().__init__() + self.comfy_attention = ComfyAttention() self.heads = heads self.head_dim = head_dim inner = heads * head_dim @@ -196,7 +197,7 @@ class Attention(nn.Module): q = AttentionTensorContainer(q.transpose(0, 1).unsqueeze(0)) k = AttentionTensorContainer(k.transpose(0, 1).unsqueeze(0)) v = AttentionTensorContainer(v.transpose(0, 1).unsqueeze(0)) - out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options) + out = optimized_attention(q, k, v, self.heads, preferred_attention=self.comfy_attention, mask=None, skip_reshape=True, transformer_options=transformer_options) return self.out_proj(out.squeeze(0)) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 4d6cf8598..7e4f94351 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -1,6 +1,7 @@ import math import sys import inspect +import json import torch import torch.nn.functional as F @@ -71,6 +72,34 @@ def get_attention_function(name: str, default: Any=...) -> Union[Callable, None] return default return REGISTERED_ATTENTION_FUNCTIONS[name] + +class ComfyAttention(nn.Module): + def __init__(self): + super().__init__() + self.config = None + self.function = None + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + self.config = None + self.function = None + metadata = state_dict.pop(prefix + "config", None) + if metadata is not None: + config = json.loads(metadata.numpy().tobytes()) + method = config.get("attention") + if method == "comfy_kitchen_int8": + self.config = config + if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE and comfy_kitchen.int8_attention_is_available(model_management.get_torch_device()): + self.function = attention_comfy_kitchen_int8 + else: + logging.warning(f"Ignoring unsupported attention method {method!r} for {prefix.rstrip('.')}") + super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) + + def _save_to_state_dict(self, destination, prefix, keep_vars): + super()._save_to_state_dict(destination, prefix, keep_vars) + if self.config is not None: + destination[prefix + "config"] = torch.tensor(list(json.dumps(self.config).encode("utf-8")), dtype=torch.uint8) + + from comfy.cli_args import args import comfy.ops ops = comfy.ops.disable_weight_init @@ -171,6 +200,7 @@ class AttentionTensorContainer: def wrap_attn(func): @functools.wraps(func) def wrapper(*args, **kwargs): + preferred_attention = kwargs.pop("preferred_attention", None) containers = None if len(args) >= 3 and isinstance(args[0], AttentionTensorContainer): if not isinstance(args[1], AttentionTensorContainer) or not isinstance(args[2], AttentionTensorContainer): @@ -191,6 +221,10 @@ def wrap_attn(func): return optimized_attention_override.container_function(*args, **kwargs) args = tuple(container.take() for container in containers) + args[3:] return optimized_attention_override(func, *args, **kwargs) + if preferred_attention is not None: + attention = preferred_attention.function + if attention is not None: + return attention(*args, **kwargs) if containers is not None: if wrapper.container_function is not None: diff --git a/comfy/ldm/qwen_image21/model.py b/comfy/ldm/qwen_image21/model.py index 3df95bf00..b4f0755e7 100644 --- a/comfy/ldm/qwen_image21/model.py +++ b/comfy/ldm/qwen_image21/model.py @@ -12,7 +12,7 @@ import comfy.rmsnorm from comfy.ldm.flux.layers import EmbedND, timestep_embedding from comfy.ldm.flux.math import apply_rope1 from comfy.ldm.lightricks.model import TimestepEmbedding -from comfy.ldm.modules.attention import optimized_attention +from comfy.ldm.modules.attention import ComfyAttention, optimized_attention from comfy.ldm.wan.model_animate2 import PoseBranchCache @@ -69,6 +69,7 @@ class SwiGLUFeedForward(nn.Module): class Attention(nn.Module): def __init__(self, dim, heads, dim_head, eps=1e-6, dtype=None, device=None, operations=None): super().__init__() + self.comfy_attention = ComfyAttention() self.heads = heads inner_dim = heads * dim_head self.to_q = operations.Linear(dim, inner_dim, bias=False, dtype=dtype, device=device) @@ -103,7 +104,7 @@ class Attention(nn.Module): q, k = comfy.quant_ops.ck.rms_rope(q, k, pe, q_scale, k_scale, self.norm_q.eps) comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_stream) comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_stream) - return self.to_out[0](attn_fn(q, k, v, self.heads)) + return self.to_out[0](attn_fn(q, k, v, self.heads, preferred_attention=self.comfy_attention)) def _split_rows(p): @@ -164,11 +165,11 @@ class LastLayer(nn.Module): def block_causal_attention(segments, transformer_options={}, cache=None, block_index=0, prefix_len=0): # segments: (start, end, mask); text segments get a causal mask, image blocks attend to everything before their end - def attn(q, k, v, heads): + def attn(q, k, v, heads, preferred_attention=None): if cache is not None: # K and V stacked on dim 1 so batch stays first and quantized rows are per token and head cache.put(block_index, torch.stack([k[:, :prefix_len], v[:, :prefix_len]], dim=1)) - outs = [optimized_attention(q[:, start:end].flatten(2), k[:, :end].flatten(2), v[:, :end].flatten(2), heads, mask=mask, transformer_options=transformer_options) + outs = [optimized_attention(q[:, start:end].flatten(2), k[:, :end].flatten(2), v[:, :end].flatten(2), heads, mask=mask, transformer_options=transformer_options, preferred_attention=preferred_attention) for start, end, mask in segments] return torch.cat(outs, dim=1) if len(outs) > 1 else outs[0] return attn @@ -176,8 +177,8 @@ def block_causal_attention(segments, transformer_options={}, cache=None, block_i def prefix_cached_attention(prefix_k, prefix_v, transformer_options={}): # target-only queries: block-causal reduces to full attention over [cached prefix, target] - def attn(q, k, v, heads): - return optimized_attention(q.flatten(2), torch.cat([prefix_k, k], dim=1).flatten(2), torch.cat([prefix_v, v], dim=1).flatten(2), heads, transformer_options=transformer_options) + def attn(q, k, v, heads, preferred_attention=None): + return optimized_attention(q.flatten(2), torch.cat([prefix_k, k], dim=1).flatten(2), torch.cat([prefix_v, v], dim=1).flatten(2), heads, transformer_options=transformer_options, preferred_attention=preferred_attention) return attn diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index bbfc9959b..9d363be93 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -874,6 +874,8 @@ class ModelPatcher: bk = self.backup.get(k, None) hbk = self.hook_backup.get(k, None) weight, set_func, convert_func = get_key_weight(self.model, k) + if not isinstance(weight, torch.Tensor): + continue if bk is not None: weight = bk.weight if hbk is not None: