metrale_model_layers/layers/ops/

w8a16_gemm_pipelined_m32.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! 2026-09-25: `w8a16_gemm_pipelined_m32`, the 32-row M-tile twin of
//! [`super::w8a16_gemm_pipelined`], and the rule for when the twin runs.
//!
//! At 1..=32 rows the 128-row tile of `w8a16_gemm_pipelined` is mostly
//! padding. The twin is a 32x32 tile with 128-K steps and caller-supplied A/C
//! row pitches (kernel header: `kernels/gb10/common/w8a16_gemm_pipelined_m32.cu`),
//! so one kernel serves the contiguous GDN in_proj_qkvz/out_proj arms
//! ([`w8a16_gemm_pipelined_by_m`]), the attention o_proj and the strided
//! multi-seq Q/K/V tier.
//!
//! Numerics: it uses the same sub-MMA windows and fold points as
//! `w8a16_gemm_pipelined`, and the `native_fp8_gdn_proj_m32_microtest` example
//! checks its output byte for byte against it. Against the scalar GEMV family
//! it reassociates the K reduction, as the `*_m16` tiers do.
//!
//! The kernel has its own module (`w8a16_gemm_pipelined_m32`). Callers resolve
//! it with `try_target_kernel`, so a target without the module gets a zero
//! handle and its callers keep to their other kernels.
//!
//! Owner: model-layers ops.
//! Invariants: none beyond the types.

use anyhow::{Result, ensure};
use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};

/// 2026-09-25: Rows one M tile covers (`PM32_M_TILE`). Above it the kernel
/// still runs, with `grid.y = ceil(M / 32)` and one weight pass per tile;
/// [`w8a16_pipelined_prefers_m32`] sends contiguous callers to the 128-tile
/// kernel past this bound.
pub const W8A16_M32_TILE_ROWS: u32 = 32;

/// 2026-09-25: N tile (`PM32_N_TILE`), the grid.x granularity.
const W8A16_M32_TILE_COLS: u32 = 32;

/// 2026-09-27: Rows the skinny body of the same kernel takes (`PM16_MAX_M`),
/// when `K` is at most [`W8A16_M16_MAX_K`]. It gives each 8-column block the
/// canonical tile's per-output arithmetic with the 128-K steps' products
/// computed in parallel by 8 warps and added in step order.
pub const W8A16_M16_MAX_ROWS: u32 = 16;

/// 2026-09-27: Largest `K` of the skinny body (`PM16_MAX_K`): one 1152-byte
/// shared slot per 128-K step within 48 KiB.
pub const W8A16_M16_MAX_K: u32 = 5120;

/// 2026-09-27: Dynamic shared memory of the 32-row tile: `PM32_SMEM_BYTES`
/// (A and BF16 B stages, the raw B stages and the E4M3 table).
const W8A16_M32_SMEM_BYTES: u32 = 44_032;

/// 2026-09-27: Dynamic shared memory of a skinny block for `k`:
/// `PM16_SMEM_BYTES(K)`, the block's 8 weight rows `k + 16` bytes apart.
pub fn w8a16_m16_smem_bytes(k: u32) -> u32 {
    8 * (k + 16)
}

/// 2026-09-25: K granularity: one FP8 scale block per K-step (`PM32_K_STEP ==
/// PM32_FP8_BLOCK`), which is why `K % 128 == 0` is required rather than
/// padded.
pub const W8A16_M32_K_BLOCK: u32 = 128;

/// 2026-09-25: The selection rule between the 128-tile `w8a16_gemm_pipelined`
/// and its 32-tile twin at a contiguous call site: the twin takes `1..=32`
/// rows when its handle is nonzero and `K` is a positive multiple of 128.
/// Above 32 rows the twin would need `grid.y > 1`, streaming the weight once
/// per 32-row tile. [`w8a16_gemm_pipelined_by_m`] applies it.
pub fn w8a16_pipelined_prefers_m32(m: u32, k: u32, m32_kernel: KernelHandle) -> bool {
    (1..=W8A16_M32_TILE_ROWS).contains(&m)
        && k >= W8A16_M32_K_BLOCK
        && k.is_multiple_of(W8A16_M32_K_BLOCK)
        && m32_kernel.0 != 0
}

/// 2026-09-25: Strided launch of `w8a16_gemm_pipelined_m32`.
///
/// `input` `[M, a_row_stride]` BF16 (first `k` of each row read), `weight`
/// `[N, K]` FP8 E4M3 with `block_scale` `[N/128, K/128]` FP32, `output`
/// `[M, c_row_stride]` BF16 (first `n` of each row written). Strides in
/// elements. The A rows feed 16-byte `cp.async` chunks, so `a_row_stride`
/// must keep them 16 B-aligned (multiple of 8).
///
/// Returns an error for `m == 0`, for `k` not a positive multiple of 128, for
/// pitches shorter than `k` / `n`, and for an `a_row_stride` that is not a
/// multiple of 8. Same argument order as `w8a16_gemv_batch{4,16}_strided`, so
/// the multi-seq Q/K/V tier holds it as one more `StridedBatchGemv` arm.
///
/// Grid: (ceil(N/32), ceil(M/32), 1)  Block: (256, 1, 1); 2026-09-27: at
/// 1..=16 rows with K <= [`W8A16_M16_MAX_K`], (ceil(N / 8), 1, 1).
#[allow(clippy::too_many_arguments)]
pub fn w8a16_gemm_pipelined_m32_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<()> {
    ensure!(
        m >= 1,
        "w8a16_gemm_pipelined_m32: m=0 has no rows to compute"
    );
    ensure!(
        k >= W8A16_M32_K_BLOCK && k.is_multiple_of(W8A16_M32_K_BLOCK),
        "w8a16_gemm_pipelined_m32: K={k} must be a positive multiple of \
         {W8A16_M32_K_BLOCK} (one FP8 scale block per K-step)"
    );
    ensure!(
        a_row_stride >= k && c_row_stride >= n,
        "w8a16_gemm_pipelined_m32: row pitches (a={a_row_stride}, c={c_row_stride}) \
         must cover the used extents (k={k}, n={n})"
    );
    ensure!(
        a_row_stride.is_multiple_of(8),
        "w8a16_gemm_pipelined_m32: a_row_stride={a_row_stride} must keep rows \
         16B-aligned (cp.async A tile)"
    );
    // 2026-09-28: The row-tile kernel under the canonical tiers (`w8a16_tc_rows.rs`).
    if super::w8a16_tc_rows_launch(
        gpu,
        input,
        weight,
        block_scale,
        output,
        m,
        n,
        k,
        a_row_stride,
        c_row_stride,
        stream,
    )? {
        return Ok(());
    }
    // 2026-09-27: 1..=16 rows (K <= 5120) run the kernel's skinny body.
    let (grid, block, smem) = if m <= W8A16_M16_MAX_ROWS && k <= W8A16_M16_MAX_K {
        ([div_ceil(n, 8), 1, 1], [256, 1, 1], w8a16_m16_smem_bytes(k))
    } else {
        (
            [
                div_ceil(n, W8A16_M32_TILE_COLS),
                div_ceil(m, W8A16_M32_TILE_ROWS),
                1,
            ],
            [256, 1, 1],
            W8A16_M32_SMEM_BYTES,
        )
    };
    KernelLaunch::new(gpu, kernel)
        .grid(grid)
        .block(block)
        .shared_mem(smem)
        .arg_ptr(input)
        .arg_ptr(weight)
        .arg_ptr(block_scale)
        .arg_ptr(output)
        .arg_u32(m)
        .arg_u32(n)
        .arg_u32(k)
        .arg_u32(a_row_stride)
        .arg_u32(c_row_stride)
        .launch(stream)
}

/// 2026-09-25: Contiguous `[M, K]` -> `[M, N]` launch of the same kernel: the
/// pitches are `k` and `n`. Same signature as [`super::w8a16_gemm_pipelined`].
#[allow(clippy::too_many_arguments)]
pub fn w8a16_gemm_pipelined_m32(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    weight: DevicePtr,
    block_scale: DevicePtr,
    output: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    stream: u64,
) -> Result<()> {
    w8a16_gemm_pipelined_m32_strided(
        gpu,
        kernel,
        input,
        weight,
        block_scale,
        output,
        m,
        n,
        k,
        k,
        n,
        stream,
    )
}

/// 2026-09-26: Rows of the `w8a16_gemm_pipelined_m64` tile.
pub const W8A16_M64_TILE_ROWS: u32 = 64;

/// 2026-09-26: 33..=64 rows take the 64-row twin `w8a16_gemm_pipelined_m64`
/// when its handle is nonzero and `K` is a positive multiple of 128: one weight
/// pass for a whole k = 2 or 3 batched verify at C = 16, where the 32-row twin
/// would stream the weight per tile and the 128-row tile is three-quarters
/// empty. The three kernels give every output the same sub-MMAs in the same
/// order (`native_fp8_gdn_proj_m32_microtest` compares them byte for byte).
pub fn w8a16_pipelined_prefers_m64(m: u32, k: u32, m64_kernel: KernelHandle) -> bool {
    (W8A16_M32_TILE_ROWS + 1..=W8A16_M64_TILE_ROWS).contains(&m)
        && k >= W8A16_M32_K_BLOCK
        && k.is_multiple_of(W8A16_M32_K_BLOCK)
        && m64_kernel.0 != 0
}

/// 2026-09-26: Strided launch of `w8a16_gemm_pipelined_m64`: the arguments
/// and checks of [`w8a16_gemm_pipelined_m32_strided`], grid
/// (ceil(N/32), ceil(M/64), 1), block 256. The multi-seq attention Q/K/V and
/// O arms take it above 32 rows under `ModelLevers::fp8_attn_m32`.
#[allow(clippy::too_many_arguments)]
pub fn w8a16_gemm_pipelined_m64_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<()> {
    ensure!(
        m >= 1 && k >= W8A16_M32_K_BLOCK && k.is_multiple_of(W8A16_M32_K_BLOCK),
        "w8a16_gemm_pipelined_m64: m={m} must be positive and K={k} a positive multiple \
         of {W8A16_M32_K_BLOCK}"
    );
    ensure!(
        a_row_stride >= k && c_row_stride >= n && a_row_stride.is_multiple_of(8),
        "w8a16_gemm_pipelined_m64: row pitches (a={a_row_stride}, c={c_row_stride}) must \
         cover k={k} / n={n}, and a must keep rows 16B-aligned"
    );
    if super::w8a16_tc_rows_launch(
        gpu,
        input,
        weight,
        block_scale,
        output,
        m,
        n,
        k,
        a_row_stride,
        c_row_stride,
        stream,
    )? {
        return Ok(());
    }
    KernelLaunch::new(gpu, kernel)
        .grid([
            div_ceil(n, W8A16_M32_TILE_COLS),
            div_ceil(m, W8A16_M64_TILE_ROWS),
            1,
        ])
        .block([256, 1, 1])
        .arg_ptr(input)
        .arg_ptr(weight)
        .arg_ptr(block_scale)
        .arg_ptr(output)
        .arg_u32(m)
        .arg_u32(n)
        .arg_u32(k)
        .arg_u32(a_row_stride)
        .arg_u32(c_row_stride)
        .launch(stream)
}

/// 2026-09-26: Contiguous launch of `w8a16_gemm_pipelined_m64` (pitches `k`
/// and `n`). Same signature as [`super::w8a16_gemm_pipelined`].
#[allow(clippy::too_many_arguments)]
pub fn w8a16_gemm_pipelined_m64(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    weight: DevicePtr,
    block_scale: DevicePtr,
    output: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    stream: u64,
) -> Result<()> {
    w8a16_gemm_pipelined_m64_strided(
        gpu,
        kernel,
        input,
        weight,
        block_scale,
        output,
        m,
        n,
        k,
        k,
        n,
        stream,
    )
}

/// 2026-09-25: Contiguous block-scaled W8A16 GEMM, tile chosen by
/// [`w8a16_pipelined_prefers_m32`]: the 32-tile twin on `m32_kernel`, else
/// (2026-09-26) [`w8a16_pipelined_prefers_m64`]'s 64-tile twin on `m64_kernel`,
/// else the 128-tile `w8a16_gemm_pipelined` on `full_kernel`. The GDN
/// `in_proj_qkvz` / `out_proj` batched-verify arms (`trait_decode_batched.rs`)
/// call it.
#[allow(clippy::too_many_arguments)]
pub fn w8a16_gemm_pipelined_by_m(
    gpu: &dyn GpuBackend,
    full_kernel: KernelHandle,
    m32_kernel: KernelHandle,
    m64_kernel: KernelHandle,
    input: DevicePtr,
    weight: DevicePtr,
    block_scale: DevicePtr,
    output: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    stream: u64,
) -> Result<()> {
    if w8a16_pipelined_prefers_m32(m, k, m32_kernel) {
        w8a16_gemm_pipelined_m32(
            gpu,
            m32_kernel,
            input,
            weight,
            block_scale,
            output,
            m,
            n,
            k,
            stream,
        )
    } else if w8a16_pipelined_prefers_m64(m, k, m64_kernel) {
        w8a16_gemm_pipelined_m64(
            gpu,
            m64_kernel,
            input,
            weight,
            block_scale,
            output,
            m,
            n,
            k,
            stream,
        )
    } else {
        super::w8a16_gemm_pipelined(
            gpu,
            full_kernel,
            input,
            weight,
            block_scale,
            output,
            m,
            n,
            k,
            stream,
        )
    }
}

#[cfg(test)]
#[path = "w8a16_gemm_pipelined_m32_tests.rs"]
mod tests;