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

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

Author: pre-commit-ci[bot]

transformer_engine/common/cast/nvfp4/quantize_nvfp4_per_token.cu

@@ -57,14 +57,14 @@ constexpr int kMaxTensorsPerKernel = 64;
 // ============================================================================
 struct NVFP4PerTokenMultiArgs {
   // K1 outputs (per-tensor pointers; one fp32 array per tensor)
-  void* row_amax_list[kMaxTensorsPerKernel];        // each: float* (M_i,)
-  void* col_amax_list[kMaxTensorsPerKernel];        // each: float* (K,)
+  void* row_amax_list[kMaxTensorsPerKernel];  // each: float* (M_i,)
+  void* col_amax_list[kMaxTensorsPerKernel];  // each: float* (K,)
 
   // K2 outputs (per-tensor pointers; FP4 codes + e4m3 inner SF)
-  void* q_row_list[kMaxTensorsPerKernel];           // each: uint8* (M_i, K/2)
-  void* s_dec_row_list[kMaxTensorsPerKernel];       // each: fp8e4m3* (M_i, K/16)
-  void* q_col_list[kMaxTensorsPerKernel];           // each: uint8* (K, M_i/2)
-  void* s_dec_col_list[kMaxTensorsPerKernel];       // each: fp8e4m3* (K, M_i/16)
+  void* q_row_list[kMaxTensorsPerKernel];      // each: uint8* (M_i, K/2)
+  void* s_dec_row_list[kMaxTensorsPerKernel];  // each: fp8e4m3* (M_i, K/16)
+  void* q_col_list[kMaxTensorsPerKernel];      // each: uint8* (K, M_i/2)
+  void* s_dec_col_list[kMaxTensorsPerKernel];  // each: fp8e4m3* (K, M_i/16)
 
   // Shared layout info
   int split_sections_range[kMaxTensorsPerKernel + 1];  // prefix sum w/ leading 0
@@ -178,13 +178,11 @@ __global__ void __launch_bounds__(kColStripThreads)
 
   int cur_tensor_id = 0;
   while (cur_tensor_id < args.num_tensors &&
-         args.split_sections_range[cur_tensor_id + 1] ==
-             args.split_sections_range[cur_tensor_id]) {
+         args.split_sections_range[cur_tensor_id + 1] == args.split_sections_range[cur_tensor_id]) {
     ++cur_tensor_id;
   }
-  int cur_tensor_end = (cur_tensor_id < args.num_tensors)
-                           ? args.split_sections_range[cur_tensor_id + 1]
-                           : 0;
+  int cur_tensor_end =
+      (cur_tensor_id < args.num_tensors) ? args.split_sections_range[cur_tensor_id + 1] : 0;
 
   // Walk M in kInnerK-row chunks. split_sections[i] % 64 == 0 implies
   // every chunk boundary aligns with a split boundary, so we never
@@ -194,10 +192,8 @@ __global__ void __launch_bounds__(kColStripThreads)
     // until cur_tensor_end > m_base. Only flush for NON-EMPTY tensors
     // (empty tensors' col_amax_list[] slots are NULL).
     while (m_base >= cur_tensor_end) {
-      if (col_in_range && my_col_amax_acc > 0.f &&
-          args.col_amax_list[cur_tensor_id] != nullptr) {
-        float* dst =
-            reinterpret_cast<float*>(args.col_amax_list[cur_tensor_id]) + my_col;
+      if (col_in_range && my_col_amax_acc > 0.f && args.col_amax_list[cur_tensor_id] != nullptr) {
+        float* dst = reinterpret_cast<float*>(args.col_amax_list[cur_tensor_id]) + my_col;
         atomicMaxFloat(dst, my_col_amax_acc);
       }
       my_col_amax_acc = 0.f;
@@ -207,8 +203,8 @@ __global__ void __launch_bounds__(kColStripThreads)
     }
     if (cur_tensor_id >= args.num_tensors) break;
 
-    // Per-element scan within the 16-row chunk (verbatim from single-tensor
-    // K1 col Pass 1).
+      // Per-element scan within the 16-row chunk (verbatim from single-tensor
+      // K1 col Pass 1).
 #pragma unroll
     for (int e = 0; e < kInnerK; ++e) {
       const int gr = m_base + e;
@@ -255,17 +251,14 @@ __global__ void __launch_bounds__(kRowwiseThreads)
   const int local_row = global_row - args.split_sections_range[tensor_id];
 
   // Read the row's outer amax (populated by K1-group rowwise).
-  const float row_amax =
-      reinterpret_cast<float*>(args.row_amax_list[tensor_id])[local_row];
+  const float row_amax = reinterpret_cast<float*>(args.row_amax_list[tensor_id])[local_row];
   const float S_enc = compute_global_encode_scaling_factor_FP4(fmaxf(row_amax, 1e-12f));
 
   // Per-tensor row base pointers.
-  uint8_t* row_out =
-      reinterpret_cast<uint8_t*>(args.q_row_list[tensor_id]) +
-      static_cast<size_t>(local_row) * (K / 2);
-  fp8e4m3* s_dec_out =
-      reinterpret_cast<fp8e4m3*>(args.s_dec_row_list[tensor_id]) +
-      static_cast<size_t>(local_row) * (K / kInnerK);
+  uint8_t* row_out = reinterpret_cast<uint8_t*>(args.q_row_list[tensor_id]) +
+                     static_cast<size_t>(local_row) * (K / 2);
+  fp8e4m3* s_dec_out = reinterpret_cast<fp8e4m3*>(args.s_dec_row_list[tensor_id]) +
+                       static_cast<size_t>(local_row) * (K / kInnerK);
 
   // === verbatim from single-tensor K1 rowwise Pass 2 ===
   const int n_blocks = K / kInnerK;
@@ -275,8 +268,7 @@ __global__ void __launch_bounds__(kRowwiseThreads)
     float bmx = 0.f;
 #pragma unroll
     for (int e = 0; e < kInnerK; e++) {
-      const float v =
-          static_cast<float>(in[static_cast<size_t>(global_row) * K + b * kInnerK + e]);
+      const float v = static_cast<float>(in[static_cast<size_t>(global_row) * K + b * kInnerK + e]);
       vals[e] = v;
       bmx = fmaxf(bmx, fabsf(v));
     }
@@ -341,12 +333,12 @@ __global__ void __launch_bounds__(kColStripThreads)
   // Per-tensor cached state. Initialize so the first chunk (b == 0, m_base == 0)
   // triggers the boundary-advance to populate these.
   float S_enc_cur = 0.f;
-  int cur_tensor_id = -1;     // -1 forces first-iteration advance
-  int cur_tensor_end = 0;     // exclusive
-  int local_block_base = 0;   // global block index of this tensor's first block
+  int cur_tensor_id = -1;    // -1 forces first-iteration advance
+  int cur_tensor_end = 0;    // exclusive
+  int local_block_base = 0;  // global block index of this tensor's first block
   uint8_t* col_out = nullptr;
   fp8e4m3* s_dec_col_out = nullptr;
-  int cur_tensor_M = 0;       // = split_sections[cur_tensor_id]
+  int cur_tensor_M = 0;  // = split_sections[cur_tensor_id]
 
   const int n_blocks_m = sum_M / kInnerK;
   for (int b = 0; b < n_blocks_m; b++) {
@@ -366,8 +358,7 @@ __global__ void __launch_bounds__(kColStripThreads)
       need_refresh = true;
     }
     if (need_refresh && col_in_range && cur_tensor_M > 0) {
-      const float col_amax =
-          reinterpret_cast<float*>(args.col_amax_list[cur_tensor_id])[my_col];
+      const float col_amax = reinterpret_cast<float*>(args.col_amax_list[cur_tensor_id])[my_col];
       S_enc_cur = compute_global_encode_scaling_factor_FP4(fmaxf(col_amax, 1e-12f));
       col_out = reinterpret_cast<uint8_t*>(args.q_col_list[cur_tensor_id]) +
                 static_cast<size_t>(my_col) * (cur_tensor_M / 2);
@@ -499,21 +490,20 @@ void populate_args(NVFP4PerTokenMultiArgs* args, std::vector<Tensor*>& outputs,
   args->split_sections_range[0] = 0;
   for (size_t i = 0; i < num_tensors; ++i) {
     Tensor* o = outputs[i];
-    NVTE_CHECK(split_sections[i] % 64 == 0, "split_sections[", i,
-               "] = ", split_sections[i], " must be a multiple of 64");
+    NVTE_CHECK(split_sections[i] % 64 == 0, "split_sections[", i, "] = ", split_sections[i],
+               " must be a multiple of 64");
     args->split_sections_range[i + 1] =
         args->split_sections_range[i] + static_cast<int>(split_sections[i]);
     // Empty splits skip pointer validation -- the kernel boundary-advance
     // loop walks PAST them in zero iterations, never touching the pointer.
     if (split_sections[i] == 0) continue;
     if (which_buffers & kBufRowAmax) {
-      NVTE_CHECK(o->amax.dptr != nullptr,
-                 "NVFP4 per-token grouped: outputs[", i, "].amax must be allocated for rowwise");
+      NVTE_CHECK(o->amax.dptr != nullptr, "NVFP4 per-token grouped: outputs[", i,
+                 "].amax must be allocated for rowwise");
       args->row_amax_list[i] = o->amax.dptr;
     }
     if (which_buffers & kBufColAmax) {
-      NVTE_CHECK(o->columnwise_amax.dptr != nullptr,
-                 "NVFP4 per-token grouped: outputs[", i,
+      NVTE_CHECK(o->columnwise_amax.dptr != nullptr, "NVFP4 per-token grouped: outputs[", i,
                  "].columnwise_amax must be allocated for columnwise");
       args->col_amax_list[i] = o->columnwise_amax.dptr;
     }
@@ -525,10 +515,9 @@ void populate_args(NVFP4PerTokenMultiArgs* args, std::vector<Tensor*>& outputs,
       args->s_dec_row_list[i] = o->scale_inv.dptr;
     }
     if (which_buffers & kBufColCast) {
-      NVTE_CHECK(
-          o->columnwise_data.dptr != nullptr && o->columnwise_scale_inv.dptr != nullptr,
-          "NVFP4 per-token grouped: outputs[", i,
-          "].columnwise_data + .columnwise_scale_inv must be allocated for columnwise cast");
+      NVTE_CHECK(o->columnwise_data.dptr != nullptr && o->columnwise_scale_inv.dptr != nullptr,
+                 "NVFP4 per-token grouped: outputs[", i,
+                 "].columnwise_data + .columnwise_scale_inv must be allocated for columnwise cast");
       args->q_col_list[i] = o->columnwise_data.dptr;
       args->s_dec_col_list[i] = o->columnwise_scale_inv.dptr;
     }
@@ -557,8 +546,8 @@ void quantize_per_token_grouped(const Tensor& input, std::vector<Tensor*>& outpu
   const int sum_M = static_cast<int>(input.flat_first_dim());
   const int K = static_cast<int>(input.flat_last_dim());
   if (sum_M == 0 || K == 0) return;
-  NVTE_CHECK(K % kInnerK == 0,
-             "NVFP4 per-token grouped: K (", K, ") must be a multiple of ", kInnerK);
+  NVTE_CHECK(K % kInnerK == 0, "NVFP4 per-token grouped: K (", K, ") must be a multiple of ",
+             kInnerK);
 
   // Amax buffer pointers must be populated whenever EITHER the K1 (writes
   // amax) or K2 (reads amax) pass runs in that direction. K2 reads
@@ -606,16 +595,16 @@ std::vector<transformer_engine::Tensor*> collect_outputs(NVTETensor* outputs, si
 }  // namespace
 
 void nvte_group_nvfp4_per_token_amax(const NVTETensor input, NVTETensor* outputs,
-                                     const size_t* split_sections, size_t num_tensors,
-                                     bool rowwise, bool columnwise, cudaStream_t stream) {
+                                     const size_t* split_sections, size_t num_tensors, bool rowwise,
+                                     bool columnwise, cudaStream_t stream) {
 #if FP4_TYPE_SUPPORTED
   NVTE_API_CALL(nvte_group_nvfp4_per_token_amax);
   using namespace transformer_engine;
   if (num_tensors == 0) return;
   const Tensor* in = convertNVTETensorCheck(input);
   std::vector<Tensor*> outs = collect_outputs(outputs, num_tensors);
-  nvfp4_per_token_group::quantize_per_token_grouped(*in, outs, split_sections, num_tensors,
-                                                    rowwise, columnwise,
+  nvfp4_per_token_group::quantize_per_token_grouped(*in, outs, split_sections, num_tensors, rowwise,
+                                                    columnwise,
                                                     /*do_amax=*/true, /*do_cast=*/false, stream);
 #else
   (void)input;
@@ -630,16 +619,16 @@ void nvte_group_nvfp4_per_token_amax(const NVTETensor input, NVTETensor* outputs
 }
 
 void nvte_group_nvfp4_per_token_cast(const NVTETensor input, NVTETensor* outputs,
-                                     const size_t* split_sections, size_t num_tensors,
-                                     bool rowwise, bool columnwise, cudaStream_t stream) {
+                                     const size_t* split_sections, size_t num_tensors, bool rowwise,
+                                     bool columnwise, cudaStream_t stream) {
 #if FP4_TYPE_SUPPORTED
   NVTE_API_CALL(nvte_group_nvfp4_per_token_cast);
   using namespace transformer_engine;
   if (num_tensors == 0) return;
   const Tensor* in = convertNVTETensorCheck(input);
   std::vector<Tensor*> outs = collect_outputs(outputs, num_tensors);
-  nvfp4_per_token_group::quantize_per_token_grouped(*in, outs, split_sections, num_tensors,
-                                                    rowwise, columnwise,
+  nvfp4_per_token_group::quantize_per_token_grouped(*in, outs, split_sections, num_tensors, rowwise,
+                                                    columnwise,
                                                     /*do_amax=*/false, /*do_cast=*/true, stream);
 #else
   (void)input;
@@ -662,8 +651,8 @@ void nvte_group_nvfp4_per_token_quantize(const NVTETensor input, NVTETensor* out
   if (num_tensors == 0) return;
   const Tensor* in = convertNVTETensorCheck(input);
   std::vector<Tensor*> outs = collect_outputs(outputs, num_tensors);
-  nvfp4_per_token_group::quantize_per_token_grouped(*in, outs, split_sections, num_tensors,
-                                                    rowwise, columnwise,
+  nvfp4_per_token_group::quantize_per_token_grouped(*in, outs, split_sections, num_tensors, rowwise,
+                                                    columnwise,
                                                     /*do_amax=*/true, /*do_cast=*/true, stream);
 #else
   (void)input;

transformer_engine/common/cast/nvfp4/quantize_nvfp4_per_token_group.cu

@@ -15,31 +15,30 @@
 extern "C" {
 #endif
 
-
 /*! \brief Composite K1+K2: per-row + per-col amax (K1) then FP4 + 1x16
  *         e4m3 SF encode (K2), back-to-back on the same stream.
  *
  *  This is the production entry point for the per-token cast on bf16 +
  *  128-aligned shapes.
  */
-void nvte_nvfp4_per_token_quantize(const NVTETensor input, const NVTETensor noop,
-                                          NVTETensor output, cudaStream_t stream);
+void nvte_nvfp4_per_token_quantize(const NVTETensor input, const NVTETensor noop, NVTETensor output,
+                                   cudaStream_t stream);
 
 /*! \brief Kernel 1 in isolation: per-row + per-col amax via TMA + atomicMax.
  *         Pre-zeroes the amax buffers and merges per-CTA partials into
  *         ``output->amax`` (size [M]) / ``output->columnwise_amax``
  *         (size [K]). Does NOT touch FP4 data / scale_inv slots.
  */
-void nvte_nvfp4_per_token_amax(const NVTETensor input, const NVTETensor noop,
-                                      NVTETensor output, cudaStream_t stream);
+void nvte_nvfp4_per_token_amax(const NVTETensor input, const NVTETensor noop, NVTETensor output,
+                               cudaStream_t stream);
 
 /*! \brief Kernel 2 in isolation: FP4 + 1x16 e4m3 SF encode given a
  *         pre-filled ``output->amax`` / ``output->columnwise_amax``. Reads
  *         the outer amax buffer(s) and writes the FP4 data / scale_inv
  *         tensors only.
  */
-void nvte_nvfp4_per_token_encode(const NVTETensor input, const NVTETensor noop,
-                                        NVTETensor output, cudaStream_t stream);
+void nvte_nvfp4_per_token_encode(const NVTETensor input, const NVTETensor noop, NVTETensor output,
+                                 cudaStream_t stream);
 
 /*! \brief Returns 1 iff the per-token kernels accept ``(M, K, dtype)``.
  *
@@ -59,8 +58,7 @@ int nvte_nvfp4_per_token_can_dispatch(size_t M, size_t K, int input_dtype_enum);
  *      d[i, j] = d[i, j] * row_amax_a[i] * row_amax_b[j]
  */
 void nvte_nvfp4_per_token_post_scale(NVTETensor d, const NVTETensor row_amax_a,
-                                     const NVTETensor row_amax_b,
-                                     cudaStream_t stream);
+                                     const NVTETensor row_amax_b, cudaStream_t stream);
 
 /* ============================================================================
  * Grouped (multi-tensor) per-token quantize.
@@ -76,8 +74,8 @@ void nvte_nvfp4_per_token_post_scale(NVTETensor d, const NVTETensor row_amax_a,
  *  \param[in]     stream         CUDA stream
  */
 void nvte_group_nvfp4_per_token_amax(const NVTETensor input, NVTETensor* outputs,
-                                     const size_t* split_sections, size_t num_tensors,
-                                     bool rowwise, bool columnwise, cudaStream_t stream);
+                                     const size_t* split_sections, size_t num_tensors, bool rowwise,
+                                     bool columnwise, cudaStream_t stream);
 
 /*! \brief Grouped per-token encode (FP4 + 1x16 e4m3 inner SF) using the
  *         row_amax / col_amax values already populated by
@@ -94,8 +92,8 @@ void nvte_group_nvfp4_per_token_amax(const NVTETensor input, NVTETensor* outputs
  *  \param[in]     stream         CUDA stream
  */
 void nvte_group_nvfp4_per_token_cast(const NVTETensor input, NVTETensor* outputs,
-                                     const size_t* split_sections, size_t num_tensors,
-                                     bool rowwise, bool columnwise, cudaStream_t stream);
+                                     const size_t* split_sections, size_t num_tensors, bool rowwise,
+                                     bool columnwise, cudaStream_t stream);
 
 /*! \brief Composite K1+K2 grouped per-token quantize. Calls the amax + cast
  *         kernels on the same stream. This is the external API

transformer_engine/common/include/transformer_engine/nvfp4_per_token.h

@@ -448,30 +448,27 @@ void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwi
                                 const at::Tensor &scale_inv_colwise, int rows, int cols,
                                 size_t start_offset);
 
-void nvfp4_per_token_quantize(const at::Tensor &input, at::Tensor q_row,
-                              at::Tensor s_dec_row, at::Tensor row_amax,
-                              at::Tensor q_col, at::Tensor s_dec_col,
+void nvfp4_per_token_quantize(const at::Tensor &input, at::Tensor q_row, at::Tensor s_dec_row,
+                              at::Tensor row_amax, at::Tensor q_col, at::Tensor s_dec_col,
                               at::Tensor col_amax, bool rowwise, bool columnwise);
 
-void nvfp4_per_token_amax(const at::Tensor &input, at::Tensor row_amax,
-                                 at::Tensor col_amax, bool rowwise, bool columnwise);
+void nvfp4_per_token_amax(const at::Tensor &input, at::Tensor row_amax, at::Tensor col_amax,
+                          bool rowwise, bool columnwise);
 
-void nvfp4_per_token_encode(const at::Tensor &input, at::Tensor q_row,
-                                   at::Tensor s_dec_row, at::Tensor row_amax,
-                                   at::Tensor q_col, at::Tensor s_dec_col,
-                                   at::Tensor col_amax, bool rowwise, bool columnwise);
+void nvfp4_per_token_encode(const at::Tensor &input, at::Tensor q_row, at::Tensor s_dec_row,
+                            at::Tensor row_amax, at::Tensor q_col, at::Tensor s_dec_col,
+                            at::Tensor col_amax, bool rowwise, bool columnwise);
 
 void nvfp4_per_token_group_quantize(
     const at::Tensor &input, const std::vector<int64_t> &split_sections,
     std::vector<at::Tensor> q_row_list, std::vector<at::Tensor> s_dec_row_list,
     std::vector<at::Tensor> row_amax_list, std::vector<at::Tensor> q_col_list,
-    std::vector<at::Tensor> s_dec_col_list, std::vector<at::Tensor> col_amax_list,
-    bool rowwise, bool columnwise);
+    std::vector<at::Tensor> s_dec_col_list, std::vector<at::Tensor> col_amax_list, bool rowwise,
+    bool columnwise);
 
 // Amax-only variant of the grouped quantize. Useful for multi-rank training
 // where amax is allReduced before the cast pass.
-void nvfp4_per_token_group_amax(const at::Tensor &input,
-                                const std::vector<int64_t> &split_sections,
+void nvfp4_per_token_group_amax(const at::Tensor &input, const std::vector<int64_t> &split_sections,
                                 std::vector<at::Tensor> row_amax_list,
                                 std::vector<at::Tensor> col_amax_list, bool rowwise,
                                 bool columnwise);