feat(kernels): explicit hd512 sm_scale and the final-logit softcap

pegainfer-kernels/csrc/shared/elementwise.cu

@@ -302,6 +302,25 @@ __global__ void gelu_tanh_mul_kernel(
   }
 }
 
+// ============================================================================
+// In-place final-logit softcap: buf[i] = bf16(cap * tanh(f32(buf[i]) / cap)).
+// Gemma 4 declares final_logit_softcapping (30.0 at every published size) and
+// applies it to the LM head output in the compute dtype; f32 internal with a
+// single rounding matches the reference's bf16 elementwise chain.
+// ============================================================================
+
+__global__ void softcap_bf16_kernel(
+    __nv_bfloat16 *__restrict__ buf,
+    float cap,
+    int n) {
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
+       idx < n;
+       idx += gridDim.x * blockDim.x) {
+    float x = __bfloat162float(buf[idx]);
+    buf[idx] = __float2bfloat16(cap * tanhf(x / cap));
+  }
+}
+
 // ============================================================================
 // In-place multiply by a host scalar: buf[i] = bf16(f32(buf[i]) * scale).
 // Serves Gemma 4's per-layer layer_scalar, a [1] weight the model reads to
@@ -714,6 +733,14 @@ CUresult scale_bf16_in_place_cuda(
   return (CUresult)cudaGetLastError();
 }
 
+CUresult softcap_bf16_in_place_cuda(
+    __nv_bfloat16 *buf, float cap, int n, cudaStream_t stream) {
+  int block = 256;
+  int grid = n / block + (n % block != 0);
+  softcap_bf16_kernel<<<grid, block, 0, stream>>>(buf, cap, n);
+  return (CUresult)cudaGetLastError();
+}
+
 CUresult embedding_decode_cuda(
     const __nv_bfloat16 *embed, const uint32_t *token_id,
     __nv_bfloat16 *out, int hidden_size, cudaStream_t stream) {

pegainfer-kernels/src/ffi/shared.rs

@@ -139,6 +139,13 @@ unsafe extern "C" {
         stream: CUstream,
     ) -> CUresult;
 
+    pub fn softcap_bf16_in_place_cuda(
+        buf: *mut Half,
+        cap: f32,
+        n: i32,
+        stream: CUstream,
+    ) -> CUresult;
+
     pub fn embedding_batched_cuda(
         embed: *const Half,
         token_ids: *const u32,

pegainfer-kernels/src/ops.rs

@@ -104,6 +104,7 @@ pub use elementwise::scaled_add_rows_token_range_into;
 pub use elementwise::silu_mul_batch;
 pub use elementwise::silu_mul_batch_into;
 pub use elementwise::silu_mul_fused_batch_into;
+pub use elementwise::softcap_bf16_in_place;
 pub use elementwise::write_vec_into;
 pub use embedding::embedding_batch;
 pub use embedding::embedding_batch_vocab_shard;

pegainfer-kernels/src/ops/attention.rs

@@ -1729,7 +1729,12 @@ pub fn paged_attention_batch_decode_hd512_into(
     meta: &Hd512DecodeMetadata,
     output: &mut HiddenStates,
     num_qo_heads: usize,
+    sm_scale: f32,
 ) -> Result<()> {
+    anyhow::ensure!(
+        sm_scale.is_finite(),
+        "paged_attention_batch_decode_hd512 sm_scale {sm_scale} must be finite"
+    );
     let num_kv_heads = layout.num_kv_heads;
     let head_dim = layout.head_dim;
     anyhow::ensure!(
@@ -1821,7 +1826,6 @@ pub fn paged_attention_batch_decode_hd512_into(
         "batch hd512 decode",
     )?;
 
-    let sm_scale = 1.0f32 / (head_dim as f32).sqrt();
     let result = unsafe {
         ffi::paged_attention_decode_cuda_hd512(
             q_ptr as *const ffi::Half,
@@ -1868,7 +1872,12 @@ pub fn paged_attention_batch_decode_via_prefill_hd512_into(
     positions_d: &CudaSlice<i32>,
     output: &mut HiddenStates,
     num_qo_heads: usize,
+    sm_scale: f32,
 ) -> Result<()> {
+    anyhow::ensure!(
+        sm_scale.is_finite(),
+        "paged_attention_batch_decode_via_prefill_hd512 sm_scale {sm_scale} must be finite"
+    );
     let num_kv_heads = layout.num_kv_heads;
     let head_dim = layout.head_dim;
     anyhow::ensure!(
@@ -2002,7 +2011,6 @@ pub fn paged_attention_batch_decode_via_prefill_hd512_into(
     let k_offset = (layer * layout.layer_stride) as i64;
     let v_offset = (layer * layout.layer_stride + layout.kv_block_len) as i64;
     let stride_page = layout.page_stride as i64;
-    let sm_scale = 1.0f32 / (head_dim as f32).sqrt();
 
     let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream);
     let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream);
@@ -2071,7 +2079,12 @@ pub fn batch_prefill_paged_hd512_into(
     plan: &PrefillPagedPlan,
     output: &mut HiddenStates,
     num_qo_heads: usize,
+    sm_scale: f32,
 ) -> Result<()> {
+    anyhow::ensure!(
+        sm_scale.is_finite(),
+        "batch_prefill_paged_hd512 sm_scale {sm_scale} must be finite"
+    );
     let num_kv_heads = layout.num_kv_heads;
     let head_dim = layout.head_dim;
     anyhow::ensure!(
@@ -2150,7 +2163,6 @@ pub fn batch_prefill_paged_hd512_into(
     let k_offset = (layer * layout.layer_stride) as i64;
     let v_offset = (layer * layout.layer_stride + layout.kv_block_len) as i64;
     let stride_page = layout.page_stride as i64;
-    let sm_scale = 1.0f32 / (head_dim as f32).sqrt();
 
     let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream);
     let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream);
@@ -2219,7 +2231,12 @@ pub fn single_prefill_hd512_into(
     num_q_heads: usize,
     num_kv_heads: usize,
     kv_len: usize,
+    sm_scale: f32,
 ) -> Result<()> {
+    anyhow::ensure!(
+        sm_scale.is_finite(),
+        "single_prefill_hd512 sm_scale {sm_scale} must be finite"
+    );
     assert_eq!(q.hidden_dim, num_q_heads * 512);
     assert_eq!(output.hidden_dim, q.hidden_dim);
     assert_eq!(output.seq_len, q.seq_len);
@@ -2252,7 +2269,7 @@ pub fn single_prefill_hd512_into(
             q_seq_len as i32,
             kv_len as i32,
             k_cache.seq_len as i32,
-            1.0f32 / (512.0f32).sqrt(),
+            sm_scale,
             crate::tensor::active_cu_stream(ctx),
         )
     };

pegainfer-kernels/src/ops/elementwise.rs

@@ -911,6 +911,35 @@ pub fn gelu_tanh_mul_batch_into(
     Ok(())
 }
 
+/// In-place final-logit softcap: `x = cap * tanh(x / cap)` in f32 with a
+/// single rounding back to bf16, matching the reference's compute-dtype
+/// application. Gemma 4 declares `final_logit_softcapping` 30.0 at every
+/// published size.
+pub fn softcap_bf16_in_place(ctx: &DeviceContext, buf: &mut HiddenStates, cap: f32) -> Result<()> {
+    if !cap.is_finite() || cap <= 0.0 {
+        return Err(anyhow!(
+            "softcap_bf16_in_place cap {cap} must be positive and finite"
+        ));
+    }
+    let n = buf.checked_extent("softcap_bf16 buf")?;
+    let n = super::checked_i32(n, "softcap_bf16 extent")?;
+    let (buf_ptr, _gb) = buf.data.device_ptr_mut(&ctx.stream);
+
+    let result = unsafe {
+        ffi::softcap_bf16_in_place_cuda(
+            buf_ptr as *mut ffi::Half,
+            cap,
+            n,
+            crate::tensor::active_cu_stream(ctx),
+        )
+    };
+    result
+        .result()
+        .map_err(|e| anyhow!("softcap_bf16_in_place_cuda failed: {e}"))?;
+
+    Ok(())
+}
+
 /// In-place multiply by a host scalar — Gemma 4's per-layer `layer_scalar`,
 /// a `[1]` weight the model reads to the host at load.
 pub fn scale_bf16_in_place(ctx: &DeviceContext, buf: &mut HiddenStates, scale: f32) -> Result<()> {

pegainfer-kernels/tests/softcap_smoke.rs

@@ -0,0 +1,62 @@
+//! Device gate for the final-logit softcap in csrc/shared/elementwise.cu.
+//!
+//! 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.
+
+mod common;
+
+use half::bf16;
+use pegainfer_kernels::ops::softcap_bf16_in_place;
+use pegainfer_kernels::tensor::HiddenStates;
+
+/// Mirrors the kernel arithmetic: f32 tanh, one rounding back to bf16.
+fn host_softcap(x: f32, cap: f32) -> f32 {
+    bf16::from_f32(cap * (x / cap).tanh()).to_f32()
+}
+
+#[test]
+fn softcap_matches_host_reference() {
+    let Some(ctx) = common::device_or_skip() else {
+        return;
+    };
+    let ctx = &ctx;
+    // Values at and past the cap are the sensitive cases: near zero the cap
+    // is a near-identity (tanh linear region), so a mutation that drops the
+    // division or the multiply only shows at |x| comparable to cap. 30.0 is
+    // the value every published Gemma 4 size declares.
+    let cap = 30.0f32;
+    let vals = [-90.0f32, -30.0, -4.0, -0.5, 0.0, 0.5, 4.0, 30.0, 90.0];
+    let host: Vec<bf16> = vals.iter().map(|&v| bf16::from_f32(v)).collect();
+    let mut buf = HiddenStates::from_host(ctx, &host, vals.len(), 1).expect("buf H2D");
+
+    softcap_bf16_in_place(ctx, &mut buf, cap).expect("softcap launch");
+
+    let got = buf.to_host(ctx).expect("buf D2H");
+    for (i, &v) in vals.iter().enumerate() {
+        let e = host_softcap(v, cap);
+        // Two bf16 ulp relative, floored near zero: host tanh and device
+        // tanhf may differ in the last f32 ulp.
+        let tol = (e.abs() * 0.008).max(0.02);
+        assert!(
+            (got[i] - e).abs() <= tol,
+            "softcap[{i}] (x {v}): got {}, expected {e} (tolerance {tol})",
+            got[i]
+        );
+    }
+}
+
+#[test]
+fn softcap_rejects_bad_cap() {
+    let Some(ctx) = common::device_or_skip() else {
+        return;
+    };
+    let ctx = &ctx;
+    let mut buf = HiddenStates::zeros(ctx, 4, 1).expect("buf alloc");
+    for bad in [0.0f32, -30.0, f32::NAN, f32::INFINITY] {
+        assert!(
+            softcap_bf16_in_place(ctx, &mut buf, bad).is_err(),
+            "cap {bad} must be rejected"
+        );
+    }
+}