@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fafe76b4af

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 04:30 UTC

[Async][Spec Decoding] Zero-bubble async scheduling + spec decoding (#32951)

Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Co-authored-by: zhrrr <43847754+izhuhaoran@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

Matthew Bonanni committed Mar 23, 2026 at 15:37 UTC fafe76b4af175d48b9b69204c995ee167b46aca1

9 files changed +488 -209

tests/v1/spec_decode/test_eagle.py

+1 -1

@@ -177,7 +177,7 @@ def test_prepare_next_token_ids():

177

178 next_token_ids_from_padded, valid_sampled_tokens_count = (

179 proposer.prepare_next_token_ids_padded(

180 - common_attn_metadata,

180 + common_attn_metadata.seq_lens_cpu,

181 sampled_token_ids_tensor,

182 mock_requests,

183 mock_input_batch,

tests/v1/spec_decode/test_extract_hidden_states.py

+1 -1

@@ -187,7 +187,7 @@ def test_prepare_next_token_ids_padded():

187 )

188

189 next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(

190 - common_attn_metadata,

190 + common_attn_metadata.seq_lens_cpu,

191 sampled_token_ids,

192 mock_requests,

193 mock_input_batch,

vllm/config/vllm.py

+13

@@ -766,6 +766,19 @@ class VllmConfig: # type: ignore[misc]

766 else:

767 self.parallel_config.disable_nccl_for_dp_synchronization = False

768

769 + if (

770 + self.speculative_config is not None

771 + and self.scheduler_config.async_scheduling

772 + and self.model_config is not None

773 + and not self.model_config.disable_cascade_attn

774 + ):

775 + logger.warning_once(

776 + "Disabling cascade attention (not yet compatible with "

777 + "async speculative decoding).",

778 + scope="local",

779 + )

780 + self.model_config.disable_cascade_attn = True

781 +

782 if (

783 self.model_config is not None

784 and self.model_config.multimodal_config is not None

vllm/v1/spec_decode/eagle.py

+4 -8

@@ -71,7 +71,6 @@ class SpecDecodeBaseProposer:

71 self.method = self.speculative_config.method

72 self.pass_hidden_states_to_model = pass_hidden_states_to_model

74 - self.runner = runner

74 self.device = device

75 self.dtype = vllm_config.model_config.dtype

76 self.max_model_len = vllm_config.model_config.max_model_len

@@ -424,8 +423,6 @@ class SpecDecodeBaseProposer:

423 )

424 )

425

427 - assert self.runner is not None

428 -

426 per_layer_attn_metadata: dict[str, object] = {}

427 for attn_group in self.draft_attn_groups:

428 attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

@@ -821,7 +818,7 @@ class SpecDecodeBaseProposer:

818

819 def prepare_next_token_ids_padded(

820 self,

824 - common_attn_metadata: CommonAttentionMetadata,

821 + seq_lens_cpu: torch.Tensor,

822 sampled_token_ids: torch.Tensor,

823 requests: dict[str, CachedRequestState],

824 gpu_input_batch: InputBatch,

@@ -836,11 +833,10 @@ class SpecDecodeBaseProposer:

833 """

834 # Precompute get_token_id for when there is no valid next token

835 num_reqs = gpu_input_batch.num_reqs

836 + seq_lens_list = seq_lens_cpu[:num_reqs].tolist()

837 self.backup_next_token_ids.np[:num_reqs] = np.array(

838 [

841 - requests[gpu_input_batch.req_ids[i]].get_token_id(

842 - common_attn_metadata.seq_lens_cpu[i].item()

843 - )

839 + requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])

840 for i in range(num_reqs)

841 ],

842 dtype=np.int32,

@@ -925,7 +921,7 @@ class SpecDecodeBaseProposer:

921 num_reqs=common_attn_metadata.num_reqs,

922 num_actual_tokens=total_num_tokens,

923 max_query_len=new_query_len_per_req.max().item(),

928 - max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(),

924 + max_seq_len=common_attn_metadata.max_seq_len,

925 block_table_tensor=common_attn_metadata.block_table_tensor,

926 slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],

927 causal=True,

vllm/v1/spec_decode/extract_hidden_states.py

+3 -4

@@ -286,7 +286,7 @@ class ExtractHiddenStatesProposer:

286

287 def prepare_next_token_ids_padded(

288 self,

289 - common_attn_metadata: CommonAttentionMetadata,

289 + seq_lens: torch.Tensor,

290 sampled_token_ids: torch.Tensor,

291 requests: dict[str, CachedRequestState],

292 gpu_input_batch: InputBatch,

@@ -303,11 +303,10 @@ class ExtractHiddenStatesProposer:

303 device = sampled_token_ids.device

304

305 # Compute backup tokens for discarded / invalid requests

306 + seq_lens_list = seq_lens[:num_reqs].tolist()

307 backup_tokens_gpu = torch.tensor(

308 [

308 - requests[gpu_input_batch.req_ids[i]].get_token_id(

309 - common_attn_metadata.seq_lens_cpu[i].item()

310 - )

309 + requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])

310 for i in range(num_reqs)

311 ],

312 dtype=torch.int32,

vllm/v1/spec_decode/utils.py

+34

@@ -3,6 +3,7 @@

3 import torch

5 from vllm.config import VllmConfig, replace

6 +from vllm.platforms import current_platform

7 from vllm.triton_utils import tl, triton

8 from vllm.v1.attention.backends.utils import (

9 CommonAttentionMetadata,

@@ -463,3 +464,36 @@ def copy_and_expand_eagle_inputs_kernel(

464 out_idx,

465 mask=is_new_token_region & in_bounds,

466 )

467 +

468 +

469 +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

470 +def update_num_computed_tokens_for_batch_change(

471 + num_computed_tokens: torch.Tensor,

472 + num_accepted_tokens: torch.Tensor,

473 + prev_positions: torch.Tensor,

474 + valid_sampled_token_count: torch.Tensor,

475 + prev_num_draft_tokens: torch.Tensor,

476 + cpu_num_computed_tokens: torch.Tensor,

477 +) -> None:

478 + """Correct num_computed_tokens for async spec decode drift.

479 +

480 + Requests that had drafts: corrected = prev_gpu + valid_count.

481 + New requests or non-draft (e.g. prefills): use CPU value directly.

482 + """

483 + # Clamp because prev_positions can be -1 for new requests

484 + gather_indices = prev_positions.clamp(min=0)

485 +

486 + valid_counts = valid_sampled_token_count[gather_indices]

487 + prev_computed = num_computed_tokens[gather_indices]

488 + prev_drafts = prev_num_draft_tokens[gather_indices]

489 +

490 + participating = (prev_positions >= 0) & (prev_drafts > 0)

491 + corrected = prev_computed + valid_counts.int()

492 +

493 + n = prev_positions.shape[0]

494 + num_computed_tokens[:n].copy_(

495 + torch.where(participating, corrected, cpu_num_computed_tokens)

496 + )

497 + num_accepted_tokens.copy_(

498 + torch.where(participating, valid_counts, num_accepted_tokens)

499 + )

vllm/v1/worker/block_table.py

+85 -64

@@ -6,7 +6,9 @@ import torch

7 from vllm.distributed import get_dcp_group, get_pcp_group

8 from vllm.logger import init_logger

9 +from vllm.triton_utils import tl, triton

10 from vllm.utils.math_utils import cdiv

11 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID

12 from vllm.v1.utils import CpuGpuBuffer

13 from vllm.v1.worker.cp_utils import get_total_cp_world_size

@@ -131,71 +133,33 @@ class BlockTable:

133 self.block_table.np[src_tgt] = self.block_table.np[tgt_src]

134

135 def compute_slot_mapping(

134 - self, req_indices: np.ndarray, positions: np.ndarray

136 + self,

137 + num_reqs: int,

138 + query_start_loc: torch.Tensor,

139 + positions: torch.Tensor,

140 ) -> None:

136 - # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

137 - # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]

138 - # where K is the max_num_blocks_per_req and the block size is 2.

139 - # NOTE(woosuk): We can't simply use `token_indices // block_size`

140 - # here because M (max_model_len) is not necessarily divisible by

141 - # block_size.

141 + num_tokens = positions.shape[0]

142 total_cp_world_size = self.pcp_world_size * self.dcp_world_size

143 total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank

144 - if total_cp_world_size > 1:

145 - # Note(hc): The DCP implement store kvcache with an interleave

146 - # style, the kvcache for the token whose token_idx is i is

147 - # always stored on the GPU whose dcp_rank equals i % cp_world_size:

148 -

149 - # Use a "virtual block" which equals to world_size * block_size

150 - # for block_table_indices calculation.

151 - virtual_block_size = self.block_size * total_cp_world_size

152 - block_table_indices = (

153 - req_indices * self.max_num_blocks_per_req

154 - + positions // virtual_block_size

155 - )

156 -

157 - block_numbers = self.block_table.np.ravel()[block_table_indices]

158 - # Use virtual_block_size for mask calculation, which marks local

159 - # tokens.

160 - virtual_block_offsets = positions % virtual_block_size

161 - mask = (

162 - virtual_block_offsets

163 - // self.cp_kv_cache_interleave_size

164 - % total_cp_world_size

165 - == total_cp_rank

166 - )

167 - # Calculate local block_offsets

168 - block_offsets = (

169 - virtual_block_offsets

170 - // (total_cp_world_size * self.cp_kv_cache_interleave_size)

171 - * self.cp_kv_cache_interleave_size

172 - + virtual_block_offsets % self.cp_kv_cache_interleave_size

173 - )

174 - # Calculate slot_mapping

175 - slot_mapping = block_numbers * self.block_size + block_offsets

176 - # Write final slots, use -1 for not-local

177 - self.slot_mapping.np[: req_indices.shape[0]] = np.where(

178 - mask, slot_mapping, -1

179 - )

180 - else:

181 - block_table_indices = (

182 - req_indices * self.max_num_blocks_per_req + positions // self.block_size

183 - )

184 -

185 - block_numbers = self.block_table.np.ravel()[block_table_indices]

186 - block_offsets = positions % self.block_size

187 - np.add(

188 - block_numbers * self.block_size,

189 - block_offsets,

190 - out=self.slot_mapping.np[: req_indices.shape[0]],

191 - )

144 + _compute_slot_mapping_kernel[(num_reqs + 1,)](

145 + num_tokens,

146 + self.max_num_batched_tokens,

147 + query_start_loc,

148 + positions,

149 + self.block_table.gpu,

150 + self.block_table.gpu.stride(0),

151 + self.block_size,

152 + self.slot_mapping.gpu,

153 + TOTAL_CP_WORLD_SIZE=total_cp_world_size,

154 + TOTAL_CP_RANK=total_cp_rank,

155 + CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,

156 + PAD_ID=PAD_SLOT_ID,

157 + BLOCK_SIZE=1024,

158 + )

159

160 def commit_block_table(self, num_reqs: int) -> None:

161 self.block_table.copy_to_gpu(num_reqs)

162

196 - def commit_slot_mapping(self, num_tokens: int) -> None:

197 - self.slot_mapping.copy_to_gpu(num_tokens)

198 -

163 def clear(self) -> None:

164 self.block_table.gpu.fill_(0)

165 self.block_table.cpu.fill_(0)

@@ -320,19 +284,18 @@ class MultiGroupBlockTable:

284 block_table.swap_row(src, tgt)

285

286 def compute_slot_mapping(

323 - self, req_indices: np.ndarray, positions: np.ndarray

287 + self,

288 + num_reqs: int,

289 + query_start_loc: torch.Tensor,

290 + positions: torch.Tensor,

291 ) -> None:

292 for block_table in self.block_tables:

326 - block_table.compute_slot_mapping(req_indices, positions)

293 + block_table.compute_slot_mapping(num_reqs, query_start_loc, positions)

294

295 def commit_block_table(self, num_reqs: int) -> None:

296 for block_table in self.block_tables:

297 block_table.commit_block_table(num_reqs)

298

332 - def commit_slot_mapping(self, num_tokens: int) -> None:

333 - for block_table in self.block_tables:

334 - block_table.commit_slot_mapping(num_tokens)

335 -

299 def clear(self) -> None:

300 for block_table in self.block_tables:

301 block_table.clear()

@@ -340,3 +303,61 @@ class MultiGroupBlockTable:

303 def __getitem__(self, idx: int) -> "BlockTable":

304 """Returns the BlockTable for the i-th KV cache group."""

305 return self.block_tables[idx]

306 +

307 +

308 +@triton.jit

309 +def _compute_slot_mapping_kernel(

310 + num_tokens,

311 + max_num_tokens,

312 + query_start_loc_ptr, # [num_reqs + 1], int32

313 + positions_ptr, # [num_tokens], int64

314 + block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat)

315 + block_table_stride, # max_num_blocks_per_req

316 + block_size,

317 + slot_mapping_ptr, # [max_num_tokens], int64

318 + TOTAL_CP_WORLD_SIZE: tl.constexpr,

319 + TOTAL_CP_RANK: tl.constexpr,

320 + CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,

321 + PAD_ID: tl.constexpr,

322 + BLOCK_SIZE: tl.constexpr,

323 +):

324 + req_idx = tl.program_id(0)

325 +

326 + if req_idx == tl.num_programs(0) - 1:

327 + # Pad remaining slots for CUDA graph compatibility.

328 + for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):

329 + offsets = i + tl.arange(0, BLOCK_SIZE)

330 + tl.store(

331 + slot_mapping_ptr + offsets,

332 + PAD_ID,

333 + mask=offsets < max_num_tokens,

334 + )

335 + return

336 +

337 + start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64)

338 + end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64)

339 +

340 + virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE

341 + row_offset = req_idx * block_table_stride

342 + for i in range(start_idx, end_idx, BLOCK_SIZE):

343 + offsets = i + tl.arange(0, BLOCK_SIZE)

344 + mask = offsets < end_idx

345 + pos = tl.load(positions_ptr + offsets, mask=mask, other=0)

346 + block_indices = pos // virtual_block_size

347 + block_numbers = tl.load(block_table_ptr + row_offset + block_indices).to(

348 + tl.int64

349 + )

350 +

351 + virtual_block_offsets = pos - block_indices * virtual_block_size

352 + is_local = (

353 + virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE

354 + ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK

355 + local_block_offsets = (

356 + virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)

357 + ) * CP_KV_CACHE_INTERLEAVE_SIZE + (

358 + virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE

359 + )

360 +

361 + slot_ids = block_numbers * block_size + local_block_offsets

362 + slot_ids = tl.where(is_local, slot_ids, PAD_ID)

363 + tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)

vllm/v1/worker/gpu_input_batch.py

+6 -4

@@ -219,7 +219,7 @@ class InputBatch:

219

220 # Speculative decoding

221 self.num_accepted_tokens_cpu_tensor = torch.ones(

222 - (max_num_reqs,), dtype=torch.int64, device="cpu", pin_memory=pin_memory

222 + (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory

223 )

224 self.num_accepted_tokens_cpu = self.num_accepted_tokens_cpu_tensor.numpy()

225

@@ -989,13 +989,15 @@ class InputBatch:

989 continue

990 num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)

991 # Also account for case where there may be a smaller number of

992 - # output placeholders (tokens can be discarded after a kv-load failure).

992 + # output placeholders (tokens can be discarded after kv-load

993 + # failure) or a larger number (async spec decode adds optimistic

994 + # placeholders that may exceed the actual acceptance count).

995 first_placeholder = req_output_token_ids.index(-1)

996 num_placeholders = len(req_output_token_ids) - first_placeholder

997 num_to_replace = min(num_sampled_ids, num_placeholders)

998 del new_ids[num_to_replace:]

997 - end_index = first_placeholder + num_to_replace

998 - req_output_token_ids[first_placeholder:end_index] = new_ids

999 + req_output_token_ids[first_placeholder:] = new_ids

1000 + # ^ Implicitly resizes to (first_placeholder + num_to_replace)

1001

1002 def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None:

1003 """

vllm/v1/worker/gpu_model_runner.py

+341 -127

@@ -7,7 +7,7 @@ import itertools

7 import threading

8 import time

9 from collections import defaultdict

10 -from collections.abc import Iterable, Iterator, Sequence

10 +from collections.abc import Callable, Iterable, Iterator, Sequence

11 from contextlib import contextmanager

12 from copy import copy, deepcopy

13 from dataclasses import dataclass, replace

@@ -172,6 +172,7 @@ from vllm.v1.spec_decode.ngram_proposer_gpu import (

172 update_scheduler_for_invalid_drafts,

173 )

174 from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer

175 +from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_change

176 from vllm.v1.structured_output.utils import apply_grammar_bitmask

177 from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext

178 from vllm.v1.worker import mamba_utils

@@ -570,6 +571,7 @@ class GPUModelRunner(

571 self.rejection_sampler = RejectionSampler(self.sampler)

572

573 self.num_spec_tokens = 0

574 + self.valid_sampled_token_count_gpu: torch.Tensor | None = None

575 if self.speculative_config:

576 self.num_spec_tokens = self.speculative_config.num_speculative_tokens

577 draft_config = self.speculative_config.draft_model_config

@@ -577,6 +579,9 @@ class GPUModelRunner(

579 self.effective_drafter_max_model_len = draft_config.max_model_len

580 else:

581 self.effective_drafter_max_model_len = self.max_model_len

582 + self.use_async_spec_decode = (

583 + self.use_async_scheduling and self.num_spec_tokens > 0

584 + )

585

586 # Request states.

587 self.requests: dict[str, CachedRequestState] = {}

@@ -659,11 +664,31 @@ class GPUModelRunner(

664

665 # Persistent buffers for CUDA graphs.

666 self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)

662 - self.positions = self._make_buffer(self.max_num_tokens, dtype=torch.int64)

667 + self.positions = torch.zeros(

668 + self.max_num_tokens, dtype=torch.int64, device=self.device

669 + )

670 self.query_start_loc = self._make_buffer(

671 self.max_num_reqs + 1, dtype=torch.int32

672 )

666 - self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

673 + self.seq_lens = torch.zeros(

674 + self.max_num_reqs, dtype=torch.int32, device=self.device

675 + )

676 + self.optimistic_seq_lens_cpu = torch.zeros(

677 + self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory

678 + )

679 + self.num_computed_tokens = torch.zeros(

680 + self.max_num_reqs, dtype=torch.int32, device=self.device

681 + )

682 + self.prev_num_draft_tokens = self._make_buffer(

683 + self.max_num_reqs, dtype=torch.int32

684 + )

685 + self.req_indices = self._make_buffer(self.max_num_tokens, dtype=torch.int64)

686 + # Maps current batch position -> previous batch position (-1 for new reqs)

687 + self.prev_positions = self._make_buffer(self.max_num_reqs, dtype=torch.int64)

688 + self.num_scheduled_tokens = self._make_buffer(

689 + self.max_num_reqs, dtype=torch.int32

690 + )

691 +

692 self.encoder_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

693 if self.dcp_world_size > 1:

694 self.dcp_local_seq_lens = self._make_buffer(

@@ -683,7 +708,7 @@ class GPUModelRunner(

708 self.max_num_reqs, dtype=torch.int32

709 )

710 self.num_accepted_tokens = self._make_buffer(

686 - self.max_num_reqs, dtype=torch.int64

711 + self.max_num_reqs, dtype=torch.int32

712 )

713

714 # Only relevant for multimodal models

@@ -722,12 +747,14 @@ class GPUModelRunner(

747 # None in the first PP rank. The rest are set after load_model.

748 self.intermediate_tensors: IntermediateTensors | None = None

749

725 - # OPTIMIZATION: Cache the tensors rather than creating them every step.

726 - # Keep in int64 to avoid overflow with long context

727 - self.arange_np = np.arange(

728 - max(self.max_num_reqs + 1, self.max_model_len, self.max_num_tokens),

729 - dtype=np.int64,

730 - )

750 + # OPTIMIZATION: Cache the arange tensors rather than creating them

751 + # every step. Keep in int64 to avoid overflow with long context.

752 + # - arange_np: immutable [0, 1, 2, ...] used as source for batched computation

753 + # - query_pos: CpuGpuBuffer for the computed batched arange result

754 + arange_size = max(self.max_num_reqs + 1, self.max_num_tokens)

755 + self.arange_np = np.arange(arange_size, dtype=np.int64)

756 + self.query_pos = self._make_buffer(arange_size, dtype=torch.int64)

757 + self._arange_scratch = np.empty(arange_size, dtype=np.int64)

758

759 # Layer pairings for cross-layer KV sharing.

760 # If an Attention layer `layer_name` is in the keys of this dict, it

@@ -812,7 +839,7 @@ class GPUModelRunner(

839 self.valid_sampled_token_count_copy_stream = torch.cuda.Stream()

840 self.valid_sampled_token_count_cpu = torch.empty(

841 self.max_num_reqs,

815 - dtype=torch.int64,

842 + dtype=torch.int32,

843 device="cpu",

844 pin_memory=self.pin_memory,

845 )

@@ -903,13 +930,13 @@ class GPUModelRunner(

930 return self.mrope_positions.gpu[:, :num_tokens]

931 if self.uses_xdrope_dim > 0:

932 return self.xdrope_positions.gpu[:, :num_tokens]

906 - return self.positions.gpu[:num_tokens]

933 + return self.positions[:num_tokens]

934 else:

935 if self.uses_mrope:

936 return self.mrope_positions.gpu[:, num_tokens]

937 if self.uses_xdrope_dim > 0:

938 return self.xdrope_positions.gpu[:, num_tokens]

912 - return self.positions.gpu[num_tokens]

939 + return self.positions[num_tokens]

940

941 def _make_buffer(

942 self, *size: int | torch.SymInt, dtype: torch.dtype, numpy: bool = True

@@ -953,7 +980,7 @@ class GPUModelRunner(

980 if len(token_type_id_requests) == 0:

981 return model_kwargs

982

956 - seq_lens = self.seq_lens.gpu[:num_reqs]

983 + seq_lens = self.seq_lens[:num_reqs]

984 token_type_ids = []

985

986 for i in range(num_reqs):

@@ -1021,7 +1048,7 @@ class GPUModelRunner(

1048 def _sync_device(self) -> None:

1049 torch.accelerator.synchronize()

1050

1024 - def _update_states(self, scheduler_output: "SchedulerOutput") -> None:

1051 + def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None:

1052 """Update the cached states and the persistent batch with the scheduler

1053 output.

1054

@@ -1086,6 +1113,8 @@ class GPUModelRunner(

1113 ngram_gpu_new_reqs: list[CachedRequestState] = []

1114

1115 reqs_to_add: list[CachedRequestState] = []

1116 + deferred_spec_decode_corrections = []

1117 +

1118 # Add new requests to the cached states.

1119 for new_req_data in scheduler_output.scheduled_new_reqs:

1120 req_id = new_req_data.req_id

@@ -1172,10 +1201,8 @@ class GPUModelRunner(

1201 scheduler_output,

1202 self.input_batch.req_id_to_index,

1203 )

1175 -

1176 - # Wait until valid_sampled_tokens_count is copied to cpu,

1177 - # then use it to update actual num_computed_tokens of each request.

1178 - valid_sampled_token_count = self._get_valid_sampled_token_count()

1204 + if self.use_async_spec_decode:

1205 + self.prev_num_draft_tokens.np.fill(0)

1206

1207 for i, req_id in enumerate(req_data.req_ids):

1208 req_state = self.requests[req_id]

@@ -1202,15 +1229,30 @@ class GPUModelRunner(

1229 if req_index is None:

1230 req_state.prev_num_draft_len = 0

1231 else:

1205 - assert self.input_batch.prev_req_id_to_index is not None

1206 - prev_req_index = self.input_batch.prev_req_id_to_index[req_id]

1207 - num_accepted = valid_sampled_token_count[prev_req_index] - 1

1208 - num_rejected = req_state.prev_num_draft_len - num_accepted

1209 - num_computed_tokens -= num_rejected

1210 - req_state.output_token_ids.extend([-1] * num_accepted)

1232 + # Optimistically assume all accepted; queue up a correction

1233 + # to be called after the model forward to preserve async

1234 + # scheduling. Corrected on GPU in _prepare_inputs.

1235 + optimistic_num_accepted = req_state.prev_num_draft_len

1236 + req_state.output_token_ids.extend([-1] * optimistic_num_accepted)

1237 +

1238 + deferred_spec_decode_corrections.append(

1239 + (req_id, optimistic_num_accepted, req_state)

1240 + )

1241 +

1242 + prev_req_index = (

1243 + self.input_batch.prev_req_id_to_index.get(req_id)

1244 + if self.input_batch.prev_req_id_to_index

1245 + else None

1246 + )

1247 + if prev_req_index is not None:

1248 + self.prev_num_draft_tokens.np[prev_req_index] = (

1249 + optimistic_num_accepted

1250 + )

1251

1212 - if is_ngram_gpu and num_accepted > 0 and req_index is not None:

1213 - self.input_batch.num_tokens_no_spec[req_index] += num_accepted

1252 + if is_ngram_gpu and optimistic_num_accepted > 0:

1253 + self.input_batch.num_tokens_no_spec[req_index] += (

1254 + optimistic_num_accepted

1255 + )

1256

1257 # Update the cached states.

1258 req_state.num_computed_tokens = num_computed_tokens

@@ -1238,7 +1280,8 @@ class GPUModelRunner(

1280 )

1281 elif num_output_tokens < len(req_state.output_token_ids):

1282 # Some output tokens were discarded due to a sync-KV-load

1241 - # failure. Align the cached state.

1283 + # failure, or output_token_ids was inflated by the optimistic

1284 + # extend above (async spec decode). Align the cached state.

1285 del req_state.output_token_ids[num_output_tokens:]

1286 if req_index is not None:

1287 end_idx = (

@@ -1326,6 +1369,40 @@ class GPUModelRunner(

1369 _pinned_val_buf=self._ngram_pinned_val_buf,

1370 )

1371

1372 + if deferred_spec_decode_corrections:

1373 +

1374 + def correct_spec_decode_token_counts():

1375 + valid_sampled_token_count = self._get_valid_sampled_token_count()

1376 + if not valid_sampled_token_count:

1377 + return

1378 + prev_req_id_to_index = self.input_batch.prev_req_id_to_index

1379 + if not prev_req_id_to_index:

1380 + return

1381 + for (

1382 + req_id,

1383 + optimistic_num_accepted,

1384 + req_state,

1385 + ) in deferred_spec_decode_corrections:

1386 + prev_req_index = prev_req_id_to_index.get(req_id)

1387 + if prev_req_index is None:

1388 + continue

1389 + num_accepted = valid_sampled_token_count[prev_req_index] - 1

1390 + correction = optimistic_num_accepted - num_accepted

1391 + req_state.num_computed_tokens -= correction

1392 + cur_req_index = self.input_batch.req_id_to_index.get(req_id)

1393 + if cur_req_index is None:

1394 + continue

1395 + self.input_batch.num_computed_tokens_cpu[cur_req_index] -= (

1396 + correction

1397 + )

1398 + if is_ngram_gpu and correction > 0:

1399 + self.input_batch.num_tokens_no_spec[cur_req_index] -= correction

1400 + self.num_tokens_no_spec_gpu[cur_req_index] -= correction

1401 +

1402 + return correct_spec_decode_token_counts

1403 + else:

1404 + return None

1405 +

1406 def _update_states_after_model_execute(

1407 self, output_token_ids: torch.Tensor, scheduler_output: "SchedulerOutput"

1408 ) -> None:

@@ -1340,6 +1417,9 @@ class GPUModelRunner(

1417 if not self.speculative_config or not self.model_config.is_hybrid:

1418 return

1419

1420 + # TODO: Remove .cpu() sync to enable fully async for hybrid model;

1421 + # Use num_computed_tokens.gpu instead of req.num_computed_tokens to

1422 + # support aligned mamba cache mode.

1423 # Find the number of accepted tokens for each sequence.

1424 num_reqs = output_token_ids.size(0)

1425 self.num_accepted_tokens.gpu[:num_reqs] = (

@@ -1486,12 +1566,14 @@ class GPUModelRunner(

1566 def _get_cumsum_and_arange(

1567 self,

1568 num_tokens: np.ndarray,

1569 + arange_out: np.ndarray,

1570 cumsum_dtype: np.dtype | None = None,

1490 - ) -> tuple[np.ndarray, np.ndarray]:

1571 + ) -> np.ndarray:

1572 """Get the cumulative sum and batched arange of the given array.

1492 - # E.g., [2, 5, 3] -> ([2, 7, 10], [0, 1, 0, 1, 2, 3, 4, 0, 1, 2])

1493 - # Equivalent to but faster than:

1494 - # np.concatenate([np.arange(n) for n in num_tokens])

1573 + E.g., [2, 5, 3] -> [2, 7, 10], arange written to

1574 + arange_out[:10] as [0, 1, 0, 1, 2, 3, 4, 0, 1, 2].

1575 + Equivalent to but faster than:

1576 + np.concatenate([np.arange(n) for n in num_tokens])

1577 """

1578 # Step 1. [2, 5, 3] -> [2, 7, 10]

1579 cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype)

@@ -1499,13 +1581,33 @@ class GPUModelRunner(

1581 # Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7]

1582 cumsums_offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens)

1583 # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

1502 - arange = self.arange_np[:total_num_tokens] - cumsums_offsets

1584 + np.subtract(

1585 + self.arange_np[:total_num_tokens],

1586 + cumsums_offsets,

1587 + out=arange_out[:total_num_tokens],

1588 + )

1589 +

1590 + return cu_num_tokens

1591 +

1592 + def _compute_prev_positions(self, num_reqs: int) -> None:

1593 + """Build prev_positions mapping: current pos -> previous pos (-1 if new).

1594 +

1595 + Populates self.prev_positions.np[:num_reqs] with the mapping.

1596 + """

1597 + prev_req_id_to_index = self.input_batch.prev_req_id_to_index

1598 + prev_positions = self.prev_positions.np[:num_reqs]

1599 +

1600 + if not prev_req_id_to_index:

1601 + prev_positions.fill(-1)

1602 + return

1603

1504 - return cu_num_tokens, arange

1604 + for i, req_id in enumerate(self.input_batch.req_ids[:num_reqs]):

1605 + prev_positions[i] = prev_req_id_to_index.get(req_id, -1)

1606

1607 def _prepare_input_ids(

1608 self,

1609 scheduler_output: "SchedulerOutput",

1610 + num_reqs: int,

1611 total_num_scheduled_tokens: int,

1612 cu_num_tokens: np.ndarray,

1613 ) -> None:

@@ -1513,7 +1615,11 @@ class GPUModelRunner(

1615

1616 Carefully handles the `prev_sampled_token_ids` which can be cached

1617 from the previous engine iteration, in which case those tokens on the

1516 - GPU need to be copied into the corresponding slots into input_ids."""

1618 + GPU need to be copied into the corresponding slots into input_ids.

1619 +

1620 + Uses self.prev_positions[:num_reqs] which maps current pos -> prev pos

1621 + (-1 for new requests).

1622 + """

1623

1624 if self.input_batch.prev_sampled_token_ids is None:

1625 # Normal scheduling case

@@ -1526,47 +1632,50 @@ class GPUModelRunner(

1632 # Async scheduling case, where some decode requests from the previous

1633 # iteration won't have entries in input_ids_cpu and need to be copied

1634 # on the GPU from prev_sampled_token_ids.

1529 - prev_req_id_to_index = self.input_batch.prev_req_id_to_index

1530 - assert prev_req_id_to_index is not None

1635 + prev_positions = self.prev_positions.np[:num_reqs]

1636 + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens

1637 sample_flattened_indices: list[int] = []

1638 spec_flattened_indices: list[int] = []

1533 - prev_common_req_indices: list[int] = []

1639 prev_draft_token_indices: list[int] = []

1535 - indices_match = True

1640 + prev_indices: list[int] = []

1641 + common_indices_match = True

1642 max_flattened_index = -1

1643 total_num_spec_tokens = 0

1538 - scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens

1644

1540 - for req_id, cur_index in self.input_batch.req_id_to_index.items():

1541 - if (prev_index := prev_req_id_to_index.get(req_id)) is not None:

1542 - prev_common_req_indices.append(prev_index)

1543 - # We need to compute the flattened input_ids index of the

1544 - # last token in each common request.

1545 - draft_len = len(scheduled_spec_tokens.get(req_id, ()))

1546 - total_num_spec_tokens += draft_len

1547 - flattened_index = cu_num_tokens[cur_index].item() - 1

1548 - # example: cu_num_tokens = [2, 5, 8], draft_tokens = [1, 2, 2]

1549 - # sample_flattened_indices = [0, 2, 5]

1550 - # spec_flattened_indices = [1, 3, 4, 6, 7]

1551 - sample_flattened_indices.append(flattened_index - draft_len)

1552 - spec_flattened_indices.extend(

1553 - range(flattened_index - draft_len + 1, flattened_index + 1)

1554 - )

1555 - start = prev_index * self.num_spec_tokens

1556 - # prev_draft_token_indices is used to find which draft_tokens_id

1557 - # should be copied to input_ids

1558 - # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]]

1559 - # flatten draft_tokens_id [1,2,3,4,5,6]

1560 - # draft_len of each request [1, 2, 1]

1561 - # then prev_draft_token_indices is [0, 2, 3, 4]

1562 - prev_draft_token_indices.extend(range(start, start + draft_len))

1563 - indices_match &= prev_index == flattened_index

1564 - max_flattened_index = max(max_flattened_index, flattened_index)

1645 + for cur_index in range(num_reqs):

1646 + prev_index = prev_positions[cur_index]

1647 + if prev_index < 0:

1648 + continue

1649 + prev_indices.append(prev_index)

1650 + req_id = self.input_batch.req_ids[cur_index]

1651 + # We need to compute the flattened input_ids index of the

1652 + # last token in each common request.

1653 + draft_len = len(scheduled_spec_tokens.get(req_id, ()))

1654 + total_num_spec_tokens += draft_len

1655 + flattened_index = cu_num_tokens[cur_index].item() - 1

1656 + # example: cu_num_tokens = [2, 5, 8], draft_tokens = [1, 2, 2]

1657 + # sample_flattened_indices = [0, 2, 5]

1658 + # spec_flattened_indices = [1, 3, 4, 6, 7]

1659 + sample_flattened_indices.append(flattened_index - draft_len)

1660 + spec_flattened_indices.extend(

1661 + range(flattened_index - draft_len + 1, flattened_index + 1)

1662 + )

1663 + start = prev_index * self.num_spec_tokens

1664 + # prev_draft_token_indices is used to find which draft_tokens_id

1665 + # should be copied to input_ids

1666 + # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]]

1667 + # flatten draft_tokens_id [1,2,3,4,5,6]

1668 + # draft_len of each request [1, 2, 1]

1669 + # then prev_draft_token_indices is [0, 2, 3, 4]

1670 + prev_draft_token_indices.extend(range(start, start + draft_len))

1671 + common_indices_match &= prev_index == flattened_index

1672 + max_flattened_index = max(max_flattened_index, flattened_index)

1673 +

1674 num_common_tokens = len(sample_flattened_indices)

1675 total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens

1676 if num_common_tokens < total_without_spec:

1677 # If not all requests are decodes from the last iteration,

1569 - # We need to copy the input_ids_cpu to the GPU first.

1678 + # we need to copy the input_ids_cpu to the GPU first.

1679 self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

1680 if self.enable_prompt_embeds:

1681 self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)

@@ -1575,7 +1684,7 @@ class GPUModelRunner(

1684 # No requests in common with the previous iteration

1685 # So input_ids.cpu will have all the input ids.

1686 return

1578 - if indices_match and max_flattened_index == (num_common_tokens - 1):

1687 + if common_indices_match and max_flattened_index == (num_common_tokens - 1):

1688 # Common-case optimization: the batch is unchanged

1689 # and no reordering happened.

1690 # The indices are both the same permutation of 0..N-1 so

@@ -1592,7 +1701,7 @@ class GPUModelRunner(

1701 sample_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory

1702 ).to(self.device, non_blocking=True)

1703 prev_common_req_indices_tensor = torch.tensor(

1595 - prev_common_req_indices, dtype=torch.int64, pin_memory=self.pin_memory

1704 + prev_indices, dtype=torch.int64, pin_memory=self.pin_memory

1705 ).to(self.device, non_blocking=True)

1706 self.input_ids.gpu.scatter_(

1707 dim=0,

@@ -1696,15 +1805,15 @@ class GPUModelRunner(

1805 req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)

1806

1807 # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]

1699 - # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

1700 - cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens)

1808 + # self.query_pos.np[:10]: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

1809 + cu_num_tokens = self._get_cumsum_and_arange(

1810 + num_scheduled_tokens, self.query_pos.np

1811 + )

1812

1813 # Get positions.

1703 - positions_np = self.positions.np[:total_num_scheduled_tokens]

1704 - np.add(

1705 - self.input_batch.num_computed_tokens_cpu[req_indices],

1706 - arange,

1707 - out=positions_np,

1814 + positions_np = (

1815 + self.input_batch.num_computed_tokens_cpu[req_indices]

1816 + + self.query_pos.np[: cu_num_tokens[-1]]

1817 )

1818

1819 # Calculate M-RoPE positions.

@@ -1782,9 +1891,6 @@ class GPUModelRunner(

1891

1892 output_idx += num_sched

1893

1785 - self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)

1786 - self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)

1787 -

1894 # Prepare the attention metadata.

1895 self.query_start_loc.np[0] = 0

1896 self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens

@@ -1794,12 +1900,21 @@ class GPUModelRunner(

1900 self.query_start_loc.copy_to_gpu()

1901 query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]

1902

1797 - self.seq_lens.np[:num_reqs] = (

1798 - self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens

1903 + # Compute optimistic seq_lens (assumes all draft tokens from previous

1904 + # iteration accepted). Store in optimistic_seq_lens_cpu for use by

1905 + # _build_attention_metadata (max_seq_len) and discard_request_mask.

1906 + # seq_lens (GPU) will be computed later using the same optimistic values.

1907 + torch.add(

1908 + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],

1909 + torch.from_numpy(num_scheduled_tokens),

1910 + out=self.optimistic_seq_lens_cpu[:num_reqs],

1911 )

1800 - # Fill unused with 0 for full cuda graph mode.

1801 - self.seq_lens.np[num_reqs:].fill(0)

1802 - self.seq_lens.copy_to_gpu()

1912 + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0)

1913 +

1914 + # Build prev_positions mapping: current pos -> prev pos (-1 if new).

1915 + # Used for gathering from previous iteration's GPU tensors.

1916 + prev_req_id_to_index = self.input_batch.prev_req_id_to_index

1917 + self._compute_prev_positions(num_reqs)

1918

1919 num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids]

1920 num_tokens_np = np.array(num_tokens, dtype=np.int32)

@@ -1807,13 +1922,78 @@ class GPUModelRunner(

1922 # Record which requests should not be sampled,

1923 # so that we could clear the sampled tokens before returning

1924 self.discard_request_mask.np[:num_reqs] = (

1810 - self.seq_lens.np[:num_reqs] < num_tokens_np

1925 + self.optimistic_seq_lens_cpu[:num_reqs].numpy() < num_tokens_np

1926 )

1927 self.discard_request_mask.copy_to_gpu(num_reqs)

1928

1929 + # Sync num_accepted_tokens from CPU (set by

1930 + # _update_states_after_model_execute for hybrid models).

1931 + if self.num_accepted_tokens_event is not None:

1932 + self.num_accepted_tokens_event.synchronize()

1933 + self.num_accepted_tokens.np[:num_reqs] = (

1934 + self.input_batch.num_accepted_tokens_cpu[:num_reqs]

1935 + )

1936 + self.num_accepted_tokens.np[num_reqs:].fill(1)

1937 + self.num_accepted_tokens.copy_to_gpu()

1938 + else:

1939 + self.num_accepted_tokens.np.fill(1)

1940 + self.num_accepted_tokens.gpu.fill_(1)

1941 +

1942 + # Update num_computed_tokens on GPU. In async spec decode,

1943 + # CPU values are optimistic (all drafts accepted). The kernel

1944 + # corrects on GPU using the previous step's

1945 + # valid_sampled_token_count_gpu. Otherwise, just copy from CPU.

1946 + if (

1947 + self.use_async_spec_decode

1948 + and self.valid_sampled_token_count_gpu is not None

1949 + and prev_req_id_to_index

1950 + ):

1951 + self.prev_positions.copy_to_gpu(num_reqs)

1952 + self.prev_num_draft_tokens.copy_to_gpu()

1953 + cpu_values = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs].to(

1954 + device=self.device, non_blocking=True

1955 + )

1956 + update_num_computed_tokens_for_batch_change(

1957 + self.num_computed_tokens,

1958 + self.num_accepted_tokens.gpu[:num_reqs],

1959 + self.prev_positions.gpu[:num_reqs],

1960 + self.valid_sampled_token_count_gpu,

1961 + self.prev_num_draft_tokens.gpu,

1962 + cpu_values,

1963 + )

1964 + else:

1965 + self.num_computed_tokens[:num_reqs].copy_(

1966 + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],

1967 + non_blocking=True,

1968 + )

1969 +

1970 + self.req_indices.np[:total_num_scheduled_tokens] = req_indices

1971 + self.req_indices.copy_to_gpu(total_num_scheduled_tokens)

1972 + req_indices_gpu = self.req_indices.gpu[:total_num_scheduled_tokens]

1973 +

1974 + self.query_pos.copy_to_gpu(total_num_scheduled_tokens)

1975 + self.num_scheduled_tokens.np[:num_reqs] = num_scheduled_tokens

1976 + self.num_scheduled_tokens.copy_to_gpu(num_reqs)

1977 + num_scheduled_tokens_gpu = self.num_scheduled_tokens.gpu[:num_reqs]

1978 + self.positions[:total_num_scheduled_tokens] = (

1979 + self.num_computed_tokens[req_indices_gpu].to(torch.int64)

1980 + + self.query_pos.gpu[:total_num_scheduled_tokens]

1981 + )

1982 + self.seq_lens[:num_reqs] = (

1983 + self.num_computed_tokens[:num_reqs] + num_scheduled_tokens_gpu

1984 + )

1985 + self.seq_lens[num_reqs:].fill_(0)

1986 +

1987 + self.input_batch.block_table.compute_slot_mapping(

1988 + num_reqs,

1989 + self.query_start_loc.gpu[: num_reqs + 1],

1990 + self.positions[:total_num_scheduled_tokens],

1991 + )

1992 +

1993 # Copy the tensors to the GPU.

1994 self._prepare_input_ids(

1995 scheduler_output,

1996 + num_reqs,

1997 total_num_scheduled_tokens,

1998 cu_num_tokens,

1999 )

@@ -1830,9 +2010,14 @@ class GPUModelRunner(

2010 self.xdrope_positions.cpu[:, :total_num_scheduled_tokens],

2011 non_blocking=True,

2012 )

1833 - else:

1834 - # Common case (1D positions)

1835 - self.positions.copy_to_gpu(total_num_scheduled_tokens)

2013 + if self.use_async_spec_decode and (self.uses_mrope or self.uses_xdrope_dim > 0):

2014 + drift = self.num_computed_tokens[req_indices_gpu].to(

2015 + torch.int64

2016 + ) - self.input_batch.num_computed_tokens_cpu_tensor[req_indices].to(

2017 + device=self.device, dtype=torch.int64, non_blocking=True

2018 + )

2019 + target = self.mrope_positions if self.uses_mrope else self.xdrope_positions

2020 + target.gpu[:, :total_num_scheduled_tokens] += drift

2021

2022 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0

2023 if not use_spec_decode:

@@ -1857,12 +2042,13 @@ class GPUModelRunner(

2042 draft_token_ids,

2043 ) in scheduler_output.scheduled_spec_decode_tokens.items():

2044 req_idx = self.input_batch.req_id_to_index[req_id]

1860 - num_draft_tokens[req_idx] = len(draft_token_ids)

2045 + draft_len = len(draft_token_ids)

2046 + num_draft_tokens[req_idx] = draft_len

2047 if (

2048 self.input_batch.num_computed_tokens_cpu[req_idx]

2049 >= self.input_batch.num_prompt_tokens[req_idx]

2050 ):

1865 - num_decode_draft_tokens[req_idx] = len(draft_token_ids)

2051 + num_decode_draft_tokens[req_idx] = draft_len

2052 spec_decode_metadata = self._calc_spec_decode_metadata(

2053 num_draft_tokens, cu_num_tokens

2054 )

@@ -1924,16 +2110,7 @@ class GPUModelRunner(

2110 # window size when capturing to make sure the correct kernel is selected.

2111 max_seq_len = self.max_model_len

2112 else:

1927 - max_seq_len = self.seq_lens.np[:num_reqs].max().item()

1928 -

1929 - if use_spec_decode:

1930 - if self.num_accepted_tokens_event is not None:

1931 - self.num_accepted_tokens_event.synchronize()

1932 - self.num_accepted_tokens.np[:num_reqs] = (

1933 - self.input_batch.num_accepted_tokens_cpu[:num_reqs]

1934 - )

1935 - self.num_accepted_tokens.np[num_reqs:].fill(1)

1936 - self.num_accepted_tokens.copy_to_gpu()

2113 + max_seq_len = self.optimistic_seq_lens_cpu.numpy()[:num_reqs].max().item()

2114

2115 kv_cache_groups = self.kv_cache_config.kv_cache_groups

2116

@@ -1963,22 +2140,29 @@ class GPUModelRunner(

2140 attn_gid = self.routed_experts_attn_gid

2141 slot_mapping_attn = slot_mappings[attn_gid]

2142 self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy()

1966 - # Compute is_prefilling: True if request is still in prefill phase

1967 - # (num_computed_tokens < num_prompt_tokens). Used by mamba backends to

1968 - # distinguish actual decodes from short extends.

2143 num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

2144 :num_reqs_padded

2145 ]

2146 num_prompt_tokens_cpu = self.input_batch.num_prompt_tokens_cpu_tensor[

2147 :num_reqs_padded

2148 ]

2149 + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs_padded]

2150 +

2151 + # is_prefilling: True if request is still in prefill phase.

2152 + # Used by mamba backends to distinguish actual decodes from

2153 + # short extends.

2154 is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu

2155

2156 + if self.use_async_spec_decode:

2157 + # GPU tensors are authoritative in async mode.

2158 + seq_lens_cpu = None

2159 + num_computed_tokens_cpu = None

2160 +

2161 cm_base = CommonAttentionMetadata(

2162 query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],

2163 query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],

1980 - seq_lens=self.seq_lens.gpu[:num_reqs_padded],

1981 - _seq_lens_cpu=self.seq_lens.cpu[:num_reqs_padded],

2164 + seq_lens=self.seq_lens[:num_reqs_padded],

2165 + _seq_lens_cpu=seq_lens_cpu,

2166 _num_computed_tokens_cpu=num_computed_tokens_cpu,

2167 num_reqs=num_reqs_padded,

2168 num_actual_tokens=num_tokens_padded,

@@ -1992,7 +2176,7 @@ class GPUModelRunner(

2176

2177 if self.dcp_world_size > 1:

2178 self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(

1995 - self.seq_lens.cpu[:num_reqs],

2179 + self.optimistic_seq_lens_cpu[:num_reqs],

2180 self.dcp_world_size,

2181 self.dcp_rank,

2182 self.parallel_config.cp_kv_cache_interleave_size,

@@ -2396,33 +2580,34 @@ class GPUModelRunner(

2580 # [4, 1, 3, 1, 2]

2581 num_sampled_tokens = num_draft_tokens + 1

2582

2399 - # Step 1. cu_num_sampled_tokens: [4, 5, 8, 9, 11]

2400 - # arange: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]

2401 - cu_num_sampled_tokens, arange = self._get_cumsum_and_arange(

2402 - num_sampled_tokens, cumsum_dtype=np.int32

2583 + # Step 1.

2584 + # cu_num_sampled_tokens: [4, 5, 8, 9, 11]

2585 + # _arange_scratch[:11]: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]

2586 + cu_num_sampled_tokens = self._get_cumsum_and_arange(

2587 + num_sampled_tokens, self._arange_scratch, cumsum_dtype=np.int32

2588 )

2589 # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]

2590 logits_indices = np.repeat(

2591 cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens

2592 )

2593 # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]

2409 - logits_indices += arange

2594 + logits_indices += self._arange_scratch[: cu_num_sampled_tokens[-1]]

2595

2596 # Compute the bonus logits indices.

2597 bonus_logits_indices = cu_num_sampled_tokens - 1

2598

2599 # Compute the draft logits indices.

2600 # cu_num_draft_tokens: [3, 3, 5, 5, 6]

2416 - # arange: [0, 1, 2, 0, 1, 0]

2417 - cu_num_draft_tokens, arange = self._get_cumsum_and_arange(

2418 - num_draft_tokens, cumsum_dtype=np.int32

2601 + # _arange_scratch[:6]: [0, 1, 2, 0, 1, 0]

2602 + cu_num_draft_tokens = self._get_cumsum_and_arange(

2603 + num_draft_tokens, self._arange_scratch, cumsum_dtype=np.int32

2604 )

2605 # [0, 0, 0, 5, 5, 9]

2606 target_logits_indices = np.repeat(

2607 cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens

2608 )

2609 # [0, 1, 2, 5, 6, 9]

2425 - target_logits_indices += arange

2610 + target_logits_indices += self._arange_scratch[: cu_num_draft_tokens[-1]]

2611

2612 # TODO: Optimize the CPU -> GPU copy.

2613 cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(

@@ -2924,7 +3109,7 @@ class GPUModelRunner(

3109 )

3110

3111 hidden_states = hidden_states[:num_scheduled_tokens]

2927 - seq_lens_cpu = self.seq_lens.cpu[:num_reqs]

3112 + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs]

3113

3114 pooling_metadata = self.input_batch.get_pooling_metadata()

3115 pooling_metadata.build_pooling_cursor(

@@ -3083,9 +3268,9 @@ class GPUModelRunner(

3268 elif self.uses_xdrope_dim > 0:

3269 positions = self.xdrope_positions.gpu[:, :num_input_tokens]

3270 else:

3086 - positions = self.positions.gpu[:num_input_tokens]

3271 + positions = self.positions[:num_input_tokens]

3272 if num_input_tokens > num_scheduled_tokens:

3088 - self.positions.gpu[num_scheduled_tokens:num_input_tokens].zero_()

3273 + self.positions[num_scheduled_tokens:num_input_tokens].zero_()

3274

3275 if is_first_rank:

3276 intermediate_tensors = None

@@ -3610,7 +3795,7 @@ class GPUModelRunner(

3795 self.synchronize_input_prep(),

3796 ):

3797 # Update persistent batch states.

3613 - self._update_states(scheduler_output)

3798 + deferred_state_corrections_fn = self._update_states(scheduler_output)

3799

3800 if has_ec_transfer() and not get_ec_transfer().is_consumer:

3801 with self.maybe_get_ec_connector_output(

@@ -3723,6 +3908,12 @@ class GPUModelRunner(

3908 pad_attn = cudagraph_mode == CUDAGraphMode.FULL

3909

3910 if self.cache_config.mamba_cache_mode == "align":

3911 + # preprocess_mamba reads req_state.num_computed_tokens (CPU)

3912 + # to decide copy operations, so we must apply deferred

3913 + # corrections before it runs.

3914 + if deferred_state_corrections_fn:

3915 + deferred_state_corrections_fn()

3916 + deferred_state_corrections_fn = None

3917 mamba_utils.preprocess_mamba(

3918 scheduler_output,

3919 self.kv_cache_config,

@@ -3734,6 +3925,14 @@ class GPUModelRunner(

3925 self.model.get_mamba_state_copy_func(),

3926 self._get_mamba_copy_bufs(),

3927 )

3928 + # preprocess_mamba resets num_accepted_tokens_cpu to 1

3929 + # for requests whose state was copied to a new block.

3930 + # Re-sync to GPU so the mamba kernel reads from the

3931 + # correct initial state slot (init_token_idx = 0).

3932 + self.num_accepted_tokens.np[:num_reqs] = (

3933 + self.input_batch.num_accepted_tokens_cpu[:num_reqs]

3934 + )

3935 + self.num_accepted_tokens.copy_to_gpu(num_reqs)

3936

3937 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0

3938 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices

@@ -3894,6 +4093,12 @@ class GPUModelRunner(

4093 slot_mappings,

4094 )

4095 self.kv_connector_output = kv_connector_output

4096 +

4097 + # Now the batch has been launched we can wait for corrections from the

4098 + # previous model forward without breaking async scheduling.

4099 + if deferred_state_corrections_fn:

4100 + deferred_state_corrections_fn()

4101 +

4102 return None

4103

4104 @torch.inference_mode

@@ -3958,6 +4163,7 @@ class GPUModelRunner(

4163

4164 self._draft_token_ids = None

4165 self._draft_token_req_ids = None

4166 + self.valid_sampled_token_count_gpu = None

4167 self.input_batch.prev_sampled_token_ids = None

4168

4169 def propose_draft_token_ids(sampled_token_ids):

@@ -4002,7 +4208,7 @@ class GPUModelRunner(

4208 assert spec_decode_common_attn_metadata is not None

4209 next_token_ids, valid_sampled_tokens_count = (

4210 self.drafter.prepare_next_token_ids_padded(

4005 - spec_decode_common_attn_metadata,

4211 + self.optimistic_seq_lens_cpu,

4212 sampled_token_ids,

4213 self.requests,

4214 self.input_batch,

@@ -4237,6 +4443,9 @@ class GPUModelRunner(

4443 counts_cpu[: counts.shape[0]].copy_(counts, non_blocking=True)

4444 self.valid_sampled_token_count_event.record()

4445

4446 + if self.use_async_spec_decode:

4447 + # Stash for GPU-side correction in _prepare_inputs.

4448 + self.valid_sampled_token_count_gpu = valid_sampled_tokens_count

4449 self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1)

4450

4451 def _get_valid_sampled_token_count(self) -> list[int]:

@@ -4366,7 +4575,7 @@ class GPUModelRunner(

4575 )

4576 next_token_ids, valid_sampled_tokens_count = (

4577 self.drafter.prepare_next_token_ids_padded(

4369 - common_attn_metadata,

4578 + self.optimistic_seq_lens_cpu,

4579 sampled_token_ids,

4580 self.requests,

4581 self.input_batch,

@@ -4405,7 +4614,7 @@ class GPUModelRunner(

4614 )

4615 next_token_ids, valid_sampled_tokens_count = (

4616 self.drafter.prepare_next_token_ids_padded(

4408 - common_attn_metadata,

4617 + self.optimistic_seq_lens_cpu,

4618 sampled_token_ids,

4619 self.requests,

4620 self.input_batch,

@@ -5148,14 +5357,19 @@ class GPUModelRunner(

5357 # In the mixed batch mode (used for FI warmup), we use

5358 # shorter sequence lengths to run faster.

5359 # TODO(luka) better system for describing dummy batches

5151 - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] # type: ignore[assignment]

5360 + seq_lens = torch.tensor( # type: ignore[assignment]

5361 + [1] * num_decode_tokens + [num_prefill_tokens + 1],

5362 + dtype=torch.int,

5363 + )

5364 else:

5365 seq_lens = max_query_len # type: ignore[assignment]

5154 - self.seq_lens.np[:num_reqs] = seq_lens

5155 - self.seq_lens.np[num_reqs:] = 0

5156 - self.seq_lens.copy_to_gpu()

5366 + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens

5367 + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0)

5368 + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True)

5369

5158 - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)

5370 + cum_num_tokens = self._get_cumsum_and_arange(

5371 + num_scheduled_tokens, self.query_pos.np

5372 + )

5373 self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens

5374 self.query_start_loc.copy_to_gpu()

5375

@@ -5201,7 +5415,7 @@ class GPUModelRunner(

5415 elif self.uses_xdrope_dim > 0:

5416 positions = self.xdrope_positions.gpu[:, :num_tokens_padded]

5417 else:

5204 - positions = self.positions.gpu[:num_tokens_padded]

5418 + positions = self.positions[:num_tokens_padded]

5419

5420 if get_pp_group().is_first_rank:

5421 intermediate_tensors = None