mistr_grouped_query_attention

A step of the mistr_native implementation. Called by mistr_mistral_block. See the file header for the source the module follows.

Usage

mistr_grouped_query_attention(
  Q,
  K,
  V,
  n_heads,
  n_kv_heads,
  mask = NULL,
  positions = NULL,
  base = 10000
)

Arguments

Q

A matrix; passed to as.matrix.

K

A matrix; passed to as.matrix.

V

A matrix; passed to as.matrix.

n_heads

A count; the body uses it as seq_len(...).

n_kv_heads

Numeric; combined arithmetically in the body.

mask

Optional; may be NULL. A matrix; indexed by row and column.

positions

Optional; may be NULL. A vector; its length is taken.

base

Passed to mistr_apply_rope. Defaults to 10000.

Value

The value of out, as built in the body.

Examples

set.seed(11)
Q <- matrix(rnorm(16), 4, 4); K <- matrix(rnorm(8), 4, 2); V <- matrix(rnorm(8), 4, 2)
dim(mistr_grouped_query_attention(Q, K, V, n_heads = 2, n_kv_heads = 1))
#> [1] 4 4