gpt-5 / cudac0b7b7

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/util/Optional.h>
#include <c10/util/BFloat16.h>

#include "kernel.h"

#include <vector>
#include <stdexcept>
#include <cmath>

namespace py = pybind11;

#define CHECK_TENSOR(t) TORCH_CHECK((t).is_cuda(), #t " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(t) TORCH_CHECK((t).is_contiguous(), #t " must be contiguous")
#define CHECK_DTYPE(t, dt) TORCH_CHECK((t).scalar_type() == (dt), #t " has wrong dtype")

static inline float default_sm_scale() {
  // 1 / sqrt(128)
  return 1.0f / std::sqrt(128.0f);
}

std::vector<torch::Tensor> run(torch::Tensor q,
                               torch::Tensor k,
                               torch::Tensor v,
                               torch::Tensor qo_indptr,
                               torch::Tensor kv_indptr,
                               c10::optional<double> sm_scale_opt) {
  // Validate device and layout
  CHECK_TENSOR(q);
  CHECK_TENSOR(k);
  CHECK_TENSOR(v);
  CHECK_TENSOR(qo_indptr);
  CHECK_TENSOR(kv_indptr);

  CHECK_CONTIGUOUS(q);
  CHECK_CONTIGUOUS(k);
  CHECK_CONTIGUOUS(v);
  CHECK_CONTIGUOUS(qo_indptr);
  CHECK_CONTIGUOUS(kv_indptr);

  CHECK_DTYPE(q, torch::kBFloat16);
  CHECK_DTYPE(k, torch::kBFloat16);
  CHECK_DTYPE(v, torch::kBFloat16);
  CHECK_DTYPE(qo_indptr, torch::kInt32);
  CHECK_DTYPE(kv_indptr, torch::kInt32);

  TORCH_CHECK(q.dim() == 3, "q must be [total_q, 32, 128]");
  TORCH_CHECK(k.dim() == 3, "k must be [total_kv, 4, 128]");
  TORCH_CHECK(v.dim() == 3, "v must be [total_kv, 4, 128]");
  TORCH_CHECK(q.size(1) == 32 && q.size(2) == 128, "q last dims must be [32, 128]");
  TORCH_CHECK(k.size(1) == 4 && k.size(2) == 128, "k last dims must be [4, 128]");
  TORCH_CHECK(v.size(1) == 4 && v.size(2) == 128, "v last dims must be [4, 128]");
  TORCH_CHECK(qo_indptr.dim() == 1, "qo_indptr must be 1D");
  TORCH_CHECK(kv_indptr.dim() == 1, "kv_indptr must be 1D");
  TORCH_CHECK(qo_indptr.size(0) == kv_indptr.size(0), "qo_indptr and kv_indptr must have the same length");

  const int64_t len_indptr = qo_indptr.size(0);
  const int64_t total_q = q.size(0);
  const int64_t total_kv = k.size(0);

  // Constraints: totals equal to last element of indptr
  auto qo_indptr_cpu = qo_indptr.cpu();
  auto kv_indptr_cpu = kv_indptr.cpu();
  const int32_t total_q_chk = qo_indptr_cpu.data_ptr<int32_t>()[len_indptr - 1];
  const int32_t total_kv_chk = kv_indptr_cpu.data_ptr<int32_t>()[len_indptr - 1];
  TORCH_CHECK(total_q == total_q_chk, "total_q must equal qo_indptr[-1]");
  TORCH_CHECK(total_kv == total_kv_chk, "total_kv must equal kv_indptr[-1]");

  float sm_scale = sm_scale_opt.has_value() ? static_cast<float>(*sm_scale_opt) : default_sm_scale();
  if (!(sm_scale > 0.0f)) sm_scale = default_sm_scale();

  // Allocate outputs
  auto options_bf16 = torch::TensorOptions().dtype(torch::kBFloat16).device(q.device());
  auto options_f32  = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());

  torch::Tensor output = torch::empty({total_q, 32, 128}, options_bf16);
  torch::Tensor lse    = torch::empty({total_q, 32}, options_f32);

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

  const __nv_bfloat16* q_ptr  = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<c10::BFloat16>());
  const __nv_bfloat16* k_ptr  = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<c10::BFloat16>());
  const __nv_bfloat16* v_ptr  = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<c10::BFloat16>());
  const int32_t* qo_ptr       = qo_indptr.data_ptr<int32_t>();
  const int32_t* kv_ptr       = kv_indptr.data_ptr<int32_t>();
  __nv_bfloat16* out_ptr      = reinterpret_cast<__nv_bfloat16*>(output.data_ptr<c10::BFloat16>());
  float* lse_ptr              = lse.data_ptr<float>();

  gqa_ragged_prefill_causal_h32_kv4_d128_launcher(
      q_ptr, k_ptr, v_ptr,
      qo_ptr, kv_ptr,
      static_cast<int>(len_indptr),
      static_cast<int>(total_q),
      static_cast<int>(total_kv),
      sm_scale,
      out_ptr, lse_ptr,
      stream.stream());

  // Check for launch errors (synchronous check)
  cudaError_t err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));

  return {output, lse};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run,
        py::arg("q"),
        py::arg("k"),
        py::arg("v"),
        py::arg("qo_indptr"),
        py::arg("kv_indptr"),
        py::arg("sm_scale") = py::none(),
        "GQA Ragged Prefill Causal Attention (h=32, kv=4, d=128) optimized kernel.");
}