diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 832177c637..6a2760b699 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -13,6 +13,7 @@ add_executable(test_operator test_qdq.cu test_cast_mxfp8.cu test_cast_mxfp8_grouped.cu + test_cast_mxfp8_grouped_swiglu.cu test_cast_nvfp4_transpose.cu test_cast_float8blockwise.cu test_cast_float8blockwise_grouped.cu diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu b/tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu new file mode 100644 index 0000000000..24dc27e3a9 --- /dev/null +++ b/tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu @@ -0,0 +1,449 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +// Only the two grouped layouts with a uniform last dim are supported: the output is +// [T, F] with F shared by every expert. +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1 +}; + +constexpr size_t SCALE_DIM_Y = 32; + +// Host mirror of mxfp8::swizzle::gemm_swizzled_scale_idx. The FC2 wgrad GEMM reads this +// operand transposed, so its scale matrix is the [cols, rows/32] transpose of the compact +// one, tiled 128x4: +// https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout +size_t gemm_swizzled_scale_idx(const size_t i, const size_t j, const size_t num_tiles_X) { + constexpr size_t TILE_DIM_X = 4; + constexpr size_t TILE_DIM_Y = 128; + constexpr size_t TILE_SIZE = TILE_DIM_X * TILE_DIM_Y; + const size_t tile_idx_X = j / TILE_DIM_X; + const size_t tile_idx_Y = i / TILE_DIM_Y; + const size_t idx_in_tile_X = j % TILE_DIM_X; + const size_t idx_in_tile_Y = i % TILE_DIM_Y; + size_t idx = (tile_idx_Y * num_tiles_X + tile_idx_X) * TILE_SIZE; + idx += (idx_in_tile_Y % 32) * 16 + (idx_in_tile_Y / 32) * 4 + idx_in_tile_X; + return idx; +} + +/** + * Reference for a single expert: (silu(act) * gate) * prob, then columnwise MXFP8. + * input : [rows, 2 * cols], last dim = [act | gate] + * prob : [rows], per-token router weight in the input dtype + * output : [rows, cols] + * scales : this expert's block of e8m0 exponents, compact or GEMM-swizzled + */ +template +void compute_ref(const InputType* input, + const InputType* prob, + OutputType* output, + fp8e8m0* scales, + const size_t rows, + const size_t cols, + const size_t scales_stride, + const bool with_gemm_swizzled_scales) { + const size_t blocks_Y = divide_round_up(rows, SCALE_DIM_Y); + // Number of 4-wide tiles along the swizzled matrix's column axis (which is rows / 32). + const size_t swizzled_tiles_X = divide_round_up(rows, scale_tensor_alignment_Y_rowwise); + const size_t input_stride = 2 * cols; + + #pragma omp parallel proc_bind(spread) + { + // Buffer to cache the weighted activation of one 32-element block + std::vector cache(SCALE_DIM_Y); + #pragma omp for schedule(static) + for (size_t block_Y = 0; block_Y < blocks_Y; ++block_Y) { + const size_t i_min = block_Y * SCALE_DIM_Y; + const size_t i_max = std::min(rows, i_min + SCALE_DIM_Y); + + for (size_t j = 0; j < cols; ++j) { + float block_amax = 0.0f; + for (size_t i = i_min; i < i_max; ++i) { + const float act_elt = static_cast(input[i * input_stride + j]); + const float gate_elt = static_cast(input[i * input_stride + cols + j]); + const float prob_elt = static_cast(prob[i]); + // Numerical truncation: the kernel rounds the weighted activation back + // through InputType before quantizing, so the reference must too. + const float elt = static_cast( + static_cast(silu(act_elt) * gate_elt * prob_elt)); + cache[i - i_min] = elt; + block_amax = std::max(block_amax, std::abs(elt)); + } + + const fp8e8m0 biased_exponent = + float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); + const size_t scale_idx = with_gemm_swizzled_scales + ? gemm_swizzled_scale_idx(j, block_Y, swizzled_tiles_X) + : block_Y * scales_stride + j; + scales[scale_idx] = biased_exponent; + + const float scale_reciprocal = exp2f_rcp(biased_exponent); + for (size_t i = i_min; i < i_max; ++i) { + output[i * cols + j] = + static_cast(cache[i - i_min] * scale_reciprocal); + } + } + } + } +} + +template +void compare_quantized_elts(const std::string& name, + const T* ref_data, + const T* test_data, + const size_t numel, + const size_t tolerable_mismatches_limit) { + size_t mismatches_num = 0; + int64_t first_mismatch_idx = -1; + + for (size_t i = 0; i < numel; ++i) { + const double t = static_cast(test_data[i]); + const double r = static_cast(ref_data[i]); + if (t == r) { + continue; + } + // Tolerate round-to-nearest picking the other side of the real value: the kernel's + // silu intrinsic and the CPU reference can disagree in the last ULP, which flips + // codes that sit on a rounding boundary. + const double mean = (t + r) / 2; + const double mean_p = mean >= 0 ? mean * (1 + 1e-6) : mean * (1 - 1e-6); + const double mean_m = mean >= 0 ? mean * (1 - 1e-6) : mean * (1 + 1e-6); + const double cast_mean_p = static_cast(static_cast(mean_p)); + const double cast_mean_m = static_cast(static_cast(mean_m)); + if (cast_mean_m == std::min(t, r) && cast_mean_p == std::max(t, r)) { + continue; + } + + mismatches_num++; + if (first_mismatch_idx == -1) { + first_mismatch_idx = static_cast(i); + } + if (mismatches_num > tolerable_mismatches_limit) { + GTEST_FAIL() << mismatches_num << " mismatch(es) in " << name + << ", more than the tolerable limit of " + << tolerable_mismatches_limit << "." << std::endl + << "First mismatch at " << first_mismatch_idx << ": " + << static_cast(test_data[first_mismatch_idx]) << " vs " + << static_cast(ref_data[first_mismatch_idx]); + } + } +} + +template +void performTest(const ShapeRepresentation shape_rep, + const size_t num_tensors, + const std::vector& rows_per_tensor, + const size_t F, + const bool with_gemm_swizzled_scales, + const bool expect_rejection) { + using namespace test; + + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + size_t T = 0; + for (size_t t = 0; t < num_tensors; ++t) { + T += rows_per_tensor[t]; + } + + const size_t in_elts = T * 2 * F; + const size_t out_elts = T * F; + const size_t scales_stride = round_up_to_nearest_multiple(F, scale_tensor_alignment_X_colwise); + + // Element offsets into the [T, F] output, and e8m0 offsets into the scale buffer. Both + // layouts share these offsets: per-expert row counts are 128-aligned, so a compact and + // a swizzled block for the same expert occupy the same number of scales. + std::vector data_offsets(num_tensors + 1, 0); + std::vector scale_offsets(num_tensors + 1, 0); + std::vector first_dims(num_tensors, 0); + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = rows_per_tensor[t]; + first_dims[t] = static_cast(M); + data_offsets[t + 1] = data_offsets[t] + static_cast(M * F); + const size_t blocks_Y = round_up_to_nearest_multiple(divide_round_up(M, SCALE_DIM_Y), + scale_tensor_alignment_Y_colwise); + scale_offsets[t + 1] = scale_offsets[t] + blocks_Y * scales_stride; + } + const size_t sfs_num = scale_offsets[num_tensors]; + + std::mt19937 gen; + std::uniform_real_distribution<> dis(-2.0, 1.0); + std::vector in_data(in_elts); + for (size_t i = 0; i < in_elts; ++i) { + in_data[i] = static_cast(dis(gen)); + } + + // prob follows TE's cuDNN fc1_prob_tensor convention: model (input) dtype. + Tensor prob("prob", std::vector{T}, itype); + fillUniform(&prob); + + const size_t in_data_size = in_elts * sizeof(InputType); + const size_t out_data_size = out_elts * sizeof(OutputType); + const size_t scales_size = sfs_num * sizeof(fp8e8m0); + + auto in_data_d = cuda_alloc(in_data_size); + auto out_data_d = cuda_alloc(out_data_size); + auto out_scales_d = cuda_alloc(scales_size); + auto first_dims_d = cuda_alloc(num_tensors * sizeof(int64_t)); + auto offsets_d = cuda_alloc((num_tensors + 1) * sizeof(int64_t)); + + NVTE_CHECK_CUDA(cudaMemcpy(in_data_d.get(), in_data.data(), in_data_size, + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(first_dims_d.get(), first_dims.data(), + num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(offsets_d.get(), data_offsets.data(), + (num_tensors + 1) * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemset(out_data_d.get(), 0, out_data_size)); + NVTE_CHECK_CUDA(cudaMemset(out_scales_d.get(), 0, scales_size)); + + std::vector in_logical_shape_vec = {T, 2 * F}; + std::vector out_logical_shape_vec = {T, F}; + std::vector scales_shape_vec = {sfs_num}; + NVTEShape in_logical_shape = nvte_make_shape(in_logical_shape_vec.data(), + in_logical_shape_vec.size()); + NVTEShape out_logical_shape = nvte_make_shape(out_logical_shape_vec.data(), + out_logical_shape_vec.size()); + NVTEShape scales_shape = nvte_make_shape(scales_shape_vec.data(), scales_shape_vec.size()); + + NVTEShape first_dims_shape; + NVTEShape offsets_shape; + first_dims_shape.ndim = 1; + offsets_shape.ndim = 1; + first_dims_shape.data[0] = num_tensors; + offsets_shape.data[0] = num_tensors + 1; + + NVTEGroupedTensor in_group_tensor = + nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, in_logical_shape); + NVTEGroupedTensor out_group_tensor = + nvte_create_grouped_tensor(NVTE_MXFP8_1D_SCALING, num_tensors, out_logical_shape); + + NVTEBasicTensor in_data_tensor = {in_data_d.get(), static_cast(itype), + in_logical_shape}; + nvte_set_grouped_tensor_param(in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, + &in_data_tensor, sizeof(in_data_tensor)); + + // Columnwise only: the MoE FC2 weight-gradient GEMM is the sole consumer. + NVTEBasicTensor out_data_tensor = {out_data_d.get(), static_cast(otype), + out_logical_shape}; + NVTEBasicTensor out_scales_tensor = {out_scales_d.get(), NVTEDType::kNVTEFloat8E8M0, + scales_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedColumnwiseData, + &out_data_tensor, sizeof(out_data_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedColumnwiseScaleInv, + &out_scales_tensor, sizeof(out_scales_tensor)); + + // The launcher derives the grouped layout from the output metadata: leaving first_dims + // unset means SAME_BOTH_DIMS, setting it means VARYING_FIRST_DIM. + if (shape_rep == VARYING_FIRST_DIM) { + NVTEBasicTensor first_dims_tensor = {first_dims_d.get(), kNVTEInt64, first_dims_shape}; + NVTEBasicTensor offsets_tensor = {offsets_d.get(), kNVTEInt64, offsets_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedFirstDims, + &first_dims_tensor, sizeof(first_dims_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, + &offsets_tensor, sizeof(offsets_tensor)); + } + + if (with_gemm_swizzled_scales) { + const uint8_t flag = 1; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedWithGEMMSwizzledScales, + &flag, sizeof(flag)); + } + + if (expect_rejection) { + EXPECT_THROW(nvte_group_swiglu_quantize(in_group_tensor, prob.data(), out_group_tensor, 0), + std::runtime_error); + nvte_destroy_grouped_tensor(in_group_tensor); + nvte_destroy_grouped_tensor(out_group_tensor); + return; + } + + // Reference (CPU), one expert at a time. + std::vector out_data_ref(out_elts, static_cast(0.0f)); + std::vector out_scales_ref(sfs_num, static_cast(0)); + const InputType* const prob_ptr = prob.rowwise_cpu_dptr(); + size_t row_base = 0; + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = rows_per_tensor[t]; + if (M == 0) { + continue; + } + // data_offsets are F-based, so the [T, 2F] input offset is twice as large. + compute_ref(in_data.data() + 2 * data_offsets[t], + prob_ptr + row_base, + out_data_ref.data() + data_offsets[t], + out_scales_ref.data() + scale_offsets[t], + M, F, scales_stride, with_gemm_swizzled_scales); + row_base += M; + } + + // GPU + nvte_group_swiglu_quantize(in_group_tensor, prob.data(), out_group_tensor, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + std::vector out_data_h(out_elts); + std::vector out_scales_h(sfs_num); + NVTE_CHECK_CUDA(cudaMemcpy(out_data_h.data(), out_data_d.get(), out_data_size, + cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(out_scales_h.data(), out_scales_d.get(), scales_size, + cudaMemcpyDeviceToHost)); + + // A last-ULP silu difference can push a block amax onto the next e8m0 exponent, so a + // few scale mismatches are tolerated; every element of such a block is then allowed to + // differ as well. + const size_t scale_diff_abs_tolerance = 0; + const double abs_tolerable_mismatches_limit = 1.0; + const double rel_tolerable_mismatches_limit = 1.0e-4; + + size_t mismatches_scales = 0; + compare_scaling_factors("colwise_scales", out_scales_h.data(), out_scales_ref.data(), + 1, sfs_num, sfs_num, mismatches_scales, scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, rel_tolerable_mismatches_limit); + + compare_quantized_elts("colwise_output", out_data_ref.data(), out_data_h.data(), + out_elts, 32 * mismatches_scales); + + nvte_destroy_grouped_tensor(in_group_tensor); + nvte_destroy_grouped_tensor(out_group_tensor); +} + +// {shape_representation, num_tensors, F, rows_of_each_expert...} +// Per-expert row counts are multiples of 128, which the kernel requires. +std::vector> input_configs = { + {SAME_BOTH_DIMS, 1, 128, 128}, + {SAME_BOTH_DIMS, 2, 256, 128, 128}, + {VARYING_FIRST_DIM, 2, 128, 128, 384}, + {VARYING_FIRST_DIM, 3, 256, 128, 384, 512}, + // Empty expert in the middle must not terminate the persistent work loop. + {VARYING_FIRST_DIM, 4, 256, 128, 384, 0, 512}, + // F is not a multiple of the 128-wide chunk, exercising the partial-tile bounds check. + {VARYING_FIRST_DIM, 4, 160, 128, 384, 512, 512}, + {VARYING_FIRST_DIM, 5, 512, 128, 256, 384, 1024, 2304}, +}; + +std::vector> input_configs_small = { + {SAME_BOTH_DIMS, 1, 128, 128}, + {VARYING_FIRST_DIM, 3, 256, 128, 384, 512}, + {VARYING_FIRST_DIM, 4, 160, 128, 384, 512, 512}, +}; + +} // namespace + +class GroupedSwigluQuantizeMXFP8TestSuite : public ::testing::TestWithParam + , // Config + bool, // GEMM-swizzled scales + transformer_engine::DType, // InputType + transformer_engine::DType // OutputType + >> {}; + +TEST_P(GroupedSwigluQuantizeMXFP8TestSuite, Test) { + // Skip tests for pre-Blackwell architectures + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const std::vector config = std::get<0>(GetParam()); + const bool with_gemm_swizzled_scales = std::get<1>(GetParam()); + const DType input_type = std::get<2>(GetParam()); + const DType output_type = std::get<3>(GetParam()); + + const ShapeRepresentation shape_rep = static_cast(config[0]); + const size_t num_tensors = config[1]; + const size_t F = config[2]; + const std::vector rows_per_tensor(config.begin() + 3, config.end()); + + // The swizzled layout tiles the scale matrix 128-wide along F, and each expert owns a + // block sized by its own token count. Configs that violate either requirement must be + // rejected by the launcher rather than silently produce a wrong layout. + const bool expect_rejection = with_gemm_swizzled_scales + && ((F % 128 != 0) + || (num_tensors > 1 && shape_rep == SAME_BOTH_DIMS)); + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(output_type, OutputType, + performTest(shape_rep, num_tensors, rows_per_tensor, F, + with_gemm_swizzled_scales, expect_rejection); + ); + ); +} + +namespace { + +std::string MakeGroupedSwigluQuantizeMXFP8TestName( + const testing::TestParamInfo& info) { + const std::vector config = std::get<0>(info.param); + + std::string name; + switch (static_cast(config[0])) { + case ShapeRepresentation::SAME_BOTH_DIMS: name = "SAME_BOTH_DIMS"; break; + case ShapeRepresentation::VARYING_FIRST_DIM: name = "VARYING_FIRST_DIM"; break; + } + + name += "_N_" + std::to_string(config[1]); + name += "_F_" + std::to_string(config[2]); + for (size_t i = 3; i < config.size(); ++i) { + name += (i == 3 ? "_ROWS_" : "X") + std::to_string(config[i]); + } + + name += std::get<1>(info.param) ? "_SWIZZLED" : "_COMPACT"; + name += "_" + test::typeName(std::get<2>(info.param)) + + "_" + test::typeName(std::get<3>(info.param)); + + return name; +} + +} // namespace + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_GroupedSwigluQuantizeMXFP8_Shapes, + GroupedSwigluQuantizeMXFP8TestSuite, + ::testing::Combine( + ::testing::ValuesIn(input_configs), + ::testing::Values(false, true), + ::testing::Values(DType::kBFloat16), + ::testing::Values(DType::kFloat8E4M3)), + MakeGroupedSwigluQuantizeMXFP8TestName); + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_GroupedSwigluQuantizeMXFP8_Dtypes, + GroupedSwigluQuantizeMXFP8TestSuite, + ::testing::Combine( + ::testing::ValuesIn(input_configs_small), + ::testing::Values(false, true), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)), + MakeGroupedSwigluQuantizeMXFP8TestName); diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index eeb6e7a394..ebbab6c86a 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -637,6 +637,57 @@ def test_group_quantize_precomputed_offsets(self, output_dbias: bool) -> None: assert torch.equal(grouped_output.rowwise_data, expected_output.rowwise_data) assert torch.equal(grouped_output.scale_inv, expected_output.scale_inv) + @pytest.mark.parametrize("optimize_for_gemm", [False, True]) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_swiglu_quantize_shapes(self, optimize_for_gemm: bool) -> None: + """Test the grouped weighted-SwiGLU MXFP8 recompute binding plumbs shapes/dtypes. + + Numerics live in tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu; this only + covers the pybind layer: a [T, 2F] input plus a [T] prob must come back as a + columnwise-MXFP8 [T, F] grouped output. + """ + num_tensors = 3 + last_dim = 256 + split_sizes_list = [128, 384, 512] + total_tokens = sum(split_sizes_list) + + input_2f = torch.randn(total_tokens, 2 * last_dim, dtype=torch.bfloat16, device="cuda") + # prob rides in the model dtype, matching TE's cuDNN fc1_prob_tensor convention. + prob = torch.rand(total_tokens, dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor(split_sizes_list, dtype=torch.int64, device="cuda") + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=False, columnwise=True) + quantizer.optimize_for_gemm = optimize_for_gemm + + grouped_output = tex.group_swiglu_quantize( + input_2f, prob, quantizer, num_tensors, first_dims + ) + + outputs = grouped_output.split_into_quantized_tensors() + assert len(outputs) == num_tensors + for rows, output in zip(split_sizes_list, outputs): + assert output.shape == (rows, last_dim) + assert output._columnwise_data.numel() == rows * last_dim + # One e8m0 exponent per 32-row block of every column. Both the compact and the + # GEMM-swizzled layout need the same number of scales. + assert output._columnwise_scale_inv.numel() == (rows // 32) * last_dim + + with pytest.raises(RuntimeError): + tex.group_swiglu_quantize(input_2f, prob.float(), quantizer, num_tensors, first_dims) + + # Both operands reach the kernel as raw pointers over a densely packed range, so a + # strided view must be rejected instead of being read as if it were contiguous. + wide = torch.randn(total_tokens, 4 * last_dim, dtype=torch.bfloat16, device="cuda") + with pytest.raises(RuntimeError): + tex.group_swiglu_quantize( + wide[:, : 2 * last_dim], prob, quantizer, num_tensors, first_dims + ) + + strided_prob = torch.rand(2 * total_tokens, dtype=torch.bfloat16, device="cuda")[::2] + with pytest.raises(RuntimeError): + tex.group_swiglu_quantize(input_2f, strided_prob, quantizer, num_tensors, first_dims) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_bgrad_group_quantize_zero_size_tensor(self) -> None: """Test bgrad_group_quantize handles zero-row input without error.""" diff --git a/transformer_engine/common/activation/swiglu_grouped.cu b/transformer_engine/common/activation/swiglu_grouped.cu index 160ab66288..43ed201d41 100644 --- a/transformer_engine/common/activation/swiglu_grouped.cu +++ b/transformer_engine/common/activation/swiglu_grouped.cu @@ -15,6 +15,15 @@ void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu stream); } +void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_swiglu_quantize); + using namespace transformer_engine; + // Weighted-SwiGLU recompute: (silu(act) * gate) * prob -> columnwise MXFP8. + dispatch::group_swiglu_quantize_fwd_helper>(input, prob, output, nullptr, + stream); +} + void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_dsilu); diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 033d464bcf..c56c11e85d 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -21,6 +21,7 @@ #include "../fp8/quantize_fp8.cuh" #include "../fp8_blockwise/group_quantize_fp8_blockwise.cuh" #include "../mxfp8/group_quantize_mxfp8.cuh" +#include "../mxfp8/group_swiglu_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" #include "../nvfp4/quantize_4over6_nvfp4.cuh" @@ -498,6 +499,46 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor } } +// Grouped weighted-SwiGLU recompute: input [T, 2F] ([act|gate]) + prob [T] +// -> columnwise MXFP8 of (silu(act) * gate) * prob. +template +void group_swiglu_quantize_fwd_helper(const NVTEGroupedTensor input, const NVTETensor prob, + NVTEGroupedTensor output, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + using namespace detail; + + NVTEScalingMode scaling_mode = nvte_grouped_tensor_scaling_mode(output); + + const GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); + const Tensor *prob_tensor = convertNVTETensorCheck(prob); + + // Quantization config + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Noop flag (graph-safe skip) + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + mxfp8::group_swiglu_quantize(input_tensor, prob_tensor, noop_tensor, + output_tensor, &quant_config_cpp, stream); + break; + } + default: + NVTE_ERROR("group_swiglu_quantize only supports NVTE_MXFP8_1D_SCALING, got: " + + to_string(scaling_mode) + "."); + } +} + template void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, NVTEGroupedTensor output, NVTEGroupedTensor dbias, diff --git a/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh new file mode 100644 index 0000000000..a7f0366788 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh @@ -0,0 +1,510 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_swiglu_quantize_mxfp8.cuh + * \brief Grouped weighted-SwiGLU fused with columnwise MXFP8 quantization. + * + * MoE backward recompute of the FC2 input, without re-running the FC1 GEMM: + * + * input : FC1 output, grouped, logical shape [T, 2F] (last dim = [act|gate]). + * prob : per-token router weight, [T], in the input dtype. + * output : columnwise-MXFP8 of (silu(act) * gate) * prob, grouped [T, F]. + * + * "SwiGLU" is TE's gated convention (same as gated_mxfp8.cuh): the first half of + * the last dim is the activation input, the second half is the gate, i.e. + * swiglu(x) = silu(x[:, :F]) * x[:, F:]. "weighted" is the per-token prob factor, + * applied after the activation. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_SWIGLU_QUANTIZE_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_GROUP_SWIGLU_QUANTIZE_MXFP8_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/cuda_runtime.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "../core/common.cuh" +#include "swizzle.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace group_swiglu_quantize_kernel { + +using namespace dispatch::common; + +// Reuse the same tiling as group_quantize_mxfp8 so the scheduler/TMA math match. +struct TunableConfig { + static constexpr uint CHUNK_DIM_Y = 128; + static constexpr uint CHUNK_DIM_X = 128; + static constexpr uint THREADS_PER_CHUNK = 128; + static constexpr uint STATIC_PERSISTENT_BLOCKS_PER_SM = 24; +}; + +constexpr size_t SCALE_DIM_Y = 32; +constexpr size_t SCALE_DIM_X = 32; + +constexpr uint PREFETCH_STAGES = 1; +constexpr uint BUFFS_NUM = PREFETCH_STAGES + 1; + +constexpr uint CHUNK_DIM_Y = TunableConfig::CHUNK_DIM_Y; +constexpr uint CHUNK_DIM_X = TunableConfig::CHUNK_DIM_X; +constexpr uint THREADS_PER_CHUNK = TunableConfig::THREADS_PER_CHUNK; + +constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; + +constexpr uint THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; +constexpr uint THREADS_Y = THREADS_PER_CHUNK / THREADS_X; + +constexpr uint BUFF_DIM_Y = THREADS_Y; +constexpr uint BUFF_DIM_X = CHUNK_DIM_X; +constexpr uint BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; +static_assert(BUFF_DIM_Y == 32); + +constexpr uint STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; +static_assert(STAGES >= 1); +static_assert(CHUNK_DIM_Y % BUFF_DIM_Y == 0); +static_assert(CHUNK_DIM_Y % SCALE_DIM_Y == 0); +static_assert(CHUNK_DIM_X % SCALE_DIM_X == 0); + +// Columnwise weighted-SwiGLU + MXFP8 quantization of one 32-row buffer slice. +// Each thread owns one column j and reduces amax over the BUFF_DIM_Y rows, then +// writes the e8m0 block scale and the scaled FP8 column. +template +__device__ __forceinline__ void process_colwise_gated_stage( + const size_t buff, const int stage, const size_t tid_X_colwise, + const size_t scales_offset_Y_colwise, const size_t scales_offset_X_colwise, + const size_t scale_stride_colwise, const size_t tensor_base_for_scales, const size_t rows, + const size_t cols, const size_t data_row_base, const IType *const __restrict__ prob_ptr, + IType *sInAct_ptr, IType *sInGate_ptr, OType *sOutColwise_ptr, e8m0_t *scales_colwise) { + using IType3D = IType[BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; + using OType3D = OType[BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; + + const auto &sInAct = *reinterpret_cast(sInAct_ptr); + const auto &sInGate = *reinterpret_cast(sInGate_ptr); + auto &sOutColwise = *reinterpret_cast(sOutColwise_ptr); + + const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; + const size_t global_scales_offset_X = scales_offset_X_colwise; + const bool colwise_scale_is_within_bounds = global_scales_offset_X < cols; + + size_t scale_idx = 0; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // The FC2 wgrad GEMM consumes this operand transposed, so its scale matrix + // is the [cols, rows/32] transpose of the compact one, tiled 128x4. Each + // expert gets its own swizzled block, sized exactly like its compact block + // because per-expert row counts are 128-aligned. + const size_t tensor_base_row = tensor_base_for_scales / cols; + const size_t tensor_scales_offset_Y_base = tensor_base_row / SCALE_DIM_Y; + const size_t tensor_scales_base = tensor_base_row * scale_stride_colwise / SCALE_DIM_Y; + const size_t local_scales_offset_Y = global_scales_offset_Y - tensor_scales_offset_Y_base; + scale_idx = tensor_scales_base + + swizzle::gemm_swizzled_scale_idx( + global_scales_offset_X, local_scales_offset_Y, + DIVUP(rows, static_cast(scale_tensor_alignment_Y_rowwise))); + } else { + scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + } + + const size_t j = tid_X_colwise; + + float rInCompute[BUFF_DIM_Y]; + float thread_amax = 0.0f; +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + const float act_elt = static_cast(sInAct[buff][i][j]); + const float gate_elt = static_cast(sInGate[buff][i][j]); + // is_job_valid guarantees every row of a valid 128-aligned block is a real + // token of this expert, so the absolute token index is always in [0, T). + // prob rides along in the input (model) dtype, matching cuDNN fc1_prob_tensor. + const float prob = static_cast(prob_ptr[data_row_base + i]); + + float elt = OP(act_elt, {}) * gate_elt * prob; + + // Match round-trip precision of the plain quantize path (cast through IType). + if constexpr (!std::is_same_v) { + elt = static_cast(static_cast(elt)); + } + thread_amax = fmaxf(thread_amax, fabsf(elt)); + rInCompute[i] = elt; + } + + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + scales_colwise[scale_idx] = + colwise_scale_is_within_bounds ? biased_exponent : static_cast(0); + + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); +#pragma unroll + for (int i = 0; i < SCALE_DIM_Y; ++i) { + sOutColwise[buff][i][j] = static_cast(rInCompute[i] * block_scale_inverse); + } +} + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8_kernel( + const __grid_constant__ CUtensorMap tensor_map_input_act_static, + const __grid_constant__ CUtensorMap tensor_map_input_gate_static, + const __grid_constant__ CUtensorMap tensor_map_output_colwise_static, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, const IType *const __restrict__ prob_ptr, + e8m0_t *const __restrict__ scales_colwise_ptr, const float *__restrict__ noop, + const size_t work_blocks_X, const size_t work_blocks_Y) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + constexpr ShapeRepresentation shape_rep = SHAPE_REP; + constexpr bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); + // The shape-rep switch instantiates this kernel for all four reps, but only the + // single-tensor ones are ever dispatched. Compile the others as no-ops. + if constexpr (!is_single_tensor) { + return; + } else { + const bool leading_thread = (threadIdx.x == 0); + + const size_t tid_X_colwise = threadIdx.x; + + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + + // shmem layout: [act input][gate input][colwise output] + extern __shared__ unsigned char dynamic_shmem[]; + unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem); + + IType *sInAct_ptr = reinterpret_cast(dshmem); + IType *sInGate_ptr = reinterpret_cast(dshmem + buff_size_aligned_in); + OType *sOutColwise_ptr = reinterpret_cast(dshmem + 2 * buff_size_aligned_in); + + // Per-buffer byte count transferred by TMA (act + gate) into one slice. + constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + const size_t total_work_blocks = work_blocks_X * work_blocks_Y; + const size_t launch_block_id = blockIdx.y * gridDim.x + blockIdx.x; + + int IN_buff_readable_parity[BUFFS_NUM] = {0}; + + if (launch_block_id >= total_work_blocks) { + return; + } + int32_t ctaid_X = static_cast(launch_block_id % work_blocks_X); + int32_t ctaid_Y = static_cast(launch_block_id / work_blocks_X); + size_t static_block_stride = gridDim.x * gridDim.y; + size_t static_next_block_id = launch_block_id + static_block_stride; + + bool job_finished = false; + + __shared__ uint64_t IN_buff_readable_mbar[BUFFS_NUM]; + initialize_barriers(IN_buff_readable_mbar, leading_thread); + + while (!job_finished) { + const JobDescriptor current_job = decode_job( + num_tensors, first_logical_dim, last_logical_dim, work_blocks_X, ctaid_X, ctaid_Y, + offsets_ptr, first_dims_ptr, last_dims_ptr); + const bool current_job_is_valid = + is_job_valid(current_job, total_work_blocks, offsets_ptr); + if (!current_job_is_valid) { + break; + } + if (!job_has_work(current_job)) { + advance_to_next_job(job_finished, ctaid_X, ctaid_Y, static_next_block_id, + static_block_stride, total_work_blocks, work_blocks_X); + continue; + } + + const size_t rows = current_job.rows; + const size_t cols = current_job.cols; + const BlockDescriptor current_block = + decode_block(current_job, offsets_ptr); + + const size_t scale_alignment_X_colwise = + static_cast(scale_tensor_alignment_X_colwise); + const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); + + // Only the swizzled layout needs the per-expert base; offsets_ptr may be null + // otherwise (SAME_BOTH_DIMS), so keep the read inside the constexpr branch. + size_t tensor_base_for_scales = 0; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + tensor_base_for_scales = (num_tensors > 1) + ? static_cast(offsets_ptr[current_job.tensor_id]) + : current_block.tensor_base; + } + + const size_t block_id_Y = current_block.block_id_Y; + const size_t block_id_X = current_block.block_id_X; + const size_t block_offset_Y = current_block.block_offset_Y; + const size_t block_offset_X = current_block.block_offset_X; + + const size_t scales_block_offset_Y_colwise = block_id_Y * CHUNK_DIM_Y / SCALE_DIM_Y; + const size_t scales_block_offset_X_colwise = block_id_X * CHUNK_DIM_X; + const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise; + const size_t scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; + + __syncthreads(); + + int buff_in = 0; + +// Prime the pipeline with the first PREFETCH_STAGES slices (act + gate). +#pragma unroll + for (int stage = 0; stage < PREFETCH_STAGES; ++stage) { + const size_t buff = stage; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t buff_offset = buff * BUFF_DIM; + uint64_t *barrier = &IN_buff_readable_mbar[buff]; + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(barrier, 2 * shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sInAct_ptr[buff_offset]), + reinterpret_cast(&tensor_map_input_act_static), global_offset_X, + global_offset_Y, barrier); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sInGate_ptr[buff_offset]), + reinterpret_cast(&tensor_map_input_gate_static), global_offset_X, + global_offset_Y, barrier); + } + } + +#pragma unroll + for (int stage = 0; stage < STAGES; ++stage) { + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + if (stage < STAGES - PREFETCH_STAGES) { + const size_t next_prefetch_buff = (buff_in + PREFETCH_STAGES) % BUFFS_NUM; + const size_t next_prefetch_stage = stage + PREFETCH_STAGES; + const size_t next_prefetch_stage_offset_Y = next_prefetch_stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + next_prefetch_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_prefetch_buff_offset = next_prefetch_buff * BUFF_DIM; + uint64_t *barrier = &IN_buff_readable_mbar[next_prefetch_buff]; + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(barrier, 2 * shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sInAct_ptr[next_prefetch_buff_offset]), + reinterpret_cast(&tensor_map_input_act_static), global_offset_X, + global_offset_Y, barrier); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&sInGate_ptr[next_prefetch_buff_offset]), + reinterpret_cast(&tensor_map_input_gate_static), global_offset_X, + global_offset_Y, barrier); + } + } + + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&IN_buff_readable_mbar[buff_in], + IN_buff_readable_parity[buff_in]); + IN_buff_readable_parity[buff_in] ^= 1; + ptx::cp_async_bulk_wait_group_read(); + + const size_t buff = buff_in; + const size_t data_row_base = block_offset_Y + stage_offset_Y; + process_colwise_gated_stage( + buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, + scale_stride_colwise, tensor_base_for_scales, rows, cols, data_row_base, prob_ptr, + sInAct_ptr, sInGate_ptr, sOutColwise_ptr, scales_colwise_ptr); + + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); + + const size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t buff_offset = buff * BUFF_DIM; + if (leading_thread) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_colwise_static), + global_offset_X, global_offset_Y, + reinterpret_cast(&sOutColwise_ptr[buff_offset])); + ptx::cp_async_bulk_commit_group(); + } + + buff_in = (buff_in + 1) % BUFFS_NUM; + } + + advance_to_next_job(job_finished, ctaid_X, ctaid_Y, static_next_block_id, static_block_stride, + total_work_blocks, work_blocks_X); + } + + destroy_barriers(IN_buff_readable_mbar, leading_thread); + } // if constexpr (is_single_tensor) +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +} // namespace group_swiglu_quantize_kernel + +// Host launcher: grouped weighted-SwiGLU -> columnwise MXFP8. +// input : GroupedTensor [T, 2F] ([act|gate]) in a floating input dtype. +// prob : Tensor [T] per-token weights, in the input (model) dtype. +// output : GroupedTensor with columnwise_data / columnwise_scale_inv for [T, F]. +template +void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const Tensor *noop, + GroupedTensor *output, const QuantizationConfig *quant_config, + cudaStream_t stream) { + using namespace group_swiglu_quantize_kernel; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + + NVTE_CHECK(output->has_columnwise_data(), + "group_swiglu_quantize requires columnwise output data to be allocated."); + NVTE_CHECK(!output->has_data(), + "group_swiglu_quantize produces a columnwise output only; " + "rowwise is not implemented."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must be same."); + NVTE_CHECK(input->has_data(), "Cannot quantize tensor without rowwise data."); + + // Determine grouped shape representation from the output metadata. + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (output->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else { + NVTE_CHECK(false, + "group_swiglu_quantize requires all experts to share the same last dim F " + "(grouped layout SAME_BOTH_DIMS or VARYING_FIRST_DIM)."); + } + + const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + + // Output logical shape drives the schedule ([T, F]); input is [T, 2F]. + const size_t first_logical_dim = output->logical_shape.data[0]; // T + const size_t out_last_logical_dim = output->logical_shape.data[1]; // F + const size_t in_last_logical_dim = input->logical_shape.data[1]; // 2F + + NVTE_CHECK(in_last_logical_dim == 2 * out_last_logical_dim, + "group_swiglu_quantize input last dim must be 2x the output last dim ([act|gate])."); + NVTE_CHECK(input->logical_shape.data[0] == first_logical_dim, + "group_swiglu_quantize input/output must share the token dimension T."); + + const size_t T = first_logical_dim; + const size_t F = out_last_logical_dim; + const size_t num_tensors = input->num_tensors; + + NVTE_CHECK(prob != nullptr && prob->data.dptr != nullptr, "prob tensor must be allocated."); + // prob follows TE's cuDNN fc1_prob_tensor convention: model (input) dtype. + NVTE_CHECK(prob->data.dtype == input->dtype(), + "prob tensor must have the same dtype as the input (model dtype)."); + NVTE_CHECK(prob->data.numel() >= T, "prob tensor must have at least T elements."); + + // Single-tensor schedule: one virtual work grid over [T, F]. + const size_t work_blocks_Y = DIVUP(T, static_cast(CHUNK_DIM_Y)); + const size_t work_blocks_X = DIVUP(F, static_cast(CHUNK_DIM_X)); + + NVTE_CHECK(T % 128 == 0, "group_swiglu_quantize requires T divisible by 128."); + + const size_t sm_num = static_cast(transformer_engine::cuda::sm_count()); + const size_t static_grid_size = sm_num * TunableConfig::STATIC_PERSISTENT_BLOCKS_PER_SM; + NVTE_CHECK(static_grid_size > 0, "Static persistent grid size must be greater than zero."); + const dim3 grid(static_grid_size); + const size_t block_size = THREADS_PER_CHUNK; + + const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(output->last_dims.dptr); + + if (with_gemm_swizzled_scales) { + // The swizzled block is tiled 128x4 over the transposed [F, rows/32] scale + // matrix, so a partial F tile would not map onto a whole number of tiles. + NVTE_CHECK(F % 128 == 0, + "group_swiglu_quantize with GEMM-swizzled scales requires the output " + "last dim (F) to be divisible by 128, got ", + F, "."); + if (num_tensors > 1) { + // Each expert owns a separate swizzled block whose extent depends on its + // own token count, so per-expert first dims and offsets are mandatory. + NVTE_CHECK(shape_rep == ShapeRepresentation::VARYING_FIRST_DIM, + "group_swiglu_quantize with GEMM-swizzled scales and multiple experts " + "requires per-expert first dims (pass first_dims / split_sections)."); + NVTE_CHECK(offsets_ptr != nullptr, + "group_swiglu_quantize with GEMM-swizzled scales requires tensor_offsets " + "to locate each expert's swizzled scale block."); + } + } + + const float *const noop_ptr = reinterpret_cast(noop->data.dptr); + e8m0_t *const scales_colwise_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + NVTE_CHECK(scales_colwise_ptr != nullptr, "Columnwise scaling tensor must be allocated"); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH( + shape_rep, SHAPE_REP, + { + alignas(64) CUtensorMap tensor_map_input_act{}; + alignas(64) CUtensorMap tensor_map_input_gate{}; + alignas(64) CUtensorMap tensor_map_output_colwise{}; + + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; + + const IType *const prob_dptr = reinterpret_cast(prob->data.dptr); + + // act half: [T, F] view of the [T, 2F] buffer, stride 2F, offset 0. + create_2D_tensor_map(tensor_map_input_act, input->data, T, F, BUFF_DIM_Y, + BUFF_DIM_X, 2 * F, 0, input_type_bit_size); + // gate half: same view, offset F. + create_2D_tensor_map(tensor_map_input_gate, input->data, T, F, BUFF_DIM_Y, + BUFF_DIM_X, 2 * F, F, input_type_bit_size); + // colwise output: [T, F] contiguous, stride F. + create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, T, F, + BUFF_DIM_Y, BUFF_DIM_X, F, 0, output_type_bit_size); + + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; + constexpr size_t output_buff_size = + (buff_elems_total * output_type_bit_size) / 8; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + // [act][gate][colwise out] + const size_t dshmem_size = + 2 * buff_size_aligned_in + buff_size_aligned_out + TMA_SHMEM_ALIGNMENT; + + auto kernel = + group_swiglu_quantize_mxfp8_kernel; + + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input_act, tensor_map_input_gate, tensor_map_output_colwise, + num_tensors, T, F, offsets_ptr, first_dims_ptr, last_dims_ptr, prob_dptr, + scales_colwise_ptr, noop_ptr, work_blocks_X, work_blocks_Y); + + NVTE_CHECK_CUDA(cudaGetLastError()); + }); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) +} + +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine +#endif // TRANSFORMER_ENGINE_GROUP_SWIGLU_QUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 4ed083740d..c80b10c947 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -85,6 +85,27 @@ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream); */ void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); +/*! \brief Grouped weighted-SwiGLU "recompute" fused with MXFP8 columnwise quantization. + * + * Computes, per token t and feature f: + * output[t, f] = ( silu(input[t, f]) * input[t, F + f] ) * prob[t] + * where the grouped input has logical shape [T, 2F] (last dim = [act | gate]) and + * the grouped output has logical shape [T, F]. Only the columnwise MXFP8 output is + * produced (it feeds the MoE FC2 weight-gradient GEMM). Restrictions: + * NVTE_MXFP8_1D_SCALING output, uniform F across experts (SAME_BOTH_DIMS / + * VARYING_FIRST_DIM), per-expert token counts divisible by 128. Scales may be + * compact or in the cuBLAS GEMM-swizzled layout; the swizzled layout additionally + * requires F divisible by 128 and, for multiple experts, VARYING_FIRST_DIM. + * + * \param[in] input Grouped input tensor [T, 2F] ([act|gate]). + * \param[in] prob Per-token weights, at least T elements, in the same + * dtype as \p input. + * \param[in,out] output Grouped output tensor [T, F] (columnwise MXFP8). + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob, + NVTEGroupedTensor output, cudaStream_t stream); + /*! \brief Computes the ReLU activation of the input. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..baebbdd370 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -352,6 +352,12 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const std::optional tensor_offsets, std::optional noop_flag); +py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional last_dims, + std::optional tensor_offsets); + py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8d77a9e349..260e430298 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -21,6 +21,7 @@ #include "common/common.h" #include "common/util/system.h" #include "pybind.h" +#include "transformer_engine/activation.h" #include "transformer_engine/multi_tensor.h" #include "transformer_engine/recipe.h" #include "transformer_engine/transformer_engine.h" @@ -397,6 +398,87 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const return py::reinterpret_borrow(grouped_output_py); } +py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional last_dims, + std::optional tensor_offsets) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + // Grouped weighted-SwiGLU recompute of the MoE FC2 input: + // input_2f : [T, 2F] ([act|gate]) in model dtype (bf16). + // prob : [T] per-token weights, model dtype (matches TE fc1_prob_tensor). + // output : columnwise MXFP8 of (silu(act) * gate) * prob, logical [T, F]. + NVTE_CHECK(input_2f.dim() == 2, "group_swiglu_quantize input must be 2D [T, 2F]."); + const auto T = static_cast(input_2f.size(0)); + const auto two_f = static_cast(input_2f.size(1)); + NVTE_CHECK(two_f % 2 == 0, "group_swiglu_quantize input last dim must be even (=2F)."); + const size_t F = two_f / 2; + + NVTE_CHECK(IsMXFP8Quantizers(quantizer.ptr()), + "group_swiglu_quantize only supports MXFP8 quantizers."); + NVTE_CHECK(input_2f.is_cuda(), "group_swiglu_quantize input must be a CUDA tensor."); + // Both operands are handed to the kernel as raw pointers over a densely packed + // range, so a strided view would be read as if it were contiguous. + NVTE_CHECK(input_2f.is_contiguous(), "group_swiglu_quantize input must be contiguous."); + NVTE_CHECK(prob.is_contiguous(), "group_swiglu_quantize prob must be contiguous."); + NVTE_CHECK(prob.device() == input_2f.device(), + "group_swiglu_quantize prob must be on the same device as the input."); + NVTE_CHECK(prob.numel() >= static_cast(T), + "group_swiglu_quantize prob must have at least T elements."); + NVTE_CHECK(prob.scalar_type() == input_2f.scalar_type(), + "group_swiglu_quantize prob must have the same dtype as the input (model dtype)."); + + // The grouped metadata is turned into offsets by a kernel on the guarded device below, + // and the fused kernel then indexes the input with those offsets. + auto check_metadata_device = [&input_2f](const std::optional &metadata, + const char *name) { + if (metadata.has_value()) { + NVTE_CHECK(metadata->device() == input_2f.device(), "group_swiglu_quantize ", name, + " must be on the same device as the input."); + } + }; + check_metadata_device(first_dims, "first_dims"); + check_metadata_device(last_dims, "last_dims"); + check_metadata_device(tensor_offsets, "tensor_offsets"); + + // Allocate the output and launch on the operands' device rather than on whatever + // torch.cuda.set_device last selected. + at::cuda::CUDAGuard device_guard(input_2f.device()); + + const bool empty_input_buffer = (T == 0 || F == 0); + + auto quantizer_cpp = convert_quantizer(quantizer); + + // Input GroupedTensor: [T, 2F]. + std::vector in_logical_shape = {T, two_f}; + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, in_logical_shape); + grouped_input_tensor.set_rowwise_data(input_2f.data_ptr(), + GetTransformerEngineDType(input_2f.scalar_type()), + std::vector{static_cast(input_2f.numel())}); + + // Output GroupedTensor: [T, F] (columnwise MXFP8). Driving logical_last_dim = F + // makes the allocated data/scales and the tensor_offsets F-based. + std::vector out_logical_shape = {T, F}; + auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( + num_tensors, out_logical_shape, GetTransformerEngineDType(input_2f.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, T, F); + + if (empty_input_buffer) { + return py::reinterpret_borrow(grouped_output_py); + } + + auto prob_te = makeTransformerEngineTensor(prob); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_swiglu_quantize(grouped_input_tensor.data(), prob_te.data(), + grouped_output_tensor_cpp.data(), at::cuda::getCurrentCUDAStream()); + }); + + return py::reinterpret_borrow(grouped_output_py); +} + py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7e9d114be8..30f4cf71c3 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -209,6 +209,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), py::arg("noop_flag") = py::none()); + m.def("group_swiglu_quantize", transformer_engine::pytorch::group_swiglu_quantize, + "Grouped weighted-SwiGLU recompute fused with columnwise MXFP8 quantization", + py::arg("input_2f"), py::arg("prob"), py::arg("quantizer"), py::arg("num_tensors"), + py::arg("first_dims") = py::none(), py::arg("last_dims") = py::none(), + py::arg("tensor_offsets") = py::none()); transformer_engine::pytorch::bind_quantize_with_amax_extensions(m); m.def("group_dequantize", transformer_engine::pytorch::group_dequantize, "Dequantize group tensor", py::arg("input"), py::arg("otype"));