Jump Forward#

class JumpForwardPlanner#

Turns a grammar-forced string into tokens that can be appended without a forward pass.

The forced text is re-encoded together with the window, the most recent tokens of the sequence, so the result is the tokenizer’s own split of the joined text rather than the forced text tokenized on its own. Window tokens at the end that the re-encoding changes are rolled back and replaced; a change to any token before them skips the jump, because those are emitted or have KV.

A tokenizer that prepends a word marker (Tokenizer::prependsSpaceMarker) only reproduces a window that starts at a word boundary, so the caller must widen the window to one; a window text that does not start with a space is skipped.

The last planned token is handed back when another vocabulary token extends its piece and the plan does not complete the grammar: the text after it is not known yet, and a canonical tokenization may merge it with what follows.

The grammar state is left unchanged.

Public Functions

JumpForwardPlanner( tokenizer::Tokenizer const &tokenizer, GuidedDecoder const &decoder, int32_t maxJumpTokens )#

Both references are borrowed and must outlive the planner.

Parameters:: maxJumpTokens – Tokens one jump may add to a sequence

JumpForwardPlan plan( int32_t slot, std::string const &forced, std::vector<int32_t> const &windowFullIds, int32_t mutableTailCount, int32_t budgetTokens ) const#

Parameters:

slot – Decoder slot whose grammar produced forced
forced – Text the grammar forces next, see GuidedDecoder::forcedContinuation
windowFullIds – Most recent tokens of the sequence, full vocabulary
mutableTailCount – Trailing window tokens that may be rolled back: not emitted, no KV
budgetTokens – Tokens the sequence may still grow by. Rolled-back tokens are replaced and do not count

Returns:

An empty plan with its skip reason when nothing can be appended

void buildWindow( std::vector<int32_t> const &tokens, int32_t generatedCount, std::vector<int32_t> &window ) const#

The window a jump re-encodes: the last token, widened backwards to a word boundary when the tokenizer prepends a word marker. Only generated tokens are candidates: prompt pieces may be template markup the tokenizer cannot re-encode, and the widening is bounded to keep the per-step host work small.

Parameters:

tokens – The sequence, prompt first
generatedCount – Tokens at the end of tokens that were generated
window – Receives the window

Public Static Attributes

static constexpr size_t kMaxWindowTokens = {16}#: Generated tokens the window may reach back over to find a word boundary.

class HostDraftTree#

A verify tree on the host at the grafting width: each slot’s draft nodes first, then the forced nodes the tree walk grafts under them. The rows are pinned, so a round uploads the batch with one copy per array.

A node’s mask row is its ancestors-or-self set and its depth the number of ancestors the tree holds, derived from its parent the way buildSpecVerificationTree derives them: a node whose parent is missing hangs off nothing. Nodes follow their parents.

Whole-batch views for the upload: [maxBatchSize, width], [maxBatchSize, width, width], [maxBatchSize].

inline Tensor const &idsTensor() const noexcept#

inline Tensor const &parentsTensor() const noexcept#

inline Tensor const &depthsTensor() const noexcept#

inline Tensor const &forcedTensor() const noexcept#

inline Tensor const &maskTensor() const noexcept#

inline Tensor const &validCountsTensor() const noexcept#

Public Functions

HostDraftTree(int32_t maxBatchSize, int32_t width)#

inline int32_t width() const noexcept#

void loadDraft( int32_t slot, int32_t const *tokens, int32_t const *parents, int32_t draftWidth )#: Load one slot’s draft tree, draftWidth nodes whose parents are lower indices or -1, and pad the rest.

void loadChain(int32_t slot, int32_t const *tokens, int32_t count)#: Load one slot’s draft chain, count nodes each hanging off the one before, and pad the rest.

void truncate(int32_t slot, int32_t count)#: Keep a slot’s first count nodes and turn the rest back into padding; count is between 1 and the node count.

int32_t insertChain( int32_t slot, int32_t attach, int32_t const *tokens, int32_t count, int8_t chainId )#

Graft count forced tokens under node attach: a child already carrying the next token is reused, the rest are inserted. Every node on the chain is flagged forced and labelled chainId.

Throws:: std::runtime_error – if the chain does not fit the width
Returns:: Nodes inserted

int32_t count(int32_t slot) const#

int32_t const *ids(int32_t slot) const#

int32_t const *parents(int32_t slot) const#

int32_t const *depths(int32_t slot) const#

int8_t const *forced(int32_t slot) const#

int8_t const *chainOf(int32_t slot) const#: Per node, the chain it is on, or -1.

int8_t const *mask(int32_t slot) const#: [width, width] ancestors-or-self rows.

int32_t forcedChildOf(int32_t slot, int32_t node) const#: The forced child of a node, or -1.

class JumpForwardCoordinator#

The decoder-independent part of a jump: the gate, the plan, and the token-list, logprob and statistics bookkeeping of applying forced tokens to a slot.

A plan whose rollbackCount is 1 replaces the slot’s sampled token, so its tokens are fed to the grammar from the state before that token; otherwise from the current state.

Public Functions

JumpForwardCoordinator( tokenizer::Tokenizer const &tokenizer, GuidedDecoder &decoder, int32_t maxJumpTokens )#

Both references are borrowed and must outlive the coordinator.

Parameters:: maxJumpTokens – Tokens one jump may add to a sequence

bool slotMayJump( DecodingInferenceContext const &context, int32_t slot ) const#: Whether a slot takes part in jumps this step: it asked for jump-forward, is unfinished, has left its reasoning block and has a live grammar. The same gate as the mask.

JumpForwardPlan plan( DecodingInferenceContext &context, int32_t slot, int32_t mutableTailCount, int32_t budget )#

Plan a jump for one slot.

Empty without touching the statistics unless the slot may jump (see slotMayJump); empty with the skip counted when nothing can be appended.

Parameters:

mutableTailCount – Trailing tokens of the slot that may be rolled back: not emitted, no KV
budget – Tokens the slot may still grow by

Applied applyPlan( DecodingInferenceContext &context, int32_t slot, JumpForwardPlan const &plan, int32_t sampledOutputSpaceId )#

Apply a plan to a slot: roll back the sampled token when the plan replaces it, feed the plan to the grammar once, append the tokens it took to the token list with synthetic logprob rows, and count.

A replacing run must spell more than the token it replaces; otherwise nothing is applied and the sampled token is re-accepted.

Parameters:: sampledOutputSpaceId – The slot’s sampled token in the output vocabulary
Returns:: What changed in the token list; both zero when nothing was applied

int32_t planDraftTree( DecodingInferenceContext &context, int32_t slot, HostDraftTree &tree, TreeWalkLimits const &limits )#

Graft forced chains into one slot’s draft tree before it is verified.

Walks the draft nodes in index order and feeds each node’s root path to the grammar; where the grammar then forces text, the text is planned behind the node and grafted under it, reusing draft children that already carry it. Skipped with their subtrees: a node the grammar refuses or terminates at, and a node beside a forced sibling, which acceptance never reaches. A node already on a chain is not planned again. A plan that would respell the node itself is refused and counted. The matcher ends where it started.

Parameters:: tree – The slot’s draft tree, loaded; receives the grafted nodes
Returns:: Chains grafted

int32_t planDraftChain( DecodingInferenceContext &context, int32_t slot, HostDraftTree &tree, TreeWalkLimits const &limits )#

Graft forced chains into a slot’s draft chain before it is verified; the result stays a chain.

Feeds the chain to the grammar from the root and plans the text the grammar forces after each node. A plan the draft after the node already spells flags those nodes forced and the walk goes on. The first plan the draft does not spell ends the walk. If it has at least as many tokens as the draft nodes it departs from that the grammar accepts in a row, it keeps the draft nodes that match it and replaces the rest of the draft with its remaining tokens; otherwise the draft stays whole and the plan is counted as draft_longer. A node on a chain is not planned again. The walk ends at a node the grammar refuses or terminates at. A plan that would respell the node itself is refused and counted. The matcher ends where it started.

Parameters:: tree – The slot’s draft chain, loaded; receives the grafted nodes
Returns:: Chains grafted

Public Static Functions

static void writeSyntheticLogprobRows( LogprobsSlot &logprobs, int32_t topK, int32_t const *tokens, int32_t count )#: Append one row per token: the token at log-probability 0, then negative-id fillers. A forced token has no logits row and the grammar gives it probability one.

static void writeSyntheticLogprobRow( LogprobsSlot &logprobs, int32_t topK, int32_t step, int32_t token )#: Rewrite row step as a forced token’s row.

struct Applied#

Public Members

int32_t rolledBack = {0}#: Sampled tokens removed from the token list: 0 or 1.

int32_t appended = {0}#: Plan tokens appended to the token list.

struct TreeWalkLimits#

Limits of one slot’s grafting round.

Public Members

int32_t maxPathLength#: Accepted path capacity of the round.

int32_t maxJumpTokens#: Tokens one chain may add.

int32_t remainingGeneration#: Tokens the slot may still generate.

class SpecJumpGrafter#

The grafting layer of a speculative decoder’s jump-forward round: grafts forced chains into the drafted verify input on the host, uploads the result, and accounts the forced tokens on the accepted paths.

A grafting round verifies jumpWidth rows per slot: the slot’s draft, the nodes grafted into it, then padding; a slot without a chain keeps its draft. Its accept step takes the forced flags and node counts held here, and a tree’s verify descriptors are prepared from the mask held here.

Per round: walk; after kGrafted, upload, verify and accept at the jump width, pass forcedPath to appendAcceptedTokens, collect logprobs, then account.

Public Types

enum class Walk#

Values:

enumerator kNoSlot#: No slot may jump; nothing was captured.

enumerator kNoChain#: Nothing was grafted; the guided decoder holds the draft as captured.

enumerator kGrafted#: Chains were grafted; upload them before verifying.

Public Functions

SpecJumpGrafter( tokenizer::Tokenizer const &tokenizer, GuidedDecoder &guidedDecoder, int32_t maxBatchSize, JumpGraftGeometry geometry, Limits const &limits )#: Both references are borrowed and must outlive the grafter.

Walk walk( DecodingInferenceContext &context, SpecVerifyInputs const &drafted )#

Capture the draft through the guided decoder, wait for it, and graft the chains the grammars plan for the slots that may jump. The GPU idles until the upload, since verification needs the grafted input.

Parameters:: drafted – The draft at width rows per slot without forced flags: parents for a tree, and node counts when the draft is ragged

void upload( Tensor &ids, Tensor *parents, Tensor *depths, cudaStream_t stream )#

After kGrafted: reshape the decoder’s verify views to jumpWidth and write the grafted input into them, and into the forced flags, node counts and, for a tree, the mask held here.

Parameters:

parents – The tree’s parent view; null for a chain
depths – The tree’s depth view; null for a chain

bool wide() const noexcept#: Whether the last walk grafted, so the round verifies jumpWidth rows per slot.

Tensor &mask()#: Device [activeBatch, jumpWidth, jumpWidth] int8 ancestors-or-self rows as uploaded. Tree geometry only.

Tensor const &mask() const#

Tensor const *forcedFlags() const noexcept#: Device [activeBatch, jumpWidth] int8 forced flags as uploaded; null unless wide.

Tensor const *validCounts() const noexcept#: Device [activeBatch] int32 node counts as uploaded, for either geometry; null unless wide.

decoder_utils::ForcedPathInfo forcedPath( Tensor const &deviceAcceptedRows )#

appendAcceptedTokens’ forced path for a wide round, which copies the accepted rows into pinned rows held here.

Parameters:: deviceAcceptedRows – [activeBatch, pathCapacity] verify rows the accept step walked

void account( DecodingInferenceContext &context, int32_t const *hostAcceptLengths, int32_t const *hostAcceptedIds, int32_t acceptCapacity )#

After a wide round’s appendAcceptedTokens and logprob collection: for each jump-forward slot, count the round, the forced tokens on its accepted path and the chains they came from, and rewrite each forced token’s logprob row, which holds the verify logits of the node whose prediction the token replaced.

Parameters:

hostAcceptLengths – [activeBatch] tokens appended, as appendAcceptedTokens left them
hostAcceptedIds – [activeBatch, acceptCapacity] accepted tokens

HostDraftTree const &hostTree() const noexcept#: The last walk’s grafted input.

struct Limits#

Public Members

int32_t width#: Draft rows per slot.

int32_t jumpWidth#: Verify rows per slot of a grafting round.

int32_t pathCapacity#: Accepted path capacity of a grafting round.

int32_t maxJumpTokens#: Tokens one chain may add.

class JumpGraftSession#

A speculative decoder’s share of jump-forward: the grafter of a deployment that can jump, whether the current round grafted, and the steps every grafting decoder performs the same way. Passive: a decoder calls each helper where its step allows, keeps the rest of its round as it is, and may use the grafter directly.

Public Types

using Walk = SpecJumpGrafter::Walk #

Public Functions

JumpGraftSession() = default#: Empty until initialize finds a deployment that can jump.

void initialize( tokenizer::Tokenizer const &tokenizer, GuidedDecoder &guidedDecoder, int32_t maxBatchSize, JumpGraftGeometry geometry, SpecDecodeConfig const &specConfig, int32_t plainWidth )#: Builds the grafter when specConfig.maxJumpTokens > 0; otherwise no round ever grafts. plainWidth is the draft rows per slot, which a decoder may size below specConfig.verifyTokenCount.

inline bool canJump() const noexcept#

bool active(DecodingInferenceContext const &context) const noexcept#: Whether this request set’s slots may jump: it asked for jump-forward on a guided deployment that can jump.

Walk walk( DecodingInferenceContext &context, SpecVerifyInputs const &drafted )#: Walks the drafted input when active and records whether the round grafted; kNoSlot when inactive.

inline void reset() noexcept#: Forgets the last walk, for a decoder whose round may skip walk.

inline bool wide() const noexcept#: Whether the current round verifies the grafted input at the jump width.

inline SpecJumpGrafter &grafter()#: The grafter; only on a deployment that can jump.

inline SpecJumpGrafter const &grafter() const#

void prepareGraftedTreeDescriptors( PipelineIO &io, Tensor const &kvCacheLengths, int32_t activeBatchSize, cudaStream_t stream ) const#: Writes a grafting round’s tree descriptors at the jump width from the grafter’s mask, over the ones the draft wrote.

inline int32_t verifyWidth(int32_t plain) const noexcept#: The current round’s values: the grafting round’s when wide, the plain ones otherwise.

inline int32_t acceptDepth(int32_t plain) const noexcept#

Tensor const *validCounts(Tensor const *plain) const noexcept#

Tensor const *forcedFlags() const noexcept#

std::optional<decoder_utils::ForcedPathInfo> forcedPath( Tensor const &deviceAcceptedRows )#: appendAcceptedTokens’ forced path of a wide round; empty otherwise.

void account( DecodingInferenceContext &context, int32_t const *hostAcceptLengths, int32_t const *hostAcceptedIds, int32_t acceptCapacity )#: After a wide round’s appendAcceptedTokens and logprob collection; nothing otherwise.

void addGraphTiers( std::vector<int32_t> &verifyWidths, std::vector<int32_t> *acceptLengths ) const#: For graph capture: adds the grafting round’s verify width past the plain one at the front, and its accepted path past the plain depths when acceptLengths is given.

struct JumpForwardSlotStats#

Per-sequence counters, reported once when the sequence ends.

Public Members

int32_t jumps = {0}#

int32_t forcedTokens = {0}#: Tokens appended without a forward pass.

int32_t rolledBackTokens = {0}#: Sampled tokens a re-encoding replaced.

int32_t chunkSteps = {0}#: Catch-up chunks or wide verify rounds this sequence took part in.

std::array<int32_t, static_cast<size_t>(JumpSkipReason::kCount)> skips = {}#

struct JumpForwardPlan#

Public Functions

inline bool empty() const noexcept#

Public Members

int32_t rollbackCount = {0}#: Window tokens to undo before appending the new ones.

std::vector<int32_t> fullVocabIds#: New tokens, full vocabulary.

std::vector<int32_t> outputSpaceIds#: The same tokens in the engine’s output vocabulary.

JumpSkipReason skipReason = {JumpSkipReason::kNone}#

struct Applied

Public Members

int32_t rolledBack = {0}: Sampled tokens removed from the token list: 0 or 1.

int32_t appended = {0}: Plan tokens appended to the token list.

struct TreeWalkLimits

Limits of one slot’s grafting round.

Public Members

int32_t maxPathLength: Accepted path capacity of the round.

int32_t maxJumpTokens: Tokens one chain may add.

int32_t remainingGeneration: Tokens the slot may still generate.

struct Limits

Public Members

int32_t width: Draft rows per slot.

int32_t jumpWidth: Verify rows per slot of a grafting round.

int32_t pathCapacity: Accepted path capacity of a grafting round.

int32_t maxJumpTokens: Tokens one chain may add.

char const *trt_edgellm::rt::jumpSkipReasonName( JumpSkipReason reason ) noexcept#

std::string trt_edgellm::rt::formatJumpForwardStats( JumpForwardSlotStats const &stats )#

std::string_view trt_edgellm::rt::completeUtf8Prefix( std::string_view text ) noexcept#: Longest prefix of text made of complete UTF-8 code points.