Package io.github.kirstenali.deepj.tensor

Interface CausalAttentionBackend

All Superinterfaces:

TensorBackend

All Known Implementing Classes:

MetalBackend

public interface CausalAttentionBackend
extends TensorBackend

Optional backend capability for memory-efficient, fused causal attention.

Method Summary

Modifier and Type	Method	Description
CausalAttentionResult	causalAttention(Tensor query, Tensor key, Tensor value, int heads, float scale)
default CausalAttentionResult	causalAttention(Tensor query, Tensor key, Tensor value, int queryHeads, int keyValueHeads, float scale)
AttentionGradients	causalAttentionBackward(Tensor query, Tensor key, Tensor value, CausalAttentionResult saved, Tensor gradient, int heads, float scale)
default AttentionGradients	causalAttentionBackward(Tensor query, Tensor key, Tensor value, CausalAttentionResult saved, Tensor gradient, int queryHeads, int keyValueHeads, float scale)
boolean	supportsCausalAttention(Tensor query, Tensor key, Tensor value, int heads, float scale)
default boolean	supportsCausalAttention(Tensor query, Tensor key, Tensor value, int queryHeads, int keyValueHeads, float scale)

Methods inherited from interface io.github.kirstenali.deepj.tensor.TensorBackend

adamWUpdate, add, addBroadcastCols, addBroadcastRows, addInPlace, addRowVector, addScalar, addScalarInPlace, batchedMatmul, cast, causalMask, causalSoftmax, clamp, copy, copyInto, crossEntropy, crossEntropy, crossEntropyGradient, crossEntropyGradient, crossEntropyLoss, crossEntropyLoss, divide, divideBroadcastCols, divideInPlace, divideScalar, divideScalarInPlace, exp, expand, expInPlace, gelu, geluBackward, geluInPlace, headPrefixMatmul, indexAdd, indexSelect, indexSelectGradient, l2Norm, layerNormBackward, log, logInPlace, materializeTensor, matmul, maxAlongRows, maximumAxes, maximumAxesBackward, meanAlongRows, meanAxes, mergeHeads, multiply, multiplyBroadcastCols, multiplyBroadcastRows, multiplyInPlace, multiplyScalar, multiplyScalarInPlace, neg, negInPlace, normalizationBackward, normalize, permute, pow, releaseResources, releaseTemporaryResources, relu, reluBackward, reluInPlace, reshape, rmsNorm, rmsNormBackward, rotary, scatterAddRows, select, sigmoid, sigmoidInPlace, slice, sliceGradient, sliceRows, softmaxAxis, softmaxAxisBackward, softmaxBackward, softmaxBackward, softmaxRows, splitHeads, sqrt, sqrtInPlace, subtract, subtractBroadcastCols, subtractInPlace, sum, sumAbs, sumAlongCols, sumAlongRows, sumAxes, sumRows, sumToShape, swiGlu, swiGluBackward, tanh, tanhInPlace, transpose, unsqueeze, usesDeviceStorage, varianceAlongRows, varianceAxes, view, writeHeadCache, zeros, zeros

Method Details

supportsCausalAttention

boolean supportsCausalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int heads,
 float scale)

causalAttention

CausalAttentionResult causalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int heads,
 float scale)

causalAttentionBackward

AttentionGradients causalAttentionBackward(Tensor query,
 Tensor key,
 Tensor value,
 CausalAttentionResult saved,
 Tensor gradient,
 int heads,
 float scale)

supportsCausalAttention

default boolean supportsCausalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int queryHeads,
 int keyValueHeads,
 float scale)

causalAttention

default CausalAttentionResult causalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int queryHeads,
 int keyValueHeads,
 float scale)

causalAttentionBackward

default AttentionGradients causalAttentionBackward(Tensor query,
 Tensor key,
 Tensor value,
 CausalAttentionResult saved,
 Tensor gradient,
 int queryHeads,
 int keyValueHeads,
 float scale)