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

DeltaNet

tileops.ops.linear_attention.deltanet.DeltaNetOp

Combined DeltaNet fwd+bwd operator with autograd support (ungated).

Wraps DeltaNetFwdKernel and DeltaNetBwdKernel in a torch.autograd.Function so that output.backward(do) automatically invokes the TileOPs backward kernels.

Layout: BHSD (batch, head, seq_len, dim).

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    beta,
)

Run deltanet forward with autograd backward support.

Parameters:

• q (Tensor) –

  Query tensor [B, H, S, DK].

• k (Tensor) –

  Key tensor [B, H, S, DK].

• v (Tensor) –

  Value tensor [B, H, S, DV].

• beta (Tensor) –

  Beta tensor [B, H, S].

Returns:

• Tensor –

  Output tensor o [B, H, S, DV] (supports .backward()).

tileops.ops.linear_attention.deltanet.DeltaNetFwdOp

DeltaNet forward operator (ungated).

Pipeline: prepare_wy_repr(k, beta) -> (Aw, Au) -> deltanet_fwd(q, k, v, beta, Aw, Au) -> o.

Layout: BHSD (batch, head, seq_len, dim).

Layout convention difference with FLA

TileOPs uses BHSD layout: q/k [B, H, S, DK], v [B, H, S, DV], beta [B, H, S].

FLA (fla.ops.delta_rule.chunk_delta_rule) uses BTHN layout: q/k [B, T, H, K], v [B, T, H, V], beta [B, T, H].

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    beta,
)

Run deltanet forward.

Parameters:

• q (Tensor) –

  Query tensor [B, H, S, DK].

• k (Tensor) –

  Key tensor [B, H, S, DK].

• v (Tensor) –

  Value tensor [B, H, S, DV].

• beta (Tensor) –

  Beta tensor [B, H, S].

Returns:

• Tensor –

  Tuple of (o, S, Aw, Au).

tileops.ops.linear_attention.deltanet.DeltaNetBwdOp

DeltaNet backward operator (ungated).

Pipeline: prepare_wy_repr -> fwd (to get Aw, Au) -> bwd kernel -> (dq, dk, dv, dbeta).

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    do,
    q,
    k,
    v,
    beta,
    S,
    Aw,
    Au,
    w,
    u,
)

Run deltanet backward.

Parameters:

• do (Tensor) –

  Gradient of output [B, H, S, DV].

• q (Tensor) –

  Query tensor [B, H, S, DK].

• k (Tensor) –

  Key tensor [B, H, S, DK].

• v (Tensor) –

  Value tensor [B, H, S, DV].

• beta (Tensor) –

  Beta tensor [B, H, S].

• S (Tensor) –

  Per-chunk boundary states from forward [B, H, NC+1, DK, DV].

• Aw (Tensor) –

  A_inv matrix from forward [B, H, S, BC].

• Au (Tensor) –

  A_inv matrix from forward [B, H, S, BC].

• w (Tensor) –

  WY w vectors from forward [B, H, S, DK].

• u (Tensor) –

  WY u vectors from forward [B, H, S, DV].

Returns:

• Tuple[Tensor, Tensor, Tensor, Tensor] –

  Tuple of (dq, dk, dv, dbeta).

tileops.ops.linear_attention.deltanet_recurrence.DeltaNetDecodeFwdOp

DeltaNet decode (single-step recurrence, ungated).

Computes one step of the delta rule (no gate): v_new = beta * (v - S @ k) o = S @ q + (q . k) * v_new S_new = S + outer(k, v_new)

Layout: BHD (batch, head, dim). Supports float32, float16, and bfloat16 with fp32 accumulation.

For fp32 dtype, dispatches to a dedicated FP32 kernel that uses element-wise matvec instead of T.gemm to avoid TF32 mantissa truncation.

__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(
    q,
    k,
    v,
    beta,
    state,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float32.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• beta (Tensor) –

  Input tensor, dtype same_as(q).

• state (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tuple[Tensor, Tensor] –

  o, new_state, as the manifest declares. Shape rules: o.shape == (B, H, DV); new_state.shape == (B, H, DK, DV).

Gated DeltaNet

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetOp

Combined Gated DeltaNet fwd+bwd operator with autograd support.

Wraps GatedDeltaNetFwdKernel and GatedDeltaNetBwdKernel in a torch.autograd.Function so that output.backward(do) automatically invokes the TileOPs backward kernels.

This makes end-to-end benchmarking against FLA straightforward::

op = GatedDeltaNetOp(chunk_size=chunk_size)
o = op(q, k, v, g, beta)   # forward
o.backward(do)              # backward via TileOPs kernels

Layout: BHSD (batch, head, seq_len, dim).

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    g,
    beta,
)

Run gated deltanet forward with autograd backward support.

Parameters:

• q (Tensor) –

  Query tensor [B, H, S, DK].

• k (Tensor) –

  Key tensor [B, H, S, DK].

• v (Tensor) –

  Value tensor [B, H, S, DV].

• g (Tensor) –

  Gate tensor [B, H, S].

• beta (Tensor) –

  Beta tensor [B, H, S].

Returns:

• Tensor –

  Output tensor o [B, H, S, DV] (supports .backward()).

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetBTHDFwdOp

Gated DeltaNet forward over token-major (BTHD) inputs.

Same operator as GatedDeltaNetBHTDFwdOp and the same four outputs, over the token-major memory order the FLA reference uses: q/k [B, S, H, DK], v [B, S, H, DV], g/beta [B, S, H]. A separate entry because the memory order is part of the signature, not a mode of one signature.

It runs the warp-specialized production pipeline, so it serves Hopper with chunk_size=64, equal K/V dimensions in {64, 128}, and float16 or bfloat16. Any other call is refused, naming what it failed.

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    g,
    beta,
)

Run the token-major forward.

Parameters:

• q (Tensor) –

  Query tensor [B, S, H, DK].

• k (Tensor) –

  Key tensor [B, S, H, DK].

• v (Tensor) –

  Value tensor [B, S, H, DV].

• g (Tensor) –

  Gate tensor [B, S, H].

• beta (Tensor) –

  Beta tensor [B, S, H].

Returns:

• Tensor –

  Tuple of (o, S, Aw, Au).

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetBHTDFwdOp

Gated DeltaNet forward operator.

Pipeline: prepare_wy_repr(k, g, beta) -> (Aw, Au) -> gated_deltanet_fwd(q, k, v, g, beta, Aw, Au) -> o.

Head-major (BHTD) inputs: q/k [B, H, S, DK], v [B, H, S, DV], g/beta [B, H, S]. GatedDeltaNetBwdOp consumes the S this returns, in the same layout. Token-major callers want GatedDeltaNetBTHDFwdOp.

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    g,
    beta,
)

Run gated deltanet forward.

Parameters:

• q (Tensor) –

  Query tensor [B, H, S, DK].

• k (Tensor) –

  Key tensor [B, H, S, DK].

• v (Tensor) –

  Value tensor [B, H, S, DV].

• g (Tensor) –

  Gate tensor [B, H, S].

• beta (Tensor) –

  Beta tensor [B, H, S].

Returns:

• Tensor –

  Tuple of (o, S, Aw, Au).

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetPrefillBTHDFwdOp

Gated DeltaNet inference prefill operator.

This is the serving-oriented zero-state prefill interface: (q, k, v, g, beta) -> (o, final_state). It intentionally does not expose backward-only training artifacts such as Aw and Au. Token-major (BTHD) inputs, the FLA/Qwen convention: q/k/v/o [B, T, H, D], g/beta [B, T, H]. Head-major callers want GatedDeltaNetPrefillBHTDFwdOp. When chunk_size is not specified, the op uses a small-stream serving default: 128 for batch * heads <= 8 when the sequence length allows it, otherwise 64.

__init__

__init__(
    chunk_size=None,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (Optional[int], default: None ) –

  Manifest params.chunk_size, int | None, default None.

• 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(
    q,
    k,
    v,
    g,
    beta,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float32.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• g (Tensor) –

  Input tensor, dtype same_as(q).

• beta (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tuple[Tensor, Tensor] –

  o, final_state, as the manifest declares. Shape rules: final_state.shape == (B, H, DK, DV).

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetPrefillBHTDFwdOp

Gated DeltaNet inference prefill over head-major (BHTD) inputs.

q/k/v/o [B, H, T, D], g/beta [B, H, T] — the TileOps convention. Same kernel and same arithmetic as GatedDeltaNetPrefillBTHDFwdOp; only the memory order the tensors arrive in differs, and memory order is part of the signature, so it is its own entry.

__init__

__init__(
    chunk_size=None,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (Optional[int], default: None ) –

  Manifest params.chunk_size, int | None, default None.

• 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(
    q,
    k,
    v,
    g,
    beta,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float32.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• g (Tensor) –

  Input tensor, dtype same_as(q).

• beta (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tuple[Tensor, Tensor] –

  o, final_state, as the manifest declares. Shape rules: final_state.shape == (B, H, DK, DV).

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetDecodeFwdOp

Gated DeltaNet decode (single-step recurrence).

Computes one step of the gated delta rule

S_t = S_{t-1} (alpha_t (I - beta_t k_t k_t^T)) + beta_t v_t k_t^T o_t = S_t q_t

Layout: BHD (batch, head, dim). Supports float32, float16, and bfloat16 with fp32 accumulation.

For fp32 dtype, dispatches to a dedicated FP32 kernel that uses element-wise matvec instead of T.gemm to avoid TF32 mantissa truncation.

__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(
    q,
    k,
    v,
    g,
    beta,
    state,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float32.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• g (Tensor) –

  Input tensor, dtype same_as(q).

• beta (Tensor) –

  Input tensor, dtype same_as(q).

• state (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tuple[Tensor, Tensor] –

  o, new_state, as the manifest declares. Shape rules: o.shape == (B, H, DV); new_state.shape == (B, H, DK, DV).

tileops.ops.linear_attention.gated_deltanet.GatedDeltaNetBwdOp

Gated DeltaNet backward operator.

Pipeline: prepare_wy_repr -> fwd (to get Aw, Au) -> bwd kernel -> (dq, dk, dv, dg, dbeta).

__init__

__init__(
    chunk_size=64,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    do,
    q,
    k,
    v,
    g,
    beta,
    S,
)

Run gated deltanet backward.

Parameters:

• do (Tensor) –

  Gradient of output [B, H, S, DV].

• q (Tensor) –

  Query tensor [B, H, S, DK].

• k (Tensor) –

  Key tensor [B, H, S, DK].

• v (Tensor) –

  Value tensor [B, H, S, DV].

• g (Tensor) –

  Gate tensor [B, H, S].

• beta (Tensor) –

  Beta tensor [B, H, S].

• S (Tensor) –

  Per-chunk boundary states from forward [B, H, NC+1, DK, DV].

Returns:

• Tuple[Tensor, Tensor, Tensor, Tensor, Tensor] –

  Tuple of (dq, dk, dv, dg, dbeta).

Gated linear attention

tileops.ops.linear_attention.gla.GLAFwdOp

GLA (Gated Linear Attention) forward operator.

Chunked GLA forward: (q, k, v, g) -> (o, final_state).

Layout: BTHD (batch, seq_len, heads, dim).

__init__

__init__(
    chunk_size=64,
    scale=-1.0,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

• scale (float, default: -1.0 ) –

  Query scale factor (default: dim_k**-0.5).

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    g,
    initial_state=None,
)

Run GLA forward.

Parameters:

• q (Tensor) –

  Query tensor [B, T, H, K].

• k (Tensor) –

  Key tensor [B, T, H, K].

• v (Tensor) –

  Value tensor [B, T, H, V].

• g (Tensor) –

  Log-space forget gates [B, T, H, K].

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

  Optional fp32 initial hidden state [B, H, K, V]; absent starts the recurrence from zeros.

Returns:

• Tuple[Tensor, Tensor] –

  Tuple of (o, final_state).

tileops.ops.linear_attention.gla.GLABwdOp

GLA (Gated Linear Attention) backward operator.

Computes gradients (dq, dk, dv, dg) given output gradient do.

Uses h_out saved from the forward pass (no recomputation needed).

Layout: BTHD (batch, seq_len, heads, dim).

__init__

__init__(
    chunk_size=64,
    scale=-1.0,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• chunk_size (int, default: 64 ) –

  Chunk size for chunked linear attention.

• scale (float, default: -1.0 ) –

  Query scale factor (default: dim_k**-0.5).

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

  Optional kernel overrides.

• tune (bool, default: False ) –

  Whether to autotune kernels.

forward

forward(
    q,
    k,
    v,
    g,
    h,
    do,
    dht,
    has_initial_state=False,
)

Run GLA backward.

Parameters:

• q (Tensor) –

  Queries [B, T, H, K].

• k (Tensor) –

  Keys [B, T, H, K].

• v (Tensor) –

  Values [B, T, H, V].

• g (Tensor) –

  Log-space forget gates [B, T, H, K].

• h (Tensor) –

  Hidden states from forward [B, NT+1, H, K, V] (fp32).

• do (Tensor) –

  Output gradient [B, T, H, V].

• dht (Tensor) –

  Final-state gradient [B, H, K, V].

• has_initial_state (bool, default: False ) –

  Whether initial_state was provided by the user.

Returns:

• Tuple[Tensor, Tensor, Tensor, Tensor] –

  Tuple of (dq, dk, dv, dg).

tileops.ops.linear_attention.gla_recurrence.GLADecodeFwdOp

GLA (Gated Linear Attention) decode (single-step recurrence).

Computes one step of the gated linear attention recurrence

S_new = diag(exp(gk)) @ S + outer(k, v) o = scale * q^T @ S_new

Layout: BHD (batch, head, dim). Supports float32, float16, and bfloat16 with fp32 accumulation.

For fp32 dtype, dispatches to a dedicated FP32 kernel that uses element-wise matvec instead of T.gemm to avoid TF32 mantissa truncation.

__init__

__init__(
    scale=-1.0,
    kernel_map=None,
    tune=False,
)

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

Parameters:

• scale (float, default: -1.0 ) –

  Manifest params.scale, float, default -1.0.

• 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(
    q,
    k,
    v,
    gk,
    state,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float32.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• gk (Tensor) –

  Input tensor, dtype same_as(q).

• state (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tuple[Tensor, Tensor] –

  o, new_state, as the manifest declares. Shape rules: o.shape == (B, H, DV); new_state.shape == (B, H, DK, DV).