From cedb988b4b61ffd97ad9588f2120452740d1a8c4 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Sun, 18 Oct 2026 08:28:29 +0000 Subject: [PATCH] fix(personaplex): round the FP16 depth input like the reference The reference builds each depth step's input as F.linear(temporal_hidden) plus the previous-token embedding in the depth transformer's dtype: in FP16 the weights, the projection output, the embedding and the sum are each rounded to Half. The depth-input helper summed FP32 copies of the BF16 tables in FP32 and left a single Half rounding to the depth engine, so the depth input differed from the reference in every FP16 depth step. The helper now rounds at those boundaries to the depth engine's compute dtype, taken per codebook from the engine's cache_k_0 dtype. FP32 depth engines keep the FP32 formula. A GPU regression test checks a binary counterexample whose reference result is 0 (the old formula gives 2^-12) and that the FP32 path is unchanged. Fixes #1788 Signed-off-by: chaofengw --- families/personaplex/runtime/CMakeLists.txt | 21 ++++ families/personaplex/runtime/pipeline.cpp | 4 +- .../personaplex/runtime/sampling_kernels.cu | 26 +++-- .../personaplex/runtime/sampling_kernels.h | 3 +- .../tests/cpp/test_depth_embedding_kernel.cpp | 101 ++++++++++++++++++ 5 files changed, 147 insertions(+), 8 deletions(-) create mode 100644 families/personaplex/tests/cpp/test_depth_embedding_kernel.cpp diff --git a/families/personaplex/runtime/CMakeLists.txt b/families/personaplex/runtime/CMakeLists.txt index 6e5a0ac750..c075a5adc6 100644 --- a/families/personaplex/runtime/CMakeLists.txt +++ b/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() diff --git a/families/personaplex/runtime/pipeline.cpp b/families/personaplex/runtime/pipeline.cpp index afd548414f..024d679d70 100644 --- a/families/personaplex/runtime/pipeline.cpp +++ b/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(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)); diff --git a/families/personaplex/runtime/sampling_kernels.cu b/families/personaplex/runtime/sampling_kernels.cu index 1508a814b5..bcb3150afc 100644 --- a/families/personaplex/runtime/sampling_kernels.cu +++ b/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); diff --git a/families/personaplex/runtime/sampling_kernels.h b/families/personaplex/runtime/sampling_kernels.h index bc80b54825..c8bae6eac7 100644 --- a/families/personaplex/runtime/sampling_kernels.h +++ b/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 diff --git a/families/personaplex/tests/cpp/test_depth_embedding_kernel.cpp b/families/personaplex/tests/cpp/test_depth_embedding_kernel.cpp new file mode 100644 index 0000000000..e2f9238901 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +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 +T* upload(const std::vector& 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 projection_values{1.0F, tiny, 1.0F, tiny}; + const std::vector embedding_values{-1.0F}; + const std::vector 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; +}