From b0cdd5af7483759dab45fde980459b6985b165b6 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 21:14:15 +0000 Subject: [PATCH] fix: extern declarations for gather_swa functions in gather_kv.cu --- dsv4/kernels/cuda/gather_kv.cu | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/dsv4/kernels/cuda/gather_kv.cu b/dsv4/kernels/cuda/gather_kv.cu index 7ea27878..6ed5cda3 100644 --- a/dsv4/kernels/cuda/gather_kv.cu +++ b/dsv4/kernels/cuda/gather_kv.cu @@ -101,6 +101,17 @@ void gather_kv_cuda( } +// Forward declarations for gather_swa.cu functions +// (both files compiled together into one .so) +extern void gather_swa_cuda( + torch::Tensor swa_fp8, torch::Tensor swa_rope, torch::Tensor swa_inv, + torch::Tensor swa_pos, torch::Tensor request_slots, torch::Tensor output, + int64_t head_dim, int64_t rope_dim); +extern void gather_all_compressed_cuda( + torch::Tensor entries_fp8, torch::Tensor entries_rope, torch::Tensor inv_scale, + torch::Tensor block_table, torch::Tensor block_lens, torch::Tensor output, + int64_t entries_per_block, int64_t rope_dim); + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("gather_kv", &gather_kv_cuda, "Gather KV entries into dense tile"); // gather_swa + gather_all_compressed are defined in gather_swa.cu