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

Batched GEMM

tileops.ops.gemm.bmm.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]\). The batch and contraction dims are checked at forward() time. A kernel is compiled on first use for each (batch, m, n, k, dtype) combination and cached.

__init__

__init__(
    kernel_map=None,
    tune=False,
)

Build the op. No shape or dtype is bound until the first call.

Parameters:

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

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.

Raises:

• ValueError –

  The operands disagree on dtype or device, either is not 3D, or their batch or contraction dims do not match.

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.ops.gemm.bmm.BmmFp8KNFwdOp

Batched FP8 GEMM over b in \([B \times K \times N]\): d[i] = (a[i] @ b[i]) * scale_a * scale_b.

This is torch.bmm's memory order. The fp8-TN WGMMA kernel wants K innermost, so this op transposes b before the call; BmmFp8NKFwdOp takes \([B \times N \times K]\) and hands it over as it stands.

__init__

__init__(
    out_dtype="bfloat16",
    kernel_map=None,
    tune=False,
)

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

Parameters:

• out_dtype (dtype | str, default: 'bfloat16' ) –

  Output tensor dtype (torch.float16 or torch.bfloat16).

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune (applied when a kernel is first built).

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, \([B \times K \times N]\), same dtype as a.

• 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.

Raises:

• ValueError –

  An input is not on CUDA, a dtype is not the one listed above, a scale is not 0-dim, the batch or contraction dims do not match, or \(K\) is not a multiple of 32 — the FP8 WGMMA K-step.

Example

op = BmmFp8KNFwdOp(out_dtype=torch.bfloat16)
d = op(a, b_kn, scale_a, scale_b)
flops, nbytes = op.eval_roofline()    # valid after the forward

tileops.ops.gemm.bmm.BmmFp8NKFwdOp

Batched FP8 GEMM over b in \([B \times N \times K]\).

K is innermost, which is the order the fp8-TN WGMMA kernel reads, so b reaches it without a transpose. Same kernel and same arithmetic as BmmFp8KNFwdOp; only the memory order b arrives in differs, and memory order is part of the signature, so it is its own entry.

b is \([B \times N \times K]\); every other argument, the return value and the errors are BmmFp8KNFwdOp.forward's.

Example

op = BmmFp8NKFwdOp(out_dtype=torch.bfloat16)
d = op(a, b_nk, scale_a, scale_b)

__init__

__init__(
    out_dtype="bfloat16",
    kernel_map=None,
    tune=False,
)

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

Parameters:

• out_dtype (dtype | str, default: 'bfloat16' ) –

  Output tensor dtype (torch.float16 or torch.bfloat16).

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune (applied when a kernel is first built).

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, \([B \times K \times N]\), same dtype as a.

• 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.

Raises:

• ValueError –

  An input is not on CUDA, a dtype is not the one listed above, a scale is not 0-dim, the batch or contraction dims do not match, or \(K\) is not a multiple of 32 — the FP8 WGMMA K-step.

Example

op = BmmFp8KNFwdOp(out_dtype=torch.bfloat16)
d = op(a, b_kn, scale_a, scale_b)
flops, nbytes = op.eval_roofline()    # valid after the forward

Dense GEMM

tileops.ops.gemm.gemm.GemmFwdOp

Dense GEMM, input-inferred and aligned to DeepGEMM's call-time JIT.

The logical dims m, n, k and the dtype are derived from the forward inputs; nothing is committed at construction. The dtype-specialized kernel is built (and cached) on first use for each (m, n, k, dtype) — mirroring DeepGEMM's compile-on-first-call + per-config cache.

The (trans_a, trans_b) pair selects one of four layouts, matching DeepGEMM's nt / nn / tn / tt:

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,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype 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).

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune (applied when a kernel is first built).

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.

Raises:

• ValueError –

  The contraction dims the two operands contribute do not match.

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.ops.gemm.gemm.GemmFp8FwdOp

Dense FP8 NT GEMM, input-inferred.

Public layout is a: \([M \times K]\) and b: \([N \times K]\). scale_a and scale_b must be either per-tensor \([1 \times 1]\) scales or block128 scales with shapes \([M \times \lceil K/128 \rceil]\) and \([N \times \lceil K/128 \rceil]\).

__init__

__init__(
    out_dtype="bfloat16",
    kernel_map=None,
    tune=False,
)

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

Parameters:

• out_dtype (dtype | str, default: 'bfloat16' ) –

  Output dtype.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

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 block128 \([M \times \lceil K/128 \rceil]\).

• scale_b (Tensor) –

  The same for b: \([1 \times 1]\) or \([N \times \lceil K/128 \rceil]\). Both scales take the same form.

• bias (Optional[Tensor], default: None ) –

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

Returns:

• Tensor –

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

Raises:

• ValueError –

  A dtype is not one of those listed above, a or b is not 2D, the contraction dims do not match, the two scales are not both per-tensor or both block128, or the bias is not \([N]\).

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.ops.gemm.gemm.GemmW4A16FwdOp

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

Public layout is activation: \([M \times K]\) and packed_weight: \([N \times K/2]\). Two unsigned INT4 values are packed per byte: the low nibble stores even K and the high nibble stores odd K. weight_scale and weight_zero are group128 metadata with shape \([N \times K/128]\). The kernel dequantizes the current W4 tile into A16 shared memory and computes activation @ W.T.

__init__

__init__(
    group_size=GROUP_SIZE,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• group_size (int, default: GROUP_SIZE ) –

  Manifest params.group_size, int, default 128.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    activation,
    packed_weight,
    weight_scale,
    weight_zero,
)

Dequantize the INT4 weight tile by tile and multiply.

Parameters:

• activation (Tensor) –

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

• packed_weight (Tensor) –

  Weights, \([N \times K/2]\), torch.uint8 — two INT4 values per byte, even \(K\) in the low nibble.

• weight_scale (Tensor) –

  Group scales, \([N \times K/128]\), torch.float32.

• weight_zero (Tensor) –

  Group zero points, \([N \times K/128]\), torch.uint8.

Returns:

• Tensor –

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

Raises:

• ValueError –

  A dtype is not one of those listed above, activation or packed_weight is not 2D, \(K\) is odd, \(K\) is not divisible by group_size, or a packed or metadata shape disagrees with \(K\).

Example

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

Grouped GEMM

tileops.ops.gemm.grouped_gemm.GroupedGemmFwdOp

Grouped GEMM with configurable transpose modes.

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}\)

__init__

__init__(
    transpose_a=False,
    transpose_b=True,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• transpose_a (bool, default: False ) –

  Manifest params.transpose_a, bool, default False.

• transpose_b (bool, default: True ) –

  Manifest params.transpose_b, bool, default True.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    a,
    b,
    batch_sizes,
    batch_offsets,
    batch_padded_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 \([K \times \mathit{batch\_sum}]\) when transpose_a.

• b (Tensor) –

  Per-group weights, \([\mathit{batch\_count} \times N \times K]\) when transpose_a is false, or \([\mathit{batch\_sum} \times N]\) when it is.

• batch_sizes (Tensor) –

  Rows per group, 1D torch.int32.

• batch_offsets (Tensor) –

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

• batch_padded_offsets (Tensor) –

  Start row of each group in the padded output, 1D torch.int32.

Returns:

• Tensor –

  The per-group products, \([\mathit{batch\_sum} \times N]\), in the dtype of

• Tensor –

  the inputs.

Raises:

• ValueError –

  a or b is not on CUDA, their dtypes differ or are neither float16 nor bfloat16, the metadata tensors are not 1D int32 of equal length, or the operand ranks and dims disagree with the layout flags.

Example

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