metrale_model_layers/forward/
quant_weights.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The `QuantWeights` trait: the matvec operations a per-layer forward
4//! needs from a quantised weight, so `super::qwen3_5` never names a concrete
5//! weight type.
6//!
7//! Owner: model-layers.
8//! Invariants: none beyond the types.
9//!
10//! `gemv` is required. The fused operations are optional: `gemv_gate_up_with`
11//! defaults to two `gemv` calls, and `gemv_silu_gate` / `gemv_silu_gate_resid`
12//! default to an error, so a caller learns that no fused kernel exists instead
13//! of getting a silent slow path.
14
15use anyhow::{Result, bail};
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
17use metrale_gpu_runtime::mlx_int8::{self, MlxInt8Weight};
18
19/// 2026-09-25: A quantised `[N, K]` weight that runs matvecs on a `GpuBackend`.
20/// The one implementation, for `MlxInt8Weight`, is in this file.
21pub trait QuantWeights: Send + Sync {
22    /// 2026-09-25: Output dimension `N` of the `[N, K]` weight.
23    fn out_features(&self) -> u32;
24
25    /// 2026-09-25: Input dimension `K` of the `[N, K]` weight.
26    fn in_features(&self) -> u32;
27
28    /// 2026-09-25: Decode matvec `y = self @ x`.
29    ///
30    /// `x` is a BF16 buffer of length `in_features()`; `y` must hold at
31    /// least `out_features()` BF16 slots.
32    fn gemv(&self, gpu: &dyn GpuBackend, x: DevicePtr, y: DevicePtr, stream: u64) -> Result<()>;
33
34    /// 2026-09-25: Dual-output GEMV over one input: `gate_y = self @ x` and
35    /// `up_y = other @ x`. The default is two `gemv` calls. The
36    /// `MlxInt8Weight` impl overrides it with the single
37    /// `mlx_int8_gemv_gate_up` launch.
38    ///
39    /// `where Self: Sized` keeps this method off the `dyn QuantWeights`
40    /// surface; the forward modules call it through a generic parameter.
41    fn gemv_gate_up_with(
42        &self,
43        other: &Self,
44        gpu: &dyn GpuBackend,
45        x: DevicePtr,
46        gate_y: DevicePtr,
47        up_y: DevicePtr,
48        stream: u64,
49    ) -> Result<()>
50    where
51        Self: Sized,
52    {
53        debug_assert_eq!(self.out_features(), other.out_features());
54        debug_assert_eq!(self.in_features(), other.in_features());
55        self.gemv(gpu, x, gate_y, stream)?;
56        other.gemv(gpu, x, up_y, stream)?;
57        Ok(())
58    }
59
60    /// 2026-09-25: Fused FFN tail `y = self @ (silu(gate) ⊙ up)`. The default
61    /// returns an error; an implementation overrides it with a fused kernel
62    /// or its own composition.
63    fn gemv_silu_gate(
64        &self,
65        _gpu: &dyn GpuBackend,
66        _gate: DevicePtr,
67        _up: DevicePtr,
68        _y: DevicePtr,
69        _stream: u64,
70    ) -> Result<()> {
71        bail!(
72            "QuantWeights::gemv_silu_gate is not implemented for this weight type — \
73             override with a fused kernel or compose silu+mul+gemv on the caller side"
74        )
75    }
76
77    /// 2026-09-25: [`Self::gemv_silu_gate`] with the residual add folded in:
78    ///   `y[n] = x_resid[n] + sum_k self[n, k] * (silu(gate[k]) ⊙ up[k])`.
79    ///
80    /// The default returns an error.
81    fn gemv_silu_gate_resid(
82        &self,
83        _gpu: &dyn GpuBackend,
84        _gate: DevicePtr,
85        _up: DevicePtr,
86        _x_resid: DevicePtr,
87        _y: DevicePtr,
88        _stream: u64,
89    ) -> Result<()> {
90        bail!(
91            "QuantWeights::gemv_silu_gate_resid is not implemented for this weight type — \
92             override with a fused kernel or compose silu+mul+gemv+add on the caller side"
93        )
94    }
95}
96
97// 2026-09-25: The impl sits beside the trait because of the orphan rule:
98// `MlxInt8Weight` belongs to metrale-gpu-runtime, which cannot see this trait.
99// Each method forwards to the weight type's own fused-kernel functions.
100
101impl QuantWeights for MlxInt8Weight {
102    fn out_features(&self) -> u32 {
103        self.out_features
104    }
105    fn in_features(&self) -> u32 {
106        self.in_features
107    }
108    fn gemv(&self, gpu: &dyn GpuBackend, x: DevicePtr, y: DevicePtr, stream: u64) -> Result<()> {
109        MlxInt8Weight::gemv(self, gpu, x, y, stream)
110    }
111    fn gemv_gate_up_with(
112        &self,
113        other: &Self,
114        gpu: &dyn GpuBackend,
115        x: DevicePtr,
116        gate_y: DevicePtr,
117        up_y: DevicePtr,
118        stream: u64,
119    ) -> Result<()> {
120        // 2026-09-25: One `mlx_int8_gemv_gate_up` launch computes both outputs.
121        mlx_int8::gemv_gate_up(gpu, self, other, x, gate_y, up_y, stream)
122    }
123    fn gemv_silu_gate(
124        &self,
125        gpu: &dyn GpuBackend,
126        gate: DevicePtr,
127        up: DevicePtr,
128        y: DevicePtr,
129        stream: u64,
130    ) -> Result<()> {
131        MlxInt8Weight::gemv_silu_gate(self, gpu, gate, up, y, stream)
132    }
133    fn gemv_silu_gate_resid(
134        &self,
135        gpu: &dyn GpuBackend,
136        gate: DevicePtr,
137        up: DevicePtr,
138        x_resid: DevicePtr,
139        y: DevicePtr,
140        stream: u64,
141    ) -> Result<()> {
142        MlxInt8Weight::gemv_silu_gate_resid(self, gpu, gate, up, x_resid, y, stream)
143    }
144}
metrale_model_layers/forward/quant_weights.rs

metrale_model_layers/forward/
quant_weights.rs
