Context
Google's Gemma 4 (26B-A4B, released March 2026) introduces a hybrid attention design with two different head dimensions within the same model:
- 26 out of 30 layers: sliding window attention with
head_dim=256 — works with FA2
- 4 out of 30 layers: global attention with
head_dim=512 — exceeds FA2's 256 limit
This is the first widely-used open model (to my knowledge) that requires head_dim=512 with multiple attention heads in production. Unlike the case in #801 (single head, head_dim=512), Gemma 4's global attention layers have num_attention_heads=8 and num_key_value_heads=4 with head_dim=512, so there is a real opportunity for FA2/FA3 speedup if the head_dim limit were raised.
Request
Support head_dim=512 in FlashAttention (FA2 and/or FA3/FA4).
Model config details
{
"head_dim": 256,
"global_head_dim": 512,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"sliding_window_pattern": 6,
"num_hidden_layers": 30
}Global attention layers at indices [3, 9, 15, 21]. Each has 8 query heads and 4 KV heads with head_dim=512.
Current error
RuntimeError: FlashAttention only supports head dimensions up to 256
Current workaround
Fall back to SDPA for all 30 layers, losing FA2 speedup even on the 26 layers that are FA2-compatible (head_dim=256).
HuggingFace Transformers could implement per-layer attention dispatch (FA2 for sliding layers, SDPA for global layers), but native FA support for head_dim=512 would be the ideal solution — it would benefit Gemma 4 and any future models adopting larger head dimensions.
Related
Why this matters now
Gemma 4 is likely the first of several models to adopt hybrid head dimensions. As context windows grow and models mix local/global attention patterns, head_dim > 256 for global layers may become common. Supporting 512 would future-proof FlashAttention for this architectural trend.
Context
Google's Gemma 4 (26B-A4B, released March 2026) introduces a hybrid attention design with two different head dimensions within the same model:
head_dim=256— works with FA2head_dim=512— exceeds FA2's 256 limitThis is the first widely-used open model (to my knowledge) that requires
head_dim=512with multiple attention heads in production. Unlike the case in #801 (single head, head_dim=512), Gemma 4's global attention layers havenum_attention_heads=8andnum_key_value_heads=4withhead_dim=512, so there is a real opportunity for FA2/FA3 speedup if the head_dim limit were raised.Request
Support
head_dim=512in FlashAttention (FA2 and/or FA3/FA4).Model config details
{ "head_dim": 256, "global_head_dim": 512, "num_attention_heads": 8, "num_key_value_heads": 4, "sliding_window_pattern": 6, "num_hidden_layers": 30 }Global attention layers at indices
[3, 9, 15, 21]. Each has 8 query heads and 4 KV heads withhead_dim=512.Current error
Current workaround
Fall back to SDPA for all 30 layers, losing FA2 speedup even on the 26 layers that are FA2-compatible (head_dim=256).
HuggingFace Transformers could implement per-layer attention dispatch (FA2 for sliding layers, SDPA for global layers), but native FA support for head_dim=512 would be the ideal solution — it would benefit Gemma 4 and any future models adopting larger head dimensions.
Related
Why this matters now
Gemma 4 is likely the first of several models to adopt hybrid head dimensions. As context windows grow and models mix local/global attention patterns,
head_dim > 256for global layers may become common. Supporting 512 would future-proof FlashAttention for this architectural trend.