pub fn fused_k_norm_rope_cache_write_bf16(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
k_in: DevicePtr,
k_norm_weight: DevicePtr,
positions: DevicePtr,
k_cache: DevicePtr,
slot_mapping: DevicePtr,
num_tokens: u32,
num_kv_heads: u32,
head_dim: u32,
rotary_dim: u32,
block_size: u32,
rms_eps: f32,
theta: f32,
stream: u64,
) -> Result<()>Expand description
2026-09-25: K norm (x * rsqrt(mean(x^2) + eps) * (1 + w)), RoPE and the paged BF16 cache
write in one kernel. K stays in FP32 from the BF16 load to the single BF16 rounding at the
store. One block per (token, KV head) with head_dim threads; the kernel supports
head_dim <= 256.