Skip to content

Rename NemotronCacheConfig to NemotronConfig and add blank penalty to the decoder #2042

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Nenad Banfic (nenad1002) merged 4 commits into from
Apr 16, 2026
Merged

Rename NemotronCacheConfig to NemotronConfig and add blank penalty to the decoder #2042

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
75e0add
Add blank penalty
nenad1002 Mar 22, 2026
1d83fe1
Rename NemotronConfig
nenad1002 Mar 22, 2026
b3846fe
Fix copilot suggestions
nenad1002 Mar 22, 2026
987c61d
Move blank penalty to search
nenad1002 Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/config.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1202,6 +1202,8 @@ struct Search_Element : JSON::Element {
v_.past_present_share_buffer = JSON::Get<bool>(value);
} else if (name == "early_stopping") {
v_.early_stopping = JSON::Get<bool>(value);
} else if (name == "blank_penalty") {
v_.blank_penalty = static_cast<float>(JSON::Get<double>(value));
} else {
throw JSON::unknown_value_error{};
}
Expand Down
1 change: 1 addition & 0 deletions src/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,6 +379,7 @@ struct Config {
bool past_present_share_buffer{}; // The past/present kv tensors are shared and allocated once to max_length (cuda only)
int random_seed{-1}; // -1 = Seed with random device, otherwise use value to seed RNG
std::optional<size_t> chunk_size; // Chunk size for prefill chunking during context processing. If present, chunking is enabled with the chunk size > 0.
float blank_penalty{}; // Penalty applied to blank token logits in CTC/RNNT decoding. Default 0 means no penalty.
} search;

struct Engine {
Expand Down
49 changes: 26 additions & 23 deletions src/models/nemotron_speech.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

namespace Generators {

void NemotronCacheConfig::PopulateFromConfig(const Config& config) {
void NemotronConfig::PopulateFromConfig(const Config& config) {
const auto& enc = config.model.encoder;
const auto& dec = config.model.decoder;
const auto& jo = config.model.joiner;
Expand Down Expand Up @@ -39,6 +39,7 @@ void NemotronCacheConfig::PopulateFromConfig(const Config& config) {
chunk_samples = config.model.chunk_samples;
blank_id = config.model.blank_id;
max_symbols_per_step = config.model.max_symbols_per_step;
blank_penalty = config.search.blank_penalty;

// Vocab size from top-level config
vocab_size = config.model.vocab_size;
Expand Down Expand Up @@ -71,7 +72,7 @@ void NemotronCacheConfig::PopulateFromConfig(const Config& config) {
dec_out_lstm_cell = dec.outputs.lstm_cell_state;
}

void NemotronEncoderCache::Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
void NemotronEncoderCache::Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
auto cache_channel_type = session_info.GetInputDataType(cfg.enc_in_cache_channel);
auto cache_time_type = session_info.GetInputDataType(cfg.enc_in_cache_time);
auto cache_channel_len_type = session_info.GetInputDataType(cfg.enc_in_cache_channel_len);
Expand All @@ -92,11 +93,11 @@ void NemotronEncoderCache::Initialize(const NemotronCacheConfig& cfg, const Sess
*cache_last_channel_len->GetTensorMutableData<int64_t>() = 0;
}

void NemotronEncoderCache::Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
void NemotronEncoderCache::Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
Initialize(cfg, session_info, allocator, device);
}

void NemotronDecoderState::Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
void NemotronDecoderState::Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
auto lstm_hidden_type = session_info.GetInputDataType(cfg.dec_in_lstm_hidden);
auto lstm_cell_type = session_info.GetInputDataType(cfg.dec_in_lstm_cell);

Expand All @@ -111,14 +112,14 @@ void NemotronDecoderState::Initialize(const NemotronCacheConfig& cfg, const Sess
last_token = cfg.blank_id; // Start with blank/SOS token
}

void NemotronDecoderState::Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
void NemotronDecoderState::Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device) {
Initialize(cfg, session_info, allocator, device);
}

NemotronSpeechModel::NemotronSpeechModel(std::unique_ptr<Config> config, OrtEnv& ort_env)
: Model{std::move(config)} {
cache_config_ = NemotronCacheConfig{};
cache_config_.PopulateFromConfig(*config_);
nemotron_config_ = NemotronConfig{};
nemotron_config_.PopulateFromConfig(*config_);

// Create session options
encoder_session_options_ = OrtSessionOptions::Create();
Expand Down Expand Up @@ -169,7 +170,7 @@ std::unique_ptr<State> NemotronSpeechModel::CreateState(DeviceSpan<int32_t> /*se
NemotronEncoderSubState::NemotronEncoderSubState(const NemotronSpeechModel& model, const GeneratorParams& params)
: State{params, model},
model_{model} {
auto& cfg = model_.cache_config_;
auto& cfg = model_.nemotron_config_;
auto& allocator = model_.allocator_cpu_;
auto& device = *model_.p_device_;

Expand Down Expand Up @@ -250,7 +251,7 @@ DeviceSpan<float> NemotronEncoderSubState::Run(int /*total_length*/, DeviceSpan<
NemotronPredictionSubState::NemotronPredictionSubState(const NemotronSpeechModel& model, const GeneratorParams& params)
: State{params, model},
model_{model} {
auto& cfg = model_.cache_config_;
auto& cfg = model_.nemotron_config_;
auto& allocator = model_.allocator_cpu_;
auto& device = *model_.p_device_;

Expand Down Expand Up @@ -304,7 +305,7 @@ DeviceSpan<float> NemotronPredictionSubState::Run(int /*total_length*/, DeviceSp
NemotronJoinerSubState::NemotronJoinerSubState(const NemotronSpeechModel& model, const GeneratorParams& params)
: State{params, model},
model_{model} {
auto& cfg = model_.cache_config_;
auto& cfg = model_.nemotron_config_;

// Register inputs
encoder_input_idx_ = inputs_.size();
Expand Down Expand Up @@ -339,15 +340,15 @@ NemotronSpeechState::NemotronSpeechState(const NemotronSpeechModel& model,
const GeneratorParams& params)
: State{params, model},
nemotron_model_{model} {
cache_config_ = model.cache_config_;
nemotron_config_ = model.nemotron_config_;

encoder_state_ = std::make_unique<NemotronEncoderSubState>(model, params);
prediction_state_ = std::make_unique<NemotronPredictionSubState>(model, params);
joiner_state_ = std::make_unique<NemotronJoinerSubState>(model, params);

// Pre-allocate encoder frame for joiner input
auto enc_out_type = model_.session_info_.GetOutputDataType(cache_config_.enc_out_encoded);
auto frame_shape = std::array<int64_t, 3>{1, 1, cache_config_.hidden_dim};
auto enc_out_type = model_.session_info_.GetOutputDataType(nemotron_config_.enc_out_encoded);
auto frame_shape = std::array<int64_t, 3>{1, 1, nemotron_config_.hidden_dim};
encoder_frame_ = OrtValue::CreateTensor(model_.allocator_cpu_, frame_shape, enc_out_type);
}

Expand All @@ -363,7 +364,7 @@ DeviceSpan<float> NemotronSpeechState::Run(int /*total_length*/,

void NemotronSpeechState::SetExtraInputs(const std::vector<ExtraInput>& extra_inputs) {
for (const auto& input : extra_inputs) {
if (input.name == Config::Defaults::AudioFeaturesName || input.name == cache_config_.enc_in_audio) {
if (input.name == Config::Defaults::AudioFeaturesName || input.name == nemotron_config_.enc_in_audio) {
current_mel_ = input.tensor;
need_encoder_run_ = true;
chunk_done_ = false;
Expand All @@ -389,11 +390,11 @@ void NemotronSpeechState::ResetStreamingState() {
auto& allocator = model_.allocator_cpu_;
auto& device = *model_.p_device_;

encoder_state_->cache_.Reset(cache_config_, model_.session_info_, allocator, device);
encoder_state_->cache_.Reset(nemotron_config_, model_.session_info_, allocator, device);
encoder_state_->UpdateCacheInputs();
encoder_state_->first_run_ = true;

prediction_state_->lstm_state_.Reset(cache_config_, model_.session_info_, allocator, device);
prediction_state_->lstm_state_.Reset(nemotron_config_, model_.session_info_, allocator, device);
prediction_state_->UpdateInputs();
prediction_state_->first_run_ = true;

Expand Down Expand Up @@ -473,7 +474,7 @@ std::span<const int32_t> NemotronSpeechState::StepToken() {
// Reshape decoder output for joiner: [1, dim] -> [1, 1, dim]
auto dec_out_shape = prediction_state_->outputs_[0]->GetTensorTypeAndShapeInfo()->GetShape();
auto decoder_frame_shape = std::array<int64_t, 3>{1, 1, dec_out_shape[1]};
auto dec_out_type = model_.session_info_.GetOutputDataType(cache_config_.dec_out_outputs);
auto dec_out_type = model_.session_info_.GetOutputDataType(nemotron_config_.dec_out_outputs);
auto decoder_frame = OrtValue::CreateTensor(allocator, decoder_frame_shape, dec_out_type);
ByteWrapTensor(*model_.p_device_, *decoder_frame)
.CopyFrom(ByteWrapTensor(*model_.p_device_, *prediction_state_->outputs_[0]));
Expand All @@ -483,21 +484,23 @@ std::span<const int32_t> NemotronSpeechState::StepToken() {
joiner_state_->Run(0, dummy_tokens);

// Argmax over logits
const float* logits_data = joiner_state_->outputs_[0]->GetTensorData<float>();
const float* logits = joiner_state_->outputs_[0]->GetTensorData<float>();
auto logits_shape = joiner_state_->outputs_[0]->GetTensorTypeAndShapeInfo()->GetShape();
int total_logits = 1;
for (auto d : logits_shape) total_logits *= static_cast<int>(d);

Comment thread
nenad1002 marked this conversation as resolved.
// Apply blank penalty virtually during argmax to avoid mutating ORT output buffer
int best_token = 0;
float best_score = logits_data[0];
float best_score = logits[0] - (nemotron_config_.blank_id == 0 ? nemotron_config_.blank_penalty : 0.0f);
for (int i = 1; i < total_logits; ++i) {
if (logits_data[i] > best_score) {
best_score = logits_data[i];
float score = (i == nemotron_config_.blank_id) ? logits[i] - nemotron_config_.blank_penalty : logits[i];
if (score > best_score) {
best_score = score;
best_token = i;
}
}

if (best_token == cache_config_.blank_id) {
if (best_token == nemotron_config_.blank_id) {
time_step_++;
symbol_step_ = 0;
continue;
Expand All @@ -511,7 +514,7 @@ std::span<const int32_t> NemotronSpeechState::StepToken() {
prediction_state_->outputs_[2] = nullptr;

symbol_step_++;
if (symbol_step_ >= cache_config_.max_symbols_per_step) {
if (symbol_step_ >= nemotron_config_.max_symbols_per_step) {
time_step_++;
symbol_step_ = 0;
}
Expand Down
15 changes: 8 additions & 7 deletions src/models/nemotron_speech.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

namespace Generators {

struct NemotronCacheConfig {
struct NemotronConfig {
// Encoder dimensions (from encoder.hidden_size / num_hidden_layers)
int num_encoder_layers{};
int hidden_dim{};
Expand All @@ -23,6 +23,7 @@ struct NemotronCacheConfig {
// Vocabulary
int vocab_size{};
int blank_id{};
float blank_penalty{};

// Streaming chunk config
int chunk_frames{};
Expand Down Expand Up @@ -76,8 +77,8 @@ struct NemotronEncoderCache {
std::unique_ptr<OrtValue> cache_last_time;
std::unique_ptr<OrtValue> cache_last_channel_len;

void Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
void Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
void Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
void Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
};

/// Holds the RNNT decoder LSTM hidden states between decoding steps.
Expand All @@ -86,8 +87,8 @@ struct NemotronDecoderState {
std::unique_ptr<OrtValue> lstm_cell_state;
int last_token{0}; // Last emitted non-blank token (for autoregressive feedback)

void Initialize(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
void Reset(const NemotronCacheConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
void Initialize(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
void Reset(const NemotronConfig& cfg, const SessionInfo& session_info, OrtAllocator& allocator, DeviceInterface& device);
};

struct NemotronSpeechModel : Model {
Expand All @@ -105,7 +106,7 @@ struct NemotronSpeechModel : Model {
std::unique_ptr<OrtSessionOptions> decoder_session_options_;
std::unique_ptr<OrtSessionOptions> joiner_session_options_;

NemotronCacheConfig cache_config_;
NemotronConfig nemotron_config_;
};

/// Sub-state for the streaming encoder.
Expand Down Expand Up @@ -201,7 +202,7 @@ struct NemotronSpeechState : State {

private:
const NemotronSpeechModel& nemotron_model_;
NemotronCacheConfig cache_config_;
NemotronConfig nemotron_config_;

std::unique_ptr<NemotronEncoderSubState> encoder_state_;
std::unique_ptr<NemotronPredictionSubState> prediction_state_;
Expand Down
26 changes: 13 additions & 13 deletions src/models/nemotron_streaming_processor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,24 +16,24 @@ NemotronStreamingProcessor::NemotronStreamingProcessor(Model& model)
throw std::runtime_error("NemotronStreamingProcessor requires a nemotron_speech model type. Got: " + model.config_->model.type);
}

cache_config_ = nemotron_model->cache_config_;
nemotron_config_ = nemotron_model->nemotron_config_;

if (cache_config_.pre_encode_cache_size <= 0) {
if (nemotron_config_.pre_encode_cache_size <= 0) {
throw std::runtime_error("NemotronStreamingProcessor requires pre_encode_cache_size > 0. Got: " +
std::to_string(cache_config_.pre_encode_cache_size));
std::to_string(nemotron_config_.pre_encode_cache_size));
}

// Initialize mel extractor from config
nemo_mel::NemoMelConfig mel_cfg{
cache_config_.num_mels, cache_config_.fft_size,
cache_config_.hop_length, cache_config_.win_length,
cache_config_.sample_rate,
cache_config_.preemph, cache_config_.log_eps};
nemotron_config_.num_mels, nemotron_config_.fft_size,
nemotron_config_.hop_length, nemotron_config_.win_length,
nemotron_config_.sample_rate,
nemotron_config_.preemph, nemotron_config_.log_eps};
mel_extractor_ = nemo_mel::NemoStreamingMelExtractor{mel_cfg};

// Initialize mel pre-encode cache (time-major ring buffer, zeros for first chunk)
mel_pre_encode_cache_.assign(
static_cast<size_t>(cache_config_.pre_encode_cache_size) * cache_config_.num_mels, 0.0f);
static_cast<size_t>(nemotron_config_.pre_encode_cache_size) * nemotron_config_.num_mels, 0.0f);
cache_pos_ = 0;
}

Expand All @@ -43,7 +43,7 @@ std::unique_ptr<NamedTensors> NemotronStreamingProcessor::Process(const float* a
// Append incoming audio to accumulation buffer
audio_buffer_.insert(audio_buffer_.end(), audio_data, audio_data + num_samples);

const size_t chunk_size = static_cast<size_t>(cache_config_.chunk_samples);
const size_t chunk_size = static_cast<size_t>(nemotron_config_.chunk_samples);

// Process the first complete chunk available
if (audio_buffer_.size() >= chunk_size) {
Expand All @@ -63,7 +63,7 @@ std::unique_ptr<NamedTensors> NemotronStreamingProcessor::Flush() {
return nullptr;
}

const size_t chunk_size = static_cast<size_t>(cache_config_.chunk_samples);
const size_t chunk_size = static_cast<size_t>(nemotron_config_.chunk_samples);
audio_buffer_.resize(chunk_size, 0.0f); // Pad with silence

auto mel = BuildMelTensor(audio_buffer_.data(), chunk_size);
Expand All @@ -79,12 +79,12 @@ std::unique_ptr<OrtValue> NemotronStreamingProcessor::BuildMelTensor(const float
// Compute mel spectrogram for this chunk: returns [num_mels, num_frames] (frequency-major)
auto [mel_data, num_frames] = mel_extractor_.Process(audio_chunk, chunk_samples);

const int cache_size = cache_config_.pre_encode_cache_size;
const int num_mels = cache_config_.num_mels;
const int cache_size = nemotron_config_.pre_encode_cache_size;
const int num_mels = nemotron_config_.num_mels;
const int total_mel_frames = cache_size + num_frames;

// Create output tensor: [1, total_mel_frames, num_mels] (time-major)
auto signal_type = model_.session_info_.GetInputDataType(cache_config_.enc_in_audio);
auto signal_type = model_.session_info_.GetInputDataType(nemotron_config_.enc_in_audio);

// TODO: Optimize for GPU/CUDA later, CPU always expects float32.
if (signal_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
Expand Down
6 changes: 3 additions & 3 deletions src/models/nemotron_streaming_processor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,12 @@ struct NemotronStreamingProcessor : StreamingProcessor {
std::unique_ptr<NamedTensors> Process(const float* audio_data, size_t num_samples) override;
std::unique_ptr<NamedTensors> Flush() override;

int GetChunkSamples() const { return cache_config_.chunk_samples; }
int GetSampleRate() const { return cache_config_.sample_rate; }
int GetChunkSamples() const { return nemotron_config_.chunk_samples; }
int GetSampleRate() const { return nemotron_config_.sample_rate; }

private:
Model& model_;
NemotronCacheConfig cache_config_;
NemotronConfig nemotron_config_;

// Log-mel feature extraction
nemo_mel::NemoStreamingMelExtractor mel_extractor_;
Expand Down
Loading