claude-opus-4-1-20250805 / cudab336ab

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

// Macro for checking CUDA errors
#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)

// Helper macros for input validation
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor run(
    torch::Tensor probs,
    torch::Tensor top_k,
    torch::Tensor top_p
) {
    // Input validation
    CHECK_INPUT(probs);
    CHECK_INPUT(top_k);
    CHECK_INPUT(top_p);
    
    // Check dimensions
    TORCH_CHECK(probs.dim() == 2, "probs must be 2D tensor, got ", probs.dim(), "D");
    TORCH_CHECK(top_k.dim() == 1, "top_k must be 1D tensor, got ", top_k.dim(), "D");
    TORCH_CHECK(top_p.dim() == 1, "top_p must be 1D tensor, got ", top_p.dim(), "D");
    
    const int batch_size = probs.size(0);
    const int vocab_size = probs.size(1);
    
    // Verify vocabulary size
    TORCH_CHECK(vocab_size == VOCAB_SIZE, 
                "vocab_size must be ", VOCAB_SIZE, ", but got ", vocab_size);
    
    // Check batch dimensions match
    TORCH_CHECK(top_k.size(0) == batch_size, 
                "top_k batch size (", top_k.size(0), ") doesn't match probs batch size (", batch_size, ")");
    TORCH_CHECK(top_p.size(0) == batch_size,
                "top_p batch size (", top_p.size(0), ") doesn't match probs batch size (", batch_size, ")");
    
    // Check dtypes
    TORCH_CHECK(probs.scalar_type() == torch::kFloat32, 
                "probs must be float32, got ", probs.scalar_type());
    TORCH_CHECK(top_k.scalar_type() == torch::kInt32,
                "top_k must be int32, got ", top_k.scalar_type());
    TORCH_CHECK(top_p.scalar_type() == torch::kFloat32,
                "top_p must be float32, got ", top_p.scalar_type());
    
    // Ensure tensors are contiguous
    probs = probs.contiguous();
    top_k = top_k.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 CUDA stream from PyTorch
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_top_k_top_p_sampling(
        probs.data_ptr<float>(),
        top_k.data_ptr<int32_t>(),
        top_p.data_ptr<float>(),
        samples.data_ptr<int64_t>(),
        batch_size,
        stream
    );
    
    // Check for kernel launch errors
    CUDA_CHECK(cudaGetLastError());
    
    // Ensure kernel completion for debugging (can be removed in production)
    CUDA_CHECK(cudaStreamSynchronize(stream));
    
    return samples;
}

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