w8a16_gemv_batch4_strided

Function w8a16_gemv_batch4_strided 

Source
pub fn w8a16_gemv_batch4_strided(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    weight: DevicePtr,
    block_scale: DevicePtr,
    output: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    a_row_stride: u32,
    c_row_stride: u32,
    stream: u64,
) -> Result<()>
Expand description

Strided sibling of w8a16_gemv_batch4 (M<=4).

WHY: the multi-sequence decode Q/K/V buffer is [n, per_seq_qkv] with Q at offset 0, K after Q and V after K inside every row, so the contiguous [M, N] writer cannot address one projection across rows. Without a strided writer the native-FP8 attention projections fell back to three scalar w8a16_gemv launches PER ROW at decode concurrency 2..=8 — the third-largest bucket in the C=4 decode profile (issue #927). This writes one projection for all M rows in ONE launch.

LAYOUT: input [M, a_row_stride] BF16, only the first k elements of each row read; weight/block_scale are the raw w8a16_gemv pointers ([N, K] FP8 E4M3 and [N/128, K/128] FP32); output [M, c_row_stride] BF16, only the first n elements of each row written. Both strides are in ELEMENTS. a_row_stride must keep each activation row 16-byte aligned (multiple of 8) — the kernel’s activation loads are uint4.

Bit-identical per row to w8a16_gemv: same template body, same K-iteration order and same reduction tree as w8a16_gemv_batch4; only the row pitches change. Verified by examples/native_fp8_qkv_batch_microtest.

Kernel: w8a16_gemv_batch4_strided (module w8a16_gemv_batch4). Grid: (ceil(N/4), 1, 1) Block: (256, 1, 1)