diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 1e3c04573b..c53dae0927 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -47,6 +47,99 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, dispatch::quantize_fwd_helper(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(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 in_list; + std::vector 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 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(out_list[i].amax.dptr); + } + + void *params_dev = nullptr; + float *temp_amax_dev = nullptr; + NVTE_CHECK_CUDA( + cudaMallocAsync(¶ms_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; + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + in_list[0].dtype(), IType, + quantize_4over6_kernel::launch_quantize_4over6_batched( + params_dev, temp_amax_dev, static_cast(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; diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 50776a3ed6..c29f1d3017 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/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 +__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(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(kBatchedAmaxTiles) * kBatchedAmaxThreads) { + local = fmaxf(local, fabsf(static_cast(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(&temp_amax[z]), __float_as_int(local)); + } + } +} + +// grid = (col_tiles, row_tiles, num_tensors). Rowwise output only. +template +__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(p.input); + fp4e2m1x2 *output = reinterpret_cast(p.output); + nvfp4_scale_t *scales = reinterpret_cast(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(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(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 +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(params_dev); + batched_amax_kernel + <<>>( + params, temp_amax_dev, rows, cols); + + const dim3 grid(DIVUP(cols, static_cast(kTileCols)), + DIVUP(rows, static_cast(kTileRows)), num_tensors); + const dim3 block(kThreads); + const size_t shmem = kPipelineStages * kStageRows * kTileCols * sizeof(IType); + auto kernel = quantize_4over6_batched_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); + kernel<<>>(params, temp_amax_dev, rows, cols, scale_stride, noop); +} + } // namespace quantize_4over6_kernel #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 4d6d24ba65..5ed803cdaa 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/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. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..f6dfc37c46 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/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 nvfp4_quantize_4over6_multi(const std::vector &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, diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_multi_quantize.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_multi_quantize.cpp new file mode 100644 index 0000000000..23acdbb135 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_multi_quantize.cpp @@ -0,0 +1,118 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../extensions.h" +#include "common.h" +#include "common/util/system.h" +#include "pybind.h" + +namespace transformer_engine::pytorch { + +namespace { + +std::vector get_tensor_shape(const TensorWrapper &tensor) { + const auto &shape = tensor.shape(); + return std::vector(shape.data, shape.data + shape.ndim); +} + +} // namespace + +std::vector nvfp4_quantize_4over6_multi(const std::vector &tensors, + py::handle quantizer, + const py::object &outputs_py) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(!tensors.empty(), "nvfp4_quantize_4over6_multi expects a non-empty tensor list."); + auto quantizer_cpp = convert_quantizer(quantizer); + NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), + "nvfp4_quantize_4over6_multi only supports NVFP4 quantizers."); + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled, + "nvfp4_quantize_4over6_multi requires a non-disabled 4over6 mode."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht && !nvfp4_quantizer_cpp->stochastic_rounding && + !nvfp4_quantizer_cpp->with_2d_quantization, + "Batched 4over6 requires non-RHT, non-stochastic-rounding, 1D quantization."); + NVTE_CHECK(nvfp4_quantizer_cpp->rowwise_usage && !nvfp4_quantizer_cpp->columnwise_usage, + "Batched 4over6 supports rowwise-only quantization."); + NVTE_CHECK(!nvfp4_quantizer_cpp->row_scaled_nvfp4, + "Batched 4over6 targets per-tensor-scaled (weight) tensors; row-scaled activations " + "should use whole-buffer quantization plus row slices."); + + const size_t num_tensors = tensors.size(); + const bool has_outputs = !outputs_py.is_none(); + py::list outs; + if (has_outputs) { + outs = py::cast(outputs_py); + NVTE_CHECK(py::len(outs) == num_tensors, + "nvfp4_quantize_4over6_multi: outputs must match tensors in length."); + } + auto stream = at::cuda::getCurrentCUDAStream(); + + std::vector keepalive; + std::vector inputs; + std::vector outputs; + std::vector out_py_list; + keepalive.reserve(num_tensors); + inputs.reserve(num_tensors); + outputs.reserve(num_tensors); + out_py_list.reserve(num_tensors); + + size_t rows = 0, cols = 0; + for (size_t i = 0; i < num_tensors; ++i) { + keepalive.push_back(tensors[i].contiguous()); + inputs.push_back(makeTransformerEngineTensor(keepalive.back())); + const bool use_given_output = has_outputs && !outs[i].is_none(); + if (use_given_output) { + // Quantize into the caller-provided workspace (weight cache path). + auto [out_cpp, out_py] = quantizer_cpp->convert_and_update_tensor(outs[i]); + outputs.push_back(std::move(out_cpp)); + out_py_list.push_back(std::move(out_py)); + } else { + const auto shape = get_tensor_shape(inputs.back()); + auto [out_cpp, out_py] = quantizer_cpp->create_tensor(shape, inputs.back().dtype()); + outputs.push_back(std::move(out_cpp)); + out_py_list.push_back(std::move(out_py)); + } + const size_t this_cols = keepalive.back().size(-1); + const size_t this_rows = keepalive.back().numel() / this_cols; + if (i == 0) { + rows = this_rows; + cols = this_cols; + } else { + NVTE_CHECK(this_rows == rows && this_cols == cols, + "nvfp4_quantize_4over6_multi requires same-shaped tensors, but tensor ", i, + " has shape (", this_rows, ", ", this_cols, ") vs (", rows, ", ", cols, ")."); + NVTE_CHECK(inputs.back().dtype() == inputs[0].dtype(), + "nvfp4_quantize_4over6_multi requires the same dtype for all tensors."); + } + } + NVTE_CHECK(cols % 16 == 0, "nvfp4_quantize_4over6_multi requires columns divisible by 16."); + + std::vector in_nvte(num_tensors), out_nvte(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + in_nvte[i] = inputs[i].data(); + out_nvte[i] = outputs[i].data(); + } + + QuantizationConfigWrapper quant_config; + quant_config.set_nvfp4_4over6_mode(nvfp4_quantizer_cpp->nvfp4_4over6_mode); + const auto err_use_fast_math = + transformer_engine::getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"); + if (err_use_fast_math) { + quant_config.set_nvfp4_4over6_err_use_fast_math(true); + } + + NVTE_SCOPED_GIL_RELEASE({ + nvte_nvfp4_quantize_4over6_multi(in_nvte.data(), out_nvte.data(), quant_config, num_tensors, + stream); + }); + return out_py_list; +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7e9d114be8..6051cfe2a8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -171,6 +171,8 @@ void init_router_bindings(pybind11::module &m) { void bind_quantize_with_amax_extensions(py::module_ &m) { m.def("nvfp4_quantize_with_amax", nvfp4_quantize_with_amax, py::arg("tensor"), py::arg("quantizer"), py::arg("rowwise_amax"), py::arg("columnwise_amax")); + m.def("nvfp4_quantize_4over6_multi", nvfp4_quantize_4over6_multi, py::arg("tensors"), + py::arg("quantizer"), py::arg("outputs") = py::none()); m.def("nvfp4_group_quantize_with_amax", nvfp4_group_quantize_with_amax, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"), diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 143fee7ba2..8c75a79ce5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -632,15 +632,37 @@ def forward( inp_view = inp.reshape(-1, in_features) inputmats: list if fp8 and not debug: - # Disable bulk allocation when CPU offloading is active: offloading skips small - # tensors (like scales), but bulk allocation shares storage across all tensors, - # so if scales can't be offloaded, nothing in the group can be offloaded. - inputmats = tex.split_quantize( - inp_view, - m_splits, - input_quantizers, - disable_bulk_allocation=cpu_offloading, - ) + q0 = input_quantizers[0] + if ( + isinstance(q0, NVFP4Quantizer) + and q0.row_scaled_nvfp4 + and not cpu_offloading + and q0.rowwise_usage + and not q0.columnwise_usage + and os.getenv("NVTE_ROW_SCALED_FUSED_QUANTIZE", "1") == "1" + ): + # Row-scaled NVFP4 quantization is row-local (per-row amax, + # block scales and 4/6 selection never cross row boundaries), + # so quantizing the whole buffer in one kernel launch is + # bit-identical to quantizing each expert split separately. + xq = q0(inp_view) + inputmats = [] + offset = 0 + for m in m_splits: + # Rowwise scale layout requires 128-row alignment (cuBLAS + # block-scaling), matching tex.split_quantize's allocation. + inputmats.append(xq.slice_rows(offset, offset + m, scale_row_multiple=128)) + offset += m + else: + # Disable bulk allocation when CPU offloading is active: offloading skips small + # tensors (like scales), but bulk allocation shares storage across all tensors, + # so if scales can't be offloaded, nothing in the group can be offloaded. + inputmats = tex.split_quantize( + inp_view, + m_splits, + input_quantizers, + disable_bulk_allocation=cpu_offloading, + ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -657,17 +679,50 @@ def forward( if fp8 or debug: weights_fp8 = [] update_ws = is_first_microbatch is None or is_first_microbatch - for i in range(num_gemms): - weight_fp8, new_workspaces[i] = quantize_weight( - tensor=weights[i], - quantizer=weight_quantizers[i], - workspace=weight_workspaces[i] if weight_workspaces else None, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - workspace_dtype=activation_dtype, - cache=cache_weight, + wq0 = weight_quantizers[0] + batched_weight_quantize = ( + fp8 + and not debug + and num_gemms > 1 + and cache_weight + and skip_fp8_weight_update is None + and isinstance(wq0, NVFP4Quantizer) + and wq0.nvfp4_use_4over6 + and not wq0.row_scaled_nvfp4 + and wq0.rowwise_usage + and not wq0.columnwise_usage + and all(w.shape == weights[0].shape for w in weights) + and all( + isinstance(q, NVFP4Quantizer) and q.nvfp4_use_4over6 for q in weight_quantizers ) - weights_fp8.append(weight_fp8) + and os.getenv("NVTE_NVFP4_BATCHED_WEIGHT_QUANTIZE", "1") == "1" + ) + if batched_weight_quantize and update_ws: + # Batched 4over6: quantize all expert weights in two kernel + # launches instead of two launches per expert. Each expert keeps + # its own scalar amax, so results are bit-identical to the loop. + ws_in = list(weight_workspaces) if weight_workspaces else [None] * num_gemms + weights_fp8 = tex.nvfp4_quantize_4over6_multi(weights, wq0, ws_in) + for i in range(num_gemms): + if ws_in[i] is None: + new_workspaces[i] = weights_fp8[i] + elif batched_weight_quantize and all( + ws is not None for ws in (weight_workspaces or []) + ): + # Cache hit: weights unchanged, reuse workspaces without quantizing. + weights_fp8 = list(weight_workspaces) + else: + for i in range(num_gemms): + weight_fp8, new_workspaces[i] = quantize_weight( + tensor=weights[i], + quantizer=weight_quantizers[i], + workspace=weight_workspaces[i] if weight_workspaces else None, + update_workspace=update_ws, + skip_update_flag=skip_fp8_weight_update, + workspace_dtype=activation_dtype, + cache=cache_weight, + ) + weights_fp8.append(weight_fp8) else: weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights] diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 09f040ba67..384eff5753 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -235,6 +235,75 @@ def restore_from_saved( self._amax_columnwise = tensors[5] return tensors[6:] + def slice_rows(self, start: int, end: int, *, scale_row_multiple: int = 0): + """Row-range view of a rowwise-quantized NVFP4 tensor (zero-copy). + + All rowwise fields are row-major, so slicing dim 0 yields valid views + sharing storage with this tensor. Only supported for rowwise-only + tensors; per-row amax (row-scaled) is sliced along, scalar amax is kept. + + If scale_row_multiple > 0 and the row count is not a multiple of it, + the rowwise scale is copied into a fresh buffer whose rows are padded + up to that multiple (TE rowwise scale layout requires 128-row + alignment for cuBLAS block-scaling); padding rows are left uninitialized + and must not be consumed, matching split_quantize's allocation behavior. + """ + if self._columnwise_data is not None: + raise RuntimeError("slice_rows only supports rowwise-only NVFP4 tensors") + num_rows = self._rowwise_data.shape[0] + if not (0 <= start <= end <= num_rows): + raise ValueError(f"slice_rows({start}, {end}) out of range for {num_rows} rows") + amax = self._amax_rowwise + if amax is not None and amax.numel() > 1: + amax = amax[start:end] + scale = self._rowwise_scale_inv[start:end] + rows = end - start + if scale_row_multiple > 0 and rows % scale_row_multiple != 0: + padded_rows = -(-rows // scale_row_multiple) * scale_row_multiple + padded = torch.empty( + (padded_rows,) + tuple(scale.shape[1:]), dtype=scale.dtype, device=scale.device + ) + padded[:rows].copy_(scale) + scale = padded + data = self._rowwise_data[start:end] + if type(self) is NVFP4TensorStorage: + return NVFP4TensorStorage( + data, + scale, + None, + None, + amax, + None, + self._fp4_dtype, + self._quantizer, + self._with_gemm_swizzled_scales, + fake_dtype=self._dtype, + row_scaled_nvfp4=self._row_scaled_nvfp4, + nvfp4_use_4over6=self._nvfp4_use_4over6, + nvfp4_e4m3_max=self._nvfp4_e4m3_max, + ) + from ..nvfp4_tensor import NVFP4Tensor + + parent_shape = self.size() + shape = (rows,) + tuple(parent_shape[1:]) + return NVFP4Tensor( + shape, + self._dtype, + rowwise_data=data, + rowwise_scale_inv=scale, + columnwise_data=None, + columnwise_scale_inv=None, + amax_rowwise=amax, + amax_columnwise=None, + fp4_dtype=self._fp4_dtype, + quantizer=self._quantizer, + with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, + row_scaled_nvfp4=self._row_scaled_nvfp4, + nvfp4_use_4over6=self._nvfp4_use_4over6, + nvfp4_e4m3_max=self._nvfp4_e4m3_max, + device=data.device, + ) + def get_data_tensors(self): """Get this Tensor's data.""" return self._rowwise_data, self._columnwise_data