#include "rmsnorm_nvfp4_m1.h"

#include <torch/extension.h>

#include <vector>

namespace {

std::vector<torch::Tensor> rmsnorm_nvfp4_m1(
    const torch::Tensor& input,
    const torch::Tensor& weight,
    double epsilon) {
  TORCH_CHECK(
      input.is_cuda() && input.scalar_type() == at::kBFloat16 &&
          input.dim() == 2 && input.size(0) == 1 && input.is_contiguous(),
      "RMSNorm input must be contiguous CUDA bfloat16 [1,K]");
  TORCH_CHECK(
      weight.is_cuda() && weight.scalar_type() == at::kBFloat16 &&
          weight.dim() == 1 && weight.is_contiguous(),
      "RMSNorm weight must be contiguous CUDA bfloat16 [K]");
  TORCH_CHECK(
      weight.device() == input.device() && weight.numel() == input.size(1),
      "RMSNorm input and weight dimensions/devices differ");
  TORCH_CHECK(
      input.size(1) > 0 && input.size(1) % 32 == 0,
      "fused RMSNorm-to-NVFP4 requires K divisible by 32");
  TORCH_CHECK(epsilon > 0.0, "RMSNorm epsilon must be positive");
  return rmsnorm_nvfp4_m1_cuda(input, weight, epsilon);
}

}  // namespace

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def(
      "quantize",
      &rmsnorm_nvfp4_m1,
      "Fused Qwen RMSNorm to packed NVFP4 at logical M=1");
}