Skip to content

Commit bdd4cde

Browse files
Krzysztof Rymskicopybara-github
authored andcommitted
Add enum to string function
PiperOrigin-RevId: 914243776
1 parent 646d91c commit bdd4cde

2 files changed

Lines changed: 26 additions & 0 deletions

File tree

gemma/configs.cc

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -736,4 +736,23 @@ AttentionImpl GetAttentionImpl(const std::string& impl_name) {
736736
return AttentionImpl::kFlash;
737737
}
738738

739+
std::string KVEncodingToString(KVEncoding encoding) {
740+
switch (encoding) {
741+
case KVEncoding::kF32:
742+
return "F32";
743+
case KVEncoding::kBF16:
744+
return "BF16";
745+
case KVEncoding::kF32TwoTranspositions:
746+
return "F32TwoTranspositions";
747+
case KVEncoding::kBF16TwoTranspositions:
748+
return "BF16TwoTranspositions";
749+
case KVEncoding::kInt8:
750+
return "Int8";
751+
case KVEncoding::kInt8TwoTranspositions:
752+
return "Int8TwoTranspositions";
753+
default:
754+
return "Unknown";
755+
}
756+
}
757+
739758
} // namespace gcpp

gemma/configs.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,13 @@ enum class KVEncoding {
9292
kInt8TwoTranspositions = 6,
9393
};
9494

95+
// Returns a string representation of the KVEncoding.
96+
// This representation will not change and can be used for
97+
// serialization or logging.
98+
// Note that no reverse function exists to convert a string back to a
99+
// KVEncoding.
100+
std::string KVEncodingToString(KVEncoding encoding);
101+
95102
enum class AttentionImpl {
96103
kFlash = 0, // Flash Attention (default)
97104
kFlashTransposedQs,

0 commit comments

Comments
 (0)