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; +}