/*************************************************************************
 * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#include <transformer_engine/transformer_engine.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <cerrno>
#include <climits>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <limits>
#include <mutex>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "common.h"
#include "common/util/cuda_runtime.h"
#include "common/util/logging.h"

namespace transformer_engine {

size_t typeToNumBits(const DType type) {
  TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(type, T,
                                     return TypeInfo<T>::size;);  // NOLINT(*)
}

size_t typeToSize(const DType type) {
  NVTE_CHECK(type != DType::kFloat4E2M1, "typeToSize() Does not support FP4 data type.");
  return typeToNumBits(type) / 8;
}

std::string to_string(const NVTEScalingMode &mode) {
  switch (mode) {
    case NVTE_DELAYED_TENSOR_SCALING:
      return "NVTE_DELAYED_TENSOR_SCALING";
    case NVTE_MXFP8_1D_SCALING:
      return "NVTE_MXFP8_1D_SCALING";
    case NVTE_BLOCK_SCALING_1D:
      return "NVTE_BLOCK_SCALING_1D";
    case NVTE_BLOCK_SCALING_2D:
      return "NVTE_BLOCK_SCALING_2D";
    case NVTE_NVFP4_1D_SCALING:
      return "NVTE_NVFP4_1D_SCALING";
    case NVTE_INVALID_SCALING:
      return "NVTE_INVALID_SCALING";
  }
  return "Invalid Scaling";
}

void CheckNoopTensor(const Tensor &t, std::string_view name) {
  if (t.data.has_data()) {
    NVTE_CHECK(t.numel() == 1, "Expected 1 element for ", name, " noop, but found ", t.numel(),
               ".");
    NVTE_CHECK(t.data.dtype == DType::kFloat32, "Found wrong dtype for ", name,
               " noop. Expected kFloat32.");
  }
}

void CheckScaleTensorShape(const Tensor &t, std::string_view name) {
  NVTE_CHECK(t.scaling_mode != NVTE_INVALID_SCALING, "Invalid scaling mode!");
  if (is_tensor_scaling(t.scaling_mode)) {
    if (is_fp8_dtype(t.dtype())) {
      // FP8 tensor with tensor scaling
      if (t.has_data()) {
        NVTE_CHECK(t.scale_inv.numel() == 1, "Tensor \"", name,
                   "\" has invalid scale_inv shape (expected 1 entry, got ", t.scale_inv.shape,
                   ")");
      }
      if (t.has_columnwise_data()) {
        NVTE_CHECK(t.columnwise_scale_inv.numel() == 1, "Tensor \"", name,
                   "\" has invalid columnwise_scale_inv shape (expected 1 entry, got ",
                   t.columnwise_scale_inv.shape, ")");
      }
    } else {
      // High-precision tensor
      if (t.has_data()) {
        NVTE_CHECK(t.scale_inv.numel() == 0, "Tensor \"", name,
                   "\" has invalid scale_inv shape (expected 0 entries, got ", t.scale_inv.shape,
                   ")");
      }
      if (t.has_columnwise_data()) {
        NVTE_CHECK(t.columnwise_scale_inv.numel() == 0, "Tensor \"", name,
                   "\" has invalid columnwise_scale_inv shape (expected 0 entries, got ",
                   t.columnwise_scale_inv.shape, ")");
      }
    }
  } else {
    if (t.scaling_mode == NVTE_MXFP8_1D_SCALING) {
      // Need (4, 128) alignment even for e8 scaling factor
      constexpr std::array<size_t, 2> block_alignment{128ul, 4ul};
      const auto [first_dim, last_dim] = t.flat_2d_dims();

      if (t.has_data()) {
        constexpr std::array<size_t, 2> block_shape{1, 32};
        const std::array<size_t, 2> expected{
            DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[0]), block_alignment[0]),
            DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[1]), block_alignment[1])};
        NVTE_CHECK(t.scale_inv.shape == expected, "Tensor \"", name,
                   "\" has invalid scale_inv shape (expected ", expected, ", got ",
                   t.scale_inv.shape, ")");
      }
      if (t.has_columnwise_data()) {
        constexpr std::array<size_t, 2> block_shape{32, 1};
        const std::array<size_t, 2> expected{
            DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[0]), block_alignment[1]),
            DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[1]), block_alignment[0])};
        NVTE_CHECK(t.columnwise_scale_inv.shape == expected, "Tensor \"", name,
                   "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ",
                   t.columnwise_scale_inv.shape, ")");
      }
    } else if (t.scaling_mode == NVTE_NVFP4_1D_SCALING) {
      const auto [first_dim, last_dim] = t.flat_2d_dims();

      if (t.has_data()) {
        constexpr std::array<size_t, 2> block_shape{1, 16};
        constexpr std::array<size_t, 2> block_alignment{128, 4};
        const std::array<size_t, 2> expected{
            DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[0]), block_alignment[0]),
            DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[1]), block_alignment[1])};
        NVTE_CHECK(t.scale_inv.shape == expected, "Tensor \"", name,
                   "\" has invalid scale_inv shape (expected ", expected, ", got ",
                   t.scale_inv.shape, ")");
      }
      if (t.has_columnwise_data()) {
        constexpr std::array<size_t, 2> block_shape{1, 16};
        constexpr std::array<size_t, 2> block_alignment{128, 4};
        const std::array<size_t, 2> expected{
            DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[0]), block_alignment[0]),
            DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[1]), block_alignment[1])};
        NVTE_CHECK(t.columnwise_scale_inv.shape == expected, "Tensor \"", name,
                   "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ",
                   t.columnwise_scale_inv.shape, ")");
      }
    }
  }
}

void CheckInputTensor(const Tensor &t, std::string_view name, bool check_scale_inv_shapes) {
  const DType type = t.dtype();
  if (is_fp8_dtype(type)) {
    // FP8 input needs to have scale_inv
    if (t.has_data()) {
      NVTE_CHECK(t.scale_inv.has_data(), "FP8 scaling factor input ", name,
                 "_scale_inverse must be allocated");
      NVTE_CHECK(t.scale_inv.dtype == DType::kFloat32 || t.scale_inv.dtype == DType::kFloat8E8M0,
                 "FP8 scaling factor input ", name,
                 "_scale_inverse has invalid dtype "
                 "(expected Float32 or Byte, got ",
                 to_string(t.scale_inv.dtype), ")");
    }
    if (t.has_columnwise_data()) {
      NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP8 scaling factor input ", name,
                 "_columnwise_scale_inverse must be allocated");
      NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat32 ||
                     t.columnwise_scale_inv.dtype == DType::kFloat8E8M0,
                 "FP8 scaling factor input ", name,
                 "_columnwise_scale_inverse has invalid dtype "
                 "(expected Float32 or Byte, got ",
                 to_string(t.columnwise_scale_inv.dtype), ")");
    }
  } else if (is_fp4_dtype(type)) {
    // TODO(ksivaman): Fix this to check for amaxes and other details.
    // For now only needed for swizzle.
    if (t.has_data()) {
      NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor input ", name,
                 "_scale_inverse must be allocated");
      NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor input ", name,
                 "_scale_inverse has invalid dtype "
                 "(expected DType::kFloat8E4M3, got ",
                 to_string(t.scale_inv.dtype), ")");
    }
    if (t.has_columnwise_data()) {
      NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor input ", name,
                 "_columnwise_scale_inverse must be allocated");
      NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3, "FP8 scaling factor input ",
                 name,
                 "_columnwise_scale_inverse has invalid dtype "
                 "(expected DType::kFloat8E4M3, got ",
                 to_string(t.columnwise_scale_inv.dtype), ")");
    }
  } else {
    NVTE_CHECK(!t.scale.has_data(), "Scale is not supported for non-FP8 input ", name);
    NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv is not supported for non-FP8 input ", name);
    NVTE_CHECK(!t.columnwise_scale_inv.has_data(), "Scale_inv is not supported for non-FP8 input ",
               name);
  }
  NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Input ", name, " is not allocated!");

  if (check_scale_inv_shapes) {
    CheckScaleTensorShape(t, name);
  }
}

void CheckOutputTensor(const Tensor &t, std::string_view name, bool allow_empty) {
  const DType type = t.dtype();
  if (is_fp8_dtype(type)) {
    // FP8 output needs to have scale, scale_inv and (if delayed scaling) amax
    if (t.scaling_mode == NVTE_DELAYED_TENSOR_SCALING && t.amax.has_data()) {
      NVTE_CHECK(t.amax.dtype == DType::kFloat32, "Invalid amax dtype (expected ",
                 to_string(DType::kFloat32), ", got ", to_string(t.amax.dtype), ")");
      NVTE_CHECK(t.amax.numel() == 1, "Invalid shape of amax in output ", name,
                 " (expected 1 entry, got shape=", t.amax.shape, ")");
    }
    if (t.has_data()) {
      NVTE_CHECK(t.scale_inv.has_data(), "FP8 scaling factor output ", name,
                 "_scale_inverse must be allocated");
      NVTE_CHECK(t.scale_inv.dtype == DType::kFloat32 || t.scale_inv.dtype == DType::kFloat8E8M0,
                 "FP8 scaling factor output ", name,
                 "_scale_inverse has invalid dtype "
                 "(expected Float32 or Float8E8M0, got ",
                 to_string(t.scale_inv.dtype), ")");
    }
    if (t.has_columnwise_data()) {
      NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP8 scaling factor output ", name,
                 "_columnwise_scale_inverse must be allocated");
      NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat32 ||
                     t.columnwise_scale_inv.dtype == DType::kFloat8E8M0,
                 "FP8 scaling factor output ", name,
                 "_columnwise_scale_inverse has invalid dtype "
                 "(expected Float32 or Float8E8M0, got ",
                 to_string(t.columnwise_scale_inv.dtype), ")");
    }
  } else if (is_fp4_dtype(type)) {
    // FP4 output needs to have the scale_inv
    if (t.has_data()) {
      NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor output ", name,
                 "_scale_inverse must be allocated");
      NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor output ", name,
                 "_scale_inverse has invalid dtype "
                 "(expected Float8E4M3, got ",
                 to_string(t.scale_inv.dtype), ")");
    }
    if (t.has_columnwise_data()) {
      NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor output ", name,
                 "_columnwise_scale_inverse must be allocated");
      NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor output ",
                 name,
                 "_columnwise_scale_inverse has invalid dtype "
                 "(expected Float8E4M3, got ",
                 to_string(t.columnwise_scale_inv.dtype), ")");
    }
  } else {
    NVTE_CHECK(!t.scale.has_data(), "Scale is not supported for non-FP8 output ", name);
    NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv is not supported for non-FP8 output ", name);
    NVTE_CHECK(!t.columnwise_scale_inv.has_data(), "Scale_inv is not supported for non-FP8 input ",
               name);
  }

  if (!allow_empty) {
    NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Output ", name, " is not allocated!");
  }

  CheckScaleTensorShape(t, name);
}

void CheckGroupedTensorShapeArrays(const GroupedTensor &t, std::string_view name) {
  NVTE_CHECK(t.num_tensors > 0, "Grouped tensor ", name, " has no tensors!");

  // Helper lambda to validate shape arrays
  // All three arrays are OPTIONAL:
  // - first_dims: empty if all tensors have same first dimension
  // - last_dims: empty if all tensors have same last dimension
  // - tensor_offsets: empty if all tensors have same shape (offsets are predictable)
  auto check_shape_array = [&](const SimpleTensor &arr, const char *arr_name) {
    if (arr.has_data()) {
      NVTE_CHECK(arr.shape.size() == 1, "Grouped tensor ", name, " ", arr_name, " must be 1D");
      NVTE_CHECK(arr.dtype == DType::kInt64, "Grouped tensor ", name, " ", arr_name,
                 " must have dtype Int64");
      NVTE_CHECK(arr.shape[0] == t.num_tensors, "Grouped tensor ", name, " ", arr_name, " size (",
                 arr.shape[0], ") must equal num_tensors (", t.num_tensors, ")");
    }
  };

  // Validate shape arrays (all optional)
  check_shape_array(t.first_dims, "first_dims");
  check_shape_array(t.last_dims, "last_dims");

  // tensor_offsets uses CSR-style prefix-sum layout with num_tensors+1 entries:
  // offsets[i] = start of tensor i, offsets[num_tensors] = total elements
  if (t.tensor_offsets.has_data()) {
    NVTE_CHECK(t.tensor_offsets.shape.size() == 1, "Grouped tensor ", name,
               " tensor_offsets must be 1D");
    NVTE_CHECK(t.tensor_offsets.dtype == DType::kInt64, "Grouped tensor ", name,
               " tensor_offsets must have dtype Int64");
    NVTE_CHECK(t.tensor_offsets.shape[0] == t.num_tensors + 1, "Grouped tensor ", name,
               " tensor_offsets size (", t.tensor_offsets.shape[0], ") must equal num_tensors+1 (",
               t.num_tensors + 1, ")");
  }

  // tensor_offsets is required if any dimension varies
  // (i.e., required unless all_same_shape())
  if (!t.all_same_shape()) {
    NVTE_CHECK(
        t.tensor_offsets.dptr != nullptr, "Grouped tensor ", name,
        " must have tensor_offsets when any dimension varies (first_dims or last_dims is set)");
  }

  // Validate logical_shape
  NVTE_CHECK(t.logical_shape.ndim == 2, "Grouped tensor ", name, " logical_shape must be 2D");
  NVTE_CHECK(t.logical_shape.data[0] > 0 && t.logical_shape.data[1] > 0, "Grouped tensor ", name,
             " logical_shape must have positive dimensions");

  // Validate all data fields are 1D (flattened)
  if (t.has_data()) {
    NVTE_CHECK(t.data.shape.size() == 1, "Grouped tensor ", name, " data must be 1D");
  }
  if (t.has_columnwise_data()) {
    NVTE_CHECK(t.columnwise_data.shape.size() == 1, "Grouped tensor ", name,
               " columnwise_data must be 1D");
  }

  // Validate data size matches logical_shape
  size_t expected_numel = t.logical_shape.data[0] * t.logical_shape.data[1];
  if (t.has_data()) {
    NVTE_CHECK(t.data.numel() == expected_numel, "Grouped tensor ", name, " data size (",
               t.data.numel(), ") must match logical_shape size (", expected_numel, ")");
  }
  if (t.has_columnwise_data()) {
    NVTE_CHECK(t.columnwise_data.numel() == expected_numel, "Grouped tensor ", name,
               " columnwise_data size (", t.columnwise_data.numel(),
               ") must match logical_shape size (", expected_numel, ")");
  }
}

// Helper function to check scale_inv for both input and output
static void CheckGroupedScaleInv(const GroupedTensor &t, std::string_view name, bool is_output) {
  const char *tensor_type = is_output ? "output" : "input";

  // Helper to check scale_inv for both rowwise and columnwise layouts
  auto check_scales = [&](DType expected_dtype) {
    if (t.has_data()) {
      NVTE_CHECK(t.scale_inv.has_data(), tensor_type, " ", name,
                 " rowwise scale_inv must be allocated");
      NVTE_CHECK(t.scale_inv.dtype == expected_dtype, tensor_type, " ", name,
                 " rowwise scale_inv has invalid dtype (expected ", to_string(expected_dtype),
                 ", got ", to_string(t.scale_inv.dtype), ")");
    }
    if (t.has_columnwise_data()) {
      NVTE_CHECK(t.columnwise_scale_inv.has_data(), tensor_type, " ", name,
                 " columnwise scale_inv must be allocated");
      NVTE_CHECK(t.columnwise_scale_inv.dtype == expected_dtype, tensor_type, " ", name,
                 " columnwise scale_inv has invalid dtype (expected ", to_string(expected_dtype),
                 ", got ", to_string(t.columnwise_scale_inv.dtype), ")");
    }
  };

  // Determine expected dtype based on data type and scaling mode
  if (is_fp8_dtype(t.dtype()) && is_tensor_scaling(t.scaling_mode)) {
    check_scales(DType::kFloat32);
  } else if (is_fp8_block_scaling(t.scaling_mode)) {
    check_scales(DType::kFloat32);
  } else if (is_mxfp8_scaling(t.scaling_mode)) {
    check_scales(DType::kFloat8E8M0);
  } else if (is_nvfp4_scaling(t.scaling_mode)) {
    check_scales(DType::kFloat8E4M3);
  } else {
    // Non-quantized types should not have scale/scale_inv
    NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv not supported for non-quantized ", tensor_type,
               " ", name);
    NVTE_CHECK(!t.columnwise_scale_inv.has_data(), "Scale_inv not supported for non-quantized ",
               tensor_type, " ", name);
  }
}

void CheckInputGroupedTensor(const GroupedTensor &t, std::string_view name) {
  NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Input grouped tensor ", name,
             " not allocated");
  CheckGroupedScaleInv(t, name, false);
  CheckGroupedTensorShapeArrays(t, name);
}

void CheckOutputGroupedTensor(const GroupedTensor &t, std::string_view name, bool allow_empty) {
  if (!allow_empty) {
    NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Output grouped tensor ", name,
               " not allocated");
  }

  // Only perform dtype-specific validation if data is allocated
  if (t.has_data() || t.has_columnwise_data()) {
    // Amax validation for delayed scaling
    if (is_fp8_dtype(t.dtype()) && t.scaling_mode == NVTE_DELAYED_TENSOR_SCALING) {
      NVTE_CHECK(t.amax.has_data(), "Output ", name, " amax must be allocated");
      NVTE_CHECK(t.amax.dtype == DType::kFloat32, "Output ", name, " amax must be Float32");
    }
    CheckGroupedScaleInv(t, name, true);
  }

  CheckGroupedTensorShapeArrays(t, name);
}

namespace {

constexpr size_t kDefaultTensorHandlePoolSizeMB = 20;
constexpr size_t kBytesPerMB = 1024 * 1024;

size_t GetTensorHandlePoolSizeMB(const char *env_var) {
  const char *env_value = std::getenv(env_var);
  if (env_value == nullptr || env_value[0] == '\0') {
    return kDefaultTensorHandlePoolSizeMB;
  }

  const std::string value(env_value);
  constexpr const char *kWhitespace = " \t\n\r\f\v";
  const size_t first = value.find_first_not_of(kWhitespace);
  const size_t last = value.find_last_not_of(kWhitespace);
  NVTE_CHECK(first != std::string::npos, env_var, " must be a positive integer.");

  const char *begin = value.c_str() + first;
  const char *expected_end = value.c_str() + last + 1;
  errno = 0;
  char *end = nullptr;
  const uint64_t pool_size_mb = std::strtoul(begin, &end, 10);

  NVTE_CHECK(value[first] >= '0' && value[first] <= '9' && end == expected_end && pool_size_mb > 0,
             env_var, " must be a positive integer, got \"", value, "\".");
  NVTE_CHECK(errno != ERANGE, env_var, " is too large.");
  NVTE_CHECK(pool_size_mb <= std::numeric_limits<size_t>::max() / kBytesPerMB, env_var,
             " is too large.");
  return static_cast<size_t>(pool_size_mb);
}

size_t GetTensorHandlePoolCapacity(size_t pool_size_mb, size_t handle_size, const char *handle_name,
                                   const char *env_var) {
  const size_t pool_size_bytes = pool_size_mb * kBytesPerMB;
  NVTE_CHECK(pool_size_bytes >= handle_size, env_var, "=", pool_size_mb,
             " MiB is too small for one ", handle_name, " handle of size ", handle_size, " bytes.");
  return pool_size_bytes / handle_size;
}

}  // namespace

class TensorAllocator {
 public:
  static TensorAllocator &instance() {
    static TensorAllocator allocator;
    return allocator;
  }

  ~TensorAllocator() {}

  void Allocate(NVTEScalingMode mode, NVTETensor *out, size_t N) {
    std::lock_guard<std::mutex> lock(mutex);
    const size_t available = free_list.size() + (memory.capacity() - memory.size());
    NVTE_CHECK(available >= N, "Cannot allocate ", N,
               " new NVTETensors. Maximum number of tensors reached: ", MAX_TENSOR_NUM, " (",
               TENSOR_HANDLE_POOL_SIZE_MB,
               " MiB handle pool). If your application legitimately needs more tensor handles, "
               "increase NVTE_TENSOR_HANDLE_POOL_SIZE_MB.");
    for (size_t i = 0; i < N; ++i) {
      uintptr_t index;
      if (!free_list.empty()) {
        index = free_list.back();
        free_list.pop_back();
      } else {
        memory.emplace_back();
        index = memory.size();
        size = index;
        memory[index - 1].nvte_tensor = reinterpret_cast<NVTETensor>(index);
      }
      memory[index - 1].scaling_mode = mode;
      out[i] = reinterpret_cast<NVTETensor>(index);
    }
    if (debug) {
      std::cout << "Allocated range of " << N << " tensors. Free list size: " << free_list.size()
                << " and capacity " << free_list.capacity() << std::endl;
    }
  }

  NVTETensor Allocate(NVTEScalingMode mode) {
    NVTETensor t;
    Allocate(mode, &t, 1);
    return t;
  }

  void Free(NVTETensor *t, size_t N) {
    std::lock_guard<std::mutex> lock(mutex);
    for (size_t i = 0; i < N; ++i) {
      uintptr_t index = reinterpret_cast<uintptr_t>(t[i]);
      if (index == 0) continue;
      NVTE_CHECK(index <= memory.size(), "Invalid tensor.");
      free_list.push_back(index);
      // Clean up
      memory[index - 1].clear();
    }
    if (debug) {
      std::cout << "Freed range of " << N << " tensors. Free list size: " << free_list.size()
                << " and capacity " << free_list.capacity() << std::endl;
    }
  }

  void Free(NVTETensor t) {
    if (reinterpret_cast<uintptr_t>(t) == 0) return;
    Free(&t, 1);
  }

  Tensor *convertNVTETensor(NVTETensor t) {
    uintptr_t index = reinterpret_cast<uintptr_t>(t);
    // 1-based indexing to enable 0-initialization of NVTETensor
    // to be invalid tensor
    static_assert(nullptr == 0);
    if (index != 0 && index <= size) {
      return &(memory[index - 1]);
    }
    return nullptr;
  }

  void setDebug(bool debug) {
    std::lock_guard<std::mutex> lock(mutex);
    this->debug = debug;
  }

 private:
  TensorAllocator() {
    std::lock_guard<std::mutex> lock(mutex);
    memory.reserve(MAX_TENSOR_NUM);
  }

  std::mutex mutex;
  std::atomic<size_t> size;
  // Should be replaced by virtual memory allocation
  const size_t TENSOR_HANDLE_POOL_SIZE_MB =
      GetTensorHandlePoolSizeMB("NVTE_TENSOR_HANDLE_POOL_SIZE_MB");
  const size_t MAX_TENSOR_NUM = GetTensorHandlePoolCapacity(
      TENSOR_HANDLE_POOL_SIZE_MB, sizeof(Tensor), "NVTETensor", "NVTE_TENSOR_HANDLE_POOL_SIZE_MB");
  std::vector<uintptr_t> free_list;
  std::vector<Tensor> memory;
  bool debug = false;
};

Tensor *convertNVTETensor(const NVTETensor t) {
  return TensorAllocator::instance().convertNVTETensor(t);
}

Tensor *convertNVTETensorCheck(const NVTETensor t) {
  Tensor *ptr = TensorAllocator::instance().convertNVTETensor(t);
  NVTE_CHECK(ptr != nullptr, "Invalid tensor.");
  return ptr;
}

// GroupedTensor allocator - similar pattern to TensorAllocator
class GroupedTensorAllocator {
 public:
  static GroupedTensorAllocator &instance() {
    static GroupedTensorAllocator allocator;
    return allocator;
  }

  ~GroupedTensorAllocator() {}

  NVTEGroupedTensor Allocate(NVTEScalingMode mode, size_t num_tensors, NVTEShape logical_shape) {
    std::lock_guard<std::mutex> lock(mutex);
    if (!free_list.empty()) {
      uintptr_t index = free_list.back();
      NVTEGroupedTensor ret = reinterpret_cast<NVTEGroupedTensor>(index);
      free_list.pop_back();
      // 1-based indexing - fully reinitialize the tensor to avoid stale data
      memory[index - 1].scaling_mode = mode;
      memory[index - 1].num_tensors = num_tensors;
      memory[index - 1].logical_shape = logical_shape;
      memory[index - 1].nvte_tensor = ret;
      return ret;
    }
    if (memory.size() < memory.capacity()) {
      memory.emplace_back(mode, num_tensors);
      GroupedTensor &t = memory.back();
      size = memory.size();
      // 1-based indexing
      uintptr_t index = memory.size();
      t.logical_shape = logical_shape;
      t.nvte_tensor = reinterpret_cast<NVTEGroupedTensor>(index);
      return reinterpret_cast<NVTEGroupedTensor>(index);
    }
    NVTE_ERROR(
        "Cannot allocate a new NVTEGroupedTensor. Maximum number of grouped tensors reached: ",
        MAX_GROUPED_TENSOR_NUM, " (", GROUPED_TENSOR_HANDLE_POOL_SIZE_MB,
        " MiB handle pool). If your application legitimately needs more grouped tensor handles, "
        "increase NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB.");
  }

  void Free(NVTEGroupedTensor t) {
    uintptr_t index = reinterpret_cast<uintptr_t>(t);
    if (index == 0) return;
    std::lock_guard<std::mutex> lock(mutex);
    NVTE_CHECK(index <= memory.size(), "Invalid grouped tensor.");
    free_list.push_back(index);
    // Clean up
    memory[index - 1].clear();
  }

  GroupedTensor *convertNVTEGroupedTensor(NVTEGroupedTensor t) {
    uintptr_t index = reinterpret_cast<uintptr_t>(t);
    // 1-based indexing to enable 0-initialization of NVTEGroupedTensor
    // to be invalid tensor
    static_assert(nullptr == 0);
    if (index != 0 && index <= size) {
      return &(memory[index - 1]);
    }
    return nullptr;
  }

 private:
  GroupedTensorAllocator() {
    std::lock_guard<std::mutex> lock(mutex);
    memory.reserve(MAX_GROUPED_TENSOR_NUM);
  }

  std::mutex mutex;
  std::atomic<size_t> size;
  const size_t GROUPED_TENSOR_HANDLE_POOL_SIZE_MB =
      GetTensorHandlePoolSizeMB("NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB");
  const size_t MAX_GROUPED_TENSOR_NUM =
      GetTensorHandlePoolCapacity(GROUPED_TENSOR_HANDLE_POOL_SIZE_MB, sizeof(GroupedTensor),
                                  "NVTEGroupedTensor", "NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB");
  std::vector<uintptr_t> free_list;
  std::vector<GroupedTensor> memory;
};

GroupedTensor *convertNVTEGroupedTensor(const NVTEGroupedTensor t) {
  return GroupedTensorAllocator::instance().convertNVTEGroupedTensor(t);
}

GroupedTensor *convertNVTEGroupedTensorCheck(const NVTEGroupedTensor t) {
  GroupedTensor *ptr = GroupedTensorAllocator::instance().convertNVTEGroupedTensor(t);
  NVTE_CHECK(ptr != nullptr, "Invalid grouped tensor.");
  return ptr;
}

}  // namespace transformer_engine

NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode) {
  NVTETensor ret = transformer_engine::TensorAllocator::instance().Allocate(scaling_mode);
  return ret;
}

void nvte_create_tensors(NVTEScalingMode scaling_mode, NVTETensor *tensors, size_t N) {
  transformer_engine::TensorAllocator::instance().Allocate(scaling_mode, tensors, N);
}

void nvte_destroy_tensor(NVTETensor tensor) {
  transformer_engine::TensorAllocator::instance().Free(tensor);
}

void nvte_destroy_tensors(NVTETensor *tensors, size_t N) {
  transformer_engine::TensorAllocator::instance().Free(tensors, N);
}

NVTEDType nvte_tensor_type(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return kNVTEFloat32;
  return static_cast<NVTEDType>(t->dtype());
}

NVTEShape nvte_make_shape(const size_t *data, size_t ndim) {
  NVTEShape ret;
  if (ndim == 0) {
    ret.ndim = 0;
    return ret;
  }
  NVTE_CHECK(ndim <= sizeof(ret.data) / sizeof(ret.data[0]),
             "Too many dims for NVTEShape (requested: ", ndim,
             ", max: ", sizeof(ret.data) / sizeof(ret.data[0]), ")");
  if (data == nullptr) {
    std::fill(ret.data, ret.data + ndim, 0);
  } else {
    std::copy(data, data + ndim, ret.data);
  }
  ret.ndim = ndim;
  return ret;
}

NVTEShape nvte_tensor_shape(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) {
    NVTE_ERROR("Invalid tensor: received null pointer in nvte_tensor_shape");
  }
  return t->shape();
}

NVTEShape nvte_tensor_columnwise_shape(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) {
    NVTE_ERROR("Invalid tensor: received null pointer in nvte_tensor_columnwise_shape");
  }
  return t->columnwise_data.shape;
}

size_t nvte_tensor_ndims(const NVTETensor tensor) { return nvte_tensor_shape(tensor).ndim; }

size_t nvte_tensor_size(const NVTETensor tensor, const size_t dim) {
  const auto &shape = nvte_tensor_shape(tensor);
  NVTE_CHECK(0 <= dim && dim < shape.ndim, "Attempted to access index ", dim,
             " in a shape array with ", shape.ndim, " entries");
  return shape.data[dim];
}

size_t nvte_tensor_numel(const NVTETensor tensor) {
  const auto &shape = nvte_tensor_shape(tensor);
  size_t numel = 1;
  for (size_t i = 0; i < shape.ndim; i++) {
    numel *= shape.data[i];
  }
  return numel;
}

size_t nvte_tensor_element_size_bits(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return 8 * sizeof(float);
  return transformer_engine::typeToNumBits(t->dtype());
}

size_t nvte_tensor_element_size(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return sizeof(float);
  NVTE_CHECK(!is_fp4_dtype(t->dtype()),
             "For FP4 type please use the nvte_tensor_element_size_bits.");
  return nvte_tensor_element_size_bits(tensor) / 8;
}

size_t nvte_tensor_size_bytes(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return 0;
  return (nvte_tensor_numel(tensor) * nvte_tensor_element_size_bits(tensor)) / 8;
}

void *nvte_tensor_data(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return nullptr;
  return t->data.dptr;
}

void *nvte_tensor_columnwise_data(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return nullptr;
  return t->columnwise_data.dptr;
}

float *nvte_tensor_amax(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return nullptr;
  NVTE_CHECK(t->amax.dtype == transformer_engine::DType::kFloat32,
             "Tensor's amax must have Float32 type!");
  return reinterpret_cast<float *>(t->amax.dptr);
}

float *nvte_tensor_scale(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return nullptr;
  NVTE_CHECK(t->scale.dtype == transformer_engine::DType::kFloat32,
             "Tensor's scale must have Float32 type!");
  return reinterpret_cast<float *>(t->scale.dptr);
}

float *nvte_tensor_scale_inv(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return nullptr;
  return reinterpret_cast<float *>(t->scale_inv.dptr);
}

void *nvte_tensor_columnwise_scale_inv(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) return nullptr;
  return t->columnwise_scale_inv.dptr;
}

NVTEShape nvte_tensor_scale_inv_shape(const NVTETensor tensor) {
  auto *t = transformer_engine::convertNVTETensor(tensor);
  if (t == nullptr) {
    return nvte_make_shape(nullptr, 1);
  }
  return nvte_make_shape(t->scale_inv.shape.data(), t->scale_inv.shape.size());
}

void nvte_set_tensor_param(NVTETensor *tensor, NVTETensorParam param_name,
                           const NVTEBasicTensor *param) {
  NVTE_CHECK(tensor != nullptr, "Tensor pointer can't be NULL.");
  auto *t = transformer_engine::convertNVTETensor(*tensor);
  NVTE_CHECK(t != nullptr, "Tensor is not allocated.");
  switch (param_name) {
    case kNVTERowwiseData:
      t->data = *param;
      break;
    case kNVTEColumnwiseData:
      t->columnwise_data = *param;
      break;
    case kNVTEScale:
      t->scale = *param;
      break;
    case kNVTEAmax:
      t->amax = *param;
      break;
    case kNVTERowwiseScaleInv:
      t->scale_inv = *param;
      break;
    case kNVTEColumnwiseScaleInv:
      t->columnwise_scale_inv = *param;
      break;
    case kNVTEColumnwiseAmax:
      t->columnwise_amax = *param;
      break;
    default:
      NVTE_ERROR("Unsupported tensor parameter (", static_cast<int>(param_name),
                 "). Consider using nvte_set_tensor_param_v2 instead.");
  }
}

NVTEBasicTensor nvte_get_tensor_param(const NVTETensor tensor, NVTETensorParam param_name) {
  if (tensor == nullptr) {
    return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 1)};
  }
  const auto &t = *transformer_engine::convertNVTETensorCheck(tensor);
  switch (param_name) {
    case kNVTERowwiseData:
      return t.data;
    case kNVTEColumnwiseData:
      return t.columnwise_data;
    case kNVTEScale:
      return t.scale;
    case kNVTEAmax:
      return t.amax;
    case kNVTERowwiseScaleInv:
      return t.scale_inv;
    case kNVTEColumnwiseScaleInv:
      return t.columnwise_scale_inv;
    case kNVTEColumnwiseAmax:
      return t.columnwise_amax;
    default:
      NVTE_ERROR("Unsupported tensor parameter (", static_cast<int>(param_name),
                 "). Consider using nvte_set_tensor_param_v2 instead.");
  }
}

void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const void *buf,
                              size_t size_in_bytes) {
  // Check attribute and buffer
  NVTE_CHECK(param < kNVTENumTensorParams, "Invalid NVTETensorParam (got ", static_cast<int>(param),
             ")");
  NVTE_CHECK(tensor != nullptr, "Tensor pointer can't be NULL.");
  auto &t = *transformer_engine::convertNVTETensorCheck(tensor);
  const auto &attr_size = transformer_engine::Tensor::attr_sizes[param];
  NVTE_CHECK(size_in_bytes >= attr_size,
             "Buffer is too small for tensor parameter "
             "(parameter ",
             static_cast<int>(param), " needs ", attr_size, " bytes, but buffer has ",
             size_in_bytes, " bytes)");
  NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)");

  // Read from buffer
  switch (param) {
    case kNVTERowwiseData: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.data = *basic_tensor;
      break;
    }
    case kNVTEColumnwiseData: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.columnwise_data = *basic_tensor;
      break;
    }
    case kNVTEScale: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.scale = *basic_tensor;
      break;
    }
    case kNVTEAmax: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.amax = *basic_tensor;
      break;
    }
    case kNVTERowwiseScaleInv: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.scale_inv = *basic_tensor;
      break;
    }
    case kNVTEColumnwiseScaleInv: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.columnwise_scale_inv = *basic_tensor;
      break;
    }
    case kNVTEColumnwiseAmax: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.columnwise_amax = *basic_tensor;
      break;
    }
    case kNVTEWithGEMMSwizzledScales:
      t.with_gemm_swizzled_scales = static_cast<bool>(*reinterpret_cast<const uint8_t *>(buf));
      break;
    case kNVTERowScaledNVFP4:
      t.row_scaled_nvfp4 = static_cast<bool>(*reinterpret_cast<const uint8_t *>(buf));
      break;
    case kNVTENVFP4E4M3Max:
      std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size);
      NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256,
                 "Unsupported NVFP4 E4M3 max (got ", t.nvfp4_e4m3_max, ")");
      break;
    default:
      NVTE_ERROR("Unsupported tensor parameter (", static_cast<int>(param), ")");
  }
}

void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, void *buf,
                              size_t size_in_bytes, size_t *size_written) {
  using namespace transformer_engine;

  // Check param
  NVTE_CHECK(param < kNVTENumTensorParams, "Invalid NVTETensorParam (got ", static_cast<int>(param),
             ")");

  // Write attribute size if provided
  const auto &attr_size = Tensor::attr_sizes[param];
  if (size_written != nullptr) {
    *size_written = attr_size;
  }

  // Return immediately if buffer is not provided
  if (buf == nullptr) {
    return;
  }

  // Check buffer size
  NVTE_CHECK(size_in_bytes >= attr_size,
             "Buffer is too small for tensor parameter "
             "(parameter ",
             static_cast<int>(param), " needs ", attr_size, " bytes, but buffer has ",
             size_in_bytes, " bytes)");

  // Get C++ tensor
  const Tensor *t = convertNVTETensor(tensor);
  std::optional<Tensor> dummy;
  if (t == nullptr) {
    // Make dummy tensor if provided tensor is invalid
    dummy.emplace();
    t = &(*dummy);
  }

  // Write to buffer
  switch (param) {
    case kNVTERowwiseData: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->data);
      break;
    }
    case kNVTEColumnwiseData: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->columnwise_data);
      break;
    }
    case kNVTEScale: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->scale);
      break;
    }
    case kNVTEAmax: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->amax);
      break;
    }
    case kNVTERowwiseScaleInv: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->scale_inv);
      break;
    }
    case kNVTEColumnwiseScaleInv: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->columnwise_scale_inv);
      break;
    }
    case kNVTEColumnwiseAmax: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->columnwise_amax);
      break;
    }
    case kNVTEWithGEMMSwizzledScales:
      *reinterpret_cast<uint8_t *>(buf) = static_cast<uint8_t>(t->with_gemm_swizzled_scales);
      break;
    case kNVTERowScaledNVFP4:
      *reinterpret_cast<uint8_t *>(buf) = static_cast<uint8_t>(t->row_scaled_nvfp4);
      break;
    case kNVTENVFP4E4M3Max:
      std::memcpy(buf, &t->nvfp4_e4m3_max, attr_size);
      break;
    default:
      NVTE_ERROR("Unsupported tensor parameter (", static_cast<int>(param), ")");
  }
}

NVTEScalingMode nvte_tensor_scaling_mode(const NVTETensor tensor) {
  if (tensor == nullptr) {
    return NVTE_DELAYED_TENSOR_SCALING;
  }
  const auto &t = *transformer_engine::convertNVTETensorCheck(tensor);
  return t.scaling_mode;
}

void nvte_tensor_pack_create(NVTETensorPack *pack) {
  transformer_engine::TensorAllocator::instance().Allocate(NVTE_DELAYED_TENSOR_SCALING,
                                                           pack->tensors, pack->MAX_SIZE);
}

void nvte_tensor_pack_destroy(NVTETensorPack *pack) {
  transformer_engine::TensorAllocator::instance().Free(pack->tensors, pack->MAX_SIZE);
}

void nvte_zero_tensor(const NVTETensor tensor, cudaStream_t stream) {
  if (tensor == nullptr) return;
  const auto &t = *transformer_engine::convertNVTETensorCheck(tensor);

  // Zero out tensor data if allocated
  if (t.data.dptr != nullptr) {
    const auto size = t.data.buffer_size_bytes();
    if (size > 0) {
      NVTE_CHECK_CUDA(cudaMemsetAsync(t.data.dptr, 0, size, stream));
    }
  }

  // Zero out amax if allocated
  if (t.amax.dptr != nullptr) {
    const auto size = t.amax.buffer_size_bytes();
    if (size > 0) {
      NVTE_CHECK_CUDA(cudaMemsetAsync(t.amax.dptr, 0, size, stream));
    }
  }
}

NVTEQuantizationConfig nvte_create_quantization_config() {
  return new transformer_engine::QuantizationConfig;
}

void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config,
                                            NVTEQuantizationConfigAttribute attr, void *buf,
                                            size_t size_in_bytes, size_t *size_written) {
  using namespace transformer_engine;

  // Write attribute size
  NVTE_CHECK(attr < kNVTEQuantizationConfigNumAttributes,
             "Invalid NVTEQuantizationConfigAttribute (got ", static_cast<int>(attr), ")");
  const auto &attr_size = QuantizationConfig::attr_sizes[attr];
  if (size_written != nullptr) {
    *size_written = attr_size;
  }

  // Return immediately if buffer is not provided
  if (buf == nullptr) {
    return;
  }

  // Check buffer size
  NVTE_CHECK(size_in_bytes >= attr_size,
             "Buffer is too small for quantization config attribute "
             "(attribute ",
             static_cast<int>(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes,
             " bytes)");

  // bool size is implementation-dependent, so we explicitly specify
  // uint8_t in the user-facing API.
  auto bool_to_uint8 = [](bool in, void *out) {
    *reinterpret_cast<uint8_t *>(out) = static_cast<uint8_t>(in);
  };

  // Write to buffer
  NVTE_CHECK(config != nullptr, "Invalid NVTEQuantizationConfig (got NULL)");
  const auto &config_ = *reinterpret_cast<const QuantizationConfig *>(config);
  switch (attr) {
    case kNVTEQuantizationConfigForcePow2Scales:
      bool_to_uint8(config_.force_pow_2_scales, buf);
      break;
    case kNVTEQuantizationConfigAmaxEpsilon:
      std::memcpy(buf, &config_.amax_epsilon, attr_size);
      break;
    case kNVTEQuantizationConfigNoopTensor:
      std::memcpy(buf, &config_.noop_tensor, attr_size);
      break;
    case kNVTEQuantizationConfigFloat8BlockScaleTensorFormat: {
      // Deprecated
      const auto invalid = Float8BlockScaleTensorFormat::INVALID;
      std::memcpy(buf, &invalid, attr_size);
      break;
    }
    case kNVTEQuantizationConfigRNGState:
      std::memcpy(buf, &config_.rng_state, attr_size);
      break;
    case kNVTEQuantizationConfigNVFP42DQuantization:
      bool_to_uint8(config_.nvfp4_2d_quantization, buf);
      break;
    case kNVTEQuantizationConfigStochasticRounding:
      bool_to_uint8(config_.stochastic_rounding, buf);
      break;
    case kNVTEQuantizationConfigUseFastMath:
      bool_to_uint8(config_.use_fast_math, buf);
      break;
    case kNVTEQuantizationConfigNVFP44Over6Mode: {
      const auto val = static_cast<uint8_t>(config_.nvfp4_4over6_mode);
      std::memcpy(buf, &val, attr_size);
      break;
    }
    case kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath:
      bool_to_uint8(config_.nvfp4_4over6_err_use_fast_math, buf);
      break;
    default:
      NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast<int>(attr), ")");
  }
}

void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config,
                                            NVTEQuantizationConfigAttribute attr, const void *buf,
                                            size_t size_in_bytes) {
  using namespace transformer_engine;

  // Check attribute and buffer
  NVTE_CHECK(attr < kNVTEQuantizationConfigNumAttributes,
             "Invalid NVTEQuantizationConfigAttribute (got ", static_cast<int>(attr), ")");
  const auto &attr_size = QuantizationConfig::attr_sizes[attr];
  NVTE_CHECK(size_in_bytes >= attr_size,
             "Buffer is too small for quantization config attribute "
             "(attribute ",
             static_cast<int>(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes,
             " bytes)");
  NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)");

  // bool size is implementation-dependent, so we explicitly specify
  // uint8_t in the user-facing API.
  auto uint8_to_bool = [](const void *in, bool &out) {
    out = static_cast<bool>(*reinterpret_cast<const uint8_t *>(in));
  };

  // Read from buffer
  NVTE_CHECK(config != nullptr, "Invalid NVTEQuantizationConfig (got NULL)");
  auto &config_ = *reinterpret_cast<QuantizationConfig *>(config);
  switch (attr) {
    case kNVTEQuantizationConfigForcePow2Scales:
      uint8_to_bool(buf, config_.force_pow_2_scales);
      break;
    case kNVTEQuantizationConfigAmaxEpsilon:
      std::memcpy(&config_.amax_epsilon, buf, attr_size);
      break;
    case kNVTEQuantizationConfigNoopTensor:
      std::memcpy(&config_.noop_tensor, buf, attr_size);
      break;
    case kNVTEQuantizationConfigFloat8BlockScaleTensorFormat:
      // Deprecated
      break;
    case kNVTEQuantizationConfigRNGState:
      std::memcpy(&config_.rng_state, buf, attr_size);
      break;
    case kNVTEQuantizationConfigNVFP42DQuantization:
      uint8_to_bool(buf, config_.nvfp4_2d_quantization);
      break;
    case kNVTEQuantizationConfigStochasticRounding:
      uint8_to_bool(buf, config_.stochastic_rounding);
      break;
    case kNVTEQuantizationConfigUseFastMath:
      uint8_to_bool(buf, config_.use_fast_math);
      break;
    case kNVTEQuantizationConfigNVFP44Over6Mode: {
      const auto val = *reinterpret_cast<const uint8_t *>(buf);
      NVTE_CHECK(val == static_cast<uint8_t>(kNVTENVFP44Over6Disabled) ||
                     val == static_cast<uint8_t>(kNVTENVFP44Over6MinMAE) ||
                     val == static_cast<uint8_t>(kNVTENVFP44Over6MinMSE),
                 "Invalid NVFP4 4over6 mode (got ", static_cast<int>(val), ")");
      config_.nvfp4_4over6_mode = static_cast<NVTENVFP44Over6Mode>(val);
      break;
    }
    case kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath:
      uint8_to_bool(buf, config_.nvfp4_4over6_err_use_fast_math);
      break;
    default:
      NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast<int>(attr), ")");
  }
}

void nvte_destroy_quantization_config(NVTEQuantizationConfig config) {
  if (config != nullptr) {
    delete reinterpret_cast<transformer_engine::QuantizationConfig *>(config);
  }
}

int nvte_is_non_tn_fp8_gemm_supported() {
  int num_devices = transformer_engine::cuda::num_devices();
  static std::vector<int> cache(num_devices, -1);
  static std::vector<std::once_flag> flags(num_devices);
  int device_id = transformer_engine::cuda::current_device();
  std::call_once(flags[device_id], [&]() {
    int deviceComputeCapability = transformer_engine::cuda::sm_arch(device_id);
    // Note: this is temporary restriction and should be lifted in the future.
    // (remove the note once it's done.)
    cache[device_id] = (deviceComputeCapability >= 100 && deviceComputeCapability < 120) ||
                       deviceComputeCapability >= 130;
  });
  return cache[device_id];
}

// Grouped Tensor C API implementations
NVTEGroupedTensor nvte_create_grouped_tensor(NVTEScalingMode scaling_mode, size_t num_tensors,
                                             NVTEShape logical_shape) {
  NVTE_CHECK(num_tensors > 0, "Number of tensors must be greater than 0");
  NVTE_CHECK(logical_shape.ndim == 2, "Logical shape must be 2D");
  // NVTE_CHECK(logical_shape.data[0] > 0 && logical_shape.data[1] > 0,
  //            "Logical shape must have positive dimensions");
  NVTEGroupedTensor ret = transformer_engine::GroupedTensorAllocator::instance().Allocate(
      scaling_mode, num_tensors, logical_shape);
  return ret;
}

void nvte_destroy_grouped_tensor(NVTEGroupedTensor tensor) {
  transformer_engine::GroupedTensorAllocator::instance().Free(tensor);
}

void nvte_set_grouped_tensor_param(NVTEGroupedTensor tensor, NVTEGroupedTensorParam param,
                                   const void *buf, size_t size_in_bytes) {
  using namespace transformer_engine;

  // Check attribute and buffer
  NVTE_CHECK(param < kNVTENumGroupedTensorParams, "Invalid NVTEGroupedTensorParam (got ",
             static_cast<int>(param), ")");
  NVTE_CHECK(tensor != nullptr, "Grouped tensor pointer can't be NULL.");
  auto &t = *convertNVTEGroupedTensorCheck(tensor);
  const auto &attr_size = GroupedTensor::attr_sizes[param];
  NVTE_CHECK(size_in_bytes >= attr_size,
             "Buffer is too small for grouped tensor parameter "
             "(parameter ",
             static_cast<int>(param), " needs ", attr_size, " bytes, but buffer has ",
             size_in_bytes, " bytes)");
  NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)");

  // Read from buffer
  switch (param) {
    case kNVTEGroupedRowwiseData: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.data = *basic_tensor;
      break;
    }
    case kNVTEGroupedColumnwiseData: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.columnwise_data = *basic_tensor;
      break;
    }
    case kNVTEGroupedScale: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.scale = *basic_tensor;
      break;
    }
    case kNVTEGroupedAmax: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.amax = *basic_tensor;
      break;
    }
    case kNVTEGroupedRowwiseScaleInv: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.scale_inv = *basic_tensor;
      break;
    }
    case kNVTEGroupedColumnwiseScaleInv: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.columnwise_scale_inv = *basic_tensor;
      break;
    }
    case kNVTEGroupedColumnwiseAmax: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.columnwise_amax = *basic_tensor;
      break;
    }
    case kNVTEGroupedFirstDims: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.first_dims = *basic_tensor;
      NVTE_CHECK(t.first_dims.dtype == DType::kInt64, "first_dims must have dtype Int64");
      break;
    }
    case kNVTEGroupedLastDims: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.last_dims = *basic_tensor;
      NVTE_CHECK(t.last_dims.dtype == DType::kInt64, "last_dims must have dtype Int64");
      break;
    }
    case kNVTEGroupedTensorOffsets: {
      const NVTEBasicTensor *basic_tensor = reinterpret_cast<const NVTEBasicTensor *>(buf);
      t.tensor_offsets = *basic_tensor;
      NVTE_CHECK(t.tensor_offsets.dtype == DType::kInt64, "tensor_offsets must have dtype Int64");
      break;
    }
    case kNVTEGroupedWithGEMMSwizzledScales:
      t.with_gemm_swizzled_scales = static_cast<bool>(*reinterpret_cast<const uint8_t *>(buf));
      break;
    default:
      NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast<int>(param), ")");
  }
}

void nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, NVTEGroupedTensorParam param,
                                   void *buf, size_t size_in_bytes, size_t *size_written) {
  using namespace transformer_engine;

  // Check param
  NVTE_CHECK(param < kNVTENumGroupedTensorParams, "Invalid NVTEGroupedTensorParam (got ",
             static_cast<int>(param), ")");

  // Write attribute size if provided
  const auto &attr_size = GroupedTensor::attr_sizes[param];
  if (size_written != nullptr) {
    *size_written = attr_size;
  }

  // Return immediately if buffer is not provided
  if (buf == nullptr) {
    return;
  }

  // Check buffer size
  NVTE_CHECK(size_in_bytes >= attr_size,
             "Buffer is too small for grouped tensor parameter "
             "(parameter ",
             static_cast<int>(param), " needs ", attr_size, " bytes, but buffer has ",
             size_in_bytes, " bytes)");

  // Get C++ grouped tensor
  const GroupedTensor *t = convertNVTEGroupedTensor(tensor);
  std::optional<GroupedTensor> dummy;
  if (t == nullptr) {
    // Make dummy grouped tensor if provided tensor is invalid
    dummy.emplace(NVTE_DELAYED_TENSOR_SCALING, 1);
    t = &(*dummy);
  }

  // Write to buffer
  switch (param) {
    case kNVTEGroupedRowwiseData: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->data);
      break;
    }
    case kNVTEGroupedColumnwiseData: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->columnwise_data);
      break;
    }
    case kNVTEGroupedScale: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->scale);
      break;
    }
    case kNVTEGroupedAmax: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->amax);
      break;
    }
    case kNVTEGroupedRowwiseScaleInv: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->scale_inv);
      break;
    }
    case kNVTEGroupedColumnwiseScaleInv: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->columnwise_scale_inv);
      break;
    }
    case kNVTEGroupedColumnwiseAmax: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->columnwise_amax);
      break;
    }
    case kNVTEGroupedFirstDims: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->first_dims);
      break;
    }
    case kNVTEGroupedLastDims: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->last_dims);
      break;
    }
    case kNVTEGroupedTensorOffsets: {
      NVTEBasicTensor *basic_tensor = reinterpret_cast<NVTEBasicTensor *>(buf);
      *basic_tensor = static_cast<NVTEBasicTensor>(t->tensor_offsets);
      break;
    }
    case kNVTEGroupedWithGEMMSwizzledScales:
      *reinterpret_cast<uint8_t *>(buf) = static_cast<uint8_t>(t->with_gemm_swizzled_scales);
      break;
    default:
      NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast<int>(param), ")");
  }
}

size_t nvte_grouped_tensor_num_tensors(const NVTEGroupedTensor tensor) {
  auto *t = transformer_engine::convertNVTEGroupedTensor(tensor);
  if (t == nullptr) return 0;
  return t->num_tensors;
}

NVTEDType nvte_grouped_tensor_type(const NVTEGroupedTensor tensor) {
  auto *t = transformer_engine::convertNVTEGroupedTensor(tensor);
  if (t == nullptr) return kNVTEFloat32;
  return static_cast<NVTEDType>(t->dtype());
}

NVTEScalingMode nvte_grouped_tensor_scaling_mode(const NVTEGroupedTensor tensor) {
  if (tensor == nullptr) {
    return NVTE_DELAYED_TENSOR_SCALING;
  }
  const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor);
  return t.scaling_mode;
}

NVTEShape nvte_get_grouped_tensor_logical_shape(const NVTEGroupedTensor tensor) {
  if (tensor == nullptr) {
    return nvte_make_shape(nullptr, 1);
  }
  const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor);
  return t.logical_shape;
}
