From 424c6111e6674e689a4924a63403f9c6cd614996 Mon Sep 17 00:00:00 2001 From: qiyuw Date: Mon, 10 Aug 2026 18:23:14 +0000 Subject: [PATCH 1/3] mxfp8: add swizzled-scale fast path for cast-only quantization Produce GEMM-ready scales directly in the specialized cast-only kernels to avoid separate scale-swizzle launches while preserving generic fallbacks for unsupported shapes. Signed-off-by: qiyuw --- tests/cpp/operator/test_cast_mxfp8.cu | 238 +++++++ .../test_mxfp8_quantize_swizzle_fusion.py | 33 + .../common/cast/mxfp8/quantize_mxfp8.cuh | 78 +-- .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 592 +++++++++++++++--- 4 files changed, 809 insertions(+), 132 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index c7c778ce1e..edeb87ebe1 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -7,10 +7,12 @@ #include #include #include +#include #include #include #include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -735,3 +737,239 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::ValuesIn(input_scenarios)), test_name_generator); + +// ============================================================================ +// Swizzled-scales cast-only tests +// +// Validate the WITH_GEMM_SWIZZLED_SCALES=true code path added by the +// CastTraitsSwizzle port. The specialized kernel dispatches to +// CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> whenever the +// output tensor has set_with_gemm_swizzled_scales(true), producing scales +// directly in GEMM-swizzled layout. +// +// Reference construction: run the well-tested linear-scale path, then apply +// nvte_swizzle_scaling_factors (independently tested by SwizzleTestSuite) to +// transform the linear scales into the swizzled layout. Byte-compare the +// direct swizzled cast output against this reference for both scale tensors +// and the FP8 output data itself. +// +// Pass criteria per test: +// 1. FP8 rowwise data byte-identical between linear and swizzled paths. +// 2. FP8 colwise data byte-identical between linear and swizzled paths. +// 3. Swizzled rowwise scale bytes match linear->swizzle reference. +// 4. Swizzled colwise scale bytes match linear->swizzle reference. +// 5. Rowwise-only FP8 data and swizzled scales match the same reference. +// 6. No CUDA errors from any launch. +// ============================================================================ + +class SwizzledScalesFusedCastMXFP8TestSuite : public ::testing::TestWithParam< + std::tuple, + transformer_engine::DType, + transformer_engine::DType>> {}; + +TEST_P(SwizzledScalesFusedCastMXFP8TestSuite, TestSwizzledCastMXFP8) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto shape = std::get<0>(GetParam()); + const DType itype = std::get<1>(GetParam()); + const DType otype = std::get<2>(GetParam()); + + // BIDIMENSIONAL scaling needs a 2D+ input. + if (shape.size() < 2) { + GTEST_SKIP(); + } + + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + const size_t out_bytes = rows * cols; // fp8 = 1 byte/elem + + // Input filled with the same values for all casts. + Tensor input("input", shape, itype); + fillUniform(&input); + + // Target: swizzled-scale cast. Dispatcher routes to + // CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> on kernel #3. + Tensor output_swizzled("output_swizzled", shape, otype, true, true, NVTE_MXFP8_1D_SCALING); + output_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_swizzled.data(), 0); + + // Rowwise-only activations take the pointer-based cast-only kernel. This + // directly exercises its WITH_GEMM_SWIZZLED_SCALES trait specialization + // for shapes whose column count is a multiple of 128. + Tensor output_rowwise_swizzled("output_rowwise_swizzled", shape, otype, + /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + output_rowwise_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_rowwise_swizzled.data(), 0); + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "swizzled-scale nvte_quantize failed"; + + // Reference construction: nvte_swizzle_scaling_factors accepts tensors with + // exactly one scale direction, so we build the rowwise and colwise references + // independently. Each is a single-direction linear cast followed by a + // single-direction swizzle transform, using the same input as the target. + // MXFP8 is a deterministic per-direction operation, so a rowwise-only cast + // produces the same rowwise fp8 bytes and scales as a rowwise+colwise cast. + + // Rowwise reference. + Tensor linear_row("linear_row", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_row.data(), 0); + Tensor ref_row_swz("ref_row_swz", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + ref_row_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_row_swz.rowwise_dptr(), linear_row.rowwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_row.data(), ref_row_swz.data(), 0); + } + + // Colwise reference. + Tensor linear_col("linear_col", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_col.data(), 0); + Tensor ref_col_swz("ref_col_swz", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + ref_col_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_col_swz.columnwise_dptr(), linear_col.columnwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_col.data(), ref_col_swz.data(), 0); + } + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "reference construction failed"; + + // ---- Comparisons ---- + + // (1) & (2): FP8 output data byte-identical to single-direction linear casts. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test_row = reinterpret_cast( + output_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref_row = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_row[i], ref_row[i]) + << "rowwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_row[i]) + << " linear=" << static_cast(ref_row[i]) << ")"; + } + const uint8_t *test_col = reinterpret_cast( + output_swizzled.columnwise_cpu_dptr()); + const uint8_t *ref_col = reinterpret_cast( + linear_col.columnwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_col[i], ref_col[i]) + << "colwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_col[i]) + << " linear=" << static_cast(ref_col[i]) << ")"; + } + } + ); + + // (3): Swizzled rowwise scale bytes — directly-written vs linear-then-transform. + { + const size_t n = product(output_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled rowwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (4): Swizzled colwise scale bytes — exercises the uint16-pair-packed flush + // path added by CACHE_COLWISE_SCALE_IN_SMEM + WITH_SWIZZLED_SCALES. + { + const size_t n = product(output_swizzled.columnwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.columnwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_col_swz.columnwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled colwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (5) & (6): The rowwise-only swizzled path matches the same independently + // constructed linear-then-swizzle reference. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test = reinterpret_cast( + output_rowwise_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only fp8 data mismatch at index " << i; + } + } + ); + { + const size_t n = product(output_rowwise_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = + output_rowwise_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only swizzled scale mismatch at offset " << i; + } + } + +} + +std::string swizzled_test_name_generator( + const testing::TestParamInfo& info) { + std::string name; + const auto &shape = std::get<0>(info.param); + for (size_t i = 0; i < shape.size(); ++i) { + if (i > 0) name += "x"; + name += std::to_string(shape[i]); + } + name += "X" + test::typeName(std::get<1>(info.param)) + + "X" + test::typeName(std::get<2>(info.param)); + return name; +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_FusedCastMXFP8_SwizzledCastOnly, + SwizzledScalesFusedCastMXFP8TestSuite, + ::testing::Values( + // 1. Aligned, small — sanity. + std::make_tuple(std::vector{128, 128}, DType::kBFloat16, DType::kFloat8E4M3), + // 2. Aligned, small — second dtype pair. + std::make_tuple(std::vector{128, 128}, DType::kFloat32, DType::kFloat8E5M2), + // 3. Aligned, medium. + std::make_tuple(std::vector{256, 384}, DType::kBFloat16, DType::kFloat8E4M3), + // 4. Odd number of scale rows (96/32 = 3) - exercises colwise flush + // odd-row scalar tail. + std::make_tuple(std::vector{96, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 5. cols/32 not a multiple of 4 - exercises rowwise flush scalar tail + // (need cols multiple of 8 for TMA-alignment on 16-bit input). + std::make_tuple(std::vector{256, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 6. Odd colwise scale rows + rowwise tail cols - combined tail paths. + std::make_tuple(std::vector{96, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 7. Small but tile-aligned. + std::make_tuple(std::vector{32, 64}, DType::kBFloat16, DType::kFloat8E4M3), + // 8. Minimum aligned size - single 32x32 tile. + std::make_tuple(std::vector{32, 32}, DType::kBFloat16, DType::kFloat8E4M3), + // 9. Multi-CTA at scale — stresses gmem writes across many CTAs. + std::make_tuple(std::vector{4096, 32768}, DType::kBFloat16, DType::kFloat8E4M3), + // 10. 4D input — matches the existing suite's rank coverage. + std::make_tuple(std::vector{16, 8, 4, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 11. Third input dtype. + std::make_tuple(std::vector{128, 128}, DType::kFloat16, DType::kFloat8E4M3), + // 12. Larger fp32. + std::make_tuple(std::vector{1024, 1024}, DType::kFloat32, DType::kFloat8E4M3) + ), + swizzled_test_name_generator); diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py index 127b487650..16a2d75de6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -129,3 +129,36 @@ def test_mxfp8_quantize_swizzle_fusion( return_rowwise=return_rowwise, return_transpose=return_transpose, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("M, N", [(96, 160), (4096, 576), (4096, 2112)]) +def test_mxfp8_bidirectional_swizzled_row_scale_padding(M: int, N: int) -> None: + """The specialized bidirectional kernel must not overwrite padded row scales.""" + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + quantizer.optimize_for_gemm = True + scale = quantizer(x)._rowwise_scale_inv.view(torch.uint8) + + scale_rows = torch.arange(M, device=scale.device, dtype=torch.int64).view(-1, 1) + scale_cols = torch.arange(N // 32, device=scale.device, dtype=torch.int64).view(1, -1) + num_tiles_x = math.ceil(N / 128) + scale_indices = ( + ((scale_rows // 128) * num_tiles_x + scale_cols // 4) * (128 * 4) + + (scale_rows % 32) * 16 + + ((scale_rows % 128) // 32) * 4 + + scale_cols % 4 + ) + valid_mask = torch.zeros(scale.numel(), dtype=torch.bool, device=scale.device) + valid_mask[scale_indices.view(-1)] = True + + torch.testing.assert_close( + scale.view(-1)[~valid_mask], + torch.zeros_like(scale.view(-1)[~valid_mask]), + atol=0, + rtol=0, + ) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 8c57f112d2..3073321691 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -648,6 +648,30 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, float *const amax_ptr = reinterpret_cast(output->amax.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); + // Clear padding before either the generic or specialized kernel writes + // directly into the GEMM-swizzled scale layout. + if (with_gemm_swizzled_scales && (cols % 128 != 0 || rows % 128 != 0)) { + constexpr size_t zero_threads = 256; + if (use_rowwise_scaling) { + const size_t size_bytes = output->scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + if (use_colwise_scaling) { + const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->columnwise_scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + } + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( @@ -669,17 +693,27 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, bidimensional_traits::blockDIM::M) <= max_grid_dim_y; const bool is_full_rowwise_chunk = (cols % 128 == 0); + const bool has_full_bidimensional_chunks = + (rows % bidimensional_traits::colChunkElems == 0) && + (cols % bidimensional_traits::rowChunkElems == 0); + // Both rowwise and bidimensional cast-only kernels select their + // scale layout from WITH_GEMM_SWIZZLED_SCALES. const bool scaling_type_has_specialized_support = (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk && rowwise_specialized_grid_fits) || (scaling_type == ScalingType::BIDIMENSIONAL && + has_full_bidimensional_chunks && bidimensional_specialized_grid_fits); - if (specialized::hasSpec() && - !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) { + // Specialized cast-only kernels do not consume the device noop flag. + // Preserve cached outputs by keeping noop-aware calls on the generic path. + if (noop_ptr == nullptr && + specialized::hasSpec() && + scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { - using traits = specialized::CastTraits; + using traits = specialized::CastTraits< + IType, OType, true, false, WITH_GEMM_SWIZZLED_SCALES>; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -698,7 +732,11 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, break; } case ScalingType::BIDIMENSIONAL: { - using traits = specialized::CastTraits; + using traits = specialized::CastTraitsSwizzle< + IType, OType, + /*NumStages=*/2, /*IterM=*/1, /*IterN=*/4, + /*kCacheColwise=*/WITH_GEMM_SWIZZLED_SCALES, + /*kSwizzled=*/WITH_GEMM_SWIZZLED_SCALES>; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -788,38 +826,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - // Zero out swizzled scales if padding is needed - /// TODO (tmoon) Handle this within the cast kernel - if (with_gemm_swizzled_scales) { - constexpr size_t TILE_DIM_X = 128; // Tile dim in data buffer - constexpr size_t TILE_DIM_Y = 128; - if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) { - // Use a noop-aware zero kernel so that the clear is skipped - // when quantization is a noop (e.g. FP8 weight caching). - constexpr size_t zero_threads = 256; - if (use_rowwise_scaling) { - const size_t size_bytes = output->scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->scale_inv.dptr), size_bytes, - noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - if (use_colwise_scaling) { - const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->columnwise_scale_inv.dptr), - size_bytes, noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - } - } - switch (scaling_type) { case ScalingType::ROWWISE: { auto kernel = quantize_mxfp8_kernel #include "../../../util/ptx.cuh" +#include "../swizzle.cuh" // gemm_swizzled_scale_idx (parent dir, GEMM scale swizzle) #include "state_counter.cuh" -#include "swizzle.cuh" +#include "swizzle.cuh" // specialized/swizzle.cuh (TMA input bank-conflict swizzle) namespace transformer_engine { namespace dispatch { @@ -24,6 +25,10 @@ namespace quantize_kernel { namespace specialized { namespace ptx = transformer_engine::ptx; + +// Bring in the GEMM-swizzled scale index helper (from ../swizzle.cuh). +// Used only when the kernel is instantiated with CastTraits::_with_swizzled_scales=true. +using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; namespace { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -122,19 +127,21 @@ struct Layout { static constexpr int32_t num = M * N; }; -template +template struct CastTraits; // 1x32 -template -struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, + _kSwizzled> { static constexpr bool isRowwise = true; static constexpr bool isColwise = false; using IType = _IType; using OType = _OType; static constexpr int32_t chunkElems = 32; - using threadLayout = Layout<1, 32>; + using threadLayout = Layout<1, THREADS_PER_WARP>; static constexpr int32_t numThreadsPerChunk = 1; static constexpr int32_t warpDimM = threadLayout::M; static constexpr int32_t warpDimN = threadLayout::N * chunkElems; @@ -151,14 +158,22 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { using iterLayout = Layout<1, 1>; static constexpr int32_t blockDimM = iterLayout::M * blockIterDimM; static constexpr int32_t blockDimN = iterLayout::N * blockIterDimN; + static constexpr int32_t rowwiseScaleStride = blockDimN / chunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 1; static constexpr int32_t numPrefetch = numStages - 1; static constexpr bool _use_cvt_4x = true; static constexpr bool _cache_rowwise_scale_in_smem = true; + static constexpr bool _with_swizzled_scales = _kSwizzled; - static constexpr int32_t numThreads = warpLayout::num * 32; + static constexpr int32_t numThreads = warpLayout::num * THREADS_PER_WARP; static constexpr size_t smem_rowwise_scale = _cache_rowwise_scale_in_smem ? (blockDimM * (blockDimN / chunkElems) * sizeof(e8m0_t)) : 0ul; @@ -498,13 +513,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re block_coords.y = blockIdx.y * CastTraits::blockDimM; block_coords.x = blockIdx.x * CastTraits::blockDimN; - constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows); @@ -514,7 +524,40 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + if constexpr (CastTraits::_with_swizzled_scales) { + // Four adjacent rowwise scale columns are contiguous in the GEMM scale + // layout, so write them as one uint32_t whenever possible. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::chunkElems; + const size_t num_tiles_x = DIVUP(cols, static_cast(128)); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / groups_per_row; + const int32_t col = (i % groups_per_row) * cols_per_group; + const uint32_t value = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + *reinterpret_cast(&scales_rowwise[idx]) = value; + } + + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / remaining_per_row; + const int32_t col = remaining_start + (i % remaining_per_row); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + scales_rowwise[idx] = sRowwiseScale[row * stride_in_smem + col]; + } + } + } else if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { using DataType = int32_t; constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; @@ -527,8 +570,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -546,8 +590,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -577,11 +622,12 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr int32_t rowChunkElems = 32; static constexpr int32_t colChunkElems = 32; - using rowThreadLayout = Layout<32, 1>; // 32x1 + using rowThreadLayout = Layout; // 32x1 using colThreadLayout = Layout; // 1x32 static_assert(rowThreadLayout::num == colThreadLayout::num, "rowThreadLayout::num must be equal to colThreadLayout::num"); - static_assert(rowThreadLayout::num == 32, "rowThreadLayout::num must be 32"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); using rowWarpDim = Layout; using colWarpDim = Layout; @@ -603,6 +649,13 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { using iterLayout = Layout<1, 4>; using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 2; @@ -636,7 +689,7 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { "It requires aligned smem pointer"); static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; - static constexpr int32_t numThreads = numWarps * 32; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); @@ -660,7 +713,9 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; // && _colwise_reduce_max != ColwiseReduceMax::Redux; static constexpr size_t smem_colwise_reduce = - _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; static constexpr size_t smem = _reuse_input_out_smem @@ -670,6 +725,146 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { smem_alignment + smem_rowwise_scale + smem_colwise_reduce); }; +// Standalone trait for the non-warp-specialized rowwise+colwise cast_only kernel. +// Exposes numStages, iterM, iterN, and the two colwise-scale features as +// caller-controllable template axes. Both colwise flags default to true, +// giving callers the swizzled + colwise-scale-cached fast path by default. +// +// This trait duck-types the same interface CastTraits<_, _, true, true> exposes, +// so it drops into quantize_mxfp8_kernel_cast_only without any +// changes to the kernel signature. Kernel #1 (rowwise-only) and Kernel #2 +// (warp-specialized row+col) don't accept it: kernel #1 requires isColwise=false, +// kernel #2 requires _use_warp_specialization=true - both are wrong here. +template +struct CastTraitsSwizzle { + static constexpr bool isRowwise = true; + static constexpr bool isColwise = true; + using IType = _IType; + using OType = _OType; + + static constexpr int32_t rowChunkElems = 32; + static constexpr int32_t colChunkElems = 32; + + using rowThreadLayout = Layout; // 32x1 + using colThreadLayout = Layout; // 1x32 + static_assert(rowThreadLayout::num == colThreadLayout::num, + "rowThreadLayout::num must be equal to colThreadLayout::num"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); + + using rowWarpDim = Layout; + using colWarpDim = Layout; + using warpDim = + Layout; + + static constexpr bool _tma_swizzle = true; + using warpLayout = Layout<1, 2>; + static_assert(_tma_swizzle ? (warpLayout::N == 2) : true); + static constexpr CUtensorMapSwizzle input_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + static constexpr CUtensorMapSwizzle output_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + using blockIterDim = Layout; + + using iterLayout = Layout<_IterM, _IterN>; + using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; + + static constexpr int32_t numStages = _NumStages; + + using inputUnitType = uint4; + static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType); + static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit; + using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; + using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; + + using colIndexSwz = swz::Swizzle<5, 0, 5>; + + using rowOutputUnitType = uint4; + static constexpr int32_t rowNumOutUnitsPerChunk = + rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType); + static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType); + + using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; + using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; + + static constexpr bool _use_cvt_4x = true; + static constexpr bool _use_warp_specialization = false; + static constexpr bool _need_wait_group = iterLayout::num > numStages; + static constexpr bool _reuse_input_out_smem = false; + static_assert(_reuse_input_out_smem == false, "Just don't use it"); + static constexpr bool _cache_rowwise_scale_in_smem = true; + + static constexpr bool _colwise_source_coming_from_rowwise = true; + static constexpr ColwiseReduceMax _colwise_reduce_max = ColwiseReduceMax::Redux; + static_assert(_colwise_reduce_max != ColwiseReduceMax::RedAsync, + "It requires aligned smem pointer"); + + // The two colwise-scale features exposed as caller-controllable trait axes. + // Both default to true so callers get the vectorized swizzled path by default. + static constexpr bool _cache_colwise_scale_in_smem = _kCacheColwise; + static constexpr bool _with_swizzled_scales = _kSwizzled; + + static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; + static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); + + static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); + static constexpr size_t smemInputPerBlock = smemInputPerWarp * warpLayout::num; + + static constexpr size_t smemRowwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemRowwiseOutputPerBlock = smemRowwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemColwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemColwiseOutputPerBlock = smemColwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemInput = smemInputPerBlock * numStages; + static constexpr size_t smemRowwiseOutput = smemRowwiseOutputPerBlock * numStages; + static constexpr size_t smemColwiseOutput = smemColwiseOutputPerBlock * numStages; + + static constexpr size_t smem_rowwise_scale = + _cache_rowwise_scale_in_smem ? (blockDIM::M * (blockDIM::N / rowChunkElems) * sizeof(e8m0_t)) + : 0ul; + + // Extra shmem for cached colwise scales - only when the flag is on. + static constexpr size_t smem_colwise_scale = + _cache_colwise_scale_in_smem + ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t) + : 0ul; + + using ColwiseReduceDataType = float; + static constexpr bool _need_smem_for_colwise_reduce = + _colwise_source_coming_from_rowwise; + static constexpr size_t smem_colwise_reduce = + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; + + static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; + static constexpr size_t smem = _reuse_input_out_smem + ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + + smem_alignment + smem_rowwise_scale + + smem_colwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + + smem_alignment + smem_rowwise_scale + + smem_colwise_scale + smem_colwise_reduce); +}; + __device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { return (x + align - 1) & ~((align)-1); } @@ -728,12 +923,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } // TODO: maybe we can assign a different barrier for each warp @@ -744,8 +939,10 @@ __global__ void quantize_mxfp8_kernel_cast_only( #pragma unroll for (int32_t i = 0; i < CastTraits::numStages; i++) { ptx::mbarrier_init(&ldg_producer[i], 1); - ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * 32); - ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * 32); + ptx::mbarrier_init(&ldg_consumer[i], + CastTraits::warpLayout::num * THREADS_PER_WARP); + ptx::mbarrier_init(&stg_producer[i], + CastTraits::warpLayout::num * THREADS_PER_WARP); ptx::mbarrier_init(&stg_consumer[i], 1); } ptx::fence_mbarrier_init_release_cluster(); @@ -1085,15 +1282,10 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { - ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u); + ptx::numbered_barrier_sync(CastTraits::warpLayout::num * THREADS_PER_WARP, 0u); - constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); @@ -1117,8 +1309,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1137,8 +1330,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1181,6 +1375,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( extern __shared__ char smem[]; char *smemAligned = reinterpret_cast( align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + // Re-assert .shared address space lost by the intptr_t round-trip in + // align_to() so NVVM InferAddressSpaces emits LDS/STS instead of LD.E/ST.E. + __builtin_assume(__isShared(smemAligned)); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -1191,15 +1388,20 @@ __global__ void quantize_mxfp8_kernel_cast_only( // colwise output will reuse input buffer OType *sColOutput; e8m0_t *sRowwiseScale = nullptr; + e8m0_t *sColwiseScale = nullptr; ColwiseReduceDataType *sColwiseReduce = nullptr; if constexpr (CastTraits::_reuse_input_out_smem) { sColOutput = reinterpret_cast(sInput); if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sRowOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1211,9 +1413,13 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1223,7 +1429,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( rowOutputUnitType *sColOutputUnit = reinterpret_cast(sColOutput); if constexpr (CastTraits::_need_smem_for_colwise_reduce) { - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } __shared__ uint64_t producer[CastTraits::numStages]; @@ -1243,7 +1449,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_colwise_source_coming_from_rowwise && CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { - ptx::mbarrier_init(colwise_reduce_barrier, 32); + ptx::mbarrier_init(colwise_reduce_barrier, THREADS_PER_WARP); } ptx::fence_mbarrier_init_release_cluster(); @@ -1263,18 +1469,35 @@ __global__ void quantize_mxfp8_kernel_cast_only( (threadIdx.x % CastTraits::rowThreadLayout::N) * (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit); - size_t rowwise_scale_base_offset = - (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) * - static_cast(scale_stride_rowwise) + + // Scale coordinates in absolute (compact) scale-tensor space. Shared by both + // compact-layout offsets and by the CastTraits::_with_swizzled_scales branches below. + const int32_t row_scale_row_base = + block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N); + const int32_t row_scale_col_base = (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) / - CastTraits::rowChunkElems; + CastTraits::rowChunkElems; + const int32_t col_scale_row_base = + (block_coords.y + warp_coords.y + + (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / + CastTraits::colChunkElems; + const int32_t col_scale_col_base = + block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N); + + size_t rowwise_scale_base_offset = + static_cast(row_scale_row_base) * static_cast(scale_stride_rowwise) + + row_scale_col_base; size_t colwise_scale_base_offset = - ((block_coords.y + warp_coords.y + - (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / - CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N)); + static_cast(col_scale_row_base) * static_cast(scale_stride_colwise) + + col_scale_col_base; + + // Precomputed swizzle constants (each swizzle tile is 128 rows x 4 cols in scale space). + // Rowwise scale tensor has DIVUP(cols, 128) tiles across; + // colwise scale tensor has DIVUP(rows, 128) tiles across (X/Y axes are transposed). + const size_t row_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(cols, static_cast(128)) : 0; + const size_t col_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(rows, static_cast(128)) : 0; constexpr int32_t rowwise_scale_stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; @@ -1401,6 +1624,14 @@ __global__ void quantize_mxfp8_kernel_cast_only( iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = + row_scale_row_base + iter_m * CastTraits::blockIterDim::M; + int32_t abs_col = row_scale_col_base + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + size_t idx = + gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = row_biased_exponent; } else { size_t rowwise_scale_offset = rowwise_scale_base_offset + @@ -1416,12 +1647,36 @@ __global__ void quantize_mxfp8_kernel_cast_only( e8m0_t col_biased_exponent = to_e8m0(col_amax); float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); sColwiseReduce[threadIdx.x] = col_scale_inverse; - size_t colwise_scale_offset = - colwise_scale_base_offset + - iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - iter_n * CastTraits::blockIterDim::N; - scales_colwise[colwise_scale_offset] = col_biased_exponent; + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + // Cache in shmem; end-of-kernel flush handles gmem indexing. + int32_t smem_row = + (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) * + CastTraits::colChunkElems) / + CastTraits::colChunkElems + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t smem_col = + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) + + iter_n * CastTraits::blockIterDim::N; + sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = + col_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = col_scale_row_base + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t abs_col = col_scale_col_base + iter_n * CastTraits::blockIterDim::N; + // Colwise scale tensor's X/Y axes are transposed vs rowwise + // (see col_swz_num_tiles_X = DIVUP(rows, 128)), so pass + // (abs_col, abs_row) - abs_col is the swizzle "row" dim. + size_t idx = + gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + scales_colwise[idx] = col_biased_exponent; + } else { + size_t colwise_scale_offset = + colwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + iter_n * CastTraits::blockIterDim::N; + scales_colwise[colwise_scale_offset] = col_biased_exponent; + } __syncwarp(); } } @@ -1546,58 +1801,203 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - scale_stride_rowwise); + DIVUP(cols, static_cast(CastTraits::rowChunkElems))); int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { - using DataType = int32_t; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + if constexpr (CastTraits::_with_swizzled_scales) { + // Swizzled flush: within a 128x4 swizzle tile, 4 consecutive column entries + // for the same row live at 4 consecutive gmem bytes. Group by 4 so each + // thread writes a uint32_t when col%4 == 0, then scalar tail for remainder. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::rowChunkElems; + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + + uint32_t val4 = *reinterpret_cast( + &sRowwiseScale[row * stride_in_smem + col]); + + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + *reinterpret_cast(&scales_rowwise[idx]) = val4; + } - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + // Tail (valid_coords.x % 4 != 0) + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_per_row; + int32_t col = remaining_start + (i % remaining_per_row); + e8m0_t val = sRowwiseScale[row * stride_in_smem + col]; + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = val; + } } } else { - using DataType = PreferredDataType; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + using PreferredDataType = typename CastTraits::PreferredDataType; - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + } + + // Cached colwise scale flush (swizzled path only). Same barrier semantics as + // the rowwise flush above: the last-iter __syncthreads already ordered every + // in-loop sColwiseScale byte store before this block reads them. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + // In GEMM swizzle, contiguous bytes run over four scale-row indices for a + // fixed logical column. A 64-row CTA owns two such rows; pack them as a + // uint16 store only when the pair stays inside the same 4-row swizzle group. + const int32_t row_pairs = valid_rows / 2; + const int32_t total_pairs = row_pairs * valid_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_pairs; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = (i / valid_cols) * 2; + int32_t col = i % valid_cols; + const int32_t abs_col = block_coords.x + col; + const int32_t abs_row = scale_row_base + row; + e8m0_t val0 = sColwiseScale[row * CastTraits::blockDIM::N + col]; + e8m0_t val1 = sColwiseScale[(row + 1) * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + if (((abs_row & 3) != 3) && + ((reinterpret_cast(&scales_colwise[idx]) & (alignof(uint16_t) - 1)) == 0)) { + uint16_t val2 = + static_cast(val0) | (static_cast(val1) << 8); + *reinterpret_cast(&scales_colwise[idx]) = val2; + } else { + scales_colwise[idx] = val0; + size_t idx1 = + gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X); + scales_colwise[idx1] = val1; + } + } + // Odd-row tail. + if ((valid_rows & 1) != 0) { + const int32_t row = valid_rows - 1; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < valid_cols; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t col = i; + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, + scale_row_base + row, + col_swz_num_tiles_X); + scales_colwise[idx] = val; + } + } + } + + // Cached colwise scale flush (non-swizzled linear layout). + // Rows in scales_colwise are indexed by (block_coords.y / colChunkElems), + // columns are logical input columns; 4 adjacent columns for the same scale + // row live at 4 consecutive gmem bytes, so pack as uint32 stores. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && !CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_cols / cols_per_group; + const int32_t total_groups = valid_rows * groups_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + uint32_t val4 = *reinterpret_cast( + &sColwiseScale[row * CastTraits::blockDIM::N + col]); + size_t idx = static_cast(scale_row_base + row) * + static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + *reinterpret_cast(&scales_colwise[idx]) = val4; + } + // Column tail (valid_cols not multiple of 4). + const int32_t remaining_start = groups_per_row * cols_per_group; + if (remaining_start < valid_cols) { + const int32_t remaining_cols = valid_cols - remaining_start; + const int32_t total_remaining = valid_rows * remaining_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_cols; + int32_t col = remaining_start + (i % remaining_cols); + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = static_cast(scale_row_base + row) * + static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + scales_colwise[idx] = val; } } } @@ -1605,7 +2005,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( ptx::cp_async_bulk_wait_group_read<0>(); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} +} // NOLINT(readability/fn_size) } // namespace specialized } // namespace quantize_kernel From 5a9f59c1644cb89abc70ecc2562f3e528b2dcfde Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:49:48 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 17 ++- .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 101 ++++++++---------- 2 files changed, 50 insertions(+), 68 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 3073321691..a3a29a5d67 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -701,8 +701,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const bool scaling_type_has_specialized_support = (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk && rowwise_specialized_grid_fits) || - (scaling_type == ScalingType::BIDIMENSIONAL && - has_full_bidimensional_chunks && + (scaling_type == ScalingType::BIDIMENSIONAL && has_full_bidimensional_chunks && bidimensional_specialized_grid_fits); // Specialized cast-only kernels do not consume the device noop flag. @@ -712,8 +711,8 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { - using traits = specialized::CastTraits< - IType, OType, true, false, WITH_GEMM_SWIZZLED_SCALES>; + using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -732,11 +731,11 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, break; } case ScalingType::BIDIMENSIONAL: { - using traits = specialized::CastTraitsSwizzle< - IType, OType, - /*NumStages=*/2, /*IterM=*/1, /*IterN=*/4, - /*kCacheColwise=*/WITH_GEMM_SWIZZLED_SCALES, - /*kSwizzled=*/WITH_GEMM_SWIZZLED_SCALES>; + using traits = + specialized::CastTraitsSwizzle; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index af03f8b599..2fd231b29b 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -133,8 +133,7 @@ struct CastTraits; // 1x32 template -struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, - _kSwizzled> { +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, _kSwizzled> { static constexpr bool isRowwise = true; static constexpr bool isColwise = false; using IType = _IType; @@ -735,12 +734,8 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { // changes to the kernel signature. Kernel #1 (rowwise-only) and Kernel #2 // (warp-specialized row+col) don't accept it: kernel #1 requires isColwise=false, // kernel #2 requires _use_warp_specialization=true - both are wrong here. -template +template struct CastTraitsSwizzle { static constexpr bool isRowwise = true; static constexpr bool isColwise = true; @@ -818,7 +813,7 @@ struct CastTraitsSwizzle { // The two colwise-scale features exposed as caller-controllable trait axes. // Both default to true so callers get the vectorized swizzled path by default. static constexpr bool _cache_colwise_scale_in_smem = _kCacheColwise; - static constexpr bool _with_swizzled_scales = _kSwizzled; + static constexpr bool _with_swizzled_scales = _kSwizzled; static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; @@ -843,26 +838,23 @@ struct CastTraitsSwizzle { // Extra shmem for cached colwise scales - only when the flag is on. static constexpr size_t smem_colwise_scale = - _cache_colwise_scale_in_smem - ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t) - : 0ul; + _cache_colwise_scale_in_smem ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t) + : 0ul; using ColwiseReduceDataType = float; - static constexpr bool _need_smem_for_colwise_reduce = - _colwise_source_coming_from_rowwise; + static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; static constexpr size_t smem_colwise_reduce = _need_smem_for_colwise_reduce ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; - static constexpr size_t smem = _reuse_input_out_smem - ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + - smem_alignment + smem_rowwise_scale + - smem_colwise_scale + smem_colwise_reduce) - : (smemInput + smemRowwiseOutput + smemColwiseOutput + - smem_alignment + smem_rowwise_scale + - smem_colwise_scale + smem_colwise_reduce); + static constexpr size_t smem = + _reuse_input_out_smem + ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce); }; __device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { @@ -939,10 +931,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( #pragma unroll for (int32_t i = 0; i < CastTraits::numStages; i++) { ptx::mbarrier_init(&ldg_producer[i], 1); - ptx::mbarrier_init(&ldg_consumer[i], - CastTraits::warpLayout::num * THREADS_PER_WARP); - ptx::mbarrier_init(&stg_producer[i], - CastTraits::warpLayout::num * THREADS_PER_WARP); + ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); + ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); ptx::mbarrier_init(&stg_consumer[i], 1); } ptx::fence_mbarrier_init_release_cluster(); @@ -1625,12 +1615,10 @@ __global__ void quantize_mxfp8_kernel_cast_only( iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; } else if constexpr (CastTraits::_with_swizzled_scales) { - int32_t abs_row = - row_scale_row_base + iter_m * CastTraits::blockIterDim::M; + int32_t abs_row = row_scale_row_base + iter_m * CastTraits::blockIterDim::M; int32_t abs_col = row_scale_col_base + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); - size_t idx = - gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); scales_rowwise[idx] = row_biased_exponent; } else { size_t rowwise_scale_offset = @@ -1649,16 +1637,13 @@ __global__ void quantize_mxfp8_kernel_cast_only( sColwiseReduce[threadIdx.x] = col_scale_inverse; if constexpr (CastTraits::_cache_colwise_scale_in_smem) { // Cache in shmem; end-of-kernel flush handles gmem indexing. - int32_t smem_row = - (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) * - CastTraits::colChunkElems) / - CastTraits::colChunkElems + - iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); - int32_t smem_col = - warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) + - iter_n * CastTraits::blockIterDim::N; - sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = - col_biased_exponent; + int32_t smem_row = (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) * + CastTraits::colChunkElems) / + CastTraits::colChunkElems + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t smem_col = warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) + + iter_n * CastTraits::blockIterDim::N; + sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = col_biased_exponent; } else if constexpr (CastTraits::_with_swizzled_scales) { int32_t abs_row = col_scale_row_base + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); @@ -1666,8 +1651,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( // Colwise scale tensor's X/Y axes are transposed vs rowwise // (see col_swz_num_tiles_X = DIVUP(rows, 128)), so pass // (abs_col, abs_row) - abs_col is the swizzle "row" dim. - size_t idx = - gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); scales_colwise[idx] = col_biased_exponent; } else { size_t colwise_scale_offset = @@ -1825,8 +1809,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( int32_t group = i % groups_per_row; int32_t col = group * cols_per_group; - uint32_t val4 = *reinterpret_cast( - &sRowwiseScale[row * stride_in_smem + col]); + uint32_t val4 = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); int32_t abs_row = block_coords.y + row; int32_t abs_col = base_col + col; @@ -1871,7 +1855,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; } } else { using DataType = PreferredDataType; @@ -1891,7 +1876,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; } } } @@ -1928,13 +1914,11 @@ __global__ void quantize_mxfp8_kernel_cast_only( size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); if (((abs_row & 3) != 3) && ((reinterpret_cast(&scales_colwise[idx]) & (alignof(uint16_t) - 1)) == 0)) { - uint16_t val2 = - static_cast(val0) | (static_cast(val1) << 8); + uint16_t val2 = static_cast(val0) | (static_cast(val1) << 8); *reinterpret_cast(&scales_colwise[idx]) = val2; } else { scales_colwise[idx] = val0; - size_t idx1 = - gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X); + size_t idx1 = gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X); scales_colwise[idx1] = val1; } } @@ -1945,8 +1929,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( i += CastTraits::warpLayout::num * THREADS_PER_WARP) { const int32_t col = i; e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; - size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, - scale_row_base + row, + size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, scale_row_base + row, col_swz_num_tiles_X); scales_colwise[idx] = val; } @@ -1977,11 +1960,11 @@ __global__ void quantize_mxfp8_kernel_cast_only( int32_t row = i / groups_per_row; int32_t group = i % groups_per_row; int32_t col = group * cols_per_group; - uint32_t val4 = *reinterpret_cast( - &sColwiseScale[row * CastTraits::blockDIM::N + col]); - size_t idx = static_cast(scale_row_base + row) * - static_cast(scale_stride_colwise) + - static_cast(block_coords.x + col); + uint32_t val4 = + *reinterpret_cast(&sColwiseScale[row * CastTraits::blockDIM::N + col]); + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); *reinterpret_cast(&scales_colwise[idx]) = val4; } // Column tail (valid_cols not multiple of 4). @@ -1994,9 +1977,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( int32_t row = i / remaining_cols; int32_t col = remaining_start + (i % remaining_cols); e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; - size_t idx = static_cast(scale_row_base + row) * - static_cast(scale_stride_colwise) + - static_cast(block_coords.x + col); + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); scales_colwise[idx] = val; } } From 79ba8c301e93f332a18b2b586f98b083e79918e8 Mon Sep 17 00:00:00 2001 From: qiyuw Date: Tue, 11 Aug 2026 00:31:16 +0000 Subject: [PATCH 3/3] initializes padded scale entries Signed-off-by: qiyuw --- .../common/cast/mxfp8/specialized/quantize_mxfp8.cuh | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index 2fd231b29b..fd528b5b6c 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -1788,8 +1788,16 @@ __global__ void quantize_mxfp8_kernel_cast_only( int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); - end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - DIVUP(cols, static_cast(CastTraits::rowChunkElems))); + if constexpr (CastTraits::_with_swizzled_scales) { + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + DIVUP(cols, static_cast(CastTraits::rowChunkElems))); + } else { + // The compact layout's padded entries are consumed by a later swizzle. + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); + } int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems);