[PyTorch] Custom kernel to compute reciprocal of a single float

transformer_engine/pytorch/csrc/extensions.h

@@ -7,6 +7,10 @@
 #ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_
 #define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_
 
+#include <cstdint>
+#include <optional>
+#include <vector>
+
 #include "common.h"
 #include "common/common.h"
 
@@ -352,6 +356,19 @@ void fused_amax_and_scale_update_after_reduction(const at::Tensor &amax_reductio
  const std::string &amax_compute_algo,
  transformer_engine::DType fp8_dtype, float margin);
 
+/* Reciprocal of a single float32
+ *
+ * /param[in] Input tensor.
+ * /param[out] Output tensor.
+ * /param[in] Offset within input tensor.
+ * /param[in] Offset within output tensor.
+ * /param[in] float32 flag indicating whether to avoid updating
+ * output.
+ */
+at::Tensor scalar_reciprocal(const at::Tensor &src, std::optional<at::Tensor> dst = std::nullopt,
+ int64_t src_offset = 0, int64_t dst_offset = 0,
+ const std::optional<at::Tensor> &noop_flag = std::nullopt);
+
 /***************************************************************************************************
  * Rotary positional embedding
  **************************************************************************************************/

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -6,6 +6,8 @@
 
 #include <pybind11/functional.h>
 
+#include <optional>
+
 #include "../comm_gemm_overlap.h"
 #include "../extensions.h"
 
@@ -137,6 +139,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fused_amax_and_scale_update_after_reduction", &fused_amax_and_scale_update_after_reduction,
  "Update amax history and FP8 scale/scale_inv after reduction",
  py::call_guard<py::gil_scoped_release>());
+ m.def("scalar_reciprocal", &scalar_reciprocal, "Reciprocal of a single float",
+ py::call_guard<py::gil_scoped_release>(), py::arg("src"), py::arg("dst") = std::nullopt,
+ py::arg("src_offset") = 0, py::arg("dst_offset") = 0, py::arg("noop_flag") = std::nullopt);
 
  // fused apply rope
  m.def("fused_rope_forward", &fused_rope_forward, "Fused Apply RoPE FWD",

transformer_engine/pytorch/csrc/extensions/recipe.cu

@@ -12,9 +12,9 @@
 #include "extensions.h"
 
 void fused_amax_and_scale_update_after_reduction(
- const at::Tensor &amax_reduction_buffer, std::vector<at::Tensor> amax_histories,
+ const at::Tensor& amax_reduction_buffer, std::vector<at::Tensor> amax_histories,
  std::vector<at::Tensor> scales, std::vector<at::Tensor> scale_invs,
- const std::string &amax_compute_algo, transformer_engine::DType fp8_dtype, float margin) {
+ const std::string& amax_compute_algo, transformer_engine::DType fp8_dtype, float margin) {
  using namespace transformer_engine;
  size_t num_tensors = amax_histories.size();
  std::vector<Tensor> t_amax_histories(num_tensors);
@@ -51,3 +51,41 @@ void fused_amax_and_scale_update_after_reduction(
  te_scale_invs, amax_compute_algo.c_str(), static_cast<NVTEDType>(fp8_dtype), margin,
  at::cuda::getCurrentCUDAStream());
 }
+
+namespace {
+
+__global__ void __launch_bounds__(1)
+ scalar_reciprocal_kernel(const float* __restrict__ src, float* __restrict__ dst,
+ const float* __restrict__ noop) {
+ if (noop != nullptr && *noop == 1.f) {
+ return;
+ }
+ *dst = __frcp_rn(*src);
+}
+
+} // namespace
+
+at::Tensor scalar_reciprocal(const at::Tensor& src, std::optional<at::Tensor> dst,
+ int64_t src_offset, int64_t dst_offset,
+ const std::optional<at::Tensor>& noop_flag) {
+ using namespace transformer_engine;
+
+ // Allocate output tensor if needed
+ NVTE_CHECK(dst || dst_offset == 0,
+ "Provided offset in output tensor without providing output tensor");
+ at::Tensor dst_val = dst ? dst.value() : allocateTorchTensor(1, DType::kFloat32);
+
+ // Get pointers
+ const float* src_ptr = reinterpret_cast<const float*>(getDataPtr(src, src_offset));
+ float* dst_ptr = reinterpret_cast<float*>(getDataPtr(dst_val, dst_offset));
+ const float* noop_ptr = nullptr;
+ if (noop_flag) {
+ noop_ptr = reinterpret_cast<const float*>(getDataPtr(noop_flag.value()));
+ }
+
+ // Launch kernel
+ scalar_reciprocal_kernel<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(src_ptr, dst_ptr,
+ noop_ptr);
+
+ return dst_val;
+}

transformer_engine/pytorch/float8_tensor.py

@@ -739,17 +739,12 @@ def cast_transpose_(
  transpose_out=transpose,
  noop_flag=noop_flag,
  )
- scale = fp8_meta.scale[fp8_meta_index : fp8_meta_index + 1]
- scale_inv = self._scale_inv
- if noop_flag is None:
- torch.reciprocal(scale, out=scale_inv)
- else:
- torch.where(
- noop_flag.bool(),
- scale_inv,
- scale.reciprocal(),
- out=scale_inv,
- )
+ tex.scalar_reciprocal(
+ fp8_meta.scale,
+ self._scale_inv,
+ src_offset=fp8_meta_index,
+ noop_flag=noop_flag,
+ )
  self._transpose_invalid = False
 
  @torch.no_grad()
@@ -853,7 +848,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
  fp8_meta_index = dst._fp8_meta_index
  scale = dst._fp8_meta[fp8_meta_key].scale[fp8_meta_index]
  amax = dst._fp8_meta[fp8_meta_key].amax_history[0][fp8_meta_index]
-dst._scale_inv.copy_(scale.detach().reciprocal())
+tex.scalar_reciprocal(scale, dst._scale_inv)
 
  # Cast to FP8
  if not dst._data.is_contiguous():

transformer_engine/pytorch/module/layernorm_linear.py

@@ -152,6 +152,7 @@ def forward(
  zero_centered_gamma,
  is_grad_enabled,
  )
+ ln_out_fp8_scale_inv = None
 
  # Column Parallel Linear
  ln_out_gathered = False
@@ -212,6 +213,12 @@ def forward(
 
  assert isinstance(weight_fp8, Float8Tensor)
 
+ # FP8 scale-inverse for GEMM input
+ ln_out_fp8_scale_inv = tex.scalar_reciprocal(
+ fp8_meta["scaling_fwd"].scale,
+ src_offset=int(tex.FP8FwdTensors.GEMM1_INPUT),
+ )
+
  if fp8_meta["recipe"].fp8_mha:
  out_index, meta_tensor, output_te_dtype, output_dtype = (
  tex.FP8FwdTensors.GEMM1_OUTPUT,
@@ -232,8 +239,8 @@ def forward(
  0,
  weight_fp8._fp8_dtype,
  ln_out_total,
-fp8_meta["scaling_fwd"].scale_inv,
-tex.FP8FwdTensors.GEMM1_INPUT,
+ln_out_fp8_scale_inv,
+0,
  fp8_dtype_forward,
  output_dtype,
  get_workspace(),
@@ -323,7 +330,7 @@ def forward(
  weight_fp8,
  weight.main_grad if cpu_offloading and fuse_wgrad_accumulation else None,
  ln_out if weight.requires_grad else None,
-fp8_meta["scaling_fwd"].scale_inv.clone() if fp8 else None,
+ln_out_fp8_scale_inv,
  )
 
  ctx.activation_dtype = activation_dtype
@@ -395,7 +402,7 @@ def backward(
  weight_fp8,
  main_grad,
  ln_out,
-fwd_scale_inverses,
+ln_out_fp8_scale_inv,
  ) = ctx.saved_tensors
 
  # Gather intermediate/activation tensors if needed
@@ -592,8 +599,8 @@ def backward(
  ln_out_total_t = tex.fp8_transpose(ln_out_total, fp8_dtype_forward)
  wgrad, _ = tex.fp8_gemm(
  ln_out_total_t,
-fwd_scale_inverses,
-tex.FP8FwdTensors.GEMM1_INPUT,
+ln_out_fp8_scale_inv,
+0,
  fp8_dtype_forward,
  (
  grad_output_t._data
@@ -618,8 +625,8 @@ def backward(
  else:
  ln_out_total_c = torch.ops.tex_ts.cast_from_fp8_ts(
  ln_out_total,
-fwd_scale_inverses,
-tex.FP8FwdTensors.GEMM1_INPUT,
+ln_out_fp8_scale_inv,
+0,
  fp8_dtype_forward,
  TE_DType[ctx.activation_dtype],
  )

transformer_engine/pytorch/module/linear.py

@@ -117,10 +117,12 @@ def forward(
  inputmat = cast_if_needed(inputmat, activation_dtype)
  inputmat_t = None
  inputmat_no_fp8 = inputmat
+ inputmat_fp8_scale_inv = None
 
  if fp8:
  fp8_dtype_forward = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True)
  if isinstance(inputmat, Float8Tensor):
+ inputmat_fp8_scale_inv = inputmat._scale_inv
  if (
  not fp8_meta["recipe"].override_linear_precision.wgrad
  and is_grad_enabled
@@ -152,6 +154,12 @@ def forward(
  fp8_dtype_forward,
  )
 
+ # FP8 scale-inverse
+ inputmat_fp8_scale_inv = tex.scalar_reciprocal(
+ fp8_meta["scaling_fwd"].scale,
+ src_offset=int(tex.FP8FwdTensors.GEMM1_INPUT),
+ )
+
  # Column Parallel Linear
  if parallel_mode == "column" and sequence_parallel:
  inputmat_total, _ = gather_along_first_dim(inputmat, tp_group)
@@ -222,8 +230,8 @@ def forward(
  if isinstance(inputmat_total, Float8Tensor)
  else inputmat_total
  ),
-fp8_meta["scaling_fwd"].scale_inv,
-tex.FP8FwdTensors.GEMM1_INPUT,
+inputmat_fp8_scale_inv,
+0,
  fp8_dtype_forward,
  proj_out_pttype,
  get_workspace(),
@@ -335,7 +343,7 @@ def forward(
  weight,
  weight_fp8,
  weight.main_grad if cpu_offloading and fuse_wgrad_accumulation else None,
-fp8_meta["scaling_fwd"].scale_inv.clone() if fp8 else None,
+inputmat_fp8_scale_inv,
  )
 
  ctx.activation_dtype = activation_dtype
@@ -388,7 +396,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
  weight,
  weight_fp8,
  main_grad,
-fwd_scale_inverses,
+inputmat_fp8_scale_inv,
  ) = ctx.saved_tensors
 
  # Gather intermediate/activation tensors if needed
@@ -545,8 +553,8 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
  if isinstance(inputmat_t_total, Float8Tensor)
  else inputmat_t_total
  ),
-fwd_scale_inverses,
-tex.FP8FwdTensors.GEMM1_INPUT,
+inputmat_fp8_scale_inv,
+0,
  fp8_dtype_forward,
  grad_output_t,
  ctx.fp8_meta["scaling_bwd"].scale_inv,