Kolors text encoder carries an unreachable torch<2 attention path built on the MPS-unsafe baddbmm(empty, beta=0) idiom

[RudraMantri123]

Describe the bug

CoreAttention.forward in pipelines/kolors/text_encoder.py computes raw attention scores with:

matmul_input_buffer = torch.empty(
    output_size[0] * output_size[1], output_size[2], output_size[3],
    dtype=query_layer.dtype, device=query_layer.device,
)
matmul_result = torch.baddbmm(matmul_input_buffer, ..., beta=0.0, alpha=(1.0 / self.norm_factor))

This relies on the documented baddbmm contract that beta=0 causes input to be ignored, so NaN/Inf in the uninitialised buffer must not propagate. The MPS backend violates that contract: pytorch/pytorch#187521 (labeled module: correctness (silent), fixed in pytorch main but not in any released torch, including 2.13.0 — the fix missed the release branch). #14438 was this exact mechanism biting SDXL through Attention.get_attention_scores, fixed by #14459. After #14459, this Kolors call site is the last remaining instance of the idiom in the repository.

Reproduction

The primitive fails at exactly the shapes this code requests. On Apple Silicon (torch 2.13.0, M-series):

import torch

for b_np, sq in [(16, 1024), (16, 2048)]:  # b*np, seq — Kolors text-encoder score shapes
    shape = (b_np, sq, sq)
    junk = torch.full(shape, float("nan"), device="mps", dtype=torch.float16)
    del junk  # freed NaN pages go back to the allocator pool
    buf = torch.empty(shape, device="mps", dtype=torch.float16)  # recycles them
    q = torch.randn(b_np, sq, 64, device="mps", dtype=torch.float16)
    k = torch.randn(b_np, sq, 64, device="mps", dtype=torch.float16)
    out = torch.baddbmm(buf, q, k.transpose(1, 2), beta=0.0, alpha=0.125)
    print(shape, "buffer had NaN:", True, "-> output has NaN:", bool(out.isnan().any()))

prints output has NaN: True for both shapes on my machine (M5 Pro).

Scope caveat, stated honestly: I have not reproduced NaN in CoreAttention's output end-to-end through the Kolors pipeline — in my attempts the allocator handed the buffer clean pages even after aggressive dirtying. In the #14438 case it is component offloading that reliably leaves large freed dirty regions behind, and the Kolors text encoder has no equivalent. So this is a latent correctness hazard (unsafe idiom + shapes confirmed vulnerable at the primitive level), not a bug I can show corrupting Kolors outputs today.

Fix in #14620: route the MPS case through a buffer-free scaled bmm, identical to the approach taken in #14459; all other devices keep the current baddbmm path unchanged.

System Info

• diffusers @ main
• torch 2.13.0 (any released torch is affected; fixed only in pytorch main)
• Apple Silicon (M5 Pro), macOS / MPS backend

Who can help?

@yiyixuxu @asomoza

---

Correction (same day): on re-verification I found the vulnerable call is unreachable on any supported torch. CoreAttention.forward branches on int(torch.__version__.split(".")[0]) >= 2 and uses scaled_dot_product_attention for all of torch 2.x; the baddbmm code above only executes on torch 1.x, and diffusers requires torch ≥ 2.6. So this is not a live bug — it is a latent hazard inside dead code (which is also why end-to-end corruption was never reproducible). #14620 has been reworked accordingly: it now deletes the entire torch < 2 branch, verified bit-identical to main across devices, dtypes, and mask branches.