Expand description
Small host-side tensor kernels used by the native SAM3 bring-up.
Compute-heavy ops go through rlx_cpu::blas::sgemm_auto so we get
Accelerate/OpenBLAS instead of naive Rust loops. The pure-Rust path
falls back to a scalar gemm with the same signature (see
rlx-cpu/src/blas.rs).
Functionsยง
- gelu_
tanh - layer_
norm - linear
out = x @ w_t + bwherew_tis[in_dim, out_dim](already transposed at load time byWeightMap::take_transposed).- matmul
out = a @ bwith default row-major no-transpose strides.- matmul_
bt out = alpha * a @ b^T(no bias).bis[n, k].- multihead_
attention - Multi-head attention forward for an upstream
nn.MultiheadAttentionwith the standardin_proj_weight [3*E, E]/out_proj.weight [E, E]parameterisation. Inputs are batch-first[B, L_q, E]/[B, L_k, E].key_padding_mask(when present) is[B, L_k]withtruemeaning the position must be ignored. - sigmoid
- softmax_
rows - In-place row-wise softmax: each
cols-long row is normalised independently. Numerically stable.