Description
pegainfer-kernels/csrc/shared/paged_attention.cu keeps its FlashInfer launch bodies as file-local templates (decode_launch, decode_split_kv_launch, prefill_paged_launch, the cta-tile resolver) and its paged-KV view builder as a file-local helper. paged_attention_hd512.cu, the Gemma 4 global-attention unit, cannot reach them, so it carries copies: the include prelude, the paged-KV builder, and the single-prefill and batch-prefill bodies with 512 written in by hand — 89 identical lines, and a second place to keep in step every time the shared body changes (the recent cta-tile override went into one copy and not the other by design, which is exactly the drift the copies invite).
The QK-norm/RoPE prep pair, prefill_attention_hd512.cu and prefill_attention_hd256_plain.cu, do the same to each other: the RoPE pair rotation, the paged element-offset arithmetic, the kernel parameter lists and the three launch wrappers are ~150 identical lines under two names.
Proposed Plan
- A
paged_launch.cuh header owns the prelude, the paged-KV builder and the launch bodies as templates over HEAD_DIM and variant; paged_attention.cu and paged_attention_hd512.cu keep only their extern "C" entry points, each a one-line instantiation. The hd512 unit stays a separate translation unit so the long template builds still compile in parallel.
- A
qk_prep.cuh header owns the rotation and offset helpers and the shared wrapper bodies for the prep pair; each unit keeps its own kernels.
- No entry point changes name, signature or behaviour; the Rust FFI layer is untouched.
Acceptance Criteria
- The set of exported symbols of each affected object and of the kernels static library is identical before and after (
nm diff empty).
- fmt, all-target Clippy at deny-warnings and the library tests stay green; the kernels crate's GPU integration suites for both prep families and the 12B HF fixture gates that drive both attention families pass unchanged.
- A 26B checkpoint answers the pinned-composition prompts with the same bytes as
main.
Description
pegainfer-kernels/csrc/shared/paged_attention.cukeeps its FlashInfer launch bodies as file-local templates (decode_launch,decode_split_kv_launch,prefill_paged_launch, the cta-tile resolver) and its paged-KV view builder as a file-local helper.paged_attention_hd512.cu, the Gemma 4 global-attention unit, cannot reach them, so it carries copies: the include prelude, the paged-KV builder, and the single-prefill and batch-prefill bodies with512written in by hand — 89 identical lines, and a second place to keep in step every time the shared body changes (the recent cta-tile override went into one copy and not the other by design, which is exactly the drift the copies invite).The QK-norm/RoPE prep pair,
prefill_attention_hd512.cuandprefill_attention_hd256_plain.cu, do the same to each other: the RoPE pair rotation, the paged element-offset arithmetic, the kernel parameter lists and the three launch wrappers are ~150 identical lines under two names.Proposed Plan
paged_launch.cuhheader owns the prelude, the paged-KV builder and the launch bodies as templates overHEAD_DIMand variant;paged_attention.cuandpaged_attention_hd512.cukeep only theirextern "C"entry points, each a one-line instantiation. The hd512 unit stays a separate translation unit so the long template builds still compile in parallel.qk_prep.cuhheader owns the rotation and offset helpers and the shared wrapper bodies for the prep pair; each unit keeps its own kernels.Acceptance Criteria
nmdiff empty).main.