claude-opus-4-1-20250805 / cuda9be3ac

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector>
#include <stdexcept>
#include "kernel.h"

// Constant for vocabulary size
constexpr int VOCAB_SIZE_CHECK = 128256;

// Main run function
torch::Tensor run(torch::Tensor probs, torch::Tensor top_p) {
    // Input validation
    TORCH_CHECK(probs.dim() == 2, "probs must be a 2D tensor");
    TORCH_CHECK(top_p.dim() == 1, "top_p must be a 1D tensor");
    TORCH_CHECK(probs.size(0) == top_p.size(0), 
                "Batch size mismatch between probs and top_p");
    TORCH_CHECK(probs.size(1) == VOCAB_SIZE_CHECK, 
                "Vocabulary size must be 128256, got ", probs.size(1));
    TORCH_CHECK(probs.scalar_type() == torch::kFloat32, 
                "probs must be float32");
    TORCH_CHECK(top_p.scalar_type() == torch::kFloat32, 
                "top_p must be float32");
    TORCH_CHECK(probs.is_cuda(), "probs must be a CUDA tensor");
    TORCH_CHECK(top_p.is_cuda(), "top_p must be a CUDA tensor");
    TORCH_CHECK(probs.is_contiguous(), "probs must be contiguous");
    TORCH_CHECK(top_p.is_contiguous(), "top_p must be contiguous");
    
    const int batch_size = probs.size(0);
    
    // Create output tensor on the same device as input
    auto options = torch::TensorOptions()
        .dtype(torch::kInt64)
        .device(probs.device());
    torch::Tensor samples = torch::empty({batch_size}, options);
    
    // Get current CUDA stream for this device
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch the optimized kernel
    launch_top_p_sampling(
        probs.data_ptr<float>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<int64_t>(),
        batch_size,
        stream
    );
    
    // Ensure kernel completion
    cudaError_t error = cudaStreamSynchronize(stream);
    if (error != cudaSuccess) {
        TORCH_CHECK(false, "CUDA error during kernel execution: ", 
                    cudaGetErrorString(error));
    }
    
    return samples;
}

// Python bindings using pybind11
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Top-p (nucleus) sampling from probability distributions",
          py::arg("probs"), py::arg("top_p"));
}