| 69 | output_scale = torch.zeros((((m + 128 - 1) // 128) * 128, (n // block_size + 4 - 1) // 4), device=device, dtype=torch.int32) |
| 70 | |
| 71 | torch.ops.lightx2v_kernel.scaled_nvfp4_quant_sm120.default(output, input, output_scale, input_global_scale) |
| 72 | output_scale = output_scale.view(torch.float8_e4m3fn) |
| 73 | return output, output_scale |
| 74 | |
| 75 | |
| 76 | def scaled_mxfp4_quant(input: torch.Tensor): |
| 77 | m, n = input.shape |
| 78 | block_size = 32 |
| 79 | device = input.device |
| 80 | |
| 81 | output = torch.empty((m, n // 2), device=device, dtype=torch.uint8) |
| 82 | output_scale = torch.zeros(((m + 128 - 1) // 128 * 128, (n // block_size + 4 - 1) // 4), device=device, dtype=torch.int32) |
| 83 | |
| 84 | torch.ops.lightx2v_kernel.scaled_mxfp4_quant_sm120.default(output, input, output_scale) |