gpt-5 / cudafd76fe

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <vector>
#include <stdexcept>
#include "kernel.h"

using torch::Tensor;

namespace {

void check_inputs(const Tensor& q_nope,
                  const Tensor& q_pe,
                  const Tensor& ckv_cache,
                  const Tensor& kpe_cache,
                  const Tensor& qo_indptr,
                  const Tensor& kv_indptr,
                  const Tensor& kv_indices) {
  TORCH_CHECK(q_nope.is_cuda(), "q_nope must be a CUDA tensor");
  TORCH_CHECK(q_pe.is_cuda(), "q_pe must be a CUDA tensor");
  TORCH_CHECK(ckv_cache.is_cuda(), "ckv_cache must be a CUDA tensor");
  TORCH_CHECK(kpe_cache.is_cuda(), "kpe_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_nope.scalar_type() == at::kBFloat16, "q_nope must be bfloat16");
  TORCH_CHECK(q_pe.scalar_type() == at::kBFloat16, "q_pe must be bfloat16");
  TORCH_CHECK(ckv_cache.scalar_type() == at::kBFloat16, "ckv_cache must be bfloat16");
  TORCH_CHECK(kpe_cache.scalar_type() == at::kBFloat16, "kpe_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_nope.dim() == 3, "q_nope must be [total_q, 16, 512]");
  TORCH_CHECK(q_pe.dim() == 3, "q_pe must be [total_q, 16, 64]");
  TORCH_CHECK(ckv_cache.dim() == 3, "ckv_cache must be [num_pages, 1, 512]");
  TORCH_CHECK(kpe_cache.dim() == 3, "kpe_cache must be [num_pages, 1, 64]");
  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");

  TORCH_CHECK(q_nope.size(1) == 16 && q_nope.size(2) == 512, "q_nope shape mismatch");
  TORCH_CHECK(q_pe.size(1) == 16 && q_pe.size(2) == 64, "q_pe shape mismatch");
  TORCH_CHECK(ckv_cache.size(1) == 1 && ckv_cache.size(2) == 512, "ckv_cache shape mismatch (page_size=1, head_dim_ckv=512)");
  TORCH_CHECK(kpe_cache.size(1) == 1 && kpe_cache.size(2) == 64, "kpe_cache shape mismatch (page_size=1, head_dim_kpe=64)");
}

} // anonymous

std::vector<Tensor> run(Tensor q_nope,
                        Tensor q_pe,
                        Tensor ckv_cache,
                        Tensor kpe_cache,
                        Tensor qo_indptr,
                        Tensor kv_indptr,
                        Tensor kv_indices,
                        double sm_scale_double) {
  check_inputs(q_nope, q_pe, ckv_cache, kpe_cache, qo_indptr, kv_indptr, kv_indices);

  // Ensure contiguity
  q_nope = q_nope.contiguous();
  q_pe = q_pe.contiguous();
  ckv_cache = ckv_cache.contiguous();
  kpe_cache = kpe_cache.contiguous();
  qo_indptr = qo_indptr.contiguous();
  kv_indptr = kv_indptr.contiguous();
  kv_indices = kv_indices.contiguous();

  const int64_t total_q = q_nope.size(0);
  const int64_t num_heads = q_nope.size(1);
  const int64_t head_dim_ckv = q_nope.size(2);
  const int64_t head_dim_kpe = q_pe.size(2);
  const int64_t num_pages = ckv_cache.size(0);
  const int64_t page_size = ckv_cache.size(1); // should be 1
  const int64_t len_indptr = qo_indptr.size(0);
  const int64_t num_kv_indices = kv_indices.size(0);

  TORCH_CHECK(num_heads == 16, "num_qo_heads must be 16");
  TORCH_CHECK(head_dim_ckv == 512, "head_dim_ckv must be 512");
  TORCH_CHECK(head_dim_kpe == 64, "head_dim_kpe must be 64");
  TORCH_CHECK(page_size == 1, "page_size must be 1");

  // Constraints
  {
    int32_t total_q_constr = qo_indptr.index({len_indptr - 1}).item<int32_t>();
    int32_t num_kv_indices_constr = kv_indptr.index({len_indptr - 1}).item<int32_t>();
    TORCH_CHECK(static_cast<int64_t>(total_q_constr) == total_q,
                "Constraint violated: total_q == qo_indptr[-1]");
    TORCH_CHECK(static_cast<int64_t>(num_kv_indices_constr) == num_kv_indices,
                "Constraint violated: num_kv_indices == kv_indptr[-1]");
  }

  // Allocate outputs
  Tensor output = torch::empty({total_q, 16, 512}, q_nope.options());
  Tensor lse = torch::empty({total_q, 16}, q_nope.options().dtype(torch::kFloat32));

  float sm_scale = static_cast<float>(sm_scale_double);

  // Raw pointers
  const __nv_bfloat16* q_nope_ptr = reinterpret_cast<const __nv_bfloat16*>(q_nope.data_ptr<at::BFloat16>());
  const __nv_bfloat16* q_pe_ptr   = reinterpret_cast<const __nv_bfloat16*>(q_pe.data_ptr<at::BFloat16>());
  const __nv_bfloat16* ckv_ptr    = reinterpret_cast<const __nv_bfloat16*>(ckv_cache.data_ptr<at::BFloat16>());
  const __nv_bfloat16* kpe_ptr    = reinterpret_cast<const __nv_bfloat16*>(kpe_cache.data_ptr<at::BFloat16>());

  const int32_t* qo_indptr_ptr = qo_indptr.data_ptr<int32_t>();
  const int32_t* kv_indptr_ptr = kv_indptr.data_ptr<int32_t>();
  const int32_t* kv_indices_ptr = kv_indices.data_ptr<int32_t>();

  __nv_bfloat16* out_ptr = reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>());
  float* lse_ptr = lse.data_ptr<float>();

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  mla_h16_ckv512_kpe64_ps1::launch_mla_paged_prefill_causal_h16_ckv512_kpe64_ps1(
      q_nope_ptr, q_pe_ptr, ckv_ptr, kpe_ptr,
      qo_indptr_ptr, kv_indptr_ptr, kv_indices_ptr,
      sm_scale,
      static_cast<int32_t>(total_q),
      static_cast<int32_t>(len_indptr),
      static_cast<int32_t>(num_pages),
      static_cast<int32_t>(num_kv_indices),
      out_ptr, lse_ptr, stream);

  return {output, lse};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run, "mla_paged_prefill_causal_h16_ckv512_kpe64_ps1 (B200-optimized)");
}