gpt-5 / cuda95c7fe

#include <torch/extension.h>
#include "kernel.h"

#include <vector>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <ATen/cuda/CUDAContext.h>

namespace py = pybind11;

static inline torch::Tensor to_device_contig(torch::Tensor t, c10::Device device, c10::ScalarType dtype) {
  if (t.device() == device && t.scalar_type() == dtype && t.is_contiguous()) {
    return t;
  }
  return t.to(device, dtype, /*non_blocking=*/false, /*copy=*/true).contiguous();
}

std::vector<torch::Tensor> run(
    torch::Tensor q,            // [B, 32, 128] bfloat16
    torch::Tensor k_cache,      // [num_pages, 1, 8, 128] bfloat16
    torch::Tensor v_cache,      // [num_pages, 1, 8, 128] bfloat16
    torch::Tensor kv_indptr,    // [B+1] int32
    torch::Tensor kv_indices,   // [num_kv_indices] int32
    double sm_scale_double      // default 1/sqrt(128)
) {
  TORCH_CHECK(q.dim() == 3, "q must be [B, 32, 128]");
  TORCH_CHECK(q.size(1) == 32 && q.size(2) == 128, "q must be [B, 32, 128]");
  TORCH_CHECK(k_cache.dim() == 4 && k_cache.size(1) == 1 && k_cache.size(2) == 8 && k_cache.size(3) == 128,
              "k_cache must be [num_pages, 1, 8, 128]");
  TORCH_CHECK(v_cache.dim() == 4 && v_cache.size(1) == 1 && v_cache.size(2) == 8 && v_cache.size(3) == 128,
              "v_cache must be [num_pages, 1, 8, 128]");
  TORCH_CHECK(kv_indptr.dim() == 1, "kv_indptr must be 1D");
  TORCH_CHECK(kv_indices.dim() == 1, "kv_indices must be 1D");

  const auto B = q.size(0);
  TORCH_CHECK(kv_indptr.size(0) == B + 1, "len_indptr must be batch_size + 1");

  // Determine target device: use CUDA
  c10::Device cuda_device = c10::Device(torch::kCUDA, at::cuda::current_device());

  // Move inputs to CUDA if needed and ensure correct dtype/contiguity
  auto q_dev         = to_device_contig(q, cuda_device, at::kBFloat16);
  auto k_cache_dev   = to_device_contig(k_cache, cuda_device, at::kBFloat16);
  auto v_cache_dev   = to_device_contig(v_cache, cuda_device, at::kBFloat16);
  auto kv_indptr_dev = to_device_contig(kv_indptr, cuda_device, at::kInt);
  auto kv_indices_dev= to_device_contig(kv_indices, cuda_device, at::kInt);

  // Allocate outputs on CUDA
  auto output_dev = torch::empty_like(q_dev, q_dev.options().dtype(at::kBFloat16));
  auto lse_dev    = torch::empty({B, 32}, q_dev.options().dtype(at::kFloat));

  // sm_scale
  float sm_scale = static_cast<float>(sm_scale_double);

  // Launch CUDA kernel
  gqa_paged_decode_h32_kv8_d128_ps1_cuda(
      q_dev, k_cache_dev, v_cache_dev, kv_indptr_dev, kv_indices_dev,
      sm_scale, output_dev, lse_dev
  );

  // If original inputs were on CPU, return CPU tensors; otherwise return CUDA tensors
  if (!q.is_cuda()) {
    auto output_cpu = output_dev.to(torch::kCPU, at::kBFloat16);
    auto lse_cpu    = lse_dev.to(torch::kCPU, at::kFloat);
    return {output_cpu, lse_cpu};
  } else {
    return {output_dev, lse_dev};
  }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  double default_sm_scale = 1.0 / std::sqrt(128.0);
  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") = default_sm_scale,
      "GQA paged decode kernel (h32, kv8, d128, page_size=1) optimized for B200"
  );
}