Pipeline Io#

struct StreamingPrefillBuffers#

Persistent copies of the prefill-time input embeddings and engine hidden_states output, used by streaming consumers that run concurrently with the base model’s decode loop.

The base model’s inputsEmbeds and outputHiddenStates tensors are reshaped to {B, 1, H} and overwritten by every decode step. This struct retains the {B, prefillLen, H} view as it stood at the end of prefill, so consumers reading these buffers do not race with decode writes.

Public Functions

void populateFromPrefill(
Tensor const &liveInputEmbeds,
Tensor const &liveEngineHiddenStates,
int32_t batch,
int32_t prefillLen,
int32_t hiddenSize,
int32_t maxBatch,
int32_t maxSeq,
cudaStream_t stream,
)#

Allocate on first call (sized to the worst case {maxBatch, maxSeq, hiddenSize}), reshape to the current request’s {batch, prefillLen, hiddenSize}, and copy from the live PipelineIO buffers on stream. Subsequent calls reuse the same allocation. Must be invoked after prefill and before the first decode step on the same stream so the copies precede any overwrite of outputHiddenStates.

Public Members

Tensor inputEmbeds#

Prefill-time layer-0 input embeddings.

Tensor engineHiddenStates#

Prefill-time engine hidden_states output.

struct PipelineIO#

All tensors flowing through the inference pipeline. POINTER STABILITY INVARIANT: After buildTensorMap() is called, this struct must not be moved, and deepstackEmbeds must not be resized. TensorMap holds Tensor* pointers into these members — any reallocation invalidates them.

Public Members

Tensor inputsEmbeds#
Tensor outputLogits#
Tensor selectTokenIndices#
Tensor contextLengths#

GPU.

Tensor hostContextLengths#

CPU (pinned, [maxBatch] INT32)

Tensor hostSelectTokenIndices#

CPU (pinned, [maxBatch, 1] INT64) — pairs with selectTokenIndices for H2D staging

Tensor visionBlockIds#

Gemma4 Unified block IDs, [batch, seq_len] INT32; empty for other models.

std::vector<Tensor> deepstackEmbeds#
Tensor mropeCosSin#
Tensor baseHiddenStates#
Tensor draftHiddenStatesIn#
Tensor draftHiddenStatesOut#
Tensor outputHiddenStates#

Engine hidden_states output. Used by the vanilla LLM path; SpecDecode routes its hidden states through baseHiddenStates instead.

StreamingPrefillBuffers streamingPrefill#

Per-request copies of inputsEmbeds / outputHiddenStates that streaming consumers (e.g. the Qwen3-Omni Talker) read while the base model’s decode loop overwrites the live buffers. Populated by LLMInferenceRuntime only when streaming output is enabled for the request; otherwise the buffers stay empty (no allocation cost).

Tensor packedAttentionMask#

Packed proposal attention mask, [batch, proposalSize, divUp(proposalSize, 32)] INT32. Written by proposal/verify input preparation kernels; consumed by the base and draft engines via the kAttentionMask binding.

Tensor specDecodePositionIds#

SpecDecode position IDs, [batch, proposalSize] INT32. Written by proposal/verify input preparation kernels; consumed by the base and draft engines via the kAttentionPosId binding.

Tensor specVerifyPhaseMarker#

Shape-only marker for hybrid MTP/DFlash base engines. The runtime binds this tensor at shape [0] for normal prefill/decode and [1] for spec verify; plugins branch on the shape, not the payload.

Tensor specTreeParentIds#

DDTree parent node ids, [batch, proposalSize] INT32. Runtime-owned metadata for tree attention and hybrid state plugin bindings.

Tensor specTreeDepths#

DDTree depth per node, [batch, proposalSize] INT32. Runtime-owned metadata for tree attention and hybrid state plugin bindings.

Public Static Functions

static PipelineIO createForLLM(
LLMEngineConfig const &cfg,
cudaStream_t stream,
)#

Build PipelineIO for the vanilla single-engine LLM runtime (basic I/O tensors, deepstack embeds, MRope cos/sin cache).

static PipelineIO createForSpecDecode(
DeploymentConfig const &bundle,
int32_t maxRuntimeBatchSize,
cudaStream_t stream,
)#

Build PipelineIO for a two-engine speculative-decoding runtime (basic I/O, hidden states, deepstack embeds, MRope cos/sin cache).

void trt_edgellm::rt::allocateBasicIO(
PipelineIO &io,
int32_t maxBatch,
int32_t maxSeq,
int32_t hiddenSize,
int32_t vocabSize,
nvinfer1::DataType dtype,
)#
void trt_edgellm::rt::allocateDeepstackEmbeds(
PipelineIO &io,
int32_t numFeatures,
int32_t maxBatch,
int32_t maxSeq,
int32_t hiddenSize,
nvinfer1::DataType dtype,
)#
void trt_edgellm::rt::allocateSpecDecodeHiddenStates(
PipelineIO &io,
int32_t maxBatch,
int32_t maxSeq,
int32_t baseHiddenDim,
int32_t draftHiddenDim,
nvinfer1::DataType dtype,
bool allocateDraftHiddenStates,
)#
void trt_edgellm::rt::allocateMRope(
PipelineIO &io,
int32_t maxBatch,
int32_t maxKVCacheCapacity,
int32_t rotaryDim,
)#
void trt_edgellm::rt::buildTensorMap(
TensorMap &map,
PipelineIO &io,
SharedResources &res,
LLMEngineConfig const &cfg,
int32_t kvCacheIndex,
)#

Populate a TensorMap from PipelineIO + SharedResources for engine binding.

This is the critical glue function that wires all allocated tensors into the name-to-pointer map consumed by TensorRegistry::bindAll().

Parameters:

  • map – Output map to populate.

  • io – Pipeline I/O tensors.

  • res – Shared resources (KV caches, RoPE pool, LoRA, zero buffer).

  • cfg – Engine configuration.

  • kvCacheIndex – Index into res.cacheManagers for the target engine.

void trt_edgellm::rt::buildTensorMapForSpecDecodeDraft(
TensorMap &map,
PipelineIO &io,
SharedResources &res,
LLMEngineConfig const &cfg,
)#

Populate a TensorMap for a SpecDecode draft engine. Delegates to buildTensorMap with kvCacheIndex=1 for the common bindings, then patches in draft-engine- specific bindings (base/draft hidden states in+out, packed proposal attention mask, proposal position IDs).

Preconditions: io must have been constructed via PipelineIO::createForSpecDecode for an EAGLE/MTP-style draft path where draftHiddenStatesIn/Out are populated alongside baseHiddenStates, packedAttentionMask, and specDecodePositionIds. DFlash uses its own draft TensorMap.

Parameters:

  • map – Output map for the draft engine’s bindings.

  • io – Pipeline I/O (must be the SpecDecode-flavoured one).

  • res – Shared resources.

  • cfg – Draft engine configuration.

void trt_edgellm::rt::buildTensorMapForGemma4MTPDraft(
TensorMap &map,
PipelineIO &io,
SharedResources &res,
DeploymentConfig const &bundle,
)#

Populate a TensorMap for a Gemma4 MTP assistant draft engine.

Unlike EAGLE/MTP draft engines, Gemma4 assistant engines do not own a draft KV cache. Their past_key_values_* bindings are zero-copy aliases to the base target KV cache selected by draftCfg.gemma4MTPKVSharingMap.