fuel_quantized/
awq_repack.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Host repack of a HuggingFace `*-AWQ` (GEMM-version) checkpoint into the
3//! buffers baracuda's W4A16 AWQ kernel reads (ROADMAP row 31).
4//!
5//! Pure host code: no CUDA, no baracuda dependency. The kernel layout is the one
6//! documented on baracuda's `Int4AwqGemmArgs` (baracuda#219, alpha.147), which its
7//! `tests/awq_reference.rs` verifies on a device against a dequantise-then-GEMM
8//! reference. Everything is OC-major, with `groups = IC / group_size` and
9//! `mult = 128 / group_size`:
10//!
11//! | buffer    | kernel layout                                                        |
12//! |-----------|----------------------------------------------------------------------|
13//! | `weight`  | `[OC, IC/8]` `u32`; nibble `ic % 8` of word `ic / 8` in row `oc`     |
14//! | `zeros`   | `[OC, zeros_w]` `u32`; nibble `g % 8` of word `g / 8` in row `oc`    |
15//! | `scales`  | `[OC, 8 * zeros_w]` `f16`; entry `g` of row `oc`                     |
16//!
17//! with `zeros_w = ceil(ceil(groups / 8) / mult) * mult`. The tail of each
18//! `zeros`/`scales` row past `groups` is padding the kernel never reads as data;
19//! it is written as zero.
20//!
21//! The HF checkpoint is the transpose and uses AutoAWQ's interleaved nibble order
22//! along `OC` ([`HF_AWQ_PACK_ORDER`]):
23//!
24//! | tensor    | HF layout                                                            |
25//! |-----------|----------------------------------------------------------------------|
26//! | `qweight` | `[IC, OC/8]` `u32`                                                   |
27//! | `qzeros`  | `[IC/group_size, OC/8]` `u32`                                        |
28//! | `scales`  | `[IC/group_size, OC]` `f16`                                          |
29//!
30//! Each HF word `col` holds logical columns `col*8 + HF_AWQ_PACK_ORDER[i]` at
31//! nibble `i`.
32
33use half::f16;
34
35/// AutoAWQ's `order_map` (`awq/modules/linear/gemm.py`, `WQLinear_GEMM.from_linear`):
36/// nibble position `i` of an HF-packed word holds the value of logical column
37/// `HF_AWQ_PACK_ORDER[i]` within its group of 8. The same order applies to `qweight`
38/// and `qzeros`. This is the HuggingFace checkpoint's order only: the op
39/// (`fuel_graph::registry::awq_matmul`) takes the sequential kernel layout this
40/// module produces.
41pub const HF_AWQ_PACK_ORDER: [usize; 8] = [0, 2, 4, 6, 1, 3, 5, 7];
42
43/// The kernel-layout buffers for one linear layer.
44#[derive(Clone, Debug, PartialEq)]
45pub struct AwqKernelBuffers {
46    /// `[OC, IC/8]` packed int4 weights.
47    pub weight: Vec<u32>,
48    /// `[OC, zeros_w]` packed int4 zero-points.
49    pub zeros: Vec<u32>,
50    /// `[OC, 8 * zeros_w]` scales.
51    pub scales: Vec<f16>,
52    /// Words per `zeros` row (`scales` rows hold `8 * zeros_w`).
53    pub zeros_w: usize,
54}
55
56/// Why a repack was refused.
57#[derive(Clone, Debug, PartialEq, Eq)]
58pub enum AwqRepackError {
59    /// `group_size` is not 64 or 128.
60    BadGroupSize(usize),
61    /// A dimension is not a multiple of what the packing requires.
62    BadDim {
63        what: &'static str,
64        value: usize,
65        multiple_of: usize,
66    },
67    /// An input slice has the wrong length.
68    BadLen {
69        what: &'static str,
70        got: usize,
71        want: usize,
72    },
73}
74
75impl std::fmt::Display for AwqRepackError {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        match self {
78            Self::BadGroupSize(g) => write!(f, "AWQ group size must be 64 or 128, got {g}"),
79            Self::BadDim {
80                what,
81                value,
82                multiple_of,
83            } => {
84                write!(
85                    f,
86                    "AWQ {what} = {value} must be a multiple of {multiple_of}"
87                )
88            }
89            Self::BadLen { what, got, want } => {
90                write!(f, "AWQ {what} has {got} elements, expected {want}")
91            }
92        }
93    }
94}
95
96impl std::error::Error for AwqRepackError {}
97
98/// Words per `zeros` row for `ic` input features at `group_size` (the kernel pads
99/// each row): `ceil(ceil(groups / 8) / mult) * mult`.
100pub fn awq_zeros_w(ic: usize, group_size: usize) -> Result<usize, AwqRepackError> {
101    if group_size != 64 && group_size != 128 {
102        return Err(AwqRepackError::BadGroupSize(group_size));
103    }
104    if !ic.is_multiple_of(group_size) {
105        return Err(AwqRepackError::BadDim {
106            what: "IC",
107            value: ic,
108            multiple_of: group_size,
109        });
110    }
111    let groups = ic / group_size;
112    let mult = 128 / group_size;
113    Ok(groups.div_ceil(8).div_ceil(mult) * mult)
114}
115
116fn check_len(what: &'static str, got: usize, want: usize) -> Result<(), AwqRepackError> {
117    if got == want {
118        Ok(())
119    } else {
120        Err(AwqRepackError::BadLen { what, got, want })
121    }
122}
123
124fn check_multiple(
125    what: &'static str,
126    value: usize,
127    multiple_of: usize,
128) -> Result<(), AwqRepackError> {
129    if value.is_multiple_of(multiple_of) {
130        Ok(())
131    } else {
132        Err(AwqRepackError::BadDim {
133            what,
134            value,
135            multiple_of,
136        })
137    }
138}
139
140/// Unpack one HF-packed row-major `[rows, oc/8]` tensor into logical `[rows, oc]`
141/// 4-bit codes.
142fn unpack_hf(words: &[u32], rows: usize, oc: usize) -> Vec<u8> {
143    let wpr = oc / 8;
144    let mut out = vec![0u8; rows * oc];
145    for r in 0..rows {
146        for col in 0..wpr {
147            let w = words[r * wpr + col];
148            for (i, &logical) in HF_AWQ_PACK_ORDER.iter().enumerate() {
149                out[r * oc + col * 8 + logical] = ((w >> (4 * i)) & 0xF) as u8;
150            }
151        }
152    }
153    out
154}
155
156/// Repack an HF `*-AWQ` layer (`qweight [IC, OC/8]`, `qzeros [IC/g, OC/8]`,
157/// `scales [IC/g, OC]`) into the kernel layout described in the module docs.
158///
159/// Checks only what the packing itself needs (`IC % 8`, `OC % 8`, `IC % group_size`,
160/// slice lengths). The kernel's own constraints (`OC % 64`, `IC % (32 * split_k)`,
161/// sm_80+) are baracuda's `can_implement`'s to enforce.
162pub fn repack_hf_awq_for_kernel(
163    qweight: &[u32],
164    qzeros: &[u32],
165    scales: &[f16],
166    ic: usize,
167    oc: usize,
168    group_size: usize,
169) -> Result<AwqKernelBuffers, AwqRepackError> {
170    let zeros_w = awq_zeros_w(ic, group_size)?;
171    check_multiple("IC", ic, 8)?;
172    check_multiple("OC", oc, 8)?;
173    let groups = ic / group_size;
174    check_len("qweight", qweight.len(), ic * (oc / 8))?;
175    check_len("qzeros", qzeros.len(), groups * (oc / 8))?;
176    check_len("scales", scales.len(), groups * oc)?;
177
178    let q = unpack_hf(qweight, ic, oc); // [IC, OC]
179    let z = unpack_hf(qzeros, groups, oc); // [groups, OC]
180    Ok(AwqKernelBuffers {
181        weight: pack_oc_major(&q, ic, oc, ic / 8),
182        zeros: pack_oc_major(&z, groups, oc, zeros_w),
183        scales: transpose_scales(scales, groups, oc, zeros_w * 8),
184        zeros_w,
185    })
186}
187
188/// Transpose logical `[rows, oc]` 4-bit codes to OC-major and pack 8 per `u32` along
189/// the row axis: output is `[oc, words_per_oc]`, word `r / 8` nibble `r % 8`.
190fn pack_oc_major(codes: &[u8], rows: usize, oc: usize, words_per_oc: usize) -> Vec<u32> {
191    let mut out = vec![0u32; oc * words_per_oc];
192    for o in 0..oc {
193        for r in 0..rows {
194            out[o * words_per_oc + r / 8] |= u32::from(codes[r * oc + o]) << (4 * (r % 8));
195        }
196    }
197    out
198}
199
200/// Transpose `[groups, oc]` scales to OC-major `[oc, sf_w]`, zero-padded past `groups`.
201fn transpose_scales(scales: &[f16], groups: usize, oc: usize, sf_w: usize) -> Vec<f16> {
202    let mut out = vec![f16::ZERO; oc * sf_w];
203    for o in 0..oc {
204        for g in 0..groups {
205            out[o * sf_w + g] = scales[g * oc + o];
206        }
207    }
208    out
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    struct Rng(u64);
216    impl Rng {
217        fn next(&mut self) -> u64 {
218            self.0 ^= self.0 << 13;
219            self.0 ^= self.0 >> 7;
220            self.0 ^= self.0 << 17;
221            self.0
222        }
223    }
224
225    /// Logical codes for one layer: `q [IC, OC]`, `z [groups, OC]`, `s [groups, OC]`.
226    struct Logical {
227        q: Vec<u8>,
228        z: Vec<u8>,
229        s: Vec<f16>,
230    }
231
232    fn logical(ic: usize, oc: usize, group: usize) -> Logical {
233        let mut r = Rng(0xD1B5_4A32_D192_ED03 ^ ((ic * 31 + oc) * 31 + group) as u64);
234        let groups = ic / group;
235        Logical {
236            q: (0..ic * oc).map(|_| (r.next() & 0xF) as u8).collect(),
237            z: (0..groups * oc).map(|_| (r.next() & 0xF) as u8).collect(),
238            // Distinct, non-zero values so a misplaced scale is visible.
239            s: (0..groups * oc)
240                .map(|i| f16::from_f32(0.01 + i as f32 * 0.001))
241                .collect(),
242        }
243    }
244
245    /// Pack `[rows, oc]` codes the way AutoAWQ's `from_linear` does (the writer
246    /// side, written independently of `unpack_hf`): word `col` of a row gets the
247    /// value of logical column `col*8 + order_map[i]` at nibble `i`.
248    fn pack_hf(codes: &[u8], rows: usize, oc: usize, order: [usize; 8]) -> Vec<u32> {
249        let mut out = vec![0u32; rows * oc / 8];
250        for r in 0..rows {
251            for col in 0..oc / 8 {
252                for i in 0..8 {
253                    out[r * (oc / 8) + col] |=
254                        u32::from(codes[r * oc + col * 8 + order[i]]) << (i * 4);
255                }
256            }
257        }
258        out
259    }
260
261    /// Read the kernel buffers back with baracuda's documented read formula and
262    /// return the logical codes it would see.
263    fn read_kernel(b: &AwqKernelBuffers, ic: usize, oc: usize, group: usize) -> Logical {
264        let groups = ic / group;
265        let sf_w = b.zeros_w * 8;
266        let mut q = vec![0u8; ic * oc];
267        let mut z = vec![0u8; groups * oc];
268        let mut s = vec![f16::ZERO; groups * oc];
269        for o in 0..oc {
270            for i in 0..ic {
271                q[i * oc + o] = ((b.weight[o * (ic / 8) + i / 8] >> (4 * (i % 8))) & 0xF) as u8;
272            }
273            for g in 0..groups {
274                z[g * oc + o] = ((b.zeros[o * b.zeros_w + g / 8] >> (4 * (g % 8))) & 0xF) as u8;
275                s[g * oc + o] = b.scales[o * sf_w + g];
276            }
277        }
278        Logical { q, z, s }
279    }
280
281    /// `(ic, oc, group)`: group 128 with `groups % 8 != 0` (padded rows), group 64
282    /// with an odd `ceil(groups/8)` (padded to `mult`), and a tidy case.
283    const SHAPES: [(usize, usize, usize); 4] = [
284        (128 * 9, 64, 128),
285        (64 * 8, 64, 64),
286        (128 * 16, 64, 128),
287        (64 * 24, 128, 64),
288    ];
289
290    #[test]
291    fn repack_round_trips_through_the_documented_kernel_read_formula() {
292        for (ic, oc, group) in SHAPES {
293            let l = logical(ic, oc, group);
294            let groups = ic / group;
295            let qw = pack_hf(&l.q, ic, oc, HF_AWQ_PACK_ORDER);
296            let qz = pack_hf(&l.z, groups, oc, HF_AWQ_PACK_ORDER);
297            let b = repack_hf_awq_for_kernel(&qw, &qz, &l.s, ic, oc, group).unwrap();
298            assert_eq!(b.weight.len(), oc * ic / 8, "{ic}x{oc} g{group}");
299            assert_eq!(b.zeros.len(), oc * b.zeros_w);
300            assert_eq!(b.scales.len(), oc * b.zeros_w * 8);
301            let back = read_kernel(&b, ic, oc, group);
302            assert!(back.q == l.q, "weights {ic}x{oc} g{group}");
303            assert!(back.z == l.z, "zeros {ic}x{oc} g{group}");
304            assert!(back.s == l.s, "scales {ic}x{oc} g{group}");
305        }
306    }
307
308    #[test]
309    fn zeros_w_matches_baracuda_padding_formula_and_rows_are_padded_with_zero() {
310        // Pinned values from baracuda's `tests/awq_reference.rs` formula.
311        assert_eq!(awq_zeros_w(128 * 9, 128).unwrap(), 2); // groups 9 -> ceil(9/8)=2, mult 1
312        assert_eq!(awq_zeros_w(64 * 8, 64).unwrap(), 2); // groups 8 -> 1, mult 2 -> 2
313        assert_eq!(awq_zeros_w(128 * 16, 128).unwrap(), 2); // groups 16 -> 2
314        assert_eq!(awq_zeros_w(64 * 24, 64).unwrap(), 4); // groups 24 -> 3, mult 2 -> 4
315        let (ic, oc, group) = SHAPES[0];
316        let l = logical(ic, oc, group);
317        let qw = pack_hf(&l.q, ic, oc, HF_AWQ_PACK_ORDER);
318        let qz = pack_hf(&l.z, ic / group, oc, HF_AWQ_PACK_ORDER);
319        let b = repack_hf_awq_for_kernel(&qw, &qz, &l.s, ic, oc, group).unwrap();
320        let groups = ic / group;
321        let sf_w = b.zeros_w * 8;
322        for o in 0..oc {
323            assert!(
324                b.scales[o * sf_w + groups..(o + 1) * sf_w]
325                    .iter()
326                    .all(|x| *x == f16::ZERO)
327            );
328            // Nibbles of the last zeros word at or past `groups % 8` are padding.
329            let last = b.zeros[o * b.zeros_w + groups / 8];
330            assert_eq!(last >> (4 * (groups % 8)), 0, "row {o} padding nibbles");
331            assert!(
332                b.zeros[o * b.zeros_w + groups / 8 + 1..(o + 1) * b.zeros_w]
333                    .iter()
334                    .all(|w| *w == 0)
335            );
336        }
337    }
338
339    /// Sabotage sibling: the interleaved order is load-bearing. A checkpoint packed
340    /// in plain sequential nibble order must NOT repack to the right codes.
341    #[test]
342    fn a_sequentially_packed_checkpoint_does_not_round_trip() {
343        let (ic, oc, group) = SHAPES[0];
344        let l = logical(ic, oc, group);
345        let seq = [0, 1, 2, 3, 4, 5, 6, 7];
346        let qw = pack_hf(&l.q, ic, oc, seq);
347        let qz = pack_hf(&l.z, ic / group, oc, seq);
348        let b = repack_hf_awq_for_kernel(&qw, &qz, &l.s, ic, oc, group).unwrap();
349        let back = read_kernel(&b, ic, oc, group);
350        assert!(back.q != l.q, "the AWQ order must matter for qweight");
351        assert!(back.z != l.z, "the AWQ order must matter for qzeros");
352    }
353
354    #[test]
355    fn bad_inputs_are_refused() {
356        let ok = |ic, oc, g, nw, nz, ns| {
357            repack_hf_awq_for_kernel(&vec![0; nw], &vec![0; nz], &vec![f16::ZERO; ns], ic, oc, g)
358        };
359        assert_eq!(
360            ok(256, 64, 32, 0, 0, 0),
361            Err(AwqRepackError::BadGroupSize(32))
362        );
363        assert!(matches!(
364            ok(200, 64, 128, 0, 0, 0),
365            Err(AwqRepackError::BadDim { what: "IC", .. })
366        ));
367        assert!(matches!(
368            ok(256, 60, 128, 0, 0, 0),
369            Err(AwqRepackError::BadDim { what: "OC", .. })
370        ));
371        assert!(matches!(
372            ok(256, 64, 128, 1, 0, 0),
373            Err(AwqRepackError::BadLen {
374                what: "qweight",
375                ..
376            })
377        ));
378        // Control: the same call with right lengths succeeds.
379        assert!(ok(256, 64, 128, 256 * 8, 2 * 8, 2 * 64).is_ok());
380    }
381}
fuel_quantized/awq_repack.rs

fuel_quantized/
awq_repack.rs
