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.linear_attention.DeltaNetChunkFwdOp ¶

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,
    *,
    target=None
)

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

Parameters:

chunk_size (int, default: 64 ) –

Chunk size for chunked linear attention.
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,
    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:

Tuple[Tensor, ...] –

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

tileops.linear_attention.DeltaNetChunkBwdOp ¶

DeltaNet backward operator (ungated).

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

init ¶

__init__(
    chunk_size=64,
    *,
    target=None
)

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

Parameters:

chunk_size (int, default: 64 ) –

Chunk size for chunked linear attention.
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(
    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.linear_attention.DeltaNetRecurrentFwdOp ¶

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__(
    *,
    target=None
)

Build the op. Shapes and dtype 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(
    q,
    k,
    v,
    beta,
    state,
)

Run one decode step.

Parameters:

q (Tensor) –

Query [B, H, DK].
k (Tensor) –

Key [B, H, DK].
v (Tensor) –

Value [B, H, DV].
beta (Tensor) –

Delta-rule step size [B, H].
state (Tensor) –

Recurrent state [B, H, DK, DV].

Returns:

Tuple[Tensor, Tensor] –

o [B, H, DV] and new_state [B, H, DK, DV].

tileops.linear_attention.DeltaNetFwdOp ¶

Inference forward for the ungated delta rule.

The input layout is [B, T, H, D]. One call covers equal-length prefill, packed-varlen prefill, and single-token decode. The recurrent state is FP32 and belongs to the caller: initial_state is optional, while (o, final_state) is always returned.

The in-tree implementations serve a 64- or 128-wide square state in float16 or bfloat16, with an even or single head count, and with or without use_qk_l2norm_in_kernel. Prefill runs on SM89 and SM90, equal-length or packed, with a sequence that is not a whole number of 64-token chunks. On SM89, shared memory holds the offsets of at most 703 packed sequences over a 128-wide state. Decode runs one token per sequence on SM80, SM89 and SM90, and does not take cu_seqlens. A key width that is neither 64 nor 128, or one that differs from the value width, has no in-tree kernel and needs an external target implementation.

beta contains the already-transformed update strength. This Op does not apply a sigmoid or another beta transform.

init ¶

__init__(
    scale=None,
    use_qk_l2norm_in_kernel=False,
    *,
    target=None
)

Configure the attention scale, Q/K normalization, and target.

Parameters:

scale (Optional[float], default: None ) –

Query scale, or None for K**-0.5.
use_qk_l2norm_in_kernel (bool, default: False ) –

Normalize Q and K internally.
target (Target, default: None ) –

Backend target, or None to resolve from the input device.

forward ¶

forward(
    q,
    k,
    v,
    beta,
    initial_state=None,
    cu_seqlens=None,
    cu_seqlens_cpu=None,
)

Run prefill or decode and return (o, final_state).

Gated DeltaNet¶

tileops.linear_attention.GDNFwdOp ¶

Gated DeltaNet (GDN): inference forward for the gated delta rule.

q and k use [B, T, H, K]. v, g, beta and the output use HV recurrent heads, where HV is a multiple of H. The recurrent state is always FP32. It is key-major [N, HV, K, V] by default and value-major [N, HV, V, K] when state_v_first=True. For equal-length inputs, N == B. Passing cu_seqlens selects packed varlen mode: the token tensors have B == 1 and N is the number of packed sequences.

By default, g contains the precomputed log-space decay and beta contains the post-sigmoid update strength. use_gate_in_kernel=True instead treats g as the raw gate and requires A_log and dt_bias; use_beta_sigmoid_in_kernel=True treats beta as raw logits. With allow_neg_eigval=True, the beta transform is 2 * sigmoid(beta) rather than sigmoid(beta).

One interface covers prefill and decode, and always returns (o, final_state). This intentionally pins FLA's output_final_state=True because inference callers need the state to continue recurrence. The target callable selects the execution path from the current inputs; in particular, T == 1 is decode rather than a separate public Op.

The in-tree kernels serve a 64- or 128-wide square state in either layout, with HV a multiple of H, an even or single head count, and any combination of the three input transforms. Prefill runs on SM89 and SM90, equal-length or packed, with a sequence that is not a whole number of 64-token chunks. On SM89, shared memory holds the offsets of at most 703 packed sequences over a 128-wide state. Decode runs one token per sequence on SM80, SM89 and SM90, and does not take cu_seqlens. Other calls need an external target implementation.

init ¶

__init__(
    scale=None,
    use_qk_l2norm_in_kernel=False,
    use_beta_sigmoid_in_kernel=False,
    allow_neg_eigval=False,
    state_v_first=False,
    use_gate_in_kernel=False,
    *,
    target=None
)

Fix recurrence semantics; tensor metadata comes from each call.

Parameters:

scale (Optional[float], default: None ) –

Query scale, or None for K**-0.5.
use_qk_l2norm_in_kernel (bool, default: False ) –

Normalize Q and K internally.
use_beta_sigmoid_in_kernel (bool, default: False ) –

Treat beta as raw logits and apply sigmoid internally.
allow_neg_eigval (bool, default: False ) –

Compute 2 * sigmoid(beta) instead of sigmoid(beta). Valid only with use_beta_sigmoid_in_kernel=True.
state_v_first (bool, default: False ) –

Use value-major recurrent state layout [N, HV, V, K] instead of key-major [N, HV, K, V].
use_gate_in_kernel (bool, default: False ) –

Treat g as a raw gate and internally compute -exp(A_log) * softplus(g + dt_bias). Otherwise, g must already contain the log-space decay.
target (Target, default: None ) –

Backend target, or None to resolve from the input device.

forward ¶

forward(
    q,
    k,
    v,
    g,
    beta,
    initial_state=None,
    cu_seqlens=None,
    cu_seqlens_cpu=None,
    A_log=None,
    dt_bias=None,
)

Run prefill or decode and return (o, final_state).

Kimi Delta Attention¶

tileops.linear_attention.KDAFwdOp ¶

Kimi Delta Attention (KDA): the gated delta rule with a per-key-channel decay.

q and k use [B, T, H, K]. v, beta and the output use HV recurrent heads, where HV is a multiple of H; g carries one log-space decay per key channel, [B, T, HV, K], which is what separates this operator from the scalar per-head gate of GDNFwdOp. The recurrent state is always FP32 and key-major [N, HV, K, V]; state_v_first=True asks for the value-major layout. For equal-length inputs N == B. Passing cu_seqlens selects packed varlen mode: the token tensors have B == 1 and N is the number of packed sequences.

By default g contains the precomputed log-space decay and beta the post-sigmoid update strength. use_gate_in_kernel=True instead treats g as the raw gate and requires A_log; use_beta_sigmoid_in_kernel=True treats beta as raw logits, and with allow_neg_eigval=True the beta transform is 2 * sigmoid(beta).

One interface covers equal-length prefill, packed varlen and single-token decode, and always returns (o, final_state): an inference caller needs the state to continue the recurrence. The target callable selects the execution path from the current inputs, so T == 1 is decode rather than a separate public Op.

The in-tree implementations cover SM90 prefill and decode over a 64- or 128-wide square state with precomputed gate and beta values. The raw-gate and raw-beta variants, the value-major state and other state widths still require an external target implementation.

init ¶

__init__(
    scale=None,
    use_qk_l2norm_in_kernel=False,
    use_beta_sigmoid_in_kernel=False,
    allow_neg_eigval=False,
    state_v_first=False,
    use_gate_in_kernel=False,
    lower_bound=None,
    *,
    target=None
)

Fix recurrence semantics; tensor metadata comes from each call.

Parameters:

scale (Optional[float], default: None ) –

Query scale, or None for K**-0.5.
use_qk_l2norm_in_kernel (bool, default: False ) –

Normalize Q and K internally.
use_beta_sigmoid_in_kernel (bool, default: False ) –

Treat beta as raw logits and apply sigmoid internally.
allow_neg_eigval (bool, default: False ) –

Compute 2 * sigmoid(beta) instead of sigmoid(beta). Valid only with use_beta_sigmoid_in_kernel=True.
state_v_first (bool, default: False ) –

Use value-major recurrent state layout [N, HV, V, K] instead of key-major [N, HV, K, V].
use_gate_in_kernel (bool, default: False ) –

Treat g as a raw gate and internally compute the log-space decay from A_log and dt_bias. Otherwise g must already contain it.
lower_bound (Optional[float], default: None ) –

Lower bound of the forget gate in log space, or None. Only used with use_gate_in_kernel=True.
target (Target, default: None ) –

Backend target, or None to resolve from the input device.

forward ¶

forward(
    q,
    k,
    v,
    g,
    beta,
    initial_state=None,
    cu_seqlens=None,
    cu_seqlens_cpu=None,
    A_log=None,
    dt_bias=None,
)

Run prefill or decode and return (o, final_state).

Gated Linear Attention¶

tileops.linear_attention.GLAChunkFwdOp ¶

Gated Linear Attention (GLA) 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,
    *,
    target=None
)

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

Parameters:

chunk_size (int, default: 64 ) –

Chunk size for chunked linear attention.
scale (float, default: -1.0 ) –

Query scale factor; a non-positive value means dim_k ** -0.5.
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,
    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.linear_attention.GLAChunkBwdOp ¶

Gated Linear Attention (GLA) 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,
    has_initial_state=False,
    *,
    target=None
)

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

Parameters:

chunk_size (int, default: 64 ) –

Chunk size for chunked linear attention.
scale (float, default: -1.0 ) –

Query scale factor; a non-positive value means dim_k ** -0.5.
has_initial_state (bool, default: False ) –

Whether the forward this backward pairs with was given an initial state.
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,
    k,
    v,
    g,
    h,
    do,
    dht,
)

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

Returns:

Tuple[Tensor, Tensor, Tensor, Tensor] –

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

tileops.linear_attention.GLARecurrentFwdOp ¶

Gated Linear Attention (GLA) 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,
    *,
    target=None
)

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

Parameters:

scale (float, default: -1.0 ) –

Query scale; a non-positive value means DK ** -0.5.
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,
    k,
    v,
    gk,
    state,
)

Run one decode step.

Parameters:

q (Tensor) –

Query [B, H, DK].
k (Tensor) –

Key [B, H, DK].
v (Tensor) –

Value [B, H, DV].
gk (Tensor) –

Log-space key gate [B, H, DK].
state (Tensor) –

Recurrent state [B, H, DK, DV].

Returns:

Tuple[Tensor, Tensor] –

o [B, H, DV] and new_state [B, H, DK, DV].

tileops.linear_attention.GLAFwdOp ¶

Gated Linear Attention (GLA) for inference, with caller-owned FP32 state.

Q, K, V and the log-space, per-key gate G use FP16/BF16 BTHD layout. One call is equal-length prefill, packed-varlen prefill, or single-token decode. An absent initial_state starts from zero; every call returns (o, final_state). The in-tree kernels serve equal-length prefill, packed-varlen prefill and decode.

init ¶

__init__(
    scale=None,
    *,
    target=None
)

Fix the query scale and optional backend target for this instance.

Parameters:

scale (Optional[float], default: None ) –

Positive query scale, or None for K**-0.5.
target (Target, default: None ) –

Backend target, or None to resolve from the input device.

forward ¶

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

Run GLA prefill or decode and return output plus final FP32 state.

Linear Attention Operators¶

DeltaNet¶

tileops.linear_attention.DeltaNetChunkFwdOp ¶

__init__ ¶

forward ¶

tileops.linear_attention.DeltaNetChunkBwdOp ¶

__init__ ¶

forward ¶

tileops.linear_attention.DeltaNetRecurrentFwdOp ¶

__init__ ¶

forward ¶

tileops.linear_attention.DeltaNetFwdOp ¶

__init__ ¶

forward ¶

Gated DeltaNet¶

tileops.linear_attention.GDNFwdOp ¶

__init__ ¶

forward ¶

Kimi Delta Attention¶

tileops.linear_attention.KDAFwdOp ¶

__init__ ¶

forward ¶

Gated Linear Attention¶

tileops.linear_attention.GLAChunkFwdOp ¶

__init__ ¶

forward ¶

tileops.linear_attention.GLAChunkBwdOp ¶

__init__ ¶

forward ¶

tileops.linear_attention.GLARecurrentFwdOp ¶

__init__ ¶

forward ¶

tileops.linear_attention.GLAFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶