ClikaRT::ops::group_query_attention

function

group_query_attention()

std::array<Tensor 3> group_query_attention(
    Tensor q,
    Tensor k,
    Tensor v,
    OptionalTensor past_key = {},
    OptionalTensor past_value = {},
    OptionalTensor kvcache_start = {},
    OptionalTensor rope_cos = {},
    OptionalTensor rope_sin = {},
    OptionalTensor position_ids = {},
    OptionalTensor attn_mask = {},
    std::optional<bool> is_causal = std::nullopt,
    ScalarOrTensor q_scale = {},
    std::optional<double> softcap = std::nullopt,
    std::optional<int64_t> sliding_window = std::nullopt,
    std::optional<bool> smooth_softmax = std::nullopt,
    std::optional<RotaryMode> rotary_mode = std::nullopt,
    std::optional<int64_t> num_heads = std::nullopt,
    std::optional<int64_t> kv_num_heads = std::nullopt,
    OptionalTensor out_present_key = {},
    OptionalTensor out_present_value = {},
    ScalarOrTensor k_scale = {},
    ScalarOrTensor v_scale = {},
    OptionalTensor head_sink = {},
    OptionalTensor q_norm_gain = {},
    OptionalTensor k_norm_gain = {},
    std::optional<double> qk_norm_eps = std::nullopt,
    OptionalTensor slot_ids = {}
)

Fused GQA with RoPE + KV cache. Bind out_present_key/out_present_value to the same buffers as past_key/past_value (a KVCache::keys/values(layer) view) + pass kvcache_start to append the new post-RoPE K/V IN PLACE into the cache (the decode perf path). q/k/v are hidden-folded [ΣS, heads*head_dim]; num_heads/kv_num_heads drive the in-op head split (GQA). head_sink is the per-head softmax sink [H_q], a virtual logit folded into the softmax denominator (attention-sink models bind one per layer), the same contract as the varlen face; it rides the parameter tail here. q_norm_gain/k_norm_gain engage the POST-rope per-head RMS norm: after the in-op rotation, every head's [head_dim] q (and new-k) vector is RMS-normalized with the gain BEFORE any cache append, so the cache holds rotated+normed keys. Each gain is rank-1 [head_dim] (one vector shared across heads; any other shape rejects), any float dtype, applied at its own dtype. The gains require the in-op rope planes (rope_cos/rope_sin), and they travel WITH qk_norm_eps: pass the model's own rms-norm epsilon alongside the gains, or neither (a gain without the epsilon, or an epsilon with no gain, rejects). Absent ⇒ the gain-less path, unchanged. A rope-free per-head norm composes qk_rms_norm instead. slot_ids ([B] Int32) names the cache row each batch row appends to and attends from on a continuous [max_seqs, H_kv, max_seq, D] cache: a sequence keeps its row while the batch composition changes around it. Absent, batch row b uses cache row b. Every entry must lie in [0, max_seqs) and no two rows may share one (each rejects). On that cache kvcache_start stays the rank-1 [B] layout selector whose values are not read: each row's write offset is cu_seqlens_k[b] - q_len[b]. A paged block table and the dense in-place form take no slot_ids (the block table is its own row map; the dense form addresses rows by batch index).

Returns: {out_attn, present_key, present_value}.

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