| 611 | |
| 612 | |
| 613 | class FP8RowwiseLinearMethod(LinearMethodBase): |
| 614 | |
| 615 | def create_weights(self, module: Linear, in_features: int, |
| 616 | out_features: int, bias: bool, dtype: torch.dtype): |
| 617 | weight_shape = (out_features, in_features) |
| 618 | |
| 619 | module.weight = Parameter(torch.empty(weight_shape, |
| 620 | dtype=torch.float8_e4m3fn), |
| 621 | requires_grad=False) |
| 622 | module.weight_scale = Parameter(torch.empty(out_features), |
| 623 | requires_grad=False) |
| 624 | # Not really used for Gemm now. |
| 625 | # Only used to quantize output of FP8 attention. |
| 626 | module.input_scale = Parameter(torch.tensor(1., dtype=torch.float32), |
| 627 | requires_grad=False) |
| 628 | module.inv_input_scale = Parameter(torch.tensor(1., |
| 629 | dtype=torch.float32), |
| 630 | requires_grad=False) |
| 631 | if bias: |
| 632 | module.bias = Parameter(torch.empty((out_features), dtype=dtype), |
| 633 | requires_grad=False) |
| 634 | else: |
| 635 | module.register_parameter("bias", None) |
| 636 | |
| 637 | def apply(self, module: Linear, input: torch.Tensor, |
| 638 | bias: Optional[torch.Tensor]): |
| 639 | # FP8 tensor inputs are from attention. Directly use ones as scale. |
| 640 | if input.dtype == torch.float8_e4m3fn: |
| 641 | qinput = input |
| 642 | cur_input_scale = torch.ones(input.shape[0], |
| 643 | device=input.device, |
| 644 | dtype=torch.float32) |
| 645 | else: |
| 646 | # Use dynamic per-token quantization for activation |
| 647 | qinput, cur_input_scale = torch.ops.tensorrt_llm.quantize_e4m3_activation( |
| 648 | input) |
| 649 | |
| 650 | # This op does not support bias now. |
| 651 | output = torch.ops.trtllm.fp8_rowwise_gemm( |
| 652 | qinput, |
| 653 | module.weight, |
| 654 | cur_input_scale.float(), |
| 655 | module.weight_scale, |
| 656 | module.dtype or input.dtype, |
| 657 | ) |
| 658 | if bias is not None: |
| 659 | output = output + bias |
| 660 | return output |
| 661 | |
| 662 | def _get_scale_name(self, weights: List[Dict]): |
| 663 | # `weight_scale_inv` for DS recipe and `weight_scale` for ModelOpt recipe. |
| 664 | # Actually they hold identical values of data_amax / 448. |
| 665 | scale_name = "weight_scale_inv" |
| 666 | if scale_name not in weights[0]: |
| 667 | scale_name = "weight_scale" |
| 668 | return scale_name |
| 669 | |
| 670 | def load_weights_vanilla(self, module: Linear, weights: List[Dict]): |