From 67f24008fed2f5d61225ccdcb301c3db9950e94b Mon Sep 17 00:00:00 2001 From: ceo Date: Sun, 18 Oct 2026 19:25:26 +0900 Subject: [PATCH] [Common, PyTorch] Fuse NVFP4 scale swizzling on SM120 and SM121 On SM120/SM121 every NVFP4 quantization with optimize_for_gemm wrote compact scales and then ran the standalone swizzle kernel: the 2D in-kernel path was gated to SM100/110, and the SM12x RHT path rejected swizzled output. Write GEMM-swizzled scales from the generic 1D, columnwise-RHT 1D and tuned 1D kernels, enable the existing 2D path on SM120/SM121, and report 128x128 alignment for these paths in gemm_swizzled_scales_alignment so eager allocation and torch.compile metadata agree. Aligned groups of four scales are written with one 32-bit store. Row-scaled NVFP4 and 4over6 keep compact scales. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ceo --- .../cpp/operator/test_cast_nvfp4_transpose.cu | 110 +++++++++++++++++- .../test_nvfp4_rht_quantize_swizzle_fusion.py | 20 ++-- .../nvfp4/test_nvfp4_sm12x_swizzle_fusion.py | 110 ++++++++++++++++++ .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 83 +++++++++---- .../quantize_transpose_nvfp4_tuned_1D.cuh | 50 ++++++-- .../common/cast/nvfp4/swizzle.cuh | 26 +++++ ...cast_col_hadamard_transform_cast_fusion.cu | 3 - transformer_engine/pytorch/csrc/quantizer.cpp | 26 ++++- 8 files changed, 373 insertions(+), 55 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_sm12x_swizzle_fusion.py diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index c2529bd419f..cafcb035347 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -1681,14 +1681,15 @@ TEST(NVFP4UE5M3ReferenceTest, Grouped) { } #endif -TEST(CastNVFP4RHTTest, RejectsSwizzledScalesOnSM12x) { +TEST(CastNVFP4RHTTest, RejectsUnalignedSwizzledScalesOnSM12x) { const int sm_arch = getDeviceComputeCapability(); if (sm_arch != 120 && sm_arch != 121) { GTEST_SKIP(); } - Tensor input("input", std::vector{128, 128}, DType::kBFloat16); - Tensor output("output", std::vector{128, 128}, DType::kFloat4E2M1, + // In-kernel GEMM-swizzled scales require 128-aligned rows and cols. + Tensor input("input", std::vector{64, 128}, DType::kBFloat16); + Tensor output("output", std::vector{64, 128}, DType::kFloat4E2M1, false, true, NVTE_NVFP4_1D_SCALING); Tensor hadamard("hadamard", std::vector{16, 16}, DType::kBFloat16); output.set_with_gemm_swizzled_scales(true); @@ -1696,9 +1697,106 @@ TEST(CastNVFP4RHTTest, RejectsSwizzledScalesOnSM12x) { try { nvte_quantize_with_hadamard_transform(input.data(), output.data(), hadamard.data(), nullptr, nullptr); - FAIL() << "Expected the SM12x RHT path to reject swizzled scales"; + FAIL() << "Expected the SM12x RHT path to reject swizzled scales for 64 rows"; } catch (const std::runtime_error& e) { - EXPECT_NE(std::string(e.what()).find("does not support GEMM-swizzled scales"), - std::string::npos); + EXPECT_NE(std::string(e.what()).find("require rows and cols to be multiples of 128"), + std::string::npos) << e.what(); + } +} + +TEST(CastNVFP4RHTTest, SwizzledScalesMatchCompactOnSM12x) { + const int sm_arch = getDeviceComputeCapability(); + if (sm_arch != 120 && sm_arch != 121) { + GTEST_SKIP(); + } + + constexpr size_t rows = 256; + constexpr size_t cols = 384; + const std::vector shape{rows, cols}; + + Tensor input("input", shape, DType::kBFloat16); + fillUniform(&input); + const bf16 *input_data = input.rowwise_cpu_dptr(); + float amax = 0.0f; + for (size_t i = 0; i < rows * cols; ++i) { + amax = fmaxf(amax, fabsf(static_cast(input_data[i]))); + } + + // Orthonormal 16-point Sylvester Hadamard matrix: H[i][j] = (-1)^popcount(i & j) / 4. + Tensor hadamard("hadamard", std::vector{16, 16}, DType::kBFloat16); + bf16 *hadamard_data = hadamard.rowwise_cpu_dptr(); + for (unsigned i = 0; i < 16; ++i) { + for (unsigned j = 0; j < 16; ++j) { + hadamard_data[i * 16 + j] = + static_cast(__builtin_popcount(i & j) % 2 ? -0.25f : 0.25f); + } + } + hadamard.from_cpu(); + + // Rowwise scales form a [rows, cols/16] matrix; columnwise scales belong to the transposed + // data and form a [cols, rows/16] matrix. With 128-aligned dims neither buffer is padded. + const auto dims = get_scale_tensor_dims(rows, cols, 1, 16); + const auto dims_t = get_scale_tensor_dims(cols, rows, 1, 16); + for (const auto &d : {dims, dims_t}) { + ASSERT_EQ(d[0], d[2]); + ASSERT_EQ(d[1], d[3]); + } + + // Permutes a compact [M, K] scale matrix into cuBLAS's block-scaling layout (128x4 tiles + // stored as 32 rows of 16 bytes), as compute_ref_swizzle<128, 4, true> in test_swizzle.cu. + const auto swizzle_ref = [](const uint8_t *scales, const size_t M, const size_t K) { + std::vector out(M * K); + for (size_t m = 0; m < M; ++m) { + for (size_t k = 0; k < K; ++k) { + const size_t tile_offset = (m / 128) * 128 * K + (k / 4) * 512; + out[tile_offset + (m % 32) * 16 + (m % 128) / 32 * 4 + k % 4] = scales[m * K + k]; + } + } + return out; + }; + const auto check_direction = [&](const std::string &name, const fp4e2m1 *data, + const fp4e2m1 *swizzled_data, const uint8_t *scales, + const uint8_t *swizzled_scales, const size_t M, + const size_t K) { + compareResults(name + "_data", reinterpret_cast(swizzled_data), + reinterpret_cast(data), rows * cols / 2); + // Buffers start zeroed, so a block skipped by both kernels would otherwise go unnoticed. + ASSERT_TRUE(std::none_of(scales, scales + M * K, [](uint8_t s) { return s == 0; })) + << name << " compact scales contain unwritten (zero) entries"; + compareResults(name + "_scales", swizzled_scales, swizzle_ref(scales, M, K).data(), + M * K); + }; + + for (const bool rowwise : {false, true}) { + SCOPED_TRACE(rowwise ? "rowwise+columnwise" : "columnwise-only"); + Tensor compact("compact", shape, DType::kFloat4E2M1, rowwise, true, + NVTE_NVFP4_1D_SCALING); + Tensor swizzled("swizzled", shape, DType::kFloat4E2M1, rowwise, true, + NVTE_NVFP4_1D_SCALING); + swizzled.set_with_gemm_swizzled_scales(true); + for (Tensor *output : {&compact, &swizzled}) { + if (rowwise) { + output->set_amax(amax); + } + // Each RHT output sums 16 terms of +-x/4, so 4 * amax bounds the columnwise amax. + output->set_tensor_amax_columnwise(4.0f * amax); + nvte_quantize_with_hadamard_transform(input.data(), output->data(), hadamard.data(), + nullptr, nullptr); + } + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + compact.to_cpu(); + swizzled.to_cpu(); + + check_direction("columnwise", compact.columnwise_cpu_dptr(), + swizzled.columnwise_cpu_dptr(), + compact.columnwise_cpu_scale_inv_ptr(), + swizzled.columnwise_cpu_scale_inv_ptr(), dims_t[0], dims_t[1]); + if (rowwise) { + check_direction("rowwise", compact.rowwise_cpu_dptr(), + swizzled.rowwise_cpu_dptr(), + compact.rowwise_cpu_scale_inv_ptr(), + swizzled.rowwise_cpu_scale_inv_ptr(), dims[0], dims[1]); + } } } diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_swizzle_fusion.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_swizzle_fusion.py index 67dc0b47b88..c4d86de4bbd 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_swizzle_fusion.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_swizzle_fusion.py @@ -208,23 +208,25 @@ def _make_swizzle_fusion_quantizer() -> NVFP4Quantizer: @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( - "M, N, expected_swizzled", + "M, N, expected_swizzled, expected_swizzled_sm12x", [ # Eligible: rows%64==0 AND cols%128==0 -> the fused RHT cast-fusion # kernel can bake the GEMM-swizzled SF in directly. - (64, 128, True), - (128, 256, True), + (64, 128, True, False), + (128, 256, True, True), # Ineligible (cols%128 != 0) -> the fused kernel is gated off, so # create_tensor reports swizzled=False (the fused kernel did not bake # it in). This is the case hit in production at irregular shapes like # (8192, 11328) where 11328%128==64. - (64, 144, False), - (128, 144, False), + (64, 144, False, False), + (128, 144, False, False), # Ineligible (rows%64 != 0) -> same gating. - (48, 128, False), + (48, 128, False, False), ], ) -def test_nvfp4_rht_swizzle_fusion_shape_gate(M: int, N: int, expected_swizzled: bool) -> None: +def test_nvfp4_rht_swizzle_fusion_shape_gate( + M: int, N: int, expected_swizzled: bool, expected_swizzled_sm12x: bool +) -> None: """The fused-kernel shape gate must reflect cast-fusion eligibility. Only ``row_cast_col_hadamard_transform_cast_fusion.cu`` can bake @@ -243,6 +245,10 @@ def test_nvfp4_rht_swizzle_fusion_shape_gate(M: int, N: int, expected_swizzled: ineligible shapes too and flips the flag back to True; see ``test_nvfp4_rht_swizzle_fusion_end_to_end_swizzled``.) """ + # The SM12x 1D path needs full 128-row scale tiles; TMEM architectures + # retain their existing 64-row eligibility. + if torch.cuda.get_device_capability() in ((12, 0), (12, 1)): + expected_swizzled = expected_swizzled_sm12x quantizer = _make_swizzle_fusion_quantizer() out = quantizer.make_empty([M, N], dtype=torch.bfloat16) assert out._with_gemm_swizzled_scales is expected_swizzled, ( diff --git a/tests/pytorch/nvfp4/test_nvfp4_sm12x_swizzle_fusion.py b/tests/pytorch/nvfp4/test_nvfp4_sm12x_swizzle_fusion.py new file mode 100644 index 00000000000..bf43bcc23c0 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_sm12x_swizzle_fusion.py @@ -0,0 +1,110 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""SM12x scale-layout fidelity and allocation gates for fused NVFP4 quantization.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import NVFP4Quantizer + +from nvfp4_utils import swizzle_nvfp4_scale + + +@pytest.fixture(autouse=True) +def require_sm12x(): + """These kernels require an SM120 or SM121 device.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if torch.cuda.get_device_capability() not in ((12, 0), (12, 1)): + pytest.skip("SM120/SM121 is required") + available, reason = te.is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + + +@pytest.mark.parametrize("shape", [(128, 128), (256, 512), (384, 256)]) +@pytest.mark.parametrize("direction", ["row", "column", "both"]) +@pytest.mark.parametrize("kind", ["1d", "2d", "rht"]) +@pytest.mark.parametrize("stochastic_rounding", [False, True]) +@pytest.mark.parametrize("scale_dtype", [te.DType.kFloat8E4M3, te.DType.kFloat8UE5M3]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_sm12x_swizzled_outputs_match_compact_reference( + shape, direction, kind, stochastic_rounding, scale_dtype, dtype +): + """Check FP4 bytes, both amaxes and an independent scale permutation.""" + if kind == "rht" and dtype != torch.bfloat16: + pytest.skip("NVFP4 RHT only supports BF16 input") + torch.manual_seed(1234) + x = torch.randn(shape, device="cuda", dtype=dtype) + compact = NVFP4Quantizer( + rowwise=direction != "column", + columnwise=direction != "row", + with_2d_quantization=kind == "2d", + with_rht=kind == "rht", + with_post_rht_amax=kind == "rht", + with_random_sign_mask=False, + stochastic_rounding=stochastic_rounding, + scale_dtype=scale_dtype, + ) + compact.optimize_for_gemm = False + fused = compact.copy() + fused.optimize_for_gemm = True + # Test allocation separately: the fallback also sets this flag after quantization. + allocated = fused.make_empty(shape, dtype=x.dtype) + assert allocated._with_gemm_swizzled_scales + + torch.manual_seed(5678) + expected = compact(x) + torch.manual_seed(5678) + actual = fused(x) + for columnwise, prefix in [(False, "rowwise"), (True, "columnwise")]: + ref_data = getattr(expected, f"_{prefix}_data") + if ref_data is None: + assert getattr(actual, f"_{prefix}_data") is None + continue + assert torch.equal(getattr(actual, f"_{prefix}_data"), ref_data) + suffix = "columnwise" if columnwise else "rowwise" + assert torch.equal(getattr(actual, f"_amax_{suffix}"), getattr(expected, f"_amax_{suffix}")) + reference_scales = swizzle_nvfp4_scale( + *shape, getattr(expected, f"_{prefix}_scale_inv"), columnwise=columnwise + ) + assert torch.equal( + getattr(actual, f"_{prefix}_scale_inv").view(torch.uint8), + reference_scales.view(torch.uint8), + ) + + +@pytest.mark.parametrize("kind", ["1d", "2d", "rht"]) +@pytest.mark.parametrize( + "shape,eligible", [((128, 128), True), ((64, 128), False), ((128, 192), False)] +) +def test_sm12x_in_kernel_swizzle_shape_gate(kind, shape, eligible): + """Partial scale tiles must retain compact allocation and the existing fallback.""" + q = NVFP4Quantizer( + with_2d_quantization=kind == "2d", + with_rht=kind == "rht", + with_post_rht_amax=kind == "rht", + ) + q.optimize_for_gemm = True + out = q.make_empty(shape, dtype=torch.bfloat16) + assert out._with_gemm_swizzled_scales is eligible + + +@pytest.mark.parametrize("with_rht", [False, True]) +@pytest.mark.parametrize("special_recipe", ["row_scaled", "4over6"]) +def test_sm12x_special_recipes_keep_compact_allocation(with_rht, special_recipe): + """Unsupported recipes must be rejected before selecting the RHT scale layout.""" + q = NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_rht=with_rht, + with_post_rht_amax=with_rht, + row_scaled_nvfp4=special_recipe == "row_scaled", + nvfp4_use_4over6=special_recipe == "4over6", + ) + q.optimize_for_gemm = True + out = q.make_empty((128, 128), dtype=torch.bfloat16) + assert not out._with_gemm_swizzled_scales diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index f83564b35fa..3e32259e432 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -667,7 +667,8 @@ __device__ __forceinline__ void quantize_columnwise_rht_stage( template + bool ROW_SCALED_NVFP4, bool RETURN_ROWWISE = true, bool APPLY_COLUMNWISE_RHT = false, + bool WITH_GEMM_SWIZZLED_SCALES = 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, @@ -679,6 +680,8 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t cols, const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(!(WITH_GEMM_SWIZZLED_SCALES && ROW_SCALED_NVFP4), + "Row-scaled NVFP4 does not support GEMM-swizzled scales."); constexpr bool NO_ACTIVATIONS_NOT_FP32_INPUT = (!COMPUTE_ACTIVATIONS) && (!std::is_same_v); @@ -1147,7 +1150,15 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t scales_offset_Y = scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; const size_t scales_offset_X = scales_offset_X_rowwise; - const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + size_t scale_idx_global; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // Write the scale directly into the cuBLAS GEMM-swizzled layout (as the 2D kernel + // does), so no separate swizzle pass is needed before the GEMM. + scale_idx_global = + swizzle::gemm_swizzled_scale_idx(scales_offset_Y, scales_offset_X, SFs_per_row); + } else { + scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + } const bool rowwise_scale_is_within_bounds_Y = (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; @@ -1230,20 +1241,30 @@ __global__ void __launch_bounds__(THREADS_NUM) // 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); - ScaleType *dst = &scales_t_ptr[scale_idx_global]; - constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(ScaleType); - 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); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // The swizzled layout scatters the contiguous columnwise scales in groups of four, so emit + // each group at its swizzled offset (same scheme as the 2D kernel). The transposed scale + // matrix has rows / SCALE_DIM columns (exact: the swizzled path requires 128-aligned dims). + const size_t col_length_t = rows / SCALE_DIM; + swizzle::store_gemm_swizzled_scales(scales_t_ptr, &out_colwise_scales_sh[scale_idx_sh], + scales_offset_Y_t, scales_offset_X_t, col_length_t, + static_cast(count)); } else { - // Safe path: element-wise store for tails or unaligned destinations - scales_vec.store_to_elts(dst, 0, count); + using ScalesVec = Vec; + 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; + ScaleType *dst = &scales_t_ptr[scale_idx_global]; + constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(ScaleType); + 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); + } } } @@ -1791,18 +1812,17 @@ __global__ void __launch_bounds__(THREADS_NUM) 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); if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - // The swizzled layout scatters the contiguous columnwise scales, so the - // vectorized store cannot be used. Emit each scale at its swizzled offset. + // The swizzled layout scatters the contiguous columnwise scales in groups + // of four, so the full-chunk vectorized store cannot be used. Emit each + // group of four at its swizzled offset with one 4-byte store. // The transposed scale matrix has `rows / SCALE_DIM` (= M/16) column tiles // (exact because the swizzled path requires 128-aligned dims). Read // SCALE_DIM by value; passing the namespace-scope constexpr by reference // (e.g. via DIVUP) would ODR-use it and fail to compile in device code. const size_t col_length_t = rows / SCALE_DIM; - for (size_t k = 0; k < count; ++k) { - const size_t off = swizzle::gemm_swizzled_scale_idx(scales_offset_Y_t, - scales_offset_X_t + k, col_length_t); - scales_t_ptr[off] = out_colwise_scales_sh[scale_idx_sh + k]; - } + swizzle::store_gemm_swizzled_scales(scales_t_ptr, &out_colwise_scales_sh[scale_idx_sh], + scales_offset_Y_t, scales_offset_X_t, col_length_t, + static_cast(count)); } else { using ScalesVec = Vec; ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); @@ -1885,12 +1905,11 @@ void quantize_transpose_impl(const Tensor &input, const Tensor *noop, Tensor *ou "Row-scaled NVFP4 does not support disabling second-level scaling."); NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), "Row-scaled NVFP4 quantization does not produce columnwise output."); - // In-kernel GEMM-swizzled scale output is only implemented on the 2D quantization - // path; every other configuration must emit compact scales. + // In-kernel GEMM-swizzled scale output is implemented by every kernel below (2D, 1D, + // columnwise-RHT 1D) and by the tuned 1D kernel; row-scaled NVFP4 must emit compact scales. const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; - NVTE_CHECK(!with_gemm_swizzled_scales || use_2d_quantization, - "In-kernel GEMM-swizzled NVFP4 scales are only supported on the 2D quantization path; " - "other paths must emit scales in compact format."); + NVTE_CHECK(!with_gemm_swizzled_scales || !row_scaled_nvfp4, + "Row-scaled NVFP4 quantization does not support GEMM-swizzled scales."); if (return_transpose) { NVTE_CHECK(output->has_columnwise_data(), "NVFP4 transposed output tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), @@ -1996,6 +2015,14 @@ void quantize_transpose_impl(const Tensor &input, const Tensor *noop, Tensor *ou ScaleType, USE_STOCHASTIC_ROUNDING, RETURN_TRANSPOSE, ROW_SCALED_NVFP4, RETURN_ROWWISE, /*APPLY_COLUMNWISE_RHT=*/true>; + if constexpr (!ROW_SCALED_NVFP4) { + if (with_gemm_swizzled_scales) { + kernel = quantize_transpose_nvfp4_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, ScaleType, USE_STOCHASTIC_ROUNDING, + RETURN_TRANSPOSE, ROW_SCALED_NVFP4, RETURN_ROWWISE, + /*APPLY_COLUMNWISE_RHT=*/true, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + } + } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); kernel<<>>( @@ -2026,6 +2053,14 @@ void quantize_transpose_impl(const Tensor &input, const Tensor *noop, Tensor *ou auto kernel = quantize_transpose_nvfp4_kernel; + if constexpr (!ROW_SCALED_NVFP4) { + if (with_gemm_swizzled_scales) { + kernel = quantize_transpose_nvfp4_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, ScaleType, USE_STOCHASTIC_ROUNDING, + RETURN_TRANSPOSE, ROW_SCALED_NVFP4, /*RETURN_ROWWISE=*/true, + /*APPLY_COLUMNWISE_RHT=*/false, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + } + } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); kernel<<>>( diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 3c24f9a4111..7e6fad618e3 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -23,6 +23,7 @@ #include "../../../util/ptx_arch_spec.cuh" #include "../../../utils.cuh" #include "../core_nvfp4.cuh" +#include "../swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -372,7 +373,7 @@ __device__ __forceinline__ void rowwise_scaling( } template + bool RETURN_TRANSPOSE, bool ROW_SCALED_NVFP4, bool WITH_GEMM_SWIZZLED_SCALES = false> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D_kernel( const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, @@ -381,6 +382,8 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D const float *const amax_colwise_ptr, const size_t rows, const size_t cols, const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(!(WITH_GEMM_SWIZZLED_SCALES && ROW_SCALED_NVFP4), + "Row-scaled NVFP4 does not support GEMM-swizzled scales."); if (noop != nullptr && noop[0] == 1.0f) { return; } @@ -650,10 +653,18 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D for (size_t row = threadIdx.x; row < TunableConfig::CHUNK_DIM_Y; row += THREADS_NUM) { const size_t row_global = scales_block_offset_Y_rowwise + row; if (row_global < rows) { - ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); - const size_t scale_idx_global = - row_global * scale_stride + scales_block_offset_X_rowwise; - scales_vec.store_to_elts(&scales_ptr[scale_idx_global], 0, count); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // Emit each group of four scales at its cuBLAS GEMM-swizzled offset (same + // scheme as the 2D kernel), so no separate swizzle pass is needed before the GEMM. + const size_t col_length = cols / SCALE_DIM; + swizzle::store_gemm_swizzled_scales(scales_ptr, sSFrowwise[row], row_global, + scales_block_offset_X_rowwise, col_length, count); + } else { + ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); + const size_t scale_idx_global = + row_global * scale_stride + scales_block_offset_X_rowwise; + scales_vec.store_to_elts(&scales_ptr[scale_idx_global], 0, count); + } } } } @@ -668,10 +679,17 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D row_tr += THREADS_NUM) { const size_t row_tr_global = scales_block_offset_Y_tr + row_tr; if (row_tr_global < cols) { - ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); - const size_t scale_idx_global = - row_tr_global * scale_stride_t + scales_block_offset_X_tr; - scales_vec.store_to_elts(&scales_t_ptr[scale_idx_global], 0, count); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // The transposed scale matrix has rows / SCALE_DIM columns. + const size_t col_length_t = rows / SCALE_DIM; + swizzle::store_gemm_swizzled_scales(scales_t_ptr, sSFcolwise[row_tr], row_tr_global, + scales_block_offset_X_tr, col_length_t, count); + } else { + ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); + const size_t scale_idx_global = + row_tr_global * scale_stride_t + scales_block_offset_X_tr; + scales_vec.store_to_elts(&scales_t_ptr[scale_idx_global], 0, count); + } } } } @@ -746,6 +764,13 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA NVTE_CHECK(cols % 32 == 0, "Number of tensor cols must be a multiple of 32"); // 16B alignment for TMA + // In-kernel GEMM-swizzled scales: the swizzled layout is tiled 128x4, so only 128-aligned + // dims avoid padded (unwritten) scale tiles. + const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + NVTE_CHECK(!with_gemm_swizzled_scales || !row_scaled_nvfp4, + "Row-scaled NVFP4 quantization does not support GEMM-swizzled scales."); + NVTE_CHECK(!with_gemm_swizzled_scales || (rows % 128 == 0 && cols % 128 == 0), + "In-kernel GEMM-swizzled NVFP4 scales require rows and cols to be multiples of 128."); const int blocks_Y = DIVUP(rows, static_cast(TunableConfig::CHUNK_DIM_Y)); const int blocks_X = DIVUP(cols, static_cast(TunableConfig::CHUNK_DIM_X)); @@ -827,6 +852,13 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, quantize_transpose_nvfp4_tuned_1D_kernel; + if constexpr (!ROW_SCALED_NVFP4) { + if (with_gemm_swizzled_scales) { + kernel = quantize_transpose_nvfp4_tuned_1D_kernel< + ScaleType, USE_STOCHASTIC_ROUNDING, USE_FAST_MATH, RETURN_TRANSPOSE, + ROW_SCALED_NVFP4, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + } + } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); diff --git a/transformer_engine/common/cast/nvfp4/swizzle.cuh b/transformer_engine/common/cast/nvfp4/swizzle.cuh index 6e82974b605..7e727845941 100644 --- a/transformer_engine/common/cast/nvfp4/swizzle.cuh +++ b/transformer_engine/common/cast/nvfp4/swizzle.cuh @@ -12,6 +12,7 @@ #define TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_SWIZZLE_CUH_ #include +#include namespace transformer_engine { namespace dispatch { @@ -50,6 +51,31 @@ __device__ __forceinline__ size_t gemm_swizzled_scale_idx(size_t row_idx, size_t return ((rb * cbg_cnt + cbg) * kRowsPerBaseBlockCol + d3) * 16 + d4 * kColsPerBaseBlockCol + d5; } +/*! \brief Store \p count consecutive compact scales of one row at their GEMM-swizzled offsets. + * + * \p src holds the scales of compact columns \p col_idx .. \p col_idx + \p count - 1 of row + * \p row_idx. In the swizzled layout the scales of columns 4c .. 4c+3 of a row are 4 contiguous bytes + * at a 4-aligned offset, so whole groups of four are written with one aligned 32-bit store: the same + * bytes at the same addresses as per-scale stores. Misaligned cases fall back to byte stores. + */ +template +__device__ __forceinline__ void store_gemm_swizzled_scales(ScaleT *dst, const ScaleT *src, + size_t row_idx, size_t col_idx, + size_t col_length, int count) { + static_assert(sizeof(ScaleT) == 1, "NVFP4 scales are 1 byte wide."); + int k = 0; + if (col_idx % 4 == 0 && reinterpret_cast(src) % 4 == 0 && + reinterpret_cast(dst) % 4 == 0) { + for (; k + 4 <= count; k += 4) { + const size_t off = gemm_swizzled_scale_idx(row_idx, col_idx + k, col_length); + *reinterpret_cast(dst + off) = *reinterpret_cast(src + k); + } + } + for (; k < count; ++k) { + dst[gemm_swizzled_scale_idx(row_idx, col_idx + k, col_length)] = src[k]; + } +} + } // namespace swizzle } // namespace nvfp4 } // namespace dispatch 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 94d395526be..7aecb2879ca 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 @@ -1396,9 +1396,6 @@ void nvte_quantize_with_hadamard_transform(const NVTETensor input, NVTETensor ou NVTE_API_CALL(nvte_quantize_with_hadamard_transform); using namespace transformer_engine; Tensor &output_cpp = *convertNVTETensorCheck(output); - const int sm_arch = transformer_engine::cuda::sm_arch(transformer_engine::cuda::current_device()); - NVTE_CHECK((sm_arch != 120 && sm_arch != 121) || !output_cpp.with_gemm_swizzled_scales, - "NVFP4 RHT quantization on SM120/SM121 does not support GEMM-swizzled scales."); QuantizationConfig quant_config_cpp; if (quant_config != nullptr) { quant_config_cpp = *reinterpret_cast(quant_config); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index ba0410bcacb..fc850a9968b 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1949,7 +1949,7 @@ NVFP4Quantizer::ShapeAlignment NVFP4Quantizer::rht_cast_fusion_alignment( } NVFP4Quantizer::ShapeAlignment NVFP4Quantizer::get_2d_swizzle_fusion_alignment(int sm_arch) { - if (sm_arch < 100 || sm_arch > 110) { + if ((sm_arch < 100 || sm_arch > 110) && sm_arch != 120 && sm_arch != 121) { return std::nullopt; } return std::pair{128, 128}; @@ -1970,16 +1970,30 @@ NVFP4Quantizer::ShapeAlignment NVFP4Quantizer::gemm_swizzled_scales_alignment( if (!optimize_for_gemm) { return std::nullopt; } + // Row-scaled and 4over6 outputs keep compact scales on every path (the in-kernel swizzled + // stores are instantiated for neither). + if (row_scaled_nvfp4 || nvfp4_use_4over6) { + return std::nullopt; + } + const bool sm12x = sm_arch == 120 || sm_arch == 121; if (with_rht) { - // The SM12x no-TMEM RHT path reuses the compact-scale 1D TMA quantizer. - if (sm_arch == 120 || sm_arch == 121) { - return std::nullopt; + const auto alignment = rht_cast_fusion_alignment(sm_arch, disable_rht_cast_fusion); + // The SM12x no-TMEM RHT path reuses the 1D TMA quantizers (the columnwise-RHT kernel, or + // the tuned 1D kernel for rowwise-only output), which write GEMM-swizzled scales in-kernel + // for 128-aligned shapes. + if (alignment && sm12x) { + return std::pair{128, 128}; } - return rht_cast_fusion_alignment(sm_arch, disable_rht_cast_fusion); + return alignment; } - if (with_2d_quantization && !row_scaled_nvfp4 && !nvfp4_use_4over6) { + if (with_2d_quantization) { return get_2d_swizzle_fusion_alignment(sm_arch); } + // Plain 1D: the tuned (BF16) and generic 1D kernels write GEMM-swizzled scales in-kernel for + // 128-aligned shapes. Enabled where validated (SM120/SM121). + if (sm12x) { + return std::pair{128, 128}; + } return std::nullopt; }