metrale_model_layers/weight_map/quantized/
transpose.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `QuantizedWeight` transposed-twin builders for the GEMM paths: plain, group-scaled, concatenated and padded.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - Every builder returns new buffers and leaves its inputs unchanged.
8//! - A transpose runs on the `transpose_u8` kernel when it resolves and
9//!   `METRALE_HOST_TRANSPOSE` is not `1`; otherwise on the host.
10
11use super::*;
12
13impl QuantizedWeight {
14    /// 2026-09-25: Resolve the `transpose_u8` GPU kernel for the load-time transpose
15    /// paths, or `None` to use the host byte-loop fallback. `None` when the
16    /// target's kernel set lacks it, or when `METRALE_HOST_TRANSPOSE=1` forces
17    /// the host path.
18    fn host_transpose_kernel(
19        gpu: &dyn GpuBackend,
20    ) -> Option<metrale_gpu_runtime::gpu::KernelHandle> {
21        if std::env::var("METRALE_HOST_TRANSPOSE").as_deref() == Ok("1") {
22            return None;
23        }
24        let k = crate::layers::try_kernel(gpu, "transpose_u8", "transpose_u8");
25        (k.0 != 0).then_some(k)
26    }
27
28    /// 2026-09-25: Transpose the packed weight `[N, K/2]` to `[K/2, N]` and the
29    /// scales `[N, K/16]` to `[K/16, N]`, into new buffers.
30    pub fn transpose_for_gemm(
31        &self,
32        gpu: &dyn GpuBackend,
33        n: usize,
34        k: usize,
35    ) -> Result<QuantizedWeight> {
36        // 2026-09-25: 16 is the NVFP4 group; native-MXFP4 experts call
37        // `transpose_for_gemm_gs` with 32.
38        self.transpose_for_gemm_gs(gpu, n, k, 16)
39    }
40
41    /// 2026-09-25: `transpose_for_gemm` with an explicit scale block size. The scale
42    /// tensor is `[N, K/group_size]`; the packed-weight transpose does not depend on it.
43    pub fn transpose_for_gemm_gs(
44        &self,
45        gpu: &dyn GpuBackend,
46        n: usize,
47        k: usize,
48        group_size: usize,
49    ) -> Result<QuantizedWeight> {
50        let half_k = k / 2;
51        let num_groups = k / group_size;
52        let packed_size = n * half_k;
53        let scale_size = n * num_groups;
54
55        // 2026-09-25: Both launches run on stream 0, which is synchronized before
56        // the new buffers are returned.
57        if let Some(tk) = Self::host_transpose_kernel(gpu) {
58            let new_weight = gpu.alloc(packed_size)?;
59            let new_scale = gpu.alloc(scale_size)?;
60            crate::layers::ops::transpose_u8(
61                gpu,
62                tk,
63                self.weight,
64                new_weight,
65                n as u32,
66                half_k as u32,
67                0,
68            )?;
69            crate::layers::ops::transpose_u8(
70                gpu,
71                tk,
72                self.weight_scale,
73                new_scale,
74                n as u32,
75                num_groups as u32,
76                0,
77            )?;
78            gpu.synchronize(0)?;
79            return Ok(QuantizedWeight {
80                weight: new_weight,
81                weight_scale: new_scale,
82                weight_scale_2: self.weight_scale_2,
83                input_scale: self.input_scale,
84                weight_scale_2_vec: self.weight_scale_2_vec,
85                act: self.act,
86            });
87        }
88
89        let mut buf = vec![0u8; packed_size];
90        gpu.copy_d2h(self.weight, &mut buf)?;
91        let mut t_buf = vec![0u8; packed_size];
92        for i in 0..n {
93            for j in 0..half_k {
94                t_buf[j * n + i] = buf[i * half_k + j];
95            }
96        }
97        let new_weight = gpu.alloc(packed_size)?;
98        gpu.copy_h2d(&t_buf, new_weight)?;
99
100        let mut sbuf = vec![0u8; scale_size];
101        gpu.copy_d2h(self.weight_scale, &mut sbuf)?;
102        let mut st_buf = vec![0u8; scale_size];
103        for i in 0..n {
104            for j in 0..num_groups {
105                st_buf[j * n + i] = sbuf[i * num_groups + j];
106            }
107        }
108        let new_scale = gpu.alloc(scale_size)?;
109        gpu.copy_h2d(&st_buf, new_scale)?;
110
111        Ok(QuantizedWeight {
112            weight: new_weight,
113            weight_scale: new_scale,
114            weight_scale_2: self.weight_scale_2,
115            input_scale: self.input_scale,
116            weight_scale_2_vec: self.weight_scale_2_vec,
117            act: self.act,
118        })
119    }
120
121    /// 2026-09-25: Transpose several weights `(weight, n)` sharing one K and
122    /// concatenate them along N into one `[K/2, N_total]` twin, so one GEMM
123    /// replaces several.
124    ///
125    /// The result carries the first part's `weight_scale_2`, `input_scale` and
126    /// `weight_scale_2_vec`, so the caller must ensure every part shares them;
127    /// nothing here checks. An empty `parts` is an error.
128    pub fn transpose_concat_for_gemm(
129        gpu: &dyn GpuBackend,
130        parts: &[(&QuantizedWeight, usize)],
131        k: usize,
132    ) -> Result<QuantizedWeight> {
133        Self::transpose_concat_for_gemm_gs(gpu, parts, k, 16)
134    }
135
136    /// 2026-09-25: `transpose_concat_for_gemm_gs` with the output row stride
137    /// padded to `align_up(n_total, align)` and the pad columns zeroed.
138    ///
139    /// Row `r` of the result starts at byte `r * stride`, and the tile GEMM reads
140    /// B with 16-byte `cp.async`, which needs 16-byte aligned rows; an N such as
141    /// an odd vocab size is not. Returns `(weight, stride)`; pass the stride to
142    /// `w4a16_gemm_n128_ldb`.
143    pub fn transpose_concat_for_gemm_padded(
144        gpu: &dyn GpuBackend,
145        parts: &[(&QuantizedWeight, usize)],
146        k: usize,
147        group_size: usize,
148        align: usize,
149    ) -> Result<(QuantizedWeight, usize)> {
150        let n_total: usize = parts.iter().map(|(_, n)| *n).sum();
151        let stride = n_total.div_ceil(align) * align;
152        Self::transpose_impl(gpu, parts, k, group_size, stride).map(|w| (w, stride))
153    }
154
155    /// 2026-09-25: `transpose_concat_for_gemm` with an explicit scale block size.
156    pub fn transpose_concat_for_gemm_gs(
157        gpu: &dyn GpuBackend,
158        parts: &[(&QuantizedWeight, usize)],
159        k: usize,
160        group_size: usize,
161    ) -> Result<QuantizedWeight> {
162        let n_total: usize = parts.iter().map(|(_, n)| *n).sum();
163        Self::transpose_impl(gpu, parts, k, group_size, n_total)
164    }
165
166    /// 2026-09-25: The implementation behind both concat builders. `stride` is the
167    /// row pitch of the output and must be at least `n_total` (a debug assertion
168    /// only); columns `n_total..stride` are zero.
169    fn transpose_impl(
170        gpu: &dyn GpuBackend,
171        parts: &[(&QuantizedWeight, usize)],
172        k: usize,
173        group_size: usize,
174        stride: usize,
175    ) -> Result<QuantizedWeight> {
176        let first = parts
177            .first()
178            .map(|(w, _)| *w)
179            .context("transpose_concat_for_gemm: empty parts")?;
180        let half_k = k / 2;
181        let num_groups = k / group_size;
182        let n_total: usize = parts.iter().map(|(_, n)| *n).sum();
183        debug_assert!(
184            stride >= n_total,
185            "transpose_impl: stride {stride} < n_total {n_total}"
186        );
187
188        // 2026-09-25: Per part, one `transpose_u8` into a contiguous `[half_k, n]`
189        // temp, then one pitched copy into the part's column window. The pad
190        // columns are zeroed up front, as the host path's staging vector is.
191        if let Some(tk) = Self::host_transpose_kernel(gpu) {
192            let new_weight = gpu.alloc(stride * half_k)?;
193            let new_scale = gpu.alloc(stride * num_groups)?;
194            if stride > n_total {
195                gpu.memset(new_weight, 0, stride * half_k)?;
196                gpu.memset(new_scale, 0, stride * num_groups)?;
197            }
198            let mut temps: Vec<DevicePtr> = Vec::with_capacity(parts.len() * 2);
199            let mut n_off = 0usize;
200            for (w, n) in parts {
201                let n = *n;
202                let t_w = gpu.alloc(n * half_k)?;
203                crate::layers::ops::transpose_u8(
204                    gpu,
205                    tk,
206                    w.weight,
207                    t_w,
208                    n as u32,
209                    half_k as u32,
210                    0,
211                )?;
212                gpu.copy_d2d_2d_async(t_w, n, new_weight.offset(n_off), stride, n, half_k, 0)?;
213                let t_s = gpu.alloc(n * num_groups)?;
214                crate::layers::ops::transpose_u8(
215                    gpu,
216                    tk,
217                    w.weight_scale,
218                    t_s,
219                    n as u32,
220                    num_groups as u32,
221                    0,
222                )?;
223                gpu.copy_d2d_2d_async(t_s, n, new_scale.offset(n_off), stride, n, num_groups, 0)?;
224                temps.push(t_w);
225                temps.push(t_s);
226                n_off += n;
227            }
228            gpu.synchronize(0)?;
229            for t in temps {
230                gpu.free(t)?;
231            }
232            return Ok(QuantizedWeight {
233                weight: new_weight,
234                weight_scale: new_scale,
235                weight_scale_2: first.weight_scale_2,
236                input_scale: first.input_scale,
237                weight_scale_2_vec: first.weight_scale_2_vec,
238                act: metrale_config::Nvfp4Act::combine(parts.iter().map(|(w, _)| w.act)),
239            });
240        }
241
242        let mut t_buf = vec![0u8; stride * half_k];
243        let mut st_buf = vec![0u8; stride * num_groups];
244        let mut n_off = 0usize;
245        for (w, n) in parts {
246            let n = *n;
247            let mut buf = vec![0u8; n * half_k];
248            gpu.copy_d2h(w.weight, &mut buf)?;
249            for i in 0..n {
250                for j in 0..half_k {
251                    t_buf[j * stride + n_off + i] = buf[i * half_k + j];
252                }
253            }
254            let mut sbuf = vec![0u8; n * num_groups];
255            gpu.copy_d2h(w.weight_scale, &mut sbuf)?;
256            for i in 0..n {
257                for j in 0..num_groups {
258                    st_buf[j * stride + n_off + i] = sbuf[i * num_groups + j];
259                }
260            }
261            n_off += n;
262        }
263
264        let new_weight = gpu.alloc(t_buf.len())?;
265        gpu.copy_h2d(&t_buf, new_weight)?;
266        let new_scale = gpu.alloc(st_buf.len())?;
267        gpu.copy_h2d(&st_buf, new_scale)?;
268
269        Ok(QuantizedWeight {
270            weight: new_weight,
271            weight_scale: new_scale,
272            weight_scale_2: first.weight_scale_2,
273            input_scale: first.input_scale,
274            weight_scale_2_vec: first.weight_scale_2_vec,
275            act: metrale_config::Nvfp4Act::combine(parts.iter().map(|(w, _)| w.act)),
276        })
277    }
278}
metrale_model_layers/weight_map/quantized/transpose.rs

metrale_model_layers/weight_map/quantized/
transpose.rs
