feat(kernels): batched decode prep entries at both head dims - #892
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: 4f15c83bef
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| let table_rows = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { | ||
| anyhow::anyhow!( | ||
| "hd256 paged decode prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" | ||
| ) | ||
| })?; |
There was a problem hiding this comment.
When cos_max_pos * rotary_dim exceeds i32::MAX, this new wrapper accepts the tables because it only checks usize overflow, but the CUDA kernel indexes them with the signed-int expression pos * rotary_dim + d; sufficiently long contexts can therefore wrap the offset and read outside the RoPE caches. The existing paged-prefill wrappers explicitly pass the table extent through checked_i32, while both new decode wrappers omit that guard (the hd512 copy is at line 3599), so apply the same validation before launch.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. Fixed both decode wrappers by applying checked_i32 to the complete RoPE table extent
| /// Per-request hd256 prep; `page_origins` is the released local-window front. | ||
| /// Invalid device metadata traps before the first paged-pool access. | ||
| #[allow(clippy::too_many_arguments)] | ||
| pub fn qkv_norm_rope_paged_decode_hd256_plain_into( |
There was a problem hiding this comment.
This introduces scheduler-facing contracts whose central invariant—local rows carry a released-front origin while global rows remain absolute—is not recorded in project documentation; a repo-wide search for #891 and both new symbol names under docs/ returns no match, and this commit changes no docs. Add or update a Gemma 4 or kernels domain document with the rationale and next scheduler step (and index it if new), as required by the repository workflow.
AGENTS.md reference: AGENTS.md:L163-L165
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This slice should only adds unconsumed kernel entries; it does not introduce scheduler or serving behavior. The local/global metadata distinction is documented on the public API and in this PR, while the durable Gemma serving contract will be documented with the first batched-serving consumer, when it becomes reachable.
4f15c83 to
c6d4721
Compare
|
@codex review please |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c6d4721f63
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| q.checked_extent("hd256 paged decode prep q")?; | ||
| q_out.checked_extent("hd256 paged decode prep q_out")?; | ||
| k.checked_extent("hd256 paged decode prep k")?; | ||
| v.checked_extent("hd256 paged decode prep v")?; |
There was a problem hiding this comment.
When a batched decode input contains more than i32::MAX elements, these calls validate only the usize extent, while the CUDA kernel computes source and Q-output offsets with signed int expressions such as token * q_dim; the overflow can therefore read or write the wrong memory. The existing paged-prefill wrappers immediately above narrow the Q/K extents through checked_i32, but both new decode wrappers omit those guards (the hd512 copy is at lines 3624–3633), so apply the same checks before either launch.
Useful? React with 👍 / 👎.
Signed-off-by: Feathbow <feathbow@gmail.com>
c6d4721 to
3c96871
Compare
|
@codex check again plz |
|
Codex Review: Didn't find any major issues. Keep them coming! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
1 participant
Description
Closes #891
A serving step that decodes one token for each of several requests needs its prep metadata per row: every request sits at its own absolute position, holds its own page table, and, for the sliding family, its own released front. The paged preps carry one request's metadata for the whole launch, so a batch can only be served one request at a time.
Both plain paged preps gain a per-token metadata form, where the page row is the window
page_indptrcuts out ofpage_indicesand a row outside that window traps rather than reading a page the request does not hold. The origin stays where it is real: the sliding family releases its front as the window moves, so its form takes a per-request origin; the global family never releases, so its absolute page index is its row index and it takes none. On top of those, one batched decode entry per head dim: at 256 it writes K and V into the pool per token, at 512 it carries the QK norm and the partial rotary for the family whose value is the key's scale-free branch.Nothing calls the new entries yet. The scheduler and the batched serving step that consume them are the next change, which is also where a numeric gate first has something to measure.
Test Env
Single GPU (sm_89, x86_64), CUDA 12.9.
Verification
cargo build/clippy -D warningsforpegainfer-kernelsand forpegainfer-gemma4 --features gemma4(the feature build is what proves the existing single-request entries kept their signatures); the CI package set with--locked.hd256_qk_rope_plain_smoke: 6 passed, including the new batched closed-form case.hd512_qk_rope_smoke: 6 passed, including the new batched case.hd256_window_prefill_reject: 1 passed.hd256_decode_csr_trapis a manual gate, compiled by CI and run by hand: a window of two pages out of a one-page array passes every host check and has to be refused by the kernel. Run separately because a trap poisons the context. Run by hand here: 1 passed, and with the same fixture narrowed to the one legal page the test fails instead, so it is the kernel's refusal it measures.pegainfer-gemma4 --features gemma4 --lib: 26 passed, 8 ignored (the checkpoint gates).Type of Change