MCPcopy Create free account
hub / github.com/pytorch/executorch / __call__

Method __call__

backends/mlx/patterns.py:1110–1175  ·  view source on GitHub ↗
(self, P, n)

Source from the content-addressed store, hash-verified

1108 return cls(head, [dequant], qdata, scale, per_tensor_scale, output_dtype)
1109
1110 def __call__(self, P, n):
1111 assert n == self.head
1112
1113 x_node, w_node = n.args[0:2]
1114 b_node = n.args[2] if len(n.args) > 2 else None
1115
1116 needs_cast = x_node.meta["val"].dtype != self.output_dtype
1117 has_bias = b_node is not None
1118 has_per_tensor_scale = True
1119
1120 _, per_tensor_scale_value = P.get_placeholder_target_and_tensor(
1121 self.per_tensor_scale
1122 )
1123 from torch._subclasses.fake_tensor import FakeTensor
1124
1125 if not isinstance(per_tensor_scale_value, FakeTensor):
1126 if per_tensor_scale_value.item() == 1.0:
1127 has_per_tensor_scale = False
1128
1129 x, w, scales, bias, per_tensor_scale = P.slot_map(
1130 [x_node, self.qdata, self.scale, b_node, self.per_tensor_scale]
1131 )
1132
1133 out = P.make_or_get_slot(n)
1134 P.emit(
1135 QuantizedMatmulNode(
1136 x=P.slot_to_tid(x),
1137 w=P.slot_to_tid(w),
1138 scales=P.slot_to_tid(scales),
1139 out=P.slot_to_tid(out),
1140 biases=None,
1141 group_size=16,
1142 bits=4,
1143 mode="nvfp4",
1144 transpose=True,
1145 )
1146 )
1147
1148 if has_per_tensor_scale:
1149 P.emit(
1150 MultiplyNode(
1151 a=P.slot_to_tid(out),
1152 b=P.slot_to_tid(per_tensor_scale),
1153 out=P.slot_to_tid(out),
1154 )
1155 )
1156
1157 if has_bias:
1158 P.emit(
1159 AddNode(
1160 a=P.slot_to_tid(out),
1161 b=P.slot_to_tid(bias),
1162 out=P.slot_to_tid(out),
1163 )
1164 )
1165
1166 if needs_cast:
1167 P.emit(

Callers

nothing calls this directly

Calls 7

AddNodeClass · 0.85
slot_mapMethod · 0.80
emitMethod · 0.80
slot_to_tidMethod · 0.80
make_or_get_slotMethod · 0.45

Tested by

no test coverage detected