Summary
When using the same mask_mod function with both flex_attention (via create_block_mask) and F.scaled_dot_product_attention (via create_mask → dense mask), the attention outputs differ numerically starting from the very first layer, and the discrepancy amplifies significantly through deeper layers.
Setup
Both backends share the identical mask_mod callable (a custom mixed causal mask with sample_id_map and modality_map). The only difference is the attention computation path:
# Flex path
block_mask = create_block_mask(mask_mod, B=B, H=None, Q_LEN=L, KV_LEN=L, BLOCK_SIZE=128)
attn_output = flex_attention(q, k, v, block_mask=block_mask)
# SDPA path
dense_mask = create_mask(mask_mod, B=1, H=H, Q_LEN=L, KV_LEN=L, device=q.device)
attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=dense_mask)
Same q, k, v inputs are passed to both.
Observations
Per-layer mismatch (same batch, atol=1e-3, rtol=1e-2):
| Layer | Mismatch % | Max Abs Diff | Max Relative Diff |
|---|---|---|---|
| 0 | 0.1% | 0.008 | 308x |
| 1 | 3.7% | 0.039 | 4512x |
| 2 | 17.5% | 0.109 | 24320x |
| 4 | 32% | 0.13 | 17536x |
| 8 | 46% | 0.28 | 45568x |
| 12 | 55% | 0.42 | 7680x |
| 23 | 40-62% | 0.37-0.53 | 6720-103424x |
The error grows monotonically from layer 0 → deeper layers, following a classic error amplification pattern in stacked transformers.
Final loss comparison (across 9 steps):
| Metric | Difference |
|---|---|
| Loss (CE) | < 0.1% relative difference |
| Grad norm | 2-8% difference at early steps, converges to ~1% later |
Key Question
Layer 0 already shows a max absolute difference of ~0.008 despite receiving identical q/k/v and mathematically equivalent masks. This suggests the discrepancy originates from the softmax computation itself, not from mask logic differences.
Given that flex_attention uses block-based computation (with create_block_mask + compiled kernel) while SDPA likely uses Flash Attention’s tiled softmax with different accumulation ordering:
-
Is this Layer-0-level discrepancy (~1e-3 abs, ~300x relative) expected due to different floating-point accumulation order in the softmax reduction? Or does it indicate a bug in how
create_block_mask/create_maskmaterialize the samemask_mod? -
Is there a recommended tolerance for comparing
flex_attentionvsSDPAoutputs when they share the same mask? The per-layer divergence makes unit-testing difficult — evenatol=1e-3fails at Layer 0, and the error snowballs to 0.3-0.5 abs diff by Layer 12. -
Is there a way to force both backends to use the same numerical reduction path (e.g., same softmax tiling strategy) so that outputs are closer to bit-exact, at least for testing purposes?
Additional Context
- We verified that
mask_modproduces identical logical masks by iteratingmask_mod(b, h, q_idx, kv_idx)for all positions — no isolated queries (all positions have visible KV tokens). - Running in
bf16precision.
Detail logs
[Loss Compare]
Text_CrossEntropyLoss Flex=6.52010918 SDPA=6.52417183 Diff=4.062653e-03 Ratio=0.999377
total_loss Flex=6.52010918 SDPA=6.52417183 Diff=4.062653e-03 Ratio=0.999377
共 24 个 attention 层
[FAIL] language_model.layers.0.self_attn: Tensor-likes are not close!
Mismatched elements: 22 / 129920 (0.0%)
Greatest absolute difference: 0.00341796875 at index (0, 78, 5, 30) (up to 0.001 allowed)
Greatest relative difference: 0.8125 at index (0, 89, 7, 5) (up to 0.01 allowed)
[FAIL] language_model.layers.1.self_attn: Tensor-likes are not close!
Mismatched elements: 8375 / 129920 (6.4%)
Greatest absolute difference: 0.0234375 at index (0, 1, 12, 23) (up to 0.001 allowed)
Greatest relative difference: 320.0 at index (0, 129, 4, 15) (up to 0.01 allowed)
[FAIL] language_model.layers.10.self_attn: Tensor-likes are not close!
Mismatched elements: 66125 / 129920 (50.9%)
Greatest absolute difference: 0.15625 at index (0, 98, 6, 1) (up to 0.001 allowed)
Greatest relative difference: 5888.0 at index (0, 113, 11, 30) (up to 0.01 allowed)
[FAIL] language_model.layers.11.self_attn: Tensor-likes are not close!
Mismatched elements: 48616 / 129920 (37.4%)
Greatest absolute difference: 0.1875 at index (0, 74, 9, 46) (up to 0.001 allowed)
Greatest relative difference: 2704.0 at index (0, 130, 12, 33) (up to 0.01 allowed)
[FAIL] language_model.layers.12.self_attn: Tensor-likes are not close!
Mismatched elements: 74103 / 129920 (57.0%)
Greatest absolute difference: 0.18359375 at index (0, 64, 0, 36) (up to 0.001 allowed)
Greatest relative difference: 1728.0 at index (0, 89, 3, 49) (up to 0.01 allowed)
[FAIL] language_model.layers.13.self_attn: Tensor-likes are not close!
Mismatched elements: 69709 / 129920 (53.7%)
Greatest absolute difference: 0.234375 at index (0, 104, 8, 33) (up to 0.001 allowed)
Greatest relative difference: 2656.0 at index (0, 58, 0, 35) (up to 0.01 allowed)
[FAIL] language_model.layers.14.self_attn: Tensor-likes are not close!
Mismatched elements: 78551 / 129920 (60.5%)
Greatest absolute difference: 0.18359375 at index (0, 63, 1, 5) (up to 0.001 allowed)
Greatest relative difference: 318.0 at index (0, 17, 11, 35) (up to 0.01 allowed)
[FAIL] language_model.layers.15.self_attn: Tensor-likes are not close!
Mismatched elements: 77262 / 129920 (59.5%)
Greatest absolute difference: 0.193359375 at index (0, 65, 2, 2) (up to 0.001 allowed)
Greatest relative difference: 1784.0 at index (0, 49, 10, 17) (up to 0.01 allowed)
[FAIL] language_model.layers.16.self_attn: Tensor-likes are not close!
Mismatched elements: 61989 / 129920 (47.7%)
Greatest absolute difference: 0.29296875 at index (0, 69, 6, 26) (up to 0.001 allowed)
Greatest relative difference: 1448.0 at index (0, 64, 10, 1) (up to 0.01 allowed)
[FAIL] language_model.layers.17.self_attn: Tensor-likes are not close!
Mismatched elements: 63597 / 129920 (49.0%)
Greatest absolute difference: 0.1484375 at index (0, 71, 0, 17) (up to 0.001 allowed)
Greatest relative difference: 1872.0 at index (0, 90, 2, 62) (up to 0.01 allowed)
[FAIL] language_model.layers.18.self_attn: Tensor-likes are not close!
Mismatched elements: 70307 / 129920 (54.1%)
Greatest absolute difference: 0.234375 at index (0, 64, 1, 1) (up to 0.001 allowed)
Greatest relative difference: 4896.0 at index (0, 16, 8, 8) (up to 0.01 allowed)
[FAIL] language_model.layers.19.self_attn: Tensor-likes are not close!
Mismatched elements: 81260 / 129920 (62.5%)
Greatest absolute difference: 0.134765625 at index (0, 65, 11, 49) (up to 0.001 allowed)
Greatest relative difference: 10496.0 at index (0, 22, 12, 22) (up to 0.01 allowed)
[FAIL] language_model.layers.2.self_attn: Tensor-likes are not close!
Mismatched elements: 24628 / 129920 (19.0%)
Greatest absolute difference: 0.0546875 at index (0, 110, 9, 62) (up to 0.001 allowed)
Greatest relative difference: 1080.0 at index (0, 58, 4, 50) (up to 0.01 allowed)
[FAIL] language_model.layers.20.self_attn: Tensor-likes are not close!
Mismatched elements: 75108 / 129920 (57.8%)
Greatest absolute difference: 0.2265625 at index (0, 64, 1, 31) (up to 0.001 allowed)
Greatest relative difference: 8192.0 at index (0, 58, 0, 50) (up to 0.01 allowed)
[FAIL] language_model.layers.21.self_attn: Tensor-likes are not close!
Mismatched elements: 80848 / 129920 (62.2%)
Greatest absolute difference: 0.359375 at index (0, 11, 9, 34) (up to 0.001 allowed)
Greatest relative difference: 8576.0 at index (0, 30, 1, 63) (up to 0.01 allowed)
[FAIL] language_model.layers.22.self_attn: Tensor-likes are not close!
Mismatched elements: 94080 / 129920 (72.4%)
Greatest absolute difference: 0.25 at index (0, 48, 5, 13) (up to 0.001 allowed)
Greatest relative difference: 4736.0 at index (0, 42, 9, 18) (up to 0.01 allowed)
[FAIL] language_model.layers.23.self_attn: Tensor-likes are not close!
Mismatched elements: 95057 / 129920 (73.2%)
Greatest absolute difference: 0.375 at index (0, 63, 1, 43) (up to 0.001 allowed)
Greatest relative difference: 1120.0 at index (0, 119, 7, 57) (up to 0.01 allowed)
[FAIL] language_model.layers.3.self_attn: Tensor-likes are not close!
Mismatched elements: 33047 / 129920 (25.4%)
Greatest absolute difference: 0.109375 at index (0, 144, 4, 61) (up to 0.001 allowed)
Greatest relative difference: 684.0 at index (0, 84, 10, 12) (up to 0.01 allowed)
[FAIL] language_model.layers.4.self_attn: Tensor-likes are not close!
Mismatched elements: 50552 / 129920 (38.9%)
Greatest absolute difference: 0.078125 at index (0, 102, 1, 14) (up to 0.001 allowed)
Greatest relative difference: 9728.0 at index (0, 69, 1, 60) (up to 0.01 allowed)
[FAIL] language_model.layers.5.self_attn: Tensor-likes are not close!
Mismatched elements: 51365 / 129920 (39.5%)
Greatest absolute difference: 0.0703125 at index (0, 102, 1, 49) (up to 0.001 allowed)
Greatest relative difference: 1752.0 at index (0, 82, 7, 51) (up to 0.01 allowed)
[FAIL] language_model.layers.6.self_attn: Tensor-likes are not close!
Mismatched elements: 51867 / 129920 (39.9%)
Greatest absolute difference: 0.09765625 at index (0, 64, 13, 38) (up to 0.001 allowed)
Greatest relative difference: 5408.0 at index (0, 106, 5, 51) (up to 0.01 allowed)
[FAIL] language_model.layers.7.self_attn: Tensor-likes are not close!
Mismatched elements: 59228 / 129920 (45.6%)
Greatest absolute difference: 0.1171875 at index (0, 115, 11, 4) (up to 0.001 allowed)
Greatest relative difference: 3552.0 at index (0, 70, 8, 44) (up to 0.01 allowed)
[FAIL] language_model.layers.8.self_attn: Tensor-likes are not close!
Mismatched elements: 66295 / 129920 (51.0%)
Greatest absolute difference: 0.13671875 at index (0, 65, 7, 45) (up to 0.001 allowed)
Greatest relative difference: 3904.0 at index (0, 133, 13, 33) (up to 0.01 allowed)
[FAIL] language_model.layers.9.self_attn: Tensor-likes are not close!
Mismatched elements: 54833 / 129920 (42.2%)
Greatest absolute difference: 0.390625 at index (0, 98, 13, 1) (up to 0.001 allowed)
Greatest relative difference: 4288.0 at index (0, 138, 12, 55) (up to 0.01 allowed)
总结: 0 PASS, 24 FAIL (atol=0.001, rtol=0.01)
============================================================
Step=20 fp=45c22128d9cf
Grad norm - Flex: 122.765945, SDPA: 123.155594
[Loss Compare]
Text_CrossEntropyLoss Flex=2.00179005 SDPA=2.00706196 Diff=5.271912e-03 Ratio=0.997373
total_loss Flex=2.00179005 SDPA=2.00706196 Diff=5.271912e-03 Ratio=0.997373
共 24 个 attention 层
[FAIL] language_model.layers.0.self_attn: Tensor-likes are not close!
Mismatched elements: 2510 / 1883392 (0.1%)
Greatest absolute difference: 0.008056640625 at index (0, 1902, 1, 6) (up to 0.001 allowed)
Greatest relative difference: 308.0 at index (0, 883, 1, 28) (up to 0.01 allowed)
[FAIL] language_model.layers.1.self_attn: Tensor-likes are not close!
Mismatched elements: 69901 / 1883392 (3.7%)
Greatest absolute difference: 0.0390625 at index (0, 435, 9, 23) (up to 0.001 allowed)
Greatest relative difference: 4512.0 at index (0, 1315, 8, 57) (up to 0.01 allowed)
[FAIL] language_model.layers.10.self_attn: Tensor-likes are not close!
Mismatched elements: 801407 / 1883392 (42.6%)
Greatest absolute difference: 0.31640625 at index (0, 1902, 6, 47) (up to 0.001 allowed)
Greatest relative difference: 12288.0 at index (0, 2089, 3, 26) (up to 0.01 allowed)
[FAIL] language_model.layers.11.self_attn: Tensor-likes are not close!
Mismatched elements: 693197 / 1883392 (36.8%)
Greatest absolute difference: 0.421875 at index (0, 97, 9, 46) (up to 0.001 allowed)
Greatest relative difference: 5888.0 at index (0, 1303, 9, 29) (up to 0.01 allowed)
[FAIL] language_model.layers.12.self_attn: Tensor-likes are not close!
Mismatched elements: 1043259 / 1883392 (55.4%)
Greatest absolute difference: 0.41796875 at index (0, 1902, 0, 11) (up to 0.001 allowed)
Greatest relative difference: 7680.0 at index (0, 2007, 9, 36) (up to 0.01 allowed)
[FAIL] language_model.layers.13.self_attn: Tensor-likes are not close!
Mismatched elements: 983369 / 1883392 (52.2%)
Greatest absolute difference: 0.2890625 at index (0, 96, 3, 1) (up to 0.001 allowed)
Greatest relative difference: 15104.0 at index (0, 1925, 13, 46) (up to 0.01 allowed)
[FAIL] language_model.layers.14.self_attn: Tensor-likes are not close!
Mismatched elements: 1000623 / 1883392 (53.1%)
Greatest absolute difference: 0.30078125 at index (0, 96, 1, 26) (up to 0.001 allowed)
Greatest relative difference: 30720.0 at index (0, 634, 12, 22) (up to 0.01 allowed)
[FAIL] language_model.layers.15.self_attn: Tensor-likes are not close!
Mismatched elements: 962511 / 1883392 (51.1%)
Greatest absolute difference: 0.28125 at index (0, 98, 2, 15) (up to 0.001 allowed)
Greatest relative difference: 93696.0 at index (0, 303, 5, 16) (up to 0.01 allowed)
[FAIL] language_model.layers.16.self_attn: Tensor-likes are not close!
Mismatched elements: 967249 / 1883392 (51.4%)
Greatest absolute difference: 0.375 at index (0, 1395, 0, 19) (up to 0.001 allowed)
Greatest relative difference: 23168.0 at index (0, 188, 4, 29) (up to 0.01 allowed)
[FAIL] language_model.layers.17.self_attn: Tensor-likes are not close!
Mismatched elements: 730961 / 1883392 (38.8%)
Greatest absolute difference: 0.09375 at index (0, 6, 13, 43) (up to 0.001 allowed)
Greatest relative difference: 7584.0 at index (0, 2047, 1, 6) (up to 0.01 allowed)
[FAIL] language_model.layers.18.self_attn: Tensor-likes are not close!
Mismatched elements: 635743 / 1883392 (33.8%)
Greatest absolute difference: 0.328125 at index (0, 1230, 3, 1) (up to 0.001 allowed)
Greatest relative difference: 3744.0 at index (0, 1123, 4, 6) (up to 0.01 allowed)
[FAIL] language_model.layers.19.self_attn: Tensor-likes are not close!
Mismatched elements: 974955 / 1883392 (51.8%)
Greatest absolute difference: 0.3125 at index (0, 1006, 3, 19) (up to 0.001 allowed)
Greatest relative difference: 6848.0 at index (0, 963, 7, 63) (up to 0.01 allowed)
[FAIL] language_model.layers.2.self_attn: Tensor-likes are not close!
Mismatched elements: 329838 / 1883392 (17.5%)
Greatest absolute difference: 0.109375 at index (0, 953, 2, 6) (up to 0.001 allowed)
Greatest relative difference: 24320.0 at index (0, 1901, 7, 1) (up to 0.01 allowed)
[FAIL] language_model.layers.20.self_attn: Tensor-likes are not close!
Mismatched elements: 941269 / 1883392 (50.0%)
Greatest absolute difference: 0.34375 at index (0, 1904, 10, 24) (up to 0.001 allowed)
Greatest relative difference: 45056.0 at index (0, 324, 6, 43) (up to 0.01 allowed)
[FAIL] language_model.layers.21.self_attn: Tensor-likes are not close!
Mismatched elements: 848215 / 1883392 (45.0%)
Greatest absolute difference: 0.34375 at index (0, 8, 9, 34) (up to 0.001 allowed)
Greatest relative difference: 18688.0 at index (0, 580, 7, 5) (up to 0.01 allowed)
[FAIL] language_model.layers.22.self_attn: Tensor-likes are not close!
Mismatched elements: 676540 / 1883392 (35.9%)
Greatest absolute difference: 0.28125 at index (0, 2078, 12, 3) (up to 0.001 allowed)
Greatest relative difference: 3344.0 at index (0, 1477, 3, 47) (up to 0.01 allowed)
[FAIL] language_model.layers.23.self_attn: Tensor-likes are not close!
Mismatched elements: 752953 / 1883392 (40.0%)
Greatest absolute difference: 0.375 at index (0, 512, 1, 32) (up to 0.001 allowed)
Greatest relative difference: 6720.0 at index (0, 95, 0, 63) (up to 0.01 allowed)
[FAIL] language_model.layers.3.self_attn: Tensor-likes are not close!
Mismatched elements: 290306 / 1883392 (15.4%)
Greatest absolute difference: 0.140625 at index (0, 1834, 2, 61) (up to 0.001 allowed)
Greatest relative difference: 2512.0 at index (0, 1891, 4, 11) (up to 0.01 allowed)
[FAIL] language_model.layers.4.self_attn: Tensor-likes are not close!
Mismatched elements: 603740 / 1883392 (32.1%)
Greatest absolute difference: 0.1328125 at index (0, 1300, 6, 59) (up to 0.001 allowed)
Greatest relative difference: 17536.0 at index (0, 843, 3, 44) (up to 0.01 allowed)
[FAIL] language_model.layers.5.self_attn: Tensor-likes are not close!
Mismatched elements: 620758 / 1883392 (33.0%)
Greatest absolute difference: 0.1171875 at index (0, 1902, 0, 63) (up to 0.001 allowed)
Greatest relative difference: 3680.0 at index (0, 1444, 4, 19) (up to 0.01 allowed)
[FAIL] language_model.layers.6.self_attn: Tensor-likes are not close!
Mismatched elements: 678255 / 1883392 (36.0%)
Greatest absolute difference: 0.1796875 at index (0, 1902, 7, 38) (up to 0.001 allowed)
Greatest relative difference: 85504.0 at index (0, 1889, 1, 39) (up to 0.01 allowed)
[FAIL] language_model.layers.7.self_attn: Tensor-likes are not close!
Mismatched elements: 874510 / 1883392 (46.4%)
Greatest absolute difference: 0.21875 at index (0, 1337, 8, 4) (up to 0.001 allowed)
Greatest relative difference: 45568.0 at index (0, 1501, 7, 7) (up to 0.01 allowed)
[FAIL] language_model.layers.8.self_attn: Tensor-likes are not close!
Mismatched elements: 863251 / 1883392 (45.8%)
Greatest absolute difference: 0.27734375 at index (0, 1903, 7, 46) (up to 0.001 allowed)
Greatest relative difference: 38912.0 at index (0, 1890, 11, 38) (up to 0.01 allowed)
[FAIL] language_model.layers.9.self_attn: Tensor-likes are not close!
Mismatched elements: 650700 / 1883392 (34.5%)
Greatest absolute difference: 0.375 at index (0, 1902, 10, 8) (up to 0.001 allowed)
Greatest relative difference: 5536.0 at index (0, 881, 2, 55) (up to 0.01 allowed)
总结: 0 PASS, 24 FAIL (atol=0.001, rtol=0.01)
============================================================