gpt-o3 / cuda14adb8

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

#include "kernel.h"

/* -------------------------------------------------------------------------- */
/*  Python-visible entry point                                                */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor probs,
                  torch::Tensor top_p)
{
    /* Sanity checks ----------------------------------------------------- */
    TORCH_CHECK(probs.is_cuda(),  "probs tensor must be on CUDA");
    TORCH_CHECK(top_p.is_cuda(),  "top_p tensor must be on CUDA");

    TORCH_CHECK(probs.dtype() == torch::kFloat32,
                "probs must be float32, got ", probs.dtype());
    TORCH_CHECK(top_p.dtype() == torch::kFloat32,
                "top_p must be float32, got ", top_p.dtype());

    TORCH_CHECK(probs.dim() == 2,
                "probs must be rank-2, got dim=", probs.dim());
    TORCH_CHECK(probs.size(1) == VOCAB_SIZE,
                "probs second dim must be ", VOCAB_SIZE,
                ", got ", probs.size(1));

    TORCH_CHECK(top_p.dim() == 1,
                "top_p must be rank-1, got dim=", top_p.dim());
    TORCH_CHECK(top_p.size(0) == probs.size(0),
                "top_p length (", top_p.size(0),
                ") must equal batch size (", probs.size(0), ")");

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

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

    auto out_opts = torch::TensorOptions()
                        .dtype(torch::kInt64)
                        .device(probs.device());
    torch::Tensor samples = torch::empty({B}, out_opts);

    /* Dispatch ---------------------------------------------------------- */
    at::cuda::CUDAGuard guard(probs.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    top_p_sampling_from_probs_v129280_launcher(
        probs.data_ptr<float>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<int64_t>(),
        B,
        stream);

    return samples;
}

/* -------------------------------------------------------------------------- */
/*  Pybind11 module                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.doc() = "Top-p (nucleus) sampling kernel for DeepSeek-V3, vocab=129 280";
    m.def("run",
          &run,
          pybind11::arg("probs"),
          pybind11::arg("top_p"),
          "Sample a token per row from a probability matrix using top-p.");
}