use crate::registry::{
BackwardKind, FusedOpEntry, FusedOpFamily, FusedOpParams, FusedOps, PatternMatch,
SubgraphPattern, decompose_via_recipe,
};
use crate::{Graph, NodeId, QuantType};
use fuel_ir::{DType, Shape};
use fuel_kernel_seam_types::{OpAttrs, OpTag, PatternNode};
pub fn entry() -> FusedOpEntry {
FusedOpEntry {
destructive_input: None,
id: FusedOps::QMATMUL,
name: "QMatMul",
family: FusedOpFamily::Quantized,
pattern: SubgraphPattern::Callable(canonical_pattern),
decompose,
backward: BackwardKind::NotDifferentiable,
shape_rule,
dtype_rule,
output_views: None,
}
}
fn shape_rule(input_shapes: &[Shape], params: &FusedOpParams) -> Shape {
debug_assert_eq!(
input_shapes.len(),
2,
"QMatMul takes 2 inputs (a, w_q_bytes)",
);
let n = match params {
FusedOpParams::QMatMul { n, .. } => *n,
_ => panic!("qmatmul::shape_rule got non-QMatMul params: {params:?}"),
};
let a_dims = input_shapes[0].dims();
let rank = a_dims.len();
debug_assert!(rank >= 2, "QMatMul activations must be rank ≥ 2");
let mut out_dims: Vec<usize> = a_dims[..rank - 1].to_vec();
out_dims.push(n);
Shape::from_dims(&out_dims)
}
fn dtype_rule(input_dtypes: &[DType], _params: &FusedOpParams) -> DType {
debug_assert_eq!(input_dtypes.len(), 2, "QMatMul takes 2 inputs",);
input_dtypes[0]
}
const QK4_0: usize = 32; const Q4_0_BYTES_PER_BLOCK: usize = 18;
fn recipe_q4_0(a_shape: &Shape, w_len_u32: usize, k: usize, n: usize) -> Option<PatternNode> {
use OpTag as T;
if k == 0 || n == 0 || !k.is_multiple_of(QK4_0) {
return None;
}
let a_dims = a_shape.dims();
if a_dims.len() < 2 || *a_dims.last()? != k {
return None;
}
let nb_per_row = k / QK4_0;
let total_blocks = n.checked_mul(nb_per_row)?;
let total_bytes = total_blocks.checked_mul(Q4_0_BYTES_PER_BLOCK)?;
if w_len_u32.checked_mul(4)? != total_bytes {
return None;
}
let four_l = total_bytes; let m_prime: usize = a_dims[..a_dims.len() - 1].iter().product();
let op = |op, attrs, operands| PatternNode::Op {
op,
attrs,
operands,
};
let bind = |i: u8| PatternNode::Bind { index: i };
let shape_attr = |dims: &[usize]| OpAttrs {
target_shape: dims.iter().map(|&d| d as i64).collect(),
..OpAttrs::default()
};
let cast_attr = |dt: DType| OpAttrs {
cast_dtype: Some(dt.as_str().to_string()),
..OpAttrs::default()
};
let sc = |v: f64| OpAttrs {
scalars: vec![v],
..OpAttrs::default()
};
let unsq = |d: u8| OpAttrs {
dims: vec![d],
..OpAttrs::default()
};
let cat = |ax: i64| OpAttrs {
axis: Some(ax),
..OpAttrs::default()
};
let slice = |ax: i64, start: u64, len: u64| OpAttrs {
axis: Some(ax),
slice_start: Some(start),
slice_len: Some(len),
..OpAttrs::default()
};
let mul_s = |v: f64, x: PatternNode| op(T::MulScalar, sc(v), vec![x]);
let add_s = |v: f64, x: PatternNode| op(T::AddScalar, sc(v), vec![x]);
let floor = |x: PatternNode| op(T::Floor, OpAttrs::default(), vec![x]);
let sub = |a: PatternNode, b: PatternNode| op(T::Sub, OpAttrs::default(), vec![a, b]);
let mul = |a: PatternNode, b: PatternNode| op(T::Mul, OpAttrs::default(), vec![a, b]);
let add = |a: PatternNode, b: PatternNode| op(T::Add, OpAttrs::default(), vec![a, b]);
let wf = op(T::Cast, cast_attr(DType::F64), vec![bind(1)]); let f1 = floor(mul_s(1.0 / 256.0, wf.clone()));
let b0 = sub(wf, mul_s(256.0, f1.clone()));
let f2 = floor(mul_s(1.0 / 256.0, f1.clone()));
let b1 = sub(f1, mul_s(256.0, f2.clone()));
let f3 = floor(mul_s(1.0 / 256.0, f2.clone()));
let b2 = sub(f2, mul_s(256.0, f3.clone()));
let b3 = f3; let lane = |x: PatternNode| op(T::Unsqueeze, unsq(1), vec![x]); let stacked = op(
T::Concat,
cat(1),
vec![lane(b0), lane(b1), lane(b2), lane(b3)],
); let bytes_f64 = op(T::Reshape, shape_attr(&[four_l]), vec![stacked]); let bytes = op(T::Cast, cast_attr(DType::F32), vec![bytes_f64]); let blocks = op(
T::Reshape,
shape_attr(&[total_blocks, Q4_0_BYTES_PER_BLOCK]),
vec![bytes],
);
let sb0 = op(T::Slice, slice(1, 0, 1), vec![blocks.clone()]); let sb1 = op(T::Slice, slice(1, 1, 1), vec![blocks.clone()]); let bits = add(sb0, mul_s(256.0, sb1)); let sign = floor(mul_s(1.0 / 32768.0, bits.clone()));
let rem = sub(bits, mul_s(32768.0, sign.clone()));
let exp = floor(mul_s(1.0 / 1024.0, rem.clone())); let mant = sub(rem, mul_s(1024.0, exp.clone())); let signed = add_s(1.0, mul_s(-2.0, sign)); let h0 = floor(mul_s(0.5, exp.clone()));
let bit0 = sub(exp.clone(), mul_s(2.0, h0.clone()));
let h1 = floor(mul_s(0.5, h0.clone()));
let bit1 = sub(h0, mul_s(2.0, h1.clone()));
let h2 = floor(mul_s(0.5, h1.clone()));
let bit2 = sub(h1, mul_s(2.0, h2.clone()));
let h3 = floor(mul_s(0.5, h2.clone()));
let bit3 = sub(h2, mul_s(2.0, h3.clone()));
let h4 = floor(mul_s(0.5, h3.clone()));
let bit4 = sub(h3, mul_s(2.0, h4));
let fac0 = add_s(1.0, bit0); let fac1 = add_s(1.0, mul_s(3.0, bit1)); let fac2 = add_s(1.0, mul_s(15.0, bit2)); let fac3 = add_s(1.0, mul_s(255.0, bit3)); let fac4 = add_s(1.0, mul_s(65535.0, bit4)); let pow2exp = mul(mul(mul(mul(fac0, fac1), fac2), fac3), fac4); let normal_mag = mul_s(2f64.powi(-25), mul(add_s(1024.0, mant.clone()), pow2exp));
let sub_mag = mul_s(2f64.powi(-24), mant);
let is_sub = op(
T::Relu,
OpAttrs::default(),
vec![add_s(1.0, mul_s(-1.0, exp))],
);
let one_minus_sub = add_s(1.0, mul_s(-1.0, is_sub.clone()));
let mag = add(mul(is_sub, sub_mag), mul(one_minus_sub, normal_mag));
let scale = mul(signed, mag);
let qbytes = op(T::Slice, slice(1, 2, 16), vec![blocks]); let high = floor(mul_s(1.0 / 16.0, qbytes.clone())); let low = sub(qbytes, mul_s(16.0, high.clone())); let x0 = add_s(-8.0, low); let x1 = add_s(-8.0, high); let codes = op(T::Concat, cat(1), vec![x0, x1]); let scale_bc = op(
T::BroadcastTo,
shape_attr(&[total_blocks, QK4_0]),
vec![scale],
); let dequant_blocks = mul(codes, scale_bc); let weight_nk = op(T::Reshape, shape_attr(&[n, k]), vec![dequant_blocks]); let weight_kn = op(T::Transpose, OpAttrs::default(), vec![weight_nk]);
let a2 = op(T::Reshape, shape_attr(&[m_prime, k]), vec![bind(0)]); let out2 = op(T::MatMul, OpAttrs::default(), vec![a2, weight_kn]); let mut out_dims: Vec<usize> = a_dims[..a_dims.len() - 1].to_vec();
out_dims.push(n);
Some(op(T::Reshape, shape_attr(&out_dims), vec![out2]))
}
pub fn decompose(graph: &mut Graph, id: NodeId, params: &FusedOpParams) -> NodeId {
let (quant_type, k, n) = match params {
FusedOpParams::QMatMul { quant_type, k, n } => (*quant_type, *k, *n),
_ => return id,
};
let (a_shape, w_len) = {
let node = graph.node(id);
if node.inputs.len() != 2 {
return id;
}
let a_shape = graph.node(node.inputs[0]).shape.clone();
let w_len = graph.node(node.inputs[1]).shape.elem_count();
(a_shape, w_len)
};
let recipe_node = match quant_type {
QuantType::Q4_0 => recipe_q4_0(&a_shape, w_len, k, n),
_ => None,
};
let recipe_node = match recipe_node {
Some(r) => r,
None => return id, };
decompose_via_recipe(graph, id, &recipe_node, Some(Vec::new()))
}
pub fn canonical_pattern(_graph: &Graph, _root: NodeId) -> Option<PatternMatch> {
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Node, Op};
fn decode_f16_bits_arith(bits: f32) -> f32 {
let sign = (bits / 32768.0).floor();
let rem = bits - 32768.0 * sign;
let exp = (rem / 1024.0).floor();
let mant = rem - 1024.0 * exp;
let signed = 1.0 - 2.0 * sign;
let h0 = (exp * 0.5).floor();
let bit0 = exp - 2.0 * h0;
let h1 = (h0 * 0.5).floor();
let bit1 = h0 - 2.0 * h1;
let h2 = (h1 * 0.5).floor();
let bit2 = h1 - 2.0 * h2;
let h3 = (h2 * 0.5).floor();
let bit3 = h2 - 2.0 * h3;
let h4 = (h3 * 0.5).floor();
let bit4 = h3 - 2.0 * h4;
let pow2exp = (1.0 + bit0)
* (1.0 + 3.0 * bit1)
* (1.0 + 15.0 * bit2)
* (1.0 + 255.0 * bit3)
* (1.0 + 65535.0 * bit4);
let normal_mag = (1024.0 + mant) * pow2exp * 2f32.powi(-25);
let sub_mag = mant * 2f32.powi(-24);
let is_sub = (1.0 - exp).max(0.0);
let mag = is_sub * sub_mag + (1.0 - is_sub) * normal_mag;
signed * mag
}
#[test]
fn f16_decode_arithmetic_matches_half_for_all_finite_bit_patterns() {
for b in 0u16..=u16::MAX {
let exp = (b >> 10) & 0x1F;
if exp == 0x1F {
continue; }
let want = half::f16::from_bits(b).to_f32();
let got = decode_f16_bits_arith(b as f32);
assert_eq!(
got.to_bits(),
want.to_bits(),
"f16 bits {b:#06x}: got {got}, want {want}",
);
}
}
fn fused_q4_0_node(g: &mut Graph, leading: &[usize], k: usize, n: usize) -> NodeId {
let mut a_dims = leading.to_vec();
a_dims.push(k);
let act = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&a_dims),
dtype: DType::F32,
});
let total_blocks = n * (k / QK4_0);
let total_bytes = total_blocks * Q4_0_BYTES_PER_BLOCK;
let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[total_bytes / 4]),
dtype: DType::U32,
});
let mut out_dims = leading.to_vec();
out_dims.push(n);
g.push(Node {
op: Op::Fused(
FusedOps::QMATMUL,
FusedOpParams::QMatMul {
quant_type: QuantType::Q4_0,
k,
n,
},
),
inputs: vec![act, w],
shape: Shape::from_dims(&out_dims),
dtype: DType::F32,
})
}
#[test]
fn qmatmul_q4_0_decompose_fires_and_is_shape_correct() {
for leading in [vec![1usize], vec![4usize], vec![2usize, 3]] {
let k = 64; let n = 5;
let params = FusedOpParams::QMatMul {
quant_type: QuantType::Q4_0,
k,
n,
};
let mut g = Graph::new();
let fused = fused_q4_0_node(&mut g, &leading, k, n);
let out_sh = g.node(fused).shape.clone();
let new_root = decompose(&mut g, fused, ¶ms);
assert_ne!(
new_root, fused,
"Q4_0 recipe decompose fires (leading={leading:?})"
);
assert_eq!(
g.node(new_root).shape,
out_sh,
"output shape matches shape_rule"
);
assert_eq!(g.node(new_root).dtype, DType::F32, "output dtype is F32");
}
}
#[test]
fn qmatmul_nonq4_0_format_is_a_surfaced_gap_fixpoint() {
let mut g = Graph::new();
let k = 64;
let n = 3;
let act = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[2, k]),
dtype: DType::F32,
});
let total_bytes = n * (k / 32) * 34; let w = g.push(Node {
op: Op::Const,
inputs: vec![],
shape: Shape::from_dims(&[total_bytes / 4]),
dtype: DType::U32,
});
let fused = g.push(Node {
op: Op::Fused(
FusedOps::QMATMUL,
FusedOpParams::QMatMul {
quant_type: QuantType::Q8_0,
k,
n,
},
),
inputs: vec![act, w],
shape: Shape::from_dims(&[2, n]),
dtype: DType::F32,
});
let before = g.len();
let out = decompose(
&mut g,
fused,
&FusedOpParams::QMatMul {
quant_type: QuantType::Q8_0,
k,
n,
},
);
assert_eq!(out, fused, "non-Q4_0 => surfaced gap => fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
#[test]
fn qmatmul_wrong_params_is_a_fixpoint_not_a_crash() {
let mut g = Graph::new();
let fused = fused_q4_0_node(&mut g, &[2], 64, 5);
let before = g.len();
let out = decompose(&mut g, fused, &FusedOpParams::Rope);
assert_eq!(out, fused, "wrong params => typed decline => fixpoint");
assert_eq!(g.len(), before, "declined before any emission");
}
}