Skip to content

v0.1.0-alpha.5 — Path C: Q4_K uint32 decode weight loads

Choose a tag to compare

@ai-hpc ai-hpc released this 15 May 06:32
· 161 commits to main since this release
748831c

Stacked on top of v0.1.0-alpha.3's Path B batched prefill. Validated on Jetson Orin Nano Super 8 GB, Qwen3-4B Q4_K_M, 25 W MAXN SUPER, GPU locked at 918 MHz.

Headline

alpha.3 alpha.5 Δ
Decode 7.5 tok/s 9.1 tok/s +21%
Prefill 15.4 tok/s 15.2 tok/s unchanged
TTFT 1181 ms ~1180 ms unchanged
Output reference bit-identical

Cumulative from the alpha.2 byte-path baseline: prefill +85%, decode +21%, TTFT −46%. Output bit-identical at every merged step.

What shipped — Path C (Q4_K uint32 weight loads)

Each lane reads four packed q-bytes as one `uint32_t` from `blk.qs` instead of one byte. Eliminates the byte-by-byte inner loop on the hot decode kernels. Same math, fewer issued load instructions, constant-per-lane scales. Output bit-identical to the byte path at FP16.

PR Kernel Cumulative decode
#25 Wo (residual-fused) 8.0 tok/s
#26 gate/up pair 8.7 tok/s
#27 QKV triple (Qwen3-4B: Q4_K + Q4_K + Q6_K) 8.9 tok/s
#28 default-on (`JLLM_Q4K_UINT32_LOADS`) 9.1 tok/s

`JLLM_Q4K_UINT32_LOADS=0` opts back into the byte path. Same env-var grammar as `JLLM_BATCHED_PREFILL` (unset = enabled; `=0` = disabled).

Also folded in (direct-to-main between alpha.3 and Path C)

  • Device-resident layer weight arena — weights copied into a per-layer device allocation at load time instead of streaming from mmap'd host memory.
  • Decode GEMV residual-add fusion — the residual `vec_add` is now inside `gemv_quant_add_typed_kernel`. One fewer launch per layer per token.
  • `gemv_rows_per_block` default lowered from 8 to 4 for the decode shapes.

Closed without merge (negative results, documented inline)

  • #23 — CUDA Graphs decode capture. Launches were already overlapping with GPU execution on the LPDDR5-bound workload; eliminating host-side overhead bought ≤0%.
  • #24 — Split-K Q4_K GEMV. SPLIT_K=4 grew block size 4× and register pressure dropped blocks-per-SM 4×, so warps-per-SM was unchanged. Bottleneck is arithmetic intensity per byte loaded, not per-warp MSHR depth.
  • #29 — Q6_K uint32 weight loads. `block_q6_K` is 210 B (not divisible by 4); consecutive blocks alternate between 4-aligned and 2-aligned positions, half the uint32 loads fault with `cudaErrorMisalignedAddress`.

Full diagnoses on the closed PRs and on the Path C umbrella issue (#19).

Known limits

  • Q6_K decode kernels (W_down ~8% of decode, output logits ~7%) are still on the byte-by-byte path. The 210-byte block size rules out uint32 vectorization without breaking GGUF compatibility. `uint16` would yield ~half the Q4_K win on ~half the decode share — deferred until a fresh `nsys` profile justifies the work.
  • Decode `cudaStreamSynchronize` calls, sampling cost, and attention/RoPE overhead haven't been re-profiled since Path C landed. The pre-Path-C profile showing "93% of decode in five K-quant GEMV kernels" is now stale (four of those kernels are vectorized).

Docs

CHANGELOG.md · README.md · docs/performance.md · docs/kernels.md