Skip to content

Prewarm native MTP batch shapes before serving requests - #438

Draft
Dampish0 wants to merge 1 commit into
theroyallab:mainfrom
Dampish0:codex/mtp-batch-prewarm
Draft

Prewarm native MTP batch shapes before serving requests#438
Dampish0 wants to merge 1 commit into
theroyallab:mainfrom
Dampish0:codex/mtp-batch-prewarm

Conversation

@Dampish0

Copy link
Copy Markdown

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_sizes setting. 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:

  • is disabled by default;
  • only runs for native MTP drafting;
  • deduplicates and sorts configured batch sizes;
  • rejects values outside 1..max_batch_size;
  • uses a fixed synthetic prompt and does not log or return generated text;
  • documents the setting in config_sample.yml.

Examples

model:
  max_batch_size: 4
  mtp_warmup_batch_sizes: [1, 2, 3, 4]

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.py
  • python -m ruff format --check backends/exllamav3/model.py common/config_models.py
  • python -m py_compile backends/exllamav3/model.py common/config_models.py
  • Pydantic config default and [1, 2, 3, 4] parsing checks
  • Runtime checks for non-MTP no-op and out-of-range rejection
  • Deterministic generation output matched the unpatched build by SHA-256
  • Live concurrency 1-4 and internal API health checks completed without CUDA or runtime errors

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant