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

#include "transformer_engine/cast.h"

#include <algorithm>
#include <array>
#include <cstdint>
#include <memory>
#include <mutex>
#include <optional>
#include <tuple>
#include <utility>
#include <vector>

#include "../extensions.h"
#include "common.h"
#include "common/common.h"
#include "common/util/cuda_runtime.h"
#include "common/util/system.h"
#include "pybind.h"
#include "transformer_engine/multi_tensor.h"
#include "transformer_engine/recipe.h"
#include "transformer_engine/transformer_engine.h"

namespace transformer_engine {
namespace pytorch {

namespace {

std::vector<size_t> get_tensor_shape(const TensorWrapper &tensor) {
  const auto &shape = tensor.shape();
  return std::vector<size_t>(shape.data, shape.data + shape.ndim);
}

void allreduce_nvfp4_amax_tensors(NVFP4Quantizer *nvfp4_quantizer_cpp,
                                  std::vector<at::Tensor> &&amax_tensors) {
  if (!nvfp4_quantizer_cpp->with_amax_reduction || amax_tensors.empty()) {
    return;
  }
  c10d::AllreduceCoalescedOptions opts;
  opts.reduceOp = c10d::ReduceOp::MAX;
  NVTE_SCOPED_GIL_RELEASE({
    nvfp4_quantizer_cpp->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait();
  });
}

}  // namespace

py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output,
                    std::optional<at::Tensor> noop_flag) {
  // Convert quantizer to C++ object
  auto quantizer_cpp = convert_quantizer(quantizer);

  // Convert input tensor to C++ object
  auto input_contiguous = tensor.contiguous();
  auto input_cpp = makeTransformerEngineTensor(input_contiguous);

  // Initialize output tensor
  TensorWrapper output_cpp;
  py::object output_py;
  if (output.is_none()) {
    const auto shape = get_tensor_shape(input_cpp);
    const auto fake_dtype = input_cpp.dtype();
    std::tie(output_cpp, output_py) = quantizer_cpp->create_tensor(shape, fake_dtype);
  } else {
    std::tie(output_cpp, output_py) = quantizer_cpp->convert_and_update_tensor(output);
  }

  // Initialize no-op flag
  std::optional<TensorWrapper> noop_flag_cpp;
  if (noop_flag.has_value()) {
    noop_flag_cpp = makeTransformerEngineTensor(*noop_flag);
  }

  // Perform quantization
  quantizer_cpp->quantize(input_cpp, output_cpp, noop_flag_cpp);

  // Post-quantize swizzle for quantizers whose kernel does not bake
  // the GEMM-swizzled scale layout in directly
  if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales()) {
    inplace_swizzle_scale_for_gemm(output_py);
  }

  return output_py;
}

py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
                                    const at::Tensor &rowwise_amax,
                                    const at::Tensor &columnwise_amax) {
  using namespace transformer_engine::pytorch::detail;
  init_extension();

  NVTE_CHECK(tensor.dim() >= 2, "Tensor must be at least 2D");
  NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(),
             "Precomputed amax tensors must be CUDA tensors.");
  NVTE_CHECK(
      rowwise_amax.scalar_type() == at::kFloat && columnwise_amax.scalar_type() == at::kFloat,
      "Precomputed amax tensors must be float32.");
  NVTE_CHECK(rowwise_amax.numel() == 1 && columnwise_amax.numel() == 1,
             "nvfp4_quantize_with_amax expects scalar rowwise and columnwise amaxes.");

  auto quantizer_cpp = convert_quantizer(quantizer);
  NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()),
             "nvfp4_quantize_with_amax only supports NVFP4 quantizers.");
  NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast<NVFP4Quantizer *>(quantizer_cpp.get());

  auto input_contiguous = tensor.contiguous();
  auto input_cpp = makeTransformerEngineTensor(input_contiguous);

  const auto shape = get_tensor_shape(input_cpp);
  const auto fake_dtype = input_cpp.dtype();
  auto [output_cpp, output_py] = quantizer_cpp->create_tensor(shape, fake_dtype);

  if (output_cpp.get_amax().data_ptr != nullptr) {
    output_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, getTensorShape(rowwise_amax));
    output_py.attr("_amax_rowwise") = py::cast(rowwise_amax);
  }
  if (output_cpp.get_columnwise_amax().data_ptr != nullptr) {
    output_cpp.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32,
                                   getTensorShape(columnwise_amax));
    output_py.attr("_amax_columnwise") = py::cast(columnwise_amax);
  }

  nvfp4_quantizer_cpp->quantize_impl(input_cpp, output_cpp, std::nullopt, false);
  if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales()) {
    inplace_swizzle_scale_for_gemm(output_py);
  }

  return output_py;
}

py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector<size_t> &shape,
                                         at::ScalarType dtype, at::Device device, bool pin_memory) {
  auto quantizer_cpp = convert_quantizer(quantizer);
  auto te_dtype = GetTransformerEngineDType(dtype);
  auto [_, output_py] = quantizer_cpp->create_tensor(shape, te_dtype, device, pin_memory);
  return output_py;
}

namespace {

// helper functions for NVFP4 grouped quantization (cuda graph safe with shapes stored in device without D2H copy)
void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor,
                               GroupedTensorWrapper &grouped_output_tensor,
                               NVFP4Quantizer *nvfp4_quantizer_cpp, cudaStream_t stream,
                               bool compute_amax) {
  size_t num_tensors = grouped_input_tensor.num_tensors();

  // assert the 2D scaling case, since 2D scaling grouped quant kernel is not ready yet
  NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization,
             "2D scaling grouped quant kernel is not ready yet");
  NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled,
             "NVFP4 4over6 quantization is not supported for grouped quantization.");
  NVTE_CHECK(nvfp4_quantizer_cpp->with_rht,
             "graph safe grouped quant kernel for non-RHT path is not ready yet");
  NVTE_CHECK(nvfp4_quantizer_cpp->with_post_rht_amax,
             "grouped NVFP4 RHT quantization expects post-RHT amax buffers.");

  auto quant_config_cpp = QuantizationConfigWrapper();

  // stochastic rounding
  bool need_stochastic_rounding = nvfp4_quantizer_cpp->stochastic_rounding;
  auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA);
  at::Tensor rng_states_tensor;  // Declare tensor outside, do not allocate yet
  TensorWrapper te_rng_state;

  if (need_stochastic_rounding) {
    // in fused kernel, one rng state will be used by the grouped kernel to generate random
    // number for different tensors in the group, so we only need to allocate one rng state
    const size_t rng_elts_per_thread = 1024 * num_tensors;
    rng_states_tensor = torch::empty({2}, opts);
    auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
        std::nullopt, at::cuda::detail::getDefaultCUDAGenerator());
    at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread);
    philox_unpack(philox_args, static_cast<int64_t *>(rng_states_tensor.data_ptr()));

    te_rng_state = makeTransformerEngineTensor(rng_states_tensor);
    quant_config_cpp.set_rng_state(te_rng_state.data());
    quant_config_cpp.set_stochastic_rounding(true);
  }

  // fast math
  const auto use_fast_math = transformer_engine::getenv<bool>("NVTE_USE_FAST_MATH");
  if (use_fast_math) {
    quant_config_cpp.set_use_fast_math(true);
  }

  if (compute_amax) {
    NVTE_SCOPED_GIL_RELEASE({
      nvte_group_hadamard_transform_amax_graph_safe(
          grouped_input_tensor.data(), grouped_output_tensor.data(), 0,
          nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream);
    });
  }

  // RHT cast fusion
  auto tile_scheduler_workspace_torch = at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32));
  auto nvte_tile_scheduler_workspace = makeTransformerEngineTensor(tile_scheduler_workspace_torch);

  auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix);
  NVTE_SCOPED_GIL_RELEASE({
    nvte_group_hadamard_transform_cast_fusion_graph_safe(
        grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(),
        quant_config_cpp, nvte_tile_scheduler_workspace.data(), stream);
  });
}

float fp8_max_for_dtype(const DType dtype) {
  if (!is_fp8_dtype(dtype)) {
    NVTE_ERROR("Expected FP8 dtype for grouped current-scaling quantization, got ",
               to_string(dtype), ".");
  }
  TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(
      dtype, T, { return transformer_engine::detail::TypeExtrema<T>::max; });
  return 0.0f;
}

// Computes per-group amax and scale/scale_inv for grouped FP8 current scaling.
//
// 1) nvte_group_compute_amax_with_config: per-group amax over the input tensor
//    (current scaling -> amax is computed from this batch's inputs, never from
//    history).
// 2) Optional NCCL allreduce on the amax buffer (when DP/TP amax sync is on).
// 3) nvte_group_compute_scale_from_amax: graph-safe per-group math kernel that
//    derives scale = max_fp8 / amax and scale_inv = 1/scale per group in a
//    single launch. It honors the config's no-op flag, so a skipped weight
//    update (skip_fp8_weight_update) preserves the cached scale instead of
//    recomputing it from a stale/uninitialized amax.
//
// After this returns, the per-group `scale` / `scale_inv` buffers are
// populated and the actual cast/transpose is performed by `nvte_group_quantize`.
void compute_grouped_fp8_current_scaling_amax_and_scale(
    const GroupedTensorWrapper &grouped_input_tensor,
    const GroupedTensorWrapper &grouped_output_tensor, const py::object &grouped_output_py,
    Float8CurrentScalingQuantizer *quantizer_cpp, const std::optional<TensorWrapper> &noop_flag) {
  QuantizationConfigWrapper quant_config;
  quant_config.set_force_pow_2_scales(quantizer_cpp->force_pow_2_scales);
  quant_config.set_amax_epsilon(quantizer_cpp->amax_epsilon);
  if (noop_flag.has_value()) {
    quant_config.set_noop_tensor(noop_flag->data());
  }

  auto stream = at::cuda::getCurrentCUDAStream();
  NVTE_SCOPED_GIL_RELEASE({
    nvte_group_compute_amax_with_config(grouped_input_tensor.data(), grouped_output_tensor.data(),
                                        quant_config, stream);
  });

  if (quantizer_cpp->with_amax_reduction) {
    // NCCL collectives require an at::Tensor; the amax buffer lives on the Python
    // object (same allocation as the grouped tensor's amax), so fetch it only on
    // this path.
    auto amax = grouped_output_py.attr("amax").cast<at::Tensor>();
    c10d::AllreduceOptions opts;
    opts.reduceOp = c10d::ReduceOp::MAX;
    std::vector<at::Tensor> tensors = {amax};
    NVTE_SCOPED_GIL_RELEASE(
        { quantizer_cpp->amax_reduction_group->allreduce(tensors, opts)->wait(); });
  }

  // Derive per-group scale/scale_inv from the (possibly reduced) amax. The same
  // quant_config carries force_pow_2_scales/amax_epsilon and, crucially, the
  // no-op flag, so the update is skipped on device when skip_fp8_weight_update is
  // set -- preserving the cached scale instead of recomputing it from a
  // stale/uninitialized amax. Buffers and the FP8 max are read from the grouped
  // tensor inside the kernel launch.
  NVTE_SCOPED_GIL_RELEASE(
      { nvte_group_compute_scale_from_amax(grouped_output_tensor.data(), quant_config, stream); });
}

}  // namespace

py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                          std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                          std::optional<at::Tensor> tensor_offsets,
                          std::optional<at::Tensor> noop_flag) {
  using namespace transformer_engine::pytorch::detail;
  init_extension();

  NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D");

  // No-op flag for CUDA graph weight caching (skip_fp8_weight_update).
  std::optional<TensorWrapper> noop_flag_cpp;
  if (noop_flag.has_value()) {
    noop_flag_cpp = makeTransformerEngineTensor(*noop_flag);
  }

  std::vector<size_t> logical_shape;
  for (const auto &d : tensor.sizes()) {
    logical_shape.push_back(d);
  }
  const auto logical_first_dim = logical_shape[0];
  const auto logical_last_dim = logical_shape[1];

  bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0;

  auto quantizer_cpp = convert_quantizer(quantizer);

  // Create input GroupedTensor.
  auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape);
  grouped_input_tensor.set_rowwise_data(tensor.data_ptr(),
                                        GetTransformerEngineDType(tensor.scalar_type()),
                                        std::vector<size_t>{static_cast<size_t>(tensor.numel())});

  // Create output GroupedTensor.
  auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor(
      num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
      py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
      logical_first_dim, logical_last_dim);

  // dispatch to scaling methods
  enum class GroupedQuantizationMode {
    FP8_CURRENT_SCALING_GROUPED_QUANTIZE,
    MXFP8_GROUPED_QUANTIZE,
    NVFP4_GROUPED_QUANTIZE,
    FP8_BLOCKWISE_GROUPED_QUANTIZE,
    INVALID_FOR_GROUPED_QUANTIZE
  };
  GroupedQuantizationMode grouped_quantization_mode =
      GroupedQuantizationMode::INVALID_FOR_GROUPED_QUANTIZE;
  if (detail::IsMXFP8Quantizers(quantizer.ptr())) {
    grouped_quantization_mode = GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE;
  } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) {
    grouped_quantization_mode = GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE;
  } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) {
    grouped_quantization_mode = GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE;
  } else if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) {
    grouped_quantization_mode = GroupedQuantizationMode::FP8_BLOCKWISE_GROUPED_QUANTIZE;
  }

  if (empty_input_buffer) {
    // early return for empty input buffer
    // just return the output tensor as is
    // no need to quantize
    return py::reinterpret_borrow<py::object>(grouped_output_py);
  }

  switch (grouped_quantization_mode) {
    case GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE: {
      // NVFP4 grouped quantization
      NVTE_CHECK(!last_dims.has_value(),
                 "group_quantize: varying last dim is not supported with NVFP4.");
      NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast<NVFP4Quantizer *>(quantizer_cpp.get());
      group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp,
                                nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), true);
      break;
    }
    case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: {
      auto *fp8_quantizer_cpp = static_cast<Float8CurrentScalingQuantizer *>(quantizer_cpp.get());
      compute_grouped_fp8_current_scaling_amax_and_scale(
          grouped_input_tensor, grouped_output_tensor_cpp, grouped_output_py, fp8_quantizer_cpp,
          noop_flag_cpp);
      QuantizationConfigWrapper quant_config_cpp;
      if (noop_flag_cpp.has_value()) {
        quant_config_cpp.set_noop_tensor(noop_flag_cpp->data());
      }
      NVTE_SCOPED_GIL_RELEASE({
        nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(),
                            quant_config_cpp, at::cuda::getCurrentCUDAStream());
      });
      break;
    }
    case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: {
      QuantizationConfigWrapper quant_config_cpp;
      if (noop_flag_cpp.has_value()) {
        quant_config_cpp.set_noop_tensor(noop_flag_cpp->data());
      }
      NVTE_SCOPED_GIL_RELEASE({
        nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(),
                            quant_config_cpp, at::cuda::getCurrentCUDAStream());
      });
      break;
    }
    case GroupedQuantizationMode::FP8_BLOCKWISE_GROUPED_QUANTIZE: {
      Float8BlockQuantizer *fp8_block_quantizer_cpp =
          static_cast<Float8BlockQuantizer *>(quantizer_cpp.get());
      QuantizationConfigWrapper quant_config_cpp;
      quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales);
      quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon);
      NVTE_SCOPED_GIL_RELEASE({
        nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(),
                            quant_config_cpp, at::cuda::getCurrentCUDAStream());
      });
      break;
    }
    case GroupedQuantizationMode::INVALID_FOR_GROUPED_QUANTIZE:
    default:
      NVTE_ERROR(
          "group_quantize: only supports MXFP8, NVFP4, Float8CurrentScalingQuantizer, or "
          "Float8Blockwise quantizer.");
      break;
  }

  return py::reinterpret_borrow<py::object>(grouped_output_py);
}

py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
                                          const size_t num_tensors,
                                          std::optional<at::Tensor> first_dims,
                                          std::optional<at::Tensor> last_dims,
                                          const at::Tensor &rowwise_amax,
                                          const at::Tensor &columnwise_amax,
                                          std::optional<at::Tensor> tensor_offsets) {
  using namespace transformer_engine::pytorch::detail;
  init_extension();

  NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D");
  NVTE_CHECK(!last_dims.has_value(),
             "nvfp4_group_quantize_with_amax: varying last dim is not supported with NVFP4.");
  NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(),
             "Precomputed amax tensors must be CUDA tensors.");
  NVTE_CHECK(
      rowwise_amax.scalar_type() == at::kFloat && columnwise_amax.scalar_type() == at::kFloat,
      "Precomputed amax tensors must be float32.");
  NVTE_CHECK(rowwise_amax.numel() == static_cast<int64_t>(num_tensors),
             "Rowwise amax must contain one value per group.");
  NVTE_CHECK(columnwise_amax.numel() == static_cast<int64_t>(num_tensors),
             "Columnwise amax must contain one value per group.");

  std::vector<size_t> logical_shape;
  for (const auto &d : tensor.sizes()) {
    logical_shape.push_back(d);
  }
  const auto logical_first_dim = logical_shape[0];
  const auto logical_last_dim = logical_shape[1];

  bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0;

  auto quantizer_cpp = convert_quantizer(quantizer);
  NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()),
             "nvfp4_group_quantize_with_amax only supports NVFP4 quantizers.");
  NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast<NVFP4Quantizer *>(quantizer_cpp.get());

  auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape);
  grouped_input_tensor.set_rowwise_data(
      tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor));

  auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor(
      num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
      py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
      logical_first_dim, logical_last_dim);

  if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) {
    grouped_output_tensor_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32,
                                       getTensorShape(rowwise_amax));
    grouped_output_py.attr("amax") = py::cast(rowwise_amax);
  }
  if (grouped_output_tensor_cpp.get_columnwise_amax().data_ptr != nullptr) {
    grouped_output_tensor_cpp.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32,
                                                  getTensorShape(columnwise_amax));
    grouped_output_py.attr("columnwise_amax") = py::cast(columnwise_amax);
  }

  std::vector<at::Tensor> amax_tensors;
  if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) {
    amax_tensors.push_back(rowwise_amax);
  }
  if (grouped_output_tensor_cpp.get_columnwise_amax().data_ptr != nullptr) {
    amax_tensors.push_back(columnwise_amax);
  }
  allreduce_nvfp4_amax_tensors(nvfp4_quantizer_cpp, std::move(amax_tensors));

  if (empty_input_buffer) {
    return py::reinterpret_borrow<py::object>(grouped_output_py);
  }

  group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, nvfp4_quantizer_cpp,
                            at::cuda::getCurrentCUDAStream(), false);

  return py::reinterpret_borrow<py::object>(grouped_output_py);
}

py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                const size_t num_tensors, std::optional<at::Tensor> first_dims,
                                std::optional<at::Tensor> last_dims,
                                std::optional<at::Tensor> tensor_offsets) {
  using namespace transformer_engine::pytorch::detail;
  init_extension();

  NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D");
  NVTE_CHECK(!last_dims.has_value(),
             "bgrad_group_quantize: varying last dim is not supported because the underlying "
             "MXFP8 dbias kernel requires a constant last dimension across grouped tensors.");

  std::vector<size_t> logical_shape;
  for (const auto &d : tensor.sizes()) {
    logical_shape.push_back(d);
  }
  const auto logical_first_dim = logical_shape[0];
  const auto logical_last_dim = logical_shape[1];

  bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0;

  NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()) ||
                 detail::IsFloat8BlockwiseQuantizers(quantizer.ptr()),
             "bgrad_group_quantize: only MXFP8 and FP8 block-scaling quantizers are supported.");

  auto quantizer_cpp = convert_quantizer(quantizer);

  // MXFP8 carries no quantization knobs; FP8 block-scaling reads scale constraints
  // off the quantizer, matching the forward group_quantize dispatch.
  QuantizationConfigWrapper quant_config_cpp;
  if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) {
    auto *fp8_block_quantizer_cpp = static_cast<Float8BlockQuantizer *>(quantizer_cpp.get());
    quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales);
    quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon);
  }

  auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape);
  grouped_input_tensor.set_rowwise_data(tensor.data_ptr(),
                                        GetTransformerEngineDType(tensor.scalar_type()),
                                        std::vector<size_t>{static_cast<size_t>(tensor.numel())});

  auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor(
      num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
      py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
      logical_first_dim, logical_last_dim);

  if (empty_input_buffer) {
    at::Tensor dbias_torch =
        at::zeros({static_cast<int64_t>(num_tensors), static_cast<int64_t>(logical_last_dim)},
                  tensor.options());
    return py::make_tuple(py::reinterpret_borrow<py::object>(grouped_output_py),
                          py::cast(std::move(dbias_torch)));
  }

  const std::vector<size_t> dbias_logical_shape = {num_tensors, logical_last_dim};
  GroupedTensorWrapper grouped_dbias(num_tensors, dbias_logical_shape, NVTE_DELAYED_TENSOR_SCALING);
  at::Tensor dbias_torch =
      at::empty({static_cast<int64_t>(num_tensors), static_cast<int64_t>(logical_last_dim)},
                tensor.options());
  grouped_dbias.set_rowwise_data(dbias_torch.data_ptr(),
                                 GetTransformerEngineDType(tensor.scalar_type()),
                                 getTensorShape(dbias_torch));
  TensorWrapper workspace_nvte;
  auto stream = at::cuda::getCurrentCUDAStream();
  NVTE_SCOPED_GIL_RELEASE({
    nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(),
                              grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp,
                              stream);
  });
  if (workspace_nvte.ndim() > 0 && workspace_nvte.numel() > 0) {
    at::Tensor workspace_torch = allocateSpace(workspace_nvte.shape(), workspace_nvte.dtype());
    workspace_nvte = makeTransformerEngineTensor(workspace_torch.data_ptr(), workspace_nvte.shape(),
                                                 workspace_nvte.dtype());
  }
  NVTE_SCOPED_GIL_RELEASE({
    nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(),
                              grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp,
                              stream);
  });
  return py::make_tuple(py::reinterpret_borrow<py::object>(grouped_output_py),
                        py::cast(std::move(dbias_torch)));
}

py::object dequantize(const py::handle &input, transformer_engine::DType otype) {
  init_extension();

  const auto none = py::none();

  const auto &input_tensor = makeTransformerEngineTensor(input, none);

  NoneQuantizer q(none);

  const auto &shape = convertShape(input_tensor.shape());

  auto [out_tensor, out] = q.create_tensor(shape, otype);

  NVTE_SCOPED_GIL_RELEASE({
    nvte_dequantize(input_tensor.data(), out_tensor.data(), at::cuda::getCurrentCUDAStream());
  });

  return out;
}

py::object group_dequantize(const py::handle &input, transformer_engine::DType otype) {
  using namespace pybind11::literals;
  init_extension();

  // Extract fields from the Python GroupedTensor.
  const auto num_tensors = input.attr("num_tensors").cast<size_t>();
  const auto logical_shape_py = input.attr("logical_shape").cast<py::tuple>();
  const auto logical_first_dim = logical_shape_py[0].cast<size_t>();
  const auto logical_last_dim = logical_shape_py[1].cast<size_t>();
  const std::vector<size_t> logical_shape = {logical_first_dim, logical_last_dim};
  const auto &quantizer = convert_quantizer(input.attr("quantizer"));

  // Extract optional tensor attributes.
  auto get_optional_tensor = [&input](const char *name) -> std::optional<at::Tensor> {
    auto attr = input.attr(name);
    if (attr.is_none()) return std::nullopt;
    return attr.cast<at::Tensor>();
  };
  auto rowwise_data = get_optional_tensor("rowwise_data");
  auto columnwise_data = get_optional_tensor("columnwise_data");
  auto rowwise_scale_inv = get_optional_tensor("scale_inv");
  auto columnwise_scale_inv = get_optional_tensor("columnwise_scale_inv");
  auto first_dims = get_optional_tensor("first_dims");
  auto last_dims = get_optional_tensor("last_dims");
  auto tensor_offsets = get_optional_tensor("tensor_offsets");

  // Early-return for empty input.
  if (logical_first_dim == 0 || logical_last_dim == 0) {
    NoneQuantizer q{py::none()};
    auto [out_cpp, out_py] =
        q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims,
                                last_dims, tensor_offsets, logical_first_dim, logical_last_dim);
    return py::reinterpret_borrow<py::object>(out_py);
  }

  // Build input GroupedTensorWrapper.
  // Data tensors are stored as flat 1D buffers; use the quantizer's dtype
  // (e.g. kFloat8E4M3) rather than the raw tensor scalar_type (uint8).
  const NVTEScalingMode scaling_mode = quantizer->get_scaling_mode();
  const bool is_block_scaling =
      (scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D);
  const bool is_nvfp4 = (scaling_mode == NVTE_NVFP4_1D_SCALING);
  const DType scale_dtype = is_block_scaling ? DType::kFloat32
                            : is_nvfp4       ? DType::kFloat8E4M3
                                             : DType::kFloat8E8M0;
  auto input_cpp = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode);
  if (rowwise_data.has_value()) {
    input_cpp.set_rowwise_data(rowwise_data->data_ptr(), quantizer->dtype,
                               std::vector<size_t>{static_cast<size_t>(rowwise_data->numel())});
    if (rowwise_scale_inv.has_value()) {
      input_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), scale_dtype,
                                      getTensorShape(*rowwise_scale_inv));
    }
  }
  if (columnwise_data.has_value()) {
    input_cpp.set_columnwise_data(
        columnwise_data->data_ptr(), quantizer->dtype,
        std::vector<size_t>{static_cast<size_t>(columnwise_data->numel())});
    if (columnwise_scale_inv.has_value()) {
      input_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), scale_dtype,
                                         getTensorShape(*columnwise_scale_inv));
    }
  }
  if (first_dims.has_value()) {
    input_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims));
  }
  if (last_dims.has_value()) {
    input_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims));
  }
  if (tensor_offsets.has_value()) {
    input_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64,
                                 getTensorShape(*tensor_offsets));
  }

  // Create output GroupedTensor using NoneQuantizer.
  NoneQuantizer q{py::none()};
  auto [out_cpp, out_py] =
      q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, last_dims,
                              tensor_offsets, logical_first_dim, logical_last_dim);

  NVTE_SCOPED_GIL_RELEASE({
    nvte_group_dequantize(input_cpp.data(), out_cpp.data(), at::cuda::getCurrentCUDAStream());
  });

  return py::reinterpret_borrow<py::object>(out_py);
}

py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
                                    const size_t num_tensors, std::optional<at::Tensor> first_dims,
                                    DType otype, std::optional<at::Tensor> tensor_offsets,
                                    bool return_dequantized) {
  init_extension();

  const bool has_rowwise =
      !grouped_x.attr("rowwise_data").is_none() && !grouped_x.attr("scale_inv").is_none();
  const bool has_columnwise = !grouped_x.attr("columnwise_data").is_none() &&
                              !grouped_x.attr("columnwise_scale_inv").is_none();
  const bool swizzled = grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>();

  NVTE_CHECK(has_rowwise, "Grouped input has no rowwise data and scales for the GEMM to consume.");

  // The tensor's own quantization must match what the op expects on every path: even a
  // pass-through hands its data straight to the GEMM. This keeps the input's format rather than
  // converting between formats.
  const auto input_quantizer = grouped_x.attr("quantizer");
  NVTE_CHECK(!input_quantizer.is_none(), "Grouped input has no quantizer.");
  NVTE_CHECK(Py_TYPE(input_quantizer.ptr()) == Py_TYPE(quantizer.ptr()),
             "Grouped input and the op disagree on quantization format.");
  NVTE_CHECK(input_quantizer.attr("dtype").cast<DType>() == quantizer.attr("dtype").cast<DType>(),
             "Grouped input and the quantizer disagree on the FP8 dtype.");

  // The columnwise copy is only worth building when a wgrad GEMM will consume it. Read this
  // before the usage is overridden below.
  const bool need_columnwise = quantizer.attr("columnwise_usage").cast<bool>();

  if (swizzled) {
    // Already GEMM-ready. Nothing can be derived from here, since dequantization requires scales
    // in compact format, so the input must already carry everything that will be consumed. No
    // quantization kernel runs, which makes this path format-agnostic.
    NVTE_CHECK(has_columnwise || !need_columnwise,
               "Grouped input has swizzled scales but no columnwise data for the wgrad GEMM. It "
               "cannot be rebuilt, because dequantization requires scales in compact format.");
    NVTE_CHECK(!return_dequantized,
               "Cannot return a dequantized tensor for an already-swizzled grouped input: "
               "dequantization requires scales in compact format.");
    return py::none();
  }

  NVTE_CHECK(!has_columnwise,
             "Grouped input already has columnwise data but unswizzled scales; requantizing "
             "from it is not supported.");

  // Everything below runs quantization kernels, so it is MXFP8-only.
  NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()),
             "Requantizing a grouped input is only supported for MXFP8.");

  const auto logical_shape = grouped_x.attr("logical_shape").cast<py::tuple>();
  const auto total_tokens = logical_shape[0].cast<size_t>();
  const auto hidden_dim = logical_shape[1].cast<size_t>();
  // Each group's token count must be a multiple of 128 too, so that every group's scales start
  // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA
  // graph capture), so that half is the caller's contract rather than an assertion.
  NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0,
             "Requantizing a grouped input requires dims that are multiples of 128, but got (",
             total_tokens, ", ", hidden_dim, ").");

  // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left
  // undefined when nothing consumes it, which skips the pass entirely.
  at::Tensor dequantized;
  if (need_columnwise || return_dequantized) {
    dequantized = group_dequantize(grouped_x, otype)
                      .attr("rowwise_data")
                      .cast<at::Tensor>()
                      .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
  }

  // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle
  // resets columnwise_scale_inv to None, which would strand the columnwise data below with a
  // null scale pointer.
  grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
  // The swizzle hands back a 2D [num_tensors * padded_m, padded_k] scale buffer, but grouped
  // tensors carry scales as a flat array indexed by element offsets (scale_inv_offsets), so
  // per-group slicing breaks unless it is flattened back.
  grouped_x.attr("scale_inv") = grouped_x.attr("scale_inv").attr("reshape")(-1);

  if (need_columnwise) {
    // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
    // data because the two directions scale along perpendicular axes. Quantizing rowwise as well
    // would redo work we already have, so that direction is switched off; the caller sets
    // optimize_for_gemm, which makes the kernel emit swizzled columnwise scales directly.
    quantizer.attr("set_usage")(py::arg("rowwise") = false, py::arg("columnwise") = true);
    auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt,
                                     tensor_offsets, std::nullopt);
    grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
    grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
  }

  if (return_dequantized) {
    return py::cast(dequantized);
  }
  return py::none();
}

namespace {

void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,
                                std::vector<py::handle> &quantizer_py_list,
                                std::vector<std::unique_ptr<Quantizer>> &quantizer_cpp_list,
                                std::vector<TensorWrapper> &output_list) {
  // Check number of tensors
  const size_t num_tensors = input_list.size();
  NVTE_CHECK(quantizer_py_list.size() == num_tensors, "Expected ", num_tensors,
             " Python quantizers, but got ", quantizer_py_list.size());
  NVTE_CHECK(quantizer_cpp_list.size() == num_tensors, "Expected ", num_tensors,
             " C++ quantizers, but got ", quantizer_cpp_list.size());
  NVTE_CHECK(output_list.size() == num_tensors, "Expected ", num_tensors,
             " output tensors, but got ", output_list.size());

  // Choose implementation
  // Note: Currently only have fused kernel for FP8 delayed scaling
  bool with_fused_kernel = true;
  for (size_t i = 0; i < num_tensors; i++) {
    if (!detail::IsFloat8Quantizers(quantizer_py_list[i].ptr())) {
      with_fused_kernel = false;
      break;
    }
    if (nvte_tensor_data(output_list[i].data()) == nullptr ||
        nvte_tensor_columnwise_data(output_list[i].data()) == nullptr) {
      with_fused_kernel = false;
      break;
    }
  }

  // Launch TE kernel
  if (with_fused_kernel) {
    // Fused kernel for multi-tensor quantize
    std::vector<NVTETensor> nvte_tensor_input_list;
    std::vector<NVTETensor> nvte_tensor_output_list;
    for (size_t i = 0; i < num_tensors; ++i) {
      nvte_tensor_input_list.push_back(input_list[i].data());
      nvte_tensor_output_list.push_back(output_list[i].data());
    }
    NVTE_SCOPED_GIL_RELEASE({
      nvte_multi_cast_transpose(nvte_tensor_input_list.size(), nvte_tensor_input_list.data(),
                                nvte_tensor_output_list.data(), at::cuda::getCurrentCUDAStream());
    });
  } else {
    // Quantize kernels individually
    for (size_t i = 0; i < num_tensors; ++i) {
      quantizer_cpp_list[i]->quantize(input_list[i], output_list[i]);
    }
  }
}

}  // namespace

std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
                                              std::vector<py::handle> quantizer_list) {
  // Check number of tensors
  const size_t num_tensors = tensor_list.size();
  NVTE_CHECK(quantizer_list.size() == num_tensors, "Expected ", num_tensors,
             " quantizers, but got ", quantizer_list.size());

  // Convert quantizers to C++ objects
  std::vector<std::unique_ptr<Quantizer>> quantizer_cpp_list;
  for (size_t i = 0; i < num_tensors; i++) {
    quantizer_cpp_list.push_back(convert_quantizer(quantizer_list[i]));
  }

  // Initialize input and output tensors
  std::vector<TensorWrapper> input_cpp_list;
  std::vector<TensorWrapper> output_cpp_list;
  std::vector<py::object> output_py_list;
  for (size_t i = 0; i < num_tensors; ++i) {
    // Convert input tensor to C++ object
    const auto &input_py = tensor_list[i];
    NVTE_CHECK(input_py.is_contiguous(), "Input tensor ", i, " is not contiguous");
    input_cpp_list.emplace_back(makeTransformerEngineTensor(input_py));
    const auto &input_cpp = input_cpp_list.back();
    const auto input_shape = input_cpp.shape();
    const auto input_dtype = GetTransformerEngineDType(input_py.scalar_type());

    // Construct output tensor
    std::vector<size_t> output_shape(input_shape.data, input_shape.data + input_shape.ndim);
    auto [output_cpp, output_py] = quantizer_cpp_list[i]->create_tensor(output_shape, input_dtype);
    output_cpp_list.emplace_back(std::move(output_cpp));
    output_py_list.emplace_back(std::move(output_py));
  }

  // Perform multi-tensor quantization
  multi_tensor_quantize_impl(input_cpp_list, quantizer_list, quantizer_cpp_list, output_cpp_list);

  return output_py_list;
}

namespace {

std::tuple<std::vector<py::object>, std::vector<TensorWrapper>> bulk_allocate_fp8_blockwise_tensors(
    std::vector<std::vector<size_t>> &shape_list, std::vector<py::handle> &quantizer_py_list,
    std::vector<Float8BlockQuantizer *> &quantizer_cpp_list) {
  init_extension();
  std::tuple<std::vector<py::object>, std::vector<TensorWrapper>> retval;
  auto &tensor_py_list = std::get<0>(retval);
  auto &tensor_cpp_list = std::get<1>(retval);

  // Number of tensors
  const size_t num_tensors = shape_list.size();
  if (num_tensors == 0) {
    return retval;
  }

  // Quantization parameters
  const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage;
  const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage;
  const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode();
  const auto is_2D_scaled = scaling_mode == NVTE_BLOCK_SCALING_2D;
  const auto fp8_dtype = quantizer_cpp_list[0]->dtype;

  // Allocate row-wise data
  std::vector<at::Tensor> rowwise_data_list, rowwise_scale_list;
  std::vector<std::vector<size_t>> rowwise_data_shapes, rowwise_scale_shapes;
  if (rowwise_usage) {
    for (size_t i = 0; i < num_tensors; ++i) {
      rowwise_data_shapes.emplace_back(shape_list[i]);
      rowwise_scale_shapes.emplace_back(
          quantizer_cpp_list[i]->get_scale_shape(shape_list[i], false));
    }

    // Bulk-allocate data and scale tensors
    std::vector<std::vector<size_t>> shapes = rowwise_data_shapes;
    std::vector<at::ScalarType> dtypes(num_tensors, torch::kUInt8);
    std::vector<size_t> alignments(num_tensors, 256);
    shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end());
    dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32);
    alignments.insert(alignments.end(), num_tensors, 16);
    auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments);

    // Split data and scale tensors
    for (size_t i = 0; i < num_tensors; ++i) {
      rowwise_data_list.emplace_back(std::move(tensors[i]));
      rowwise_scale_list.emplace_back(std::move(tensors[num_tensors + i]));
    }
  }

  // Allocate column-wise data
  std::vector<at::Tensor> columnwise_data_list, columnwise_scale_list;
  std::vector<std::vector<size_t>> columnwise_data_shapes, columnwise_scale_shapes;
  if (columnwise_usage) {
    for (size_t i = 0; i < num_tensors; ++i) {
      columnwise_data_shapes.emplace_back();
      auto &shape = columnwise_data_shapes.back();
      shape.push_back(shape_list[i].back());
      for (size_t j = 0; j < shape_list[i].size() - 1; ++j) {
        shape.push_back(shape_list[i][j]);
      }
      columnwise_scale_shapes.emplace_back(
          quantizer_cpp_list[i]->get_scale_shape(shape_list[i], true));
    }

    // Bulk-allocate data and scale tensors
    std::vector<std::vector<size_t>> shapes = columnwise_data_shapes;
    std::vector<at::ScalarType> dtypes(num_tensors, torch::kUInt8);
    std::vector<size_t> alignments(num_tensors, 256);
    shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end());
    dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32);
    alignments.insert(alignments.end(), num_tensors, 16);
    auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments);

    // Split data and scale tensors
    for (size_t i = 0; i < num_tensors; ++i) {
      columnwise_data_list.push_back(tensors[i]);
      columnwise_scale_list.push_back(tensors[num_tensors + i]);
    }
  }

  // Construct FP8 block-wise tensors
  py::handle Float8BlockwiseQTensorClass(
      reinterpret_cast<PyObject *>(Float8BlockwiseQTensorStoragePythonClass));
  for (size_t i = 0; i < num_tensors; ++i) {
    // Create tensor objects with proper reference counting
    py::object rowwise_data = rowwise_usage ? py::cast(rowwise_data_list[i]) : py::none();
    py::object rowwise_scale = rowwise_usage ? py::cast(rowwise_scale_list[i]) : py::none();
    py::object columnwise_data =
        (columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none());
    py::object columnwise_scale =
        (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none());

    // Construct Python tensor.
    tensor_py_list.emplace_back(Float8BlockwiseQTensorClass(
        rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, MakePythonDType(fp8_dtype),
        quantizer_py_list[i], is_2D_scaled));

    // Construct C++ tensor
    tensor_cpp_list.emplace_back(makeTransformerEngineTensor(
        rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr,
        columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr,
        rowwise_usage ? rowwise_data_shapes[i] : std::vector<size_t>{0},
        columnwise_usage ? columnwise_data_shapes[i] : std::vector<size_t>{0}, fp8_dtype, nullptr,
        nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr,
        columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr,
        rowwise_usage ? rowwise_scale_shapes[i] : std::vector<size_t>{0},
        columnwise_usage ? columnwise_scale_shapes[i] : std::vector<size_t>{0}, scaling_mode));
  }

  return retval;
}

std::tuple<std::vector<py::object>, std::vector<TensorWrapper>> bulk_allocate_mxfp8_tensors(
    std::vector<std::vector<size_t>> &shape_list, std::vector<py::handle> &quantizer_py_list,
    std::vector<MXFP8Quantizer *> &quantizer_cpp_list) {
  init_extension();
  std::tuple<std::vector<py::object>, std::vector<TensorWrapper>> retval;
  auto &tensor_py_list = std::get<0>(retval);
  auto &tensor_cpp_list = std::get<1>(retval);

  // Number of tensors
  const size_t num_tensors = shape_list.size();
  if (num_tensors == 0) {
    return retval;
  }

  // Quantization parameters
  const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage;
  const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage;
  const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode();
  const auto fp8_dtype = quantizer_cpp_list[0]->dtype;
  const bool with_gemm_swizzled_scales = quantizer_cpp_list[0]->optimize_for_gemm;

  // Allocate row-wise data
  std::vector<at::Tensor> rowwise_data_list, rowwise_scale_list;
  std::vector<std::vector<size_t>> rowwise_data_shapes, rowwise_scale_shapes;
  if (rowwise_usage) {
    for (size_t i = 0; i < num_tensors; ++i) {
      rowwise_data_shapes.emplace_back(shape_list[i]);
      rowwise_scale_shapes.emplace_back(
          quantizer_cpp_list[i]->get_scale_shape(shape_list[i], false));
    }

    // Bulk-allocate data and scale tensors
    std::vector<std::vector<size_t>> shapes = rowwise_data_shapes;
    std::vector<at::ScalarType> dtypes(num_tensors, torch::kUInt8);
    std::vector<size_t> alignments(num_tensors, 256);
    shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end());
    dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8);
    alignments.insert(alignments.end(), num_tensors, 16);
    auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments);

    // Split data and scale tensors
    for (size_t i = 0; i < num_tensors; ++i) {
      rowwise_data_list.emplace_back(std::move(tensors[i]));
      rowwise_scale_list.emplace_back(std::move(tensors[num_tensors + i]));
    }
  }

  // Allocate column-wise data
  std::vector<at::Tensor> columnwise_data_list, columnwise_scale_list;
  std::vector<std::vector<size_t>> columnwise_data_shapes, columnwise_scale_shapes;
  if (columnwise_usage) {
    for (size_t i = 0; i < num_tensors; ++i) {
      // For MXFP8, the columnwise data doesn't need transpose
      // because of TN, NT, NN layout support in SM100
      columnwise_data_shapes.emplace_back(shape_list[i]);
      columnwise_scale_shapes.emplace_back(
          quantizer_cpp_list[i]->get_scale_shape(shape_list[i], true));
    }

    // Bulk-allocate data and scale tensors
    std::vector<std::vector<size_t>> shapes = columnwise_data_shapes;
    std::vector<at::ScalarType> dtypes(num_tensors, torch::kUInt8);
    std::vector<size_t> alignments(num_tensors, 256);
    shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end());
    dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8);
    alignments.insert(alignments.end(), num_tensors, 16);
    auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments);

    // Split data and scale tensors
    for (size_t i = 0; i < num_tensors; ++i) {
      columnwise_data_list.push_back(tensors[i]);
      columnwise_scale_list.push_back(tensors[num_tensors + i]);
    }
  }

  // Construct mxfp8 tensors
  py::handle MXFP8TensorClass(reinterpret_cast<PyObject *>(MXFP8TensorStoragePythonClass));
  for (size_t i = 0; i < num_tensors; ++i) {
    // Create tensor objects with proper reference counting
    py::object rowwise_data = rowwise_usage ? py::cast(rowwise_data_list[i]) : py::none();
    py::object rowwise_scale = rowwise_usage ? py::cast(rowwise_scale_list[i]) : py::none();
    py::object columnwise_data =
        (columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none());
    py::object columnwise_scale =
        (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none());

    // Construct Python tensor.
    tensor_py_list.emplace_back(MXFP8TensorClass(rowwise_data, rowwise_scale, columnwise_data,
                                                 columnwise_scale, MakePythonDType(fp8_dtype),
                                                 quantizer_py_list[i], with_gemm_swizzled_scales));

    // Construct C++ tensor
    tensor_cpp_list.emplace_back(makeTransformerEngineTensor(
        rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr,
        columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr,
        rowwise_usage ? rowwise_data_shapes[i] : std::vector<size_t>{0},
        columnwise_usage ? columnwise_data_shapes[i] : std::vector<size_t>{0}, fp8_dtype, nullptr,
        nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr,
        columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr,
        rowwise_usage ? rowwise_scale_shapes[i] : std::vector<size_t>{0},
        columnwise_usage ? columnwise_scale_shapes[i] : std::vector<size_t>{0}, scaling_mode));
    tensor_cpp_list.back().set_with_gemm_swizzled_scales(with_gemm_swizzled_scales);
  }

  return retval;
}

// allocate fp4 data, fp8 scalings, and amax values
// layout: [fp4_data0, ..., fp4_dataN, fp8_scaling0, ..., fp8_scalingN, amax0, ..., amaxN]
// amax buffer will be zeroed out by later amax kernels, so we can use empty to allocate
std::tuple<std::vector<py::object>, std::vector<TensorWrapper>, bool> bulk_allocate_nvfp4_tensors(
    std::vector<std::vector<size_t>> &shape_list, std::vector<py::handle> &quantizer_py_list,
    std::vector<NVFP4Quantizer *> &quantizer_cpp_list) {
  init_extension();
  std::tuple<std::vector<py::object>, std::vector<TensorWrapper>, bool> retval;
  auto &tensor_py_list = std::get<0>(retval);
  auto &tensor_cpp_list = std::get<1>(retval);
  auto &contiguous_data_and_scale = std::get<2>(retval);
  contiguous_data_and_scale = true;

  // Number of tensors
  const size_t num_tensors = shape_list.size();
  if (num_tensors == 0) {
    return retval;
  }

  const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage;
  const bool row_scaled_nvfp4 = quantizer_cpp_list[0]->row_scaled_nvfp4;
  const bool nvfp4_use_4over6 =
      quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;
  const int nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max;
  const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage;
  if (row_scaled_nvfp4) {
    NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage.");
    NVTE_CHECK(!columnwise_usage,
               "Row-scaled NVFP4 bulk allocation does not support columnwise usage.");
  }
  const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode();
  const auto fp4_dtype = quantizer_cpp_list[0]->dtype;

  // with_gemm_swizzled_scales is a single group-wide boolean baked
  // into every output tensor. We can safely request it only when
  // (a) every quantizer in the group has optimize_for_gemm and
  // with_rht set, and (b) every tensor's shape qualifies for RHT
  // cast-fusion. Disagreement among quantizers would silently give
  // some outputs a layout that their own quantizer did not request;
  // the NVTE_CHECK loop below turns that into a loud error. The
  // final flag ANDs all three predicates (the shape-eligibility one
  // is computed by the loop further below).
  const bool group_optimize_for_gemm = quantizer_cpp_list[0]->optimize_for_gemm;
  const bool group_with_rht = quantizer_cpp_list[0]->with_rht;
  for (size_t i = 1; i < num_tensors; ++i) {
    NVTE_CHECK(quantizer_cpp_list[i]->optimize_for_gemm == group_optimize_for_gemm,
               "NVFP4 bulk allocation requires all quantizers in the group to share "
               "the same optimize_for_gemm value (tensor 0=",
               group_optimize_for_gemm, ", tensor ", i, "=",
               quantizer_cpp_list[i]->optimize_for_gemm, ").");
    NVTE_CHECK(quantizer_cpp_list[i]->with_rht == group_with_rht,
               "NVFP4 bulk allocation requires all quantizers in the group to share "
               "the same with_rht value (tensor 0=",
               group_with_rht, ", tensor ", i, "=", quantizer_cpp_list[i]->with_rht, ").");
  }
  bool all_tensors_rht_cast_fusion_eligible = true;
  for (size_t i = 0; i < num_tensors; ++i) {
    if (!NVFP4Quantizer::is_eligible_for_rht_cast_fusion(shape_list[i],
                                                         /*for_grouped_kernel=*/true)) {
      all_tensors_rht_cast_fusion_eligible = false;
      break;
    }
  }
  const bool with_gemm_swizzled_scales =
      group_optimize_for_gemm && group_with_rht && all_tensors_rht_cast_fusion_eligible;

  // Helper function to get size of byte buffer holding FP4 data (last dim divided by 2)
  auto fp4_byte_shape = [](const std::vector<size_t> &shape) -> std::vector<size_t> {
    NVTE_CHECK(!shape.empty());
    NVTE_CHECK(shape.back() % 2 == 0);
    std::vector<size_t> out(shape.begin(), shape.end());
    out.back() /= 2;
    return out;
  };

  // Helper function to get size of amax buffer
  auto amax_shape = [](const std::vector<size_t> &shape,
                       bool row_scaled = false) -> std::vector<size_t> {
    if (row_scaled) {
      const auto [rows, _] = get_2d_dims(shape);
      return {rows};
    }
    return {1};
  };

  // Allocate row-wise data
  std::vector<at::Tensor> rowwise_data_list, rowwise_scale_list, amax_rowwise_list;
  std::vector<std::vector<size_t>> rowwise_data_shapes, rowwise_scale_shapes;
  if (rowwise_usage) {
    for (size_t i = 0; i < num_tensors; ++i) {
      rowwise_data_shapes.emplace_back(shape_list[i]);
      rowwise_scale_shapes.emplace_back(
          quantizer_cpp_list[i]->get_scale_shape(shape_list[i], false));
    }

    // Check whether data and scales can be packed in contiguous
    // buffer. Amaxes are not contiguous since they are aligned to
    // 16B.
    for (size_t i = 0; i < num_tensors; ++i) {
      if (product(rowwise_data_shapes[i]) / 2 % 256 != 0) {
        contiguous_data_and_scale = false;
      }
      if (product(rowwise_scale_shapes[i]) % 16 != 0) {
        contiguous_data_and_scale = false;
      }
    }

    // Bulk-allocate tensors data, scale, and amax tensors
    std::vector<std::vector<size_t>> shapes;
    for (size_t i = 0; i < num_tensors; ++i) {
      shapes.emplace_back(fp4_byte_shape(rowwise_data_shapes[i]));
    }
    std::vector<at::ScalarType> dtypes(num_tensors, torch::kUInt8);
    std::vector<size_t> alignments(num_tensors, 256);
    shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end());
    dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8);
    alignments.insert(alignments.end(), num_tensors, 16);
    for (size_t i = 0; i < num_tensors; ++i) {
      shapes.emplace_back(amax_shape(rowwise_data_shapes[i], row_scaled_nvfp4));
    }
    dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32);
    alignments.insert(alignments.end(), num_tensors, 16);
    auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments);

    // Split data, scale, and amax tensors
    for (size_t i = 0; i < num_tensors; ++i) {
      rowwise_data_list.push_back(tensors[i]);
      rowwise_scale_list.push_back(tensors[num_tensors + i]);
      amax_rowwise_list.push_back(tensors[2 * num_tensors + i]);
    }
  }

  // Allocate column-wise data
  std::vector<at::Tensor> columnwise_data_list, columnwise_scale_list, amax_columnwise_list;
  std::vector<std::vector<size_t>> columnwise_data_shapes, columnwise_scale_shapes;
  if (columnwise_usage) {
    for (size_t i = 0; i < num_tensors; ++i) {
      // push the transposed shape into NVFP4 columnwise shape
      // NVFP4 on SM100 is TN only
      columnwise_data_shapes.emplace_back();
      auto &shape = columnwise_data_shapes.back();
      shape.push_back(shape_list[i].back());
      for (size_t j = 0; j < shape_list[i].size() - 1; ++j) {
        shape.push_back(shape_list[i][j]);
      }
      columnwise_scale_shapes.emplace_back(
          quantizer_cpp_list[i]->get_scale_shape(shape_list[i], true));
    }

    // Check whether data and scales can be packed in contiguous
    // buffer. Amaxes are not contiguous since they are aligned to
    // 16B.
    for (size_t i = 0; i < num_tensors; ++i) {
      if (product(columnwise_data_shapes[i]) / 2 % 256 != 0) {
        contiguous_data_and_scale = false;
      }
      if (product(columnwise_scale_shapes[i]) % 16 != 0) {
        contiguous_data_and_scale = false;
      }
    }

    // Bulk-allocate tensors data, scale, and amax tensors
    std::vector<std::vector<size_t>> shapes;
    for (size_t i = 0; i < num_tensors; ++i) {
      shapes.emplace_back(fp4_byte_shape(columnwise_data_shapes[i]));
    }
    std::vector<at::ScalarType> dtypes(num_tensors, torch::kUInt8);
    std::vector<size_t> alignments(num_tensors, 256);
    shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end());
    dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8);
    alignments.insert(alignments.end(), num_tensors, 16);
    for (size_t i = 0; i < num_tensors; ++i) {
      shapes.emplace_back(amax_shape(columnwise_data_shapes[i]));
    }
    dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32);
    alignments.insert(alignments.end(), num_tensors, 16);
    auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments);

    // Split data, scale, and amax tensors
    for (size_t i = 0; i < num_tensors; ++i) {
      columnwise_data_list.push_back(tensors[i]);
      columnwise_scale_list.push_back(tensors[num_tensors + i]);
      amax_columnwise_list.push_back(tensors[2 * num_tensors + i]);
    }
  }

  // Construct nvfp4 tensors
  py::handle NVFP4TensorClass(reinterpret_cast<PyObject *>(NVFP4TensorStoragePythonClass));
  for (size_t i = 0; i < num_tensors; ++i) {
    // Create tensor objects with proper reference counting
    py::object rowwise_data = rowwise_usage ? py::cast(rowwise_data_list[i]) : py::none();
    py::object rowwise_scale = rowwise_usage ? py::cast(rowwise_scale_list[i]) : py::none();
    py::object columnwise_data =
        (columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none());
    py::object columnwise_scale =
        (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none());
    py::object amax_rowwise = rowwise_usage ? py::cast(amax_rowwise_list[i]) : py::none();
    py::object amax_columnwise = columnwise_usage ? py::cast(amax_columnwise_list[i]) : py::none();

    // Construct Python tensor.
    tensor_py_list.emplace_back(NVFP4TensorClass(
        rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, amax_rowwise,
        amax_columnwise, MakePythonDType(fp4_dtype), quantizer_py_list[i],
        with_gemm_swizzled_scales, py::arg("row_scaled_nvfp4") = row_scaled_nvfp4,
        py::arg("nvfp4_use_4over6") = nvfp4_use_4over6,
        py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max));

    // Construct C++ tensor
    // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor,
    // then set the amax and amax_columnwise values.
    {
      auto tensor_wrapper = makeTransformerEngineTensor(
          rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr,
          columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr,
          rowwise_usage ? rowwise_data_shapes[i] : std::vector<size_t>{0},
          columnwise_usage ? columnwise_data_shapes[i] : std::vector<size_t>{0}, fp4_dtype,
          /*amax_ptr=*/nullptr,
          /*scale_ptr=*/nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr,
          columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr,
          rowwise_usage ? rowwise_scale_shapes[i] : std::vector<size_t>{0},
          columnwise_usage ? columnwise_scale_shapes[i] : std::vector<size_t>{0}, scaling_mode);
      tensor_wrapper.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales);
      tensor_wrapper.set_row_scaled_nvfp4(row_scaled_nvfp4);
      tensor_wrapper.set_nvfp4_e4m3_max(nvfp4_e4m3_max);

      // Set the amax rowwise and amax columnwise if available
      if (rowwise_usage) {
        tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32,
                                getTensorShape(amax_rowwise_list[i]));
      }
      if (columnwise_usage) {
        tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32,
                                           std::vector<size_t>{1});
      }

      tensor_cpp_list.emplace_back(std::move(tensor_wrapper));
    }
  }

  return retval;
}

// Owns all allocations/wrappers backing quant_config_list[*].set_rng_state(...).
struct StochasticRngStateResources {
  at::Tensor rng_states_tensor;          // [2 * num_tensors], int64, CUDA
  at::Tensor rng_states_tensor_colwise;  // optional, same shape/dtype/device
  std::vector<TensorWrapper> te_rng_state_list;
  std::vector<TensorWrapper> te_rng_state_list_colwise;

  bool enabled{false};
  bool need_separate_rng_states{false};
  bool with_bulk_generate_rng_states{false};
};

// Populates quant_config_list (+ optional colwise list) with rng_state pointers and stochastic flag.
static StochasticRngStateResources setup_stochastic_rounding_rng_states_helper(
    size_t num_tensors, bool stochastic_rounding, bool with_bulk_generate_rng_states,
    bool need_separate_rng_states,
    std::vector<QuantizationConfigWrapper> &quant_config_list_rowwise,
    std::vector<QuantizationConfigWrapper> &quant_config_list_colwise) {
  // the return object will be used to keep rng states alive
  StochasticRngStateResources res;
  res.enabled = stochastic_rounding;
  res.need_separate_rng_states = need_separate_rng_states;
  res.with_bulk_generate_rng_states = with_bulk_generate_rng_states;

  if (!stochastic_rounding) return res;

  // Basic sanity: caller usually pre-sizes these to num_tensors.
  TORCH_CHECK(quant_config_list_rowwise.size() == num_tensors,
              "quant_config_list_rowwise must be sized to num_tensors");
  if (need_separate_rng_states) {
    TORCH_CHECK(quant_config_list_colwise.size() == num_tensors,
                "quant_config_list_colwise must be sized to num_tensors when "
                "need_separate_rng_states=true");
  }

  const size_t rng_elts_per_thread =
      res.with_bulk_generate_rng_states ? (1024 * num_tensors) : 1024;

  auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA);
  res.rng_states_tensor = torch::empty({static_cast<int64_t>(2 * num_tensors)}, opts);
  if (need_separate_rng_states) {
    res.rng_states_tensor_colwise = torch::empty({static_cast<int64_t>(2 * num_tensors)}, opts);
  }

  res.te_rng_state_list.reserve(num_tensors);
  if (need_separate_rng_states) res.te_rng_state_list_colwise.reserve(num_tensors);

  for (size_t i = 0; i < num_tensors; ++i) {
    auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
        std::nullopt, at::cuda::detail::getDefaultCUDAGenerator());

    // Rowwise RNG state
    at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread);
    int64_t *rng_state_ptr = static_cast<int64_t *>(res.rng_states_tensor.data_ptr()) + i * 2;
    philox_unpack(philox_args, rng_state_ptr);

    res.te_rng_state_list.push_back(makeTransformerEngineTensor(
        static_cast<void *>(rng_state_ptr), std::vector<size_t>{2}, DType::kInt64));
    quant_config_list_rowwise[i].set_rng_state(res.te_rng_state_list[i].data());
    quant_config_list_rowwise[i].set_stochastic_rounding(true);

    // Colwise RNG state (only if you truly need a different sequence)
    if (need_separate_rng_states) {
      // re-initialize philox_args for colwise RNG state
      at::PhiloxCudaState philox_args_col = init_philox_state(gen, rng_elts_per_thread);
      int64_t *rng_state_ptr_colwise =
          static_cast<int64_t *>(res.rng_states_tensor_colwise.data_ptr()) + i * 2;

      philox_unpack(philox_args_col, rng_state_ptr_colwise);

      res.te_rng_state_list_colwise.push_back(makeTransformerEngineTensor(
          static_cast<void *>(rng_state_ptr_colwise), std::vector<size_t>{2}, DType::kInt64));
      quant_config_list_colwise[i].set_rng_state(res.te_rng_state_list_colwise[i].data());
      quant_config_list_colwise[i].set_stochastic_rounding(true);
    }

    // break the loop if we are using bulk generate rng states
    if (res.with_bulk_generate_rng_states) break;
  }

  return res;
}

// Implements split-quantize NVFP4 with Row/Column-wise Hadamard Transform (RHT)
void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input,
                                               const std::vector<TensorWrapper> &input_list,
                                               std::vector<TensorWrapper> &output_list,
                                               const std::vector<size_t> &split_sections,
                                               const std::vector<NVFP4Quantizer *> &quantizers,
                                               cudaStream_t stream) {
  const size_t num_tensors = split_sections.size();
  const auto &quantizer = *quantizers.front();
  const bool nvfp4_use_4over6 = quantizer.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;
  NVTE_CHECK(!nvfp4_use_4over6,
             "NVFP4 4over6 quantization is not supported with RHT split quantization.");

  std::vector<NVTETensor> nvte_tensor_input_list;
  std::vector<NVTETensor> nvte_tensor_output_list;
  for (size_t i = 0; i < num_tensors; ++i) {
    nvte_tensor_input_list.push_back(input_list[i].data());
    nvte_tensor_output_list.push_back(output_list[i].data());
  }

  // trigger the row-col fusion when the split-sections shapes are all 128 aligned for max performance
  bool all_aligned_token_dim =
      std::all_of(split_sections.begin(), split_sections.end(),
                  [](size_t split_section) { return split_section % 128 == 0; });

  // in the case when rowwise and colwise cannot be fused, we have to generate the RNG states twice
  // so that rowwise and colwise will have different random numbers
  bool need_separate_rng_states =
      (!all_aligned_token_dim) && quantizer.rowwise_usage && quantizer.columnwise_usage;

  // Objects for TE C API
  std::vector<QuantizationConfigWrapper> quant_config_list;
  std::vector<QuantizationConfigWrapper> quant_config_list_colwise;
  for (size_t i = 0; i < num_tensors; ++i) {
    quant_config_list.emplace_back(QuantizationConfigWrapper());
    quant_config_list_colwise.emplace_back(QuantizationConfigWrapper());
  }

  // this is true because we have already built grouped kernels for rowwise and colwise quantization with RHT
  bool with_bulk_generate_rng_states = true;

  // Stochastic rounding
  bool need_stochastic_rounding = quantizer.stochastic_rounding;
  auto stochastic_rng_state_resources = setup_stochastic_rounding_rng_states_helper(
      num_tensors, need_stochastic_rounding, with_bulk_generate_rng_states,
      need_separate_rng_states, quant_config_list, quant_config_list_colwise);

  for (auto &config : quant_config_list) {
    config.set_nvfp4_4over6_mode(quantizer.nvfp4_4over6_mode);
  }
  for (auto &config : quant_config_list_colwise) {
    config.set_nvfp4_4over6_mode(quantizer.nvfp4_4over6_mode);
  }

  // Enable NVFP4 kernels to use math operations that sacrifice
  // accuracy for performance. These optimizations are experimental
  // and inconsistently implemented.
  // What math is accelerated? Only the high precision math, so numerical impact is minimal
  // 1. replace 1 / x by reciprocal_approximate_ftz(x)
  // 2. when RHT cast fusion is available, fusion allows cast to be performed on FP32 data,
  //    this will essentially remove a round trip between FP32 to BF16 then FP32
  // NVFP4 4over6 candidate error math is controlled separately by
  // NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH.
  const auto use_fast_math = transformer_engine::getenv<bool>("NVTE_USE_FAST_MATH");
  if (use_fast_math && !nvfp4_use_4over6) {
    for (auto &config : quant_config_list) {
      config.set_use_fast_math(true);
    }
    for (auto &config : quant_config_list_colwise) {
      config.set_use_fast_math(true);
    }
  }

  const auto use_4over6_err_use_fast_math =
      transformer_engine::getenv<bool>("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH");
  if (use_4over6_err_use_fast_math) {
    for (auto &config : quant_config_list) {
      config.set_nvfp4_4over6_err_use_fast_math(true);
    }
    for (auto &config : quant_config_list_colwise) {
      config.set_nvfp4_4over6_err_use_fast_math(true);
    }
  }

  auto &quant_config_list_colwise_to_use =
      need_separate_rng_states ? quant_config_list_colwise : quant_config_list;

  // Compute amaxes
  if (quantizer.with_post_rht_amax) {
    // We need:
    // 1. Rowwise amax = amax for input
    // 2. Columnwise amax = amax for RHT(input.t)
    nvte_group_hadamard_transform_amax(
        input.data(), reinterpret_cast<NVTETensor *>(nvte_tensor_output_list.data()),
        split_sections.data(), num_tensors, 0, quantizer.rht_matrix_random_sign_mask_t, stream);
  } else {
    // RHT is enabled, but amax is pre-RHT amax
    NVTE_ERROR("NVFP4 split-quantize does not yet support pre-RHT amax");
  }

  // Check that RHT matrix is available
  NVTE_CHECK(quantizer.rht_matrix.defined() && quantizer.rht_matrix.numel() > 0,
             "RHT matrix is not available.");
  auto rht_matrix_nvte = makeTransformerEngineTensor(quantizer.rht_matrix);

  if (all_aligned_token_dim) {
    // allocate a tile scheduler workspace
    auto tile_scheduler_workspace_torch =
        at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32));
    auto nvte_tile_scheduler_workspace =
        makeTransformerEngineTensor(tile_scheduler_workspace_torch);
    // call the fully-fused grouped kernel for rowwise quantization & colwise RHT quantization transpose
    nvte_group_hadamard_transform_cast_fusion(
        input.data(), reinterpret_cast<NVTETensor *>(nvte_tensor_output_list.data()),
        rht_matrix_nvte.data(), split_sections.data(), num_tensors, quant_config_list[0],
        nvte_tile_scheduler_workspace.data(), stream);
  } else {
    // Separate quantization for rowwise usage and columnwise usage
    // Rowwise quantization fusion with grouped version
    if (quantizer.rowwise_usage) {
      std::vector<TensorWrapper> out_identity_list;
      std::vector<NVTETensor> nvte_tensor_out_identity_list;
      for (size_t i = 0; i < num_tensors; i++) {
        bool is_empty_split = input_list[i].numel() == 0;
        TensorWrapper out_identity(output_list[i].scaling_mode());
        auto out_identity_data = output_list[i].get_rowwise_data();
        auto out_identity_scale_inv = output_list[i].get_rowwise_scale_inv();
        auto out_identity_amax = output_list[i].get_amax();
        if (!is_empty_split) {
          out_identity.set_rowwise_data(out_identity_data.data_ptr,
                                        static_cast<DType>(out_identity_data.dtype),
                                        out_identity_data.shape);
          out_identity.set_rowwise_scale_inv(out_identity_scale_inv.data_ptr,
                                             static_cast<DType>(out_identity_scale_inv.dtype),
                                             out_identity_scale_inv.shape);
          out_identity.set_amax(out_identity_amax.data_ptr,
                                static_cast<DType>(out_identity_amax.dtype),
                                out_identity_amax.shape);
        }
        out_identity_list.emplace_back(std::move(out_identity));
        nvte_tensor_out_identity_list.push_back(out_identity_list.back().data());
      }
      nvte_group_nvfp4_quantize_with_amax(input.data(), nvte_tensor_out_identity_list.data(),
                                          split_sections.data(), num_tensors, quant_config_list[0],
                                          stream);
    }

    // Columnwise RHT quantization fusion with grouped version
    if (quantizer.columnwise_usage) {
      std::vector<TensorWrapper> out_transpose_list;
      std::vector<NVTETensor> nvte_tensor_out_transpose_list;
      for (size_t i = 0; i < num_tensors; i++) {
        bool is_empty_split = input_list[i].numel() == 0;
        auto out_columnwise_data = output_list[i].get_columnwise_data();
        auto out_columnwise_scale_inv = output_list[i].get_columnwise_scale_inv();
        auto out_columnwise_amax = output_list[i].get_columnwise_amax();

        // Create a wrapper for the columnwise output, as the rowwise output. Input is in transposed layout.
        TensorWrapper out_transpose(output_list[i].scaling_mode());
        if (!is_empty_split) {
          auto [cw_first, cw_last] = get_2d_dims(out_columnwise_data.shape, true);
          std::vector<size_t> colwise_data_shape_2d = {cw_first, cw_last};

          out_transpose.set_rowwise_data(out_columnwise_data.data_ptr,
                                         static_cast<DType>(out_columnwise_data.dtype),
                                         colwise_data_shape_2d);
          out_transpose.set_rowwise_scale_inv(out_columnwise_scale_inv.data_ptr,
                                              static_cast<DType>(out_columnwise_scale_inv.dtype),
                                              out_columnwise_scale_inv.shape);
          out_transpose.set_amax(out_columnwise_amax.data_ptr,
                                 static_cast<DType>(out_columnwise_amax.dtype),
                                 out_columnwise_amax.shape);
        }
        out_transpose_list.emplace_back(std::move(out_transpose));
        nvte_tensor_out_transpose_list.push_back(out_transpose_list.back().data());
      }
      nvte_group_hadamard_transform_cast_fusion_columnwise(
          input.data(), reinterpret_cast<NVTETensor *>(nvte_tensor_out_transpose_list.data()),
          rht_matrix_nvte.data(), split_sections.data(), num_tensors,
          quant_config_list_colwise_to_use[0], stream);
    }
  }
}

void split_quantize_nvfp4_impl_helper(const TensorWrapper &input,
                                      const std::vector<TensorWrapper> &input_list,
                                      std::vector<TensorWrapper> &output_list,
                                      const std::vector<size_t> &split_sections,
                                      const std::vector<NVFP4Quantizer *> &quantizers,
                                      cudaStream_t stream) {
  const size_t num_tensors = input_list.size();
  const auto &quantizer = *quantizers.front();
  const bool nvfp4_use_4over6 = quantizer.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;
  NVTE_CHECK(!nvfp4_use_4over6 || !quantizer.stochastic_rounding,
             "NVFP4 4over6 quantization does not support stochastic rounding.");

  std::vector<NVTETensor> nvte_tensor_input_list;
  std::vector<NVTETensor> nvte_tensor_output_list;
  for (size_t i = 0; i < num_tensors; ++i) {
    nvte_tensor_input_list.push_back(input_list[i].data());
    nvte_tensor_output_list.push_back(output_list[i].data());
  }

  // In this case without RHT, the rowwise and colwise quantization are fused
  // we don't need separate rng states for rowwise and colwise
  bool need_separate_rng_states = false;

  // Objects for TE C API
  std::vector<QuantizationConfigWrapper> quant_config_list;
  for (size_t i = 0; i < num_tensors; ++i) {
    quant_config_list.emplace_back(QuantizationConfigWrapper());
  }

  // TODO: this is only true because the non-RHT path doesn't have grouped kernels yet, which we can be optimized
  // so that we can generate all rng states at once
  bool with_bulk_generate_rng_states = false;

  bool need_stochastic_rounding = quantizer.stochastic_rounding;

  // place holder for colwise rng states, which are not needed in this case
  std::vector<QuantizationConfigWrapper> dummy_quant_config_list_colwise;

  auto stochastic_rng_state_resources = setup_stochastic_rounding_rng_states_helper(
      num_tensors, need_stochastic_rounding, with_bulk_generate_rng_states,
      need_separate_rng_states, quant_config_list,
      dummy_quant_config_list_colwise);  // colwise rng states are not needed in this case

  for (auto &config : quant_config_list) {
    config.set_nvfp4_4over6_mode(quantizer.nvfp4_4over6_mode);
  }

  // NVFP4 4over6 candidate error math is controlled separately by
  // NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH.
  const auto use_fast_math = transformer_engine::getenv<bool>("NVTE_USE_FAST_MATH");
  if (use_fast_math && !nvfp4_use_4over6) {
    for (auto &config : quant_config_list) {
      config.set_use_fast_math(true);
    }
  }

  const auto use_4over6_err_use_fast_math =
      transformer_engine::getenv<bool>("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH");
  if (use_4over6_err_use_fast_math) {
    for (auto &config : quant_config_list) {
      config.set_nvfp4_4over6_err_use_fast_math(true);
    }
  }

  // We need:
  // 1. Rowwise amax = amax for input
  // 2. Columnwise amax = amax for input too
  // Columnwise amax will be filled with a fused D2D copy from rowwise amax
  // Note that the multi compute amax API expects rowwise amax pointer to be not null
  // So we need to set the pointer accordingly to make colwise-only quantization work
  std::vector<void *> orig_amax_ptr_list;
  for (size_t i = 0; i < num_tensors; i++) {
    auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr;
    orig_amax_ptr_list.push_back(rowwise_amax_ptr);
    auto columnwise_amax_ptr = output_list[i].get_columnwise_amax().data_ptr;
    void *amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr;
    NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer");
    output_list[i].set_amax(amax_ptr, DType::kFloat32, std::vector<size_t>{1});
  }
  nvte_group_amax(input.data(), reinterpret_cast<NVTETensor *>(nvte_tensor_output_list.data()),
                  split_sections.data(), num_tensors, stream);
  for (size_t i = 0; i < num_tensors; i++) {
    output_list[i].set_amax(orig_amax_ptr_list[i], DType::kFloat32, std::vector<size_t>{1});
  }

  // Quantize tensors individually
  for (size_t i = 0; i < num_tensors; i++) {
    // skip this round if input is empty
    if (input_list[i].numel() == 0) {
      continue;
    }
    nvte_quantize_v2(input_list[i].data(), output_list[i].data(), quant_config_list[i], stream);
  }
}

void split_quantize_nvfp4_impl(const TensorWrapper &input,
                               const std::vector<TensorWrapper> &input_list,
                               std::vector<TensorWrapper> &output_list,
                               const std::vector<size_t> &split_sections,
                               const std::vector<NVFP4Quantizer *> &quantizers) {
  // Check tensor lists
  const size_t num_tensors = split_sections.size();
  NVTE_CHECK(input_list.size() == num_tensors, "Expected ", num_tensors, " input tensors, but got ",
             input_list.size(), ".");
  NVTE_CHECK(output_list.size() == num_tensors, "Expected ", num_tensors,
             " output tensors, but got ", output_list.size(), ".");
  NVTE_CHECK(quantizers.size() == num_tensors, "Expected ", num_tensors,
             " NVFP4 quantizers, but got ", quantizers.size(), ".");

  // sanity check all the quantizers have the same scaling mode
  bool all_same_scaling_mode =
      std::all_of(quantizers.begin(), quantizers.end(), [&](const NVFP4Quantizer *quantizer) {
        return quantizer->get_scaling_mode() == quantizers.front()->get_scaling_mode();
      });
  NVTE_CHECK(all_same_scaling_mode, "All quantizers must have the same scaling mode");

  // Trivial cases
  if (num_tensors == 0) {
    return;
  }
  if (input.numel() == 0) {
    for (const auto &tensor : input_list) {
      NVTE_CHECK(tensor.numel() == 0,
                 "Input tensor has zero elements but got split with non-zero elements");
    }
    return;
  }

  // Assume all quantizers have identical config
  const auto &quantizer = *quantizers.front();
  NVTE_CHECK(!quantizer.with_2d_quantization,
             "NVFP4 split-quantize does not support 2D quantization");
  NVTE_CHECK(!quantizer.with_amax_reduction,
             "NVFP4 split-quantize does not support amax reduction");
  if (quantizer.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled) {
    NVTE_CHECK(!quantizer.with_rht, "NVFP4 4over6 quantization does not support RHT.");
    NVTE_CHECK(!quantizer.stochastic_rounding,
               "NVFP4 4over6 quantization does not support stochastic rounding.");
  }

  // Check input tensor shape
  const size_t input_last_dim = input.ndim() > 0 ? input.size(input.ndim() - 1) : 1;
  NVTE_CHECK(input_last_dim % 128 == 0,
             "NVFP4 multi-quantize requires inner dim to be multiple of 128.");

  // CUDA stream
  auto stream = at::cuda::getCurrentCUDAStream();

  // The grouped Hadamard transform kernels are implemented for the SM100 family
  // only. On other architectures, where
  // NVFP4Quantizer::is_eligible_for_rht_cast_fusion is false as well, quantize
  // each split on its own instead. That takes the generic unfused RHT path.
  const int sm = transformer_engine::cuda::sm_arch();
  const bool grouped_rht_supported = sm >= 100 && sm <= 110;

  // Perform multi-tensor quantization
  NVTE_SCOPED_GIL_RELEASE({
    if (quantizer.with_rht) {  // Quantize row-wise data, RHT+quantize column-wise data
      // Check that config is supported
      NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input");
      if (grouped_rht_supported) {
        // Fuse the rowwise and colwise into one when the kernel is ready
        split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections,
                                                  quantizers, stream);
      } else {
        for (size_t i = 0; i < num_tensors; ++i) {
          if (input_list[i].numel() == 0) {
            continue;
          }
          quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt);
        }
      }
    } else {  // NVFP4 quantize
      // Fuse the rowwise and colwise into one when the kernel is ready
      split_quantize_nvfp4_impl_helper(input, input_list, output_list, split_sections, quantizers,
                                       stream);
    }
  });
}

}  // namespace

std::vector<py::object> split_quantize(const at::Tensor &tensor,
                                       const std::vector<size_t> &split_sections,
                                       std::vector<py::handle> quantizer_list,
                                       bool disable_bulk_allocation) {
  init_extension();

  // Check number of tensors
  const size_t num_splits = split_sections.size();
  NVTE_CHECK(quantizer_list.size() == num_splits, "Expected ", num_splits, " quantizers, but got ",
             quantizer_list.size());
  if (num_splits == 0) {
    return {};
  }

  // Input tensor properties
  auto input_py = tensor.contiguous();
  uint8_t *input_dptr = reinterpret_cast<uint8_t *>(input_py.data_ptr());
  auto input_dtype = GetTransformerEngineDType(input_py.scalar_type());
  std::vector<size_t> input_shape;
  size_t input_size = 1;
  for (const auto &d : input_py.sizes()) {
    input_shape.push_back(d);
    input_size *= d;
  }
  NVTE_CHECK(input_shape.size() > 0, "Input tensor has 0 dims");

  // Split input tensor along dim 0
  std::vector<TensorWrapper> input_list;
  std::vector<std::vector<size_t>> split_shapes;
  size_t dim0_offset = 0;
  const size_t dim0_stride =
      input_shape[0] == 0 ? 0 : input_py.element_size() * input_size / input_shape[0];
  for (size_t i = 0; i < num_splits; ++i) {
    NVTE_CHECK(dim0_offset + split_sections[i] <= input_shape[0],
               "Attempted to split tensor with shape=", input_shape,
               " along dim 0 with split_sections=", split_sections);
    split_shapes.push_back(input_shape);
    auto &split_shape = split_shapes.back();
    split_shape[0] = split_sections[i];
    void *split_dptr = static_cast<void *>(input_dptr + dim0_offset * dim0_stride);
    input_list.emplace_back(makeTransformerEngineTensor(split_dptr, split_shape, input_dtype));
    dim0_offset += split_sections[i];
  }

  // Convert quantizers to C++ objects
  std::vector<std::unique_ptr<Quantizer>> quantizer_cpp_list;
  for (size_t i = 0; i < num_splits; i++) {
    quantizer_cpp_list.push_back(convert_quantizer(quantizer_list[i]));
  }

  // Choose implementation for allocating and populating tensors
  enum class AllocationMethod { UNFUSED, BULK_FP8_BLOCKWISE, BULK_MXFP8, BULK_NVFP4 };
  enum class QuantizationMethod { UNFUSED, FUSED_NVFP4 };
  AllocationMethod allocation_method = AllocationMethod::UNFUSED;
  QuantizationMethod quantization_method = QuantizationMethod::UNFUSED;
  if (!disable_bulk_allocation) {
    if (std::all_of(quantizer_list.begin(), quantizer_list.end(),
                    [](const py::handle &quantizer) -> bool {
                      return detail::IsFloat8BlockwiseQuantizers(quantizer.ptr());
                    })) {
      allocation_method = AllocationMethod::BULK_FP8_BLOCKWISE;
    } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(),
                           [](const py::handle &quantizer) -> bool {
                             return detail::IsMXFP8Quantizers(quantizer.ptr());
                           })) {
      allocation_method = AllocationMethod::BULK_MXFP8;
    } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(),
                           [](const py::handle &quantizer) -> bool {
                             return detail::IsNVFP4Quantizers(quantizer.ptr());
                           })) {
      allocation_method = AllocationMethod::BULK_NVFP4;
      const bool has_row_scaled_nvfp4 =
          std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(),
                      [](const std::unique_ptr<Quantizer> &quantizer) {
                        return static_cast<NVFP4Quantizer *>(quantizer.get())->row_scaled_nvfp4;
                      });
      if (has_row_scaled_nvfp4) {
        quantization_method = QuantizationMethod::UNFUSED;
      } else {
        quantization_method = QuantizationMethod::FUSED_NVFP4;
      }
    }
  }

  // Allocate output tensors
  std::vector<TensorWrapper> output_cpp_list;
  std::vector<py::object> output_py_list;
  switch (allocation_method) {
    case AllocationMethod::BULK_FP8_BLOCKWISE: {
      // Bulk allocation for FP8 block-scaling tensors
      std::vector<Float8BlockQuantizer *> blockwise_quantizers;
      for (auto &quantizer : quantizer_cpp_list) {
        blockwise_quantizers.push_back(static_cast<Float8BlockQuantizer *>(quantizer.get()));
      }
      std::tie(output_py_list, output_cpp_list) =
          bulk_allocate_fp8_blockwise_tensors(split_shapes, quantizer_list, blockwise_quantizers);
      break;
    }
    case AllocationMethod::BULK_MXFP8: {
      // Bulk allocation for MXFP8 tensors
      std::vector<MXFP8Quantizer *> mxfp8_quantizers;
      for (auto &quantizer : quantizer_cpp_list) {
        mxfp8_quantizers.push_back(static_cast<MXFP8Quantizer *>(quantizer.get()));
      }
      std::tie(output_py_list, output_cpp_list) =
          bulk_allocate_mxfp8_tensors(split_shapes, quantizer_list, mxfp8_quantizers);
      break;
    }
    case AllocationMethod::BULK_NVFP4: {
      // Bulk allocation for NVFP4 tensors
      std::vector<NVFP4Quantizer *> nvfp4_quantizers;
      for (auto &quantizer : quantizer_cpp_list) {
        nvfp4_quantizers.push_back(static_cast<NVFP4Quantizer *>(quantizer.get()));
      }
      bool contiguous_data_and_scale = false;
      std::tie(output_py_list, output_cpp_list, contiguous_data_and_scale) =
          bulk_allocate_nvfp4_tensors(split_shapes, quantizer_list, nvfp4_quantizers);
      if (quantization_method == QuantizationMethod::FUSED_NVFP4 && !input_shape.empty() &&
          input_shape.back() % 128 != 0) {
        static std::once_flag once_unfused_nvfp4_fallback_warning;
        std::call_once(once_unfused_nvfp4_fallback_warning, []() {
          NVTE_WARN(
              "Unfused NVFP4 quantization fallback is triggered because the input tensor inner "
              "dimension is not a multiple of 128, disabling NVFP4 grouped kernel fusion. "
              "NVFP4 might bring performance regressions for this input tensor shape.");
        });
        quantization_method = QuantizationMethod::UNFUSED;
      }
      if (!contiguous_data_and_scale) {
        // Avoid fused quantize kernel if data is not contiguous
        quantization_method = QuantizationMethod::UNFUSED;
      }
      break;
    }
    default: {
      // Allocate output tensors individually
      for (size_t i = 0; i < num_splits; ++i) {
        auto [output_cpp, output_py] =
            quantizer_cpp_list[i]->create_tensor(split_shapes[i], input_dtype);
        output_cpp_list.emplace_back(std::move(output_cpp));
        output_py_list.emplace_back(std::move(output_py));
      }
    }
  }

  // Quantize into output tensors
  switch (quantization_method) {
    case QuantizationMethod::FUSED_NVFP4: {
      // Fused NVFP4 quantize kernel
      auto input_nvte = makeTransformerEngineTensor(input_dptr, input_shape, input_dtype);
      std::vector<NVFP4Quantizer *> nvfp4_quantizers;
      for (auto &quantizer : quantizer_cpp_list) {
        nvfp4_quantizers.push_back(static_cast<NVFP4Quantizer *>(quantizer.get()));
      }
      split_quantize_nvfp4_impl(input_nvte, input_list, output_cpp_list, split_sections,
                                nvfp4_quantizers);
      break;
    }
    default:
      // General multi-tensor quantization
      multi_tensor_quantize_impl(input_list, quantizer_list, quantizer_cpp_list, output_cpp_list);
  }

  return output_py_list;
}

}  // namespace pytorch
}  // namespace transformer_engine
