Cpu optimized/qwen3 decode path -> downstream - #39
Conversation
Performance kernels layered on the standardized CPU flash attention: - RawInterleavedKvCache: an f16, head-major raw KV cache that halves the bytes streamed per decode step and lays K/V out so each kv-head reads contiguously; grows on demand. - causal_decode_f16kv_interleaved / causal_prefill_f16kv_headmajor: decode and prefill kernels that read the interleaved f16 cache directly via raw slices (no per-step dequantization, no separate f32 KV copy), with software prefetch of the next K/V row. - FLASH_DECODE_POOL: a dedicated rayon pool sized for the q_len=1 decode case (saturates well below the core count), parallelizing over kv-heads. cpu_flash_attn tests pass (9/9).
DrJesseGlass
changed the title
Adds the int8 SDOT-based quantized GEMM path used for prefill and decode on aarch64 (NEON dotprod) and x86 (AVX2): - Row-tiled GEMM in k_quants::matmul: tile activation rows so each weight column streams once per tile, with rayon-parallel columns; the per-row activation quantization is parallelized over the pool. - GgmlType::vec_dot_multi trait hook + a multi-row NEON kernel that unpacks each Q4_K weight superblock once and dots it against R activation rows, amortizing the nibble-unpack. Bit-identical to the single-row vec_dot. - QStorage/QTensor::cat_rows for bit-exact row concatenation (used by fused projection pairs), plus a QStorage owned-data roundtrip test. All 38 quantized_tests pass (bit-exact vs the scalar reference).
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b83be2efd5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| pool.install(|| { | ||
| if m == 1 { | ||
| let lhs_row = &lhs_b[..k_in_blocks]; | ||
| dst.into_par_iter() |
There was a problem hiding this comment.
When m == 1, this iterates over the entire dst slice even though the function only requires m * n outputs and explicitly allows longer buffers by only rejecting dst.len() < m * n. Any caller that passes a scratch buffer with extra capacity, e.g. matmul((1, k, n), ..., &mut dst) where dst.len() > n, will compute col_idx >= n and slice rhs_t past the available columns, panicking instead of leaving the tail untouched as the previous row-sliced implementation did.
Useful? React with 👍 / 👎.
| fn forward(&self, ids: &Tensor) -> Result<Tensor> { | ||
| let _enter = self.span.enter(); | ||
| let dims = ids.dims().to_vec(); | ||
| let ids = ids.flatten_all()?.to_vec1::<u32>()?; |
There was a problem hiding this comment.
This CPU quantized embedding replaces the normal Embedding path, but to_vec1::<u32>() rejects index tensors with dtypes that index_select previously accepted, notably I64 (common for token IDs) and U8. In those cases quantized Qwen3 now fails before lookup on CPU while the dense/GPU path still works, so the gather should either mirror index_select's accepted index dtypes or normalize them before this call.
Useful? React with 👍 / 👎.
3356a33 to
f6f89ac
Compare
Add vdotq_s32_acc (accumulate-chain SDOT) and use it in both vec_dot_q4k_q8k (decode GEMV) and vec_dot_q4k_q8k_xr (multi-row prefill): both sub-block-half dots accumulate into one register via SDOT's native accumulate, then one reduction - instead of vdupq_n_s32(0) + a vaddq per pair. Integer accumulation is associative, so this is bit-identical to the prior form (quantized_tests q4k pass). It brings candle's Q4_K kernels to instruction parity with ggml's hand-tuned C: ~+9.7% cumulative 1-thread decode and ~+4-5% prefill on Graviton2/N1.
When m=1 and the decode pool has a single thread (the Lambda 1-vCPU tier), run the GEMV serially instead of pool.install + into_par_iter - that path is pure rayon split/join overhead per call (decode issues ~200 matmuls/token) with no parallelism to gain. Bit-identical (same per-column vec_dot); gated CANDLE_DECODE_SERIAL (default on, =0 restores the rayon path). ~+2.5% N1 1-thread decode.
RoPE is ~13% of single-thread decode and almost all of it is per-Tensor-op framework overhead - ~150 tiny ops/token (narrow / to_dtype / contiguous / apply_op3, x2, per layer). rope_neox_f32 does the neox rotation directly on the CPU f32 slice using the already-present zero-alloc cos_sin_at(pos), building one output tensor - no intermediate cos/sin tensors, no apply_op3 dispatch. Bit-identical to candle_nn::rotary_emb::rope by construction (same op order: dst[i1]=a*cos-bb*sin, dst[i2]=a*sin+bb*cos). Gated by CANDLE_ROPE_FUSED (default on); only taken on CPU f32, else the original path runs. Measured ~+3.4% M1 / ~+1.1% N1 1-thread decode.
causal_decode_f16kv_interleaved issues only ~h_kv tiny tasks; on a single-thread decode pool (the Lambda 1-vCPU tier) the rayon split/join is pure per-call overhead with no parallelism to gain. Factor the per-kv-head work into a shared closure and run it serially when the pool has <=1 thread. Bit-identical to the parallel path (same closure); gated CANDLE_FLASH_SERIAL (default on, =0 forces rayon). ~+0.6% N1 1-thread decode.
Add dot_f16_f16 (aarch64 fmla .8h with two accumulators + a portable reference) and route the decode flash score through it under the opt-in f16-attn-dot cargo feature: each kv-head narrows its rk query rows to f16 once (amortized over kv_len), so the inner q.k is a pure f16.f16 dot. Halves the dot's FMA count and drops the in-loop fcvtl; ~+1.4% N1 1-thread decode. The f16 accumulator loses ~1e-3 rel (softmax + greedy argmax absorb it; greedy text stays coherent), so this is NOT bit-exact and is default-OFF - the exact f32-accumulating dot_f32_f16 runs unless the feature is enabled.
… deduplicated pool builders (simplification/reuse)
vec_dot_multi was fixed to vec_dot_q4k_q8k_xr::<2> (the Cortex-X925 tuning, where R=4 spilled). On Neoverse-N1 a weight column needs only 4 nibble vectors, so wider row groups fit the register file and amortize the unpack further. Replace the fixed-R=2 loop with a descending 8->4->2->1 ladder capped by CANDLE_Q4K_XR_R (clamped 2..=8, default 4 = the validated N1 win), and raise the _xr kernel's group cap from 4 to 8. Bit-identical to vec_dot (quantized_tests q4k pass at R=2/4/8). ~+5-9% prefill on N1; the env knob keeps per-arch tuning available.
…3-integration-base
|
@codex review |
f6f89ac to
61c3d61
Compare
|
@codex review |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…embedding Wire the CPU quantized Qwen3 model onto the serverless kernels: - FusedPairProj: fuse q/k and gate/up into one quantized matmul via QTensor::cat_rows when the two weights share a dtype and inner dim. Rows are quantized independently in every ggml format, so this is a bit-exact row concatenation that halves dispatch/fork-join cost and shares one activation quantization. CPU only; other devices keep the split projections. attn_v stays separate (often a wider quant). - f16 raw-KV flash path: prefill writes the f16 RawInterleavedKvCache and runs causal_prefill_f16kv_headmajor directly over it, the same cache decode reads via causal_decode_f16kv_interleaved. No separate f32 KV copy and no per-layer cat of the full cache. Drops the old InterleavedKvCache prefill staging. - KV prealloc is the initial reservation only (the cache grows on demand); size it for serverless via CANDLE_KV_PREALLOC to cut idle RAM (~0.11 MB/position for Qwen3-0.6B). Default 1024. - QuantizedEmbedding / EmbedTokens: on CPU keep the token embedding quantized and dequantize only the rows each forward needs, instead of materializing the full f32 table (~1 GB RSS at 128k vocab). The tied lm_head shares the same quantized bytes. - Skip causal-mask materialization only when the CPU flash path actually runs (b == 1); batched CPU falls back to standard attention.
414cbca to
bdd7d79
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bdd7d799da
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let mut rows = Vec::with_capacity(ids.len() * self.row_bytes); | ||
| for &id in &ids { | ||
| let off = id as usize * self.row_bytes; | ||
| rows.extend_from_slice(&data[off..off + self.row_bytes]); |
There was a problem hiding this comment.
When the CPU quantized embedding path receives an id that is >= the embedding vocab size, this direct slice indexes past data and panics, whereas the previous dense Embedding path routed through index_select and returned a Result error. This can abort a CPU Qwen3 request for a tokenizer/model mismatch or otherwise invalid input id; check id against the row count before slicing and bail instead of panicking.
Useful? React with 👍 / 👎.
1 participant
No description provided.