Should attention_sage auto-cast fp32 inputs to bf16 to recover sage perf when third-party nodes leak fp32?

[Booyaka101]

Hi! Filing this as a discussion rather than a cold PR because the change touches a perf-vs-correctness tradeoff in core and I'd like to gauge whether you'd accept this approach before doing the PR work.

The observation

When running with --use-sage-attention, several third-party Flux 2 custom nodes (comfyui-flux2fun-controlnet, RES4LYF, others) silently fall back to pytorch attention because their attention paths leak float32 tensors into attention_sage.

This shows up as:

Error running sage attention: Input tensors must be in dtype of torch.float16 or torch.bfloat16, using pytorch attention instead.

Repeated per attention call per sampling step. On a Flux 2 Dev keyframe render with Fun-Controlnet (4 control layers × 8 sampling steps × multiple attention calls per layer), this fires ~32+ times per render and costs an estimated ~10-15% wall-time per render.

Why fp32 leaks in

The model itself is bf16 throughout. The fp32 promotion is incidental:

• nn.LayerNorm / nn.RMSNorm upcast to fp32 internally for variance computation and return fp32 unless explicitly cast back.
• Rotary embedding application that uses fp32 freqs_cis returns fp32 q/k.
• Third-party authors don't always cast back to the model's native dtype before downstream attention.

Comfy core's stock Flux paths handle this carefully (e.g. apply_rope1 ends with .type_as(x) to cast back). Third-party nodes don't always.

The proposed change

Cast fp32 q/k/v to bf16 before calling sageattn(...):

if q.dtype == torch.float32:
    q = q.to(torch.bfloat16)
if k.dtype == torch.float32:
    k = k.to(torch.bfloat16)
if v.dtype == torch.float32:
    v = v.to(torch.bfloat16)
out = sageattn(q, k, v, ...)

Quality impact: zero. The model's native dtype is bf16; we're matching that.

Perf impact: 10-15% per render whenever an attention path leaks fp32 (eliminates the pytorch fallback).

Concerns I'd raise myself

1. Masking third-party bugs. Could be argued this enables sloppy custom_nodes by making them "work" even though they're leaking dtypes. Counter: custom_node authors aren't going to fix it; users get the perf hit either way.

2. Future fp32-attention models. If a model genuinely needs fp32 attention precision, this would silently downcast. Counter: no current diffusion model in the wild does this; could be gated behind --strict-attn-dtype.

3. Where the cast should live. Could go in attention_sage (proposed), or in a wrapper, or as a registered attention function with an opt-in flag.

Question

Would you accept a PR adding this auto-cast, and if so:

• (A) as default behavior in attention_sage?
• (B) gated behind a new flag like --sage-auto-cast?
• (C) prefer to file fix-request issues against each leaking third-party node instead?

Happy to do whichever direction you'd prefer.

Reproduction

• ComfyUI 0.19.3, py3.12.8, torch 2.9.1+cu128, sageattention 2.0.1, RTX 4090 Win11
• Workflow: Flux 2 Dev FP8 + PuLID-Flux2 + comfyui-flux2fun-controlnet
• Without cast: ~32 sage fallback log lines per render
• With cast: zero fallbacks, sage runs on every attention call

Source: came up while building https://github.com//autoseries (local AI video pipeline). Happy to share full benchmarks if useful.

Thanks for ComfyUI!

[Booyaka101]

Following up on my own discussion - I attempted the proposed patch locally and it produced broken renders, so I am withdrawing the suggestion as written. Sharing the failure mode here in case anyone else considers a similar approach.

What broke

I cast q, k, v independently to torch.bfloat16 before sageattn(...):

if q.dtype == torch.float32: q = q.to(torch.bfloat16)
if k.dtype == torch.float32: k = k.to(torch.bfloat16)
if v.dtype == torch.float32: v = v.to(torch.bfloat16)
out = sageattn(q, k, v, ...)

In a Flux 2 Dev workflow with Fun-Controlnet active, this produced horizontal-banding noise renders that decode to garbage. Reverted and the renders are clean again with the original pytorch-attention fallback path.

Likely root cause

In a mixed-precision pipeline some attention paths feed v as fp16 while q/k come through fp32 (RMSNorm output). My patch produced q=bf16, k=bf16, v=fp16 mixed-dtype input. Sage doesn't appear to validate dtype consistency cleanly — instead of erroring (which would trigger the existing pytorch-attention fallback), it processes the mixed-dtype tensors and emits garbage.

A correct version would homogenize to a target dtype

If anyone wants to revisit this, the safe approach is something like:

target_dtype = v.dtype if v.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
if q.dtype != target_dtype: q = q.to(target_dtype)
if k.dtype != target_dtype: k = k.to(target_dtype)
if v.dtype != target_dtype: v = v.to(target_dtype)

I have not validated this version - flagging it as a starting point only. The cleanest fix is probably for sage itself to validate dtype consistency, or for the third-party nodes that leak fp32 (Fun-Controlnet etc) to cast back to model dtype before attention - filed those separately at bryanmcguire/comfyui-flux2fun-controlnet#12.

Apologies for the noise. Treating this as the rule confirmed: never propose a Comfy core change without end-to-end validation across multiple sampling steps. Lesson absorbed.