Flex Attention Silently Incorrect with bfloat16 and Masks

[zaptrem]

PyTorch flex_attention Numerical Instability with bfloat16 and Masks

🐛 Describe the bug

flex_attention produces significantly different results than F.scaled_dot_product_attention when using bfloat16 dtype with cross-attention masks with significant masked padding. The error grows to >3% of the output magnitude, which I believe is beyond acceptable numerical precision differences.

Key findings:

• Bug triggers: When attending to ≤58% of positions with bfloat16
• Bug disappears: When attending to ≥59% of positions
• Error magnitude: Up to 3.4% relative error (vs 0.3% expected tolerance)
• Does NOT occur with float32: Same code with float32 has <0.001% error
• Comparison baseline: Physical array slicing with F.scaled_dot_product_attention provides ground truth

This affects production models using bfloat16 for memory efficiency, particularly in autoregressive generation scenarios where cross-attention masks are commonly used to handle variable-length conditioning.

Minimal Reproduction

import torch
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
import torch.nn.functional as F

B, H, Q, D, KV = 1, 1, 1, 64, 1000

# Test showing error threshold: masking <=58% causes bug, >=59% is OK
print("PyTorch flex_attention bfloat16 bug demonstration:")
print("-" * 50)

for mask_amount in [580, 590]:
    mask = create_block_mask(
        lambda b, h, q, kv: kv < mask_amount,
        B, H, Q, KV, "cuda"
    )

    torch.manual_seed(0)
    q = torch.randn(B, H, Q, D, dtype=torch.bfloat16, device="cuda") * 5
    k = torch.randn(B, H, KV, D, dtype=torch.bfloat16, device="cuda") * 5
    v = torch.randn(B, H, KV, D, dtype=torch.bfloat16, device="cuda") * 5

    flex_out = flex_attention(q, k, v, block_mask=mask)
    ref_out = F.scaled_dot_product_attention(q, k[:, :, :mask_amount], v[:, :, :mask_amount])

    error = (flex_out - ref_out).abs().max() / ref_out.abs().max()
    status = "BUG!" if error > 0.01 else "OK"
    print(f"Mask {mask_amount}/{KV} ({mask_amount/KV:.0%}): Error {error:.1%} ({status})")

print("\nBug triggers when masking ≤58% of positions with bfloat16")

Output:

PyTorch flex_attention bfloat16 bug demonstration:
--------------------------------------------------
Mask 580/1000 (58%): Error 3.4% (BUG!)
Mask 590/1000 (59%): Error 0.3% (OK)

Bug triggers when masking ≤58% of positions with bfloat16

Expected behavior

flex_attention with a mask should produce nearly identical results to F.scaled_dot_product_attention with physical array slicing (within normal bfloat16 precision tolerance of <1%).

Actual behavior

Large numerical divergence (>3% error) when using bfloat16 with masks that exclude more than ~42% of positions.

Versions

Collecting environment information...
PyTorch version: 2.8.0+cu126
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.13 (main, Jun  4 2025, 08:57:29) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.5.0-1023-aws-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-80GB
Nvidia driver version: 535.183.01
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             24
On-line CPU(s) list:                0-23
Vendor ID:                          GenuineIntel
Model name:                         Intel Xeon Processor (SapphireRapids)
CPU family:                         6
Model:                              143
Thread(s) per core:                 1
Core(s) per socket:                 24
Socket(s):                          1
Stepping:                           4
BogoMIPS:                           4200.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd arat vnmi avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b fsrm md_clear serialize tsxldtrk avx512_fp16 arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          768 KiB (24 instances)
L1i cache:                          768 KiB (24 instances)
L2 cache:                           96 MiB (24 instances)
L3 cache:                           16 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-23
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Unknown: No mitigations
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI Syscall hardening, KVM SW loop
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.8.0+cu126
[pip3] torchaudio==2.8.0+cu126
[pip3] torchvision==0.23.0+cu126
[pip3] triton==3.4.0
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-nvrtc-cu12==12.6.77
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-nccl-cu12==2.27.3
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvidia-nvtx-cu12==12.6.77
[conda] Could not collect

cc @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @Chillee @drisspg @yanboliang @BoyuanFeng

[drisspg]

So the eager baseline is going be pretty different in implementation compared to what SDPA w/ masking will run especially in bfloat16, can you run w/ compiled

[ezyang]

I can repro with the script and the divergence disappears with torch.compile

Mask 580/1000 (58%): Error 0.0% (OK)
Mask 590/1000 (59%): Error 0.0% (OK)

But Driss, doesn't this just mean we have to try harder to have a more accurate eager implementation?

[drisspg]

I think there are a few things here;
Our eager implementation is very very different to flex-attention's compiled implementation. We are a vmapped, unfused impl. This means that we don't actually do iterative softmax which can have a significant impact on numerics, in fact this in theory is a benefit for the eager impl see internal note Ed: https://fburl.com/workplace/heavwoh5

Besides, and what I suspect is showcasing in this impl is that we don't actually skip work in the eager impl but instead we emulate it by masking all attention scores to 0 where they would be skipped by the kernel for a non None block-mask.

So making a more accurate impl in eager seems like quite a large lift IMO. That being said there may be a real issue for eager but I that would require some more review of whats going on here.

Does NOT occur with float32: Same code with float32 has <0.001% error

leads me to believe that is not the case and this is quark of FP

[zaptrem]

I think there are a few things here;
Our eager implementation is very very different to flex-attention's compiled implementation. We are a vmapped, unfused impl. This means that we don't actually do iterative softmax which can have a significant impact on numerics, in fact this in theory is a benefit for the eager impl see internal note Ed: https://fburl.com/workplace/heavwoh5

Besides, and what I suspect is showcasing in this impl is that we don't actually skip work in the eager impl but instead we emulate it by masking all attention scores to 0 where they would be skipped by the kernel for a non None block-mask.

So making a more accurate impl in eager seems like quite a large lift IMO. That being said there may be a real issue for eager but I that would require some more review of whats going on here.

Does NOT occur with float32: Same code with float32 has <0.001% error

leads me to believe that is not the case and this is quark of FP

1. I believe the relative difference continues to grow with magnitude to sometimes >6-10%, which seems a little too much to attribute only to fp weirdness. I guess it could be okay if the error is normally distributed or something? I found this while trying to figure out why my model sounded weird and it hurt my trust in flex attention.

2. Can you explain the rationale in the internal doc for us outsiders?

Thanks!

[drisspg]

1.) I think that this is a more accurate representation of the situation

assert_close_with_ulp is a small utility that runs is_close which is what is run underhood by PT for assert_close and for all failure prints the ULP distance from expected(reference) to what was calculated.

import torch
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
import torch.nn.functional as F
from transformer_nuggets.numerics import assert_close_with_ulp

torch.manual_seed(42)
B, H, Q, D, KV = 1, 1, 1, 64, 1000

# Test showing error threshold: masking <=58% causes bug, >=59% is OK
print("PyTorch flex_attention bfloat16 bug demonstration:")
print("-" * 50)

initial_dtype = torch.float64
dtype = torch.bfloat16
torch.manual_seed(0)
for mask_amount in [580, 590]:
    print(f"Mask {mask_amount}/{KV} ({mask_amount/KV:.0%})")
    mask = create_block_mask(
        lambda b, h, q, kv: kv < mask_amount,
        B, H, Q, KV, "cuda"
    )

    q = torch.randn(B, H, Q, D, dtype=initial_dtype, device="cuda") * 5
    k = torch.randn(B, H, KV, D, dtype=initial_dtype, device="cuda") * 5
    v = torch.randn(B, H, KV, D, dtype=initial_dtype, device="cuda") * 5

    q_lp, k_lp, v_lp = q.to(dtype), k.to(dtype), v.to(dtype)

    flex_out = flex_attention(q_lp, k_lp, v_lp, block_mask=mask)
    # flex_out = torch.compile(flex_attention)(q, k, v, block_mask=mask)
    ref_out = F.scaled_dot_product_attention(q, k[:, :, :mask_amount], v[:, :, :mask_amount])
    ref_low_p = F.scaled_dot_product_attention(q_lp, k_lp[:, :, :mask_amount], v_lp[:, :, :mask_amount])
    # assert_close_with_ulp(flex_out, ref_out)
    error_ref = (ref_low_p - ref_out).abs().max() / ref_out.abs().max()
    error_flex = (flex_out - ref_out).abs().max() / ref_out.abs().max()
    print(f"  Reference error: {error_ref:.2%})")
    print(f"  Flex error: {error_flex:.1%})")

    try:
        assert_close_with_ulp(ref_low_p, ref_out.to(dtype))
    except Exception as E:
        print("REF ERROR STATS".center(100,"-"))
        print(E)
    try:
        assert_close_with_ulp(flex_out, ref_out.to(dtype))
    except Exception as E:
        print("FLEX ERROR STATS".center(100,"-"))
        print(E)

WhenI run this w/ initial_dtype of bf16

initial_dtype = torch.bfloat16
dtype = torch.bfloat16

We get:

❯ python flex/eager_numbers.py
PyTorch flex_attention bfloat16 bug demonstration:
--------------------------------------------------
Mask 580/1000 (58%)
/home/dev/meta/pytorch/torch/nn/attention/flex_attention.py:1479: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.

SOLUTION: Use torch.compile(flex_attention)(...)

If you want to debug your score_mod/mask_mod, you can set:
torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True

This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
  _warn_once(
  Reference error: 0.00%)
  Flex error: 3.4%)
------------------------------------------FLEX ERROR STATS------------------------------------------
Tensor-likes are not close!

Mismatched elements: 62 / 64 (96.9%)
Greatest absolute difference: 0.375 at index (0, 0, 0, 11) (up to 1e-08 allowed)
Greatest relative difference: 1.50781 at index (0, 0, 0, 54) (up to 1e-05 allowed)

============================================================
ULP Analysis of Failures:
============================================================

Total failures: 62
ULP distances: min=1, max=176, mean=26.0

Top 10 failures by absolute difference:
  #  | Index                     | Abs Diff    | Rel Diff    | ULP  | Expected     | Actual
----------------------------------------------------------------------------------------------------
   1 | (0, 0, 0, 11)             | 3.750000e-01 | 6.738281e-02 |   12 |    -5.562500 |    -5.187500
   2 | (0, 0, 0, 54)             | 3.164062e-01 | 1.507812e+00 |  176 |     0.209961 |     0.527344
   3 | (0, 0, 0, 7)              | 3.125000e-01 | 4.394531e-02 |   10 |    -7.125000 |    -6.812500
   4 | (0, 0, 0, 33)             | 3.125000e-01 | 5.029297e-02 |   10 |     6.218750 |     5.906250
   5 | (0, 0, 0, 43)             | 3.125000e-01 | 2.819824e-02 |    5 |   -11.062500 |   -10.750000
   6 | (0, 0, 0, 1)              | 3.046875e-01 | 5.117188e-01 |  131 |     0.597656 |     0.292969
   7 | (0, 0, 0, 53)             | 3.046875e-01 | 1.777344e-01 |   39 |    -1.718750 |    -1.414062
   8 | (0, 0, 0, 25)             | 2.851562e-01 | 4.804688e-01 |   73 |     0.593750 |     0.878906
   9 | (0, 0, 0, 55)             | 2.734375e-01 | 5.664062e-01 |  160 |     0.482422 |     0.209961
  10 | (0, 0, 0, 58)             | 2.695312e-01 | 2.871094e-01 |   69 |    -0.937500 |    -0.667969

Note: Maximum absolute and relative errors occur at different locations
  Max abs diff location (0, 0, 0, 11): 12 ULP
  Max rel diff location (0, 0, 0, 54): 176 ULP
Mask 590/1000 (59%)
  Reference error: 0.00%)
  Flex error: 1.1%)
------------------------------------------FLEX ERROR STATS------------------------------------------
Tensor-likes are not close!

Mismatched elements: 56 / 64 (87.5%)
Greatest absolute difference: 0.125 at index (0, 0, 0, 30) (up to 1e-08 allowed)
Greatest relative difference: 0.147461 at index (0, 0, 0, 61) (up to 1e-05 allowed)

============================================================
ULP Analysis of Failures:
============================================================

Total failures: 56
ULP distances: min=1, max=27, mean=2.9

Top 10 failures by absolute difference:
  #  | Index                     | Abs Diff    | Rel Diff    | ULP  | Expected     | Actual
----------------------------------------------------------------------------------------------------
   1 | (0, 0, 0, 30)             | 1.250000e-01 | 1.086426e-02 |    2 |    11.500000 |    11.375000
   2 | (0, 0, 0, 17)             | 9.375000e-02 | 1.757812e-02 |    3 |    -5.343750 |    -5.250000
   3 | (0, 0, 0, 34)             | 9.375000e-02 | 1.226807e-02 |    3 |    -7.625000 |    -7.531250
   4 | (0, 0, 0, 40)             | 9.375000e-02 | 1.586914e-02 |    3 |    -5.906250 |    -5.812500
   5 | (0, 0, 0, 45)             | 9.375000e-02 | 2.233887e-02 |    3 |    -4.187500 |    -4.093750
   6 | (0, 0, 0, 1)              | 7.031250e-02 | 5.834961e-02 |    9 |    -1.203125 |    -1.132812
   7 | (0, 0, 0, 2)              | 6.250000e-02 | 6.072998e-03 |    1 |   -10.312500 |   -10.250000
   8 | (0, 0, 0, 4)              | 6.250000e-02 | 8.605957e-03 |    2 |    -7.250000 |    -7.187500
   9 | (0, 0, 0, 12)             | 6.250000e-02 | 7.507324e-03 |    1 |    -8.312500 |    -8.250000
  10 | (0, 0, 0, 15)             | 6.250000e-02 | 1.324463e-02 |    2 |    -4.718750 |    -4.656250

Note: Maximum absolute and relative errors occur at different locations
  Max abs diff location (0, 0, 0, 30): 2 ULP
  Max rel diff location (0, 0, 0, 61): 27 ULP
~/me/my_scripts main ?4 ❯       

which I think is a reasonable delta, I suspect that the pseudo phase shift is because w/ 580 you are masking less and thus more V entries are contributing to your final output

This is what I get if I set the initial_dtype to float.64 (which will under the hood also enable the math backend since it is the only one that supports fp64):

❯ python flex/eager_numbers.py
PyTorch flex_attention bfloat16 bug demonstration:
--------------------------------------------------
Mask 580/1000 (58%)
/home/dev/meta/pytorch/torch/nn/attention/flex_attention.py:1479: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.

SOLUTION: Use torch.compile(flex_attention)(...)

If you want to debug your score_mod/mask_mod, you can set:
torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True

This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
  _warn_once(
  Reference error: 0.36%)
  Flex error: 0.5%)
------------------------------------------REF ERROR STATS-------------------------------------------
Tensor-likes are not close!

Mismatched elements: 13 / 64 (20.3%)
Greatest absolute difference: 0.0625 at index (0, 0, 0, 3) (up to 1e-08 allowed)
Greatest relative difference: 0.00747681 at index (0, 0, 0, 26) (up to 1e-05 allowed)

============================================================
ULP Analysis of Failures:
============================================================

Total failures: 13
ULP distances: min=1, max=1, mean=1.0

Top 10 failures by absolute difference:
  #  | Index                     | Abs Diff    | Rel Diff    | ULP  | Expected     | Actual
----------------------------------------------------------------------------------------------------
   1 | (0, 0, 0, 3)              | 6.250000e-02 | 6.896973e-03 |    1 |     9.062500 |     9.000000
   2 | (0, 0, 0, 7)              | 3.125000e-02 | 6.408691e-03 |    1 |    -4.875000 |    -4.906250
   3 | (0, 0, 0, 2)              | 3.125000e-02 | 6.591797e-03 |    1 |     4.750000 |     4.718750
   4 | (0, 0, 0, 26)             | 3.125000e-02 | 7.476807e-03 |    1 |    -4.187500 |    -4.218750
   5 | (0, 0, 0, 44)             | 3.125000e-02 | 5.767822e-03 |    1 |    -5.406250 |    -5.375000
   6 | (0, 0, 0, 10)             | 3.125000e-02 | 4.089355e-03 |    1 |     7.656250 |     7.687500
   7 | (0, 0, 0, 55)             | 1.562500e-02 | 6.256104e-03 |    1 |    -2.500000 |    -2.515625
   8 | (0, 0, 0, 5)              | 1.562500e-02 | 4.516602e-03 |    1 |     3.468750 |     3.453125
   9 | (0, 0, 0, 23)             | 1.562500e-02 | 6.256104e-03 |    1 |     2.500000 |     2.484375
  10 | (0, 0, 0, 52)             | 7.812500e-03 | 4.638672e-03 |    1 |    -1.687500 |    -1.679688

Note: Maximum absolute and relative errors occur at different locations
  Max abs diff location (0, 0, 0, 3): 1 ULP
  Max rel diff location (0, 0, 0, 26): 1 ULP
------------------------------------------FLEX ERROR STATS------------------------------------------
Tensor-likes are not close!

Mismatched elements: 33 / 64 (51.6%)
Greatest absolute difference: 0.0625 at index (0, 0, 0, 3) (up to 1e-08 allowed)
Greatest relative difference: 0.0175781 at index (0, 0, 0, 21) (up to 1e-05 allowed)

============================================================
ULP Analysis of Failures:
============================================================

Total failures: 33
ULP distances: min=1, max=4, mean=1.1

Top 10 failures by absolute difference:
  #  | Index                     | Abs Diff    | Rel Diff    | ULP  | Expected     | Actual
----------------------------------------------------------------------------------------------------
   1 | (0, 0, 0, 3)              | 6.250000e-02 | 6.896973e-03 |    1 |     9.062500 |     9.000000
   2 | (0, 0, 0, 16)             | 6.250000e-02 | 4.730225e-03 |    1 |    13.187500 |    13.125000
   3 | (0, 0, 0, 24)             | 6.250000e-02 | 5.920410e-03 |    1 |   -10.562500 |   -10.500000
   4 | (0, 0, 0, 29)             | 6.250000e-02 | 6.225586e-03 |    1 |    10.062500 |    10.000000
   5 | (0, 0, 0, 41)             | 6.250000e-02 | 6.805420e-03 |    1 |     9.187500 |     9.125000
   6 | (0, 0, 0, 2)              | 3.125000e-02 | 6.591797e-03 |    1 |     4.750000 |     4.718750
   7 | (0, 0, 0, 4)              | 3.125000e-02 | 4.455566e-03 |    1 |    -7.000000 |    -6.968750
   8 | (0, 0, 0, 6)              | 3.125000e-02 | 6.652832e-03 |    1 |     4.687500 |     4.656250
   9 | (0, 0, 0, 15)             | 3.125000e-02 | 6.164551e-03 |    1 |     5.062500 |     5.031250
  10 | (0, 0, 0, 31)             | 3.125000e-02 | 4.364014e-03 |    1 |     7.156250 |     7.125000

Note: Maximum absolute and relative errors occur at different locations
  Max abs diff location (0, 0, 0, 3): 1 ULP
  Max rel diff location (0, 0, 0, 21): 4 ULP
Mask 590/1000 (59%)
  Reference error: 0.24%)
  Flex error: 0.2%)
------------------------------------------REF ERROR STATS-------------------------------------------
Tensor-likes are not close!

Mismatched elements: 2 / 64 (3.1%)
Greatest absolute difference: 0.00390625 at index (0, 0, 0, 37) (up to 1e-08 allowed)
Greatest relative difference: 0.00518799 at index (0, 0, 0, 44) (up to 1e-05 allowed)

============================================================
ULP Analysis of Failures:
============================================================

Total failures: 2
ULP distances: min=1, max=1, mean=1.0

Top 2 failures by absolute difference:
  #  | Index                     | Abs Diff    | Rel Diff    | ULP  | Expected     | Actual
----------------------------------------------------------------------------------------------------
   1 | (0, 0, 0, 37)             | 3.906250e-03 | 4.119873e-03 |    1 |    -0.949219 |    -0.945312
   2 | (0, 0, 0, 44)             | 1.953125e-03 | 5.187988e-03 |    1 |     0.376953 |     0.378906

Note: Maximum absolute and relative errors occur at different locations
  Max abs diff location (0, 0, 0, 37): 1 ULP
  Max rel diff location (0, 0, 0, 44): 1 ULP
------------------------------------------FLEX ERROR STATS------------------------------------------
Tensor-likes are not close!

Mismatched elements: 2 / 64 (3.1%)
Greatest absolute difference: 0.00390625 at index (0, 0, 0, 37) (up to 1e-08 allowed)
Greatest relative difference: 0.00518799 at index (0, 0, 0, 44) (up to 1e-05 allowed)

============================================================
ULP Analysis of Failures:
============================================================

Total failures: 2
ULP distances: min=1, max=1, mean=1.0

Top 2 failures by absolute difference:
  #  | Index                     | Abs Diff    | Rel Diff    | ULP  | Expected     | Actual
----------------------------------------------------------------------------------------------------
   1 | (0, 0, 0, 37)             | 3.906250e-03 | 4.119873e-03 |    1 |    -0.949219 |    -0.945312
   2 | (0, 0, 0, 44)             | 1.953125e-03 | 5.187988e-03 |    1 |     0.376953 |     0.378906

Note: Maximum absolute and relative errors occur at different locations
  Max abs diff location (0, 0, 0, 37): 1 ULP
  Max rel diff location (0, 0, 0, 44): 1 ULP

With this we can see that both low precision runs are within 1 ulp of the True reference. So I am proposing that using bfloat16 as the measuring stick isn't correct.

Side note the original repro does not error if you don't scale input by 5. Otherwise the random BF16 draws stay near ±1, the softmax sums never mix meaningfully different exponents, and the kernel’s cancellation error stays hidden. Boosting the input magnitude spreads the activations across several BF16 exponent bins, which is what makes the precision loss show up

[drisspg]

For 2.) Its roughly that iterative softmax's error significant and can be as large as using lower precision inputs

[v0i0]

i was comparing various attn backends the other day, and eager flex is indeed especially bad. this is because it downcasts scores, which other implementations do not, i believe. see #163986