69723ee31c
Summary of sweep (all with CUDA graph capture):
| Batch | flex tps | vLLM tps | flex / vLLM | flex mem |
|------:|---------:|---------:|------------:|---------:|
| 8 | 680 | 1900 | 35.8 % | 44 GB |
| 16 | 1626 | 3698 | 44.0 % | 44 GB |
| 32 | 3134 | 6318 | 49.6 % | 44 GB |
| 64 | 5474 | 10459 | **52.3 %** | 44 GB |
| 128 | 5565 | 14996 | 37.1 % | 81 GB |
| 256 | 5812 | 21170 | 27.5 % | 154 GB |
Canonical GRPO (batch 64 + LoRA rank 32):
- vLLM: 7775 tok/s / 156 GB
- flex: **5616 tok/s / 44 GB** = **72% of vLLM at 3.5x less memory**
Up from 9 % (transformers CB) at the start of this work.
Best FlexKernelOptions after sweep:
decode: PRESCALE_QK, USE_TMA, BLOCKS_ARE_CONTIGUOUS, num_warps=8, num_stages=3
prefill: FORCE_USE_FLEX_ATTENTION, PRESCALE_QK, USE_TMA
Biggest single win: `num_warps=8` (+28% at batch 64). Inductor's default
picks 4 on small Triton blocks; 8 is better for our decode shapes.
`BLOCKS_ARE_CONTIGUOUS` adds +10% (safe in our setup because
PageTable.reserve allocates pages sequentially on a fresh batch).
TMA adds 2-3%.
Items documented that broke correctness or didn't help:
- ROWS_GUARANTEED_SAFE=true NaNs the softmax on padded batch slots that
only attend to reserved page 0 (mask returns False for every kv_idx).
- BACKEND="TRITON_DECODE" from the docs raises
NameError('TRITON_DECODE is not defined') inside Inductor.
- USE_TMA + torch.compile(call_model_with_flex_kwargs) -> misaligned
address at runtime (compile breaks TMA alignment assumptions).
- torch.compile(flex_attention, mode="max-autotune") nests cudagraph_trees
inside our raw CUDA graph -> "Cannot prepare for replay during
capturing stage". max-autotune-no-cudagraphs works but same throughput
as default mode.
- compile on call_model_with_flex_kwargs: same as eager walker (CUDA
graph capture already fuses every op in the walker).
- num_warps=4 / 16 both slower than num_warps=8.
CLI surface added to qwen3_flex_inference.py:
--decode_kernel_options JSON (FlexKernelOptions for decode)
--prefill_kernel_options JSON (same for prefill)
--compile_model_forward MODE (optional torch.compile on the walker)
25 lines
No EOL
1.4 KiB
JSON
25 lines
No EOL
1.4 KiB
JSON
{
|
|
"backend": "qwen3_flex",
|
|
"capture_cudagraph": true,
|
|
"lora_adapter": null,
|
|
"n_prompts": 16,
|
|
"n_decoded_tokens": 7446,
|
|
"wall_times_s": [
|
|
6.069258489005733,
|
|
4.580210059997626,
|
|
5.2929606509860605,
|
|
5.689690829021856,
|
|
5.988061942043714
|
|
],
|
|
"median_wall_s": 5.689690829021856,
|
|
"best_wall_s": 4.580210059997626,
|
|
"decode_tps_median": 1308.6827076823927,
|
|
"decode_tps_best": 1625.6896304891004,
|
|
"max_new_tokens": 512,
|
|
"peak_memory_gb": 43.700210094451904,
|
|
"sample_completions": [
|
|
" To solve this problem, we need to determine how many ways we can divide a \\(20 \\times 24\\) rectangle into \\(4 \\times 5\\) rectangles. We will consider rotations and reflections as distinct.\n\nFirst, let's calculate the area of the \\(20 \\times 24\\) rectangle:\n\\[\n20 \\times 24 = 480\n",
|
|
" To solve this problem, we need to find the area of the region inside the larger circle \\( C \\) with radius 30 and outside the six smaller congruent circles that form a ring and are each internally tangent to \\( C \\).\n\nFirst, let's denote the radius of each of the six smaller circles as \\( r \\). Since the six smaller circles form a ring and are each externally",
|
|
" \nA 10-digit palindrome has the form \\( \\overline{abcdefghij} \\) where \\( a = j \\), \\( b = i \\), \\( c = h \\), \\( d = g \\), \\( e = f \\), and \\( f = e \\). This means the number can be written as \\( \\overline{abcdeedcba} \\).\n\nTo determine"
|
|
]
|
|
} |