diff --git a/comfy/cli_args.py b/comfy/cli_args.py index d02b92d0ae9..407e138386e 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -147,6 +147,8 @@ def from_string(cls, value: str): attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.") attn_group.add_argument("--use-quad-cross-attention", action="store_true", help="Use the sub-quadratic cross attention optimization . Ignored when xformers is used.") attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.") +attn_group.add_argument("--use-kitchen-bf16-attention", action="store_true", help="Use Comfy Kitchen FP16/BF16 attention on supported AMD ROCm/HIP GPUs.") +attn_group.add_argument("--use-kitchen-int8-attention", action="store_true", help="Use Comfy Kitchen INT8 attention on supported AMD ROCm/HIP GPUs.") attn_group.add_argument("--use-sage-attention", action="store_true", help="Use sage attention.") attn_group.add_argument("--use-flash-attention", action="store_true", help="Use FlashAttention.") attn_group.add_argument("--use-ck-attention", action="store_true", help="Use Comfy Kitchen attention.") diff --git a/comfy/ldm/lumina/model.py b/comfy/ldm/lumina/model.py index cdf03b2b5e3..43fa29e972c 100644 --- a/comfy/ldm/lumina/model.py +++ b/comfy/ldm/lumina/model.py @@ -7,6 +7,7 @@ import torch.nn.functional as F import comfy.ldm.common_dit import comfy.model_management +import comfy.model_prefetch import comfy.ops import comfy.quant_ops @@ -17,6 +18,120 @@ import comfy.patcher_extension import comfy.utils from comfy.ldm.chroma_radiance.layers import NerfEmbedder +from comfy.quant_ops import QuantizedTensor, TensorWiseINT8Layout + +_FUSED_RMS_MODULATED = getattr(TensorWiseINT8Layout, "fused_rms_modulated", None) +_FUSED_SWIGLU_FFN = getattr(TensorWiseINT8Layout, "fused_swiglu_ffn", None) + + +def _int8_convrot_weight(weight) -> bool: + """Return whether a weight uses the supported INT8 ConvRot layout.""" + return ( + isinstance(weight, QuantizedTensor) + and weight._layout_cls == "TensorWiseINT8Layout" + and getattr(weight._params, "convrot", False) + ) + + +def _fused_rms_modulated_linear(x, linear, norm, modulation_scale): + """Run fused RMS modulation and INT8 projection when supported.""" + weight = getattr(linear, "weight", None) + if ( + comfy.model_management.in_training + or not callable(_FUSED_RMS_MODULATED) + or not _int8_convrot_weight(weight) + ): + return None + with ( + comfy.ops.CastBiasWeightContext(norm, x, offloadable=True) as (norm_weight, _), + comfy.ops.CastBiasWeightContext(linear, x, offloadable=True) as (weight, bias), + ): + scale = modulation_scale.to(device=x.device, dtype=x.dtype) + fused = _FUSED_RMS_MODULATED( + x, weight, bias, norm_weight, norm.eps, scale, + ) + return None if fused is NotImplemented else fused + + +def _fused_swiglu_ffn_postnorm(x, feed_forward, norm): + """Run the post-normalized fused INT8 SwiGLU FFN when supported.""" + if ( + comfy.model_management.in_training + or not callable(_FUSED_SWIGLU_FFN) + ): + return None + w1 = getattr(feed_forward.w1, "weight", None) + w2 = getattr(feed_forward.w2, "weight", None) + w3 = getattr(feed_forward.w3, "weight", None) + if ( + feed_forward.w1.bias is not None + or feed_forward.w2.bias is not None + or feed_forward.w3.bias is not None + or not all(_int8_convrot_weight(w) for w in (w1, w2, w3)) + ): + return None + normed = norm(x) + with ( + comfy.ops.CastBiasWeightContext(feed_forward.w1, normed, offloadable=True) as (w1, b1), + comfy.ops.CastBiasWeightContext(feed_forward.w3, normed, offloadable=True) as (w3, b3), + comfy.ops.CastBiasWeightContext(feed_forward.w2, normed, offloadable=True) as (w2, b2), + ): + fused = _FUSED_SWIGLU_FFN( + normed, w1, w3, w2, b1, b3, b2, + ) + return None if fused is NotImplemented else fused + + +def _fused_swiglu_ffn(x, feed_forward, norm, modulation_scale): + """Run fused RMS modulation and an INT8 SwiGLU FFN when supported.""" + w1 = getattr(feed_forward.w1, "weight", None) + w2 = getattr(feed_forward.w2, "weight", None) + w3 = getattr(feed_forward.w3, "weight", None) + if ( + comfy.model_management.in_training + or not callable(_FUSED_SWIGLU_FFN) + or feed_forward.w1.bias is not None + or feed_forward.w2.bias is not None + or feed_forward.w3.bias is not None + or not all(_int8_convrot_weight(w) for w in (w1, w2, w3)) + ): + return None + with ( + comfy.ops.CastBiasWeightContext(feed_forward.w1, x, offloadable=True) as (w1, b1), + comfy.ops.CastBiasWeightContext(feed_forward.w3, x, offloadable=True) as (w3, b3), + comfy.ops.CastBiasWeightContext(feed_forward.w2, x, offloadable=True) as (w2, b2), + comfy.ops.CastBiasWeightContext(norm, x, offloadable=True) as (norm_weight, _), + ): + scale = modulation_scale.to(device=x.device, dtype=x.dtype) + fused = _FUSED_SWIGLU_FFN( + x, w1, w3, w2, b1, b3, b2, + norm_weight=norm_weight, norm_eps=norm.eps, modulation_scale=scale, + ) + return None if fused is NotImplemented else fused + + +def _fused_rms_gated_residual(activation, norm, residual, gate): + """Run fused RMS normalization, gating, and residual addition.""" + ck = getattr(comfy.quant_ops, "ck", None) + if ck is None or not callable(getattr(ck, "rms_gated_residual", None)): + return None + if ( + comfy.model_management.in_training + or activation.dtype != torch.bfloat16 + or residual.shape != activation.shape + or gate.ndim != 2 + or gate.shape[0] != 1 + or gate.shape[1] != activation.shape[-1] + ): + return None + norm_weight, _, offload_stream = comfy.ops.cast_bias_weight(norm, activation, offloadable=True) + try: + gate_vec = gate[0].to(device=activation.device, dtype=activation.dtype) + return ck.rms_gated_residual( + activation, norm_weight, residual, gate_vec, norm.eps, + ) + finally: + comfy.ops.uncast_bias_weight(norm, norm_weight, None, offload_stream) def invert_slices(slices, length): @@ -129,6 +244,7 @@ def forward( x_mask: torch.Tensor, freqs_cis: torch.Tensor, transformer_options={}, + qkv: torch.Tensor | None = None, ) -> torch.Tensor: """ @@ -142,8 +258,11 @@ def forward( """ bsz, seqlen, _ = x.shape + if qkv is None: + qkv = self.qkv(x) + xq, xk, xv = torch.split( - self.qkv(x), + qkv, [ self.n_local_heads * self.head_dim, self.n_local_kv_heads * self.head_dim, @@ -334,20 +453,66 @@ def forward( if self.modulation: assert adaln_input is not None scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).chunk(4, dim=1) - - x = x + apply_gate(gate_msa.unsqueeze(1).tanh(), self.attention_norm2( - clamp_fp16(self.attention( - modulate(self.attention_norm1(x), scale_msa, timestep_zero_index=timestep_zero_index), - x_mask, - freqs_cis, - transformer_options=transformer_options, - ))), timestep_zero_index=timestep_zero_index - ) - x = x + apply_gate(gate_mlp.unsqueeze(1).tanh(), self.ffn_norm2( - clamp_fp16(self.feed_forward( - modulate(self.ffn_norm1(x), scale_mlp, timestep_zero_index=timestep_zero_index), - ))), timestep_zero_index=timestep_zero_index - ) + gate_msa = gate_msa.tanh() + gate_mlp = gate_mlp.tanh() + gate_msa_t = gate_msa.unsqueeze(1) + gate_mlp_t = gate_mlp.unsqueeze(1) + + if timestep_zero_index is None and not comfy.model_management.in_training: + qkv = _fused_rms_modulated_linear( + x, self.attention.qkv, self.attention_norm1, scale_msa, + ) + if qkv is not None: + attn_out = clamp_fp16(self.attention( + x, x_mask, freqs_cis, transformer_options=transformer_options, qkv=qkv, + )) + fused_x = _fused_rms_gated_residual( + attn_out, self.attention_norm2, x, gate_msa, + ) + if fused_x is not None: + x = fused_x + else: + x = x + apply_gate(gate_msa_t, self.attention_norm2(attn_out)) + else: + x = x + apply_gate(gate_msa_t, self.attention_norm2( + clamp_fp16(self.attention( + modulate(self.attention_norm1(x), scale_msa, timestep_zero_index=timestep_zero_index), + x_mask, + freqs_cis, + transformer_options=transformer_options, + )) + )) + + ffn_out = _fused_swiglu_ffn(x, self.feed_forward, self.ffn_norm1, scale_mlp) + if ffn_out is not None: + ffn_out = clamp_fp16(ffn_out) + fused_x = _fused_rms_gated_residual( + ffn_out, self.ffn_norm2, x, gate_mlp, + ) + if fused_x is not None: + x = fused_x + else: + x = x + apply_gate(gate_mlp_t, self.ffn_norm2(ffn_out)) + else: + x = x + apply_gate(gate_mlp_t, self.ffn_norm2( + clamp_fp16(self.feed_forward( + modulate(self.ffn_norm1(x), scale_mlp, timestep_zero_index=timestep_zero_index), + )) + )) + else: + x = x + apply_gate(gate_msa_t, self.attention_norm2( + clamp_fp16(self.attention( + modulate(self.attention_norm1(x), scale_msa, timestep_zero_index=timestep_zero_index), + x_mask, + freqs_cis, + transformer_options=transformer_options, + )) + ), timestep_zero_index=timestep_zero_index) + x = x + apply_gate(gate_mlp_t, self.ffn_norm2( + clamp_fp16(self.feed_forward( + modulate(self.ffn_norm1(x), scale_mlp, timestep_zero_index=timestep_zero_index), + )) + ), timestep_zero_index=timestep_zero_index) else: assert adaln_input is None x = x + self.attention_norm2( @@ -358,11 +523,15 @@ def forward( transformer_options=transformer_options, )) ) - x = x + self.ffn_norm2( - self.feed_forward( - self.ffn_norm1(x), + ffn_out = _fused_swiglu_ffn_postnorm(x, self.feed_forward, self.ffn_norm1) + if ffn_out is not None: + x = x + self.ffn_norm2(clamp_fp16(ffn_out)) + else: + x = x + self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x), + ) ) - ) return x @@ -626,6 +795,18 @@ def __init__( self.dim = dim self.n_heads = n_heads + def get_dynamic_vram__units(self): + """Return model units in Dynamic VRAM execution order.""" + units = list(self.context_refiner) + if self.siglip_refiner is not None: + units.extend(self.siglip_refiner) + units.extend(self.noise_refiner) + units.extend(self.layers) + final_unit = getattr(self, "final_layer", None) + if final_unit is None: + final_unit = getattr(self, "dec_net", None) + return units, [] if final_unit is None else [final_unit] + def unpatchify( self, x: torch.Tensor, img_size: List[Tuple[int, int]], cap_size: List[int], return_tensor=False ) -> List[torch.Tensor]: @@ -717,6 +898,7 @@ def embed_all(self, x, cap_feats=None, siglip_feats=None, offset=0, omni=False, def patchify_and_embed( self, x: torch.Tensor, cap_feats: torch.Tensor, cap_mask: torch.Tensor, t: torch.Tensor, num_tokens, ref_latents=[], ref_contexts=[], siglip_feats=[], transformer_options={} ) -> Tuple[torch.Tensor, torch.Tensor, List[Tuple[int, int]], List[int], torch.Tensor]: + """Embed and refine caption, reference, and image tokens.""" bsz = x.shape[0] cap_mask = None # TODO? main_siglip = None @@ -771,8 +953,11 @@ def patchify_and_embed( # refine context cap_feats = torch.cat(embeds[0], dim=1) cap_freqs_cis = torch.cat(freqs_cis[0], dim=1) + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.context_refiner), cap_feats.device, transformer_options) for layer in self.context_refiner: + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, cap_feats.device, layer) cap_feats = layer(cap_feats, cap_mask, cap_freqs_cis, transformer_options=transformer_options) + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, cap_feats.device, None) feats = (cap_feats,) fc = (cap_freqs_cis,) @@ -782,8 +967,11 @@ def patchify_and_embed( siglip_feats_combined = torch.cat(embeds[1], dim=1) siglip_feats_freqs_cis = torch.cat(freqs_cis[1], dim=1) if self.siglip_refiner is not None: + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.siglip_refiner), siglip_feats_combined.device, transformer_options) for layer in self.siglip_refiner: + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, siglip_feats_combined.device, layer) siglip_feats_combined = layer(siglip_feats_combined, siglip_mask, siglip_feats_freqs_cis, transformer_options=transformer_options) + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, siglip_feats_combined.device, None) feats += (siglip_feats_combined,) fc += (siglip_feats_freqs_cis,) @@ -796,13 +984,16 @@ def patchify_and_embed( timestep_zero_index = None x_input = x + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.noise_refiner), x.device, transformer_options) for i, layer in enumerate(self.noise_refiner): + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer) x = layer(x, padded_img_mask, fc_x, t, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options) if "noise_refiner" in patches: for p in patches["noise_refiner"]: out = p({"img": x, "img_input": x_input, "txt": cap_feats, "pe": fc_x, "vec": t, "x": orig_x, "block_index": i, "transformer_options": transformer_options, "block_type": "noise_refiner"}) if "img" in out: x = out["img"] + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None) padded_full_embed = torch.cat(feats + (x,), dim=1) if timestep_zero_index is not None: @@ -815,6 +1006,7 @@ def patchify_and_embed( return padded_full_embed, mask, img_sizes, l_effective_cap_len, torch.cat(fc + (fc_x,), dim=1), timestep_zero_index def forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs): + """Execute the denoiser through registered model wrappers.""" return comfy.patcher_extension.WrapperExecutor.new_class_executor( self._forward, self, @@ -823,6 +1015,7 @@ def forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwar # def forward(self, x, t, cap_feats, cap_mask): def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_latents=[], ref_contexts=[], siglip_feats=[], transformer_options={}, **kwargs): + """Run the NextDiT denoising forward pass.""" omni = len(ref_latents) > 0 if omni: timesteps = torch.cat([timesteps * 0, timesteps], dim=0) @@ -832,11 +1025,6 @@ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_l cap_mask = attention_mask bs, c, h, w = x.shape x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size)) - """ - Forward pass of NextDiT. - t: (N,) tensor of diffusion timesteps - y: (N,) tensor of text tokens/features - """ t = self.t_embedder(t * self.time_scale, dtype=x.dtype) # (N, D) adaln_input = t @@ -858,7 +1046,9 @@ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_l transformer_options["total_blocks"] = len(self.layers) transformer_options["block_type"] = "double" img_input = img + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), img.device, transformer_options) for i, layer in enumerate(self.layers): + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, img.device, layer) transformer_options["block_index"] = i img = layer(img, mask, freqs_cis, adaln_input, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options) if "double_block" in patches: @@ -868,6 +1058,7 @@ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_l img[:, cap_size[0]:] = out["img"] if "txt" in out: img[:, :cap_size[0]] = out["txt"] + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, img.device, None) img = self.final_layer(img, adaln_input, timestep_zero_index=timestep_zero_index) img = self.unpatchify(img, img_size, cap_size, return_tensor=x_is_tensor)[:, :, :h, :w] @@ -1043,6 +1234,7 @@ def __init__( # with the pixel-space dec_net decoder. # ------------------------------------------------------------------ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_latents=[], ref_contexts=[], siglip_feats=[], transformer_options={}, **kwargs): + """Run the pixel-space NextDiT denoising forward pass.""" omni = len(ref_latents) > 0 if omni: timesteps = torch.cat([timesteps * 0, timesteps], dim=0) @@ -1089,7 +1281,9 @@ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_l transformer_options["total_blocks"] = len(self.layers) transformer_options["block_type"] = "double" img_input = img + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), img.device, transformer_options) for i, layer in enumerate(self.layers): + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, img.device, layer) transformer_options["block_index"] = i img = layer(img, mask, freqs_cis, adaln_input, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options) if "double_block" in patches: @@ -1099,6 +1293,7 @@ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_l img[:, cap_size[0]:] = out["img"] if "txt" in out: img[:, :cap_size[0]] = out["txt"] + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, img.device, None) # ---- pixel-space decoder (replaces final_layer + unpatchify) ---- # img may have padding tokens beyond N; only the first N are real image patches @@ -1120,6 +1315,7 @@ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, ref_l return -img_out def forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs): + """Execute the pixel-space denoiser through model wrappers.""" # _forward returns neg_x0 = -x0 (negated decoder output). # # Reference inference (working_inference_reference.py): diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index b22d03d7755..c549b1ed90a 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -73,6 +73,7 @@ def get_attention_function(name: str, default: Any=...) -> Union[Callable, None] from comfy.cli_args import args import comfy.ops +import comfy.quant_ops ops = comfy.ops.disable_weight_init FORCE_UPCAST_ATTENTION_DTYPE = model_management.force_upcast_attention_dtype() @@ -585,6 +586,76 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha ).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head) return out +_KITCHEN_ATTENTION_ROUTE_LOGGED = set() + +def _attention_kitchen(kernel_name, support_name, q, k, v, heads, mask, + attn_precision, skip_reshape, skip_output_reshape, + **kwargs): + """Dispatch supported calls to Kitchen HIP or PyTorch SDPA.""" + kitchen = getattr(comfy.quant_ops, "ck", None) + support = None if kitchen is None else getattr(kitchen, support_name, None) + supported = ( + model_management.is_amd() + and callable(support) + and mask is None and skip_reshape + and not kwargs.get("enable_gqa", False) + and q.ndim == 4 and q.shape[1] == heads + and support(q, k, v) + ) + route_key = (kernel_name, "hip" if supported else "sdpa", None if q is None else tuple(q.shape), None if q is None else str(q.dtype)) + if route_key not in _KITCHEN_ATTENTION_ROUTE_LOGGED: + _KITCHEN_ATTENTION_ROUTE_LOGGED.add(route_key) + logging.info( + "Kitchen %s route=%s shape=%s dtype=%s mask=%s skip_reshape=%s", + kernel_name, + "hip" if supported else "pytorch-sdpa", + None if q is None else tuple(q.shape), + None if q is None else q.dtype, + mask is not None, + skip_reshape, + ) + if not supported: + return attention_pytorch( + q, k, v, heads, mask=mask, attn_precision=attn_precision, + skip_reshape=skip_reshape, + skip_output_reshape=skip_output_reshape, **kwargs, + ) + + out = getattr(kitchen, kernel_name)(q, k, v, kwargs.get("scale")) + if not skip_output_reshape: + out = out.transpose(1, 2).reshape( + q.shape[0], q.shape[2], heads * q.shape[3] + ) + return out + +@wrap_attn +def attention_kitchen_bf16(q, k, v, heads, mask=None, attn_precision=None, + skip_reshape=False, skip_output_reshape=False, + **kwargs): + """Run Kitchen BF16 attention with PyTorch SDPA fallback.""" + return _attention_kitchen( + "hip_attention", "hip_attention_is_supported", + q, k, v, heads, mask, attn_precision, + skip_reshape, skip_output_reshape, **kwargs, + ) + +@wrap_attn +def attention_kitchen_int8(q, k, v, heads, mask=None, attn_precision=None, + skip_reshape=False, skip_output_reshape=False, + **kwargs): + """Run Kitchen INT8 attention when low-precision attention is allowed.""" + if kwargs.get("low_precision_attention", True) is False: + return attention_kitchen_bf16( + q, k, v, heads, mask=mask, attn_precision=attn_precision, + skip_reshape=skip_reshape, + skip_output_reshape=skip_output_reshape, **kwargs, + ) + return _attention_kitchen( + "hip_int8_attention", "hip_int8_attention_is_supported", + q, k, v, heads, mask, attn_precision, + skip_reshape, skip_output_reshape, **kwargs, + ) + def _comfy_kitchen_int8_inputs(q, k, v, heads, mask, skip_reshape, enable_gqa): dim_head = q.shape[-1] if skip_reshape else q.shape[-1] // heads b = q.shape[0] @@ -852,7 +923,16 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape optimized_attention = attention_basic -if model_management.sage_attention_enabled(): +if (args.use_kitchen_bf16_attention or args.use_kitchen_int8_attention) and not model_management.is_amd(): + logging.warning("Comfy Kitchen HIP attention is AMD ROCm-only; ignoring the selected Kitchen attention flag") + +if args.use_kitchen_bf16_attention and model_management.is_amd(): + logging.info("Using Comfy Kitchen FP16/BF16 attention") + optimized_attention = attention_kitchen_bf16 +elif args.use_kitchen_int8_attention and model_management.is_amd(): + logging.info("Using Comfy Kitchen INT8 attention") + optimized_attention = attention_kitchen_int8 +elif model_management.sage_attention_enabled(): logging.info("Using sage attention") optimized_attention = attention_sage elif model_management.flash_attention_enabled(): @@ -884,6 +964,8 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape # register core-supported attention functions +register_attention_function("kitchen_bf16", attention_kitchen_bf16) +register_attention_function("kitchen_int8", attention_kitchen_int8) if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE: register_attention_function("comfy_kitchen_int8", attention_comfy_kitchen_int8) if SAGE_ATTENTION_IS_AVAILABLE: diff --git a/comfy/model_base.py b/comfy/model_base.py index ba59d270511..27b1ce8013b 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -1505,6 +1505,10 @@ def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.lumina.model.NextDiT) self.memory_usage_factor_conds = ("ref_latents",) + def get_dynamic_vram__units(self): + """Return the underlying Lumina model's Dynamic VRAM units.""" + return self.diffusion_model.get_dynamic_vram__units() + def extra_conds(self, **kwargs): out = super().extra_conds(**kwargs) attention_mask = kwargs.get("attention_mask", None) diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index f0fdf1aa5bd..0fa9f8772cc 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -89,6 +89,7 @@ def __init__(self, device="cpu", max_length=77, freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=comfy.clip_model.CLIPTextModel, special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, zero_out_masked=False, return_projected_pooled=True, return_attention_masks=False, model_options={}): # clip-vit-base-patch32 + """Initialize a text encoder with device-aware quantized operations.""" super().__init__() if textmodel_json_config is None: @@ -111,7 +112,18 @@ def __init__(self, device="cpu", max_length=77, if operations is None: if quant_config is not None: - operations = comfy.ops.mixed_precision_ops(quant_config, dtype, full_precision_mm=True) + load_device = model_options.get("load_device", model_management.text_encoder_device()) + disabled = {"float8_e4m3fn", "float8_e5m2", "nvfp4", "mxfp8"} + if model_management.is_device_cuda(load_device): + if model_management.supports_fp8_compute(load_device): + disabled.difference_update(("float8_e4m3fn", "float8_e5m2")) + if model_management.supports_nvfp4_compute(load_device): + disabled.remove("nvfp4") + if model_management.supports_mxfp8_compute(load_device): + disabled.remove("mxfp8") + operations = comfy.ops.mixed_precision_ops( + quant_config, dtype, disabled=disabled, + ) logging.info("Using MixedPrecisionOps for text encoder") else: operations = comfy.ops.manual_cast diff --git a/tests-unit/comfy_test/lumina_kitchen_test.py b/tests-unit/comfy_test/lumina_kitchen_test.py new file mode 100644 index 00000000000..3cd2b7bd7e3 --- /dev/null +++ b/tests-unit/comfy_test/lumina_kitchen_test.py @@ -0,0 +1,114 @@ +"""Lumina Comfy Kitchen integration regression tests.""" + +from types import SimpleNamespace + +import pytest +import torch + +from comfy.cli_args import args + +original_cpu = args.cpu +if not torch.cuda.is_available(): + args.cpu = True + +import comfy.ldm.lumina.model as lumina_model # noqa: E402 +import comfy.ldm.modules.attention as attention # noqa: E402 +import comfy.model_base as model_base # noqa: E402 + +args.cpu = original_cpu + + +@pytest.mark.parametrize("length", [160, 4096, 4256]) +def test_kitchen_bf16_routes_every_supported_length(monkeypatch, length): + """Supported BF16 calls of every model length must use Kitchen HIP.""" + calls = [] + kitchen = SimpleNamespace( + hip_attention_is_supported=lambda q, k, v: True, + hip_attention=lambda q, k, v, scale: calls.append((q.shape, scale)) or q, + ) + monkeypatch.setattr(attention.comfy.quant_ops, "ck", kitchen) + monkeypatch.setattr(attention.model_management, "is_amd", lambda: True) + + q = torch.empty((1, 1, length, 128), dtype=torch.bfloat16) + out = attention.attention_kitchen_bf16( + q, q, q, 1, skip_reshape=True, skip_output_reshape=True, + ) + + assert out is q + assert calls == [(q.shape, None)] + + +def test_kitchen_int8_respects_explicit_low_precision_opt_out(monkeypatch): + """Explicit low-precision opt-out must select Kitchen BF16, not INT8.""" + calls = [] + kitchen = SimpleNamespace( + hip_attention_is_supported=lambda q, k, v: True, + hip_int8_attention_is_supported=lambda q, k, v: True, + hip_attention=lambda q, k, v, scale: calls.append("bf16") or q, + hip_int8_attention=lambda q, k, v, scale: calls.append("int8") or q, + ) + monkeypatch.setattr(attention.comfy.quant_ops, "ck", kitchen) + monkeypatch.setattr(attention.model_management, "is_amd", lambda: True) + + q = torch.empty((1, 1, 1024, 128), dtype=torch.bfloat16) + assert attention.attention_kitchen_int8( + q, q, q, 1, skip_reshape=True, skip_output_reshape=True, + ) is q + assert attention.attention_kitchen_int8( + q, q, q, 1, skip_reshape=True, skip_output_reshape=True, + low_precision_attention=False, + ) is q + + assert calls == ["int8", "bf16"] + + +def test_unsupported_kitchen_attention_uses_pytorch(monkeypatch): + """Unsupported Kitchen calls must retain the PyTorch fallback.""" + fallback = object() + kitchen = SimpleNamespace( + hip_attention_is_supported=lambda q, k, v: True, + hip_attention=lambda q, k, v, scale: pytest.fail("HIP kernel called"), + ) + monkeypatch.setattr(attention.comfy.quant_ops, "ck", kitchen) + monkeypatch.setattr(attention.model_management, "is_amd", lambda: True) + monkeypatch.setattr(attention, "attention_pytorch", lambda *args, **kwargs: fallback) + + q = torch.empty((1, 1, 160, 128), dtype=torch.bfloat16) + mask = torch.ones((1, 1, 160, 160), dtype=torch.bool) + + assert attention.attention_kitchen_bf16( + q, q, q, 1, mask=mask, skip_reshape=True, skip_output_reshape=True, + ) is fallback + + +def test_missing_layout_fusions_fall_back(monkeypatch): + """Missing Kitchen layout methods must retain the unfused paths.""" + monkeypatch.setattr(lumina_model, "_FUSED_RMS_MODULATED", None) + monkeypatch.setattr(lumina_model, "_FUSED_SWIGLU_FFN", None) + + linear = SimpleNamespace(weight=None) + ffn_layer = SimpleNamespace(weight=None, bias=None) + feed_forward = SimpleNamespace(w1=ffn_layer, w2=ffn_layer, w3=ffn_layer) + + assert lumina_model._fused_rms_modulated_linear(None, linear, None, None) is None + assert lumina_model._fused_swiglu_ffn_postnorm(None, feed_forward, None) is None + assert lumina_model._fused_swiglu_ffn(None, feed_forward, None, None) is None + + +@pytest.mark.parametrize("wrapper_type", [model_base.Lumina2, model_base.ZImagePixelSpace]) +def test_lumina_wrapper_delegates_dynamic_vram_units(wrapper_type): + """Lumina wrappers must expose their diffusion model's unit ordering.""" + expected = ([object(), object()], [object()]) + + class DynamicUnitsModel(torch.nn.Module): + """Minimal diffusion model exposing Dynamic VRAM units.""" + + def get_dynamic_vram__units(self): + """Return the sentinel execution order.""" + return expected + + wrapper = wrapper_type.__new__(wrapper_type) + torch.nn.Module.__init__(wrapper) + wrapper.diffusion_model = DynamicUnitsModel() + + assert wrapper.get_dynamic_vram__units() is expected