diff --git a/families/personaplex/runtime/CMakeLists.txt b/families/personaplex/runtime/CMakeLists.txt index 6e5a0ac750..85ef2c268b 100644 --- a/families/personaplex/runtime/CMakeLists.txt +++ b/families/personaplex/runtime/CMakeLists.txt @@ -81,4 +81,8 @@ if(TRTMC_BUILD_TESTS) ) add_test(NAME ${test_name} COMMAND ${test_name}) endforeach() + # The embedding plan uses the CUDA host conversions to round Half additions. + target_include_directories(test_personaplex_speech_temporal_embed_plan SYSTEM PRIVATE + ${TRTMC_CUDA_INCLUDE_DIR} + ) endif() diff --git a/families/personaplex/runtime/pipeline.cpp b/families/personaplex/runtime/pipeline.cpp index afd548414f..d3d7ca3dbe 100644 --- a/families/personaplex/runtime/pipeline.cpp +++ b/families/personaplex/runtime/pipeline.cpp @@ -540,6 +540,12 @@ namespace { void compute_text_prompt_frame_embed(const SpeechConfig& cfg, int32_t text_token_id, int32_t hidden, float* summed_embed) { + if (cfg.temporal_embed_half) { + const int32_t bos = speech_clamp_token(cfg.codebook_size, cfg.audio_vocab_size); + const std::vector audio_tokens(static_cast(cfg.num_codebooks), bos); + compute_half_summed_embed(cfg, hidden, audio_tokens, text_token_id, summed_embed); + return; + } std::fill(summed_embed, summed_embed + hidden, 0.0F); const int32_t text_tok = speech_clamp_token(text_token_id, cfg.temporal_text_vocab); const auto text_offset = static_cast(text_tok) * hidden; diff --git a/families/personaplex/runtime/plugin.cpp b/families/personaplex/runtime/plugin.cpp index bc11a8db57..cd68ea50e8 100644 --- a/families/personaplex/runtime/plugin.cpp +++ b/families/personaplex/runtime/plugin.cpp @@ -125,6 +125,9 @@ extern "C" trtmc::ITask* trtmc_create_family(const trtmc::FamilyContext& context if (temporal_shape.size() < 2 || temporal_shape[1] <= 0 || depth_shape.size() < 2 || depth_shape[1] <= 0) throw std::runtime_error("PersonaPlex engine cache geometry is invalid"); + // An FP16 temporal transformer (per-layer FP32 promotions keep the FP16 + // base) consumes Half input embeddings, summed as the reference does. + config.temporal_embed_half = temporal.module->tensor_dtype("cache_k_0") == DType::kFloat16; auto temporal_state = std::make_unique( config.temporal_num_layers, config.mimi_max_frames, static_cast(temporal_shape[1]), options.stream, diff --git a/families/personaplex/runtime/speech_config.h b/families/personaplex/runtime/speech_config.h index bbe7247193..8dea9922e4 100644 --- a/families/personaplex/runtime/speech_config.h +++ b/families/personaplex/runtime/speech_config.h @@ -34,6 +34,9 @@ struct SpeechConfig { std::vector temporal_text_embedding; int32_t temporal_text_vocab{0}; + // The temporal transformer runs in FP16: sum the input embeddings as the + // reference's Half embed_codes does (see speech_temporal_embed_plan.h). + bool temporal_embed_half{false}; int32_t text_padding_id{3}; std::vector depth_text_embedding; diff --git a/families/personaplex/runtime/speech_temporal_embed_plan.h b/families/personaplex/runtime/speech_temporal_embed_plan.h index 0947c29cd6..386b611105 100644 --- a/families/personaplex/runtime/speech_temporal_embed_plan.h +++ b/families/personaplex/runtime/speech_temporal_embed_plan.h @@ -11,6 +11,7 @@ #include #include #include +#include #include namespace trtmc { @@ -26,6 +27,63 @@ inline void add_speech_embedding_row(const std::vector& table, std::size_ } } +// Round an FP32 value to FP16 (round to nearest even) and back: the result of +// one PyTorch Half operation, which computes in FP32 and rounds once. Host-side +// cuda_fp16 conversion; no device code. +inline float personaplex_round_to_half(float value) { + return __half2float(__float2half_rn(value)); +} + +// Accumulate one embedding row as a Half tensor add: the row is rounded to +// FP16 and the first row starts the sum. +inline void accumulate_half_embedding_row(const std::vector& table, std::size_t offset, + int32_t hidden_size, bool& started, float* out_embed) { + if (offset + hidden_size > table.size()) { + return; + } + const float* row = table.data() + offset; + for (int32_t d = 0; d < hidden_size; ++d) { + const float value = personaplex_round_to_half(row[d]); + out_embed[d] = started ? personaplex_round_to_half(out_embed[d] + value) : value; + } + started = true; +} + +inline std::size_t temporal_text_offset(const SpeechConfig& cfg, int32_t hidden_size, + int32_t text_token) { + const int32_t ttok = clamp_speech_depth_token(text_token, cfg.temporal_text_vocab); + return static_cast(ttok) * hidden_size; +} + +inline std::size_t temporal_audio_offset(const SpeechConfig& cfg, int32_t hidden_size, + int32_t emb_codebook_idx, int32_t token) { + const int32_t tok = clamp_speech_depth_token(token, cfg.audio_vocab_size); + return static_cast(emb_codebook_idx) * cfg.audio_vocab_size * hidden_size + + static_cast(tok) * hidden_size; +} + +// The reference's FP16 embed_codes: audio codebooks in order (model stream, +// then user stream), then text, each add a Half tensor add. +inline void compute_half_summed_embed(const SpeechConfig& cfg, int32_t hidden_size, + const std::vector& audio_tokens, int32_t text_token, + float* out_embed) { + std::fill(out_embed, out_embed + hidden_size, 0.0F); + bool started = false; + if (cfg.audio_vocab_size > 0) { + for (std::size_t cb = 0; cb < audio_tokens.size(); ++cb) { + const auto offset = + temporal_audio_offset(cfg, hidden_size, static_cast(cb), audio_tokens[cb]); + accumulate_half_embedding_row(cfg.audio_embeddings, offset, hidden_size, started, + out_embed); + } + } + if (!cfg.temporal_text_embedding.empty() && cfg.temporal_text_vocab > 0) { + accumulate_half_embedding_row(cfg.temporal_text_embedding, + temporal_text_offset(cfg, hidden_size, text_token), + hidden_size, started, out_embed); + } +} + inline void add_temporal_text_embedding(const SpeechConfig& cfg, int32_t hidden_size, int32_t text_token, float* out_embed) { if (cfg.temporal_text_embedding.empty() || cfg.temporal_text_vocab <= 0) { @@ -54,6 +112,12 @@ inline void compute_dual_stream_summed_embed(const SpeechConfig& cfg, int32_t hi int32_t stream_cb, const int32_t* moshi_tokens, const int32_t* user_tokens, int32_t text_token, float* out_embed) { + if (cfg.temporal_embed_half) { + std::vector audio_tokens(moshi_tokens, moshi_tokens + stream_cb); + audio_tokens.insert(audio_tokens.end(), user_tokens, user_tokens + stream_cb); + compute_half_summed_embed(cfg, hidden_size, audio_tokens, text_token, out_embed); + return; + } std::fill(out_embed, out_embed + hidden_size, 0.0F); add_temporal_text_embedding(cfg, hidden_size, text_token, out_embed); for (int32_t cb = 0; cb < stream_cb; ++cb) { diff --git a/families/personaplex/tests/cpp/test_speech_temporal_embed_plan.cpp b/families/personaplex/tests/cpp/test_speech_temporal_embed_plan.cpp index 1382afe81b..92257e3237 100644 --- a/families/personaplex/tests/cpp/test_speech_temporal_embed_plan.cpp +++ b/families/personaplex/tests/cpp/test_speech_temporal_embed_plan.cpp @@ -64,6 +64,49 @@ void test_dual_stream_embed_sums_text_moshi_and_user_streams() { check_close(out[1], 4.0F + 13.0F + 25.0F + 31.0F + 45.0F, 1e-6F, "dual stream embed value 1"); } +void test_round_to_half_matches_ieee_half_rounding() { + check(trtmc::personaplex_round_to_half(1.0F + 0x1.0p-11F) == 1.0F, "tie rounds to even (down)"); + check(trtmc::personaplex_round_to_half(1.0F + 0x1.8p-10F) == 1.0F + 0x1.0p-9F, + "tie rounds to even (up)"); + check(trtmc::personaplex_round_to_half(1.0F + 0x1.0p-10F) == 1.0F + 0x1.0p-10F, + "half values are exact"); + check(trtmc::personaplex_round_to_half(65519.0F) == 65504.0F, "largest finite half"); + check(std::isinf(trtmc::personaplex_round_to_half(65520.0F)), "overflow rounds to infinity"); + check(trtmc::personaplex_round_to_half(0x1.0p-25F) == 0.0F, + "subnormal tie rounds to even zero"); + check(trtmc::personaplex_round_to_half(0x1.8p-24F) == 0x1.0p-23F, + "subnormal tie rounds to even"); + check(trtmc::personaplex_round_to_half(-0x1.0p-24F) == -0x1.0p-24F, + "smallest subnormal keeps sign"); + check(trtmc::personaplex_round_to_half(0.1F) == 0.0999755859375F, + "0.1 rounds to its half value"); +} + +void test_half_embed_sums_audio_then_text_with_half_adds() { + auto cfg = make_test_config(); + cfg.temporal_embed_half = true; + // 2048 + 1 + 1 in Half: each add rounds (2048 + 1 -> 2048), so the order matters. + cfg.audio_embeddings = {2048.0F, 2048.0F, 1.0F, 1.0F, 0.0F, 0.0F, + 1.0F, 1.0F, 0.0F, 0.0F, 0.0F, 0.0F}; + cfg.temporal_text_embedding = {0.0F, 0.0F, 1.0F, 1.0F, 0.0F, 0.0F, 0.0F, 0.0F}; + const int32_t moshi_tokens[] = {0}; + const int32_t user_tokens[] = {0}; + float out[2] = {0.0F, 0.0F}; + trtmc::compute_dual_stream_summed_embed(cfg, 2, 1, moshi_tokens, user_tokens, 1, out); + check(out[0] == 2048.0F && out[1] == 2048.0F, "Half adds round after every row"); + + cfg.temporal_embed_half = false; + trtmc::compute_dual_stream_summed_embed(cfg, 2, 1, moshi_tokens, user_tokens, 1, out); + check(out[0] == 2050.0F, "FP32 path keeps the exact sum"); + + // Audio first, text last: 2048 + -2048 + 1 = 1; text first in Half would give 0. + cfg.temporal_embed_half = true; + cfg.audio_embeddings[6] = -2048.0F; + cfg.audio_embeddings[7] = -2048.0F; + trtmc::compute_dual_stream_summed_embed(cfg, 2, 1, moshi_tokens, user_tokens, 1, out); + check(out[0] == 1.0F && out[1] == 1.0F, "Half sum adds the text embedding last"); +} + void test_temporal_embed_clamps_tokens_to_vocab() { auto cfg = make_test_config(); float out[2] = {0.0F, 0.0F}; @@ -94,6 +137,8 @@ void test_hidden_helpers_pad_or_truncate_logits() { int main() { test_dual_stream_embed_sums_text_moshi_and_user_streams(); + test_round_to_half_matches_ieee_half_rounding(); + test_half_embed_sums_audio_then_text_with_half_adds(); test_temporal_embed_clamps_tokens_to_vocab(); test_hidden_helpers_pad_or_truncate_logits();