Skip to main content

ClikaRT::ops::qk_rms_norm

function

qk_rms_norm()​

std::array<Tensor 3> qk_rms_norm(
    Tensor q,
    OptionalTensor k = {},
    OptionalTensor v = {},
    OptionalTensor w_q = {},
    OptionalTensor w_k = {},
    std::int64_t head_dim = 0,
    std::optional<double> eps = std::nullopt
)

Per-head RMS norm over packed attention projections: q, k and v in ONE call, no reshapes.

Each contiguous head_dim run of q (and k, when present) is its own normalization group: out = x / sqrt(mean(x^2) + eps) * w. v passes through untouched (it rides along so one call serves the projection triplet). Equivalent to reshaping [S, heads*head_dim] to [S, heads, head_dim], applying rms_norm, and reshaping back, with none of those steps.

Parameters

  • q: the query projection, [.., heads*head_dim].
  • k: optional key projection, same trailing structure.
  • v: optional value projection; passed through unchanged.
  • w_q: optional per-head gain for q, shape [head_dim].
  • w_k: optional per-head gain for k, shape [head_dim].
  • head_dim: the per-head run length; must divide the trailing extent.
  • eps: stability floor; std::nullopt selects 1e-5.

Returns: {q_normed, k_normed, v}; absent inputs come back absent.

Throws

  • ClikaRT::Error: (INVALID_ARGUMENT) when head_dim does not divide the trailing extent or a gain is mis-shaped.

Declared in ClikaRT/compute/ops.h, line 3536