[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

docker_build_and_test.sh

@@ -82,4 +82,4 @@ docker run --gpus all -it --rm \
     echo "=== Running operator tests ==="
     cd /workspace/TransformerEngine/tests/cpp
     ./build/operator/test_operator "$@"
-  ' _ "${TEST_ARGS[@]}"
+  ' _ "${TEST_ARGS[@]}"

patch_swizzle.py

@@ -8,8 +8,8 @@
 template <int SF_TILE_DIM_M, int SF_TILE_DIM_K>
 __global__ void __launch_bounds__(TB_DIM* TB_DIM)
     grouped_swizzle_scaling_variable_shape_kernel(
-        const void* input, 
-        void* output, 
+        const void* input,
+        void* output,
         const int64_t* m_array,
         const int64_t* k_array,
         const int* block_offsets,
@@ -42,23 +42,23 @@
   if (tensor_id == -1) return;
 
   int local_block_id = linear_block_id - block_offsets[tensor_id];
-  
+
   size_t M = rowwise ? m_array[tensor_id] : k_array[tensor_id];
   size_t K = rowwise ? k_array[tensor_id] : m_array[tensor_id];
-  
+
   size_t padded_m = round_up_to_multiple(M, 128);
   size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast<size_t>(MXFP8_BLOCK_SIZE)), 4);
-  
+
   int num_tiles_m = padded_m / SF_TILE_DIM_M;
   int num_tiles_k = padded_k / SF_TILE_DIM_K;
-  
+
   int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));
   if (vec_load_size == 3) vec_load_size = 1;
   int n_tiles_in_tb = TB_DIM * vec_load_size;
 
   int grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM);
   int grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size);
-  
+
   int block_x = local_block_id % grid_dim_x;
   int block_y = local_block_id / grid_dim_x;
 
@@ -71,29 +71,29 @@
   if (rowwise) {
       if (vec_load_size == 4) {
           swizzle_row_scaling_kernel_impl<int4, SF_TILE_DIM_M, SF_TILE_DIM_K>(
-              input_base, output_base, padded_m, padded_k, original_M, original_K, 
+              input_base, output_base, padded_m, padded_k, original_M, original_K,
               block_x, block_y, grid_dim_x, grid_dim_y);
       } else if (vec_load_size == 2) {
           swizzle_row_scaling_kernel_impl<int2, SF_TILE_DIM_M, SF_TILE_DIM_K>(
-              input_base, output_base, padded_m, padded_k, original_M, original_K, 
+              input_base, output_base, padded_m, padded_k, original_M, original_K,
               block_x, block_y, grid_dim_x, grid_dim_y);
       } else {
           swizzle_row_scaling_kernel_impl<int, SF_TILE_DIM_M, SF_TILE_DIM_K>(
-              input_base, output_base, padded_m, padded_k, original_M, original_K, 
+              input_base, output_base, padded_m, padded_k, original_M, original_K,
               block_x, block_y, grid_dim_x, grid_dim_y);
       }
   } else {
       if (vec_load_size == 4) {
           swizzle_col_scaling_kernel_impl<int4, SF_TILE_DIM_M, SF_TILE_DIM_K>(
-              input_base, output_base, padded_m, padded_k, original_M, original_K, 
+              input_base, output_base, padded_m, padded_k, original_M, original_K,
               block_x, block_y, grid_dim_x, grid_dim_y);
       } else if (vec_load_size == 2) {
           swizzle_col_scaling_kernel_impl<int2, SF_TILE_DIM_M, SF_TILE_DIM_K>(
-              input_base, output_base, padded_m, padded_k, original_M, original_K, 
+              input_base, output_base, padded_m, padded_k, original_M, original_K,
               block_x, block_y, grid_dim_x, grid_dim_y);
       } else {
           swizzle_col_scaling_kernel_impl<int, SF_TILE_DIM_M, SF_TILE_DIM_K>(
-              input_base, output_base, padded_m, padded_k, original_M, original_K, 
+              input_base, output_base, padded_m, padded_k, original_M, original_K,
               block_x, block_y, grid_dim_x, grid_dim_y);
       }
   }
@@ -113,34 +113,34 @@
   if (blockIdx.x == 0 && threadIdx.x == 0) {
     int current_block_offset = 0;
     size_t current_scale_offset = 0;
-    
+
     for (size_t i = 0; i < num_tensors; ++i) {
       block_offsets[i] = current_block_offset;
       scale_offsets[i] = current_scale_offset;
-      
+
       size_t m = rowwise ? m_array[i] : k_array[i];
       size_t k = rowwise ? k_array[i] : m_array[i];
-      
+
       size_t padded_m = round_up_to_multiple(m, 128);
       size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast<size_t>(MXFP8_BLOCK_SIZE)), 4);
-      
+
       int num_tiles_m = padded_m / 128;
       int num_tiles_k = padded_k / 4;
-      
+
       int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1));
       if (vec_load_size == 3) vec_load_size = 1;
-      
+
       int blocks_m = num_tiles_m;
       int blocks_k = DIVUP(num_tiles_k, TB_DIM * vec_load_size);
       if (!rowwise) {
           blocks_m = DIVUP(num_tiles_m, vec_load_size);
           blocks_k = DIVUP(num_tiles_k, TB_DIM);
       }
-      
+
       current_block_offset += blocks_m * blocks_k;
       current_scale_offset += padded_m * padded_k * scale_elem_size;
     }
-    
+
     block_offsets[num_tensors] = current_block_offset;
     scale_offsets[num_tensors] = current_scale_offset;
     *total_blocks = current_block_offset;
@@ -150,7 +150,10 @@
 
 namespace transformer_engine {
 """
-content = content.replace("namespace transformer_engine {\n\nvoid swizzle_grouped_scaling_factors", kernels_code + "\nvoid swizzle_grouped_scaling_factors")
+content = content.replace(
+    "namespace transformer_engine {\n\nvoid swizzle_grouped_scaling_factors",
+    kernels_code + "\nvoid swizzle_grouped_scaling_factors",
+)
 
 # 2. Modify swizzle_grouped_scaling_factors
 old_func = """void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output,
@@ -206,7 +209,7 @@
     auto launch_grouped_swizzle_variable = [&](bool rowwise) {
       const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype)
                                              : typeToSize(input->columnwise_scale_inv.dtype);
-      
+
       compute_grouped_swizzle_setup<<<1, 1, 0, stream>>>(
           m_array, k_array, d_block_offsets, d_scale_offsets, d_total_blocks,
           d_global_counter, num_tensors, rowwise, scale_elem_size);
@@ -215,7 +218,7 @@
           grouped_swizzle_scaling_variable_shape_kernel<SF_TILE_DIM_M, SF_TILE_DIM_K>,
           cudaFuncAttributeMaxDynamicSharedMemorySize, max_slm_size));
 
-      int persistent_blocks = 108 * 8; 
+      int persistent_blocks = 108 * 8;
       dim3 num_blocks(persistent_blocks);
 
       const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr;
@@ -257,4 +260,3 @@
 
 with open("transformer_engine/common/swizzle/swizzle.cu", "w") as f:
     f.write(content)
-

patch_swizzle_cpp.py

@@ -10,7 +10,7 @@
   });"""
 
 new_code = """  swizzle_output.set_with_gemm_swizzled_scales(true);
-  
+
   size_t num_tensors = input.num_tensors();
   size_t workspace_size = (num_tensors + 2) * sizeof(int) + (num_tensors + 1) * sizeof(size_t);
   workspace_size = roundup(workspace_size, 256);
@@ -40,4 +40,3 @@
 
 with open("transformer_engine/pytorch/csrc/extensions/swizzle.cpp", "w") as f:
     f.write(content)
-

transformer_engine/common/util/vectorized_pointwise.h

@@ -228,7 +228,8 @@ __launch_bounds__(unary_kernel_threads) __global__
     loader.load(tid, size);
 #pragma unroll
     for (int i = 0; i < nvec; ++i) {
-      const size_t global_idx = (aligned ? (tid * nvec + i) : (tid * nvec + i - loader.alignment()));
+      const size_t global_idx =
+          (aligned ? (tid * nvec + i) : (tid * nvec + i - loader.alignment()));
       if (global_idx >= size) continue;
 
       ComputeType val = static_cast<ComputeType>(loader.separate()[i]);
@@ -332,7 +333,8 @@ __launch_bounds__(unary_kernel_threads) __global__
     grad_loader.load(tid, size);
 #pragma unroll
     for (int i = 0; i < nvec; ++i) {
-      const size_t global_idx = (aligned ? (tid * nvec + i) : (tid * nvec + i - loader.alignment()));
+      const size_t global_idx =
+          (aligned ? (tid * nvec + i) : (tid * nvec + i - loader.alignment()));
       if (global_idx >= size) continue;
 
       ComputeType val = static_cast<ComputeType>(loader.separate()[i]);
@@ -466,19 +468,19 @@ void VectorizedUnaryKernelLauncher(const InputType *input, const fp32 *noop, Out
     switch (align) {
       case Alignment::SAME_ALIGNED:
         unary_kernel<nvec, true, fp32, Param, OP><<<grid, threads, 0, stream>>>(
-            input, noop, output, scale, amax, scale_inv, params, N, num_aligned_elements,
-            offsets, first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
+            input, noop, output, scale, amax, scale_inv, params, N, num_aligned_elements, offsets,
+            first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
         break;
       case Alignment::SAME_UNALIGNED:
         unary_kernel<nvec, false, fp32, Param, OP><<<grid, threads, 0, stream>>>(
-            input, noop, output, scale, amax, scale_inv, params, N, num_aligned_elements,
-            offsets, first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
+            input, noop, output, scale, amax, scale_inv, params, N, num_aligned_elements, offsets,
+            first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
         break;
       case Alignment::DIFFERENT: {
         // If the pointers are aligned differently we cannot vectorize
         unary_kernel<1, true, fp32, Param, OP><<<grid, threads, 0, stream>>>(
-            input, noop, output, scale, amax, scale_inv, params, N, N,
-            offsets, first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
+            input, noop, output, scale, amax, scale_inv, params, N, N, offsets, first_dims,
+            last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
         break;
       }
     }
@@ -508,19 +510,19 @@ void VectorizedUnaryGradKernelLauncher(const InputTypeGrad *grad, const InputTyp
     switch (align) {
       case Alignment::SAME_ALIGNED:
         unary_grad_kernel<nvec, true, fp32, Param, OP><<<grid, threads, 0, stream>>>(
-            grad, input, output, scale, amax, scale_inv, params, N, num_aligned_elements,
-            offsets, first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
+            grad, input, output, scale, amax, scale_inv, params, N, num_aligned_elements, offsets,
+            first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
         break;
       case Alignment::SAME_UNALIGNED:
         unary_grad_kernel<nvec, false, fp32, Param, OP><<<grid, threads, 0, stream>>>(
-            grad, input, output, scale, amax, scale_inv, params, N, num_aligned_elements,
-            offsets, first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
+            grad, input, output, scale, amax, scale_inv, params, N, num_aligned_elements, offsets,
+            first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
         break;
       case Alignment::DIFFERENT: {
         // If the pointers are aligned differently we cannot vectorize
         unary_grad_kernel<1, true, fp32, Param, OP><<<grid, threads, 0, stream>>>(
-            grad, input, output, scale, amax, scale_inv, params, N, N,
-            offsets, first_dims, last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
+            grad, input, output, scale, amax, scale_inv, params, N, N, offsets, first_dims,
+            last_dims, num_tensors, scale_numel, scale_inv_numel, amax_numel);
         break;
       }
     }