(self, P: MLXProgramBuilder, n: Node)
| 577 | return cls(head, [w], qdata, scale, per_tensor_scale, output_dtype) |
| 578 | |
| 579 | def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: |
| 580 | assert n == self.head |
| 581 | w_node, x_node = n.args[0:2] |
| 582 | |
| 583 | has_per_tensor_scale = True |
| 584 | _, per_tensor_scale_value = P.get_placeholder_target_and_tensor( |
| 585 | self.per_tensor_scale |
| 586 | ) |
| 587 | from torch._subclasses.fake_tensor import FakeTensor |
| 588 | |
| 589 | if not isinstance(per_tensor_scale_value, FakeTensor): |
| 590 | if per_tensor_scale_value.item() == 1.0: |
| 591 | has_per_tensor_scale = False |
| 592 | |
| 593 | x_dtype = x_node.meta["val"].dtype |
| 594 | needs_cast = self.output_dtype != x_dtype |
| 595 | |
| 596 | x, scales_slot, per_tensor_scale, qdata_slot = P.slot_map( |
| 597 | [x_node, self.scale, self.per_tensor_scale, self.qdata] |
| 598 | ) |
| 599 | |
| 600 | ids_index = IntOrVidOrTid.from_tid(P.slot_to_tid(x)) |
| 601 | |
| 602 | # Gather quantized weights by indices |
| 603 | _, wq_sel = P.make_tmp_slot() |
| 604 | P.emit( |
| 605 | TakeNode( |
| 606 | x=P.slot_to_tid(qdata_slot), |
| 607 | index=ids_index, |
| 608 | out=P.slot_to_tid(wq_sel), |
| 609 | axis=0, |
| 610 | ) |
| 611 | ) |
| 612 | |
| 613 | # Gather scales by indices |
| 614 | _, sc_sel = P.make_tmp_slot() |
| 615 | P.emit( |
| 616 | TakeNode( |
| 617 | x=P.slot_to_tid(scales_slot), |
| 618 | index=ids_index, |
| 619 | out=P.slot_to_tid(sc_sel), |
| 620 | axis=0, |
| 621 | ) |
| 622 | ) |
| 623 | |
| 624 | # Dequantize the gathered slices |
| 625 | out = P.make_or_get_slot(n) |
| 626 | P.emit( |
| 627 | DequantizeNode( |
| 628 | w=P.slot_to_tid(wq_sel), |
| 629 | scales=P.slot_to_tid(sc_sel), |
| 630 | out=P.slot_to_tid(out), |
| 631 | biases=None, |
| 632 | group_size=16, |
| 633 | bits=4, |
| 634 | mode="nvfp4", |
| 635 | dtype=torch_dtype_to_scalar_type(self.output_dtype), |
| 636 | ) |
nothing calls this directly
no test coverage detected