[Async][Spec Decoding] Zero-bubble async scheduling + spec decoding

vllm/config/vllm.py

@@ -634,6 +634,30 @@ def __post_init__(self):
             else:
                 self.parallel_config.disable_nccl_for_dp_synchronization = False
 
+        if (
+            self.speculative_config is not None
+            and self.scheduler_config.async_scheduling
+        ):
+            if (
+                self.model_config is not None
+                and not self.model_config.disable_cascade_attn
+            ):
+                logger.warning_once(
+                    "Disabling cascade attention (not yet compatible with "
+                    "async speculative decoding).",
+                    scope="local",
+                )
+                self.model_config.disable_cascade_attn = True
+
+            if self.cache_config.mamba_cache_mode == "align":
+                logger.warning_once(
+                    "Mamba cache mode 'align' is not compatible with async speculative "
+                    "decoding, as it relies on cpu num_computed_tokens.",
+                    scope="local",
+                )
+                self.cache_config.mamba_cache_mode = "none"
+                self.cache_config.enable_prefix_caching = False
+
         from vllm.platforms import current_platform
 
         if (

vllm/v1/executor/multiproc_executor.py

@@ -428,7 +428,7 @@ def check_health(self) -> None:
     def max_concurrent_batches(self) -> int:
         # PP requires PP-size concurrent batches to fill the pipeline.
         pp_size = self.parallel_config.pipeline_parallel_size
-        return 2 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size
+        return 3 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size
 
     def _get_output_rank(self) -> int:
         # Only returns ModelRunnerOutput from TP rank=0 and PP rank=-1

vllm/v1/spec_decode/eagle.py

@@ -648,11 +648,18 @@ def prepare_next_token_ids_padded(
         """
         # Precompute get_token_id for when there is no valid next token
         num_reqs = gpu_input_batch.num_reqs
+        num_computed_tokens_cpu = gpu_input_batch.num_computed_tokens_cpu_tensor[
+            :num_reqs
+        ]
+        query_lens_cpu = (
+            common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1]
+            - common_attn_metadata.query_start_loc_cpu[:num_reqs]
+        )
+        seq_lens_cpu = num_computed_tokens_cpu + query_lens_cpu
+        seq_lens_list = seq_lens_cpu.tolist()
         self.backup_next_token_ids.np[:num_reqs] = np.array(
             [
-                requests[gpu_input_batch.req_ids[i]].get_token_id(
-                    common_attn_metadata.seq_lens_cpu[i].item()
-                )
+                requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])
                 for i in range(num_reqs)
             ],
             dtype=np.int32,
@@ -737,7 +744,7 @@ def prepare_inputs_padded(
             num_reqs=common_attn_metadata.num_reqs,
             num_actual_tokens=total_num_tokens,
             max_query_len=new_query_len_per_req.max().item(),
-            max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(),
+            max_seq_len=common_attn_metadata.max_seq_len,
             block_table_tensor=common_attn_metadata.block_table_tensor,
             slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],
             causal=True,

vllm/v1/worker/block_table.py

@@ -131,7 +131,7 @@ def swap_row(self, src: int, tgt: int) -> None:
         self.block_table.np[src_tgt] = self.block_table.np[tgt_src]
 
     def compute_slot_mapping(
-        self, req_indices: np.ndarray, positions: np.ndarray
+        self, req_indices: torch.Tensor, positions: torch.Tensor
     ) -> None:
         # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
         # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
@@ -154,7 +154,7 @@ def compute_slot_mapping(
                 + positions // virtual_block_size
             )
 
-            block_numbers = self.block_table.np.ravel()[block_table_indices]
+            block_numbers = self.block_table.gpu.view(-1)[block_table_indices]
             # Use virtual_block_size for mask calculation, which marks local
             # tokens.
             virtual_block_offsets = positions % virtual_block_size
@@ -174,20 +174,20 @@ def compute_slot_mapping(
             # Calculate slot_mapping
             slot_mapping = block_numbers * self.block_size + block_offsets
             # Write final slots, use -1 for not-local
-            self.slot_mapping.np[: req_indices.shape[0]] = np.where(
+            self.slot_mapping.gpu[: req_indices.shape[0]] = torch.where(
                 mask, slot_mapping, -1
             )
         else:
             block_table_indices = (
                 req_indices * self.max_num_blocks_per_req + positions // self.block_size
             )
 
-            block_numbers = self.block_table.np.ravel()[block_table_indices]
+            block_numbers = self.block_table.gpu.view(-1)[block_table_indices]
             block_offsets = positions % self.block_size
-            np.add(
+            torch.add(
                 block_numbers * self.block_size,
                 block_offsets,
-                out=self.slot_mapping.np[: req_indices.shape[0]],
+                out=self.slot_mapping.gpu[: req_indices.shape[0]],
             )
 
     def commit_block_table(self, num_reqs: int) -> None:
@@ -320,7 +320,7 @@ def swap_row(self, src: int, tgt: int) -> None:
             block_table.swap_row(src, tgt)
 
     def compute_slot_mapping(
-        self, req_indices: np.ndarray, positions: np.ndarray
+        self, req_indices: torch.Tensor, positions: torch.Tensor
     ) -> None:
         for block_table in self.block_tables:
             block_table.compute_slot_mapping(req_indices, positions)

vllm/v1/worker/gpu_input_batch.py

@@ -207,7 +207,7 @@ def __init__(
 
         # Speculative decoding
         self.num_accepted_tokens_cpu_tensor = torch.ones(
-            (max_num_reqs,), dtype=torch.int64, device="cpu", pin_memory=pin_memory
+            (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
         )
         self.num_accepted_tokens_cpu = self.num_accepted_tokens_cpu_tensor.numpy()
 
@@ -938,6 +938,9 @@ def update_async_output_token_ids(self) -> None:
         In async scheduling case, update output_token_ids in sampling metadata
         from prior steps sampled token ids once they've finished copying to CPU.
         This is called right before they are needed by the logits processors.
+
+        For async spec decode mode: extends output_token_ids with actual tokens.
+        For non-async spec decode mode: replaces -1 placeholders with actual tokens.
         """
         output_token_ids = self.sampling_metadata.output_token_ids
         if self.sampled_token_ids_cpu is None or not output_token_ids:
@@ -951,10 +954,6 @@ def update_async_output_token_ids(self) -> None:
             if prev_index is None:
                 continue
             req_output_token_ids = output_token_ids[index]
-            if not req_output_token_ids or req_output_token_ids[-1] != -1:
-                # Final output id is not a placeholder, some tokens must have
-                # been discarded after a kv-load failure.
-                continue
             if sampled_token_ids is None:
                 assert self.async_copy_ready_event is not None
                 self.async_copy_ready_event.synchronize()
@@ -964,14 +963,28 @@ def update_async_output_token_ids(self) -> None:
             if not new_ids:
                 continue
             num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
-            # Also account for case where there may be a smaller number of
-            # output placeholders (tokens can be discarded after a kv-load failure).
-            first_placeholder = req_output_token_ids.index(-1)
-            num_placeholders = len(req_output_token_ids) - first_placeholder
-            num_to_replace = min(num_sampled_ids, num_placeholders)
-            del new_ids[num_to_replace:]
-            end_index = first_placeholder + num_to_replace
-            req_output_token_ids[first_placeholder:end_index] = new_ids
+            if num_sampled_ids == 0:
+                continue
+
+            # Check if there are placeholders to replace.
+            has_placeholders = req_output_token_ids and req_output_token_ids[-1] == -1
+
+            if has_placeholders:
+                # Replace placeholder token id(s) with actual sampled id(s).
+                # Also account for case where there may be a smaller number of
+                # output placeholders (tokens can be discarded after kv-load failure).
+                first_placeholder = req_output_token_ids.index(-1)
+                num_placeholders = len(req_output_token_ids) - first_placeholder
+                num_to_replace = min(num_sampled_ids, num_placeholders)
+                end_index = first_placeholder + num_to_replace
+                req_output_token_ids[first_placeholder:end_index] = new_ids[
+                    :num_to_replace
+                ]
+            else:
+                # No placeholders - extend with actual sampled tokens.
+                # This is used in async spec decode mode where we don't add
+                # placeholders upfront.
+                req_output_token_ids.extend(new_ids[:num_sampled_ids])
 
     def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None:
         """