| 1 | # Standard attention decode 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 seq_len grid (decode: q_len=1) ---- |
| 52 | # Small grid for quick iteration. Uncomment for full sweep. |
| 53 | |
| 54 | # Batch size 1 |
| 55 | - "q1s1k" |
| 56 | - "q1s512" |
| 57 | - "q1s2k" |
| 58 | - "q1s4k" |
| 59 | - "q1s8k" |
| 60 | - "q1s16k" |
| 61 | - "q1s32k" |
| 62 | |
| 63 | # Batch size 2 |
| 64 | - "2q1s512" |
| 65 | - "2q1s1k" |
| 66 | - "2q1s2k" |
| 67 | - "2q1s4k" |
| 68 | - "2q1s8k" |
| 69 | - "2q1s16k" |
| 70 | - "2q1s32k" |
| 71 | |
| 72 | # Batch size 4 |
| 73 | - "4q1s512" |
| 74 | - "4q1s1k" |
| 75 | - "4q1s2k" |
| 76 | - "4q1s4k" |
| 77 | - "4q1s8k" |
| 78 | - "4q1s16k" |
| 79 | - "4q1s32k" |
| 80 | |
| 81 | # Batch size 8 |
| 82 | - "8q1s1k" |
| 83 | - "8q1s512" |
| 84 | - "8q1s2k" |
| 85 | - "8q1s4k" |
| 86 | - "8q1s8k" |
| 87 | - "8q1s16k" |
| 88 | - "8q1s32k" |
| 89 | |
| 90 | # Batch size 16 |
| 91 | - "16q1s512" |
| 92 | - "16q1s1k" |
| 93 | - "16q1s2k" |
| 94 | - "16q1s4k" |
| 95 | - "16q1s8k" |
| 96 | - "16q1s16k" |
| 97 | - "16q1s32k" |
| 98 | |
| 99 | # Batch size 32 |
| 100 | - "32q1s512" |
| 101 | - "32q1s1k" |
| 102 | - "32q1s2k" |
| 103 | - "32q1s4k" |
| 104 | - "32q1s8k" |
| 105 | - "32q1s16k" |
| 106 | - "32q1s32k" |
| 107 | |
| 108 | # Batch size 64 |
| 109 | - "64q1s1k" |
| 110 | - "64q1s512" |
| 111 | - "64q1s2k" |
| 112 | - "64q1s4k" |
| 113 | - "64q1s8k" |
| 114 | - "64q1s16k" |
| 115 | - "64q1s32k" |
| 116 | |
| 117 | # Batch size 128 |
| 118 | - "128q1s512" |
| 119 | - "128q1s1k" |
| 120 | - "128q1s2k" |
| 121 | - "128q1s4k" |
| 122 | - "128q1s8k" |
| 123 | - "128q1s16k" |
| 124 | - "128q1s32k" |
| 125 | |
| 126 | # Batch size 256 |
| 127 | - "256q1s1k" |
| 128 | - "256q1s512" |
| 129 | - "256q1s2k" |
| 130 | - "256q1s4k" |
| 131 | - "256q1s8k" |
| 132 | - "256q1s16k" |
| 133 | - "256q1s32k" |
| 134 | |
| 135 | # Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER |
| 136 | backends: |
| 137 | - FLASH_ATTN |
| 138 | - TRITON_ATTN |
| 139 | - FLASHINFER |
| 140 | |
| 141 | device: "cuda:0" |
| 142 | profile_memory: false |