Repository navigation
fast-forward tokens in logits post processor #2365
Description
I've been working on an OpenAI-compatible REST server, utilizing TensorRT-LLM but not Triton, similar to openai_server.py but in Rust and generally production-ready. Its main feature is support for constrained decoding, using Guidance or more specifically the Rust low-level Guidance library. This allows for enforcing JSON schemas on the output (similar to OpenAI's "Structured Output" feature, though much more general and with no initial latency), as well as arbitrary context-free grammars.
I've been utilizing the batched logits post processor - I wrote a cuda kernel that applies boolean mask to logits (ie., sets logits to -inf where mask is false). Because I needed to support different temperature per token, the kernel also scales logits. There is an additional hack in there to support temperature==0.0 (greedy sampling). It would be nice it logits post processor allowed for changing sampling parameters.
However, the main feature that is missing and that would help performance is the ability to append more than one token to the sequence in the output phase - fast-forward tokens. This can be thought of an "speculation that is always right".
An example, where fast forward tokens are useful is generating data adhering to a certain JSON schema. The logit processor forces {"name":" to be generated, then the model generates John", the processor forces ,\n"age":, model generates 42, and so on. Another example is chain-of-thought reasoning, where after the model generated a sentence, the controller forces more instructions for the model, the model generates more text, and so on. If used, these greatly speed up generation process.
The sglang implements this though it was first done in Guidance.
One way to allow this functionality is to extend logit post-processor with the following callback that for each sequence is given the sampled token and returns a vector of tokens (typically it would be just the sampled token in the vector but could be something else, not necessarily starting with the sampled token).
using PostSamplingCallback = std::function< std::vector<std::vector<std::vector<TokenIdType>>> (std::vector<IdType> const &seq_ids,
std::vector<std::vector<TokenIdType>> const &sampled_tokens,
std::vector<std::optional<IdType>> const &user_seq_ids)>;The callback would take request ids, sampled tokens - a vector of num_requests of vectors of beam size of sampled tokens, and user req ids. It would return a vector of num_requests, of vectors of beam size of vectors of tokens to append to each beam.
For example:
- model generates
John, post-processor returns {John} - model generates
", post-processor returns {",,\n,",age,":} (assuming such tokens exists - note how"was replaced by",)
The signature above looks somewhat complicated due to beams and requests, for a single beam, single request it's essentially:
vector<Token> post_sampling(Token t) {
return {t}; // default
}
Activity
MartinMarciniszyn commentedon Oct 23, 2024
Hi @mmoskal, thank you for submitting this suggestion. We are in the process of improving the capabilities for generating structured output in TRT-LLM. We shall take your proposal into consideration. I will get back to you once we have a clearer design. CC @juney-nvidia , @tqchen
mmoskal commentedon Oct 23, 2024
Thank you @MartinMarciniszyn ! If you are interested in integrating the llguidance library directly into TRT-LLM let me know. The library has a C interface that allows for regex, JSON schema, and Guidance constraints.
lupickup commentedon Nov 14, 2024
@MartinMarciniszyn Adding to the interest in enabling fast-forward tokens. Azure AI has been collaborating with @mmoskal and we're very interested in combining the best-in-class inference performance of trt-llm with the best-in-class constrained decoding features guidance and it's ecosystem of libraries offers!
mmoskal commentedon Jan 22, 2025
We have just released a large JSON Schema benchmark and a paper. The paper in particular has data on how much token can be typically fast-forwarded (see difference between LLM-only and Guidance in the plot below)
There is also isolated mask-generation benchmarks - comparing LLGuidance, Outlines, XGrammar and llama.cpp grammars.
MartinMarciniszyn commentedon Jan 22, 2025
Thank you for sharing the benchmark.
Note that XGrammer has been integrated into TRT-LLM in the meanwhile. Here is an example.
We are in the process of evaluating how to implement the ability of appending more than one token to the sequence. Instead of extending the interface of
LogitsPostprocesser, one could also follow a target-draft model approach or a modified version of lookahead decoding. We will share updates here.mmoskal commentedon Jan 22, 2025
I love the integration of constrained decoding directly into TRT-LLM!
But, I would like to see LLGuidance integrated as well, as it's significantly faster than XGrammar and has much more JSON feature coverage (eg.
allOf,min/maxItems,minimum/maximum,multipleOf,additionalProperties, sibling keys; note that implementing some of these makes the grammar more complex, which is why "XGrammar.cpp" which is using llama.cpp JSON->Grammar converter is so much slower than regular XGrammar).See MaskBench for details.
karljang commentedon Oct 21, 2025
This old post:Fast JSON Decoding for Local LLMs with Compressed Finite State Machine recently caught my attention, and it seems this ticket was exploring a similar approach.
karljang commentedon Oct 21, 2025
@mmoskal ,
Copying my answer in #8504 here.
Currently, TRT-LLM doesn’t support jump-forward or fast-forward decoding.
However, for low-latency scenarios, guided decoding combined with speculative decoding can offer similar speedups.
Would you be able to give that approach a try?
Thank you!
github-actions commentedon Nov 5, 2025
Issue has not received an update in over 14 days. Adding stale label.
karljang commentedon Nov 14, 2025
Closing this issue as stale.
If you're still exploring speedups, consider trying guided decoding combined with speculative decoding, they may offer similar performance improvements.
Thank you!