claude-opus-4-1-20250805 / cudafdf15e

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

// Expected vocabulary size constant
constexpr int EXPECTED_VOCAB_SIZE = 151936;

// Main run function that interfaces with PyTorch
torch::Tensor run(torch::Tensor probs, torch::Tensor top_p) {
    // Validate input dimensions
    TORCH_CHECK(probs.dim() == 2, 
                "probs must be a 2D tensor, got ", probs.dim(), " dimensions");
    TORCH_CHECK(top_p.dim() == 1, 
                "top_p must be a 1D tensor, got ", top_p.dim(), " dimensions");
    
    // Get dimensions
    int64_t batch_size = probs.size(0);
    int64_t vocab_size = probs.size(1);
    
    // Validate sizes
    TORCH_CHECK(vocab_size == EXPECTED_VOCAB_SIZE, 
                "vocab_size must be ", EXPECTED_VOCAB_SIZE, ", got ", vocab_size);
    TORCH_CHECK(top_p.size(0) == batch_size, 
                "Batch size mismatch: probs has ", batch_size, 
                " batches, top_p has ", top_p.size(0));
    
    // Validate data types
    TORCH_CHECK(probs.scalar_type() == torch::kFloat32, 
                "probs must be float32");
    TORCH_CHECK(top_p.scalar_type() == torch::kFloat32, 
                "top_p must be float32");
    
    // Ensure tensors are on CUDA
    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(), 
                "All tensors must be on the same device");
    
    // Make tensors contiguous if needed
    probs = probs.contiguous();
    top_p = top_p.contiguous();
    
    // Create 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>(),
        static_cast<int>(batch_size),
        stream
    );
    
    // Synchronize to ensure kernel completion
    cudaError_t err = cudaStreamSynchronize(stream);
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA stream synchronization failed: ", cudaGetErrorString(err));
    }
    
    // Final error check
    err = cudaGetLastError();
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel execution failed: ", cudaGetErrorString(err));
    }
    
    return samples;
}

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