Skip to main content

ClikaRT::ops::scaled_dot_product_attention_varlen

function

scaled_dot_product_attention_varlen()​

Tensor scaled_dot_product_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 attn_mask = {},
    bool is_causal = false,
    ScalarOrTensor q_scale = {},
    ScalarOrTensor k_scale = {},
    ScalarOrTensor v_scale = {}
)

Variable-length (packed) scaled dot-product attention: ragged batches ride one token-packed tensor plus prefix-sum offsets, no padding.

Same math as scaled_dot_product_attention; the batch structure moves into cu_seqlens_*.

Parameters

  • q: packed queries [Sum(S_q), H_q, D] (token-major).
  • k: packed keys [Sum(S_kv), H_kv, D].
  • v: packed values [Sum(S_kv), H_kv, Dv].
  • cu_seqlens_q: [B+1] Int32 prefix sums of the per-sequence query lengths.
  • cu_seqlens_k: [B+1] Int32 prefix sums of the per-sequence KV lengths.
  • max_seqlen_q: optional host bound on the longest query run (an integer, or a HOST-resident 0-D Int32/Int64 tensor).
  • max_seqlen_k: optional host bound on the longest KV run, same forms.
  • attn_mask: optional, broadcastable to [B, H_q, S_q, S_kv].
  • is_causal: apply the causal mask per sequence. Default false.
  • q_scale: as in scaled_dot_product_attention.
  • k_scale: as in scaled_dot_product_attention.
  • v_scale: as in scaled_dot_product_attention.

Throws

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

Returns: packed [Sum(S_q), H_q, Dv], dtype of q.

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