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

GateConv

tileops.sequence_modeling.EngramGateConvFwdOp

Engram GateConv forward operator (post-projection fusion).

Assumes k = E @ W_K and v = E @ W_V have been computed externally via standard GEMM. This op fuses the remaining memory-bound stages: RMSNorm gating -> causal DWConv1D -> SiLU + residual

Returns Y plus saved intermediates for backward (strategy B): vhat, alpha, rrms_h, rrms_k, rrms_v

__init__

__init__(
    M,
    seq_len,
    d,
    eps=1e-06,
    *,
    target=None
)

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

Parameters:

• M (int) –

  Batch size.

• seq_len (int) –

  Sequence length.

• d (int) –

  Model hidden dimension.

• eps (float, default: 1e-06 ) –

  RMSNorm epsilon (default 1e-6).

• 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(
    H,
    k,
    v,
    rms_w_h,
    rms_w_v,
    conv_w,
)

Parameters:

• H (Tensor) –

  (M, seq_len, d) — hidden states from preceding layers.

• k (Tensor) –

  (M, seq_len, d) — key projection output (E @ W_K).

• v (Tensor) –

  (M, seq_len, d) — value projection output (E @ W_V).

• rms_w_h (Tensor) –

  (d,) — RMSNorm weight for H and k.

• rms_w_v (Tensor) –

  (d,) — RMSNorm weight for gated value (conv input).

• conv_w (Tensor) –

  (4, d) — depthwise causal conv1d weights.

Returns:

• List[Tensor] –

  List of [Y, vhat, alpha, rrms_h, rrms_k, rrms_v]: Y: (M, seq_len, d) — output to add as residual to H. vhat: (M, seq_len, d) — saved for backward. alpha: (M, seq_len) — scalar gate, saved for backward. rrms_h: (M, seq_len) — RMSNorm reciprocal rms of H. rrms_k: (M, seq_len) — RMSNorm reciprocal rms of k. rrms_v: (M, seq_len) — RMSNorm reciprocal rms of v_hat.

tileops.sequence_modeling.EngramGateConvBwdOp

Engram GateConv backward operator.

Given dY and saved intermediates from forward (vhat, alpha, rrms_*), computes gradients for H, k, v, rms_w_h, rms_w_v, conv_w.

The caller is responsible for the projection backward (GEMM): dE = dk @ W_K^T + dv @ W_V^T dW_K = E^T @ dk dW_V = E^T @ dv

__init__

__init__(
    M,
    seq_len,
    d,
    eps=1e-06,
    *,
    target=None
)

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

Parameters:

• M (int) –

  Batch size.

• seq_len (int) –

  Sequence length.

• d (int) –

  Model hidden dimension.

• eps (float, default: 1e-06 ) –

  RMSNorm epsilon (default 1e-6).

• 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(
    dY,
    H,
    k,
    v,
    rms_w_h,
    rms_w_v,
    conv_w,
    vhat,
    alpha,
    rrms_h,
    rrms_k,
    rrms_v,
)

Parameters:

• dY (Tensor) –

  (M, T, d) — gradient of output Y.

• H (Tensor) –

  (M, T, d) — hidden states (forward input).

• k (Tensor) –

  (M, T, d) — key projection (forward input, or recomputed).

• v (Tensor) –

  (M, T, d) — value projection (forward input).

• rms_w_h (Tensor) –

  (d,) — RMSNorm weight for H and k.

• rms_w_v (Tensor) –

  (d,) — RMSNorm weight for v_hat.

• conv_w (Tensor) –

  (4, d) — conv weights.

• vhat (Tensor) –

  (M, T, d) — saved from forward.

• alpha (Tensor) –

  (M, T) — saved from forward.

• rrms_h (Tensor) –

  (M, T) — saved from forward.

• rrms_k (Tensor) –

  (M, T) — saved from forward.

• rrms_v (Tensor) –

  (M, T) — saved from forward.

Returns:

• List[Tensor] –

  List of [dH, dk, dv, drms_w_h, drms_w_v, dconv_w]: dH: (M, T, d) dk: (M, T, d) dv: (M, T, d) drms_w_h: (d,) — fp32 drms_w_v: (d,) — fp32 dconv_w: (4, d) — fp32

Decode

tileops.sequence_modeling.EngramDecodeFwdOp

Engram fused decode operator — single-token inference.

Fuses GEMV projection + gating + dilated conv + SiLU into one kernel call.

Like MHA decode compiles with max_seqlen_kv and accepts real_seqlen_kv at runtime, this op compiles with max_conv_len and accepts the actual conv_state length at runtime without recompilation.

__init__

__init__(
    batch,
    d_mem,
    d,
    max_conv_len,
    conv_kernel_size,
    dilation,
    eps=1e-06,
    *,
    target=None
)

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

Parameters:

• batch (int) –

  Batch size.

• d_mem (int) –

  Memory embedding dimension.

• d (int) –

  Model hidden dimension.

• max_conv_len (int) –

  Max conv cache capacity (compile-time).

• conv_kernel_size (int) –

  Number of conv taps w (model param, e.g. 4).

• dilation (int) –

  Dilation factor δ (model param, e.g. max N-gram order).

• eps (float, default: 1e-06 ) –

  RMSNorm epsilon (default 1e-6).

• 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(
    e_t,
    h_t,
    conv_state,
    W_K,
    W_V,
    rms_w_h,
    rms_w_v,
    conv_w,
)

Parameters:

• e_t (Tensor) –

  (B, d_mem) — gathered N-gram embedding for current token.

• h_t (Tensor) –

  (B, d) — hidden state for current token.

• conv_state (Tensor) –

  (B, L, d) — conv history, L <= max_conv_len. Left-padded to max_conv_len internally if L < max_conv_len.

• W_K (Tensor) –

  (d_mem, d) — key projection weight.

• W_V (Tensor) –

  (d_mem, d) — value projection weight.

• rms_w_h (Tensor) –

  (d,) — RMSNorm weight for h and k.

• rms_w_v (Tensor) –

  (d,) — RMSNorm weight for v_hat.

• conv_w (Tensor) –

  (w, d) — depthwise conv weights (w = conv_kernel_size).

Returns:

• List[Tensor] –

  [y_t, new_conv_state]: y_t: (B, d) — output to add as residual. new_conv_state: (B, max_conv_len, d) — updated state for next step.