Write TiledAttention outputs in token-major row order.

[copybara-service]

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.