metrale_model_layers/layers/ops/
w8a16_gemv_ncol.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the N-column-blocked W8A16 batched GEMV
4//! (`kernels/hopper/common/w8a16_gemv_ncol.cu`, module `w8a16_gemv_ncol`).
5//!
6//! It computes what `w8a16_gemv_batch16` computes, with the same per-row
7//! reduction order, but one thread owns `N_COLS` adjacent output columns, so
8//! the activation loads and BF16->FP32 converts are shared by `N_COLS` weight
9//! bytes. The `native_fp8_attn_decode_batch_microtest` example compares the
10//! outputs bit for bit.
11//!
12//! The four entry points keep the `ContiguousBatchGemv` / `StridedBatchGemv`
13//! signatures of the `w8a16_gemv_batch{4,16}` launchers, so a decode tier picks
14//! a rung by swapping a function pointer and a handle.
15//!
16//! Owner: model-layers ops.
17//! Invariants: none beyond the types.
18
19use anyhow::{Result, ensure};
20use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
21use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
22
23/// 2026-09-25: Output columns one thread owns. The kernel is instantiated for
24/// both; `METRALE_ATTN_NCOL_WIDTH` picks one (`attn_ncol_gemv::NcolWidth`).
25const COLS_PER_THREAD_2: u32 = 2;
26const COLS_PER_THREAD_4: u32 = 4;
27
28/// 2026-09-25: Output groups per block: `w8a16_gemv_ncol.cu`'s
29/// `N_GROUPS_PER_BLOCK`, 64 threads each, as in `w8a16_gemv_batch4.cu`. A block
30/// covers `GROUPS_PER_BLOCK * n_cols` columns, the grid divisor below.
31const GROUPS_PER_BLOCK: u32 = 4;
32
33/// 2026-09-25: `N_COLS=2`, contiguous A `[M, K]` and C `[M, N]`.
34///
35/// Kernel: `w8a16_gemv_batch16_ncol2` (module `w8a16_gemv_ncol`).
36/// Grid: (ceil(N/8), 1, 1)  Block: (256, 1, 1)
37#[allow(clippy::too_many_arguments)]
38pub fn w8a16_gemv_batch16_ncol2(
39    gpu: &dyn GpuBackend,
40    kernel: KernelHandle,
41    input: DevicePtr,
42    weight: DevicePtr,
43    block_scale: DevicePtr,
44    output: DevicePtr,
45    m: u32,
46    n: u32,
47    k: u32,
48    stream: u64,
49) -> Result<()> {
50    ncol_launch(
51        gpu,
52        kernel,
53        input,
54        weight,
55        block_scale,
56        output,
57        m,
58        n,
59        k,
60        k,
61        n,
62        COLS_PER_THREAD_2,
63        false,
64        stream,
65    )
66}
67
68/// 2026-09-25: `N_COLS=4`, contiguous A and C. Grid: (ceil(N/16), 1, 1)
69#[allow(clippy::too_many_arguments)]
70pub fn w8a16_gemv_batch16_ncol4(
71    gpu: &dyn GpuBackend,
72    kernel: KernelHandle,
73    input: DevicePtr,
74    weight: DevicePtr,
75    block_scale: DevicePtr,
76    output: DevicePtr,
77    m: u32,
78    n: u32,
79    k: u32,
80    stream: u64,
81) -> Result<()> {
82    ncol_launch(
83        gpu,
84        kernel,
85        input,
86        weight,
87        block_scale,
88        output,
89        m,
90        n,
91        k,
92        k,
93        n,
94        COLS_PER_THREAD_4,
95        false,
96        stream,
97    )
98}
99
100/// 2026-09-25: `N_COLS=2`, explicit A and C row pitches in elements, for the
101/// multi-seq `[n, per_seq_qkv]` decode buffer, as `w8a16_gemv_batch16_strided`.
102///
103/// Kernel: `w8a16_gemv_batch16_ncol2_strided` (module `w8a16_gemv_ncol`).
104#[allow(clippy::too_many_arguments)]
105pub fn w8a16_gemv_batch16_ncol2_strided(
106    gpu: &dyn GpuBackend,
107    kernel: KernelHandle,
108    input: DevicePtr,
109    weight: DevicePtr,
110    block_scale: DevicePtr,
111    output: DevicePtr,
112    m: u32,
113    n: u32,
114    k: u32,
115    a_row_stride: u32,
116    c_row_stride: u32,
117    stream: u64,
118) -> Result<()> {
119    ncol_launch(
120        gpu,
121        kernel,
122        input,
123        weight,
124        block_scale,
125        output,
126        m,
127        n,
128        k,
129        a_row_stride,
130        c_row_stride,
131        COLS_PER_THREAD_2,
132        true,
133        stream,
134    )
135}
136
137/// 2026-09-25: `N_COLS=4`, explicit row pitches.
138#[allow(clippy::too_many_arguments)]
139pub fn w8a16_gemv_batch16_ncol4_strided(
140    gpu: &dyn GpuBackend,
141    kernel: KernelHandle,
142    input: DevicePtr,
143    weight: DevicePtr,
144    block_scale: DevicePtr,
145    output: DevicePtr,
146    m: u32,
147    n: u32,
148    k: u32,
149    a_row_stride: u32,
150    c_row_stride: u32,
151    stream: u64,
152) -> Result<()> {
153    ncol_launch(
154        gpu,
155        kernel,
156        input,
157        weight,
158        block_scale,
159        output,
160        m,
161        n,
162        k,
163        a_row_stride,
164        c_row_stride,
165        COLS_PER_THREAD_4,
166        true,
167        stream,
168    )
169}
170
171/// 2026-09-25: Shared launch body. `strided` selects the argument tail: the
172/// contiguous entry points pass `a_row_stride = K` / `c_row_stride = N` inside
173/// the kernel and take neither.
174///
175/// Returns an error for `m` outside `1..=16`, `k` not a multiple of 16, pitches
176/// shorter than `k` / `n`, or an `a_row_stride` that is not a multiple of 8.
177#[allow(clippy::too_many_arguments)]
178fn ncol_launch(
179    gpu: &dyn GpuBackend,
180    kernel: KernelHandle,
181    input: DevicePtr,
182    weight: DevicePtr,
183    block_scale: DevicePtr,
184    output: DevicePtr,
185    m: u32,
186    n: u32,
187    k: u32,
188    a_row_stride: u32,
189    c_row_stride: u32,
190    n_cols: u32,
191    strided: bool,
192    stream: u64,
193) -> Result<()> {
194    // 2026-09-25: The kernel is `w8a16_gemv_ncol_impl<16, N_COLS>`: above 16
195    // rows it would leave rows 16.. unwritten rather than fail.
196    ensure!(
197        (1..=16).contains(&m),
198        "w8a16_gemv_batch16_ncol{n_cols}: m={m} outside 1..=16 (kernel MAX_M)"
199    );
200    ensure!(
201        k.is_multiple_of(16),
202        "w8a16_gemv_batch16_ncol{n_cols}: K={k} not a multiple of 16 (uint4 loads)"
203    );
204    ensure!(
205        a_row_stride >= k && c_row_stride >= n,
206        "w8a16_gemv_batch16_ncol{n_cols}: row pitches (a={a_row_stride}, c={c_row_stride}) \
207         must cover the used extents (k={k}, n={n})"
208    );
209    ensure!(
210        a_row_stride.is_multiple_of(8),
211        "w8a16_gemv_batch16_ncol{n_cols}: a_row_stride={a_row_stride} must keep rows \
212         16B-aligned (uint4 activation loads)"
213    );
214    let mut launch = KernelLaunch::new(gpu, kernel)
215        .grid([div_ceil(n, GROUPS_PER_BLOCK * n_cols), 1, 1])
216        .block([256, 1, 1])
217        .arg_ptr(input)
218        .arg_ptr(weight)
219        .arg_ptr(block_scale)
220        .arg_ptr(output)
221        .arg_u32(m)
222        .arg_u32(n)
223        .arg_u32(k);
224    if strided {
225        launch = launch.arg_u32(a_row_stride).arg_u32(c_row_stride);
226    }
227    launch.launch(stream)
228}
metrale_model_layers/layers/ops/w8a16_gemv_ncol.rs

metrale_model_layers/layers/ops/
w8a16_gemv_ncol.rs
