diff --git a/tests/test_qkv_format_to_string.cpp b/tests/test_qkv_format_to_string.cpp new file mode 100644 index 0000000000..bc6656f091 --- /dev/null +++ b/tests/test_qkv_format_to_string.cpp @@ -0,0 +1,14 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include "transformer_engine/fused_attn.h" + +TEST(QKVFormatToString, BHSD) { + EXPECT_EQ(transformer_engine::to_string(NVTE_QKV_Format::NVTE_BHSD), "NVTE_BHSD"); +} diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 1ac7a36383..307b3c5ec6 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -79,6 +79,8 @@ std::string to_string(NVTE_QKV_Format format) { return "NVTE_SBHD"; case NVTE_BSHD: return "NVTE_BSHD"; + case NVTE_BHSD: + return "NVTE_BHSD"; case NVTE_THD: return "NVTE_THD"; case NVTE_BSHD_2SBHD: