prompt_kernel.hpp - mozsearch

Enable keyboard shortcuts

#pragma once

#include "model_loader.hpp"

#include <vector>

namespace pk {

// Post-encoder prompt conditioning for multilingual nemotron models.

//

// Mirrors NeMo EncDecRNNTBPEModelWithPrompt.forward():

//   encoded[T, D] = transpose(encoder_out[D, T])

//   onehot[T, P]  = one_hot(prompt_index, num_prompts) broadcast over T

//   out[T, D]     = prompt_kernel(cat([encoded, onehot]))  // Linear->ReLU->Linear

// then transposed back to channels-first [D, T].

//

// The one-hot is constant over time (one language per utterance), so this is a

// concat + two matmuls + ReLU. Weights: prompt_kernel.0.{weight,bias} (D+P->2D),

// prompt_kernel.2.{weight,bias} (2D->D). present() is false for non-prompt

// models (callers skip apply()).

class PromptKernel {

public:

    explicit PromptKernel(const ModelLoader& ml);

    bool present() const { return present_; }

    // Apply the prompt projection to a channels-first encoder output

    // enc_out[d_model, T] for the given prompt_index, writing channels-first

    // out[d_model, T]. If present()==false this is a no-op copy (out = enc_out).

    void apply(const std::vector<float>& enc_out, int d_model, int T,

               int prompt_index, std::vector<float>& out) const;

private:

    const ModelLoader& ml_;

    bool present_ = false;

    int  num_prompts_ = 0;

    int  d_model_ = 0;

};

} // namespace pk