40 lines
1.4 KiB
Plaintext
40 lines
1.4 KiB
Plaintext
// torch binding for gqa_decode_attn
|
|
// kernel defined in gqa_decode_attn.cuh
|
|
#include "gqa_decode_attn.cuh"
|
|
#include <torch/extension.h>
|
|
|
|
torch::Tensor gqa_decode_attn(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor mask
|
|
) {
|
|
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && mask.is_cuda());
|
|
TORCH_CHECK(q.dtype() == torch::kBFloat16);
|
|
TORCH_CHECK(k.dtype() == torch::kBFloat16);
|
|
TORCH_CHECK(v.dtype() == torch::kBFloat16);
|
|
TORCH_CHECK(mask.dtype() == torch::kBool);
|
|
TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1");
|
|
|
|
int B = q.size(0), n_heads = q.size(1), n_kv = k.size(1);
|
|
int seq_len = k.size(2), hd = q.size(3);
|
|
TORCH_CHECK(hd % 32 == 0, "head_dim must be multiple of 32");
|
|
int group_size = n_heads / n_kv;
|
|
auto out = torch::empty_like(q);
|
|
|
|
size_t smem = DC_CHUNK * hd * sizeof(bf16);
|
|
dim3 block(32, group_size);
|
|
dim3 grid(B * n_kv);
|
|
|
|
gqa_decode_attn_kernel<<<grid, block, smem>>>(
|
|
reinterpret_cast<const bf16*>(q.data_ptr()),
|
|
reinterpret_cast<const bf16*>(k.data_ptr()),
|
|
reinterpret_cast<const bf16*>(v.data_ptr()),
|
|
mask.data_ptr<bool>(),
|
|
reinterpret_cast<bf16*>(out.data_ptr()),
|
|
B, n_heads, n_kv, seq_len, hd
|
|
);
|
|
return out;
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode v2 (per-KV-head, shared K)");
|
|
}
|