gpt-5 / cuda21ea96

#include "kernel.h"

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <vector>
#include <cmath>
#include <stdexcept>
#include <iostream>

static inline void check_inputs(
    const torch::Tensor& q,
    const torch::Tensor& k_cache,
    const torch::Tensor& v_cache,
    const torch::Tensor& qo_indptr,
    const torch::Tensor& kv_indptr,
    const torch::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(qo_indptr.is_cuda(), "qo_indptr 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.scalar_type() == at::kBFloat16, "q must be bfloat16");
  TORCH_CHECK(k_cache.scalar_type() == at::kBFloat16, "k_cache must be bfloat16");
  TORCH_CHECK(v_cache.scalar_type() == at::kBFloat16, "v_cache must be bfloat16");
  TORCH_CHECK(qo_indptr.scalar_type() == at::kInt, "qo_indptr must be int32");
  TORCH_CHECK(kv_indptr.scalar_type() == at::kInt, "kv_indptr must be int32");
  TORCH_CHECK(kv_indices.scalar_type() == at::kInt, "kv_indices must be int32");

  TORCH_CHECK(q.dim() == 3, "q must have shape [total_q, 32, 128]");
  TORCH_CHECK(k_cache.dim() == 4, "k_cache must have shape [num_pages, 1, 8, 128]");
  TORCH_CHECK(v_cache.dim() == 4, "v_cache must have shape [num_pages, 1, 8, 128]");

  TORCH_CHECK(q.is_contiguous(), "q must be contiguous");
  TORCH_CHECK(k_cache.is_contiguous(), "k_cache must be contiguous");
  TORCH_CHECK(v_cache.is_contiguous(), "v_cache must be contiguous");
  TORCH_CHECK(qo_indptr.is_contiguous(), "qo_indptr must be contiguous");
  TORCH_CHECK(kv_indptr.is_contiguous(), "kv_indptr must be contiguous");
  TORCH_CHECK(kv_indices.is_contiguous(), "kv_indices must be contiguous");

  const int64_t total_q_sz = q.size(0);
  const int64_t num_qo_heads = q.size(1);
  const int64_t head_dim = q.size(2);
  TORCH_CHECK(num_qo_heads == NUM_QO_HEADS, "num_qo_heads must be 32");
  TORCH_CHECK(head_dim == HEAD_DIM, "head_dim must be 128");

  const int64_t num_pages = k_cache.size(0);
  const int64_t page_size = k_cache.size(1);
  const int64_t num_kv_heads = k_cache.size(2);
  const int64_t head_dim_k = k_cache.size(3);
  (void)num_pages;
  TORCH_CHECK(page_size == PAGE_SIZE, "page_size must be 1");
  TORCH_CHECK(num_kv_heads == NUM_KV_HEADS, "num_kv_heads must be 8");
  TORCH_CHECK(head_dim_k == HEAD_DIM, "KV head_dim must be 128");
  TORCH_CHECK(v_cache.size(0) == k_cache.size(0) &&
              v_cache.size(1) == k_cache.size(1) &&
              v_cache.size(2) == k_cache.size(2) &&
              v_cache.size(3) == k_cache.size(3), "v_cache must match k_cache shape");

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

  const int64_t len_indptr = qo_indptr.size(0);
  TORCH_CHECK(kv_indptr.size(0) == len_indptr, "qo_indptr and kv_indptr must have same length");

  // Constraints
  // total_q == qo_indptr[-1] and num_kv_indices == kv_indptr[-1]
  auto qo_indptr_cpu = qo_indptr.cpu();
  auto kv_indptr_cpu = kv_indptr.cpu();
  const int32_t* qo_ptr_h = qo_indptr_cpu.data_ptr<int32_t>();
  const int32_t* kv_ptr_h = kv_indptr_cpu.data_ptr<int32_t>();

  int32_t total_q_from_indptr = qo_ptr_h[len_indptr - 1];
  TORCH_CHECK(total_q_sz == static_cast<int64_t>(total_q_from_indptr),
              "Constraint failed: total_q must equal qo_indptr[-1]");

  int32_t num_kv_indices_from_indptr = kv_ptr_h[len_indptr - 1];
  TORCH_CHECK(kv_indices.size(0) == static_cast<int64_t>(num_kv_indices_from_indptr),
              "Constraint failed: num_kv_indices must equal kv_indptr[-1]");
}

std::vector<torch::Tensor> run(
    torch::Tensor q,             // [total_q, 32, 128], bf16, CUDA
    torch::Tensor k_cache,       // [num_pages, 1, 8, 128], bf16, CUDA
    torch::Tensor v_cache,       // [num_pages, 1, 8, 128], bf16, CUDA
    torch::Tensor qo_indptr,     // [len_indptr], int32, CUDA
    torch::Tensor kv_indptr,     // [len_indptr], int32, CUDA
    torch::Tensor kv_indices,    // [num_kv_indices], int32, CUDA
    double sm_scale_double       // scalar
) {
  check_inputs(q, k_cache, v_cache, qo_indptr, kv_indptr, kv_indices);

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

  // Prepare outputs
  auto opts_out = q.options();
  torch::Tensor output = torch::empty({total_q, NUM_QO_HEADS, HEAD_DIM}, opts_out);
  torch::Tensor lse = torch::empty({total_q, NUM_QO_HEADS},
                                   q.options().dtype(at::kFloat)); // float32

  // Build mapping arrays q_seq and q_pos on CPU, then move to GPU (same device)
  torch::Tensor q_seq_cpu = torch::empty({total_q}, torch::TensorOptions().dtype(at::kInt).device(torch::kCPU));
  torch::Tensor q_pos_cpu = torch::empty({total_q}, torch::TensorOptions().dtype(at::kInt).device(torch::kCPU));

  // Copy indptr to CPU for generating mapping
  auto qo_indptr_cpu = qo_indptr.cpu();
  const int32_t* qo_indptr_h = qo_indptr_cpu.data_ptr<int32_t>();

  int32_t* q_seq_h = q_seq_cpu.data_ptr<int32_t>();
  int32_t* q_pos_h = q_pos_cpu.data_ptr<int32_t>();

  for (int64_t b = 0; b < len_indptr - 1; ++b) {
    const int32_t start = qo_indptr_h[b];
    const int32_t end = qo_indptr_h[b + 1];
    for (int32_t qg = start; qg < end; ++qg) {
      q_seq_h[qg] = static_cast<int32_t>(b);
      q_pos_h[qg] = static_cast<int32_t>(qg - start);
    }
  }

  // Move mapping to device (same device as q)
  torch::Tensor q_seq = q_seq_cpu.to(q.device(), /*non_blocking=*/true);
  torch::Tensor q_pos = q_pos_cpu.to(q.device(), /*non_blocking=*/true);

  // Launch kernel
  float sm_scale = static_cast<float>(sm_scale_double);
  launch_gqa_paged_prefill_causal_h32_kv8_d128_ps1(
      q, k_cache, v_cache,
      qo_indptr, kv_indptr, kv_indices,
      q_seq, q_pos,
      sm_scale,
      output, lse
  );

  // Return results (stay on device)
  return {output, lse};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run,
        "gqa_paged_prefill_causal_h32_kv8_d128_ps1 (B200-optimized)",
        pybind11::arg("q"),
        pybind11::arg("k_cache"),
        pybind11::arg("v_cache"),
        pybind11::arg("qo_indptr"),
        pybind11::arg("kv_indptr"),
        pybind11::arg("kv_indices"),
        pybind11::arg("sm_scale") = 1.0 / std::sqrt(static_cast<double>(HEAD_DIM)));
}