Skip to content

v1.3.0 — split-K decode attention (beats llama.cpp at depth)

Choose a tag to compare

@ai-hpc ai-hpc released this 21 Jun 15:31
· 4 commits to main since this release
624e90a

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 bakes n_splits at
    capture; the kernel derives each split's KV range from *d_pos at replay).
  • Handles FP16 + INT8 KV, Gemma's cache_head_dim sliding 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