From e883fd898ef8adcdecdba359c8369397e5f0a4f7 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Thu, 15 Oct 2026 03:57:03 -0700 Subject: [PATCH] Write TiledAttention outputs in token-major row order. LocalAttentionForAllHeadsTokensAndBatch reads Q from row `token_idx * qbatch.Size() + qi`, and the caller adds `att_out` to `x` in that same token-major order. However, it wrote `att_out`, `softmax_d` and `softmax_max` to row `qi * num_tokens + token_idx`. These differ whenever there are several queries and several tokens per query. The dot-product (attention.cc) and flash (flash_attention.cc) paths already use token-major rows. No current caller is affected: prefill uses one query per QBatch and decode one token. Batched embeddings will be the first caller with both. Effect on Gemma 3 4B (sfp) when that case is exercised: 4 different 64-token prompts run through all layers as one QBatch vs. one at a time, comparing the next-token argmax at every position: attention batched == single cos(last-token logits) flash 61-64 / 64 >= 0.99993 flash_transposed_qs 5 / 64 0.70 - 0.87 (before) flash_transposed_qs 63-64 / 64 >= 0.99994 (after) With one query, both row orders coincide, hence single-prompt results are unchanged. The goldens of the multi-token tests (2 tokens, qbatch 2) had encoded the query-major order; their rows 1 and 2 are swapped accordingly. The test fails without the fix. PiperOrigin-RevId: 1000029819 --- gemma/tiled_attention.cc | 3 +- gemma/tiled_attention_test.cc | 66 ++++++++++++++++++----------------- 2 files changed, 36 insertions(+), 33 deletions(-) diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index 6aff4b7c..dd2c00d3 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -1254,8 +1254,9 @@ void LocalAttentionForAllHeadsTokensAndBatch( size_t token_idx = div_heads_per_kv_head.Divide(q_idx); size_t head_in_group_idx = div_heads_per_kv_head.Remainder(q_idx); + // Rows are interleaved token-major, as for the Q read and KV write. const size_t batch_index = - current_qbatch_idx * num_query_tokens + token_idx; + token_idx * qbatch.Size() + current_qbatch_idx; const size_t q_head_idx = kv_head_idx * heads_per_kv_head + head_in_group_idx; const size_t activations_att_out_start_idx = q_head_idx * qkv_dim; diff --git a/gemma/tiled_attention_test.cc b/gemma/tiled_attention_test.cc index dc6259ea..8c605b2e 100644 --- a/gemma/tiled_attention_test.cc +++ b/gemma/tiled_attention_test.cc @@ -433,6 +433,8 @@ void TestLocalAttentionForAllHeadsTokensAndBatch() { } } +// Expected `att_out` of the tests below for 2 tokens and qbatch 2. Rows are +// token-major, i.e. row `token_idx * qbatch_size + qi`, as in `TiledAttention`. const std::vector AttentionMultipleTokensAttentionGoldens = { 34.7414, 34.7717, 34.8022, 34.8327, 34.8631, 34.8936, 34.9241, 34.9545, 34.985, 35.0156, 35.046, 35.0765, 35.1068, 35.1373, 35.1678, 35.1982, @@ -466,38 +468,6 @@ const std::vector AttentionMultipleTokensAttentionGoldens = { 38.1612, 38.1903, 38.2197, 38.249, 38.2781, 38.3075, 38.3368, 38.366, 38.3955, 38.4248, 38.4539, 38.4834, 38.5127, 38.5419, 38.5714, 38.6008, 38.63, 38.6595, 38.6889, 38.7181, 38.7477, 38.777, 38.8063, 38.8358, - 39.0984, 39.1479, 39.1976, 39.2475, 39.297, 39.3468, 39.3967, 39.4463, - 39.4961, 39.546, 39.5957, 39.6455, 39.695, 39.7447, 39.7946, 39.8441, - 39.8939, 39.9438, 39.9934, 40.0431, 40.0931, 40.1427, 40.1925, 40.2425, - 40.2921, 40.342, 40.3915, 40.4412, 40.4911, 40.5407, 40.5904, 40.6403, - 40.6899, 40.7397, 40.7897, 40.8393, 40.8892, 40.9387, 40.9884, 41.0382, - 41.0878, 41.1375, 41.1874, 41.237, 41.2868, 41.3367, 41.3863, 41.4361, - 41.4861, 41.5358, 41.5856, 41.6351, 41.6849, 41.7347, 41.7843, 41.834, - 41.884, 41.9336, 41.9834, 42.0333, 42.083, 42.1328, 42.1823, 42.232, - 38.9699, 39.0188, 39.068, 39.1173, 39.1663, 39.2155, 39.2648, 39.3138, - 39.3631, 39.4124, 39.4615, 39.5108, 39.5597, 39.6089, 39.6581, 39.7071, - 39.7563, 39.8056, 39.8546, 39.9039, 39.9532, 40.0023, 40.0515, 40.1009, - 40.15, 40.1993, 40.2483, 40.2974, 40.3467, 40.3957, 40.4449, 40.4942, - 40.5433, 40.5925, 40.6419, 40.691, 40.7402, 40.7892, 40.8383, 40.8876, - 40.9366, 40.9857, 41.035, 41.0841, 41.1333, 41.1826, 41.2317, 41.2809, - 41.3303, 41.3794, 41.4287, 41.4777, 41.5268, 41.5761, 41.6251, 41.6743, - 41.7237, 41.7727, 41.8219, 41.8713, 41.9204, 41.9697, 42.0186, 42.0677, - 43.4945, 43.5425, 43.5902, 43.6376, 43.6856, 43.7334, 43.7808, 43.8289, - 43.8766, 43.9241, 43.9722, 44.02, 44.0675, 44.1157, 44.1635, 44.2111, - 44.2583, 44.3062, 44.3538, 44.4011, 44.449, 44.4966, 44.544, 44.5919, - 44.6396, 44.6869, 44.735, 44.7826, 44.8301, 44.8781, 44.9258, 44.9733, - 45.0213, 45.0691, 45.1166, 45.1647, 45.2125, 45.26, 45.3081, 45.356, - 45.4035, 45.4508, 45.4987, 45.5462, 45.5936, 45.6415, 45.6891, 45.7364, - 45.7844, 45.832, 45.8794, 45.9274, 45.9751, 46.0225, 46.0705, 46.1183, - 46.1657, 46.2138, 46.2615, 46.309, 46.3571, 46.4049, 46.4525, 46.5006, - 43.4125, 43.4603, 43.5077, 43.5549, 43.6027, 43.6502, 43.6974, 43.7453, - 43.7928, 43.84, 43.8879, 43.9355, 43.9828, 44.0307, 44.0783, 44.1256, - 44.1726, 44.2203, 44.2676, 44.3147, 44.3624, 44.4098, 44.4569, 44.5046, - 44.552, 44.5992, 44.6469, 44.6944, 44.7416, 44.7894, 44.8369, 44.8841, - 44.9319, 44.9795, 45.0267, 45.0746, 45.1222, 45.1694, 45.2173, 45.265, - 45.3123, 45.3593, 45.407, 45.4543, 45.5014, 45.5491, 45.5965, 45.6436, - 45.6913, 45.7387, 45.7859, 45.8336, 45.8811, 45.9283, 45.9761, 46.0236, - 46.0708, 46.1186, 46.1661, 46.2134, 46.2613, 46.3088, 46.3561, 46.404, 34.7729, 34.8035, 34.8341, 34.8648, 34.8953, 34.9259, 34.9567, 34.9872, 35.0179, 35.0486, 35.0792, 35.1098, 35.1404, 35.171, 35.2016, 35.2322, 35.2628, 35.2935, 35.324, 35.3547, 35.3854, 35.416, 35.4466, 35.4774, @@ -530,6 +500,38 @@ const std::vector AttentionMultipleTokensAttentionGoldens = { 38.2159, 38.2451, 38.2747, 38.3042, 38.3334, 38.363, 38.3925, 38.4218, 38.4514, 38.4809, 38.5102, 38.5398, 38.5693, 38.5986, 38.6283, 38.6578, 38.6872, 38.7168, 38.7464, 38.7757, 38.8054, 38.835, 38.8644, 38.8941, + 39.0984, 39.1479, 39.1976, 39.2475, 39.297, 39.3468, 39.3967, 39.4463, + 39.4961, 39.546, 39.5957, 39.6455, 39.695, 39.7447, 39.7946, 39.8441, + 39.8939, 39.9438, 39.9934, 40.0431, 40.0931, 40.1427, 40.1925, 40.2425, + 40.2921, 40.342, 40.3915, 40.4412, 40.4911, 40.5407, 40.5904, 40.6403, + 40.6899, 40.7397, 40.7897, 40.8393, 40.8892, 40.9387, 40.9884, 41.0382, + 41.0878, 41.1375, 41.1874, 41.237, 41.2868, 41.3367, 41.3863, 41.4361, + 41.4861, 41.5358, 41.5856, 41.6351, 41.6849, 41.7347, 41.7843, 41.834, + 41.884, 41.9336, 41.9834, 42.0333, 42.083, 42.1328, 42.1823, 42.232, + 38.9699, 39.0188, 39.068, 39.1173, 39.1663, 39.2155, 39.2648, 39.3138, + 39.3631, 39.4124, 39.4615, 39.5108, 39.5597, 39.6089, 39.6581, 39.7071, + 39.7563, 39.8056, 39.8546, 39.9039, 39.9532, 40.0023, 40.0515, 40.1009, + 40.15, 40.1993, 40.2483, 40.2974, 40.3467, 40.3957, 40.4449, 40.4942, + 40.5433, 40.5925, 40.6419, 40.691, 40.7402, 40.7892, 40.8383, 40.8876, + 40.9366, 40.9857, 41.035, 41.0841, 41.1333, 41.1826, 41.2317, 41.2809, + 41.3303, 41.3794, 41.4287, 41.4777, 41.5268, 41.5761, 41.6251, 41.6743, + 41.7237, 41.7727, 41.8219, 41.8713, 41.9204, 41.9697, 42.0186, 42.0677, + 43.4945, 43.5425, 43.5902, 43.6376, 43.6856, 43.7334, 43.7808, 43.8289, + 43.8766, 43.9241, 43.9722, 44.02, 44.0675, 44.1157, 44.1635, 44.2111, + 44.2583, 44.3062, 44.3538, 44.4011, 44.449, 44.4966, 44.544, 44.5919, + 44.6396, 44.6869, 44.735, 44.7826, 44.8301, 44.8781, 44.9258, 44.9733, + 45.0213, 45.0691, 45.1166, 45.1647, 45.2125, 45.26, 45.3081, 45.356, + 45.4035, 45.4508, 45.4987, 45.5462, 45.5936, 45.6415, 45.6891, 45.7364, + 45.7844, 45.832, 45.8794, 45.9274, 45.9751, 46.0225, 46.0705, 46.1183, + 46.1657, 46.2138, 46.2615, 46.309, 46.3571, 46.4049, 46.4525, 46.5006, + 43.4125, 43.4603, 43.5077, 43.5549, 43.6027, 43.6502, 43.6974, 43.7453, + 43.7928, 43.84, 43.8879, 43.9355, 43.9828, 44.0307, 44.0783, 44.1256, + 44.1726, 44.2203, 44.2676, 44.3147, 44.3624, 44.4098, 44.4569, 44.5046, + 44.552, 44.5992, 44.6469, 44.6944, 44.7416, 44.7894, 44.8369, 44.8841, + 44.9319, 44.9795, 45.0267, 45.0746, 45.1222, 45.1694, 45.2173, 45.265, + 45.3123, 45.3593, 45.407, 45.4543, 45.5014, 45.5491, 45.5965, 45.6436, + 45.6913, 45.7387, 45.7859, 45.8336, 45.8811, 45.9283, 45.9761, 46.0236, + 46.0708, 46.1186, 46.1661, 46.2134, 46.2613, 46.3088, 46.3561, 46.404, 39.1594, 39.2093, 39.2593, 39.3095, 39.3594, 39.4094, 39.4597, 39.5096, 39.5597, 39.61, 39.6599, 39.7101, 39.7599, 39.8099, 39.8601, 39.91, 39.96, 40.0102, 40.0601, 40.1102, 40.1605, 40.2104, 40.2605, 40.3108,