Merge remote-tracking branch 'origin/nvdlff-inspect-support' into pgadzinski/debugtools-cppqtensor

pggPL · pggPL · commit a89c44f28223 · 2025-01-30T03:04:44.000-08:00
diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
@@ -11,7 +11,7 @@ pytest -v -s $TE_PATH/tests/pytorch/test_sanity.py
 pytest -v -s $TE_PATH/tests/pytorch/test_recipe.py
 pytest -v -s $TE_PATH/tests/pytorch/test_deferred_init.py
 PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 pytest -v -s $TE_PATH/tests/pytorch/test_numerics.py
-PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 pytest -v -s $TE_PATH/tests/pytorch/test_cuda_graphs.py
+NVTE_CUDNN_MXFP8_NORM=0 PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 pytest -v -s $TE_PATH/tests/pytorch/test_cuda_graphs.py
 pytest -v -s $TE_PATH/tests/pytorch/test_jit.py
 pytest -v -s $TE_PATH/tests/pytorch/test_fused_rope.py
 pytest -v -s $TE_PATH/tests/pytorch/test_float8tensor.py
diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh
@@ -11,5 +11,5 @@ pytest -v -s $TE_PATH/tests/pytorch/distributed/test_numerics.py
 pytest -v -s $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py
 pytest -v -s $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py
 pytest -v -s $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py
-pytest -v -s $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
+# pytest -v -s $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py  ### TODO Debug UB support with te.Sequential
 pytest -v -s $TE_PATH/tests/pytorch/fused_attn/test_fused_attn_with_cp.py
diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h
@@ -423,10 +423,8 @@ struct is_fp8<fp8e5m2> : std::true_type {};
 size_t typeToSize(const DType type);
 
 void CheckNoopTensor(const Tensor &t, const std::string &name);
-void CheckInputTensor(const Tensor &t, const std::string &name,
-                      bool check_scale_inv_alignment = false);
-void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty = false,
-                       bool check_scale_inv_alignment = false);
+void CheckInputTensor(const Tensor &t, const std::string &name);
+void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty = false);
 
 bool is_fp8_dtype(const DType t);
 
diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu
@@ -210,8 +210,8 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s
     return;
   }
 
-  CheckInputTensor(*input, "scaling_factor_input", true);
-  CheckInputTensor(*output, "scaling_factor_output", true);
+  CheckInputTensor(*input, "scaling_factor_input");
+  CheckInputTensor(*output, "scaling_factor_output");
 
   auto& scaling_mode = input->scaling_mode;
 
diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp
@@ -65,7 +65,7 @@ void CheckNoopTensor(const Tensor &t, const std::string &name) {
   }
 }
 
-void CheckScaleTensorShape(const Tensor &t, bool check_scale_inv_alignment) {
+void CheckScaleTensorShape(const Tensor &t) {
   NVTE_CHECK(t.scaling_mode != NVTE_INVALID_SCALING, "Invalid scaling mode!");
   if (is_tensor_scaling(t.scaling_mode)) {
     // per-tensor scaling
@@ -80,7 +80,6 @@ void CheckScaleTensorShape(const Tensor &t, bool check_scale_inv_alignment) {
     }
   } else {
     if (t.scaling_mode == NVTE_MXFP8_1D_SCALING) {
-      if (!check_scale_inv_alignment) return;
       // Need (4, 128) alignment even for e8 scaling factor
       auto block_alignment = std::vector<size_t>{128ul / typeToSize(t.scale_inv.dtype),
                                                  4ul / typeToSize(t.scale_inv.dtype)};
@@ -111,7 +110,7 @@ void CheckScaleTensorShape(const Tensor &t, bool check_scale_inv_alignment) {
   }
 }
 
-void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale_inv_alignment) {
+void CheckInputTensor(const Tensor &t, const std::string &name) {
   const DType type = t.dtype();
   if (is_fp8_dtype(type)) {
     // FP8 input needs to have scale_inv
@@ -143,11 +142,10 @@ void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale
   }
   NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Input ", name, " is not allocated!");
 
-  CheckScaleTensorShape(t, check_scale_inv_alignment);
+  CheckScaleTensorShape(t);
 }
 
-void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty,
-                       bool check_scale_inv_alignment) {
+void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty) {
   const DType type = t.dtype();
   if (is_fp8_dtype(type)) {
     // FP8 output needs to have scale, scale_inv and (if delayed scaling) amax
@@ -189,7 +187,7 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt
     NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Output ", name, " is not allocated!");
   }
 
-  CheckScaleTensorShape(t, check_scale_inv_alignment);
+  CheckScaleTensorShape(t);
 }
 
 }  // namespace transformer_engine
diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py
@@ -11,25 +11,13 @@
 import transformer_engine_torch as tex
 from ..constants import TE_DType
 from ..utils import assert_dim_for_fp8_exec, get_sm_count
-from ..tensor.quantized_tensor import QuantizedTensor
-from ..tensor.float8_tensor import Float8Tensor, Float8TensorBase
-from ..tensor.mxfp8_tensor import MXFP8Tensor, MXFP8TensorBase
 
 from ..tensor.quantized_tensor import Quantizer
+from ..tensor.float8_tensor import Float8Tensor
+from ..tensor.mxfp8_tensor import MXFP8Tensor
 from ..tensor._internal.float8_tensor_base import Float8TensorBase
 from ..tensor._internal.mxfp8_tensor_base import MXFP8TensorBase
 
-
-
-from ..tensor.quantized_tensor import (
-    QuantizedTensor,
-    Quantizer,
-    prepare_for_saving,
-    restore_from_saved,
-)
-
-from ..tensor._internal.mxfp8_tensor_base import MXFP8TensorBase
-
 __all__ = [
     "general_gemm",
     "general_grouped_gemm",
diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
@@ -223,4 +223,9 @@ std::vector<size_t> convertShape(const NVTEShape& shape) {
   return std::vector<size_t>(shape.data, shape.data + shape.ndim);
 }
 
+int roundup(const int value, const int multiple) {
+  assert(multiple > 0);
+  return ((value + multiple - 1) / multiple) * multiple;
+}
+
 }  // namespace transformer_engine::pytorch
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
@@ -59,19 +59,6 @@ class FP8TensorMeta {
   at::Tensor amax_history;
 };
 
-// FP8TensorMeta for block scaling, this structure allows
-// indexing into it the same way (i.e. using FP8FwdTensors
-// and FP8BwdTensors) for both hopper and blackwell recipes.
-// TODO(ksivaman): check perf with this design; should be ok
-// since there are no amax reductions, or bulk amax/scale
-// updates for block scaling.
-class MXFP8TensorMeta {
- public:
-  std::vector<at::Tensor> scale;
-  std::vector<at::Tensor> scale_inv;
-  std::vector<at::Tensor> amax_history;
-};
-
 // Used as named indices on the `scale`, `scale_inv`,
 // and `amax` tensors in the `FP8TensorMeta` class.
 enum FP8FwdTensors {
@@ -265,6 +252,8 @@ void* getDataPtr(at::Tensor tensor, int offset = 0);
 
 std::vector<size_t> convertShape(const NVTEShape& shape);
 
+int roundup(const int value, const int multiple);
+
 }  // namespace transformer_engine::pytorch
 
 namespace std {
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
@@ -361,8 +361,6 @@ at::Tensor rowwise_swizzle(at::Tensor input, at::Tensor scale_inv);
 
 at::Tensor columnwise_swizzle(at::Tensor input, at::Tensor scale_inv);
 
-at::Tensor pad_scale_inv(at::Tensor scale_inv, bool rowwise);
-
 /***************************************************************************************************
  * Comm+GEMM Overlap Wrappers
  **************************************************************************************************/
diff --git a/transformer_engine/pytorch/csrc/extensions/quantizer.cpp b/transformer_engine/pytorch/csrc/extensions/quantizer.cpp
@@ -174,28 +174,34 @@ std::pair<TensorWrapper, py::object> MXFP8Quantizer::create_tensor(
   opts = opts.dtype(torch::kUInt8).device(torch::kCUDA);
   auto last_dim = torch_shape.back();
 
+  NVTE_CHECK(last_dim % MXFP8_BLOCK_SIZE == 0 && (numel / last_dim) % MXFP8_BLOCK_SIZE == 0,
+             "MXFP8 requires tensor dims that are divisble by ", MXFP8_BLOCK_SIZE,
+             " (got shape=", torch_shape, ")");
+
   at::Tensor data;
   if (rowwise_usage) {
     if (rowwise_data.has_value()) {
       data = std::move(*rowwise_data);
     } else {
       data = at::empty(torch_shape, opts);
     }
-    rowwise_scale_inv = at::empty({numel / last_dim, last_dim / MXFP8_BLOCK_SIZE}, opts);
+    auto sinv0 = roundup(numel / last_dim, 128);
+    auto sinv1 = roundup(last_dim / MXFP8_BLOCK_SIZE, 4);
+    rowwise_scale_inv = at::zeros({sinv0, sinv1}, opts);
     tensor.set_rowwise_data(data.data_ptr(), this->dtype, shape);
-    tensor.set_rowwise_scale_inv(
-        rowwise_scale_inv.data_ptr(), DType::kFloat8E8M0,
-        std::vector<size_t>{numel / last_dim, last_dim / MXFP8_BLOCK_SIZE});
-  } else {
+    tensor.set_rowwise_scale_inv(rowwise_scale_inv.data_ptr(), DType::kFloat8E8M0,
+                                 std::vector<size_t>{sinv0, sinv1});
   }
+
   if (columnwise_usage) {
+    auto sinv0 = roundup(numel / (last_dim * MXFP8_BLOCK_SIZE), 4);
+    auto sinv1 = roundup(last_dim, 128);
     columnwise_data = at::empty(torch_shape, opts);
-    columnwise_scale_inv = at::empty({numel / (last_dim * MXFP8_BLOCK_SIZE), last_dim}, opts);
+    columnwise_scale_inv = at::zeros({sinv0, sinv1}, opts);
 
     tensor.set_columnwise_data(columnwise_data.data_ptr(), this->dtype, shape);
-    tensor.set_columnwise_scale_inv(
-        columnwise_scale_inv.data_ptr(), DType::kFloat8E8M0,
-        std::vector<size_t>{numel / (last_dim * MXFP8_BLOCK_SIZE), last_dim});
+    tensor.set_columnwise_scale_inv(columnwise_scale_inv.data_ptr(), DType::kFloat8E8M0,
+                                    std::vector<size_t>{sinv0, sinv1});
   }
   this->set_quantization_params(&tensor);
 
diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
@@ -65,24 +65,11 @@ void swizzle_scaling_factors(transformer_engine::TensorWrapper& input, bool roww
   }
 }
 
-at::Tensor pad_scale_inv(at::Tensor scale_inv, bool rowwise) {
-  size_t dim_1_mod = (rowwise) ? 128 : 4;
-  size_t dim_2_mod = (rowwise) ? 4 : 128;
-  size_t dim_1_pad = (dim_1_mod - scale_inv.sizes()[0] % dim_1_mod) % dim_1_mod;
-  size_t dim_2_pad = (dim_2_mod - scale_inv.sizes()[1] % dim_2_mod) % dim_2_mod;
-  if (dim_1_pad == 0 && dim_2_pad == 0) {
-    return scale_inv;
-  }
-  return at::constant_pad_nd(scale_inv, {0, dim_2_pad, 0, dim_1_pad}, 0.0);
-}
-
-at::Tensor rowwise_swizzle(at::Tensor input, at::Tensor _scale_inv) {
+at::Tensor rowwise_swizzle(at::Tensor input, at::Tensor scale_inv) {
   using namespace transformer_engine::pytorch;
 
   NVTE_CHECK(input.element_size() == 1, "8-bit input required for swizzling scaling factors.");
 
-  auto scale_inv = pad_scale_inv(_scale_inv, true);
-
   auto options = at::TensorOptions().dtype(scale_inv.dtype()).device(torch::kCUDA);
   auto swizzled_scale_inv = at::empty_like(scale_inv, options);
 
@@ -102,13 +89,11 @@ at::Tensor rowwise_swizzle(at::Tensor input, at::Tensor _scale_inv) {
   return swizzled_scale_inv;
 }
 
-at::Tensor columnwise_swizzle(at::Tensor input, at::Tensor _scale_inv) {
+at::Tensor columnwise_swizzle(at::Tensor input, at::Tensor scale_inv) {
   using namespace transformer_engine::pytorch;
 
   NVTE_CHECK(input.element_size() == 1, "8-bit input required for swizzling scaling factors.");
 
-  auto scale_inv = pad_scale_inv(_scale_inv, false);
-
   auto options = at::TensorOptions().dtype(scale_inv.dtype()).device(torch::kCUDA);
   auto swizzled_scale_inv = at::empty_like(scale_inv, options);
 
diff --git a/transformer_engine/pytorch/fp8.py b/transformer_engine/pytorch/fp8.py
@@ -102,7 +102,6 @@ class FP8GlobalStateManager:
     skip_fp8_weight_update_tensor = None
     mxfp8_available = None
     reason_for_no_mxfp8 = ""
-    debug_enabled = False
 
     @classmethod
     def reset(cls) -> None:
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
@@ -4,6 +4,7 @@
 
 """Internal function used by multiple modules."""
 
+import os
 from typing import Any, List, Optional, Tuple, Union, Callable
 from dataclasses import dataclass
 
@@ -12,6 +13,10 @@
 from .. import cpp_extensions as tex
 from ..constants import TE_DType
 from ..utils import get_default_init_method
+from ..tensor.mxfp8_tensor import MXFP8Quantizer
+
+
+_use_cudnn_mxfp8_norm = bool(int(os.getenv("NVTE_CUDNN_MXFP8_NORM", "0")))
 
 
 def _get_normalization_func(normalization: str, forward: bool):
@@ -46,17 +51,25 @@ def apply_normalization(
 
     inputs = (inputmat, ln_weight) if ln_bias is None else (inputmat, ln_weight, ln_bias)
 
+    split_mxfp8_cast = False
+    if not _use_cudnn_mxfp8_norm and isinstance(output_quantizer, MXFP8Quantizer):
+        split_mxfp8_cast = True
+
     output = normalization_func(
         *inputs,
         eps,
-        ln_out,
-        output_quantizer,
+        None if split_mxfp8_cast else ln_out,
+        None if split_mxfp8_cast else output_quantizer,
         TE_DType[output_dtype] if output_dtype in TE_DType else output_dtype,
         fwd_ln_sm_margin,
         zero_centered_gamma,
     )
 
-    return output
+    return (
+        (output_quantizer.quantize(output[0], out=ln_out), *output[1:])
+        if split_mxfp8_cast
+        else output
+    )
 
 
 class _NoopCatFunc(torch.autograd.Function):
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
@@ -898,11 +898,9 @@ def grad_output_preprocess(
 
         # FP8 with all-gather: unfused bgrad, fused cast + transpose
         if gather_grad_output:
+            grad_bias = None
             if ctx.use_bias:
-                # TODO: We know it creates spike in memory usage, we should WAR that
                 grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)
-            else:
-                grad_bias = None
             if ctx.ub_overlap_ag:
                 # TODO: Implement
                 raise NotImplementedError(
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
@@ -101,15 +101,12 @@ def forward(
         inputmats = []
 
         weight_requires_grad = weights[0].requires_grad
-        backward_needs_input = is_grad_enabled and weight_requires_grad  # #TODO
 
         if input_quantizers[0] is not None:
             for input_quantizer in input_quantizers:
                 input_quantizer.set_usage(
                     rowwise=True,
-                    columnwise=(
-                        is_grad_enabled and weight_requires_grad
-                    ),  # TODO: and not sequence parallel?
+                    columnwise=(is_grad_enabled and weight_requires_grad),
                 )
             columnwise_usage = is_grad_enabled and inp.requires_grad
             if not columnwise_usage:
@@ -312,11 +309,8 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
                 # Deallocate input tensor
                 clear_tensor_data(*inputmats)
 
-                # clear_tensor_data(*weights) # TODO: 2 cases - own and do not won weight
-
-                # TODO - handle it later
-                """def handle_custom_ddp_from_mcore(w, wgrad):
-                    if w.requires_grad:
+                def handle_custom_ddp_from_mcore(w, wgrad):
+                    if ctx.weights_requires_grad:
                         if ctx.fuse_wgrad_accumulation and hasattr(w, "grad_added_to_main_grad"):
                             w.grad_added_to_main_grad = True
                             if getattr(w, "zero_out_wgrad", False):
@@ -341,7 +335,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
 
                 wgrad_list = [
                     handle_custom_ddp_from_mcore(w, wgrad) for w, wgrad in zip(weights, wgrad_list)
-                ]"""
+                ]
             else:
                 wgrad_list = [None] * ctx.num_gemms
 
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -601,7 +601,10 @@ def backward(
         # Scatter fp8 weight buffers
         # if ctx.fp8 and not isinstance(weight, QuantizedTensor):
         #    _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8)
+<<<<<<< HEAD
 
+=======
+>>>>>>> origin/nvdlff-inspect-support
 
         return (
             dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -748,7 +748,7 @@ def backward(
                 if fc2_bias_grad is None:
                     fc2_bias_grad = fc2_bias_grad_
                 del fc2_bias_grad_
-            clear_tensor_data(act_out, grad_output)
+            clear_tensor_data(act_out)
 
             # bias computation
             fc1_bias_grad = None
@@ -896,10 +896,10 @@ def backward(
                 fc1_dgrad_work = None
 
             # Residual gradient
-            #dgrad = fc1_dgrad.view(inputmat.shape)
-            #if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered:
-            #    dgrad = dgrad + grad_outputs[1].view_as(dgrad)
-            dgrad = fc1_dgrad
+            dgrad = fc1_dgrad.view(inputmat.shape)
+            if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered:
+                dgrad = dgrad + grad_outputs[1].view_as(dgrad)
+
             # Norm gradient
             dgamma = None
             dbeta = None
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py
diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py
diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py
diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py

Original file line number	Diff line number	Diff line change
`@@ -210,8 +210,8 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s`
`210`	`210`	`return;`
`211`	`211`	`}`
`212`	`212`
`213`		`- CheckInputTensor(*input, "scaling_factor_input", true);`
`214`		`- CheckInputTensor(*output, "scaling_factor_output", true);`
	`213`	`+ CheckInputTensor(*input, "scaling_factor_input");`
	`214`	`+ CheckInputTensor(*output, "scaling_factor_output");`
`215`	`215`
`216`	`216`	`auto& scaling_mode = input->scaling_mode;`
`217`	`217`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ void CheckNoopTensor(const Tensor &t, const std::string &name) {`
`65`	`65`	`}`
`66`	`66`	`}`
`67`	`67`
`68`		`-void CheckScaleTensorShape(const Tensor &t, bool check_scale_inv_alignment) {`
	`68`	`+void CheckScaleTensorShape(const Tensor &t) {`
`69`	`69`	`NVTE_CHECK(t.scaling_mode != NVTE_INVALID_SCALING, "Invalid scaling mode!");`
`70`	`70`	`if (is_tensor_scaling(t.scaling_mode)) {`
`71`	`71`	`// per-tensor scaling`
`@@ -80,7 +80,6 @@ void CheckScaleTensorShape(const Tensor &t, bool check_scale_inv_alignment) {`
`80`	`80`	`}`
`81`	`81`	`} else {`
`82`	`82`	`if (t.scaling_mode == NVTE_MXFP8_1D_SCALING) {`
`83`		`- if (!check_scale_inv_alignment) return;`
`84`	`83`	`// Need (4, 128) alignment even for e8 scaling factor`
`85`	`84`	`auto block_alignment = std::vector<size_t>{128ul / typeToSize(t.scale_inv.dtype),`
`86`	`85`	`4ul / typeToSize(t.scale_inv.dtype)};`
`@@ -111,7 +110,7 @@ void CheckScaleTensorShape(const Tensor &t, bool check_scale_inv_alignment) {`
`111`	`110`	`}`
`112`	`111`	`}`
`113`	`112`
`114`		`-void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale_inv_alignment) {`
	`113`	`+void CheckInputTensor(const Tensor &t, const std::string &name) {`
`115`	`114`	`const DType type = t.dtype();`
`116`	`115`	`if (is_fp8_dtype(type)) {`
`117`	`116`	`// FP8 input needs to have scale_inv`
`@@ -143,11 +142,10 @@ void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale`
`143`	`142`	`}`
`144`	`143`	`NVTE_CHECK(t.has_data() \|\| t.has_columnwise_data(), "Input ", name, " is not allocated!");`
`145`	`144`
`146`		`- CheckScaleTensorShape(t, check_scale_inv_alignment);`
	`145`	`+ CheckScaleTensorShape(t);`
`147`	`146`	`}`
`148`	`147`
`149`		`-void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty,`
`150`		`- bool check_scale_inv_alignment) {`
	`148`	`+void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty) {`
`151`	`149`	`const DType type = t.dtype();`
`152`	`150`	`if (is_fp8_dtype(type)) {`
`153`	`151`	`// FP8 output needs to have scale, scale_inv and (if delayed scaling) amax`
`@@ -189,7 +187,7 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt`
`189`	`187`	`NVTE_CHECK(t.has_data() \|\| t.has_columnwise_data(), "Output ", name, " is not allocated!");`
`190`	`188`	`}`
`191`	`189`
`192`		`- CheckScaleTensorShape(t, check_scale_inv_alignment);`
	`190`	`+ CheckScaleTensorShape(t);`
`193`	`191`	`}`
`194`	`192`
`195`	`193`	`} // namespace transformer_engine`
Original file line number	Diff line number	Diff line change
`@@ -223,4 +223,9 @@ std::vector<size_t> convertShape(const NVTEShape& shape) {`
`223`	`223`	`return std::vector<size_t>(shape.data, shape.data + shape.ndim);`
`224`	`224`	`}`
`225`	`225`
	`226`	`+int roundup(const int value, const int multiple) {`
	`227`	`+ assert(multiple > 0);`
	`228`	`+ return ((value + multiple - 1) / multiple) * multiple;`
	`229`	`+}`
	`230`	`+`
`226`	`231`	`} // namespace transformer_engine::pytorch`