cpu-optimized/flash_kv_kernels -> downstream - #35
Conversation
![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: 846459312e
ℹ️ 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".
| /// Decode with interleaved KV cache. No ALiBi, no softcap. | ||
| #[allow(clippy::too_many_arguments)] | ||
| pub fn causal_decode_f32_interleaved( | ||
| pub fn causal_decode_f16kv_interleaved( |
There was a problem hiding this comment.
In the default candle-transformers build, this rename/signature change leaves existing downstream callers broken: I checked cargo check -p candle-transformers --all-targets, and quantized_qwen3.rs plus smol/quantized_smollm3.rs still import/call causal_decode_f32_interleaved, so the lib and lib-test targets fail to compile before reaching the new f16-KV implementation. Please either keep a compatible wrapper or update those callers to the new function name and head_stride/f16-cache contract.
Useful? React with 👍 / 👎.
| let mut res = T::zero(); | ||
| // SAFETY: `a` and `b` are both at least `a.len()` long and `res` is a valid | ||
| // out pointer, pre-zeroed for the scalar fallback that accumulates into it. | ||
| unsafe { T::vec_dot(a.as_ptr(), b.as_ptr(), &mut res, a.len()) }; | ||
| res.to_f64() as f32 |
There was a problem hiding this comment.
When T is f16/bf16, VecOps::vec_dot accumulates internally in f32 but then stores through res: *mut Self, so this helper rounds or saturates the attention logit to half before converting it back to f32. That regresses the f16 varlen path, which previously widened q to f32 and kept f32 dots: B>1 f16 attention or the new f16-KV kernels can produce inf/NaN or distorted softmax weights for large-but-valid f32 dot products. Use a f32 output accumulator for half dtypes instead.
Useful? React with 👍 / 👎.
9ff701c to
63236ce
Compare
DrJesseGlass
changed the title
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).
5751ed3 to
4909bfd
Compare
DrJesseGlass
changed the title
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.
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! 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". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! 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". |
8ec998b to
69eb61e
Compare
|
@codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. 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". |
1 participant
CPU flash attention: f16 KV cache, head-major decode/prefill kernels
Performance kernels layered on the standardized CPU flash attention:
bytes streamed per decode step and lays K/V out so each kv-head reads
contiguously; grows on demand.
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.
case (saturates well below the core count), parallelizing over kv-heads.