diff --git a/CMakeLists.txt b/CMakeLists.txt index 80dcdd5b..de0f5b84 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -327,6 +327,7 @@ add_library(engine_core OBJECT src/framework/modules/speech_encoders/whisper_embedding.cpp src/framework/modules/speech_encoders/whisper_frontend.cpp src/framework/modules/speech_encoders/campplus_encoder.cpp + src/framework/modules/codecs/nemo_nano_codec.cpp src/framework/modules/pitch_extractors/rmvpe_pitch_extractor.cpp src/framework/modules/vocoders/bigvgan_vocoder.cpp src/framework/modules/vocoders/hift_vocoder.cpp diff --git a/include/engine/framework/codecs/mimi_codec_runtime.h b/include/engine/framework/codecs/mimi_codec_runtime.h index 7c855893..899b625d 100644 --- a/include/engine/framework/codecs/mimi_codec_runtime.h +++ b/include/engine/framework/codecs/mimi_codec_runtime.h @@ -203,6 +203,8 @@ class MimiEncoderRuntime { ~MimiEncoderRuntime(); std::vector encode(const runtime::AudioBuffer & audio); + void reset_streaming(); + std::vector encode_streaming(const runtime::AudioBuffer & audio, bool flush); private: struct Impl; diff --git a/include/engine/framework/modules/activation_modules.h b/include/engine/framework/modules/activation_modules.h index fcca633e..0b57350f 100644 --- a/include/engine/framework/modules/activation_modules.h +++ b/include/engine/framework/modules/activation_modules.h @@ -21,6 +21,23 @@ class ReluModule { static const core::ModuleSchema & static_schema() noexcept; }; +struct LeakyReluConfig { + float negative_slope = 0.01F; +}; + +class LeakyReluModule { +public: + explicit LeakyReluModule(LeakyReluConfig config = {}); + + const LeakyReluConfig & config() const noexcept; + const core::ModuleSchema & schema() const noexcept; + core::TensorValue build(core::ModuleBuildContext & ctx, const core::TensorValue & input) const; + static const core::ModuleSchema & static_schema() noexcept; + +private: + LeakyReluConfig config_; +}; + class SigmoidModule { public: const core::ModuleSchema & schema() const noexcept; diff --git a/include/engine/framework/modules/attention/cross_attention.h b/include/engine/framework/modules/attention/cross_attention.h index 221e487a..26113aee 100644 --- a/include/engine/framework/modules/attention/cross_attention.h +++ b/include/engine/framework/modules/attention/cross_attention.h @@ -3,8 +3,15 @@ #include "engine/framework/core/module.h" #include "engine/framework/modules/attention/types.h" +#include + namespace engine::modules { +struct CrossAttentionKeyValue { + core::TensorValue key; + core::TensorValue value; +}; + class CrossAttentionModule { public: explicit CrossAttentionModule(AttentionConfig config); @@ -18,6 +25,29 @@ class CrossAttentionModule { const core::TensorValue & memory, const AttentionWeights & weights) const; + core::TensorValue build( + core::ModuleBuildContext & ctx, + const core::TensorValue & query, + const core::TensorValue & memory, + const AttentionWeights & weights, + const core::TensorValue & memory_mask, + const core::TensorValue * attention_prior = nullptr, + core::TensorValue * last_attention = nullptr) const; + + core::TensorValue build_cached( + core::ModuleBuildContext & ctx, + const core::TensorValue & query, + const CrossAttentionKeyValue & key_value, + const AttentionWeights & weights, + const core::TensorValue & memory_mask, + const core::TensorValue * attention_prior = nullptr, + core::TensorValue * last_attention = nullptr) const; + + CrossAttentionKeyValue build_key_value( + core::ModuleBuildContext & ctx, + const core::TensorValue & memory, + const AttentionWeights & weights) const; + static const core::ModuleSchema & static_schema() noexcept; private: diff --git a/include/engine/framework/modules/attention/feed_forward.h b/include/engine/framework/modules/attention/feed_forward.h index 7354b6ad..6d27710e 100644 --- a/include/engine/framework/modules/attention/feed_forward.h +++ b/include/engine/framework/modules/attention/feed_forward.h @@ -2,6 +2,7 @@ #include "engine/framework/core/module.h" #include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" #include "engine/framework/modules/linear_module.h" #include @@ -45,6 +46,20 @@ struct GatedFeedForwardWeights { LinearWeights down_proj; }; +struct ConvFeedForwardConfig { + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t kernel_size = 0; + bool causal = false; + bool use_bias = false; + GeluApproximation gelu_approximation = GeluApproximation::Tanh; +}; + +struct ConvFeedForwardWeights { + Conv1dWeights proj; + Conv1dWeights out; +}; + class FeedForwardModule { public: explicit FeedForwardModule(FeedForwardConfig config); @@ -99,4 +114,22 @@ class GatedFeedForwardModule { GatedFeedForwardConfig config_; }; +class ConvFeedForwardModule { +public: + explicit ConvFeedForwardModule(ConvFeedForwardConfig config); + + const ConvFeedForwardConfig & config() const noexcept; + const core::ModuleSchema & schema() const noexcept; + + core::TensorValue build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvFeedForwardWeights & weights) const; + + static const core::ModuleSchema & static_schema() noexcept; + +private: + ConvFeedForwardConfig config_; +}; + } // namespace engine::modules diff --git a/include/engine/framework/modules/attention/self_attention.h b/include/engine/framework/modules/attention/self_attention.h index 72567508..357b3eb3 100644 --- a/include/engine/framework/modules/attention/self_attention.h +++ b/include/engine/framework/modules/attention/self_attention.h @@ -2,6 +2,7 @@ #include "engine/framework/core/module.h" #include "engine/framework/modules/attention/types.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" #include @@ -19,6 +20,22 @@ class SelfAttentionModule { const core::TensorValue & input, const AttentionWeights & weights) const; + core::TensorValue build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const AttentionWeights & weights, + const std::optional & attention_mask) const; + + StreamingAttentionOutputs build_cached_tail( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const AttentionWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask, + FastKVSetRowsMode set_rows_mode = FastKVSetRowsMode::BackendViewOptimized) const; + static const core::ModuleSchema & static_schema() noexcept; private: diff --git a/include/engine/framework/modules/attention/types.h b/include/engine/framework/modules/attention/types.h index 6c801d99..1857447c 100644 --- a/include/engine/framework/modules/attention/types.h +++ b/include/engine/framework/modules/attention/types.h @@ -18,6 +18,12 @@ struct AttentionConfig { ggml_prec projection_precision = GGML_PREC_DEFAULT; ggml_prec attention_precision = GGML_PREC_DEFAULT; AttentionPrefixCacheLayout prefix_cache_layout = AttentionPrefixCacheLayout::SequenceHeads; + bool use_packed_qkv = false; + bool causal = false; + bool use_packed_kv = false; + int64_t key_value_size = 0; + int64_t attention_size = 0; + int64_t head_dim = 0; }; struct RelativeAttentionConfig { diff --git a/include/engine/framework/modules/codecs/nemo_nano_codec.h b/include/engine/framework/modules/codecs/nemo_nano_codec.h new file mode 100644 index 00000000..e0005191 --- /dev/null +++ b/include/engine/framework/modules/codecs/nemo_nano_codec.h @@ -0,0 +1,53 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include + +namespace engine::modules { + +struct NemoNanoCodecConfig { + int64_t sample_rate = 22050; + int64_t input_dim = 32; + int64_t base_channels = 864; + int64_t audio_codebooks = 0; + std::vector upsample_rates; + std::vector resblock_kernel_sizes; + std::vector resblock_dilation_sizes; + std::vector fsq_num_levels; + std::vector fsq_dim_base_index; +}; + +struct NemoNanoCodecRuntimeOptions { + size_t graph_arena_bytes = 1024ull * 1024ull * 1024ull; + size_t weight_context_bytes = 2048ull * 1024ull * 1024ull; + assets::TensorStorageType weight_storage_type = assets::TensorStorageType::Native; +}; + +class NemoNanoCodecRuntime { +public: + NemoNanoCodecRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + NemoNanoCodecConfig config, + NemoNanoCodecRuntimeOptions options); + ~NemoNanoCodecRuntime(); + + NemoNanoCodecRuntime(const NemoNanoCodecRuntime &) = delete; + NemoNanoCodecRuntime & operator=(const NemoNanoCodecRuntime &) = delete; + + runtime::AudioBuffer decode_codes(const std::vector & codes); + void release_runtime_graph(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::modules diff --git a/include/engine/framework/modules/streaming_conv_modules.h b/include/engine/framework/modules/streaming_conv_modules.h index 3cf24464..43a947d3 100644 --- a/include/engine/framework/modules/streaming_conv_modules.h +++ b/include/engine/framework/modules/streaming_conv_modules.h @@ -60,6 +60,12 @@ enum class StreamingPadMode { Replicate, }; +enum class StreamingConv1dPaddingMode { + StreamingSame, + StrictCausal, + Explicit, +}; + struct StreamingConv1dConfig { int64_t in_channels = 0; int64_t out_channels = 0; @@ -68,6 +74,9 @@ struct StreamingConv1dConfig { int dilation = 1; bool use_bias = true; StreamingPadMode pad_mode = StreamingPadMode::Constant; + StreamingConv1dPaddingMode padding_mode = StreamingConv1dPaddingMode::StreamingSame; + int64_t explicit_left = 0; + int64_t explicit_right = 0; }; using StreamingConv1dWeights = Conv1dWeights; @@ -84,6 +93,12 @@ class StreamingConv1dModule { StreamingConv1dConfig config_; }; +using CausalConv1dPadMode = StreamingPadMode; +using CausalConv1dPaddingMode = StreamingConv1dPaddingMode; +using CausalConv1dConfig = StreamingConv1dConfig; +using CausalConv1dWeights = StreamingConv1dWeights; +using CausalConv1dModule = StreamingConv1dModule; + struct DepthwiseConvTranspose1dConfig { int64_t channels = 0; int64_t kernel_size = 0; diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index 66b05e79..70a834e9 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -26,6 +26,7 @@ struct QwenCausalDecodeRuntimeConfig { QwenCausalDecodeOutputMode output_mode = QwenCausalDecodeOutputMode::Logits; bool return_hidden = false; std::optional readback_round_type; + std::vector logits_readback_token_ids; }; struct QwenCausalDecodeRuntimeWeights { @@ -41,6 +42,12 @@ struct QwenCausalPrefillResult { runtime::TransformerKVState state; }; +struct QwenCausalBatchedPrefillResult { + std::vector logits; + std::vector hidden; + runtime::TransformerBatchedKVState state; +}; + struct QwenCausalDecodeStepResult { std::vector logits; std::vector hidden; @@ -60,11 +67,31 @@ class QwenCausalDecodeRuntime { QwenCausalPrefillResult prefill_tokens(const std::vector & token_ids); QwenCausalPrefillResult prefill_embeddings(const std::vector & embeddings, int64_t steps); + QwenCausalBatchedPrefillResult prefill_tokens_batched( + const std::vector & token_ids, + int64_t batch_size, + int64_t steps); + QwenCausalBatchedPrefillResult prefill_embeddings_batched( + const std::vector & embeddings, + int64_t batch_size, + int64_t steps); + void start_decode_tokens(const runtime::TransformerKVState & state, int64_t required_cache_steps); void start_decode_embeddings(const runtime::TransformerKVState & state, int64_t required_cache_steps); QwenCausalDecodeStepResult decode_token(int32_t token); QwenCausalDecodeStepResult decode_embedding(const std::vector & embedding); + void start_decode_tokens_batched( + const runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps); + void start_decode_embeddings_batched( + const runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps); + QwenCausalDecodeStepResult decode_tokens_batched(const std::vector & tokens); + QwenCausalDecodeStepResult decode_embeddings_batched( + const std::vector & embeddings, + int64_t batch_size); + int64_t decode_cache_steps() const noexcept; int64_t decode_current_end() const noexcept; int64_t decode_valid_steps() const noexcept; diff --git a/include/engine/framework/modules/transformers/qwen_causal_decoder.h b/include/engine/framework/modules/transformers/qwen_causal_decoder.h index 829d652f..4231fcae 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decoder.h @@ -22,6 +22,7 @@ enum class QwenCausalDecoderLogitsMode { struct QwenDecoderHiddenConfig { QwenDecoderStackConfig stack; QwenCausalDecoderLogitsMode hidden_mode = QwenCausalDecoderLogitsMode::LastStep; + ggml_type static_cache_type = GGML_TYPE_F32; }; struct QwenCausalDecoderConfig { @@ -31,6 +32,7 @@ struct QwenCausalDecoderConfig { bool use_lm_head_bias = false; ggml_prec lm_head_precision = GGML_PREC_DEFAULT; std::optional lm_head_input_type; + ggml_type static_cache_type = GGML_TYPE_F32; }; struct QwenCausalDecoderWeights { @@ -56,6 +58,12 @@ struct QwenDecoderHiddenStaticCacheOutputs { runtime::TransformerKVCache cache; }; +struct QwenDecoderHiddenBatchedStaticCacheOutputs { + core::TensorValue sequence; + core::TensorValue hidden; + runtime::TransformerBatchedKVCache cache; +}; + struct QwenCausalDecoderOutputs { core::TensorValue sequence; core::TensorValue hidden; @@ -70,6 +78,13 @@ struct QwenCausalDecoderStaticCacheOutputs { runtime::TransformerKVCache cache; }; +struct QwenCausalDecoderBatchedStaticCacheOutputs { + core::TensorValue sequence; + core::TensorValue hidden; + core::TensorValue logits; + runtime::TransformerBatchedKVCache cache; +}; + class QwenDecoderHiddenModule { public: explicit QwenDecoderHiddenModule(QwenDecoderHiddenConfig config); @@ -94,6 +109,16 @@ class QwenDecoderHiddenModule { const core::TensorValue & attention_mask, const std::optional & cache_slot = std::nullopt) const; + QwenDecoderHiddenBatchedStaticCacheOutputs build_static_cache_tail_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderHiddenWeights & weights, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot) const; + private: QwenDecoderHiddenConfig config_; }; @@ -122,6 +147,16 @@ class QwenCausalDecoderModule { const core::TensorValue & attention_mask, const std::optional & cache_slot = std::nullopt) const; + QwenCausalDecoderBatchedStaticCacheOutputs build_static_cache_tail_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenCausalDecoderWeights & weights, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot) const; + private: QwenCausalDecoderConfig config_; }; @@ -147,4 +182,12 @@ void write_qwen_cached_step_mask( int64_t visible_prefix_steps, int64_t current_slot); +void write_qwen_batched_cached_step_mask( + ggml_tensor * tensor, + std::vector & scratch, + int64_t batch_size, + int64_t mask_steps, + int64_t visible_prefix_steps, + int64_t current_slot); + } // namespace engine::modules diff --git a/include/engine/framework/modules/transformers/qwen_decoder.h b/include/engine/framework/modules/transformers/qwen_decoder.h index 56b85603..3b43a744 100644 --- a/include/engine/framework/modules/transformers/qwen_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_decoder.h @@ -160,6 +160,17 @@ class QwenDecoderLayerModule { const std::optional & cache_slot, const core::TensorValue & attention_mask) const; + QwenDecoderLayerOutputs build_with_static_cache_tail_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask) const; + static const core::ModuleSchema & static_schema() noexcept; private: diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index 81156423..db55f117 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/include/engine/framework/runtime/kv_cache.h @@ -70,6 +70,62 @@ class TransformerKVCache { std::vector layers_; }; +struct BatchedKVLayerState { + int64_t valid_steps = 0; + std::vector key; + std::vector value; +}; + +struct TransformerBatchedKVState { + int64_t batch_size = 0; + int64_t current_end = 0; + std::vector layers; +}; + +class TransformerBatchedKVCache { +public: + TransformerBatchedKVCache() = default; + TransformerBatchedKVCache( + int64_t cache_steps, + int64_t batch_size, + int64_t row_elems, + std::vector keys, + std::vector values); + TransformerBatchedKVCache( + int64_t cache_steps, + int64_t batch_size, + int64_t row_elems, + std::vector keys, + std::vector values, + TransformerKVCacheOptions options); + + void import_state(const TransformerBatchedKVState & state); + TransformerBatchedKVState export_state() const; + + void advance_after_direct_append(int64_t steps); + + int64_t batch_size() const noexcept; + int64_t valid_steps() const noexcept; + int64_t current_end() const noexcept; + int64_t cache_steps() const noexcept; + +private: + struct LayerCache { + core::TensorValue key_tensor; + core::TensorValue value_tensor; + std::vector import_key_scratch; + std::vector import_value_scratch; + }; + + int64_t cache_steps_ = 0; + int64_t batch_size_ = 0; + int64_t row_elems_ = 0; + int64_t valid_steps_ = 0; + int64_t current_end_ = 0; + TransformerKVCacheOptions options_; + std::vector layers_; +}; + core::TensorValue view_transformer_kv_cache_steps( core::ModuleBuildContext & ctx, const core::TensorValue & cache, diff --git a/include/engine/framework/sampling/hf_sampler.h b/include/engine/framework/sampling/hf_sampler.h index 4fcf00a3..d4c7ebb5 100644 --- a/include/engine/framework/sampling/hf_sampler.h +++ b/include/engine/framework/sampling/hf_sampler.h @@ -23,6 +23,8 @@ struct HfTorchSamplingState { const TorchCudaSamplingPolicy * policy = nullptr; uint64_t seed = 0; uint64_t call_index = 0; + uint64_t offset_blocks = 0; + bool use_offset_blocks = false; }; class HfSamplerScratch { diff --git a/include/engine/framework/sampling/torch_random.h b/include/engine/framework/sampling/torch_random.h index e652ffa2..f4fc808c 100644 --- a/include/engine/framework/sampling/torch_random.h +++ b/include/engine/framework/sampling/torch_random.h @@ -81,4 +81,12 @@ float torch_cuda_tensor_iterator_exponential_element( int64_t multiprocessor_count, int64_t max_threads_per_multiprocessor); +float torch_cuda_tensor_iterator_exponential_element_at_offset( + uint64_t seed, + uint64_t total_elements, + uint64_t element_index, + uint64_t offset_blocks, + int64_t multiprocessor_count, + int64_t max_threads_per_multiprocessor); + } // namespace engine::sampling diff --git a/src/framework/codecs/mimi_codec_runtime.cpp b/src/framework/codecs/mimi_codec_runtime.cpp index fdbbf337..58c7618b 100644 --- a/src/framework/codecs/mimi_codec_runtime.cpp +++ b/src/framework/codecs/mimi_codec_runtime.cpp @@ -2028,6 +2028,9 @@ struct MimiEncoderRuntime::Impl { int threads = 1; size_t graph_arena_bytes = 0; RuntimeCache cache; + std::optional streaming_state; + bool streaming_transformer_initialized = false; + std::vector streaming_partial; }; MimiEncoderRuntime::MimiEncoderRuntime( @@ -2065,4 +2068,46 @@ std::vector MimiEncoderRuntime::encode(const runtime::AudioBuffer & aud return impl_->encode_streaming_with_state(mono, state, transformer_initialized); } +void MimiEncoderRuntime::reset_streaming() { + impl_->streaming_state = make_mimi_encoder_state(impl_->config); + impl_->streaming_transformer_initialized = false; + impl_->streaming_partial.clear(); +} + +std::vector MimiEncoderRuntime::encode_streaming(const runtime::AudioBuffer & audio, bool flush) { + if (!impl_->streaming_state.has_value()) { + reset_streaming(); + } + if (audio.sample_rate <= 0 || audio.channels <= 0) { + throw std::runtime_error("Mimi codec streaming encoder received invalid audio format"); + } + auto mono = engine::audio::convert_interleaved_audio_to_mono_linear_resampled( + audio.samples, + audio.sample_rate, + audio.channels, + impl_->config.sample_rate); + impl_->streaming_partial.insert(impl_->streaming_partial.end(), mono.begin(), mono.end()); + const int64_t available = static_cast(impl_->streaming_partial.size()); + int64_t encode_samples = (available / kMimiFrameSamples) * kMimiFrameSamples; + if (flush && available > encode_samples) { + encode_samples += kMimiFrameSamples; + } + if (encode_samples == 0) { + return {}; + } + std::vector encode_input( + impl_->streaming_partial.begin(), + impl_->streaming_partial.begin() + static_cast( + std::min(available, encode_samples))); + encode_input.resize(static_cast(encode_samples), 0.0F); + impl_->streaming_partial.erase( + impl_->streaming_partial.begin(), + impl_->streaming_partial.begin() + static_cast( + std::min(available, encode_samples))); + return impl_->encode_streaming_with_state( + encode_input, + *impl_->streaming_state, + impl_->streaming_transformer_initialized); +} + } // namespace engine::codecs diff --git a/src/framework/modules/activation_modules.cpp b/src/framework/modules/activation_modules.cpp index 5ff0294f..c8f1928c 100644 --- a/src/framework/modules/activation_modules.cpp +++ b/src/framework/modules/activation_modules.cpp @@ -16,14 +16,24 @@ const core::ModulePortSpec kActivationOutputs[] = { const core::ModuleSchema kReluSchema = { "ReLU", - "nn.activation", - kActivationInputs, - 1, + "nn.activation", + kActivationInputs, + 1, kActivationOutputs, 1, "Applies rectified linear activation elementwise.", }; +const core::ModuleSchema kLeakyReluSchema = { + "LeakyReLU", + "nn.activation", + kActivationInputs, + 1, + kActivationOutputs, + 1, + "Applies leaky rectified linear activation elementwise.", +}; + const core::ModuleSchema kSigmoidSchema = { "Sigmoid", "nn.activation", @@ -317,6 +327,33 @@ const core::ModuleSchema & ReluModule::static_schema() noexcept { return kReluSchema; } +LeakyReluModule::LeakyReluModule(LeakyReluConfig config) : config_(config) { +} + +const LeakyReluConfig & LeakyReluModule::config() const noexcept { + return config_; +} + +const core::ModuleSchema & LeakyReluModule::schema() const noexcept { + return static_schema(); +} + +core::TensorValue LeakyReluModule::build(core::ModuleBuildContext & ctx, const core::TensorValue & input) const { + if (ctx.ggml == nullptr) { + throw std::runtime_error("ModuleBuildContext.ggml is null"); + } + core::validate_rank_between(input, 1, core::kMaxTensorRank, "input"); + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + return core::wrap_tensor( + ggml_leaky_relu(ctx.ggml, contiguous.tensor, config_.negative_slope, false), + input.shape, + GGML_TYPE_F32); +} + +const core::ModuleSchema & LeakyReluModule::static_schema() noexcept { + return kLeakyReluSchema; +} + const core::ModuleSchema & SigmoidModule::schema() const noexcept { return static_schema(); } diff --git a/src/framework/modules/attention/attention_internal.h b/src/framework/modules/attention/attention_internal.h index 60879f1d..616a5dca 100644 --- a/src/framework/modules/attention/attention_internal.h +++ b/src/framework/modules/attention/attention_internal.h @@ -67,6 +67,16 @@ inline const core::ModuleSchema kGatedFeedForwardSchema = { "Applies a gated feed-forward block using gate, up, and down projections.", }; +inline const core::ModuleSchema kConvFeedForwardSchema = { + "ConvFeedForward", + "nn.block", + kInputOutputInputs, + 1, + kSingleOutput, + 1, + "Applies a convolutional feed-forward block to [batch, frames, hidden] inputs.", +}; + inline const core::ModuleSchema kSelfAttentionSchema = { "SelfAttention", "nn.attention", diff --git a/src/framework/modules/attention/cross_attention.cpp b/src/framework/modules/attention/cross_attention.cpp index f3e83005..13888656 100644 --- a/src/framework/modules/attention/cross_attention.cpp +++ b/src/framework/modules/attention/cross_attention.cpp @@ -4,6 +4,150 @@ namespace engine::modules { using namespace attention::internal; +namespace { + +int64_t cross_key_value_size(const AttentionConfig & config) { + return config.key_value_size > 0 ? config.key_value_size : config.hidden_size; +} + +int64_t cross_attention_size(const AttentionConfig & config) { + return config.attention_size > 0 ? config.attention_size : config.hidden_size; +} + +int64_t cross_head_dim(const AttentionConfig & config) { + const int64_t attention_size = cross_attention_size(config); + return config.head_dim > 0 ? config.head_dim : attention_size / config.num_heads; +} + +LinearWeights require_packed_kv_weights(const AttentionWeights & weights, bool use_bias) { + if (!weights.qkv_weight.has_value()) { + throw std::runtime_error("CrossAttentionModule packed KV path requires qkv_weight"); + } + if (use_bias && !weights.qkv_bias.has_value()) { + throw std::runtime_error("CrossAttentionModule packed KV path requires qkv_bias when bias is enabled"); + } + return {*weights.qkv_weight, weights.qkv_bias}; +} + +void validate_cross_query(const core::TensorValue & query, const AttentionConfig & config) { + validate_sequence_input(query, config.hidden_size, "query"); + const int64_t attention_size = cross_attention_size(config); + const int64_t head_dim = cross_head_dim(config); + if (attention_size <= 0 || head_dim <= 0 || config.num_heads <= 0 || + attention_size != config.num_heads * head_dim) { + throw std::runtime_error("CrossAttentionModule attention_size must equal num_heads * head_dim"); + } +} + +void validate_cross_memory(const core::TensorValue & memory, const AttentionConfig & config) { + validate_sequence_input(memory, cross_key_value_size(config), "memory"); +} + +void validate_cross_cache( + const CrossAttentionKeyValue & key_value, + const core::TensorValue & query, + const core::TensorValue & memory_mask, + const AttentionConfig & config) { + const int64_t head_dim = cross_head_dim(config); + if (key_value.key.shape.rank != 4 || key_value.value.shape.rank != 4 || + key_value.key.shape.dims[0] != query.shape.dims[0] || + key_value.value.shape.dims[0] != query.shape.dims[0] || + key_value.key.shape.dims[1] != config.num_heads || + key_value.value.shape.dims[1] != config.num_heads || + key_value.key.shape.dims[2] != memory_mask.shape.dims[1] || + key_value.value.shape.dims[2] != memory_mask.shape.dims[1] || + key_value.key.shape.dims[3] != head_dim || + key_value.value.shape.dims[3] != head_dim) { + throw std::runtime_error("CrossAttentionModule cached KV shape is invalid"); + } +} + +core::TensorValue build_cross_query( + core::ModuleBuildContext & ctx, + const core::TensorValue & query, + const AttentionConfig & config, + const AttentionWeights & weights) { + const int64_t attention_size = cross_attention_size(config); + const int64_t head_dim = cross_head_dim(config); + auto projected = LinearModule({ + config.hidden_size, + attention_size, + config.use_bias, + config.projection_precision, + }).build(ctx, query, make_linear_weights(weights.q_weight, weights.q_bias)); + projected = core::reshape_tensor( + ctx, + ensure_contiguous_layout(ctx, projected), + core::TensorShape::from_dims({projected.shape.dims[0], projected.shape.dims[1], config.num_heads, head_dim})); + return permute_tensor(ctx, projected, {0, 2, 1, 3}); +} + +core::TensorValue build_cross_probabilities( + core::ModuleBuildContext & ctx, + const core::TensorValue & query_heads, + const core::TensorValue & key_heads, + const core::TensorValue & memory_mask, + int64_t head_dim, + const core::TensorValue * attention_prior, + core::TensorValue * last_attention) { + auto kt = permute_tensor(ctx, key_heads, {0, 1, 3, 2}); + auto scores = MatMulModule().build(ctx, query_heads, kt); + scores = core::wrap_tensor( + ggml_scale(ctx.ggml, scores.tensor, 1.0F / std::sqrt(static_cast(head_dim))), + scores.shape, + GGML_TYPE_F32); + auto key_mask = core::reshape_tensor( + ctx, + ensure_contiguous_layout(ctx, core::wrap_tensor(ggml_cast(ctx.ggml, memory_mask.tensor, GGML_TYPE_F32), memory_mask.shape, GGML_TYPE_F32)), + core::TensorShape::from_dims({memory_mask.shape.dims[0], 1, 1, memory_mask.shape.dims[1]})); + auto score_mask = core::wrap_tensor( + ggml_scale_bias(ctx.ggml, key_mask.tensor, 1.0e30F, -1.0e30F), + key_mask.shape, + GGML_TYPE_F32); + scores = AddModule().build(ctx, scores, RepeatModule({scores.shape}).build(ctx, score_mask)); + auto probs = SoftmaxModule().build(ctx, scores); + probs = MulModule().build(ctx, probs, RepeatModule({probs.shape}).build(ctx, key_mask)); + if (attention_prior != nullptr) { + auto prior_repeated = core::wrap_tensor(ggml_repeat(ctx.ggml, attention_prior->tensor, probs.tensor), probs.shape, GGML_TYPE_F32); + probs = MulModule().build(ctx, probs, prior_repeated); + auto normalizer = core::wrap_tensor( + ggml_repeat(ctx.ggml, ggml_sum_rows(ctx.ggml, probs.tensor), probs.tensor), + probs.shape, + GGML_TYPE_F32); + probs = core::wrap_tensor(ggml_div(ctx.ggml, probs.tensor, normalizer.tensor), probs.shape, GGML_TYPE_F32); + } + if (last_attention != nullptr) { + *last_attention = SliceModule({2, query_heads.shape.dims[2] - 1, 1}).build(ctx, probs); + *last_attention = ensure_contiguous_layout(ctx, *last_attention); + } + return probs; +} + +core::TensorValue build_cross_output( + core::ModuleBuildContext & ctx, + const core::TensorValue & query, + const core::TensorValue & probabilities, + const core::TensorValue & value_heads, + const AttentionConfig & config, + const AttentionWeights & weights) { + const int64_t attention_size = cross_attention_size(config); + auto context = MatMulModule().build(ctx, probabilities, value_heads); + context = permute_tensor(ctx, context, {0, 2, 1, 3}); + context = ensure_contiguous_layout(ctx, context); + context = core::reshape_tensor( + ctx, + context, + core::TensorShape::from_dims({query.shape.dims[0], query.shape.dims[1], attention_size})); + return LinearModule({ + attention_size, + config.hidden_size, + config.use_bias, + config.projection_precision, + }).build(ctx, context, make_linear_weights(weights.out_weight, weights.out_bias)); +} + +} // namespace + CrossAttentionModule::CrossAttentionModule(AttentionConfig config) : config_(config) { validate_attention_config(config_); } @@ -21,9 +165,86 @@ core::TensorValue CrossAttentionModule::build( const core::TensorValue & query, const core::TensorValue & memory, const AttentionWeights & weights) const { + if (config_.use_packed_kv) { + throw std::runtime_error("CrossAttentionModule packed KV path requires memory_mask"); + } return build_attention_impl(ctx, query, memory, config_, require_attention_weights(weights, config_.use_bias)); } +core::TensorValue CrossAttentionModule::build( + core::ModuleBuildContext & ctx, + const core::TensorValue & query, + const core::TensorValue & memory, + const AttentionWeights & weights, + const core::TensorValue & memory_mask, + const core::TensorValue * attention_prior, + core::TensorValue * last_attention) const { + if (!config_.use_packed_kv) { + throw std::runtime_error("CrossAttentionModule masked path currently requires packed KV"); + } + validate_cross_query(query, config_); + validate_cross_memory(memory, config_); + const auto key_value = build_key_value(ctx, memory, weights); + return build_cached(ctx, query, key_value, weights, memory_mask, attention_prior, last_attention); +} + +core::TensorValue CrossAttentionModule::build_cached( + core::ModuleBuildContext & ctx, + const core::TensorValue & query, + const CrossAttentionKeyValue & key_value, + const AttentionWeights & weights, + const core::TensorValue & memory_mask, + const core::TensorValue * attention_prior, + core::TensorValue * last_attention) const { + if (!config_.use_packed_kv) { + throw std::runtime_error("CrossAttentionModule cached path requires packed KV"); + } + validate_cross_query(query, config_); + validate_cross_cache(key_value, query, memory_mask, config_); + auto query_heads = build_cross_query(ctx, query, config_, weights); + auto probs = build_cross_probabilities( + ctx, + query_heads, + key_value.key, + memory_mask, + cross_head_dim(config_), + attention_prior, + last_attention); + return build_cross_output(ctx, query, probs, key_value.value, config_, weights); +} + +CrossAttentionKeyValue CrossAttentionModule::build_key_value( + core::ModuleBuildContext & ctx, + const core::TensorValue & memory, + const AttentionWeights & weights) const { + if (!config_.use_packed_kv) { + throw std::runtime_error("CrossAttentionModule build_key_value requires packed KV"); + } + validate_cross_memory(memory, config_); + const int64_t attention_size = cross_attention_size(config_); + const int64_t head_dim = cross_head_dim(config_); + auto kv = LinearModule({ + cross_key_value_size(config_), + 2 * attention_size, + config_.use_bias, + config_.projection_precision, + }).build(ctx, memory, require_packed_kv_weights(weights, config_.use_bias)); + auto key = SliceModule({2, 0, attention_size}).build(ctx, kv); + key = core::reshape_tensor( + ctx, + ensure_contiguous_layout(ctx, key), + core::TensorShape::from_dims({key.shape.dims[0], key.shape.dims[1], config_.num_heads, head_dim})); + auto value = SliceModule({2, attention_size, attention_size}).build(ctx, kv); + value = core::reshape_tensor( + ctx, + ensure_contiguous_layout(ctx, value), + core::TensorShape::from_dims({value.shape.dims[0], value.shape.dims[1], config_.num_heads, head_dim})); + return { + ensure_contiguous_layout(ctx, permute_tensor(ctx, key, {0, 2, 1, 3})), + ensure_contiguous_layout(ctx, permute_tensor(ctx, value, {0, 2, 1, 3})), + }; +} + const core::ModuleSchema & CrossAttentionModule::static_schema() noexcept { return kCrossAttentionSchema; } diff --git a/src/framework/modules/attention/feed_forward.cpp b/src/framework/modules/attention/feed_forward.cpp index 15789597..9d32602e 100644 --- a/src/framework/modules/attention/feed_forward.cpp +++ b/src/framework/modules/attention/feed_forward.cpp @@ -1,4 +1,5 @@ #include "attention_internal.h" +#include "engine/framework/modules/streaming_conv_modules.h" namespace engine::modules { @@ -83,4 +84,83 @@ const core::ModuleSchema & GatedFeedForwardModule::static_schema() noexcept { return kGatedFeedForwardSchema; } +namespace { + +core::TensorValue build_conv_feed_forward_conv( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const Conv1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t kernel_size, + bool causal, + bool use_bias) { + if (causal) { + return StreamingConv1dModule({ + in_channels, + out_channels, + kernel_size, + 1, + 1, + use_bias, + StreamingPadMode::Constant, + StreamingConv1dPaddingMode::StrictCausal, + }).build(ctx, input, weights); + } + return Conv1dModule({in_channels, out_channels, kernel_size, 1, static_cast(kernel_size / 2), 1, use_bias}) + .build(ctx, input, weights); +} + +} // namespace + +ConvFeedForwardModule::ConvFeedForwardModule(ConvFeedForwardConfig config) : config_(config) { + validate_hidden_positive(config_.hidden_size, "ConvFeedForwardConfig.hidden_size"); + validate_hidden_positive(config_.intermediate_size, "ConvFeedForwardConfig.intermediate_size"); + if (config_.kernel_size <= 0) { + throw std::runtime_error("ConvFeedForwardConfig.kernel_size must be positive"); + } +} + +const ConvFeedForwardConfig & ConvFeedForwardModule::config() const noexcept { + return config_; +} + +const core::ModuleSchema & ConvFeedForwardModule::schema() const noexcept { + return static_schema(); +} + +core::TensorValue ConvFeedForwardModule::build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvFeedForwardWeights & weights) const { + core::validate_rank_between(input, 3, 3, "input"); + core::validate_last_dim(input, config_.hidden_size, "input"); + + auto x = TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, input); + x = build_conv_feed_forward_conv( + ctx, + x, + weights.proj, + config_.hidden_size, + config_.intermediate_size, + config_.kernel_size, + config_.causal, + config_.use_bias); + x = GeluModule({config_.gelu_approximation}).build(ctx, x); + x = build_conv_feed_forward_conv( + ctx, + x, + weights.out, + config_.intermediate_size, + config_.hidden_size, + config_.kernel_size, + config_.causal, + config_.use_bias); + return TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); +} + +const core::ModuleSchema & ConvFeedForwardModule::static_schema() noexcept { + return kConvFeedForwardSchema; +} + } // namespace engine::modules diff --git a/src/framework/modules/attention/self_attention.cpp b/src/framework/modules/attention/self_attention.cpp index f0e6e4e3..9fc5791a 100644 --- a/src/framework/modules/attention/self_attention.cpp +++ b/src/framework/modules/attention/self_attention.cpp @@ -1,9 +1,78 @@ #include "attention_internal.h" +#include "engine/framework/modules/attention/grouped_query_attention.h" + namespace engine::modules { using namespace attention::internal; +namespace { + +void require_packed_self_attention_config(const AttentionConfig & config) { + if (!config.use_packed_qkv) { + throw std::runtime_error("SelfAttentionModule packed QKV path requires use_packed_qkv"); + } +} + +LinearWeights require_packed_qkv_weights(const AttentionWeights & weights, bool use_bias) { + if (!weights.qkv_weight.has_value()) { + throw std::runtime_error("SelfAttentionModule packed QKV path requires qkv_weight"); + } + if (use_bias && !weights.qkv_bias.has_value()) { + throw std::runtime_error("SelfAttentionModule packed QKV path requires qkv_bias when bias is enabled"); + } + return {*weights.qkv_weight, weights.qkv_bias}; +} + +struct PackedQkvHeads { + core::TensorValue q; + core::TensorValue k; + core::TensorValue v; +}; + +PackedQkvHeads build_packed_qkv_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const AttentionConfig & config, + const AttentionWeights & weights) { + const int64_t head_dim = config.hidden_size / config.num_heads; + auto qkv = LinearModule({ + config.hidden_size, + 3 * config.hidden_size, + config.use_bias, + config.projection_precision, + }).build(ctx, input, require_packed_qkv_weights(weights, config.use_bias)); + + auto q = SliceModule({2, 0, config.hidden_size}).build(ctx, qkv); + q = reshape_heads(ctx, ensure_contiguous_layout(ctx, q), config.num_heads, head_dim); + auto k = SliceModule({2, config.hidden_size, config.hidden_size}).build(ctx, qkv); + k = reshape_heads(ctx, ensure_contiguous_layout(ctx, k), config.num_heads, head_dim); + auto v = SliceModule({2, 2 * config.hidden_size, config.hidden_size}).build(ctx, qkv); + v = reshape_heads(ctx, ensure_contiguous_layout(ctx, v), config.num_heads, head_dim); + return {q, k, v}; +} + +core::TensorValue project_attention_output( + core::ModuleBuildContext & ctx, + const core::TensorValue & context, + const core::TensorShape & input_shape, + const AttentionConfig & config, + const AttentionWeights & weights) { + auto output = ensure_contiguous_layout(ctx, context); + output = core::reshape_tensor( + ctx, + output, + core::TensorShape::from_dims({input_shape.dims[0], input_shape.dims[1], config.hidden_size})); + return LinearModule({ + config.hidden_size, + config.hidden_size, + config.use_bias, + config.projection_precision, + }).build(ctx, output, make_linear_weights(weights.out_weight, weights.out_bias)); +} + +} // namespace + SelfAttentionModule::SelfAttentionModule(AttentionConfig config) : config_(config) { validate_attention_config(config_); } @@ -20,7 +89,86 @@ core::TensorValue SelfAttentionModule::build( core::ModuleBuildContext & ctx, const core::TensorValue & input, const AttentionWeights & weights) const { - return build_attention_impl(ctx, input, input, config_, require_attention_weights(weights, config_.use_bias)); + return build(ctx, input, weights, std::nullopt); +} + +core::TensorValue SelfAttentionModule::build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const AttentionWeights & weights, + const std::optional & attention_mask) const { + if (!config_.use_packed_qkv) { + if (attention_mask.has_value()) { + throw std::runtime_error("SelfAttentionModule attention_mask requires packed QKV path"); + } + return build_attention_impl(ctx, input, input, config_, require_attention_weights(weights, config_.use_bias)); + } + require_packed_self_attention_config(config_); + validate_sequence_input(input, config_.hidden_size, "input"); + const int64_t head_dim = config_.hidden_size / config_.num_heads; + const auto qkv = build_packed_qkv_heads(ctx, input, config_, weights); + auto q_heads = permute_tensor(ctx, qkv.q, {0, 2, 1, 3}); + q_heads = ensure_contiguous_layout(ctx, q_heads); + auto k_heads = permute_tensor(ctx, qkv.k, {0, 2, 1, 3}); + auto v_heads = permute_tensor(ctx, qkv.v, {0, 2, 1, 3}); + core::TensorValue context; + if (attention_mask.has_value()) { + context = GroupedQueryAttentionModule({ + head_dim, + GroupedQueryAttentionLowering::FlashGroupedViewKV, + config_.attention_precision, + config_.causal ? AttentionCausality::Causal : AttentionCausality::NonCausal, + }).build(ctx, q_heads, k_heads, v_heads, attention_mask); + } else { + auto kt = permute_tensor(ctx, k_heads, {0, 1, 3, 2}); + auto scores = MatMulModule().build(ctx, q_heads, kt); + if (config_.causal) { + scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, 0), scores.shape, GGML_TYPE_F32); + } + scores = core::wrap_tensor(ggml_scale(ctx.ggml, scores.tensor, 1.0F / std::sqrt(static_cast(head_dim))), scores.shape, GGML_TYPE_F32); + auto probs = SoftmaxModule().build(ctx, scores); + context = MatMulModule().build(ctx, probs, v_heads); + context = permute_tensor(ctx, context, {0, 2, 1, 3}); + } + return project_attention_output(ctx, context, input.shape, config_, weights); +} + +StreamingAttentionOutputs SelfAttentionModule::build_cached_tail( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const AttentionWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask, + FastKVSetRowsMode set_rows_mode) const { + require_packed_self_attention_config(config_); + if (!config_.causal) { + throw std::runtime_error("SelfAttentionModule cached packed tail requires causal attention"); + } + validate_sequence_input(input, config_.hidden_size, "input"); + const int64_t head_dim = config_.hidden_size / config_.num_heads; + auto qkv = build_packed_qkv_heads(ctx, input, config_, weights); + qkv.k = ensure_contiguous_layout(ctx, qkv.k); + qkv.v = ensure_contiguous_layout(ctx, qkv.v); + const FastKVSetRowsModule set_rows({set_rows_mode}); + auto updated_key = set_rows.build(ctx, cache_key, qkv.k, cache_slot); + auto updated_value = set_rows.build(ctx, cache_value, qkv.v, cache_slot); + auto q_heads = permute_tensor(ctx, qkv.q, {0, 2, 1, 3}); + q_heads = ensure_contiguous_layout(ctx, q_heads); + auto k_heads = permute_tensor(ctx, updated_key, {0, 2, 1, 3}); + auto v_heads = permute_tensor(ctx, updated_value, {0, 2, 1, 3}); + auto context = GroupedQueryAttentionModule({ + head_dim, + GroupedQueryAttentionLowering::FlashGroupedViewKV, + config_.attention_precision, + AttentionCausality::Causal, + }).build(ctx, q_heads, k_heads, v_heads, attention_mask); + return { + project_attention_output(ctx, context, input.shape, config_, weights), + updated_key, + updated_value, + }; } const core::ModuleSchema & SelfAttentionModule::static_schema() noexcept { diff --git a/src/framework/modules/codecs/nemo_nano_codec.cpp b/src/framework/modules/codecs/nemo_nano_codec.cpp new file mode 100644 index 00000000..2b2b1610 --- /dev/null +++ b/src/framework/modules/codecs/nemo_nano_codec.cpp @@ -0,0 +1,581 @@ +#include "engine/framework/modules/codecs/nemo_nano_codec.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::modules { +namespace { + +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct CodecResidualBlockWeights { + Snake1dWeights input_snake; + Conv1dWeights input_conv; + Snake1dWeights skip_snake; + Conv1dWeights skip_conv; +}; + +struct CodecStageWeights { + Snake1dWeights up_snake; + std::vector upsample_groups; + std::vector residuals; +}; + +struct CodecWeights { + std::shared_ptr store; + Conv1dWeights pre_conv; + std::vector stages; + Snake1dWeights post_snake; + Conv1dWeights post_conv; +}; + +std::vector read_exact_f32_tensor(ggml_tensor * tensor, size_t count, const char * name) { + auto out = core::read_tensor_f32(tensor); + if (out.size() != count) { + throw std::runtime_error( + std::string(name) + " readback element count mismatch: expected " + + std::to_string(count) + ", got " + std::to_string(out.size())); + } + return out; +} + +void validate_config(const NemoNanoCodecConfig & config) { + if (config.sample_rate <= 0 || config.input_dim <= 0 || config.base_channels <= 0 || config.audio_codebooks <= 0) { + throw std::runtime_error("NeMo nano codec config requires positive sample_rate, input_dim, base_channels, and audio_codebooks"); + } + if (config.upsample_rates.empty() || config.resblock_kernel_sizes.empty() || config.resblock_dilation_sizes.empty()) { + throw std::runtime_error("NeMo nano codec config requires upsample and residual block settings"); + } + if (config.fsq_num_levels.empty() || config.fsq_num_levels.size() != config.fsq_dim_base_index.size()) { + throw std::runtime_error("NeMo nano codec config requires matching FSQ levels and base indices"); + } + if (config.input_dim != config.audio_codebooks * static_cast(config.fsq_num_levels.size())) { + throw std::runtime_error("NeMo nano codec input_dim must match audio_codebooks times FSQ dimensions per group"); + } +} + +std::vector fold_weight_norm( + const std::vector & g, + const std::vector & v, + int64_t outer, + int64_t inner, + int64_t kernel) { + std::vector out(v.size()); + for (int64_t o = 0; o < outer; ++o) { + double sum = 0.0; + for (int64_t i = 0; i < inner; ++i) { + for (int64_t k = 0; k < kernel; ++k) { + const float value = v[static_cast((o * inner + i) * kernel + k)]; + sum += static_cast(value) * static_cast(value); + } + } + const float scale = g[static_cast(o)] / static_cast(std::sqrt(sum)); + for (int64_t i = 0; i < inner; ++i) { + for (int64_t k = 0; k < kernel; ++k) { + const size_t index = static_cast((o * inner + i) * kernel + k); + out[index] = v[index] * scale; + } + } + } + return out; +} + +std::vector> split_grouped_transpose_conv1d_weight( + const std::vector & weight, + int64_t in_channels, + int64_t out_channels, + int64_t kernel) { + if (static_cast(weight.size()) != in_channels * kernel) { + throw std::runtime_error("NeMo nano codec grouped ConvTranspose1d folded weight shape mismatch"); + } + const int64_t inputs_per_group = in_channels / out_channels; + if (inputs_per_group <= 0 || inputs_per_group * out_channels != in_channels) { + throw std::runtime_error("NeMo nano codec grouped ConvTranspose1d channel ratio is invalid"); + } + std::vector> groups(static_cast(out_channels)); + for (auto & group : groups) { + group.resize(static_cast(inputs_per_group * kernel), 0.0F); + } + for (int64_t group = 0; group < out_channels; ++group) { + const int64_t input_start = group * inputs_per_group; + for (int64_t input_offset = 0; input_offset < inputs_per_group; ++input_offset) { + const int64_t in_channel = input_start + input_offset; + for (int64_t tap = 0; tap < kernel; ++tap) { + groups[static_cast(group)][static_cast(input_offset * kernel + tap)] = + weight[static_cast(in_channel * kernel + tap)]; + } + } + } + return groups; +} + +std::vector load_weight_norm_grouped_convtranspose1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t in_channels, + int64_t out_channels, + int64_t kernel_size, + bool use_bias) { + const auto g = source.require_f32(prefix + ".parametrizations.weight.original0", {in_channels, 1, 1}); + const auto v = source.require_f32(prefix + ".parametrizations.weight.original1", {in_channels, 1, kernel_size}); + const int64_t inputs_per_group = in_channels / out_channels; + if (inputs_per_group <= 0 || inputs_per_group * out_channels != in_channels) { + throw std::runtime_error("NeMo nano codec grouped ConvTranspose1d channel ratio is invalid"); + } + const auto folded = fold_weight_norm(g, v, in_channels, 1, kernel_size); + const auto groups = split_grouped_transpose_conv1d_weight(folded, in_channels, out_channels, kernel_size); + const auto bias = use_bias ? source.require_f32(prefix + ".bias", {out_channels}) : std::vector{}; + std::vector weights; + weights.reserve(static_cast(out_channels)); + for (int64_t group = 0; group < out_channels; ++group) { + ConvTranspose1dWeights item; + item.weight = store.make_from_f32( + core::TensorShape::from_dims({inputs_per_group, 1, kernel_size}), + storage_type, + groups[static_cast(group)]); + if (use_bias) { + item.bias = store.make_f32(core::TensorShape::from_dims({1}), {bias[static_cast(group)]}); + } + weights.push_back(std::move(item)); + } + return weights; +} + +Snake1dWeights load_half_snake_alpha( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + int64_t channels) { + const int64_t snake_channels = channels / 2; + const auto values = source.require_f32(name, {1, snake_channels, 1}); + return {store.make_f32(core::TensorShape::from_dims({snake_channels}), values)}; +} + +CodecWeights load_codec_weights( + const assets::TensorSource & source, + const NemoNanoCodecConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + const NemoNanoCodecRuntimeOptions & options) { + CodecWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "framework.nemo_nano_codec.weights", + options.weight_context_bytes); + weights.pre_conv = binding::weight_norm_conv1d_from_source( + *weights.store, + source, + "audio_decoder.pre_conv.conv", + options.weight_storage_type, + config.base_channels, + config.input_dim, + 7, + true); + int64_t in_channels = config.base_channels; + weights.stages.reserve(config.upsample_rates.size()); + for (size_t stage = 0; stage < config.upsample_rates.size(); ++stage) { + const int64_t rate = config.upsample_rates[stage]; + const int64_t out_channels = in_channels / 2; + CodecStageWeights stage_weights; + stage_weights.up_snake = load_half_snake_alpha( + *weights.store, + source, + "audio_decoder.activations." + std::to_string(stage) + ".activation.snake_act.alpha", + in_channels); + stage_weights.upsample_groups = load_weight_norm_grouped_convtranspose1d( + *weights.store, + source, + "audio_decoder.up_sample_conv_layers." + std::to_string(stage) + ".conv", + options.weight_storage_type, + in_channels, + out_channels, + rate * 2, + true); + for (size_t kernel_index = 0; kernel_index < config.resblock_kernel_sizes.size(); ++kernel_index) { + const int64_t kernel = config.resblock_kernel_sizes[kernel_index]; + for (size_t dilation_index = 0; dilation_index < config.resblock_dilation_sizes.size(); ++dilation_index) { + const std::string prefix = + "audio_decoder.res_layers." + std::to_string(stage) + + ".res_blocks." + std::to_string(kernel_index) + + ".res_blocks." + std::to_string(dilation_index); + CodecResidualBlockWeights block; + block.input_snake = load_half_snake_alpha( + *weights.store, + source, + prefix + ".input_activation.activation.snake_act.alpha", + out_channels); + block.input_conv = binding::weight_norm_conv1d_from_source( + *weights.store, + source, + prefix + ".input_conv.conv", + options.weight_storage_type, + out_channels, + out_channels, + kernel, + true); + block.skip_snake = load_half_snake_alpha( + *weights.store, + source, + prefix + ".skip_activation.activation.snake_act.alpha", + out_channels); + block.skip_conv = binding::weight_norm_conv1d_from_source( + *weights.store, + source, + prefix + ".skip_conv.conv", + options.weight_storage_type, + out_channels, + out_channels, + kernel, + true); + stage_weights.residuals.push_back(std::move(block)); + } + } + weights.stages.push_back(std::move(stage_weights)); + in_channels = out_channels; + } + weights.post_snake = load_half_snake_alpha( + *weights.store, + source, + "audio_decoder.post_activation.activation.snake_act.alpha", + in_channels); + weights.post_conv = binding::weight_norm_conv1d_from_source( + *weights.store, + source, + "audio_decoder.post_conv.conv", + options.weight_storage_type, + 1, + in_channels, + 3, + true); + weights.store->upload(); + return weights; +} + +core::TensorValue causal_grouped_convtranspose1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const std::vector & weights, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + int64_t stride) { + const int64_t inputs_per_group = in_channels / out_channels; + if (static_cast(weights.size()) != out_channels) { + throw std::runtime_error("NeMo nano codec ConvTranspose1d group count mismatch"); + } + core::TensorValue out; + for (int64_t group = 0; group < out_channels; ++group) { + const int64_t input_start = group * inputs_per_group; + auto input_slice = SliceModule({1, input_start, inputs_per_group}).build(ctx, input); + const auto & group_weights = weights[static_cast(group)]; + auto group_out = ConvTranspose1dModule({ + inputs_per_group, + 1, + kernel, + static_cast(stride), + 0, + 1, + group_weights.bias.has_value(), + }).build(ctx, input_slice, group_weights); + out = out.valid() ? ConcatModule({1}).build(ctx, out, group_out) : group_out; + } + const int64_t trim_right = kernel - stride; + if (trim_right <= 0) { + return out; + } + return SliceModule({2, 0, out.shape.dims[2] - trim_right}).build(ctx, out); +} + +core::TensorValue half_snake( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const Snake1dWeights & weights) { + const int64_t snake_channels = input.shape.dims[1] / 2; + auto left = SliceModule({1, 0, snake_channels}).build(ctx, input); + auto right = SliceModule({1, snake_channels, input.shape.dims[1] - snake_channels}).build(ctx, input); + left = Snake1dModule({snake_channels}).build(ctx, left, weights); + right = LeakyReluModule({0.01F}).build(ctx, right); + return ConcatModule({1}).build(ctx, left, right); +} + +core::TensorValue codec_residual( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const CodecResidualBlockWeights & weights, + int64_t channels, + int64_t kernel, + int64_t dilation) { + auto x = half_snake(ctx, input, weights.input_snake); + x = CausalConv1dModule({ + channels, + channels, + kernel, + 1, + static_cast(dilation), + true, + CausalConv1dPadMode::Constant, + CausalConv1dPaddingMode::StrictCausal, + }).build(ctx, x, weights.input_conv); + x = half_snake(ctx, x, weights.skip_snake); + x = CausalConv1dModule({ + channels, + channels, + kernel, + 1, + 1, + true, + CausalConv1dPadMode::Constant, + CausalConv1dPaddingMode::StrictCausal, + }).build(ctx, x, weights.skip_conv); + return ResidualAddModule().build(ctx, input, x); +} + +core::TensorValue codec_residual_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const CodecStageWeights & weights, + const NemoNanoCodecConfig & config, + int64_t channels) { + core::TensorValue summed; + size_t residual_index = 0; + for (const int64_t kernel : config.resblock_kernel_sizes) { + auto branch = input; + for (const int64_t dilation : config.resblock_dilation_sizes) { + branch = codec_residual( + ctx, + branch, + weights.residuals[residual_index++], + channels, + kernel, + dilation); + } + summed = summed.valid() ? AddModule().build(ctx, summed, branch) : branch; + } + return core::wrap_tensor( + ggml_scale( + ctx.ggml, + summed.tensor, + 1.0F / static_cast(config.resblock_kernel_sizes.size())), + summed.shape, + summed.type); +} + +core::TensorValue build_codec_decoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const NemoNanoCodecConfig & config, + const CodecWeights & weights) { + auto x = CausalConv1dModule({ + config.input_dim, + config.base_channels, + 7, + 1, + 1, + true, + CausalConv1dPadMode::Constant, + CausalConv1dPaddingMode::StrictCausal, + }).build(ctx, input, weights.pre_conv); + int64_t channels = config.base_channels; + for (size_t stage = 0; stage < weights.stages.size(); ++stage) { + const int64_t rate = config.upsample_rates[stage]; + const int64_t out_channels = channels / 2; + x = half_snake(ctx, x, weights.stages[stage].up_snake); + x = causal_grouped_convtranspose1d(ctx, x, weights.stages[stage].upsample_groups, channels, out_channels, rate * 2, rate); + x = codec_residual_layer(ctx, x, weights.stages[stage], config, out_channels); + channels = out_channels; + } + x = half_snake(ctx, x, weights.post_snake); + x = CausalConv1dModule({ + channels, + 1, + 3, + 1, + 1, + true, + CausalConv1dPadMode::Constant, + CausalConv1dPaddingMode::StrictCausal, + }).build(ctx, x, weights.post_conv); + return core::wrap_tensor(ggml_clamp(ctx.ggml, x.tensor, -1.0F, 1.0F), x.shape, GGML_TYPE_F32); +} + +} // namespace + +struct NemoNanoCodecRuntime::Impl { + struct Graph { + Graph( + const Impl & owner, + int64_t input_frames) + : frames(input_frames), + owner_backend(owner.backend) { + ggml_init_params params{owner.options.graph_arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("NeMo nano codec failed to create graph context"); + } + core::ModuleBuildContext build{ctx.get(), "framework.nemo_nano_codec", owner.backend_type}; + input = core::make_tensor(build, GGML_TYPE_F32, core::TensorShape::from_dims({1, owner.config.input_dim, frames})); + output = build_codec_decoder(build, input, owner.config, *owner.weights); + output = core::ensure_backend_addressable_layout(build, output); + graph = ggml_new_graph_custom(ctx.get(), 262144, false); + ggml_set_output(output.tensor); + ggml_build_forward_expand(graph, output.tensor); + gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(owner.backend)); + if (gallocr == nullptr) { + throw std::runtime_error("NeMo nano codec failed to create graph allocator"); + } + if (!ggml_gallocr_reserve(gallocr, graph) || !ggml_gallocr_alloc_graph(gallocr, graph)) { + throw std::runtime_error("NeMo nano codec failed to allocate graph"); + } + } + + ~Graph() { + if (owner_backend != nullptr && graph != nullptr) { + core::release_backend_graph_resources(owner_backend, graph); + } + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + gallocr = nullptr; + } + } + + Graph(const Graph &) = delete; + Graph & operator=(const Graph &) = delete; + + int64_t frames = 0; + ggml_backend_t owner_backend = nullptr; + std::unique_ptr ctx; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t gallocr = nullptr; + core::TensorValue input; + core::TensorValue output; + }; + + Impl( + std::shared_ptr source, + core::ExecutionContext & execution, + NemoNanoCodecConfig input_config, + NemoNanoCodecRuntimeOptions input_options) + : config(std::move(input_config)), + backend(execution.backend()), + backend_type(execution.backend_type()), + options(input_options) { + validate_config(config); + if (source == nullptr) { + throw std::runtime_error("NeMo nano codec runtime requires tensor source"); + } + weights = std::make_shared( + load_codec_weights(*source, config, backend, backend_type, options)); + } + + std::vector fsq_decode(const std::vector & codes) const { + const int64_t frames = static_cast(codes.size()) / config.audio_codebooks; + if (frames <= 0 || static_cast(codes.size()) != frames * config.audio_codebooks) { + throw std::runtime_error("NeMo nano codec code shape is invalid"); + } + const int64_t dims_per_group = static_cast(config.fsq_num_levels.size()); + std::vector out(static_cast(config.input_dim * frames), 0.0F); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t group = 0; group < config.audio_codebooks; ++group) { + const int32_t index = codes[static_cast(frame * config.audio_codebooks + group)]; + for (int64_t d = 0; d < dims_per_group; ++d) { + const int32_t base = config.fsq_dim_base_index[static_cast(d)]; + const int32_t levels = config.fsq_num_levels[static_cast(d)]; + const int32_t nonnegative = (index / base) % levels; + const int32_t scale = levels / 2; + const float value = static_cast(nonnegative - scale) / static_cast(scale); + const int64_t channel = group * dims_per_group + d; + out[static_cast(channel * frames + frame)] = value; + } + } + } + return out; + } + + Graph & graph_for_frames(int64_t frames) { + if (graph == nullptr || graph->frames != frames) { + graph = std::make_unique(*this, frames); + } + return *graph; + } + + runtime::AudioBuffer decode_codes(const std::vector & codes) { + const int64_t frames = static_cast(codes.size()) / config.audio_codebooks; + auto dequantized = fsq_decode(codes); + auto & graph_ref = graph_for_frames(frames); + ggml_backend_tensor_set( + graph_ref.input.tensor, + dequantized.data(), + 0, + dequantized.size() * sizeof(float)); + const auto start = Clock::now(); + core::compute_backend_graph(backend, graph_ref.graph); + debug::timing_log_scalar("nemo_nano_codec.graph.compute_ms", debug::elapsed_ms(start)); + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(config.sample_rate); + audio.channels = 1; + audio.samples = read_exact_f32_tensor( + graph_ref.output.tensor, + static_cast(graph_ref.output.shape.dims[2]), + "NeMo nano codec output"); + return audio; + } + + NemoNanoCodecConfig config; + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + NemoNanoCodecRuntimeOptions options; + std::shared_ptr weights; + std::unique_ptr graph; +}; + +NemoNanoCodecRuntime::NemoNanoCodecRuntime( + std::shared_ptr source, + core::ExecutionContext & execution, + NemoNanoCodecConfig config, + NemoNanoCodecRuntimeOptions options) + : impl_(std::make_unique(std::move(source), execution, std::move(config), options)) {} + +NemoNanoCodecRuntime::~NemoNanoCodecRuntime() = default; + +runtime::AudioBuffer NemoNanoCodecRuntime::decode_codes(const std::vector & codes) { + return impl_->decode_codes(codes); +} + +void NemoNanoCodecRuntime::release_runtime_graph() { + impl_->graph.reset(); +} + +} // namespace engine::modules diff --git a/src/framework/modules/streaming_conv_modules.cpp b/src/framework/modules/streaming_conv_modules.cpp index 26b51ea4..ff75beb4 100644 --- a/src/framework/modules/streaming_conv_modules.cpp +++ b/src/framework/modules/streaming_conv_modules.cpp @@ -6,6 +6,7 @@ #include #include +#include namespace engine::modules { @@ -73,6 +74,36 @@ core::TensorValue repeat_first_frame( return RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], prefix_frames})}).build(ctx, first); } +core::TensorValue zeros_like_suffix( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t suffix_frames) { + auto suffix = RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], suffix_frames})}) + .build(ctx, SliceModule({2, input.shape.dims[2] - 1, 1}).build(ctx, input)); + auto suffix_contiguous = tensor_layout::ensure_contiguous_layout_if_needed(ctx, suffix); + return core::wrap_tensor(ggml_scale(ctx.ggml, suffix_contiguous.tensor, 0.0f), suffix.shape, GGML_TYPE_F32); +} + +core::TensorValue repeat_last_frame( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t suffix_frames) { + auto last = SliceModule({2, input.shape.dims[2] - 1, 1}).build(ctx, input); + return RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], suffix_frames})}).build(ctx, last); +} + +std::pair streaming_conv1d_padding(const StreamingConv1dConfig & config, int64_t effective_kernel) { + switch (config.padding_mode) { + case StreamingConv1dPaddingMode::StreamingSame: + return {effective_kernel - config.stride, 0}; + case StreamingConv1dPaddingMode::StrictCausal: + return {effective_kernel - 1, 0}; + case StreamingConv1dPaddingMode::Explicit: + return {config.explicit_left, config.explicit_right}; + } + throw std::runtime_error("StreamingConv1dModule unknown padding mode"); +} + } DepthwiseConv1dModule::DepthwiseConv1dModule(DepthwiseConv1dConfig config) : config_(config) { @@ -175,6 +206,9 @@ StreamingConv1dModule::StreamingConv1dModule(StreamingConv1dConfig config) : con if (config_.stride <= 0 || config_.dilation <= 0) { throw std::runtime_error("StreamingConv1d stride and dilation must be positive"); } + if (config_.explicit_left < 0 || config_.explicit_right < 0) { + throw std::runtime_error("StreamingConv1d explicit padding must be non-negative"); + } } core::TensorValue StreamingConv1dModule::build( @@ -182,9 +216,9 @@ core::TensorValue StreamingConv1dModule::build( const core::TensorValue & input, const StreamingConv1dWeights & weights) const { const int64_t effective_kernel = (config_.kernel_size - 1) * config_.dilation + 1; - const int64_t left_pad = effective_kernel - config_.stride; - if (left_pad < 0) { - throw std::runtime_error("StreamingConv1dModule requires effective_kernel >= stride"); + const auto [left_pad, right_pad] = streaming_conv1d_padding(config_, effective_kernel); + if (left_pad < 0 || right_pad < 0) { + throw std::runtime_error("StreamingConv1dModule computed negative padding"); } if (input.shape.dims[2] <= 0) { throw std::runtime_error("StreamingConv1dModule input must have frames"); @@ -197,6 +231,12 @@ core::TensorValue StreamingConv1dModule::build( : zeros_like_prefix(ctx, input, left_pad); padded = ConcatModule({2}).build(ctx, prefix, input); } + if (right_pad > 0) { + core::TensorValue suffix = config_.pad_mode == StreamingPadMode::Replicate + ? repeat_last_frame(ctx, input, right_pad) + : zeros_like_suffix(ctx, input, right_pad); + padded = ConcatModule({2}).build(ctx, padded, suffix); + } return Conv1dModule({ config_.in_channels, config_.out_channels, diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index b8ec4f17..2dd7afb6 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -42,6 +42,16 @@ void validate_runtime_config(const QwenCausalDecodeRuntimeConfig & config) { if (config.readback_round_type.has_value() && *config.readback_round_type != GGML_TYPE_BF16) { throw std::runtime_error("QwenCausalDecodeRuntime readback rounding currently supports only bf16"); } + if (!config.logits_readback_token_ids.empty()) { + if (config.output_mode != QwenCausalDecodeOutputMode::Logits) { + throw std::runtime_error("QwenCausalDecodeRuntime compact logits readback requires logits mode"); + } + for (const int32_t token : config.logits_readback_token_ids) { + if (token < 0 || token >= config.decoder.logits_size) { + throw std::runtime_error("QwenCausalDecodeRuntime compact logits token id is outside logits size"); + } + } + } } core::TensorValue token_embedding_input( @@ -62,10 +72,26 @@ core::TensorValue token_embedding_input( core::TensorShape::from_dims({1, steps, config.stack.hidden_size})); } +core::TensorValue token_embedding_input_batched( + core::ModuleBuildContext & ctx, + const QwenCausalDecodeRuntimeWeights & weights, + const QwenCausalDecoderConfig & config, + ggml_tensor * token_ids, + int64_t batch_size, + int64_t steps) { + auto ids = core::wrap_tensor( + token_ids, + core::TensorShape::from_dims({batch_size, steps}), + GGML_TYPE_I32); + return EmbeddingModule({weights.token_embedding.shape.dims[0], config.stack.hidden_size}) + .build(ctx, ids, weights.token_embedding); +} + QwenDecoderHiddenConfig hidden_config_from_runtime(const QwenCausalDecodeRuntimeConfig & config) { QwenDecoderHiddenConfig out; out.stack = config.decoder.stack; out.hidden_mode = config.decoder.logits_mode; + out.static_cache_type = config.decoder.static_cache_type; return out; } @@ -96,6 +122,75 @@ void round_readback(std::vector & values, const QwenCausalDecodeRuntimeCo } } +core::TensorValue compact_logits_readback( + core::ModuleBuildContext & ctx, + const QwenCausalDecodeRuntimeConfig & config, + const core::TensorValue & logits, + const core::TensorValue & token_ids) { + if (logits.shape.last_dim() != config.decoder.logits_size) { + throw std::runtime_error("QwenCausalDecodeRuntime compact logits requires logits on the last dimension"); + } + + const int64_t rows = logits.shape.prefix_elements(); + const int64_t compact_size = static_cast(config.logits_readback_token_ids.size()); + const auto flat = core::reshape_tensor( + ctx, + logits, + core::TensorShape::from_dims({rows, config.decoder.logits_size})); + const auto transposed = core::wrap_tensor( + ggml_transpose(ctx.ggml, flat.tensor), + core::TensorShape::from_dims({config.decoder.logits_size, rows}), + flat.type); + const auto source = core::wrap_tensor( + ggml_cont(ctx.ggml, transposed.tensor), + transposed.shape, + transposed.type); + const auto selected = core::wrap_tensor( + ggml_get_rows(ctx.ggml, source.tensor, token_ids.tensor), + core::TensorShape::from_dims({compact_size, rows}), + GGML_TYPE_F32); + const auto selected_t = core::wrap_tensor( + ggml_transpose(ctx.ggml, selected.tensor), + core::TensorShape::from_dims({rows, compact_size}), + GGML_TYPE_F32); + const auto contiguous = core::wrap_tensor( + ggml_cont(ctx.ggml, selected_t.tensor), + selected_t.shape, + selected_t.type); + return core::reshape_tensor(ctx, contiguous, logits.shape.with_last_dim(compact_size)); +} + +ggml_tensor * make_logits_readback_token_ids( + ggml_context * ctx, + const QwenCausalDecodeRuntimeConfig & config) { + if (config.logits_readback_token_ids.empty()) { + return nullptr; + } + return ggml_new_tensor_1d( + ctx, + GGML_TYPE_I32, + static_cast(config.logits_readback_token_ids.size())); +} + +core::TensorValue wrap_logits_readback_token_ids( + ggml_tensor * tensor, + const QwenCausalDecodeRuntimeConfig & config) { + return core::wrap_tensor( + tensor, + core::TensorShape::from_dims({static_cast(config.logits_readback_token_ids.size())}), + GGML_TYPE_I32); +} + +void upload_logits_readback_token_ids( + ggml_tensor * tensor, + const QwenCausalDecodeRuntimeConfig & config) { + ggml_backend_tensor_set( + tensor, + config.logits_readback_token_ids.data(), + 0, + config.logits_readback_token_ids.size() * sizeof(int32_t)); +} + QwenCausalDecoderOutputs build_causal_prefill( core::ModuleBuildContext & ctx, const QwenCausalDecodeRuntimeConfig & config, @@ -167,6 +262,48 @@ QwenCausalDecoderStaticCacheOutputs build_causal_decode( }; } +QwenCausalDecoderBatchedStaticCacheOutputs build_causal_decode_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const QwenCausalDecodeRuntimeConfig & config, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenCausalDecodeRuntimeWeights & weights, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot) { + if (config.output_mode == QwenCausalDecodeOutputMode::Logits) { + auto causal_weights = causal_decoder_weights(weights); + return QwenCausalDecoderModule(config.decoder) + .build_static_cache_tail_batched( + ctx, + graph, + input, + positions, + causal_weights, + cache_steps, + attention_mask, + cache_slot); + } + + auto hidden_out = QwenDecoderHiddenModule(hidden_config_from_runtime(config)) + .build_static_cache_tail_batched( + ctx, + graph, + input, + positions, + hidden_weights_from_runtime(weights), + cache_steps, + attention_mask, + cache_slot); + return { + std::move(hidden_out.sequence), + hidden_out.hidden, + {}, + std::move(hidden_out.cache), + }; +} + } // namespace class QwenCausalDecodeRuntime::Impl { @@ -220,6 +357,45 @@ class QwenCausalDecodeRuntime::Impl { return run_prefill(); } + QwenCausalBatchedPrefillResult prefill_tokens_batched( + const std::vector & token_ids, + int64_t batch_size, + int64_t steps) { + if (batch_size <= 0 || steps <= 0) { + throw std::runtime_error("QwenCausalDecodeRuntime batched prefill requires positive batch and steps"); + } + if (token_ids.size() != static_cast(batch_size * steps)) { + throw std::runtime_error("QwenCausalDecodeRuntime batched prefill token size mismatch"); + } + ensure_batched_prefill_token_graph(batch_size, steps); + ggml_backend_tensor_set( + batched_prefill_input_, + token_ids.data(), + 0, + token_ids.size() * sizeof(int32_t)); + return run_batched_prefill(); + } + + QwenCausalBatchedPrefillResult prefill_embeddings_batched( + const std::vector & embeddings, + int64_t batch_size, + int64_t steps) { + if (batch_size <= 0 || steps <= 0) { + throw std::runtime_error("QwenCausalDecodeRuntime batched prefill requires positive batch and steps"); + } + const size_t expected = static_cast(batch_size * steps * config_.decoder.stack.hidden_size); + if (embeddings.size() != expected) { + throw std::runtime_error("QwenCausalDecodeRuntime batched prefill embedding size mismatch"); + } + ensure_batched_prefill_embedding_graph(batch_size, steps); + ggml_backend_tensor_set( + batched_prefill_input_, + embeddings.data(), + 0, + embeddings.size() * sizeof(float)); + return run_batched_prefill(); + } + void start_decode_tokens(const runtime::TransformerKVState & state, int64_t required_cache_steps) { if (required_cache_steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime decode requires positive cache capacity"); @@ -257,6 +433,55 @@ class QwenCausalDecodeRuntime::Impl { return run_decode_step(); } + void start_decode_tokens_batched( + const runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps) { + if (required_cache_steps <= 0) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode requires positive cache capacity"); + } + ensure_batched_decode_token_graph(required_cache_steps, state.batch_size); + batched_decode_cache_.import_state(state); + } + + void start_decode_embeddings_batched( + const runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps) { + if (required_cache_steps <= 0) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode requires positive cache capacity"); + } + ensure_batched_decode_embedding_graph(required_cache_steps, state.batch_size); + batched_decode_cache_.import_state(state); + } + + QwenCausalDecodeStepResult decode_tokens_batched(const std::vector & tokens) { + ensure_batched_decode_started(); + if (batched_decode_input_kind_ != InputKind::Token) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode graph expects tokens"); + } + if (tokens.size() != static_cast(batched_decode_batch_size_)) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode token size mismatch"); + } + ggml_backend_tensor_set(batched_decode_input_, tokens.data(), 0, tokens.size() * sizeof(int32_t)); + return run_batched_decode_step(); + } + + QwenCausalDecodeStepResult decode_embeddings_batched( + const std::vector & embeddings, + int64_t batch_size) { + ensure_batched_decode_started(); + if (batched_decode_input_kind_ != InputKind::Embedding) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode graph expects embeddings"); + } + if (batch_size != batched_decode_batch_size_) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode embedding batch mismatch"); + } + if (embeddings.size() != static_cast(batch_size * config_.decoder.stack.hidden_size)) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode embedding size mismatch"); + } + ggml_backend_tensor_set(batched_decode_input_, embeddings.data(), 0, embeddings.size() * sizeof(float)); + return run_batched_decode_step(); + } + int64_t decode_cache_steps() const noexcept { return decode_cache_steps_; } @@ -272,6 +497,8 @@ class QwenCausalDecodeRuntime::Impl { void release_runtime_graphs() { release_prefill_graph(); release_decode_graph(); + release_batched_prefill_graph(); + release_batched_decode_graph(); } private: @@ -332,6 +559,7 @@ class QwenCausalDecodeRuntime::Impl { core::TensorShape::from_dims({1, 1, steps, steps}), GGML_TYPE_F16); auto decoder_out = build_causal_prefill(ctx, config_, x, positions, weights_, attention_mask); + prefill_logits_readback_token_ids_ = make_logits_readback_token_ids(prefill_ctx_.get(), config_); for (const auto & layer : decoder_out.state.layers) { if (!layer.key.has_value() || !layer.value.has_value()) { throw std::runtime_error("QwenCausalDecodeRuntime prefill decoder did not return K/V state"); @@ -350,10 +578,18 @@ class QwenCausalDecodeRuntime::Impl { prefill_values_.push_back(value); } if (config_.output_mode == QwenCausalDecodeOutputMode::Logits) { + auto logits = decoder_out.logits; + if (prefill_logits_readback_token_ids_ != nullptr) { + logits = compact_logits_readback( + ctx, + config_, + logits, + wrap_logits_readback_token_ids(prefill_logits_readback_token_ids_, config_)); + } prefill_logits_ = ggml_cpy( prefill_ctx_.get(), - decoder_out.logits.tensor, - ggml_dup_tensor(prefill_ctx_.get(), decoder_out.logits.tensor)); + logits.tensor, + ggml_dup_tensor(prefill_ctx_.get(), logits.tensor)); ggml_set_output(prefill_logits_); } if (config_.return_hidden || config_.output_mode == QwenCausalDecodeOutputMode::Hidden) { @@ -394,6 +630,9 @@ class QwenCausalDecodeRuntime::Impl { mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); + if (prefill_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(prefill_logits_readback_token_ids_, config_); + } prefill_input_kind_ = input_kind; prefill_steps_ = steps; debug::timing_log_scalar( @@ -436,6 +675,201 @@ class QwenCausalDecodeRuntime::Impl { return out; } + void ensure_batched_prefill_token_graph(int64_t batch_size, int64_t steps) { + if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Token && + batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { + debug::timing_log_scalar(config_.trace_name + ".batched_prefill.graph.build_ms", 0.0); + return; + } + release_batched_prefill_graph(); + build_batched_prefill_graph(InputKind::Token, batch_size, steps); + } + + void ensure_batched_prefill_embedding_graph(int64_t batch_size, int64_t steps) { + if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Embedding && + batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { + debug::timing_log_scalar(config_.trace_name + ".batched_prefill.graph.build_ms", 0.0); + return; + } + release_batched_prefill_graph(); + build_batched_prefill_graph(InputKind::Embedding, batch_size, steps); + } + + void build_batched_prefill_graph(InputKind input_kind, int64_t batch_size, int64_t steps) { + const auto build_start = Clock::now(); + ggml_init_params params{config_.prefill_graph_arena_bytes, nullptr, true}; + batched_prefill_ctx_.reset(ggml_init(params)); + if (batched_prefill_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize QwenCausalDecodeRuntime batched prefill graph context"); + } + core::ModuleBuildContext ctx{batched_prefill_ctx_.get(), config_.trace_name.c_str(), backend_type_}; + core::TensorValue x; + if (input_kind == InputKind::Token) { + auto input = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({batch_size, steps})); + batched_prefill_input_ = input.tensor; + x = token_embedding_input_batched( + ctx, + weights_, + config_.decoder, + batched_prefill_input_, + batch_size, + steps); + } else { + auto input = core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({batch_size, steps, config_.decoder.stack.hidden_size})); + batched_prefill_input_ = input.tensor; + x = input; + } + batched_prefill_positions_ = ggml_new_tensor_1d(batched_prefill_ctx_.get(), GGML_TYPE_I32, steps); + auto positions = core::wrap_tensor( + batched_prefill_positions_, + core::TensorShape::from_dims({steps}), + GGML_TYPE_I32); + auto attention = core::make_tensor( + ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({batch_size, 1, steps, steps})); + batched_prefill_attention_mask_ = attention.tensor; + auto decoder_out = build_causal_prefill(ctx, config_, x, positions, weights_, attention); + batched_prefill_logits_readback_token_ids_ = + make_logits_readback_token_ids(batched_prefill_ctx_.get(), config_); + for (const auto & layer : decoder_out.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("QwenCausalDecodeRuntime batched prefill decoder did not return K/V state"); + } + auto * key = ggml_cpy( + batched_prefill_ctx_.get(), + layer.key->tensor, + ggml_dup_tensor(batched_prefill_ctx_.get(), layer.key->tensor)); + auto * value = ggml_cpy( + batched_prefill_ctx_.get(), + layer.value->tensor, + ggml_dup_tensor(batched_prefill_ctx_.get(), layer.value->tensor)); + ggml_set_output(key); + ggml_set_output(value); + batched_prefill_keys_.push_back(key); + batched_prefill_values_.push_back(value); + } + if (config_.output_mode == QwenCausalDecodeOutputMode::Logits) { + auto logits = decoder_out.logits; + if (batched_prefill_logits_readback_token_ids_ != nullptr) { + logits = compact_logits_readback( + ctx, + config_, + logits, + wrap_logits_readback_token_ids(batched_prefill_logits_readback_token_ids_, config_)); + } + batched_prefill_logits_ = ggml_cpy( + batched_prefill_ctx_.get(), + logits.tensor, + ggml_dup_tensor(batched_prefill_ctx_.get(), logits.tensor)); + ggml_set_output(batched_prefill_logits_); + } + if (config_.return_hidden || config_.output_mode == QwenCausalDecodeOutputMode::Hidden) { + batched_prefill_hidden_ = ggml_cpy( + batched_prefill_ctx_.get(), + decoder_out.hidden.tensor, + ggml_dup_tensor(batched_prefill_ctx_.get(), decoder_out.hidden.tensor)); + ggml_set_output(batched_prefill_hidden_); + } + batched_prefill_graph_ = ggml_new_graph_custom(batched_prefill_ctx_.get(), 65536, false); + for (auto * key : batched_prefill_keys_) { + ggml_build_forward_expand(batched_prefill_graph_, key); + } + for (auto * value : batched_prefill_values_) { + ggml_build_forward_expand(batched_prefill_graph_, value); + } + if (batched_prefill_logits_ != nullptr) { + ggml_build_forward_expand(batched_prefill_graph_, batched_prefill_logits_); + } + if (batched_prefill_hidden_ != nullptr) { + ggml_build_forward_expand(batched_prefill_graph_, batched_prefill_hidden_); + } + batched_prefill_gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (batched_prefill_gallocr_ == nullptr || + !ggml_gallocr_reserve(batched_prefill_gallocr_, batched_prefill_graph_) || + !ggml_gallocr_alloc_graph(batched_prefill_gallocr_, batched_prefill_graph_)) { + throw std::runtime_error("failed to allocate QwenCausalDecodeRuntime batched prefill graph"); + } + const auto positions_values = qwen_position_ids(steps); + ggml_backend_tensor_set( + batched_prefill_positions_, + positions_values.data(), + 0, + positions_values.size() * sizeof(int32_t)); + const auto mask = qwen_causal_prefill_mask_values(batch_size, steps); + ggml_backend_tensor_set( + batched_prefill_attention_mask_, + mask.data(), + 0, + mask.size() * sizeof(ggml_fp16_t)); + if (batched_prefill_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(batched_prefill_logits_readback_token_ids_, config_); + } + batched_prefill_input_kind_ = input_kind; + batched_prefill_batch_size_ = batch_size; + batched_prefill_steps_ = steps; + debug::timing_log_scalar( + config_.trace_name + ".batched_prefill.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + QwenCausalBatchedPrefillResult run_batched_prefill() { + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph(backend_, batched_prefill_graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("QwenCausalDecodeRuntime batched prefill graph compute failed"); + } + QwenCausalBatchedPrefillResult out; + if (batched_prefill_logits_ != nullptr) { + out.logits.resize(static_cast(ggml_nelements(batched_prefill_logits_))); + ggml_backend_tensor_get( + batched_prefill_logits_, + out.logits.data(), + 0, + out.logits.size() * sizeof(float)); + } + if (batched_prefill_hidden_ != nullptr) { + out.hidden.resize(static_cast(ggml_nelements(batched_prefill_hidden_))); + ggml_backend_tensor_get( + batched_prefill_hidden_, + out.hidden.data(), + 0, + out.hidden.size() * sizeof(float)); + round_readback(out.hidden, config_); + } + out.state.batch_size = batched_prefill_batch_size_; + out.state.current_end = batched_prefill_steps_; + out.state.layers.resize(batched_prefill_keys_.size()); + const size_t layer_values = static_cast( + batched_prefill_batch_size_ * + batched_prefill_steps_ * + config_.decoder.stack.num_key_value_heads * + config_.decoder.stack.head_dim); + for (size_t layer = 0; layer < batched_prefill_keys_.size(); ++layer) { + auto & state = out.state.layers[layer]; + state.valid_steps = batched_prefill_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get( + batched_prefill_keys_[layer], + state.key.data(), + 0, + state.key.size() * sizeof(float)); + ggml_backend_tensor_get( + batched_prefill_values_[layer], + state.value.data(), + 0, + state.value.size() * sizeof(float)); + round_readback(state.key, config_); + round_readback(state.value, config_); + } + return out; + } + void ensure_decode_token_graph(int64_t cache_steps) { if (decode_graph_ != nullptr && decode_input_kind_ == InputKind::Token && decode_cache_steps_ >= cache_steps) { debug::timing_log_scalar(config_.trace_name + ".decode.graph.build_ms", 0.0); @@ -497,11 +931,20 @@ class QwenCausalDecodeRuntime::Impl { attention_mask, cache_slot); decode_cache_ = std::move(decoder_out.cache); + decode_logits_readback_token_ids_ = make_logits_readback_token_ids(decode_ctx_.get(), config_); if (config_.output_mode == QwenCausalDecodeOutputMode::Logits) { + auto logits = decoder_out.logits; + if (decode_logits_readback_token_ids_ != nullptr) { + logits = compact_logits_readback( + ctx, + config_, + logits, + wrap_logits_readback_token_ids(decode_logits_readback_token_ids_, config_)); + } decode_logits_ = ggml_cpy( decode_ctx_.get(), - decoder_out.logits.tensor, - ggml_dup_tensor(decode_ctx_.get(), decoder_out.logits.tensor)); + logits.tensor, + ggml_dup_tensor(decode_ctx_.get(), logits.tensor)); ggml_set_output(decode_logits_); } if (config_.return_hidden || config_.output_mode == QwenCausalDecodeOutputMode::Hidden) { @@ -521,6 +964,9 @@ class QwenCausalDecodeRuntime::Impl { if (decode_buffer_ == nullptr) { throw std::runtime_error("failed to allocate QwenCausalDecodeRuntime decode graph"); } + if (decode_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(decode_logits_readback_token_ids_, config_); + } decode_attention_mask_values_.assign( static_cast(cache_steps), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); @@ -532,12 +978,138 @@ class QwenCausalDecodeRuntime::Impl { debug::trace_log_scalar(config_.trace_name + ".decode.cache_steps", cache_steps); } + void ensure_batched_decode_token_graph(int64_t cache_steps, int64_t batch_size) { + if (batched_decode_graph_ != nullptr && batched_decode_input_kind_ == InputKind::Token && + batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size) { + debug::timing_log_scalar(config_.trace_name + ".batched_decode.graph.build_ms", 0.0); + return; + } + release_batched_decode_graph(); + build_batched_decode_graph(InputKind::Token, batch_size, cache_steps); + } + + void ensure_batched_decode_embedding_graph(int64_t cache_steps, int64_t batch_size) { + if (batched_decode_graph_ != nullptr && batched_decode_input_kind_ == InputKind::Embedding && + batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size) { + debug::timing_log_scalar(config_.trace_name + ".batched_decode.graph.build_ms", 0.0); + return; + } + release_batched_decode_graph(); + build_batched_decode_graph(InputKind::Embedding, batch_size, cache_steps); + } + + void build_batched_decode_graph(InputKind input_kind, int64_t batch_size, int64_t cache_steps) { + if (batch_size <= 0) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode requires positive batch size"); + } + if (config_.decoder.stack.runtime.static_cache.update_mode != QwenDecoderStaticCacheUpdateMode::DirectSetRows) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode supports only DirectSetRows cache update"); + } + const auto build_start = Clock::now(); + ggml_init_params params{config_.decode_graph_arena_bytes, nullptr, true}; + batched_decode_ctx_.reset(ggml_init(params)); + if (batched_decode_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize QwenCausalDecodeRuntime batched decode graph context"); + } + core::ModuleBuildContext ctx{batched_decode_ctx_.get(), config_.trace_name.c_str(), backend_type_}; + core::TensorValue x; + if (input_kind == InputKind::Token) { + auto input = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({batch_size, 1})); + batched_decode_input_ = input.tensor; + x = token_embedding_input_batched(ctx, weights_, config_.decoder, batched_decode_input_, batch_size, 1); + } else { + auto input = core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({batch_size, 1, config_.decoder.stack.hidden_size})); + batched_decode_input_ = input.tensor; + x = input; + } + batched_decode_positions_ = ggml_new_tensor_1d(batched_decode_ctx_.get(), GGML_TYPE_I32, 1); + auto positions = core::wrap_tensor( + batched_decode_positions_, + core::TensorShape::from_dims({1}), + GGML_TYPE_I32); + auto slot = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({batch_size})); + batched_decode_cache_slot_ = slot.tensor; + auto attention = core::make_tensor( + ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({batch_size, 1, 1, cache_steps})); + batched_decode_attention_mask_ = attention.tensor; + batched_decode_graph_ = ggml_new_graph_custom(batched_decode_ctx_.get(), 65536, false); + auto decoder_out = build_causal_decode_batched( + ctx, + batched_decode_graph_, + config_, + x, + positions, + weights_, + cache_steps, + attention, + slot); + batched_decode_cache_ = std::move(decoder_out.cache); + batched_decode_logits_readback_token_ids_ = + make_logits_readback_token_ids(batched_decode_ctx_.get(), config_); + if (config_.output_mode == QwenCausalDecodeOutputMode::Logits) { + auto logits = decoder_out.logits; + if (batched_decode_logits_readback_token_ids_ != nullptr) { + logits = compact_logits_readback( + ctx, + config_, + logits, + wrap_logits_readback_token_ids(batched_decode_logits_readback_token_ids_, config_)); + } + batched_decode_logits_ = ggml_cpy( + batched_decode_ctx_.get(), + logits.tensor, + ggml_dup_tensor(batched_decode_ctx_.get(), logits.tensor)); + ggml_set_output(batched_decode_logits_); + } + if (config_.return_hidden || config_.output_mode == QwenCausalDecodeOutputMode::Hidden) { + batched_decode_hidden_ = ggml_cpy( + batched_decode_ctx_.get(), + decoder_out.hidden.tensor, + ggml_dup_tensor(batched_decode_ctx_.get(), decoder_out.hidden.tensor)); + ggml_set_output(batched_decode_hidden_); + } + if (batched_decode_logits_ != nullptr) { + ggml_build_forward_expand(batched_decode_graph_, batched_decode_logits_); + } + if (batched_decode_hidden_ != nullptr) { + ggml_build_forward_expand(batched_decode_graph_, batched_decode_hidden_); + } + batched_decode_buffer_ = ggml_backend_alloc_ctx_tensors(batched_decode_ctx_.get(), backend_); + if (batched_decode_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate QwenCausalDecodeRuntime batched decode graph"); + } + if (batched_decode_logits_readback_token_ids_ != nullptr) { + upload_logits_readback_token_ids(batched_decode_logits_readback_token_ids_, config_); + } + batched_decode_attention_mask_values_.assign( + static_cast(batch_size * cache_steps), + ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + batched_decode_cache_slots_.resize(static_cast(batch_size)); + batched_decode_batch_size_ = batch_size; + batched_decode_cache_steps_ = cache_steps; + batched_decode_input_kind_ = input_kind; + debug::timing_log_scalar( + config_.trace_name + ".batched_decode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + void ensure_decode_started() const { if (decode_graph_ == nullptr) { throw std::runtime_error("QwenCausalDecodeRuntime decode graph has not been started"); } } + void ensure_batched_decode_started() const { + if (batched_decode_graph_ == nullptr) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode graph has not been started"); + } + } + QwenCausalDecodeStepResult run_decode_step() { if (decode_cache_.valid_steps() >= decode_cache_steps_) { throw std::runtime_error("QwenCausalDecodeRuntime decode cache exhausted"); @@ -572,6 +1144,57 @@ class QwenCausalDecodeRuntime::Impl { return out; } + QwenCausalDecodeStepResult run_batched_decode_step() { + if (batched_decode_cache_.valid_steps() >= batched_decode_cache_steps_) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode cache exhausted"); + } + const int32_t position = static_cast(batched_decode_cache_.current_end()); + ggml_backend_tensor_set(batched_decode_positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(batched_decode_cache_.valid_steps()); + for (int64_t batch = 0; batch < batched_decode_batch_size_; ++batch) { + batched_decode_cache_slots_[static_cast(batch)] = + static_cast(batch * batched_decode_cache_steps_ + cache_slot); + } + ggml_backend_tensor_set( + batched_decode_cache_slot_, + batched_decode_cache_slots_.data(), + 0, + batched_decode_cache_slots_.size() * sizeof(int32_t)); + write_qwen_batched_cached_step_mask( + batched_decode_attention_mask_, + batched_decode_attention_mask_values_, + batched_decode_batch_size_, + batched_decode_cache_steps_, + batched_decode_cache_.valid_steps(), + cache_slot); + core::set_backend_threads(backend_, threads_); + const ggml_status status = core::compute_backend_graph(backend_, batched_decode_graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("QwenCausalDecodeRuntime batched decode graph compute failed"); + } + QwenCausalDecodeStepResult out; + if (batched_decode_logits_ != nullptr) { + out.logits.resize(static_cast(ggml_nelements(batched_decode_logits_))); + ggml_backend_tensor_get( + batched_decode_logits_, + out.logits.data(), + 0, + out.logits.size() * sizeof(float)); + } + if (batched_decode_hidden_ != nullptr) { + out.hidden.resize(static_cast(ggml_nelements(batched_decode_hidden_))); + ggml_backend_tensor_get( + batched_decode_hidden_, + out.hidden.data(), + 0, + out.hidden.size() * sizeof(float)); + round_readback(out.hidden, config_); + } + batched_decode_cache_.advance_after_direct_append(1); + return out; + } + void release_prefill_graph() { if (prefill_graph_ != nullptr) { core::release_backend_graph_resources(backend_, prefill_graph_); @@ -584,6 +1207,7 @@ class QwenCausalDecodeRuntime::Impl { prefill_input_ = nullptr; prefill_positions_ = nullptr; prefill_attention_mask_ = nullptr; + prefill_logits_readback_token_ids_ = nullptr; prefill_logits_ = nullptr; prefill_hidden_ = nullptr; prefill_keys_.clear(); @@ -593,6 +1217,29 @@ class QwenCausalDecodeRuntime::Impl { prefill_input_kind_ = InputKind::None; } + void release_batched_prefill_graph() { + if (batched_prefill_graph_ != nullptr) { + core::release_backend_graph_resources(backend_, batched_prefill_graph_); + } + if (batched_prefill_gallocr_ != nullptr) { + ggml_gallocr_free(batched_prefill_gallocr_); + batched_prefill_gallocr_ = nullptr; + } + batched_prefill_ctx_.reset(); + batched_prefill_input_ = nullptr; + batched_prefill_positions_ = nullptr; + batched_prefill_attention_mask_ = nullptr; + batched_prefill_logits_readback_token_ids_ = nullptr; + batched_prefill_logits_ = nullptr; + batched_prefill_hidden_ = nullptr; + batched_prefill_keys_.clear(); + batched_prefill_values_.clear(); + batched_prefill_graph_ = nullptr; + batched_prefill_batch_size_ = 0; + batched_prefill_steps_ = 0; + batched_prefill_input_kind_ = InputKind::None; + } + void release_decode_graph() { if (decode_graph_ != nullptr) { core::release_backend_graph_resources(backend_, decode_graph_); @@ -606,6 +1253,7 @@ class QwenCausalDecodeRuntime::Impl { decode_positions_ = nullptr; decode_cache_slot_ = nullptr; decode_attention_mask_ = nullptr; + decode_logits_readback_token_ids_ = nullptr; decode_logits_ = nullptr; decode_hidden_ = nullptr; decode_graph_ = nullptr; @@ -615,6 +1263,31 @@ class QwenCausalDecodeRuntime::Impl { decode_input_kind_ = InputKind::None; } + void release_batched_decode_graph() { + if (batched_decode_graph_ != nullptr) { + core::release_backend_graph_resources(backend_, batched_decode_graph_); + } + if (batched_decode_buffer_ != nullptr) { + ggml_backend_buffer_free(batched_decode_buffer_); + batched_decode_buffer_ = nullptr; + } + batched_decode_ctx_.reset(); + batched_decode_input_ = nullptr; + batched_decode_positions_ = nullptr; + batched_decode_cache_slot_ = nullptr; + batched_decode_attention_mask_ = nullptr; + batched_decode_logits_readback_token_ids_ = nullptr; + batched_decode_logits_ = nullptr; + batched_decode_hidden_ = nullptr; + batched_decode_graph_ = nullptr; + batched_decode_cache_ = runtime::TransformerBatchedKVCache(); + batched_decode_attention_mask_values_.clear(); + batched_decode_cache_slots_.clear(); + batched_decode_batch_size_ = 0; + batched_decode_cache_steps_ = 0; + batched_decode_input_kind_ = InputKind::None; + } + ggml_backend_t backend_ = nullptr; core::BackendType backend_type_ = core::BackendType::Cpu; int threads_ = 1; @@ -625,6 +1298,7 @@ class QwenCausalDecodeRuntime::Impl { ggml_tensor * prefill_input_ = nullptr; ggml_tensor * prefill_positions_ = nullptr; ggml_tensor * prefill_attention_mask_ = nullptr; + ggml_tensor * prefill_logits_readback_token_ids_ = nullptr; ggml_tensor * prefill_logits_ = nullptr; ggml_tensor * prefill_hidden_ = nullptr; std::vector prefill_keys_; @@ -634,11 +1308,27 @@ class QwenCausalDecodeRuntime::Impl { int64_t prefill_steps_ = 0; InputKind prefill_input_kind_ = InputKind::None; + std::unique_ptr batched_prefill_ctx_; + ggml_tensor * batched_prefill_input_ = nullptr; + ggml_tensor * batched_prefill_positions_ = nullptr; + ggml_tensor * batched_prefill_attention_mask_ = nullptr; + ggml_tensor * batched_prefill_logits_readback_token_ids_ = nullptr; + ggml_tensor * batched_prefill_logits_ = nullptr; + ggml_tensor * batched_prefill_hidden_ = nullptr; + std::vector batched_prefill_keys_; + std::vector batched_prefill_values_; + ggml_cgraph * batched_prefill_graph_ = nullptr; + ggml_gallocr_t batched_prefill_gallocr_ = nullptr; + int64_t batched_prefill_batch_size_ = 0; + int64_t batched_prefill_steps_ = 0; + InputKind batched_prefill_input_kind_ = InputKind::None; + std::unique_ptr decode_ctx_; ggml_tensor * decode_input_ = nullptr; ggml_tensor * decode_positions_ = nullptr; ggml_tensor * decode_cache_slot_ = nullptr; ggml_tensor * decode_attention_mask_ = nullptr; + ggml_tensor * decode_logits_readback_token_ids_ = nullptr; ggml_tensor * decode_logits_ = nullptr; ggml_tensor * decode_hidden_ = nullptr; ggml_cgraph * decode_graph_ = nullptr; @@ -647,6 +1337,23 @@ class QwenCausalDecodeRuntime::Impl { runtime::TransformerKVCache decode_cache_; int64_t decode_cache_steps_ = 0; InputKind decode_input_kind_ = InputKind::None; + + std::unique_ptr batched_decode_ctx_; + ggml_tensor * batched_decode_input_ = nullptr; + ggml_tensor * batched_decode_positions_ = nullptr; + ggml_tensor * batched_decode_cache_slot_ = nullptr; + ggml_tensor * batched_decode_attention_mask_ = nullptr; + ggml_tensor * batched_decode_logits_readback_token_ids_ = nullptr; + ggml_tensor * batched_decode_logits_ = nullptr; + ggml_tensor * batched_decode_hidden_ = nullptr; + ggml_cgraph * batched_decode_graph_ = nullptr; + ggml_backend_buffer_t batched_decode_buffer_ = nullptr; + std::vector batched_decode_attention_mask_values_; + std::vector batched_decode_cache_slots_; + runtime::TransformerBatchedKVCache batched_decode_cache_; + int64_t batched_decode_batch_size_ = 0; + int64_t batched_decode_cache_steps_ = 0; + InputKind batched_decode_input_kind_ = InputKind::None; }; QwenCausalDecodeRuntime::QwenCausalDecodeRuntime( @@ -667,6 +1374,20 @@ QwenCausalPrefillResult QwenCausalDecodeRuntime::prefill_embeddings( return impl_->prefill_embeddings(embeddings, steps); } +QwenCausalBatchedPrefillResult QwenCausalDecodeRuntime::prefill_tokens_batched( + const std::vector & token_ids, + int64_t batch_size, + int64_t steps) { + return impl_->prefill_tokens_batched(token_ids, batch_size, steps); +} + +QwenCausalBatchedPrefillResult QwenCausalDecodeRuntime::prefill_embeddings_batched( + const std::vector & embeddings, + int64_t batch_size, + int64_t steps) { + return impl_->prefill_embeddings_batched(embeddings, batch_size, steps); +} + void QwenCausalDecodeRuntime::start_decode_tokens( const runtime::TransformerKVState & state, int64_t required_cache_steps) { @@ -687,6 +1408,28 @@ QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_embedding(const std:: return impl_->decode_embedding(embedding); } +void QwenCausalDecodeRuntime::start_decode_tokens_batched( + const runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps) { + impl_->start_decode_tokens_batched(state, required_cache_steps); +} + +void QwenCausalDecodeRuntime::start_decode_embeddings_batched( + const runtime::TransformerBatchedKVState & state, + int64_t required_cache_steps) { + impl_->start_decode_embeddings_batched(state, required_cache_steps); +} + +QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_tokens_batched(const std::vector & tokens) { + return impl_->decode_tokens_batched(tokens); +} + +QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_embeddings_batched( + const std::vector & embeddings, + int64_t batch_size) { + return impl_->decode_embeddings_batched(embeddings, batch_size); +} + int64_t QwenCausalDecodeRuntime::decode_cache_steps() const noexcept { return impl_->decode_cache_steps(); } diff --git a/src/framework/modules/transformers/qwen_causal_decoder.cpp b/src/framework/modules/transformers/qwen_causal_decoder.cpp index 1a819cbe..94ca0050 100644 --- a/src/framework/modules/transformers/qwen_causal_decoder.cpp +++ b/src/framework/modules/transformers/qwen_causal_decoder.cpp @@ -21,6 +21,10 @@ void validate_config(const QwenCausalDecoderConfig & config) { if (config.stack.layers <= 0) { throw std::runtime_error("QwenCausalDecoderConfig requires a positive layer count"); } + if (config.static_cache_type != GGML_TYPE_F32 && config.static_cache_type != GGML_TYPE_F16 && + config.static_cache_type != GGML_TYPE_BF16) { + throw std::runtime_error("QwenCausalDecoderConfig static cache type must be f32, f16, or bf16"); + } } void validate_hidden_config(const QwenDecoderHiddenConfig & config) { @@ -30,6 +34,10 @@ void validate_hidden_config(const QwenDecoderHiddenConfig & config) { if (config.stack.layers <= 0) { throw std::runtime_error("QwenDecoderHiddenConfig requires a positive layer count"); } + if (config.static_cache_type != GGML_TYPE_F32 && config.static_cache_type != GGML_TYPE_F16 && + config.static_cache_type != GGML_TYPE_BF16) { + throw std::runtime_error("QwenDecoderHiddenConfig static cache type must be f32, f16, or bf16"); + } } void validate_steps(int64_t steps, const char * label) { @@ -38,6 +46,13 @@ void validate_steps(int64_t steps, const char * label) { } } +runtime::TransformerKVCacheOptions transformer_cache_options(ggml_type type) { + runtime::TransformerKVCacheOptions out; + out.allow_f16_storage = type == GGML_TYPE_F16; + out.allow_bf16_storage = type == GGML_TYPE_BF16; + return out; +} + core::TensorValue select_hidden_steps( core::ModuleBuildContext & ctx, const core::TensorValue & hidden_sequence, @@ -53,6 +68,7 @@ QwenDecoderHiddenConfig hidden_config_from_causal(const QwenCausalDecoderConfig QwenDecoderHiddenConfig out; out.stack = config.stack; out.hidden_mode = config.logits_mode; + out.static_cache_type = config.static_cache_type; return out; } @@ -127,11 +143,11 @@ QwenDecoderHiddenStaticCacheOutputs QwenDecoderHiddenModule::build_static_cache_ for (const auto & layer : weights.stack.layers) { cache_keys.push_back(core::make_tensor( ctx, - GGML_TYPE_F32, + config_.static_cache_type, core::TensorShape::from_dims({1, cache_steps, config_.stack.num_key_value_heads, config_.stack.head_dim}))); cache_values.push_back(core::make_tensor( ctx, - GGML_TYPE_F32, + config_.static_cache_type, core::TensorShape::from_dims({1, cache_steps, config_.stack.num_key_value_heads, config_.stack.head_dim}))); auto out = layer_module.build_with_static_cache_tail( ctx, @@ -151,7 +167,81 @@ QwenDecoderHiddenStaticCacheOutputs QwenDecoderHiddenModule::build_static_cache_ return { std::move(x), hidden, - runtime::TransformerKVCache(cache_steps, step_elems, std::move(cache_keys), std::move(cache_values)), + runtime::TransformerKVCache( + cache_steps, + step_elems, + std::move(cache_keys), + std::move(cache_values), + transformer_cache_options(config_.static_cache_type)), + }; +} + +QwenDecoderHiddenBatchedStaticCacheOutputs QwenDecoderHiddenModule::build_static_cache_tail_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderHiddenWeights & weights, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot) const { + if (graph == nullptr) { + throw std::runtime_error("QwenDecoderHiddenModule batched static-cache build requires a graph"); + } + validate_steps(cache_steps, "QwenDecoderHiddenModule batched static-cache build"); + if (input.shape.rank != 3 || input.shape.dims[0] <= 0 || input.shape.dims[1] != 1 || + input.shape.dims[2] != config_.stack.hidden_size) { + throw std::runtime_error("QwenDecoderHiddenModule batched static-cache input shape must be [batch, 1, hidden]"); + } + if (static_cast(weights.stack.layers.size()) != config_.stack.layers) { + throw std::runtime_error("QwenDecoderHiddenWeights layer count does not match config"); + } + + const int64_t batch_size = input.shape.dims[0]; + const int64_t row_elems = config_.stack.num_key_value_heads * config_.stack.head_dim; + std::vector cache_keys; + std::vector cache_values; + cache_keys.reserve(weights.stack.layers.size()); + cache_values.reserve(weights.stack.layers.size()); + + auto x = input; + const QwenDecoderLayerModule layer_module(qwen_decoder_layer_config_from_stack(config_.stack)); + for (const auto & layer : weights.stack.layers) { + cache_keys.push_back(core::make_tensor( + ctx, + config_.static_cache_type, + core::TensorShape::from_dims( + {batch_size, cache_steps, config_.stack.num_key_value_heads, config_.stack.head_dim}))); + cache_values.push_back(core::make_tensor( + ctx, + config_.static_cache_type, + core::TensorShape::from_dims( + {batch_size, cache_steps, config_.stack.num_key_value_heads, config_.stack.head_dim}))); + auto out = layer_module.build_with_static_cache_tail_batched( + ctx, + graph, + x, + positions, + layer, + cache_keys.back(), + cache_values.back(), + cache_slot, + attention_mask); + x = out.output; + } + + auto hidden = RMSNormModule({config_.stack.hidden_size, config_.stack.rms_norm_eps, true, false}) + .build(ctx, x, weights.final_norm); + return { + std::move(x), + hidden, + runtime::TransformerBatchedKVCache( + cache_steps, + batch_size, + row_elems, + std::move(cache_keys), + std::move(cache_values), + transformer_cache_options(config_.static_cache_type)), }; } @@ -253,6 +343,56 @@ QwenCausalDecoderStaticCacheOutputs QwenCausalDecoderModule::build_static_cache_ }; } +QwenCausalDecoderBatchedStaticCacheOutputs QwenCausalDecoderModule::build_static_cache_tail_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenCausalDecoderWeights & weights, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot) const { + if (graph == nullptr) { + throw std::runtime_error("QwenCausalDecoderModule batched static-cache build requires a graph"); + } + validate_steps(cache_steps, "QwenCausalDecoderModule batched static-cache build"); + if (input.shape.rank != 3 || input.shape.dims[0] <= 0 || input.shape.dims[1] != 1 || + input.shape.dims[2] != config_.stack.hidden_size) { + throw std::runtime_error("QwenCausalDecoderModule batched static-cache input shape must be [batch, 1, hidden]"); + } + + auto hidden_out = QwenDecoderHiddenModule(hidden_config_from_causal(config_)) + .build_static_cache_tail_batched( + ctx, + graph, + input, + positions, + hidden_weights_from_causal(weights), + cache_steps, + attention_mask, + cache_slot); + auto logits_input = hidden_out.hidden; + if (config_.lm_head_input_type.has_value() && logits_input.type != *config_.lm_head_input_type) { + logits_input = core::wrap_tensor( + ggml_cast(ctx.ggml, logits_input.tensor, *config_.lm_head_input_type), + logits_input.shape, + *config_.lm_head_input_type); + } + const auto logits = LinearModule({ + config_.stack.hidden_size, + config_.logits_size, + config_.use_lm_head_bias, + config_.lm_head_precision, + }) + .build(ctx, logits_input, weights.lm_head); + return { + std::move(hidden_out.sequence), + hidden_out.hidden, + logits, + std::move(hidden_out.cache), + }; +} + std::vector qwen_position_ids(int64_t steps, int64_t offset) { validate_steps(steps, "qwen_position_ids"); std::vector out(static_cast(steps), 0); @@ -357,4 +497,45 @@ void write_qwen_cached_step_mask( ggml_backend_tensor_set(tensor, scratch.data(), 0, scratch.size() * sizeof(ggml_fp16_t)); } +void write_qwen_batched_cached_step_mask( + ggml_tensor * tensor, + std::vector & scratch, + int64_t batch_size, + int64_t mask_steps, + int64_t visible_prefix_steps, + int64_t current_slot) { + if (tensor == nullptr) { + throw std::runtime_error("write_qwen_batched_cached_step_mask requires a tensor"); + } + if (batch_size <= 0) { + throw std::runtime_error("write_qwen_batched_cached_step_mask requires positive batch size"); + } + validate_steps(mask_steps, "write_qwen_batched_cached_step_mask"); + if (visible_prefix_steps < 0 || visible_prefix_steps > mask_steps) { + throw std::runtime_error("write_qwen_batched_cached_step_mask visible prefix is out of range"); + } + if (current_slot < 0 || current_slot >= mask_steps) { + throw std::runtime_error("write_qwen_batched_cached_step_mask current slot is out of range"); + } + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + const size_t row_size = static_cast(mask_steps); + const size_t total_size = static_cast(batch_size) * row_size; + if (scratch.size() != total_size) { + scratch.resize(total_size); + } + for (int64_t batch = 0; batch < batch_size; ++batch) { + const size_t offset = static_cast(batch) * row_size; + std::fill( + scratch.begin() + static_cast(offset), + scratch.begin() + static_cast(offset + row_size), + masked); + for (int64_t i = 0; i < visible_prefix_steps; ++i) { + scratch[offset + static_cast(i)] = visible; + } + scratch[offset + static_cast(current_slot)] = visible; + } + ggml_backend_tensor_set(tensor, scratch.data(), 0, scratch.size() * sizeof(ggml_fp16_t)); +} + } // namespace engine::modules diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index 98a9525d..a50b7191 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -812,6 +812,166 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( return {output, stored_key, stored_value}; } +QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_batched( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask) const { + (void)graph; + validate_sequence_input(input, config_.hidden_size, "input"); + if (input.shape.dims[0] <= 0 || input.shape.dims[1] != 1) { + throw std::runtime_error("Qwen decoder batched static-cache update requires [batch, 1, hidden] input"); + } + if (config_.runtime.static_cache.update_mode != QwenDecoderStaticCacheUpdateMode::DirectSetRows) { + throw std::runtime_error("Qwen decoder batched static-cache update supports only DirectSetRows"); + } + const int64_t dim = require_head_dim(config_); + const int64_t kv_repeats = config_.num_attention_heads / config_.num_key_value_heads; + + auto x_norm = RMSNormModule({config_.hidden_size, config_.rms_norm_eps, true, false}) + .build(ctx, input, weights.input_norm); + if (config_.activation_cast.enabled && config_.activation_cast.after_input_norm) { + x_norm = activation_cast(ctx, x_norm, config_.activation_cast); + } + auto qkv = build_qkv_projections(ctx, x_norm, weights, config_, dim); + if (config_.activation_cast.enabled && config_.activation_cast.after_qkv_projection) { + qkv.q = activation_cast(ctx, qkv.q, config_.activation_cast); + qkv.k = activation_cast(ctx, qkv.k, config_.activation_cast); + qkv.v = activation_cast(ctx, qkv.v, config_.activation_cast); + } + + auto q = reshape_qwen_heads(ctx, qkv.q, config_.num_attention_heads, dim); + auto k = reshape_qwen_heads(ctx, qkv.k, config_.num_key_value_heads, dim); + if (config_.use_qk_norm) { + q = RMSNormModule({dim, config_.rms_norm_eps, true, false}).build(ctx, q, weights.q_norm); + k = RMSNormModule({dim, config_.rms_norm_eps, true, false}).build(ctx, k, weights.k_norm); + if (config_.activation_cast.enabled && config_.activation_cast.after_qk_norm) { + q = activation_cast(ctx, q, config_.activation_cast); + k = activation_cast(ctx, k, config_.activation_cast); + } + } + auto v = reshape_qwen_heads(ctx, qkv.v, config_.num_key_value_heads, dim); + + if (config_.position_encoding == QwenDecoderPositionEncoding::Rotary) { + const core::TensorValue * rope_factors = weights.rope_frequency_factors.has_value() + ? &*weights.rope_frequency_factors + : nullptr; + q = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, q, positions, rope_factors); + k = RoPEModule({dim, config_.rope_type, config_.rope_theta}).build(ctx, k, positions, rope_factors); + if (config_.activation_cast.enabled && config_.activation_cast.after_rope) { + q = activation_cast(ctx, q, config_.activation_cast); + k = activation_cast(ctx, k, config_.activation_cast); + } + } + k = core::ensure_backend_addressable_layout(ctx, k); + v = core::ensure_backend_addressable_layout(ctx, v); + + const FastKVSetRowsModule set_rows({ + config_.runtime.static_cache.set_rows_mode == QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized + ? FastKVSetRowsMode::BackendViewOptimized + : FastKVSetRowsMode::Exact, + }); + auto attention_key_cache = set_rows.build(ctx, cache_key, k, cache_slot); + auto attention_value_cache = set_rows.build(ctx, cache_value, v, cache_slot); + if (config_.activation_cast.enabled && config_.activation_cast.after_static_cache_update) { + attention_key_cache = activation_cast(ctx, attention_key_cache, config_.activation_cast); + attention_value_cache = activation_cast(ctx, attention_value_cache, config_.activation_cast); + } + + auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); + auto k_heads = TransposeModule({{0, 2, 1, 3}, attention_key_cache.shape.rank}).build(ctx, attention_key_cache); + auto v_heads = TransposeModule({{0, 2, 1, 3}, attention_value_cache.shape.rank}).build(ctx, attention_value_cache); + core::TensorValue context; + const bool use_grouped_query = + config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery && + config_.runtime.attention.grouped_query_min_steps > 0 && + cache_key.shape.dims[1] >= config_.runtime.attention.grouped_query_min_steps && + kv_repeats > 1; + if (use_grouped_query) { + k_heads = core::ensure_backend_addressable_layout(ctx, k_heads); + v_heads = core::ensure_backend_addressable_layout(ctx, v_heads); + context = attention_from_grouped_query_heads( + ctx, + q_heads, + k_heads, + v_heads, + dim, + config_.num_attention_heads, + config_.num_key_value_heads, + attention_mask); + } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || + config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { + k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); + v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); + context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); + } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { + context = flash_attention_from_grouped_heads_view_kv( + ctx, + q_heads, + k_heads, + v_heads, + dim, + attention_mask, + config_.attention_precision); + } else { + k_heads = core::wrap_tensor(ggml_cont(ctx.ggml, k_heads.tensor), k_heads.shape, k_heads.type); + v_heads = core::wrap_tensor(ggml_cont(ctx.ggml, v_heads.tensor), v_heads.shape, v_heads.type); + context = flash_attention_from_grouped_heads( + ctx, + q_heads, + k_heads, + v_heads, + dim, + attention_mask, + config_.attention_precision); + } + if (config_.activation_cast.enabled && config_.activation_cast.after_attention) { + context = activation_cast(ctx, context, config_.activation_cast); + } + context = core::ensure_backend_addressable_layout(ctx, context); + context = core::reshape_tensor( + ctx, + context, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config_.num_attention_heads * dim})); + + auto attn_out = LinearModule( + { + config_.num_attention_heads * dim, + config_.hidden_size, + weights.self_attention.out_bias.has_value(), + config_.projection_precision, + }) + .build( + ctx, + context, + {weights.self_attention.out_weight, weights.self_attention.out_bias}); + if (config_.activation_cast.enabled && config_.activation_cast.after_attention_output) { + attn_out = activation_cast(ctx, attn_out, config_.activation_cast); + } + auto x = AddModule{}.build(ctx, input, attn_out); + if (config_.activation_cast.enabled && config_.activation_cast.after_residual) { + x = activation_cast(ctx, x, config_.activation_cast); + } + + auto ff_in = RMSNormModule({config_.hidden_size, config_.rms_norm_eps, true, false}) + .build(ctx, x, weights.post_norm); + if (config_.activation_cast.enabled && config_.activation_cast.after_ffn_norm) { + ff_in = activation_cast(ctx, ff_in, config_.activation_cast); + } + auto ff = build_mlp(ctx, ff_in, config_, weights.mlp); + auto output = AddModule{}.build(ctx, x, ff); + if (config_.activation_cast.enabled && config_.activation_cast.after_output) { + output = activation_cast(ctx, output, config_.activation_cast); + } + return {output, k, v}; +} + const core::ModuleSchema & QwenDecoderLayerModule::static_schema() noexcept { return kQwenDecoderLayerSchema; } diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 555cfd44..567aac4d 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -216,6 +216,152 @@ void TransformerKVCache::trace_log_state(const std::string & name, int64_t num_h } } +TransformerBatchedKVCache::TransformerBatchedKVCache( + int64_t cache_steps, + int64_t batch_size, + int64_t row_elems, + std::vector keys, + std::vector values) + : TransformerBatchedKVCache(cache_steps, batch_size, row_elems, std::move(keys), std::move(values), {}) {} + +TransformerBatchedKVCache::TransformerBatchedKVCache( + int64_t cache_steps, + int64_t batch_size, + int64_t row_elems, + std::vector keys, + std::vector values, + TransformerKVCacheOptions options) + : cache_steps_(std::max(0, cache_steps)), + batch_size_(std::max(0, batch_size)), + row_elems_(std::max(0, row_elems)), + options_(options) { + if (cache_steps_ <= 0 || batch_size_ <= 0 || row_elems_ <= 0) { + throw std::runtime_error("TransformerBatchedKVCache requires positive cache_steps, batch_size, and row_elems"); + } + if (keys.size() != values.size()) { + throw std::runtime_error("TransformerBatchedKVCache key/value layer counts must match"); + } + const size_t cache_elems = static_cast(batch_size_ * cache_steps_ * row_elems_); + layers_.reserve(keys.size()); + for (size_t layer = 0; layer < keys.size(); ++layer) { + validate_cache_tensor(keys[layer], options_); + validate_cache_tensor(values[layer], options_); + layers_.push_back(LayerCache{ + std::move(keys[layer]), + std::move(values[layer]), + std::vector(cache_elems, 0.0F), + std::vector(cache_elems, 0.0F), + }); + } +} + +void TransformerBatchedKVCache::import_state(const TransformerBatchedKVState & state) { + if (state.batch_size != batch_size_) { + throw std::runtime_error("TransformerBatchedKVCache state batch size does not match cache batch size"); + } + current_end_ = state.current_end; + if (layers_.empty()) { + valid_steps_ = 0; + return; + } + if (state.layers.size() != layers_.size()) { + throw std::runtime_error("TransformerBatchedKVCache state layer count does not match cache layer count"); + } + const int64_t state_steps = state.layers.empty() ? 0 : state.layers.front().valid_steps; + if (state_steps > cache_steps_) { + throw std::runtime_error("TransformerBatchedKVCache state valid_steps exceeds cache capacity"); + } + valid_steps_ = state_steps; + const size_t copy_elems = static_cast(state_steps * row_elems_); + for (size_t layer = 0; layer < layers_.size(); ++layer) { + auto & cache = layers_[layer]; + const auto & source = state.layers[layer]; + if (source.valid_steps != state_steps) { + throw std::runtime_error("TransformerBatchedKVCache requires consistent valid_steps across all layers"); + } + const size_t state_elems = static_cast(batch_size_) * copy_elems; + if (source.key.size() != source.value.size() || source.key.size() != state_elems) { + throw std::runtime_error("TransformerBatchedKVCache source tensors do not match batch * valid_steps * row_elems"); + } + std::fill(cache.import_key_scratch.begin(), cache.import_key_scratch.end(), 0.0F); + std::fill(cache.import_value_scratch.begin(), cache.import_value_scratch.end(), 0.0F); + for (int64_t batch = 0; batch < batch_size_; ++batch) { + const size_t src_offset = static_cast(batch) * copy_elems; + const size_t dst_offset = static_cast(batch * cache_steps_ * row_elems_); + std::copy( + source.key.begin() + static_cast(src_offset), + source.key.begin() + static_cast(src_offset + copy_elems), + cache.import_key_scratch.begin() + static_cast(dst_offset)); + std::copy( + source.value.begin() + static_cast(src_offset), + source.value.begin() + static_cast(src_offset + copy_elems), + cache.import_value_scratch.begin() + static_cast(dst_offset)); + } + write_cache_tensor(cache.key_tensor, cache.import_key_scratch, options_); + write_cache_tensor(cache.value_tensor, cache.import_value_scratch, options_); + } +} + +TransformerBatchedKVState TransformerBatchedKVCache::export_state() const { + TransformerBatchedKVState state; + state.batch_size = batch_size_; + state.current_end = current_end_; + state.layers.resize(layers_.size()); + const size_t copy_elems = static_cast(valid_steps_ * row_elems_); + const size_t state_elems = static_cast(batch_size_) * copy_elems; + for (size_t layer = 0; layer < layers_.size(); ++layer) { + auto & out = state.layers[layer]; + out.valid_steps = valid_steps_; + out.key.resize(state_elems); + out.value.resize(state_elems); + if (copy_elems == 0) { + continue; + } + const auto key_values = read_cache_tensor(layers_[layer].key_tensor, options_); + const auto value_values = read_cache_tensor(layers_[layer].value_tensor, options_); + for (int64_t batch = 0; batch < batch_size_; ++batch) { + const size_t src_offset = static_cast(batch * cache_steps_ * row_elems_); + const size_t dst_offset = static_cast(batch) * copy_elems; + std::copy( + key_values.begin() + static_cast(src_offset), + key_values.begin() + static_cast(src_offset + copy_elems), + out.key.begin() + static_cast(dst_offset)); + std::copy( + value_values.begin() + static_cast(src_offset), + value_values.begin() + static_cast(src_offset + copy_elems), + out.value.begin() + static_cast(dst_offset)); + } + } + return state; +} + +void TransformerBatchedKVCache::advance_after_direct_append(int64_t steps) { + if (steps <= 0) { + return; + } + if (valid_steps_ + steps > cache_steps_) { + throw std::runtime_error("TransformerBatchedKVCache direct append exceeds cache capacity"); + } + valid_steps_ += steps; + current_end_ += steps; +} + +int64_t TransformerBatchedKVCache::batch_size() const noexcept { + return batch_size_; +} + +int64_t TransformerBatchedKVCache::valid_steps() const noexcept { + return valid_steps_; +} + +int64_t TransformerBatchedKVCache::current_end() const noexcept { + return current_end_; +} + +int64_t TransformerBatchedKVCache::cache_steps() const noexcept { + return cache_steps_; +} + core::TensorValue view_transformer_kv_cache_steps( core::ModuleBuildContext & ctx, const core::TensorValue & cache, diff --git a/src/framework/sampling/hf_sampler.cpp b/src/framework/sampling/hf_sampler.cpp index 321a709c..f6cb8cf6 100644 --- a/src/framework/sampling/hf_sampler.cpp +++ b/src/framework/sampling/hf_sampler.cpp @@ -386,13 +386,21 @@ int32_t HfTokenSampler::sample_from_processed_scores( int32_t best_token = -1; for (size_t index = 0; index < scratch.candidates_.size(); ++index) { const int32_t token = scratch.candidates_[index]; - const float exponential = torch_cuda_tensor_iterator_exponential_element( - torch_state->seed, - static_cast(scores.size()), - static_cast(token), - torch_state->call_index, - torch_state->policy->multiprocessor_count, - torch_state->policy->max_threads_per_multiprocessor); + const float exponential = torch_state->use_offset_blocks + ? torch_cuda_tensor_iterator_exponential_element_at_offset( + torch_state->seed, + static_cast(scores.size()), + static_cast(token), + torch_state->offset_blocks, + torch_state->policy->multiprocessor_count, + torch_state->policy->max_threads_per_multiprocessor) + : torch_cuda_tensor_iterator_exponential_element( + torch_state->seed, + static_cast(scores.size()), + static_cast(token), + torch_state->call_index, + torch_state->policy->multiprocessor_count, + torch_state->policy->max_threads_per_multiprocessor); const double rank = scratch.weights_[index] / static_cast(exponential); if (rank > best_rank) { best_rank = rank; diff --git a/src/framework/sampling/torch_random.cpp b/src/framework/sampling/torch_random.cpp index 82367d9a..9bcf0f47 100644 --- a/src/framework/sampling/torch_random.cpp +++ b/src/framework/sampling/torch_random.cpp @@ -444,6 +444,29 @@ float torch_cuda_tensor_iterator_exponential_element( uint64_t call_index, int64_t multiprocessor_count, int64_t max_threads_per_multiprocessor) { + const uint64_t offset_blocks = + call_index * torch_cuda_tensor_iterator_offset_blocks(total_elements, TorchCudaSamplingPolicy{ + multiprocessor_count, + max_threads_per_multiprocessor, + false, + 0, + }); + return torch_cuda_tensor_iterator_exponential_element_at_offset( + seed, + total_elements, + element_index, + offset_blocks, + multiprocessor_count, + max_threads_per_multiprocessor); +} + +float torch_cuda_tensor_iterator_exponential_element_at_offset( + uint64_t seed, + uint64_t total_elements, + uint64_t element_index, + uint64_t offset_blocks, + int64_t multiprocessor_count, + int64_t max_threads_per_multiprocessor) { if (total_elements == 0 || element_index >= total_elements) { throw std::invalid_argument("torch CUDA TensorIterator exponential element index is out of range"); } @@ -465,8 +488,9 @@ float torch_cuda_tensor_iterator_exponential_element( const int component = static_cast(chunk % unroll_factor); const uint64_t loop_index = chunk / unroll_factor; const uint64_t sequence = element_index % stride; - const uint64_t offset_blocks = call_index * (counter_offset / unroll_factor) + loop_index; - const float uniform = torch_cuda_uniform_tensor_iterator_element(seed, sequence, offset_blocks, component); + (void)counter_offset; + const float uniform = + torch_cuda_uniform_tensor_iterator_element(seed, sequence, offset_blocks + loop_index, component); return -std::log(uniform); }