From 01fda5a757c6a4f067193709224e52dc85855e46 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Tue, 4 Aug 2026 11:14:32 +0100 Subject: [PATCH] feat(kernels): explicit hd512 sm_scale and the final-logit softcap Signed-off-by: Feathbow --- pegainfer-kernels/csrc/shared/elementwise.cu | 27 +++++++++ pegainfer-kernels/src/ffi/shared.rs | 7 +++ pegainfer-kernels/src/ops.rs | 1 + pegainfer-kernels/src/ops/attention.rs | 25 ++++++-- pegainfer-kernels/src/ops/elementwise.rs | 29 +++++++++ pegainfer-kernels/tests/softcap_smoke.rs | 62 ++++++++++++++++++++ 6 files changed, 147 insertions(+), 4 deletions(-) create mode 100644 pegainfer-kernels/tests/softcap_smoke.rs diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index 48a3f303d..c18de2e4e 100644 --- a/pegainfer-kernels/csrc/shared/elementwise.cu +++ b/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<<>>(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) { diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index bc8d47500..d3b84b47d 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/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, diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index 2e7d796eb..225bb6d3c 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/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; diff --git a/pegainfer-kernels/src/ops/attention.rs b/pegainfer-kernels/src/ops/attention.rs index 35af89ef9..2cbb51793 100644 --- a/pegainfer-kernels/src/ops/attention.rs +++ b/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, 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), ) }; diff --git a/pegainfer-kernels/src/ops/elementwise.rs b/pegainfer-kernels/src/ops/elementwise.rs index 05f9fb53c..a15c92868 100644 --- a/pegainfer-kernels/src/ops/elementwise.rs +++ b/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<()> { diff --git a/pegainfer-kernels/tests/softcap_smoke.rs b/pegainfer-kernels/tests/softcap_smoke.rs new file mode 100644 index 000000000..6d31eaf9d --- /dev/null +++ b/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 = 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" + ); + } +}