Function blockwise_quant_fp8_e4m3

tests/python/contrib/test_cutlass_gemm.py:178–245

(shape: tuple[int, int], block_size: tuple[int, int], dtype: str)

def blockwise_quant_fp8_e4m3(shape: tuple[int, int], block_size: tuple[int, int], dtype: str):
    w_full_np = (np.random.rand(*shape) * 2 - 1).astype(dtype)
    w_scale_shape = (
        *shape[:-2],
        (shape[-2] + block_size[0] - 1) // block_size[0],
        (shape[-1] + block_size[1] - 1) // block_size[1],
    )
    # For each (block_size[0], block_size[1]) block, compute the max abs value of `w_full_np`
    w_max_abs_np = np.zeros(w_scale_shape, dtype="float32")
    for i in range(w_scale_shape[-2]):
        for j in range(w_scale_shape[-1]):
            block_shape = (
                *shape[:-2],
                min(block_size[0], shape[-2] - i * block_size[0]),
                min(block_size[1], shape[-1] - j * block_size[1]),
            )
            w_max_abs_np[..., i, j] = np.max(
                np.abs(
                    w_full_np[
                        ...,
                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
                    ]
                ).reshape(*shape[:-2], block_shape[-2] * block_shape[-1]),
                axis=-1,
            )
    # Scale is the `w_max_abs_np` divided by the max value of quant_dtype in ml_dtypes
    fp8_max = float(ml_dtypes.finfo("float8_e4m3fn").max)
    w_scale_np = w_max_abs_np / fp8_max
    # `w_np` is the `w_full_np` divided by the `w_scale_np` (with block awareness),
    # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
    w_np = np.zeros_like(w_full_np, dtype="float8_e4m3fn")
    if len(w_scale_shape) == 2:
        for i in range(w_scale_shape[-2]):
            for j in range(w_scale_shape[-1]):
                w_np[
                    i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
                    j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
                ] = np.clip(
                    w_full_np[
                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
                    ]
                    / w_scale_np[..., i, j],
                    -fp8_max,
                    fp8_max,
                )
    else:
        for e in range(w_scale_shape[0]):
            for i in range(w_scale_shape[-2]):
                for j in range(w_scale_shape[-1]):
                    w_np[
                        e,
                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
                    ] = np.clip(
                        w_full_np[
                            e,