[xegpu] Performance enhancements for fused attention kernel.

[charithaintc]

This PR adds several perf improvements for the xegpu fused attention example.

1. Simplify the code generation for fused attention loop: Remove the Load/DPAS interleaving approach for K/V loads. Original approach is best for minimal register usage but complicates the prefetch insertion. For now sticking with simpler single K/V loads instead of smaller interleaved loads.
2. Add <fastmath> flag for math.exp for faster exponentiation.
3. Add prefetch for K/V loads.

Performance

Intel Data Center GPU Max 1100, Z=2 H=8 n_ctx=4096 n_head=64 f16, wg_rows=128 sg_rows=16 inner_loop_tile_size=64. Mean of 50 runs after 20 warmup (median of 3
such runs). FLOPs = 4 * Z * H * n_ctx^2 * n_head = 68.72 GFLOP, i.e. the two
matmuls only, as in the flash attention tutorials.

#	Configuration	time (us)	TFLOPS	vs main
1	main — K/V split into 4 chunks interleaved with 8 DPAS, f16 softmax	1447	47.5	1.00x
2	This PR without fastmath<fast> and without prefetch	2429	28.3	0.60x
3	This PR without prefetch	944	72.8	1.53x
4	This PR	674	101.9	2.15x

Rows 2-4 all use a single full-tile K/V load with 2 DPAS ops and an f32 softmax /
f32 accumulators; row 2 adds nothing else, row 3 adds fastmath<fast> on the two
math.exp ops, row 4 adds the K/V xegpu.insert_prefetch (nb_prefetch=1).
None of rows 1-4 spill; the intermediate f16-softmax version of row 2 spilled
3584 B per thread inside the loop, which is what moving the softmax to f32 fixed.

Reproduce:

python examples/xegpu/fused_attention.py --n-ctx=4096 --nruns=50 --nwarmup=20   # row 4
python examples/xegpu/fused_attention.py --n-ctx=4096 --nb-prefetch=0 ...        # row 3

Assisted by: Claude

[charithaintc]

cc @tkarna @rengolin

[tkarna]

Looks good to me!

[tkarna]

Since you only need to add prefetch ops to the loads inside the loop, you could first match the scf.for and then match load ops within the scf.for. The more semantics we can bake into the matching the more robust the schedules become.

[charithaintc]

makes sense. I fixed it.

[charithaintc]

@tkarna please merge if no further comments.