Uploaded using `kernel-builder`.

d934615 verified 6 months ago

10.9 kB

"""Minimal MX / NVFP4 quantization + scale swizzling utilities.
Ported from torchao (BSD-3) to avoid the runtime dependency:
  torchao/prototype/mx_formats/{mx_tensor, nvfp4_tensor, utils, constants}.py
  torchao/prototype/custom_fp_utils.py
  torchao/prototype/mx_formats/kernels.py
All quantizers are pure-PyTorch. Use the `to_mx_compiled` / `to_mxfp4_compiled` /
`to_nvfp4_compiled` module-level handles if you want torch.compile-generated
Triton kernels (much faster on big tensors; one-time compile overhead).
Only the FLOOR scaling mode is ported (torchao's default for MX formats).
"""

import torch

F8E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max  # 448.0
F8E4M3_MAX_POW2 = 8
E8M0_EXPONENT_BIAS = 127
E8M0_EXPONENT_NAN_VAL = 255
F32_EXP_BIAS = 127
F32_MIN_NORMAL = 2 ** (-F32_EXP_BIAS + 1)  # 2**-126
MBITS_F32 = 23
EBITS_F32 = 8

# FP4 E2M1 constants
F4_E2M1_MAX = 6.0
F4_E2M1_MAX_POW2 = 2
F4_E2M1_MAX_INT = 7  # 3-bit magnitude mask
EBITS_F4_E2M1, MBITS_F4_E2M1 = 2, 1

E4M3_EPS = torch.finfo(torch.float8_e4m3fn).tiny


def _n_ones(n: int) -> int:
    return (1 << n) - 1


def to_mx(data_hp: torch.Tensor, block_size: int = 32):
    """MXFP8-e4m3 quantization with FLOOR scaling.
    Args:
        data_hp: (..., K) bf16 or fp32 tensor, contiguous, K % block_size == 0.
    Returns:
        qdata: (..., K) float8_e4m3fn
        scale: (..., K // block_size) float8_e8m0fnu
    """
    assert data_hp.dtype in (torch.bfloat16, torch.float32)
    assert data_hp.shape[-1] % block_size == 0
    assert data_hp.is_contiguous()

    orig_shape = data_hp.shape
    data_hp = data_hp.reshape(*orig_shape[:-1], orig_shape[-1] // block_size, block_size)
    max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1)

    data_hp = data_hp.to(torch.float32)
    max_abs = max_abs.to(torch.float32)

    # FLOOR scaling: extract biased exponent of max_abs via bit-shift
    max_abs_int32 = max_abs.view(torch.int32)
    extracted_pow2 = ((torch.bitwise_right_shift(max_abs_int32, MBITS_F32)) & 0xFF) - F32_EXP_BIAS
    scale_e8m0_unbiased = extracted_pow2 - F8E4M3_MAX_POW2
    scale_e8m0_unbiased = torch.clamp(
        scale_e8m0_unbiased, min=-E8M0_EXPONENT_BIAS, max=E8M0_EXPONENT_BIAS + 1
    )
    scale_e8m0_biased = (scale_e8m0_unbiased + E8M0_EXPONENT_BIAS).to(torch.uint8)
    # restore NaN sentinel (uint8 cast drops NaN)
    scale_e8m0_biased = torch.where(torch.isnan(max_abs), E8M0_EXPONENT_NAN_VAL, scale_e8m0_biased)

    # reconstruct fp32 scale from biased exponent
    scale_fp32 = (torch.bitwise_left_shift(scale_e8m0_biased.to(torch.int32), MBITS_F32)).view(
        torch.float32
    )
    # avoid 2**-127 being flushed to 0 (pytorch #125557)
    scale_fp32 = torch.clamp(scale_fp32, min=F32_MIN_NORMAL)

    data_lp = data_hp / scale_fp32
    # eager fp8 cast is unsaturated; clamp explicitly
    if not torch._dynamo.is_compiling():
        data_lp = torch.clamp(data_lp, min=-F8E4M3_MAX, max=F8E4M3_MAX)

    qdata = data_lp.to(torch.float8_e4m3fn).reshape(orig_shape)
    scale = scale_e8m0_biased.view(torch.float8_e8m0fnu).squeeze(-1)
    return qdata, scale


def _f32_to_floatx_unpacked(x: torch.Tensor, ebits: int, mbits: int) -> torch.Tensor:
    """FP32 -> sub-byte float (uint8, code in low bits). Verbatim from torchao.
    Round-to-nearest-even via magic-adder; saturation on overflow; no NaN.
    """
    assert x.dtype == torch.float
    assert 1 + ebits + mbits <= 8
    exp_bias = _n_ones(ebits - 1)
    max_int = _n_ones(ebits + mbits)
    sign_mask = 1 << (ebits + mbits)
    magic_adder = _n_ones(MBITS_F32 - mbits - 1)
    max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2**mbits))
    min_normal = 2 ** (1 - exp_bias)
    denorm_exp = (F32_EXP_BIAS - exp_bias) + (MBITS_F32 - mbits) + 1
    denorm_mask_int = denorm_exp << MBITS_F32
    denorm_mask_float = torch.tensor(denorm_mask_int, dtype=torch.int32).view(torch.float32)

    x = x.view(torch.int32)
    sign = x & 0x80000000
    x = x ^ sign
    x = x.view(torch.float)
    saturate_mask = x >= max_normal
    denormal_mask = torch.logical_and(torch.logical_not(saturate_mask), x < min_normal)
    normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask))
    denormal_x = x + denorm_mask_float
    denormal_x = denormal_x.view(torch.int32)
    denormal_x -= denorm_mask_int
    denormal_x = denormal_x.to(torch.uint8)
    normal_x = x.view(torch.int32)
    mant_odd = (normal_x >> (MBITS_F32 - mbits)) & 1
    val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder
    normal_x += val_to_add
    normal_x += mant_odd
    normal_x = normal_x >> (MBITS_F32 - mbits)
    normal_x = normal_x.to(torch.uint8)
    x = torch.full_like(x, max_int, dtype=torch.uint8)
    x = torch.where(denormal_mask, denormal_x, x)
    x = torch.where(normal_mask, normal_x, x)
    sign_lp = sign >> (MBITS_F32 + EBITS_F32 - mbits - ebits)
    sign_lp = sign_lp.to(torch.uint8)
    sign_lp = sign_lp & sign_mask
    x = x | sign_lp
    return x.to(torch.uint8)


def _pack_uint4(uint8_data: torch.Tensor) -> torch.Tensor:
    """Pack 4-bit uint8 values in pairs: pair (a,b) -> byte (b<<4 | a)."""
    shape = uint8_data.shape
    assert shape[-1] % 2 == 0
    uint8_data = uint8_data.contiguous().view(-1)
    return (uint8_data[::2] | uint8_data[1::2] << 4).view(*shape[:-1], shape[-1] // 2)


def _compute_e8m0_scale_floor(max_abs: torch.Tensor, target_max_pow2: int) -> torch.Tensor:
    """Return biased E8M0 scale (uint8) for FLOOR-mode MX quantization."""
    max_abs_int32 = max_abs.view(torch.int32)
    extracted_pow2 = ((torch.bitwise_right_shift(max_abs_int32, MBITS_F32)) & 0xFF) - F32_EXP_BIAS
    scale_unbiased = extracted_pow2 - target_max_pow2
    scale_unbiased = torch.clamp(
        scale_unbiased, min=-E8M0_EXPONENT_BIAS, max=E8M0_EXPONENT_BIAS + 1
    )
    scale_biased = (scale_unbiased + E8M0_EXPONENT_BIAS).to(torch.uint8)
    scale_biased = torch.where(torch.isnan(max_abs), E8M0_EXPONENT_NAN_VAL, scale_biased)
    return scale_biased


def to_mxfp4(x: torch.Tensor, block_size: int = 32):
    """MXFP4 quantization: E2M1 data + E8M0 per-block scales, FLOOR scaling.
    Args:
        x: (..., K) bf16/fp16/fp32, contiguous, K % block_size == 0.
    Returns:
        qdata_packed: uint8, shape (..., K // 2). Two FP4 values per byte
                      (first -> low nibble, second -> high nibble).
        scale: float8_e8m0fnu, shape (..., K // block_size).
    """
    assert x.dtype in (torch.bfloat16, torch.float16, torch.float32)
    assert x.shape[-1] % block_size == 0
    assert x.is_contiguous()

    orig_shape = x.shape
    data_hp = x.reshape(*orig_shape[:-1], orig_shape[-1] // block_size, block_size)
    max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1)
    data_hp = data_hp.to(torch.float32)
    max_abs = max_abs.to(torch.float32)

    scale_biased = _compute_e8m0_scale_floor(max_abs, F4_E2M1_MAX_POW2)
    scale_fp32 = (torch.bitwise_left_shift(scale_biased.to(torch.int32), MBITS_F32)).view(
        torch.float32
    )
    scale_fp32 = torch.clamp(scale_fp32, min=F32_MIN_NORMAL)

    data_lp = data_hp / scale_fp32
    data_lp = data_lp.reshape(orig_shape)
    data_lp = _f32_to_floatx_unpacked(data_lp.float(), EBITS_F4_E2M1, MBITS_F4_E2M1)
    data_lp = _pack_uint4(data_lp)

    scale = scale_biased.view(torch.float8_e8m0fnu).squeeze(-1)
    return data_lp, scale


def nvfp4_per_tensor_scale(amax: torch.Tensor) -> torch.Tensor:
    """NVFP4 per-tensor scale: amax / (F8E4M3_MAX * F4_E2M1_MAX) = amax / 2688."""
    return amax.to(torch.float32) / (F8E4M3_MAX * F4_E2M1_MAX)


def to_nvfp4(x: torch.Tensor, block_size: int = 16, per_tensor_scale=None):
    """NVFP4 quantization: E2M1 data + E4M3 per-block scales + optional fp32 per-tensor scale.
    Args:
        x: (..., K) bf16/fp32, contiguous, K % 16 == 0.
        block_size: must be 16.
        per_tensor_scale: scalar fp32 tensor, or None (uses 1.0 / returns unit).
    Returns:
        qdata_packed: uint8, shape (..., K // 2)
        scale: float8_e4m3fn, shape (..., K // 16)
        per_tensor_scale: scalar fp32 tensor (1.0 if None was passed)
    """
    assert x.dtype in (torch.bfloat16, torch.float32)
    assert x.shape[-1] % block_size == 0
    assert x.is_contiguous()
    assert block_size == 16, "NVFP4 requires block_size=16"

    orig_shape = x.shape
    data_hp = x.float().reshape(*orig_shape[:-1], orig_shape[-1] // block_size, block_size)
    max_abs = torch.amax(torch.abs(data_hp), dim=-1)
    block_scale = max_abs / F4_E2M1_MAX

    if per_tensor_scale is None:
        block_scale_fp8 = torch.clamp(block_scale, min=E4M3_EPS, max=F8E4M3_MAX).to(
            torch.float8_e4m3fn
        )
        recip = 1.0 / block_scale_fp8.to(torch.float32)
        returned_pts = torch.tensor(1.0, dtype=torch.float32, device=x.device)
    else:
        scaled = block_scale.to(torch.float32) / per_tensor_scale
        block_scale_fp8 = torch.clamp(scaled, min=E4M3_EPS, max=F8E4M3_MAX).to(torch.float8_e4m3fn)
        recip = (1.0 / per_tensor_scale) / block_scale_fp8.to(torch.float32)
        returned_pts = per_tensor_scale.to(torch.float32)

    data_scaled = data_hp * recip.unsqueeze(-1)
    data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX)
    data_scaled = data_scaled.view(orig_shape)
    data_lp = _f32_to_floatx_unpacked(data_scaled.float(), EBITS_F4_E2M1, MBITS_F4_E2M1)
    data_lp = _pack_uint4(data_lp)
    return data_lp, block_scale_fp8, returned_pts


# ---------------------------------------------------------------------------
# torch.compile-wrapped fast paths. Generates fused Triton quant kernels via
# Inductor. dynamic=True avoids recompilation on shape changes.
# ---------------------------------------------------------------------------
to_mx_compiled = torch.compile(to_mx, dynamic=True)
to_mxfp4_compiled = torch.compile(to_mxfp4, dynamic=True)
to_nvfp4_compiled = torch.compile(to_nvfp4, dynamic=True)


def _ceil_div(a, b):
    return (a + b - 1) // b


def to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
    """Swizzle a (H, W) e8m0 scale tensor into the 128x4 blocked layout
    cuBLAS expects for MXFP8 _scaled_mm. Returns a 1-D flat tensor of size
    32*ceil(H/128) * 16*ceil(W/4)."""
    rows, cols = input_matrix.shape
    n_row_blocks = _ceil_div(rows, 128)
    n_col_blocks = _ceil_div(cols, 4)
    padded_rows = n_row_blocks * 128
    padded_cols = n_col_blocks * 4

    padded = input_matrix
    if torch.compiler.is_compiling() or (rows, cols) != (padded_rows, padded_cols):
        padded = torch.zeros(
            (padded_rows, padded_cols),
            device=input_matrix.device,
            dtype=input_matrix.dtype,
        )
        padded[:rows, :cols] = input_matrix

    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    return rearranged.flatten()

	"""Minimal MX / NVFP4 quantization + scale swizzling utilities.

	Ported from torchao (BSD-3) to avoid the runtime dependency:
	torchao/prototype/mx_formats/{mx_tensor, nvfp4_tensor, utils, constants}.py
	torchao/prototype/custom_fp_utils.py
	torchao/prototype/mx_formats/kernels.py

	All quantizers are pure-PyTorch. Use the `to_mx_compiled` / `to_mxfp4_compiled` /
	`to_nvfp4_compiled` module-level handles if you want torch.compile-generated
	Triton kernels (much faster on big tensors; one-time compile overhead).

	Only the FLOOR scaling mode is ported (torchao's default for MX formats).
	"""

	import torch

	F8E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max # 448.0
	F8E4M3_MAX_POW2 = 8
	E8M0_EXPONENT_BIAS = 127
	E8M0_EXPONENT_NAN_VAL = 255
	F32_EXP_BIAS = 127
	F32_MIN_NORMAL = 2 (-F32_EXP_BIAS + 1) # 2-126
	MBITS_F32 = 23
	EBITS_F32 = 8

	# FP4 E2M1 constants
	F4_E2M1_MAX = 6.0
	F4_E2M1_MAX_POW2 = 2
	F4_E2M1_MAX_INT = 7 # 3-bit magnitude mask
	EBITS_F4_E2M1, MBITS_F4_E2M1 = 2, 1

	E4M3_EPS = torch.finfo(torch.float8_e4m3fn).tiny


	def _n_ones(n: int) -> int:
	return (1 << n) - 1


	def to_mx(data_hp: torch.Tensor, block_size: int = 32):
	"""MXFP8-e4m3 quantization with FLOOR scaling.

	Args:
	data_hp: (..., K) bf16 or fp32 tensor, contiguous, K % block_size == 0.
	Returns:
	qdata: (..., K) float8_e4m3fn
	scale: (..., K // block_size) float8_e8m0fnu
	"""
	assert data_hp.dtype in (torch.bfloat16, torch.float32)
	assert data_hp.shape[-1] % block_size == 0
	assert data_hp.is_contiguous()

	orig_shape = data_hp.shape
	data_hp = data_hp.reshape(*orig_shape[:-1], orig_shape[-1] // block_size, block_size)
	max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1)

	data_hp = data_hp.to(torch.float32)
	max_abs = max_abs.to(torch.float32)

	# FLOOR scaling: extract biased exponent of max_abs via bit-shift
	max_abs_int32 = max_abs.view(torch.int32)
	extracted_pow2 = ((torch.bitwise_right_shift(max_abs_int32, MBITS_F32)) & 0xFF) - F32_EXP_BIAS
	scale_e8m0_unbiased = extracted_pow2 - F8E4M3_MAX_POW2
	scale_e8m0_unbiased = torch.clamp(
	scale_e8m0_unbiased, min=-E8M0_EXPONENT_BIAS, max=E8M0_EXPONENT_BIAS + 1
	)
	scale_e8m0_biased = (scale_e8m0_unbiased + E8M0_EXPONENT_BIAS).to(torch.uint8)
	# restore NaN sentinel (uint8 cast drops NaN)
	scale_e8m0_biased = torch.where(torch.isnan(max_abs), E8M0_EXPONENT_NAN_VAL, scale_e8m0_biased)

	# reconstruct fp32 scale from biased exponent
	scale_fp32 = (torch.bitwise_left_shift(scale_e8m0_biased.to(torch.int32), MBITS_F32)).view(
	torch.float32
	)
	# avoid 2**-127 being flushed to 0 (pytorch #125557)
	scale_fp32 = torch.clamp(scale_fp32, min=F32_MIN_NORMAL)

	data_lp = data_hp / scale_fp32
	# eager fp8 cast is unsaturated; clamp explicitly
	if not torch._dynamo.is_compiling():
	data_lp = torch.clamp(data_lp, min=-F8E4M3_MAX, max=F8E4M3_MAX)

	qdata = data_lp.to(torch.float8_e4m3fn).reshape(orig_shape)
	scale = scale_e8m0_biased.view(torch.float8_e8m0fnu).squeeze(-1)
	return qdata, scale


	def _f32_to_floatx_unpacked(x: torch.Tensor, ebits: int, mbits: int) -> torch.Tensor:
	"""FP32 -> sub-byte float (uint8, code in low bits). Verbatim from torchao.

	Round-to-nearest-even via magic-adder; saturation on overflow; no NaN.
	"""
	assert x.dtype == torch.float
	assert 1 + ebits + mbits <= 8
	exp_bias = _n_ones(ebits - 1)
	max_int = _n_ones(ebits + mbits)
	sign_mask = 1 << (ebits + mbits)
	magic_adder = _n_ones(MBITS_F32 - mbits - 1)
	max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2**mbits))
	min_normal = 2 ** (1 - exp_bias)
	denorm_exp = (F32_EXP_BIAS - exp_bias) + (MBITS_F32 - mbits) + 1
	denorm_mask_int = denorm_exp << MBITS_F32
	denorm_mask_float = torch.tensor(denorm_mask_int, dtype=torch.int32).view(torch.float32)

	x = x.view(torch.int32)
	sign = x & 0x80000000
	x = x ^ sign
	x = x.view(torch.float)
	saturate_mask = x >= max_normal
	denormal_mask = torch.logical_and(torch.logical_not(saturate_mask), x < min_normal)
	normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask))
	denormal_x = x + denorm_mask_float
	denormal_x = denormal_x.view(torch.int32)
	denormal_x -= denorm_mask_int
	denormal_x = denormal_x.to(torch.uint8)
	normal_x = x.view(torch.int32)
	mant_odd = (normal_x >> (MBITS_F32 - mbits)) & 1
	val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder
	normal_x += val_to_add
	normal_x += mant_odd
	normal_x = normal_x >> (MBITS_F32 - mbits)
	normal_x = normal_x.to(torch.uint8)
	x = torch.full_like(x, max_int, dtype=torch.uint8)
	x = torch.where(denormal_mask, denormal_x, x)
	x = torch.where(normal_mask, normal_x, x)
	sign_lp = sign >> (MBITS_F32 + EBITS_F32 - mbits - ebits)
	sign_lp = sign_lp.to(torch.uint8)
	sign_lp = sign_lp & sign_mask
	x = x \| sign_lp
	return x.to(torch.uint8)


	def _pack_uint4(uint8_data: torch.Tensor) -> torch.Tensor:
	"""Pack 4-bit uint8 values in pairs: pair (a,b) -> byte (b<<4 \| a)."""
	shape = uint8_data.shape
	assert shape[-1] % 2 == 0
	uint8_data = uint8_data.contiguous().view(-1)
	return (uint8_data[::2] \| uint8_data[1::2] << 4).view(*shape[:-1], shape[-1] // 2)


	def _compute_e8m0_scale_floor(max_abs: torch.Tensor, target_max_pow2: int) -> torch.Tensor:
	"""Return biased E8M0 scale (uint8) for FLOOR-mode MX quantization."""
	max_abs_int32 = max_abs.view(torch.int32)
	extracted_pow2 = ((torch.bitwise_right_shift(max_abs_int32, MBITS_F32)) & 0xFF) - F32_EXP_BIAS
	scale_unbiased = extracted_pow2 - target_max_pow2
	scale_unbiased = torch.clamp(
	scale_unbiased, min=-E8M0_EXPONENT_BIAS, max=E8M0_EXPONENT_BIAS + 1
	)
	scale_biased = (scale_unbiased + E8M0_EXPONENT_BIAS).to(torch.uint8)
	scale_biased = torch.where(torch.isnan(max_abs), E8M0_EXPONENT_NAN_VAL, scale_biased)
	return scale_biased


	def to_mxfp4(x: torch.Tensor, block_size: int = 32):
	"""MXFP4 quantization: E2M1 data + E8M0 per-block scales, FLOOR scaling.

	Args:
	x: (..., K) bf16/fp16/fp32, contiguous, K % block_size == 0.
	Returns:
	qdata_packed: uint8, shape (..., K // 2). Two FP4 values per byte
	(first -> low nibble, second -> high nibble).
	scale: float8_e8m0fnu, shape (..., K // block_size).
	"""
	assert x.dtype in (torch.bfloat16, torch.float16, torch.float32)
	assert x.shape[-1] % block_size == 0
	assert x.is_contiguous()

	orig_shape = x.shape
	data_hp = x.reshape(*orig_shape[:-1], orig_shape[-1] // block_size, block_size)
	max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1)
	data_hp = data_hp.to(torch.float32)
	max_abs = max_abs.to(torch.float32)

	scale_biased = _compute_e8m0_scale_floor(max_abs, F4_E2M1_MAX_POW2)
	scale_fp32 = (torch.bitwise_left_shift(scale_biased.to(torch.int32), MBITS_F32)).view(
	torch.float32
	)
	scale_fp32 = torch.clamp(scale_fp32, min=F32_MIN_NORMAL)

	data_lp = data_hp / scale_fp32
	data_lp = data_lp.reshape(orig_shape)
	data_lp = _f32_to_floatx_unpacked(data_lp.float(), EBITS_F4_E2M1, MBITS_F4_E2M1)
	data_lp = _pack_uint4(data_lp)

	scale = scale_biased.view(torch.float8_e8m0fnu).squeeze(-1)
	return data_lp, scale


	def nvfp4_per_tensor_scale(amax: torch.Tensor) -> torch.Tensor:
	"""NVFP4 per-tensor scale: amax / (F8E4M3_MAX * F4_E2M1_MAX) = amax / 2688."""
	return amax.to(torch.float32) / (F8E4M3_MAX * F4_E2M1_MAX)


	def to_nvfp4(x: torch.Tensor, block_size: int = 16, per_tensor_scale=None):
	"""NVFP4 quantization: E2M1 data + E4M3 per-block scales + optional fp32 per-tensor scale.

	Args:
	x: (..., K) bf16/fp32, contiguous, K % 16 == 0.
	block_size: must be 16.
	per_tensor_scale: scalar fp32 tensor, or None (uses 1.0 / returns unit).
	Returns:
	qdata_packed: uint8, shape (..., K // 2)
	scale: float8_e4m3fn, shape (..., K // 16)
	per_tensor_scale: scalar fp32 tensor (1.0 if None was passed)
	"""
	assert x.dtype in (torch.bfloat16, torch.float32)
	assert x.shape[-1] % block_size == 0
	assert x.is_contiguous()
	assert block_size == 16, "NVFP4 requires block_size=16"

	orig_shape = x.shape
	data_hp = x.float().reshape(*orig_shape[:-1], orig_shape[-1] // block_size, block_size)
	max_abs = torch.amax(torch.abs(data_hp), dim=-1)
	block_scale = max_abs / F4_E2M1_MAX

	if per_tensor_scale is None:
	block_scale_fp8 = torch.clamp(block_scale, min=E4M3_EPS, max=F8E4M3_MAX).to(
	torch.float8_e4m3fn
	)
	recip = 1.0 / block_scale_fp8.to(torch.float32)
	returned_pts = torch.tensor(1.0, dtype=torch.float32, device=x.device)
	else:
	scaled = block_scale.to(torch.float32) / per_tensor_scale
	block_scale_fp8 = torch.clamp(scaled, min=E4M3_EPS, max=F8E4M3_MAX).to(torch.float8_e4m3fn)
	recip = (1.0 / per_tensor_scale) / block_scale_fp8.to(torch.float32)
	returned_pts = per_tensor_scale.to(torch.float32)

	data_scaled = data_hp * recip.unsqueeze(-1)
	data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX)
	data_scaled = data_scaled.view(orig_shape)
	data_lp = _f32_to_floatx_unpacked(data_scaled.float(), EBITS_F4_E2M1, MBITS_F4_E2M1)
	data_lp = _pack_uint4(data_lp)
	return data_lp, block_scale_fp8, returned_pts


	# ---------------------------------------------------------------------------
	# torch.compile-wrapped fast paths. Generates fused Triton quant kernels via
	# Inductor. dynamic=True avoids recompilation on shape changes.
	# ---------------------------------------------------------------------------
	to_mx_compiled = torch.compile(to_mx, dynamic=True)
	to_mxfp4_compiled = torch.compile(to_mxfp4, dynamic=True)
	to_nvfp4_compiled = torch.compile(to_nvfp4, dynamic=True)


	def _ceil_div(a, b):
	return (a + b - 1) // b


	def to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
	"""Swizzle a (H, W) e8m0 scale tensor into the 128x4 blocked layout
	cuBLAS expects for MXFP8 _scaled_mm. Returns a 1-D flat tensor of size
	32ceil(H/128) 16*ceil(W/4)."""
	rows, cols = input_matrix.shape
	n_row_blocks = _ceil_div(rows, 128)
	n_col_blocks = _ceil_div(cols, 4)
	padded_rows = n_row_blocks * 128
	padded_cols = n_col_blocks * 4

	padded = input_matrix
	if torch.compiler.is_compiling() or (rows, cols) != (padded_rows, padded_cols):
	padded = torch.zeros(
	(padded_rows, padded_cols),
	device=input_matrix.device,
	dtype=input_matrix.dtype,
	)
	padded[:rows, :cols] = input_matrix

	blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
	rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
	return rearranged.flatten()