std1.0.0-beta-7 SafeC docs

std :: ml::gpu_rocm

Function std::rocm_linear_attention

#include <std/ml/gpu_rocm.h> source · definition

int rocm_linear_attention(
    struct RocmBuf q,
    struct RocmBuf k,
    struct RocmBuf v,
    struct RocmBuf out,
    unsigned long heads,
    unsigned long seqLen,
    unsigned long kvLen,
    unsigned long headDim,
    int causal,
)

Linear attention with phi(x) = elu(x) + 1, as attention.h's linear_attention_forward: out_i = phi(q_i) S / phi(q_i) z with S, z the sums of phi(k_j)^T v_j and phi(k_j). Chunked (128 rows on matrix cores, 32 otherwise): per-chunk states, a prefix over chunks, then each chunk's rows -- so causal (j <= i, which needs kvLen == seqLen) runs in parallel too. Cost O(n d^2) instead of O(n^2 d). The chunk states live in a scratch buffer this library keeps and grows; during graph capture a call that would need to grow it fails.