Skip to main content

ClikaRT::ops::qk_rms_norm

function

qk_rms_norm()​

std::array<Tensor 3> qk_rms_norm(
    Tensor query,
    OptionalTensor key = {},
    OptionalTensor value = {},
    OptionalTensor query_weight = {},
    OptionalTensor key_weight = {},
    std::int64_t head_dim = 0,
    std::optional<double> eps = std::nullopt
)

Per-head RMS norm over packed attention projections: query, key and value in ONE call, no reshapes.

Each contiguous head_dim run of query (and key, when present) is its own normalization group: out = x / sqrt(mean(x^2) + eps) * w. value 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

  • query: the query projection, [.., heads*head_dim].
  • key: optional key projection, same trailing structure.
  • value: optional value projection; passed through unchanged.
  • query_weight: optional per-head gain for query, shape [head_dim].
  • key_weight: optional per-head gain for key, 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 3556