claude-opus-4-1-20250805 / cudaea19e3

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

// Macro for CUDA error checking
#define CUDA_CHECK(call) \
    do { \
        cudaError_t error = call; \
        if (error != cudaSuccess) { \
            throw std::runtime_error(std::string("CUDA error at ") + __FILE__ + ":" + \
                std::to_string(__LINE__) + " - " + cudaGetErrorString(error)); \
        } \
    } while(0)

torch::Tensor run(torch::Tensor probs, torch::Tensor top_p) {
    // Input validation
    TORCH_CHECK(probs.dim() == 2, "probs must be 2-dimensional");
    TORCH_CHECK(top_p.dim() == 1, "top_p must be 1-dimensional");
    
    const int batch_size = probs.size(0);
    const int vocab_size = probs.size(1);
    
    TORCH_CHECK(vocab_size == VOCAB_SIZE, 
                "vocab_size must be ", VOCAB_SIZE, ", got ", vocab_size);
    TORCH_CHECK(top_p.size(0) == batch_size, 
                "batch size mismatch between probs and top_p");
    
    // Ensure CUDA tensors
    TORCH_CHECK(probs.is_cuda(), "probs must be on CUDA device");
    TORCH_CHECK(top_p.is_cuda(), "top_p must be on CUDA device");
    TORCH_CHECK(probs.device() == top_p.device(), 
                "probs and top_p must be on the same device");
    
    // Convert to float32 if necessary
    if (probs.scalar_type() != torch::kFloat32) {
        probs = probs.to(torch::kFloat32);
    }
    if (top_p.scalar_type() != torch::kFloat32) {
        top_p = top_p.to(torch::kFloat32);
    }
    
    // Ensure contiguous memory layout
    probs = probs.contiguous();
    top_p = top_p.contiguous();
    
    // Allocate output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kInt64)
        .device(probs.device());
    torch::Tensor samples = torch::empty({batch_size}, options);
    
    // Get current CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_top_p_sampling(
        probs.data_ptr<float>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<int64_t>(),
        batch_size,
        stream
    );
    
    // Check for kernel launch errors
    CUDA_CHECK(cudaGetLastError());
    
    return samples;
}

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