Linear Algebra 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(...).

Dense GEMM¶

tileops.gemm.GemmFwdOp ¶

Dense GEMM. Nothing is committed at construction: m, n, k and the dtype come from the forward inputs, so eval_roofline() is valid only after a call.

The (trans_a, trans_b) pair selects one of four layouts:

Flags	Layout	Product
`(False, True)`	NT, the default	\(d = a \mathbin{@} b^{\top}\)
`(False, False)`	NN	\(d = a \mathbin{@} b\)
`(True, False)`	TN	\(d = a^{\top} \mathbin{@} b\)
`(True, True)`	TT	\(d = a^{\top} \mathbin{@} b^{\top}\)

init ¶

__init__(
    trans_a=False,
    trans_b=True,
    *,
    target=None
)

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

Parameters:

trans_a (bool, default: False ) –

Whether a is stored transposed (\([K \times M]\)).
trans_b (bool, default: True ) –

Whether b is stored transposed (\([N \times K]\)). Default True (NT).
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(
    a,
    b,
)

Multiply the two matrices under the layout the constructor selected.

Parameters:

a (Tensor) –

Left operand, \([M \times K]\), or \([K \times M]\) when trans_a.
b (Tensor) –

Right operand, \([N \times K]\) under the default NT layout, or \([K \times N]\) when trans_b is false.

Returns:

Tensor –

The product, \([M \times N]\), in the dtype of the inputs.

Example

op = GemmFwdOp()                      # NT by default
d = op(a, b)                          # a=[M,K], b=[N,K] -> d=[M,N]
flops, nbytes = op.eval_roofline()    # valid after the forward

tileops.gemm.GemmFP8FwdOp ¶

Dense FP8 NT GEMM, input-inferred: \(d = (a \cdot s_a) \mathbin{@} (b \cdot s_b)^{\top} + \text{bias}\).

a is \([M \times K]\) and b is \([N \times K]\), the operand torch._scaled_mm receives as b.T. scale_a and scale_b are both per-tensor \([1 \times 1]\) scales, or scale_a is per 1x128 block along K, \([M \times \lceil K/128 \rceil]\), with scale_b per 1x128 block, \([N \times \lceil K/128 \rceil]\), or per 128x128 block, \([\lceil N/128 \rceil \times \lceil K/128 \rceil]\).

init ¶

__init__(
    out_dtype=torch.bfloat16,
    *,
    target=None
)

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

Parameters:

out_dtype (dtype, default: bfloat16 ) –

Output dtype, torch.bfloat16 or torch.float16.
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(
    a,
    b,
    scale_a,
    scale_b,
    bias=None,
)

Multiply the two FP8 matrices, apply the scales, and add the bias.

Parameters:

a (Tensor) –

Left operand, \([M \times K]\), torch.float8_e4m3fn.
b (Tensor) –

Right operand, \([N \times K]\), same dtype as a.
scale_a (Tensor) –

torch.float32 scales for a: per-tensor \([1 \times 1]\), or per 1x128 block \([M \times \lceil K/128 \rceil]\).
scale_b (Tensor) –

torch.float32 scales for b: \([1 \times 1]\) when scale_a is, else per 1x128 block \([N \times \lceil K/128 \rceil]\) or per 128x128 block \([\lceil N/128 \rceil \times \lceil K/128 \rceil]\).
bias (Optional[Tensor], default: None ) –

Optional bias, \([N]\), in out_dtype.

Returns:

Tensor –

The scaled product plus bias, \([M \times N]\), in out_dtype.

Example

op = GemmFP8FwdOp(out_dtype=torch.bfloat16)
d = op(a, b, scale_a, scale_b)        # per-tensor scales
flops, nbytes = op.eval_roofline()    # valid after the forward

tileops.gemm.GemmW4A16FwdOp ¶

Dense W4A16 NT GEMM with group-wise affine weight dequantization.

Inputs are activation [M, K], prepacked weight [N, K/2], and scale and zero-point tensors [N, K/group_size]. Weight (n, k) is (q - zero[n, g]) * scale[n, g] with g = k // group_size and q its INT4 value. repack converts a row-major packed weight once at load time; the two layouts have the same shape and dtype and cannot be distinguished at runtime. The output is activation @ W.T with shape [M, N]. The in-tree kernel serves group_size = 128 only and refuses any other value when it is built.

init ¶

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

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

Parameters:

group_size (int, default: 128 ) –

Weights per dequantization 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(
    activation,
    packed_weight,
    weight_scale,
    weight_zero,
)

Multiply the activations by the dequantized INT4 weight.

Parameters:

activation (Tensor) –

Activations, \([M \times K]\), torch.float16.
packed_weight (Tensor) –

Weights, \([N \times K/2]\), torch.uint8, in the order repack returns.
weight_scale (Tensor) –

Group scales, \([N \times K/\text{group\_size}]\), in the activation dtype.
weight_zero (Tensor) –

Group zero points, \([N \times K/\text{group\_size}]\), torch.uint8.

Returns:

Tensor –

The product, \([M \times N]\), in torch.float16.

Example

op = GemmW4A16FwdOp()
d = op(activation, packed_weight, weight_scale, weight_zero)
flops, nbytes = op.eval_roofline()    # valid after the forward

Batched GEMM¶

tileops.gemm.BmmFwdOp ¶

Batched dense GEMM: \(d_i = a_i \mathbin{@} b_i \quad \text{for } i \in [0, B)\).

Shapes are strictly 3D: a is \([B \times M \times K]\), b is \([B \times K \times N]\), and d is \([B \times M \times N]\). A kernel is compiled on first use for each (batch, m, n, k, dtype) combination and cached. The in-tree kernels need \(K\) to be a multiple of 16 and refuse other calls when built.

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(
    a,
    b,
)

Multiply the two batches, one GEMM per batch item.

Parameters:

a (Tensor) –

Left operand, \([B \times M \times K]\).
b (Tensor) –

Right operand, \([B \times K \times N]\).

Returns:

Tensor –

The product, \([B \times M \times N]\), in the dtype of the inputs.

Example

op = BmmFwdOp()
d = op(a, b)                          # a=[B,M,K], b=[B,K,N] -> d=[B,M,N]
flops, nbytes = op.eval_roofline()    # valid after the forward

tileops.gemm.BmmFP8FwdOp ¶

Batched FP8 GEMM: d[i] = (a[i] @ b[i]) * scale_a * scale_b.

trans_b states which axis order b arrives in: False is torch.bmm's \([B \times K \times N]\), True is \([B \times N \times K]\).

FP8 WGMMA reads b K-innermost and has no transposed operand mode, so a b that does not already lie that way is transposed into a new buffer, costing one extra read and write of it. Which calls pay is decided by b's strides, not by trans_b: passing b K-innermost avoids the copy under either value of the flag, and is the faster call. The in-tree kernel needs \(K\) to be a multiple of 32 and refuses other calls when built.

init ¶

__init__(
    out_dtype=torch.bfloat16,
    trans_b=False,
    *,
    target=None
)

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

Parameters:

out_dtype (dtype, default: bfloat16 ) –

Output tensor dtype (torch.float16 or torch.bfloat16).
trans_b (bool, default: False ) –

Whether b's axes are \([B \times N \times K]\) rather than \([B \times K \times N]\). Which of the two is faster is decided by b's strides, not by this flag.
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(
    a,
    b,
    scale_a,
    scale_b,
)

Multiply the two FP8 batches and apply the two scales.

Parameters:

a (Tensor) –

Left operand, \([B \times M \times K]\), torch.float8_e4m3fn.
b (Tensor) –

Right operand, same dtype as a: \([B \times K \times N]\), or \([B \times N \times K]\) when trans_b.
scale_a (Tensor) –

Per-tensor scale for a, a 0-dim torch.float32 tensor.
scale_b (Tensor) –

Per-tensor scale for b, a 0-dim torch.float32 tensor.

Returns:

Tensor –

The scaled product, \([B \times M \times N]\), in out_dtype.

Example

op = BmmFP8FwdOp(out_dtype=torch.bfloat16)      # b as [B, K, N]
d = op(a, b_kn, scale_a, scale_b)
flops, nbytes = op.eval_roofline()    # valid after the forward

Grouped GEMM¶

tileops.gemm.GroupedGemmFwdOp ¶

Grouped GEMM with configurable transpose modes. Nothing is committed at construction: the extents and the dtype come from the forward inputs.

The (transpose_a, transpose_b) pair selects one of four layouts:

Flags	Layout	Product
`(False, True)`	NT	\(C = A \mathbin{@} B^{\top}\)
`(False, False)`	NN	\(C = A \mathbin{@} B\)
`(True, False)`	TN	\(C = A^{\top} \mathbin{@} B\)
`(True, True)`	TT	\(C = A^{\top} \mathbin{@} B^{\top}\)

The metadata values are the caller's obligation and are not checked, since checking them would synchronise: batch_sizes sums to the packed row count, and batch_offsets is its exclusive prefix sum.

init ¶

__init__(
    transpose_a=False,
    transpose_b=True,
    *,
    target=None
)

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

Parameters:

transpose_a (bool, default: False ) –

Whether the groups split the contraction: a is \([\mathit{batch\_sum} \times N]\) and the output keeps a group axis.
transpose_b (bool, default: True ) –

Whether the per-group operand is stored transposed. Default True (NT).
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(
    a,
    b,
    batch_sizes,
    batch_offsets,
)

Run one GEMM per group, with the groups packed along a single axis.

Parameters:

a (Tensor) –

Activations for every group, \([\mathit{batch\_sum} \times K]\), or \([\mathit{batch\_sum} \times N]\) when transpose_a.
b (Tensor) –

Per-group weights, \([\mathit{batch\_count} \times N \times K]\) under the default NT layout, \([\mathit{batch\_count} \times K \times N]\) under NN; with transpose_a, \([K \times \mathit{batch\_sum}]\) or \([\mathit{batch\_sum} \times K]\).
batch_sizes (Tensor) –

Rows per group, 1D torch.int32.
batch_offsets (Tensor) –

Start row of each group in a, 1D torch.int32.

Returns:

Tensor –

The per-group products in the dtype of the inputs:
Tensor –

\([\mathit{batch\_sum} \times N]\), or
Tensor –

\([\mathit{batch\_count} \times N \times K]\) when transpose_a.

Example

op = GroupedGemmFwdOp()               # NT by default
d = op(a, b, batch_sizes, batch_offsets)

Linear Algebra Operators¶

Dense GEMM¶

tileops.gemm.GemmFwdOp ¶

__init__ ¶

forward ¶

tileops.gemm.GemmFP8FwdOp ¶

__init__ ¶

forward ¶

tileops.gemm.GemmW4A16FwdOp ¶

__init__ ¶

forward ¶

Batched GEMM¶

tileops.gemm.BmmFwdOp ¶

__init__ ¶

forward ¶

tileops.gemm.BmmFP8FwdOp ¶

__init__ ¶

forward ¶

Grouped GEMM¶

tileops.gemm.GroupedGemmFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶