v1.3.0 — split-K decode attention (beats llama.cpp at depth)
Split-K (flash-decoding) decode attention — fixes genie's decode-at-depth
collapse and pushes decode past llama.cpp at every context depth.
The decode-attention kernels launched grid=(n_heads) — only 8 blocks on the
8-SM Orin (ncu: ~11 % occupancy, 3 % throughput), and each per-head block
walked the whole KV serially, so decode throughput fell with depth. Split-K
splits the KV range across n_splits blocks per head, each computing an
online-softmax partial; a reduce kernel combines them per head. Default-on
(JLLM_SPLITK=0 reverts).
Decode-at-depth vs llama.cpp (same Orin, same gemma-4-E2B-it-Q4_K_M, MAXN_SUPER)
| Context depth | v1.2.1 (no split-K) | v1.3.0 | llama.cpp tg |
v1.3.0 vs llama |
|---|---|---|---|---|
| ~512 | 25.7 | 31.8 | 20.1 | 1.58× |
| ~1024 | 19.7 | 30.5 | 20.8 | 1.47× |
| ~2048 | 14.2 | 26.8 | 20.4 | 1.31× |
| ~4096 | 8.6 | 21.0 | 19.1 | 1.10× |
Without split-K genie was 2.2× behind llama.cpp at 4096 (8.6 vs 19.1); with
it genie leads at every depth. Shallow decode is unchanged (the single-block
kernel is kept below 64 tokens).
Details
flash_decode_splitk_partial_kernel+flash_decode_reduce_kernel, wired into
both the per-step and CUDA-graph decode paths (the graph bakesn_splitsat
capture; the kernel derives each split's KV range from*d_posat replay).- Handles FP16 + INT8 KV, Gemma's
cache_head_dimsliding layers, and the
sliding-window mask. Numerically exact vs the single-block kernel (standalone
tests/test_flash_decode_splitk.cu: max abs diff 0.0, 3.2–7.3× at seq 512–4096).
Full Changelog: v1.2.1...v1.3.0