| 1 | # Standard attention prefill benchmark configuration |
| 2 | # Sweeps num_q_heads and num_kv_heads to isolate effects of: |
| 3 | # 1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads) |
| 4 | # 2. Absolute head count (fixed 4:1 ratio, vary scale) |
| 5 | |
| 6 | model: |
| 7 | num_layers: 32 |
| 8 | num_q_heads: 32 # Base value, overridden by sweep |
| 9 | num_kv_heads: 8 # Base value, overridden by sweep |
| 10 | head_dim: 128 |
| 11 | block_size: 16 |
| 12 | |
| 13 | # Head count sweep: each entry overrides num_q_heads, num_kv_heads, and |
| 14 | # head_dim where it differs from the base (128). Head counts are per-GPU |
| 15 | # (i.e. after TP sharding). |
| 16 | # |
| 17 | # Group A — vary GQA ratio (fixed q=32, head_dim=128): |
| 18 | # 32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA) |
| 19 | # |
| 20 | # Groups B-E — real model configs at various TP degrees: |
| 21 | # Model head_dim Full TP2 TP4 TP8 |
| 22 | # Llama 3 8B 128 32:8 16:4 8:2 4:1 |
| 23 | # Llama 3 70B 128 64:8 32:4 16:2 8:1 |
| 24 | # GPT-OSS 120B 64 64:8 32:4 16:2 8:1 |
| 25 | # Llama 3 405B 128 128:8 64:4 32:2 16:1 |
| 26 | model_parameter_sweep: |
| 27 | values: |
| 28 | # --- head_dim=128 (Llama 3 family) --- |
| 29 | - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1 |
| 30 | - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1 |
| 31 | - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8 |
| 32 | - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4 |
| 33 | - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2 |
| 34 | - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1 |
| 35 | - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8 |
| 36 | - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4 |
| 37 | - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1 |
| 38 | - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1 |
| 39 | - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8 |
| 40 | - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4 |
| 41 | - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2 |
| 42 | - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1 |
| 43 | # --- head_dim=64 (GPT-OSS 120B) --- |
| 44 | - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8 |
| 45 | - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4 |
| 46 | - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2 |
| 47 | - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1 |
| 48 | label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}" |
| 49 | |
| 50 | batch_specs: |
| 51 | # ---- batch_size x prefill_len grid (prefill: q_len == seq_len) ---- |
| 52 | # Total tokens = batch_size * prefill_len, and prefill compute scales with |
| 53 | # prefill_len^2, so the largest cells are expensive. Trim batch sizes or |
| 54 | # lengths for quick iteration. |
| 55 | |
| 56 | # Batch size 1 |
| 57 | - "q512" |
| 58 | - "q1k" |
| 59 | - "q2k" |
| 60 | - "q4k" |
| 61 | - "q8k" |
| 62 | - "q16k" |
| 63 | - "q32k" |
| 64 | |
| 65 | # Batch size 2 |
| 66 | - "2q512" |
| 67 | - "2q1k" |
| 68 | - "2q2k" |
| 69 | - "2q4k" |
| 70 | - "2q8k" |
| 71 | - "2q16k" |
| 72 | - "2q32k" |
| 73 | |
| 74 | # Batch size 4 |
| 75 | - "4q512" |
| 76 | - "4q1k" |
| 77 | - "4q2k" |
| 78 | - "4q4k" |
| 79 | - "4q8k" |
| 80 | - "4q16k" |
| 81 | - "4q32k" |
| 82 | |
| 83 | # Batch size 8 |
| 84 | - "8q512" |
| 85 | - "8q1k" |
| 86 | - "8q2k" |
| 87 | - "8q4k" |
| 88 | - "8q8k" |
| 89 | - "8q16k" |
| 90 | - "8q32k" |
| 91 | |
| 92 | # Batch size 16 |
| 93 | - "16q512" |
| 94 | - "16q1k" |
| 95 | - "16q2k" |
| 96 | - "16q4k" |
| 97 | - "16q8k" |
| 98 | - "16q16k" |
| 99 | - "16q32k" |
| 100 | |
| 101 | # Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER |
| 102 | backends: |
| 103 | - FLASH_ATTN |
| 104 | - TRITON_ATTN |
| 105 | - FLASHINFER |
| 106 | |
| 107 | device: "cuda:0" |
| 108 | profile_memory: false |