spark_model/layers/ops/
w8a16_gemv_ncol.rs

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

spark_model/layers/ops/
w8a16_gemv_ncol.rs
