Add scaled_dot_product_attention with explicit head axis

[CarloLucibello]

Summary

Introduces a new attention interface, scaled_dot_product_attention, taking inputs with an explicit head axis (head_dim, nheads, seq_len, batch...) as in PyTorch, and deprecates the old packed-head dot_product_attention.

The number of heads is inferred from the tensor shape (size(q, 2)), so there is no nheads keyword anymore.

What's new

• scaled_dot_product_attention(q, k, v, [bias]; fdrop, mask, scale, is_causal) — returns only the attention output (v_head_dim, nheads, q_len, batch...).
• scaled_dot_product_attention_scores(q, k, [bias]; ...) — returns the attention weights (kv_len, q_len, nheads, batch...).
• Grouped-query attention (GQA): key/value may have fewer heads than the query (number of query heads must be divisible by the number of kv heads).
• New scale (denominator, default √head_dim) and is_causal keywords.
• make_causal_mask default changed to dims=3 to match the new layout's sequence axis.

Deprecation

The old dot_product_attention / dot_product_attention_scores keep working with a depwarn, forwarding through the new implementation (heads split out of the feature dim, then joined back on output). Output and scores come from a single pass, so a stochastic fdrop is not applied twice.

Notes / open questions

• This is a API change with deprecations (new exported names, old ones deprecated).
• The kernel is unchanged (batched_mul + softmax + batched_mul) — this PR is about the interface, not a fused/flash kernel. It sets up the nheads/is_causal surface a future NNkernels/cuDNN backend could dispatch on.
• scale is kept as the denominator (matching the previous behavior) rather than PyTorch's multiplier convention — open to changing if preferred.

Tests

• test/attention.jl rewritten for the new API (80 tests pass): batch shapes, GQA, is_causal, scale, scores↔output consistency, gradients through the GQA path.
• test/ext_amdgpu/attention.jl updated to the new layout.

🤖 Generated with Claude Code