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

Multi-head compression

tileops.ops.sequence_modeling.mhc.MHCPreFwdOp

Layout: BSHD

__init__

__init__(
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from 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(
    phi,
    x,
    b,
    alpha_pre,
    alpha_post,
    alpha_res,
    sinkhorn_repeat,
    sinkhorn_eps=0.02,
)

Run the op on the inputs the manifest declares.

Parameters:

• phi (Tensor) –

  Input tensor, dtype float32.

• x (Tensor) –

  Input tensor, dtype bfloat16.

• b (Tensor) –

  Input tensor, dtype float32.

Returns:

• Tensor –

  x_res, x_layer, as the manifest declares.

tileops.ops.sequence_modeling.mhc.MHCPostFwdOp

Layout: BSHD

__init__

__init__(
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from 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(
    x_layer_out,
    h_post,
    x_res,
)

Run the op on the inputs the manifest declares.

Parameters:

• x_layer_out (Tensor) –

  Input tensor, dtype bfloat16.

• h_post (Tensor) –

  Input tensor, dtype float32.

• x_res (Tensor) –

  Input tensor, dtype bfloat16.

Returns:

• Tensor –

  x_out, as the manifest declares.