From a8b1160451353a47a66883eae832d8ed4e05dc87 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Sat, 17 Oct 2026 10:02:21 +0000 Subject: [PATCH] fix(personaplex): preserve Half embedding additions Sum audio embeddings before text and round every FP16 addition as the Native temporal input does. Keep the existing FP32 path and infer the temporal base dtype from the engine cache contract. Add host-only rounding and summation regressions. This fixes an input defect; complete audio parity still requires separate validation. Refs: #1778 Signed-off-by: chaofengw --- families/personaplex/runtime/CMakeLists.txt | 4 ++ families/personaplex/runtime/pipeline.cpp | 6 ++ families/personaplex/runtime/plugin.cpp | 3 + families/personaplex/runtime/speech_config.h | 3 + .../runtime/speech_temporal_embed_plan.h | 64 +++++++++++++++++++ .../cpp/test_speech_temporal_embed_plan.cpp | 45 +++++++++++++ 6 files changed, 125 insertions(+) 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();