Guided Decoder#

class GuidedDecoder#

Grammar-constrained decoding backed by XGrammar.

Owns the tokenizer info, the compiled-grammar cache, and one grammar matcher per active batch slot.

Lifetime mirrors the runtime, but the per-slot matchers are request-local and must be reindexed by compactSlots on every batch eviction, in lockstep with the per-slot vectors in DecodingInferenceContext.

Public Functions

GuidedDecoder()#

~GuidedDecoder()#: Out-of-line: Impl is incomplete here.

GuidedDecoder(GuidedDecoder const&) = delete#

GuidedDecoder &operator=(GuidedDecoder const&) = delete#

void initialize( int32_t maxBatchSize, int32_t maxRowsPerSlot, int32_t outputVocabSize, int32_t fullVocabSize, Tokenizer const *tokenizer, Tensor const &reducedToFullVocabMap, cudaStream_t stream )#

Allocate reusable buffers and record the vocabulary geometry.

The tokenizer info is built on the first guided request instead: it costs one idToPiece call per output-vocabulary entry.

Parameters:

maxRowsPerSlot – Logits rows each slot can own in one step: 1 for vanilla decode, the speculative verify size otherwise. Sizes the bitmask buffers.
outputVocabSize – Width of the engine’s logits, i.e. the reduced vocabulary when the engine prunes, else the full vocabulary
fullVocabSize – Tokenizer vocabulary size; only used to size the full-to-output reverse map, which exists only when the engine prunes
tokenizer – Borrowed; must outlive this object
reducedToFullVocabMap – Output-space index to full token ID; pass an unallocated tensor when the engine does not prune
stream – Used for the one-off device-to-host copy of that map

bool prepareSlot( int32_t slot, GuidedDecodingParams const &params, std::string &failReason )#

Compile a grammar and install a fresh matcher for one slot.

Runs before any GPU work for the request, so a failure marks just this slot. This is the only input-driven step here, hence the only one that catches exceptions.

Parameters:: failReason – [out] Set only when the call returns false
Returns:: False when the guide could not be compiled; the slot is left unconstrained

void reset()#: Drop every matcher.

void compactSlots(std::vector<int32_t> const &batchMapping)#: Reindex per-slot state after a batch eviction. batchMapping[i] is the new index of old slot i, or -1 when it was evicted. Must run with the same mapping, and at the same point, as the DecodingInferenceContext vectors.

bool hasAnyGrammar() const noexcept#

bool hasGrammar(int32_t slot) const noexcept#

bool isTerminated(int32_t slot) const noexcept#: True once the matcher has accepted the stop token. Filling a mask past that point is a hard error in XGrammar, so callers must gate on this.

bool advance(int32_t slot, int32_t outputSpaceToken)#

Advance a slot’s grammar by one accepted token.

Parameters:: outputSpaceToken – Token index in the engine’s output vocabulary, i.e. captured before mapReducedVocabToFullVocab
Returns:: False when the grammar rejects the token, which the caller turns into kError

bool advanceCommitted( int32_t slot, int32_t const *committedFullSpace, int32_t count, int8_t &reasoningEnded )#

Advance a slot’s grammar over every token the step committed.

Speculative decoding commits a variable number of tokens per step, so the vanilla one-token advance does not apply. Tokens up to and including the reasoning-end marker are skipped: the marker is a separator, not constrained output, and consuming it is what flips reasoningEnded.

Parameters:

committedFullSpace – Tokens committed this step, in the full vocabulary
reasoningEnded – [inout] Latched to 1 by the marker; constrains from the next token on

Returns:

False when the grammar rejected a token, which the caller turns into kError

void fillMasks( int32_t activeBatchSize, int32_t rowsPerSlot, std::vector<int8_t> const &maskSuppressedPerSlot, std::vector<int32_t> &unsatisfiableSlots, cudaStream_t stream )#

Fill and upload one bitmask row per slot, at each slot’s current grammar state.

Vanilla decode only. Speculative verification needs a different mask per verify row and must call fillMasksForDraftTree: filling rowsPerSlot rows from one un-advanced matcher would just repeat the same mask, which is exactly what a grammar must not do.

Parameters:

maskSuppressedPerSlot – Slots to leave unconstrained this step
unsatisfiableSlots – [out] Slots whose mask came out all-zero, i.e. the grammar cannot be satisfied in this engine’s vocabulary

void captureDraftChains( Tensor const &draftChainIds, int32_t activeBatchSize, int32_t rowsPerSlot, cudaStream_t stream )#

Copy this step’s draft chains to the host and mark when they are there.

The matchers live on the host but the chains are produced on the device, so the mask fill has to wait for them. Recording an event instead of synchronizing the stream is what lets the caller enqueue the verify forward first and then wait, so the GPU is never idle while the host walks the grammar.

Parameters:: draftChainIds – [activeBatchSize, rowsPerSlot] device tensor of full-vocabulary IDs, node 0 being the token the previous step committed

void captureDraftTree( Tensor const &nodeTokenIds, Tensor const &parentIds, OptionalInputTensor const &validCounts, int32_t activeBatchSize, int32_t rowsPerSlot, cudaStream_t stream )#

Copy this step’s draft tree, topology included, to the host.

Same contract as captureDraftChains, plus the shape of the tree. A chain leaves the topology unset and the walk assumes node i descends from node i - 1.

Parameters:

nodeTokenIds – [activeBatchSize, rowsPerSlot] full-vocabulary IDs, node 0 the root
parentIds – [activeBatchSize, rowsPerSlot] parent node index; -1 at the root and in padding. The builder appends a node only once its parent is in the tree, so parentIds[i] < i always holds
validCounts – [activeBatchSize] nodes actually built; the rest are padding. Absent where the builder always fills the whole tree, as EAGLE’s does

int32_t const *hostDraftTokens() const#: Draft tokens captured by captureDraftChains or captureDraftTree, valid after waitForDraftTopology. Both geometries share this buffer.

int32_t const *hostDraftParentIds() const#: Tree topology captured by captureDraftTree; null after captureDraftChains.

int32_t const *hostDraftValidCounts() const#

void waitForDraftTopology()#: Block until the copy started by captureDraftChains or captureDraftTree has landed. Long is normal: the wait covers the drafting forwards, during which the verify forward is already running.

void fillMasksForDraftTree( int32_t activeBatchSize, int32_t rowsPerSlot, int32_t const *draftTokensFullSpace, int32_t const *parentIds, int32_t const *validCounts, std::vector<int8_t> const &slotSuppressed, std::vector<int8_t> const &reasoningEndedPerSlot, std::vector<int32_t> &unsatisfiableSlots, cudaStream_t stream )#

Fill and upload one bitmask row per draft-tree node, for every slot.

Walks each slot’s tree depth-first, feeding the grammar the token on the way down and rewinding it one step on the way back up, so that row i carries the mask for the state after the root-to-node-i path. The matcher ends where it started. Node 0 is the token the previous step already committed, so it is not fed to the matcher; the caller must have advanced past it.

A node whose token the grammar refuses prunes its subtree: those rows keep rowNeedsMask == 0 and are never applied. They are unreachable, because acceptance is path-based and the refused node’s parent row was masked. Sibling branches are unaffected, which is the whole reason the walk is a DFS and not a sweep over row order.

Parameters:

draftTokensFullSpace – [activeBatchSize, rowsPerSlot] node token IDs in the full vocabulary, as the verify tree carries them
parentIds – [activeBatchSize, rowsPerSlot] parent node index, or null for a chain, which is the degenerate tree with parentIds[i] == i - 1. A non-root node may report -1 when its parent missed the verify selection; it hangs off nothing, so it is unreachable and left unmasked
validCounts – [activeBatchSize] nodes actually built, or null for a chain, where every row is a node
slotSuppressed – Slots to leave entirely unconstrained this step
reasoningEndedPerSlot – Per-slot reasoning state at the start of the step. Read only: the flag is latched when the step’s tokens are committed, in advanceCommitted
unsatisfiableSlots – [out] Slots whose row 0 came out all-zero

void applyMask( Tensor &logits, int32_t activeBatchSize, int32_t rowsPerSlot, cudaStream_t stream )#: Apply the uploaded bitmask to logits, shaped [activeBatchSize, rowsPerSlot, vocab].

int64_t cacheSizeBytes() const#: Compiled-grammar cache footprint in bytes.

struct Impl#

Public Functions

inline int32_t toOutputSpace(int32_t fullId) const#: Full ID -> output-space index. Identity, plus a bounds check, when the engine does not prune; the caller therefore has a single code path.

inline bool isReasoningEnd(int32_t fullId) const#

inline bool rowHasAnyToken(int32_t rowId) const#: True when the row’s mask admits at least one token.

inline bool rowAllows(int32_t rowId, int32_t outputToken) const#

void ensureCompiler()#

xgrammar::CompiledGrammar compile(GuidedDecodingParams const &params)#

Public Members

Tokenizer const *tokenizer = {nullptr}#

int32_t maxBatchSize = {0}#

int32_t maxRows = {0}#: Bitmask rows allocated up front = maxBatchSize * maxRowsPerSlot.

int32_t outputVocabSize = {0}#

int32_t bitmaskSize = {0}#: int32 words per row, i.e. xgrammar::GetBitmaskSize(outputVocabSize)

Tensor outputToFullVocab#: Output-space index -> full tokenizer ID, [outputVocabSize]. Unset when the engine does not prune, in which case the mapping is the identity.

bool hasReducedVocab = {false}#

int32_t fullVocabSize = {0}#: Full tokenizer ID -> output-space index, or -1 for a token this engine cannot emit. Only built when the engine prunes; the speculative path needs this direction because the verify tree carries full-vocabulary IDs and the matchers live in output space.

std::vector<int32_t> fullToOutput#

std::vector<int32_t> reasoningEnd#: Reasoning-end markers in the full vocabulary; -1 entries never match.

std::vector<int8_t> warnedUnmappable#: One warning per request when a draft token has no output-space image, which means the engine’s reduced vocabulary is not a superset of the draft’s. Silently dropping every draft node instead would look like an unexplained acceptance-rate collapse.

std::optional<xgrammar::TokenizerInfo> tokenizerInfo#: Built on first use: one idToPiece call per output-vocabulary entry is too expensive to pay for runs that never use guided decoding.

std::optional<xgrammar::GrammarCompiler> compiler#

std::vector<std::optional<xgrammar::GrammarMatcher>> matchers#

Tensor deviceBitmask#: [maxRows, bitmaskSize] the mask the apply kernel reads

Tensor deviceRowNeedsMask#: [maxRows] per-row flag; the kernel leaves a zeroed row alone

Tensor hostBitmask#: [maxRows, bitmaskSize] where XGrammar fills; only flags are cleared per step

Tensor hostRowNeedsMask#: [maxRows] staging for the flags above

Tensor hostDraftTokens#

[batch, rowsPerSlot] node tokens, full vocabulary

Speculative decoding only, allocated when a slot owns more than one verify row. Holds this step’s draft geometry: the matchers live on the host, so the walk waits for this copy, and the event is blocking-sync because that wait spans the drafting forwards.

Tensor hostDraftParentIds#: Tree only: [batch, rowsPerSlot] parent index; -1 at root/padding.

Tensor hostDraftValidCounts#: Tree only: [batch] nodes actually built; later rows are padding.

bool draftIsTree = {false}#: Whether this step’s draft geometry is a tree.

bool draftHasValidCounts = {false}#: Whether the tree builder reports a node count at all.

cudaEvent_t draftCopyReady = {nullptr}#: Signals that the copy has landed.

bool draftCopyPending = {false}#: A copy is in flight; guards against a stale event.

std::vector<int32_t> firstChild#: [rowsPerSlot] lowest-numbered child, or -1 for a leaf

std::vector<int32_t> nextSibling#: [rowsPerSlot] next child of the same parent, or -1

class Tokenizer#

Tokenizer class for encoding and decoding text.

Provides tokenization functionality including pretokenization, encoding, and decoding. Supports loading from HuggingFace model directories.

Public Functions

Tokenizer() noexcept#

~Tokenizer() noexcept = default#

std::vector<Rank> encode( std::string const &text, bool addBos = false, bool addEos = false ) const#

Encode text to token IDs.

Parameters:

text – Input text to encode
addBos – Whether to add beginning-of-sequence token
addEos – Whether to add end-of-sequence token

Throws:

std::runtime_error – if tokenization encounters an error

Returns:

Vector of token IDs

std::string decode( std::vector<Rank> const &tokens, bool skipSpecialTokens = false ) const#

Decode token IDs back to text.

Parameters:

tokens – Vector of token IDs
skipSpecialTokens – Whether to skip special tokens in output

Returns:

Decoded text string (well-formed UTF-8; invalid byte sequences are replaced with U+FFFD via sanitizeUtf8Streaming/Flush)

std::string idToPiece(Rank token, bool skipSpecialTokens = true) const#

Single-token piece lookup for the streaming hot path.

Forwards to TokenEncoder::getRankToken with the skip-special policy applied. Returns an empty string when the token is a special token and skipSpecialTokens=true, or when rank is unknown (matches the silent-skip semantics of decode()).

Parameters:

token – Token ID (Rank).
skipSpecialTokens – Skip special tokens (BOS/EOS/etc).

Returns:

Raw piece bytes (possibly not independently valid UTF-8) or “”.

bool loadFromHF(std::filesystem::path const &modelDir)#

Load tokenizer from HuggingFace model directory.

Parameters:: modelDir – Path to the model directory containing tokenizer files
Returns:: true if directory exists, tokenizer.json is found and parsed successfully, pretokenizer and encoder are created successfully; false if directory doesn’t exist, tokenizer.json is missing/corrupt, or initialization fails

inline int getNumVocab() const noexcept#

Get total vocabulary size.

Returns:: Number of tokens in vocabulary

inline Rank getBosId() const noexcept#

Get beginning-of-sequence token ID.

Returns:: BOS token ID

inline Rank getEosId() const noexcept#

Get end-of-sequence token ID.

Returns:: EOS token ID

inline std::vector<Rank> getEosIds() const#

Get all configured end-of-sequence token IDs.

Returns:: EOS token IDs, including IDs injected from engine config when present.

inline bool isEosId(Rank token) const noexcept#

Check whether a token is one of the configured EOS IDs.

Parameters:: token – Token ID to check
Returns:: true if the token is an EOS token

inline bool isEosToken(Rank tokenId) const noexcept#

Check if a token is an end-of-sequence token.

Parameters:: tokenId – Token ID to check
Returns:: true if the token is any EOS token (primary or additional)

inline void setAdditionalEosIds(std::vector<Rank> const &ids)#

Set additional end-of-sequence token IDs.

Parameters:: ids – Vector of EOS token IDs from engine config.

inline Rank getPadId() const noexcept#

Get padding token ID.

Returns:: PAD token ID (returns EOS if PAD is not set)

inline Rank getUnkId() const noexcept#

Get unknown token ID.

Returns:: UNK token ID

inline Rank getTokenId(std::string const &token) const noexcept#

Get token ID by token string.

Parameters:: token – Token string to look up
Returns:: Token ID, or -1 if not found

bool isInitialized() const noexcept#

Check if tokenizer is properly initialized.

Returns:: true if initialized, false otherwise

inline TokenToRanks const &getSpecialTokensEncoder() const noexcept#: Special token string → id map (e.g. Alpamayo trajectory placeholder expansion).

bool trt_edgellm::rt::hasGuidedDecoding( LLMGenerationRequest const &request ) noexcept#

std::vector<int32_t> trt_edgellm::rt::reasoningStartMarkers( Tokenizer const &tokenizer )#: Token IDs that open a reasoning block, or -1 where the model has no such marker.

std::vector<int32_t> trt_edgellm::rt::reasoningEndMarkers( Tokenizer const &tokenizer )#: Token IDs that close a reasoning block, or -1 where the model has no such marker. Single source of truth: the grammar gate and the runtime’s thinking bookkeeping must agree on which token ends the block, or the constraint starts at the wrong moment.

bool trt_edgellm::rt::reasoningClosedInPrompt( std::vector<int32_t> const &promptTokens, std::vector<int32_t> const &startMarkers, std::vector<int32_t> const &endMarkers, bool thinkingEnabled ) noexcept#

Whether the prompt leaves the reasoning block closed, i.e. whether guided decoding may constrain from the very first generated token.

The most recent marker wins. A tokenizer carrying none of these markers has no reasoning phase, and is reported as closed.

Parameters:

startMarkers – Ids opening a reasoning block; entries below zero are absent from the tokenizer and ignored
endMarkers – Ids closing one, same convention
thinkingEnabled – Whether the request allows a reasoning phase. When the prompt carries no marker, a disabled reasoning phase is closed while an enabled phase may still be opened by the model.

void trt_edgellm::rt::applyGuidedDecodingMask( GuidedDecoder &decoder, DecodingInferenceContext &context, Tensor &logits, int32_t activeBatchSize, int32_t rowsPerSlot, cudaStream_t stream )#

Constrain one step’s logits to the grammar-legal tokens.

A slot whose grammar admits nothing ends with FinishReason::kError, keeping the text it generated so far. Slots still inside their thinking block are left unconstrained, so a JSON grammar cannot mask away the opening <think>.

Parameters:: rowsPerSlot – Logits rows owned by each slot; 1 for vanilla decode

void trt_edgellm::rt::advanceGuidedDecoding( GuidedDecoder &decoder, DecodingInferenceContext &context, int32_t const *outputSpaceIds, int32_t activeBatchSize )#

Advance each slot’s grammar by the token it just produced.

Must run before thinkingDone is updated for this token; that ordering is what keeps the </think> separator itself out of the grammar.

Parameters:: outputSpaceIds – Sampled indices in the engine’s output vocabulary, captured before the reduced-vocabulary remap

void trt_edgellm::rt::applyGuidedDecodingMaskForDraftTree( GuidedDecoder &decoder, DecodingInferenceContext &context, Tensor &logits, int32_t activeBatchSize, int32_t rowsPerSlot, cudaStream_t stream )#

Constrain one speculative verification step’s logits, one mask per verify row.

Waits for the draft tokens captured earlier in the step, walks each slot’s draft tree to build its masks, and applies them. The wait sits here rather than next to the copy so that the verify forward is already enqueued and the grammar walk overlaps it.

Parameters:: rowsPerSlot – Verify rows each slot owns, i.e. the deployment’s verify size

void trt_edgellm::rt::advanceGuidedDecodingForCommitted( GuidedDecoder &decoder, DecodingInferenceContext &context, int32_t const *hostAcceptedTokenIds, int32_t const *hostAcceptLengths, int32_t maxAcceptDepth, int32_t activeBatchSize )#

Advance each slot’s grammar over the tokens this speculative step committed.

Parameters:

hostAcceptedTokenIds – [activeBatchSize, maxAcceptDepth] full-vocabulary IDs
hostAcceptLengths – Per-slot count after appendAcceptedTokens trimmed it at EOS or max length; using the raw acceptance would advance past tokens that were discarded

bool trt_edgellm::rt::validateGuidedDecodingParams( GuidedDecodingParams const &params, std::string &failReason )#

Reject guides XGrammar would accept but not honour, and oversized ones.

Some JSON Schema keywords compile cleanly and then do nothing, so the output would violate the schema while the API claims it cannot. Compilation cannot catch that, which makes this the only line of defence. The blacklist is chosen by measurement against the pinned XGrammar version and must be re-checked when that pin moves.

Parameters:: failReason – [out] Set only when the call returns false
Returns:: False when the request must be rejected outright