feat(kernels): pool-write prep for Gemma 4 serving at both head dims

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;

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<int64_t>(page_id) * stride_page
+        + block_offset_elems
+        + static_cast<int64_t>(offset_in_page) * num_kv_heads * HD256_PLAIN
+        + static_cast<int64_t>(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<<<prep_grid, THREADS_HD256_PLAIN, 0, stream>>>(
+        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"

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,