c++ bulk alloc worked, still draft version

Signed-off-by: zhongboz <[email protected]>

Author: zhongbozhu

transformer_engine/pytorch/csrc/common.h

@@ -194,6 +194,8 @@ class Float8BlockQuantizer : public Quantizer {
   std::pair<TensorWrapper, py::object> create_tensor(
       const std::vector<size_t>& shape, DType dtype,
       std::optional<at::Tensor> rowwise_data = std::nullopt) const override;
+
+  std::pair<size_t, size_t> get_scale_shape(const std::vector<size_t>& shape, bool columnwise) const;
 };
 
 class MXFP8Quantizer : public Quantizer {

transformer_engine/pytorch/csrc/extensions.h

@@ -107,6 +107,11 @@ namespace transformer_engine::pytorch {
  * Transpose
  **************************************************************************************************/
 
+std::vector<py::object> fused_bulk_alloc_outputs(at::Tensor inpput_view, std::vector<int> m_splits, 
+                                                std::vector<py::handle> quantizer_list);
+
+py::object simple_sanity_check(at::Tensor input, py::handle quantizer);
+
 std::vector<py::object> fused_multi_quantize(std::vector<at::Tensor> input_list,
                                              std::optional<std::vector<py::object>> output_list,
                                              std::vector<py::handle> quantizer_list,

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -18,6 +18,8 @@
 #include "../extensions.h"
 #include "common.h"
 
+#include <iostream>
+
 namespace transformer_engine::pytorch {
 
 PyTypeObject *Float8TensorPythonClass = nullptr;  /// TODO Remove
@@ -199,6 +201,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("ln_out"), py::arg("quantizer"), py::arg("otype"), py::arg("sm_margin"),
         py::arg("zero_centered_gamma"));
   m.def("rmsnorm_bwd", &rmsnorm_bwd, "Backward of RMSNorm");
+  m.def("fused_bulk_alloc_outputs", &transformer_engine::pytorch::fused_bulk_alloc_outputs, "Fused Bulk Alloc Outputs",
+        py::arg("input_view"), py::arg("m_splits"), py::arg("quantizer_list"));
+  m.def("simple_sanity_check", &transformer_engine::pytorch::simple_sanity_check, "foo",
+        py::arg("input"), py::arg("quantizer"));
   m.def("fused_multi_quantize", &transformer_engine::pytorch::fused_multi_quantize,
         "Fused Multi-tensor Cast + Transpose", py::arg("input_list"), py::arg("output_list"),
         py::arg("quantizer_list"), py::arg("otype"));

transformer_engine/pytorch/csrc/extensions/quantizer.cpp

@@ -387,6 +387,57 @@ std::pair<TensorWrapper, py::object> Float8BlockQuantizer::create_tensor(
   return {std::move(tensor), std::move(ret)};
 }
 
+std::pair<size_t, size_t> Float8BlockQuantizer::get_scale_shape(const std::vector<size_t>& shape, bool columnwise) const {
+  using namespace pybind11::literals;
+  std::vector<int64_t> torch_shape;
+  size_t numel = 1;
+  for (auto s : shape) {
+    torch_shape.emplace_back(static_cast<int64_t>(s));
+    numel *= s;
+  }
+
+  size_t k_dim = torch_shape.size() == 0 ? 1u : torch_shape.back();
+  size_t m_dim = numel / k_dim;
+  constexpr size_t kBlockLen = 128;
+
+  std::pair<size_t, size_t> scale_shape;
+
+  if (columnwise) {
+    size_t sinv0 = 0;
+    size_t sinv1 = 0;
+    if (block_scaling_dim == 2) {
+      sinv0 = (k_dim + kBlockLen - 1) / kBlockLen;
+      sinv1 = roundup((m_dim + kBlockLen - 1) / kBlockLen, 4);
+    } else if (block_scaling_dim == 1) {
+      sinv0 = (m_dim + kBlockLen - 1) / kBlockLen;
+      sinv1 = roundup(k_dim, 4);
+    } else {
+      NVTE_CHECK(false,
+                 "Unsupported block_scaling_dim in create_tensor columnwise."
+                 "Expected 1 or 2. Got ",
+                 block_scaling_dim);
+    }
+    scale_shape = {sinv0, sinv1};
+  }else {
+    size_t sinv0 = 0;
+    size_t sinv1 = 0;
+    if (block_scaling_dim == 2) {
+      sinv0 = (m_dim + kBlockLen - 1) / kBlockLen;
+      sinv1 = roundup((k_dim + kBlockLen - 1) / kBlockLen, 4);
+    } else if (block_scaling_dim == 1) {
+      sinv0 = (k_dim + kBlockLen - 1) / kBlockLen;
+      sinv1 = roundup(m_dim, 4);
+    } else {
+      NVTE_CHECK(false,
+                 "Unsupported block_scaling_dim in create_tensor rowwise."
+                 "Expected 1 or 2. Got ",
+                 block_scaling_dim);
+    }
+    scale_shape = {sinv0, sinv1};
+  }
+  return scale_shape;
+}
+
 MXFP8Quantizer::MXFP8Quantizer(const py::handle& quantizer) : Quantizer(quantizer) {
   this->dtype = quantizer.attr("dtype").cast<DType>();
 }

transformer_engine/pytorch/csrc/extensions/transpose.cpp

@@ -5,12 +5,179 @@
  ************************************************************************/
 
 #include <optional>
+#include <pybind.h>
 
 #include "extensions.h"
 #include "pybind.h"
 
+#include <iostream>
+
 namespace transformer_engine::pytorch {
 
+std::vector<py::object> fused_bulk_alloc_outputs(at::Tensor input_view, std::vector<int> m_splits, 
+                                                std::vector<py::handle> quantizer_list) {
+  init_extension();
+  using namespace pybind11::literals;  // For operator""_a
+
+  int num_splits = m_splits.size();
+
+  // convert all the quantizers
+  std::vector<std::unique_ptr<Quantizer>> quantizers;
+  for (int i = 0; i < num_splits; i++) {
+    quantizers.push_back(convert_quantizer(quantizer_list[i]));
+  }
+
+  bool rowwise_usage = quantizers[0]->rowwise_usage;
+  bool columnwise_usage = quantizers[0]->columnwise_usage;
+  size_t hidden_dim = input_view.size(1);
+
+  std::vector<py::object> output_list;
+
+  if (detail::IsFloat8BlockwiseQuantizers(quantizer_list[0].ptr())) {
+    // implement the fuse bulk alloc for blockwise quantizer
+    // downcast quantizers, resorces are owned by the unique_ptr, so use raw ptr here just to get the attributes
+    std::vector<Float8BlockQuantizer*> blockwise_quantizers;
+    for (size_t i = 0; i < quantizers.size(); i++) {
+      Quantizer* raw_ptr = quantizers[i].get();
+      Float8BlockQuantizer* blockwise_quantizer = static_cast<Float8BlockQuantizer*>(raw_ptr);
+      blockwise_quantizers.push_back(blockwise_quantizer);
+    }
+
+    bool is_2D_scaled = blockwise_quantizers[0]->get_scaling_mode() == NVTE_BLOCK_SCALING_2D;
+    transformer_engine::DType fp8_dtype = blockwise_quantizers[0]->dtype;
+
+    size_t fp8_elem_size = 1;
+    size_t scale_elem_size = 4;
+
+    std::vector<std::pair<size_t, size_t>> rowwise_data_shapes;
+    std::vector<std::pair<size_t, size_t>> rowwise_scale_shapes;
+    std::vector<size_t> rowwise_data_sizes;
+    std::vector<size_t> rowwise_scale_sizes;
+    std::vector<std::pair<size_t, size_t>> columnwise_data_shapes;
+    std::vector<std::pair<size_t, size_t>> columnwise_scale_shapes;
+    std::vector<size_t> columnwise_data_sizes;
+    std::vector<size_t> columnwise_scale_sizes;
+    for (int i = 0; i < num_splits; i++) {
+      std::pair<size_t, size_t> input_view_i_shape = std::make_pair((size_t)m_splits[i], (size_t)hidden_dim);
+      if (rowwise_usage) {
+        rowwise_data_shapes.emplace_back(input_view_i_shape);
+        rowwise_scale_shapes.emplace_back(blockwise_quantizers[i]->get_scale_shape({input_view_i_shape.first, input_view_i_shape.second}, false));
+        rowwise_data_sizes.emplace_back(input_view_i_shape.first * input_view_i_shape.second * fp8_elem_size);
+        rowwise_scale_sizes.emplace_back(rowwise_scale_shapes.back().first * rowwise_scale_shapes.back().second * scale_elem_size);
+      }
+      if (columnwise_usage) {
+        columnwise_data_shapes.emplace_back(std::make_pair(input_view_i_shape.second, input_view_i_shape.first));
+        columnwise_scale_shapes.emplace_back(blockwise_quantizers[i]->get_scale_shape({input_view_i_shape.first, input_view_i_shape.second}, true));
+        columnwise_data_sizes.emplace_back(input_view_i_shape.first * input_view_i_shape.second * fp8_elem_size);
+        columnwise_scale_sizes.emplace_back(columnwise_scale_shapes.back().first * columnwise_scale_shapes.back().second * scale_elem_size);
+      }
+    }
+
+    size_t total_size_rowwise_data = std::accumulate(rowwise_data_sizes.begin(), rowwise_data_sizes.end(), 0);
+    size_t total_size_rowwise_scale = std::accumulate(rowwise_scale_sizes.begin(), rowwise_scale_sizes.end(), 0);
+    size_t total_size_columnwise_data = std::accumulate(columnwise_data_sizes.begin(), columnwise_data_sizes.end(), 0);
+    size_t total_size_columnwise_scale = std::accumulate(columnwise_scale_sizes.begin(), columnwise_scale_sizes.end(), 0);
+
+    size_t total_size_rowwise = total_size_rowwise_data + total_size_rowwise_scale;
+    size_t total_size_columnwise = total_size_columnwise_data + total_size_columnwise_scale;
+
+    std::vector<at::Tensor> rowwise_data_list;
+    std::vector<at::Tensor> rowwise_scale_list;
+    std::vector<at::Tensor> columnwise_data_list;
+    std::vector<at::Tensor> columnwise_scale_list;
+
+    at::Tensor rowwise_full_tensor;
+    at::Tensor columnwise_full_tensor;
+
+    if (rowwise_usage) {
+      rowwise_full_tensor = at::empty({(int64_t)total_size_rowwise}, at::device(input_view.device()).dtype(torch::kUInt8));
+      // use raw pointer math + from blob, avoid torch slice to reduce cpu overhead
+      uint8_t* rowwise_data_ptr = rowwise_full_tensor.data_ptr<uint8_t>();
+      uint8_t* rowwise_scale_ptr = rowwise_full_tensor.data_ptr<uint8_t>() + total_size_rowwise_data;
+      // use from_blob to construct rowwise_data_list and rowwise_scale_list
+      for (int i = 0; i < num_splits; i++) {
+        rowwise_data_list.emplace_back(at::from_blob(rowwise_data_ptr, {static_cast<int64_t>(rowwise_data_shapes[i].first), static_cast<int64_t>(rowwise_data_shapes[i].second)}, at::device(input_view.device()).dtype(torch::kUInt8)));
+        rowwise_scale_list.emplace_back(at::from_blob(rowwise_scale_ptr, {static_cast<int64_t>(rowwise_scale_shapes[i].first), static_cast<int64_t>(rowwise_scale_shapes[i].second)}, at::device(input_view.device()).dtype(torch::kFloat32)));
+        rowwise_data_ptr += rowwise_data_sizes[i];
+        rowwise_scale_ptr += rowwise_scale_sizes[i];
+      }
+    }
+
+    if (columnwise_usage) {
+      columnwise_full_tensor = at::empty({(int64_t)total_size_columnwise}, at::device(input_view.device()).dtype(torch::kUInt8));
+      uint8_t* columnwise_data_ptr = columnwise_full_tensor.data_ptr<uint8_t>();
+      uint8_t* columnwise_scale_ptr = columnwise_full_tensor.data_ptr<uint8_t>() + total_size_columnwise_data;
+      for (int i = 0; i < num_splits; i++) {
+        columnwise_data_list.emplace_back(at::from_blob(columnwise_data_ptr, {static_cast<int64_t>(columnwise_data_shapes[i].first), static_cast<int64_t>(columnwise_data_shapes[i].second)}, at::device(input_view.device()).dtype(torch::kUInt8)));
+        columnwise_scale_list.emplace_back(at::from_blob(columnwise_scale_ptr, {static_cast<int64_t>(columnwise_scale_shapes[i].first), static_cast<int64_t>(columnwise_scale_shapes[i].second)}, at::device(input_view.device()).dtype(torch::kFloat32)));
+        columnwise_data_ptr += columnwise_data_sizes[i];
+        columnwise_scale_ptr += columnwise_scale_sizes[i];
+      }
+    }
+        
+    for (int i = 0; i < num_splits; i++) {
+
+      py::handle Float8BlockwiseQTensorClass(
+        reinterpret_cast<PyObject*>(Float8BlockwiseQTensorBasePythonClass));
+
+      // Create the tensor object with proper reference counting
+      py::object rowwise_data = rowwise_usage ? py::cast(rowwise_data_list[i]) : py::none();
+      py::object columnwise_data = columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none();
+      py::object rowwise_scale = rowwise_usage ? py::cast(rowwise_scale_list[i]) : py::none();
+      py::object columnwise_scale = columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none();
+
+      py::object ret = Float8BlockwiseQTensorClass(
+          "rowwise_data"_a = rowwise_data,
+          "columnwise_data"_a = columnwise_data, 
+          "rowwise_scale_inv"_a = rowwise_scale,
+          "columnwise_scale_inv"_a = columnwise_scale,
+          "fp8_dtype"_a = fp8_dtype,
+          "quantizer"_a = quantizer_list[i],
+          "is_2D_scaled"_a = is_2D_scaled);
+
+      output_list.emplace_back(std::move(ret));
+    }
+
+    py::handle Float8BlockwiseQTensorClass(
+        reinterpret_cast<PyObject*>(Float8BlockwiseQTensorBasePythonClass));
+
+    // put the two full tensor into a python class to maintain their life cycle
+    py::object ret = Float8BlockwiseQTensorClass(
+        "rowwise_data"_a = rowwise_full_tensor,
+        "columnwise_data"_a = columnwise_full_tensor,
+        "rowwise_scale_inv"_a = py::none(),
+        "columnwise_scale_inv"_a = py::none(),
+        "fp8_dtype"_a = transformer_engine::DType::kFloat8E4M3, "quantizer"_a = py::none(), "is_2D_scaled"_a = true);
+    
+    output_list.emplace_back(std::move(ret));
+
+  }else{
+    NVTE_ERROR("Fused bulk alloc is not supported for this quantizer type");
+  }
+
+  return output_list;
+}
+
+py::object simple_sanity_check(at::Tensor input, py::handle quantizer){
+  init_extension();
+  using namespace pybind11::literals;  // For operator""_a
+  py::handle Float8BlockwiseQTensorClass(
+    reinterpret_cast<PyObject*>(Float8BlockwiseQTensorBasePythonClass));
+
+  py::object ret = Float8BlockwiseQTensorClass(
+      "rowwise_data"_a = input,
+      "columnwise_data"_a = input,
+      "rowwise_scale_inv"_a = input,
+      "columnwise_scale_inv"_a = input,
+      "fp8_dtype"_a = transformer_engine::DType::kFloat8E4M3, "quantizer"_a = quantizer, "is_2D_scaled"_a = true);
+
+  // py::handle Float8TensorClass(reinterpret_cast<PyObject*>(Float8TensorBasePythonClass));
+  // py::object ret = Float8TensorClass("data"_a = py::none(), "fp8_scale_inv"_a = py::none(),
+  //                         "fp8_dtype"_a = transformer_engine::DType::kFloat8E4M3, "data_transpose"_a = py::none(),
+  //                         "quantizer"_a = py::none());
+  return ret;
+}
+
 std::vector<py::object> fused_multi_quantize(std::vector<at::Tensor> input_list,
                                              std::optional<std::vector<py::object>> output_list,
                                              std::vector<py::handle> quantizer_list,

transformer_engine/pytorch/module/grouped_linear.py

@@ -50,7 +50,7 @@
     restore_from_saved,
 )
 
-from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer, bulk_alloc_float8_blockwise_tensor
+from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer
 
 __all__ = ["GroupedLinear"]
 
@@ -124,12 +124,22 @@ def forward(
                 output_quantizer.set_usage(rowwise=True, columnwise=False)
 
         fprop_gemm_use_split_accumulator = _2X_ACC_FPROP
+        full_buffer_rowwise = None
+        full_buffer_columnwise = None
         if fp8:
             recipe = FP8GlobalStateManager.get_fp8_recipe()
             if hasattr(recipe, "fp8_gemm_fprop"):
                 fprop_gemm_use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator
-            # TODO(zhongbo): make bulk alloc available for all quantizers
-            output_list = bulk_alloc_float8_blockwise_tensor(inp_view, m_splits, input_quantizers) if isinstance(input_quantizers[0], Float8BlockQuantizer) else None
+
+            alloc_output = tex.fused_bulk_alloc_outputs(inp_view, m_splits, input_quantizers) if isinstance(input_quantizers[0], Float8BlockQuantizer) else None
+            # alloc_output = tex.simple_sanity_check(inp_view, input_quantizers[0]) if isinstance(input_quantizers[0], Float8BlockQuantizer) else None
+            output_list = None
+            if alloc_output is not None:
+                # last element if the full buffer, all the previous tensor are view of the full buffer
+                output_list = alloc_output[:-1]
+                full_buffer_rowwise = alloc_output[-1]._rowwise_data
+                full_buffer_columnwise = alloc_output[-1]._columnwise_data
+ 
             inputmats = tex.fused_multi_quantize(
                 inputmats_no_fp8, output_list, input_quantizers, TE_DType[activation_dtype]
             )
@@ -204,6 +214,7 @@ def forward(
                 for inputmat in inputmats:
                     if isinstance(inputmat, QuantizedTensorBase):
                         inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)
+                        full_tensor_rowwise = None
             if inp.requires_grad:
                 for weight in weights_fp8:
                     if isinstance(weight, QuantizedTensorBase):
@@ -216,6 +227,7 @@ def forward(
                 *biases,
             )
             ctx.save_for_backward(*tensors_to_save)
+            ctx.full_buffer_columnwise = full_buffer_columnwise
             ctx.tensor_objects = tensor_objects
 
             ctx.weights_requires_grad = weights[0].requires_grad
@@ -260,6 +272,8 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
             biases = saved_tensors[3 * N : 4 * N]
             main_grads = ctx.main_grads
 
+            full_buffer_columnwise = ctx.full_buffer_columnwise
+
             if ctx.cpu_offloading and ctx.fuse_wgrad_accumulation:  # TOSO
                 for i in ctx.num_gemms:
                     w = torch.nn.Parameter(weights[i], weights[i].requires_grad)
@@ -275,6 +289,9 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
             )
             grad_output = [None] * ctx.num_gemms
             grad_biases = [None] * ctx.num_gemms
+
+            full_buffer_rowwise_dy = None
+            full_buffer_columnwise_dy = None
             if ctx.fp8:
                 if ctx.use_bias:
                     # unfuse bgrad for now until cast_transpose + dgrad calculation is ready
@@ -289,7 +306,12 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
                                 grad_output_mats[i], ctx.grad_output_quantizers[i]
                             )
                 else:
-                    output_list = bulk_alloc_float8_blockwise_tensor(grad_output_view, ctx.m_splits, ctx.grad_output_quantizers) if isinstance(ctx.grad_output_quantizers[0], Float8BlockQuantizer) else None
+                    alloc_output = tex.fused_bulk_alloc_outputs(grad_output_view, ctx.m_splits, ctx.grad_output_quantizers) if isinstance(ctx.grad_output_quantizers[0], Float8BlockQuantizer) else None
+                    output_list = None
+                    if alloc_output is not None:
+                        output_list = alloc_output[:-1]
+                        full_buffer_rowwise_dy = alloc_output[-1]._rowwise_data
+                        full_buffer_columnwise_dy = alloc_output[-1]._columnwise_data
                     grad_output = tex.fused_multi_quantize(
                         grad_output_mats,
                         output_list,