spark_model/forward/
quant_weights.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Per-quant weight abstraction.
3//!
4//! Plug-in trait the vendor-agnostic forward modules (e.g.
5//! [`super::qwen3_5`]) call instead of reaching into a concrete weight
6//! type like `MlxInt8Weight`. Each backend's weight loader implements
7//! `QuantWeights` and overrides whichever fused variants it ships
8//! kernels for; the rest fall back to default impls that compose the
9//! unfused primitives or fail loudly.
10//!
11//! Convention: `gemv` is mandatory (every backend ships one). Fused
12//! variants (`gemv_silu_gate`, `gemv_silu_gate_resid`,
13//! `gemv_gate_up_with`) are advisory — backends override only the
14//! ones they have fused kernels for. Defaults either fall back to a
15//! correct-but-slower composition (`gemv_gate_up_with`) or error
16//! loudly so the caller can choose between hard-requiring the fused
17//! kernel and degrading to a manual silu+mul+gemv pipeline.
18
19use anyhow::{Result, bail};
20use spark_runtime::gpu::{DevicePtr, GpuBackend};
21use spark_runtime::weights::mlx_int8::{self, MlxInt8Weight};
22
23/// A quantised weight tensor that can drive matvec / matmul ops on a
24/// `GpuBackend`. Implementations live with each backend's weight
25/// loader (Metal: `MlxInt8Weight`; future CUDA: `Nvfp4Weight`,
26/// `Fp8DenseWeight`, …).
27pub trait QuantWeights: Send + Sync {
28    /// Output dimension `N` of the underlying `[N, K]` weight.
29    fn out_features(&self) -> u32;
30
31    /// Input dimension `K` of the underlying `[N, K]` weight.
32    fn in_features(&self) -> u32;
33
34    /// Decode-path matvec: `y = self @ x`.
35    ///
36    /// `x` is a BF16 buffer of length `in_features()`; `y` must hold at
37    /// least `out_features()` BF16 slots.
38    fn gemv(&self, gpu: &dyn GpuBackend, x: DevicePtr, y: DevicePtr, stream: u64) -> Result<()>;
39
40    /// Dual-output GEMV with shared input: `gate_y = self @ x`,
41    /// `up_y = other @ x`. The default impl is two serial `gemv`
42    /// calls — correct on any backend, just slower than the fused
43    /// kernel some backends ship (e.g. Metal's
44    /// `mlx_int8_gemv_gate_up`). Backends with a fused dual-output
45    /// path override this to halve x-side memory bandwidth and
46    /// remove a launch.
47    ///
48    /// `where Self: Sized` keeps this method off the dyn-trait surface
49    /// (it's only callable through generic-parameter dispatch, which
50    /// is what the forward modules use anyway).
51    fn gemv_gate_up_with(
52        &self,
53        other: &Self,
54        gpu: &dyn GpuBackend,
55        x: DevicePtr,
56        gate_y: DevicePtr,
57        up_y: DevicePtr,
58        stream: u64,
59    ) -> Result<()>
60    where
61        Self: Sized,
62    {
63        debug_assert_eq!(self.out_features(), other.out_features());
64        debug_assert_eq!(self.in_features(), other.in_features());
65        self.gemv(gpu, x, gate_y, stream)?;
66        other.gemv(gpu, x, up_y, stream)?;
67        Ok(())
68    }
69
70    /// Fused FFN tail: `y = self @ (silu(gate) ⊙ up)`.
71    ///
72    /// Default impl errors — backends that ship the fused kernel
73    /// override; backends that don't can either error out (forcing
74    /// the caller to do the unfused dance) or override with their
75    /// own composition.
76    fn gemv_silu_gate(
77        &self,
78        _gpu: &dyn GpuBackend,
79        _gate: DevicePtr,
80        _up: DevicePtr,
81        _y: DevicePtr,
82        _stream: u64,
83    ) -> Result<()> {
84        bail!(
85            "QuantWeights::gemv_silu_gate is not implemented for this weight type — \
86             override with a fused kernel or compose silu+mul+gemv on the caller side"
87        )
88    }
89
90    /// Same as [`Self::gemv_silu_gate`] but additionally folds the
91    /// layer-output residual addition into the same kernel:
92    ///   `y[n] = x_resid[n] + sum_k self[n, k] * (silu(gate[k]) ⊙ up[k])`.
93    ///
94    /// Default impl errors. Backends that ship a `_resid` variant of
95    /// the fused kernel override.
96    fn gemv_silu_gate_resid(
97        &self,
98        _gpu: &dyn GpuBackend,
99        _gate: DevicePtr,
100        _up: DevicePtr,
101        _x_resid: DevicePtr,
102        _y: DevicePtr,
103        _stream: u64,
104    ) -> Result<()> {
105        bail!(
106            "QuantWeights::gemv_silu_gate_resid is not implemented for this weight type — \
107             override with a fused kernel or compose silu+mul+gemv+add on the caller side"
108        )
109    }
110}
111
112// ── Backend impls ────────────────────────────────────────────────────
113//
114// Lives next to the trait so the orphan rule applies (the trait is
115// foreign to spark-runtime, but native here). Each impl is a thin
116// forwarding shim onto the concrete weight type's inherent methods,
117// keeping the optimised fused-kernel paths intact.
118
119impl QuantWeights for MlxInt8Weight {
120    fn out_features(&self) -> u32 {
121        self.out_features
122    }
123    fn in_features(&self) -> u32 {
124        self.in_features
125    }
126    fn gemv(&self, gpu: &dyn GpuBackend, x: DevicePtr, y: DevicePtr, stream: u64) -> Result<()> {
127        MlxInt8Weight::gemv(self, gpu, x, y, stream)
128    }
129    fn gemv_gate_up_with(
130        &self,
131        other: &Self,
132        gpu: &dyn GpuBackend,
133        x: DevicePtr,
134        gate_y: DevicePtr,
135        up_y: DevicePtr,
136        stream: u64,
137    ) -> Result<()> {
138        // Atlas Metal ships a fused dual-output kernel
139        // (`mlx_int8_gemv_gate_up`); use it instead of two serial
140        // gemvs to halve x-side bandwidth and remove a launch.
141        mlx_int8::gemv_gate_up(gpu, self, other, x, gate_y, up_y, stream)
142    }
143    fn gemv_silu_gate(
144        &self,
145        gpu: &dyn GpuBackend,
146        gate: DevicePtr,
147        up: DevicePtr,
148        y: DevicePtr,
149        stream: u64,
150    ) -> Result<()> {
151        MlxInt8Weight::gemv_silu_gate(self, gpu, gate, up, y, stream)
152    }
153    fn gemv_silu_gate_resid(
154        &self,
155        gpu: &dyn GpuBackend,
156        gate: DevicePtr,
157        up: DevicePtr,
158        x_resid: DevicePtr,
159        y: DevicePtr,
160        stream: u64,
161    ) -> Result<()> {
162        MlxInt8Weight::gemv_silu_gate_resid(self, gpu, gate, up, x_resid, y, stream)
163    }
164}
spark_model/forward/quant_weights.rs

spark_model/forward/
quant_weights.rs
