claude-opus-4-1-20250805 / cuda29eefb

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

torch::Tensor run(torch::Tensor probs, torch::Tensor top_k) {
    // Input validation
    TORCH_CHECK(probs.dim() == 2, "probs must be a 2D tensor");
    TORCH_CHECK(top_k.dim() == 1, "top_k must be a 1D tensor");
    TORCH_CHECK(probs.size(1) == 129280, "vocab_size must be 129280, got ", probs.size(1));
    TORCH_CHECK(probs.size(0) == top_k.size(0), 
                "Batch size mismatch: probs has ", probs.size(0), 
                " samples, top_k has ", top_k.size(0));
    
    // Ensure CUDA tensors
    TORCH_CHECK(probs.is_cuda(), "probs must be on CUDA device");
    TORCH_CHECK(top_k.is_cuda(), "top_k must be on CUDA device");
    TORCH_CHECK(probs.device() == top_k.device(), 
                "probs and top_k must be on the same device");
    
    // Type conversion if needed
    torch::Tensor probs_float = probs;
    if (probs.dtype() != torch::kFloat32) {
        probs_float = probs.to(torch::kFloat32);
    }
    
    torch::Tensor top_k_int = top_k;
    if (top_k.dtype() != torch::kInt32) {
        top_k_int = top_k.to(torch::kInt32);
    }
    
    // Ensure contiguous tensors
    probs_float = probs_float.contiguous();
    top_k_int = top_k_int.contiguous();
    
    int batch_size = probs_float.size(0);
    
    // Allocate output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kInt64)
        .device(probs_float.device())
        .requires_grad(false);
    torch::Tensor samples = torch::empty({batch_size}, options);
    
    // Get current CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_top_k_sampling(
        probs_float.data_ptr<float>(),
        top_k_int.data_ptr<int>(),
        samples.data_ptr<int64_t>(),
        batch_size,
        stream
    );
    
    // Synchronize to ensure kernel completion
    cudaStreamSynchronize(stream);
    
    // Check for any CUDA errors
    cudaError_t error = cudaGetLastError();
    if (error != cudaSuccess) {
        throw std::runtime_error(
            std::string("CUDA error after kernel execution: ") + cudaGetErrorString(error)
        );
    }
    
    return samples;
}

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