fix(personaplex): round the FP16 depth input like the reference

families/personaplex/runtime/CMakeLists.txt

@@ -81,4 +81,25 @@ if(TRTMC_BUILD_TESTS)
     )
     add_test(NAME ${test_name} COMMAND ${test_name})
   endforeach()
+
+  add_executable(test_personaplex_depth_embedding_kernel
+    ${PROJECT_SOURCE_DIR}/families/personaplex/tests/cpp/test_depth_embedding_kernel.cpp
+    ${PROJECT_SOURCE_DIR}/families/personaplex/runtime/sampling_kernels.cu
+  )
+  target_include_directories(test_personaplex_depth_embedding_kernel PRIVATE
+    ${PROJECT_SOURCE_DIR}
+    ${PROJECT_SOURCE_DIR}/core/runtime/include
+  )
+  target_include_directories(test_personaplex_depth_embedding_kernel SYSTEM PRIVATE
+    ${TRTMC_CUDA_INCLUDE_DIR}
+  )
+  target_link_libraries(test_personaplex_depth_embedding_kernel PRIVATE
+    ${TRTMC_CUDART_LIBRARY}
+  )
+  add_test(NAME test_personaplex_depth_embedding_kernel
+    COMMAND test_personaplex_depth_embedding_kernel)
+  set_tests_properties(test_personaplex_depth_embedding_kernel PROPERTIES
+    SKIP_RETURN_CODE 77
+    LABELS gpu
+  )
 endif()

families/personaplex/runtime/pipeline.cpp

@@ -393,6 +393,8 @@ void SpeechPipeline::prepare_depth_input(const SpeechTemporalDeviceOutput& tempo
                                     forced_audio_provided &&
                                     forced_audio_provided[previous_codebook] != 0;
     const int32_t forced_previous = previous_is_forced ? forced_audio_tokens[previous_codebook] : 0;
+    // The depth transformer computes in its KV-cache dtype; its input is built in that dtype.
+    const DType depth_dtype = depth_engine_for_codebook(codebook).tensor_dtype("cache_k_0");
     personaplex_prepare_depth_embedding(
         static_cast<float*>(device_workspace_->depth_embed.data()), temporal_output.hidden,
         temporal_output.hidden_dtype,
@@ -404,7 +406,7 @@ void SpeechPipeline::prepare_depth_input(const SpeechTemporalDeviceOutput& tempo
         device_workspace_->depth_audio_embeddings.numel(), selected_tokens, codebook, text_token,
         text_token_is_forced, forced_previous, previous_is_forced, config_.depth_hidden_size,
         config_.temporal_hidden_size, config_.depth_text_vocab, config_.audio_vocab_size,
-        config_.num_depformer_emb, stream_);
+        config_.num_depformer_emb, depth_dtype, stream_);
     if (const auto error = cudaGetLastError(); error != cudaSuccess) {
         throw std::runtime_error(std::string("SpeechPipeline depth embedding: ") +
                                  cudaGetErrorString(error));

families/personaplex/runtime/sampling_kernels.cu

@@ -178,6 +178,16 @@ __device__ float load_hidden<__nv_bfloat16>(const __nv_bfloat16* hidden, int32_t
     return __bfloat162float(hidden[index]);
 }
 
+// The reference depth input is F.linear(temporal_hidden) + embedding(previous_token) in the depth
+// transformer's dtype: operands, the projection output and the sum are each rounded to that dtype.
+__device__ float round_to_depth_dtype(float value, DType depth_dtype) {
+    if (depth_dtype == DType::kFloat16)
+        return __half2float(__float2half_rn(value));
+    if (depth_dtype == DType::kBFloat16)
+        return __bfloat162float(__float2bfloat16_rn(value));
+    return value;
+}
+
 __device__ float text_seed(const float* embedding, int64_t elements, const int32_t* selected_tokens,
                            int32_t forced_token, bool token_is_forced, int32_t vocab_size,
                            int32_t hidden_size, int32_t row) {
@@ -210,7 +220,8 @@ __global__ void prepare_depth_embedding_kernel(
     int64_t depth_audio_embedding_elements, const int32_t* selected_tokens, int32_t codebook,
     int32_t forced_text_token, bool text_token_is_forced, int32_t forced_previous_token,
     bool previous_token_is_forced, int32_t depth_hidden, int32_t temporal_hidden_dim,
-    int32_t depth_text_vocab, int32_t audio_vocab_size, int32_t num_depformer_embeddings) {
+    int32_t depth_text_vocab, int32_t audio_vocab_size, int32_t num_depformer_embeddings,
+    DType depth_dtype) {
     __shared__ float partial_sums[kProjectionThreads];
     const int32_t row = blockIdx.x;
     if (row >= depth_hidden)
@@ -224,8 +235,8 @@ __global__ void prepare_depth_embedding_kernel(
     if (projection_is_valid) {
         for (int32_t column = threadIdx.x; column < temporal_hidden_dim;
              column += kProjectionThreads) {
-            projected +=
-                depth_projection[row_offset + column] * load_hidden(temporal_hidden, column);
+            projected += round_to_depth_dtype(depth_projection[row_offset + column], depth_dtype) *
+                         round_to_depth_dtype(load_hidden(temporal_hidden, column), depth_dtype);
         }
     }
     partial_sums[threadIdx.x] = projected;
@@ -244,7 +255,9 @@ __global__ void prepare_depth_embedding_kernel(
                                             selected_tokens, codebook, forced_previous_token,
                                             previous_token_is_forced, audio_vocab_size,
                                             depth_hidden, num_depformer_embeddings, row);
-        output[row] = seed + partial_sums[0];
+        const float projection = round_to_depth_dtype(partial_sums[0], depth_dtype);
+        output[row] =
+            round_to_depth_dtype(round_to_depth_dtype(seed, depth_dtype) + projection, depth_dtype);
     }
 }
 
@@ -282,7 +295,8 @@ void personaplex_prepare_depth_embedding(
     const int32_t* selected_tokens, int32_t codebook, int32_t forced_text_token,
     bool text_token_is_forced, int32_t forced_previous_token, bool previous_token_is_forced,
     int32_t depth_hidden, int32_t temporal_hidden_dim, int32_t depth_text_vocab,
-    int32_t audio_vocab_size, int32_t num_depformer_embeddings, cudaStream_t stream) {
+    int32_t audio_vocab_size, int32_t num_depformer_embeddings, DType depth_dtype,
+    cudaStream_t stream) {
     if (!output || !selected_tokens || depth_hidden <= 0)
         return;
     const auto launch = personaplex_depth_projection_launch_plan(depth_hidden);
@@ -293,7 +307,7 @@ void personaplex_prepare_depth_embedding(
         depth_audio_embeddings, depth_audio_embedding_elements, selected_tokens, codebook,         \
         forced_text_token, text_token_is_forced, forced_previous_token, previous_token_is_forced,  \
         depth_hidden, temporal_hidden_dim, depth_text_vocab, audio_vocab_size,                     \
-        num_depformer_embeddings)
+        num_depformer_embeddings, depth_dtype)
     switch (temporal_hidden_dtype) {
     case DType::kFloat16:
         TRTMC_LAUNCH_DEPTH_EMBED(__half);

families/personaplex/runtime/sampling_kernels.h

@@ -31,6 +31,7 @@ void personaplex_prepare_depth_embedding(
     const int32_t* selected_tokens, int32_t codebook, int32_t forced_text_token,
     bool text_token_is_forced, int32_t forced_previous_token, bool previous_token_is_forced,
     int32_t depth_hidden, int32_t temporal_hidden_dim, int32_t depth_text_vocab,
-    int32_t audio_vocab_size, int32_t num_depformer_embeddings, cudaStream_t stream);
+    int32_t audio_vocab_size, int32_t num_depformer_embeddings, DType depth_dtype,
+    cudaStream_t stream);
 
 } // namespace trtmc

families/personaplex/tests/cpp/test_depth_embedding_kernel.cpp

@@ -0,0 +1,101 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// =============================================================================
+// ISO 26262 Traceability
+// =============================================================================
+// Trace ID:       UT-AUD-CPP-17
+// Architecture:   ARCH-FAC-001
+// Unit Design:    UD-AUD-01
+// Intent:         Depth input embedding rounds at the reference's depth-dtype boundaries
+// Preconditions:  A CUDA device (the test is skipped without one)
+// Postconditions: An FP16 depth gets Half(Half(projection) + Half(embedding)); an FP32 depth keeps
+// the FP32 sum
+// =============================================================================
+
+#include "families/personaplex/runtime/sampling_kernels.h"
+
+#include <cmath>
+#include <cstdint>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <iostream>
+#include <vector>
+
+namespace {
+
+int g_failures = 0;
+
+void check_equal(float actual, float expected, const char* name) {
+    if (actual != expected) {
+        std::cerr << "FAIL: " << name << " actual=" << actual << " expected=" << expected << '\n';
+        ++g_failures;
+    }
+}
+
+template <typename T>
+T* upload(const std::vector<T>& values) {
+    T* device = nullptr;
+    if (cudaMalloc(&device, values.size() * sizeof(T)) != cudaSuccess ||
+        cudaMemcpy(device, values.data(), values.size() * sizeof(T), cudaMemcpyHostToDevice) !=
+            cudaSuccess)
+        return nullptr;
+    return device;
+}
+
+// One depth row, temporal hidden [1, 1] in Half, projection row [1, 2^-12] (exact in BF16 and
+// Half) and embedding -1 for both the text (codebook 0) and audio (codebook 1) seeds. The
+// reference computes F.linear in Half (1 + 2^-12 rounds to 1) and then adds -1 in Half, giving 0;
+// a single rounding of the FP32 sum gives 2^-12.
+float depth_input(int32_t codebook, trtmc::DType depth_dtype) {
+    const std::vector<__half> hidden_values{__float2half(1.0F), __float2half(1.0F)};
+    const float tiny = std::ldexp(1.0F, -12);
+    const std::vector<float> projection_values{1.0F, tiny, 1.0F, tiny};
+    const std::vector<float> embedding_values{-1.0F};
+    const std::vector<int32_t> token_values{0, 0};
+    auto* hidden = upload(hidden_values);
+    auto* projection = upload(projection_values);
+    auto* text_embedding = upload(embedding_values);
+    auto* audio_embedding = upload(embedding_values);
+    auto* tokens = upload(token_values);
+    float* output = nullptr;
+    float result = NAN;
+    if (hidden && projection && text_embedding && audio_embedding && tokens &&
+        cudaMalloc(&output, sizeof(float)) == cudaSuccess) {
+        trtmc::personaplex_prepare_depth_embedding(output, hidden, trtmc::DType::kFloat16,
+                                                   projection, 4, text_embedding, 1,
+                                                   audio_embedding, 1, tokens, codebook, 0, false,
+                                                   0, false, 1, 2, 1, 1, 1, depth_dtype, nullptr);
+        if (cudaGetLastError() != cudaSuccess ||
+            cudaMemcpy(&result, output, sizeof(float), cudaMemcpyDeviceToHost) != cudaSuccess)
+            result = NAN;
+    }
+    cudaFree(output);
+    cudaFree(tokens);
+    cudaFree(audio_embedding);
+    cudaFree(text_embedding);
+    cudaFree(projection);
+    cudaFree(hidden);
+    return result;
+}
+
+} // namespace
+
+int main() {
+    int device_count = 0;
+    if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0)
+        return 77;
+    const float tiny = std::ldexp(1.0F, -12);
+    check_equal(depth_input(0, trtmc::DType::kFloat16), 0.0F, "half depth text seed");
+    check_equal(depth_input(1, trtmc::DType::kFloat16), 0.0F, "half depth audio seed");
+    check_equal(depth_input(0, trtmc::DType::kFloat32), tiny, "fp32 depth text seed");
+    check_equal(depth_input(1, trtmc::DType::kFloat32), tiny, "fp32 depth audio seed");
+    if (g_failures != 0) {
+        std::cerr << g_failures << " check(s) failed\n";
+        return 1;
+    }
+    std::cout << "depth embedding kernel: all checks passed\n";
+    return 0;
+}