实验性运算符

注意力运算符

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)

带 BF16/INT4 KV 的解码分组查询注意力 Split-K。

解码分组查询注意力 (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 K 的数量（控制上下文长度维度 (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) 头和）