fused CPU/f32 neox RoPE - #48
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: 70445d9078
ℹ️ 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".
| for bh in 0..b * h { | ||
| let chunk = bh * t * d; | ||
| for it in 0..t { | ||
| let (cos, sin) = self.cos_sin_at(offset + it); |
There was a problem hiding this comment.
When the CPU/f32 fast path is enabled, offset + it is passed to cos_sin_at, which slices self.cos_f32[start..end] without bounds checks returning a Result. For inputs whose prompt/decode position exceeds max_position_embeddings, the previous path failed through self.cos.narrow(...)?, but this default path will panic and crash the caller instead of propagating an error; add the same range validation before entering the fused loop or make cos_sin_at fallible.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1bcdbc82ea
ℹ️ 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".
| /// Apply RoPE (q, k shape: B x H x L x D) | ||
| pub fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> { | ||
| // CPU f32 fast path: fused neox on raw slices, bit-identical to the op path. | ||
| if *FUSED_ROPE && q.device().is_cpu() && q.dtype() == DType::F32 { |
There was a problem hiding this comment.
This condition also catches prompt/prefill calls (seq_len > 1) on CPU f32 tensors, not just single-token decode. The new rope_neox_f32 implementation below iterates b*h*t*d in a single thread, whereas the existing candle_nn::rotary_emb::rope CPU path uses par_chunks for t > 1, so long CPU prefills now bypass the parallel implementation and can become much slower. Please either gate this branch on seq_len == 1 or parallelize it before replacing the op path.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e94b8827e0
ℹ️ 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".
| fn rope_neox_f32(&self, x: &Tensor, offset: usize) -> Result<Tensor> { | ||
| let (b, h, t, d) = x.dims4()?; | ||
| let half = d / 2; | ||
| debug_assert_eq!(half, self.half_d); |
There was a problem hiding this comment.
When this fast path is used with a CPU/F32 decode input whose last dimension does not match the RotaryEmbedding head dim (for example an inconsistent GGUF key_length/projection shape, or a direct caller of the public apply), this debug_assert is compiled out in release. The old candle_nn::rotary_emb::rope path returned a shape error, but the fused path can either index past cos[j]/sin[j] and panic or return a tensor using only part of the RoPE table, so make this a runtime bail! before slicing.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4c6d4d4fec
ℹ️ 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".
| // candle_nn::rotary_emb::rope (same op order) but on raw slices, no apply_op3. | ||
| fn rope_neox_f32(&self, x: &Tensor, offset: usize) -> Result<Tensor> { | ||
| let (b, h, t, d) = x.dims4()?; | ||
| let half = d / 2; |
There was a problem hiding this comment.
When the CPU/f32 decode fast path runs with an odd head dimension, half is floor(d / 2), so the new check still passes for a matching RotaryEmbedding and the loop only writes the paired coordinates, leaving the final coordinate as zero; with d == 1 it reaches the max_pos division with half_d == 0 and panics. The previous rotary_emb::rope path rejected these shapes via its cos_n_embd * 2 != n_embd validation, so this changes an edge/invalid config from a clean error into corrupted q/k values or a panic; bail when d is odd before continuing.
Useful? React with 👍 / 👎.
|
@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". |
1 participant
No description provided.