gpt-5 / cuda00b2dd

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cmath>
#include "kernel.h"

namespace py = pybind11;

static inline void check_inputs(
    const at::Tensor& q,
    const at::Tensor& k_cache,
    const at::Tensor& v_cache,
    const at::Tensor& kv_indptr,
    const at::Tensor& kv_indices) {

  TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor");
  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be a CUDA tensor");
  TORCH_CHECK(v_cache.is_cuda(), "v_cache must be a CUDA tensor");
  TORCH_CHECK(kv_indptr.is_cuda(), "kv_indptr must be a CUDA tensor");
  TORCH_CHECK(kv_indices.is_cuda(), "kv_indices must be a CUDA tensor");

  TORCH_CHECK(q.dtype() == at::kBFloat16, "q must be bfloat16");
  TORCH_CHECK(k_cache.dtype() == at::kBFloat16, "k_cache must be bfloat16");
  TORCH_CHECK(v_cache.dtype() == at::kBFloat16, "v_cache must be bfloat16");
  TORCH_CHECK(kv_indptr.dtype() == at::kInt, "kv_indptr must be int32");
  TORCH_CHECK(kv_indices.dtype() == at::kInt, "kv_indices must be int32");

  TORCH_CHECK(q.dim() == 3, "q must have shape [B, 32, 128]");
  TORCH_CHECK(q.size(1) == GQA_NUM_QO_HEADS, "q.num_qo_heads must be 32");
  TORCH_CHECK(q.size(2) == GQA_HEAD_DIM, "q.head_dim must be 128");

  TORCH_CHECK(k_cache.dim() == 4, "k_cache must have shape [num_pages, 1, 4, 128]");
  TORCH_CHECK(v_cache.dim() == 4, "v_cache must have shape [num_pages, 1, 4, 128]");
  TORCH_CHECK(k_cache.size(1) == GQA_PAGE_SIZE, "page_size must be 1 for k_cache");
  TORCH_CHECK(v_cache.size(1) == GQA_PAGE_SIZE, "page_size must be 1 for v_cache");
  TORCH_CHECK(k_cache.size(2) == GQA_NUM_KV_HEADS, "k_cache.num_kv_heads must be 4");
  TORCH_CHECK(v_cache.size(2) == GQA_NUM_KV_HEADS, "v_cache.num_kv_heads must be 4");
  TORCH_CHECK(k_cache.size(3) == GQA_HEAD_DIM, "k_cache.head_dim must be 128");
  TORCH_CHECK(v_cache.size(3) == GQA_HEAD_DIM, "v_cache.head_dim must be 128");

  TORCH_CHECK(kv_indptr.dim() == 1, "kv_indptr must be 1D");
  TORCH_CHECK(kv_indices.dim() == 1, "kv_indices must be 1D");

  int64_t batch_size = q.size(0);
  TORCH_CHECK(kv_indptr.size(0) == batch_size + 1, "len_indptr must equal batch_size + 1");

  // Verify num_kv_indices == kv_indptr[-1].
  // Copy small vector to CPU to avoid tricky GPU indexing in C++ frontend.
  at::Tensor kv_indptr_cpu = kv_indptr.to(torch::kCPU);
  int32_t last_offset = kv_indptr_cpu.data_ptr<int32_t>()[kv_indptr_cpu.numel() - 1];
  TORCH_CHECK(kv_indices.size(0) == static_cast<int64_t>(last_offset),
              "num_kv_indices must equal kv_indptr[-1]");
}

std::tuple<at::Tensor, at::Tensor> run(
    at::Tensor q,             // [B, 32, 128] bfloat16
    at::Tensor k_cache,       // [P, 1, 4, 128] bfloat16
    at::Tensor v_cache,       // [P, 1, 4, 128] bfloat16
    at::Tensor kv_indptr,     // [B+1] int32
    at::Tensor kv_indices,    // [kv_indptr[-1]] int32
    c10::optional<double> sm_scale_opt // optional float
) {
  check_inputs(q, k_cache, v_cache, kv_indptr, kv_indices);

  // Make tensors contiguous for predictable addressing
  q = q.contiguous();
  k_cache = k_cache.contiguous();
  v_cache = v_cache.contiguous();
  kv_indptr = kv_indptr.contiguous();
  kv_indices = kv_indices.contiguous();

  int64_t batch_size = q.size(0);
  int64_t num_pages = k_cache.size(0);
  // Default sm_scale = 1/sqrt(128)
  float sm_scale = sm_scale_opt.has_value() ? static_cast<float>(sm_scale_opt.value())
                                            : static_cast<float>(1.0 / std::sqrt(static_cast<double>(GQA_HEAD_DIM)));

  // Allocate outputs
  auto out = torch::empty({batch_size, (int64_t)GQA_NUM_QO_HEADS, (int64_t)GQA_HEAD_DIM},
                          q.options()); // bfloat16
  auto lse = torch::empty({batch_size, (int64_t)GQA_NUM_QO_HEADS},
                          q.options().dtype(torch::kFloat32)); // float32

  // Launch kernel
  auto stream = at::cuda::getCurrentCUDAStream();

  const __nv_bfloat16* q_ptr = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>());
  const __nv_bfloat16* k_ptr = reinterpret_cast<const __nv_bfloat16*>(k_cache.data_ptr<at::BFloat16>());
  const __nv_bfloat16* v_ptr = reinterpret_cast<const __nv_bfloat16*>(v_cache.data_ptr<at::BFloat16>());
  const int32_t* indptr_ptr = kv_indptr.data_ptr<int32_t>();
  const int32_t* indices_ptr = kv_indices.data_ptr<int32_t>();
  __nv_bfloat16* out_ptr = reinterpret_cast<__nv_bfloat16*>(out.data_ptr<at::BFloat16>());
  float* lse_ptr = lse.data_ptr<float>();

  gqa_paged_decode_h32_kv4_d128_ps1_launcher(
      q_ptr,
      k_ptr,
      v_ptr,
      indptr_ptr,
      indices_ptr,
      static_cast<int>(batch_size),
      static_cast<int>(num_pages),
      sm_scale,
      out_ptr,
      lse_ptr,
      stream);

  // Ensure kernel launch was successful
  auto err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess, "Kernel launch failed: ", cudaGetErrorString(err));

  return {out, lse};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run,
        py::arg("q"),
        py::arg("k_cache"),
        py::arg("v_cache"),
        py::arg("kv_indptr"),
        py::arg("kv_indices"),
        py::arg("sm_scale") = py::none(),
        "GQA paged decode kernel for (h=32, kv=4, d=128, page_size=1).");
}