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
Fused K-path: rms_norm → RoPE → BF16 paged cache write in one kernel.
Replaces the chained ops::rms_norm + ops::rope + ops::reshape_and_cache
sequence for the K projection. Keeps K in FP32 between the three
operations and BF16-rounds ONLY at cache write — vLLM-equivalent
precision regime. Eliminates the two intermediate BF16 rounding steps
that previously compounded at deep attention layers (L35-L39) where K
magnitudes peak ~18× vs L0, causing the documented BF16-KV cliff.