Skip to content

Remove the unreachable torch<2 attention path in the Kolors text encoder #14620

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
RudraMantri123 wants to merge 2 commits into
base: main
Choose a base branch
from
+90 −93
Open

Remove the unreachable torch<2 attention path in the Kolors text encoder #14620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6a503ba
Avoid the unsafe baddbmm beta=0 idiom on MPS in the Kolors text encoder
RudraMantri123 Aug 27, 2026
fe10d47
Remove the unreachable torch<2 attention path instead of patching it
RudraMantri123 Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 16 additions & 93 deletions src/diffusers/pipelines/kolors/text_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,101 +130,24 @@ def __init__(self, config: ChatGLMConfig, layer_number):
self.attention_dropout = torch.nn.Dropout(config.attention_dropout)

def forward(self, query_layer, key_layer, value_layer, attention_mask):
pytorch_major_version = int(torch.__version__.split(".")[0])
if pytorch_major_version >= 2:
query_layer, key_layer, value_layer = [
k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]
]
if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
context_layer = torch.nn.functional.scaled_dot_product_attention(
query_layer, key_layer, value_layer, is_causal=True
)
else:
if attention_mask is not None:
attention_mask = ~attention_mask
context_layer = torch.nn.functional.scaled_dot_product_attention(
query_layer, key_layer, value_layer, attention_mask
)
context_layer = context_layer.permute(2, 0, 1, 3)
new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
context_layer = context_layer.reshape(*new_context_layer_shape)
else:
# Raw attention scores

# [b, np, sq, sk]
output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))

# [sq, b, np, hn] -> [sq, b * np, hn]
query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
# [sk, b, np, hn] -> [sk, b * np, hn]
key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)

# preallocting input tensor: [b * np, sq, sk]
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,
)

# Raw attention scores. [b * np, sq, sk]
matmul_result = torch.baddbmm(
matmul_input_buffer,
query_layer.transpose(0, 1), # [b * np, sq, hn]
key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
beta=0.0,
alpha=(1.0 / self.norm_factor),
# diffusers requires torch >= 2.6, so scaled_dot_product_attention is always
# available; the pre-2.0 manual attention path this module originally carried
# was unreachable (and used an MPS-unsafe baddbmm idiom, see
# https://github.com/huggingface/diffusers/issues/14624).
query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
context_layer = torch.nn.functional.scaled_dot_product_attention(
query_layer, key_layer, value_layer, is_causal=True
)

# change view to [b, np, sq, sk]
attention_scores = matmul_result.view(*output_size)

# ===========================
# Attention probs and dropout
# ===========================

# attention scores and attention mask [b, np, sq, sk]
if self.attention_softmax_in_fp32:
attention_scores = attention_scores.float()
if self.coeff is not None:
attention_scores = attention_scores * self.coeff
if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
attention_mask = torch.ones(
output_size[0], 1, output_size[2], output_size[3], device=attention_scores.device, dtype=torch.bool
)
attention_mask.tril_()
attention_mask = ~attention_mask
else:
if attention_mask is not None:
attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
attention_probs = F.softmax(attention_scores, dim=-1)
attention_probs = attention_probs.type_as(value_layer)

# This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper.
attention_probs = self.attention_dropout(attention_probs)
# =========================
# Context layer. [sq, b, hp]
# =========================

# value_layer -> context layer.
# [sk, b, np, hn] --> [b, np, sq, hn]

# context layer shape: [b, np, sq, hn]
output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
# change view [sk, b * np, hn]
value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
# change view [b * np, sq, sk]
attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
# matmul: [b * np, sq, hn]
context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
# change view [b, np, sq, hn]
context_layer = context_layer.view(*output_size)
# [b, np, sq, hn] --> [sq, b, np, hn]
context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
# [sq, b, np, hn] --> [sq, b, hp]
new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
context_layer = context_layer.view(*new_context_layer_shape)
attention_mask = ~attention_mask
context_layer = torch.nn.functional.scaled_dot_product_attention(
query_layer, key_layer, value_layer, attention_mask
)
context_layer = context_layer.permute(2, 0, 1, 3)
new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
context_layer = context_layer.reshape(*new_context_layer_shape)

return context_layer

Expand Down
74 changes: 74 additions & 0 deletions tests/pipelines/kolors/test_kolors_text_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest
import torch

from diffusers.pipelines.kolors.text_encoder import ChatGLMConfig, CoreAttention

from ...testing_utils import torch_device


class TestKolorsCoreAttention:
"""
`CoreAttention` computes attention through `scaled_dot_product_attention`; the module
used to carry a second, manual attention path gated on torch < 2 that was unreachable
on any supported torch (https://github.com/huggingface/diffusers/issues/14624). These
tests pin the behaviour of the remaining path on every backend so its removal, and any
future rework, stay observable.
"""

def get_attention(self):
config = ChatGLMConfig(
hidden_size=256,
num_attention_heads=8,
kv_channels=32,
multi_query_attention=False,
attention_softmax_in_fp32=True,
)
return CoreAttention(config, layer_number=1)

def get_inputs(self, device, dtype=torch.float32, seq_len=64, batch=2, heads=8, head_dim=32):
torch.manual_seed(0)
shape = (seq_len, batch, heads, head_dim)
return tuple(torch.randn(shape, dtype=dtype).to(device) for _ in range(3))

@pytest.mark.parametrize("masked", [False, True])
def test_output_matches_cpu_reference(self, masked):
# The device path must stay numerically equivalent to the CPU path, for both the
# causal (mask=None) branch and the explicit-mask branch of forward.
attention = self.get_attention()
query, key, value = self.get_inputs("cpu")
seq_len, batch = query.shape[0], query.shape[1]

if masked:
torch.manual_seed(1)
# ChatGLM convention: True marks positions that must NOT be attended to.
attention_mask = torch.rand(batch, 1, seq_len, seq_len) < 0.25
attention_mask[..., 0] = False # keep at least one visible key per query row
else:
attention_mask = None

with torch.no_grad():
expected = attention(query, key, value, attention_mask)
actual = attention(
query.to(torch_device),
key.to(torch_device),
value.to(torch_device),
attention_mask.to(torch_device) if attention_mask is not None else None,
)

assert torch.isfinite(expected).all()
torch.testing.assert_close(actual.cpu(), expected, atol=1e-4, rtol=1e-4)
Loading