From 6a503ba846a0f83b7b3ac07f32916aa82a9f8c69 Mon Sep 17 00:00:00 2001 From: Rudra Mantri <189433012+RudraMantri123@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:47:32 +0530 Subject: [PATCH 1/2] Avoid the unsafe baddbmm beta=0 idiom on MPS in the Kolors text encoder CoreAttention computes raw attention scores with baddbmm(input=torch.empty(...), beta=0), relying on beta=0 causing input to be ignored. MPS does not honour that contract (pytorch#187521): NaN in the uninitialised buffer propagates into the output. Verified directly at the shapes this code requests -- a freed NaN block of (b*np, sq, sk) is handed back by torch.empty and survives baddbmm. This is the last remaining instance of the idiom in the repository after #14459 fixed both copies of Attention.get_attention_scores. Use the same buffer-free scaled bmm on MPS, which relies on no contract and skips the scores-sized allocation; all other devices keep the existing baddbmm path unchanged. Adds a CPU-equivalence test that runs on every backend and an MPS-gated test that dirties the allocator before computing scores. --- .../pipelines/kolors/text_encoder.py | 44 +++++++---- .../kolors/test_kolors_text_encoder.py | 74 +++++++++++++++++++ 2 files changed, 102 insertions(+), 16 deletions(-) create mode 100644 tests/pipelines/kolors/test_kolors_text_encoder.py diff --git a/src/diffusers/pipelines/kolors/text_encoder.py b/src/diffusers/pipelines/kolors/text_encoder.py index 434f4fed6fbb..b7a08b3d37b6 100644 --- a/src/diffusers/pipelines/kolors/text_encoder.py +++ b/src/diffusers/pipelines/kolors/text_encoder.py @@ -159,23 +159,35 @@ def forward(self, query_layer, key_layer, value_layer, attention_mask): # [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), - ) + if query_layer.device.type == "mps": + # `baddbmm(input=torch.empty(...), beta=0)` relies on beta=0 causing `input` to be + # ignored. MPS does not honour that contract — NaN/Inf in the uninitialised buffer + # propagate into the scores. Use a buffer-free scaled bmm instead, which also skips + # the scores-sized allocation. See + # https://github.com/huggingface/diffusers/pull/14459 and + # https://github.com/pytorch/pytorch/issues/187521. + matmul_result = torch.bmm( + query_layer.transpose(0, 1) * (1.0 / self.norm_factor), # [b * np, sq, hn] + key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] + ) + else: + # 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, + ) + + 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), + ) # change view to [b, np, sq, sk] attention_scores = matmul_result.view(*output_size) diff --git a/tests/pipelines/kolors/test_kolors_text_encoder.py b/tests/pipelines/kolors/test_kolors_text_encoder.py new file mode 100644 index 000000000000..e4afd1465cf3 --- /dev/null +++ b/tests/pipelines/kolors/test_kolors_text_encoder.py @@ -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 raw attention scores. On MPS it must not route them through + `baddbmm(input=torch.empty(...), beta=0)`, because MPS does not honour the documented + "input is ignored when beta=0" contract and NaN from the uninitialised buffer can reach + the scores. See https://github.com/huggingface/diffusers/pull/14459. + """ + + 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=256, 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)) + + def test_scores_match_cpu_reference(self): + # The device-specific score path must stay numerically equivalent to the CPU path. + attention = self.get_attention() + query, key, value = self.get_inputs("cpu") + + with torch.no_grad(): + expected = attention(query, key, value, None) + actual = attention(query.to(torch_device), key.to(torch_device), value.to(torch_device), None) + + torch.testing.assert_close(actual.cpu(), expected, atol=1e-4, rtol=1e-4) + + @pytest.mark.skipif(torch_device != "mps", reason="guards an MPS-specific baddbmm contract violation") + def test_scores_finite_with_dirty_allocator(self): + # Free NaN-filled blocks of exactly the scores shape so the allocator can hand those + # pages to the score computation, then assert the output stays finite. + attention = self.get_attention() + seq_len, batch, heads = 256, 2, 8 + query, key, value = self.get_inputs(torch_device, dtype=torch.float16, seq_len=seq_len) + + scores_shape = (batch * heads, seq_len, seq_len) + for _ in range(4): + junk = torch.full(scores_shape, float("nan"), device=torch_device, dtype=torch.float16) + del junk + + with torch.no_grad(): + output = attention(query, key, value, None) + + assert torch.isfinite(output).all(), "NaN reached the Kolors attention output on MPS" From fe10d4747c23b98962ade9547b8258f11609559f Mon Sep 17 00:00:00 2001 From: Rudra Mantri <189433012+RudraMantri123@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:47:14 +0530 Subject: [PATCH 2/2] Remove the unreachable torch<2 attention path instead of patching it CoreAttention.forward gates on int(torch.__version__.split(".")[0]) >= 2 and diffusers requires torch >= 2.6, so the manual attention branch -- the one carrying the MPS-unsafe baddbmm(torch.empty(...), beta=0) idiom -- cannot execute on any supported install. Every forward goes through scaled_dot_product_attention. That makes the previous commit's approach wrong twice over: the bug it guarded against is unreachable, and the bmm branch it added was itself dead code that nothing could execute. Delete the whole torch<2 path instead, which removes the unsafe idiom truthfully. Verified the deletion changes nothing observable: outputs are bit-identical to main across cpu/mps, fp32/fp16, and both mask branches. The test file now pins CPU-equivalence of the remaining SDPA path (masked and causal) on every backend; the dirty-allocator test is dropped because it exercised SDPA all along and could never fail. --- .../pipelines/kolors/text_encoder.py | 121 +++--------------- .../kolors/test_kolors_text_encoder.py | 54 ++++---- 2 files changed, 43 insertions(+), 132 deletions(-) diff --git a/src/diffusers/pipelines/kolors/text_encoder.py b/src/diffusers/pipelines/kolors/text_encoder.py index b7a08b3d37b6..9eeeda20b5f6 100644 --- a/src/diffusers/pipelines/kolors/text_encoder.py +++ b/src/diffusers/pipelines/kolors/text_encoder.py @@ -130,113 +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) + # 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 + ) 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) - - # Raw attention scores. [b * np, sq, sk] - if query_layer.device.type == "mps": - # `baddbmm(input=torch.empty(...), beta=0)` relies on beta=0 causing `input` to be - # ignored. MPS does not honour that contract — NaN/Inf in the uninitialised buffer - # propagate into the scores. Use a buffer-free scaled bmm instead, which also skips - # the scores-sized allocation. See - # https://github.com/huggingface/diffusers/pull/14459 and - # https://github.com/pytorch/pytorch/issues/187521. - matmul_result = torch.bmm( - query_layer.transpose(0, 1) * (1.0 / self.norm_factor), # [b * np, sq, hn] - key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] - ) - else: - # 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, - ) - - 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), - ) - - # 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 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 diff --git a/tests/pipelines/kolors/test_kolors_text_encoder.py b/tests/pipelines/kolors/test_kolors_text_encoder.py index e4afd1465cf3..f08909590900 100644 --- a/tests/pipelines/kolors/test_kolors_text_encoder.py +++ b/tests/pipelines/kolors/test_kolors_text_encoder.py @@ -23,10 +23,11 @@ class TestKolorsCoreAttention: """ - `CoreAttention` computes raw attention scores. On MPS it must not route them through - `baddbmm(input=torch.empty(...), beta=0)`, because MPS does not honour the documented - "input is ignored when beta=0" contract and NaN from the uninitialised buffer can reach - the scores. See https://github.com/huggingface/diffusers/pull/14459. + `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): @@ -39,36 +40,35 @@ def get_attention(self): ) return CoreAttention(config, layer_number=1) - def get_inputs(self, device, dtype=torch.float32, seq_len=256, batch=2, heads=8, head_dim=32): + 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)) - def test_scores_match_cpu_reference(self): - # The device-specific score path must stay numerically equivalent to the CPU path. + @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] - with torch.no_grad(): - expected = attention(query, key, value, None) - actual = attention(query.to(torch_device), key.to(torch_device), value.to(torch_device), None) - - torch.testing.assert_close(actual.cpu(), expected, atol=1e-4, rtol=1e-4) - - @pytest.mark.skipif(torch_device != "mps", reason="guards an MPS-specific baddbmm contract violation") - def test_scores_finite_with_dirty_allocator(self): - # Free NaN-filled blocks of exactly the scores shape so the allocator can hand those - # pages to the score computation, then assert the output stays finite. - attention = self.get_attention() - seq_len, batch, heads = 256, 2, 8 - query, key, value = self.get_inputs(torch_device, dtype=torch.float16, seq_len=seq_len) - - scores_shape = (batch * heads, seq_len, seq_len) - for _ in range(4): - junk = torch.full(scores_shape, float("nan"), device=torch_device, dtype=torch.float16) - del junk + 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(): - output = attention(query, key, value, None) + 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(output).all(), "NaN reached the Kolors attention output on MPS" + assert torch.isfinite(expected).all() + torch.testing.assert_close(actual.cpu(), expected, atol=1e-4, rtol=1e-4)