fused_k_norm_rope_cache_write_bf16

Function fused_k_norm_rope_cache_write_bf16 

Source
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.