Allow model files to contain which attention should be used for each block. (#16419)

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 `<module path>.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.

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))
 
 

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:

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
 
 

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: