| 82 | |
| 83 | |
| 84 | def scaled_mxfp8_quant(input: torch.Tensor): |
| 85 | m, n = input.shape |
| 86 | block_size = 32 |
| 87 | device = input.device |
| 88 | |
| 89 | output = torch.empty((m, n), device=device, dtype=torch.uint8) |
| 90 | output_scale = torch.empty(((m + 128 - 1) // 128 * 128, (n // block_size + 4 - 1) // 4), device=device, dtype=torch.int32) |
| 91 | |
| 92 | torch.ops.lightx2v_kernel.scaled_mxfp8_quant_sm120.default(output, input, output_scale) |
| 93 | output_scale = output_scale.view(torch.float8_e8m0fnu) |
| 94 | return output, output_scale |
| 95 | |
| 96 | |
| 97 | def scaled_mxfp8_gelu_quant(input: torch.Tensor): |