[Common][PyTorch] Fuse NVFP4 quantization launches for MoE experts

transformer_engine/common/cast/cast.cu

@@ -47,6 +47,99 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output,
   dispatch::quantize_fwd_helper<IS_ACT, Empty, nullptr>(input, output, quant_config, stream);
 }
 
+void nvte_nvfp4_quantize_4over6_multi(const NVTETensor *inputs, NVTETensor *outputs,
+                                      const NVTEQuantizationConfig quant_config,
+                                      const size_t num_tensors, cudaStream_t stream) {
+  NVTE_API_CALL(nvte_nvfp4_quantize_4over6_multi);
+  using namespace transformer_engine;
+  using namespace transformer_engine::dispatch::nvfp4;
+
+  NVTE_CHECK(inputs != nullptr && outputs != nullptr,
+             "Multi 4over6 quantization requires non-null tensor lists.");
+  NVTE_CHECK(num_tensors > 0, "Multi 4over6 quantization requires a non-empty tensor list.");
+
+  QuantizationConfig quant_config_cpp;
+  if (quant_config != nullptr) {
+    quant_config_cpp = *reinterpret_cast<const QuantizationConfig *>(quant_config);
+  }
+  NVTE_CHECK(quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled,
+             "Multi 4over6 quantization requires a non-disabled 4over6 mode.");
+  NVTE_CHECK(!quant_config_cpp.stochastic_rounding,
+             "Multi 4over6 quantization does not support stochastic rounding.");
+
+  std::vector<Tensor> in_list;
+  std::vector<Tensor> out_list;
+  in_list.reserve(num_tensors);
+  out_list.reserve(num_tensors);
+  size_t rows = 0, cols = 0;
+  for (size_t i = 0; i < num_tensors; ++i) {
+    in_list.push_back(*convertNVTETensorCheck(inputs[i]));
+    out_list.push_back(*convertNVTETensorCheck(outputs[i]));
+    const auto &in = in_list.back();
+    const auto &out = out_list.back();
+    NVTE_CHECK(out.has_data() && !out.has_columnwise_data(),
+               "Multi 4over6 supports rowwise-only outputs.");
+    NVTE_CHECK(out.scale_inv.dptr != nullptr, "Multi 4over6 requires allocated scaling tensors.");
+    NVTE_CHECK(out.amax.dptr != nullptr, "Multi 4over6 requires allocated amax tensors.");
+    NVTE_CHECK(is_fp4_dtype(out.data.dtype), "Multi 4over6 output must have FP4 type.");
+    NVTE_CHECK(!out.row_scaled_nvfp4,
+               "Multi 4over6 targets per-tensor-scaled tensors, not row-scaled ones.");
+    NVTE_CHECK(!out.with_gemm_swizzled_scales, "Multi 4over6 requires compact scale layout.");
+    if (i == 0) {
+      rows = in.flat_first_dim();
+      cols = in.flat_last_dim();
+    } else {
+      NVTE_CHECK(in.flat_first_dim() == rows && in.flat_last_dim() == cols,
+                 "Multi 4over6 requires same-shaped input tensors.");
+      NVTE_CHECK(in.dtype() == in_list[0].dtype(), "Multi 4over6 requires the same input dtype.");
+      NVTE_CHECK(out.scale_inv.shape == out_list[0].scale_inv.shape,
+                 "Multi 4over6 requires the same scaling tensor shape.");
+    }
+  }
+  NVTE_CHECK(cols % quantize_4over6_kernel::kGroupSize == 0,
+             "Multi 4over6 quantization requires columns divisible by ",
+             quantize_4over6_kernel::kGroupSize, ".");
+  const size_t scale_stride = out_list[0].scale_inv.shape[1];
+
+  using quantize_4over6_kernel::BatchedQuantizeParams;
+  std::vector<BatchedQuantizeParams> host_params(num_tensors);
+  for (size_t i = 0; i < num_tensors; ++i) {
+    host_params[i].input = in_list[i].data.dptr;
+    host_params[i].output = out_list[i].data.dptr;
+    host_params[i].scales = out_list[i].scale_inv.dptr;
+    host_params[i].amax = static_cast<float *>(out_list[i].amax.dptr);
+  }
+
+  void *params_dev = nullptr;
+  float *temp_amax_dev = nullptr;
+  NVTE_CHECK_CUDA(
+      cudaMallocAsync(&params_dev, num_tensors * sizeof(BatchedQuantizeParams), stream));
+  NVTE_CHECK_CUDA(cudaMallocAsync(&temp_amax_dev, num_tensors * sizeof(float), stream));
+  NVTE_CHECK_CUDA(cudaMemcpyAsync(params_dev, host_params.data(),
+                                  num_tensors * sizeof(BatchedQuantizeParams),
+                                  cudaMemcpyHostToDevice, stream));
+  NVTE_CHECK_CUDA(cudaMemsetAsync(temp_amax_dev, 0, num_tensors * sizeof(float), stream));
+
+  const float *noop_ptr = nullptr;
+  TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH(
+      out_list[0].nvfp4_e4m3_max, E4M3_MAX,
+      TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH(
+          quant_config_cpp.nvfp4_4over6_mode, MODE,
+          TRANSFORMER_ENGINE_SWITCH_CONDITION(
+              quant_config_cpp.nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, {
+                using Cfg = quantize_4over6_kernel::Config<MODE, ERR_USE_FAST_MATH>;
+                TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(
+                    in_list[0].dtype(), IType,
+                    quantize_4over6_kernel::launch_quantize_4over6_batched<Cfg, E4M3_MAX, IType>(
+                        params_dev, temp_amax_dev, static_cast<int>(num_tensors), rows, cols,
+                        scale_stride, noop_ptr, stream););
+              });););
+
+  NVTE_CHECK_CUDA(cudaFreeAsync(params_dev, stream));
+  NVTE_CHECK_CUDA(cudaFreeAsync(temp_amax_dev, stream));
+  NVTE_CHECK_CUDA(cudaGetLastError());
+}
+
 void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
   NVTE_API_CALL(nvte_dequantize);
   using namespace transformer_engine;

transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh

@@ -653,6 +653,146 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out
   });
 }
 
+// ---------------------------------------------------------------------------
+// Batched (multi-tensor) 4over6 quantization.
+//
+// Quantizes many same-shaped tensors in two kernel launches (one batched amax
+// pass + one batched quantize pass) instead of 2 launches per tensor. Each
+// tensor keeps its own scalar amax / global scale, so results are bit-identical
+// to calling the single-tensor path in a loop: per-tensor amax is an exact max
+// (order independent), and every 1x16 block's scales and 4/6 selection are
+// block-local.
+// ---------------------------------------------------------------------------
+
+struct BatchedQuantizeParams {
+  const void *input;  // IType*, [rows, cols]
+  void *output;       // fp4e2m1x2*, [rows, cols/2]
+  void *scales;       // nvfp4_scale_t*, [rows, scale_stride]
+  float *amax;        // float*, scalar output amax for this tensor
+};
+
+constexpr int kBatchedAmaxThreads = 256;
+constexpr int kBatchedAmaxTiles = 8;
+
+// Computes one scalar amax per tensor. grid = (kBatchedAmaxTiles, 1, num_tensors).
+// Partial block maxima are combined with atomicMax on the int view of the float
+// (exact for non-negative floats), after the caller zero-initializes temp_amax.
+template <typename IType>
+__global__ void __launch_bounds__(kBatchedAmaxThreads)
+    batched_amax_kernel(const BatchedQuantizeParams *params, float *temp_amax, const size_t rows,
+                        const size_t cols) {
+  const int z = blockIdx.z;
+  const IType *input = reinterpret_cast<const IType *>(params[z].input);
+  const size_t total = rows * cols;
+  float local = 0.0f;
+  for (size_t idx = blockIdx.x * kBatchedAmaxThreads + threadIdx.x; idx < total;
+       idx += static_cast<size_t>(kBatchedAmaxTiles) * kBatchedAmaxThreads) {
+    local = fmaxf(local, fabsf(static_cast<float>(input[idx])));
+  }
+#pragma unroll
+  for (int offset = kWarpThreads / 2; offset > 0; offset /= 2) {
+    local = fmaxf(local, __shfl_down_sync(0xffffffffu, local, offset));
+  }
+  __shared__ float smem[kBatchedAmaxThreads / kWarpThreads];
+  if ((threadIdx.x & (kWarpThreads - 1)) == 0) {
+    smem[threadIdx.x / kWarpThreads] = local;
+  }
+  __syncthreads();
+  if (threadIdx.x < kWarpThreads) {
+    constexpr int num_warps = kBatchedAmaxThreads / kWarpThreads;
+    local = threadIdx.x < num_warps ? smem[threadIdx.x] : 0.0f;
+#pragma unroll
+    for (int offset = num_warps / 2; offset > 0; offset /= 2) {
+      local = fmaxf(local, __shfl_down_sync(0xffffffffu, local, offset));
+    }
+    if (threadIdx.x == 0) {
+      atomicMax(reinterpret_cast<int *>(&temp_amax[z]), __float_as_int(local));
+    }
+  }
+}
+
+// grid = (col_tiles, row_tiles, num_tensors). Rowwise output only.
+template <typename Cfg, int E4M3_MAX, typename IType>
+__global__ void __launch_bounds__(kThreads)
+    quantize_4over6_batched_kernel(const BatchedQuantizeParams *params_arr, const float *temp_amax,
+                                   const size_t rows, const size_t cols, const size_t scale_stride,
+                                   const float *noop) {
+#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+  if (noop != nullptr && noop[0] == 1.0f) {
+    return;
+  }
+  const int z = blockIdx.z;
+  const BatchedQuantizeParams &p = params_arr[z];
+  const IType *input = reinterpret_cast<const IType *>(p.input);
+  fp4e2m1x2 *output = reinterpret_cast<fp4e2m1x2 *>(p.output);
+  nvfp4_scale_t *scales = reinterpret_cast<nvfp4_scale_t *>(p.scales);
+  const float *amax = temp_amax + z;
+
+  // Publish this tensor's scalar amax into its own output buffer.
+  if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0) {
+    *p.amax = temp_amax[z];
+  }
+
+  extern __shared__ char dynamic_shmem[];
+  auto *tiles = reinterpret_cast<IType *>(dynamic_shmem);
+  const size_t tile_col = blockIdx.x * kTileCols;
+  const size_t tile_row = blockIdx.y * kTileRows;
+
+  IType *stage_tiles[kPipelineStages];
+#pragma unroll
+  for (int stage = 0; stage < kPipelineStages; ++stage) {
+    stage_tiles[stage] = &tiles[stage * kStageRows * kTileCols];
+  }
+
+  load_stage_to_shared_async(input, stage_tiles[0], rows, cols, tile_row, tile_col);
+  cp_async_commit_group();
+  cp_async_wait_group<0>();
+  __syncthreads();
+
+  for (int stage = 0; stage < kPipelineStages; ++stage) {
+    const int next_stage = stage + 1;
+    if (next_stage < kPipelineStages) {
+      const size_t next_stage_row = tile_row + next_stage * kStageRows;
+      load_stage_to_shared_async(input, stage_tiles[next_stage], rows, cols, next_stage_row,
+                                 tile_col);
+      cp_async_commit_group();
+    }
+
+    const size_t stage_row = tile_row + stage * kStageRows;
+    IType *stage_tile = stage_tiles[stage];
+
+    quantize_stage_rowwise</*USE_2D_QUANTIZATION=*/false, /*ROW_SCALED_NVFP4=*/false, Cfg,
+                           E4M3_MAX>(stage_tile, output, scales, amax, rows, cols, stage_row,
+                                     tile_col, scale_stride);
+
+    if (next_stage < kPipelineStages) {
+      cp_async_wait_group<0>();
+      __syncthreads();
+    }
+  }
+#else
+  NVTE_DEVICE_ERROR("sm_100 or higher is required.");
+#endif
+}
+
+template <typename Cfg, int E4M3_MAX, typename IType>
+void launch_quantize_4over6_batched(const void *params_dev, float *temp_amax_dev, int num_tensors,
+                                    size_t rows, size_t cols, size_t scale_stride,
+                                    const float *noop, cudaStream_t stream) {
+  const auto *params = reinterpret_cast<const BatchedQuantizeParams *>(params_dev);
+  batched_amax_kernel<IType>
+      <<<dim3(kBatchedAmaxTiles, 1, num_tensors), kBatchedAmaxThreads, 0, stream>>>(
+          params, temp_amax_dev, rows, cols);
+
+  const dim3 grid(DIVUP(cols, static_cast<size_t>(kTileCols)),
+                  DIVUP(rows, static_cast<size_t>(kTileRows)), num_tensors);
+  const dim3 block(kThreads);
+  const size_t shmem = kPipelineStages * kStageRows * kTileCols * sizeof(IType);
+  auto kernel = quantize_4over6_batched_kernel<Cfg, E4M3_MAX, IType>;
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem);
+  kernel<<<grid, block, shmem, stream>>>(params, temp_amax_dev, rows, cols, scale_stride, noop);
+}
+
 }  // namespace quantize_4over6_kernel
 
 #endif  // FP4_TYPE_SUPPORTED

transformer_engine/common/include/transformer_engine/cast.h

@@ -125,6 +125,24 @@ void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor no
 void nvte_quantize_v2(const NVTETensor input, NVTETensor output,
                       const NVTEQuantizationConfig quant_config, cudaStream_t stream);
 
+/*! \brief Batched NVFP4 4over6 quantization of multiple same-shaped tensors.
+ *
+ *  Quantizes `num_tensors` tensors in two kernel launches (one batched amax
+ *  pass and one batched quantize pass) instead of two launches per tensor.
+ *  All tensors must share the same shape, dtype and NVFP4 4over6 configuration;
+ *  each output keeps its own scalar amax / global scale, so results are
+ *  bit-identical to calling the single-tensor path in a loop.
+ *
+ *  \param[in]      inputs          List of input tensors to be cast.
+ *  \param[in,out]  outputs         List of output quantized tensors.
+ *  \param[in]      quant_config    Quantization configuration.
+ *  \param[in]      num_tensors     Number of input and output tensors.
+ *  \param[in]      stream          CUDA stream used for the operation.
+ */
+void nvte_nvfp4_quantize_4over6_multi(const NVTETensor *inputs, NVTETensor *outputs,
+                                      const NVTEQuantizationConfig quant_config,
+                                      const size_t num_tensors, cudaStream_t stream);
+
 /*! \brief Casts input tensor to MXFP8. Additionally, reduces the input along columns.
  *         If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING,
  *         the block quantization (MXFP8) of the specified shape of the block will be used.

transformer_engine/pytorch/csrc/extensions.h

@@ -345,6 +345,10 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz
                                     const at::Tensor &rowwise_amax,
                                     const at::Tensor &columnwise_amax);
 
+std::vector<py::object> nvfp4_quantize_4over6_multi(const std::vector<at::Tensor> &tensors,
+                                                    py::handle quantizer,
+                                                    const py::object &outputs_py);
+
 py::object dequantize(const py::handle &input, DType otype);
 
 py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,