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)