Skip to main content

ClikaRT::ops::group_query_attention_varlen

function

group_query_attention_varlen()

std::array<Tensor 3> group_query_attention_varlen(
    Tensor q,
    Tensor k,
    Tensor v,
    Tensor cu_seqlens_q,
    Tensor cu_seqlens_k,
    ScalarOrTensor max_seqlen_q = {},
    ScalarOrTensor max_seqlen_k = {},
    OptionalTensor past_key = {},
    OptionalTensor past_value = {},
    OptionalTensor kvcache_start = {},
    OptionalTensor rope_cos = {},
    OptionalTensor rope_sin = {},
    OptionalTensor position_ids = {},
    OptionalTensor attn_mask = {},
    OptionalTensor head_sink = {},
    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 q_norm_gain = {},
    OptionalTensor k_norm_gain = {},
    std::optional<double> qk_norm_eps = std::nullopt,
    OptionalTensor slot_ids = {}
)

The varlen face of the fused GQA above, with the same qk-norm tail: q_norm_gain/k_norm_gain (rank-1 [head_dim], post-rope, applied before the K/V append) travel with qk_norm_eps (the model's rms-norm epsilon) and require the in-op rope planes; absent ⇒ unchanged. The same slot_ids contract: on a continuous cache it names each batch row's cache row ([B] Int32, in range, pairwise distinct; absent = row b for batch row b), and kvcache_start's values stay unread there (the write offsets derive from cu_seqlens_k).

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