fix: add missing NVTE_BHSD_BHSD_BHSD to layout to_string

[andrewwhitecdw]

This PR addresses the following issue in transformer_engine/common/fused_attn/fused_attn.cpp: add missing NVTE_BHSD_BHSD_BHSD to layout to_string.

Changes

• transformer_engine/common/fused_attn/fused_attn.cpp: add missing NVTE_BHSD_BHSD_BHSD to layout to_string.

Details

--- a/transformer_engine/common/fused_attn/fused_attn.cpp
+++ b/transformer_engine/common/fused_attn/fused_attn.cpp
@@ -1,27 +1,31 @@
-    case NVTE_Paged_KV_THD_SBHD_SBHD:
-      return "NVTE_Paged_KV_THD_SBHD_SBHD";
-    default:
-      return "UNKNOWN_QKV_LAYOUT(" + std::to_string(static_cast<int>(layout)) + ")";
-  }
-}
-
-std::string to_string(NVTE_QKV_Format format) {
-  switch (format) {
-    case NVTE_SBHD:
-      return "NVTE_SBHD";
-    case NVTE_BSHD:
-      return "NVTE_BSHD";
-    case NVTE_THD:
-      return "NVTE_THD";
-    case NVTE_BSHD_2SBHD:
-      return "NVTE_BSHD_2SBHD";
-    case NVTE_SBHD_2BSHD:
-      return "NVTE_SBHD_2BSHD";
-    case NVTE_THD_2BSHD:
-      return "NVTE_THD_2BSHD";
-    case NVTE_THD_2SBHD:
-      return "NVTE_THD_2SBHD";
-    default:
-      return "UNKNOWN_QKV_FORMAT(" + std::to_string(static_cast<int>(format)) + ")";
-  }
-}
+    case NVTE_Paged_KV_THD_SBHD_SBHD:
+      return "NVTE_Paged_KV_THD_SBHD_SBHD";
+    case NVTE_BHSD_BHSD_BHSD:
+      return "NVTE_BHSD_BHSD_BHSD";
+    default:
+      return "UNKNOWN_QKV_LAYOUT(" + std::to_string(static_cast<int>(layout)) + ")";
+  }
+}
+
+std::string to_string(NVTE_QKV_Format format) {
+  switch (format) {
+    case NVTE_SBHD:
+      return "NVTE_SBHD";
+    case NVTE_BSHD:
+      return "NVTE_BSHD";
+    case NVTE_THD:
+      return "NVTE_THD";
+    case NVTE_BSHD_2SBHD:
+      return "NVTE_BSHD_2SBHD";
+    case NVTE_SBHD_2BSHD:
+      return "NVTE_SBHD_2BSHD";
+    case NVTE_THD_2BSHD:
+      return "NVTE_THD_2BSHD";
+    case NVTE_THD_2SBHD:
+      return "NVTE_THD_2SBHD";
+    case NVTE_BHSD:
+      return "NVTE_BHSD";
+    default:
+      return "UNKNOWN_QKV_FORMAT(" + std::to_string(static_cast<int>(format)) + ")";
+  }
+}

Tests

• tests/cpp/test_fused_attn.cpp

--- a/tests/cpp/test_fused_attn.cpp
+++ b/tests/cpp/test_fused_attn.cpp
@@ -0,0 +1,16 @@
+#include <gtest/gtest.h>
+#include "transformer_engine/fused_attn.h"
+
+namespace transformer_engine {
+namespace test {
+
+TEST(FusedAttnUtils, QkvLayoutToStringBhsd) {
+  EXPECT_EQ(to_string(NVTE_BHSD_BHSD_BHSD), "NVTE_BHSD_BHSD_BHSD");
+}
+
+TEST(FusedAttnUtils, QkvFormatToStringBhsd) {
+  EXPECT_EQ(to_string(NVTE_BHSD), "NVTE_BHSD");
+}
+
+}  // namespace test
+}  // namespace transformer_engine

[greptile-apps]

Greptile Summary

This PR completes fused-attention diagnostic string conversion for the existing BHSD layout and format enum values.

• Adds NVTE_BHSD_BHSD_BHSD handling to to_string(NVTE_QKV_Layout).
• Adds NVTE_BHSD handling to to_string(NVTE_QKV_Format).

Confidence Score: 5/5

The PR appears safe to merge because the added cases match the existing enum declarations and BHSD mapping behavior.

The change only completes diagnostic string conversion for two existing, uniquely valued enums and does not alter fused-attention dispatch, layout conversion, or tensor processing.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds correct, non-conflicting string representations for the existing BHSD layout and format enum values; no actionable issue identified.

_{Reviews (1): Last reviewed commit: "fix: add missing NVTE_BHSD_BHSD_BHSD to ..." | Re-trigger Greptile}

[cyanguwa]

Please consolidate 3361 with 3362, and make sure the enum order is consistent style-wise (for example, BSHD, SBHD, and BHSD probably should sit together, instead of with the BSHD_BS2HD ones). Thanks.