Visual Token Pruner#

class VisualTokenPruner#

Abstract prefill-time visual-token pruner (batch 1).

The non-virtual pruneForPrefill() owns everything algorithms must not diverge on: the enablement guards (minVisualTokens, target-is-a-reduction), modality partitioning, and the target computation. Subclasses implement prune() — with full freedom over what “pruning” means (subset selection, token merging, per-image quotas, …) — and typically finish by calling compactToKeepList(), which owns the invariant-laden buffer compaction: all text tokens are always kept; kept tokens retain their original absolute RoPE positions; decode continues at the position after the unpruned sequence (matching the HF DART reference).

Instances are created through createVisualTokenPruner() and reused across requests, so per-request device work buffers should be preallocated in the constructor. The caller is responsible for the runtime gates (fresh KV cache, mRoPE engine, no spec decode, …) and for shrinking the context lengths to the returned pruned length.

Subclassed by trt_edgellm::rt::DartPruner

Public Functions

virtual ~VisualTokenPruner() = default#

virtual char const *name() const noexcept = 0#: Algorithm name (matches the registry key).

int32_t pruneForPrefill( std::vector<int32_t> const &hostTokenIds, PipelineIO &io, int32_t origLen, cudaStream_t stream )#

Prune the assembled prefill inputs for a batch-1 request.

May synchronize stream (algorithm-dependent); on return the compaction kernels are enqueued on stream and the PipelineIO tensors are reshaped to the pruned length.

Parameters:

hostTokenIds – The request’s expanded token ids (image placeholders already inserted); visual positions are tokenId == imageTokenId.
io – Pipeline buffers; inputsEmbeds must currently be [1, origLen, hiddenSize].
origLen – The unpruned prefill length (== hostTokenIds.size()).
stream – CUDA stream all device work runs on.

Returns:

The pruned length P (< origLen), or origLen when pruning is skipped (no/too-few visual tokens, or the target keep count is not a reduction).

inline VisualPrunerConfig const &config() const noexcept#

struct VisualPrunerConfig#

Visual-token pruning configuration (runtime-side; nothing is required in the exported engine config — the prune operates on runtime buffers only).

Public Members

bool enabled = {false}#

std::string algorithm = {"dart"}#

Pruning algorithm name, resolved through the pruner registry. Built-in: “dart” (default, duplication-aware; paper: “Stop Looking for Important Tokens

in Multimodal Language Models: Duplication Matters More”).

float reductionRatio = {0.25F}#: Fraction of visual tokens to remove (0.25 = keep 75%).

int32_t minVisualTokens = {16}#: Skip pruning when the request has fewer visual tokens than this (accuracy and break-even guard: tiny images gain nothing from pruning).

int32_t pivotImageTokens = {4}#

int32_t pivotTextTokens = {4}#

struct ImageSpan#

One contiguous run of visual tokens (one image, or one video-frame block) with its per-image retention quota. Pruning within each span independently — rather than over one global candidate pool — guarantees every image keeps its proportional share of tokens (at least one), so a low-information image can never be starved by the others.

Public Members

int32_t begin = {0}#: First sequence position of the span (inclusive)

int32_t end = {0}#: One past the last sequence position of the span (exclusive)

int32_t targetTokens = {0}#: Visual tokens to retain from this span (1 <= target <= end - begin)

struct PruneRequest#

Guard-checked, modality-partitioned view of one batch-1 prefill request, prepared by VisualTokenPruner::pruneForPrefill and handed to the algorithm hook.

Public Members

Tensor const *embeds = {nullptr}#: Non-owning [origLen, hiddenSize] FP16 GPU view of the assembled input embeddings.

std::vector<int32_t> const *imagePositions = {nullptr}#: Positions of visual tokens in the sequence (ascending, non-empty).

std::vector<int32_t> const *textPositions = {nullptr}#: Positions of all non-visual tokens in the sequence (ascending).

std::vector<ImageSpan> const *imageSpans = {nullptr}#: Contiguous visual spans (one per image) with per-span retention quotas.

int32_t targetImageTokens = {0}#: Total number of visual tokens to retain — the sum of the per-span quotas (1 <= targetImageTokens < imagePositions->size()).

int32_t origLen = {0}#: Unpruned prefill length (== imagePositions->size() + textPositions->size()).

void trt_edgellm::rt::registerVisualPruner( std::string const &name, VisualPrunerFactory factory )#: Register a pruning algorithm under name (case-sensitive). The built-in “dart” is registered automatically; call this to plug in a custom algorithm before constructing the runtime. Re-registering a name replaces the previous factory.

std::unique_ptr<VisualTokenPruner> trt_edgellm::rt::createVisualTokenPruner( VisualPrunerConfig const &config, LLMEngineConfig const &engineConfig )#

Instantiate the pruner named by config.algorithm.

Throws:: std::runtime_error – if the name is not registered or the config is invalid.

std::vector<std::string> trt_edgellm::rt::registeredVisualPrunerNames( )#: Names of all registered pruning algorithms (for CLI help / validation).