claude-opus-4-1 / cuda16cd03

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <ATen/cuda/CUDAContext.h>
#include "kernel.h"

torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
    // Input validation
    TORCH_CHECK(hidden_states.device().is_cuda(), "hidden_states must be a CUDA tensor");
    TORCH_CHECK(weight.device().is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");
    TORCH_CHECK(hidden_states.scalar_type() == torch::kBFloat16, "hidden_states must be bfloat16");
    TORCH_CHECK(weight.scalar_type() == torch::kBFloat16, "weight must be bfloat16");
    
    // Check dimensions
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D tensor");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1D tensor");
    
    const int64_t batch_size = hidden_states.size(0);
    const int64_t hidden_size = hidden_states.size(1);
    
    TORCH_CHECK(hidden_size == 7168, "hidden_size must be 7168, got ", hidden_size);
    TORCH_CHECK(weight.size(0) == hidden_size, "weight size must match hidden_size");
    
    // Allocate output tensor
    torch::Tensor output = torch::empty_like(hidden_states);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_rmsnorm_h7168(
        hidden_states.data_ptr(),
        weight.data_ptr(),
        output.data_ptr(),
        static_cast<int>(batch_size),
        stream
    );
    
    // Ensure kernel completes
    cudaError_t err = cudaStreamSynchronize(stream);
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel execution error: ", cudaGetErrorString(err));
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "RMSNorm forward pass for hidden_size=7168",
          py::arg("hidden_states"), py::arg("weight"));
}