Skip to content

Commit 8930052

Browse files
Oleg-Goncharovpre-commit-ci[bot]
Oleg-Goncharov
and
pre-commit-ci[bot]
authored
[Common] Improved performance of Group MXFP8 kernels (#3337)
* Improved performance of Group MXFP8 kernels Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Addressed PR review comments Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Fix Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Changed kernel launch configs from 3D grid to 1D grid Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Fixed TMA synchronization. Restored the number of tensors check Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed max() instantiation issue Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 7695bbd commit 8930052

3 files changed

Lines changed: 553 additions & 182 deletions

File tree

‎tests/cpp/operator/test_cast_mxfp8_grouped.cu‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,42 @@ std::vector<std::vector<size_t>> input_config_small = {
700700
{VARYING_FIRST_DIM, 4, 512,160, 128,0,0,256},
701701
};
702702

703+
// Focused coverage for chunks that span multiple X and Y stages. Keep this separate from the
704+
// general shape matrix to avoid multiplying these larger cases by every operation and dtype.
705+
std::vector<std::vector<size_t>> input_config_multichunk = {
706+
{SAME_BOTH_DIMS, 2, 768,384},
707+
{VARYING_FIRST_DIM, 4, 1024,384, 128,256,384,256},
708+
{VARYING_LAST_DIM, 3, 384,1152, 128,384,640},
709+
{VARYING_BOTH_DIMS, 3, 1,(128*128)+(256*384)+(384*640),
710+
128,256,384, 128,384,640},
711+
};
712+
713+
std::vector<std::vector<size_t>> input_config_same_both_multichunk = {
714+
{SAME_BOTH_DIMS, 2, 768,384},
715+
};
716+
717+
std::vector<size_t> make_max_descriptor_input_config() {
718+
constexpr size_t num_tensors = 64;
719+
std::vector<size_t> first_dims(num_tensors, 128);
720+
std::vector<size_t> last_dims(num_tensors, 128);
721+
first_dims.front() = 256;
722+
last_dims.back() = 256;
723+
724+
size_t elts_num = 0;
725+
for (size_t t = 0; t < num_tensors; ++t) {
726+
elts_num += first_dims[t] * last_dims[t];
727+
}
728+
729+
std::vector<size_t> config = {VARYING_BOTH_DIMS, num_tensors, 1, elts_num};
730+
config.insert(config.end(), first_dims.begin(), first_dims.end());
731+
config.insert(config.end(), last_dims.begin(), last_dims.end());
732+
return config;
733+
}
734+
735+
std::vector<std::vector<size_t>> input_config_max_descriptors = {
736+
make_max_descriptor_input_config(),
737+
};
738+
703739
class GroupedFusedCastMXFP8TestSuite : public ::testing::TestWithParam
704740
<std::tuple<ProcessingMethod,
705741
ActivationKind,
@@ -946,3 +982,51 @@ INSTANTIATE_TEST_SUITE_P(
946982
::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16),
947983
::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)),
948984
MakeGroupedFusedCastMXFP8TestName);
985+
986+
INSTANTIATE_TEST_SUITE_P(
987+
OperatorTest_GroupedFusedCastMXFP8_MultiChunk,
988+
GroupedFusedCastMXFP8TestSuite,
989+
::testing::Combine(
990+
::testing::Values(ProcessingMethod::CAST_ONLY),
991+
::testing::Values(ActivationKind::Identity),
992+
::testing::ValuesIn(scaling_directions),
993+
::testing::ValuesIn(input_config_multichunk),
994+
::testing::Values(DType::kBFloat16),
995+
::testing::Values(DType::kFloat8E4M3)),
996+
MakeGroupedFusedCastMXFP8TestName);
997+
998+
INSTANTIATE_TEST_SUITE_P(
999+
OperatorTest_GroupedFusedCastMXFP8_MultiChunkDBias,
1000+
GroupedFusedCastMXFP8TestSuite,
1001+
::testing::Combine(
1002+
::testing::Values(ProcessingMethod::CAST_DBIAS),
1003+
::testing::Values(ActivationKind::Identity),
1004+
::testing::ValuesIn(scaling_directions),
1005+
::testing::ValuesIn(input_config_same_both_multichunk),
1006+
::testing::Values(DType::kBFloat16),
1007+
::testing::Values(DType::kFloat8E4M3)),
1008+
MakeGroupedFusedCastMXFP8TestName);
1009+
1010+
INSTANTIATE_TEST_SUITE_P(
1011+
OperatorTest_GroupedFusedCastMXFP8_MultiChunkDBiasDAct,
1012+
GroupedFusedCastMXFP8TestSuite,
1013+
::testing::Combine(
1014+
::testing::Values(ProcessingMethod::CAST_DBIAS_DACT),
1015+
::testing::Values(ActivationKind::GeLU),
1016+
::testing::ValuesIn(scaling_directions),
1017+
::testing::ValuesIn(input_config_same_both_multichunk),
1018+
::testing::Values(DType::kBFloat16),
1019+
::testing::Values(DType::kFloat8E4M3)),
1020+
MakeGroupedFusedCastMXFP8TestName);
1021+
1022+
INSTANTIATE_TEST_SUITE_P(
1023+
OperatorTest_GroupedFusedCastMXFP8_MaxDescriptors,
1024+
GroupedFusedCastMXFP8TestSuite,
1025+
::testing::Combine(
1026+
::testing::Values(ProcessingMethod::CAST_ONLY),
1027+
::testing::Values(ActivationKind::Identity),
1028+
::testing::Values(ScalingDirection::BOTH),
1029+
::testing::ValuesIn(input_config_max_descriptors),
1030+
::testing::Values(DType::kBFloat16),
1031+
::testing::Values(DType::kFloat8E4M3)),
1032+
MakeGroupedFusedCastMXFP8TestName);

‎transformer_engine/common/cast/core/grouped_tma.cuh‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ struct alignas(128) TensorMapStorage {
3838
alignas(128) CUtensorMap act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
3939
alignas(128) CUtensorMap output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
4040
alignas(128) CUtensorMap output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
41+
size_t rows[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
42+
size_t cols[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
43+
size_t offsets[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
4144
};
4245

4346
// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs.
@@ -115,6 +118,9 @@ __global__ void __launch_bounds__(1)
115118
const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr);
116119

117120
const size_t offset_elts = offsets_ptr[tensor_id];
121+
g_tensor_maps.rows[tensor_id] = rows;
122+
g_tensor_maps.cols[tensor_id] = cols;
123+
g_tensor_maps.offsets[tensor_id] = offset_elts;
118124

119125
// Zero-sized groups: skip TMA descriptor update. The main kernel already returns
120126
// early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension

0 commit comments

Comments
 (0)