[PyTorch] Fix linter warnings from unused args (#816)

* Fix linter warnings from unused args
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Update .gitignore
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

.gitignore

@@ -26,4 +26,4 @@ docs/_build
 .ipynb_checkpoints
 docs/doxygen
 *.log
-CMakeFiles/CMakeSystem.cmake
+CMakeFiles/CMakeSystem.cmake
\ No newline at end of file

transformer_engine/pytorch/attention.py

@@ -1757,11 +1757,12 @@ class _PrepareQKVForFA(torch.autograd.Function):
        to separate contiguous q, k, v tensors in (b, s, ...) layout."""
 
     @staticmethod
-    def forward(ctx,
-                query_layer: torch.Tensor,
-                key_layer: torch.Tensor,
-                value_layer: torch.Tensor
-    ) -> torch.Tensor:
+    def forward(
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         # All inputs received are non-contiguous tensors.
         # The `query_layer` tensor is used to access the
         # full memory region of the QKV tensor.
@@ -1773,10 +1774,11 @@ class _PrepareQKVForFA(torch.autograd.Function):
         return query_layer, key_layer, value_layer
 
     @staticmethod
-    def backward(ctx,
-                 dq: torch.Tensor,
-                 dk: torch.Tensor,
-                 dv: torch.Tensor
+    def backward(
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
+        dq: torch.Tensor,
+        dk: torch.Tensor,
+        dv: torch.Tensor
     ) -> Tuple[Union[torch.Tensor, None], ...]:
         dqkv = tex.fa_prepare_bwd(dq, dk, dv)
         dq, dk, dv = split_tensor_along_dim(dqkv, -1, 3)

transformer_engine/pytorch/float8_tensor.py

@@ -46,7 +46,7 @@ class _FromFloat8Func(torch.autograd.Function):
     """Cast from FP8 to other dtype"""
     @staticmethod
     def forward(
-        ctx,
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
         tensor: Float8Tensor,
         dtype: Optional[torch.dtype] = None,
     ) -> torch.Tensor:
@@ -63,7 +63,10 @@ class _FromFloat8Func(torch.autograd.Function):
         return out
 
     @staticmethod
-    def backward(ctx, grad):
+    def backward(
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
+        grad: torch.Tensor,
+    ) -> Tuple[Optional[torch.Tensor], ...]:
         # Assume that we want gradients in full precision
         return grad, None
 
@@ -97,7 +100,7 @@ class _ToFloat8Func(torch.autograd.Function):
     """Cast to FP8 from other dtype"""
     @staticmethod
     def forward(
-        ctx,
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
         tensor: torch.Tensor,
         fp8_meta: Optional[Dict[str, Any]] = None,
         fp8_meta_forward: bool = True,
@@ -106,7 +109,7 @@ class _ToFloat8Func(torch.autograd.Function):
         scale: Optional[torch.Tensor] = None,
         amax: Optional[torch.Tensor] = None,
         scale_inv: Optional[torch.Tensor] = None,
-    ):
+    ) -> Float8Tensor:
 
         # Manually compute scale-inverse if needed
         if scale is not None and scale_inv is None:
@@ -189,7 +192,10 @@ class _ToFloat8Func(torch.autograd.Function):
         )
 
     @staticmethod
-    def backward(ctx, grad):
+    def backward(
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
+        grad: torch.Tensor,
+    ) -> Tuple[Optional[torch.Tensor], ...]:
         # Assume that we want gradients in full precision
         return grad, None, None, None, None, None, None, None