diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index a65010ef02..0bc1b6c337 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -9,6 +9,8 @@ # Due to the structure of NVFP4Quantizer, we need to test the RHT functionality # together with the quantization functionality. +import os + import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer @@ -127,6 +129,34 @@ def check_quantization_nvfp4_versus_reference( ref_quantizer._apply_rht(x.t().contiguous()) if with_rht else x.t().contiguous() ) ref_amax_colwise_t = torch.max(torch.abs(x_t_for_amax)).to(torch.float32).view(1) + + # SM120/121 uses TE's native single-K=16 MMA Hadamard arithmetic. cuBLAS may + # choose a different reduction order (and applies the random-sign matrix in + # the ATen reference orientation), so use the unfused TE kernel as the exact + # reference for the fused TE kernel on these architectures. + if torch.cuda.get_device_capability() in ((12, 0), (12, 1)): + env_name = "NVTE_NVFP4_DISABLE_RHT_CAST_FUSION" + old_value = os.environ.get(env_name) + os.environ[env_name] = "1" + try: + native_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=False, + columnwise=True, + with_amax_reduction=False, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + native = native_quantizer(x) + finally: + if old_value is None: + os.environ.pop(env_name) + else: + os.environ[env_name] = old_value + qx_t_ref = unpack_fp4(native._columnwise_data.view(dtype=torch.uint8)) + sx_t_ref = native._columnwise_scale_inv.view(dtype=torch.uint8) + ref_amax_colwise_t = native._amax_columnwise else: qx_t_ref = None sx_t_ref = None diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py new file mode 100644 index 0000000000..9152c5dc13 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py @@ -0,0 +1,174 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Regression tests for the no-TMEM fused NVFP4 RHT path on SM120/SM121.""" + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import NVFP4Quantizer + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def _is_sm12x(device: int = 0) -> bool: + return torch.cuda.get_device_capability(device) in ((12, 0), (12, 1)) + + +def _native_unfused_columnwise(x: torch.Tensor, with_random_sign_mask: bool): + """Run TE's native K=16 RHT with cast fusion disabled.""" + + env_name = "NVTE_NVFP4_DISABLE_RHT_CAST_FUSION" + old_value = os.environ.get(env_name) + os.environ[env_name] = "1" + try: + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=False, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + return quantizer(x) + finally: + if old_value is None: + os.environ.pop(env_name) + else: + os.environ[env_name] = old_value + + +def _unpack_fp4(x: torch.Tensor) -> torch.Tensor: + unpacked = x.view(torch.uint8).repeat_interleave(2, dim=-1) + unpacked[..., 0::2] &= 0x0F + unpacked[..., 1::2] >>= 4 + return unpacked + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("with_random_sign_mask", [False, True]) +def test_sm12x_post_rht_amax_matches_native_k16(with_random_sign_mask: bool) -> None: + """The fused post-RHT amax must match TE's native K=16 MMA RHT.""" + + if not _is_sm12x(): + pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") + + torch.manual_seed(1234) + x = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16) + torch.manual_seed(5678) + expected = _native_unfused_columnwise(x, with_random_sign_mask) + + torch.manual_seed(5678) + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=False, + columnwise=True, + with_amax_reduction=False, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + out = quantizer(x) + + torch.testing.assert_close(out._amax_columnwise, expected._amax_columnwise, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("shape", [(128, 128), (256, 256)]) +@pytest.mark.parametrize("rowwise", [False, True]) +@pytest.mark.parametrize("with_random_sign_mask", [False, True]) +@pytest.mark.parametrize("seed", [1234, 2026]) +def test_sm12x_fused_rht_codes_and_scales_match_native_k16( + shape: tuple[int, int], rowwise: bool, with_random_sign_mask: bool, seed: int +) -> None: + """Fused RHT codes/scales must match TE's native K=16 MMA RHT path.""" + + if not _is_sm12x(): + pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") + + torch.manual_seed(seed) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(5678) + fused_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=True, + with_amax_reduction=False, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + fused = fused_quantizer(x) + + torch.manual_seed(5678) + expected = _native_unfused_columnwise(x, with_random_sign_mask) + + torch.testing.assert_close( + _unpack_fp4(fused._columnwise_data), + _unpack_fp4(expected._columnwise_data), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + fused._columnwise_scale_inv, + expected._columnwise_scale_inv, + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + fused._amax_columnwise, expected._amax_columnwise, atol=0.0, rtol=0.0 + ) + + +def _distributed_amax_worker(rank: int, world_size: int, init_file: str) -> None: + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + torch.manual_seed(4321 + rank) + x = torch.randn((128, 128), device=f"cuda:{rank}", dtype=torch.bfloat16) + x.mul_(rank + 1) + torch.manual_seed(5678) + expected_amax = _native_unfused_columnwise(x, with_random_sign_mask=True)._amax_columnwise + + torch.manual_seed(5678) + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=False, + columnwise=True, + with_amax_reduction=True, + amax_reduction_group=dist.group.WORLD, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=True, + ) + out = quantizer(x) + + dist.all_reduce(expected_amax, op=dist.ReduceOp.MAX) + torch.testing.assert_close(out._amax_columnwise, expected_amax, atol=0.0, rtol=0.0) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two CUDA devices are required") +def test_sm12x_post_rht_amax_reduction_is_global(tmp_path) -> None: + """ATen post-RHT amax must be computed before the distributed MAX reduction.""" + + if not all(_is_sm12x(device) for device in range(2)): + pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") + + init_file = os.fspath(tmp_path / "nvfp4_rht_amax_init") + mp.spawn(_distributed_amax_worker, args=(2, init_file), nprocs=2, join=True) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..532c760bf0 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -19,6 +19,7 @@ #include #include "../../common.h" +#include "../../hadamard_transform/hadamard_transform_utils.cuh" #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" @@ -30,6 +31,16 @@ namespace transformer_engine { namespace dispatch { namespace nvfp4 { +__device__ __forceinline__ void load_matrix_b_16x16_from_shared(uint32_t &b0, uint32_t &b1, + uint32_t &b2, uint32_t &b3, + const void *addr, uint32_t stride) { + asm volatile( + "wmma.load.b.sync.aligned.row.m16n16k16.shared::cta.bf16 " + "{%0,%1,%2,%3}, [%4], %5;\n" + : "=r"(b0), "=r"(b1), "=r"(b2), "=r"(b3) + : "l"(addr), "r"(stride)); +} + namespace rowwise_amax_kernel { using namespace ptx; @@ -317,11 +328,12 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / template + bool ROW_SCALED_NVFP4, bool RETURN_ROWWISE = true, bool APPLY_COLUMNWISE_RHT = false> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, const __grid_constant__ CUtensorMap tensor_map_output_t, + const __grid_constant__ CUtensorMap tensor_map_rht, nvfp4_scale_t *const scales_ptr, nvfp4_scale_t *const scales_t_ptr, const float *noop, const float *const amax_rowwise_ptr, @@ -408,6 +420,8 @@ __global__ void __launch_bounds__(THREADS_NUM) 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; + constexpr size_t out_mem_colwise_scales = + (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM * sizeof(nvfp4_scale_t); extern __shared__ char dynamic_shmem[]; uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); @@ -425,6 +439,9 @@ __global__ void __launch_bounds__(THREADS_NUM) 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 *rht_sh = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + + out_mem_rowwise_scales + out_mem_colwise_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; @@ -448,9 +465,16 @@ __global__ void __launch_bounds__(THREADS_NUM) // Initialize shared memory barrier with the number of threads participating in the barrier. #pragma nv_diag_suppress static_var_with_dynamic_init __shared__ alignas(8) uint64_t mbar[STAGES]; + __shared__ alignas(8) uint64_t mbar_rht[1]; initialize_barriers(mbar, is_master_thread); + if constexpr (APPLY_COLUMNWISE_RHT) { + initialize_barriers<1, THREADS_NUM>(mbar_rht, is_master_thread); + copy_2d_to_shared(rht_sh, &tensor_map_rht, 0, 0, 16 * 16 * sizeof(IType), &mbar_rht[0], + is_master_thread); + } + copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, shmem_buff_size, &mbar[0], is_master_thread); @@ -483,11 +507,79 @@ __global__ void __launch_bounds__(THREADS_NUM) // Wait for the data to have arrived ptx::mbarrier_wait_parity(&mbar[stage], 0); + if constexpr (APPLY_COLUMNWISE_RHT) { + ptx::mbarrier_wait_parity(&mbar_rht[0], 0); + + // SM120/121 have legacy warp MMA but no TMEM. Form sixteen 16x16 tiles from + // the TMA input stage and compute A @ H in registers. The col-major A load + // consumes the transposed TMA layout directly, avoiding an operand staging copy. + constexpr int kWarps = THREADS_NUM / THREADS_PER_WARP; + constexpr int kTilesX = BUFF_IN_DIM_X / SCALE_DIM; + constexpr int kTilesY = BUFF_DIM_Y / SCALE_DIM; + constexpr int kMmaTiles = kTilesX * kTilesY; + const int warp = threadIdx.x / THREADS_PER_WARP; + uint32_t b_frag[4]; + load_matrix_b_16x16_from_shared(b_frag[0], b_frag[1], b_frag[2], b_frag[3], rht_sh, + SCALE_DIM); + + for (int tile = warp; tile < kMmaTiles; tile += kWarps) { + const int tile_y = tile / kTilesX; + const int tile_x = tile % kTilesX; + uint32_t a_frag[4]; + uint32_t c_frag[4]; + uint32_t unused_amax = 0; + IType *tile_in = + &in_sh[buff_offset_in + tile_y * SCALE_DIM * BUFF_IN_DIM_X + tile_x * SCALE_DIM]; + load_matrix_16x16_from_shared(a_frag[0], a_frag[1], a_frag[2], a_frag[3], tile_in, + BUFF_IN_DIM_X); + mma_m16_n16_k16_b16_b16_b16_noacc( + a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag[0], b_frag[1], b_frag[2], b_frag[3], + c_frag[0], c_frag[1], c_frag[2], c_frag[3], unused_amax); + + // A WMMA lane owns two adjacent pairs in each of two rows. Quantize + // those fragments in registers and write both packed FP4 pairs + // directly to the transposed-output staging buffer. + const int lane = threadIdx.x % THREADS_PER_WARP; + const int lane_in_quad = lane & 3; + const int row_in_half = lane >> 2; +#pragma unroll + for (int row_half = 0; row_half < 2; ++row_half) { + const uint32_t c_lo = c_frag[row_half * 2]; + const uint32_t c_hi = c_frag[row_half * 2 + 1]; + __nv_bfloat162 lo = *reinterpret_cast(&c_lo); + __nv_bfloat162 hi = *reinterpret_cast(&c_hi); + float row_amax = + fmaxf(fmaxf(fabsf(__bfloat162float(lo.x)), fabsf(__bfloat162float(lo.y))), + fmaxf(fabsf(__bfloat162float(hi.x)), fabsf(__bfloat162float(hi.y)))); + row_amax = fmaxf(row_amax, __shfl_xor_sync(0xffffffff, row_amax, 1)); + row_amax = fmaxf(row_amax, __shfl_xor_sync(0xffffffff, row_amax, 2)); + const nvfp4_scale_t scale = compute_decoding_scaling_factor(row_amax, S_enc_colwise); + const float scale_inverse = fminf(1.0f / (static_cast(scale) * S_dec_colwise), + detail::TypeExtrema::max); + const float2 scale_inverse_2x{scale_inverse, scale_inverse}; + const uint64_t values = static_cast(c_lo) | (static_cast(c_hi) << 32); + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + const fp4e2m1x4 packed = + ptx::mul_cvt_bf16_to_fp4_4x(values, scale_inverse_2x, rbits); + const uint16_t packed_bits = *reinterpret_cast(&packed); + const int output_row = tile_x * SCALE_DIM + row_in_half + row_half * 8; + uint8_t *output_bytes = reinterpret_cast(out_t_data_sh) + buff_offset_out_t + + output_row * BUFF_OUT_T_DIM_X + tile_y * (SCALE_DIM / 2); + output_bytes[lane_in_quad] = static_cast(packed_bits); + output_bytes[lane_in_quad + 4] = static_cast(packed_bits >> 8); + if (lane_in_quad == 0) { + const size_t scale_idx = + output_row * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + tile_y; + out_colwise_scales_sh[scale_idx] = scale; + } + } + } + } float block_amax = 0.0f; // COLWISE scaling - if constexpr (RETURN_TRANSPOSE) { + if constexpr (RETURN_TRANSPOSE && !APPLY_COLUMNWISE_RHT) { #pragma unroll for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; @@ -562,18 +654,36 @@ __global__ void __launch_bounds__(THREADS_NUM) // 3. Scale elements fp4e2m1x4 regs[SCALE_DIM / 4]; + if constexpr (APPLY_COLUMNWISE_RHT && NO_ACTIVATIONS_NOT_FP32_INPUT) { + uint32_t *regs_8x = reinterpret_cast(regs); #pragma unroll - for (int e = 0; e < SCALE_DIM / 4; ++e) { - const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); - regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( - elts, block_scale_inverse_2x, rbits); - } else { - const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); - const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); - regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( - in01, in23, block_scale_inverse_2x, rbits); + for (int e = 0; e < SCALE_DIM / 8; ++e) { + const uint64_t elts03 = *reinterpret_cast(&in_colwise_IType[8 * e]); + const uint64_t elts47 = *reinterpret_cast(&in_colwise_IType[8 * e + 4]); + if constexpr (USE_STOCHASTIC_ROUNDING) { + const uint32_t rbits03 = get_rbits(rng, random_uint4, rnd_idx); + const uint32_t rbits47 = get_rbits(rng, random_uint4, rnd_idx); + regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_stochastic_rounding( + elts03, elts47, block_scale_inverse, rbits03, rbits47); + } else { + regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest(elts03, elts47, + block_scale_inverse); + } + } + } else { +#pragma unroll + for (int e = 0; e < SCALE_DIM / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); + regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); + const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); + regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } } } @@ -601,7 +711,7 @@ __global__ void __launch_bounds__(THREADS_NUM) } // ROWWISE scaling - { + if constexpr (RETURN_ROWWISE) { const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; #pragma unroll for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { @@ -818,9 +928,11 @@ __global__ void __launch_bounds__(THREADS_NUM) 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])); + if constexpr (RETURN_ROWWISE) { + 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])); + } if constexpr (RETURN_TRANSPOSE) { ptx::cp_async_bulk_tensor_2d_shared_to_global( @@ -853,6 +965,9 @@ __global__ void __launch_bounds__(THREADS_NUM) } destroy_barriers(mbar, is_master_thread); + if constexpr (APPLY_COLUMNWISE_RHT) { + destroy_barriers<1>(mbar_rht, is_master_thread); + } #else NVTE_DEVICE_ERROR("sm_100 or higher is required."); #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -1418,9 +1533,10 @@ __global__ void __launch_bounds__(THREADS_NUM) #endif // FP4_TYPE_SUPPORTED } // namespace quantize_transpose_kernel -template +template void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { + const QuantizationConfig *quant_config, cudaStream_t stream, + const Tensor *rht_matrix = nullptr) { #if FP4_TYPE_SUPPORTED using namespace quantize_transpose_kernel; using namespace ptx; @@ -1437,6 +1553,14 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, // Columnwise-only (no rowwise output) is supported on the optimized 2D path; the rowwise pass // and its store are gated out via the RETURN_ROWWISE template bool. const bool return_rowwise = output->has_data(); + if constexpr (apply_columnwise_rht) { + NVTE_CHECK(rht_matrix != nullptr, "Columnwise RHT requires an RHT matrix."); + NVTE_CHECK(return_transpose, "Columnwise RHT requires columnwise output."); + NVTE_CHECK(rht_matrix->dtype() == DType::kBFloat16, "Columnwise RHT matrix must be BF16."); + NVTE_CHECK( + rht_matrix->dim() == 2 && rht_matrix->shape()[0] == 16 && rht_matrix->shape()[1] == 16, + "Columnwise RHT matrix must have shape [16, 16]."); + } if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); @@ -1453,7 +1577,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, CheckOutputTensor(*output, "output", false); NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - NVTE_CHECK(return_rowwise || (return_transpose && use_2d_quantization), + NVTE_CHECK(return_rowwise || (return_transpose && (use_2d_quantization || apply_columnwise_rht)), "NVFP4 optimized kernel supports rowwise output (1D or 2D), or columnwise-only output " "with 2D quantization."); if (return_rowwise) { @@ -1523,6 +1647,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_output{}; alignas(64) CUtensorMap tensor_map_output_transpose{}; + alignas(64) CUtensorMap tensor_map_rht{}; create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, sizeof(IType) * 8); @@ -1535,6 +1660,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); } + if constexpr (apply_columnwise_rht) { + create_2D_tensor_map(tensor_map_rht, rht_matrix->data, 16, 16, 16, 16, 16, 0, + sizeof(IType) * 8); + } constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; constexpr size_t buff_size_aligned_in = @@ -1542,6 +1671,8 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, 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 rht_matrix_mem = 16 * 16 * sizeof(IType); + constexpr size_t rht_mem = DIVUP_TO_MULTIPLE(rht_matrix_mem, TMA_SHMEM_ALIGNMENT); constexpr size_t in_mem = buff_size_aligned_in; @@ -1551,7 +1682,8 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, 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; + const size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + + (apply_columnwise_rht ? rht_mem : 0) + TMA_SHMEM_ALIGNMENT; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, @@ -1559,13 +1691,23 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_rowwise, RETURN_ROWWISE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - // The 1D kernel always produces rowwise output (no RETURN_ROWWISE); the dispatch only - // routes columnwise-only requests here when use_2d_quantization is true. - auto kernel = quantize_transpose_nvfp4_kernel; - - if constexpr (use_2d_quantization) { + if constexpr (apply_columnwise_rht) { + auto kernel = + quantize_transpose_nvfp4_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, tensor_map_rht, + scales_ptr, scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, + rows, cols, scale_stride, scale_stride_transpose, rng_state); + } else if constexpr (use_2d_quantization) { + auto kernel = + quantize_transpose_nvfp4_2D_kernel; if (with_gemm_swizzled_scales) { kernel = quantize_transpose_nvfp4_2D_kernel< COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, @@ -1575,13 +1717,23 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; } + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, + scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, + scale_stride, scale_stride_transpose, rng_state); + } else { + auto kernel = quantize_transpose_nvfp4_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, tensor_map_rht, + scales_ptr, scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, + rows, cols, scale_stride, scale_stride_transpose, rng_state); } - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - kernel<<>>( - tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, - scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, - scale_stride, scale_stride_transpose, rng_state); }); }); });); diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 479922a9bf..36cabd9658 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -17,6 +17,8 @@ #include #include +#include "common/cast/core/common.cuh" +#include "common/cast/nvfp4/quantize_transpose_nvfp4.cuh" #include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" @@ -1314,6 +1316,22 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, NVTE_CHECK(m % hadamard_dimension == 0, "num_rows must be divisible by hadamard_dimension"); + // SM120/121 do not provide TMEM. Reuse the NVFP4 1D TMA pipeline and perform the + // 16-point columnwise RHT in registers, while leaving the SM100/110 UMMA/TMEM path below intact. + const int sm_arch = transformer_engine::cuda::sm_arch(transformer_engine::cuda::current_device()); + if (sm_arch == 120 || sm_arch == 121) { + Tensor noop; + if (output_.columnwise_data.dptr != nullptr) { + dispatch::nvfp4::quantize_transpose(input_, &noop, &output_, &quant_config, + stream, &hadamard_matrix_); + } else { + // RHT only affects the columnwise result. Keep rowwise-only quantization on + // the regular 1D scaling path rather than accidentally selecting 2D scaling. + dispatch::nvfp4::quantize_transpose(input_, &noop, &output_, &quant_config, stream); + } + return; + } + int k_tile_size = 1024; // Honor the output tensor's GEMM-swizzled-scales flag: when set, emit diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4fff3f92de..865d0f9432 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1908,10 +1908,15 @@ void NVFP4Quantizer::set_quantization_params(TensorWrapper* tensor) const { bool NVFP4Quantizer::is_eligible_for_rht_cast_fusion(const std::vector& shape, bool for_grouped_kernel) { + if (transformer_engine::getenv("NVTE_NVFP4_DISABLE_RHT_CAST_FUSION")) { + return false; + } const auto [rows, cols] = get_2d_dims(shape); const size_t row_align = for_grouped_kernel ? 128 : 64; - return rows % row_align == 0 && cols % 128 == 0 && transformer_engine::cuda::sm_arch() >= 100 && - transformer_engine::cuda::sm_arch() <= 110; + const int sm_arch = transformer_engine::cuda::sm_arch(); + const bool supported_arch = (sm_arch >= 100 && sm_arch <= 110) || + (!for_grouped_kernel && (sm_arch == 120 || sm_arch == 121)); + return rows % row_align == 0 && cols % 128 == 0 && supported_arch; } bool NVFP4Quantizer::is_eligible_for_2d_swizzle_fusion(const std::vector& shape) { @@ -1930,6 +1935,12 @@ bool nvfp4_emits_gemm_swizzled_scales(const NVFP4Quantizer& q, const std::vector return false; } if (q.with_rht) { + const int sm_arch = transformer_engine::cuda::sm_arch(); + // The SM12x no-TMEM RHT path reuses the compact-scale 1D TMA quantizer. + // Keep the existing post-quantize swizzle until that path gains an in-kernel swizzle epilogue. + if (sm_arch == 120 || sm_arch == 121) { + return false; + } return NVFP4Quantizer::is_eligible_for_rht_cast_fusion(shape); } // Plain (non-RHT) 2D quantize kernel can bake the swizzled layout for aligned shapes.