From 5ae2ae831f1d9fc15f107f24778772e729fb7294 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Tue, 4 Aug 2026 07:29:25 -0700 Subject: [PATCH 1/7] [Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel Add nvte_group_swiglu_quantize, which fuses the weighted SwiGLU activation with columnwise MXFP8 quantization over grouped (MoE) tensors. Signed-off-by: Cael Ling --- tests/cpp/operator/CMakeLists.txt | 1 + .../test_cast_mxfp8_grouped_swiglu.cu | 449 +++++++++++++++ tests/pytorch/test_grouped_tensor.py | 39 ++ .../common/activation/swiglu_grouped.cu | 9 + .../common/cast/dispatch/quantize.cuh | 41 ++ .../mxfp8/group_swiglu_quantize_mxfp8.cuh | 509 ++++++++++++++++++ .../include/transformer_engine/activation.h | 21 + transformer_engine/pytorch/csrc/extensions.h | 6 + .../pytorch/csrc/extensions/cast.cpp | 60 +++ .../pytorch/csrc/extensions/pybind.cpp | 5 + 10 files changed, 1140 insertions(+) create mode 100644 tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu create mode 100644 transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh 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..43875797cf 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -637,6 +637,45 @@ 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) + @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..8848181976 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh @@ -0,0 +1,509 @@ +/************************************************************************* + * 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..d63c1fe349 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -6,6 +6,8 @@ #include "transformer_engine/cast.h" +#include "transformer_engine/activation.h" + #include #include #include @@ -397,6 +399,64 @@ 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(prob.is_cuda(), "group_swiglu_quantize prob must be a CUDA tensor."); + 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)."); + + 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")); From 29798819fc4f40a6b96ba7f8d63341bfbe3032a0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:37:11 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../mxfp8/group_swiglu_quantize_mxfp8.cuh | 275 +++++++++--------- .../pytorch/csrc/extensions/cast.cpp | 9 +- 2 files changed, 142 insertions(+), 142 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh index 8848181976..a7f0366788 100644 --- a/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh @@ -176,168 +176,170 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8 } else { const bool leading_thread = (threadIdx.x == 0); - const size_t tid_X_colwise = threadIdx.x; + 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); + 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); + // 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); + 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; + // 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; + 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}; + 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; + 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 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 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; + 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(); + __syncthreads(); - int buff_in = 0; + 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; + 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 next_prefetch_buff_offset = next_prefetch_buff * BUFF_DIM; - uint64_t *barrier = &IN_buff_readable_mbar[next_prefetch_buff]; + 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[next_prefetch_buff_offset]), + 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[next_prefetch_buff_offset]), + reinterpret_cast(&sInGate_ptr[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(); +#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); + } + } - 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::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(); - ptx::fence_proxy_async_shared_cta(); - __syncthreads(); + 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); - 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(); + 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; } - 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); } - 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) @@ -383,9 +385,9 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const 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 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 + 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])."); @@ -448,7 +450,8 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const TRANSFORMER_ENGINE_SWITCH_CONDITION( with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH( - shape_rep, SHAPE_REP, { + 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{}; @@ -456,8 +459,7 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const 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); + 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, @@ -471,8 +473,7 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const 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 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 = @@ -498,9 +499,9 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const NVTE_CHECK_CUDA(cudaGetLastError()); }); // NOLINT(*) - ); // NOLINT(*) - ); // NOLINT(*) - ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) } } // namespace mxfp8 diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index d63c1fe349..decd519c03 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -6,8 +6,6 @@ #include "transformer_engine/cast.h" -#include "transformer_engine/activation.h" - #include #include #include @@ -23,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" @@ -432,9 +431,9 @@ py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &p // 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())}); + 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. From e0928903b78b6296abb74f8dcfc3034f4b6e9e13 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Tue, 4 Aug 2026 20:51:16 -0700 Subject: [PATCH 3/7] [PyTorch] Require contiguous, colocated operands in group_swiglu_quantize Signed-off-by: Cael Ling --- tests/pytorch/test_grouped_tensor.py | 12 ++++++++++++ transformer_engine/pytorch/csrc/extensions/cast.cpp | 12 +++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 43875797cf..ebbab6c86a 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -676,6 +676,18 @@ def test_group_swiglu_quantize_shapes(self, optimize_for_gemm: bool) -> None: 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/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index decd519c03..70357073a4 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -418,12 +418,22 @@ py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &p NVTE_CHECK(IsMXFP8Quantizers(quantizer.ptr()), "group_swiglu_quantize only supports MXFP8 quantizers."); - NVTE_CHECK(prob.is_cuda(), "group_swiglu_quantize prob must be a CUDA tensor."); + 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)."); + // 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); From 5bfd7913811b0d62fb374838c22e0dff4c5e02b3 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Tue, 4 Aug 2026 22:53:55 -0700 Subject: [PATCH 4/7] [PyTorch] Require grouped metadata tensors on the current CUDA device Signed-off-by: Cael Ling --- transformer_engine/pytorch/csrc/quantizer.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4308464c52..c05d1903c2 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -94,6 +94,13 @@ void check_grouped_metadata_tensor(const at::Tensor& metadata_tensor, const char NVTE_CHECK(metadata_tensor.is_contiguous(), metadata_name, " must be contiguous."); NVTE_CHECK(static_cast(metadata_tensor.numel()) == expected_len, metadata_name, " must have length ", expected_len, "."); + // The offsets are built and consumed by kernels on the current device's stream, and the + // grouped data is allocated there too, so metadata on another device would be read + // through a foreign pointer. + const int current_device = static_cast(at::cuda::current_device()); + const int metadata_device = static_cast(metadata_tensor.get_device()); + NVTE_CHECK(metadata_device == current_device, metadata_name, " must be on the current CUDA ", + "device (cuda:", current_device, "), but is on cuda:", metadata_device, "."); } std::optional build_grouped_tensor_offsets(const size_t num_tensors, From 8572e3aa93c56d12b4c8f665f72ab3ef459ab02e Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Tue, 4 Aug 2026 23:57:51 -0700 Subject: [PATCH 5/7] [PyTorch] Scope the grouped metadata device check to group_swiglu_quantize Signed-off-by: Cael Ling --- transformer_engine/pytorch/csrc/extensions/cast.cpp | 13 +++++++++++++ transformer_engine/pytorch/csrc/quantizer.cpp | 7 ------- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 70357073a4..260e430298 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -430,6 +430,19 @@ py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &p 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()); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index c05d1903c2..4308464c52 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -94,13 +94,6 @@ void check_grouped_metadata_tensor(const at::Tensor& metadata_tensor, const char NVTE_CHECK(metadata_tensor.is_contiguous(), metadata_name, " must be contiguous."); NVTE_CHECK(static_cast(metadata_tensor.numel()) == expected_len, metadata_name, " must have length ", expected_len, "."); - // The offsets are built and consumed by kernels on the current device's stream, and the - // grouped data is allocated there too, so metadata on another device would be read - // through a foreign pointer. - const int current_device = static_cast(at::cuda::current_device()); - const int metadata_device = static_cast(metadata_tensor.get_device()); - NVTE_CHECK(metadata_device == current_device, metadata_name, " must be on the current CUDA ", - "device (cuda:", current_device, "), but is on cuda:", metadata_device, "."); } std::optional build_grouped_tensor_offsets(const size_t num_tensors, From a5f6859c44a847b3f0b15e5708fe94dfda58c1e3 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Sat, 8 Aug 2026 06:02:13 -0700 Subject: [PATCH 6/7] [Common/PyTorch] Rename to group_scaled_swiglu and speed up the kernel lines up with TE's existing ScaledSwiGLU op, so nvte_group_swiglu_quantize becomes nvte_group_scaled_swiglu, and the kernel header, dispatch, bindings and tests follow. Also makes the kernel actually faster than the unfused path it replaces. Per-token scales are staged through shared memory once per chunk, the output buffer is single-buffered to raise occupancy, and SiLU uses an approximate exp and divide. The activation is quantized to MXFP8 immediately, so the approximation stays far below one FP8 ULP. Signed-off-by: Cael Ling --- tests/cpp/operator/CMakeLists.txt | 2 +- ... test_cast_mxfp8_grouped_scaled_swiglu.cu} | 28 ++-- tests/pytorch/test_grouped_tensor.py | 16 +- .../common/activation/swiglu_grouped.cu | 12 +- .../common/cast/dispatch/quantize.cuh | 18 +- ...xfp8.cuh => group_scaled_swiglu_mxfp8.cuh} | 155 ++++++++++++------ .../include/transformer_engine/activation.h | 6 +- transformer_engine/pytorch/csrc/extensions.h | 10 +- .../pytorch/csrc/extensions/cast.cpp | 36 ++-- .../pytorch/csrc/extensions/pybind.cpp | 4 +- 10 files changed, 171 insertions(+), 116 deletions(-) rename tests/cpp/operator/{test_cast_mxfp8_grouped_swiglu.cu => test_cast_mxfp8_grouped_scaled_swiglu.cu} (95%) rename transformer_engine/common/cast/mxfp8/{group_swiglu_quantize_mxfp8.cuh => group_scaled_swiglu_mxfp8.cuh} (76%) diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 6a2760b699..1994d98bbc 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -13,7 +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_mxfp8_grouped_scaled_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_scaled_swiglu.cu similarity index 95% rename from tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu rename to tests/cpp/operator/test_cast_mxfp8_grouped_scaled_swiglu.cu index 24dc27e3a9..13b18e2ce9 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped_scaled_swiglu.cu @@ -74,7 +74,7 @@ void compute_ref(const InputType* input, #pragma omp parallel proc_bind(spread) { - // Buffer to cache the weighted activation of one 32-element block + // Buffer to cache the scaled 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) { @@ -87,7 +87,7 @@ void compute_ref(const InputType* input, 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 + // Numerical truncation: the kernel rounds the scaled 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)); @@ -280,7 +280,7 @@ void performTest(const ShapeRepresentation shape_rep, } if (expect_rejection) { - EXPECT_THROW(nvte_group_swiglu_quantize(in_group_tensor, prob.data(), out_group_tensor, 0), + EXPECT_THROW(nvte_group_scaled_swiglu(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); @@ -307,7 +307,7 @@ void performTest(const ShapeRepresentation shape_rep, } // GPU - nvte_group_swiglu_quantize(in_group_tensor, prob.data(), out_group_tensor, 0); + nvte_group_scaled_swiglu(in_group_tensor, prob.data(), out_group_tensor, 0); NVTE_CHECK_CUDA(cudaDeviceSynchronize()); auto err = cudaGetLastError(); ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); @@ -360,14 +360,14 @@ std::vector> input_configs_small = { } // namespace -class GroupedSwigluQuantizeMXFP8TestSuite : public ::testing::TestWithParam +class GroupedScaledSwigluMXFP8TestSuite : public ::testing::TestWithParam , // Config bool, // GEMM-swizzled scales transformer_engine::DType, // InputType transformer_engine::DType // OutputType >> {}; -TEST_P(GroupedSwigluQuantizeMXFP8TestSuite, Test) { +TEST_P(GroupedScaledSwigluMXFP8TestSuite, Test) { // Skip tests for pre-Blackwell architectures if (getDeviceComputeCapability() < blackwellComputeCapability) { GTEST_SKIP(); @@ -403,8 +403,8 @@ TEST_P(GroupedSwigluQuantizeMXFP8TestSuite, Test) { namespace { -std::string MakeGroupedSwigluQuantizeMXFP8TestName( - const testing::TestParamInfo& info) { +std::string MakeGroupedScaledSwigluMXFP8TestName( + const testing::TestParamInfo& info) { const std::vector config = std::get<0>(info.param); std::string name; @@ -429,21 +429,21 @@ std::string MakeGroupedSwigluQuantizeMXFP8TestName( } // namespace INSTANTIATE_TEST_SUITE_P( - OperatorTest_GroupedSwigluQuantizeMXFP8_Shapes, - GroupedSwigluQuantizeMXFP8TestSuite, + OperatorTest_GroupedScaledSwigluMXFP8_Shapes, + GroupedScaledSwigluMXFP8TestSuite, ::testing::Combine( ::testing::ValuesIn(input_configs), ::testing::Values(false, true), ::testing::Values(DType::kBFloat16), ::testing::Values(DType::kFloat8E4M3)), - MakeGroupedSwigluQuantizeMXFP8TestName); + MakeGroupedScaledSwigluMXFP8TestName); INSTANTIATE_TEST_SUITE_P( - OperatorTest_GroupedSwigluQuantizeMXFP8_Dtypes, - GroupedSwigluQuantizeMXFP8TestSuite, + OperatorTest_GroupedScaledSwigluMXFP8_Dtypes, + GroupedScaledSwigluMXFP8TestSuite, ::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); + MakeGroupedScaledSwigluMXFP8TestName); diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index ebbab6c86a..1f30d0477e 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -639,10 +639,10 @@ def test_group_quantize_precomputed_offsets(self, output_dbias: bool) -> None: @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. + def test_group_scaled_swiglu_shapes(self, optimize_for_gemm: bool) -> None: + """Test the grouped scaled SwiGLU MXFP8 recompute binding plumbs shapes/dtypes. - Numerics live in tests/cpp/operator/test_cast_mxfp8_grouped_swiglu.cu; this only + Numerics live in tests/cpp/operator/test_cast_mxfp8_grouped_scaled_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. """ @@ -660,9 +660,7 @@ def test_group_swiglu_quantize_shapes(self, optimize_for_gemm: bool) -> None: 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 - ) + grouped_output = tex.group_scaled_swiglu(input_2f, prob, quantizer, num_tensors, first_dims) outputs = grouped_output.split_into_quantized_tensors() assert len(outputs) == num_tensors @@ -674,19 +672,19 @@ def test_group_swiglu_quantize_shapes(self, optimize_for_gemm: bool) -> None: 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) + tex.group_scaled_swiglu(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( + tex.group_scaled_swiglu( 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) + tex.group_scaled_swiglu(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: diff --git a/transformer_engine/common/activation/swiglu_grouped.cu b/transformer_engine/common/activation/swiglu_grouped.cu index 43ed201d41..1d72d26ec6 100644 --- a/transformer_engine/common/activation/swiglu_grouped.cu +++ b/transformer_engine/common/activation/swiglu_grouped.cu @@ -15,13 +15,13 @@ 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); +void nvte_group_scaled_swiglu(const NVTEGroupedTensor input, const NVTETensor prob, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_scaled_swiglu); using namespace transformer_engine; - // Weighted-SwiGLU recompute: (silu(act) * gate) * prob -> columnwise MXFP8. - dispatch::group_swiglu_quantize_fwd_helper>(input, prob, output, nullptr, - stream); + // Scaled SwiGLU recompute: (silu(act) * gate) * prob -> columnwise MXFP8. + dispatch::group_scaled_swiglu_fwd_helper>(input, prob, output, nullptr, + stream); } void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index c56c11e85d..f00645b718 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -21,7 +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/group_scaled_swiglu_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" #include "../nvfp4/quantize_4over6_nvfp4.cuh" @@ -499,13 +499,13 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor } } -// Grouped weighted-SwiGLU recompute: input [T, 2F] ([act|gate]) + prob [T] +// Grouped scaled 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) { +void group_scaled_swiglu_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); @@ -529,12 +529,12 @@ void group_swiglu_quantize_fwd_helper(const NVTEGroupedTensor input, const NVTET switch (scaling_mode) { case NVTE_MXFP8_1D_SCALING: { - mxfp8::group_swiglu_quantize(input_tensor, prob_tensor, noop_tensor, - output_tensor, &quant_config_cpp, stream); + mxfp8::group_scaled_swiglu(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: " + + NVTE_ERROR("group_scaled_swiglu only supports NVTE_MXFP8_1D_SCALING, got: " + to_string(scaling_mode) + "."); } } diff --git a/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh similarity index 76% rename from transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh rename to transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh index a7f0366788..16ed07ef91 100644 --- a/transformer_engine/common/cast/mxfp8/group_swiglu_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh @@ -4,8 +4,8 @@ * See LICENSE for license information. ************************************************************************/ -/*! \file group_swiglu_quantize_mxfp8.cuh - * \brief Grouped weighted-SwiGLU fused with columnwise MXFP8 quantization. +/*! \file group_scaled_swiglu_mxfp8.cuh + * \brief Grouped scaled SwiGLU fused with columnwise MXFP8 quantization. * * MoE backward recompute of the FC2 input, without re-running the FC1 GEMM: * @@ -15,12 +15,12 @@ * * "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, + * swiglu(x) = silu(x[:, :F]) * x[:, F:]. "scaled" 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_ +#ifndef TRANSFORMER_ENGINE_GROUP_SCALED_SWIGLU_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_GROUP_SCALED_SWIGLU_MXFP8_CUH_ #include #include @@ -38,7 +38,7 @@ namespace transformer_engine { namespace dispatch { namespace mxfp8 { -namespace group_swiglu_quantize_kernel { +namespace group_scaled_swiglu_kernel { using namespace dispatch::common; @@ -56,6 +56,13 @@ constexpr size_t SCALE_DIM_X = 32; constexpr uint PREFETCH_STAGES = 1; constexpr uint BUFFS_NUM = PREFETCH_STAGES + 1; +// Holding both the act and the gate input slice already costs this kernel 1.7x the +// shared memory of the plain quantize kernel, which is what caps resident blocks per +// SM. Single-buffering the (1 byte per element) output slice buys one more block back +// at the cost of overlapping the TMA store with the next stage's compute. +constexpr uint OUT_BUFFS_NUM = 1; +static_assert(OUT_BUFFS_NUM >= 1 && OUT_BUFFS_NUM <= BUFFS_NUM); + 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; @@ -76,19 +83,40 @@ 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. +// silu(x) = x / (1 + exp(-x)), evaluated directly on the MUFU units. +// +// Both operations are deliberately approximate. The generic path pays a full-precision +// expf and a correctly rounded reciprocal, each a range-reduction / Newton-refinement +// chain costing an order of magnitude more than the MUFU op it wraps, and none of that +// precision survives rounding to an MXFP8 output whose mantissa is at most 3 bits wide. +// +// The exp is written as inline PTX rather than `__expf` to reach the ftz variant of +// ex2. Because the non-ftz variant has to return a denormal when the result underflows, +// nvcc guards every MUFU.EX2 with a compare plus two predicated multiplies that +// evaluate ex2(arg/2)^2 in that range. That guard cannot affect this expression: +// exp(-x) only underflows for x greater than about 87, where it is below 1e-38 and is +// then added to 1.0f, whose fp32 ulp is 1.2e-7. The sum is exactly 1.0f either way, so +// flushing to zero is bit-identical here and drops three instructions per element. +__device__ __forceinline__ float silu_approx(const float x) { + constexpr float LOG2_E = 1.4426950408889634f; + float exp_neg_x; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(exp_neg_x) : "f"(-x * LOG2_E)); + return __fdividef(x, 1.0f + exp_neg_x); +} + +// Columnwise scaled 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 buff, const size_t out_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) { + const size_t cols, const float *const sProb, 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]; + using OType3D = OType[OUT_BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; const auto &sInAct = *reinterpret_cast(sInAct_ptr); const auto &sInGate = *reinterpret_cast(sInGate_ptr); @@ -124,12 +152,19 @@ __device__ __forceinline__ void process_colwise_gated_stage( 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]); + // Staged in shared memory for the whole chunk by the caller: every thread needs + // all rows, so reading it from global here would issue one broadcast load per + // row per warp on the critical path. + const float prob = sProb[stage * BUFF_DIM_Y + i]; + + float act_x; + if constexpr (OP == &silu) { + act_x = silu_approx(act_elt); + } else { + act_x = OP(act_elt, {}); + } - float elt = OP(act_elt, {}) * gate_elt * prob; + float elt = act_x * gate_elt * prob; // Match round-trip precision of the plain quantize path (cast through IType). if constexpr (!std::is_same_v) { @@ -147,13 +182,13 @@ __device__ __forceinline__ void process_colwise_gated_stage( 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); + sOutColwise[out_buff][i][j] = static_cast(rInCompute[i] * block_scale_inverse); } } template -__global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8_kernel( +__global__ void __launch_bounds__(THREADS_PER_CHUNK) group_scaled_swiglu_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, @@ -182,16 +217,21 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8 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 out_buff_elems_total = OUT_BUFFS_NUM * buff_elems; constexpr size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + DIVUP_TO_MULTIPLE(out_buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + constexpr size_t prob_buff_size = + DIVUP_TO_MULTIPLE(CHUNK_DIM_Y * sizeof(float), TMA_SHMEM_ALIGNMENT); - // shmem layout: [act input][gate input][colwise output] + // shmem layout: [act input][gate input][colwise output][prob] 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); + float *sProb_ptr = + reinterpret_cast(dshmem + 2 * buff_size_aligned_in + buff_size_aligned_out); // Per-buffer byte count transferred by TMA (act + gate) into one slice. constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -257,6 +297,14 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8 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; + // Stage this chunk's per-token prob once. 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. + for (size_t row = threadIdx.x; row < CHUNK_DIM_Y; row += THREADS_PER_CHUNK) { + sProb_ptr[row] = static_cast(prob_ptr[block_offset_Y + row]); + } + __syncthreads(); int buff_in = 0; @@ -310,26 +358,31 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8 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(); + // Wait until the store groups still holding an output slice have drained. Only + // the leading thread commits those groups, so the wait is a no-op on the other + // threads and the barrier is what stops them from overwriting a slice the TMA + // unit has not finished reading. + ptx::cp_async_bulk_wait_group_read(); + __syncthreads(); const size_t buff = buff_in; - const size_t data_row_base = block_offset_Y + stage_offset_Y; + const size_t out_buff = buff_in % OUT_BUFFS_NUM; 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); + buff, out_buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, + scale_stride_colwise, tensor_base_for_scales, rows, cols, sProb_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; + const size_t out_buff_offset = out_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])); + reinterpret_cast(&sOutColwise_ptr[out_buff_offset])); ptx::cp_async_bulk_commit_group(); } @@ -345,25 +398,25 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_swiglu_quantize_mxfp8 #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } -} // namespace group_swiglu_quantize_kernel +} // namespace group_scaled_swiglu_kernel -// Host launcher: grouped weighted-SwiGLU -> columnwise MXFP8. +// Host launcher: grouped scaled 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; +void group_scaled_swiglu(const GroupedTensor *input, const Tensor *prob, const Tensor *noop, + GroupedTensor *output, const QuantizationConfig *quant_config, + cudaStream_t stream) { + using namespace group_scaled_swiglu_kernel; checkCuDriverContext(stream); CheckNoopTensor(*noop, "cast_noop"); NVTE_CHECK(output->has_columnwise_data(), - "group_swiglu_quantize requires columnwise output data to be allocated."); + "group_scaled_swiglu requires columnwise output data to be allocated."); NVTE_CHECK(!output->has_data(), - "group_swiglu_quantize produces a columnwise output only; " + "group_scaled_swiglu 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, @@ -378,7 +431,7 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; } else { NVTE_CHECK(false, - "group_swiglu_quantize requires all experts to share the same last dim F " + "group_scaled_swiglu requires all experts to share the same last dim F " "(grouped layout SAME_BOTH_DIMS or VARYING_FIRST_DIM)."); } @@ -390,9 +443,9 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const 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])."); + "group_scaled_swiglu 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."); + "group_scaled_swiglu input/output must share the token dimension T."); const size_t T = first_logical_dim; const size_t F = out_last_logical_dim; @@ -408,7 +461,7 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const 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."); + NVTE_CHECK(T % 128 == 0, "group_scaled_swiglu 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; @@ -424,17 +477,17 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const // 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 " + "group_scaled_swiglu 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 " + "group_scaled_swiglu 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 " + "group_scaled_swiglu with GEMM-swizzled scales requires tensor_offsets " "to locate each expert's swizzled scale block."); } } @@ -474,20 +527,24 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const 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 out_buff_elems_total = OUT_BUFFS_NUM * buff_elems; constexpr size_t output_buff_size = - (buff_elems_total * output_type_bit_size) / 8; + (out_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; + constexpr size_t prob_buff_size = + DIVUP_TO_MULTIPLE(CHUNK_DIM_Y * sizeof(float), TMA_SHMEM_ALIGNMENT); + + // [act][gate][colwise out][prob] + const size_t dshmem_size = 2 * buff_size_aligned_in + buff_size_aligned_out + + prob_buff_size + TMA_SHMEM_ALIGNMENT; auto kernel = - group_swiglu_quantize_mxfp8_kernel; + group_scaled_swiglu_mxfp8_kernel; NVTE_CHECK_CUDA(cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); @@ -507,4 +564,4 @@ void group_swiglu_quantize(const GroupedTensor *input, const Tensor *prob, const } // namespace mxfp8 } // namespace dispatch } // namespace transformer_engine -#endif // TRANSFORMER_ENGINE_GROUP_SWIGLU_QUANTIZE_MXFP8_CUH_ +#endif // TRANSFORMER_ENGINE_GROUP_SCALED_SWIGLU_MXFP8_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index c80b10c947..f3d3fd3737 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -85,7 +85,7 @@ 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. +/*! \brief Grouped scaled 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] @@ -103,8 +103,8 @@ void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu * \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); +void nvte_group_scaled_swiglu(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, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index baebbdd370..b9bb018f9e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -352,11 +352,11 @@ 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 group_scaled_swiglu(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, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 260e430298..a7a4eb4fc5 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -398,44 +398,44 @@ 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) { +py::object group_scaled_swiglu(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: + // Grouped scaled 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]."); + NVTE_CHECK(input_2f.dim() == 2, "group_scaled_swiglu 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)."); + NVTE_CHECK(two_f % 2 == 0, "group_scaled_swiglu 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."); + "group_scaled_swiglu only supports MXFP8 quantizers."); + NVTE_CHECK(input_2f.is_cuda(), "group_scaled_swiglu 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(input_2f.is_contiguous(), "group_scaled_swiglu input must be contiguous."); + NVTE_CHECK(prob.is_contiguous(), "group_scaled_swiglu prob must be contiguous."); NVTE_CHECK(prob.device() == input_2f.device(), - "group_swiglu_quantize prob must be on the same device as the input."); + "group_scaled_swiglu 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."); + "group_scaled_swiglu 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)."); + "group_scaled_swiglu 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, + NVTE_CHECK(metadata->device() == input_2f.device(), "group_scaled_swiglu ", name, " must be on the same device as the input."); } }; @@ -472,8 +472,8 @@ py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &p 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()); + nvte_group_scaled_swiglu(grouped_input_tensor.data(), prob_te.data(), + grouped_output_tensor_cpp.data(), at::cuda::getCurrentCUDAStream()); }); return py::reinterpret_borrow(grouped_output_py); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 30f4cf71c3..2392758a8f 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -209,8 +209,8 @@ 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", + m.def("group_scaled_swiglu", transformer_engine::pytorch::group_scaled_swiglu, + "Grouped scaled 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()); From 0b94fcb3c83957979b071a0f3aa09d73747639f7 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Sat, 8 Aug 2026 06:03:47 -0700 Subject: [PATCH 7/7] [PyTorch] Add a benchmark for group_scaled_swiglu Times the fused kernel against the unfused path it replaces, with the activation half computed three ways: eager PyTorch, torch.compile, and TE's existing ScaledSwiGLU op. Signed-off-by: Cael Ling --- benchmarks/benchmark_group_scaled_swiglu.py | 571 ++++++++++++++++++++ 1 file changed, 571 insertions(+) create mode 100644 benchmarks/benchmark_group_scaled_swiglu.py diff --git a/benchmarks/benchmark_group_scaled_swiglu.py b/benchmarks/benchmark_group_scaled_swiglu.py new file mode 100644 index 0000000000..d0cff0cb08 --- /dev/null +++ b/benchmarks/benchmark_group_scaled_swiglu.py @@ -0,0 +1,571 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark for the grouped scaled-SwiGLU MXFP8 kernel (nvte_group_scaled_swiglu). + +Compares the fused kernel against the unfused path it replaces when recomputing +the MoE FC2 weight-gradient input: + + fused : [T, 2F] bf16 --(scaled SwiGLU + columnwise MXFP8)--> [T, F] fp8 + scales + unfused : [T, 2F] bf16 --(scaled SwiGLU)--> [T, F] bf16 --(group_quantize)--> fp8 + +Both paths are bandwidth bound, so the number that explains the speedup is DRAM +traffic. Per output element the unfused path moves 4 bytes reading [T, 2F], 2 +writing the bf16 activation, 2 reading it back, and 1 writing FP8; the fused +kernel keeps the activation in registers and moves 4 + 1. That is 9 vs 5 bytes +per element, so ~1.8x is the ceiling for the fused kernel. + +Every unfused variant ends in the same ``tex.group_quantize`` call with the same +quantizer the fused path uses (columnwise only), so the variants differ *only* in how +the activation is computed. That choice changes what the speedup means: + + unfused-eager plain PyTorch ops, one kernel and one DRAM round trip per + elementwise op. Much too loose to quote. It is kept because it is + the only way to see that torch.compile actually fused: a silent + Inductor fallback would drag unfused-compiled toward this number + and quietly inflate the reported speedup. + unfused-compiled the same expression through torch.compile, which fuses it into one + elementwise kernel. The tightest unfused implementation, so this is + the number to quote: against it, the fused kernel's remaining + advantage can only come from fusing the *quantization* in. + unfused-te-op the same two steps assembled from TE's existing components, i.e. + what a user falls back on today without this kernel. It is slower + than unfused-compiled for a structural reason rather than just + Python overhead: ScaledSwiGLU runs tex.swiglu and then a *separate* + kernel to apply the per-token scale, so the bf16 intermediate makes + one extra DRAM round trip. Operation-fuser overhead adds to that. + +Shapes must respect the kernel's restrictions: every expert's token count is +divisible by 128, and the GEMM-swizzled scale layout also needs F divisible by +128. The default total token count mirrors benchmark_group_quantize_current_scaling.py. + +Example: + python benchmarks/benchmark_group_scaled_swiglu.py + python benchmarks/benchmark_group_scaled_swiglu.py --hidden 4096 --num-groups 8 64 +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import dataclass +from typing import Callable, List, Optional + +# IMPORTANT: import transformer_engine before torch to avoid cublasLt symbol-resolution +# issues caused by torch's bundled CUDA libs. +import transformer_engine.pytorch # noqa: F401 - registers extension +from transformer_engine.pytorch import MXFP8Quantizer +import transformer_engine.pytorch.ops as te_ops +import transformer_engine_torch as tex +import torch + + +BF16_BYTES = 2 +FP8_BYTES = 1 +# One e8m0 exponent per 32-row block of every column. +SCALE_BLOCK_ROWS = 32 +# The kernel schedules 128-row blocks, so every expert's token count must be a multiple. +TOKEN_ALIGNMENT = 128 +# The swizzled scale layout tiles the transposed scale matrix 128-wide along F. +SWIZZLE_F_ALIGNMENT = 128 + +VARIANTS = ("fused", "unfused-eager", "unfused-compiled", "unfused-te-op") + + +@dataclass +class CaseResult: + variant: str + tokens: int + hidden: int + num_groups: int + swizzled_scales: bool + loop: str + iters: int + per_iter_us: float + min_bytes: int + bw_TBps: float + speedup_vs_fused: Optional[float] = None + + +def _distribute_blocks(blocks: int, num_groups: int, imbalance: str) -> List[int]: + """Split ``blocks`` 128-row blocks across ``num_groups`` experts, each >= 1.""" + if imbalance == "uniform": + if blocks % num_groups != 0: + raise SystemExit( + f"{blocks} blocks of {TOKEN_ALIGNMENT} tokens do not divide evenly across" + f" {num_groups} experts; pick a --tokens that is a multiple of" + f" {TOKEN_ALIGNMENT * num_groups} or use --imbalance zipf." + ) + return [blocks // num_groups] * num_groups + + if imbalance == "mild": + weights = [0.8 + 0.4 * i / max(1, num_groups - 1) for i in range(num_groups)] + elif imbalance == "zipf": + weights = [1.0 / ((i + 1) ** 0.7) for i in range(num_groups)] + else: + raise SystemExit(f"unknown imbalance={imbalance}") + + total_weight = sum(weights) + counts = [max(1, int(round(w * blocks / total_weight))) for w in weights] + + # Fix up rounding so the blocks sum back to the requested total, always keeping + # every expert at >= 1 block. + order = sorted(range(num_groups), key=lambda i: counts[i], reverse=True) + idx = 0 + while sum(counts) != blocks: + target = order[idx % num_groups] + if sum(counts) < blocks: + counts[target] += 1 + elif counts[target] > 1: + counts[target] -= 1 + idx += 1 + return counts + + +def _make_first_dims(tokens: int, num_groups: int, imbalance: str) -> List[int]: + if tokens % TOKEN_ALIGNMENT != 0: + raise SystemExit(f"--tokens must be a multiple of {TOKEN_ALIGNMENT}, got {tokens}") + blocks = _distribute_blocks(tokens // TOKEN_ALIGNMENT, num_groups, imbalance) + return [b * TOKEN_ALIGNMENT for b in blocks] + + +def _make_quantizer(swizzled_scales: bool) -> MXFP8Quantizer: + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + # The FC2 wgrad GEMM only consumes the columnwise operand. + quantizer.set_usage(rowwise=False, columnwise=True) + quantizer.optimize_for_gemm = swizzled_scales + return quantizer + + +def _scaled_swiglu_bf16(input_2f: torch.Tensor, prob: torch.Tensor, hidden: int) -> torch.Tensor: + act = input_2f[:, :hidden] + gate = input_2f[:, hidden:] + return torch.nn.functional.silu(act) * gate * prob.unsqueeze(1) + + +def _fused_bytes(tokens: int, hidden: int) -> int: + read_input = tokens * 2 * hidden * BF16_BYTES + write_fp8 = tokens * hidden * FP8_BYTES + write_scales = (tokens // SCALE_BLOCK_ROWS) * hidden + return read_input + write_fp8 + write_scales + + +def _unfused_bytes(tokens: int, hidden: int) -> int: + # Activation kernel: read [T, 2F] bf16, write the [T, F] bf16 intermediate. + activation = tokens * 2 * hidden * BF16_BYTES + tokens * hidden * BF16_BYTES + # Quantize kernel: read that intermediate back, write FP8 plus scales. + quantize = ( + tokens * hidden * BF16_BYTES + + tokens * hidden * FP8_BYTES + + (tokens // SCALE_BLOCK_ROWS) * hidden + ) + return activation + quantize + + +def _compile_activation() -> Optional[Callable]: + try: + return torch.compile(_scaled_swiglu_bf16, dynamic=False) + except Exception as exc: # torch.compile is optional for this benchmark + print(f" (torch.compile unavailable, skipping unfused-compiled: {exc})") + return None + + +def _te_op_activation() -> Optional[Callable]: + """TE's own ScaledSwiGLU as the activation half of the unfused path. + + ``glu_interleave_size=None`` keeps the contiguous ``[act | gate]`` layout the + fused kernel expects; the fused grouped MLP instead runs this op with 32-wide + interleaving, which would make the comparison a layout difference rather than a + fusion one. The op routes through the operation fuser and autograd, so it also + carries framework overhead that is not kernel time. + """ + try: + op = te_ops.ScaledSwiGLU(glu_interleave_size=None) + except Exception as exc: + print(f" (te_ops.ScaledSwiGLU unavailable, skipping unfused-te-op: {exc})") + return None + + def activation(input_2f: torch.Tensor, prob: torch.Tensor, hidden: int) -> torch.Tensor: + del hidden # the op infers F from the input + with torch.no_grad(): + return op(input_2f, prob) + + return activation + + +def _make_runner( + variant: str, + quantizer: MXFP8Quantizer, + hidden: int, + num_groups: int, + first_dims: Optional[torch.Tensor], + compiled_activation: Optional[Callable], + te_op_activation: Optional[Callable], +) -> Callable[[torch.Tensor, torch.Tensor], object]: + if variant == "fused": + return lambda x, prob: tex.group_scaled_swiglu(x, prob, quantizer, num_groups, first_dims) + + activation = { + "unfused-eager": _scaled_swiglu_bf16, + "unfused-compiled": compiled_activation, + "unfused-te-op": te_op_activation, + }[variant] + + def run(x: torch.Tensor, prob: torch.Tensor): + intermediate = activation(x, prob, hidden) + return tex.group_quantize(intermediate, quantizer, num_groups, first_dims) + + return run + + +def _time_eager(runner, inputs, probs, iters: int) -> float: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for it in range(iters): + i = it % len(inputs) + runner(inputs[i], probs[i]) + end.record() + end.synchronize() + return start.elapsed_time(end) + + +def _time_graph(runner, inputs, probs, iters: int, calls_per_replay: int = 16): + """Capture ``calls_per_replay`` calls into one CUDA graph and replay it. + + Removes Python and launch overhead, which is what makes a single memory-bound + kernel look slower than it is. + """ + static_x, static_prob = inputs[0], probs[0] + + # Warmup on a side stream before capture, per the torch CUDA-graph docs. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + runner(static_x, static_prob) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(calls_per_replay): + runner(static_x, static_prob) + + replays = max(1, iters // calls_per_replay) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(replays): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end), replays * calls_per_replay + + +def _check_fused_matches_unfused( + quantizer: MXFP8Quantizer, + input_2f: torch.Tensor, + prob: torch.Tensor, + hidden: int, + num_groups: int, + first_dims: Optional[torch.Tensor], +) -> None: + """Guard against timing a kernel that is not computing the right thing. + + This is a smoke test, not a numerics test: correctness lives in + tests/cpp/operator/test_cast_mxfp8_grouped_scaled_swiglu.cu, whose CPU reference + mirrors the kernel's arithmetic order exactly. Here the reference is plain + PyTorch, which rounds to bf16 three times (after silu, after the gate multiply, + after the prob multiply) where the kernel rounds once, for ~0.2% of relative + spread on the pre-quantization values. + + MXFP8 turns that spread into whole-block disagreements: each 32-row block shares + one e8m0 scale, and e8m0 is a pure power of two, so a block whose amax sits within + ~0.2% of a power-of-two boundary can pick a different exponent in the two paths + and shift all 32 of its codes at once. That affects on the order of + 0.002/ln(2) ~ 0.3% of blocks, hence a similar fraction of elements. The budget + below sits above that but far below the ~100% a wrong formula would produce. + """ + fused = tex.group_scaled_swiglu(input_2f, prob, quantizer, num_groups, first_dims) + reference = tex.group_quantize( + _scaled_swiglu_bf16(input_2f, prob, hidden), quantizer, num_groups, first_dims + ) + if fused.columnwise_data.numel() != reference.columnwise_data.numel(): + raise RuntimeError( + f"fused output has {fused.columnwise_data.numel()} elements but the reference has" + f" {reference.columnwise_data.numel()}; the benchmark is comparing different shapes." + ) + if int(fused.columnwise_data.view(torch.uint8).max().item()) == 0: + raise RuntimeError("fused output is entirely zero; the kernel produced nothing.") + + # Coarse code-level comparison. Signs agree between the two paths in practice, so + # treating the FP8 bytes as integers is good enough to spot a gross mismatch. + fused_codes = fused.columnwise_data.view(torch.uint8).to(torch.int16) + reference_codes = reference.columnwise_data.view(torch.uint8).to(torch.int16) + mismatch = (fused_codes - reference_codes).abs() > 1 + mismatch_rate = float(mismatch.sum().item()) / max(1, mismatch.numel()) + if mismatch_rate > 2e-2: + raise RuntimeError( + "fused output disagrees with the unfused reference on" + f" {100.0 * mismatch_rate:.3f}% of FP8 codes by more than 1 ULP, which is too much" + " to be block-scale rounding; the benchmark would be timing the wrong computation." + ) + + +def run_case( + variant: str, + *, + tokens: int, + hidden: int, + num_groups: int, + swizzled_scales: bool, + same_shape: bool, + imbalance: str, + num_buffers: int, + warmup: int, + iters: int, + loop: str, + compiled_activation: Optional[Callable], + te_op_activation: Optional[Callable], +) -> Optional[CaseResult]: + quantizer = _make_quantizer(swizzled_scales) + first_dims = None + if not same_shape: + first_dims = torch.tensor( + _make_first_dims(tokens, num_groups, imbalance), dtype=torch.int64, device="cuda" + ) + + inputs = [ + torch.randn(tokens, 2 * hidden, dtype=torch.bfloat16, device="cuda") + for _ in range(num_buffers) + ] + probs = [torch.rand(tokens, dtype=torch.bfloat16, device="cuda") for _ in range(num_buffers)] + + runner = _make_runner( + variant, quantizer, hidden, num_groups, first_dims, compiled_activation, te_op_activation + ) + + for it in range(warmup): + i = it % len(inputs) + runner(inputs[i], probs[i]) + torch.cuda.synchronize() + + if loop == "graph": + try: + elapsed_ms, actual_iters = _time_graph(runner, inputs, probs, iters) + except Exception as exc: # capture can fail e.g. for a compiled activation + print(f" skipping {variant} under CUDA-graph capture: {exc}") + return None + else: + elapsed_ms = _time_eager(runner, inputs, probs, iters) + actual_iters = iters + + min_bytes = ( + _fused_bytes(tokens, hidden) if variant == "fused" else _unfused_bytes(tokens, hidden) + ) + per_iter_us = elapsed_ms * 1000.0 / actual_iters + bw_TBps = min_bytes / (per_iter_us * 1.0e-6) / 1.0e12 + + return CaseResult( + variant=variant, + tokens=tokens, + hidden=hidden, + num_groups=num_groups, + swizzled_scales=swizzled_scales, + loop=loop, + iters=actual_iters, + per_iter_us=per_iter_us, + min_bytes=min_bytes, + bw_TBps=bw_TBps, + ) + + +def _print_table(results: List[CaseResult]) -> None: + header = ( + f"{'T x F':>14s} {'experts':>7s} {'scales':>8s} {'loop':>5s} " + f"{'variant':16s} {'per_iter_us':>11s} {'min_GB/iter':>11s} " + f"{'BW_TB/s':>8s} {'vs fused':>9s}" + ) + print() + print("=" * len(header)) + print(header) + print("-" * len(header)) + for r in results: + shape = f"{r.tokens}x{r.hidden}" + scales = "swizzled" if r.swizzled_scales else "compact" + speedup = f"{r.speedup_vs_fused:.2f}x" if r.speedup_vs_fused is not None else "-" + print( + f"{shape:>14s} {r.num_groups:7d} {scales:>8s} {r.loop:>5s} " + f"{r.variant:16s} {r.per_iter_us:11.2f} {r.min_bytes / 1e9:11.3f} " + f"{r.bw_TBps:8.2f} {speedup:>9s}" + ) + print("-" * len(header)) + print( + "min_GB/iter = minimum DRAM traffic the path must move (reads + writes);" + " BW = min_GB/iter / per_iter_us." + ) + print( + "min_GB/iter assumes one activation kernel plus one quantize kernel, so the" + " variants that use more than that read below their real bandwidth:" + " unfused-eager launches one kernel per elementwise op, and unfused-te-op" + " applies the per-token scale in a separate kernel (one extra round trip of" + " the bf16 intermediate) on top of op-fuser overhead." + ) + print( + "unfused-compiled is the tightest baseline; quote the speedup against it." + " unfused-te-op is what TE's existing components give you today." + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--tokens", + type=int, + default=98304, + help="Total tokens T summed over experts (multiple of 128). Default 98304.", + ) + parser.add_argument( + "--hidden", + type=int, + default=2048, + help="MoE intermediate size F; the input is [T, 2F]. Default 2048.", + ) + parser.add_argument("--num-groups", type=int, nargs="+", default=[16, 64]) + parser.add_argument( + "--imbalance", + choices=("uniform", "mild", "zipf"), + default="uniform", + help="Token distribution across experts. Default uniform.", + ) + parser.add_argument( + "--same-shape", + action="store_true", + help=( + "Use the SAME_BOTH_DIMS layout (no first_dims). The GEMM-swizzled scale" + " layout does not support it with more than one expert." + ), + ) + parser.add_argument( + "--scales", + choices=("compact", "swizzled", "both"), + default="both", + help="Scale layout to benchmark. Default both.", + ) + parser.add_argument( + "--variants", + nargs="+", + default=list(VARIANTS), + help=f"Subset of {VARIANTS}.", + ) + parser.add_argument( + "--loop", + choices=("eager", "graph", "both"), + default="both", + help="Python loop, replayed CUDA graph, or both. Default both.", + ) + parser.add_argument("--num-buffers", type=int, default=4) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--json-out", default=None) + args = parser.parse_args() + + for variant in args.variants: + if variant not in VARIANTS: + raise SystemExit(f"unknown variant={variant}") + if args.tokens % TOKEN_ALIGNMENT != 0: + raise SystemExit(f"--tokens must be a multiple of {TOKEN_ALIGNMENT}, got {args.tokens}") + + scale_layouts = { + "compact": [False], + "swizzled": [True], + "both": [False, True], + }[args.scales] + loop_modes = ("eager", "graph") if args.loop == "both" else (args.loop,) + + print(f"GPU: {torch.cuda.get_device_name(0)}") + print( + f"Config: T={args.tokens}, F={args.hidden}, experts={args.num_groups}," + f" imbalance={args.imbalance}," + f" layout={'SAME_BOTH_DIMS' if args.same_shape else 'VARYING_FIRST_DIM'}," + f" iters={args.iters}, warmup={args.warmup}" + ) + + compiled_activation = _compile_activation() if "unfused-compiled" in args.variants else None + te_op_activation = _te_op_activation() if "unfused-te-op" in args.variants else None + + # One correctness check up front: a benchmark of a wrong kernel is worthless. + check_tokens = min(args.tokens, 4096) + _check_fused_matches_unfused( + _make_quantizer(False), + torch.randn(check_tokens, 2 * args.hidden, dtype=torch.bfloat16, device="cuda"), + torch.rand(check_tokens, dtype=torch.bfloat16, device="cuda"), + args.hidden, + num_groups=1, + first_dims=None, + ) + + results: List[CaseResult] = [] + for num_groups in args.num_groups: + for swizzled_scales in scale_layouts: + if swizzled_scales and args.hidden % SWIZZLE_F_ALIGNMENT != 0: + print( + f" skipping swizzled scales: F={args.hidden} is not a multiple of" + f" {SWIZZLE_F_ALIGNMENT}" + ) + continue + if swizzled_scales and args.same_shape and num_groups > 1: + print( + " skipping swizzled scales with SAME_BOTH_DIMS and multiple experts" + " (unsupported: each expert needs its own swizzled block)" + ) + continue + for loop in loop_modes: + fused_us = None + group_results: List[CaseResult] = [] + for variant in args.variants: + if variant == "unfused-compiled" and compiled_activation is None: + continue + if variant == "unfused-te-op" and te_op_activation is None: + continue + result = run_case( + variant, + tokens=args.tokens, + hidden=args.hidden, + num_groups=num_groups, + swizzled_scales=swizzled_scales, + same_shape=args.same_shape, + imbalance=args.imbalance, + num_buffers=args.num_buffers, + warmup=args.warmup, + iters=args.iters, + loop=loop, + compiled_activation=compiled_activation, + te_op_activation=te_op_activation, + ) + if result is None: + continue + if variant == "fused": + fused_us = result.per_iter_us + group_results.append(result) + + if fused_us: + for result in group_results: + result.speedup_vs_fused = result.per_iter_us / fused_us + results.extend(group_results) + + _print_table(results) + + if args.json_out: + with open(args.json_out, "w") as f: + json.dump([r.__dict__ for r in results], f, indent=2) + print(f"Wrote {args.json_out}") + + +if __name__ == "__main__": + main()