fix(personaplex): preserve Half embedding additions

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()

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<int32_t> audio_tokens(static_cast<std::size_t>(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<std::size_t>(text_tok) * hidden;

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<PersonaplexKvCache>(
         config.temporal_num_layers, config.mimi_max_frames,
         static_cast<std::int32_t>(temporal_shape[1]), options.stream,

families/personaplex/runtime/speech_config.h

@@ -34,6 +34,9 @@ struct SpeechConfig {
 
     std::vector<float> 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<float> depth_text_embedding;

families/personaplex/runtime/speech_temporal_embed_plan.h

@@ -11,6 +11,7 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <cuda_fp16.h>
 #include <vector>
 
 namespace trtmc {
@@ -26,6 +27,63 @@ inline void add_speech_embedding_row(const std::vector<float>& 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<float>& 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<std::size_t>(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<std::size_t>(emb_codebook_idx) * cfg.audio_vocab_size * hidden_size +
+           static_cast<std::size_t>(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<int32_t>& 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<int32_t>(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<int32_t> 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) {

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();