From 53d590c59a46c4b485fc3b27b59eddc16ce50141 Mon Sep 17 00:00:00 2001 From: David Kogan Date: Tue, 4 Aug 2026 22:09:21 -0400 Subject: [PATCH 1/3] Fuse the RHT into grouped NVFP4 quantize on non-SM100 architectures Outside the SM100 family split_quantize with RHT-enabled NVFP4 quantizers quantizes each split on its own, which runs two launches per split plus a separate Hadamard transform that materializes the transformed tensor to global memory. Complete the transposed output path in group_quantize_transpose with per-split direct stores, fold the 16-point random Hadamard transform into the columnwise read, and add nvte_group_quantize_with_colwise_rht so the dispatch can use one grouped launch pair per chunk. Chunk launch pairs by input bytes so the amax pass prefetches for the quantize pass. Drop the transposed staging buffers in shared memory, which the direct stores make dead. The transposed path requires 128-row-aligned splits and is reachable only through the columnwise RHT entry point, which refuses on the SM100 family, so generic grouped quantize keeps its existing refusal everywhere. Signed-off-by: David Kogan --- .../nvfp4/test_nvfp4_group_quantize.py | 74 +++++ transformer_engine/common/cast/cast.cu | 28 ++ .../nvfp4/group_quantize_transpose_nvfp4.cuh | 255 ++++++++++-------- transformer_engine/common/common.h | 5 + .../common/include/transformer_engine/cast.h | 22 ++ .../pytorch/csrc/extensions/cast.cpp | 97 ++++++- 6 files changed, 372 insertions(+), 109 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 0c3d3127d4..d3677a5f8f 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -274,3 +274,77 @@ def test_rht_split_quantize_matches_per_tensor_reference(quantize_mode: str) -> split_sections=split_sections, with_rht=True, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("N", [512, 2048]) +@pytest.mark.parametrize( + "split_sections", + [ + [128, 128, 128, 128], + [256, 512, 128, 384], + [0, 256, 256, 0], + [0, 128, 0, 256, 128, 0, 384, 128], + [192, 320], + [64, 64, 192, 192], + ], + ids=[ + "aligned", + "mixed_aligned", + "empty_ends", + "empty_mixed", + "unaligned", + "small_unaligned", + ], +) +def test_rht_split_quantize_grouped_matches_unfused(monkeypatch, split_sections, N) -> None: + # split_quantize folds the RHT into the columnwise pass in one grouped launch when + # every split is a multiple of 128 rows, and quantizes each split on its own + # otherwise. Empty splits are dropped from the grouped launch. Both routes have to + # produce the same bytes. + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + + def run(disable_grouped: bool): + monkeypatch.setenv("NVTE_NVFP4_DISABLE_GROUPED_RHT", "1" if disable_grouped else "0") + quantizers = [ + NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + ) + for _ in split_sections + ] + return [ + { + "rowwise_data": out._rowwise_data.view(dtype=torch.uint8).clone(), + "columnwise_data": out._columnwise_data.view(dtype=torch.uint8).clone(), + "amax_rowwise": out._amax_rowwise.clone(), + "amax_columnwise": out._amax_columnwise.clone(), + "rowwise_scale_inv": out._rowwise_scale_inv.clone(), + "columnwise_scale_inv": out._columnwise_scale_inv.clone(), + } + for out in tex.split_quantize(x, split_sections, quantizers) + ] + + torch.manual_seed(0) + unfused = run(True) + fused = run(False) + + x_splits = torch.split(x, split_sections) + for i, rows in enumerate(split_sections): + if rows == 0: + continue + for key in ("rowwise_data", "columnwise_data", "amax_rowwise", "amax_columnwise"): + torch.testing.assert_close(fused[i][key], unfused[i][key], atol=0.0, rtol=0.0) + # Scale buffers are allocated with padded shapes and neither route writes the + # padding, so compare only the region both of them define. + for key, columnwise in (("rowwise_scale_inv", False), ("columnwise_scale_inv", True)): + valid = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, columnwise) + torch.testing.assert_close( + fused[i][key][: valid[0], : valid[1]], + unfused[i][key][: valid[0], : valid[1]], + atol=0.0, + rtol=0.0, + ) diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 1e3c04573b..e959141db9 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -47,6 +47,34 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, dispatch::quantize_fwd_helper(input, output, quant_config, stream); } +void nvte_group_quantize_with_colwise_rht(const NVTETensor input, const NVTETensor *outputs, + const size_t *split_sections, const size_t num_tensors, + const uint16_t rht_sign_mask_t, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_with_colwise_rht); + using namespace transformer_engine; + + QuantizationConfig config_with_rht; + if (quant_config != nullptr) { + config_with_rht = *reinterpret_cast(quant_config); + } + config_with_rht.nvfp4_colwise_rht = true; + config_with_rht.nvfp4_rht_sign_mask_t = rht_sign_mask_t; + + auto *input_tensor = convertNVTETensorCheck(input); + std::vector output_list(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + output_list[i] = convertNVTETensorCheck(outputs[i]); + } + Tensor dummy_noop; + Tensor *noop = config_with_rht.noop_tensor != nullptr + ? convertNVTETensorCheck(config_with_rht.noop_tensor) + : &dummy_noop; + dispatch::nvfp4::group_quantize_transpose( + *input_tensor, noop, output_list, split_sections, num_tensors, &config_with_rht, stream); +} + 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/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index 91c6af26b5..13b95d25c6 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -17,6 +17,7 @@ #include #include "../../common.h" +#include "../../util/cuda_runtime.h" #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" @@ -166,15 +167,46 @@ constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4 * 8) / 4; // 256 // Number of threads (rowwise scaling) that span 32 banks (4-byte banks) of shared memory constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / 16 +// In-register 16-point Walsh-Hadamard transform with a random sign mask: +// v' = 0.25 * FWHT(v * signs). Same transform the RHT cast-fusion kernels +// apply, evaluated on the 16 rows a thread already holds for the columnwise +// block, so no separate transform pass over global memory is needed. +__device__ __forceinline__ void group_rht_16pt(float (&f)[16], const uint16_t sign_mask) { +#pragma unroll + for (int i = 0; i < 16; ++i) { + if ((sign_mask >> i) & 1) { + f[i] = -f[i]; + } + } +#pragma unroll + for (int stride = 1; stride < 16; stride <<= 1) { +#pragma unroll + for (int i = 0; i < 16; ++i) { + if ((i & stride) == 0) { + const float a = f[i]; + const float b = f[i + stride]; + f[i] = a + b; + f[i + stride] = a - b; + } + } + } +#pragma unroll + for (int i = 0; i < 16; ++i) { + f[i] *= 0.25f; + } +} + template + typename IType, bool USE_STOCHASTIC_ROUNDING, bool RETURN_TRANSPOSE, + bool COLWISE_RHT = false> __global__ void __launch_bounds__(THREADS_NUM) group_quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, nvfp4_scale_t *const scales_ptr, const float *noop, const size_t rows, const size_t cols, const size_t scale_stride, const size_t *rng_state, - MultiAmaxCastTransposeFusionArgs kernel_args) { + MultiAmaxCastTransposeFusionArgs kernel_args, + const uint16_t rht_sign_mask_t) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) constexpr bool NO_ACTIVATIONS_NOT_FP32_INPUT = (!COMPUTE_ACTIVATIONS) && (!std::is_same_v); @@ -217,7 +249,6 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; const size_t tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; const size_t tid_X_colwise = threadIdx.x; - const size_t tid_Y_t = tid_X_colwise; // const size_t tid_X_t = 0; const size_t thread_offset_Y_rowwise = tid_Y_rowwise; @@ -252,15 +283,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t buff_size_aligned_in = DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); constexpr size_t in_mem = buff_size_aligned_in; - constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; - constexpr size_t out_mem_colwise_data = buff_size_aligned_out; - constexpr size_t out_mem_rowwise_scales = 0; - extern __shared__ char dynamic_shmem[]; uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding @@ -271,12 +296,6 @@ __global__ void __launch_bounds__(THREADS_NUM) // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); fp4e2m1x2 *out_data_sh = reinterpret_cast(dshmem + in_mem); - fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - - nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -302,9 +321,22 @@ __global__ void __launch_bounds__(THREADS_NUM) float S_dec_rowwise = 1.0f; UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); - // TODO (zhongbo): colwise scaling disabled for now because of transpose float S_enc_colwise = 1.0f; float S_dec_colwise = 1.0f; + // Columnwise output bindings. The transposed data of a split lands in that + // split's own buffers, so these follow the tensor id rather than a shared + // output tensor map. + uint8_t *split_colwise_data_ptr = nullptr; + nvfp4_scale_t *split_colwise_scale_ptr = nullptr; + size_t split_colwise_scale_stride = 0; + if constexpr (RETURN_TRANSPOSE) { + amax_colwise_ptr = reinterpret_cast(kernel_args.colwise_amax_list[tensor_id]); + split_colwise_data_ptr = + reinterpret_cast(kernel_args.output_colwise_data_list[tensor_id]); + split_colwise_scale_ptr = + reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); + split_colwise_scale_stride = kernel_args.output_colwise_scale_stride[tensor_id]; + } if (amax_colwise_ptr != nullptr) { UpdateEncodeDecodeScaleFP32(amax_colwise_ptr, &S_enc_colwise, &S_dec_colwise); } else { @@ -331,7 +363,6 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t buff_offset_in = buff * BUFF_IN_SIZE; const size_t buff_offset_out = buff * BUFF_OUT_SIZE; - const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE; // for stages from 1 to STAGES - 1, we need to update the tensor id // skip updating tensor id if it's the last CTA, and some stages will be out of bounds @@ -345,8 +376,15 @@ __global__ void __launch_bounds__(THREADS_NUM) UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); split_rowwise_scale_ptr = reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); - // TODO (zhongbo): colwise scaling disabled for now because of transpose - // Skip fetching colwise amax pointer and scaling factor updates + if constexpr (RETURN_TRANSPOSE) { + amax_colwise_ptr = reinterpret_cast(kernel_args.colwise_amax_list[tensor_id]); + UpdateEncodeDecodeScaleFP32(amax_colwise_ptr, &S_enc_colwise, &S_dec_colwise); + split_colwise_data_ptr = + reinterpret_cast(kernel_args.output_colwise_data_list[tensor_id]); + split_colwise_scale_ptr = reinterpret_cast( + kernel_args.output_colwise_scale_inv_list[tensor_id]); + split_colwise_scale_stride = kernel_args.output_colwise_scale_stride[tensor_id]; + } } } @@ -379,13 +417,8 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; const size_t in_thread_offset_X = thread_offset_X_colwise; - const size_t out_t_thread_offset_Y = thread_offset_X_colwise; - const size_t out_t_thread_offset_X = 0 + it * BUFF_OUT_IT_OFFSET; - const size_t shmem_offset_base_colwise_in = buff_offset_in + in_thread_offset_Y * BUFF_IN_DIM_X + in_thread_offset_X; - const size_t shmem_offset_base_colwise_out_t = - buff_offset_out_t + out_t_thread_offset_Y * BUFF_OUT_T_DIM_X + out_t_thread_offset_X; block_amax = 0.0f; float in_compute_colwise[SCALE_DIM]; @@ -393,11 +426,29 @@ __global__ void __launch_bounds__(THREADS_NUM) // 1. Read/Compute elements. Find NVFP4-block AMAX if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { IType block_amax_f16 = static_cast(0.0f); + if constexpr (COLWISE_RHT) { + // Transform the 16-row column chunk in fp32, round back to IType so + // the values match what the unfused path materializes, and take the + // block amax after the transform. + float v[SCALE_DIM]; #pragma unroll - for (int i = 0; i < SCALE_DIM; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise_in + i * BUFF_IN_DIM_X; - in_colwise_IType[i] = in_sh[shmem_offset_colwise]; - block_amax_f16 = __hmax(block_amax_f16, __habs(in_colwise_IType[i])); + for (int i = 0; i < SCALE_DIM; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + i * BUFF_IN_DIM_X; + v[i] = static_cast(in_sh[shmem_offset_colwise]); + } + group_rht_16pt(v, rht_sign_mask_t); +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + in_colwise_IType[i] = static_cast(v[i]); + block_amax_f16 = __hmax(block_amax_f16, __habs(in_colwise_IType[i])); + } + } else { +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + i * BUFF_IN_DIM_X; + in_colwise_IType[i] = in_sh[shmem_offset_colwise]; + block_amax_f16 = __hmax(block_amax_f16, __habs(in_colwise_IType[i])); + } } block_amax = static_cast(block_amax_f16); } else { @@ -434,10 +485,20 @@ __global__ void __launch_bounds__(THREADS_NUM) const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_colwise); - // Store scaling factors through SHMEM - const size_t scale_idx_sh = - tid_Y_t * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + it; - out_colwise_scales_sh[scale_idx_sh] = S_dec_b_fp8; + // Transposed coordinates: this thread's column is a row of the split's + // columnwise output, and the 16-row chunk is 16 consecutive elements + // (8 bytes of FP4) along that row. + const size_t global_row_start = block_offset_Y + stage_offset_Y + it * SCALE_DIM; + const size_t local_row_start = global_row_start - split_start; + const size_t split_rows_t = split_end - split_start; + const bool colwise_in_bounds = + !col_out_of_bounds_colwise && (global_row_start + SCALE_DIM <= rows); + + if (colwise_in_bounds && split_colwise_scale_ptr != nullptr) { + const size_t scale_idx_t = + col_base_colwise * split_colwise_scale_stride + local_row_start / SCALE_DIM; + split_colwise_scale_ptr[scale_idx_t] = S_dec_b_fp8; + } // Compute "correct" per-block encoding scaling factor constexpr float float_max = detail::TypeExtrema::max; @@ -463,26 +524,14 @@ __global__ void __launch_bounds__(THREADS_NUM) } } - const int group = thread_lane / 16; - uint32_t val[2]; - uint32_t *regs_4x = reinterpret_cast(regs); - - // Helps reducing bank conflicts - switch (group) { - case 0: - val[0] = regs_4x[0]; - val[1] = regs_4x[1]; - break; - case 1: - val[0] = regs_4x[1]; - val[1] = regs_4x[0]; - - break; + // 16 FP4 values are 8 contiguous bytes of the split's columnwise data at + // (row = col_base_colwise, cols = [local_row_start, local_row_start+16)). + // 128-row-aligned splits keep this store 8-byte aligned. + if (colwise_in_bounds && split_colwise_data_ptr != nullptr) { + const size_t byte_offset = (col_base_colwise * split_rows_t + local_row_start) / 2; + *reinterpret_cast(split_colwise_data_ptr + byte_offset) = + *reinterpret_cast(regs); } - uint32_t *out_t_data_sh_as_uint32_t = - reinterpret_cast(&out_t_data_sh[shmem_offset_base_colwise_out_t]); - out_t_data_sh_as_uint32_t[group] = val[0]; // idx1 = (group + 0) % 2; - out_t_data_sh_as_uint32_t[(group + 1) & 1] = val[1]; // idx2 = (group + 1) % 2; } } @@ -688,46 +737,15 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t global_offset_Y = block_offset_Y + stage_offset_Y; const size_t global_offset_X = block_offset_X; - // TODO(zhongbo): add back when transpose is supported - // const size_t global_offset_Y_t = block_offset_Y_t; - // const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; - ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, reinterpret_cast(&out_data_sh[buff_offset_out])); - // TODO(zhongbo): add back when transpose is supported - // if constexpr (RETURN_TRANSPOSE) { - // ptx::cp_async_bulk_tensor_2d_shared_to_global( - // reinterpret_cast(&tensor_map_output_t), global_offset_X_t, - // global_offset_Y_t, reinterpret_cast(&out_t_data_sh[buff_offset_out_t])); - // } - // Create a "bulk async-group" out of the previous bulk copy operation. ptx::cp_async_bulk_commit_group(); } } // end of stages - // TODO(zhongbo): add back when transpose is supported - // Vectorized store scaling factors through SHMEM - // if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { - // using ScalesVec = Vec; - // const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; - // ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); - // const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; - // const size_t count = // number of scales in Y dimension of this chunk - // (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); - // nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - // constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); - // if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { - // // Fast path: vectorized store when destination is properly aligned - // scales_vec.store_to(dst); - // } else { - // // Safe path: element-wise store for tails or unaligned destinations - // scales_vec.store_to_elts(dst, 0, count); - // } - // } - destroy_barriers(mbar, is_master_thread); #else NVTE_DEVICE_ERROR("sm_100 or higher is required."); @@ -767,8 +785,28 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to // return the transposed data. bool return_transpose = output->has_columnwise_data(); - // forbid return transpose for now because group quantize transpose is not supported yet - NVTE_CHECK(!return_transpose, "Return transpose is not supported for group quantize transpose."); + const bool colwise_rht = quant_config ? quant_config->nvfp4_colwise_rht : false; + const uint16_t rht_sign_mask_t = quant_config ? quant_config->nvfp4_rht_sign_mask_t : 0; + NVTE_CHECK(!colwise_rht || return_transpose, + "Columnwise RHT fusion requires a columnwise (transposed) output."); + if (return_transpose) { + // Only the columnwise RHT entry point writes the transposed output. Generic grouped + // quantize keeps the refusal it has today, on every architecture. + NVTE_CHECK(colwise_rht, "Return transpose is not supported for group quantize transpose."); + // The SM100 family has its own RHT cast-fusion kernels and never reaches this path. + const int sm = transformer_engine::cuda::sm_arch(); + NVTE_CHECK(sm < 100 || sm > 110, + "Columnwise RHT group quantize is not supported on this architecture."); + // Each split writes its own transposed buffer, and a split boundary is only + // resolved once per 128-row chunk, so splits must be 128-row aligned. + for (size_t i = 0; i < num_tensors; ++i) { + NVTE_CHECK(split_sections[i] % CHUNK_DIM_Y == 0, + "Group quantize with transposed output requires splits that are multiples of ", + CHUNK_DIM_Y, ", but split ", i, " has ", split_sections[i], " rows."); + } + NVTE_CHECK(!output->row_scaled_nvfp4, + "Group quantize with transposed output does not support row-scaled NVFP4."); + } // output_List is contiguous in memory, so take the first tensor as the contiguous output auto output_contiguous = output->data; @@ -802,6 +840,16 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, reinterpret_cast(output_list[i]->amax.dptr); kernel_args.output_rowwise_scale_inv_list[kernel_args.num_tensors] = reinterpret_cast(output_list[i]->scale_inv.dptr); + if (return_transpose) { + kernel_args.colwise_amax_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->columnwise_amax.dptr); + kernel_args.output_colwise_data_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->columnwise_data.dptr); + kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->columnwise_scale_inv.dptr); + kernel_args.output_colwise_scale_stride[kernel_args.num_tensors] = + static_cast(output_list[i]->columnwise_scale_inv.shape[1]); + } // kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i]; kernel_args.split_sections_range[kernel_args.num_tensors + 1] = kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i]; @@ -860,37 +908,34 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); - constexpr size_t in_mem = buff_size_aligned_in; constexpr size_t out_data_mem = buff_size_aligned_out; - constexpr size_t out_data_transpose_mem = buff_size_aligned_out; - constexpr size_t out_scales_transpose_mem = buff_size_scales; - - constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; - constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + constexpr size_t dshmem_size = in_mem + out_data_mem + TMA_SHMEM_ALIGNMENT; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - auto kernel = - group_quantize_transpose_nvfp4_kernel; + TRANSFORMER_ENGINE_SWITCH_CONDITION( + return_transpose, RETURN_TRANSPOSE, - if constexpr (use_2d_quantization) { - NVTE_ERROR("2D quantization is not supported for group quantize transpose."); - } + TRANSFORMER_ENGINE_SWITCH_CONDITION(colwise_rht, COLWISE_RHT, { + auto kernel = group_quantize_transpose_nvfp4_kernel; + + if constexpr (use_2d_quantization) { + NVTE_ERROR("2D quantization is not supported for group quantize transpose."); + } - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - kernel<<>>(tensor_map_input, tensor_map_output, - scales_ptr, noop_ptr, rows, cols, - scale_stride, rng_state, kernel_args); - NVTE_CHECK_CUDA(cudaGetLastError()); - });); + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>( + tensor_map_input, tensor_map_output, scales_ptr, noop_ptr, rows, cols, scale_stride, + rng_state, kernel_args, rht_sign_mask_t); + NVTE_CHECK_CUDA(cudaGetLastError()); + }););); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..b2b7dc87d3 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -606,6 +606,11 @@ struct QuantizationConfig { bool use_fast_math = false; NVTENVFP44Over6Mode nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; bool nvfp4_4over6_err_use_fast_math = false; + // Not exposed through the attribute API. Folds the random Hadamard transform + // into the columnwise NVFP4 quantization pass, for architectures without the + // SM100-family RHT cast-fusion kernels. + bool nvfp4_colwise_rht = false; + uint16_t nvfp4_rht_sign_mask_t = 0; static constexpr size_t attr_sizes[] = { sizeof(uint8_t), // force_pow_2_scales diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 4d6d24ba65..ae184d1964 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -125,6 +125,28 @@ 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 Grouped NVFP4 quantization with the random Hadamard transform folded + * into the columnwise pass. + * + * One launch quantizes up to 64 row-splits of a contiguous input. Rowwise + * outputs are quantized from the raw input, columnwise outputs from + * RHT(split^T). Every output must already hold its rowwise amax and its + * post-RHT columnwise amax, and every split must be a multiple of 128 rows. + * + * \param[in] input Input BF16 tensor, splits stacked on rows. + * \param[in,out] outputs Output NVFP4 tensors, one per split. + * \param[in] split_sections Row count of each split. + * \param[in] num_tensors Number of splits. + * \param[in] rht_sign_mask_t Random sign mask of the transposed RHT. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_with_colwise_rht(const NVTETensor input, const NVTETensor *outputs, + const size_t *split_sections, const size_t num_tensors, + const uint16_t rht_sign_mask_t, + const NVTEQuantizationConfig quant_config, + 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/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 882481e0f4..868e3b8027 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1649,11 +1649,100 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections, quantizers, stream); } else { - for (size_t i = 0; i < num_tensors; ++i) { - if (input_list[i].numel() == 0) { - continue; + // Grouped path: one grouped post-RHT amax launch plus one grouped + // RHT-quantize launch per chunk, instead of two launches per split. It needs + // both usages, per-split post-RHT amaxes, and splits that are multiples of + // 128 rows, so ragged or rowwise-only cases take the per-split loop below. + constexpr size_t kMaxTensorsPerLaunch = 64; + // Input bytes per launch pair. Sized against L2 so the quantize pass reuses + // what the amax pass just pulled in. + constexpr size_t kChunkBytes = 8u << 20; + bool use_grouped = quantizer.with_post_rht_amax && quantizer.rowwise_usage && + quantizer.columnwise_usage && !quantizer.with_2d_quantization && + !quantizer.row_scaled_nvfp4 && + quantizer.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled && + !transformer_engine::getenv("NVTE_NVFP4_DISABLE_GROUPED_RHT"); + for (size_t i = 0; use_grouped && i < num_tensors; ++i) { + use_grouped = split_sections[i] % 128 == 0; + } + + // Empty splits hold no allocation, so they stay out of the grouped launch. + // They occupy no rows either, so dropping them does not move any offset. + std::vector grouped_idx; + if (use_grouped) { + grouped_idx.reserve(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + if (split_sections[i] != 0) { + grouped_idx.push_back(i); + } + } + } + + if (use_grouped && !grouped_idx.empty()) { + const size_t cols = input.size(input.ndim() - 1); + auto *input_base = reinterpret_cast(input.get_rowwise_data().data_ptr); + + QuantizationConfigWrapper grouped_config; + TensorWrapper te_rng_state; + if (quantizer.stochastic_rounding) { + grouped_config.set_stochastic_rounding(true); + constexpr size_t rng_elts_per_thread = 1024; + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); + auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); + auto rng_state = torch::empty({2}, opts); + philox_unpack(philox_args, static_cast(rng_state.data_ptr())); + te_rng_state = makeTransformerEngineTensor(rng_state); + grouped_config.set_rng_state(te_rng_state.data()); + } + + const size_t row_bytes = cols * sizeof(uint16_t); + const size_t rows_per_chunk = + std::max(1, kChunkBytes / std::max(1, row_bytes)); + + size_t chunk_start = 0; + size_t chunk_row0 = 0; + while (chunk_start < grouped_idx.size()) { + size_t chunk_n = 0; + size_t budget = 0; + while (chunk_start + chunk_n < grouped_idx.size() && chunk_n < kMaxTensorsPerLaunch) { + const size_t next = split_sections[grouped_idx[chunk_start + chunk_n]]; + if (chunk_n > 0 && budget + next > rows_per_chunk) { + break; + } + budget += next; + ++chunk_n; + } + size_t chunk_rows = 0; + std::vector chunk_outputs(chunk_n); + std::vector chunk_splits(chunk_n); + for (size_t i = 0; i < chunk_n; ++i) { + const size_t idx = grouped_idx[chunk_start + i]; + chunk_outputs[i] = output_list[idx].data(); + chunk_splits[i] = split_sections[idx]; + chunk_rows += chunk_splits[i]; + } + TensorWrapper input_chunk; + input_chunk.set_rowwise_data(input_base + chunk_row0 * cols * sizeof(uint16_t), + DType::kBFloat16, std::vector{chunk_rows, cols}); + + nvte_group_hadamard_transform_amax(input_chunk.data(), chunk_outputs.data(), + chunk_splits.data(), chunk_n, 0, + quantizer.rht_matrix_random_sign_mask_t, stream); + nvte_group_quantize_with_colwise_rht( + input_chunk.data(), chunk_outputs.data(), chunk_splits.data(), chunk_n, + quantizer.rht_matrix_random_sign_mask_t, grouped_config, stream); + chunk_start += chunk_n; + chunk_row0 += chunk_rows; + } + } else { + for (size_t i = 0; i < num_tensors; ++i) { + if (input_list[i].numel() == 0) { + continue; + } + quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt); } - quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt); } } } else { // NVFP4 quantize From 4158554fa5a2f94d4d718a73baaa7616cecc37e9 Mon Sep 17 00:00:00 2001 From: David Kogan Date: Wed, 5 Aug 2026 00:07:52 -0400 Subject: [PATCH 2/3] Require quantizer agreement before the grouped RHT launch The grouped launch configures itself from quantizers.front() alone: with_post_rht_amax, rowwise_usage, columnwise_usage, row_scaled_nvfp4, nvfp4_4over6_mode, stochastic_rounding and the RHT sign mask. Grouped callers build one independent NVFP4Quantizer per split, and the Python-side cross-expert validator does not cover these fields for NVFP4Quantizer, so a heterogeneous list could silently take the grouped path and apply the first split's settings to the rest. Require every quantizer to agree with the first on these fields before taking the grouped path, mirroring the scaling-mode check already in this function. A mismatch falls through to the existing per-split loop, which reads each quantizer on its own. Added a launch-based regression test: two quantizers differing only in columnwise_usage do not produce byte-different rowwise output on the shapes tested, since rowwise computation does not depend on columnwise_usage and TE allocates every split's columnwise buffer regardless of that flag, so the test checks which kernel launches rather than the output bytes. Signed-off-by: David Kogan --- .../nvfp4/test_nvfp4_group_quantize.py | 54 +++++++++++++++++++ .../pytorch/csrc/extensions/cast.cpp | 21 ++++++-- 2 files changed, 72 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index d3677a5f8f..b89e5447e7 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -348,3 +348,57 @@ def run(disable_grouped: bool): atol=0.0, rtol=0.0, ) + + +def _grouped_kernel_names(split_sections, N, columnwise_per_split): + # Kernel names, not output bytes: TE allocates every split's columnwise buffer + # regardless of that split's own columnwise_usage flag, and rowwise output does + # not depend on columnwise_usage, so a heterogeneous list that wrongly takes the + # grouped path can still produce byte-identical rowwise output to the per-split + # path on some shapes. Which kernel launched is the property that actually + # distinguishes the two routes; see the launch-count table in the PR body. + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + quantizers = [ + NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=columnwise_per_split[i], + with_rht=True, + with_post_rht_amax=True, + ) + for i in range(len(split_sections)) + ] + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + tex.split_quantize(x, split_sections, quantizers) + torch.cuda.synchronize() + return {e.key for e in prof.key_averages() if e.count > 0} + + +def _launched_grouped_kernel(names): + return any("GroupHadamardAmaxTma" in n or "group_quantize_transpose_nvfp4_kernel" in n + for n in names) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rht_split_quantize_grouped_kernel_engages_when_uniform() -> None: + # Positive control for the test below: with every split asking for the same + # usage, the grouped path is eligible and must be the one that runs. + split_sections = [128, 128, 128, 128] + names = _grouped_kernel_names(split_sections, 256, [True] * len(split_sections)) + assert _launched_grouped_kernel(names) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_rht_split_quantize_declines_grouped_on_mismatched_quantizers() -> None: + # The grouped launch reads its usage flags, with_post_rht_amax, row_scaled_nvfp4, + # nvfp4_4over6_mode, stochastic_rounding and the RHT sign mask from quantizers[0] + # alone, and GroupedLinear builds one independent NVFP4Quantizer per expert. + # rowwise_usage and columnwise_usage are the one axis TE's own cross-expert + # validator explicitly allows to differ, so this is the realistic case: split 1 + # asks for rowwise only, the rest ask for both. The grouped path must decline to + # the per-split loop, which reads each split's own quantizer. + split_sections = [128, 128, 128, 128] + columnwise = [i != 1 for i in range(len(split_sections))] + names = _grouped_kernel_names(split_sections, 256, columnwise) + assert not _launched_grouped_kernel(names) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 868e3b8027..56641f9b55 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1657,9 +1657,24 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, // Input bytes per launch pair. Sized against L2 so the quantize pass reuses // what the amax pass just pulled in. constexpr size_t kChunkBytes = 8u << 20; - bool use_grouped = quantizer.with_post_rht_amax && quantizer.rowwise_usage && - quantizer.columnwise_usage && !quantizer.with_2d_quantization && - !quantizer.row_scaled_nvfp4 && + // The grouped launch configures every split from quantizers.front() alone. + // GroupedLinear builds one independent NVFP4Quantizer per expert, and the + // Python-side cross-expert validator does not cover these fields for NVFP4, + // so require them to agree here or take the per-split loop below instead of + // silently applying the first split's settings to the rest. + const bool grouped_quantizers_uniform = + std::all_of(quantizers.begin(), quantizers.end(), [&](const NVFP4Quantizer *q) { + return q->with_post_rht_amax == quantizer.with_post_rht_amax && + q->rowwise_usage == quantizer.rowwise_usage && + q->columnwise_usage == quantizer.columnwise_usage && + q->row_scaled_nvfp4 == quantizer.row_scaled_nvfp4 && + q->nvfp4_4over6_mode == quantizer.nvfp4_4over6_mode && + q->stochastic_rounding == quantizer.stochastic_rounding && + q->rht_matrix_random_sign_mask_t == quantizer.rht_matrix_random_sign_mask_t; + }); + bool use_grouped = grouped_quantizers_uniform && quantizer.with_post_rht_amax && + quantizer.rowwise_usage && quantizer.columnwise_usage && + !quantizer.with_2d_quantization && !quantizer.row_scaled_nvfp4 && quantizer.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled && !transformer_engine::getenv("NVTE_NVFP4_DISABLE_GROUPED_RHT"); for (size_t i = 0; use_grouped && i < num_tensors; ++i) { From ccb687ed711a47f8cbe1f87d4271f70f4739bd83 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 04:09:02 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/nvfp4/test_nvfp4_group_quantize.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index b89e5447e7..4c216ed93d 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -376,8 +376,9 @@ def _grouped_kernel_names(split_sections, N, columnwise_per_split): def _launched_grouped_kernel(names): - return any("GroupHadamardAmaxTma" in n or "group_quantize_transpose_nvfp4_kernel" in n - for n in names) + return any( + "GroupHadamardAmaxTma" in n or "group_quantize_transpose_nvfp4_kernel" in n for n in names + ) @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)