Merge branch 'main' into fix-mtp-dummy-run-assertion

[Bugfix] Fix assertion error in _dummy_run for MTP speculative decoding
Fix an assertion error during model warmup when using MTP speculative decoding with parallel drafting. The issue occurred because the compile range is extended for speculative decoding to accommodate drafter batches, but the assertion in _dummy_run wasn't updated to match. Root cause: PR #32887 added compile range extension in _set_compile_ranges for speculative decoding. This causes warmup sizes to exceed max_num_batched_tokens, triggering the assertion in _dummy_run. Fix: Extend the assertion bound in _dummy_run to match the extended compile range when parallel drafting is enabled. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
2026-02-23 16:54:20 -08:00 · 2026-02-13 17:51:39 -05:00
1 changed files with 18 additions and 3 deletions
@@ -4689,7 +4689,21 @@ class GPUModelRunner(
        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
        # for dummy run with LoRA so that the num_reqs collectively
        # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        # For speculative decoding, the compile range is extended:
+        # - Sequential: + 1 * max_num_seqs (one draft token per iteration)
+        # - Parallel draft: + num_speculative_tokens * max_num_seqs
+        max_dummy_run_tokens = self.scheduler_config.max_num_batched_tokens
+        if self.speculative_config is not None and (
+            self.speculative_config.uses_draft_model()
+            or self.speculative_config.use_eagle()
+        ):
+            multiplier = (
+                self.speculative_config.num_speculative_tokens
+                if self.speculative_config.parallel_drafting
+                else 1
+            )
+            max_dummy_run_tokens += multiplier * self.scheduler_config.max_num_seqs
+        assert num_tokens <= max_dummy_run_tokens
        max_num_reqs = self.scheduler_config.max_num_seqs
        if create_mixed_batch:
            assert not uniform_decode
@@ -4822,8 +4836,9 @@ class GPUModelRunner(
            remove_lora,
            num_active_loras,
        ):
-            # Make sure padding doesn't exceed max_num_tokens
-            assert num_tokens_padded <= self.max_num_tokens
+            # Make sure padding doesn't exceed max_num_tokens (extended for
+            # parallel drafting warmup)
+            assert num_tokens_padded <= max_dummy_run_tokens
            model_kwargs = self._init_model_kwargs()
            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)