cuda_nvfp4_grouped_naive_g4_n2048_k2048

#include "kernel.h"
#include <torch/extension.h>

torch::Tensor run(torch::Tensor A_fp4, torch::Tensor A_scale, torch::Tensor B_fp4,
                  torch::Tensor B_scale, torch::Tensor m_indptr, torch::Tensor alpha) {
    auto A = A_fp4.contiguous(); auto As = A_scale.contiguous();
    auto B = B_fp4.contiguous(); auto Bs = B_scale.contiguous();
    auto mi = m_indptr.to(torch::kInt32).contiguous();
    auto al = alpha.to(torch::kFloat32).contiguous();
    int G = B.size(0);
    const auto Mtot = A.size(0); const auto N = B.size(1);
    auto C = torch::empty({Mtot, N}, torch::dtype(at::kBFloat16).device(A.device()));
    launch_nvfp4_grouped(A, As, B, Bs, mi, al, C, G);
    return C;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "grouped nvfp4 gemm",
          py::arg("A_fp4"), py::arg("A_scale"), py::arg("B_fp4"), py::arg("B_scale"),
          py::arg("m_indptr"), py::arg("alpha"));
}