Prewarm native MTP batch shapes before serving requests - #438
Draft
Dampish0 wants to merge 1 commit into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Is your pull request related to a problem? Please describe.
Native Qwen3.5 MTP/recurrent CUDA graph shapes are captured lazily for each batch size. Warming a single request does not prepare later concurrency levels, so the first request at an unseen batch size can pay the full JIT and graph-capture cost.
In profiling, the first unseen batch-1 and batch-2 shapes generated at roughly 4 tokens/s. Once captured, the same model reached about 190 tokens/s at batch 1, 275-281 aggregate tokens/s at batch 2, and about 357 aggregate tokens/s at batch 4. The warmed scaling showed that steady-state kernels were healthy; the user-facing latency came from capture occurring after the API was already ready.
Why should this feature be added?
This adds an optional
model.mtp_warmup_batch_sizessetting. During generator creation, TabbyAPI runs short, internal greedy jobs for each configured batch size before startup completes. This prepares both native MTP draft and target-verification shapes before real traffic arrives.The implementation:
1..max_batch_size;config_sample.yml.Examples
With batch sizes 1-4 configured, startup graph capture took approximately 49 seconds in the tested 150k-context Q8-cache/Q6-MTP configuration. After readiness, first external requests at concurrency levels 1-4 no longer encountered the capture stall.
Additional context
Validation performed:
python -m ruff check backends/exllamav3/model.py common/config_models.pypython -m ruff format --check backends/exllamav3/model.py common/config_models.pypython -m py_compile backends/exllamav3/model.py common/config_models.py[1, 2, 3, 4]parsing checks