实验性算子

注意力算子

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 – 序列位置（包含每个词元的实际长度）；形状 = (B)

• qk_scale – 在 QK^T 之后应用的缩放因子

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

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

• use_tensor_cores – 是否使用张量核心 wmma 指令以实现快速计算

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

返回:

一个元组，包含组合后的 split-K 输出、未组合的 split-K 输出以及 split-K 元数据（包含最大 QK^T 值和 softmax(QK^T) 头和）