diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 0c3d3127d4..4c216ed93d 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -274,3 +274,132 @@ 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, + ) + + +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/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..56641f9b55 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1649,11 +1649,115 @@ 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; + // 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) { + 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