diff --git a/cactus-engine/CMakeLists.txt b/cactus-engine/CMakeLists.txt index eee348095..3c0f5096f 100644 --- a/cactus-engine/CMakeLists.txt +++ b/cactus-engine/CMakeLists.txt @@ -122,8 +122,14 @@ set(_CACTUS_ENGINE_ARCHIVES set(_CACTUS_ENGINE_DEPS cactus_engine cactus_graph cactus_kernels) if(APPLE) + find_program(CACTUS_APPLE_LIBTOOL + NAMES libtool + PATHS /usr/bin + NO_DEFAULT_PATH + REQUIRED + ) add_custom_target(cactus_engine_bundle ALL - COMMAND libtool -static -o ${_CACTUS_ENGINE_BUNDLE} ${_CACTUS_ENGINE_ARCHIVES} + COMMAND ${CACTUS_APPLE_LIBTOOL} -static -o ${_CACTUS_ENGINE_BUNDLE} ${_CACTUS_ENGINE_ARCHIVES} DEPENDS ${_CACTUS_ENGINE_DEPS} BYPRODUCTS ${_CACTUS_ENGINE_BUNDLE} ) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 1d19964ee..5e59ffdf9 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -147,13 +147,19 @@ struct Config { float rope_scaling_factor = 1.0f; float rope_mscale_all_dim = 0.0f; - enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18}; + enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18, NEMOTRON_ASR = 19}; uint32_t predictor_hidden_dim = 0; uint32_t predictor_num_layers = 0; uint32_t tdt_joint_dim = 0; uint32_t tdt_num_durations = 0; uint32_t tdt_blank_id = 0; std::vector tdt_durations; + uint32_t rnnt_joint_dim = 0; + uint32_t rnnt_blank_id = 0; + uint32_t prompt_dim = 0; + uint32_t default_prompt_id = 0; + uint32_t max_symbols_per_step = 10; + std::unordered_map prompt_dictionary; ModelType model_type = ModelType::GEMMA4; @@ -672,11 +678,28 @@ class Model { double raw_decode_ms = 0.0; }; + struct RnntStreamState { + bool initialized = false; + uint32_t last_token = 0; + size_t time_index = 0; + std::vector> dec_state; + std::vector pending; + float confirmed_sec = 0.0f; + size_t decoded_tokens = 0; + double raw_decode_ms = 0.0; + }; + std::vector transcribe_parakeet_tdt(const std::vector& audio_features, ParakeetTdtStreamState* stream = nullptr, bool is_final = true, size_t end_frame = 0, const std::atomic* should_stop = nullptr); + std::vector transcribe_nemotron_asr(const std::vector& audio_features, + const std::string& target_lang = "auto", + RnntStreamState* stream = nullptr, + bool is_final = true, + size_t end_frame = 0, + const std::atomic* should_stop = nullptr); std::vector transcribe_whisper_seq2seq(const std::vector& audio_features, const std::vector& decoder_prompt_tokens, size_t max_tokens, @@ -724,7 +747,9 @@ class Model { size_t last_prefill_tail_padding_tokens() const { return last_prefill_tail_padding_tokens_; } bool load_npu_audio_encoder(const std::string& model_path, const std::string& compute_units = ""); + bool load_npu_audio_encoders(const std::vector& model_paths, const std::string& compute_units = ""); bool has_npu_audio_encoder() const { return npu_audio_encoder_ != nullptr; } + npu::NPUEncoder* select_npu_audio_encoder(size_t feature_dim, size_t frame_count) const; bool load_npu_vision_encoder(const std::string& model_path); bool has_npu_vision_encoder() const { return npu_vision_encoder_ != nullptr; } @@ -884,11 +909,14 @@ class Model { std::string family_; std::string npu_audio_encoder_mlpackage_; + std::vector npu_audio_encoder_mlpackages_; std::string npu_audio_compute_units_; + size_t npu_audio_prewarm_frames_ = 0; std::string npu_vision_encoder_mlpackage_; std::string npu_source_encoder_mlpackage_; std::unique_ptr npu_audio_encoder_; + std::vector> npu_audio_encoders_; std::unique_ptr npu_vision_encoder_; std::unique_ptr npu_source_encoder_; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index a990becba..caf2a69b3 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -44,6 +44,32 @@ std::vector parse_config_uint_list(const std::string& value) { return numbers; } +std::unordered_map parse_config_string_uint_map(const std::string& value) { + std::unordered_map result; + std::stringstream ss(value); + std::string item; + while (std::getline(ss, item, ',')) { + size_t sep = item.find(':'); + if (sep == std::string::npos) continue; + std::string key = item.substr(0, sep); + std::string number = item.substr(sep + 1); + auto trim = [](std::string& s) { + size_t first = s.find_first_not_of(" \t"); + if (first == std::string::npos) { + s.clear(); + return; + } + size_t last = s.find_last_not_of(" \t"); + s = s.substr(first, last - first + 1); + }; + trim(key); + trim(number); + if (key.empty() || number.empty()) continue; + result[key] = static_cast(std::stoul(number)); + } + return result; +} + float read_scalar_value(Precision precision, const uint8_t* data, size_t index) { const uint8_t* ptr = data + PrecisionTraits::byte_offset_of(precision, index); switch (precision) { @@ -615,7 +641,16 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, cache_max_seq_len_ = context_size; - if (!npu_audio_encoder_mlpackage_.empty()) { + if (!npu_audio_encoder_mlpackages_.empty()) { + std::vector full_paths; + full_paths.reserve(npu_audio_encoder_mlpackages_.size()); + for (const std::string& rel_path : npu_audio_encoder_mlpackages_) { + full_paths.push_back(bundle_dir + "/" + rel_path); + } + if (!load_npu_audio_encoders(full_paths, npu_audio_compute_units_)) { + CACTUS_LOG_WARN("model", "NPU audio encoder variants failed to load; falling back to CPU"); + } + } else if (!npu_audio_encoder_mlpackage_.empty()) { std::string full_path = bundle_dir + "/" + npu_audio_encoder_mlpackage_; if (!load_npu_audio_encoder(full_path, npu_audio_compute_units_)) { CACTUS_LOG_WARN("model", "NPU audio encoder load failed for " << full_path << "; falling back to CPU"); @@ -670,9 +705,23 @@ bool Model::load_manifest() { if (obj.count("npu_audio_encoder") && obj.at("npu_audio_encoder").is()) { npu_audio_encoder_mlpackage_ = obj.at("npu_audio_encoder").get(); } + npu_audio_encoder_mlpackages_.clear(); + if (obj.count("npu_audio_encoders") && obj.at("npu_audio_encoders").is()) { + for (const auto& v : obj.at("npu_audio_encoders").get()) { + if (v.is()) npu_audio_encoder_mlpackages_.push_back(v.get()); + } + } if (obj.count("npu_audio_compute_units") && obj.at("npu_audio_compute_units").is()) { npu_audio_compute_units_ = obj.at("npu_audio_compute_units").get(); } + npu_audio_prewarm_frames_ = 0; + if (obj.count("inputs") && obj.at("inputs").is()) { + const auto& inputs = obj.at("inputs").get(); + auto it = inputs.find("active_feature_frames"); + if (it != inputs.end() && it->second.is() && it->second.get() > 0.0) { + npu_audio_prewarm_frames_ = static_cast(it->second.get()); + } + } if (obj.count("npu_vision_encoder") && obj.at("npu_vision_encoder").is()) { npu_vision_encoder_mlpackage_ = obj.at("npu_vision_encoder").get(); } @@ -742,6 +791,7 @@ bool Model::setup_tokenizer() { std::string merges = bundle_dir_ + "/merges.txt"; std::string cfg = bundle_dir_ + "/tokenizer_config.txt"; if (!fs::exists(vocab)) return false; + if (!fs::exists(cfg)) cfg = bundle_dir_ + "/config.txt"; auto rt = load_tokenizer_runtime_config(cfg); bool use_bpe = rt.tokenizer_type == TokenizerRuntimeConfig::TokenizerType::BPE || (rt.tokenizer_type == TokenizerRuntimeConfig::TokenizerType::UNKNOWN @@ -2931,6 +2981,7 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a std::vector<__fp16> npu_hidden_storage; bool used_npu = false; size_t npu_hidden_T = 0; + size_t npu_hidden_D = 0; if (has_npu_audio_encoder()) { const std::vector in_shape = npu_audio_encoder_->get_input_shape(); const std::vector out_shape = npu_audio_encoder_->get_output_shape(); @@ -2941,6 +2992,7 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a const size_t window_frames = static_cast(in_shape[1]); const size_t window_hidden = static_cast(out_shape[1]); const size_t hidden_dim_npu = static_cast(out_shape[2]); + npu_hidden_D = hidden_dim_npu; const size_t chunk_input_elems = window_frames * expected_mels; const size_t chunk_output_elems = window_hidden * hidden_dim_npu; const size_t num_chunks = (copy_frames + window_frames - 1) / window_frames; @@ -3002,13 +3054,22 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } size_t hidden_node = static_cast(audio_enc->output_node_ids[hidden_idx]); const auto& hidden_desc = audio_enc->graph->get_output_buffer(hidden_node); + if (hidden_desc.shape.size() < 3) { + CACTUS_LOG_ERROR("model", "encoder_hidden_states must be 3D [B, T, D]"); + return emitted; + } + if (used_npu && npu_hidden_D != hidden_desc.shape[2]) { + throw std::runtime_error( + "NPU audio encoder hidden dimension " + std::to_string(npu_hidden_D) + + " does not match decoder graph hidden dimension " + std::to_string(hidden_desc.shape[2])); + } const uint8_t* hidden_ptr; if (used_npu) { hidden_ptr = reinterpret_cast(npu_hidden_storage.data()); } else { hidden_ptr = static_cast(audio_enc->graph->get_output(hidden_node)); } - if (hidden_desc.shape.size() < 3 || hidden_ptr == nullptr) { + if (hidden_ptr == nullptr) { CACTUS_LOG_ERROR("model", "encoder_hidden_states must be 3D [B, T, D]"); return emitted; } @@ -3090,13 +3151,10 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a }; } - size_t commit_to = T; - if (stream) { - size_t valid_hidden = T; - if (!used_npu && expected_frames > 0) - valid_hidden = std::min(T, (copy_frames * T) / expected_frames); - commit_to = (end_frame > 0) ? std::min(end_frame, valid_hidden) : valid_hidden; - } + size_t valid_hidden = T; + if (!used_npu && expected_frames > 0 && copy_frames > 0) + valid_hidden = std::min(T, (copy_frames * T + expected_frames - 1) / expected_frames); + size_t commit_to = (stream && end_frame > 0) ? std::min(end_frame, valid_hidden) : valid_hidden; Tokenizer* stream_tok = stream ? get_tokenizer() : nullptr; const auto& tdt_vocab_bias = get_vocab_bias(); const float frame_sec = (160.0f / 16000.0f) * @@ -3221,6 +3279,360 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a return emitted; } +std::vector Model::transcribe_nemotron_asr(const std::vector& audio_features, + const std::string& target_lang, + RnntStreamState* stream, + bool is_final, + size_t end_frame, + const std::atomic* should_stop) { + std::vector emitted; + double raw_decode_ms = 0.0; + + reset_handoff_probe_rollout(); + + Component* audio_enc = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; + Component* dec = components_.count("decoder") ? &components_.at("decoder") : nullptr; + if (!audio_enc || !dec) { + CACTUS_LOG_ERROR("model", "Nemotron ASR bundle missing audio_encoder or decoder component"); + return emitted; + } + if (!bind_runtime_buffers(*audio_enc)) return emitted; + if (!bind_runtime_buffers(*dec)) return emitted; + + int feat_idx = input_index(*audio_enc, "input_features"); + if (feat_idx < 0) { + CACTUS_LOG_ERROR("model", "audio_encoder has no input_features input"); + return emitted; + } + auto& feat_buf = audio_enc->input_buffers[feat_idx]; + size_t feat_node = static_cast(audio_enc->runtime_input_node_ids[feat_idx]); + const auto& feat_desc = audio_enc->graph->get_output_buffer(feat_node); + if (feat_desc.shape.size() != 3) { + CACTUS_LOG_ERROR("model", "audio_encoder expects [1, frames, mels] input shape"); + return emitted; + } + const size_t expected_frames = feat_desc.shape[1]; + const size_t expected_mels = feat_desc.shape[2]; + const size_t source_frames = expected_mels > 0 ? audio_features.size() / expected_mels : 0; + const size_t copy_frames = std::min(source_frames, expected_frames); + std::vector transposed(expected_frames * expected_mels, 0.0f); + for (size_t t = 0; t < copy_frames; ++t) { + for (size_t m = 0; m < expected_mels; ++m) { + transposed[t * expected_mels + m] = audio_features[m * source_frames + t]; + } + } + write_typed_buffer(feat_buf, feat_desc.precision, transposed.data(), + transposed.size() * sizeof(float), Precision::FP32); + + std::vector<__fp16> npu_hidden_storage; + bool used_npu = false; + size_t npu_hidden_T = 0; + size_t npu_hidden_D = 0; + npu::NPUEncoder* npu_audio_encoder = select_npu_audio_encoder(expected_mels, source_frames); + if (npu_audio_encoder) { + const std::vector in_shape = npu_audio_encoder->get_input_shape(); + const std::vector out_shape = npu_audio_encoder->get_output_shape(); + if (in_shape.size() >= 3 && out_shape.size() >= 3 && + in_shape[1] > 0 && in_shape[2] > 0 && out_shape[1] > 0 && out_shape[2] > 0 && + static_cast(in_shape[2]) == expected_mels && + source_frames <= static_cast(in_shape[1])) { + const size_t window_frames = static_cast(in_shape[1]); + const size_t window_hidden = static_cast(out_shape[1]); + const size_t hidden_dim_npu = static_cast(out_shape[2]); + npu_hidden_D = hidden_dim_npu; + npu_hidden_storage.assign(window_hidden * hidden_dim_npu, __fp16(0)); + const bool want_probe = handoff_probe_loaded_ && + static_cast(handoff_probe_feat_dim_) == hidden_dim_npu; + std::vector<__fp16> input_fp16(window_frames * expected_mels, __fp16(0)); + if (should_stop && should_stop->load()) return emitted; + for (size_t t = 0; t < source_frames; ++t) { + const size_t local = t * expected_mels; + for (size_t m = 0; m < expected_mels; ++m) { + input_fp16[local + m] = static_cast<__fp16>(audio_features[m * source_frames + t]); + } + } + size_t written = npu_audio_encoder->encode( + input_fp16.data(), npu_hidden_storage.data(), in_shape, "x", "encoded"); + if (written > 0) { + used_npu = true; + npu_hidden_T = (source_frames * window_hidden + window_frames - 1) / window_frames; + if (npu_hidden_T > window_hidden) npu_hidden_T = window_hidden; + CACTUS_LOG_INFO("model", "Nemotron audio encoder ran on NPU (" + << window_frames << " frame capacity, " << npu_hidden_T << " valid hidden frames)"); + if (want_probe) { + const size_t probe_elems = npu_hidden_T * hidden_dim_npu; + handoff_probe_hidden_.reserve(handoff_probe_hidden_.size() + probe_elems); + const __fp16* pp = npu_hidden_storage.data(); + for (size_t i = 0; i < probe_elems; ++i) { + handoff_probe_hidden_.push_back(static_cast(pp[i])); + } + CACTUS_LOG_INFO("cloud_handoff", "Captured " << npu_hidden_T + << " NPU probe frames for the handoff probe"); + } + } else { + CACTUS_LOG_WARN("model", "NPU audio encoder failed; falling back to CPU graph"); + } + } + } + if (!used_npu) { + if (source_frames > expected_frames) { + CACTUS_LOG_WARN("model", "Nemotron audio has " << source_frames + << " frames but CPU graph capacity is " << expected_frames + << " and NPU audio encoding is unavailable; falling back with truncation"); + } + if (should_stop && should_stop->load()) return emitted; + audio_enc->graph->execute(); + maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states"); + } + + int hidden_idx = output_index(*audio_enc, "encoder_hidden_states"); + if (hidden_idx < 0) { + CACTUS_LOG_ERROR("model", "audio_encoder has no encoder_hidden_states output"); + return emitted; + } + size_t hidden_node = static_cast(audio_enc->output_node_ids[hidden_idx]); + const auto& hidden_desc = audio_enc->graph->get_output_buffer(hidden_node); + if (hidden_desc.shape.size() < 3) { + CACTUS_LOG_ERROR("model", "encoder_hidden_states must be 3D [B, T, D]"); + return emitted; + } + if (used_npu && npu_hidden_D != hidden_desc.shape[2]) { + throw std::runtime_error( + "NPU audio encoder hidden dimension " + std::to_string(npu_hidden_D) + + " does not match decoder graph hidden dimension " + std::to_string(hidden_desc.shape[2])); + } + const uint8_t* hidden_ptr; + if (used_npu) { + hidden_ptr = reinterpret_cast(npu_hidden_storage.data()); + } else { + hidden_ptr = static_cast(audio_enc->graph->get_output(hidden_node)); + } + if (hidden_ptr == nullptr) { + CACTUS_LOG_ERROR("model", "encoder_hidden_states must be 3D [B, T, D]"); + return emitted; + } + const size_t T = used_npu ? npu_hidden_T : hidden_desc.shape[1]; + const size_t D = hidden_desc.shape[2]; + const Precision hidden_precision = used_npu ? Precision::FP16 : hidden_desc.precision; + const size_t hidden_elem = PrecisionTraits::size_of(hidden_precision); + const size_t frame_bytes = D * hidden_elem; + + const int ef_idx = input_index(*dec, "encoder_frame"); + const int prompt_idx = input_index(*dec, "language_prompt"); + const int tok_in_idx = input_index(*dec, "token_ids"); + const int logits_idx = output_index(*dec, "step_logits"); + if (ef_idx < 0 || prompt_idx < 0 || tok_in_idx < 0 || logits_idx < 0) { + CACTUS_LOG_ERROR("model", "decoder missing encoder_frame / language_prompt / token_ids / step_logits ports"); + return emitted; + } + + auto& ef_buf = dec->input_buffers[ef_idx]; + const auto& ef_desc = dec->graph->get_output_buffer(static_cast(dec->runtime_input_node_ids[ef_idx])); + auto& prompt_buf = dec->input_buffers[prompt_idx]; + const auto& prompt_desc = dec->graph->get_output_buffer(static_cast(dec->runtime_input_node_ids[prompt_idx])); + auto& tok_buf = dec->input_buffers[tok_in_idx]; + const Precision tok_prec = dec->graph->get_output_buffer(static_cast(dec->runtime_input_node_ids[tok_in_idx])).precision; + void* tok_data = tok_buf.data(); + const size_t logits_node = static_cast(dec->output_node_ids[logits_idx]); + const auto& logits_desc = dec->graph->get_output_buffer(logits_node); + const Precision logits_prec = logits_desc.precision; + const size_t token_class_count = logits_desc.shape.empty() ? 0 : logits_desc.shape.back(); + if (token_class_count == 0) return emitted; + uint32_t effective_blank = config_.rnnt_blank_id; + if (effective_blank >= token_class_count) effective_blank = static_cast(token_class_count - 1); + + size_t prompt_elems = prompt_desc.byte_size / std::max(1, PrecisionTraits::size_of(prompt_desc.precision)); + if (prompt_elems == 0 && !prompt_desc.shape.empty()) prompt_elems = prompt_desc.shape.back(); + std::vector prompt(prompt_elems, 0.0f); + std::string lang = target_lang.empty() ? "auto" : target_lang; + uint32_t prompt_id = config_.default_prompt_id; + auto prompt_it = config_.prompt_dictionary.find(lang); + if (prompt_it != config_.prompt_dictionary.end()) prompt_id = prompt_it->second; + if (prompt_id >= prompt_elems) prompt_id = std::min(config_.default_prompt_id, prompt_elems ? static_cast(prompt_elems - 1) : 0); + if (!prompt.empty()) prompt[prompt_id] = 1.0f; + write_typed_buffer(prompt_buf, prompt_desc.precision, prompt.data(), prompt.size() * sizeof(float), Precision::FP32); + + std::vector state_names; + for (uint32_t i = 0; i < config_.predictor_num_layers; ++i) { + state_names.push_back("state_h_" + std::to_string(i)); + state_names.push_back("state_c_" + std::to_string(i)); + } + auto zero_state = [&](const std::string& name) -> bool { + int idx = input_index(*dec, name); + if (idx < 0) { + CACTUS_LOG_ERROR("model", "decoder missing recurrent state input " << name); + return false; + } + auto& buf = dec->input_buffers[idx]; + std::memset(buf.data(), 0, buf.size()); + return true; + }; + if (stream && stream->initialized) { + if (stream->dec_state.size() != state_names.size()) { + CACTUS_LOG_ERROR("model", "decoder recurrent stream state count mismatch"); + return emitted; + } + for (size_t i = 0; i < state_names.size(); ++i) { + int idx = input_index(*dec, state_names[i]); + if (idx < 0) { + CACTUS_LOG_ERROR("model", "decoder missing recurrent state input " << state_names[i]); + return emitted; + } + auto& buf = dec->input_buffers[idx]; + if (buf.size() != stream->dec_state[i].size()) { + CACTUS_LOG_ERROR("model", "decoder recurrent stream state size mismatch for " << state_names[i]); + return emitted; + } + std::memcpy(buf.data(), stream->dec_state[i].data(), + buf.size()); + } + } else { + for (const auto& state_name : state_names) { + if (!zero_state(state_name)) return emitted; + } + } + + struct StateCopy { void* in_data; const void* out_ptr; size_t bytes; }; + std::vector state_copies; + state_copies.reserve(state_names.size()); + for (const auto& state_name : state_names) { + int out_idx = output_index(*dec, state_name); + int in_idx = input_index(*dec, state_name); + if (out_idx < 0 || in_idx < 0) { + CACTUS_LOG_ERROR("model", "decoder missing recurrent state port " << state_name); + return emitted; + } + size_t out_node = static_cast(dec->output_node_ids[out_idx]); + const auto& out_desc = dec->graph->get_output_buffer(out_node); + auto& in_buf = dec->input_buffers[in_idx]; + const void* out_ptr = dec->graph->get_output(out_node); + if (!out_ptr || out_desc.byte_size != in_buf.size()) { + CACTUS_LOG_ERROR("model", "decoder recurrent state buffer mismatch for " << state_name); + return emitted; + } + state_copies.push_back({ + in_buf.data(), + out_ptr, + in_buf.size() + }); + } + if (state_copies.size() != state_names.size()) { + CACTUS_LOG_ERROR("model", "decoder recurrent state port count mismatch"); + return emitted; + } + + uint32_t last_token = (stream && stream->initialized) ? stream->last_token : effective_blank; + size_t time_index = (stream && stream->initialized) ? stream->time_index : 0; + size_t valid_hidden = T; + if (!used_npu && expected_frames > 0 && copy_frames > 0) + valid_hidden = std::min(T, (copy_frames * T + expected_frames - 1) / expected_frames); + size_t commit_to = (stream && end_frame > 0) ? std::min(end_frame, valid_hidden) : valid_hidden; + Tokenizer* stream_tok = stream ? get_tokenizer() : nullptr; + const auto& vocab_bias = get_vocab_bias(); + const float frame_sec = (160.0f / 16000.0f) * + static_cast(std::max(1, config_.subsampling_factor)); + const uint32_t max_symbols = std::max(1, config_.max_symbols_per_step); + + auto snapshot_state = [&]() { + std::vector> snap(state_copies.size()); + for (size_t s = 0; s < state_copies.size(); ++s) { + const uint8_t* p = static_cast(state_copies[s].in_data); + snap[s].assign(p, p + state_copies[s].bytes); + } + return snap; + }; + + std::vector> snap_state; + uint32_t snap_last_token = last_token; + size_t snap_time = time_index; + size_t confirmed_count = 0; + if (stream) snap_state = snapshot_state(); + + while (time_index < commit_to) { + if (should_stop && should_stop->load()) break; + const uint8_t* frame_ptr = hidden_ptr + time_index * frame_bytes; + write_typed_buffer(ef_buf, ef_desc.precision, frame_ptr, frame_bytes, hidden_precision); + + uint32_t symbols_added = 0; + bool advanced = false; + while (symbols_added < max_symbols) { + write_int_element(static_cast(tok_data), tok_prec, 0, static_cast(last_token)); + + const auto exec_t0 = std::chrono::steady_clock::now(); + dec->graph->execute(); + raw_decode_ms += std::chrono::duration( + std::chrono::steady_clock::now() - exec_t0).count(); + + const void* logits_ptr = dec->graph->get_output(logits_node); + auto get_logit = [&](size_t i) -> float { + if (logits_prec == Precision::FP32) return reinterpret_cast(logits_ptr)[i]; + if (logits_prec == Precision::FP16) return static_cast(reinterpret_cast(logits_ptr)[i]); + return static_cast(reinterpret_cast(logits_ptr)[i]); + }; + + uint32_t next_token = 0; + float best_score = -std::numeric_limits::infinity(); + for (size_t i = 0; i < token_class_count; ++i) { + float v = get_logit(i); + if (stream && !vocab_bias.empty()) { + auto it = vocab_bias.find(static_cast(i)); + if (it != vocab_bias.end()) v += it->second; + } + if (v > best_score) { + best_score = v; + next_token = static_cast(i); + } + } + + ++symbols_added; + if (next_token == effective_blank) { + ++time_index; + advanced = true; + break; + } + + if (stream && !is_final && stream_tok && !emitted.empty()) { + std::string piece = stream_tok->decode({next_token}); + if (!piece.empty() && piece[0] == ' ') { + snap_state = snapshot_state(); + snap_last_token = last_token; + snap_time = time_index; + confirmed_count = emitted.size(); + } + } + emitted.push_back(next_token); + last_token = next_token; + for (const auto& copy : state_copies) { + std::memcpy(copy.in_data, copy.out_ptr, copy.bytes); + } + } + + if (!advanced) ++time_index; + } + + if (stream) { + stream->initialized = true; + stream->decoded_tokens = emitted.size(); + stream->raw_decode_ms = raw_decode_ms; + stream->pending.clear(); + if (is_final) { + stream->dec_state = snapshot_state(); + stream->last_token = last_token; + stream->time_index = time_index; + stream->confirmed_sec = static_cast(time_index) * frame_sec; + } else { + stream->dec_state = std::move(snap_state); + stream->last_token = snap_last_token; + stream->time_index = snap_time; + stream->confirmed_sec = static_cast(snap_time) * frame_sec; + stream->pending.assign(emitted.begin() + confirmed_count, emitted.end()); + emitted.resize(confirmed_count); + } + } + + return emitted; +} + uint32_t Model::decode_with_images(const std::vector& tokens, const std::vector& /*image_paths*/, float temperature, float top_p, size_t top_k, const std::string& profile_file, float* out_entropy, float min_p, float repetition_penalty) { @@ -3857,6 +4269,7 @@ bool Config::from_json(const std::string& config_path) { else if (mt == "lfm2") model_type = ModelType::LFM2; else if (mt == "whisper") model_type = ModelType::WHISPER; else if (mt == "parakeet_tdt" || mt == "parakeet-tdt") model_type = ModelType::PARAKEET_TDT; + else if (mt == "nemotron_asr" || mt == "nemotron-asr") model_type = ModelType::NEMOTRON_ASR; else if (mt == "youtu") model_type = ModelType::YOUTU; else if (mt == "needle") model_type = ModelType::NEEDLE; else if (mt == "bert" || mt == "nomic") model_type = ModelType::NOMIC; @@ -3932,6 +4345,12 @@ bool Config::from_json(const std::string& config_path) { else if (key == "tdt_joint_dim") tdt_joint_dim = static_cast(std::stoul(value)); else if (key == "tdt_num_durations") tdt_num_durations = static_cast(std::stoul(value)); else if (key == "tdt_blank_id") tdt_blank_id = static_cast(std::stoul(value)); + else if (key == "rnnt_joint_dim") rnnt_joint_dim = static_cast(std::stoul(value)); + else if (key == "rnnt_blank_id") rnnt_blank_id = static_cast(std::stoul(value)); + else if (key == "prompt_dim") prompt_dim = static_cast(std::stoul(value)); + else if (key == "default_prompt_id") default_prompt_id = static_cast(std::stoul(value)); + else if (key == "max_symbols_per_step") max_symbols_per_step = static_cast(std::stoul(value)); + else if (key == "prompt_dictionary") prompt_dictionary = parse_config_string_uint_map(value); else if (key == "tdt_durations") { tdt_durations.clear(); std::stringstream ss(value); diff --git a/cactus-engine/src/model_npu.cpp b/cactus-engine/src/model_npu.cpp index 8c421dee1..42f812a36 100644 --- a/cactus-engine/src/model_npu.cpp +++ b/cactus-engine/src/model_npu.cpp @@ -5,21 +5,113 @@ #include #include #include +#include #include namespace cactus { namespace engine { +static void prewarm_npu_audio_encoder(npu::NPUEncoder& encoder) { + const std::vector input_shape = encoder.get_input_shape(); + const std::vector output_shape = encoder.get_output_shape(); + if (input_shape.empty() || output_shape.empty()) return; + size_t input_elems = 1; + for (int dim : input_shape) { + if (dim <= 0) return; + input_elems *= static_cast(dim); + } + size_t output_elems = 1; + for (int dim : output_shape) { + if (dim <= 0) return; + output_elems *= static_cast(dim); + } + std::vector<__fp16> input(input_elems, __fp16(0)); + std::vector<__fp16> output(output_elems, __fp16(0)); + if (!encoder.preallocate(input_shape, "x", "encoded")) return; + if (encoder.encode(input.data(), output.data(), input_shape, "x", "encoded") == 0) { + CACTUS_LOG_WARN("model", "NPU audio encoder warmup failed"); + } +} + bool Model::load_npu_audio_encoder(const std::string& model_path, const std::string& compute_units) { auto encoder = npu::create_encoder(); if (!encoder) return false; if (!encoder->load(model_path, compute_units)) return false; if (!encoder->is_available()) return false; + prewarm_npu_audio_encoder(*encoder); npu_audio_encoder_ = std::move(encoder); CACTUS_LOG_INFO("model", "NPU audio encoder loaded from: " << model_path); return true; } +bool Model::load_npu_audio_encoders(const std::vector& model_paths, const std::string& compute_units) { + npu_audio_encoders_.clear(); + bool loaded = false; + for (const std::string& model_path : model_paths) { + auto encoder = npu::create_encoder(); + if (!encoder) continue; + if (!encoder->load(model_path, compute_units) || !encoder->is_available()) { + CACTUS_LOG_WARN("model", "NPU audio encoder load failed for " << model_path << "; skipping"); + continue; + } + npu_audio_encoders_.push_back(std::move(encoder)); + CACTUS_LOG_INFO("model", "NPU audio encoder loaded from: " << model_path); + loaded = true; + } + std::sort(npu_audio_encoders_.begin(), npu_audio_encoders_.end(), + [](const std::unique_ptr& a, const std::unique_ptr& b) { + const auto as = a ? a->get_input_shape() : std::vector{}; + const auto bs = b ? b->get_input_shape() : std::vector{}; + const int af = as.size() >= 2 ? as[1] : 0; + const int bf = bs.size() >= 2 ? bs[1] : 0; + return af < bf; + }); + if (npu_audio_prewarm_frames_ > 0) { + size_t feature_dim = 0; + for (const auto& encoder : npu_audio_encoders_) { + const auto shape = encoder ? encoder->get_input_shape() : std::vector{}; + if (shape.size() >= 3 && shape[2] > 0) { + feature_dim = static_cast(shape[2]); + break; + } + } + if (feature_dim > 0) { + npu::NPUEncoder* encoder = select_npu_audio_encoder(feature_dim, npu_audio_prewarm_frames_); + const auto selected_shape = encoder ? encoder->get_input_shape() : std::vector{}; + const size_t selected_frames = selected_shape.size() >= 2 && selected_shape[1] > 0 + ? static_cast(selected_shape[1]) : 0; + for (const auto& candidate : npu_audio_encoders_) { + if (!candidate) continue; + const auto shape = candidate->get_input_shape(); + if (shape.size() < 2 || shape[1] <= 0) continue; + if (selected_frames > 0 && static_cast(shape[1]) <= selected_frames) { + prewarm_npu_audio_encoder(*candidate); + } + } + } + } + return loaded; +} + +npu::NPUEncoder* Model::select_npu_audio_encoder(size_t feature_dim, size_t frame_count) const { + if (frame_count == 0) return nullptr; + npu::NPUEncoder* best = nullptr; + size_t best_frames = std::numeric_limits::max(); + auto consider = [&](npu::NPUEncoder* encoder) { + if (!encoder || !encoder->is_available()) return; + const std::vector shape = encoder->get_input_shape(); + if (shape.size() < 3 || shape[1] <= 0 || shape[2] <= 0) return; + const size_t frames = static_cast(shape[1]); + const size_t mels = static_cast(shape[2]); + if (mels != feature_dim || frame_count > frames || frames >= best_frames) return; + best = encoder; + best_frames = frames; + }; + for (const auto& encoder : npu_audio_encoders_) consider(encoder.get()); + consider(npu_audio_encoder_.get()); + return best; +} + bool Model::load_npu_vision_encoder(const std::string& model_path) { auto encoder = npu::create_encoder(); if (!encoder) return false; diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm index 99638e2e0..8499859f8 100644 --- a/cactus-engine/src/npu_ane.mm +++ b/cactus-engine/src/npu_ane.mm @@ -527,7 +527,7 @@ static size_t copyMLArrayToFP16(MLMultiArray* array, __fp16* output) { return true; }; - if (false && is_contiguous()) { + if (is_contiguous()) { if (array.dataType == MLMultiArrayDataTypeFloat16) { if (output != (__fp16*)array.dataPointer) { memcpy(output, array.dataPointer, count * sizeof(__fp16)); @@ -616,7 +616,7 @@ static void copyFP16ToMLArray(const __fp16* data, size_t count, MLMultiArray* ar return true; }; - if (false && is_contiguous()) { + if (is_contiguous()) { if (array.dataType == MLMultiArrayDataTypeFloat16) { memcpy(array.dataPointer, data, count * sizeof(__fp16)); } else { diff --git a/cactus-engine/src/stream.cpp b/cactus-engine/src/stream.cpp index 182d81e4e..fb4b3a542 100644 --- a/cactus-engine/src/stream.cpp +++ b/cactus-engine/src/stream.cpp @@ -40,13 +40,16 @@ struct StreamTranscribe { CactusModelHandle* model = nullptr; std::string options_json; bool is_parakeet = false; + bool is_nemotron = false; std::vector samples; + std::vector full_samples; size_t samples_decoded_up_to = 0; size_t silence_run = 0; bool cold_restart = false; cactus::engine::Model::ParakeetTdtStreamState pstate; + cactus::engine::Model::RnntStreamState rstate; std::vector committed_tokens; std::string emitted_text; std::string previous_pending; @@ -87,6 +90,18 @@ std::string strip_annotations(const std::string& text) { return trim(out); } +std::string strip_language_tags(const std::string& text) { + std::string out = std::regex_replace(text, std::regex(R"(\s*<[a-z]{2}(?:-[A-Z]{2})?>\s*)"), " "); + out = std::regex_replace(out, std::regex(R"(\s+)"), " "); + return trim(out); +} + +std::string target_lang_option(const std::string& json) { + std::string target = json_string_field(json, "target_lang"); + if (target.empty()) target = json_string_field(json, "language"); + return target.empty() ? "auto" : target; +} + std::vector parse_segments(const std::string& json) { std::vector segs; for (const std::string& obj : split_json_array(json_array_field(json, "segments"))) { @@ -119,14 +134,17 @@ std::vector whisper_transcribe(StreamTranscribe* s, const std return parse_segments(json); } -std::vector window_features(std::vector window, size_t mel_bins) { +std::vector window_features(std::vector window, size_t mel_bins, bool is_nemotron) { if (window.empty()) return {}; auto cfg = cactus::audio::get_parakeet_spectrogram_config(); + if (is_nemotron) cfg.mel_floor_additive = true; const size_t waveform_samples = window.size(); cactus::audio::apply_preemphasis(window, 0.97f); + const int mel_norm_type = is_nemotron ? 1 : 0; + const int mel_scale_type = is_nemotron ? 2 : 0; std::vector features = cactus::audio::compute_spectrogram_graph( - window, cfg, mel_bins, 0.0f, 8000.0f, cactus::audio::WHISPER_SAMPLE_RATE, 0, 0); - cactus::audio::normalize_parakeet_log_mel(features, mel_bins); + window, cfg, mel_bins, 0.0f, 8000.0f, cactus::audio::WHISPER_SAMPLE_RATE, mel_norm_type, mel_scale_type); + if (!is_nemotron) cactus::audio::normalize_parakeet_log_mel(features, mel_bins); size_t valid_frames = waveform_samples / cfg.hop_length; if (valid_frames == 0) valid_frames = 1; cactus::audio::trim_mel_frames(features, mel_bins, valid_frames); @@ -140,33 +158,47 @@ std::string parakeet_decode_window(StreamTranscribe* s, size_t window_start, siz auto* model = s->model->model.get(); const size_t mel_bins = std::max(1, static_cast(model->get_config().num_mel_bins)); std::vector features = window_features( - std::vector(s->samples.begin() + window_start, s->samples.begin() + window_end), mel_bins); + std::vector(s->samples.begin() + window_start, s->samples.begin() + window_end), + mel_bins, + s->is_nemotron); if (features.empty()) return ""; - s->pstate.time_index = decode_start_frame; + if (s->is_nemotron) s->rstate.time_index = decode_start_frame; + else s->pstate.time_index = decode_start_frame; const auto t0 = std::chrono::steady_clock::now(); - std::vector tokens = model->transcribe_parakeet_tdt( - features, &s->pstate, is_final, is_final ? 0 : decode_end_frame); + std::vector tokens = s->is_nemotron + ? model->transcribe_nemotron_asr( + features, target_lang_option(s->options_json), &s->rstate, is_final, + is_final ? 0 : decode_end_frame) + : model->transcribe_parakeet_tdt( + features, &s->pstate, is_final, is_final ? 0 : decode_end_frame); stats.total_time_ms += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); - stats.decode_tokens += s->pstate.decoded_tokens; + stats.decode_tokens += s->is_nemotron ? s->rstate.decoded_tokens : s->pstate.decoded_tokens; auto* tokenizer = model->get_tokenizer(); if (!tokenizer) return ""; + const bool strip_lang_tags = s->is_nemotron && json_bool_field_or(s->options_json, "strip_lang_tags", true); s->committed_tokens.insert(s->committed_tokens.end(), tokens.begin(), tokens.end()); std::string full = tokenizer->decode(s->committed_tokens); + if (strip_lang_tags) full = strip_language_tags(full); std::string delta = full.size() > s->emitted_text.size() ? full.substr(s->emitted_text.size()) : std::string(); s->emitted_text = full; - if (pending_text && !s->pstate.pending.empty()) { + const std::vector& pending_tokens = s->is_nemotron ? s->rstate.pending : s->pstate.pending; + if (pending_text && !pending_tokens.empty()) { std::vector combined = s->committed_tokens; - combined.insert(combined.end(), s->pstate.pending.begin(), s->pstate.pending.end()); + combined.insert(combined.end(), pending_tokens.begin(), pending_tokens.end()); std::string with_pending = tokenizer->decode(combined); - if (with_pending.size() > full.size()) *pending_text = with_pending.substr(full.size()); + if (strip_lang_tags) with_pending = strip_language_tags(with_pending); + if (with_pending.size() > full.size()) { + *pending_text = with_pending.substr(full.size()); + } } - if (!is_final && !tokens.empty() && s->pstate.confirmed_sec > 0.0f) { - s->samples_decoded_up_to = window_start + static_cast(s->pstate.confirmed_sec * kSampleRateF); + const float confirmed_sec = s->is_nemotron ? s->rstate.confirmed_sec : s->pstate.confirmed_sec; + if (!is_final && !tokens.empty() && confirmed_sec > 0.0f) { + s->samples_decoded_up_to = window_start + static_cast(confirmed_sec * kSampleRateF); } return delta; } @@ -197,7 +229,10 @@ std::string parakeet_process(StreamTranscribe* s, std::string& pending, StreamSt const size_t window_end = decode_to + kRightContextSamples; const size_t decode_start_frame = (s->samples_decoded_up_to - window_start) / spf; const size_t decode_end_frame = decode_start_frame + (decode_to - s->samples_decoded_up_to) / spf; - if (cold) s->pstate = {}; + if (cold) { + if (s->is_nemotron) s->rstate = {}; + else s->pstate = {}; + } std::string pend; const size_t prev_cursor = s->samples_decoded_up_to; @@ -205,9 +240,10 @@ std::string parakeet_process(StreamTranscribe* s, std::string& pending, StreamSt confirmed += parakeet_decode_window(s, window_start, window_end, decode_start_frame, decode_end_frame, false, &pend, stats); pending = pend; - if (s->pstate.decoded_tokens > 0) { s->cold_restart = false; s->silence_run = 0; } + const size_t decoded_tokens = s->is_nemotron ? s->rstate.decoded_tokens : s->pstate.decoded_tokens; + if (decoded_tokens > 0) { s->cold_restart = false; s->silence_run = 0; } if (s->samples_decoded_up_to > prev_cursor) continue; - if (s->pstate.decoded_tokens == 0) { + if (decoded_tokens == 0) { s->silence_run += decode_to - s->samples_decoded_up_to; s->samples_decoded_up_to = decode_to; if (s->silence_run >= kSilenceResetSamples) s->cold_restart = true; @@ -236,16 +272,31 @@ std::string parakeet_process(StreamTranscribe* s, std::string& pending, StreamSt std::string parakeet_flush(StreamTranscribe* s, StreamStats& stats) { std::lock_guard lock(s->model->model_mutex); + std::string previous_emitted; + if (s->is_nemotron) { + previous_emitted = s->emitted_text; + if (!s->full_samples.empty()) s->samples = s->full_samples; + s->samples_decoded_up_to = 0; + s->rstate = {}; + s->committed_tokens.clear(); + s->emitted_text.clear(); + } const size_t total = s->samples.size(); if (total <= s->samples_decoded_up_to) return ""; const size_t spf = parakeet_spf(s); const bool cold = s->samples_decoded_up_to == 0 || s->cold_restart; - if (cold) s->pstate = {}; + if (cold) { + if (s->is_nemotron) s->rstate = {}; + else s->pstate = {}; + } const size_t window_start = cold ? s->samples_decoded_up_to : (s->samples_decoded_up_to > kLeftContextSamples ? s->samples_decoded_up_to - kLeftContextSamples : 0); const size_t decode_start_frame = (s->samples_decoded_up_to - window_start) / spf; std::string confirmed = parakeet_decode_window(s, window_start, total, decode_start_frame, 0, true, nullptr, stats); + if (s->is_nemotron && !previous_emitted.empty() && confirmed.rfind(previous_emitted, 0) == 0) { + confirmed = trim(confirmed.substr(previous_emitted.size())); + } stats.finalize(); return confirmed; } @@ -334,14 +385,16 @@ cactus_stream_transcribe_t cactus_stream_transcribe_start(cactus_model_t model, const auto type = handle->model->get_config().model_type; const bool is_whisper = type == cactus::engine::Config::ModelType::WHISPER; const bool is_parakeet = type == cactus::engine::Config::ModelType::PARAKEET_TDT; - if (!is_whisper && !is_parakeet) { - last_error_message = "stream_transcribe_start: only Whisper and Parakeet models support streaming"; + const bool is_nemotron = type == cactus::engine::Config::ModelType::NEMOTRON_ASR; + if (!is_whisper && !is_parakeet && !is_nemotron) { + last_error_message = "stream_transcribe_start: only Whisper, Parakeet, and Nemotron ASR models support streaming"; CACTUS_LOG_ERROR("stream_transcribe_start", last_error_message); return nullptr; } auto s = std::make_unique(); s->model = handle; s->is_parakeet = is_parakeet; + s->is_nemotron = is_nemotron; if (options_json && options_json[0] != '\0') s->options_json = options_json; CACTUS_LOG_INFO("stream_transcribe_start", "streaming session opened"); return static_cast(s.release()); @@ -366,11 +419,13 @@ int cactus_stream_transcribe_process(cactus_stream_transcribe_t stream, if (pcm_buffer && pcm_buffer_size >= sizeof(int16_t)) { std::vector news = cactus::audio::pcm_buffer_to_float_samples(pcm_buffer, pcm_buffer_size); s->samples.insert(s->samples.end(), news.begin(), news.end()); + if (s->is_nemotron) s->full_samples.insert(s->full_samples.end(), news.begin(), news.end()); } StreamStats stats; std::string pending; - std::string confirmed = s->is_parakeet ? parakeet_process(s, pending, stats) - : whisper_process(s, pending, stats); + std::string confirmed = (s->is_parakeet || s->is_nemotron) + ? parakeet_process(s, pending, stats) + : whisper_process(s, pending, stats); return write_result(response_buffer, buffer_size, confirmed, pending, stats); } catch (const std::exception& e) { last_error_message = std::string("stream_transcribe_process: ") + e.what(); @@ -392,7 +447,12 @@ int cactus_stream_transcribe_stop(cactus_stream_transcribe_t stream, int result = 0; try { StreamStats stats; - std::string confirmed = s->is_parakeet ? parakeet_flush(s, stats) : whisper_flush(s, stats); + std::string confirmed = (s->is_parakeet || s->is_nemotron) + ? parakeet_flush(s, stats) + : whisper_flush(s, stats); + if (s->is_nemotron && json_bool_field_or(s->options_json, "strip_lang_tags", true)) { + confirmed = strip_language_tags(confirmed); + } result = write_result(response_buffer, buffer_size, confirmed, "", stats); } catch (const std::exception& e) { last_error_message = std::string("stream_transcribe_stop: ") + e.what(); diff --git a/cactus-engine/src/transcribe.cpp b/cactus-engine/src/transcribe.cpp index cafd98c67..89db686ed 100644 --- a/cactus-engine/src/transcribe.cpp +++ b/cactus-engine/src/transcribe.cpp @@ -9,6 +9,7 @@ #include #include #include +#include using namespace cactus::ffi; @@ -50,6 +51,18 @@ std::vector parse_whisper_timestamp_segments( return segments; } +void strip_language_tags(std::string& text) { + text = std::regex_replace(text, std::regex(R"(\s*<[a-z]{2}(?:-[A-Z]{2})?>\s*)"), " "); + text = std::regex_replace(text, std::regex(R"(\s+)"), " "); + size_t start = text.find_first_not_of(" \t\r\n"); + if (start == std::string::npos) { + text.clear(); + return; + } + size_t end = text.find_last_not_of(" \t\r\n"); + text = text.substr(start, end - start + 1); +} + } // namespace extern "C" { @@ -102,14 +115,18 @@ CACTUS_FFI_EXPORT int cactus_preprocess_audio_features( auto prepared = cactus::audio::preprocess_audio_for_gemma4(std::move(audio_samples), cfg); features = std::move(prepared.features); frames = prepared.num_frames; - } else if (lowered.find("parakeet") != std::string::npos) { + } else if (lowered.find("parakeet") != std::string::npos || lowered.find("nemotron") != std::string::npos) { + const bool is_nemotron_frontend = lowered.find("nemotron") != std::string::npos; auto cfg = cactus::audio::get_parakeet_spectrogram_config(); + if (is_nemotron_frontend) cfg.mel_floor_additive = true; const size_t waveform_samples = audio_samples.size(); cactus::audio::apply_preemphasis(audio_samples, 0.97f); + const int mel_norm_type = is_nemotron_frontend ? 1 : 0; + const int mel_scale_type = is_nemotron_frontend ? 2 : 0; features = cactus::audio::compute_spectrogram_graph( audio_samples, cfg, bins, 0.0f, 8000.0f, - cactus::audio::WHISPER_SAMPLE_RATE, 0, 0); - cactus::audio::normalize_parakeet_log_mel(features, bins); + cactus::audio::WHISPER_SAMPLE_RATE, mel_norm_type, mel_scale_type); + if (!is_nemotron_frontend) cactus::audio::normalize_parakeet_log_mel(features, bins); size_t valid_frames = waveform_samples / cfg.hop_length; if (valid_frames == 0) valid_frames = 1; cactus::audio::trim_mel_frames(features, bins, valid_frames); @@ -173,8 +190,9 @@ int cactus_transcribe( const auto model_type = handle->model->get_config().model_type; const bool is_whisper = model_type == cactus::engine::Config::ModelType::WHISPER; const bool is_parakeet = model_type == cactus::engine::Config::ModelType::PARAKEET_TDT; + const bool is_nemotron = model_type == cactus::engine::Config::ModelType::NEMOTRON_ASR; - if (!is_whisper && !is_parakeet) { + if (!is_whisper && !is_parakeet && !is_nemotron) { const uint8_t* pcm_data = pcm_buffer; size_t pcm_size = pcm_buffer_size; std::vector file_pcm; @@ -240,15 +258,18 @@ int cactus_transcribe( const size_t mel_bins = std::max(1, static_cast(handle->model->get_config().num_mel_bins)); std::vector audio_features; std::vector handoff_audio_samples; - if (is_parakeet) { + if (is_parakeet || is_nemotron) { auto cfg = cactus::audio::get_parakeet_spectrogram_config(); + if (is_nemotron) cfg.mel_floor_additive = true; const size_t waveform_samples = audio_samples.size(); handoff_audio_samples = audio_samples; cactus::audio::apply_preemphasis(audio_samples, 0.97f); + const int mel_norm_type = is_nemotron ? 1 : 0; + const int mel_scale_type = is_nemotron ? 2 : 0; audio_features = cactus::audio::compute_spectrogram_graph( audio_samples, cfg, mel_bins, 0.0f, 8000.0f, - cactus::audio::WHISPER_SAMPLE_RATE, 0, 0); - cactus::audio::normalize_parakeet_log_mel(audio_features, mel_bins); + cactus::audio::WHISPER_SAMPLE_RATE, mel_norm_type, mel_scale_type); + if (is_parakeet) cactus::audio::normalize_parakeet_log_mel(audio_features, mel_bins); size_t valid_frames = waveform_samples / cfg.hop_length; if (valid_frames == 0) valid_frames = 1; cactus::audio::trim_mel_frames(audio_features, mel_bins, valid_frames); @@ -351,6 +372,25 @@ int cactus_transcribe( callback(piece.c_str(), tok, user_data); } } + } else if (is_nemotron) { + std::string target_lang = json_string_field(options_json ? options_json : "", "target_lang"); + if (target_lang.empty()) target_lang = json_string_field(options_json ? options_json : "", "language"); + if (target_lang.empty()) target_lang = "auto"; + generated_tokens = handle->model->transcribe_nemotron_asr( + audio_features, target_lang, nullptr, true, 0, &handle->should_stop); + auto t_first = std::chrono::high_resolution_clock::now(); + time_to_first_token = + std::chrono::duration_cast(t_first - start_time).count() / 1000.0; + final_text = tokenizer->decode(generated_tokens); + if (json_bool_field_or(options_json ? options_json : "", "strip_lang_tags", true)) { + strip_language_tags(final_text); + } + if (callback) { + for (uint32_t tok : generated_tokens) { + std::string piece = tokenizer->decode({tok}); + callback(piece.c_str(), tok, user_data); + } + } } else if (is_whisper) { generated_tokens = handle->model->transcribe_whisper_seq2seq( audio_features, diff --git a/cactus-engine/src/utils.h b/cactus-engine/src/utils.h index 95e2368da..03d88e3bf 100644 --- a/cactus-engine/src/utils.h +++ b/cactus-engine/src/utils.h @@ -610,6 +610,17 @@ inline std::string json_string_field(const std::string& json, const std::string& return {}; } +inline bool json_bool_field_or(const std::string& json, const std::string& key, bool fallback) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return fallback; + size_t i = pos + pattern.size(); + while (i < json.size() && std::isspace(static_cast(json[i]))) i++; + if (json.compare(i, 4, "true") == 0) return true; + if (json.compare(i, 5, "false") == 0) return false; + return fallback; +} + inline std::string json_array_field(const std::string& json, const std::string& key) { std::string pattern = "\"" + key + "\":"; size_t pos = json.find(pattern); diff --git a/cactus-engine/tests/transcribe.cpp b/cactus-engine/tests/transcribe.cpp index 84d22b0e9..0038b9497 100644 --- a/cactus-engine/tests/transcribe.cpp +++ b/cactus-engine/tests/transcribe.cpp @@ -96,10 +96,6 @@ static std::string whisper_prompt(const std::string& model_path, const std::stri return ""; } -static void print_token(const char* token, uint32_t, void*) { - std::cout << token << std::flush; -} - static int transcribe_file(cactus_model_t model, const std::string& audio_path, const std::string& model_path, const std::string& language) { std::string prompt = whisper_prompt(model_path, language); @@ -110,7 +106,7 @@ static int transcribe_file(cactus_model_t model, const std::string& audio_path, int rc = cactus_transcribe( model, audio_path.c_str(), prompt.empty() ? nullptr : prompt.c_str(), response_buffer.data(), response_buffer.size(), - options.c_str(), print_token, nullptr, nullptr, 0); + options.c_str(), nullptr, nullptr, nullptr, 0); double total_s = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); if (rc < 0) { std::cerr << "\nTranscription failed: " << cactus_get_last_error() << "\n"; @@ -118,6 +114,7 @@ static int transcribe_file(cactus_model_t model, const std::string& audio_path, } std::string json(response_buffer.data()); + std::cout << json_string_value(json, "response"); double model_ms = json_number_value(json, "total_time_ms"); std::ostringstream stats; stats << std::fixed << std::setprecision(2) << "[processed: " << total_s << "s"; @@ -203,7 +200,7 @@ static int run_live_transcription(cactus_model_t model, const std::string& langu std::string options = "{\"language\":\"" + language + "\"}"; cactus_stream_transcribe_t stream = cactus_stream_transcribe_start(model, options.c_str()); if (!stream) { - std::cerr << "Failed to start streaming (live mode needs a Whisper or Parakeet TDT model): " + std::cerr << "Failed to start streaming (live mode needs a Whisper, Parakeet TDT, or Nemotron ASR model): " << cactus_get_last_error() << "\n"; SDL_CloseAudioDevice(device); SDL_Quit(); @@ -359,7 +356,8 @@ int main(int argc, char** argv) { std::string model_path = argv[1]; std::string audio_file; - std::string language = "en"; + std::string language; + bool language_explicit = false; for (int i = 2; i < argc; ++i) { std::string arg = argv[i]; if (arg == "-h" || arg == "--help") { @@ -371,6 +369,7 @@ int main(int argc, char** argv) { return 1; } language = argv[++i]; + language_explicit = true; } else if (!arg.empty() && arg[0] == '-') { std::cerr << "Error: unknown option '" << arg << "'\n"; print_usage(argv[0]); @@ -379,6 +378,11 @@ int main(int argc, char** argv) { audio_file = arg; } } + if (!language_explicit) { + std::string lowered = model_path; + std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) { return (char)std::tolower(c); }); + language = lowered.find("nemotron") != std::string::npos ? "auto" : "en"; + } std::cout << "Loading model from " << model_path << "...\n"; cactus_model_t model = cactus_init(model_path.c_str(), nullptr, false); diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index b7fff474d..27aa1967f 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -52,6 +52,51 @@ namespace { return filename; } + void attach_mapped_weight_metadata(BufferDesc& buffer, const GraphFile::MappedFile& mapped_file) { + Precision precision = mapped_file.precision(); + const auto& shape = mapped_file.shape(); + if (PrecisionTraits::is_cq(precision) && mapped_file.group_size() > 0) { + buffer.group_size = mapped_file.group_size(); + buffer.num_groups = mapped_file.num_groups(); + + const char* scales_base = static_cast(mapped_file.scales_data()); + uint32_t bits = PrecisionTraits::cq_bits(precision); + uint32_t cb_size = 1u << bits; + uint32_t gs = static_cast(mapped_file.group_size()); + uint32_t K = gs * static_cast(mapped_file.num_groups()); + uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); + + size_t off = 0; + buffer.cq_codebook = reinterpret_cast(scales_base + off); + off += cb_size * sizeof(__fp16); + buffer.cq_input_scale = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_norms = reinterpret_cast(scales_base + off); + off += static_cast(N) * mapped_file.num_groups() * sizeof(__fp16); + + if (mapped_file.is_orthogonal_rotation()) { + buffer.cq_rotation = reinterpret_cast(scales_base + off); + buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else { + buffer.cq_left_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_right_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_permutation = reinterpret_cast(scales_base + off); + buffer.cq_flags = 0; + } + if (mapped_file.is_interleaved_4row()) { + buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + } + } else if (precision == Precision::INT8 && mapped_file.group_size() > 0) { + buffer.group_size = mapped_file.group_size(); + buffer.num_groups = mapped_file.num_groups(); + buffer.set_activation_scales(const_cast(mapped_file.scales_data()), shape.empty() ? 0 : shape[0]); + } + } + inline void write_u32(std::ostream& out, uint32_t v) { out.write(reinterpret_cast(&v), sizeof(v)); } @@ -385,46 +430,7 @@ size_t CactusGraph::mmap_embeddings(const std::string& filename) { set_external_input(node_id, const_cast(mapped_file->data()), precision); auto& buffer = nodes_[node_index_map_.at(node_id)]->output_buffer; - if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { - buffer.group_size = mapped_file->group_size(); - buffer.num_groups = mapped_file->num_groups(); - - const char* scales_base = static_cast(mapped_file->scales_data()); - uint32_t bits = PrecisionTraits::cq_bits(precision); - uint32_t cb_size = 1u << bits; - uint32_t gs = static_cast(mapped_file->group_size()); - uint32_t K = gs * static_cast(mapped_file->num_groups()); - uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); - - size_t off = 0; - buffer.cq_codebook = reinterpret_cast(scales_base + off); - off += cb_size * sizeof(__fp16); - buffer.cq_input_scale = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_norms = reinterpret_cast(scales_base + off); - off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); - - if (mapped_file->is_orthogonal_rotation()) { - buffer.cq_rotation = reinterpret_cast(scales_base + off); - buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; - } else { - buffer.cq_left_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_right_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_permutation = reinterpret_cast(scales_base + off); - buffer.cq_flags = 0; - } - if (mapped_file->is_interleaved_4row()) { - buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; - } - } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) { - buffer.group_size = mapped_file->group_size(); - buffer.num_groups = mapped_file->num_groups(); - buffer.set_activation_scales(const_cast(mapped_file->scales_data()), shape.empty() ? 0 : shape[0]); - } + attach_mapped_weight_metadata(buffer, *mapped_file); size_t file_idx = mapped_files_.size(); mapped_files_.push_back(std::move(mapped_file)); @@ -448,44 +454,8 @@ size_t CactusGraph::mmap_weights(const std::string& filename) { size_t node_id = input(shape, precision); set_external_input(node_id, const_cast(mapped_file->data()), precision); - if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { - auto& buffer = nodes_[node_index_map_.at(node_id)]->output_buffer; - buffer.group_size = mapped_file->group_size(); - buffer.num_groups = mapped_file->num_groups(); - - - const char* scales_base = static_cast(mapped_file->scales_data()); - uint32_t bits = PrecisionTraits::cq_bits(precision); - uint32_t cb_size = 1u << bits; - uint32_t gs = static_cast(mapped_file->group_size()); - uint32_t K = gs * static_cast(mapped_file->num_groups()); - uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); - - size_t off = 0; - buffer.cq_codebook = reinterpret_cast(scales_base + off); - off += cb_size * sizeof(__fp16); - buffer.cq_input_scale = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_norms = reinterpret_cast(scales_base + off); - off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); - - if (mapped_file->is_orthogonal_rotation()) { - buffer.cq_rotation = reinterpret_cast(scales_base + off); - buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; - } else { - buffer.cq_left_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_right_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_permutation = reinterpret_cast(scales_base + off); - buffer.cq_flags = 0; - } - if (mapped_file->is_interleaved_4row()) { - buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; - } - } + auto& buffer = nodes_[node_index_map_.at(node_id)]->output_buffer; + attach_mapped_weight_metadata(buffer, *mapped_file); size_t file_idx = mapped_files_.size(); mapped_files_.push_back(std::move(mapped_file)); @@ -533,46 +503,7 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) buffer.cq_rotation = nullptr; buffer.cq_flags = 0; - if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { - buffer.group_size = mapped_file->group_size(); - buffer.num_groups = mapped_file->num_groups(); - - const char* scales_base = static_cast(mapped_file->scales_data()); - uint32_t bits = PrecisionTraits::cq_bits(precision); - uint32_t cb_size = 1u << bits; - uint32_t gs = static_cast(mapped_file->group_size()); - uint32_t K = gs * static_cast(mapped_file->num_groups()); - uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); - - size_t off = 0; - buffer.cq_codebook = reinterpret_cast(scales_base + off); - off += cb_size * sizeof(__fp16); - buffer.cq_input_scale = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_norms = reinterpret_cast(scales_base + off); - off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); - - if (mapped_file->is_orthogonal_rotation()) { - buffer.cq_rotation = reinterpret_cast(scales_base + off); - buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; - } else { - buffer.cq_left_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_right_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_permutation = reinterpret_cast(scales_base + off); - buffer.cq_flags = 0; - } - if (mapped_file->is_interleaved_4row()) { - buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; - } - } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) { - buffer.group_size = mapped_file->group_size(); - buffer.num_groups = mapped_file->num_groups(); - buffer.set_activation_scales(const_cast(mapped_file->scales_data()), shape.empty() ? 0 : shape[0]); - } + attach_mapped_weight_metadata(buffer, *mapped_file); size_t file_idx = mapped_files_.size(); mapped_files_.push_back(std::move(mapped_file)); diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index 23afc2f20..cbb63e4d6 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -569,7 +569,7 @@ void compute_rel_pos_bias_node(GraphNode& node, const std::vector 0) { + node.output_buffer.group_size = input_buffer.group_size; + node.output_buffer.num_groups = input_buffer.num_groups; + auto* scales = static_cast(input_buffer.activation_scales_data); + node.output_buffer.set_activation_scales(scales + slice_start * input_buffer.num_groups * sizeof(__fp16), slice_length); + } return; } diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp index 9173b7324..f090e6963 100644 --- a/cactus-graph/src/param_io.cpp +++ b/cactus-graph/src/param_io.cpp @@ -142,6 +142,28 @@ enum class ParamField : uint32_t { CachedVScalesPtr, MaxCacheSeqLen, CacheSinkSize, + HopLength, + Power, + Center, + MelFloor, + Dither, + PreemphasisCoef, + RemoveDcOffset, + LogMelMode, + PadModeType, + NumMelFilters, + SamplingRate, + MinFrequency, + MaxFrequency, + MelNormType, + MelScaleType, + PatchSize, + RescaleFactor, + ImageMean, + ImageStd, + TargetWidth, + TargetHeight, + ImageChannels, }; enum class FieldPersistence { @@ -167,6 +189,7 @@ const Schema& op_schema(OpType op_type) { {OpType::SCALAR_DIVIDE, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::SCALAR_NOT_EQUAL, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::CLAMP, {{ParamField::Scalar, FieldPersistence::Persistent}, {ParamField::Scale, FieldPersistence::Persistent}}}, + {OpType::LEAKY_RELU, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::SOFTMAX, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::SUM, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::MEAN, {{ParamField::Axis, FieldPersistence::Persistent}}}, @@ -177,6 +200,7 @@ const Schema& op_schema(OpType op_type) { {OpType::INDEX, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::IndexValue, FieldPersistence::Persistent}}}, {OpType::CONCAT, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::CAT, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::GLU, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::VIEW, {{ParamField::NewShape, FieldPersistence::Persistent}}}, {OpType::RESHAPE, {{ParamField::NewShape, FieldPersistence::Persistent}}}, {OpType::FLATTEN, {{ParamField::NewShape, FieldPersistence::Persistent}}}, @@ -230,15 +254,49 @@ const Schema& op_schema(OpType op_type) { {OpType::GATED_DELTANET_DECODE, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::NumKvHeads, FieldPersistence::Persistent}}}, {OpType::GATED_DELTANET_PREFILL, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::NumKvHeads, FieldPersistence::Persistent}, {ParamField::ChunkSize, FieldPersistence::Persistent}}}, {OpType::STFT, {{ParamField::Stride, FieldPersistence::Persistent}, {ParamField::NumFftBins, FieldPersistence::Persistent}}}, + {OpType::MEL_FILTER_BANK, { + {ParamField::NumMelFilters, FieldPersistence::Persistent}, + {ParamField::MinFrequency, FieldPersistence::Persistent}, + {ParamField::MaxFrequency, FieldPersistence::Persistent}, + {ParamField::SamplingRate, FieldPersistence::Persistent}, + {ParamField::MelNormType, FieldPersistence::Persistent}, + {ParamField::MelScaleType, FieldPersistence::Persistent}, + }}, + {OpType::SPECTROGRAM, { + {ParamField::NumFftBins, FieldPersistence::Persistent}, + {ParamField::HopLength, FieldPersistence::Persistent}, + {ParamField::Stride, FieldPersistence::Persistent}, + {ParamField::Power, FieldPersistence::Persistent}, + {ParamField::Center, FieldPersistence::Persistent}, + {ParamField::PadModeType, FieldPersistence::Persistent}, + {ParamField::MelFloor, FieldPersistence::Persistent}, + {ParamField::LogMelMode, FieldPersistence::Persistent}, + {ParamField::Dither, FieldPersistence::Persistent}, + {ParamField::PreemphasisCoef, FieldPersistence::Persistent}, + {ParamField::RemoveDcOffset, FieldPersistence::Persistent}, + }}, + {OpType::IMAGE_PREPROCESS, { + {ParamField::DstWidth, FieldPersistence::Persistent}, + {ParamField::DstHeight, FieldPersistence::Persistent}, + {ParamField::TargetWidth, FieldPersistence::Persistent}, + {ParamField::TargetHeight, FieldPersistence::Persistent}, + {ParamField::PatchSize, FieldPersistence::Persistent}, + {ParamField::ImageChannels, FieldPersistence::Persistent}, + {ParamField::RescaleFactor, FieldPersistence::Persistent}, + {ParamField::ImageMean, FieldPersistence::Persistent}, + {ParamField::ImageStd, FieldPersistence::Persistent}, + }}, {OpType::BILINEAR_INTERPOLATION, {{ParamField::DstHeight, FieldPersistence::Persistent}, {ParamField::DstWidth, FieldPersistence::Persistent}, {ParamField::AlignCorners, FieldPersistence::Persistent}}}, {OpType::SCATTER_TOPK, {{ParamField::NumClasses, FieldPersistence::Persistent}}}, {OpType::ALTUP_PREDICT, {{ParamField::NumAltupInputs, FieldPersistence::Persistent}}}, {OpType::ALTUP_CORRECT, {{ParamField::NumAltupInputs, FieldPersistence::Persistent}}}, + {OpType::GAUSSIAN_TOPK, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::MAXPOOL1D, {{ParamField::KernelSize, FieldPersistence::Persistent}, {ParamField::Stride, FieldPersistence::Persistent}}}, {OpType::CONV1D_CAUSAL, {{ParamField::Dilation, FieldPersistence::Persistent}}}, {OpType::CONV1D_K3, {{ParamField::Stride, FieldPersistence::Persistent}}}, {OpType::CONV1D_K7S3, {{ParamField::Stride, FieldPersistence::Persistent}}}, {OpType::CONV1D, {{ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::DENSE_MLP_TQ_FUSED, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::SAMPLE, {{ParamField::Temperature, FieldPersistence::Persistent}, {ParamField::TopP, FieldPersistence::Persistent}, {ParamField::MinP, FieldPersistence::Persistent}, {ParamField::RepetitionPenalty, FieldPersistence::Persistent}, {ParamField::TopK, FieldPersistence::Persistent}, {ParamField::RandomSeed, FieldPersistence::Persistent}, {ParamField::BiasIndices, FieldPersistence::Persistent}, {ParamField::BiasValues, FieldPersistence::Persistent}}}, }; @@ -300,6 +358,32 @@ void write_field(std::ostream& out, ParamField field, const OpParams& params) { throw std::runtime_error("Attempted to serialize runtime-only pointer field"); case ParamField::MaxCacheSeqLen: write_u64(out, static_cast(params.max_cache_seq_len)); break; case ParamField::CacheSinkSize: write_u64(out, static_cast(params.cache_sink_size)); break; + case ParamField::HopLength: write_u64(out, static_cast(params.hop_length)); break; + case ParamField::Power: write_f32(out, params.power); break; + case ParamField::Center: write_u32(out, params.center ? 1u : 0u); break; + case ParamField::MelFloor: write_f32(out, params.mel_floor); break; + case ParamField::Dither: write_f32(out, params.dither); break; + case ParamField::PreemphasisCoef: write_f32(out, params.preemphasis_coef); break; + case ParamField::RemoveDcOffset: write_u32(out, params.remove_dc_offset ? 1u : 0u); break; + case ParamField::LogMelMode: write_i32(out, static_cast(params.log_mel_mode)); break; + case ParamField::PadModeType: write_i32(out, static_cast(params.pad_mode_type)); break; + case ParamField::NumMelFilters: write_u64(out, static_cast(params.num_mel_filters)); break; + case ParamField::SamplingRate: write_u64(out, static_cast(params.sampling_rate)); break; + case ParamField::MinFrequency: write_f32(out, params.min_frequency); break; + case ParamField::MaxFrequency: write_f32(out, params.max_frequency); break; + case ParamField::MelNormType: write_i32(out, static_cast(params.mel_norm_type)); break; + case ParamField::MelScaleType: write_i32(out, static_cast(params.mel_scale_type)); break; + case ParamField::PatchSize: write_i32(out, static_cast(params.patch_size)); break; + case ParamField::RescaleFactor: write_f32(out, params.rescale_factor); break; + case ParamField::ImageMean: + for (float value : params.image_mean) write_f32(out, value); + break; + case ParamField::ImageStd: + for (float value : params.image_std) write_f32(out, value); + break; + case ParamField::TargetWidth: write_i32(out, static_cast(params.target_width)); break; + case ParamField::TargetHeight: write_i32(out, static_cast(params.target_height)); break; + case ParamField::ImageChannels: write_i32(out, static_cast(params.image_channels)); break; } } @@ -364,6 +448,32 @@ void read_field(std::istream& in, ParamField field, OpParams& params) { throw std::runtime_error("Graph file corrupted: runtime-only field serialized"); case ParamField::MaxCacheSeqLen: params.max_cache_seq_len = static_cast(read_u64(in)); break; case ParamField::CacheSinkSize: params.cache_sink_size = static_cast(read_u64(in)); break; + case ParamField::HopLength: params.hop_length = static_cast(read_u64(in)); break; + case ParamField::Power: params.power = read_f32(in); break; + case ParamField::Center: params.center = (read_u32(in) != 0); break; + case ParamField::MelFloor: params.mel_floor = read_f32(in); break; + case ParamField::Dither: params.dither = read_f32(in); break; + case ParamField::PreemphasisCoef: params.preemphasis_coef = read_f32(in); break; + case ParamField::RemoveDcOffset: params.remove_dc_offset = (read_u32(in) != 0); break; + case ParamField::LogMelMode: params.log_mel_mode = static_cast(read_i32(in)); break; + case ParamField::PadModeType: params.pad_mode_type = static_cast(read_i32(in)); break; + case ParamField::NumMelFilters: params.num_mel_filters = static_cast(read_u64(in)); break; + case ParamField::SamplingRate: params.sampling_rate = static_cast(read_u64(in)); break; + case ParamField::MinFrequency: params.min_frequency = read_f32(in); break; + case ParamField::MaxFrequency: params.max_frequency = read_f32(in); break; + case ParamField::MelNormType: params.mel_norm_type = static_cast(read_i32(in)); break; + case ParamField::MelScaleType: params.mel_scale_type = static_cast(read_i32(in)); break; + case ParamField::PatchSize: params.patch_size = static_cast(read_i32(in)); break; + case ParamField::RescaleFactor: params.rescale_factor = read_f32(in); break; + case ParamField::ImageMean: + for (float& value : params.image_mean) value = read_f32(in); + break; + case ParamField::ImageStd: + for (float& value : params.image_std) value = read_f32(in); + break; + case ParamField::TargetWidth: params.target_width = static_cast(read_i32(in)); break; + case ParamField::TargetHeight: params.target_height = static_cast(read_i32(in)); break; + case ParamField::ImageChannels: params.image_channels = static_cast(read_i32(in)); break; } } diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp index c2843f090..ef95d58ee 100644 --- a/cactus-graph/tests/test_nn.cpp +++ b/cactus-graph/tests/test_nn.cpp @@ -111,6 +111,52 @@ static void write_test_cq_weights( if (!file) throw std::runtime_error("failed writing test CQ weights"); } +static void write_test_int8_grouped_weights( + const std::filesystem::path& path, + const std::vector& shape, + size_t group_size, + const std::vector& data, + const std::vector<__fp16>& scales) { + constexpr uint32_t CACTUS_MAGIC = 0x54434143; + constexpr uint32_t FLAG_HAS_SCALES = 1 << 0; + constexpr size_t HEADER_SIZE = 84; + constexpr uint32_t alignment = 32; + + size_t rows = shape.empty() ? 1 : shape[0]; + size_t cols = data.size() / rows; + size_t num_groups = cols / group_size; + + std::ofstream file(path, std::ios::binary); + if (!file) throw std::runtime_error("cannot open test INT8 weights"); + + auto write_u32 = [&](uint32_t v) { file.write(reinterpret_cast(&v), sizeof(v)); }; + auto write_u64 = [&](uint64_t v) { file.write(reinterpret_cast(&v), sizeof(v)); }; + auto write_padding = [&](size_t bytes) { + char zero = 0; + for (size_t i = 0; i < bytes; ++i) file.write(&zero, 1); + }; + + write_u32(CACTUS_MAGIC); + write_u32(FLAG_HAS_SCALES); + write_u32(alignment); + write_u32(static_cast(shape.size())); + for (size_t i = 0; i < 4; ++i) write_u64(i < shape.size() ? shape[i] : 0); + write_u32(static_cast(Precision::INT8)); + write_u64(data.size()); + write_u64(scales.size() * sizeof(__fp16)); + write_u32(static_cast(group_size)); + write_u32(static_cast(num_groups)); + write_u64(rows); + + size_t aligned_header = align_offset_test(HEADER_SIZE, alignment); + write_padding(aligned_header - HEADER_SIZE); + file.write(reinterpret_cast(scales.data()), scales.size() * sizeof(__fp16)); + size_t data_start = align_offset_test(aligned_header + scales.size() * sizeof(__fp16), alignment); + write_padding(data_start - (aligned_header + scales.size() * sizeof(__fp16))); + file.write(reinterpret_cast(data.data()), data.size()); + if (!file) throw std::runtime_error("failed writing test INT8 weights"); +} + bool test_matmul_cq() { // Test graph-level CQ matmul dispatch for every supported bit width. const size_t M = 2, K = 128, N = 8; @@ -193,6 +239,155 @@ bool test_matmul_cq() { return true; } +bool test_mmap_int8_conv1d_pointwise_scales() { + const size_t C_in = 32; + const size_t C_out = 2; + std::filesystem::path path = std::filesystem::temp_directory_path() / "cactus_test_int8_conv1d_pointwise.weights"; + + std::vector q(C_out * C_in); + for (size_t c = 0; c < C_in; ++c) { + q[c] = 2; + q[C_in + c] = -3; + } + std::vector<__fp16> scales = {static_cast<__fp16>(0.25f), static_cast<__fp16>(0.5f)}; + write_test_int8_grouped_weights(path, {C_out, C_in, 1}, C_in, q, scales); + + CactusGraph g; + size_t input = g.input({1, 1, C_in}, Precision::FP16); + size_t weight = g.mmap_weights(path.string()); + size_t output = g.conv1d_pointwise(input, weight); + + std::vector<__fp16> x(C_in); + float sum = 0.0f; + for (size_t c = 0; c < C_in; ++c) { + float v = static_cast(static_cast(c % 5) - 2); + x[c] = static_cast<__fp16>(v); + sum += v; + } + + g.set_input(input, x.data(), Precision::FP16); + g.execute(); + auto* y = static_cast<__fp16*>(g.get_output(output)); + std::filesystem::remove(path); + + float y0 = static_cast(y[0]); + float y1 = static_cast(y[1]); + return std::abs(y0 - sum * 0.5f) < 1e-2f && std::abs(y1 - sum * -1.5f) < 1e-2f; +} + +bool test_axis0_slice_preserves_int8_scales() { + const size_t C = 2; + const size_t K = 9; + std::filesystem::path path = std::filesystem::temp_directory_path() / "cactus_test_int8_depthwise_slice.weights"; + + std::vector q(C * K); + for (size_t k = 0; k < K; ++k) { + q[k] = 2; + q[K + k] = -3; + } + std::vector<__fp16> scales = {static_cast<__fp16>(0.25f), static_cast<__fp16>(0.5f)}; + write_test_int8_grouped_weights(path, {C, 1, K}, K, q, scales); + + CactusGraph g; + size_t input = g.input({1, 1, K}, Precision::FP16); + size_t weight = g.mmap_weights(path.string()); + size_t sliced = g.slice(weight, 0, 1, 1); + size_t output = g.conv1d(input, sliced, 1); + + std::vector<__fp16> x(K, static_cast<__fp16>(1.0f)); + g.set_input(input, x.data(), Precision::FP16); + g.execute(); + auto* y = static_cast<__fp16*>(g.get_output(output)); + std::filesystem::remove(path); + + return std::abs(static_cast(y[0]) - (-13.5f)) < 1e-2f; +} + +bool test_glu_layout_invariant() { + const size_t L = 16; + const size_t C2 = 18; + + CactusGraph g; + size_t input = g.input({1, L, C2}, Precision::FP16); + size_t direct = g.glu(input, -1); + size_t transposed = g.transposeN(input, {0, 2, 1}); + size_t via_ncl = g.transposeN(g.glu(transposed, 1), {0, 2, 1}); + + std::vector<__fp16> x(L * C2); + for (size_t i = 0; i < x.size(); ++i) { + x[i] = static_cast<__fp16>((static_cast(i % 23) - 11) / 5.0f); + } + g.set_input(input, x.data(), Precision::FP16); + g.execute(); + + auto* a = static_cast<__fp16*>(g.get_output(direct)); + auto* b = static_cast<__fp16*>(g.get_output(via_ncl)); + for (size_t i = 0; i < L * (C2 / 2); ++i) { + if (std::abs(static_cast(a[i]) - static_cast(b[i])) > 1e-5f) { + return false; + } + } + return true; +} + +bool test_int8_causal_depthwise_matches_sliced_generic_conv1d() { + const size_t L = 16; + const size_t C = 2; + const size_t K = 9; + std::filesystem::path path = std::filesystem::temp_directory_path() / "cactus_test_int8_depthwise_causal.weights"; + + std::vector q(C * K); + for (size_t k = 0; k < K; ++k) { + q[k] = static_cast(k - 4); + q[K + k] = static_cast(4 - static_cast(k)); + } + std::vector<__fp16> scales = {static_cast<__fp16>(0.25f), static_cast<__fp16>(0.5f)}; + write_test_int8_grouped_weights(path, {C, 1, K}, K, q, scales); + + std::vector<__fp16> input_nlc(L * C); + for (size_t t = 0; t < L; ++t) { + for (size_t c = 0; c < C; ++c) { + input_nlc[t * C + c] = static_cast<__fp16>(static_cast((t + 1) * (c + 1)) / 16.0f); + } + } + + CactusGraph causal_graph; + size_t causal_input = causal_graph.input({1, L, C}, Precision::FP16); + size_t causal_weight = causal_graph.mmap_weights(path.string()); + size_t causal_output = causal_graph.conv1d_causal(causal_input, causal_weight, K, 1); + causal_graph.set_input(causal_input, input_nlc.data(), Precision::FP16); + causal_graph.execute(); + auto* causal = static_cast<__fp16*>(causal_graph.get_output(causal_output)); + + bool ok = true; + for (size_t c = 0; c < C; ++c) { + CactusGraph generic_graph; + size_t generic_input = generic_graph.input({1, 1, L + K - 1}, Precision::FP16); + size_t generic_weight = generic_graph.mmap_weights(path.string()); + size_t sliced_weight = generic_graph.slice(generic_weight, 0, c, 1); + size_t generic_output = generic_graph.conv1d(generic_input, sliced_weight, 1); + + std::vector<__fp16> padded(L + K - 1, static_cast<__fp16>(0)); + for (size_t t = 0; t < L; ++t) { + padded[t + K - 1] = input_nlc[t * C + c]; + } + generic_graph.set_input(generic_input, padded.data(), Precision::FP16); + generic_graph.execute(); + auto* generic = static_cast<__fp16*>(generic_graph.get_output(generic_output)); + + for (size_t t = 0; t < L; ++t) { + float a = static_cast(causal[t * C + c]); + float b = static_cast(generic[t]); + if (std::abs(a - b) > 1e-5f) { + ok = false; + break; + } + } + } + std::filesystem::remove(path); + return ok; +} + bool test_attention_int8_hybrid() { const size_t b = 1, s = 1, h = 2, kv = 2, d = 16; const size_t cache_len = 4; @@ -683,6 +878,10 @@ int main() { runner.run_test("Matrix Multiplication", test_matrix_multiplication()); runner.run_test("MatMul CQ", test_matmul_cq()); + runner.run_test("MMap INT8 Conv1D Pointwise", test_mmap_int8_conv1d_pointwise_scales()); + runner.run_test("Axis0 Slice INT8 Scales", test_axis0_slice_preserves_int8_scales()); + runner.run_test("GLU Layout Invariant", test_glu_layout_invariant()); + runner.run_test("INT8 Causal Depthwise", test_int8_causal_depthwise_matches_sliced_generic_conv1d()); runner.run_test("Transpose", test_transpose()); runner.run_test("Reshape", test_reshape()); runner.run_test("RMS Norm", test_rms_norm()); diff --git a/cactus-kernels/src/conv.cpp b/cactus-kernels/src/conv.cpp index 18990605b..146edf31d 100644 --- a/cactus-kernels/src/conv.cpp +++ b/cactus-kernels/src/conv.cpp @@ -20,71 +20,12 @@ constexpr size_t ACCELERATE_K_THRESHOLD = 32; constexpr size_t ACCELERATE_L_THRESHOLD = 128; #endif -#ifdef __APPLE__ -static void conv1d_causal_depthwise_f16_accelerate( - const __fp16* input, - const __fp16* weight, - __fp16* output, - size_t N, size_t L, size_t C, size_t K, size_t dilation) -{ - const size_t in_bs = L * C; - const size_t out_bs = L * C; - CactusThreading::parallel_for_2d(N, C, CactusThreading::Thresholds::ATTENTION, [&](size_t n, size_t c) { - const __fp16* Xb = input + n * in_bs; - __fp16* Yb = output + n * out_bs; - const __fp16* Wc = weight + c * K; - - if (dilation == 1) { - - size_t i = 0; - for(; i + 7 < L; i += 8){ - float16x8_t acc = vdupq_n_f16(0.0f); - - for (size_t k = 0; k < K; ++k) { - float16x8_t input_vec = vld1q_f16(&Xb[(i + k) * C + c]); - float16x8_t weight_vec = vdupq_n_f16(Wc[k]); - acc = vfmaq_f16(acc, input_vec, weight_vec); - } - vst1q_f16(&Yb[i * C + c], acc); - } - - for (; i < L; ++i) { - float acc = 0.0f; - for (size_t k = 0; k < K; ++k) { - acc += float(Xb[(i + k) * C + c]) * float(Wc[k]); - } - Yb[i * C + c] = (__fp16)acc; - } - } else { - - for (size_t i = 0; i < L; ++i) { - float acc = 0.0f; - - for (size_t k = 0; k < K; ++k) { - size_t idx = i + k * dilation; - acc += float(Xb[idx * C + c]) * float(Wc[k]); - } - - Yb[i * C + c] = (__fp16)acc; - } - } - }); -} -#endif - void cactus_conv1d_causal_depthwise_f16( const __fp16* input, const __fp16* weight, __fp16* output, size_t N, size_t L, size_t C, size_t K, size_t dilation) { -#ifdef __APPLE__ - if (K >= ACCELERATE_K_THRESHOLD && L >= ACCELERATE_L_THRESHOLD) { - conv1d_causal_depthwise_f16_accelerate(input, weight, output, N, L, C, K, dilation); - return; - } -#endif - const size_t in_bs = L * C; const size_t out_bs = L * C; @@ -92,9 +33,9 @@ void cactus_conv1d_causal_depthwise_f16( const __fp16* Xb = input + n * in_bs; __fp16* Yb = output + n * out_bs; - std::vector wrev(K); const __fp16* Wc = weight + c * K; - for (size_t k = 0; k < K; ++k) wrev[k] = (float)Wc[K - 1 - k]; + std::vector wrev(K); + for (size_t k = 0; k < K; ++k) wrev[k] = static_cast(Wc[K - 1 - k]); for (size_t t0 = 0; t0 < L; t0 += T_TILE_F16) { const size_t t1 = std::min(t0 + 1, L - 1); @@ -104,58 +45,57 @@ void cactus_conv1d_causal_depthwise_f16( size_t k = 0; for (; k + 8 <= K; k += 8) { - - float x0_0=0, x1_0=0, x2_0=0, x3_0=0; - float x0_1=0, x1_1=0, x2_1=0, x3_1=0; + float x0_0 = 0, x1_0 = 0, x2_0 = 0, x3_0 = 0; + float x0_1 = 0, x1_1 = 0, x2_1 = 0, x3_1 = 0; { - ptrdiff_t a0=(ptrdiff_t)t0-(ptrdiff_t)((k+0)*dilation); - ptrdiff_t a1=(ptrdiff_t)t0-(ptrdiff_t)((k+1)*dilation); - ptrdiff_t a2=(ptrdiff_t)t0-(ptrdiff_t)((k+2)*dilation); - ptrdiff_t a3=(ptrdiff_t)t0-(ptrdiff_t)((k+3)*dilation); - if (a0>=0) x0_0 = (float)Xb[(size_t)a0*C + c]; - if (a1>=0) x1_0 = (float)Xb[(size_t)a1*C + c]; - if (a2>=0) x2_0 = (float)Xb[(size_t)a2*C + c]; - if (a3>=0) x3_0 = (float)Xb[(size_t)a3*C + c]; - - ptrdiff_t b0=(ptrdiff_t)t1-(ptrdiff_t)((k+0)*dilation); - ptrdiff_t b1=(ptrdiff_t)t1-(ptrdiff_t)((k+1)*dilation); - ptrdiff_t b2=(ptrdiff_t)t1-(ptrdiff_t)((k+2)*dilation); - ptrdiff_t b3=(ptrdiff_t)t1-(ptrdiff_t)((k+3)*dilation); - if (b0>=0) x0_1 = (float)Xb[(size_t)b0*C + c]; - if (b1>=0) x1_1 = (float)Xb[(size_t)b1*C + c]; - if (b2>=0) x2_1 = (float)Xb[(size_t)b2*C + c]; - if (b3>=0) x3_1 = (float)Xb[(size_t)b3*C + c]; + ptrdiff_t a0 = static_cast(t0) - static_cast((k + 0) * dilation); + ptrdiff_t a1 = static_cast(t0) - static_cast((k + 1) * dilation); + ptrdiff_t a2 = static_cast(t0) - static_cast((k + 2) * dilation); + ptrdiff_t a3 = static_cast(t0) - static_cast((k + 3) * dilation); + if (a0 >= 0) x0_0 = static_cast(Xb[static_cast(a0) * C + c]); + if (a1 >= 0) x1_0 = static_cast(Xb[static_cast(a1) * C + c]); + if (a2 >= 0) x2_0 = static_cast(Xb[static_cast(a2) * C + c]); + if (a3 >= 0) x3_0 = static_cast(Xb[static_cast(a3) * C + c]); + + ptrdiff_t b0 = static_cast(t1) - static_cast((k + 0) * dilation); + ptrdiff_t b1 = static_cast(t1) - static_cast((k + 1) * dilation); + ptrdiff_t b2 = static_cast(t1) - static_cast((k + 2) * dilation); + ptrdiff_t b3 = static_cast(t1) - static_cast((k + 3) * dilation); + if (b0 >= 0) x0_1 = static_cast(Xb[static_cast(b0) * C + c]); + if (b1 >= 0) x1_1 = static_cast(Xb[static_cast(b1) * C + c]); + if (b2 >= 0) x2_1 = static_cast(Xb[static_cast(b2) * C + c]); + if (b3 >= 0) x3_1 = static_cast(Xb[static_cast(b3) * C + c]); } - float32x4_t xv0 = {x0_0,x1_0,x2_0,x3_0}; - float32x4_t yv0 = {x0_1,x1_1,x2_1,x3_1}; - float32x4_t wv0 = {wrev[k+0],wrev[k+1],wrev[k+2],wrev[k+3]}; + float32x4_t xv0 = {x0_0, x1_0, x2_0, x3_0}; + float32x4_t yv0 = {x0_1, x1_1, x2_1, x3_1}; + float32x4_t wv0 = {wrev[k + 0], wrev[k + 1], wrev[k + 2], wrev[k + 3]}; vacc0 = vfmaq_f32(vacc0, xv0, wv0); vacc1 = vfmaq_f32(vacc1, yv0, wv0); - float a0_0=0, a1_0=0, a2_0=0, a3_0=0; - float a0_1=0, a1_1=0, a2_1=0, a3_1=0; + float a0_0 = 0, a1_0 = 0, a2_0 = 0, a3_0 = 0; + float a0_1 = 0, a1_1 = 0, a2_1 = 0, a3_1 = 0; { - ptrdiff_t a0i=(ptrdiff_t)t0-(ptrdiff_t)((k+4)*dilation); - ptrdiff_t a1i=(ptrdiff_t)t0-(ptrdiff_t)((k+5)*dilation); - ptrdiff_t a2i=(ptrdiff_t)t0-(ptrdiff_t)((k+6)*dilation); - ptrdiff_t a3i=(ptrdiff_t)t0-(ptrdiff_t)((k+7)*dilation); - if (a0i>=0) a0_0 = (float)Xb[(size_t)a0i*C + c]; - if (a1i>=0) a1_0 = (float)Xb[(size_t)a1i*C + c]; - if (a2i>=0) a2_0 = (float)Xb[(size_t)a2i*C + c]; - if (a3i>=0) a3_0 = (float)Xb[(size_t)a3i*C + c]; - - ptrdiff_t b0i=(ptrdiff_t)t1-(ptrdiff_t)((k+4)*dilation); - ptrdiff_t b1i=(ptrdiff_t)t1-(ptrdiff_t)((k+5)*dilation); - ptrdiff_t b2i=(ptrdiff_t)t1-(ptrdiff_t)((k+6)*dilation); - ptrdiff_t b3i=(ptrdiff_t)t1-(ptrdiff_t)((k+7)*dilation); - if (b0i>=0) a0_1 = (float)Xb[(size_t)b0i*C + c]; - if (b1i>=0) a1_1 = (float)Xb[(size_t)b1i*C + c]; - if (b2i>=0) a2_1 = (float)Xb[(size_t)b2i*C + c]; - if (b3i>=0) a3_1 = (float)Xb[(size_t)b3i*C + c]; + ptrdiff_t a0i = static_cast(t0) - static_cast((k + 4) * dilation); + ptrdiff_t a1i = static_cast(t0) - static_cast((k + 5) * dilation); + ptrdiff_t a2i = static_cast(t0) - static_cast((k + 6) * dilation); + ptrdiff_t a3i = static_cast(t0) - static_cast((k + 7) * dilation); + if (a0i >= 0) a0_0 = static_cast(Xb[static_cast(a0i) * C + c]); + if (a1i >= 0) a1_0 = static_cast(Xb[static_cast(a1i) * C + c]); + if (a2i >= 0) a2_0 = static_cast(Xb[static_cast(a2i) * C + c]); + if (a3i >= 0) a3_0 = static_cast(Xb[static_cast(a3i) * C + c]); + + ptrdiff_t b0i = static_cast(t1) - static_cast((k + 4) * dilation); + ptrdiff_t b1i = static_cast(t1) - static_cast((k + 5) * dilation); + ptrdiff_t b2i = static_cast(t1) - static_cast((k + 6) * dilation); + ptrdiff_t b3i = static_cast(t1) - static_cast((k + 7) * dilation); + if (b0i >= 0) a0_1 = static_cast(Xb[static_cast(b0i) * C + c]); + if (b1i >= 0) a1_1 = static_cast(Xb[static_cast(b1i) * C + c]); + if (b2i >= 0) a2_1 = static_cast(Xb[static_cast(b2i) * C + c]); + if (b3i >= 0) a3_1 = static_cast(Xb[static_cast(b3i) * C + c]); } - float32x4_t xv1 = {a0_0,a1_0,a2_0,a3_0}; - float32x4_t yv1 = {a0_1,a1_1,a2_1,a3_1}; - float32x4_t wv1 = {wrev[k+4],wrev[k+5],wrev[k+6],wrev[k+7]}; + float32x4_t xv1 = {a0_0, a1_0, a2_0, a3_0}; + float32x4_t yv1 = {a0_1, a1_1, a2_1, a3_1}; + float32x4_t wv1 = {wrev[k + 4], wrev[k + 5], wrev[k + 6], wrev[k + 7]}; vacc0 = vfmaq_f32(vacc0, xv1, wv1); vacc1 = vfmaq_f32(vacc1, yv1, wv1); } @@ -164,14 +104,14 @@ void cactus_conv1d_causal_depthwise_f16( float acc1 = vaddvq_f32(vacc1); for (; k < K; ++k) { - ptrdiff_t a=(ptrdiff_t)t0-(ptrdiff_t)(k*dilation); - if (a>=0) acc0 += wrev[k] * (float)Xb[(size_t)a*C + c]; - ptrdiff_t b=(ptrdiff_t)t1-(ptrdiff_t)(k*dilation); - if (b>=0) acc1 += wrev[k] * (float)Xb[(size_t)b*C + c]; + ptrdiff_t a = static_cast(t0) - static_cast(k * dilation); + if (a >= 0) acc0 += wrev[k] * static_cast(Xb[static_cast(a) * C + c]); + ptrdiff_t b = static_cast(t1) - static_cast(k * dilation); + if (b >= 0) acc1 += wrev[k] * static_cast(Xb[static_cast(b) * C + c]); } - Yb[t0*C + c] = (__fp16)acc0; - if (t0 + 1 < L) Yb[t1*C + c] = (__fp16)acc1; + Yb[t0 * C + c] = static_cast<__fp16>(acc0); + if (t0 + 1 < L) Yb[t1 * C + c] = static_cast<__fp16>(acc1); } }); } @@ -960,5 +900,3 @@ void cactus_conv1d_same_depthwise_f16_k9( } }); } - - diff --git a/cactus-kernels/src/nn.cpp b/cactus-kernels/src/nn.cpp index c56de09db..5bb2cce86 100644 --- a/cactus-kernels/src/nn.cpp +++ b/cactus-kernels/src/nn.cpp @@ -61,7 +61,7 @@ void cactus_silu_f16(const __fp16* input, __fp16* output, size_t num_elements) { for (size_t i = vectorized_end; i < end_idx; ++i) { float x_f32 = static_cast(input[i]); - output[i] = static_cast<__fp16>(x_f32 / (1.0f + expf(-x_f32))); + output[i] = static_cast<__fp16>(x_f32 * fast_sigmoid_f32(x_f32)); } }); } @@ -187,7 +187,7 @@ void cactus_sigmoid_f16(const __fp16* input, __fp16* output, size_t num_elements for (size_t i = vectorized_end; i < end_idx; ++i) { float x_f32 = static_cast(input[i]); - output[i] = static_cast<__fp16>(1.0f / (1.0f + expf(-x_f32))); + output[i] = static_cast<__fp16>(fast_sigmoid_f32(x_f32)); } }); } @@ -256,7 +256,7 @@ void cactus_glu_f16( } for (; i < inner_size; ++i) { - const float gate = 1.0f / (1.0f + expf(-static_cast(b[i]))); + const float gate = fast_sigmoid_f32(static_cast(b[i])); y[i] = static_cast<__fp16>(static_cast(a[i]) * gate); } } @@ -294,7 +294,7 @@ void cactus_glu_f32( } for (; i < inner_size; ++i) { - const float gate = 1.0f / (1.0f + expf(-b[i])); + const float gate = fast_sigmoid_f32(b[i]); y[i] = a[i] * gate; } } diff --git a/cactus-kernels/src/threading.h b/cactus-kernels/src/threading.h index e5795329f..67d8167d9 100644 --- a/cactus-kernels/src/threading.h +++ b/cactus-kernels/src/threading.h @@ -167,6 +167,10 @@ inline float32x4_t fast_sigmoid_f32x4(float32x4_t x) { return vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(x)))); } +inline float fast_sigmoid_f32(float x) { + return vgetq_lane_f32(fast_sigmoid_f32x4(vdupq_n_f32(x)), 0); +} + template inline float16x8_t apply_f32_op_on_f16x8(float16x8_t v, F32x4Op op) { float32x4_t lo, hi; diff --git a/cactus-kernels/tests/test_conv.cpp b/cactus-kernels/tests/test_conv.cpp index 42b916c1e..558d934b2 100644 --- a/cactus-kernels/tests/test_conv.cpp +++ b/cactus-kernels/tests/test_conv.cpp @@ -39,6 +39,54 @@ bool test_conv1d_causal_depthwise() { return true; } +bool test_conv1d_causal_depthwise_matches_reference() { + const size_t N = 1, L = 16, C = 8, K = 9; + std::vector<__fp16> input(N * L * C), weight(C * K), causal(N * L * C); + std::vector<__fp16> padded_channel(N * (L + K - 1)), generic_channel(N * L); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + + cactus_conv1d_causal_depthwise_f16(input.data(), weight.data(), causal.data(), N, L, C, K, 1); + + for (size_t n = 0; n < N; ++n) { + for (size_t t = 0; t < L; ++t) { + for (size_t c = 0; c < C; ++c) { + const size_t causal_idx = (n * L + t) * C + c; + float expected = 0.0f; + for (size_t k = 0; k < K; ++k) { + ptrdiff_t src = static_cast(t) + static_cast(k) - static_cast(K - 1); + if (src >= 0) { + expected += static_cast(input[(n * L + static_cast(src)) * C + c]) * static_cast(weight[c * K + k]); + } + } + const __fp16 expected_h = static_cast<__fp16>(expected); + if (static_cast(causal[causal_idx]) != static_cast(expected_h)) { + std::cerr << " causal depthwise mismatch [t=" << t << ",c=" << c << "]: " + << static_cast(causal[causal_idx]) << " vs " << static_cast(expected_h) << "\n"; + return false; + } + } + } + } + for (size_t c = 0; c < C; ++c) { + std::fill(padded_channel.begin(), padded_channel.end(), static_cast<__fp16>(0)); + for (size_t t = 0; t < L; ++t) { + padded_channel[t + K - 1] = input[t * C + c]; + } + cactus_conv1d_f16(padded_channel.data(), weight.data() + c * K, nullptr, generic_channel.data(), N, L + K - 1, 1, 1, K, 1); + for (size_t t = 0; t < L; ++t) { + const size_t causal_idx = t * C + c; + if (std::abs(static_cast(causal[causal_idx]) - static_cast(generic_channel[t])) > 1e-5f) { + std::cerr << " causal depthwise generic mismatch [t=" << t << ",c=" << c << "]: " + << static_cast(causal[causal_idx]) << " vs " + << static_cast(generic_channel[t]) << "\n"; + return false; + } + } + } + return true; +} + bool test_stft_complex() { const size_t N = 1, L = 128, C_in = 1, K = 64, stride = 32; const size_t num_fft_bins = K / 2; @@ -143,6 +191,7 @@ int main() { TestRunner runner("Convolution & Pooling"); runner.run_test("conv1d_k3", test_conv1d_k3()); runner.run_test("conv1d_causal_depthwise", test_conv1d_causal_depthwise()); + runner.run_test("conv1d_causal_depthwise_ref", test_conv1d_causal_depthwise_matches_reference()); runner.run_test("stft_complex", test_stft_complex()); runner.run_test("maxpool1d", test_maxpool1d()); runner.print_benchmarks_header(); diff --git a/docs/cactus_engine.md b/docs/cactus_engine.md index 0b8e851e8..8adbadfa2 100644 --- a/docs/cactus_engine.md +++ b/docs/cactus_engine.md @@ -523,11 +523,11 @@ int cactus_benchmark_tokens( Returns the number of bytes written to `response_buffer` on success, `-1` on error. ### `cactus_transcribe` -Transcribes audio to text using Whisper or Parakeet TDT models. Supports both file-based and buffer-based audio input. +Transcribes audio to text using Whisper, Parakeet TDT, or Nemotron ASR models. Supports both file-based and buffer-based audio input. ```c int cactus_transcribe( - cactus_model_t model, // Model handle (Whisper or Parakeet TDT model) + cactus_model_t model, // Model handle (Whisper, Parakeet TDT, or Nemotron ASR model) const char* audio_file_path, // Path to WAV file (16-bit PCM) - can be NULL if using pcm_buffer const char* prompt, // Optional prompt to guide transcription (can be NULL) char* response_buffer, // Buffer for response JSON @@ -555,9 +555,9 @@ int cactus_transcribe( | Option | Type | Default | Description | |--------|------|---------|-------------| -| `max_tokens` | int | auto | Maximum tokens to generate. When unset, it defaults to the larger of 100 and an audio-length estimate (`audio_sec × 20` for Whisper, `audio_sec × 30` for Parakeet). For Whisper the result is then capped so the prompt tokens plus generated tokens fit the decoder's 448-position limit. | -| `language` | string | model default | Whisper only. Two-letter language code (e.g. `en`, `es`, `de`) substituted into the decoder prompt's language token. Ignored by Parakeet and when an explicit `prompt` is supplied. | -| `timestamps` | bool | false | Whisper only. Decodes timestamp tokens and populates `segments` with `{start, end, text}` entries (seconds). Empty otherwise, including all Parakeet transcription. | +| `max_tokens` | int | auto | Whisper maximum tokens to generate. When unset, it defaults to the larger of 100 and an audio-length estimate (`audio_sec × 20`) and is capped so the prompt tokens plus generated tokens fit the decoder's 448-position limit. | +| `language` / `target_lang` | string | model default | Whisper: two-letter language code substituted into the decoder prompt's language token. Nemotron: prompt language (`auto` by default). Ignored by Parakeet and by Whisper when an explicit `prompt` is supplied. | +| `timestamps` | bool | false | Whisper only. Decodes timestamp tokens and populates `segments` with `{start, end, text}` entries (seconds). Empty otherwise, including Parakeet and Nemotron transcription. | **Response Format:** ```json @@ -582,7 +582,7 @@ int cactus_transcribe( ``` - `response`: Full transcription text -- `segments`: `{start, end, text}` entries (seconds), populated only for Whisper when the `timestamps` option is set; empty otherwise, including all Parakeet transcription +- `segments`: `{start, end, text}` entries (seconds), populated only for Whisper when the `timestamps` option is set; empty otherwise, including Parakeet and Nemotron transcription - `cloud_handoff`: Always false for transcription - `confidence_threshold`: `-1.0` (unset) — transcription does not resolve a cloud-handoff threshold @@ -610,26 +610,26 @@ int result = cactus_transcribe(whisper, NULL, NULL, ``` ### Streaming transcription -Transcribe continuously while audio is still being captured, instead of waiting for a complete recording. You push PCM as it arrives and read back text as soon as it stabilizes. Supported for the dedicated speech models (Whisper and Parakeet TDT). +Transcribe continuously while audio is still being captured, instead of waiting for a complete recording. You push PCM as it arrives and read back text as soon as it stabilizes. Supported for the dedicated speech models (Whisper, Parakeet TDT, and Nemotron ASR). The session emits text in two parts on every call: - **`confirmed`** — newly finalized words. Append them to your running transcript; they never change. - **`pending`** — the current best guess for the still-changing tail. Replace it on every call (do not append it); it is for live display only. -Text is confirmed once two successive re-transcriptions of the audio agree (LocalAgreement, compared per **segment** for Whisper) or once a following token starts a new word (Parakeet TDT). Confirmed audio is dropped from the front of the buffer as the window advances, so memory stays bounded and the active window never approaches the model's fixed input length. +Text is confirmed once two successive re-transcriptions of the audio agree (LocalAgreement, compared per **segment** for Whisper) or once a following token starts a new word (Parakeet TDT and Nemotron ASR). Confirmed audio is dropped from the active decode window as it advances. Nemotron v1 keeps the full session audio for final flush, supports the Cactus streaming contract up to the emitted encoder window set, and does not yet use NVIDIA-style cache-aware encoder-cache tensors. #### `cactus_stream_transcribe_start` Opens a streaming session bound to an already-initialized speech model. ```c cactus_stream_transcribe_t cactus_stream_transcribe_start( - cactus_model_t model, // Whisper or Parakeet TDT model handle + cactus_model_t model, // Whisper, Parakeet TDT, or Nemotron ASR model handle const char* options_json // Optional (can be NULL); forwarded to cactus_transcribe ); ``` **Returns:** an opaque session handle, or `NULL` on error (see `cactus_get_last_error`). Free it with `cactus_stream_transcribe_stop`. -`options_json` (optional) is forwarded to the underlying `cactus_transcribe` call **for Whisper only** (e.g. `language`); the Parakeet TDT path ignores it. Chunking and segmentation are handled internally with no user-facing tunables. +`options_json` (optional) is forwarded to the underlying `cactus_transcribe` call for Whisper and Nemotron (e.g. `language` / `target_lang`); the Parakeet TDT path ignores it. Chunking and segmentation are handled internally with no user-facing tunables. #### `cactus_stream_transcribe_process` Feeds the next slice of audio. Input is 16-bit signed PCM, **16 kHz, mono** — the same format as `cactus_transcribe`'s `pcm_buffer`. Feed reasonably small chunks (≈ 0.1–2 s) for low latency. diff --git a/docs/cactus_quants.md b/docs/cactus_quants.md index 83ec3b424..77e5ad2bc 100644 --- a/docs/cactus_quants.md +++ b/docs/cactus_quants.md @@ -110,7 +110,7 @@ M is the modality-tower bit width and L is the LLM bit width: At 2-bit modality tower with 4-bit LLM, AWQ and HQQ collapse on ChartQA from ~49 to ~14 while CQ retains 48.44. -### Transcription (Whisper and Parakeet) +### Transcription (Whisper, Parakeet, and Nemotron) On speech models, CQ is the best method on Parakeet at every bit width. The standout: 2-bit Parakeet-1.1B where CQ holds 0.147 WER while every baseline collapses to ≥ 1.0. @@ -122,7 +122,8 @@ On speech models, CQ is the best method on Parakeet at every bit width. The stan | Parakeet-1.1B 2-bit WER | **0.147** | 1.043 | 1.084 | 1.043 | On Whisper, CQ is competitive but doesn't consistently lead — it's within 1% of the -best method at most settings. +best method at most settings. Nemotron ASR uses the same speech conversion/runtime +surface as Parakeet, with RNNT decoding and Cactus streaming API parity in v1. ### Mixed-Precision Allocation diff --git a/docs/cactus_transpiler.md b/docs/cactus_transpiler.md index a0dbd7773..53e6ca687 100644 --- a/docs/cactus_transpiler.md +++ b/docs/cactus_transpiler.md @@ -124,6 +124,7 @@ the model in a thin adapter that exposes a stable, task-specific interface: | `ctc_logits` | `input_features` | CTC logits | | `seq2seq_transcription` | `input_features` | transcription token IDs | | `tdt_transcription` | `input_features` | TDT token IDs | +| `rnnt_transcription` | `input_features` | RNNT token IDs | This is where model-family-specific knowledge lives (Gemma4, Qwen, LFM2, Parakeet, etc.). The adapter also injects `use_cache=False` and `return_dict=False` so the export @@ -229,7 +230,7 @@ cactus transpile google/gemma-4-E2B-it \ --weights-dir ./gemma4-weights ``` -### Speech Models (Parakeet TDT) +### Speech Models (Parakeet TDT and Nemotron ASR) ```bash cactus convert nvidia/parakeet-tdt-0.6b-v3 ./parakeet-weights --bits 4 @@ -239,6 +240,16 @@ cactus transpile nvidia/parakeet-tdt-0.6b-v3 \ --weights-dir ./parakeet-weights ``` +```bash +cactus convert nvidia/nemotron-3.5-asr-streaming-0.6b ./nemotron-weights --bits 4 +cactus transpile nvidia/nemotron-3.5-asr-streaming-0.6b \ + --audio-file recording.wav \ + --task rnnt_transcription \ + --weights-dir ./nemotron-weights +``` + +Nemotron ASR bundles provide the same Cactus one-shot and streaming APIs as Parakeet. The v1 stream path does not yet use NVIDIA-style cache-aware encoder-cache tensors. + --- ## Running Saved Bundles @@ -426,7 +437,7 @@ cactus transpile [options] | Option | Description | |--------|-------------| | `--weights-dir ` | Path to converted CQ weights (default: `weights/`) | -| `--task ` | Force task type (default: `auto` — inferred from model config). Choices: `causal_lm_logits`, `multimodal_causal_lm_logits`, `ctc_logits`, `encoder_hidden_states`, `seq2seq_transcription`, `tdt_transcription` | +| `--task ` | Force task type (default: `auto` — inferred from model config). Choices: `causal_lm_logits`, `multimodal_causal_lm_logits`, `ctc_logits`, `encoder_hidden_states`, `seq2seq_transcription`, `tdt_transcription`, `rnnt_transcription` | | `--prompt ` | Representative prompt for shape capture | | `--image-file ` | Representative image (repeatable) | | `--audio-file ` | Representative audio file | @@ -446,6 +457,8 @@ cactus transpile [options] NPU emission (CoreML `.mlpackage`s for Apple Silicon audio and vision encoders) is available through the `--npu`, `--npu-quantize`, `--npu-audio-quantize`, and `--npu-vision-quantize` flags on `cactus transpile`. +Nemotron ASR captures 1024, 2048, 4096, and 6144 input-frame variants by default; +override the set with `CACTUS_NEMOTRON_ASR_FRAMES`. ### `cactus run` @@ -490,6 +503,7 @@ tested adapters for: | Whisper | encoder | no | | Parakeet CTC | encoder | no | | Parakeet TDT | transcription | yes (encoder + decoder) | +| Nemotron ASR | transcription | yes (encoder + decoder) | Adding a new model family means writing a small adapter in `model_adapters.py` that wraps the HF forward into the canonical task interface. The rest of the pipeline diff --git a/python/README.md b/python/README.md index a259aa8ce..1ac49ac7c 100644 --- a/python/README.md +++ b/python/README.md @@ -187,7 +187,7 @@ result = cactus_complete(model, completion_messages, None, tools, None) ### Transcription -Returns a `dict` with the `response` field (transcribed text) and a `segments` array of `{start, end, text}` objects. `segments` is populated only for Whisper models when the `timestamps` option is set (`{"timestamps": True}`); it is empty otherwise, including for all Parakeet transcription. +Returns a `dict` with the `response` field (transcribed text) and a `segments` array of `{start, end, text}` objects. `segments` is populated only for Whisper models when the `timestamps` option is set (`{"timestamps": True}`); it is empty otherwise, including for Parakeet and Nemotron transcription. ```python result = cactus_transcribe( @@ -219,7 +219,7 @@ for seg in result["segments"]: ### Streaming transcription -Transcribe continuously while audio is still being captured (Whisper and Parakeet TDT). Open a session, push 16 kHz mono 16-bit PCM chunks, and read text back as it stabilizes: `confirmed` words are final (append them to your transcript), `pending` is the volatile tail (replace it each call, for live display only). +Transcribe continuously while audio is still being captured (Whisper, Parakeet TDT, and Nemotron ASR). Open a session, push 16 kHz mono 16-bit PCM chunks, and read text back as it stabilizes: `confirmed` words are final (append them to your transcript), `pending` is the volatile tail (replace it each call, for live display only). ```python stream = cactus_stream_transcribe_start(model: int, options: dict | str | None) -> int @@ -227,7 +227,7 @@ result = cactus_stream_transcribe_process(stream: int, pcm_data: bytes) -> dict result = cactus_stream_transcribe_stop(stream: int) -> dict # {"confirmed": str, "pending": ""}; destroys the session ``` -`options` is forwarded to `cactus_transcribe` for **Whisper only** (e.g. `language`, `max_tokens`); the Parakeet TDT path ignores it. Chunking is handled internally. +`options` is forwarded to `cactus_transcribe` for Whisper and Nemotron. Whisper accepts options such as `language` and `max_tokens`; Nemotron accepts `language` / `target_lang`. The Parakeet TDT path ignores it. Chunking is handled internally. ```python stream = cactus_stream_transcribe_start(model, {"language": "en"}) diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index aa40fc578..eb6cafe86 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -1041,14 +1041,15 @@ def cactus_transcribe(model, audio_path, prompt=None, options=None, callback=Non audio_path: Path to a WAV audio file. prompt: Optional prompt to guide transcription. options: Optional dict of transcription options. Set ``timestamps: True`` (Whisper - only) to populate the ``segments`` array. + only) to populate the ``segments`` array. Nemotron accepts + ``language`` or ``target_lang`` and defaults to ``auto``. callback: Optional function(text, token_id) for streaming tokens. pcm_data: Optional raw PCM audio bytes (alternative to audio_path). Returns: A dict with the transcribed text under ``response`` and a ``segments`` array of ``{start, end, text}`` objects, populated only for Whisper when ``timestamps`` is set - (empty otherwise, including all Parakeet transcription). + (empty otherwise, including Parakeet and Nemotron transcription). """ buf = ctypes.create_string_buffer(1 << 20) cb = _make_token_callback(callback) @@ -1063,13 +1064,14 @@ def cactus_transcribe(model, audio_path, prompt=None, options=None, callback=Non def cactus_stream_transcribe_start(model, options=None): - """Open a streaming transcription session on a Whisper or Parakeet TDT model. + """Open a streaming transcription session on a Whisper, Parakeet TDT, or Nemotron ASR model. Args: model: Model handle. - options: Optional dict forwarded to the underlying transcribe call - (e.g. {"language": "en", "max_tokens": 256}); chunking is - handled internally. + options: Optional dict forwarded to Whisper and Nemotron paths + (e.g. {"language": "en", "target_lang": "auto"}); chunking is + handled internally. Nemotron uses Cactus streaming API parity; + cache-aware encoder-cache tensors are future work. Returns: An opaque stream handle. Feed audio with cactus_stream_transcribe_process diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index c0bc5ba34..8275fba9d 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -107,7 +107,7 @@ def create_parser(): cactus transcribe [model] live microphone transcription with a model --file audio file to transcribe (WAV) - --language language code (default: en) + --language language code (default: model-specific) --bits 1|2|3|4 CQ quantization (default: 4) --platform {_PLATFORM_PIPE:<22} target accelerator (default: auto) --token HuggingFace token (gated models) @@ -269,8 +269,8 @@ def create_parser(): help=f"HuggingFace model id (default: {DEFAULT_TRANSCRIPTION_MODEL_ID})") transcribe_parser.add_argument("--file", dest="audio_file", default=None, help="Audio file to transcribe (WAV)") - transcribe_parser.add_argument("--language", default="en", - help="Language code (default: en)") + transcribe_parser.add_argument("--language", default=None, + help="Language code (default: model-specific)") transcribe_parser.add_argument("--bits", type=int, choices=[1, 2, 3, 4], default=4, help="CQ quantization (default: 4)") transcribe_parser.add_argument("--platform", choices=_PLATFORM_CHOICES, default="auto", @@ -346,7 +346,8 @@ def create_parser(): transpile_parser.add_argument("--task", default="auto", choices=["auto", "causal_lm_logits", "multimodal_causal_lm_logits", "ctc_logits", "encoder_hidden_states", - "seq2seq_transcription", "tdt_transcription"], + "seq2seq_transcription", "tdt_transcription", + "rnnt_transcription"], help="Transpile task (default: auto, from model config)") transpile_parser.add_argument("--prompt", help="Prompt for causal/multimodal shape capture") diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index 1f01c49ac..dd640ae35 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -83,12 +83,12 @@ def _infer_transpile_spec(*, task, plan): task=task, needs_image=task == "multimodal_causal_lm_logits", needs_audio=task in { - "tdt_transcription", "seq2seq_transcription", + "tdt_transcription", "rnnt_transcription", "seq2seq_transcription", "ctc_logits", "encoder_hidden_states", "multimodal_causal_lm_logits", }, force_component_pipeline=task in { - "tdt_transcription", "seq2seq_transcription", + "tdt_transcription", "rnnt_transcription", "seq2seq_transcription", "multimodal_causal_lm_logits", }, ) @@ -153,7 +153,7 @@ def _has_transpiled_bundle(path): _AUDIO_TASKS = frozenset({ - "tdt_transcription", "seq2seq_transcription", + "tdt_transcription", "rnnt_transcription", "seq2seq_transcription", "ctc_logits", "encoder_hidden_states", }) diff --git a/python/cactus/convert/cactus_adapters/config_utils.py b/python/cactus/convert/cactus_adapters/config_utils.py index b17b7c229..c813d4271 100644 --- a/python/cactus/convert/cactus_adapters/config_utils.py +++ b/python/cactus/convert/cactus_adapters/config_utils.py @@ -18,6 +18,17 @@ def cfg_get(c, key, default=None): return default +def prompt_dim_from_dictionary(prompt_dictionary, *candidates): + values = [] + if isinstance(prompt_dictionary, dict): + for value in prompt_dictionary.values(): + values.append(int(value)) + dim = max((int(value or 0) for value in candidates), default=0) + if values: + dim = max(dim, max(values) + 1) + return max(1, dim) + + def detect_model_type(cfg, config, output_dir=None): """Detect the model architecture type from config.""" model_type_str = str(cfg_get(cfg, 'model_type', cfg_get(config, 'model_type', '')) or '').lower().strip() @@ -26,10 +37,21 @@ def detect_model_type(cfg, config, output_dir=None): loss_cfg = cfg_get(cfg, 'loss', cfg_get(config, 'loss', None)) loss_name = str(cfg_get(loss_cfg, 'loss_name', '')).lower() - # NeMo Parakeet-TDT configs often do not expose HF-style model_type names. if decoding_model_type == 'tdt' or loss_name == 'tdt': return 'parakeet_tdt' + target = str(cfg_get(cfg, 'target', cfg_get(config, 'target', '')) or '').lower() + model_defaults = cfg_get(cfg, 'model_defaults', cfg_get(config, 'model_defaults', {})) + has_prompt = ( + 'withprompt' in target + or 'prompt' in target + or cfg_get(model_defaults, 'prompt_dictionary', None) is not None + or int(cfg_get(cfg, 'num_prompts', cfg_get(config, 'num_prompts', 0)) or 0) > 0 + ) + + if has_prompt and ('rnnt' in target or cfg_get(cfg_get(cfg, 'joint', cfg_get(config, 'joint', {})), 'num_classes', None) is not None): + return 'nemotron_asr' + if model_type_str in {'nomic_bert', 'nomic-bert'} or 'nomic' in model_type_str: return 'nomic' if 'gemma4' in model_type_str or 'tinyllama' in model_type_str: @@ -57,6 +79,8 @@ def detect_model_type(cfg, config, output_dir=None): return 'bert' elif 'whisper' in model_type_str: return 'whisper' + elif 'nemotron_asr' in model_type_str or 'nemotron-asr' in model_type_str: + return 'nemotron_asr' elif 'parakeet_tdt' in model_type_str or 'parakeet-tdt' in model_type_str: return 'parakeet_tdt' elif 'parakeet' in model_type_str: @@ -397,6 +421,82 @@ def extract_parakeet_tdt_config(config): } +def extract_nemotron_asr_config(config): + """Extract Cactus runtime config for prompt-conditioned Nemotron RNNT ASR checkpoints.""" + encoder_cfg = cfg_get(config, 'encoder', cfg_get(config, 'encoder_config', None)) + if encoder_cfg is None: + raise ValueError("Nemotron ASR conversion requires encoder config") + + preprocessor_cfg = cfg_get(config, 'preprocessor', {}) + decoder_cfg = cfg_get(config, 'decoder', {}) + prediction_cfg = cfg_get(decoder_cfg, 'prediction', cfg_get(decoder_cfg, 'prednet', {})) + joint_cfg = cfg_get(config, 'joint', {}) + jointnet_cfg = cfg_get(joint_cfg, 'jointnet', {}) + model_defaults_cfg = cfg_get(config, 'model_defaults', {}) + + hidden_dim = int(cfg_get(encoder_cfg, 'd_model', cfg_get(encoder_cfg, 'hidden_size', 0))) + attention_heads = int(cfg_get(encoder_cfg, 'n_heads', cfg_get(encoder_cfg, 'num_attention_heads', 0))) + ff_intermediate = int(cfg_get(encoder_cfg, 'ffn_hidden_size', cfg_get(encoder_cfg, 'intermediate_size', 0))) + if ff_intermediate == 0: + ff_intermediate = int(round(hidden_dim * float(cfg_get(encoder_cfg, 'ff_expansion_factor', 4.0)))) + + labels = cfg_get(config, 'labels', cfg_get(joint_cfg, 'vocabulary', [])) + if not isinstance(labels, (list, tuple)): + labels = [] + vocab_size = int(cfg_get(decoder_cfg, 'vocab_size', cfg_get(config, 'vocab_size', len(labels)))) + blank_id = int(cfg_get(config, 'rnnt_blank_id', vocab_size)) + prompt_dictionary = cfg_get(config, 'prompt_dictionary', cfg_get(model_defaults_cfg, 'prompt_dictionary', {})) + if not isinstance(prompt_dictionary, dict): + prompt_dictionary = {} + prompt_dim = prompt_dim_from_dictionary( + prompt_dictionary, + cfg_get(config, 'prompt_dim', 0), + cfg_get(config, 'num_prompts', 0), + cfg_get(model_defaults_cfg, 'num_prompts', 0), + ) + default_prompt_id = int(cfg_get(config, 'default_prompt_id', prompt_dictionary.get('auto', 0))) + prompt_dictionary_text = ",".join( + f"{str(key)}:{int(value)}" + for key, value in sorted(prompt_dictionary.items(), key=lambda item: str(item[0])) + ) + greedy_cfg = cfg_get(cfg_get(config, 'decoding', {}), 'greedy', {}) + max_symbols = cfg_get(greedy_cfg, 'max_symbols', 10) if isinstance(greedy_cfg, dict) else 10 + + return { + 'vocab_size': vocab_size + 1, + 'hidden_dim': hidden_dim, + 'num_layers': int(cfg_get(encoder_cfg, 'n_layers', cfg_get(encoder_cfg, 'num_hidden_layers', 0))), + 'attention_heads': attention_heads, + 'attention_kv_heads': int(cfg_get(encoder_cfg, 'n_kv_heads', attention_heads)), + 'attention_head_dim': int(hidden_dim // max(1, attention_heads)), + 'ffn_intermediate_dim': ff_intermediate, + 'context_length': int(cfg_get(encoder_cfg, 'pos_emb_max_len', cfg_get(encoder_cfg, 'max_position_embeddings', 0))), + 'layer_norm_eps': float(cfg_get(encoder_cfg, 'layer_norm_eps', cfg_get(encoder_cfg, 'norm_eps', 1e-5)) or 1e-5), + 'rope_theta': float(cfg_get(encoder_cfg, 'rope_theta', 0.0) or 0.0), + 'conv_kernel_size': int(cfg_get(encoder_cfg, 'conv_kernel_size', 9)), + 'subsampling_conv_kernel_size': int(cfg_get(encoder_cfg, 'subsampling_conv_kernel_size', 3)), + 'subsampling_conv_stride': int(cfg_get(encoder_cfg, 'subsampling_conv_stride', 2)), + 'subsampling_conv_channels': int(cfg_get(encoder_cfg, 'subsampling_conv_channels', 256)), + 'subsampling_factor': int(cfg_get(encoder_cfg, 'subsampling_factor', 8)), + 'num_mel_bins': int(cfg_get(preprocessor_cfg, 'features', cfg_get(encoder_cfg, 'feat_in', cfg_get(encoder_cfg, 'num_mel_bins', 128)))), + 'pad_token_id': int(cfg_get(config, 'pad_token_id', 0)), + 'encoder_hidden_act': cfg_get(encoder_cfg, 'activation', cfg_get(encoder_cfg, 'hidden_act', 'silu')), + 'predictor_hidden_dim': int(cfg_get(prediction_cfg, 'pred_hidden', cfg_get(config, 'decoder_hidden_size', 0))), + 'predictor_num_layers': int(cfg_get(prediction_cfg, 'pred_rnn_layers', cfg_get(config, 'num_decoder_layers', 0))), + 'rnnt_joint_dim': int(cfg_get(jointnet_cfg, 'joint_hidden', cfg_get(config, 'decoder_hidden_size', 0))), + 'rnnt_blank_id': blank_id, + 'prompt_dim': prompt_dim, + 'default_prompt_id': default_prompt_id, + 'prompt_dictionary': prompt_dictionary_text, + 'max_symbols_per_step': int(max_symbols), + 'sample_rate': int(cfg_get(preprocessor_cfg, 'sample_rate', cfg_get(config, 'sample_rate', 16000))), + 'tokenizer_type': 'bpe', + 'vocab_format': 'id_tab_token', + 'normalizer': 'metaspace', + 'decoder': 'replace_metaspace', + } + + def extract_complex_gemma_config(cfg, root_config): """Extract configuration parameters for Gemma3n and Gemma4 models.""" altup_num_inputs = int(cfg_get(cfg, 'altup_num_inputs', cfg_get(root_config, 'altup_num_inputs', 4))) diff --git a/python/cactus/convert/cactus_adapters/weight_patterns.py b/python/cactus/convert/cactus_adapters/weight_patterns.py index b71869fa3..ff41a0b89 100644 --- a/python/cactus/convert/cactus_adapters/weight_patterns.py +++ b/python/cactus/convert/cactus_adapters/weight_patterns.py @@ -179,6 +179,32 @@ (('joint.joint_net.2.bias', 'joint.joint_net.0.bias', 'joint.head.bias'), 'tdt_joint_out.bias'), ] +NEMOTRON_ASR_GLOBAL_WEIGHTS = [ + (('encoder.pre_encode.conv.0.weight', 'encoder.subsampling.layers.0.weight'), 'subsampling_conv0_weight.weights'), + (('encoder.pre_encode.conv.0.bias', 'encoder.subsampling.layers.0.bias'), 'subsampling_conv0_bias.bias'), + (('encoder.pre_encode.conv.2.weight', 'encoder.subsampling.layers.2.weight'), 'subsampling_depthwise1_weight.weights'), + (('encoder.pre_encode.conv.2.bias', 'encoder.subsampling.layers.2.bias'), 'subsampling_depthwise1_bias.bias'), + (('encoder.pre_encode.conv.3.weight', 'encoder.subsampling.layers.3.weight'), 'subsampling_pointwise1_weight.weights'), + (('encoder.pre_encode.conv.3.bias', 'encoder.subsampling.layers.3.bias'), 'subsampling_pointwise1_bias.bias'), + (('encoder.pre_encode.conv.5.weight', 'encoder.subsampling.layers.5.weight'), 'subsampling_depthwise2_weight.weights'), + (('encoder.pre_encode.conv.5.bias', 'encoder.subsampling.layers.5.bias'), 'subsampling_depthwise2_bias.bias'), + (('encoder.pre_encode.conv.6.weight', 'encoder.subsampling.layers.6.weight'), 'subsampling_pointwise2_weight.weights'), + (('encoder.pre_encode.conv.6.bias', 'encoder.subsampling.layers.6.bias'), 'subsampling_pointwise2_bias.bias'), + (('encoder.pre_encode.out.weight', 'encoder.subsampling.linear.weight'), 'subsampling_linear_weight.weights'), + (('encoder.pre_encode.out.bias', 'encoder.subsampling.linear.bias'), 'subsampling_linear_bias.bias'), + (('decoder.prediction.embed.weight', 'decoder.embedding.weight'), 'rnnt_predictor_embed.weights'), + (('joint.enc.weight', 'encoder_projector.weight'), 'rnnt_joint_enc.weights'), + (('joint.enc.bias', 'encoder_projector.bias'), 'rnnt_joint_enc.bias'), + (('joint.pred.weight', 'decoder.decoder_projector.weight'), 'rnnt_joint_pred.weights'), + (('joint.pred.bias', 'decoder.decoder_projector.bias'), 'rnnt_joint_pred.bias'), + (('joint.joint_net.2.weight', 'joint.joint_net.0.weight', 'joint.head.weight'), 'rnnt_joint_out.weights'), + (('joint.joint_net.2.bias', 'joint.joint_net.0.bias', 'joint.head.bias'), 'rnnt_joint_out.bias'), + (('prompt_kernel.0.weight',), 'rnnt_prompt_linear1.weights'), + (('prompt_kernel.0.bias',), 'rnnt_prompt_linear1.bias'), + (('prompt_kernel.2.weight',), 'rnnt_prompt_linear2.weights'), + (('prompt_kernel.2.bias',), 'rnnt_prompt_linear2.bias'), +] + PARAKEET_LAYER_WEIGHTS = [ (('feed_forward1.linear1.weight',), 'ff1_linear1.weights'), (('feed_forward1.linear1.bias',), 'ff1_linear1.bias'), diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 20f8dc94e..f48ab0b0c 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -31,7 +31,7 @@ from .export.validate import validate_qdq from .model_adapters.detection import SUPPORTED_FAMILIES, detect_family from .model_adapters.adapters import adapter_for_family -from .model_adapters.nemo import ensure_parakeet_tdt_nemo_source +from .model_adapters.nemo import ensure_nemo_asr_source from .quantization.cq import quantize_hadamard, quantize_orthogonal, write_cq_tensor from .compat import patch_transformers_import_compat @@ -50,7 +50,7 @@ def _load_hf(model_id_or_path: str, device: str): warnings.filterwarnings("ignore", message=".*You are using a model of type.*") for note in patch_transformers_import_compat(): print(f"note={note}") - nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) + nemo_export = ensure_nemo_asr_source(model_id_or_path, cache_dir=_hf_cache_dir()) if nemo_export is not None: model_id_or_path = nemo_export try: @@ -74,7 +74,7 @@ def _load_hf(model_id_or_path: str, device: str): processor = adapter.load_processor(model_id_or_path) model_cls = adapter.model_class(cfg) model_type = str(cfg_get(cfg, "model_type", "") or "").lower() - if isinstance(cfg, dict) and model_type == "parakeet_tdt": + if isinstance(cfg, dict) and model_type in {"parakeet_tdt", "nemotron_asr"}: return cfg, processor, None try: model = model_cls.from_pretrained( @@ -132,7 +132,7 @@ def _bits_for_component(component: str, args: argparse.Namespace) -> int: def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: - nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) + nemo_export = ensure_nemo_asr_source(model_id_or_path, cache_dir=_hf_cache_dir()) if nemo_export is not None: model_id_or_path = nemo_export root = Path(model_id_or_path) @@ -207,7 +207,7 @@ def _save_fallback_tensor(tensor, out_path: Path, precision: str, family: str) - if "conv_pointwise" in out_path.name and len(shape) == 3 and shape[2] == 1: save_pointwise_conv1d_int8_with_header(tensor, out_path) return - if "conv_depthwise.weights" in out_path.name and len(shape) == 3 and shape[1] == 1: + if "conv_depthwise" in out_path.name and len(shape) == 3 and shape[1] == 1: save_depthwise_conv_int8_with_header(tensor, out_path) return save_tensor_with_header(tensor, out_path, precision="INT8", model_type=family, allow_int8_bias=True) @@ -398,7 +398,7 @@ def convert(args: argparse.Namespace) -> None: out_dir.mkdir(parents=True, exist_ok=True) cfg, processor, model = _load_hf(args.model, args.device) - runtime_source = ensure_parakeet_tdt_nemo_source(args.model, cache_dir=_hf_cache_dir()) or args.model + runtime_source = ensure_nemo_asr_source(args.model, cache_dir=_hf_cache_dir()) or args.model family = detect_family(cfg, args.model_family) adapter = adapter_for_family(family) checkpoint_state = _load_checkpoint_state_dict(args.model) diff --git a/python/cactus/convert/export/files.py b/python/cactus/convert/export/files.py index 4529a4db8..e764727f7 100644 --- a/python/cactus/convert/export/files.py +++ b/python/cactus/convert/export/files.py @@ -12,6 +12,7 @@ "tokenizer.json", "tokenizer.model", "vocab.txt", + "tokenizer_config.txt", "tokenizer_config.json", "special_tokens_map.json", "added_tokens.json", diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index a29592d6a..59549c639 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -18,6 +18,7 @@ extract_complex_gemma_config, extract_parakeet_config, extract_parakeet_tdt_config, + extract_nemotron_asr_config, extract_vision_config, extract_whisper_config, ) @@ -432,6 +433,16 @@ def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: return NormalizedState(augmented, provenance) +class NemotronASRAdapter(ParakeetTDTAdapter): + family = "nemotron_asr" + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return extract_nemotron_asr_config(cfg) + + def runtime_model_type(self) -> str: + return "nemotron_asr" + + class Lfm2Adapter(FamilyAdapter): family = "lfm2" @@ -514,6 +525,7 @@ def scale_factor(self, output_name: str) -> float: "whisper": WhisperAdapter(), "parakeet": ParakeetAdapter(), "parakeet_tdt": ParakeetTDTAdapter(), + "nemotron_asr": NemotronASRAdapter(), } diff --git a/python/cactus/convert/model_adapters/detection.py b/python/cactus/convert/model_adapters/detection.py index 69cb2aa09..e90b697cf 100644 --- a/python/cactus/convert/model_adapters/detection.py +++ b/python/cactus/convert/model_adapters/detection.py @@ -3,7 +3,7 @@ from ..cactus_adapters.config_utils import cfg_get, detect_model_type -SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "generic"} +SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "nemotron_asr", "moonshine", "nomic", "needle", "generic"} def detect_family(config, requested: str = "auto") -> str: @@ -17,6 +17,6 @@ def detect_family(config, requested: str = "auto") -> str: if "gemma3" in raw: return "gemma3" return "generic" - if detected in {"gemma4", "qwen", "lfm2", "whisper", "moonshine", "parakeet", "parakeet_tdt", "nomic", "needle"}: + if detected in {"gemma4", "qwen", "lfm2", "whisper", "moonshine", "parakeet", "parakeet_tdt", "nemotron_asr", "nomic", "needle"}: return detected return "generic" diff --git a/python/cactus/convert/model_adapters/naming.py b/python/cactus/convert/model_adapters/naming.py index b44a4a072..bed43f25e 100644 --- a/python/cactus/convert/model_adapters/naming.py +++ b/python/cactus/convert/model_adapters/naming.py @@ -328,6 +328,11 @@ def _global_match(name: str, family: str) -> str | None: found = _candidate_table_match(name, table) if found: return found + elif family == "nemotron_asr": + for table in (wp.PARAKEET_GLOBAL_WEIGHTS, wp.NEMOTRON_ASR_GLOBAL_WEIGHTS): + found = _candidate_table_match(name, table) + if found: + return found if name in {"wte.weight", "word_embeddings.weight"}: return "token_embeddings.weights" if name in {"wpe.weight", "position_embeddings.weight"}: @@ -357,12 +362,16 @@ def _global_match(name: str, family: str) -> str | None: predictor_name = _parakeet_tdt_predictor_match(name) if predictor_name: return predictor_name + if family == "nemotron_asr": + predictor_name = _parakeet_tdt_predictor_match(name) + if predictor_name: + return predictor_name.replace("tdt_", "rnnt_", 1) return None def component_for_name(name: str, output_name: str | None = None) -> str: joined = f"{name} {output_name or ''}".lower() - if "parakeet" in joined or "tdt_" in joined or "ctc_" in joined: + if "parakeet" in joined or "nemotron" in joined or "tdt_" in joined or "rnnt_" in joined or "ctc_" in joined: return "transcription" if "audio" in joined or "encoder.conv" in joined or "subsample" in joined: return "audio" @@ -376,7 +385,7 @@ def component_for_name(name: str, output_name: str | None = None) -> str: def _component_for_family(name: str, output_name: str | None, family: str) -> str: - if family in {"parakeet", "parakeet_tdt"}: + if family in {"parakeet", "parakeet_tdt", "nemotron_asr"}: return "transcription" if family == "whisper": out = output_name or "" diff --git a/python/cactus/convert/model_adapters/nemo.py b/python/cactus/convert/model_adapters/nemo.py index af828a65c..08ab78e45 100644 --- a/python/cactus/convert/model_adapters/nemo.py +++ b/python/cactus/convert/model_adapters/nemo.py @@ -6,6 +6,8 @@ from pathlib import Path from typing import Any +from ..cactus_adapters.config_utils import prompt_dim_from_dictionary + _NEMO_EXPORT_FILES = { "model_config.yaml", @@ -14,7 +16,7 @@ } -def ensure_parakeet_tdt_nemo_source( +def ensure_nemo_asr_source( model_id_or_path: str, *, token: str | None = None, @@ -25,7 +27,12 @@ def ensure_parakeet_tdt_nemo_source( return None out = nemo_path.parent / f"{nemo_path.stem}-cactus-hf" - if (out / "config.json").exists() and (out / "pytorch_model.bin").exists() and (out / "vocab.txt").exists(): + if ( + (out / "config.json").exists() + and (out / "pytorch_model.bin").exists() + and (out / "vocab.txt").exists() + and (out / "tokenizer_config.txt").exists() + ): return str(out) tmp = out.with_name(f"{out.name}.tmp") @@ -49,10 +56,10 @@ def ensure_parakeet_tdt_nemo_source( if not weights.exists() or not config: raise RuntimeError(f"{nemo_path} is missing model_config.yaml or model_weights.ckpt") - hf_config = _parakeet_tdt_config_from_nemo(config) + hf_config = _asr_config_from_nemo(config, model_id_or_path=model_id_or_path) (tmp / "config.json").write_text(json.dumps(hf_config, indent=2, sort_keys=True), encoding="utf-8") weights.rename(tmp / "pytorch_model.bin") - _write_parakeet_tdt_tokenizer_files(tmp, hf_config) + _write_asr_tokenizer_files(tmp, hf_config) if out.exists(): shutil.rmtree(out) @@ -60,6 +67,15 @@ def ensure_parakeet_tdt_nemo_source( return str(out) +def ensure_parakeet_tdt_nemo_source( + model_id_or_path: str, + *, + token: str | None = None, + cache_dir: str | None = None, +) -> str | None: + return ensure_nemo_asr_source(model_id_or_path, token=token, cache_dir=cache_dir) + + def _find_single_nemo(model_id_or_path: str, *, token: str | None, cache_dir: str | None) -> Path | None: path = Path(model_id_or_path) if path.suffix == ".nemo" and path.is_file(): @@ -92,7 +108,7 @@ def _read_yaml(path: Path) -> dict[str, Any]: try: import yaml except Exception as exc: - raise RuntimeError("PyYAML is required to import Parakeet TDT .nemo checkpoints") from exc + raise RuntimeError("PyYAML is required to import NeMo ASR .nemo checkpoints") from exc if not path.exists(): return {} loaded = yaml.safe_load(path.read_text(encoding="utf-8")) @@ -104,6 +120,33 @@ def _get(mapping: dict[str, Any], key: str, default: Any = None) -> Any: return default if value is None else value +def _asr_config_from_nemo(root: dict[str, Any], *, model_id_or_path: str = "") -> dict[str, Any]: + if _is_nemotron_asr_config(root, model_id_or_path=model_id_or_path): + return _nemotron_asr_config_from_nemo(root) + return _parakeet_tdt_config_from_nemo(root) + + +def _is_nemotron_asr_config(root: dict[str, Any], *, model_id_or_path: str = "") -> bool: + lowered_id = model_id_or_path.lower() + target = str(_get(root, "target", "") or "").lower() + model_defaults = _get(root, "model_defaults", {}) or {} + has_prompt = ( + "nemotron" in lowered_id + or "withprompt" in target + or "prompt" in target + or bool(_get(model_defaults, "prompt_dictionary", None)) + or int(_get(root, "num_prompts", _get(model_defaults, "num_prompts", 0)) or 0) > 0 + ) + loss = _get(root, "loss", {}) or {} + decoding = _get(root, "decoding", {}) or {} + is_tdt = ( + str(_get(loss, "loss_name", "") or "").lower() == "tdt" + or str(_get(decoding, "model_type", "") or "").lower() == "tdt" + or bool(_get(model_defaults, "tdt_durations", None)) + ) + return has_prompt and not is_tdt + + def _parakeet_tdt_config_from_nemo(root: dict[str, Any]) -> dict[str, Any]: decoder = _get(root, "decoder", {}) or {} prednet = _get(decoder, "prednet", _get(decoder, "prediction", {})) or {} @@ -119,7 +162,39 @@ def _parakeet_tdt_config_from_nemo(root: dict[str, Any]) -> dict[str, Any]: return config -def _write_parakeet_tdt_tokenizer_files(root: Path, config: dict[str, Any]) -> None: +def _nemotron_asr_config_from_nemo(root: dict[str, Any]) -> dict[str, Any]: + decoder = _get(root, "decoder", {}) or {} + prednet = _get(decoder, "prednet", _get(decoder, "prediction", {})) or {} + joint = _get(root, "joint", {}) or {} + jointnet = _get(joint, "jointnet", {}) or {} + model_defaults = _get(root, "model_defaults", {}) or {} + labels = [str(token) for token in (_get(joint, "vocabulary", _get(root, "labels", [])) or [])] + decoder_vocab = int(_get(decoder, "vocab_size", len(labels))) + prompt_dictionary = _get(model_defaults, "prompt_dictionary", {}) or {} + prompt_dim = prompt_dim_from_dictionary( + prompt_dictionary, + _get(root, "prompt_dim", 0), + _get(root, "num_prompts", 0), + _get(model_defaults, "num_prompts", 0), + ) + default_prompt_id = int(prompt_dictionary.get("auto", 0)) if isinstance(prompt_dictionary, dict) else 0 + + config = dict(root) + config.update( + architectures=["NemotronASRForRNNT"], + model_type="nemotron_asr", + labels=labels, + prompt_dictionary={str(k): int(v) for k, v in prompt_dictionary.items()} if isinstance(prompt_dictionary, dict) else {}, + prompt_dim=prompt_dim, + default_prompt_id=default_prompt_id, + rnnt_blank_id=decoder_vocab, + ) + config["decoder"] = {**decoder, "prediction": dict(prednet), "vocab_size": decoder_vocab} + config["joint"] = {**joint, "jointnet": dict(jointnet), "vocabulary": labels} + return config + + +def _write_asr_tokenizer_files(root: Path, config: dict[str, Any]) -> None: tokenizer_model = next(iter(sorted(root.glob("*_tokenizer.model"))), None) if tokenizer_model is not None and not (root / "tokenizer.model").exists(): tokenizer_model.rename(root / "tokenizer.model") @@ -134,7 +209,26 @@ def _write_parakeet_tdt_tokenizer_files(root: Path, config: dict[str, Any]) -> N if nemo_vocab is not None and not (root / "vocab.txt").exists(): nemo_vocab.rename(root / "vocab.txt") + (root / "merges.txt").write_text("#version: 0.2\n", encoding="utf-8") + (root / "tokenizer_config.txt").write_text( + "\n".join(( + "tokenizer_type=bpe", + "vocab_format=id_tab_token", + "normalizer=metaspace", + "decoder=replace_metaspace", + "byte_fallback=false", + "", + )), + encoding="utf-8", + ) (root / "tokenizer_config.json").write_text( - json.dumps({"model_type": "parakeet_tdt", "tokenizer_class": "SentencePieceProcessor"}, indent=2), + json.dumps({ + "model_type": config.get("model_type", "parakeet_tdt"), + "tokenizer_class": "SentencePieceProcessor", + "tokenizer_type": "bpe", + "vocab_format": "id_tab_token", + "normalizer": "metaspace", + "decoder": "replace_metaspace", + }, indent=2), encoding="utf-8", ) diff --git a/python/cactus/convert/model_adapters/policy.py b/python/cactus/convert/model_adapters/policy.py index 578218aef..6dbf88a75 100644 --- a/python/cactus/convert/model_adapters/policy.py +++ b/python/cactus/convert/model_adapters/policy.py @@ -27,24 +27,22 @@ def policy_for_tensor(match: NameMatch, shape: tuple[int, ...], user_bits: int, return TensorPolicy("fallback", "FP16", None, component, False, "none", "norm tensor") if family == "gemma4" and name == "model.embed_vision.embedding_projection.weight": return TensorPolicy("fallback", "FP16", None, component, False, "none", "vision embedding projection scale-sensitive") - if family in {"parakeet", "parakeet_tdt"} and out.endswith(".bias") and ( + if family in {"parakeet", "parakeet_tdt", "nemotron_asr"} and out.endswith(".bias") and ( "conv_" in out or out.startswith("subsampling_") or out.startswith("ctc_head_") ): return TensorPolicy("fallback", "FP16", None, component, False, "none", "conv bias tensor") - if family in {"parakeet", "parakeet_tdt"} and "lstm" in out.lower(): + if family in {"parakeet", "parakeet_tdt", "nemotron_asr"} and "lstm" in out.lower(): return TensorPolicy("fallback", "FP16", None, component, False, "none", "lstm recurrent tensor") if family == "parakeet_tdt" and out.startswith("tdt_"): return TensorPolicy("fallback", "FP16", None, component, False, "none", "tdt decoder tensor") - if family in {"parakeet", "parakeet_tdt"} and "self_attn_bias_" in out: + if family == "nemotron_asr" and out.startswith("rnnt_predictor_embed."): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "rnnt predictor embedding tensor") + if family in {"parakeet", "parakeet_tdt", "nemotron_asr"} and "self_attn_bias_" in out: return TensorPolicy("fallback", "FP16", None, component, False, "none", "relative attention bias tensor") - if family in {"parakeet", "parakeet_tdt"} and "conv_pointwise" in out and len(shape) == 3 and shape[2] == 1: + if family in {"parakeet", "parakeet_tdt", "nemotron_asr"} and "conv_pointwise" in out and len(shape) == 3 and shape[2] == 1: return TensorPolicy("fallback", "INT8", 8, component, False, "none", "pointwise conv tensor") - if "conv_depthwise.weights" in out and len(shape) == 3 and shape[1] == 1: + if ("conv_depthwise.weights" in out or "conv_depthwise_conv.weights" in out) and len(shape) == 3 and shape[1] == 1: return TensorPolicy("fallback", "INT8", 8, component, False, "none", "depthwise conv tensor") - if family in {"parakeet", "parakeet_tdt"} and out.startswith("layer_") and ( - "conv_pointwise" in out or "conv_depthwise" in out - ): - return TensorPolicy("fallback", "FP16", None, component, False, "none", "conformer conv tensor") if family == "whisper" and out.startswith("encoder.layer_"): return TensorPolicy("fallback", "FP16", None, component, False, "none", "whisper encoder tensor") if family == "gemma4" and component == "audio" and (name.endswith(".bias") or out.endswith(".bias") or ".bias." in out): diff --git a/python/cactus/convert/tests/test_cq.py b/python/cactus/convert/tests/test_cq.py index 2cd816e25..bc87f9d7e 100644 --- a/python/cactus/convert/tests/test_cq.py +++ b/python/cactus/convert/tests/test_cq.py @@ -6,6 +6,7 @@ import numpy as np import torch +from cactus.convert.cli import _save_fallback_tensor from cactus.convert.cactus_adapters.tensor_io import ( FLAG_HAS_SCALES, GROUP_SIZE, @@ -133,6 +134,26 @@ def test_depthwise_conv_int8_preserves_kernel_shape(tmp_path): assert num_groups == 1 +def test_fallback_tensor_routes_nemotron_depthwise_conv_name(tmp_path): + weight = torch.randn(4, 1, 9) + out = tmp_path / "layer_0_conv_depthwise_conv.weights" + _save_fallback_tensor(weight, out, "INT8", "nemotron_asr") + raw = out.read_bytes() + magic, flags, alignment, ndim = struct.unpack_from("<4sIII", raw, 0) + dims = struct.unpack_from(""], + "model_defaults": { + "num_prompts": 128, + "prompt_dictionary": {"auto": 101, "en-US": 0}, + "pred_hidden": 640, + "joint_hidden": 640, + }, + "preprocessor": {"sample_rate": 16000, "features": 128}, + "encoder": { + "d_model": 1024, + "n_layers": 24, + "n_heads": 8, + "ff_expansion_factor": 4.0, + "conv_kernel_size": 9, + "subsampling_conv_channels": 256, + "subsampling_factor": 8, + "activation": "silu", + }, + "decoder": { + "vocab_size": 3, + "prediction": {"pred_hidden": 640, "pred_rnn_layers": 2}, + }, + "joint": { + "num_classes": 3, + "jointnet": {"joint_hidden": 640}, + "vocabulary": ["a", "b", ""], + }, + "decoding": {"greedy": {"max_symbols": 10}}, + } + assert detect_family(cfg) == "nemotron_asr" + extracted = extract_nemotron_asr_config(cfg) + assert extracted["vocab_size"] == 4 + assert extracted["num_mel_bins"] == 128 + assert extracted["predictor_hidden_dim"] == 640 + assert extracted["predictor_num_layers"] == 2 + assert extracted["rnnt_blank_id"] == 3 + assert extracted["prompt_dim"] == 128 + assert extracted["default_prompt_id"] == 101 + assert "auto:101" in extracted["prompt_dictionary"] + assert "en-US:0" in extracted["prompt_dictionary"] + plan = infer_component_plan_from_config(cfg, model_id="nvidia/nemotron-3.5-asr-streaming-0.6b") + assert plan is not None + assert plan.task == "rnnt_transcription" + assert plan.components == ("audio_encoder", "decoder") + + +def test_prompt_tdt_config_detection_stays_parakeet_tdt(): + cfg = { + "target": "nemo.collections.asr.models.rnnt_bpe_models_prompt.EncDecRNNTBPEModelWithPrompt", + "model_defaults": {"prompt_dictionary": {"auto": 101}, "tdt_durations": [0, 1, 2]}, + "loss": {"loss_name": "tdt"}, + "decoder": {"vocab_size": 3}, + "joint": {"num_classes": 3, "vocabulary": ["a", "b", ""]}, + } + assert detect_family(cfg) == "parakeet_tdt" + + +def test_nemotron_prompt_dim_preserves_sparse_prompt_ids(): + cfg = { + "model_type": "nemotron_asr", + "labels": ["a", "b"], + "model_defaults": {"prompt_dictionary": {"auto": 101, "en-US": 0}}, + "preprocessor": {"sample_rate": 16000, "features": 128}, + "encoder": {"d_model": 1024, "n_layers": 1, "n_heads": 8}, + "decoder": {"vocab_size": 2, "prediction": {"pred_hidden": 640, "pred_rnn_layers": 2}}, + "joint": {"num_classes": 2, "jointnet": {"joint_hidden": 640}, "vocabulary": ["a", "b"]}, + } + extracted = extract_nemotron_asr_config(cfg) + assert extracted["prompt_dim"] == 102 + assert extracted["default_prompt_id"] == 101 + + def test_whisper_hf_config_extraction(): cfg = { "model_type": "whisper", @@ -226,12 +358,17 @@ def test_parakeet_tdt_lstm_bias_normalization_records_provenance(): assert provenance.transform == "parakeet_tdt_lstm_bias_sum" -def test_policy_parakeet_pointwise_conv_int8_and_conv_bias_fp16(): +def test_policy_parakeet_conformer_conv_int8_and_conv_bias_fp16(): weight = cactus_name_for_tensor("encoder.layers.0.conv.pointwise_conv1.weight", "parakeet_tdt", 24) weight_policy = policy_for_tensor(weight, (2048, 1024, 1), 4, "parakeet_tdt") assert weight_policy.precision == "INT8" assert weight_policy.fallback_reason == "pointwise conv tensor" + depthwise = cactus_name_for_tensor("encoder.layers.0.conv.depthwise_conv.weight", "parakeet_tdt", 24) + depthwise_policy = policy_for_tensor(depthwise, (1024, 1, 9), 4, "parakeet_tdt") + assert depthwise_policy.precision == "INT8" + assert depthwise_policy.fallback_reason == "depthwise conv tensor" + bias = cactus_name_for_tensor("encoder.layers.0.conv.pointwise_conv1.bias", "parakeet_tdt", 24) bias_policy = policy_for_tensor(bias, (2048,), 4, "parakeet_tdt") assert bias_policy.precision == "FP16" diff --git a/python/cactus/server.py b/python/cactus/server.py index c8f9206a6..56f60ee37 100644 --- a/python/cactus/server.py +++ b/python/cactus/server.py @@ -30,7 +30,7 @@ LOGGER = logging.getLogger(__name__) LLM_MODEL_TYPES = {"gemma", "gemma3n", "gemma4", "lfm2", "qwen", "qwen3p5", "needle", "youtu"} -STT_MODEL_TYPES = {"whisper", "parakeet_tdt", "parakeet-tdt"} +STT_MODEL_TYPES = {"whisper", "parakeet_tdt", "parakeet-tdt", "nemotron_asr", "nemotron-asr"} @dataclass(frozen=True) diff --git a/python/cactus/transpile/README.md b/python/cactus/transpile/README.md index a0a059658..dff611bed 100644 --- a/python/cactus/transpile/README.md +++ b/python/cactus/transpile/README.md @@ -118,6 +118,7 @@ Today the main task shapes are: - `ctc_logits` - `encoder_hidden_states` - `tdt_transcription` +- `rnnt_transcription` ## 2. Prepare Example Inputs @@ -137,7 +138,7 @@ if task == "causal_lm_logits": prepared = _prepare_text_inputs(...) elif task == "multimodal_causal_lm_logits": prepared = _prepare_gemma4_multimodal_inputs(...) -elif task == "tdt_transcription": +elif task in {"tdt_transcription", "rnnt_transcription"}: prepared = _prepare_audio_inputs(...) else: prepared = _prepare_audio_inputs(...) @@ -928,7 +929,7 @@ cactus transpile \ --weights-dir /path/to/converted_weights ``` -### Parakeet TDT +### Parakeet TDT and Nemotron ASR ```bash cactus transpile \ @@ -937,6 +938,13 @@ cactus transpile \ --weights-dir /path/to/converted_weights ``` +```bash +cactus transpile \ + --audio-file /path/to/sample.wav \ + --task rnnt_transcription \ + --weights-dir /path/to/converted_weights +``` + If you want the full direct script help: ```bash @@ -1000,7 +1008,7 @@ As of the current code: - Saved bundle execution is implemented for `causal_lm_logits`, `multimodal_causal_lm_logits`, `encoder_hidden_states`, and - `tdt_transcription`. + `tdt_transcription` / `rnnt_transcription`. - `ctc_logits` can be transpiled, but the runtime does not currently implement a saved-bundle executor for that task. - The component pipeline is model-family-specific; it is not a universal split-IR diff --git a/python/cactus/transpile/audio_preprocess.py b/python/cactus/transpile/audio_preprocess.py index 5e6d0aeea..671100088 100644 --- a/python/cactus/transpile/audio_preprocess.py +++ b/python/cactus/transpile/audio_preprocess.py @@ -178,7 +178,8 @@ def generic_log_mel_features( frame_length: int, preemphasis: float | None = None, mel_floor: np.float32 = _PARAKEET_LOG_FLOOR, - normalize_active_frames_only: bool = True, + mel_floor_additive: bool = False, + normalize_active_frames_only: bool | None = True, ) -> tuple[np.ndarray, int]: waveform_tensor = torch.from_numpy(waveform).to(torch.float32).unsqueeze(0) if preemphasis is not None and abs(preemphasis) > 0.0: @@ -208,9 +209,14 @@ def generic_log_mel_features( feature_length = max(1, int(waveform.shape[0] // hop_length)) mel_filters = torch.from_numpy(_mel_filter_bank(num_mels, sample_rate, n_fft)).to(torch.float32) mel_spec = power.transpose(1, 2) @ mel_filters - mel_spec = torch.log(torch.clamp(mel_spec, min=float(mel_floor))) + if mel_floor_additive: + mel_spec = torch.log(mel_spec + float(mel_floor)) + else: + mel_spec = torch.log(torch.clamp(mel_spec, min=float(mel_floor))) - if normalize_active_frames_only: + if normalize_active_frames_only is None: + pass + elif normalize_active_frames_only: attention_mask = ( torch.arange(mel_spec.shape[1], dtype=torch.long, device=mel_spec.device).unsqueeze(0) < feature_length diff --git a/python/cactus/transpile/canonicalize/cleanup.py b/python/cactus/transpile/canonicalize/cleanup.py index b7f4de2aa..90308bb10 100644 --- a/python/cactus/transpile/canonicalize/cleanup.py +++ b/python/cactus/transpile/canonicalize/cleanup.py @@ -128,6 +128,7 @@ "attention", "attention_block", "conv_module", + "causal_layer_norm_conv_module", "dense_mlp_tq_fused", "rope", } diff --git a/python/cactus/transpile/component_partition.py b/python/cactus/transpile/component_partition.py index a0a76dd32..14674fc23 100644 --- a/python/cactus/transpile/component_partition.py +++ b/python/cactus/transpile/component_partition.py @@ -141,7 +141,7 @@ def _task_default_component(*, family: str, task: str) -> str: if task in {"ctc_logits", "encoder_hidden_states"}: return COMPONENT_AUDIO_ENCODER profile = profile_for_family(family) - if profile is not None and profile.default_task == "tdt_transcription": + if profile is not None and profile.default_task in {"tdt_transcription", "rnnt_transcription"}: return COMPONENT_AUDIO_ENCODER return COMPONENT_UNSPECIFIED diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py index 11f51cd94..fc8a3a3d2 100644 --- a/python/cactus/transpile/component_pipeline.py +++ b/python/cactus/transpile/component_pipeline.py @@ -33,6 +33,7 @@ class ComponentModuleSpec: npu_runtime_input_count: int = 1 npu_reparam: Callable[[torch.nn.Module], AbstractContextManager] | None = None npu_example_inputs: tuple[torch.Tensor, ...] | None = None + npu_variants: tuple[dict[str, Any], ...] = field(default_factory=tuple) @dataclass diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index c9e63cabd..4898db26b 100644 --- a/python/cactus/transpile/component_plan.py +++ b/python/cactus/transpile/component_plan.py @@ -122,6 +122,16 @@ def _is_tdt_config(config: Mapping[str, object], model_type: str, lowered_id: st return model_type == "parakeet_tdt" or "parakeet-tdt" in lowered_id +def _is_rnnt_asr_config(config: Mapping[str, object], model_type: str, lowered_id: str) -> bool: + if model_type == "nemotron_asr" or "nemotron-3.5-asr" in lowered_id or "nemotron-asr" in lowered_id: + return True + target = str(config.get("target", "") or "").lower() + model_defaults = config.get("model_defaults") + if isinstance(model_defaults, Mapping) and model_defaults.get("prompt_dictionary"): + return True + return "rnnt" in target and "prompt" in target + + def _looks_like_vision_language_model( *, model_type: str, @@ -171,6 +181,14 @@ def infer_component_plan_from_config( force_component_pipeline=True, ) + if _is_rnnt_asr_config(config, model_type, lowered_id): + return ComponentPlan( + task="rnnt_transcription", + components=("audio_encoder", "decoder"), + needs_audio=True, + force_component_pipeline=True, + ) + if model_type == "whisper" or "whisper" in lowered_id or any("whisper" in value for value in architectures): return ComponentPlan( task="seq2seq_transcription", diff --git a/python/cactus/transpile/fusion/__init__.py b/python/cactus/transpile/fusion/__init__.py index 463478d1c..63ceee3ce 100644 --- a/python/cactus/transpile/fusion/__init__.py +++ b/python/cactus/transpile/fusion/__init__.py @@ -1,4 +1,6 @@ from cactus.transpile.fusion.conv import ConvModuleMatch +from cactus.transpile.fusion.conv import CausalLayerNormConvModuleMatch +from cactus.transpile.fusion.conv import match_causal_layer_norm_conv_module from cactus.transpile.fusion.conv import match_conv_module from cactus.transpile.fusion.attention import AttentionBlockMatch from cactus.transpile.fusion.attention import AttentionMatch @@ -23,6 +25,7 @@ __all__ = [ "ConvModuleMatch", + "CausalLayerNormConvModuleMatch", "AttentionBlockMatch", "AttentionMatch", "SelfAttentionBlockMatch", @@ -37,6 +40,7 @@ "match_attention_block", "match_self_attention_block", "match_conv_module", + "match_causal_layer_norm_conv_module", "match_gated_deltanet", "match_gated_mlp", "match_linear", diff --git a/python/cactus/transpile/fusion/conv.py b/python/cactus/transpile/fusion/conv.py index 265501f9a..a6b12e77f 100644 --- a/python/cactus/transpile/fusion/conv.py +++ b/python/cactus/transpile/fusion/conv.py @@ -28,6 +28,22 @@ class ConvModuleMatch: node_ids: tuple[str, ...] +@dataclass(frozen=True) +class CausalLayerNormConvModuleMatch: + input_value_id: str + pointwise1_weight_value_id: str + pointwise1_bias_value_id: str | None + depthwise_weight_value_id: str + depthwise_bias_value_id: str | None + layer_norm_weight_value_id: str + layer_norm_bias_value_id: str + pointwise2_weight_value_id: str + pointwise2_bias_value_id: str | None + eps: float + depthwise_kernel_size: int + node_ids: tuple[str, ...] + + def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None: if node.op != "permute" or len(node.inputs) != 1 or len(node.outputs) != 1: return None @@ -69,6 +85,11 @@ def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None: depthwise_weight_shape = graph.values.get(depthwise.inputs[1]).shape if depthwise.inputs[1] in graph.values else None if depthwise_weight_shape is None or len(depthwise_weight_shape) != 3: return None + if ( + int(depthwise.attrs.get("groups", 1)) != int(depthwise_weight_shape[0]) + or int(depthwise_weight_shape[1]) != 1 + ): + return None depthwise_kernel_size = int(depthwise_weight_shape[2]) return ConvModuleMatch( @@ -99,6 +120,108 @@ def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None: ) +def match_causal_layer_norm_conv_module(graph: IRGraph, node: IRNode) -> CausalLayerNormConvModuleMatch | None: + if node.op != "permute" or len(node.inputs) != 1 or len(node.outputs) != 1: + return None + if tuple(node.attrs.get("permutation", ())) != (0, 2, 1): + return None + + pointwise2 = producer(graph, strip_passthrough(graph, node.inputs[0])) + if not _is_pointwise_conv1d(pointwise2): + return None + + silu = producer(graph, strip_passthrough(graph, pointwise2.inputs[0])) + if silu is None or silu.op != "silu" or len(silu.inputs) != 1: + return None + + layer_norm_input_permute = producer(graph, strip_passthrough(graph, silu.inputs[0])) + if ( + layer_norm_input_permute is None + or layer_norm_input_permute.op != "permute" + or tuple(layer_norm_input_permute.attrs.get("permutation", ())) != (0, 2, 1) + or len(layer_norm_input_permute.inputs) != 1 + ): + return None + + layer_norm = producer(graph, strip_passthrough(graph, layer_norm_input_permute.inputs[0])) + if layer_norm is None or layer_norm.op != "layer_norm" or len(layer_norm.inputs) != 3: + return None + + depthwise_output_permute = producer(graph, strip_passthrough(graph, layer_norm.inputs[0])) + if ( + depthwise_output_permute is None + or depthwise_output_permute.op != "permute" + or tuple(depthwise_output_permute.attrs.get("permutation", ())) != (0, 2, 1) + or len(depthwise_output_permute.inputs) != 1 + ): + return None + + depthwise = producer(graph, strip_passthrough(graph, depthwise_output_permute.inputs[0])) + if not _is_causal_depthwise_conv1d(depthwise): + return None + + pad = producer(graph, strip_passthrough(graph, depthwise.inputs[0])) + if pad is None or pad.op != "pad" or len(pad.inputs) != 1: + return None + if pad.attrs.get("mode") != "constant" or float(pad.attrs.get("value", 0.0)) != 0.0: + return None + + depthwise_weight_shape = graph.values.get(depthwise.inputs[1]).shape if depthwise.inputs[1] in graph.values else None + if depthwise_weight_shape is None or len(depthwise_weight_shape) != 3: + return None + if ( + int(depthwise.attrs.get("groups", 1)) != int(depthwise_weight_shape[0]) + or int(depthwise_weight_shape[1]) != 1 + ): + return None + depthwise_kernel_size = int(depthwise_weight_shape[2]) + if tuple(pad.attrs.get("pads", ())) != (depthwise_kernel_size - 1, 0): + return None + + glu = producer(graph, strip_passthrough(graph, pad.inputs[0])) + if glu is None or glu.op != "glu" or len(glu.inputs) != 1: + return None + if int(glu.attrs.get("axis", -1)) not in {1, -2}: + return None + + pointwise1 = producer(graph, strip_passthrough(graph, glu.inputs[0])) + if not _is_pointwise_conv1d(pointwise1): + return None + + input_permute = producer(graph, strip_passthrough(graph, pointwise1.inputs[0])) + if input_permute is None or input_permute.op != "permute" or len(input_permute.inputs) != 1: + return None + if tuple(input_permute.attrs.get("permutation", ())) != (0, 2, 1): + return None + + return CausalLayerNormConvModuleMatch( + input_value_id=strip_passthrough(graph, input_permute.inputs[0]), + pointwise1_weight_value_id=pointwise1.inputs[1], + pointwise1_bias_value_id=pointwise1.inputs[2] if len(pointwise1.inputs) > 2 else None, + depthwise_weight_value_id=depthwise.inputs[1], + depthwise_bias_value_id=depthwise.inputs[2] if len(depthwise.inputs) > 2 else None, + layer_norm_weight_value_id=layer_norm.inputs[1], + layer_norm_bias_value_id=layer_norm.inputs[2], + pointwise2_weight_value_id=pointwise2.inputs[1], + pointwise2_bias_value_id=pointwise2.inputs[2] if len(pointwise2.inputs) > 2 else None, + eps=float(layer_norm.attrs.get("eps", 1e-5)), + depthwise_kernel_size=depthwise_kernel_size, + node_ids=collect_node_ids( + node, + pointwise2, + silu, + layer_norm_input_permute, + layer_norm, + depthwise_output_permute, + depthwise, + pad, + glu, + pointwise1, + input_permute, + ), + ) + + def _is_pointwise_conv1d(node: IRNode | None) -> bool: return ( node is not None @@ -121,3 +244,9 @@ def _is_same_depthwise_conv1d(node: IRNode | None) -> bool: if stride != 1 or dilation != 1: return False return groups > 1 and padding >= 0 + + +def _is_causal_depthwise_conv1d(node: IRNode | None) -> bool: + if not _is_same_depthwise_conv1d(node): + return False + return int(node.attrs.get("padding", 0)) == 0 diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index 4151ccb4a..c33c3d3ff 100644 --- a/python/cactus/transpile/hf_model.py +++ b/python/cactus/transpile/hf_model.py @@ -26,7 +26,7 @@ from cactus.transpile.runtime_compat import Graph from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header -from cactus.convert.model_adapters.nemo import ensure_parakeet_tdt_nemo_source +from cactus.convert.model_adapters.nemo import ensure_nemo_asr_source from cactus.transpile.audio_preprocess import generic_log_mel_features as _generic_log_mel_features from cactus.transpile.audio_preprocess import load_audio_waveform as _load_audio_waveform from cactus.transpile.audio_preprocess import prepare_cactus_audio_features @@ -59,6 +59,10 @@ from cactus.transpile.tdt_runtime import greedy_decode_parakeet_tdt_token_ids from cactus.transpile.tdt_runtime import load_tdt_local_model from cactus.transpile.tdt_runtime import prepare_parakeet_tdt_audio_features +from cactus.transpile.rnnt_runtime import greedy_decode_nemotron_asr_token_ids +from cactus.transpile.rnnt_runtime import load_nemotron_asr_local_model +from cactus.transpile.rnnt_runtime import _valid_hidden_frame_count +from cactus.transpile.rnnt_runtime import prepare_nemotron_asr_audio_features from cactus.transpile.weight_compat import ensure_binding_compatible _DEFAULT_CAUSAL_PROMPT = "The capital of France is" @@ -79,7 +83,7 @@ def _ensure_transformers_supports_model_type(model_type: str) -> str | None: def _resolve_local_snapshot(model_id_or_path: str) -> str | None: explicit = Path(model_id_or_path) if explicit.exists(): - nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path) + nemo_export = ensure_nemo_asr_source(model_id_or_path) if nemo_export is not None: return nemo_export return str(explicit) @@ -98,7 +102,7 @@ def _resolve_local_snapshot(model_id_or_path: str) -> str | None: if not snapshots: return None snapshot = str(snapshots[-1]) - nemo_export = ensure_parakeet_tdt_nemo_source(snapshot) + nemo_export = ensure_nemo_asr_source(snapshot) if nemo_export is not None: return nemo_export return snapshot @@ -377,7 +381,7 @@ def _write_component_bundle( component_io_signatures: dict[str, dict[str, tuple[str, ...]]] | None = None, component_metadata: dict[str, dict[str, object]] | None = None, graph_filename: str = "graph.cactus", - npu_encoder_mlpackages: dict[str, str] | None = None, + npu_encoder_mlpackages: dict[str, object] | None = None, ) -> Path: bundle_dir = artifact_dir / "components" component_order = [ @@ -506,8 +510,11 @@ def _write_component_bundle( for manifest_key, mlpackage_path in (npu_encoder_mlpackages or {}).items(): if mlpackage_path: manifest_payload[manifest_key] = mlpackage_path - if family == "parakeet_tdt" and "npu_audio_encoder" in manifest_payload: + has_npu_audio = "npu_audio_encoder" in manifest_payload or "npu_audio_encoders" in manifest_payload + if family == "parakeet_tdt" and has_npu_audio: manifest_payload["npu_audio_compute_units"] = "CPU_AND_NE" + elif family == "nemotron_asr" and has_npu_audio: + manifest_payload["npu_audio_compute_units"] = "CPU_AND_GPU" _write_json(manifest_path, manifest_payload) return manifest_path @@ -561,6 +568,11 @@ def _run_parakeet_tdt_component_decode( initial_store=_named_tensor_store(prepared), ) encoder_hidden_states = np.asarray(store["encoder_hidden_states"]) + active_frames = prepared.metadata.get("active_feature_frames") + if isinstance(active_frames, int): + total_frames = int(prepared.tensors[0].shape[1]) + valid_hidden = _valid_hidden_frame_count(active_frames, total_frames, int(encoder_hidden_states.shape[1])) + encoder_hidden_states = encoder_hidden_states[:, :valid_hidden, :] batch_size = int(encoder_hidden_states.shape[0]) if batch_size != 1: raise ValueError("Parakeet TDT component decode currently expects batch size 1") @@ -608,6 +620,71 @@ def _step( } +def _run_nemotron_asr_component_decode( + *, + component_graphs, + model, + prepared: PreparedInputs, +) -> dict[str, object]: + store, _ = execute_component_pipeline( + [component_graphs["audio_encoder"]], + initial_store=_named_tensor_store(prepared), + ) + encoder_hidden_states = np.asarray(store["encoder_hidden_states"]) + batch_size = int(encoder_hidden_states.shape[0]) + if batch_size != 1: + raise ValueError("Nemotron ASR component decode currently expects batch size 1") + + hidden_dtype = prepared.tensors[0].dtype + initial_states = model.initial_decoder_state( + batch_size=batch_size, + device=torch.device("cpu"), + dtype=hidden_dtype, + ) + initial_state_arrays = tuple(state.detach().cpu().numpy() for state in initial_states) + language_prompt = model.language_prompt( + "auto", + batch_size=batch_size, + device=torch.device("cpu"), + dtype=hidden_dtype, + ).detach().cpu().numpy() + decoder_component = component_graphs["decoder"] + + def _step( + frame: np.ndarray, + token_id: int, + state_values: tuple[np.ndarray, ...], + ) -> tuple[np.ndarray, tuple[np.ndarray, ...]]: + input_store: dict[str, object] = { + "encoder_frame": np.ascontiguousarray(frame), + "language_prompt": np.ascontiguousarray(language_prompt), + "token_ids": np.full((batch_size,), token_id, dtype=np.int64), + } + for index in range(model.config.predictor_num_layers): + input_store[f"state_h_{index}"] = np.ascontiguousarray(state_values[index * 2]) + input_store[f"state_c_{index}"] = np.ascontiguousarray(state_values[index * 2 + 1]) + + runtime_inputs = [input_store[key] for key in decoder_component.input_keys] + decoder_component.transpiled_graph.set_inputs(runtime_inputs) + outputs = decoder_component.transpiled_graph.execute() + logits = outputs[0].numpy().astype(np.float32, copy=False) + next_states = tuple(output.numpy() for output in outputs[1:]) + return logits, next_states + + emitted = greedy_decode_nemotron_asr_token_ids( + config=model.config, + encoder_hidden_states=encoder_hidden_states, + initial_states=initial_state_arrays, + step=_step, + ) + + return { + "token_ids": emitted, + "transcript": model.decode_token_ids(emitted), + "encoder_hidden_shape": list(encoder_hidden_states.shape), + } + + def _run_component_pipeline_transpile( *, args, @@ -638,7 +715,7 @@ def _run_component_pipeline_transpile( npu_quantize = getattr(args, "npu_quantize", None) npu_audio_quantize = getattr(args, "npu_audio_quantize", None) npu_vision_quantize = getattr(args, "npu_vision_quantize", None) - npu_encoder_mlpackages: dict[str, str] = {} + npu_encoder_mlpackages: dict[str, object] = {} if npu_enabled and artifact_dir is not None: from .npu import run_encoder_pipeline npu_encoder_mlpackages = run_encoder_pipeline( @@ -994,6 +1071,55 @@ def _run_component_pipeline_transpile( print(f"saved_result={artifact_dir / 'result.json'}") return 0 + if task == "rnnt_transcription": + print("execute_begin=true") + transpiled_decode = _run_nemotron_asr_component_decode( + component_graphs=captured_components, + model=model, + prepared=prepared, + ) + print("execute_done=true") + print(f"transpiled_transcript={transpiled_decode['transcript']!r}") + result_payload = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "raw_ir_nodes": raw_ir_nodes, + "optimized_ir_nodes": optimized_ir_nodes, + "weight_bindings": binding_count, + "transpiled_token_ids": transpiled_decode["token_ids"], + "transpiled_transcript": transpiled_decode["transcript"], + "encoder_hidden_shape": transpiled_decode["encoder_hidden_shape"], + } + if args.skip_reference_compare: + result_payload["reference_compare_skipped"] = True + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + + print("reference_begin=true") + active_frames = prepared.metadata.get("active_feature_frames") + reference_token_ids = model.greedy_decode_token_ids( + prepared.tensors[0], + language="auto", + active_frames=active_frames if isinstance(active_frames, int) else None, + ) + reference_transcript = model.decode_token_ids(reference_token_ids) + print("reference_done=true") + print(f"reference_transcript={reference_transcript!r}") + print(f"transcript_match={reference_transcript == transpiled_decode['transcript']}") + result_payload["reference_token_ids"] = reference_token_ids + result_payload["reference_transcript"] = reference_transcript + result_payload["transcript_match"] = bool(reference_transcript == transpiled_decode["transcript"]) + result_payload["token_id_match"] = bool(reference_token_ids == transpiled_decode["token_ids"]) + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + raise NotImplementedError(f"component pipeline execution is not implemented for task={task}") @@ -1270,6 +1396,8 @@ def _infer_task_from_config(model_id_or_path: str) -> str: loss_cfg = config.get("loss") if isinstance(loss_cfg, dict) and str(loss_cfg.get("loss_name", "") or "").lower() == "tdt": return "tdt_transcription" + if model_type in {"nemotron_asr", "nemotron-asr"}: + return "rnnt_transcription" if "nomic" in model_type or any("NomicBert" in value for value in architectures): return "text_embedding" @@ -1287,6 +1415,8 @@ def _infer_task_from_config(model_id_or_path: str) -> str: return "text_embedding" if "parakeet-tdt" in lowered_id: return "tdt_transcription" + if "nemotron-3.5-asr" in lowered_id or "nemotron-asr" in lowered_id: + return "rnnt_transcription" if "whisper" in lowered_id: return "seq2seq_transcription" if "ctc" in lowered_id: @@ -1300,6 +1430,7 @@ def _infer_task_from_config(model_id_or_path: str) -> str: "Pass one explicitly with --task, for example:\n" " --task causal_lm_logits\n" " --task tdt_transcription\n" + " --task rnnt_transcription\n" " --task ctc_logits\n" " --task encoder_hidden_states\n" " --task seq2seq_transcription\n" @@ -1328,6 +1459,8 @@ def _infer_fallback_audio_input_names(config: dict[str, object], task: str) -> t return ("input_features",) if task == "tdt_transcription": return ("input_features",) + if task == "rnnt_transcription": + return ("input_features",) audio_cfg = config.get("audio_config") if isinstance(audio_cfg, dict) and any(key in audio_cfg for key in ("features", "input_feat_size", "num_mel_bins")): return ("input_features", "attention_mask") if task == "ctc_logits" else ("input_features",) @@ -1449,12 +1582,20 @@ def _prepare_fallback_audio_inputs( model_type = str(config.get("model_type", "") or "").lower() is_parakeet_like = ( task == "tdt_transcription" + or task == "rnnt_transcription" or "parakeet" in model_type + or "nemotron" in model_type or "parakeet" in type(model).__module__.lower() + or "nemotron" in type(model).__module__.lower() ) if is_parakeet_like: expected_frames = _infer_expected_input_feature_frames(model) - parakeet_features, _ = prepare_parakeet_tdt_audio_features( + prepare_native_features = ( + prepare_nemotron_asr_audio_features + if task == "rnnt_transcription" or "nemotron" in model_type + else prepare_parakeet_tdt_audio_features + ) + parakeet_features, active_frames = prepare_native_features( audio_file, expected_frames=expected_frames, expected_mels=num_mels, @@ -1469,6 +1610,7 @@ def _prepare_fallback_audio_inputs( "sample_rate": sample_rate, "fallback_audio_preprocessor": True, "native_parakeet_audio_preprocessor": True, + "active_feature_frames": int(active_frames), "input_shapes": { name: list(tensor.shape) for name, tensor in zip(input_names, tensors) @@ -2691,6 +2833,10 @@ def _load_transformers_bundle( model = load_tdt_local_model(model_source, torch_dtype=torch_dtype).eval() return model_source, None, model, config + if task == "rnnt_transcription": + model = load_nemotron_asr_local_model(model_source, torch_dtype=torch_dtype).eval() + return model_source, None, model, config + if task == "text_embedding": # nomic-bert weights are produced from the model's remote modeling code # (fused Wqkv / experts.mlp.w1,w2); the transformers-native nomic_bert has a @@ -3029,6 +3175,7 @@ def main() -> int: "encoder_hidden_states", "seq2seq_transcription", "tdt_transcription", + "rnnt_transcription", "text_embedding", ), help="Transpile task. Use auto to infer from config/model id.", @@ -3305,7 +3452,7 @@ def main() -> int: prime_static_features = getattr(canonical.module, "prime_static_multimodal_features", None) if callable(prime_static_features): prime_static_features(*prepared.tensors) - elif task == "tdt_transcription": + elif task in {"tdt_transcription", "rnnt_transcription"}: if not args.audio_file: raise RuntimeError(f"--audio-file is required for task={task}") prepared = _prepare_audio_inputs( diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index ceeee03b1..fb94a04a6 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -151,7 +151,13 @@ def transpile_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: flush=True, ) - outputs = [env[value_id] for value_id in ir.outputs] + outputs = [] + for value_id in ir.outputs: + output = env[value_id] + value = ir.values.get(value_id) + if isinstance(output, Tensor) and value is not None and value.users: + output = g.reshape(output, tuple(int(dim) for dim in output.shape)) + outputs.append(output) seen_bound_constant_ids = {int(tensor.id) for tensor in bound_constants} for tensor in getattr(g, "_transpile_materialized_constants", []): tensor_id = int(tensor.id) @@ -1835,6 +1841,60 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> current = g.conv1d_pointwise(current, pointwise2_weight, bias=pointwise2_bias) return [current] + if op == "causal_layer_norm_conv_module": + input_index = 0 + x_nlc = _tensor(env, node.inputs[input_index]) + input_index += 1 + + pointwise1_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + pointwise1_bias = None + if bool(node.attrs.get("has_pointwise1_bias", False)): + pointwise1_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + depthwise_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + depthwise_bias = None + if bool(node.attrs.get("has_depthwise_bias", False)): + depthwise_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + layer_norm_weight = _tensor(env, node.inputs[input_index]) + layer_norm_bias = _tensor(env, node.inputs[input_index + 1]) + input_index += 2 + + pointwise2_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + pointwise2_bias = None + if bool(node.attrs.get("has_pointwise2_bias", False)): + pointwise2_bias = _tensor(env, node.inputs[input_index]) + + if len(x_nlc.shape) != 3: + raise NotImplementedError(f"causal_layer_norm_conv_module expects rank-3 NLC input, got {x_nlc.shape}") + kernel_size = int(node.attrs.get("depthwise_kernel_size", 0)) + if len(depthwise_weight.shape) != 3 or kernel_size != int(depthwise_weight.shape[2]): + raise NotImplementedError( + f"causal_layer_norm_conv_module got kernel_size={kernel_size} weight_shape={depthwise_weight.shape}" + ) + + current = g.conv1d_pointwise(x_nlc, pointwise1_weight, bias=pointwise1_bias) + current = g.glu(current, axis=-1) + current = g.conv1d_causal(current, depthwise_weight, kernel_size=kernel_size, dilation=1) + if depthwise_bias is not None: + bias_reshaped = g.reshape(depthwise_bias, (1, 1, int(depthwise_weight.shape[0]))) + current, bias_reshaped = _legalize_elementwise_binary_inputs(g, current, bias_reshaped) + current = g.add(current, bias_reshaped) + current = g.layer_norm( + current, + layer_norm_weight, + bias=layer_norm_bias, + eps=float(node.attrs.get("eps", 1e-5)), + ) + current = g.silu(current) + current = g.conv1d_pointwise(current, pointwise2_weight, bias=pointwise2_bias) + return [current] + if op == "conv1d": x = _tensor(env, node.inputs[0]) weight = _tensor(env, node.inputs[1]) @@ -2000,7 +2060,12 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> if dilation != 1: raise NotImplementedError(f"conv2d with dilation != 1 is unsupported: {dilation}") - if stride == 2 and padding == 0 and groups == 1 and weight.shape[2:] == (3, 3): + is_depthwise_k3 = ( + groups == x.shape[1] == weight.shape[0] + and weight.shape[1] == 1 + and weight.shape[2:] == (3, 3) + ) + if stride == 2 and padding == 0 and weight.shape[2:] == (3, 3) and (groups == 1 or is_depthwise_k3): batch_size, channels, height, width = (int(dim) for dim in x.shape) pad_dtype = _torch_dtype_for_graph_dtype(x.dtype) top = _materialize_constant_tensor( @@ -2021,7 +2086,10 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> torch.zeros((batch_size, channels, height + 2, 1), dtype=pad_dtype), ) x = g.cat([left, x, right], axis=3) - y = g.conv2d_k3s2p1(x, weight, bias=bias) + if is_depthwise_k3: + y = g.conv2d_depthwise_k3s2p1(x, weight, bias=bias) + else: + y = g.conv2d_k3s2p1(x, weight, bias=bias) output_height = ((height - 3) // 2) + 1 output_width = ((width - 3) // 2) + 1 y = g.slice(y, axis=2, start=1, length=output_height) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index daac679ce..f3c2670f1 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -6164,6 +6164,14 @@ def build_component_module_specs( named_tensors=named_tensors, weights_dir=weights_dir, ) + if family == "nemotron_asr" and task == "rnnt_transcription": + from cactus.transpile.rnnt_runtime import build_nemotron_asr_component_specs + + return build_nemotron_asr_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + ) if family == "whisper" and task == "seq2seq_transcription": return _build_whisper_seq2seq_component_specs( model, diff --git a/python/cactus/transpile/model_profiles.py b/python/cactus/transpile/model_profiles.py index 75580eb47..22cbd7b9a 100644 --- a/python/cactus/transpile/model_profiles.py +++ b/python/cactus/transpile/model_profiles.py @@ -153,6 +153,27 @@ class ModelProfile: ) +NEMOTRON_ASR_PROFILE = ModelProfile( + family="nemotron_asr", + model_types=("nemotron_asr",), + model_id_aliases=( + ("nemotron-asr", "nvidia/nemotron-3.5-asr-streaming-0.6b"), + ("nemotron-3.5-asr", "nvidia/nemotron-3.5-asr-streaming-0.6b"), + ), + model_id_markers=("nemotron-3.5-asr", "nemotron-asr"), + avoid_native_loader=True, + default_task="rnnt_transcription", + default_components=("audio_encoder", "decoder"), + needs_audio=True, + force_component_pipeline=True, + aliases=PARAKEET_TDT_PROFILE.aliases + ( + ("joint.joint_net.2.weight", "joint.joint_net.0.weight"), + ("joint.joint_net.2.bias", "joint.joint_net.0.bias"), + ), + regex_aliases=PARAKEET_TDT_PROFILE.regex_aliases, +) + + WHISPER_PROFILE = ModelProfile( family="whisper", model_types=("whisper",), @@ -227,6 +248,7 @@ class ModelProfile: GEMMA4_PROFILE, LFM2_PROFILE, PARAKEET_TDT_PROFILE, + NEMOTRON_ASR_PROFILE, WHISPER_PROFILE, QWEN_PROFILE, NEEDLE_PROFILE, diff --git a/python/cactus/transpile/npu/pipeline.py b/python/cactus/transpile/npu/pipeline.py index e15b25cce..ac2e71973 100644 --- a/python/cactus/transpile/npu/pipeline.py +++ b/python/cactus/transpile/npu/pipeline.py @@ -23,8 +23,8 @@ def run_encoder_pipeline( quantize_bits: int | None = None, audio_quantize_bits: int | None = None, vision_quantize_bits: int | None = None, -) -> dict[str, str]: - results: dict[str, str] = {} +) -> dict[str, object]: + results: dict[str, object] = {} if not enabled or not component_specs: return results @@ -60,6 +60,30 @@ def _resolve(per_component: int | None, default: int | None) -> int | None: qdesc = f"int{qbits}" if qbits else "fp16" print(f"npu.pipeline: emitting {component} quant={qdesc}") reparam = getattr(spec, "npu_reparam", None) + variants = tuple(getattr(spec, "npu_variants", ()) or ()) + if component == "audio_encoder" and variants: + emitted_paths: list[str] = [] + for variant in variants: + variant_inputs = tuple(variant.get("example_inputs") or example_inputs) + if not variant_inputs: + continue + variant_module = variant.get("module") or module + variant_filename = str(variant.get("filename") or filename) + with (reparam(variant_module) if reparam else nullcontext()): + emitted = emit_fn( + variant_module, + bundle_root, + example_input=variant_inputs[0], + baked_inputs=variant_inputs[1:], + filename=variant_filename, + quantize_bits=qbits, + ) + if emitted: + emitted_paths.append(f"components/{emitted}") + if emitted_paths: + results["npu_audio_encoders"] = emitted_paths + results["npu_audio_encoder"] = emitted_paths[-1] + continue with (reparam(module) if reparam else nullcontext()): if component == "source_encoder": emitted = emit_fn( diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py index e1a870f8b..0351591ce 100644 --- a/python/cactus/transpile/ops.py +++ b/python/cactus/transpile/ops.py @@ -183,6 +183,18 @@ class OpSchema: ), is_fusible=True, ), + OpSchema( + "causal_layer_norm_conv_module", + num_inputs=None, + attrs=( + "eps", + "has_pointwise1_bias", + "has_depthwise_bias", + "has_pointwise2_bias", + "depthwise_kernel_size", + ), + is_fusible=True, + ), OpSchema( "gated_deltanet_prefill", num_inputs=None, diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index ea2714d0d..4196a95b5 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/python/cactus/transpile/optimize_graph.py @@ -11,6 +11,7 @@ from cactus.transpile.canonicalize.utils import rebuild_graph from cactus.transpile.fusion import match_attention from cactus.transpile.fusion import match_attention_block +from cactus.transpile.fusion import match_causal_layer_norm_conv_module from cactus.transpile.fusion import match_conv_module from cactus.transpile.fusion import match_gated_deltanet from cactus.transpile.fusion import match_gated_mlp @@ -103,8 +104,11 @@ def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig changed = True if config.enable_dense_mlp_tq_fused and fuse_dense_mlp_tq(graph): changed = True - if config.enable_conv_module and fuse_conv_modules(graph): - changed = True + if config.enable_conv_module: + if fuse_conv_modules(graph): + changed = True + if fuse_causal_layer_norm_conv_modules(graph): + changed = True if config.enable_add_clipped and fuse_add_clipped(graph): changed = True if not changed: @@ -1223,6 +1227,54 @@ def fuse_conv_modules(graph: IRGraph) -> bool: return changed +def fuse_causal_layer_norm_conv_modules(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_causal_layer_norm_conv_module(graph, node) + if match is None: + continue + + inputs = [ + match.input_value_id, + match.pointwise1_weight_value_id, + ] + if match.pointwise1_bias_value_id is not None: + inputs.append(match.pointwise1_bias_value_id) + inputs.append(match.depthwise_weight_value_id) + if match.depthwise_bias_value_id is not None: + inputs.append(match.depthwise_bias_value_id) + inputs.extend( + [ + match.layer_norm_weight_value_id, + match.layer_norm_bias_value_id, + match.pointwise2_weight_value_id, + ] + ) + if match.pointwise2_bias_value_id is not None: + inputs.append(match.pointwise2_bias_value_id) + + node.op = "causal_layer_norm_conv_module" + node.inputs = inputs + node.attrs = { + "eps": float(match.eps), + "has_pointwise1_bias": bool(match.pointwise1_bias_value_id is not None), + "has_depthwise_bias": bool(match.depthwise_bias_value_id is not None), + "has_pointwise2_bias": bool(match.pointwise2_bias_value_id is not None), + "depthwise_kernel_size": int(match.depthwise_kernel_size), + } + node.kind = "semantic" + node.meta["conv_module_nodes"] = match.node_ids + changed = True + + if changed: + rebuild_graph(graph) + return changed + + def _materialize_shifted_constant(graph: IRGraph, value_id: str, offset: float, *, suffix: str) -> str: base = graph.constants[value_id] if not isinstance(base, torch.Tensor): diff --git a/python/cactus/transpile/rnnt_runtime.py b/python/cactus/transpile/rnnt_runtime.py new file mode 100644 index 000000000..8c6e44107 --- /dev/null +++ b/python/cactus/transpile/rnnt_runtime.py @@ -0,0 +1,513 @@ +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from dataclasses import replace +import json +import os +import re +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from cactus.transpile.component_pipeline import ComponentModuleSpec +from cactus.transpile.tdt_runtime import ParakeetTDTConfig +from cactus.transpile.tdt_runtime import ParakeetTDTANEEncoder +from cactus.transpile.tdt_runtime import ParakeetTDTDecoderPrediction +from cactus.transpile.tdt_runtime import ParakeetTDTEncoder +from cactus.transpile.tdt_runtime import ParakeetTDTJoint +from cactus.transpile.tdt_runtime import _add_tdt_derived_aliases +from cactus.transpile.tdt_runtime import _cfg_get +from cactus.transpile.tdt_runtime import _copy_linear_weight +from cactus.transpile.tdt_runtime import _load_parakeet_vocabulary +from cactus.transpile.tdt_runtime import _parse_attention_context_size +from cactus.transpile.audio_preprocess import _PARAKEET_FRAME_LENGTH +from cactus.transpile.audio_preprocess import _PARAKEET_HOP_LENGTH +from cactus.transpile.audio_preprocess import _PARAKEET_LOG_FLOOR +from cactus.transpile.audio_preprocess import _PARAKEET_N_FFT +from cactus.transpile.audio_preprocess import _PARAKEET_PREEMPHASIS +from cactus.transpile.audio_preprocess import _PARAKEET_SAMPLE_RATE +from cactus.transpile.audio_preprocess import generic_log_mel_features +from cactus.transpile.audio_preprocess import load_audio_waveform +from cactus.transpile.model_profiles import NEMOTRON_ASR_PROFILE +from cactus.transpile.model_profiles import add_tensor_aliases + + +def nemotron_asr_frame_values() -> tuple[int, ...]: + raw = os.environ.get("CACTUS_NEMOTRON_ASR_FRAMES") + if raw is None: + raw = os.environ.get("CACTUS_NPU_NEMOTRON_FRAMES", "1024,2048,4096,6144") + values = tuple(sorted({int(value.strip()) for value in raw.split(",") if value.strip()})) + if not values or any(value <= 0 for value in values): + raise ValueError("CACTUS_NEMOTRON_ASR_FRAMES values must be positive frame counts") + return values + + +def _valid_hidden_frame_count(active_frames: int | None, total_frames: int, hidden_frames: int) -> int: + if active_frames is None or active_frames <= 0 or total_frames <= 0: + return hidden_frames + return min(hidden_frames, (int(active_frames) * hidden_frames + total_frames - 1) // total_frames) + + +@contextmanager +def _temporary_float32(module: torch.nn.Module): + dtypes = { + param.dtype + for param in module.parameters() + if param.is_floating_point() + } + dtype = next(iter(dtypes)) if len(dtypes) == 1 else None + if dtype is None or dtype == torch.float32: + yield + return + module.to(dtype=torch.float32) + try: + yield + finally: + module.to(dtype=dtype) + + +@dataclass +class NemotronASRConfig: + model_source: str + sample_rate: int + num_mel_bins: int + hidden_dim: int + num_layers: int + attention_heads: int + attention_head_dim: int + attention_scale: float + ff_intermediate_dim: int + conv_kernel_size: int + conv_norm_type: str + conv_is_causal: bool + subsampling_factor: int + subsampling_conv_channels: int + causal_downsampling: bool + predictor_hidden_dim: int + predictor_num_layers: int + joint_dim: int + blank_id: int + vocabulary: tuple[str, ...] + encoder_hidden_act: str + prompt_dim: int + prompt_dictionary: dict[str, int] + default_prompt_id: int + max_symbols_per_step: int + attention_context_style: str + attention_context_size: tuple[int, int] + + def as_parakeet_config(self) -> ParakeetTDTConfig: + return ParakeetTDTConfig( + model_source=self.model_source, + sample_rate=self.sample_rate, + num_mel_bins=self.num_mel_bins, + hidden_dim=self.hidden_dim, + num_layers=self.num_layers, + attention_heads=self.attention_heads, + attention_head_dim=self.attention_head_dim, + attention_scale=self.attention_scale, + ff_intermediate_dim=self.ff_intermediate_dim, + conv_kernel_size=self.conv_kernel_size, + conv_norm_type=self.conv_norm_type, + conv_is_causal=self.conv_is_causal, + subsampling_factor=self.subsampling_factor, + subsampling_conv_channels=self.subsampling_conv_channels, + causal_downsampling=self.causal_downsampling, + predictor_hidden_dim=self.predictor_hidden_dim, + predictor_num_layers=self.predictor_num_layers, + joint_dim=self.joint_dim, + num_tdt_durations=0, + tdt_durations=(), + blank_id=self.blank_id, + vocabulary=self.vocabulary, + encoder_hidden_act=self.encoder_hidden_act, + attention_context_style=self.attention_context_style, + attention_context_size=self.attention_context_size, + ) + + +def _load_tensor_state_dict(model_source: str) -> dict[str, torch.Tensor]: + root = Path(model_source) + safetensors_path = root / "model.safetensors" + if safetensors_path.exists(): + from safetensors.torch import load_file + + return _with_nemotron_asr_aliases(dict(load_file(str(safetensors_path)))) + + bin_path = root / "pytorch_model.bin" + if bin_path.exists(): + loaded = torch.load(bin_path, map_location="cpu") + if isinstance(loaded, dict) and "state_dict" in loaded and isinstance(loaded["state_dict"], dict): + loaded = loaded["state_dict"] + if isinstance(loaded, dict): + return _with_nemotron_asr_aliases({ + str(key): value + for key, value in loaded.items() + if isinstance(value, torch.Tensor) + }) + raise RuntimeError(f"unsupported Nemotron ASR checkpoint format in {model_source}") + + +def _with_nemotron_asr_aliases(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + add_tensor_aliases( + state_dict, + NEMOTRON_ASR_PROFILE, + derived_aliases=_add_tdt_derived_aliases, + ) + return state_dict + + +def _prompt_dim_from_dictionary(prompt_dictionary: dict[str, int], *candidates: int) -> int: + dim = max((int(value or 0) for value in candidates), default=0) + if prompt_dictionary: + dim = max(dim, max(int(value) for value in prompt_dictionary.values()) + 1) + return max(1, dim) + + +def prepare_nemotron_asr_audio_features( + audio_file: str | Path, + *, + expected_frames: int | None, + expected_mels: int, + torch_dtype: torch.dtype, +) -> tuple[torch.Tensor, int]: + if expected_frames is None or expected_frames <= 0: + expected_frames = max(nemotron_asr_frame_values()) + waveform = load_audio_waveform(audio_file, target_sample_rate=_PARAKEET_SAMPLE_RATE) + features, feature_length = generic_log_mel_features( + waveform, + sample_rate=_PARAKEET_SAMPLE_RATE, + num_mels=expected_mels, + n_fft=_PARAKEET_N_FFT, + hop_length=_PARAKEET_HOP_LENGTH, + frame_length=_PARAKEET_FRAME_LENGTH, + preemphasis=_PARAKEET_PREEMPHASIS, + mel_floor=_PARAKEET_LOG_FLOOR, + mel_floor_additive=True, + normalize_active_frames_only=None, + ) + active_frames = feature_length + if isinstance(expected_frames, int) and expected_frames > 0: + active_frames = min(active_frames, expected_frames) + features = features[:active_frames, :] + if features.shape[1] != expected_mels: + raise ValueError( + f"feature mel dimension mismatch: expected {expected_mels}, got {features.shape[1]}" + ) + if isinstance(expected_frames, int) and expected_frames > active_frames: + features = np.pad(features, ((0, expected_frames - active_frames), (0, 0)), mode="constant") + tensor = torch.from_numpy(np.ascontiguousarray(features)).unsqueeze(0).to(dtype=torch_dtype) + return tensor, active_frames + + +def load_nemotron_asr_config(model_source: str) -> NemotronASRConfig: + root_path = Path(model_source) + config_path = root_path / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"missing config.json for Nemotron ASR: {config_path}") + root = json.loads(config_path.read_text()) + encoder = root.get("encoder") or root.get("encoder_config") or {} + decoder = root.get("decoder") or {} + prediction = decoder.get("prediction") or decoder.get("prednet") or {} + joint = root.get("joint") or {} + jointnet = joint.get("jointnet") or {} + model_defaults = root.get("model_defaults") or {} + preprocessor = root.get("preprocessor") or {} + decoding = root.get("decoding") or {} + greedy = decoding.get("greedy") or {} + + hidden_dim = int(_cfg_get(root, "hidden_dim", _cfg_get(encoder, "d_model", _cfg_get(encoder, "hidden_size", 0)))) + attention_heads = int(_cfg_get(encoder, "n_heads", _cfg_get(encoder, "num_attention_heads", 0))) + attention_head_dim = hidden_dim // max(attention_heads, 1) + vocabulary = tuple(str(value) for value in _cfg_get(joint, "vocabulary", _cfg_get(root, "labels", ()))) + if not vocabulary: + vocabulary = _load_parakeet_vocabulary(root_path) + decoder_vocab_size = int(_cfg_get(root, "rnnt_blank_id", _cfg_get(decoder, "vocab_size", len(vocabulary)))) + prompt_dictionary = _cfg_get(root, "prompt_dictionary", _cfg_get(model_defaults, "prompt_dictionary", {})) + if not isinstance(prompt_dictionary, dict): + prompt_dictionary = {} + prompt_dictionary = {str(key): int(value) for key, value in prompt_dictionary.items()} + prompt_dim = _prompt_dim_from_dictionary( + prompt_dictionary, + int(_cfg_get(root, "prompt_dim", 0) or 0), + int(_cfg_get(root, "num_prompts", 0) or 0), + int(_cfg_get(model_defaults, "num_prompts", 0) or 0), + ) + + return NemotronASRConfig( + model_source=model_source, + sample_rate=int(_cfg_get(preprocessor, "sample_rate", _cfg_get(root, "sample_rate", 16000))), + num_mel_bins=int(_cfg_get(root, "num_mel_bins", _cfg_get(preprocessor, "features", _cfg_get(encoder, "feat_in", _cfg_get(encoder, "num_mel_bins", 128))))), + hidden_dim=hidden_dim, + num_layers=int(_cfg_get(root, "num_layers", _cfg_get(encoder, "n_layers", _cfg_get(encoder, "num_hidden_layers", 0)))), + attention_heads=attention_heads, + attention_head_dim=attention_head_dim, + attention_scale=float(attention_head_dim ** -0.5 if attention_head_dim > 0 else 1.0), + ff_intermediate_dim=int( + _cfg_get( + root, + "ffn_intermediate_dim", + _cfg_get( + encoder, + "ffn_hidden_size", + _cfg_get(encoder, "intermediate_size", round(hidden_dim * float(_cfg_get(encoder, "ff_expansion_factor", 4.0)))), + ), + ) + ), + conv_kernel_size=int(_cfg_get(root, "conv_kernel_size", _cfg_get(encoder, "conv_kernel_size", 9))), + conv_norm_type=str(_cfg_get(root, "conv_norm_type", _cfg_get(encoder, "conv_norm_type", "batch_norm"))).lower(), + conv_is_causal=str(_cfg_get(root, "conv_context_size", _cfg_get(encoder, "conv_context_size", ""))).lower() == "causal", + subsampling_factor=int(_cfg_get(root, "subsampling_factor", _cfg_get(encoder, "subsampling_factor", 8))), + subsampling_conv_channels=int(_cfg_get(root, "subsampling_conv_channels", _cfg_get(encoder, "subsampling_conv_channels", 256))), + causal_downsampling=bool(_cfg_get(root, "causal_downsampling", _cfg_get(encoder, "causal_downsampling", False))), + predictor_hidden_dim=int(_cfg_get(root, "predictor_hidden_dim", _cfg_get(root, "decoder_hidden_size", _cfg_get(prediction, "pred_hidden", _cfg_get(model_defaults, "pred_hidden", 640))))), + predictor_num_layers=int(_cfg_get(root, "predictor_num_layers", _cfg_get(root, "num_decoder_layers", _cfg_get(prediction, "pred_rnn_layers", 1)))), + joint_dim=int(_cfg_get(root, "rnnt_joint_dim", _cfg_get(jointnet, "joint_hidden", _cfg_get(model_defaults, "joint_hidden", 640)))), + blank_id=int(_cfg_get(root, "rnnt_blank_id", decoder_vocab_size)), + vocabulary=vocabulary, + encoder_hidden_act=str(_cfg_get(root, "encoder_hidden_act", _cfg_get(encoder, "activation", _cfg_get(encoder, "hidden_act", "silu")))).lower(), + prompt_dim=prompt_dim, + prompt_dictionary=prompt_dictionary, + default_prompt_id=int(_cfg_get(root, "default_prompt_id", prompt_dictionary.get("auto", 0))), + max_symbols_per_step=int(_cfg_get(root, "max_symbols_per_step", _cfg_get(greedy, "max_symbols", 10))), + attention_context_style=str(_cfg_get(encoder, "att_context_style", "regular")), + attention_context_size=_parse_attention_context_size(_cfg_get(encoder, "att_context_size", None)), + ) + + +class NemotronASRPromptKernel(nn.Module): + def __init__(self, config: NemotronASRConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + linear1_weight = state_dict["prompt_kernel.0.weight"] + linear2_weight = state_dict["prompt_kernel.2.weight"] + self.linear1 = nn.Linear(int(linear1_weight.shape[1]), int(linear1_weight.shape[0])) + self.linear2 = nn.Linear(int(linear2_weight.shape[1]), int(linear2_weight.shape[0])) + _copy_linear_weight(self.linear1, state_dict["prompt_kernel.0.weight"], bias=state_dict["prompt_kernel.0.bias"]) + _copy_linear_weight(self.linear2, state_dict["prompt_kernel.2.weight"], bias=state_dict["prompt_kernel.2.bias"]) + + def forward(self, encoder_frame: torch.Tensor, language_prompt: torch.Tensor) -> torch.Tensor: + if language_prompt.ndim == 1: + language_prompt = language_prompt.unsqueeze(0) + x = torch.cat((encoder_frame, language_prompt.to(dtype=encoder_frame.dtype, device=encoder_frame.device)), dim=-1) + return self.linear2(F.relu(self.linear1(x))) + + +class NemotronASRDecoderStep(nn.Module): + def __init__(self, config: NemotronASRConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + parakeet_config = config.as_parakeet_config() + self.prompt_kernel = NemotronASRPromptKernel(config, state_dict) + self.prediction = ParakeetTDTDecoderPrediction(parakeet_config, state_dict) + self.joint = ParakeetTDTJoint(parakeet_config, state_dict) + self.config = config + + def forward(self, encoder_frame: torch.Tensor, language_prompt: torch.Tensor, token_ids: torch.Tensor, *state_tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: + if len(state_tensors) != self.config.predictor_num_layers * 2: + raise ValueError( + f"expected {self.config.predictor_num_layers * 2} predictor state tensors, got {len(state_tensors)}" + ) + state_h = tuple(state_tensors[0::2]) + state_c = tuple(state_tensors[1::2]) + prompted_frame = self.prompt_kernel(encoder_frame, language_prompt) + predictor_hidden, next_h, next_c = self.prediction(token_ids, state_h, state_c) + logits = self.joint(prompted_frame, predictor_hidden) + outputs: list[torch.Tensor] = [logits] + for h, c in zip(next_h, next_c, strict=True): + outputs.append(h) + outputs.append(c) + return tuple(outputs) + + +class NemotronASRLocalModel(nn.Module): + def __init__(self, config: NemotronASRConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.name_or_path = config.model_source + self.family = "nemotron_asr" + self.config = config + self.encoder = ParakeetTDTEncoder(config.as_parakeet_config(), state_dict) + self.decoder_step = NemotronASRDecoderStep(config, state_dict) + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + return self.encoder(input_features) + + def initial_decoder_state(self, *, batch_size: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, ...]: + state: list[torch.Tensor] = [] + for _ in range(self.config.predictor_num_layers): + state.append(torch.zeros((batch_size, self.config.predictor_hidden_dim), device=device, dtype=dtype)) + state.append(torch.zeros((batch_size, self.config.predictor_hidden_dim), device=device, dtype=dtype)) + return tuple(state) + + def language_prompt(self, language: str = "auto", *, batch_size: int = 1, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + prompt_id = int(self.config.prompt_dictionary.get(language, self.config.default_prompt_id)) + prompt_id = max(0, min(prompt_id, self.config.prompt_dim - 1)) + prompt = torch.zeros((batch_size, self.config.prompt_dim), device=device, dtype=dtype) + prompt[:, prompt_id] = 1.0 + return prompt + + def greedy_decode_token_ids(self, input_features: torch.Tensor, *, language: str = "auto", active_frames: int | None = None) -> list[int]: + with torch.no_grad(): + encoder_hidden = self.encoder(input_features) + valid_hidden = _valid_hidden_frame_count(active_frames, int(input_features.shape[1]), int(encoder_hidden.shape[1])) + encoder_hidden = encoder_hidden[:, :valid_hidden, :] + batch = int(encoder_hidden.shape[0]) + if batch != 1: + raise ValueError("Nemotron ASR local greedy decode currently expects batch size 1") + states = self.initial_decoder_state(batch_size=batch, device=encoder_hidden.device, dtype=encoder_hidden.dtype) + language_prompt = self.language_prompt(language, batch_size=batch, device=encoder_hidden.device, dtype=encoder_hidden.dtype) + state_arrays = tuple(state.detach().cpu().numpy() for state in states) + encoder_hidden_np = encoder_hidden.detach().cpu().numpy() + prompt_np = language_prompt.detach().cpu().numpy() + + def _step(frame: np.ndarray, token_id: int, state_values: tuple[np.ndarray, ...]) -> tuple[np.ndarray, tuple[np.ndarray, ...]]: + frame_tensor = torch.from_numpy(frame).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) + prompt_tensor = torch.from_numpy(prompt_np).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) + token_tensor = torch.tensor([token_id], device=encoder_hidden.device, dtype=torch.long) + state_tensors = tuple( + torch.from_numpy(value).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) + for value in state_values + ) + outputs = self.decoder_step(frame_tensor, prompt_tensor, token_tensor, *state_tensors) + logits = outputs[0].detach().cpu().numpy() + next_states = tuple(output.detach().cpu().numpy() for output in outputs[1:]) + return logits, next_states + + return greedy_decode_nemotron_asr_token_ids( + config=self.config, + encoder_hidden_states=encoder_hidden_np, + initial_states=state_arrays, + step=_step, + ) + + def decode_token_ids(self, token_ids: list[int], *, strip_lang_tags: bool = True) -> str: + pieces: list[str] = [] + for token_id in token_ids: + if token_id < 0 or token_id >= len(self.config.vocabulary): + continue + pieces.append(self.config.vocabulary[token_id]) + text = "".join(pieces).replace("▁", " ") + text = re.sub(r"\s+", " ", text).strip() + if strip_lang_tags: + text = re.sub(r"\s*<[a-z]{2}(?:-[A-Z]{2})?>\s*", " ", text) + text = re.sub(r"\s+", " ", text).strip() + return text + + +def greedy_decode_nemotron_asr_token_ids( + *, + config: NemotronASRConfig, + encoder_hidden_states: np.ndarray, + initial_states: tuple[np.ndarray, ...], + step, +) -> list[int]: + hidden = np.ascontiguousarray(np.asarray(encoder_hidden_states)) + if hidden.ndim != 3 or hidden.shape[0] != 1: + raise ValueError(f"expected encoder_hidden_states with shape [1, T, D], got {tuple(hidden.shape)}") + if len(initial_states) != int(config.predictor_num_layers) * 2: + raise ValueError(f"expected {int(config.predictor_num_layers) * 2} initial state tensors, got {len(initial_states)}") + + states = tuple(np.ascontiguousarray(np.asarray(state)) for state in initial_states) + last_token = int(config.blank_id) + emitted: list[int] = [] + + for time_index in range(int(hidden.shape[1])): + frame = np.ascontiguousarray(hidden[:, time_index, :]) + symbols_added = 0 + while symbols_added < int(config.max_symbols_per_step): + logits, next_states = step(frame, last_token, states) + logits_array = np.asarray(logits, dtype=np.float32) + next_token = int(np.argmax(logits_array[0])) + symbols_added += 1 + if next_token == int(config.blank_id): + break + emitted.append(next_token) + last_token = next_token + states = tuple(np.ascontiguousarray(np.asarray(state).copy()) for state in next_states) + return emitted + + +def load_nemotron_asr_local_model(model_source: str, *, torch_dtype: torch.dtype) -> NemotronASRLocalModel: + config = load_nemotron_asr_config(model_source) + state_dict = _load_tensor_state_dict(model_source) + predictor_vocab_size = int(state_dict["decoder.prediction.embed.weight"].shape[0]) + if config.blank_id < 0 or config.blank_id >= predictor_vocab_size: + config = replace(config, blank_id=predictor_vocab_size - 1) + model = NemotronASRLocalModel(config, state_dict).eval() + model.to(dtype=torch_dtype) + return model + + +def build_nemotron_asr_component_specs( + model: NemotronASRLocalModel, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, +) -> list[ComponentModuleSpec]: + input_features = named_tensors["input_features"] + example_hidden = model.encoder(input_features) + batch_size = int(example_hidden.shape[0]) + initial_states = model.initial_decoder_state(batch_size=batch_size, device=example_hidden.device, dtype=example_hidden.dtype) + example_prompt = model.language_prompt("auto", batch_size=batch_size, device=example_hidden.device, dtype=example_hidden.dtype) + example_frame = example_hidden[:, :1, :].reshape(batch_size, example_hidden.shape[-1]) + example_token_id = torch.full((batch_size,), int(model.config.blank_id), device=example_hidden.device, dtype=torch.long) + + state_input_keys: list[str] = [] + state_output_keys: list[str] = [] + for index in range(model.config.predictor_num_layers): + state_input_keys.extend((f"state_h_{index}", f"state_c_{index}")) + state_output_keys.extend((f"state_h_{index}", f"state_c_{index}")) + + common_graph_meta = { + "weights_dir": weights_dir, + "task": "rnnt_transcription", + "adapter_family": "nemotron_asr", + } + + factor = max(1, int(model.config.subsampling_factor)) + npu_frame_values = nemotron_asr_frame_values() + npu_variants: list[dict[str, object]] = [] + for npu_frames in npu_frame_values: + if npu_frames <= 0 or npu_frames % factor != 0: + raise ValueError(f"CACTUS_NEMOTRON_ASR_FRAMES values must be positive multiples of {factor}, got {npu_frames}") + npu_input_features = torch.zeros( + (int(input_features.shape[0]), npu_frames, int(input_features.shape[-1])), + device=input_features.device, + dtype=input_features.dtype, + ) + valid_frames = min(int(input_features.shape[1]), npu_frames) + npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] + with torch.no_grad(): + ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) + npu_variants.append({ + "module": ParakeetTDTANEEncoder(model.encoder, ane_seq_len), + "example_inputs": (npu_input_features.to(dtype=torch.float32),), + "filename": f"audio_encoder_{npu_frames}.mlpackage", + }) + + return [ + ComponentModuleSpec( + component="audio_encoder", + module=model.encoder, + example_inputs=(input_features,), + input_keys=("input_features",), + output_keys=("encoder_hidden_states",), + graph_meta={**common_graph_meta, "component": "audio_encoder"}, + metadata={"family": "nemotron_asr", "task": "rnnt_transcription"}, + npu_runtime_input_count=1, + npu_reparam=_temporary_float32, + npu_variants=tuple(npu_variants), + ), + ComponentModuleSpec( + component="decoder", + module=model.decoder_step, + example_inputs=(example_frame, example_prompt, example_token_id, *initial_states), + input_keys=("encoder_frame", "language_prompt", "token_ids", *state_input_keys), + output_keys=("step_logits", *state_output_keys), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata={"family": "nemotron_asr", "task": "rnnt_transcription"}, + ), + ] diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 01275d7ec..2f4a7ce72 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -101,8 +101,11 @@ class ParakeetTDTConfig: attention_scale: float ff_intermediate_dim: int conv_kernel_size: int + conv_norm_type: str + conv_is_causal: bool subsampling_factor: int subsampling_conv_channels: int + causal_downsampling: bool predictor_hidden_dim: int predictor_num_layers: int joint_dim: int @@ -111,6 +114,18 @@ class ParakeetTDTConfig: blank_id: int vocabulary: tuple[str, ...] encoder_hidden_act: str + attention_context_style: str = "regular" + attention_context_size: tuple[int, int] = (-1, -1) + + +def _parse_attention_context_size(raw: Any) -> tuple[int, int]: + if isinstance(raw, (list, tuple)) and raw: + first = raw[0] + if isinstance(first, (list, tuple)) and len(first) >= 2: + return (int(first[0]), int(first[1])) + if len(raw) >= 2: + return (int(raw[0]), int(raw[1])) + return (-1, -1) def prepare_parakeet_tdt_audio_features( @@ -290,10 +305,13 @@ def load_parakeet_tdt_config(model_source: str) -> ParakeetTDTConfig: ) ), conv_kernel_size=int(_cfg_get(root, "conv_kernel_size", _cfg_get(encoder, "conv_kernel_size", 9))), + conv_norm_type=str(_cfg_get(root, "conv_norm_type", _cfg_get(encoder, "conv_norm_type", "batch_norm"))).lower(), + conv_is_causal=str(_cfg_get(root, "conv_context_size", _cfg_get(encoder, "conv_context_size", ""))).lower() == "causal", subsampling_factor=int(_cfg_get(root, "subsampling_factor", _cfg_get(encoder, "subsampling_factor", 8))), subsampling_conv_channels=int( _cfg_get(root, "subsampling_conv_channels", _cfg_get(encoder, "subsampling_conv_channels", 256)) ), + causal_downsampling=bool(_cfg_get(root, "causal_downsampling", _cfg_get(encoder, "causal_downsampling", False))), predictor_hidden_dim=int( _cfg_get(root, "predictor_hidden_dim", _cfg_get(root, "decoder_hidden_size", _cfg_get(prediction, "pred_hidden", _cfg_get(model_defaults, "pred_hidden", 640)))) ), @@ -306,6 +324,8 @@ def load_parakeet_tdt_config(model_source: str) -> ParakeetTDTConfig: blank_id=blank_id, vocabulary=vocabulary, encoder_hidden_act=str(_cfg_get(root, "encoder_hidden_act", _cfg_get(encoder, "activation", _cfg_get(encoder, "hidden_act", "silu")))).lower(), + attention_context_style=str(_cfg_get(encoder, "att_context_style", "regular")), + attention_context_size=_parse_attention_context_size(_cfg_get(encoder, "att_context_size", None)), ) @@ -406,14 +426,12 @@ def _relative_position_embeddings(*, seq_len: int, hidden_dim: int, device: torc def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: batch, heads, seq_len, _ = query.shape scores = torch.matmul(query, relative_key.transpose(-1, -2)) - rel_index = ( - torch.arange(seq_len, device=query.device).view(seq_len, 1) - - torch.arange(seq_len, device=query.device).view(1, seq_len) - + (seq_len - 1) - ) - rel_index = rel_index.view(1, 1, seq_len, seq_len).expand(batch, heads, seq_len, seq_len) - gathered = scores.gather(-1, rel_index) - return gathered * float(scale) + pos_len = 2 * seq_len - 1 + shifted = F.pad(scores, (1, 0)) + shifted = shifted.reshape(batch, heads, pos_len + 1, seq_len) + shifted = shifted[:, :, 1:, :] + shifted = shifted.reshape(batch, heads, seq_len, pos_len) + return shifted[:, :, :, :seq_len] * float(scale) def _relative_position_bias_static(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: @@ -427,6 +445,30 @@ def _relative_position_bias_static(query: torch.Tensor, relative_key: torch.Tens return x[:, :, :, :seq_len] * float(scale) +def _attention_context_bias(config: ParakeetTDTConfig, seq_len: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None: + left, right = config.attention_context_size + if left < 0 and right < 0: + return None + style = config.attention_context_style + positions = torch.arange(seq_len, device=device) + if style == "chunked_limited" and right >= 0: + chunk_size = int(right) + 1 + query_chunks = torch.div(positions, chunk_size, rounding_mode="trunc").view(seq_len, 1) + key_chunks = torch.div(positions, chunk_size, rounding_mode="trunc").view(1, seq_len) + diff_chunks = query_chunks - key_chunks + left_chunks = (int(left) // chunk_size) if left >= 0 else 10000 + allowed = (diff_chunks <= left_chunks) & (diff_chunks >= 0) + else: + query_positions = positions.view(seq_len, 1) + key_positions = positions.view(1, seq_len) + allowed = torch.ones((seq_len, seq_len), dtype=torch.bool, device=device) + if left >= 0: + allowed &= key_positions >= query_positions - int(left) + if right >= 0: + allowed &= key_positions <= query_positions + int(right) + return (~allowed).to(dtype=dtype).view(1, 1, seq_len, seq_len) * -1e4 + + class ParakeetTDTFeedForward(nn.Module): def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): super().__init__() @@ -499,6 +541,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: rel_k_heads, scale=self.config.attention_scale, ) + context_bias = _attention_context_bias(self.config, seq_len, device=x.device, dtype=rel_bias.dtype) + if context_bias is not None: + rel_bias = rel_bias + context_bias attn = F.scaled_dot_product_attention( q_u_heads, k_heads, @@ -526,11 +571,15 @@ def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, hidden_dim, hidden_dim, kernel_size=kernel, - padding=kernel // 2, + padding=0 if config.conv_is_causal else kernel // 2, groups=hidden_dim, bias=f"{prefix}.depthwise_conv.bias" in state_dict, ) - self.batch_norm = nn.BatchNorm1d(hidden_dim, eps=1e-5, affine=True, track_running_stats=True) + self.norm_type = config.conv_norm_type + if self.norm_type == "layer_norm": + self.batch_norm = nn.LayerNorm(hidden_dim) + else: + self.batch_norm = nn.BatchNorm1d(hidden_dim, eps=1e-5, affine=True, track_running_stats=True) self.pointwise_conv2 = nn.Conv1d( hidden_dim, hidden_dim, @@ -554,12 +603,13 @@ def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, ) self.batch_norm.weight.data.copy_(state_dict[f"{prefix}.batch_norm.weight"].to(dtype=self.batch_norm.weight.dtype)) self.batch_norm.bias.data.copy_(state_dict[f"{prefix}.batch_norm.bias"].to(dtype=self.batch_norm.bias.dtype)) - self.batch_norm.running_mean.data.copy_( - state_dict[f"{prefix}.batch_norm.running_mean"].to(dtype=self.batch_norm.running_mean.dtype) - ) - self.batch_norm.running_var.data.copy_( - state_dict[f"{prefix}.batch_norm.running_var"].to(dtype=self.batch_norm.running_var.dtype) - ) + if self.norm_type != "layer_norm": + self.batch_norm.running_mean.data.copy_( + state_dict[f"{prefix}.batch_norm.running_mean"].to(dtype=self.batch_norm.running_mean.dtype) + ) + self.batch_norm.running_var.data.copy_( + state_dict[f"{prefix}.batch_norm.running_var"].to(dtype=self.batch_norm.running_var.dtype) + ) self.config = config def forward(self, x: torch.Tensor, pad_mask: torch.Tensor | None = None) -> torch.Tensor: @@ -568,8 +618,13 @@ def forward(self, x: torch.Tensor, pad_mask: torch.Tensor | None = None) -> torc x = F.glu(x, dim=1) if pad_mask is not None: x = x * pad_mask + if self.config.conv_is_causal: + x = F.pad(x, (self.config.conv_kernel_size - 1, 0)) x = self.depthwise_conv(x) - x = self.batch_norm(x) + if self.norm_type == "layer_norm": + x = self.batch_norm(x.transpose(1, 2)).transpose(1, 2) + else: + x = self.batch_norm(x) x = _apply_activation(x, self.config.encoder_hidden_act) x = self.pointwise_conv2(x) return x.transpose(1, 2) @@ -606,14 +661,16 @@ class ParakeetTDTPreEncode(nn.Module): def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): super().__init__() channels = config.subsampling_conv_channels + self.causal_downsampling = config.causal_downsampling + conv_padding = 0 if config.causal_downsampling else 1 self.conv = nn.ModuleList( [ - nn.Conv2d(1, channels, kernel_size=3, stride=2, padding=1), + nn.Conv2d(1, channels, kernel_size=3, stride=2, padding=conv_padding), nn.Identity(), - nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1, groups=channels), + nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=conv_padding, groups=channels), nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0), nn.Identity(), - nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1, groups=channels), + nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=conv_padding, groups=channels), nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0), ] ) @@ -624,7 +681,8 @@ def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor _copy_conv2d_weight(self.conv[6], state_dict["encoder.pre_encode.conv.6.weight"], bias=state_dict["encoder.pre_encode.conv.6.bias"]) projected_width = config.num_mel_bins for _ in range(3): - projected_width = (projected_width + 2 - 3) // 2 + 1 + all_padding = 3 if config.causal_downsampling else 2 + projected_width = (projected_width + all_padding - 3) // 2 + 1 self.out = nn.Linear(channels * projected_width, config.hidden_dim) _copy_linear_weight( self.out, @@ -636,8 +694,14 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: if input_features.ndim != 3: raise ValueError(f"expected input_features [batch, frames, mels], got {tuple(input_features.shape)}") x = input_features.unsqueeze(1) + if self.causal_downsampling: + x = F.pad(x, (2, 1, 2, 1)) x = F.relu(self.conv[0](x)) + if self.causal_downsampling: + x = F.pad(x, (2, 1, 2, 1)) x = F.relu(self.conv[3](self.conv[2](x))) + if self.causal_downsampling: + x = F.pad(x, (2, 1, 2, 1)) x = F.relu(self.conv[6](self.conv[5](x))) x = x.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[2], -1) return self.out(x) @@ -679,7 +743,7 @@ def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): with torch.no_grad(): rel_k = attn.linear_pos(rel_pos).view( 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim - ).flip(1) + ) self.register_buffer("_rel_k", rel_k.detach(), persistent=False) def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torch.Tensor: @@ -700,6 +764,9 @@ def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torc rel_k_heads, scale=self.config.attention_scale, ) + context_bias = _attention_context_bias(self.config, seq_len, device=x.device, dtype=rel_bias.dtype) + if context_bias is not None: + rel_bias = rel_bias + context_bias if key_bias is not None: rel_bias = rel_bias + key_bias attn = F.scaled_dot_product_attention( diff --git a/python/tests/test_graph.py b/python/tests/test_graph.py index c8b05f88e..aebb4f9f9 100644 --- a/python/tests/test_graph.py +++ b/python/tests/test_graph.py @@ -110,6 +110,30 @@ def test_composed_inference_graph(self): np.testing.assert_allclose(out, expected, atol=1e-2) +class TestGraphAttentionOps(unittest.TestCase): + + def test_rel_pos_bias_matches_pytorch_indexing(self): + g = Graph() + q = g.input((1, 3, 2, 4)) + r = g.input((1, 5, 2, 4)) + y = g.rel_pos_bias(q, r, 0.25) + + q_data = (np.arange(24, dtype=np.float32).reshape(1, 3, 2, 4) / 10).astype(np.float16) + r_data = (np.arange(40, dtype=np.float32).reshape(1, 5, 2, 4) / 20).astype(np.float16) + g.set_input(q, q_data) + g.set_input(r, r_data) + g.execute() + + scores = np.einsum("bthd,brhd->bhtr", q_data.astype(np.float32), r_data.astype(np.float32)) + rel_index = ( + np.arange(3, dtype=np.int64).reshape(3, 1) + - np.arange(3, dtype=np.int64).reshape(1, 3) + + 2 + ) + expected = np.take_along_axis(scores, rel_index.reshape(1, 1, 3, 3), axis=-1) * 0.25 + np.testing.assert_allclose(y.numpy(), expected.astype(np.float16), atol=1e-2) + + class TestGraphTensorOps(unittest.TestCase): def test_view(self): @@ -571,6 +595,176 @@ def test_save_load_roundtrip_softmax_cat(self): self.assertEqual(info["precision"], expected_info["precision"]) np.testing.assert_allclose(out, expected, atol=5e-2) + def test_save_load_roundtrip_glu_axis(self): + g = Graph() + + x = g.input((1, 4, 2)) + out_node = g.glu(x, axis=1) + + x_data = np.array( + [[[1.0, 2.0], [3.0, 4.0], [-2.0, 1.0], [0.5, -1.0]]], + dtype=np.float16, + ) + + g.set_input(x, x_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "glu_axis_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_x = self._rebind_tensor(loaded, x) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_x, x_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_save_load_roundtrip_audio_preprocess_params(self): + g = Graph() + + waveform = g.input((160,), dtype=Graph.FP32) + mel_filters = g.mel_filter_bank( + num_frequency_bins=33, + num_mel_filters=8, + min_frequency=20.0, + max_frequency=7600.0, + sampling_rate=16000, + norm_type=1, + scale_type=1, + ) + out_node = g.spectrogram( + waveform, + mel_filters, + frame_length=64, + hop_length=16, + fft_length=64, + power=1.0, + center=False, + pad_mode=1, + mel_floor=1e-5, + log_mel_mode=1, + dither=0.0, + preemphasis=0.97, + remove_dc_offset=True, + ) + + waveform_data = np.linspace(-0.4, 0.6, 160, dtype=np.float32) + + g.set_input(waveform, waveform_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "audio_preprocess_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_waveform = self._rebind_tensor(loaded, waveform) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_waveform, waveform_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-5) + + def test_save_load_roundtrip_image_preprocess_params(self): + g = Graph() + + pixels = g.input((2, 2, 3), dtype=Graph.FP32) + out_node = g.image_preprocess( + pixels, + src_width=2, + src_height=2, + target_width=2, + target_height=2, + patch_size=1, + channels=3, + rescale_factor=0.25, + mean=[0.1, 0.2, 0.3], + std_dev=[1.0, 2.0, 4.0], + ) + + pixel_data = np.array( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + dtype=np.float32, + ) + + g.set_input(pixels, pixel_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "image_preprocess_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_pixels = self._rebind_tensor(loaded, pixels) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_pixels, pixel_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-6) + + def test_save_load_roundtrip_gaussian_topk_scalar(self): + g = Graph() + + x = g.input((2, 4)) + out_node = g.gaussian_topk(x, ppf=0.8) + + x_data = np.array( + [[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], + dtype=np.float16, + ) + + g.set_input(x, x_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "gaussian_topk_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_x = self._rebind_tensor(loaded, x) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_x, x_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-2) + def test_save_load_roundtrip_matmul(self): g = Graph() diff --git a/python/tests/test_transpile_aten_ops.py b/python/tests/test_transpile_aten_ops.py index b0f0ad7f4..249a1172d 100644 --- a/python/tests/test_transpile_aten_ops.py +++ b/python/tests/test_transpile_aten_ops.py @@ -9,8 +9,10 @@ from cactus.transpile.graph_ir import IRGraph from cactus.transpile.graph_ir import IRNode from cactus.transpile.graph_ir import IRValue +from cactus.transpile.lower import transpile_preoptimized_ir from cactus.transpile.normalize import normalize_target from cactus.transpile.optimize_graph import fuse_dense_mlp_tq +from cactus.transpile.optimize_graph import optimize_graph class OddlyNamedLinearBlock(torch.nn.Module): @@ -33,6 +35,55 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.silu(self.not_a_transformer_layer["banana"](x)) +class ManualLSTMCell(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight_ih = torch.nn.Parameter(torch.randn(16, 4) * 0.2) + self.weight_hh = torch.nn.Parameter(torch.randn(16, 4) * 0.2) + self.bias = torch.nn.Parameter(torch.randn(16) * 0.1) + + def forward( + self, + x: torch.Tensor, + h: torch.Tensor, + c: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + gates = torch.nn.functional.linear(x, self.weight_ih) + torch.nn.functional.linear(h, self.weight_hh) + self.bias + i, f, g, o = torch.chunk(gates, 4, dim=-1) + c_next = torch.sigmoid(f) * c + torch.sigmoid(i) * torch.tanh(g) + h_next = torch.sigmoid(o) * torch.tanh(c_next) + return h_next, c_next + + +class ReusedLSTMStateOutput(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.cell0 = ManualLSTMCell() + self.cell1 = ManualLSTMCell() + + def forward( + self, + x: torch.Tensor, + h0: torch.Tensor, + c0: torch.Tensor, + h1: torch.Tensor, + c1: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + h0_next, c0_next = self.cell0(x, h0, c0) + h1_next, _ = self.cell1(h0_next, h1, c1) + return h1_next, h0_next, c0_next + + +class DepthwiseConv2dStride2Padding0(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(4, 1, 3, 3) * 0.1) + self.bias = torch.nn.Parameter(torch.randn(4) * 0.1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.conv2d(x, self.weight, self.bias, stride=2, padding=0, groups=4) + + def _op_counts(module: torch.nn.Module) -> Counter[str]: example = torch.randn(2, 4) captured = capture_model(module, (example,)) @@ -64,6 +115,44 @@ def test_import_records_canonical_aten_op_metadata() -> None: assert all("module_paths" not in meta for meta in node_meta) +def test_reused_lstm_state_output_is_materialized() -> None: + torch.manual_seed(3) + module = ReusedLSTMStateOutput().eval().to(dtype=torch.float16) + inputs = tuple(torch.randn(1, 4, dtype=torch.float16) for _ in range(5)) + with torch.no_grad(): + expected = module(*inputs) + + captured = capture_model(module, inputs) + optimized = optimize_graph(captured.ir_graph) + lowered = transpile_preoptimized_ir(optimized) + lowered.set_inputs([value.numpy() for value in inputs]) + actual = [value.numpy() for value in lowered.execute()] + + for got, want in zip(actual, expected, strict=True): + torch.testing.assert_close( + torch.from_numpy(got), + want.detach().cpu(), + atol=2e-3, + rtol=2e-3, + ) + + +def test_depthwise_conv2d_stride2_padding0_lowers() -> None: + torch.manual_seed(11) + module = DepthwiseConv2dStride2Padding0().eval() + x = torch.randn(1, 4, 9, 11) + with torch.no_grad(): + expected = module(x) + + captured = capture_model(module, (x,)) + optimized = optimize_graph(captured.ir_graph) + lowered = transpile_preoptimized_ir(optimized) + lowered.set_inputs([x.numpy()]) + actual = torch.from_numpy(lowered.execute()[0].numpy()) + + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-3, rtol=2e-3) + + def test_dense_mlp_fusion_uses_topology_not_layer_names() -> None: graph = IRGraph( values={