Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
100 lines
3.8 KiB
C++
100 lines
3.8 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/models/nemotron_asr/assets.h"
|
|
#include "engine/models/nemotron_asr/frontend.h"
|
|
#include "engine/models/nemotron_asr/weights.h"
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
namespace engine::models::nemotron_asr {
|
|
|
|
struct NemotronEncodedAudio {
|
|
std::vector<float> values;
|
|
std::vector<int32_t> durations;
|
|
int64_t frames = 0;
|
|
int64_t valid_frames = 0;
|
|
int64_t hidden_size = 0;
|
|
};
|
|
|
|
struct NemotronEncoderStreamState {
|
|
int64_t attention_seen_frames = 0;
|
|
int64_t attention_cached_frames = 0;
|
|
bool first_chunk = true;
|
|
bool backend_cache_valid = false;
|
|
const void * backend_cache_owner = nullptr;
|
|
};
|
|
|
|
// One speaker's cache-aware stream in NeMo's pad_and_drop_preencoded layout
|
|
// (masked speaker-tagged ASR): the K/V and convolution caches live in backend
|
|
// tensors owned by this state, so several speakers can share one graph.
|
|
struct NemotronEncoderSpeakerCaches;
|
|
struct NemotronEncoderSpeakerState {
|
|
std::shared_ptr<NemotronEncoderSpeakerCaches> caches;
|
|
int64_t attention_seen_frames = 0;
|
|
int64_t attention_cached_frames = 0;
|
|
};
|
|
|
|
class NemotronFastConformerEncoderRuntime {
|
|
public:
|
|
NemotronFastConformerEncoderRuntime(
|
|
std::shared_ptr<const NemotronASRAssets> assets,
|
|
std::shared_ptr<const NemotronWeights> weights,
|
|
engine::core::ExecutionContext & execution_context,
|
|
size_t graph_arena_bytes);
|
|
~NemotronFastConformerEncoderRuntime();
|
|
|
|
void prepare_capacity(int64_t input_frames, int64_t feature_dim, int64_t lookahead_tokens);
|
|
void prepare_streaming_capacity(int64_t feature_dim, int64_t lookahead_tokens);
|
|
void release_offline_graph();
|
|
NemotronEncodedAudio encode(const NemotronFrontendFeatures & features, int64_t prompt_id, int64_t lookahead_tokens);
|
|
NemotronEncoderStreamState make_stream_state() const;
|
|
NemotronEncodedAudio encode_stream_chunk(
|
|
const NemotronFrontendFeatures & features,
|
|
int64_t prompt_id,
|
|
int64_t lookahead_tokens,
|
|
NemotronEncoderStreamState & state);
|
|
// One pad_and_drop step for several speakers in one graph (batch = rows). Each row
|
|
// holds pre-encode cache frames followed by the new frames; the subsampling runs
|
|
// statelessly over all of them and the first drop_extra_pre_encoded encoder frames
|
|
// are dropped. Rows share the frame count and have their own caches and masks.
|
|
std::vector<NemotronEncodedAudio> encode_pad_and_drop_batch(
|
|
const std::vector<const NemotronFrontendFeatures *> & features,
|
|
int64_t prompt_id,
|
|
int64_t lookahead_tokens,
|
|
const std::vector<NemotronEncoderSpeakerState *> & states);
|
|
int64_t pad_and_drop_cache_frames() const;
|
|
|
|
private:
|
|
struct Graph;
|
|
|
|
void ensure_graph(int64_t input_frames, int64_t feature_dim, int64_t lookahead_tokens);
|
|
Graph & ensure_stream_graph(
|
|
int64_t input_frames,
|
|
int64_t feature_dim,
|
|
int64_t lookahead_tokens,
|
|
int64_t prefix_frames,
|
|
bool first_chunk,
|
|
bool raw_cache = false,
|
|
int64_t rows = 1);
|
|
const std::vector<float> & relative_positional_encoding(int64_t frames);
|
|
|
|
std::shared_ptr<const NemotronASRAssets> assets_;
|
|
std::shared_ptr<const NemotronWeights> weights_;
|
|
engine::core::ExecutionContext * execution_context_ = nullptr;
|
|
size_t graph_arena_bytes_ = 0;
|
|
std::unique_ptr<Graph> graph_;
|
|
std::vector<std::unique_ptr<Graph>> stream_graphs_;
|
|
std::vector<float> input_scratch_;
|
|
std::vector<float> output_scratch_;
|
|
std::vector<float> prompt_scratch_;
|
|
std::vector<int32_t> mask_scratch_;
|
|
std::vector<float> attention_mask_scratch_;
|
|
std::unordered_map<int64_t, std::vector<float>> relative_positional_encoding_cache_;
|
|
};
|
|
|
|
} // namespace engine::models::nemotron_asr
|