cc033fee19
Wired up two suggestions from the FlashAttention-4 blog + attention-gym: 1. `--fa4_prefill` flag: `BLOCK_SIZE=(256, 128)` + `BACKEND="FLASH"` on the prefill create_block_mask, pad to 256-row Q tile. Confirmed FA4 kernel fires on Blackwell (torch 2.11 + flash-attn CuTeDSL). Output is coherent but 4617 tok/s vs 5744 baseline at batch 64 + LoRA. Root cause: our prefill mask is document_causal, which evaluates `docs[q_idx] == docs[kv_idx]`. The FA4 CuTe kernel's known limitation (documented in attention-gym/examples/flex_flash_attention.py) is that "Indexing by kv_idx is a large perf hit". The doc mask hits that slow path directly. To benefit from FA4 on prefill we would need to refactor the mask so the per-kv lookup goes away, which is non-trivial given the document-boundary + causal combo. 2. flex_autotune_replay.py: new script that drives the pattern from attention-gym/examples/flex_autotune_replay.py -- sets `TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE` + runs with `mode="max-autotune-no-cudagraphs"`, parses the JSON log (handling symbolic dims like `s40`), picks the decode-shape entry (Q_LEN=1), and writes best fwd_* kernel options as JSON. Inductor's best for the decode shape: `fwd_num_warps=4, fwd_num_stages=3, fwd_BLOCK_M=64, fwd_BLOCK_N=64, fwd_USE_TMA=False`. Applied end-to-end: 4827 tok/s vs 5744 manual baseline. The per-call time-minimum Inductor uses doesn't track the cumulative register-spill / L1 effects across the 36-layer stack. Kept `--fa4_prefill` and flex_autotune_replay.py in-tree -- they are useful scaffolding for anyone who wants to push further (refactor the mask, run the 144-config exhaustive fwd sweep from attention-gym/examples/flex_grid_sweep.py, etc.). Default config is unchanged. Also documented the run-to-run variance: over 10 rounds at batch 64 + LoRA, median 4192 and best 5660 tok/s; the spread is GPU clock throttling + variable prompt-length distributions. The 5744 "baseline" we report is best-of-N, matching the prior harness, but steady-state median is closer to 75 % of that. Writeup update in scripts/benchmarks/results/flex_vs_vllm.md.
25 lines
No EOL
1.4 KiB
JSON
25 lines
No EOL
1.4 KiB
JSON
{
|
|
"backend": "qwen3_flex",
|
|
"capture_cudagraph": true,
|
|
"lora_adapter": "outputs/lora_rank32_fresh",
|
|
"n_prompts": 64,
|
|
"n_decoded_tokens": 29601,
|
|
"wall_times_s": [
|
|
11.575495404948015,
|
|
6.411582662025467,
|
|
6.848826738016214,
|
|
7.4589985449565575,
|
|
6.9681332929758355
|
|
],
|
|
"median_wall_s": 6.9681332929758355,
|
|
"best_wall_s": 6.411582662025467,
|
|
"decode_tps_median": 4248.053066068501,
|
|
"decode_tps_best": 4616.800805723188,
|
|
"max_new_tokens": 512,
|
|
"peak_memory_gb": 44.22274446487427,
|
|
"sample_completions": [
|
|
"First, let's find the sum of the numbers in Amanda's list. The sum of the first n even numbers is given by the formula n(n+1). In this case, n = 50 (since there are 50 even numbers from 2 to 100). So, the sum of Amanda's list is 50(50+1) = ",
|
|
" \nTo find the area of the smaller square, we need to determine its side length. Let's denote the side length of the smaller square as \\( s \\).\n\nFrom the diagram, we can see that the larger square has a side length of 6. The smaller square is inscribed within the larger square such that its vertices touch the midpoints of the sides of the larger square. This means",
|
|
" To solve the problem, we need to find the number of ordered pairs \\((x, y)\\) of positive integers that satisfy the inequalities \\(x \\le 2y \\le 60\\) and \\(y \\le 2x \\le 60\\).\n\nFirst, let's rewrite the inequalities in a more convenient form:\n1. \\(x \\le 2y \\le"
|
|
]
|
|
} |