Prewarm native MTP batch shapes before serving requests

backends/exllamav3/model.py

@@ -22,6 +22,7 @@
     Tokenizer,
 )
 from exllamav3.cache import CacheLayer_quant
+from exllamav3.generator.sampler import GreedySampler
 from backends.exllamav3.grammar import ExLlamaV3Grammar
 
 from backends.exllamav3.sampler import ExllamaV3SamplerBuilder
@@ -130,6 +131,7 @@ class ExllamaV3Container:
     max_batch_size: Optional[int] = None
     draft_num_tokens: Optional[int] = None
     ngram_match_min: int = 0
+    mtp_warmup_batch_sizes: List[int] = []
 
     def __init__(self):
         # Mutable state must be created per instance. A class-level default
@@ -177,6 +179,7 @@ async def create(cls, model_directory: pathlib.Path, hf_model: HFModel, **kwargs
         self.config = Config.from_directory(str(model_directory.resolve()))
         self.model = Model.from_config(self.config)
         self.tokenizer = Tokenizer.from_config(self.config)
+        self.mtp_warmup_batch_sizes = list(unwrap(kwargs.get("mtp_warmup_batch_sizes"), []))
 
         # Prepare vision model if requested in config
         self.vision_model = None
@@ -675,6 +678,63 @@ def load_model_sync(self, progress_callback=None):
             if value:
                 yield value
 
+    async def _warmup_mtp_batches(self):
+        """Capture recurrent/MTP CUDA graph shapes before the API accepts traffic."""
+
+        if not self.mtp_warmup_batch_sizes:
+            return
+        if not self.use_draft_model or not self.draft_model.caps.get("mtp_draft", False):
+            xlogger.warning(
+                "mtp_warmup_batch_sizes was configured, but the active draft model "
+                "is not native MTP."
+            )
+            return
+
+        batch_sizes = sorted(set(self.mtp_warmup_batch_sizes))
+        invalid = [size for size in batch_sizes if size < 1 or size > self.max_batch_size]
+        if invalid:
+            raise ValueError(
+                "mtp_warmup_batch_sizes must be between 1 and max_batch_size "
+                f"({self.max_batch_size}); invalid values: {invalid}"
+            )
+
+        # More than three decode rounds are required to select and capture the fused GDN graph.
+        # A predictable prompt keeps the MTP window fully accepted while exercising draft Bx1 and
+        # target-verification Bx(draft+1) shapes. No generated text is logged or returned.
+        input_ids = self.tokenizer.encode(
+            "Warm the inference graph by continuing this sequence: 1 2 3 4 5 6 7 8",
+            encode_special_tokens=False,
+        )
+
+        async def drain(job: AsyncJob):
+            final = None
+            async for result in job:
+                final = result
+            return final
+
+        xlogger.info(f"Prewarming native MTP batch shapes: {batch_sizes}")
+        started = asyncio.get_running_loop().time()
+        for batch_size in batch_sizes:
+            jobs = [
+                AsyncJob(
+                    self.generator,
+                    input_ids=input_ids.clone(),
+                    max_new_tokens=32,
+                    sampler=GreedySampler(),
+                    seed=0,
+                    identifier=f"mtp-startup-warmup-b{batch_size}-{index}",
+                )
+                for index in range(batch_size)
+            ]
+            results = await asyncio.gather(*(drain(job) for job in jobs))
+            if self.generator.error is not None or any(result is None for result in results):
+                raise RuntimeError(f"Native MTP batch-{batch_size} startup warmup failed")
+            torch.cuda.synchronize()
+            xlogger.info(f"Prewarmed native MTP batch shape {batch_size}")
+
+        elapsed = asyncio.get_running_loop().time() - started
+        xlogger.info(f"Native MTP batch prewarm complete in {elapsed:.2f} seconds")
+
     async def create_generator(self):
         """Create and save a Exllama generator class."""
 
@@ -700,9 +760,10 @@ async def create_generator(self):
                 ngram_match_min=self.ngram_match_min,
             )
 
-            # Update the state of the container var
             if self.max_batch_size is None:
                 self.max_batch_size = self.generator.generator.max_batch_size
+
+            await self._warmup_mtp_batches()
         finally:
             # This means the generator is being recreated
             # The load lock is already released in the load function

common/config_models.py

@@ -293,6 +293,15 @@ class ModelConfig(BaseConfigModel):
         ),
         ge=1,
     )
+    mtp_warmup_batch_sizes: List[int] = Field(
+        default_factory=list,
+        description=(
+            "Precompile Qwen3.5 MTP and recurrent CUDA graph shapes during model startup.\n"
+            "Each listed batch size is warmed before the API becomes ready, preventing the first\n"
+            "request at a new concurrency level from paying JIT and graph-capture latency.\n"
+            "Example: [1, 2, 3, 4]. Only applies to native MTP drafting."
+        ),
+    )
     prompt_template: Optional[str] = Field(
         None,
         description=(

config_sample.yml

@@ -157,6 +157,12 @@ model:
   # can reduce it further to minimize VRAM overhead.
   max_batch_size:
 
+  # Precompile native MTP batch shapes during model startup (default: []).
+  # This prevents the first request at each configured concurrency level from
+  # paying the recurrent CUDA graph capture cost. Values must not exceed
+  # max_batch_size. Only applies to native MTP drafting.
+  mtp_warmup_batch_sizes: []
+
   # Set the prompt template for this model. (default: None)
   # If empty, attempts to look for the model's chat template.
   # If a model contains multiple templates in its tokenizer_config.json,