0c60ed8784
Qwen3.6: num_heads=24, num_kv_heads=4, gqa_ratio=6, head_dim=256 Before (expand GQA then bmm): k_flat: [100K, 4, 256] → expand to [100K, 24, 256] → contiguous Memory: 100K × 24 × 256 × 2B = 1.2GB allocated per decode step Then: [24, 256, 100K] @ [24, 1, 256]^T → scores After (broadcast without materializing): k_kv: [100K, 4, 256] → [4, 256, 100K] (no expansion) q: [24, 256] → [4, 6, 1, 256] scores: matmul([4, 6, 1, 256], [4, 1, 256, 100K]) → [4, 6, 100K] Broadcasting handles GQA — K stays at kv_heads size. Memory: 100K × 4 × 256 × 2B = 200MB (6x reduction) For 100K context generating 1000 tokens: Old: 1000 × 1.2GB = 1.2TB total memory traffic for GQA expansion alone New: 1000 × 200MB = 200GB total (saved 1TB of unnecessary data movement) V weighted sum still needs GQA expansion (V @ scores requires matching dims), but the dominant cost (Q @ K^T) is now broadcast.