Function rowwise_quant_fp8_e4m3

tests/python/contrib/test_cutlass_gemm.py:140–175

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

def rowwise_quant_fp8_e4m3(shape: tuple[int, int], block_size: tuple[int, int], dtype: str):
    x_full_np = (np.random.rand(*shape) * 2 - 1).astype(dtype)
    x_scale_shape = (
        *shape[:-1],
        (shape[-1] + block_size[1] - 1) // block_size[1],
    )
    # For each (block_size[1]) block, compute the max abs value of `w_full_np`
    x_max_abs_np = np.zeros(x_scale_shape, dtype="float32")
    for i in range(x_scale_shape[-1]):
        x_max_abs_np[..., i] = np.max(
            np.abs(x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])]),
            axis=-1,
        )[0]
    # Scale is the `x_max_abs_np` divided by the max value of quant_dtype in ml_dtypes
    fp8_max = float(ml_dtypes.finfo("float8_e4m3fn").max)
    x_scale_np = x_max_abs_np / fp8_max
    # `x_np` is the `x_full_np` divided by the `x_scale_np` (with block awareness),
    # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
    x_np = np.zeros_like(x_full_np, dtype="float8_e4m3fn")
    for i in range(x_scale_shape[-1]):
        x_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])] = np.clip(
            x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])]
            / x_scale_np[..., i : i + 1],
            -fp8_max,
            fp8_max,
        )

    x_scale_np = np.random.rand(*x_scale_np.shape).astype("float32") / fp8_max
    for i in range(x_scale_shape[-1]):
        x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])] = (
            x_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])].astype(
                x_scale_np.dtype
            )
            * x_scale_np[..., i : i + 1]
        )
    return x_np, x_scale_np


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