Repository navigation
QSA block top-k runs a 1M-token window on one program per row #566
Description
Before you start
- I have read the FAQ and my problem is not answered there.
- I have read the Roadmap and this is not already planned there.
- I have searched existing issues and found no duplicate.
- I am on the latest release, or on a freshly rebuilt
mainwhen building from source.
What happened
qsa_block_topk splits a wide row so its radix select runs on many SMs, but _split_plan gives up when the chunk it needs exceeds _MAX_RESIDENT (8,192): the merge holds every split's k candidates resident, so at most 16 splits of k = 512. A 1,048,576-token window at index_ratio 4 scores 262,144 blocks per row, which needs 16,384-wide chunks, so the plan is None and the whole row goes through one program: a single SM running seven radix passes over up to 262,144 scores, in every QSA layer, for every decoded token.
| window | columns | _split_plan |
|---|---|---|
| 262,144 | 65,536 | (4096, 16) |
| 524,288 | 131,072 | (8192, 16) |
| 1,048,576 | 262,144 | None |
test_block_topk_split_path_cost_tracks_live_blocks uses exactly this geometry (1 row, 262,144 columns, k 512) under a name that says split path, but never asserts that it takes one, so it has been timing the fallback.
Device time of one decode step's indexer per layer (score + top-k + expand), 1 row, Flash-Next's geometry (4 index heads of 128, 64-token pages, ratio 4, k 512), a CUDA graph of 20 replays, best of 7:
| depth | top-k | indexer | x12 QSA layers |
|---|---|---|---|
| 4,096 | 11.1 us | 16.2 us | 0.19 ms |
| 125,769 | 110.2 us | 123.2 us | 1.48 ms |
| 299,105 | 299.7 us | 330.5 us | 3.97 ms |
| 524,288 | 433.9 us | 459.0 us | 5.51 ms |
| 985,621 | 845.4 us | 868.8 us | 10.43 ms |
End to end, a server holding a 985,577-token prompt decodes at 19.19-19.41 t/s on short greedy answers and 25.47 t/s over a 511-token one; with the row split it decodes the same tokens at 23.49-23.76 and 33.65 t/s.
How did you install FreeToken
Built from source
FreeToken version
0d652e7 (kernel/triton/qsa/topk.py is unchanged since v0.1.3, where I measured it)
OS
Windows 11
OS details
Windows 11 Pro 26200, Python 3.13.3, torch 2.11.0+cu130, triton-windows 3.6.0. The kernel is Triton, so nothing here is platform-specific.
GPU and driver
RTX 3090 Ti 24 GB, driver 616.92
CPU and system RAM
AMD Ryzen 9 9950X3D, 256 GB
Checkpoint
RadixArk/Qwen3.8-Flash-Next-NVFP4
Command
python -c "from freetoken.kernel.triton.qsa.topk import _split_plan; print(_split_plan(262144, 512))"Full log
None
Anything else
Serving a 1M window takes more than this fix (#568); the kernel numbers above need none of it. I have a fix with tests and an A/B and will open a PR linked here.
Activity