spark_model/layers/ops/
fp8_gemv_batch.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! FP8-weight dual-GEMV (batch=2) dispatch.
4//!
5//! `dense_gemv_fp8w_batch2` computes two output rows from one pass over the
6//! FP8 weight matrix — the batch=2 sibling of `dense_gemv_fp8w`. It halves
7//! FP8 weight bandwidth vs two M=1 GEMV launches and is bit-identical to
8//! running `dense_gemv_fp8w` twice (per-token reduction order unchanged).
9//! Used by the K=2 MTP verify path where the two verify positions share
10//! weights but have distinct activations (lm_head, attention Q/K/V/O, SSM
11//! out_proj).
12
13use anyhow::{Result, ensure};
14use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
16
17use crate::weight_map::Fp8DenseWeight;
18
19/// Register-tiled batched row-scaled FP8 GEMV (M<=8, T=2 outputs/thread) —
20/// the FP8 twin of `w4a16_gemv_batch8_rt2`, for the DFlash drafter PROPOSE
21/// path. `input` `[M, K]` BF16, `output` `[M, N]` BF16; per-row f32 scale
22/// applied at write-out inside the kernel. Replaces the prefill-class tile
23/// GEMMs (`fp8_gemm_t_row_scaled` M64-tile / `_m16`) that pad 87%/50% of
24/// their M-tile at M=8 (~100 GB/s measured vs 180+ for the rt family).
25/// Drafter-side numerics: correctness-free under strict-argmax accept.
26/// Kernel: `fp8_gemv_rowscale_batch8_rt2` (module `fp8_gemv_rt`).
27/// Grid: (ceil(N/8), 1, 1)  Block: (256, 1, 1). Requires K % 16 == 0.
28#[allow(clippy::too_many_arguments)]
29pub fn fp8_gemv_rowscale_batch8_rt2(
30    gpu: &dyn GpuBackend,
31    kernel: KernelHandle,
32    input: DevicePtr,
33    weight: &Fp8DenseWeight,
34    output: DevicePtr,
35    m: u32,
36    n: u32,
37    k: u32,
38    stream: u64,
39) -> Result<()> {
40    ensure!(
41        (1..=8).contains(&m),
42        "fp8_gemv_rowscale_batch8_rt2: m={m} outside 1..=8 (kernel MAX_M)"
43    );
44    ensure!(
45        k.is_multiple_of(16),
46        "fp8_gemv_rowscale_batch8_rt2: K={k} not a multiple of 16"
47    );
48    KernelLaunch::new(gpu, kernel)
49        .grid([div_ceil(n, 8), 1, 1])
50        .block([256, 1, 1])
51        .arg_ptr(input)
52        .arg_ptr(weight.weight)
53        .arg_ptr(weight.row_scale)
54        .arg_ptr(output)
55        .arg_u32(m)
56        .arg_u32(n)
57        .arg_u32(k)
58        .launch(stream)
59}
60
61/// MAX_M=16 sibling of [`fp8_gemv_rowscale_batch8_rt2`] for the γ>8 DFlash
62/// propose window (flags 9..17). Same template, same launch geometry; added
63/// 2026-08-29 after STEP_TIMING measured propose 18.2ms (flag 8, rt2) vs
64/// 38.0ms (flag 9, tile fallback) — the whole γ>8 step tax.
65/// Kernel: `fp8_gemv_rowscale_batch16_rt2` (module `fp8_gemv_rt`).
66#[allow(clippy::too_many_arguments)]
67pub fn fp8_gemv_rowscale_batch16_rt2(
68    gpu: &dyn GpuBackend,
69    kernel: KernelHandle,
70    input: DevicePtr,
71    weight: &Fp8DenseWeight,
72    output: DevicePtr,
73    m: u32,
74    n: u32,
75    k: u32,
76    stream: u64,
77) -> Result<()> {
78    ensure!(
79        (1..=16).contains(&m),
80        "fp8_gemv_rowscale_batch16_rt2: m={m} outside 1..=16 (kernel MAX_M)"
81    );
82    ensure!(
83        k.is_multiple_of(16),
84        "fp8_gemv_rowscale_batch16_rt2: K={k} not a multiple of 16"
85    );
86    KernelLaunch::new(gpu, kernel)
87        .grid([div_ceil(n, 8), 1, 1])
88        .block([256, 1, 1])
89        .arg_ptr(input)
90        .arg_ptr(weight.weight)
91        .arg_ptr(weight.row_scale)
92        .arg_ptr(output)
93        .arg_u32(m)
94        .arg_u32(n)
95        .arg_u32(k)
96        .launch(stream)
97}
98
99/// FP8-weight dual-GEMV. `input` is `[2, K]` BF16, `output` is `[2, N]` BF16.
100/// Grid: (ceil(N/4), 1, 1)  Block: (256, 1, 1)
101pub fn dense_gemv_fp8w_batch2(
102    gpu: &dyn GpuBackend,
103    kernel: KernelHandle,
104    input: DevicePtr,
105    weight: &Fp8DenseWeight,
106    output: DevicePtr,
107    n: u32,
108    k: u32,
109    stream: u64,
110) -> Result<()> {
111    KernelLaunch::new(gpu, kernel)
112        .grid([div_ceil(n, 4), 1, 1])
113        .block([256, 1, 1])
114        .arg_ptr(input)
115        .arg_ptr(weight.weight)
116        .arg_ptr(weight.row_scale)
117        .arg_ptr(output)
118        .arg_u32(n)
119        .arg_u32(k)
120        .launch(stream)
121}
122
123/// Block-scaled FP8 batched GEMV (M<=4). `input` is `[M, K]` BF16, `output` is
124/// `[M, N]` BF16; `weight`/`block_scale` are the raw `w8a16_gemv` pointers (2D
125/// block-scaled FP8). One pass over the FP8 weight serves all M rows — the M=4
126/// sibling of `w8a16_gemv`, replacing `w8a16_gemm_pipelined` for n<=4 batched
127/// decode (which pads M to a 128-row MMA tile). Bit-identical per-row to
128/// `w8a16_gemv`. Grid: (ceil(N/4), 1, 1)  Block: (256, 1, 1)
129#[allow(clippy::too_many_arguments)]
130pub fn w8a16_gemv_batch4(
131    gpu: &dyn GpuBackend,
132    kernel: KernelHandle,
133    input: DevicePtr,
134    weight: DevicePtr,
135    block_scale: DevicePtr,
136    output: DevicePtr,
137    m: u32,
138    n: u32,
139    k: u32,
140    stream: u64,
141) -> Result<()> {
142    KernelLaunch::new(gpu, kernel)
143        .grid([div_ceil(n, 4), 1, 1])
144        .block([256, 1, 1])
145        .arg_ptr(input)
146        .arg_ptr(weight)
147        .arg_ptr(block_scale)
148        .arg_ptr(output)
149        .arg_u32(m)
150        .arg_u32(n)
151        .arg_u32(k)
152        .launch(stream)
153}
154
155/// Block-scaled FP8 dual-GEMV (batch=2). `input` is `[2, K]` BF16, `output` is
156/// `[2, N]` BF16; `weight`/`block_scale` are the raw `w8a16_gemv` pointers.
157/// Grid: (ceil(N/4), 1, 1)  Block: (256, 1, 1)
158#[allow(clippy::too_many_arguments)]
159pub fn w8a16_gemv_batch2(
160    gpu: &dyn GpuBackend,
161    kernel: KernelHandle,
162    input: DevicePtr,
163    weight: DevicePtr,
164    block_scale: DevicePtr,
165    output: DevicePtr,
166    n: u32,
167    k: u32,
168    stream: u64,
169) -> Result<()> {
170    KernelLaunch::new(gpu, kernel)
171        .grid([div_ceil(n, 4), 1, 1])
172        .block([256, 1, 1])
173        .arg_ptr(input)
174        .arg_ptr(weight)
175        .arg_ptr(block_scale)
176        .arg_ptr(output)
177        .arg_u32(n)
178        .arg_u32(k)
179        .launch(stream)
180}
spark_model/layers/ops/fp8_gemv_batch.rs

spark_model/layers/ops/
fp8_gemv_batch.rs
