Server: openai-style lookup decoding #12127
Open
+228
−5
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This is a proposal for lookup decoding in the server example, similar to the "predicted outputs" feature in OpenAI API. The idea is that the user includes a string in the request that will be used for lookup decoding. This gives major speedups to code refactoring or text editing tasks where most of the output will be the same. See https://platform.openai.com/docs/guides/predicted-outputs#code-refactoring-example for examples.
I'm aware of the open PR #6828, and maybe that could be used here as well, but for now I implemented a custom lookup algorithm for this. I think the main difference is that in this OAI-style API, the user is in control of when to use lookup, so if they choose not to, there will be no performance impact in any direction.
I implemented the feature based on the existing speculative decoding code, and to keep it simple I made it so that lookup is not used if speculative decoding is enabled. It uses the
n_max
parameter from speculative, but for best results this should be much higher than when using a draft model.The algorithm itself is probably not optimal, but gives good results anyway:
For the adaptive window size I just set it to double the window size every time, starting from a window size of 16. I'm not sure what actually governs the performance here, but I get pretty nice speedups with these values. If there's some more clever algorithm already in the lookup example maybe that could be substituted here.
One limitation to OpenAI compatibility: rejected token count is not included in the response (not sure what this would mean exactly, but I figure accepted tokens is more important info anyway)