| 1412 | module.register_parameter("bias", None) |
| 1413 | |
| 1414 | def apply(self, module: Linear, input: torch.Tensor, |
| 1415 | bias: Optional[torch.Tensor]): |
| 1416 | fp8_input, input_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( |
| 1417 | input) |
| 1418 | input_scale = input_scale.to(torch.float32) |
| 1419 | nrows = fp4_utils.pad_up(input.shape[0], 128) |
| 1420 | ncols = fp4_utils.pad_up(input.shape[1] // module.scaling_vector_size, |
| 1421 | 4) |
| 1422 | # 01111111 is 2^(127 - 127) = 1 in E8M0 |
| 1423 | module.fake_act_scale = torch.empty( |
| 1424 | [nrows * ncols], dtype=torch.uint8, |
| 1425 | device=fp8_input.device).fill_(127).view(fp4_utils.float4_sf_dtype) |
| 1426 | output = torch.ops.trtllm.w4a8_mxfp4_fp8_gemm(fp8_input, module.weight, |
| 1427 | module.fake_act_scale, |
| 1428 | module.weight_scale, |
| 1429 | input_scale, module.dtype) |
| 1430 | if bias is not None: |
| 1431 | output = output + bias |
| 1432 | return output |
| 1433 | |
| 1434 | def load_weight_scales(self, |
| 1435 | weights: List[Dict], |