gpt-o3 / cudac24d60

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

/* -------------------------------------------------------------------------- */
/* Python-visible wrapper                                                     */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor probs,   /* [B, 128256] - float32 */
                  torch::Tensor top_p)   /* [B]         - float32 */
{
  TORCH_CHECK(probs.is_cuda(), "probs must reside on CUDA");
  TORCH_CHECK(top_p.is_cuda(), "top_p must reside on CUDA");
  TORCH_CHECK(probs.dtype()  == torch::kFloat32, "probs must be float32");
  TORCH_CHECK(top_p.dtype()  == torch::kFloat32, "top_p must be float32");
  TORCH_CHECK(probs.dim() == 2, "probs must be 2-D (batch, vocab)");
  TORCH_CHECK(probs.size(1) == VOCAB_SIZE,
              "vocab dimension must be ", VOCAB_SIZE);

  probs = probs.contiguous();
  top_p = top_p.contiguous();

  const int batch_size = static_cast<int>(probs.size(0));

  auto out_opts  = probs.options().dtype(torch::kInt64);
  torch::Tensor samples = torch::empty({batch_size}, out_opts);

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

  top_p_sampling_from_probs_v128256_launcher(
      probs.data_ptr<float>(),
      top_p.data_ptr<float>(),
      samples.data_ptr<int64_t>(),
      batch_size,
      stream);

  return samples;
}

/* -------------------------------------------------------------------------- */
/* PyBind11 binding                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
  m.def("run", &run,
        "top_p_sampling_from_probs_v128256 (CUDA - optimised for B200)");
}