Small attention optimization: pre-scale q tile with qk_scale

[AmesingFlank]

Stacked PRs:

• ->Small attention optimization: pre-scale q tile with qk_scale #2157
• Remove redundant m_i update in example attention kernel #2156

---

Small optimization, discovered by @cota. Here, q is of size (block_q, head_dim) which is smaller than qk, which is (block_q, block_k). Typically, head_dim is 64/128 which is likely smaller than block_k. So this results in slightly less ALU pressure.

On TPUs, with

export LIBTPU_INIT_ARGS="--xla_tpu_dvfs_p_state=7 --xla_tpu_scoped_vmem_limit_kib=65536"

B, H, S, D = 8, 32, 8192, 128,

block_sizes = [2, 1024, 1024]

This optimization provides a 1~2 TFLOPs increasing, from a base line of 290 TFLOPs.

It's not much, but it's something.