fused_k_norm_rope_cache_write_fp8_kv

Function fused_k_norm_rope_cache_write_fp8_kv 

Source
pub fn fused_k_norm_rope_cache_write_fp8_kv(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    k_in: DevicePtr,
    value: DevicePtr,
    k_norm_weight: DevicePtr,
    positions: DevicePtr,
    k_cache: DevicePtr,
    v_cache: DevicePtr,
    slot_mapping: DevicePtr,
    num_tokens: u32,
    num_kv_heads: u32,
    head_dim: u32,
    rotary_dim: u32,
    block_size: u32,
    k_scale: f32,
    v_scale: f32,
    key_stride: u32,
    value_stride: u32,
    cache_stride: u64,
    rms_eps: f32,
    theta: f32,
    stream: u64,
) -> Result<()>
Expand description

Decode-path fusion of rms_norm + rope_forward + reshape_and_cache_flash_fp8 for an FP8 KV cache. Writes BOTH K (normed + rotated + quantized) and V (quantized) in one launch, so the FP8 decode chain drops from rms_norm(K) + rope(Q,K) + reshape_and_cache_fp8(K,V) to rope(Q only) + this kernel.

BIT-IDENTICAL to the chain it replaces — the kernel reproduces every intermediate BF16 rounding step rather than keeping intermediates in FP32. See the header of reshape_and_cache_fused_k_fp8.cu; do not “improve” the precision here, it would change committed cache bytes.

k_scale/v_scale are the DEQUANT scales; the kernel takes the reciprocal itself so the lowering matches reshape_and_cache_flash_fp8 exactly.

Grid: (num_tokens, num_kv_heads, 1) Block: (head_dim, 1, 1)