| 56 | |
| 57 | |
| 58 | def scaled_mxfp4_quant(input: torch.Tensor): |
| 59 | m, n = input.shape |
| 60 | block_size = 32 |
| 61 | device = input.device |
| 62 | |
| 63 | output = torch.empty((m, n // 2), device=device, dtype=torch.uint8) |
| 64 | output_scale = torch.zeros(((m + 128 - 1) // 128 * 128, (n // block_size + 4 - 1) // 4), device=device, dtype=torch.int32) |
| 65 | |
| 66 | torch.ops.lightx2v_kernel.scaled_mxfp4_quant_sm120.default(output, input, output_scale) |
| 67 | output_scale = output_scale.view(torch.float8_e8m0fnu) |
| 68 | return output, output_scale |
| 69 | |
| 70 | |
| 71 | def scaled_mxfp6_quant(input: torch.Tensor): |