From f42cb4cc84bc6a55747ce68d4f1105dda897175f Mon Sep 17 00:00:00 2001 From: Feathbow Date: Tue, 4 Aug 2026 20:45:10 +0100 Subject: [PATCH] feat(kernels): pool-write prep for Gemma 4 serving at both head dims Signed-off-by: Feathbow --- pegainfer-core/src/ops.rs | 2 + .../shared/prefill_attention_hd256_plain.cu | 244 ++++++++- .../csrc/shared/prefill_attention_hd512.cu | 18 +- pegainfer-kernels/src/ffi/shared.rs | 36 +- pegainfer-kernels/src/ops.rs | 1 + pegainfer-kernels/src/ops/attention.rs | 505 ++++++++++++++---- .../tests/hd256_qk_rope_plain_smoke.rs | 125 ++++- .../tests/hd512_qk_rope_smoke.rs | 12 +- 8 files changed, 809 insertions(+), 134 deletions(-) diff --git a/pegainfer-core/src/ops.rs b/pegainfer-core/src/ops.rs index 4fe395798..b4e7a1416 100644 --- a/pegainfer-core/src/ops.rs +++ b/pegainfer-core/src/ops.rs @@ -63,9 +63,11 @@ pub use pegainfer_kernels::ops::lora_decode_fused_delta_into; pub use pegainfer_kernels::ops::pack_lora_b_rows_into; pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd256_into; pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd512_into; +pub use pegainfer_kernels::ops::qk_norm_partial_rope_paged_prefill_hd512_into; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::qk_norm_rope_batch_decode_into; pub use pegainfer_kernels::ops::qk_norm_rope_prefill_hd256_plain_into; +pub use pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into; pub use pegainfer_kernels::ops::rms_norm; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::rms_norm_batch_into; diff --git a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu b/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu index 7af76279d..f38a7b7fa 100644 --- a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu +++ b/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu @@ -1,10 +1,10 @@ // QK-norm + RoPE prep at head_dim 256 with the plain-w norm (Gemma 4 local // layers). The hd256 sibling under csrc/qwen35/ is not this kernel: it // computes the 1+w offset norm, assumes a gated Q layout twice as wide, and -// writes K into a paged pool. All three are qwen35's choices, not the -// family's; Gemma 4 multiplies by the plain weight and ships no gate. Q and -// K land in contiguous token-major buffers; the attention entry's cache -// layout (single_prefill reads HND) stays the caller's concern. +// ships no v_norm. +// +// Two entry points share the math: the contiguous oracle form and the +// paged serving form. // // rotary_dim is a runtime argument, checked at the launcher for positive, // even and <= HD256. Gemma 4 local layers rotate the full head (256); at @@ -14,7 +14,8 @@ // // Positions derive from host-known start_pos, so the launcher rejects any // out-of-range window before launch; the device trap before the cos read is -// a second layer, not the contract. +// a second layer, not the contract. Page ids are device data, so for them +// the paged form's trap is the only check. #include "common.cuh" #include "ffi_guard.cuh" @@ -133,6 +134,152 @@ __global__ void qk_norm_rope_prefill_hd256_plain_kernel( } } +__device__ __forceinline__ int64_t paged_kv_offset_hd256_plain( + int page_id, + int64_t block_offset_elems, + int64_t stride_page, + int page_size, + int num_kv_heads, + int pos, + int kv_head, + int d) { + int offset_in_page = pos % page_size; + return static_cast(page_id) * stride_page + + block_offset_elems + + static_cast(offset_in_page) * num_kv_heads * HD256_PLAIN + + static_cast(kv_head) * HD256_PLAIN + + d; +} + +// Paged serving prep. grid.y carries three bands: [0, num_q_heads) Q, +// then num_kv_heads K, then num_kv_heads V. Q and K are plain-w normed +// and rotated; V is weightless-normed over its own head vector (v_proj +// output — a separate reduction, unlike the hd512 K=V fork) and never +// rotated. K and V write straight into the pool's per-layer K/V blocks. +__global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( + const __nv_bfloat16* __restrict__ q_batch, // [q_dim, seq_len] + const __nv_bfloat16* __restrict__ k_batch, // [kv_dim, seq_len] + const __nv_bfloat16* __restrict__ v_batch, // [kv_dim, seq_len] + const __nv_bfloat16* __restrict__ q_norm_weight, // [HD256_PLAIN] + const __nv_bfloat16* __restrict__ k_norm_weight, // [HD256_PLAIN] + const __nv_bfloat16* __restrict__ cos_cache, // [max_seq * rotary_dim] + const __nv_bfloat16* __restrict__ sin_cache, + __nv_bfloat16* __restrict__ q_batch_out, // [q_dim, seq_len] + __nv_bfloat16* __restrict__ kv_data, // paged KV pool + int64_t k_offset_elems, + int64_t v_offset_elems, + const int* __restrict__ page_indices, // request page list + int num_q_heads, + int num_kv_heads, + int start_pos, // host base position + int cos_max_pos, // rows in cos/sin tables + int rotary_dim, + float rms_eps, + int page_size, + int num_pages, // pool capacity in pages + int64_t stride_page +) { + int token = blockIdx.x; + int band = blockIdx.y; + int d = threadIdx.x; + + bool is_q = band < num_q_heads; + bool is_k = !is_q && band < num_q_heads + num_kv_heads; + int head_local = is_q ? band + : is_k ? band - num_q_heads + : band - num_q_heads - num_kv_heads; + int q_dim = num_q_heads * HD256_PLAIN; + int kv_dim = num_kv_heads * HD256_PLAIN; + + int src_offset = is_q + ? token * q_dim + head_local * HD256_PLAIN + d + : token * kv_dim + head_local * HD256_PLAIN + d; + __nv_bfloat16 x = is_q ? q_batch[src_offset] + : is_k ? k_batch[src_offset] + : v_batch[src_offset]; + + float sq = __bfloat162float(x); + sq *= sq; + float sq_sum = warp_reduce_sum(sq); + + int warp_id = d / WARP_SIZE; + int lane_id = d % WARP_SIZE; + __shared__ float warp_sums[NUM_WARPS_HD256_PLAIN]; + __shared__ float inv_rms; + __shared__ __nv_bfloat16 smem[HD256_PLAIN]; + + if (lane_id == 0) warp_sums[warp_id] = sq_sum; + __syncthreads(); + + if (d == 0) { + float total = 0.0f; + for (int i = 0; i < NUM_WARPS_HD256_PLAIN; i++) total += warp_sums[i]; + inv_rms = 1.0f / sqrtf(total / HD256_PLAIN + rms_eps); + } + __syncthreads(); + + int pos = start_pos + token; + // Reject before reading the cos/sin tables or the page list. + if (pos < 0 || pos >= cos_max_pos) __trap(); + // Check the device-resident page id before the first pool write. Q + // blocks never touch the pool. + int page_id = -1; + if (!is_q) { + page_id = page_indices[pos / page_size]; + if (page_id < 0 || page_id >= num_pages) __trap(); + } + + if (!is_q && !is_k) { + // V band: weightless norm, no RoPE — the whole block exits here. + int64_t dst = paged_kv_offset_hd256_plain( + page_id, v_offset_elems, stride_page, page_size, + num_kv_heads, pos, head_local, d); + kv_data[dst] = __float2bfloat16(__bfloat162float(x) * inv_rms); + return; + } + + smem[d] = rms_norm_elem_hd256_plain( + x, inv_rms, is_q ? q_norm_weight[d] : k_norm_weight[d]); + __syncthreads(); + + int half_rotary = rotary_dim / 2; + + if (d < half_rotary) { + __nv_bfloat16 lo = smem[d]; + __nv_bfloat16 hi = smem[d + half_rotary]; + apply_rope_pair_hd256_plain( + lo, + hi, + cos_cache[pos * rotary_dim + d], + sin_cache[pos * rotary_dim + d] + ); + + if (is_q) { + int dst = token * q_dim + head_local * HD256_PLAIN; + q_batch_out[dst + d] = lo; + q_batch_out[dst + d + half_rotary] = hi; + } else { + int64_t dst = paged_kv_offset_hd256_plain( + page_id, k_offset_elems, stride_page, page_size, + num_kv_heads, pos, head_local, d); + kv_data[dst] = lo; + kv_data[dst + half_rotary] = hi; + } + } + + if (d >= rotary_dim) { + if (is_q) { + int dst = token * q_dim + head_local * HD256_PLAIN; + q_batch_out[dst + d] = smem[d]; + } else { + int64_t dst = paged_kv_offset_hd256_plain( + page_id, k_offset_elems, stride_page, page_size, + num_kv_heads, pos, head_local, d); + kv_data[dst] = smem[d]; + } + } +} + extern "C" { int qk_norm_rope_prefill_hd256_plain_cuda( @@ -205,4 +352,91 @@ int qk_norm_rope_prefill_hd256_plain_cuda( PEGAINFER_FFI_GUARD_END(-1) } +int qkv_norm_rope_paged_prefill_hd256_plain_cuda( + const __nv_bfloat16* q_batch, + const __nv_bfloat16* k_batch, + const __nv_bfloat16* v_batch, + const __nv_bfloat16* q_norm_weight, + const __nv_bfloat16* k_norm_weight, + const __nv_bfloat16* cos_cache, + const __nv_bfloat16* sin_cache, + __nv_bfloat16* q_batch_out, + __nv_bfloat16* kv_data, + int64_t k_offset_elems, + int64_t v_offset_elems, + const int* page_indices, + int num_q_heads, + int num_kv_heads, + int seq_len, + int start_pos, + int cos_max_pos, + int rotary_dim, + float rms_eps, + int page_size, + int num_pages, + int64_t stride_page, + cudaStream_t stream +) { + PEGAINFER_FFI_GUARD_BEGIN + if (rotary_dim <= 0 || (rotary_dim & 1) != 0 || rotary_dim > HD256_PLAIN) { + pegainfer_ffi_set_last_error( + "qkv_norm_rope_paged_prefill_hd256_plain_cuda: rotary_dim must be " + "positive, even and <= 256"); + return -1; + } + if (q_batch == nullptr || k_batch == nullptr || v_batch == nullptr || + q_norm_weight == nullptr || k_norm_weight == nullptr || + cos_cache == nullptr || sin_cache == nullptr || + q_batch_out == nullptr || kv_data == nullptr || + page_indices == nullptr) { + pegainfer_ffi_set_last_error( + "qkv_norm_rope_paged_prefill_hd256_plain_cuda: null pointer argument"); + return -1; + } + if (num_q_heads <= 0 || num_kv_heads <= 0 || seq_len <= 0 || + page_size <= 0 || num_pages <= 0) { + pegainfer_ffi_set_last_error( + "qkv_norm_rope_paged_prefill_hd256_plain_cuda: num_q_heads, " + "num_kv_heads, seq_len, page_size and num_pages must be positive"); + return -1; + } + if (start_pos < 0 || start_pos + seq_len > cos_max_pos) { + pegainfer_ffi_set_last_error( + "qkv_norm_rope_paged_prefill_hd256_plain_cuda: start_pos + seq_len " + "must be <= cos_max_pos"); + return -1; + } + dim3 prep_grid(seq_len, num_q_heads + 2 * num_kv_heads); + qkv_norm_rope_paged_prefill_hd256_plain_kernel<<>>( + q_batch, + k_batch, + v_batch, + q_norm_weight, + k_norm_weight, + cos_cache, + sin_cache, + q_batch_out, + kv_data, + k_offset_elems, + v_offset_elems, + page_indices, + num_q_heads, + num_kv_heads, + start_pos, + cos_max_pos, + rotary_dim, + rms_eps, + page_size, + num_pages, + stride_page + ); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + pegainfer_ffi_set_last_error(cudaGetErrorString(err)); + return -1; + } + return 0; + PEGAINFER_FFI_GUARD_END(-1) +} + } // extern "C" diff --git a/pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu b/pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu index 7589ab80d..961ed6f32 100644 --- a/pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu +++ b/pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu @@ -1,13 +1,14 @@ // QK-norm + partial RoPE prep for head_dim 512 (Gemma 4 global layers). // Differs from the hd256 sibling in three ways that are choices, not -// oversights: plain w rather than the 1+w offset, no gate, no V. +// oversights: plain w rather than the 1+w offset, no gate, and no separate +// V input — V is the weightless RMS of the same raw row K reduces, so the +// kernel reuses inv_rms and writes V = x * inv_rms into the pool's V block +// alongside K. // // rotary_dim is a runtime argument, checked at the launcher for positive, // even and <= HD512. Evenness is load-bearing: with half_rotary floored, an // odd value leaves index rotary_dim - 1 written by neither branch. -// -// K is written into the paged pool, which feeds batch_prefill_paged; -// single_prefill wants a contiguous cache instead. + // // Positions and page ids are trapped on device — checking either on the // host would require a D2H synchronization. @@ -63,6 +64,7 @@ __global__ void qk_norm_partial_rope_paged_prefill_hd512_kernel( __nv_bfloat16* __restrict__ q_batch_out, // [q_dim, seq_len] __nv_bfloat16* __restrict__ kv_data, // paged KV pool int64_t k_offset_elems, + int64_t v_offset_elems, const int* __restrict__ page_indices, // request page list int num_q_heads, int num_kv_heads, @@ -123,6 +125,12 @@ __global__ void qk_norm_partial_rope_paged_prefill_hd512_kernel( if (!is_q) { page_id = page_indices[pos / page_size]; if (page_id < 0 || page_id >= num_pages) __trap(); + // V is the K=V fork: the weightless norm of the same raw vector, + // sharing inv_rms. No RoPE, no weight. + int64_t v_dst = paged_kv_offset_hd512( + page_id, v_offset_elems, stride_page, page_size, + num_kv_heads, pos, head_local, d); + kv_data[v_dst] = __float2bfloat16(__bfloat162float(x) * inv_rms); } int half_rotary = rotary_dim / 2; @@ -268,6 +276,7 @@ int qk_norm_partial_rope_paged_prefill_hd512_cuda( __nv_bfloat16* q_batch_out, __nv_bfloat16* kv_data, int64_t k_offset_elems, + int64_t v_offset_elems, const int* page_indices, int num_q_heads, int num_kv_heads, @@ -319,6 +328,7 @@ int qk_norm_partial_rope_paged_prefill_hd512_cuda( q_batch_out, kv_data, k_offset_elems, + v_offset_elems, page_indices, num_q_heads, num_kv_heads, diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index d3b84b47d..5b0c8dc89 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/pegainfer-kernels/src/ffi/shared.rs @@ -1018,8 +1018,8 @@ unsafe extern "C" { // on the Rust wrapper in ops::attention; the entry returns 0 on success, -1 // with a diagnostic on failure. unsafe extern "C" { - // Q and K both land in contiguous buffers shaped like their inputs; there - // is no paged pool because no Gemma 4 KV cache consumer exists yet. + // Oracle form: Q and K both land in contiguous buffers shaped like + // their inputs; the paged serving form is the qkv_ entry below. pub fn qk_norm_rope_prefill_hd256_plain_cuda( q_batch: *const Half, k_batch: *const Half, @@ -1038,6 +1038,35 @@ unsafe extern "C" { rms_eps: f32, stream: CUstream, ) -> i32; + + // Paged serving form: Q → contiguous q_batch_out; K (normed + rotated) + // and V (weightless-normed, never rotated) → straight into the paged KV + // pool at k_offset_elems / v_offset_elems. + pub fn qkv_norm_rope_paged_prefill_hd256_plain_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + num_q_heads: i32, + num_kv_heads: i32, + seq_len: i32, + start_pos: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; } // hd512 QK-norm + partial RoPE prep (Gemma 4 global layers): @@ -1047,6 +1076,8 @@ unsafe extern "C" { unsafe extern "C" { // Prefill: Q → contiguous q_batch_out; K → straight into the paged KV // pool at k_offset_elems (feeds batch_prefill_paged, not single_prefill). + // V is the K=V fork — the weightless norm of the same raw K, sharing + // its denominator — written to v_offset_elems in the same pass. pub fn qk_norm_partial_rope_paged_prefill_hd512_cuda( q_batch: *const Half, k_batch: *const Half, @@ -1057,6 +1088,7 @@ unsafe extern "C" { q_batch_out: *mut Half, kv_data: *mut Half, k_offset_elems: i64, + v_offset_elems: i64, page_indices: *const i32, num_q_heads: i32, num_kv_heads: i32, diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index 225bb6d3c..35b2df4e1 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/pegainfer-kernels/src/ops.rs @@ -38,6 +38,7 @@ pub use attention::qk_norm_partial_rope_batched_decode_hd512_into; pub use attention::qk_norm_partial_rope_paged_prefill_hd512_into; pub use attention::qk_norm_rope_batch_decode_into; pub use attention::qk_norm_rope_prefill_hd256_plain_into; +pub use attention::qkv_norm_rope_paged_prefill_hd256_plain_into; pub use attention::single_decode_nhd_into; pub use attention::single_prefill_hd256_into; pub use attention::single_prefill_hd512_into; diff --git a/pegainfer-kernels/src/ops/attention.rs b/pegainfer-kernels/src/ops/attention.rs index 2cbb51793..2c455a8e7 100644 --- a/pegainfer-kernels/src/ops/attention.rs +++ b/pegainfer-kernels/src/ops/attention.rs @@ -2282,15 +2282,16 @@ pub fn single_prefill_hd512_into( Ok(()) } -/// Plain-w QK RMSNorm + RoPE prep at head_dim 256 (Gemma 4 local layers): -/// Q and K each land in a contiguous token-major output shaped like its -/// input. No gate, no V, and the plain weight multiply — not the (1+w) -/// offset of the qwen35 hd256 sibling. `single_prefill` consumes a -/// contiguous **HND** cache, so these outputs are reassembled per head -/// before attention, not fed directly. `rotary_dim` must be positive, even -/// and <= 256 (launcher-enforced, surfaced here as `Err`); Gemma 4 local -/// layers pass 256, and a smaller even width keeps the hd512 sibling's -/// partial-RoPE tail semantics. +/// Plain-w QK RMSNorm + RoPE prep at head_dim 256 (Gemma 4 local layers), +/// oracle form: Q and K each land in a contiguous token-major output shaped +/// like its input; the serving form is +/// `qkv_norm_rope_paged_prefill_hd256_plain_into`. No gate, no V, and the +/// plain weight multiply — not the (1+w) offset of the qwen35 hd256 +/// sibling. `single_prefill` consumes a contiguous **HND** cache, so these +/// outputs are reassembled per head before attention, not fed directly. +/// `rotary_dim` must be positive, even and <= 256 (launcher-enforced, +/// surfaced here as `Err`); Gemma 4 local layers pass 256, and a smaller +/// even width keeps the hd512 sibling's partial-RoPE tail semantics. #[allow(clippy::too_many_arguments)] pub fn qk_norm_rope_prefill_hd256_plain_into( ctx: &DeviceContext, @@ -2449,6 +2450,313 @@ pub fn qk_norm_rope_prefill_hd256_plain_into( Ok(()) } +/// Everything a paged-pool prep launch derives from a caller-supplied +/// `PagedKvLayout`, re-derived with overflow checks and narrowed through +/// checked conversions: the layout's fields are public, so a contradictory +/// or wrapping layout must fail here, before anything reaches an unsafe +/// launch. +struct PagedPrepGeometry { + k_offset_elems: i64, + v_offset_elems: i64, + page_size: i32, + num_pages: i32, + stride_page: i64, +} + +fn checked_paged_prep_geometry( + what: &str, + layout: &PagedKvLayout, + pool_len: usize, + layer: usize, + head_dim: usize, + num_kv_heads: usize, +) -> Result { + anyhow::ensure!( + layout.head_dim == head_dim, + "{what} layout.head_dim {} != {head_dim}", + layout.head_dim + ); + anyhow::ensure!( + layout.num_kv_heads == num_kv_heads, + "{what} layout.num_kv_heads {} != num_kv_heads {num_kv_heads}", + layout.num_kv_heads + ); + anyhow::ensure!( + layout.page_size > 0 && layout.num_kv_heads > 0 && layout.num_layers > 0, + "{what} layout dims must be positive: page_size {} num_kv_heads {} num_layers {}", + layout.page_size, + layout.num_kv_heads, + layout.num_layers + ); + let kv_block_len = layout + .page_size + .checked_mul(layout.num_kv_heads) + .and_then(|x| x.checked_mul(layout.head_dim)) + .ok_or_else(|| { + anyhow::anyhow!( + "{what} page_size {} * num_kv_heads {} * head_dim {} overflows", + layout.page_size, + layout.num_kv_heads, + layout.head_dim + ) + })?; + anyhow::ensure!( + layout.kv_block_len == kv_block_len, + "{what} layout.kv_block_len {} != page_size * num_kv_heads * head_dim {kv_block_len}", + layout.kv_block_len + ); + let layer_stride = kv_block_len + .checked_mul(2) + .ok_or_else(|| anyhow::anyhow!("{what} 2 * kv_block_len {kv_block_len} overflows"))?; + anyhow::ensure!( + layout.layer_stride == layer_stride, + "{what} layout.layer_stride {} != 2 * kv_block_len {kv_block_len}", + layout.layer_stride + ); + let page_stride = layout.num_layers.checked_mul(layer_stride).ok_or_else(|| { + anyhow::anyhow!( + "{what} num_layers {} * layer_stride {layer_stride} overflows", + layout.num_layers + ) + })?; + anyhow::ensure!( + layout.page_stride == page_stride, + "{what} layout.page_stride {} != num_layers {} * layer_stride {layer_stride}", + layout.page_stride, + layout.num_layers + ); + anyhow::ensure!( + layer < layout.num_layers, + "{what} layer {layer} >= layout.num_layers {}", + layout.num_layers + ); + anyhow::ensure!( + pool_len.is_multiple_of(page_stride), + "{what} kv_pool.len {pool_len} is not a multiple of layout.page_stride {page_stride}" + ); + let num_pages = pool_len / page_stride; + anyhow::ensure!( + num_pages >= 1, + "{what} kv_pool.len {pool_len} holds no whole page (page_stride {page_stride})" + ); + let k_offset = layer.checked_mul(layer_stride).ok_or_else(|| { + anyhow::anyhow!("{what} layer {layer} * layer_stride {layer_stride} overflows") + })?; + let v_offset = k_offset.checked_add(kv_block_len).ok_or_else(|| { + anyhow::anyhow!("{what} K offset {k_offset} + kv_block_len {kv_block_len} overflows") + })?; + Ok(PagedPrepGeometry { + k_offset_elems: i64::try_from(k_offset) + .map_err(|_| anyhow::anyhow!("{what} K offset {k_offset} does not fit i64"))?, + v_offset_elems: i64::try_from(v_offset) + .map_err(|_| anyhow::anyhow!("{what} V offset {v_offset} does not fit i64"))?, + page_size: crate::ops::checked_i32(layout.page_size, "paged prep page_size")?, + num_pages: crate::ops::checked_i32(num_pages, "paged prep num_pages")?, + stride_page: i64::try_from(page_stride) + .map_err(|_| anyhow::anyhow!("{what} page_stride {page_stride} does not fit i64"))?, + }) +} + +/// Plain-w QKV prep at head_dim 256 (Gemma 4 local layers), serving form. +/// Q is normalised + rotated into a contiguous `q_out`; K is normalised + +/// rotated straight into the paged KV pool at layer `layer`'s K block, and +/// V — a separate v_proj head vector, unlike the hd512 K=V fork — is +/// weightless-normalised (never rotated) into the layer's V block, all in +/// one kernel with no intermediate scatter. +/// +/// The kernel `__trap()`s on any out-of-range pos or page id as the second +/// layer of the host validation's defence. Slot positions derive from +/// `start_pos + token`: absolute and cache-relative coordinates coincide +/// below the sliding window, which is this entry's contract — +/// window-crossing callers need a separate slot mapping first. +#[allow(clippy::too_many_arguments)] +pub fn qkv_norm_rope_paged_prefill_hd256_plain_into( + ctx: &DeviceContext, + q: &HiddenStates, + k: &HiddenStates, + v: &HiddenStates, + q_out: &mut HiddenStates, + kv_pool: &CudaSlice, + layout: &PagedKvLayout, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + layer: usize, + page_indices: &CudaSlice, + start_pos: usize, + cos_max_pos: usize, + num_q_heads: usize, + num_kv_heads: usize, + rotary_dim: usize, + rms_eps: f32, +) -> Result<()> { + let seq_len = q.seq_len; + let q_dim = num_q_heads.checked_mul(256).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep num_q_heads {num_q_heads} * 256 overflows") + })?; + let kv_dim = num_kv_heads.checked_mul(256).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep num_kv_heads {num_kv_heads} * 256 overflows") + })?; + anyhow::ensure!( + q.hidden_dim == q_dim, + "hd256 paged prep q.hidden_dim {} != num_q_heads {num_q_heads} * 256", + q.hidden_dim + ); + anyhow::ensure!( + q_out.hidden_dim == q.hidden_dim, + "hd256 paged prep q_out.hidden_dim {} != q.hidden_dim {}", + q_out.hidden_dim, + q.hidden_dim + ); + anyhow::ensure!( + q_out.seq_len == seq_len, + "hd256 paged prep q_out.seq_len {} != q.seq_len {seq_len}", + q_out.seq_len + ); + anyhow::ensure!( + k.hidden_dim == kv_dim, + "hd256 paged prep k.hidden_dim {} != num_kv_heads {num_kv_heads} * 256", + k.hidden_dim + ); + anyhow::ensure!( + v.hidden_dim == k.hidden_dim, + "hd256 paged prep v.hidden_dim {} != k.hidden_dim {}", + v.hidden_dim, + k.hidden_dim + ); + anyhow::ensure!( + k.seq_len == seq_len, + "hd256 paged prep k.seq_len {} != q.seq_len {seq_len}", + k.seq_len + ); + anyhow::ensure!( + v.seq_len == seq_len, + "hd256 paged prep v.seq_len {} != q.seq_len {seq_len}", + v.seq_len + ); + let geometry = checked_paged_prep_geometry( + "hd256 paged prep", + layout, + kv_pool.len(), + layer, + 256, + num_kv_heads, + )?; + ensure_vec_backed(q_norm_weight, "hd256 paged prep q_norm_weight")?; + ensure_vec_backed(k_norm_weight, "hd256 paged prep k_norm_weight")?; + ensure_vec_backed(cos_cache, "hd256 paged prep cos_cache")?; + ensure_vec_backed(sin_cache, "hd256 paged prep sin_cache")?; + anyhow::ensure!( + q_norm_weight.len == 256, + "hd256 paged prep q_norm_weight len {} != 256", + q_norm_weight.len + ); + anyhow::ensure!( + k_norm_weight.len == 256, + "hd256 paged prep k_norm_weight len {} != 256", + k_norm_weight.len + ); + let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep start_pos {start_pos} + seq_len {seq_len} overflows") + })?; + let covered = page_indices + .len() + .checked_mul(layout.page_size) + .ok_or_else(|| { + anyhow::anyhow!( + "hd256 paged prep page_indices len {} * page_size {} overflows", + page_indices.len(), + layout.page_size + ) + })?; + anyhow::ensure!( + covered >= end_pos, + "hd256 paged prep pages cover {covered} tokens, need start_pos {start_pos} + seq_len {seq_len}" + ); + anyhow::ensure!( + end_pos <= cos_max_pos, + "hd256 paged prep start_pos {start_pos} + seq_len {seq_len} > cos_max_pos {cos_max_pos}" + ); + let table_len = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { + anyhow::anyhow!( + "hd256 paged prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" + ) + })?; + anyhow::ensure!( + cos_cache.len >= table_len, + "hd256 paged prep cos_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", + cos_cache.len + ); + anyhow::ensure!( + sin_cache.len >= table_len, + "hd256 paged prep sin_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", + sin_cache.len + ); + let q_elems = q.checked_extent("hd256 paged prep q")?; + q_out.checked_extent("hd256 paged prep q_out")?; + let k_elems = k.checked_extent("hd256 paged prep k")?; + v.checked_extent("hd256 paged prep v")?; + crate::ops::checked_i32(q_elems, "hd256 paged prep q extent")?; + crate::ops::checked_i32(k_elems, "hd256 paged prep k extent")?; + crate::ops::checked_i32(table_len, "hd256 paged prep rope table extent")?; + + let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd256 paged prep num_q_heads")?; + let num_kv_heads_i32 = crate::ops::checked_i32(num_kv_heads, "hd256 paged prep num_kv_heads")?; + let seq_len_i32 = crate::ops::checked_i32(seq_len, "hd256 paged prep seq_len")?; + let start_pos_i32 = crate::ops::checked_i32(start_pos, "hd256 paged prep start_pos")?; + let cos_max_pos_i32 = crate::ops::checked_i32(cos_max_pos, "hd256 paged prep cos_max_pos")?; + let rotary_dim_i32 = crate::ops::checked_i32(rotary_dim, "hd256 paged prep rotary_dim")?; + + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let (k_ptr, _gk) = k.data.device_ptr(&ctx.stream); + let (v_ptr, _gv) = v.data.device_ptr(&ctx.stream); + let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); + // The KV state owns mutation; this call borrows its shared pool handle. + let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); + + let result = unsafe { + ffi::qkv_norm_rope_paged_prefill_hd256_plain_cuda( + q_ptr as *const ffi::Half, + k_ptr as *const ffi::Half, + v_ptr as *const ffi::Half, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + qo_ptr as *mut ffi::Half, + pool_ptr as *mut ffi::Half, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + num_q_heads_i32, + num_kv_heads_i32, + seq_len_i32, + start_pos_i32, + cos_max_pos_i32, + rotary_dim_i32, + rms_eps, + geometry.page_size, + geometry.num_pages, + geometry.stride_page, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "qkv_norm_rope_paged_prefill_hd256_plain_cuda failed with error \ + {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + Ok(()) +} + /// Causal single-sequence prefill at head_dim 256 over a contiguous **HND** /// K/V cache — `k[head, pos, dim]` with `k_cache.seq_len` allocated rows per /// head — matching `single_prefill_cuda_hd256`'s stride contract. Q and the @@ -2582,18 +2890,12 @@ pub fn single_prefill_hd256_into( /// global layers). Q is normalised + partially rotated into a contiguous /// `q_out`; K is normalised + partially rotated straight into the paged KV /// pool at layer `layer`'s K block (feeds `batch_prefill_paged`, not -/// `single_prefill`). No gate, no V; plain-w RMSNorm (not the (1+w) offset -/// of hd256). +/// `single_prefill`); V — the K=V fork — is the weightless RMS norm of the +/// same raw K, written to the layer's V block in the same pass. No gate; +/// plain-w RMSNorm. /// -/// Validated here: layout geometry consistency (the derived fields are pub -/// and hand-constructible), `layer < layout.num_layers`, `kv_pool` holding -/// whole pages (its capacity in pages becomes the kernel's page bound), -/// `page_indices` covering `[start_pos, start_pos + seq_len)`, -/// `start_pos + seq_len <= cos_max_pos` (host-side layer of the position -/// defence; the kernel `__trap()`s on any out-of-range pos or page id as -/// the second), table lengths, and 512-element norm weights. `rotary_dim` -/// must be positive, even and <= 512 — enforced by the launcher, which -/// returns -1; propagated here as `Err`. +/// The kernel `__trap()`s on any out-of-range pos or page id as the +/// second layer of the host validation's defence. #[allow(clippy::too_many_arguments)] pub fn qk_norm_partial_rope_paged_prefill_hd512_into( ctx: &DeviceContext, @@ -2616,11 +2918,16 @@ pub fn qk_norm_partial_rope_paged_prefill_hd512_into( rms_eps: f32, ) -> Result<()> { let seq_len = q.seq_len; + let q_dim = num_q_heads.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep num_q_heads {num_q_heads} * 512 overflows") + })?; + let kv_dim = num_kv_heads.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep num_kv_heads {num_kv_heads} * 512 overflows") + })?; anyhow::ensure!( - q.hidden_dim == num_q_heads * 512, - "hd512 prefill prep q.hidden_dim {} != num_q_heads {} * 512", - q.hidden_dim, - num_q_heads + q.hidden_dim == q_dim, + "hd512 prefill prep q.hidden_dim {} != num_q_heads {num_q_heads} * 512", + q.hidden_dim ); anyhow::ensure!( q_out.hidden_dim == q.hidden_dim, @@ -2634,73 +2941,23 @@ pub fn qk_norm_partial_rope_paged_prefill_hd512_into( q_out.seq_len ); anyhow::ensure!( - k.hidden_dim == num_kv_heads * 512, - "hd512 prefill prep k.hidden_dim {} != num_kv_heads {} * 512", - k.hidden_dim, - num_kv_heads + k.hidden_dim == kv_dim, + "hd512 prefill prep k.hidden_dim {} != num_kv_heads {num_kv_heads} * 512", + k.hidden_dim ); anyhow::ensure!( k.seq_len == seq_len, "hd512 prefill prep k.seq_len {} != q.seq_len {seq_len}", k.seq_len ); - anyhow::ensure!( - layout.head_dim == 512, - "hd512 prefill prep layout.head_dim {} != 512", - layout.head_dim - ); - anyhow::ensure!( - layout.num_kv_heads == num_kv_heads, - "hd512 prefill prep layout.num_kv_heads {} != num_kv_heads {num_kv_heads}", - layout.num_kv_heads - ); - // Keep the pool-capacity division below defined for public layouts. - anyhow::ensure!( - layout.page_size > 0 && layout.num_kv_heads > 0, - "hd512 prefill prep layout.page_size {} / num_kv_heads {} must be positive", - layout.page_size, - layout.num_kv_heads - ); - let kv_block_len = layout.page_size * layout.num_kv_heads * layout.head_dim; - anyhow::ensure!( - layout.kv_block_len == kv_block_len, - "hd512 prefill prep layout.kv_block_len {} != page_size {} * num_kv_heads \ - {} * head_dim {}", - layout.kv_block_len, - layout.page_size, - layout.num_kv_heads, - layout.head_dim - ); - anyhow::ensure!( - layout.layer_stride == 2 * kv_block_len, - "hd512 prefill prep layout.layer_stride {} != 2 * kv_block_len {kv_block_len}", - layout.layer_stride - ); - anyhow::ensure!( - layout.page_stride == layout.num_layers * layout.layer_stride, - "hd512 prefill prep layout.page_stride {} != num_layers {} * layer_stride {}", - layout.page_stride, - layout.num_layers, - layout.layer_stride - ); - anyhow::ensure!( - layer < layout.num_layers, - "hd512 prefill prep layer {layer} >= layout.num_layers {}", - layout.num_layers - ); - let num_pages = kv_pool.len() / layout.page_stride; - anyhow::ensure!( - kv_pool.len().is_multiple_of(layout.page_stride), - "hd512 prefill prep kv_pool.len {} is not a multiple of layout.page_stride {}", - kv_pool.len(), - layout.page_stride - ); - anyhow::ensure!( - num_pages >= 1, - "hd512 prefill prep kv_pool.len {} holds no whole page (page_stride {})", + let geometry = checked_paged_prep_geometry( + "hd512 prefill prep", + layout, kv_pool.len(), - layout.page_stride - ); + layer, + 512, + num_kv_heads, + )?; anyhow::ensure!( q_norm_weight.len == 512, "hd512 prefill prep q_norm_weight len {} != 512", @@ -2711,43 +2968,60 @@ pub fn qk_norm_partial_rope_paged_prefill_hd512_into( "hd512 prefill prep k_norm_weight len {} != 512", k_norm_weight.len ); + let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} overflows") + })?; + let covered = page_indices + .len() + .checked_mul(layout.page_size) + .ok_or_else(|| { + anyhow::anyhow!( + "hd512 prefill prep page_indices len {} * page_size {} overflows", + page_indices.len(), + layout.page_size + ) + })?; anyhow::ensure!( - page_indices.len() * layout.page_size >= start_pos + seq_len, - "hd512 prefill prep pages cover {} tokens (len {} * page_size {}), need \ - start_pos {start_pos} + seq_len {seq_len}", - page_indices.len() * layout.page_size, - page_indices.len(), - layout.page_size + covered >= end_pos, + "hd512 prefill prep pages cover {covered} tokens, need start_pos {start_pos} + seq_len {seq_len}" ); anyhow::ensure!( - start_pos + seq_len <= cos_max_pos, - "hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} > \ - cos_max_pos {cos_max_pos}" + end_pos <= cos_max_pos, + "hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} > cos_max_pos {cos_max_pos}" ); + let table_len = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { + anyhow::anyhow!( + "hd512 prefill prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" + ) + })?; anyhow::ensure!( - cos_cache.len >= cos_max_pos * rotary_dim, - "hd512 prefill prep cos_cache len {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", + cos_cache.len >= table_len, + "hd512 prefill prep cos_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", cos_cache.len ); anyhow::ensure!( - sin_cache.len >= cos_max_pos * rotary_dim, - "hd512 prefill prep sin_cache len {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", + sin_cache.len >= table_len, + "hd512 prefill prep sin_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", sin_cache.len ); - q.checked_extent("hd512 prefill prep q")?; + let q_elems = q.checked_extent("hd512 prefill prep q")?; q_out.checked_extent("hd512 prefill prep q_out")?; - k.checked_extent("hd512 prefill prep k")?; + let k_elems = k.checked_extent("hd512 prefill prep k")?; + crate::ops::checked_i32(q_elems, "hd512 prefill prep q extent")?; + crate::ops::checked_i32(k_elems, "hd512 prefill prep k extent")?; + crate::ops::checked_i32(table_len, "hd512 prefill prep rope table extent")?; ensure_vec_backed(q_norm_weight, "hd512 prefill prep q_norm_weight")?; ensure_vec_backed(k_norm_weight, "hd512 prefill prep k_norm_weight")?; ensure_vec_backed(cos_cache, "hd512 prefill prep cos_cache")?; ensure_vec_backed(sin_cache, "hd512 prefill prep sin_cache")?; - // Derive raw offsets from the validated layout. - let k_offset_elems = (layer * layout.layer_stride) as i64; - let page_size = layout.page_size; - let stride_page = layout.page_stride as i64; + let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd512 prefill prep num_q_heads")?; + let num_kv_heads_i32 = + crate::ops::checked_i32(num_kv_heads, "hd512 prefill prep num_kv_heads")?; + let seq_len_i32 = crate::ops::checked_i32(seq_len, "hd512 prefill prep seq_len")?; + let start_pos_i32 = crate::ops::checked_i32(start_pos, "hd512 prefill prep start_pos")?; + let cos_max_pos_i32 = crate::ops::checked_i32(cos_max_pos, "hd512 prefill prep cos_max_pos")?; + let rotary_dim_i32 = crate::ops::checked_i32(rotary_dim, "hd512 prefill prep rotary_dim")?; let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); let (k_ptr, _gk) = k.data.device_ptr(&ctx.stream); @@ -2770,18 +3044,19 @@ pub fn qk_norm_partial_rope_paged_prefill_hd512_into( sin_ptr as *const ffi::Half, qo_ptr as *mut ffi::Half, pool_ptr as *mut ffi::Half, - k_offset_elems, + geometry.k_offset_elems, + geometry.v_offset_elems, pi_ptr as *const i32, - num_q_heads as i32, - num_kv_heads as i32, - seq_len as i32, - start_pos as i32, - cos_max_pos as i32, - rotary_dim as i32, + num_q_heads_i32, + num_kv_heads_i32, + seq_len_i32, + start_pos_i32, + cos_max_pos_i32, + rotary_dim_i32, rms_eps, - page_size as i32, - num_pages as i32, - stride_page, + geometry.page_size, + geometry.num_pages, + geometry.stride_page, crate::tensor::active_cu_stream(ctx), ) }; diff --git a/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs b/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs index d47ee47e6..13c6c705e 100644 --- a/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs +++ b/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs @@ -2,15 +2,15 @@ //! //! Manual gate: CI compiles this but never runs it. Run on a GPU box with //! PEGAINFER_REQUIRE_GPU=1, which turns a missing device into a failure -//! rather than a skip. Unlike the hd512 sibling there are no trap binaries: -//! this entry has no device-resident indices — the only out-of-range axis -//! (the position window) is host-known and rejected by the wrapper, so the -//! kernel's position trap is unreachable through the public path. +//! rather than a skip. mod common; +use cudarc::driver::CudaSlice; use half::bf16; use pegainfer_kernels::ops::qk_norm_rope_prefill_hd256_plain_into; +use pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into; +use pegainfer_kernels::paged_kv::PagedKvLayout; use pegainfer_kernels::tensor::DeviceContext; use pegainfer_kernels::tensor::DeviceVec; use pegainfer_kernels::tensor::HiddenStates; @@ -29,6 +29,10 @@ const NUM_KV_HEADS: usize = 8; // 8 entries; the pattern is aperiodic under h -> 2h, h % 8 and h ± 1. const Q_BASE: f32 = 1.0; const K_BASE: f32 = 3.0; +// The weightless V norm erases magnitude (x * inv_rms(x) is ±1 up to eps), +// so only signs distinguish V from a mis-read source; the global flip +// against Q/K makes a V-input mix-up land a wrong sign in every slot. +const V_BASE: f32 = -5.0; const HEAD_SIGNS: [f32; NUM_Q_HEADS] = [ 1.0, 1.0, -1.0, 1.0, -1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, 1.0, -1.0, 1.0, -1.0, -1.0, ]; @@ -188,6 +192,119 @@ fn run_prep(ctx: &DeviceContext, rotary_dim: usize) -> (Vec, Vec) { (qo, ko) } +// Pool geometry for the serving-form case. Positions 1..=4 map through +// PAGE_INDICES to pages 3, 7, 5; page id 9 is an out-of-range sentinel the +// kernel must never dereference. The 8-page pool leaves in-range pages +// unreferenced, so stray writes have somewhere visible to land. +const PAGE_SIZE: usize = 2; +const NUM_LAYERS: usize = 2; +const PAGE_INDICES: [i32; 4] = [3, 7, 5, 9]; +const POOL_PAGES: usize = 8; + +/// K and V blocks at their layout-derived offsets; everything else stays +/// 0.0. V is the weightless norm of the v input — never rotated, no weight. +fn expected_pool(layout: &PagedKvLayout, layer: usize, kw: &[bf16]) -> Vec { + let mut exp = vec![0.0f32; layout.page_stride * POOL_PAGES]; + let layer_offset = (layer * layout.layer_stride) as i64; + for t in 0..SEQ_LEN { + let pos = START_POS + t; + let page = PAGE_INDICES[pos / PAGE_SIZE] as i64; + for h in 0..NUM_KV_HEADS { + let k_x = signed(K_BASE, h, t); + let k_inv = inv_rms(k_x); + let v_x = signed(V_BASE, h, t); + let v_val = bf16::from_f32(v_x * inv_rms(v_x)).to_f32(); + let base = page * layout.page_stride as i64 + + layer_offset + + (pos % PAGE_SIZE) as i64 * KV_DIM as i64 + + h as i64 * HD as i64; + for d in 0..HD { + exp[(base + d as i64) as usize] = expected_prep(k_x, kw, k_inv, d, pos, HD); + exp[(base + layout.kv_block_len as i64 + d as i64) as usize] = v_val; + } + } + } + exp +} + +/// Exact zero is the assertion, not sloppiness: it marks a slot the kernel +/// must never have written — unreferenced pages, the other layer, and the +/// slots outside the request's positions. +#[allow(clippy::float_cmp)] +fn assert_pool(got: &[f32], expected: &[f32]) { + assert_eq!(got.len(), expected.len()); + for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { + if e == 0.0 { + assert_eq!(g, 0.0, "pool[{i}]: expected untouched, got {g}"); + } else { + assert!( + (g - e).abs() < 0.02, + "pool[{i}]: got {g}, expected {e} (tolerance 0.02)" + ); + } + } +} + +#[test] +fn pool_write_matches_closed_form_and_touches_nothing_else() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let ctx = &ctx; + let qw = q_norm_weights(); + let kw = k_norm_weights(); + let layer = 1; + let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); + let q = hidden_input(ctx, Q_BASE, NUM_Q_HEADS); + let k = hidden_input(ctx, K_BASE, NUM_KV_HEADS); + let v = hidden_input(ctx, V_BASE, NUM_KV_HEADS); + let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); + let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, HD); + let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); + let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride * POOL_PAGES) + .expect("pool alloc"); + let page_indices: CudaSlice = ctx + .stream + .clone_htod(&PAGE_INDICES) + .expect("page_indices H2D"); + + qkv_norm_rope_paged_prefill_hd256_plain_into( + ctx, + &q, + &k, + &v, + &mut q_out, + &pool, + &layout, + &qn, + &kn, + &cos_dev, + &sin_dev, + layer, + &page_indices, + START_POS, + COS_MAX_POS, + NUM_Q_HEADS, + NUM_KV_HEADS, + HD, + EPS, + ) + .expect("pool prep launch"); + + let qo = q_out.to_host(ctx).expect("q_out D2H"); + assert_close( + &qo, + &expected_full(Q_BASE, &qw, Q_DIM, HD), + "pool-write Q pairing", + ); + let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); + let pool_f: Vec = pool_host.iter().map(|x| x.to_f32()).collect(); + assert_pool(&pool_f, &expected_pool(&layout, layer, &kw)); +} + /// rotary_dim = 256 is the Gemma 4 local-layer case: the full head rotates /// and the pass-through tail is empty. #[test] diff --git a/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs b/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs index 9d55bde7a..e81152cc4 100644 --- a/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs +++ b/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs @@ -116,11 +116,14 @@ fn expected_full( full } -/// K side only; everything else stays 0.0. The layer offset is derived -/// from the layout, so oracle and kernel share no hand-picked raw offset. +/// K and V blocks; everything else stays 0.0. The offsets are derived from +/// the layout, so oracle and kernel share no hand-picked raw offset. V is +/// the K=V fork: the same raw vector under the shared inv_rms, weightless +/// and un-rotated. fn expected_pool(x: f32, w: &[bf16], inv: f32, layout: &PagedKvLayout, layer: usize) -> Vec { let mut exp = vec![0.0f32; POOL_LEN]; let layer_offset = (layer * layout.layer_stride) as i64; + let v_val = bf16::from_f32(x * inv).to_f32(); for t in 0..SEQ_LEN { let pos = START_POS + t; let page = PAGE_INDICES[pos / PAGE_SIZE] as i64; @@ -131,6 +134,7 @@ fn expected_pool(x: f32, w: &[bf16], inv: f32, layout: &PagedKvLayout, layer: us + h as i64 * HD as i64; for d in 0..HD { exp[(base + d as i64) as usize] = expected_prep(x, w, inv, d, pos); + exp[(base + layout.kv_block_len as i64 + d as i64) as usize] = v_val; } } } @@ -148,8 +152,8 @@ fn assert_close(got: &[f32], expected: &[f32], what: &str) { } /// Exact zero is the assertion, not sloppiness: it marks a slot the kernel -/// must never have written. That covers the V blocks too, so they get no -/// separate check. +/// must never have written — unreferenced pages, the other layer, and the +/// slots outside the request's positions. #[allow(clippy::float_cmp)] fn assert_pool(got: &[f32], expected: &[f32]) { assert_eq!(got.len(), expected.len());