(self, P: MLXProgramBuilder, n: Node)
| 884 | ) |
| 885 | |
| 886 | def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: |
| 887 | assert n == self.head |
| 888 | |
| 889 | x_node, w_node = n.args[0:2] |
| 890 | b_node = n.args[2] if len(n.args) > 2 else None |
| 891 | |
| 892 | qdata_target, qdata = P.get_placeholder_target_and_tensor(self.qdata) |
| 893 | zero_point_target, zero_point = P.get_placeholder_target_and_tensor( |
| 894 | self.zero_point |
| 895 | ) |
| 896 | _, scale = P.get_placeholder_target_and_tensor(self.scale) |
| 897 | |
| 898 | x_slot, scale_slot, b_slot = P.slot_map([x_node, self.scale, b_node]) |
| 899 | |
| 900 | Q, B = to_mlx_qparams(qdata, scale, zero_point, self.bits) |
| 901 | w = P.make_or_get_constant(f"{qdata_target}_to_packed", Q) |
| 902 | biases = emit_quantized_biases( |
| 903 | P, zero_point_target, scale, zero_point, self.bits, B, scale_slot |
| 904 | ) |
| 905 | |
| 906 | out = P.make_or_get_slot(n) |
| 907 | has_bias = b_node is not None |
| 908 | x_dtype = x_node.meta["val"].dtype |
| 909 | needs_cast = self.out_dtype != x_dtype |
| 910 | |
| 911 | P.emit( |
| 912 | QuantizedMatmulNode( |
| 913 | x=P.slot_to_tid(x_slot), |
| 914 | w=P.slot_to_tid(w), |
| 915 | scales=P.slot_to_tid(scale_slot), |
| 916 | out=P.slot_to_tid(out), |
| 917 | biases=P.slot_to_tid(biases), |
| 918 | group_size=self.group_size, |
| 919 | bits=self.bits, |
| 920 | mode="affine", |
| 921 | transpose=True, |
| 922 | ) |
| 923 | ) |
| 924 | |
| 925 | if has_bias: |
| 926 | P.emit( |
| 927 | AddNode( |
| 928 | a=P.slot_to_tid(out), |
| 929 | b=P.slot_to_tid(b_slot), |
| 930 | out=P.slot_to_tid(out), |
| 931 | ) |
| 932 | ) |
| 933 | |
| 934 | if needs_cast: |
| 935 | P.emit( |
| 936 | AsTypeNode( |
| 937 | x=P.slot_to_tid(out), |
| 938 | out=P.slot_to_tid(out), |
| 939 | scalar_type=torch_dtype_to_scalar_type(self.out_dtype), |
| 940 | ) |
| 941 | ) |
| 942 | |
| 943 | return out |
nothing calls this directly
no test coverage detected