实验性算子

Attention 算子

std::tuple<at::Tensor, at::Tensor, at::Tensor> gqa_attn_splitk(const at::Tensor &XQ, const at::Tensor &cache_K, const at::Tensor &cache_V, const at::Tensor &seq_positions, const double qk_scale, const int64_t num_split_ks, const int64_t kv_cache_quant_num_groups, const bool use_tensor_cores, const int64_t cache_logical_dtype_int)

解码分组查询注意力 Split-K，使用 BF16/INT4 KV。

解码分组查询注意力 (GQA) 的 CUDA 实现，支持 BF16 和 INT4 KV 缓存以及 BF16 输入查询。目前仅支持最大上下文长度为 16384，固定头维度为 128，并且只有一个 KV 缓存头。它支持任意数量的查询头。

参数:

• XQ – 输入查询；形状 = (B, 1, H_Q, D)，其中 B = 批大小，H_Q = 查询头数，D = 头维度（固定为 128）

• cache_K – K 缓存；形状 = (B, MAX_T, H_KV, D)，其中 MAX_T = 最大上下文长度（固定为 16384），H_KV = KV 缓存头数（固定为 1）

• cache_V – V 缓存；形状 = (B, MAX_T, H_KV, D)

• seq_positions – 序列位置（包含每个 token 的实际长度）；形状 = (B)

• qk_scale – 应用于 QK^T 之后的缩放

• num_split_ks – split Ks 的数量（控制上下文长度维度 (MAX_T) 中的并行量）

• kv_cache_quant_num_groups – 用于每个 KV token 的组式 INT4 和 FP8 量化的组数（每组使用相同的缩放和偏差进行量化）。FP8 目前仅支持单组。

• use_tensor_cores – 是否使用 tensor core wmma 指令来实现快速实现

• cache_logical_dtype_int – 指定 kv_cache 的量化数据类型：{BF16:0 , FP8:1, INT4:2}

返回:

组合的 split-K 输出、非组合的 split-K 输出和 split-K 元组（包含最大 QK^T 和 softmax(QK^T) 头总和）