tract_core/ops/matmul/
de_block_quant.rs

1use tract_linalg::block_quant::{BlockQuant, BlockQuantFact, BlockQuantStorage, Q4_0};
2
3use crate::internal::*;
4use crate::ops::einsum::einsum_matmul::EinSumMatMul;
5use crate::ops::konst::Const;
6use crate::transform::ModelTransform;
7
8#[derive(Debug)]
9pub struct BlockQuantTransform;
10
11impl ModelTransform for BlockQuantTransform {
12    fn name(&self) -> StaticName {
13        "block_quant".into()
14    }
15
16    fn transform(&self, model: &mut TypedModel) -> TractResult<()> {
17        crate::ops::einsum::einsum_matmul::detect_all(model)?;
18        Rewriter::<()>::default()
19            .with_rule_for("block_quant_einsum_weights", block_quant_einsum_weights)
20            .rewrite(&(), model)?;
21        crate::ops::einsum::einsum_matmul::flatten_all(model)?;
22        Ok(())
23    }
24}
25
26fn block_quant_einsum_weights(
27    _ctx: &(),
28    model: &TypedModel,
29    node: &TypedNode,
30    prefix: &str,
31    op: &EinSumMatMul,
32) -> TractResult<Option<TypedModelPatch>> {
33    rule_if!(node.inputs.len() == 2);
34    for (slot, fact) in model.node_input_facts(node.id)?.iter().enumerate() {
35        let Some(a) = fact.konst.as_ref() else { continue };
36        if a.rank() != 2 {
37            continue;
38        };
39        if op.k_axis().inputs[slot][0] == 0 {
40            let mut patch = TypedModelPatch::default();
41            let mut taps = patch.taps(model, &node.inputs)?;
42            taps[slot] = patch.wire_node(
43                format!("{}.t_{}", node.name, slot),
44                AxisOp::Move(1, 0),
45                &[taps[slot]],
46            )?[0];
47            let mut new_op = op.clone();
48            new_op.op.axes = op
49                .op
50                .axes
51                .clone()
52                .remove_axis_occurency(InOut::In(slot), 0)?
53                .with_extra_axis_occurency(op.k_axis, InOut::In(slot), 1)?;
54            let output = patch.wire_node(prefix, new_op, &taps)?;
55            patch.shunt_outside(model, node.id.into(), output[0])?;
56            return Ok(Some(patch));
57        }
58        let format = Q4_0;
59        let mut patch = TypedModelPatch::default();
60        let weights = if a.datum_type() == f16::datum_type() {
61            format.quant_f16(a.try_as_plain()?.as_slice::<f16>()?)?
62        } else {
63            format.quant_f32(a.cast_to::<f32>()?.try_as_plain()?.as_slice::<f32>()?)?
64        };
65        let act_slot = 1 - slot;
66        let name = &model.node(node.inputs[slot].node).name;
67        let m = a.shape()[0];
68        let k = a.shape()[1];
69        let bqs = BlockQuantStorage::new(Box::new(format), m, k, Arc::new(weights))?;
70        let fact =
71            Box::new(BlockQuantFact::new(dyn_clone::clone_box(bqs.format()), tvec!(1, m, k)));
72        let weights = patch.wire_node(
73            format!("{name}.bq"),
74            Const::new_with_exotic_fact(
75                Arc::new(bqs.into_tensor_with_shape(a.datum_type(), &[1, m, k])),
76                fact,
77            )?,
78            &[],
79        )?;
80        let tap = patch.tap_model(model, node.inputs[act_slot])?;
81        // Block-quant tensor is rank 3 [G=1, M, K]; add a group dim to the weight's axes
82        let mut new_op = op.op.clone();
83        new_op.axes = new_op.axes.with_extra_axis('G', InOut::In(slot), 0)?;
84        let inputs = if slot == 0 { [weights[0], tap] } else { [tap, weights[0]] };
85        let wire = patch.wire_node(prefix, new_op, &inputs)?;
86        patch.shunt_outside(model, node.id.into(), wire[0])?;
87        return Ok(Some(patch));
88    }
89    Ok(None)
90}
91
92#[cfg(test)]
93mod test {
94    use super::*;
95    use crate::ops::einsum::EinSum;
96
97    // Deterministic varied fill so quantization is actually exercised.
98    fn fill(shape: &[usize], seed: usize) -> Tensor {
99        let n: usize = shape.iter().product();
100        let data: Vec<f32> =
101            (0..n).map(|i| (((i * 13 + seed * 7) % 29) as f32 - 14.0) / 14.0).collect();
102        Tensor::from_shape(shape, &data).unwrap()
103    }
104
105    fn build(axes: &str, x_shape: &[usize], w: &Tensor) -> TractResult<TypedModel> {
106        let mut model = TypedModel::default();
107        let x = model.add_source("x", f32::fact(x_shape))?;
108        let w = model.wire_node("w", Const::new(w.clone().into_arc_tensor())?, &[])?[0];
109        let out = model.wire_node(
110            "mm",
111            EinSum { axes: axes.parse()?, operating_dt: f32::datum_type(), q_params: None },
112            &[x, w],
113        )?;
114        model.select_output_outlets(&out)?;
115        model.into_decluttered()
116    }
117
118    fn eval(model: TypedModel, x: &Tensor) -> TractResult<Tensor> {
119        let out = model.into_runnable()?.run(tvec!(x.clone().into_tvalue()))?;
120        Ok(out[0].clone().into_tensor())
121    }
122
123    // `X @ W` with W a rank-2 const (`w_k_axis` = contraction axis within W).
124    // Asserts BlockQuantTransform both applies AND computes `X @ Q4_0(W)` correctly:
125    // reference uses the same Q4_0-dequantized W, isolating operand wiring from quant error.
126    fn check(axes: &str, x_shape: &[usize], w_shape: &[usize], w_k_axis: usize) -> TractResult<()> {
127        let x = fill(x_shape, 1);
128        let w = fill(w_shape, 2);
129
130        // Q4_0 blocks along the last axis, so dequantize with k moved last, then back.
131        let last = w.rank() - 1;
132        let w_deq = Q4_0
133            .simulate_precision_loss(w.clone().move_axis(w_k_axis, last)?, last)?
134            .move_axis(last, w_k_axis)?;
135        let reference = eval(build(axes, x_shape, &w_deq)?, &x)?;
136
137        let mut quant = build(axes, x_shape, &w)?;
138        BlockQuantTransform.transform(&mut quant)?;
139        let got = eval(quant, &x)?;
140
141        got.close_enough(&reference, Approximation::Approximate)
142    }
143
144    #[test]
145    fn block_quant_xw_rank2() -> TractResult<()> {
146        // plain 2D, ONNX X@W orientation: X[m,k] @ W[k,n], k is W axis 0
147        check("mk,kn->mn", &[7, 256], &[256, 256], 0)
148    }
149
150    #[test]
151    fn block_quant_xw_batched() -> TractResult<()> {
152        // BERT-style: X[b,m,k] @ W[k,n] -> [b,m,n]
153        check("bmk,kn->bmn", &[2, 7, 256], &[256, 256], 0)
154    }
155
156    #[test]
157    fn block_quant_weights_already_nk() -> TractResult<()> {
158        // canonical orientation (k inner on W): X[m,k] @ W[n,k], k is W axis 1
159        check("mk,nk->mn", &[7, 256], &[256, 256], 1)
160    }
161}
tract_core/ops/matmul/de_block_quant.rs

tract_core/ops/matmul/
de_block_quant.rs
