Skip to main content

ClikaRT::ops::scaled_dot_product_attention

function

scaled_dot_product_attention()​

Tensor scaled_dot_product_attention(
    Tensor query,
    Tensor key,
    Tensor value,
    OptionalTensor attn_mask = {},
    bool is_causal = false,
    ScalarOrTensor q_scale = {},
    ScalarOrTensor k_scale = {},
    ScalarOrTensor v_scale = {}
)

Scaled dot-product attention over dense head-major tensors.

out=softmax ⁣(s (qk⊤)+mask) v\mathrm{out} = \mathrm{softmax}\!\bigl(s\,(q k^\top) + \mathrm{mask}\bigr)\,v

s defaults to 1/D1/\sqrt{D} (the q/k head size) and is replaced wholesale by q_scale when given. Layout is head-major, D innermost.

Parameters

  • query: queries [B, H_q, S_q, D].
  • key: keys [B, H_kv, S_kv, D] (H_kv divides H_q; grouped KV heads are broadcast to their query group).
  • value: values [B, H_kv, S_kv, Dv].
  • attn_mask: optional, broadcastable to [B, H_q, S_q, S_kv]; Bool = keep-mask, float = additive logits.
  • is_causal: apply the causal (lower-triangular) mask. Default false.
  • q_scale: optional softmax pre-scale: a scalar / 0-D replaces the 1/sqrt(D) default; a [H_q] tensor gives one PER-HEAD scale that carries the WHOLE pre-scale (fold 1/sqrt(D) in yourself). Absent = the default.
  • k_scale: optional dequantization scale for quantized keys, folded into the logits together with q_scale.
  • v_scale: optional dequantization scale for quantized values; it multiplies the value side.

Throws

  • ClikaRT::Error: when the input's dtype/shape is not served for this operation (the machine-readable reason rides code_name()).

Returns: [B, H_q, S_q, Dv], dtype of query.

auto out = ClikaRT::ops::scaled_dot_product_attention(query, key, value, {}, /*is_causal=*/true);

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