Quantization and Dequantization Operators¶

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

INT8 quantization¶

tileops.quantization.INT8QuantPerTensorFwdOp ¶

Symmetric INT8 quantization of a 2-D tensor against one scale for the whole tensor.

scale is amax(|x|) / 127 in float32, or 1.0 when x is all zero, and is the dequantization multiplier: x ~= q * scale. q is x / scale rounded half to even and clamped to [-127, 127]. Both are bit-equal to the torch expression.

A NaN or an infinity in x reaches scale as it does in torch; q is then unspecified.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
)

Quantize x against its absolute maximum.

Parameters:

x (Tensor) –

\([M \times K]\), float16, bfloat16 or float32.

Returns:

Tuple[Tensor, Tensor] –

q \([M \times K]\) in int8 and scale \([1]\) in float32.

tileops.quantization.INT8QuantPerChannelFwdOp ¶

Symmetric INT8 quantization of a weight with one scale per output channel (row).

scale[n] is amax(|w[n, :]|) / 127 in float32, or 1.0 for an all-zero row, and is the dequantization multiplier: w[n, :] ~= q[n, :] * scale[n]. q is the row divided by its scale, rounded half to even and clamped to [-127, 127].

Both are bit-equal to the torch expression.

A NaN or an infinity in a row reaches that row's scale as it does in torch; the row's q is then unspecified, as it is when the row's amax is so small (below about 8.8e-44, reachable only in float32) that its scale rounds to zero.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    w,
)

Quantize each row of w against its own absolute maximum.

Parameters:

w (Tensor) –

\([N \times K]\), float16, bfloat16 or float32.

Returns:

Tuple[Tensor, Tensor] –

q \([N \times K]\) in int8 and scale \([N]\) in float32.

tileops.quantization.INT8QuantPerBlockFwdOp ¶

Symmetric INT8 quantization with one scale per 128 contiguous elements along K.

Each row of x is split into blocks of 128 elements along K; the last block of a row holds K % 128 elements when K is not a multiple of 128. A block's scale is amax(|block|) / 127 in float32, or 1.0 for an all-zero block, and is the dequantization multiplier. q is each element divided by its block's scale, rounded half to even and clamped to [-127, 127].

Both are bit-equal to the torch expression.

A NaN or an infinity in a block reaches that block's scale as it does in torch; the block's q is then unspecified, as it is when the block's amax is so small (below about 8.8e-44, reachable only in float32) that its scale rounds to zero.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
)

Quantize each 128-element block of x against its own absolute maximum.

Parameters:

x (Tensor) –

\([M \times K]\), float16, bfloat16 or float32.

Returns:

Tuple[Tensor, Tensor] –

q \([M \times K]\) in int8 and scale \([M \times \lceil K / 128 \rceil]\) in float32.

tileops.quantization.SmoothQuantFwdOp ¶

SmoothQuant activation quantization: per-channel smoothing, then symmetric INT8 per row.

Each column of x is divided by its smoothing factor, xs = x / smooth in float32, as in SmoothQuant (Xiao et al., 2023). scale[m] is amax(|xs[m, :]|) / 127, or 1.0 for an all-zero row, and is the dequantization multiplier of xs. q is xs[m, :] / scale[m] rounded half to even and clamped to [-127, 127].

Both are bit-equal to the torch expression.

A NaN or an infinity in xs (a zero smoothing factor gives one) reaches that row's scale as it does in torch; the row's q is then unspecified, as it is when the row's amax is so small (below about 8.8e-44) that its scale rounds to zero.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
    smooth,
)

Smooth x by channel, then quantize each row against its absolute maximum.

Parameters:

x (Tensor) –

\([M \times K]\), float16 or bfloat16.
smooth (Tensor) –

\([K]\), float32: the per-channel smoothing factors.

Returns:

Tuple[Tensor, Tensor] –

q \([M \times K]\) in int8 and scale \([M]\) in float32.

FP8 quantization¶

tileops.quantization.FP8QuantFwdOp ¶

Quantize each row of an index tensor to float8_e4m3fn against its own maximum.

scale_tensor is the row's absolute maximum, floored at 1e-4, over 448, to within one float32 ulp. output_tensor is the row scaled by the reciprocal of that scale, so for a finite row every element is within one float8_e4m3fn code of clamp(input / scale, -448, 448) rather than equal to it.

A row holding an infinity or a NaN has no defined result: whether the row maximum and the clamp carry the non-finite value depends on the kernel serving the call.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from the first call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    input_tensor,
)

Quantize input_tensor row by row.

Parameters:

input_tensor (Tensor) –

\([B \times S \times G \times D]\), float16, bfloat16 or float32.

Returns:

Tensor –

scale_tensor \([B \times S \times G]\) in float32 and output_tensor
Tensor –

\([B \times S \times G \times D]\) in float8_e4m3fn.

tileops.quantization.FP8QuantPerBlockFwdOp ¶

Block-scaled float8_e4m3fn quantization of a weight, one scale per 128x128 tile.

This is the block-scaled FP8 weight format of DeepSeek-V3 checkpoints, where scale is stored as weight_scale_inv. Edge tiles cover the N % 128 rows and K % 128 columns that remain. A tile's scale is amax(|tile|) / 448 in float32, or 1.0 for an all-zero tile, and is the dequantization multiplier. q is each element divided by its tile's scale, clamped to [-448, 448] and rounded to the nearest float8_e4m3fn.

Both are bit-equal to the torch expression for every input, a tile holding an infinity or a NaN and a tile whose scale is subnormal or rounds to zero included.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    w,
)

Quantize each 128x128 tile of w against its own absolute maximum.

Parameters:

w (Tensor) –

\([N \times K]\), bfloat16, float16 or float32.

Returns:

Tuple[Tensor, Tensor] –

q \([N \times K]\) in float8_e4m3fn and scale \([\lceil N / 128 \rceil \times \lceil K / 128 \rceil]\) in float32.

INT4 quantization¶

tileops.quantization.INT4QuantPerGroupFwdOp ¶

DeepSpeed's asymmetric INT4 quantization, grouped along K.

For a group's minimum lo and maximum hi, the quantization multiplier is 16 / (hi - lo) (1 for a constant group), and its offset is (hi + lo) / 2. Codes are clamp(round((w - offset) * multiplier), -8, 7), with ties to even. Adjacent signed codes occupy the high and low nibbles of each int8 byte. params stores the inverse multiplier and offset as float32, one row per group. Floating-point division may move a code at a rounding boundary by one step.

The in-tree kernels serve power-of-two group sizes from 32 to 1024 or multiples of 128 up to 65536, with N * K <= 2**31 - 1. Nonfinite groups are unspecified.

init ¶

__init__(
    group_size=128,
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

group_size (int, default: 128 ) –

Elements per group along K (default 128).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    w,
)

Quantize and pack the groups of w.

Parameters:

w (Tensor) –

\([N \times K]\), float16.

Returns:

Tuple[Tensor, Tensor] –

packed_weight \([N \times K / 2]\) in int8 and params \([(N K / group\_size) \times 2]\) in float32.

INT8 dequantization¶

tileops.quantization.INT8DequantPerTensorFwdOp ¶

Dequantize an INT8 matrix with one scale: x = (q.float() * scale).to(out_dtype).

The multiply is in float32 and its result is cast once to out_dtype, so x is bit-identical to that expression.

init ¶

__init__(
    out_dtype,
    *,
    target=None
)

Build the op. Shapes are taken from each call.

Parameters:

out_dtype (dtype) –

Dtype of x: torch.float16, torch.bfloat16 or torch.float32.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    q,
    scale,
)

Dequantize q.

Parameters:

q (Tensor) –

\([M \times K]\), int8.
scale (Tensor) –

\([1]\), float32.

Returns:

Tensor –

x \([M \times K]\) in out_dtype.

tileops.quantization.INT8DequantPerChannelFwdOp ¶

Dequantize an INT8 matrix with one scale per row: x[m, k] = q[m, k] * scale[m].

The multiply is in float32 and its result is cast once to out_dtype, so x is bit-identical to (q.float() * scale[:, None]).to(out_dtype).

init ¶

__init__(
    out_dtype,
    *,
    target=None
)

Build the op. Shapes are taken from each call.

Parameters:

out_dtype (dtype) –

Dtype of x: torch.float16, torch.bfloat16 or torch.float32.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    q,
    scale,
)

Dequantize q row by row.

Parameters:

q (Tensor) –

\([M \times K]\), int8.
scale (Tensor) –

\([M]\), float32.

Returns:

Tensor –

x \([M \times K]\) in out_dtype.

tileops.quantization.INT8DequantPerBlockFwdOp ¶

Dequantize an INT8 matrix with one scale per 128 contiguous elements of a row.

x[m, k] = q[m, k] * scale[m, k // 128]; the last block of a row may be partial. The multiply is in float32 and its result is cast once to out_dtype, so x is bit-identical to (q.float() * scale.repeat_interleave(128, dim=1)[:, :K]).to(out_dtype).

init ¶

__init__(
    out_dtype,
    *,
    target=None
)

Build the op. Shapes are taken from each call.

Parameters:

out_dtype (dtype) –

Dtype of x: torch.float16, torch.bfloat16 or torch.float32.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    q,
    scale,
)

Dequantize q block by block.

Parameters:

q (Tensor) –

\([M \times K]\), int8.
scale (Tensor) –

\([M \times \lceil K / 128 \rceil]\), float32.

Returns:

Tensor –

x \([M \times K]\) in out_dtype.

Quantization and Dequantization Operators¶

INT8 quantization¶

tileops.quantization.INT8QuantPerTensorFwdOp ¶

__init__ ¶

forward ¶

tileops.quantization.INT8QuantPerChannelFwdOp ¶

__init__ ¶

forward ¶

tileops.quantization.INT8QuantPerBlockFwdOp ¶

__init__ ¶

forward ¶

tileops.quantization.SmoothQuantFwdOp ¶

__init__ ¶

forward ¶

FP8 quantization¶

tileops.quantization.FP8QuantFwdOp ¶

__init__ ¶

forward ¶

tileops.quantization.FP8QuantPerBlockFwdOp ¶

__init__ ¶

forward ¶

INT4 quantization¶

tileops.quantization.INT4QuantPerGroupFwdOp ¶

__init__ ¶

forward ¶

INT8 dequantization¶

tileops.quantization.INT8DequantPerTensorFwdOp ¶

__init__ ¶

forward ¶

tileops.quantization.INT8DequantPerChannelFwdOp ¶

__init__ ¶

forward ¶

tileops.quantization.INT8DequantPerBlockFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶