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.3 KiB
JSON
25 lines
No EOL
1.3 KiB
JSON
{
|
|
"backend": "qwen3_flex",
|
|
"capture_cudagraph": true,
|
|
"lora_adapter": null,
|
|
"n_prompts": 128,
|
|
"n_decoded_tokens": 56415,
|
|
"wall_times_s": [
|
|
12.628456736041699,
|
|
11.507176020997576,
|
|
10.393205604981631,
|
|
10.137171553040389,
|
|
11.370744699030183
|
|
],
|
|
"median_wall_s": 11.370744699030183,
|
|
"best_wall_s": 10.137171553040389,
|
|
"decode_tps_median": 4961.416467719275,
|
|
"decode_tps_best": 5565.161811144426,
|
|
"max_new_tokens": 512,
|
|
"peak_memory_gb": 80.88226366043091,
|
|
"sample_completions": [
|
|
"Let $h$ be the height of the tetrahedron. Then, the volume of the tetrahedron is $\\frac{1}{3} \\cdot 120 \\cdot h = 40h$.<end_working_out><SOLUTION>400</SOLUTION>",
|
|
" \nTo solve this problem, we will use the concept of mass points and the properties of similar triangles. \n\nFirst, let's assign masses to the points based on the given information. Since $M$ is the midpoint of $BC$, we can assign a mass of 1 to both $B$ and $C$. This means that the mass at $M$ is 2 (since $",
|
|
" To solve this problem, we need to find the value of \\( n \\) that minimizes the sum \\( \\sum_{i=1}^{n} f(i) \\) under the given conditions. Let's break down the problem step by step.\n\n1. **Understanding the Constraints:**\n - \\( f \\) is a non-negative valued function on \\( \\{1, 2"
|
|
]
|
|
} |