Cross-Attention kernel #22

Description

Cross-attention for conditioning on different modalities (text→image, context→predictor).

Operation

# Q from one modality, K/V from another
# query: (B, Tq, D), key/value: (B, Tkv, D)
output = cross_attention(query, key, value, mask=None)

Implementation

• Flash-attention style memory efficiency
• Support for causal and padding masks
• Multi-head with head_dim parameter
• Fused QKV projection option

Tests

• Matches F.scaled_dot_product_attention
• Correct with different Q/KV sequence lengths
• Mask properly applied
• Gradient flows to Q, K, V
• gradcheck passes
• Memory scales O(Tq) not O(Tq * Tkv)

Used In

Perceiver, Flamingo, JEPA predictor, Stable Diffusion