fix(flex): support GQA/MQA in attention (q_heads % kv_heads == 0)

[AnayGarodia]

Fixes #4930.

Problem

attention_impl (flash path) and attention_naive_impl asserted strict head-count equality between Q and K/V, panicking on ONNX-compliant Grouped-Query / Multi-Query Attention where kv_heads < q_heads.

Fix

Relax the check to require q_heads % kv_heads == 0 (matching the ONNX Attention spec and burn-ndarray) and map each query head to its shared K/V head kv_h = h / (q_heads / kv_heads). K/V batch strides now use kv_heads. Standard MHA is the q_per_kv == 1 special case; mask/bias broadcasting is unchanged (still keyed on q-heads). Applied identically to both the flash and naive paths.

Testing

Added test_attention_gqa_matches_repeated_kv (4 q-heads / 2 kv-heads) and test_attention_mqa_matches_repeated_kv (4 q-heads / 1 kv-head); both assert grouped output equals plain MHA with K/V heads repeated q_heads/kv_heads times, exercising both loops.

cargo test -p burn-flex attention                                 # 14 passed
cargo clippy -p burn-flex --all-targets -- -D warnings            # clean
cargo fmt --check                                                 # clean

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.