From 055f2fff5daf18e4176b5d92a06d889fb925e31f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 8 Oct 2026 18:45:06 +0000 Subject: [PATCH 1/3] Use 32-row flex attention blocks for head_dim > 256 The default flex block sizes fault at head_dim 512 on Blackwell (misaligned address), which breaks attn_implementation="flex_attention" for Gemma 4, and are 1.7x slower than 32-row blocks on Hopper. --- .../integrations/flex_attention.py | 4 ++ tests/integrations/test_flex_attention.py | 55 +++++++++++++++++++ 2 files changed, 59 insertions(+) create mode 100644 tests/integrations/test_flex_attention.py diff --git a/src/transformers/integrations/flex_attention.py b/src/transformers/integrations/flex_attention.py index bf2aada97eb3..524f4314766d 100644 --- a/src/transformers/integrations/flex_attention.py +++ b/src/transformers/integrations/flex_attention.py @@ -309,6 +309,10 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): enable_gqa = False kernel_options = kwargs.get("kernel_options") + # The default block sizes fault at head_dim 512 on Blackwell ("misaligned address") and are 1.7x slower than + # 32-row blocks on Hopper (e.g. Gemma 4 global layers) + if kernel_options is None and query.shape[-1] > 256 and query.device.type == "cuda": + kernel_options = {"BLOCK_M": 32, "BLOCK_N": 32, "num_stages": 2, "num_warps": 4} # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it return_lse = query.device.type != "cpu" diff --git a/tests/integrations/test_flex_attention.py b/tests/integrations/test_flex_attention.py new file mode 100644 index 000000000000..f1abd3271de5 --- /dev/null +++ b/tests/integrations/test_flex_attention.py @@ -0,0 +1,55 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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 unittest + +from transformers.testing_utils import ( + is_torch_available, + require_flex_attention, + require_torch_accelerator, + torch_device, +) + + +if is_torch_available(): + import torch + from torch.nn.attention.flex_attention import create_block_mask + + from transformers.integrations.flex_attention import flex_attention_forward + + +@require_flex_attention +@require_torch_accelerator +class FlexAttentionForwardTest(unittest.TestCase): + def test_head_dim_512_forward_and_backward(self): + # Gemma 4 global layers use head_dim 512, where the default flex block sizes fault on Blackwell + torch.manual_seed(0) + batch_size, num_heads, num_kv_heads, seq_len, head_dim = 1, 8, 2, 1024, 512 + query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=torch_device, dtype=torch.bfloat16) + key = torch.randn(batch_size, num_kv_heads, seq_len, head_dim, device=torch_device, dtype=torch.bfloat16) + value = torch.randn(batch_size, num_kv_heads, seq_len, head_dim, device=torch_device, dtype=torch.bfloat16) + query.requires_grad_() + block_mask = create_block_mask(lambda b, h, q, kv: q >= kv, None, None, seq_len, seq_len, device=torch_device) + + output, _ = flex_attention_forward(torch.nn.Module(), query, key, value, block_mask) + output.float().sum().backward() + + reference_query = query.detach().float().requires_grad_() + reference = torch.nn.functional.scaled_dot_product_attention( + reference_query, key.float(), value.float(), is_causal=True, enable_gqa=True + ) + reference.sum().backward() + + torch.testing.assert_close(output.float(), reference.transpose(1, 2), rtol=2e-2, atol=2e-2) + torch.testing.assert_close(query.grad.float(), reference_query.grad, rtol=2e-2, atol=2e-2) From e0011f7ba4efac997bad1fe040d0332d9b289a1d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 9 Oct 2026 01:39:21 +0000 Subject: [PATCH 2/3] Point the flex block-size comment at pytorch/pytorch#200348 --- src/transformers/integrations/flex_attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/integrations/flex_attention.py b/src/transformers/integrations/flex_attention.py index 524f4314766d..2cac8f87777f 100644 --- a/src/transformers/integrations/flex_attention.py +++ b/src/transformers/integrations/flex_attention.py @@ -309,8 +309,8 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): enable_gqa = False kernel_options = kwargs.get("kernel_options") - # The default block sizes fault at head_dim 512 on Blackwell ("misaligned address") and are 1.7x slower than - # 32-row blocks on Hopper (e.g. Gemma 4 global layers) + # The default block sizes fault at head_dim 512 on B300 (pytorch/pytorch#200348) and are 1.7x slower than 32-row + # blocks on H100 (e.g. Gemma 4 global layers) if kernel_options is None and query.shape[-1] > 256 and query.device.type == "cuda": kernel_options = {"BLOCK_M": 32, "BLOCK_N": 32, "num_stages": 2, "num_warps": 4} # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it From 2aca36584c231c15e307b2079f56439a5fa6353b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Mon, 9 Nov 2026 21:42:54 +0000 Subject: [PATCH 3/3] Only use the 32-row flex blocks on Hopper and datacenter Blackwell, which have the shared memory for them --- src/transformers/integrations/flex_attention.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/transformers/integrations/flex_attention.py b/src/transformers/integrations/flex_attention.py index 16cd1ff4b5f6..a2532c891f3a 100644 --- a/src/transformers/integrations/flex_attention.py +++ b/src/transformers/integrations/flex_attention.py @@ -306,8 +306,15 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): kernel_options = kwargs.get("kernel_options") # The default block sizes fault at head_dim 512 on B300 (pytorch/pytorch#200348) and are 1.7x slower than 32-row - # blocks on H100 (e.g. Gemma 4 global layers) - if kernel_options is None and query.shape[-1] > 256 and query.device.type == "cuda": + # blocks on H100 (e.g. Gemma 4 global layers). In fp32, 32-row blocks need ~200 KB of shared memory, which only + # Hopper and datacenter Blackwell have, so other GPUs keep the default blocks + if ( + kernel_options is None + and query.shape[-1] > 256 + and query.device.type == "cuda" + and torch.version.hip is None + and torch.cuda.get_device_capability(query.device)[0] in (9, 10) + ): kernel_options = {"BLOCK_M": 32, "BLOCK_N": 32, "num_stages": 2, "num_warps": 4} # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it return_lse = query.device.type != "cpu"