metrale_model_layers/layers/ops/
gemm_quant_gemv.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the dense GEMVs: BF16 for one row, two rows and M
4//! rows, and FP8-weight for one row.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - [`dense_gemv_batchm`] returns an error, without launching, when `m` is
9//!   outside `1..=DENSE_GEMV_BATCHM_MAX_M`.
10
11use super::*;
12
13/// 2026-09-25: Dense BF16 GEMV for one row, `C = A @ B^T`: A `[1, K]`, B
14/// `[N, K]`, C `[1, N]`. Each 256-thread block computes 4 outputs, 64 threads
15/// per output (kernel `dense_gemv_bf16(A, B, C, N, K)`).
16pub fn dense_gemv(
17    gpu: &dyn GpuBackend,
18    kernel: KernelHandle,
19    input: DevicePtr,
20    weight: &DenseWeight,
21    output: DevicePtr,
22    n: u32,
23    k: u32,
24    stream: u64,
25) -> Result<()> {
26    KernelLaunch::new(gpu, kernel)
27        .grid([div_ceil(n, 4), 1, 1])
28        .block([256, 1, 1])
29        .arg_ptr(input)
30        .arg_ptr(weight.weight)
31        .arg_ptr(output)
32        .arg_u32(n)
33        .arg_u32(k)
34        .launch(stream)
35}
36
37/// 2026-09-25: Dense BF16 GEMV for two rows in one pass over the weight. Each
38/// row follows `dense_gemv_bf16`'s K order and reduction (gb10 kernels).
39///
40/// `input`: `[2, K]` BF16, contiguous; `output`: two rows at
41/// `output + t * out_stride` (BF16 elements). The batched SSM decode uses it
42/// for a BF16 `in_proj_qkvz` (`qwen3_ssm/trait_decode_batched.rs`).
43///
44/// Kernel: `dense_gemv_bf16_batch2(A, B, C, N, K, out_stride)`
45#[allow(clippy::too_many_arguments)]
46pub fn dense_gemv_batch2(
47    gpu: &dyn GpuBackend,
48    kernel: KernelHandle,
49    input: DevicePtr,
50    weight: &DenseWeight,
51    output: DevicePtr,
52    n: u32,
53    k: u32,
54    out_stride: u32,
55    stream: u64,
56) -> Result<()> {
57    KernelLaunch::new(gpu, kernel)
58        .grid([div_ceil(n, 4), 1, 1])
59        .block([256, 1, 1])
60        .arg_ptr(input)
61        .arg_ptr(weight.weight)
62        .arg_ptr(output)
63        .arg_u32(n)
64        .arg_u32(k)
65        .arg_u32(out_stride)
66        .launch(stream)
67}
68
69#[allow(clippy::too_many_arguments)]
70/// 2026-09-25: Rows `dense_gemv_bf16_batchm` computes: `MAX_M` in
71/// `kernels/gb10/common/dense_gemv_bf16_batchm.cu`. The kernel clamps a
72/// larger M, so [`dense_gemv_batchm`] refuses above this value. Each row's
73/// accumulation does not depend on M, so a row's result is the same at every
74/// batch width.
75///
76/// Other readers: the row thresholds of the GLM-5.3 DSA, KDA and MLP dispatch
77/// (`glm5next_dsa/layer.rs`, `glm5next_kda/mod.rs`, `glm5next_mlp/forward.rs`,
78/// `glm5next_layer/levers.rs`), the `verify_k` workspace size
79/// (`glm5_next_load/loader.rs`) and the clamp in
80/// `target_defaults::resolve_batchm_max`.
81pub const DENSE_GEMV_BATCHM_MAX_M: u32 = 16;
82
83/// 2026-09-25: The widest batch that the MTP row dispatch
84/// (`mtp_head/row_dispatch.rs`) and the BF16 arm of the model's
85/// `lm_head_batched` (`impl_a3_lm_head.rs`) hand to `dense_gemv_bf16_batchm`;
86/// `target_defaults::BASELINE_BATCHM_MAX` is set from it. Above it those paths
87/// choose other kernels, some of which accumulate in a different order, so
88/// this value decides which bits a decode of a given width produces.
89pub const DENSE_GEMV_BATCHM_DECODE_MAX_M: u32 = 8;
90
91/// 2026-09-25: Dense BF16 GEMV for M rows, `C[t] = A[t] @ B^T` for `t` in
92/// `[0, M)`, reading the weight once for all rows. Each row follows
93/// `dense_gemv_bf16`'s K order and reduction (gb10 kernels).
94///
95/// `input`: `[M, K]` BF16, contiguous; `output`: M rows at
96/// `output + t * out_stride` (BF16 elements). Refuses `m` outside
97/// `1..=DENSE_GEMV_BATCHM_MAX_M`.
98///
99/// Kernel: `dense_gemv_bf16_batchm(A, B, C, M, N, K, out_stride)`
100pub fn dense_gemv_batchm(
101    gpu: &dyn GpuBackend,
102    kernel: KernelHandle,
103    input: DevicePtr,
104    weight: &DenseWeight,
105    output: DevicePtr,
106    m: u32,
107    n: u32,
108    k: u32,
109    out_stride: u32,
110    stream: u64,
111) -> Result<()> {
112    // 2026-09-25: The kernel clamps M to its MAX_M instead of failing, which
113    // would leave the rows past it unwritten.
114    ensure!(
115        (1..=DENSE_GEMV_BATCHM_MAX_M).contains(&m),
116        "dense_gemv_batchm: m={m} outside 1..={DENSE_GEMV_BATCHM_MAX_M} \
117         (kernel MAX_M clamps silently; use dense_gemm_tc for wider batches)"
118    );
119    KernelLaunch::new(gpu, kernel)
120        .grid([div_ceil(n, 4), 1, 1])
121        .block([256, 1, 1])
122        .arg_ptr(input)
123        .arg_ptr(weight.weight)
124        .arg_ptr(output)
125        .arg_u32(m)
126        .arg_u32(n)
127        .arg_u32(k)
128        .arg_u32(out_stride)
129        .launch(stream)
130}
131
132/// 2026-09-27: [`dense_gemv_batchm`] over `m` rows split into `y_blocks` block
133/// rows of `ceil(m / y_blocks)` rows each (at most `DENSE_GEMV_BATCHM_MAX_M`):
134/// more blocks for a narrow weight such as the MoE router, and one launch past
135/// 16 rows. Each row's result is the same as in any other split.
136#[allow(clippy::too_many_arguments)]
137pub fn dense_gemv_batchm_split(
138    gpu: &dyn GpuBackend,
139    kernel: KernelHandle,
140    input: DevicePtr,
141    weight: &DenseWeight,
142    output: DevicePtr,
143    m: u32,
144    y_blocks: u32,
145    n: u32,
146    k: u32,
147    out_stride: u32,
148    stream: u64,
149) -> Result<()> {
150    ensure!(
151        m >= 1 && (1..=m).contains(&y_blocks) && m.div_ceil(y_blocks) <= DENSE_GEMV_BATCHM_MAX_M,
152        "dense_gemv_batchm_split: m={m} over y_blocks={y_blocks} must give 1..={DENSE_GEMV_BATCHM_MAX_M} rows per block row"
153    );
154    KernelLaunch::new(gpu, kernel)
155        .grid([div_ceil(n, 4), y_blocks, 1])
156        .block([256, 1, 1])
157        .arg_ptr(input)
158        .arg_ptr(weight.weight)
159        .arg_ptr(output)
160        .arg_u32(m)
161        .arg_u32(n)
162        .arg_u32(k)
163        .arg_u32(out_stride)
164        .launch(stream)
165}
166
167/// 2026-09-25: FP8-weight GEMV for one row,
168/// `C = A @ (dequant(B_fp8) * row_scale)^T`: A `[1, K]` BF16, B `[N, K]` FP8
169/// E4M3, row_scale `[N]` f32, C `[1, N]` BF16. One byte per weight, against
170/// two for [`dense_gemv`]. Each 256-thread block computes 4 outputs, 64
171/// threads per output (kernel `dense_gemv_fp8w(A, B, row_scale, C, N, K)`).
172pub fn dense_gemv_fp8w(
173    gpu: &dyn GpuBackend,
174    kernel: KernelHandle,
175    input: DevicePtr,
176    weight: &Fp8DenseWeight,
177    output: DevicePtr,
178    n: u32,
179    k: u32,
180    stream: u64,
181) -> Result<()> {
182    KernelLaunch::new(gpu, kernel)
183        .grid([div_ceil(n, 4), 1, 1])
184        .block([256, 1, 1])
185        .arg_ptr(input)
186        .arg_ptr(weight.weight)
187        .arg_ptr(weight.row_scale)
188        .arg_ptr(output)
189        .arg_u32(n)
190        .arg_u32(k)
191        .launch(stream)
192}
metrale_model_layers/layers/ops/gemm_quant_gemv.rs

metrale_model_layers/layers/ops/
gemm_quant_gemv.rs
