Code Predictor#

class CodePredictor#

Residual codebook prediction, one engine behind the talker.

One frame is one sequence: the engine runs a two-row prefill and then either numRvqLayers - 1 decode steps or, with speculation on, draft/verify rounds. The KV cache holds a single frame, so the talker resets it at the start of every frame and no state crosses a frame boundary.

This owns its engine, its weights and every workspace it touches. Two tensors belong to the talker’s codec vocabulary and are only borrowed: the code_0 embedding table and the per-frame codec-hidden buffer whose middle rows are filled here for the talker’s residual. Both arrive in Binding.

The small_to_mtp projection needs the CuTe DSL GEMM module, which the talker loads before preallocate() and unloads on destruction.

Public Functions

void preallocate( std::string const &artifactDir, std::string const &checkpointDir, int32_t maxBatch, int32_t specVerifySize, cudaStream_t stream )#: Load the engine and claim every workspace. checkpointDir may be empty. specVerifySize is the verify window the deployment asked for; 0, or anything this engine cannot honour, leaves the autoregressive loop in place. Throws on failure, with a message saying what failed.

void bind(Binding binding)#: Declare the talker’s half of the contract, which this checks against the artifacts.

inline size_t requiredContextMemory() const noexcept#

bool setContextMemory(rt::Tensor &memory)#

Hand the engine an arena the talker owns, releasing the one claimed here.

Legal only because both engines run on the talker’s stream: TensorRT forbids giving one context’s device memory to a context executing concurrently.

Returns:: false when memory is smaller than requiredContextMemory().

void resetKvCache(int32_t batchSize, cudaStream_t stream)#: Zero the KV lengths for batchSize sequences. One frame is one sequence.

bool runFrame( FramePlan const &plan, std::vector<std::vector<int32_t>> &outputCodesPerBatch, cudaStream_t stream )#

Generate codes 1..numRvqLayers for one codec frame.

Parameters:: outputCodesPerBatch – [activeBatchSize][numCodesPerFrame], code_0 first. Rows keep their capacity across frames, so reuse the same vector.
Returns:: false when an engine call fails; the talker must then abandon the frame.

bool captureDecodeGraphs(cudaStream_t stream)#: Capture one decode graph per batch size. Leaves the KV cache empty, ready for a frame.

void logAndResetSpecStats()#: Report and reset the speculation counters accumulated since the last call. No-op while speculation is off or no round has run.

inline int32_t numRvqLayers() const noexcept#

inline int32_t codebookSize() const noexcept#

rt::Tensor const &codecEmbedTable(int32_t layer) const#: Residual codebook table layer, [codebookSize, talkerHiddenSize], 0-based over the numRvqLayers residual codebooks. Borrowed; the talker reads them for its own lookups.

struct Binding#

What the talker lends, checked once at load.

Deployment-scoped, not per-request: the two tensors are borrowed and must stay alive, with at least this capacity, for as long as frames are driven.

Public Members

int32_t talkerHiddenSize = {0}#

int32_t numCodesPerFrame = {0}#

int32_t maxBatchSize = {0}#

rt::Tensor const *codeZeroEmbedTable = {nullptr}#: [codecVocabSize, talkerHiddenSize] table code_0 is looked up in. Which table that is depends on the checkpoint, which is why the talker picks it.

rt::Tensor *codecHiddens = {nullptr}#: [maxBatchSize, numCodesPerFrame, talkerHiddenSize]. Rows 1..numRvqLayers-1 are written here; rows 0 and numRvqLayers stay the talker’s to fill.

struct FramePlan#

One codec frame’s inputs. Everything is borrowed for the duration of the call.

Public Members

int32_t activeBatchSize = {0}#

int32_t frameIdx = {-1}#: Frames decoded since the batch’s prefill, the same for every row; keys the offsets of the frame’s non-speculative draws.

rt::Tensor const *codeZero = {nullptr}#: Pinned host INT32, at least [activeBatchSize]: the talker’s code_0, uploaded asynchronously.

rt::Tensor const *talkerLastHidden = {nullptr}#: [activeBatchSize, talkerHiddenSize] FP16 device view of the talker’s last hidden state.

SamplingParams const *sampling = {nullptr}#: Only temperature/topK/topP are read; batch and vocabulary are re-derived per call.

struct Binding

What the talker lends, checked once at load.

Deployment-scoped, not per-request: the two tensors are borrowed and must stay alive, with at least this capacity, for as long as frames are driven.

Public Members

int32_t talkerHiddenSize = {0}

int32_t numCodesPerFrame = {0}

int32_t maxBatchSize = {0}

rt::Tensor const *codeZeroEmbedTable = {nullptr}: [codecVocabSize, talkerHiddenSize] table code_0 is looked up in. Which table that is depends on the checkpoint, which is why the talker picks it.

rt::Tensor *codecHiddens = {nullptr}: [maxBatchSize, numCodesPerFrame, talkerHiddenSize]. Rows 1..numRvqLayers-1 are written here; rows 0 and numRvqLayers stay the talker’s to fill.

struct FramePlan

One codec frame’s inputs. Everything is borrowed for the duration of the call.

Public Members

int32_t activeBatchSize = {0}

int32_t frameIdx = {-1}: Frames decoded since the batch’s prefill, the same for every row; keys the offsets of the frame’s non-speculative draws.

rt::Tensor const *codeZero = {nullptr}: Pinned host INT32, at least [activeBatchSize]: the talker’s code_0, uploaded asynchronously.

rt::Tensor const *talkerLastHidden = {nullptr}: [activeBatchSize, talkerHiddenSize] FP16 device view of the talker’s last hidden state.

SamplingParams const *sampling = {nullptr}: Only temperature/topK/topP are read; batch and vocabulary are re-derived per call.

void trt_edgellm::rt::qwen3_omni::prepareRaggedStep( rt::RaggedExecutionBatch &batch, rt::PipelineIO &io, rt::SharedResources &resources, rt::LLMEngineConfig const &config, std::vector<int32_t> const &queryLengths, std::vector<int32_t> const &pastLengths, int32_t batchSize, int32_t queryWidth, rt::ExecutionPhase phase, cudaStream_t stream )#

Publish one step’s token-major layout to the engine’s ragged bindings.

Every execute() reads the offsets, past lengths and state indices this writes; they are not step-invariant, so binding them once per prepare() leaves the attention kernels addressing whatever the previous step left behind. The talker steps its own engine through this too.