class

GS::MASt3R::FusedOps

Fused operations dispatcher for MASt3R

Constants

KERNEL_SOURCE_PATH = "src/metal/kernels/mastr_fused.metal"

instance

new

apply_rope_2d!(x : Tensor, freqs : Tensor, height : Int32, width : Int32, num_heads : Int32)

Apply RoPE 2D (for image patches)

attention_scores(q : Tensor, k : Tensor, num_heads : Int32, rope_cos : Tensor | Nil = nil, rope_sin : Tensor | Nil = nil) : Tensor

Fused attention scores with optional RoPE Returns: scores [batch, heads, seq, seq]

layernorm_qkv(x : Tensor, ln_weight : Tensor, ln_bias : Tensor, qkv_weight : Tensor, qkv_bias : Tensor, eps : Float32 = 1e-5_f32) : Tensor

Fused LayerNorm + QKV projection Input: x [batch, seq, embed_dim] Output: qkv [batch, seq, 3 * embed_dim]

mlp(x : Tensor, w1 : Tensor, b1 : Tensor, w2 : Tensor, b2 : Tensor) : Tensor

Fused MLP: GELU(x @ W1 + b1) @ W2 + b2

residual_layernorm(x : Tensor, residual : Tensor, weight : Tensor, bias : Tensor, eps : Float32 = 1e-5_f32) : Tensor

Fused residual + LayerNorm

softmax_inplace!(scores : Tensor)

Softmax (in-place, row-wise)