gpt-o3 / cuda7f55e3

#include "kernel.h"
#include <torch/extension.h>

/* -------------------------------------------------------------------------- */
/*                            Python exposed entry                            */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor probs,
                  torch::Tensor top_k,
                  torch::Tensor top_p)
{
    TORCH_CHECK(probs.is_cuda(), "probs must reside on CUDA");
    TORCH_CHECK(top_k.is_cuda(), "top_k 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 tensor");

    const int batch_size = probs.size(0);
    auto samples = torch::empty({batch_size},
                                torch::dtype(torch::kInt64).device(probs.device()));

    top_k_top_p_sampling_from_probs_v128256_launcher(
        probs.contiguous(),
        top_k.contiguous(),
        top_p.contiguous(),
        samples);

    return samples;
}

/* -------------------------------------------------------------------------- */
/*                            pybind11 registration                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run", &run, "top-k / top-p sampling kernel (B200 optimised)");
}