| 1 | # Study 4: What is optimal reorder_batch_threshold for MLA backends supporting query length > 1? |
| 2 | # Question: At what query length does prefill pipeline become faster than decode pipeline? |
| 3 | # Methodology: For each query length, compare decode vs prefill performance to find crossover point |
| 4 | # Applies to: FlashAttn MLA, FlashMLA |
| 5 | |
| 6 | description: "Decode vs Prefill pipeline crossover analysis" |
| 7 | |
| 8 | # Test FlashAttn MLA |
| 9 | backend: FLASH_ATTN_MLA |
| 10 | |
| 11 | # Mode: decode_vs_prefill comparison (special sweep mode) |
| 12 | # For each batch spec, we'll test both decode and prefill pipelines |
| 13 | mode: "decode_vs_prefill" |
| 14 | |
| 15 | # Query lengths to test (from old benchmark_mla_threshold.py methodology) |
| 16 | # Each query length will be tested with BOTH decode and prefill pipelines: |
| 17 | # - decode: threshold >= query_length (forces decode pipeline) |
| 18 | # - prefill: threshold < query_length (forces prefill pipeline) |
| 19 | # |
| 20 | # We use q<N>s1k format which creates q_len=N, seq_len=1024 requests |
| 21 | # This tests different query lengths with fixed sequence length context |
| 22 | # |
| 23 | # Using batch_spec_ranges for automatic generation: |
| 24 | batch_spec_ranges: |
| 25 | - template: "q{q_len}s1k" |
| 26 | q_len: |
| 27 | start: 1 |
| 28 | stop: 16 |
| 29 | step: 1 |
| 30 | end_inclusive: false |
| 31 | - template: "q{q_len}s1k" |
| 32 | q_len: |
| 33 | start: 16 |
| 34 | stop: 64 |
| 35 | step: 2 |
| 36 | end_inclusive: false |
| 37 | - template: "q{q_len}s1k" |
| 38 | q_len: |
| 39 | start: 64 |
| 40 | stop: 1024 |
| 41 | step: 4 |
| 42 | end_inclusive: true |
| 43 | |
| 44 | # Batch sizes to test (from old script) |
| 45 | batch_sizes: |
| 46 | - 1 |
| 47 | - 2 |
| 48 | - 4 |
| 49 | - 8 |
| 50 | - 16 |
| 51 | - 32 |
| 52 | - 64 |
| 53 | - 128 |
| 54 | - 256 |
| 55 | |
| 56 | # Model configuration (DeepSeek V2/V3 defaults) |
| 57 | model: |
| 58 | num_layers: 10 |
| 59 | head_dim: 576 |
| 60 | num_q_heads: 128 |
| 61 | num_kv_heads: 1 |
| 62 | block_size: 128 |
| 63 | |
| 64 | # Benchmark settings |
| 65 | device: "cuda:0" |
| 66 | profile_memory: false |
| 67 | |
| 68 | # Output |
| 69 | output: |
| 70 | csv: "reorder_threshold_results.csv" |
| 71 | json: "reorder_threshold_results.json" |
| 72 | |
| 73 | # Expected outcome (reproduces old benchmark_mla_threshold.py study): |
| 74 | # - For each batch size, find the crossover point where prefill becomes faster than decode |
| 75 | # - Show decode vs prefill performance across all query lengths |
| 76 | # - Determine optimal reorder_batch_threshold based on last query length where decode is faster |
| 77 | # - Understand how crossover point varies with batch size |
| 78 | # - Provide data-driven guidance for default threshold value |
| 79 | # |
| 80 | # Methodology (from old script): |
| 81 | # - Each query length tested with BOTH pipelines: |
| 82 | # * decode: threshold >= query_length (forces decode pipeline) |
| 83 | # * prefill: threshold < query_length (forces prefill pipeline) |
| 84 | # - Compare which is faster to find crossover point |
| 85 | # |