diff --git a/arctic_inference/vllm/args.py b/arctic_inference/vllm/args.py index 8407e4bec..9cba50abd 100644 --- a/arctic_inference/vllm/args.py +++ b/arctic_inference/vllm/args.py @@ -20,7 +20,7 @@ from vllm.config import ParallelConfig from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs -from vllm.utils import FlexibleArgumentParser +from vllm.utils.argparse_utils import FlexibleArgumentParser from arctic_inference.patching import ArcticPatch from arctic_inference.vllm.config import ArcticParallelConfig @@ -50,7 +50,6 @@ class EngineArgsPatch(ArcticPatch[EngineArgs]): _orig_add_cli_args = EngineArgs.add_cli_args _orig_from_cli_args = EngineArgs.__dict__["from_cli_args"].__wrapped__ _orig_create_engine_config = EngineArgs.create_engine_config - _orig_is_v1_supported_oracle = EngineArgs._is_v1_supported_oracle def __new__(cls, *args, **kwargs): # Override __new__ to return an ArcticEngineArgs instead of an @@ -109,6 +108,11 @@ def create_engine_config(self, *args, **kwargs): if (self.ulysses_sequence_parallel_size > 1 and self.distributed_executor_backend is None): self.distributed_executor_backend = "mp" + + # Store ulysses_sequence_parallel_size for access during config initialization + from arctic_inference.vllm import ulysses + ulysses._ulysses_sp_size = self.ulysses_sequence_parallel_size + vllm_config = self._orig_create_engine_config(*args, **kwargs) # Recreate the parallel config with Arctic parameters since they might # not be passed to the parallel config __init__ when first initialized. @@ -121,21 +125,6 @@ def create_engine_config(self, *args, **kwargs): vllm_config.parallel_config = ArcticParallelConfig(**kwargs) return vllm_config - def _is_v1_supported_oracle(self, *args, **kwargs): - orig_speculative_config = self.speculative_config - - # Since Arctic Inference is only compatible with v1 and we already - # check it earlier, we can just disable this check altogether. - if (self.speculative_config is not None and - self.speculative_config.get("method") in ("arctic", "suffix")): - self.speculative_config = None - - res = self._orig_is_v1_supported_oracle(*args, **kwargs) - - self.speculative_config = orig_speculative_config - - return res - class AsyncEngineArgsPatch(ArcticPatch[AsyncEngineArgs]): diff --git a/arctic_inference/vllm/config.py b/arctic_inference/vllm/config.py index aac94303b..b310e1f4f 100644 --- a/arctic_inference/vllm/config.py +++ b/arctic_inference/vllm/config.py @@ -13,9 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass +from pydantic.dataclasses import dataclass import logging +import vllm from vllm.config import ParallelConfig, SpeculativeConfig, VllmConfig from vllm.transformers_utils.configs.mlp_speculator import MLPSpeculatorConfig @@ -55,6 +56,7 @@ def world_size(self, value: int) -> None: @dataclass class ArcticSpeculativeConfig(SpeculativeConfig): + method: str | None = None enable_suffix_decoding: bool = False suffix_cache_max_depth: int = 64 suffix_speculative_tokens: int = 0 @@ -80,21 +82,18 @@ class SpeculativeConfigPatch(ArcticPatch[SpeculativeConfig]): _orig_post_init = SpeculativeConfig.__post_init__ def __new__(cls, *args, **kwargs): - # Override __new__ to return an ArcticSpeculativeConfig instead of a - # SpeculativeConfig when creating a new instance of the class. if cls is SpeculativeConfig: return ArcticSpeculativeConfig.__new__(ArcticSpeculativeConfig, *args, **kwargs) return super(SpeculativeConfig, cls).__new__(cls) def __post_init__(self): - use_suffix = (self.method - == "suffix") or (self.method is None - and self.enable_suffix_decoding) - use_hybrid = (self.method == "arctic" - and self.enable_suffix_decoding) - if (use_suffix or self.method == "arctic") and \ - self.disable_by_batch_size is None: + is_arctic_method = self.method in ("arctic", "mlp_speculator") + use_suffix = (self.method == "suffix") or (self.method is None + and self.enable_suffix_decoding) + use_hybrid = (self.method == "arctic" and self.enable_suffix_decoding) + + if (use_suffix or is_arctic_method) and self.disable_by_batch_size is None: logger.info("Defaulting disable_by_batch_size to 64") self.disable_by_batch_size = 64 @@ -104,8 +103,30 @@ def __post_init__(self): if use_suffix: self.method = "suffix" self.enable_suffix_decoding = True - self.num_speculative_tokens = self.suffix_cache_max_depth + # Use suffix_speculative_tokens if explicitly set, otherwise + # default to 16 (not suffix_cache_max_depth which can be very + # large and makes every step process 1+N tokens even when the + # suffix cache has no matches). + # NOTE: num_speculative_tokens defaults to None (not 0). + if self.suffix_speculative_tokens > 0: + self.num_speculative_tokens = self.suffix_speculative_tokens + elif self.num_speculative_tokens is None: + self.num_speculative_tokens = 16 self._verify_args() + return + + if is_arctic_method: + actual_draft_model = getattr(self, "draft_model", None) + + self.draft_model = None + + try: + self._orig_post_init() + finally: + self.draft_model = actual_draft_model + + if self.num_speculative_tokens == 0: + self.num_speculative_tokens = getattr(self, "num_lookahead_slots", 1) else: self._orig_post_init() @@ -113,6 +134,11 @@ def __post_init__(self): class VllmConfigPatch(ArcticPatch[VllmConfig]): _orig_str = VllmConfig.__str__ + _orig_post_init = VllmConfig.__post_init__ + + from typing import Literal + OldEagleModelTypes = vllm.config.speculative.EagleModelTypes + NewEagleModelTypes = Literal["arctic", "suffix", OldEagleModelTypes] def __str__(self, *args, **kwargs): string = self._orig_str(*args, **kwargs) @@ -121,6 +147,24 @@ def __str__(self, *args, **kwargs): string += f", shift_parallel_threshold={self.parallel_config.shift_parallel_threshold}" return string + def __post_init__(self, *args, **kwargs): + # if self.speculative_config is not None: + # if self.speculative_config.method not in get_args(EagleModelTypes): + # raise ValueError( + # "Currently, async scheduling is only supported " + # "with EAGLE/MTP kind of speculative decoding" + # ) + import sys + from typing import Literal + target_module = sys.modules[VllmConfig.__module__] + original_types = getattr(target_module, "EagleModelTypes") + NewEagleModelTypes = Literal["mlp_speculator", "suffix", original_types] + setattr(target_module, "EagleModelTypes", NewEagleModelTypes) + try: + self._orig_post_init(*args, **kwargs) + finally: + setattr(target_module, "EagleModelTypes", original_types) + class MLPSpeculatorConfigPatch(ArcticPatch[MLPSpeculatorConfig]): @@ -129,3 +173,16 @@ class MLPSpeculatorConfigPatch(ArcticPatch[MLPSpeculatorConfig]): def __init__(self, *args, **kwargs): self.base_model_arch = kwargs.pop("base_model_arch", "") self._orig_init(*args, **kwargs) + + # Inject dummy attributes required by vLLM's ModelArchConfigConvertor + # The convertor tries to calculate head_size = hidden_size // num_attention_heads + if not hasattr(self, "num_attention_heads"): + self.num_attention_heads = 1 + + if not hasattr(self, "hidden_size"): + # Fallback to n_embd if present, otherwise default to a safe dummy value + self.hidden_size = getattr(self, "n_embd", 1024) + + # Ensure hidden_size is an integer to prevent TypeError during division + if hasattr(self, "hidden_size"): + self.hidden_size = int(self.hidden_size) diff --git a/arctic_inference/vllm/model_runner.py b/arctic_inference/vllm/model_runner.py index 733ead1e1..2f2fb5284 100644 --- a/arctic_inference/vllm/model_runner.py +++ b/arctic_inference/vllm/model_runner.py @@ -27,7 +27,7 @@ import vllm.envs as envs from vllm.compilation.counter import compilation_counter from vllm.compilation.monitor import set_cudagraph_capturing_enabled -from vllm.config import CUDAGraphMode, CompilationLevel, VllmConfig +from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.kv_transfer import (get_kv_transfer_group, has_kv_transfer_group) from vllm.distributed.parallel_state import (get_pp_group, get_tp_group, @@ -35,10 +35,10 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.model_executor.model_loader import get_model from vllm.sequence import IntermediateTensors -from vllm.utils import round_up, cdiv -from vllm.v1.attention.backends.utils import CommonAttentionMetadata -from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput) +from vllm.utils.math_utils import round_up, cdiv +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput, + SamplerOutput) from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.rejection_sampler import MAX_SPEC_LEN, RejectionSampler from vllm.v1.spec_decode.metadata import SpecDecodeMetadata @@ -48,6 +48,7 @@ logger, AsyncGPUModelRunnerOutput, ) +from vllm.v1.structured_output.utils import apply_grammar_bitmask if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -98,21 +99,21 @@ def is_shift_parallel_mode() -> bool: class GPUModelRunnerPatch(ArcticPatch[GPUModelRunner]): """ - Extension of vLLM's GPUModelRunner that adds: - - Ulysses sequence parallel + shift-parallel dual-model routing - - Arctic / suffix speculative decoding hooks - - SwiftKV metadata propagation - - Async scheduling correctness (fences + async output wrapper) + Rebased GPUModelRunnerPatch for vLLM v14. """ - _orig_initialize_kv_cache = GPUModelRunner.initialize_kv_cache _orig_capture_cudagraphs = GPUModelRunner._capture_cudagraphs - _orig_prepare_inputs = GPUModelRunner._prepare_inputs _orig_profile_run = GPUModelRunner.profile_run _orig_load_model = GPUModelRunner.load_model _orig_propose_draft_token_ids = GPUModelRunner.propose_draft_token_ids _orig_dummy_run = GPUModelRunner._dummy_run _orig_init = GPUModelRunner.__init__ + _orig_build_attention_metadata = GPUModelRunner._build_attention_metadata + _orig_execute_model = GPUModelRunner.execute_model + _orig_bookkeeping_sync = GPUModelRunner._bookkeeping_sync + _orig_sample_tokens = GPUModelRunner.sample_tokens + _orig_initialize_kv_cache = GPUModelRunner.initialize_kv_cache + # _orig_pad_for_sequence_parallelism = GPUModelRunner._pad_for_sequence_parallelism def __init__( self, @@ -122,7 +123,7 @@ def __init__( if vllm_config.parallel_config.ulysses_sequence_parallel_size > 1: self.use_ulysses = True pass_config = vllm_config.compilation_config.pass_config - if pass_config.enable_sequence_parallelism: + if pass_config.enable_sp: raise ValueError( "Ulysses sequence parallelism is incompatible with native " "sequence parallelism. Set enable_sequence_parallelism " @@ -131,33 +132,62 @@ def __init__( else: self.use_ulysses = False - if (vllm_config.speculative_config is not None and - vllm_config.speculative_config.method in ("arctic", "suffix", "mlp_speculator")): + arctic_methods = ("arctic", "suffix", "mlp_speculator") + is_arctic_spec = (vllm_config.speculative_config is not None and + vllm_config.speculative_config.method in arctic_methods) + + arctic_speculative_config = None + if is_arctic_spec: arctic_speculative_config = vllm_config.speculative_config vllm_config.speculative_config = None - else: - arctic_speculative_config = None self._orig_init(vllm_config, device) + + self._suffix_cache: Optional[SuffixDecodingCache] = None - self._suffix_cache = None - if arctic_speculative_config is not None: + if is_arctic_spec: self.vllm_config.speculative_config = arctic_speculative_config self.speculative_config = arctic_speculative_config + self.num_spec_tokens = getattr(self.speculative_config, + "num_speculative_tokens", 0) + self.uniform_decode_query_len = 1 + self.num_spec_tokens + + if not hasattr(self, "draft_token_ids_cpu") or self.draft_token_ids_cpu is None: + self.draft_token_ids_event = torch.Event() + self.draft_token_ids_copy_stream = torch.cuda.Stream() + self.draft_token_ids_cpu = torch.empty( + (self.max_num_reqs, self.num_spec_tokens), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + + if (self.use_async_scheduling + and self.speculative_config.method in ("arctic", "mlp_speculator", "suffix")): + if not hasattr(self, "valid_sampled_token_count_cpu") or self.valid_sampled_token_count_cpu is None: + self.valid_sampled_token_count_event = torch.Event() + self.valid_sampled_token_count_copy_stream = torch.cuda.Stream() + self.valid_sampled_token_count_cpu = torch.empty( + self.max_num_reqs, + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) if get_pp_group().is_last_rank: if self.speculative_config.method in ("arctic", "mlp_speculator"): self.drafter = ArcticProposer(self.vllm_config) - elif self.speculative_config.method != "suffix": - raise ValueError("Unknown speculative decoding method: " - f"{self.speculative_config.method}") - else: + elif self.speculative_config.method == "suffix": self.drafter = SuffixProposer() - self.rejection_sampler = RejectionSampler() + else: + raise ValueError(f"Unknown speculative decoding method: {self.speculative_config.method}") + + self.rejection_sampler = RejectionSampler(self.sampler) - if (self.speculative_config is not None and - self.speculative_config.enable_suffix_decoding): - if self.speculative_config.method not in ("arctic", "suffix", "mlp_speculator"): + if (self.speculative_config is not None and + getattr(self.speculative_config, "enable_suffix_decoding", False)): + + if self.speculative_config.method not in arctic_methods: raise ValueError( "Suffix decoding is only supported with the 'arctic', " "'mlp_speculator' or 'suffix' spec decoding methods." @@ -167,39 +197,181 @@ def __init__( max_tree_depth=spec_cfg.suffix_cache_max_depth, max_cached_requests=spec_cfg.suffix_cache_max_requests ) - + + # Async suffix decoding infrastructure: a dedicated CUDA stream and + # pinned buffer for copying sampled token IDs to CPU *without* + # serialising behind Arctic GPU drafting work on the default stream. + if self._suffix_cache is not None and self.use_async_scheduling: + self.suffix_copy_stream = torch.cuda.Stream() + self.suffix_copy_done_event = torch.Event() + max_gen_len = 1 + self.num_spec_tokens + self.suffix_sampled_ids_pinned = torch.empty( + (self.max_num_reqs, max_gen_len), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + # Pinned buffer for suffix merge results. Using pinned memory + # for H2C copies avoids the implicit default-stream + # synchronisation that cudaMemcpyAsync performs with pageable + # (non-pinned) source memory. Without this, the merge step + # blocks the CPU until ALL pending GPU work (including Arctic + # drafting) completes, destroying the async overlap. + self._suffix_merge_pinned = torch.zeros( + (self.max_num_reqs, self.num_spec_tokens), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + + # Pre-allocated GPU buffer for the merged draft tensor. + # Avoids per-step F.pad / torch.zeros allocations in the async + # Arctic drafting path (propose_draft_token_ids + suffix merge). + # Shape: [max_num_reqs, num_spec_tokens], int64 (matches + # draft_token_ids_cpu for zero-cost _copy_draft_token_ids_to_cpu). + if (self.speculative_config is not None + and self.use_async_scheduling + and self.speculative_config.method + in ("arctic", "mlp_speculator")): + self._draft_merged_gpu = torch.zeros( + (self.max_num_reqs, self.num_spec_tokens), + dtype=torch.int64, device=self.device, + ) + + # Pre-allocated pinned index buffer for suffix merge overlay. + # Avoids per-step torch.tensor(...).pin_memory() allocations. + if self._suffix_cache is not None and self.use_async_scheduling: + self._suffix_index_pinned = torch.empty( + self.max_num_reqs, dtype=torch.long, + device="cpu", pin_memory=self.pin_memory, + ) + + # Per-request response tokens for suffix pattern building in async + # mode. In async scheduling, _bookkeeping_sync writes -1 placeholders + # to token_ids_cpu instead of real values, corrupting the pattern that + # propose_suffix_draft_token_ids reads. We keep a clean copy here. + self._suffix_response_tokens: dict[str, list[int]] = {} + + # Actual draft lengths per request from the previous step. Used + # by execute_model to trim the scheduler's spec token allocation + # down to the real draft width, and communicated back to the + # scheduler (via scheduler_output._actual_draft_lens) so + # _update_after_schedule can set dynamic placeholder counts. + self._prev_actual_draft_lens: dict[str, int] = {} + + # Backup-token buffer used by suffix-only async rejection sampling. + # The arctic proposer has its own buffer; this one covers the case + # where no arctic drafter is present. + self._suffix_backup_tokens_gpu: Optional[torch.Tensor] = None + if (self._suffix_cache is not None + and self.use_async_scheduling + and self.speculative_config.method not in ("arctic", + "mlp_speculator")): + self._suffix_backup_tokens_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device, + ) + + def _suffix_only_rejection_sample( + self, + sampled_token_ids: torch.Tensor, + common_attn_metadata: "CommonAttentionMetadata", + ) -> None: + """Rejection-sample accepted tokens for suffix-only async scheduling. + + EAGLE / arctic do this inside propose_draft_token_ids via + prepare_next_token_ids_padded. For suffix-only there is no + drafter with that method, so we call the same Triton kernel + directly and feed the results into _copy_valid_sampled_token_count. + """ + from vllm.triton_utils import triton + from vllm.v1.spec_decode.utils import ( + eagle_prepare_next_token_padded_kernel, + ) + + num_reqs = self.input_batch.num_reqs + batch_size, num_tokens = sampled_token_ids.shape + device = sampled_token_ids.device + + # Compute backup tokens (last accepted token per request) on CPU, + # then copy to GPU in one shot to avoid per-element synchronisation. + backup = self._suffix_backup_tokens_gpu + assert backup is not None + backup_np = np.empty(num_reqs, dtype=np.int32) + for i in range(num_reqs): + req_id = self.input_batch.req_ids[i] + seq_len = int(common_attn_metadata.seq_lens_cpu[i].item()) + backup_np[i] = self.requests[req_id].get_token_id(seq_len) + # Copy directly from CPU numpy-backed tensor to GPU; avoids + # creating an intermediate GPU tensor via .to(device). + backup[:num_reqs].copy_( + torch.from_numpy(backup_np), non_blocking=True, + ) + + next_token_ids = torch.empty(batch_size, dtype=torch.int32, + device=device) + valid_counts = torch.empty(batch_size, dtype=torch.int32, + device=device) + + BLOCK_SIZE_TOKENS = triton.next_power_of_2(num_tokens) + eagle_prepare_next_token_padded_kernel[(batch_size,)]( + sampled_token_ids, + self.discard_request_mask.gpu, + backup, + next_token_ids, + valid_counts, + self.model_config.get_vocab_size(), + num_tokens, + batch_size, + sampled_token_ids.stride(0), + BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS, + ) + + self._copy_valid_sampled_token_count(next_token_ids, valid_counts) + + def _build_attention_metadata(self, *args, **kwargs): + attn_metadata, spec_decode_common_attn_metadata = \ + self._orig_build_attention_metadata(*args, **kwargs) + + logits_indices = kwargs.get("logits_indices", None) + if logits_indices is not None: + if isinstance(attn_metadata, list): + for ub in attn_metadata: + for meta in ub.values(): + meta.swiftkv_logits_indices = logits_indices + else: + for meta in attn_metadata.values(): + meta.swiftkv_logits_indices = logits_indices + + return attn_metadata, spec_decode_common_attn_metadata + + # set padding for SP here + def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: + + sp_size = self.parallel_config.ulysses_sequence_parallel_size + num_input_tokens = round_up(num_scheduled_tokens, sp_size) + + #if torch.distributed.get_rank() == 0: + # print(f"padding num_scheduled_tokens {num_scheduled_tokens} -> num_input_tokens {num_input_tokens}") + + return num_input_tokens + def profile_run(self) -> None: self._orig_profile_run() if getattr(self, "shift_model", None) is not None: orig_model, self.model = self.model, self.shift_model + cc = self.vllm_config.compilation_config + base_ctx = cc.static_forward_context + shift_ctx = getattr(self, 'shift_forward_context', None) try: + if shift_ctx is not None: + cc.static_forward_context = shift_ctx with set_shift_parallel_mode(True): self._dummy_run(self.max_num_tokens, is_profile=True) finally: self.model = orig_model + cc.static_forward_context = base_ctx - def _prepare_inputs(self, *args, **kwargs): - """ - Forward to upstream _prepare_inputs, then add SwiftKV-specific metadata. - """ - (attn_metadata, logits_indices, spec_decode_metadata, - num_scheduled_tokens_np, spec_decode_common_attn_metadata, - max_query_len, ubatch_slices, num_tokens_after_padding) = ( - self._orig_prepare_inputs(*args, **kwargs) - ) - - if isinstance(attn_metadata, list): - for ubatch_attn_metadata in attn_metadata: - for meta in ubatch_attn_metadata.values(): - meta.swiftkv_logits_indices = logits_indices - else: - for meta in attn_metadata.values(): - meta.swiftkv_logits_indices = logits_indices - - return (attn_metadata, logits_indices, spec_decode_metadata, - num_scheduled_tokens_np, spec_decode_common_attn_metadata, - max_query_len, ubatch_slices, num_tokens_after_padding) def monkeypatch_forward(self: GPUModelRunner): """ @@ -226,7 +398,7 @@ def ulysses_forward(*args, **kwargs): output = model_forward(*args, **kwargs) if output.size(0) == N_ulysses: - model_output = torch.empty((N, self.hidden_size), + model_output = torch.empty((N, output.shape[1]), dtype=output.dtype, device=output.device) torch.distributed.all_gather_into_tensor(model_output, @@ -243,7 +415,7 @@ def ulysses_forward(*args, **kwargs): def _dummy_run( self, num_tokens: int, - cudagraph_runtime_mode: Optional[CUDAGraphMode] = None, + cudagraph_runtime_mode: CUDAGraphMode | None = None, force_attention: bool = False, uniform_decode: bool = False, allow_microbatching: bool = True, @@ -251,71 +423,25 @@ def _dummy_run( is_profile: bool = False, create_mixed_batch: bool = False, remove_lora: bool = True, + activate_lora: bool = False, + is_graph_capturing: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Run a dummy forward pass to warm up/profile run or capture the - CUDA graph for the model. - - Args: - num_tokens: Number of tokens to run the dummy forward pass. - cudagraph_runtime_mode: used to control the behavior. - - if not set will determine the cudagraph mode based on using - the self.cudagraph_dispatcher. - - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run - - CUDAGraphMode.PIECEWISE: Piecewise cudagraph. - - CUDAGraphMode.FULL: Full cudagraph, attention metadata is - needed. - force_attention: If True, always create attention metadata. Used to - warm up attention backend when mode is NONE. - uniform_decode: If True, the batch is a uniform decode batch. - skip_eplb: If True, skip EPLB state update. - is_profile: If True, this is a profile run. - create_mixed_batch: If True, create a mixed batch with both decode - (1 token) and prefill (multiple tokens) requests. - remove_lora: If False, dummy LoRAs are not destroyed after the run - """ - assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in { - CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL - } - - # If cudagraph_mode.decode_mode() == FULL and - # cudagraph_mode.separate_routine(). This means that we are using - # different graphs and/or modes for mixed prefill-decode batches vs. - # uniform decode batches. A uniform decode batch means that all - # requests have identical query length, except a potential virtual - # request (shorter) in the batch account for padding. - # Uniform decode batch could either be common pure decode, where - # max_query_len == 1, or speculative decode, where - # max_query_len == 1 + num_spec_decode_tokens. - - # When setting max_query_len = 1, we switch to and capture the optimized - # routine of FA2 for pure decode, i.e., Flashdecode + an optimization - # for GQA/MQA. - max_query_len = self.uniform_decode_query_len if uniform_decode else \ - num_tokens - - # 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 + + from vllm.v1.worker.gpu_model_runner import supports_mm_encoder_only + if supports_mm_encoder_only(self.model): + return torch.tensor([]), torch.tensor([]) + + max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens max_num_reqs = self.scheduler_config.max_num_seqs + if create_mixed_batch: - assert not uniform_decode - # Create mixed batch: - # first half decode tokens, second half one prefill - num_decode_tokens = num_tokens // 2 + num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2) num_prefill_tokens = num_tokens - num_decode_tokens num_reqs = num_decode_tokens + 1 - - # Create decode requests (1 token each) followed by prefill request - num_scheduled_tokens_list = [1] * num_decode_tokens + [ - num_prefill_tokens - ] - # Note: Overriding max_query_len to be the prefill tokens + num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens] max_query_len = num_prefill_tokens elif uniform_decode: - assert not create_mixed_batch - num_reqs = cdiv(num_tokens, max_query_len) + num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len)) num_scheduled_tokens_list = [max_query_len] * num_reqs if num_tokens % max_query_len != 0: num_scheduled_tokens_list[-1] = num_tokens % max_query_len @@ -325,226 +451,159 @@ def _dummy_run( num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs num_scheduled_tokens_list[-1] += num_tokens % num_reqs - assert sum(num_scheduled_tokens_list) == num_tokens - assert len(num_scheduled_tokens_list) == num_reqs - num_scheduled_tokens = np.array(num_scheduled_tokens_list, - dtype=np.int32) - total_num_scheduled_tokens = int(num_scheduled_tokens.sum()) - - num_scheduled_tokens_for_logits = num_scheduled_tokens[num_scheduled_tokens > 0] - if num_scheduled_tokens_for_logits.size == 0: - # Handle edge case: all requests have 0 tokens. - logits_indices = torch.empty(0, dtype=torch.long, device=self.device) - # Use numpy for the final return indexing - logit_indices_np = np.array([], dtype=int) - else: - logits_indices_cpu = np.cumsum(num_scheduled_tokens_for_logits) - 1 - logits_indices = torch.from_numpy(logits_indices_cpu).to(self.device) - # Use numpy for the final return indexing - logit_indices_np = np.cumsum(num_scheduled_tokens) - 1 - - ubatch_slices = None - num_tokens_after_padding = None - - # We currently only microbatch if the number of tokens is - # over a certain threshold. - if self.parallel_config.enable_dbo and allow_microbatching: - ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split( - num_scheduled_tokens, - total_num_scheduled_tokens, - total_num_scheduled_tokens, - uniform_decode=uniform_decode, - vllm_config=self.vllm_config, + num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_tokens_unpadded = int(num_scheduled_tokens.sum()) + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + + if torch.distributed.get_rank() == 0: + print(f"num_tokens_unpadded: {num_tokens_unpadded}, num_reqs: {num_reqs}") + + _cg_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = ( + self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens, + max_num_scheduled_tokens=max_query_len, + use_cascade_attn=False, + allow_microbatching=allow_microbatching, + force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE), + force_uniform_decode=uniform_decode, + force_has_lora=activate_lora, ) - # Currently when DBO is enabled `ubatch_split` returns - # the num_tokens_after_padding for a single ubatch, but we have 2 - # TODO(sage,lucas): this is cruft that should be addressed in the - # padding refactor. - if ubatch_num_tokens_after_padding is not None: - num_tokens_after_padding = ubatch_num_tokens_after_padding * 2 - - # If we failed to microbatch, currently need to resynchronize - # TODO(lucas,sage): we should be able to avoid this second sync by - # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into - # a single `coordinate_batch_across_dp` function. - if num_tokens_after_padding is None: - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens_after_padding = num_tokens + num_pad - else: - num_tokens_across_dp = num_tokens_after_padding - num_tokens_after_padding = int(num_tokens_after_padding[0].item()) + ) - attn_metadata: Optional[PerLayerAttnMetadata] = None + if cudagraph_runtime_mode is None: + cudagraph_runtime_mode = _cg_mode + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + + from vllm.v1.worker.gpu_model_runner import maybe_create_ubatch_slices + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens, + num_tokens_padded, + num_reqs_padded, + self.vllm_config.parallel_config.num_ubatches, + ) - # If force_attention is True, we always capture attention. Otherwise, - # it only happens for cudagraph_runtime_mode=FULL. - if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: - attn_metadata = {} - if ubatch_slices is not None: - attn_metadata = [dict() for _ in range(len(ubatch_slices))] + logits_indices_cpu = np.cumsum(num_scheduled_tokens) - 1 + logits_indices = torch.from_numpy(logits_indices_cpu).to(self.device) + attn_metadata = None + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: if create_mixed_batch: - # In the mixed batch mode (used for FI warmup), we use - # shorter sequence lengths to run faster. - # TODO(luka) better system for describing dummy batches - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + seq_lens_list = [1] * num_decode_tokens + [num_prefill_tokens + 1] else: - seq_lens = max_query_len - self.seq_lens.np[:num_reqs] = seq_lens + seq_lens_list = [max_query_len] * num_reqs # simplified + + self.seq_lens.np[:num_reqs] = seq_lens_list self.seq_lens.np[num_reqs:] = 0 self.seq_lens.copy_to_gpu() - cum_num_tokens, _ = self._get_cumsum_and_arange( - num_scheduled_tokens) - self.query_start_loc.np[1:num_reqs + 1] = cum_num_tokens + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() - for kv_cache_group_id, kv_cache_group_spec in enumerate( - self.kv_cache_config.kv_cache_groups): - common_attn_metadata = CommonAttentionMetadata( - query_start_loc=self.query_start_loc.gpu[:num_reqs + 1], - query_start_loc_cpu=self.query_start_loc.cpu[:num_reqs + - 1], - seq_lens=self.seq_lens.gpu[:num_reqs], - seq_lens_cpu=self.seq_lens.cpu[:num_reqs], - num_computed_tokens_cpu=self.input_batch. - num_computed_tokens_cpu_tensor[:num_reqs], - num_reqs=num_reqs, - num_actual_tokens=num_tokens, - max_query_len=max_query_len, - max_seq_len=self.max_model_len, - block_table_tensor=self.input_batch. - block_table[kv_cache_group_id].get_device_tensor(num_reqs), - slot_mapping=self.input_batch.block_table[ - kv_cache_group_id].slot_mapping.gpu[:num_tokens], - causal=True) - for attn_group in self.attn_groups[kv_cache_group_id]: - if ubatch_slices is not None: - common_attn_metadata_list = split_attn_metadata( - ubatch_slices, common_attn_metadata) - for ubid, common_attn_metadata in enumerate( - common_attn_metadata_list): - assert common_attn_metadata.max_query_len == 1 - attn_metadata_i = (attn_group\ - .get_metadata_builder(ubatch_id=ubid)\ - .build_for_cudagraph_capture(common_attn_metadata)) - for layer_name in attn_group.layer_names: - assert type(attn_metadata) is list - attn_metadata[ubid][ - layer_name] = attn_metadata_i - else: - assert type(attn_metadata) is dict - attn_metadata_i = attn_group.get_metadata_builder()\ - .build_for_cudagraph_capture(common_attn_metadata) - for layer_name in attn_group.layer_names: - attn_metadata[layer_name] = attn_metadata_i + pad_attn = (cudagraph_runtime_mode == CUDAGraphMode.FULL) + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + ) if attn_metadata is not None: if isinstance(attn_metadata, list): - for ubatch_attn_metadata in attn_metadata: - for meta in ubatch_attn_metadata.values(): + for ub_meta in attn_metadata: + for meta in ub_meta.values(): meta.swiftkv_logits_indices = logits_indices else: for meta in attn_metadata.values(): meta.swiftkv_logits_indices = logits_indices - with self.maybe_dummy_run_with_lora(self.lora_config, - num_scheduled_tokens, remove_lora): - model_kwargs = self._init_model_kwargs(num_tokens) - if (self.supports_mm_inputs - and not self.model_config.is_encoder_decoder): - input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens] - model_kwargs = { - **model_kwargs, - **self._dummy_mm_kwargs(num_reqs), - } + with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, + num_sampled_tokens, activate_lora, remove_lora): + + 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) + model_kwargs.update(self._dummy_mm_kwargs(num_reqs)) elif self.enable_prompt_embeds: input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens] - model_kwargs = self._init_model_kwargs(num_tokens) + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: - input_ids = self.input_ids.gpu[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None + positions = self.positions.gpu[:num_tokens_padded] if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_tokens] - else: - positions = self.positions.gpu[:num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens_padded] - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: + intermediate_tensors = None + if not get_pp_group().is_first_rank: if self.intermediate_tensors is None: - self.intermediate_tensors = ( - self.model.make_empty_intermediate_tensors( - batch_size=self.max_num_tokens, - dtype=self.model_config.dtype, - device=self.device)) - - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_tokens, None, False) - - # filter out the valid batch descriptor - _cg_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch( - BatchDescriptor(num_tokens=num_tokens_after_padding, - uniform_decode=uniform_decode)) \ - if not is_profile else (CUDAGraphMode.NONE, None) - if cudagraph_runtime_mode is not None: - # we allow forcing NONE when the dispatcher disagrees to support - # warm ups for cudagraph capture - assert cudagraph_runtime_mode == CUDAGraphMode.NONE or \ - cudagraph_runtime_mode == _cg_mode, ( - f"Cudagraph runtime mode mismatch at dummy_run. " - f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}.") - else: - cudagraph_runtime_mode = _cg_mode + self.intermediate_tensors = self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, dtype=self.model_config.dtype, device=self.device) + intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False) - if ubatch_slices is not None: - # Adjust values to reflect a single ubatch. - # TODO(sage,lucas): this is cruft that should be addressed in - # the padding refactor. - num_tokens_after_padding = ubatch_slices[0].num_tokens + target_num_tokens = num_tokens_padded + if ubatch_slices_padded is not None: + target_num_tokens = ubatch_slices_padded[0].num_tokens if num_tokens_across_dp is not None: - num_tokens_across_dp[:] = num_tokens_after_padding - - with self.maybe_randomize_inputs(input_ids), set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens_after_padding, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, - ubatch_slices=ubatch_slices): - outputs = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - ) + num_tokens_across_dp[:] = target_num_tokens - if self.use_aux_hidden_state_outputs: - hidden_states, _ = outputs - else: - hidden_states = outputs + with self.maybe_randomize_inputs(input_ids, inputs_embeds), set_forward_context( + attn_metadata, self.vllm_config, num_tokens=target_num_tokens, + num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded): + + outputs = self.model(input_ids=input_ids, positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, **model_kwargs) + + hidden_states = outputs[0] if self.use_aux_hidden_state_outputs else outputs if self.speculative_config and self.speculative_config.use_eagle(): - assert isinstance(self.drafter, EagleProposer) - self.drafter.dummy_run(num_tokens) - - # This is necessary to avoid blocking DP. - # For dummy runs, we typically skip EPLB since we don't have any real - # requests to process. - # However, in DP settings, there may be cases when some DP ranks do - # not have any requests to process, so they're executing dummy batches. - # In such cases, we still have to trigger EPLB to make sure - # ranks execute the rearrangement in synchronization. + self.drafter.dummy_run(num_tokens, use_cudagraphs=False, is_graph_capturing=is_graph_capturing) + if not skip_eplb: self.eplb_step(is_dummy=True, is_profile=is_profile) - logit_indices = np.cumsum(num_scheduled_tokens) - 1 - return hidden_states, hidden_states[logit_indices] + return hidden_states, hidden_states[logits_indices] + + # ------------------------------------------------------------------ + # _sample: inline the base GPUModelRunner._sample logic here because + # the class is monkey-patched at runtime, making both super() and + # GPUModelRunner._sample(self, ...) resolve back to this method. + # ------------------------------------------------------------------ + def _sample( + self, + logits: torch.Tensor | None, + spec_decode_metadata: SpecDecodeMetadata | None, + ) -> SamplerOutput: + sampling_metadata = self.input_batch.sampling_metadata + self.input_batch.update_async_output_token_ids() + if spec_decode_metadata is None: + return self.sampler( + logits=logits, sampling_metadata=sampling_metadata) + + if (self.use_async_scheduling + and self._draft_token_req_ids is not None): + draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu() + self.input_batch.update_async_spec_token_ids( + draft_token_ids_cpu) + + sampler_output = self.rejection_sampler( + spec_decode_metadata, + None, # draft_probs + logits, + sampling_metadata, + ) + self._update_states_after_model_execute( + sampler_output.sampled_token_ids) + return sampler_output @torch.inference_mode() def execute_model( @@ -554,299 +613,349 @@ def execute_model( ) -> Union[ ModelRunnerOutput, AsyncGPUModelRunnerOutput, IntermediateTensors ]: - """ - This override preserves the upstream execution order but: - - wraps input prep with the async fence (fixes hangs), - - returns AsyncGPUModelRunnerOutput when async scheduling is enabled, - - routes forward() to shift_model when under the threshold, - - computes logits with the same model used for forward(). - """ - with self.synchronize_input_prep(): - self._update_states(scheduler_output) - - if not scheduler_output.total_num_scheduled_tokens: - if not has_kv_transfer_group(): - return EMPTY_MODEL_RUNNER_OUTPUT - return self.kv_connector_no_forward(scheduler_output, - self.vllm_config) - - (attn_metadata, logits_indices, spec_decode_metadata, - num_scheduled_tokens_np, spec_decode_common_attn_metadata, - max_query_len, ubatch_slices, num_tokens_after_padding) = ( - self._prepare_inputs(scheduler_output) - ) - - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + num_scheduled_tokens = getattr(scheduler_output, "total_num_scheduled_tokens", None) + if num_scheduled_tokens is None: + try: + num_scheduled_tokens = int( + sum(scheduler_output.num_scheduled_tokens.values()) + ) + except Exception: + num_scheduled_tokens = 0 use_shift_model = ( getattr(self, "use_ulysses", False) and getattr(self, "shift_model", None) is not None - and num_scheduled_tokens <= self.shift_parallel_threshold + and num_scheduled_tokens <= int(getattr(self, "shift_parallel_threshold", 0)) ) - if self.use_ulysses and not use_shift_model: - sp_size = self.parallel_config.ulysses_sequence_parallel_size - num_input_tokens = round_up(num_scheduled_tokens, sp_size) - if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and num_input_tokens // sp_size <= self.cudagraph_batch_sizes[-1]): - num_input_tokens = ( - self.vllm_config.pad_for_cudagraph(num_input_tokens // sp_size) * sp_size - ) - elif (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]): - num_input_tokens = self.vllm_config.pad_for_cudagraph( - num_scheduled_tokens + if not use_shift_model: + return self._orig_execute_model(scheduler_output, intermediate_tensors) + + orig_model = self.model + cc = self.vllm_config.compilation_config + base_ctx = cc.static_forward_context + shift_ctx = getattr(self, 'shift_forward_context', None) + try: + self.model = self.shift_model + if shift_ctx is not None: + cc.static_forward_context = shift_ctx + with set_shift_parallel_mode(True), \ + self._use_shift_cudagraph_tables(): + result = self._orig_execute_model(scheduler_output, intermediate_tensors) + finally: + self.model = orig_model + cc.static_forward_context = base_ctx + return result + + @torch.inference_mode + def sample_tokens(self, grammar_output): + """Wrapper around base sample_tokens for arctic async spec decode. + + Saves execute_model_state before the base clears it, then handles + the 'not-fits-in-drafter' case that the base only handles for Eagle. + """ + _arctic_saved_state = None + if (self.execute_model_state is not None + and self.speculative_config is not None + and self.speculative_config.method + in ("arctic", "mlp_speculator", "suffix") + and self.use_async_scheduling): + _arctic_saved_state = ( + self.execute_model_state.scheduler_output, + self.execute_model_state.spec_decode_common_attn_metadata, ) - else: - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - if (self.compilation_config.pass_config.enable_sequence_parallelism - and tp_size > 1): - num_input_tokens = round_up(num_scheduled_tokens, tp_size) - else: - num_input_tokens = num_scheduled_tokens - num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) - num_input_tokens += num_pad + result = self._orig_sample_tokens(grammar_output) - if self.supports_mm_inputs: - self._execute_mm_encoder(scheduler_output) - mm_embeds = self._gather_mm_embeddings(scheduler_output) - else: - mm_embeds = [] - - if self.supports_mm_inputs and get_pp_group().is_first_rank: - inputs_embeds_scheduled = self.model.get_input_embeddings( - input_ids=self.input_ids.gpu[:num_scheduled_tokens], - multimodal_embeddings=mm_embeds or None, - ) - self.inputs_embeds.gpu[:num_scheduled_tokens].copy_( - inputs_embeds_scheduled - ) - input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = { - **self._init_model_kwargs(num_scheduled_tokens), - **self._extract_mm_kwargs(scheduler_output), - } - else: - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = None - model_kwargs = self._init_model_kwargs(num_input_tokens) + # If _arctic_async_sampled_tensor was stashed by _bookkeeping_sync + # but never consumed by propose_draft_token_ids, this is the + # not-fits-in-drafter case. Mirror Eagle's handling: call + # _copy_valid_sampled_token_count and set draft tokens to zeros. + stashed = getattr(self, '_arctic_async_sampled_tensor', None) + if stashed is not None: + del self._arctic_async_sampled_tensor + if _arctic_saved_state is not None: + scheduler_output, common_attn_meta = _arctic_saved_state + self._arctic_handle_not_fits( + stashed, scheduler_output, common_attn_meta) - if (self.model_config.is_encoder_decoder - and scheduler_output.scheduled_encoder_inputs): - encoder_inputs = self._extract_encoder_inputs(scheduler_output) - model_kwargs.update(encoder_inputs) + return result - if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_input_tokens] - else: - positions = self.positions.gpu[:num_input_tokens] - - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True + def _arctic_handle_not_fits( + self, + sampled_token_ids: torch.Tensor, + scheduler_output: "SchedulerOutput", + common_attn_metadata, + ) -> None: + """Mirror Eagle's not-fits-in-drafter path for arctic async. + + When the input is too long for the drafter but spec decode is + active, Eagle still calls prepare_next_token_ids_padded / + _copy_valid_sampled_token_count and sets draft tokens to zeros. + Without this, _get_valid_sampled_token_count returns stale counts + and _prepare_input_ids scatters -1 placeholders into the + embedding layer. + """ + if (hasattr(self, 'drafter') + and hasattr(self.drafter, 'prepare_next_token_ids_padded') + and common_attn_metadata is not None): + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + common_attn_metadata, + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count) + else: + # Fallback for drafters without prepare_next_token_ids_padded + # (e.g. suffix-only). Compute valid counts with PyTorch ops. + mask = sampled_token_ids != -1 + valid_counts = mask.sum(dim=1) + batch_size = sampled_token_ids.shape[0] + col_indices = torch.arange( + sampled_token_ids.shape[1], + device=sampled_token_ids.device, + ).unsqueeze(0).expand_as(sampled_token_ids) + last_valid_col = ( + col_indices.masked_fill(~mask, -1).max(dim=1).values) + last_valid_col = last_valid_col.clamp(min=0) + next_token_ids = sampled_token_ids[ + torch.arange(batch_size, + device=sampled_token_ids.device), + last_valid_col, + ] + self._copy_valid_sampled_token_count( + next_token_ids, valid_counts) - uniform_decode = ( - max_query_len == self.uniform_decode_query_len - and num_scheduled_tokens == self.input_batch.num_reqs * max_query_len - ) - batch_descriptor = BatchDescriptor( - num_tokens=num_input_tokens, uniform_decode=uniform_decode - ) - cudagraph_runtime_mode, batch_descriptor = \ - self.cudagraph_dispatcher.dispatch(batch_descriptor) - - if ubatch_slices is not None: - num_input_tokens = ubatch_slices[0].num_tokens - - with set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_input_tokens, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, - ubatch_slices=ubatch_slices, - ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output: - model = self.shift_model if use_shift_model else self.model - with set_shift_parallel_mode(use_shift_model): - model_output = model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - ) + # Zero draft tokens -- same as Eagle's not-fits path. + self._draft_token_ids = torch.zeros( + 1, device=self.device, dtype=torch.int32, + ).expand(len(self.input_batch.req_ids), self.num_spec_tokens) + self._copy_draft_token_ids_to_cpu( + scheduler_output, zeros_only=True) - if self.use_aux_hidden_state_outputs: - hidden_states, aux_hidden_states = model_output - else: - hidden_states = model_output - aux_hidden_states = None + def _bookkeeping_sync( + self, + scheduler_output: "SchedulerOutput", + sampler_output: SamplerOutput, + logits: torch.Tensor | None, + hidden_states: torch.Tensor, + num_scheduled_tokens: int, + spec_decode_metadata: SpecDecodeMetadata | None, + ): + """Wrap base _bookkeeping_sync to handle arctic async spec decode. + + In the base vLLM code, only Eagle-style drafters run *before* + bookkeeping (setting prev_sampled_token_ids via + _copy_valid_sampled_token_count). Arctic/suffix drafting runs + *after* bookkeeping, so prev_sampled_token_ids is still None + when bookkeeping checks ``assert sampled_token_ids.shape[-1] == 1``. + + We fix this by: + 1. Saving the GPU sampled tensor for propose_draft_token_ids. + 2. Setting prev_sampled_token_ids to a placeholder so the + assertion is skipped. The real value will be written by + _copy_valid_sampled_token_count inside propose_draft_token_ids + (fits case) or sample_tokens (not-fits case). + """ + sampled_token_ids = sampler_output.sampled_token_ids + if (self.use_async_scheduling + and self.speculative_config is not None + and self.speculative_config.method + in ("arctic", "mlp_speculator", "suffix") + and spec_decode_metadata is not None + and sampled_token_ids.shape[-1] > 1 + and self.input_batch.prev_sampled_token_ids is None): + # Stash the full GPU tensor so propose_draft_token_ids can + # pick it up later (it normally only receives an empty list + # in the post-bookkeeping path). + self._arctic_async_sampled_tensor = sampled_token_ids + # Placeholder: first column only (bonus token per request). + # Prevents the assertion from firing; the correct value will + # be overwritten by _copy_valid_sampled_token_count shortly. + self.input_batch.prev_sampled_token_ids = ( + sampled_token_ids[:, :1]) + + return self._orig_bookkeeping_sync( + scheduler_output, sampler_output, logits, hidden_states, + num_scheduled_tokens, spec_decode_metadata) - broadcast_pp_output = ( - self.parallel_config.distributed_executor_backend == "external_launcher" - and len(get_pp_group().ranks) > 0 + def propose_draft_token_ids( + self, + scheduler_output: "SchedulerOutput", + sampled_token_ids: torch.Tensor | list[list[int]], + sampling_metadata: SamplingMetadata, + hidden_states: torch.Tensor, + sample_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + spec_decode_metadata: SpecDecodeMetadata | None, + common_attn_metadata: CommonAttentionMetadata, + ) -> list[list[int]] | torch.Tensor: + # In async mode, the base vLLM dispatches arctic to the + # post-bookkeeping path which passes valid_sampled_token_ids + # (an empty list for async). Recover the stashed GPU tensor + # so the fast async drafting path below can activate. + if (isinstance(sampled_token_ids, list) + and len(sampled_token_ids) == 0 + and hasattr(self, '_arctic_async_sampled_tensor')): + sampled_token_ids = self._arctic_async_sampled_tensor + del self._arctic_async_sampled_tensor + + # Compute the maximum number of requests to draft for. + # When disable_by_batch_size is set and the batch exceeds it, + # we still draft for the first N requests instead of disabling + # entirely. This avoids the stale-data crash that occurs when + # drafting is fully disabled one step and re-enabled the next. + batch_size = len(self.input_batch.req_ids) + draft_limit = batch_size # default: draft for all + if ( + self.speculative_config + and self.speculative_config.disable_by_batch_size + and batch_size > self.speculative_config.disable_by_batch_size + ): + draft_limit = self.speculative_config.disable_by_batch_size + + use_async_path = ( + self.speculative_config.method in ("arctic", "mlp_speculator") + and isinstance(sampled_token_ids, torch.Tensor) + and self.use_async_scheduling + and common_attn_metadata is not None ) - if not get_pp_group().is_last_rank: - assert isinstance(hidden_states, IntermediateTensors) - if not broadcast_pp_output: - hidden_states.kv_connector_output = kv_connector_output - return hidden_states - get_pp_group().send_tensor_dict( - hidden_states.tensors, - all_gather_group=get_tp_group(), - ) - logits = None - else: - if self.input_batch.pooling_params: - output = self._pool( - hidden_states, num_scheduled_tokens, num_scheduled_tokens_np + if use_async_path: + assert isinstance(sampled_token_ids, torch.Tensor) + + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + common_attn_metadata, + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, ) - output.kv_connector_output = kv_connector_output - return output - - sample_hidden_states = hidden_states[logits_indices] - with set_shift_parallel_mode(False): - logits = self.model.compute_logits(sample_hidden_states) - - if broadcast_pp_output: - model_output_broadcast_data = {} - if logits is not None: - model_output_broadcast_data["logits"] = logits.contiguous() - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( - model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 ) - assert model_output_broadcast_data is not None - logits = model_output_broadcast_data["logits"] - - if scheduler_output.grammar_bitmask is not None: - from vllm.v1.structured_output.utils import apply_grammar_bitmask - apply_grammar_bitmask(scheduler_output, self.input_batch, - logits, self.device) - - with record_function_or_nullcontext("Sample"): - sampler_output = self._sample(logits, spec_decode_metadata) - - with record_function_or_nullcontext("Bookkeep"): - ( - num_nans_in_logits, - logprobs_lists, - valid_sampled_token_ids, - prompt_logprobs_dict, - req_ids_output_copy, - req_id_to_index_output_copy, - invalid_req_indices, - ) = self._bookkeeping_sync( - scheduler_output, sampler_output, logits, hidden_states, num_scheduled_tokens + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count ) - if self._suffix_cache is not None: - self._update_suffix_cache(valid_sampled_token_ids) + target_hidden_states = self.drafter.prepare_hidden_states( + sample_hidden_states=sample_hidden_states, + sampled_token_ids=sampled_token_ids, + spec_decode_metadata=spec_decode_metadata, + ) - sampling_metadata = self.input_batch.sampling_metadata + # Only draft for the first draft_limit requests. + raw_draft = self.drafter.propose( + context_token_ids=next_token_ids[:draft_limit], + previous_hidden_states=target_hidden_states[:draft_limit], + num_predict_tokens=self.drafter.model.n_predict, + ) - if not self.speculative_config: - self._draft_token_ids = None + # Use pre-allocated GPU buffer when available. This avoids + # per-step F.pad + torch.zeros allocations. The buffer is + # [max_num_reqs, num_spec_tokens] so a single zero_() + + # copy_() handles both width and batch padding in one shot. + merged_buf = getattr(self, '_draft_merged_gpu', None) + if merged_buf is not None: + draft = merged_buf[:batch_size] + draft.zero_() + rd_rows, rd_cols = raw_draft.shape + draft[:rd_rows, :rd_cols].copy_(raw_draft) + else: + draft = raw_draft + if draft.shape[1] < self.num_spec_tokens: + draft = torch.nn.functional.pad( + draft, + (0, self.num_spec_tokens - draft.shape[1]), + value=0, + ) + if draft_limit < batch_size: + full_draft = torch.zeros( + batch_size, draft.shape[1], + dtype=draft.dtype, device=draft.device, + ) + full_draft[:draft_limit] = draft + draft = full_draft + return draft + + if isinstance(sampled_token_ids, torch.Tensor): + vocab_size = self.model_config.get_vocab_size() + sampled_token_ids_list = [ + [t for t in seq if t != -1 and t < vocab_size] + for seq in sampled_token_ids.tolist() + ] + sampled_token_ids_tensor = sampled_token_ids else: - assert spec_decode_common_attn_metadata is not None - self._draft_token_ids = self.propose_draft_token_ids( - scheduler_output, - valid_sampled_token_ids, - sampler_output.sampled_token_ids, - sampling_metadata, - hidden_states, - hidden_states[logits_indices], - aux_hidden_states, - spec_decode_metadata, - spec_decode_common_attn_metadata, - ) + sampled_token_ids_list = sampled_token_ids + sampled_token_ids_tensor = None - self.eplb_step() + arctic_spec_token_ids = None + suffix_spec_token_ids = None - output = ModelRunnerOutput( - req_ids=req_ids_output_copy, - req_id_to_index=req_id_to_index_output_copy, - sampled_token_ids=valid_sampled_token_ids, - logprobs=logprobs_lists, - prompt_logprobs_dict=prompt_logprobs_dict, - pooler_output=[], - ) + if self.speculative_config.method in ("arctic", "mlp_speculator"): + if sampled_token_ids_tensor is None: + import numpy as np + sampled_token_ids_tensor = torch.tensor(sampled_token_ids_list, device=self.device) - if self.use_async_scheduling: - return AsyncGPUModelRunnerOutput( - model_runner_output=output, - sampled_token_ids=sampler_output.sampled_token_ids, - invalid_req_indices=invalid_req_indices, - async_output_copy_stream=self.async_output_copy_stream, + previous_hidden_states = self.drafter.prepare_hidden_states( + sample_hidden_states=sample_hidden_states, + sampled_token_ids=sampled_token_ids_tensor, + spec_decode_metadata=spec_decode_metadata, ) - return output - def propose_draft_token_ids( - self, - scheduler_output: "SchedulerOutput", - sampled_token_ids: list[list[int]], - original_sampled_token_ids: Union[np.ndarray, torch.Tensor], - sampling_metadata: SamplingMetadata, - hidden_states: torch.Tensor, - sample_hidden_states: torch.Tensor, - aux_hidden_states: Optional[torch.Tensor], - spec_decode_metadata: Optional[SpecDecodeMetadata], - common_attn_metadata: CommonAttentionMetadata, - ) -> list[list[int]]: - disable_spec_decode = ( - self.speculative_config - and self.speculative_config.disable_by_batch_size - and len(self.input_batch.req_ids) > self.speculative_config.disable_by_batch_size - ) - if disable_spec_decode: - return [[] for _ in sampled_token_ids] + next_token_ids = self.drafter.prepare_next_token_ids_cpu( + sampled_token_ids_list, + self.requests, + self.input_batch, + scheduler_output.num_scheduled_tokens, + ) + + # Only draft for the first draft_limit requests. + arctic_output_tensor = self.drafter.propose( + context_token_ids=next_token_ids[:draft_limit], + previous_hidden_states=previous_hidden_states[:draft_limit], + num_predict_tokens=self.drafter.model.n_predict, + ) + + arctic_spec_token_ids = arctic_output_tensor.tolist() + # Pad with empty lists for requests beyond draft_limit. + if draft_limit < batch_size: + arctic_spec_token_ids.extend( + [] for _ in range(batch_size - draft_limit) + ) - suffix_spec_token_ids = None - new_sampled_token_ids = sampled_token_ids.copy() if self._suffix_cache is not None: - results = self.propose_suffix_draft_token_ids(new_sampled_token_ids) + self._update_suffix_cache(sampled_token_ids_list) + results = self.propose_suffix_draft_token_ids(sampled_token_ids_list) + suffix_spec_token_ids = [] min_score = 0 if self.speculative_config.method == "suffix" \ - else self.speculative_config.num_speculative_tokens - for i, result in enumerate(results): + else self.drafter.model.n_predict + + for result in results: if result.score >= min_score: - new_sampled_token_ids[i] = [] suffix_spec_token_ids.append(result.token_ids) else: suffix_spec_token_ids.append([]) spec_token_ids = None - if self.speculative_config.method == "suffix": - pass - elif self.speculative_config.method in ("arctic", "mlp_speculator"): - assert isinstance(self.drafter, ArcticProposer) - if isinstance(original_sampled_token_ids, np.ndarray): - original_sampled_token_ids_tensor = torch.from_numpy( - original_sampled_token_ids - ).to(self.device) - else: - original_sampled_token_ids_tensor = original_sampled_token_ids - - previous_hidden_states = self.drafter.prepare_hidden_states( - sample_hidden_states=sample_hidden_states, - sampled_token_ids=original_sampled_token_ids_tensor, - spec_decode_metadata=spec_decode_metadata, - ) - spec_token_ids = self.propose_arctic_draft_token_ids( - scheduler_output, - new_sampled_token_ids, - previous_hidden_states=previous_hidden_states, - ) + if suffix_spec_token_ids is not None and arctic_spec_token_ids is not None: + spec_token_ids = [ + s_tokens if s_tokens else a_tokens + for s_tokens, a_tokens in zip(suffix_spec_token_ids, arctic_spec_token_ids) + ] + elif suffix_spec_token_ids is not None: + spec_token_ids = suffix_spec_token_ids + elif arctic_spec_token_ids is not None: + spec_token_ids = arctic_spec_token_ids else: spec_token_ids = self._orig_propose_draft_token_ids( scheduler_output, - new_sampled_token_ids, + sampled_token_ids_list, sampling_metadata, hidden_states, sample_hidden_states, @@ -856,87 +965,62 @@ def propose_draft_token_ids( ) if spec_token_ids is None: - spec_token_ids = suffix_spec_token_ids - elif suffix_spec_token_ids is not None: - spec_token_ids = [ - suffix_spec_token_ids[i] or spec_token_ids[i] - for i in range(len(suffix_spec_token_ids)) - ] - - return spec_token_ids - - def propose_arctic_draft_token_ids( - self, - scheduler_output: "SchedulerOutput", - sampled_token_ids: list[list[int]], - previous_hidden_states: Optional[torch.Tensor] = None, - ) -> list[list[int]]: - last_tokens = [] - max_spec_tokens = self.speculative_config.num_speculative_tokens - for i, sampled_ids in enumerate(sampled_token_ids): - num_sampled_ids = len(sampled_ids) - - if num_sampled_ids == 0: - if self.speculative_config.enable_suffix_decoding: - return [[]] * len(sampled_token_ids) - req_id = self.input_batch.req_ids[i] - req_state = self.requests[req_id] - seq_len = (req_state.num_computed_tokens + - scheduler_output.num_scheduled_tokens[req_id]) - sampled_ids = [req_state.get_token_id(seq_len)] - num_sampled_ids = len(sampled_ids) - - start_idx = self.input_batch.num_tokens_no_spec[i] - end_idx = start_idx + num_sampled_ids - - current_max_spec_tokens = min( - max_spec_tokens, - self.max_model_len - end_idx - 1, - ) - if current_max_spec_tokens <= 0: - last_tokens.append(None) - continue - - if num_sampled_ids > 0: - self.input_batch.token_ids_cpu[i, start_idx:end_idx] = sampled_ids - last_tokens.append( - self.input_batch.token_ids_cpu[i, end_idx - 1].item() - ) - else: - last_tokens.append(None) - continue - - valid_last_tokens = [t for t in last_tokens if t is not None] - if not valid_last_tokens: - return [[] for _ in sampled_token_ids] - - if previous_hidden_states is not None: - indices = [i for i, t in enumerate(last_tokens) if t is not None] - previous_hidden_states = previous_hidden_states[indices] - - final_max_spec_tokens = min( - max_spec_tokens, - self.max_model_len - self.input_batch.num_tokens_no_spec.max() - 1, - ) - if final_max_spec_tokens <= 0: - return [[] for _ in sampled_token_ids] - - drafter_output = self.drafter.propose( - valid_last_tokens, - previous_hidden_states=previous_hidden_states, - num_predict_tokens=final_max_spec_tokens, - ) - - draft_token_ids_list = drafter_output.tolist() - final_draft_token_ids = [] - draft_iter = iter(draft_token_ids_list) - for t in last_tokens: - if t is not None: - final_draft_token_ids.append(next(draft_iter)) + spec_token_ids = [[] for _ in range(len(self.input_batch.req_ids))] + + # For async scheduling the base _prepare_input_ids asserts that + # _draft_token_ids is a torch.Tensor and uses it to scatter draft + # tokens into the input. If we reached here (non-async code path) + # while async scheduling is active we must: + # 1. Convert the list-of-lists draft tokens to a padded tensor. + # 2. Call _copy_valid_sampled_token_count so the next step's + # _get_valid_sampled_token_count returns correct counts. + if (self.use_async_scheduling + and isinstance(spec_token_ids, list) + and isinstance(sampled_token_ids, torch.Tensor)): + # --- _copy_valid_sampled_token_count --- + if (hasattr(self, 'drafter') + and hasattr(self.drafter, 'prepare_next_token_ids_padded') + and common_attn_metadata is not None): + next_tok, valid_cnt = ( + self.drafter.prepare_next_token_ids_padded( + common_attn_metadata, + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + )) + self._copy_valid_sampled_token_count(next_tok, valid_cnt) else: - final_draft_token_ids.append([]) + # Manual fallback (suffix-only drafter). + mask = sampled_token_ids != -1 + valid_cnt = mask.sum(dim=1) + _bs = sampled_token_ids.shape[0] + cols = torch.arange( + sampled_token_ids.shape[1], + device=sampled_token_ids.device, + ).unsqueeze(0).expand_as(sampled_token_ids) + last_col = cols.masked_fill(~mask, -1).max(dim=1).values + last_col = last_col.clamp(min=0) + next_tok = sampled_token_ids[ + torch.arange(_bs, device=sampled_token_ids.device), + last_col, + ] + self._copy_valid_sampled_token_count(next_tok, valid_cnt) + + # --- Convert list[list[int]] -> padded tensor --- + padded = torch.zeros( + batch_size, self.num_spec_tokens, + dtype=torch.int32, device=self.device, + ) + for i, tokens in enumerate(spec_token_ids): + length = min(len(tokens), self.num_spec_tokens) + if length > 0: + padded[i, :length] = torch.tensor( + tokens[:length], dtype=torch.int32, + device=self.device) + spec_token_ids = padded - return final_draft_token_ids + return spec_token_ids def _update_suffix_cache(self, sampled_token_ids: list[list[int]]) -> None: seen_req_ids = set() @@ -948,18 +1032,24 @@ def _update_suffix_cache(self, sampled_token_ids: list[list[int]]) -> None: continue index = self.input_batch.req_id_to_index[req_id] - if req_id not in self._suffix_cache.active_requests: + is_new = req_id not in self._suffix_cache.active_requests + if is_new: if req_id in self._suffix_cache.cached_requests: self._suffix_cache.evict_cached_response(req_id) num_prompt_tokens = self.input_batch.num_prompt_tokens[index] prompt_token_ids = self.input_batch.token_ids_cpu[index, :num_prompt_tokens] self._suffix_cache.start_request(req_id, prompt_token_ids.tolist()) + self._suffix_response_tokens[req_id] = [] self._suffix_cache.add_active_response(req_id, sampled_ids) + self._suffix_response_tokens[req_id].extend(sampled_ids) + stopped_ids = [] for req_id in list(self._suffix_cache.active_requests): if req_id not in seen_req_ids: self._suffix_cache.stop_request(req_id) + self._suffix_response_tokens.pop(req_id, None) + stopped_ids.append(req_id) def propose_suffix_draft_token_ids( self, @@ -974,23 +1064,46 @@ def propose_suffix_draft_token_ids( continue req_id = self.input_batch.req_ids[i] - if req_id in self.input_batch.spec_decode_unsupported_reqs: - results.append(SuffixDecodingDraft()) - continue - - num_tokens = self.input_batch.num_tokens_no_spec[i] - if num_tokens >= self.max_model_len: - results.append(SuffixDecodingDraft()) - continue + index = self.input_batch.req_id_to_index[req_id] - start = max(0, num_tokens - config.suffix_cache_max_depth) - pattern = self.input_batch.token_ids_cpu[i, start:num_tokens].tolist() + # In async mode, token_ids_cpu contains -1 placeholders at + # decoded positions (written by _bookkeeping_sync). Build the + # pattern from the clean _suffix_response_tokens instead. + if (self.use_async_scheduling + and req_id in self._suffix_response_tokens): + response = self._suffix_response_tokens[req_id] + num_prompt = int( + self.input_batch.num_prompt_tokens[index]) + num_tokens = num_prompt + len(response) + if num_tokens >= self.max_model_len: + results.append(SuffixDecodingDraft()) + continue + # Take up to suffix_cache_max_depth tokens from the tail. + depth = config.suffix_cache_max_depth + if len(response) >= depth: + pattern = response[-depth:] + else: + need = depth - len(response) + prompt_start = max(0, num_prompt - need) + prompt_part = self.input_batch.token_ids_cpu[ + index, prompt_start:num_prompt].tolist() + pattern = prompt_part + response + else: + num_tokens = self.input_batch.num_tokens_no_spec[i] + if num_tokens >= self.max_model_len: + results.append(SuffixDecodingDraft()) + continue + start = max(0, num_tokens - config.suffix_cache_max_depth) + pattern = self.input_batch.token_ids_cpu[ + i, start:num_tokens].tolist() + + max_spec = min( + MAX_SPEC_LEN, self.max_model_len - num_tokens - 1 + ) result = self._suffix_cache.speculate( req_id, pattern, - max_spec_tokens=min( - MAX_SPEC_LEN, self.max_model_len - num_tokens - 1 - ), + max_spec_tokens=max_spec, max_spec_factor=config.suffix_max_spec_factor, max_spec_offset=config.suffix_max_spec_offset, min_token_prob=config.suffix_min_token_prob, @@ -1000,6 +1113,505 @@ def propose_suffix_draft_token_ids( return results + + def _start_suffix_copy( + self, + sampled_token_ids: torch.Tensor, + ) -> None: + """Initiate an async D2H copy of sampled token IDs for suffix decoding. + + Copies ``sampled_token_ids`` to a pinned CPU buffer on a dedicated + CUDA stream (``suffix_copy_stream``) that only waits for prior work + on the default stream (i.e. sampling). + + **This MUST be called BEFORE launching Arctic GPU work on the default + stream** so that the copy is not ordered behind Arctic kernels. The + resulting timeline is:: + + Default stream: [sample] -> [arctic prepare / propose ...] + Suffix stream : [sample] -> [D2H copy] -> [event] + CPU : wait -> suffix logic + + The companion ``_finish_suffix_copy`` synchronises on the copy event + and returns the materialised CPU list. + """ + n_rows = sampled_token_ids.shape[0] + n_cols = sampled_token_ids.shape[-1] + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(self.suffix_copy_stream): + self.suffix_copy_stream.wait_stream(default_stream) + self.suffix_sampled_ids_pinned[:n_rows, :n_cols].copy_( + sampled_token_ids, non_blocking=True, + ) + self.suffix_copy_done_event.record() + self._suffix_copy_shape = (n_rows, n_cols) + + def _finish_suffix_copy(self) -> list[list[int]]: + """Wait for the suffix copy and return sampled token IDs as CPU lists. + + Synchronises on the copy event recorded by ``_start_suffix_copy``, + then parses the pinned buffer into ``list[list[int]]``, applying + rejection-sampling for spec-decode batches and masking discarded + (still-in-prefill) requests. + """ + self.suffix_copy_done_event.synchronize() + n_rows, n_cols = self._suffix_copy_shape + pinned = self.suffix_sampled_ids_pinned[:n_rows, :n_cols] + + num_reqs = self.input_batch.num_reqs + discard_indices = np.nonzero( + self.discard_request_mask.np[:num_reqs] + )[0] + + if n_cols == 1: + result = pinned.tolist() + for i in discard_indices: + result[int(i)].clear() + else: + result, _ = RejectionSampler.parse_output( + pinned, + self.input_batch.vocab_size, + discard_indices, + ) + + return result + + @torch.inference_mode + def sample_tokens( + self, grammar_output: "GrammarOutput | None" + ) -> Union[ModelRunnerOutput, AsyncGPUModelRunnerOutput, IntermediateTensors]: + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + + if self.execute_model_state is None: + # Nothing to do (PP non-final rank case), output isn't used. + if not kv_connector_output: + return None # noqa + + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output + if kv_connector_output.is_empty(): + return EMPTY_MODEL_RUNNER_OUTPUT + + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.kv_connector_output = kv_connector_output + return output + + # Unpack ephemeral state. + ( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + cudagraph_stats, + ) = self.execute_model_state + # Clear ephemeral state. + self.execute_model_state = None + + # Apply structured output bitmasks if present. + if grammar_output is not None: + apply_grammar_bitmask( + scheduler_output, grammar_output, self.input_batch, logits + ) + + with record_function_or_nullcontext("gpu_model_runner: sample"): + sampler_output = self._sample(logits, spec_decode_metadata) + + self._draft_token_ids = None + self._draft_token_req_ids = None + self.input_batch.prev_sampled_token_ids = None + + def propose_draft_token_ids( + sampled_token_ids: torch.Tensor | list[np.ndarray], + ) -> None: + assert spec_decode_common_attn_metadata is not None + with record_function_or_nullcontext("gpu_model_runner: draft"): + self._draft_token_ids = self.propose_draft_token_ids( + scheduler_output, + sampled_token_ids, + self.input_batch.sampling_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + ) + self._copy_draft_token_ids_to_cpu(scheduler_output) + + # --- Draft proposal orchestration --- + # + # There are three drafting modes depending on async scheduling and + # which drafters are active: + # + # A) Arctic-only + async: + # Pre-bookkeeping. Arctic runs entirely on GPU tensors + # (the existing EAGLE-like async path). + # + # B) Suffix (with or without arctic) + async: + # Pre-bookkeeping. Uses a *dedicated* suffix_copy_stream so + # that the D2H transfer of sampled token IDs only waits on + # "sampling complete" -- NOT on Arctic GPU work. Timeline: + # + # Default stream: [sample] -> [arctic GPU work ----------->] + # Suffix stream: [sample] -> [D2H copy] -> [event] + # CPU : wait -> suffix + # + # At merge time, if any requests need arctic fallback, we + # sync on the arctic result (.tolist()); by then the GPU + # kernels have had the full suffix-CPU window to finish. + # + # C) Non-async (any combination): + # Post-bookkeeping. propose_draft_token_ids handles everything + # using CPU valid_sampled_token_ids from bookkeeping. + + has_suffix = self._suffix_cache is not None + is_arctic_method = ( + self.speculative_config is not None + and self.speculative_config.method in ("arctic", "mlp_speculator") + ) + input_fits_in_drafter = spec_decode_common_attn_metadata is not None + sampled_token_ids = sampler_output.sampled_token_ids + + # Determine if we should use async pre-bookkeeping drafting. + use_async_spec = ( + self.use_async_scheduling + and self.speculative_config is not None + and (is_arctic_method or has_suffix) + and input_fits_in_drafter + ) + + if use_async_spec: + if is_arctic_method and not has_suffix: + # (A) Arctic-only async: use existing closure which calls + # propose_draft_token_ids (the method) with a GPU tensor + # and triggers the async GPU path internally. + propose_draft_token_ids(sampled_token_ids) + + # Track actual draft lengths for next step's allocation. + _n_predict = self.drafter.model.n_predict + _batch_size = len(self.input_batch.req_ids) + _disable_bs = ( + self.speculative_config.disable_by_batch_size + if self.speculative_config else None + ) + _draft_limit = _batch_size + if _disable_bs and _batch_size > _disable_bs: + _draft_limit = _disable_bs + self._prev_actual_draft_lens = { + req_id: _n_predict if i < _draft_limit else 0 + for i, req_id in enumerate(self.input_batch.req_ids) + } + scheduler_output._actual_draft_lens = ( + self._prev_actual_draft_lens + ) + + elif is_arctic_method: + # (B) Arctic + suffix async. + # D2H copy of sampled tokens runs on a dedicated + # suffix_copy_stream that only waits for sampling, + # NOT for Arctic GPU kernels enqueued afterwards. + # Suffix CPU work then overlaps with Arctic GPU. + + # Step 1: Initiate async D2H copy before Arctic GPU work. + self._start_suffix_copy(sampled_token_ids) + + # Step 2: Launch arctic on the default stream. + with record_function_or_nullcontext( + "gpu_model_runner: draft (arctic)" + ): + arctic_draft = self.propose_draft_token_ids( + scheduler_output, + sampled_token_ids, + self.input_batch.sampling_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + ) + + # Step 3: Wait for D2H copy, then run suffix on CPU. + # This overlaps with remaining Arctic GPU work. + with record_function_or_nullcontext( + "gpu_model_runner: draft (suffix)" + ): + sampled_cpu = self._finish_suffix_copy() + self._update_suffix_cache(sampled_cpu) + suffix_results = self.propose_suffix_draft_token_ids( + sampled_cpu + ) + + min_score = self.drafter.model.n_predict + suffix_draft = [ + result.token_ids if result.score >= min_score + else [] + for result in suffix_results + ] + + # Step 4: Merge arctic + suffix results. + # Suffix takes priority when available. + # + # The merged result MUST be a torch.Tensor on GPU + # because upstream _prepare_input_ids (next iteration) + # asserts isinstance(self._draft_token_ids, torch.Tensor) + # and uses it for on-device scatter. + + # Collect suffix rows that have results. + suffix_indices = [ + i for i, s in enumerate(suffix_draft) if s + ] + + # arctic_draft is already a [batch, num_spec_tokens] + # GPU tensor (from the pre-allocated _draft_merged_gpu + # buffer when available, or F.pad fallback). + if isinstance(arctic_draft, torch.Tensor): + merged = arctic_draft + else: + # Shouldn't happen in async path, but handle + # gracefully. + k = self.num_spec_tokens + n = len(arctic_draft) + pin = self._suffix_merge_pinned[:n, :k] + pin_np = pin.numpy() + pin_np[:] = 0 + for i, row in enumerate(arctic_draft): + t = row[:k] + if t: + pin_np[i, :len(t)] = t + merged = pin.to( + device=self.device, non_blocking=True) + + if merged.shape[1] < self.num_spec_tokens: + merged = torch.nn.functional.pad( + merged, + (0, self.num_spec_tokens - merged.shape[1]), + value=0, + ) + + # Batch-overwrite rows that have suffix results. + # CRITICAL: use the pre-allocated *pinned* merge + # buffer so that cudaMemcpyAsync does NOT synchronise + # the default stream. With pageable (non-pinned) + # memory CUDA must sync the stream first, blocking the + # CPU until Arctic GPU work finishes -- destroying + # the async overlap. + width = merged.shape[1] + if suffix_indices: + n_sfx = len(suffix_indices) + overlay_pin = \ + self._suffix_merge_pinned[:n_sfx, :width] + overlay_np = overlay_pin.numpy() + overlay_np[:] = 0 + for j, idx in enumerate(suffix_indices): + s = suffix_draft[idx] + slen = min(len(s), width) + overlay_np[j, :slen] = s[:slen] + # Pinned H2C -- truly non-blocking. + overlay_t = overlay_pin.to( + device=self.device, non_blocking=True) + # Use pre-allocated pinned index buffer when + # available to avoid per-step allocation. + idx_pinned = getattr( + self, '_suffix_index_pinned', None) + if idx_pinned is not None: + idx_pin = idx_pinned[:n_sfx] + idx_pin[:] = torch.tensor( + suffix_indices, dtype=torch.long) + else: + idx_pin = torch.tensor( + suffix_indices, dtype=torch.long + ).pin_memory() + idx_t = idx_pin.to( + device=self.device, non_blocking=True) + merged.index_copy_(0, idx_t, overlay_t) + + self._draft_token_ids = merged + self._copy_draft_token_ids_to_cpu(scheduler_output) + + # Track actual draft lengths for next step's allocation. + _n_predict = self.drafter.model.n_predict + _batch_size = len(self.input_batch.req_ids) + _disable_bs = ( + self.speculative_config.disable_by_batch_size + if self.speculative_config else None + ) + _draft_limit = _batch_size + if _disable_bs and _batch_size > _disable_bs: + _draft_limit = _disable_bs + _actual_lens: dict[str, int] = {} + for _i, _req_id in enumerate(self.input_batch.req_ids): + _s = (len(suffix_draft[_i]) + if _i < len(suffix_draft) else 0) + _a = _n_predict if _i < _draft_limit else 0 + # Suffix takes priority when available. + _actual_lens[_req_id] = _s if _s > 0 else _a + self._prev_actual_draft_lens = _actual_lens + scheduler_output._actual_draft_lens = _actual_lens + + else: + # (B2) Suffix-only async. + # No Arctic GPU work to overlap with, so skip the + # copy stream (its overhead exceeds the marginal + # overlap with the lightweight rejection kernel). + # Instead: rejection sample → parse tokens directly + # → suffix CPU work → build GPU tensor via pinned buf. + self._suffix_only_rejection_sample( + sampled_token_ids, + spec_decode_common_attn_metadata, + ) + + # Parse sampled tokens to CPU lists. .cpu() syncs the + # default stream (waits for the rejection kernel, which + # is lightweight), then parse_output performs CPU-side + # rejection to extract accepted token IDs. + with record_function_or_nullcontext( + "gpu_model_runner: draft (suffix)" + ): + num_reqs = self.input_batch.num_reqs + discard_indices = np.nonzero( + self.discard_request_mask.np[:num_reqs] + )[0] + n_cols = sampled_token_ids.shape[-1] + if n_cols == 1: + sampled_cpu = sampled_token_ids.tolist() + for idx in discard_indices: + sampled_cpu[int(idx)].clear() + else: + sampled_cpu, _ = RejectionSampler.parse_output( + sampled_token_ids.cpu(), + self.input_batch.vocab_size, + discard_indices, + ) + + self._update_suffix_cache(sampled_cpu) + suffix_results = self.propose_suffix_draft_token_ids( + sampled_cpu + ) + suffix_draft = [ + result.token_ids if result.score >= 0 + else [] + for result in suffix_results + ] + + # Build GPU tensor from suffix lists via pinned buffer. + k = self.num_spec_tokens + n = len(suffix_draft) + pin = self._suffix_merge_pinned[:n, :k] + pin_np = pin.numpy() + pin_np[:] = 0 + for i, s in enumerate(suffix_draft): + if s: + slen = min(len(s), k) + pin_np[i, :slen] = s[:slen] + self._draft_token_ids = pin.to( + device=self.device, non_blocking=True) + self._copy_draft_token_ids_to_cpu(scheduler_output) + + # Track actual draft lengths for next step's allocation. + _actual_lens_b2: dict[str, int] = {} + for _i, _req_id in enumerate(self.input_batch.req_ids): + _s = (len(suffix_draft[_i]) + if _i < len(suffix_draft) else 0) + _actual_lens_b2[_req_id] = _s + self._prev_actual_draft_lens = _actual_lens_b2 + scheduler_output._actual_draft_lens = _actual_lens_b2 + + # --- Bookkeeping --- + with record_function_or_nullcontext("gpu_model_runner: bookkeep"): + ( + num_nans_in_logits, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) = self._bookkeeping_sync( + scheduler_output, + sampler_output, + logits, + hidden_states, + scheduler_output.total_num_scheduled_tokens, + spec_decode_metadata, + ) + + # (C) Non-async drafting: run after bookkeeping. + # Pass the raw GPU sampled_token_ids tensor (not the cleaned + # valid_sampled_token_ids list) so that propose_draft_token_ids + # gets a properly shaped 2-D tensor with -1 markers for rejected + # tokens. The method already handles tensors correctly: it + # filters out -1 to build the CPU list for + # prepare_next_token_ids_cpu and keeps the raw tensor for + # prepare_hidden_states. + if ( + self.speculative_config is not None + and not use_async_spec + and input_fits_in_drafter + ): + propose_draft_token_ids(sampled_token_ids) + + with record_function_or_nullcontext("gpu_model_runner: eplb"): + self.eplb_step() + + with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + if self.model_config.enable_return_routed_experts: + capturer = RoutedExpertsCapturer.get_instance() + if capturer is not None: + capturer.save_captured_experts(indices=self.slot_mapping) # noqa + else: + logger.error("RoutedExpertsCapturer not initialized.") + + output = ModelRunnerOutput( + req_ids=req_ids_output_copy, + req_id_to_index=req_id_to_index_output_copy, + sampled_token_ids=valid_sampled_token_ids, + logprobs=logprobs_lists, + prompt_logprobs_dict=prompt_logprobs_dict, + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output + if self.supports_mm_inputs + else None, + num_nans_in_logits=num_nans_in_logits, + cudagraph_stats=cudagraph_stats, + ) + # Attach actual draft lengths to the ModelRunnerOutput so the + # scheduler can read them reliably in update_from_output. + # This survives the async pipeline (scheduler_output attrs + # may not due to object lifecycle in the batch queue). + output._actual_draft_lens = getattr( + scheduler_output, '_actual_draft_lens', None) + + if not self.use_async_scheduling: + return output + + with record_function_or_nullcontext( + "gpu_model_runner: AsyncGPUModelRunnerOutput" + ): + async_output = AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=sampler_output.sampled_token_ids, + logprobs_tensors=sampler_output.logprobs_tensors, + invalid_req_indices=invalid_req_indices, + async_output_copy_stream=self.async_output_copy_stream, + vocab_size=self.input_batch.vocab_size, + ) + with record_function_or_nullcontext( + "gpu_model_runner: set_async_sampled_token_ids" + ): + # Save ref of sampled_token_ids CPU tensor if the batch contains + # any requests with sampling params that require output ids. + self.input_batch.set_async_sampled_token_ids( + async_output.sampled_token_ids_cpu, + async_output.async_copy_ready_event, + ) + + return async_output + + def load_model(self, eep_scale_up: bool = False) -> None: load_shift_model = ( self.vllm_config.parallel_config.enable_shift_parallel @@ -1023,6 +1635,10 @@ def load_model(self, eep_scale_up: bool = False) -> None: self.shift_parallel_threshold = ( shift_config.parallel_config.shift_parallel_threshold ) + self.shift_forward_context = ( + shift_config.compilation_config.static_forward_context + ) + if "SwiftKV" in self.model.__class__.__name__: if hasattr(self.model, "model") and hasattr(self.model.model, "decode_runner"): self.model.model.decode_runner = self.shift_model.model.decode_runner @@ -1032,107 +1648,182 @@ def load_model(self, eep_scale_up: bool = False) -> None: else: self.shift_model = None self.shift_parallel_threshold = 0 + self.shift_forward_context = None + - def _capture_cudagraphs(self, compilation_cases: list[int], - cudagraph_runtime_mode: CUDAGraphMode, - uniform_decode: bool): + def initialize_kv_cache(self, kv_cache_config) -> None: + self._orig_initialize_kv_cache(kv_cache_config) + shift_ctx = getattr(self, 'shift_forward_context', None) + if shift_ctx is None: + return + base_ctx = self.compilation_config.static_forward_context + bound = 0 + for name, shift_attn in shift_ctx.items(): + base_attn = base_ctx.get(name) + if base_attn is not None and hasattr(base_attn, 'kv_cache'): + shift_attn.kv_cache = base_attn.kv_cache + bound += 1 + if is_global_first_rank(): + logger.info("Bound KV cache to %d shift model attention layers", + bound) + + from vllm.forward_context import BatchDescriptor + def _case_bs(self, case) -> int: + # vLLM can pass ints, tuples, or sometimes BatchDescriptor-like objects + if isinstance(case, int): + return case + if isinstance(case, BatchDescriptor): + return int(case.num_tokens) + if isinstance(case, tuple): + return int(case[0]) + # last resort + return int(getattr(case, "num_tokens")) + + def _with_bs(self, case, new_bs: int): + if isinstance(case, tuple): + return (new_bs, *case[1:]) + if isinstance(case, BatchDescriptor): + # Best-effort reconstruction; adjust if your BatchDescriptor signature differs. + return BatchDescriptor( + num_tokens=new_bs, + num_reqs=case.num_reqs, + uniform=case.uniform, + has_lora=case.has_lora, + ) + return new_bs + + + @contextlib.contextmanager + def _shift_graph_capture_context(self): + """Disable custom all-reduce on the _SP_TP group during shift model + graph capture so it falls back to pynccl (NCCL), which handles graph + capture natively. The _SP_TP group's ca_comm was never set up through + the normal vLLM graph_capture() path, so using it inside a CUDA graph + would crash.""" + from vllm.distributed.device_communicators.cuda_communicator import ( + CudaCommunicator, + ) + sp_tp = parallel_state._SP_TP + ca_comm = None + if sp_tp is not None and sp_tp.device_communicator is not None: + assert isinstance(sp_tp.device_communicator, CudaCommunicator) + ca_comm = sp_tp.device_communicator.ca_comm + saved_disabled = None + if ca_comm is not None: + saved_disabled = ca_comm.disabled + ca_comm.disabled = True + try: + yield + finally: + if ca_comm is not None and saved_disabled is not None: + ca_comm.disabled = saved_disabled + + @contextlib.contextmanager + def _use_shift_cudagraph_tables(self): + """Temporarily swap compilation_config sizes to the shift (unscaled) + lookup table so that vLLM internals (dispatcher, pad_for_cudagraph, + bounds checks) all see the shift model's sizes.""" + cc = self.compilation_config + saved_sizes = cc.cudagraph_capture_sizes + saved_max = cc.max_cudagraph_capture_size + saved_table = cc.bs_to_padded_graph_size + + shift_sizes = self.vllm_config._shift_cudagraph_capture_sizes + shift_max = self.vllm_config._shift_max_cudagraph_capture_size + shift_table = self.vllm_config._shift_bs_to_padded_graph_size + + cc.cudagraph_capture_sizes = shift_sizes + cc.max_cudagraph_capture_size = shift_max + cc.bs_to_padded_graph_size = shift_table + try: + yield + finally: + cc.cudagraph_capture_sizes = saved_sizes + cc.max_cudagraph_capture_size = saved_max + cc.bs_to_padded_graph_size = saved_table + + def _capture_cudagraphs( + self, + compilation_cases: list[tuple[int, bool]], + cudagraph_runtime_mode: CUDAGraphMode, + uniform_decode: bool, + ): """ Capture CUDA graphs for both base (SP) and shift (TP) variants, splitting shapes by threshold so both models have required graphs captured. """ - # make sure it is a valid mode assert cudagraph_runtime_mode != CUDAGraphMode.NONE and \ cudagraph_runtime_mode in [CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE] sp_size = parallel_state._SP.world_size tp_size = parallel_state._TP.world_size - - # capture shapes for the base model - compilation_cases_base = [ - min(shape * sp_size, self.max_num_tokens) - for shape in compilation_cases - if shape * sp_size > self.shift_parallel_threshold - ] - # filter again for FULL mode - if cudagraph_runtime_mode == CUDAGraphMode.FULL: - compilation_cases_base = [shape for shape in compilation_cases_base if - shape <= self.scheduler_config.max_num_seqs] + threshold = int(getattr(self, "shift_parallel_threshold", 0)) + has_shift = getattr(self, "shift_model", None) is not None + is_swiftkv = "SwiftKV" in self.model.__class__.__name__ + + # --- Base model (Ulysses SP): uses the scaled lookup table (default) --- + # Exclude sizes at or below the shift threshold -- those batches + # are routed to the shift model at runtime. Capturing them for the + # base model would deadlock because the Ulysses all-to-all collectives + # diverge across ranks at small batch sizes. + if has_shift and not is_swiftkv: + compilation_cases_base = [ + case for case in compilation_cases + if self._case_bs(case) > threshold + ] + else: + compilation_cases_base = list(compilation_cases) if is_global_first_rank(): - logger.info(f"base model (SP={sp_size}, TP={tp_size}) shapes {compilation_cases_base}, cudagraph mode {cudagraph_runtime_mode}") + logger.info( + "base model (SP=%s, TP=%s) cudagraph mode %s shapes %s", + sp_size, tp_size, cudagraph_runtime_mode, + [self._case_bs(c) for c in compilation_cases_base], + ) - # capture the base model graphs if compilation_cases_base: self._orig_capture_cudagraphs( compilation_cases_base, cudagraph_runtime_mode, uniform_decode ) - if getattr(self, "shift_model", None) is not None: - - # capture shapes for the shift model + # --- Shift model (SP*TP fused as TP-only): uses the unscaled lookup table --- + # The incoming compilation_cases contain *scaled* base sizes (e.g. + # [4, 8, ..., 2048] with sp_size=4). The shift model needs the + # *unscaled* sizes from its own capture list (e.g. [1, 2, ..., 512]). + # We rebuild the cases from _shift_cudagraph_capture_sizes, copying + # the non-bs fields (like has_lora) from the first matching base case. + if has_shift: + shift_sizes = self.vllm_config._shift_cudagraph_capture_sizes + # Use the first base case as a template for non-bs fields + template = compilation_cases[0] if compilation_cases else None compilation_cases_shift = [ - shape for shape in compilation_cases - if shape <= self.shift_parallel_threshold - or "SwiftKV" in self.model.__class__.__name__ + self._with_bs(template, bs) if template is not None else bs + for bs in reversed(shift_sizes) ] - # filter again for FULL mode - if cudagraph_runtime_mode == CUDAGraphMode.FULL: - compilation_cases_shift = [shape for shape in compilation_cases_shift if - shape <= self.scheduler_config.max_num_seqs] if is_global_first_rank(): - logger.info(f"shift model (SPxTP={sp_size * tp_size}) shapes {compilation_cases_shift}, cudagraph mode {cudagraph_runtime_mode}") + logger.info( + "shift model (SPxTP=%s) shapes %s", + sp_size * tp_size, + [self._case_bs(c) for c in compilation_cases_shift], + ) - # capture the shift model graphs if compilation_cases_shift: orig_model, self.model = self.model, self.shift_model + cc = self.vllm_config.compilation_config + base_ctx = cc.static_forward_context + shift_ctx = getattr(self, 'shift_forward_context', None) try: - with set_shift_parallel_mode(True): - # capture the shift model graphs + if shift_ctx is not None: + cc.static_forward_context = shift_ctx + with set_shift_parallel_mode(True), \ + self._use_shift_cudagraph_tables(), \ + self._shift_graph_capture_context(): self._orig_capture_cudagraphs( - compilation_cases_shift, cudagraph_runtime_mode, uniform_decode + compilation_cases_shift, + cudagraph_runtime_mode, + uniform_decode, ) finally: self.model = orig_model - - def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: - """ - Initialize KV cache for base model, then bind the same buffers into - the shift-parallel model so they share cache/state. - """ - self._orig_initialize_kv_cache(kv_cache_config) - - if getattr(self, "shift_model", None) is not None: - forward_context = self.vllm_config.compilation_config.static_forward_context - - try: - from vllm.attention import Attention - for mod in self.shift_model.modules(): - if isinstance(mod, Attention): - if hasattr(mod, "layer_name"): - ln = mod.layer_name - if ln in forward_context: - mod.kv_cache = forward_context[ln].kv_cache - else: - logger.warning(f"Could not find {ln} in forward_context for shift_model.") - else: - logger.warning("Could not bind KV cache for shift_model: " - "Attention module missing 'layer_name'.") - except ImportError: - logger.warning("Could not import Attention to bind KV cache for shift_model.") - - try: - from vllm.model_executor.layers.mamba.abstract import MambaBase - for mod in self.shift_model.modules(): - if isinstance(mod, MambaBase): - if hasattr(mod, "layer_name"): - ln = mod.layer_name - if ln in forward_context: - mod.state = forward_context[ln].state - else: - logger.warning(f"Could not find {ln} in forward_context for shift_model.") - else: - logger.warning("Could not bind Mamba state for shift_model: " - "Mamba module missing 'layer_name'.") - except ImportError: - pass \ No newline at end of file + cc.static_forward_context = base_ctx diff --git a/arctic_inference/vllm/patches.py b/arctic_inference/vllm/patches.py index 78f2a5a8a..0d13baed7 100644 --- a/arctic_inference/vllm/patches.py +++ b/arctic_inference/vllm/patches.py @@ -13,9 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os import vllm from vllm.logger import init_logger +from vllm.v1.core.sched.async_scheduler import AsyncScheduler +from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.engine.core import EngineCoreProc from vllm.v1.worker.worker_base import WorkerBase @@ -35,6 +36,193 @@ logger = init_logger(__name__) +class AsyncSchedulerPatch(ArcticPatch[AsyncScheduler]): + """Patch AsyncScheduler to: + 1. Respect ``disable_by_batch_size`` when allocating spec token + placeholders (the worker only drafts for the first N requests). + 2. Use the previous step's actual draft length for dynamic placeholder + allocation, avoiding wasted verification compute when the real draft + width (e.g. Arctic n_predict=3) is much smaller than + num_speculative_tokens (e.g. 12). + 3. Store ``_scheduled_spec_count`` so that the post-fix in + ``update_from_output`` can compensate for worker-side trimming. + """ + + _orig_update_after_schedule = AsyncScheduler._update_after_schedule + + def _update_after_schedule(self, scheduler_output): + # Call the base Scheduler._update_after_schedule (NOT the + # AsyncScheduler override which we are replacing). + Scheduler._update_after_schedule(self, scheduler_output) + + has_structured_output_requests = False + pending_structured_output_tokens = False + spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens + + # Respect disable_by_batch_size: only add spec token placeholders + # for the first N decode requests (matching the worker's draft_limit + # in propose_draft_token_ids). + spec_config = getattr(self.vllm_config, 'speculative_config', None) + disable_bs = ( + spec_config.disable_by_batch_size if spec_config else None + ) + decode_with_spec_count = 0 + for req_id in scheduler_output.num_scheduled_tokens: + request = self.requests[req_id] + has_structured_output_requests |= request.use_structured_output + pending_structured_output_tokens |= ( + request.use_structured_output + and request.num_output_placeholders > 0 + ) + cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ())) + # Store the originally-scheduled spec count so that + # update_from_output can compensate for worker-side trimming. + request._scheduled_spec_count = cur_num_spec_tokens + + if ( + request.num_computed_tokens + == request.num_tokens + + request.num_output_placeholders + + cur_num_spec_tokens + ): + # The request will generate a new token + spec tokens. + request.num_output_placeholders += 1 + cur_num_spec_tokens + + # Check if beyond the disable_by_batch_size limit. + decode_with_spec_count += 1 + if disable_bs and decode_with_spec_count > disable_bs: + # Beyond limit: no spec token placeholders. + request.spec_token_ids = [] + continue + + # Use previous step's actual draft length to size + # placeholders. When suffix had a good match (actual + # > n_predict), allocate the full width so the next + # step can verify all suffix tokens. When suffix + # didn't match (actual = n_predict from arctic), + # allocate only that many to avoid wasting attention + # compute on zero-padded positions. + # Cold start: allocate full width (generous). + prev_actual = getattr( + request, '_prev_actual_draft_len', None) + if prev_actual is not None: + num_placeholders = min( + max(prev_actual, 1), self.num_spec_tokens) + else: + num_placeholders = self.num_spec_tokens + + request.spec_token_ids = [-1] * num_placeholders + + scheduler_output.has_structured_output_requests = ( + has_structured_output_requests) + scheduler_output.pending_structured_output_tokens = ( + pending_structured_output_tokens) + + def update_from_output(self, scheduler_output, model_runner_output): + """Wrap Scheduler.update_from_output to store actual draft counts. + + We infer the drafter's real capability from the acceptance + results so that the next ``_update_after_schedule`` can size + placeholders correctly. This works even when the worker runs + in a separate process (where scheduler_output._actual_draft_lens + set by the worker doesn't survive serialisation back to the + scheduler). + + Strategy: + 1. **Primary path**: read ``_actual_draft_lens`` from the + ``model_runner_output`` object (attached by the model runner + to the ``ModelRunnerOutput`` dataclass, which reliably + survives the async pipeline). + 2. **Legacy path**: read from ``scheduler_output._actual_draft_lens`` + (works in same-process non-async mode). + 3. **Fallback**: infer from acceptance results with exponential + growth — when all drafted tokens are accepted, double the + allocation so suffix decoding reaches full capacity in + O(log n) steps instead of O(n). + """ + sampled_token_ids = model_runner_output.sampled_token_ids + req_id_to_index = model_runner_output.req_id_to_index + + result = Scheduler.update_from_output( + self, scheduler_output, model_runner_output) + + # Primary path: read from model_runner_output (most reliable + # for async scheduling — the ModelRunnerOutput object is + # returned by get_output() and guaranteed to survive). + actual_lens = getattr( + model_runner_output, '_actual_draft_lens', None) + + # Legacy path: read from scheduler_output (works for non-async + # or same-process setups where the attribute is preserved). + if not actual_lens: + actual_lens = getattr( + scheduler_output, '_actual_draft_lens', None) + + if actual_lens: + for req_id, actual_len in actual_lens.items(): + request = self.requests.get(req_id) + if request is not None: + request._prev_actual_draft_len = actual_len + return result + + # Fallback: infer from acceptance results (multi-process case). + # Uses exponential growth when all drafted tokens are accepted + # (indicating the drafter / suffix cache can handle more), so + # the allocation converges to num_spec_tokens in O(log n) + # steps: + # step 0: 1 position → accept 1/1 → prev = 2 + # step 1: 2 positions → accept 2/2 → prev = 4 + # step 2: 4 positions → accept 4/4 → prev = 8 + # ... + # When not all are accepted (normal drafter), linear growth: + # step 0: 1 position → accept 1 → prev = 2 + # step 1: 2 positions → accept 2 → prev = 3 + # step 2: 3 positions → steady state (n_predict = 3) + if not sampled_token_ids: + return result + + for req_id in scheduler_output.num_scheduled_tokens: + scheduled_spec = ( + scheduler_output.scheduled_spec_decode_tokens.get(req_id)) + if not scheduled_spec: + continue + + req_index = req_id_to_index.get(req_id) + if req_index is None: + continue + + request = self.requests.get(req_id) + if request is None: + continue + + generated = sampled_token_ids[req_index] + num_accepted = (len(generated) - 1) if generated else 0 + num_draft = len(scheduled_spec) + + prev = getattr(request, '_prev_actual_draft_len', None) + + if num_accepted > 0: + if num_accepted >= num_draft and num_draft > 0: + # All drafted tokens accepted — the drafter (or + # suffix cache) could produce more if given room. + # Double the allocation for exponential ramp-up. + new_val = min( + num_draft * 2, self.num_spec_tokens) + else: + # Partial acceptance: grow linearly. + new_val = min( + num_accepted + 1, self.num_spec_tokens) + request._prev_actual_draft_len = max( + prev or 0, new_val) + elif prev is None: + # First step for this request, zero acceptance. + # Seed with 1 so _update_after_schedule doesn't + # fall back to the full num_spec_tokens next time. + request._prev_actual_draft_len = 1 + + return result + + class EngineCoreProcPatch(ArcticPatch[EngineCoreProc]): _orig_run_engine_core = EngineCoreProc.run_engine_core @@ -95,6 +283,10 @@ def apply_arctic_patches(): EngineCoreProcPatch.apply_patch() WorkerBasePatch.apply_patch() + # Async scheduler patches for spec decode (disable_by_batch_size + # interaction + dynamic draft width allocation). + AsyncSchedulerPatch.apply_patch() + # Patches to vLLM arguments and configuration objects. EngineArgsPatch.apply_patch() AsyncEngineArgsPatch.apply_patch() diff --git a/arctic_inference/vllm/spec_dec/arctic_proposer.py b/arctic_inference/vllm/spec_dec/arctic_proposer.py index 193b84294..4f471f66d 100644 --- a/arctic_inference/vllm/spec_dec/arctic_proposer.py +++ b/arctic_inference/vllm/spec_dec/arctic_proposer.py @@ -13,12 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Union +from typing import Optional, Union, List from vllm.config import VllmConfig from vllm.model_executor.model_loader import get_model +from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker.gpu_model_runner import logger +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch +from vllm.v1.utils import CpuGpuBuffer +from vllm.utils.platform_utils import is_pin_memory_available import numpy as np import torch @@ -39,6 +43,9 @@ def __init__( self.model = None self.device = None + self.max_batch_size = vllm_config.scheduler_config.max_num_seqs + self.backup_next_token_ids = None # type: Optional[CpuGpuBuffer] + def load_model( self, model: Union[ArcticMLPSpeculator, ArcticLSTMSpeculator], @@ -101,63 +108,186 @@ def load_model( quant_config=draft_config_quant_config, parallel_config=draft_config_parallel_config, scheduler_config=self.vllm_config.scheduler_config, - speculative_config=self.vllm_config.speculative_config, + speculative_config=self.speculative_config, load_config=self.vllm_config.load_config, device_config=self.vllm_config.device_config, ) self.model = get_model(vllm_config=draft_worker_config) - self.device = next(model.parameters()).device + self.device = next(self.model.parameters()).device self.input_hidden_dim = self.model.input_hidden_dim if isinstance( self.model, ArcticLSTMSpeculator) else self.model.emb_dim + self.backup_next_token_ids = CpuGpuBuffer( + self.max_batch_size, + dtype=torch.int32, + pin_memory=is_pin_memory_available(), + device=self.device, + with_numpy=True, + ) + def prepare_hidden_states( self, sample_hidden_states: torch.Tensor, - sampled_token_ids: Union[np.ndarray, list[list[int]]], + sampled_token_ids: Union[torch.Tensor, np.ndarray, List[List[int]]], spec_decode_metadata: SpecDecodeMetadata, ) -> torch.Tensor: - if sample_hidden_states is not None: - assert sample_hidden_states.shape[-1] == self.input_hidden_dim, \ - f"hidden_states shape mismatch: {sample_hidden_states.shape[-1]} != {self.input_hidden_dim}. \ - Please make sure spec model is trained using the same base model." - - # TODO(Ye): fuse into a single kernel + assert sample_hidden_states is not None, "sample_hidden_states must be provided" + + if isinstance(sampled_token_ids, np.ndarray): + sampled_token_ids = torch.as_tensor( + sampled_token_ids, device=sample_hidden_states.device, dtype=torch.long + ) + elif isinstance(sampled_token_ids, list): + sampled_token_ids = torch.as_tensor( + sampled_token_ids, device=sample_hidden_states.device, dtype=torch.long + ) + elif sampled_token_ids.device != sample_hidden_states.device: + sampled_token_ids = sampled_token_ids.to(sample_hidden_states.device, non_blocking=True) + max_gen_len = sampled_token_ids.shape[-1] - if max_gen_len == 1: + num_requests = sampled_token_ids.shape[0] + if max_gen_len == 1 and sample_hidden_states.shape[0] == num_requests: + # Fast path: one row per request, no index-select needed. return sample_hidden_states assert spec_decode_metadata is not None - valid_mask = sampled_token_ids != -1 - gen_lens = valid_mask.sum(dim=1) - num_sampled_tokens = np.array(spec_decode_metadata.num_draft_tokens) - num_sampled_tokens = torch.tensor(num_sampled_tokens, - device=gen_lens.device) + 1 - hidden_states_idx = (gen_lens - 1) + torch.cumsum( - num_sampled_tokens, 0) - num_sampled_tokens - previous_hidden_states = sample_hidden_states[hidden_states_idx] + if hasattr(spec_decode_metadata, "cu_num_draft_tokens") and spec_decode_metadata.cu_num_draft_tokens is not None: + cu = spec_decode_metadata.cu_num_draft_tokens + num_draft_tokens_gpu = torch.cat([cu[0:1], cu[1:] - cu[:-1]]) + else: + num_draft_tokens_gpu = torch.as_tensor( + spec_decode_metadata.num_draft_tokens, + device=sample_hidden_states.device, + dtype=torch.int64 + ) + + num_processed_tokens_per_req = num_draft_tokens_gpu + 1 + + offsets = torch.cumsum(num_processed_tokens_per_req, dim=0) - num_processed_tokens_per_req + + vocab_size = self.vllm_config.model_config.get_vocab_size() + valid_mask = (sampled_token_ids != -1) & (sampled_token_ids < vocab_size) + gen_lens = valid_mask.sum(dim=1).to(dtype=torch.int64) + + last_valid = torch.clamp(gen_lens - 1, min=0) + hidden_states_idx = offsets + last_valid + + previous_hidden_states = sample_hidden_states.index_select( + dim=0, index=hidden_states_idx + ) + + assert previous_hidden_states.size(-1) == self.input_hidden_dim, ( + f"hidden_states dim {previous_hidden_states.size(-1)} != speculator expected {self.input_hidden_dim}. " + "Make sure the spec model is trained with the same base model." + ) + return previous_hidden_states def propose( self, - context_token_ids: np.ndarray, - previous_hidden_states: torch.Tensor, + context_token_ids: Union[torch.Tensor, np.ndarray, List[int]], + previous_hidden_states: Optional[torch.Tensor], num_predict_tokens: int, - ) -> Optional[np.ndarray]: - assert num_predict_tokens > 0, \ - f"num_predict_tokens must be greater than 0, got {num_predict_tokens}." - - input_ids = torch.tensor(context_token_ids, device=self.device) + ) -> Optional[torch.Tensor]: + assert num_predict_tokens > 0 + if isinstance(context_token_ids, torch.Tensor): + if context_token_ids.device != self.device: + input_ids = context_token_ids.to(self.device, non_blocking=True) + else: + input_ids = context_token_ids + else: + input_ids = torch.as_tensor(context_token_ids, device=self.device, dtype=torch.long) next_tokens = self.model.generate_proposals( input_ids=input_ids, previous_hidden_states=previous_hidden_states, num_predict_tokens=num_predict_tokens, ) + return next_tokens + + # Borrow from eagle + def prepare_next_token_ids_cpu( + self, + sampled_token_ids: list[list[int]], + requests: dict[str, CachedRequestState], + gpu_input_batch: InputBatch, + num_scheduled_tokens: dict[str, int], + ) -> torch.Tensor: + req_ids = gpu_input_batch.req_ids + next_token_ids: list[int] = [] + for i, token_ids in enumerate(sampled_token_ids): + if token_ids: + # Common case. + next_token_id = token_ids[-1] + else: + # Partial prefill (rare case). + # Get the next token id from the request state. + req_id = req_ids[i] + req_state = requests[req_id] + seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id] + next_token_id = req_state.get_token_id(seq_len) + next_token_ids.append(next_token_id) + next_token_ids = torch.tensor( + next_token_ids, dtype=torch.int32, device=self.device + ) + return next_token_ids + + + def prepare_next_token_ids_padded( + self, + common_attn_metadata: CommonAttentionMetadata, + sampled_token_ids: torch.Tensor, + requests: dict[str, CachedRequestState], + gpu_input_batch: InputBatch, + discard_request_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.triton_utils import triton + from vllm.v1.spec_decode.utils import eagle_prepare_next_token_padded_kernel + + num_reqs = gpu_input_batch.num_reqs + 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() + ) + for i in range(num_reqs) + ], + dtype=np.int32, + ) + self.backup_next_token_ids.copy_to_gpu(num_reqs) + backup_tokens_gpu = self.backup_next_token_ids.gpu + + batch_size, num_tokens = sampled_token_ids.shape + device = sampled_token_ids.device + + assert discard_request_mask.dtype == torch.bool + assert backup_tokens_gpu.dtype == torch.int32 + + next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device) + valid_sampled_tokens_count = next_token_ids.new_empty(batch_size) + + # Kernel grid: one program per request (row) + grid = (batch_size,) + + # Find the next power of 2 for block sizes + BLOCK_SIZE_TOKENS = triton.next_power_of_2(num_tokens) + eagle_prepare_next_token_padded_kernel[grid]( + sampled_token_ids, + discard_request_mask, + backup_tokens_gpu, + next_token_ids, + valid_sampled_tokens_count, + gpu_input_batch.vocab_size, + num_tokens, + batch_size, + sampled_token_ids.stride(0), + BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS, + ) - return next_tokens.cpu().numpy() + return next_token_ids, valid_sampled_tokens_count class SuffixProposer: diff --git a/arctic_inference/vllm/spec_dec/arctic_speculator.py b/arctic_inference/vllm/spec_dec/arctic_speculator.py index 9414f59b4..22c4a09d2 100644 --- a/arctic_inference/vllm/spec_dec/arctic_speculator.py +++ b/arctic_inference/vllm/spec_dec/arctic_speculator.py @@ -255,6 +255,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: else vllm_config.scheduler_config.max_num_seqs ) self.cuda_graph_max_batch_size = padding_size(disable_by_batch_size) + self.cuda_graph_max_batch_size = padding_size( + vllm_config.scheduler_config.max_num_seqs) self.static_cuda_buffers = { "last_tokens": torch.empty(self.cuda_graph_max_batch_size, @@ -339,6 +341,7 @@ def generate_token_ids( argidx = torch.argmax(vals, -1).reshape(batch_size, -1) last_tokens = torch.gather(indices, -1, argidx) + last_tokens.clamp_(0, self.vocab_size - 1) if next_tokens_tensors[head_index] == None: next_tokens_tensors[head_index] = last_tokens else: @@ -445,11 +448,24 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.n_predict = config.n_predict self.vocab_size = config.vocab_size self.input_hidden_dim = config.input_hidden_dim - config.inner_dim = [int(i) for i in config.inner_dim.split(".")] + + def _parse_dim(value): + """Helper to normalize dimension config into a list of ints.""" + if isinstance(value, str): + return [int(i) for i in value.split(".")] + elif isinstance(value, int): + return [value] + elif isinstance(value, list): + return [int(i) for i in value] + return value + + config.inner_dim = _parse_dim(config.inner_dim) self.inner_dim = config.inner_dim - config.emb_dim = [int(i) for i in config.emb_dim.split(".")] - self.emb_dim = config.emb_dim - config.proj_dim = [int(i) for i in config.proj_dim.split(".")] + + config.emb_dim = _parse_dim(config.emb_dim) + self.emb_dim = config.emb_dim + + config.proj_dim = _parse_dim(config.proj_dim) self.proj_dim = config.proj_dim self.max_speculative_tokens = config.num_lookahead_tokens @@ -617,6 +633,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: else vllm_config.scheduler_config.max_num_seqs ) self.cuda_graph_max_batch_size = padding_size(disable_by_batch_size) + self.static_cuda_buffers = { "last_tokens": torch.empty(self.cuda_graph_max_batch_size, 1, dtype=torch.long), @@ -778,6 +795,7 @@ def generate_token_ids( next_tokens_tensors: List[torch.Tensor], cell_states: torch.Tensor = None, ) -> torch.Tensor: + last_tokens.clamp_(0, self.vocab_size - 1) for head_index in range(num_predict_tokens): if self.method == "sum_lstm": states, cell_states = self.generate_states( @@ -810,6 +828,7 @@ def generate_token_ids( argidx = torch.argmax(vals, -1).reshape(batch_size, -1) last_tokens = torch.gather(indices, -1, argidx) + last_tokens.clamp_(0, self.vocab_size - 1) if next_tokens_tensors[head_index] == None: next_tokens_tensors[head_index] = last_tokens else: diff --git a/arctic_inference/vllm/spec_dec/logits_processor_opt.py b/arctic_inference/vllm/spec_dec/logits_processor_opt.py index 4516a8656..b90c9ba86 100644 --- a/arctic_inference/vllm/spec_dec/logits_processor_opt.py +++ b/arctic_inference/vllm/spec_dec/logits_processor_opt.py @@ -44,8 +44,7 @@ def __init__(self, self.soft_cap = soft_cap # Whether to use gather or all-gather to gather the logits. - self.use_gather = not current_platform.is_tpu( - ) and not envs.VLLM_USE_V1 + self.use_gather = False self.skip_last_gather = skip_last_gather diff --git a/arctic_inference/vllm/spec_dec/vocab_parallel_embedding.py b/arctic_inference/vllm/spec_dec/vocab_parallel_embedding.py index ce59e5859..0b766229d 100644 --- a/arctic_inference/vllm/spec_dec/vocab_parallel_embedding.py +++ b/arctic_inference/vllm/spec_dec/vocab_parallel_embedding.py @@ -421,7 +421,12 @@ def forward(self, input_): self.shard_indices.added_vocab_start_index, self.shard_indices.added_vocab_end_index) else: - masked_input = input_ + # For tp_size==1 there is no masking, so clamp to the + # valid embedding range. Async scheduling + spec decode + # can leave -1 sentinel tokens in input_ids that would + # otherwise crash F.embedding. + masked_input = input_.clamp( + 0, self.num_embeddings_per_partition - 1) # Get the embeddings. output_parallel = self.quant_method.embedding(self, masked_input.long()) diff --git a/arctic_inference/vllm/swiftkv/llama_swiftkv.py b/arctic_inference/vllm/swiftkv/llama_swiftkv.py index a89a097ff..675dc4aa8 100644 --- a/arctic_inference/vllm/swiftkv/llama_swiftkv.py +++ b/arctic_inference/vllm/swiftkv/llama_swiftkv.py @@ -20,10 +20,12 @@ from torch import nn import vllm.distributed.parallel_state as parallel_state -from vllm.attention.backends.abstract import AttentionType +from vllm.v1.attention.backend import AttentionType from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.forward_context import ForwardContext, get_forward_context +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import (BatchDescriptor, ForwardContext, + get_forward_context) from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import (ColumnParallelLinear, @@ -47,7 +49,7 @@ try: from vllm.v1.attention.backends.flashinfer import FlashInferMetadata FLASHINFER_AVAILABLE = True -except ImportError: +except (ImportError, RuntimeError): FLASHINFER_AVAILABLE = False FlashInferMetadata = None @@ -76,8 +78,6 @@ def __init__( hidden_size: int, num_heads: int, num_kv_heads: int, - rope_theta: float = 10000, - rope_scaling: Optional[dict[str, Any]] = None, max_position_embeddings: int = 8192, quant_config: Optional[QuantizationConfig] = None, bias: bool = False, @@ -91,8 +91,6 @@ def __init__( hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, - rope_theta=rope_theta, - rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, quant_config=quant_config, bias=bias, @@ -148,16 +146,8 @@ def __init__( ) -> None: super().__init__() self.hidden_size = config.hidden_size - rope_theta = getattr(config, "rope_theta", 10000) - rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is not None and getattr( - config, "original_max_position_embeddings", None): - rope_scaling["original_max_position_embeddings"] = ( - config.original_max_position_embeddings) max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False) self.self_attn = LlamaSwiftKVAttention( @@ -166,8 +156,6 @@ def __init__( num_heads=config.num_attention_heads, num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), - rope_theta=rope_theta, - rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, quant_config=quant_config, bias=attention_bias, @@ -376,6 +364,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + def _init_prefill_runner(self, vllm_config: VllmConfig): vllm_config.compilation_config = copy.copy( vllm_config.compilation_config) @@ -705,6 +696,28 @@ def forward( k_states, v_states)) + # When swiftkv_select filters tokens (mixed prefill-decode batches), + # the decode runner processes fewer tokens than the original batch. + # Piecewise CUDA graphs captured for the original batch size cannot + # be replayed with modified attention metadata (stale FA3 scheduler + # metadata, changed query_start_loc, etc.), so we fall back to eager + # compiled execution for the decode runner on these batches. + # For decode-only batches all tokens survive, so CUDA graphs are + # used normally -- preserving decode throughput. + fwd_ctx = get_forward_context() + saved_batch_descriptor = fwd_ctx.batch_descriptor + saved_cudagraph_mode = fwd_ctx.cudagraph_runtime_mode + decode_num_tokens = hidden_states.shape[0] + if (saved_batch_descriptor is not None + and saved_batch_descriptor.num_tokens != decode_num_tokens): + fwd_ctx.batch_descriptor = BatchDescriptor( + num_tokens=decode_num_tokens, + num_reqs=saved_batch_descriptor.num_reqs, + uniform=saved_batch_descriptor.uniform, + has_lora=saved_batch_descriptor.has_lora, + ) + fwd_ctx.cudagraph_runtime_mode = CUDAGraphMode.NONE + with model_runner.set_shift_parallel_mode(True): hidden_states = self.decode_runner( hidden_states, @@ -714,6 +727,9 @@ def forward( v_states, ) + fwd_ctx.batch_descriptor = saved_batch_descriptor + fwd_ctx.cudagraph_runtime_mode = saved_cudagraph_mode + attn_metadata = get_attn_metadata_for_swiftkv() if attn_metadata is not None: logits_indices = attn_metadata.swiftkv_logits_indices @@ -842,6 +858,9 @@ def _init_model(self, def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.get_input_embeddings(input_ids) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + def forward( self, input_ids: torch.Tensor, diff --git a/arctic_inference/vllm/ulysses.py b/arctic_inference/vllm/ulysses.py index f4f526dff..7de12536a 100644 --- a/arctic_inference/vllm/ulysses.py +++ b/arctic_inference/vllm/ulysses.py @@ -16,14 +16,17 @@ import threading import weakref from contextlib import contextmanager -from concurrent.futures import ThreadPoolExecutor -from typing import Optional, Any +from concurrent.futures import Future +from collections import deque +from collections.abc import Callable +from typing import Optional, cast +import time import torch import vllm.distributed.parallel_state as parallel_state import vllm.envs as envs from vllm.attention.layer import Attention -from vllm.config import ModelConfig, ParallelConfig, CUDAGraphMode +from vllm.config import ModelConfig, ParallelConfig, CUDAGraphMode, VllmConfig from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.distributed.parallel_state import (init_model_parallel_group, get_world_group, @@ -31,18 +34,22 @@ destroy_distributed_environment) from vllm.v1.executor.multiproc_executor import ( set_multiprocessing_worker_envs) -from vllm.utils import get_distributed_init_method, get_open_port, get_loopback_ip +from vllm.utils.network_utils import get_distributed_init_method, get_open_port, get_loopback_ip +from vllm.utils.system_utils import get_mp_context from vllm.v1.executor.abstract import FailureCallback from vllm.v1.executor.multiproc_executor import (MultiprocExecutor, WorkerProc, - UnreadyWorkerProcHandle) -from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator -from vllm.model_executor.layers.fused_moe import FusedMoE + UnreadyWorkerProcHandle, + FutureWrapper) from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher -from vllm.forward_context import BatchDescriptor +from vllm.config.compilation import CompilationConfig +from vllm.v1.engine.core import EngineCore, EngineCoreOutputs +from vllm.v1.outputs import ModelRunnerOutput from arctic_inference.patching import ArcticPatch +# global variable to hack compilation config +_ulysses_sp_size = 1 def apply_shift_parallel_patches(): UlyssesModelConfig.apply_patch() @@ -51,6 +58,9 @@ def apply_shift_parallel_patches(): UlyssesMultiprocExecutor.apply_patch() UlyssesAttention.apply_patch() UlyssesCudagraphDispatcher.apply_patch() + UlyssesCompilationConfig.apply_patch() + UlyssesVllmConfig.apply_patch() + UlyssesEngineCore.apply_patch() class UlyssesModelConfig(ArcticPatch[ModelConfig]): @@ -99,17 +109,19 @@ class UlyssesParallelState(ArcticPatch[parallel_state]): def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, + prefill_context_model_parallel_size: int = 1, decode_context_model_parallel_size: Optional[int] = 1, backend: Optional[str] = None, ) -> None: - - from vllm.distributed.parallel_state import _DP, _EP, _PP, _TP - # Get world size and rank. Ensure some consistencies. + + from vllm.distributed.parallel_state import _DP, _EP, _PP, _TP, _DCP, _PCP + assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() rank = torch.distributed.get_rank() backend = backend or torch.distributed.get_backend( - get_world_group().device_group) + get_world_group().device_group + ) data_parallel_size = 1 from vllm.config import get_current_vllm_config @@ -117,114 +129,180 @@ def initialize_model_parallel( if config is not None: data_parallel_size = config.parallel_config.data_parallel_size - sequence_parallel_size = \ - config.parallel_config.ulysses_sequence_parallel_size + sequence_parallel_size = config.parallel_config.ulysses_sequence_parallel_size - all_ranks = torch.arange(world_size).reshape( - -1, data_parallel_size, pipeline_model_parallel_size, - sequence_parallel_size, tensor_model_parallel_size) # noqa + # vLLM types allow None, but group building needs an int + if decode_context_model_parallel_size is None: + # treat "no DCP" as DCP==TP (common interpretation) + decode_context_model_parallel_size = tensor_model_parallel_size - # Build the tensor model-parallel groups. - assert _TP is None, ("tensor model parallel group is already initialized") + # Layout order (extended from vLLM's): ExternalDP x DP x PP x PCP x SP x TP + all_ranks = torch.arange(world_size).reshape( + -1, + data_parallel_size, + pipeline_model_parallel_size, + prefill_context_model_parallel_size, + sequence_parallel_size, + tensor_model_parallel_size, + ) + + assert _TP is None, "tensor model parallel group is already initialized" group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0) group_ranks = [x.tolist() for x in group_ranks] TP_group_ranks = group_ranks - # message queue broadcaster is only used in tensor model parallel group - _TP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - use_message_queue_broadcaster=True, - group_name="tp") - - # Build the pipeline model-parallel groups. - assert _PP is None, ( - "pipeline model parallel group is already initialized") - group_ranks = all_ranks.transpose(2, 4).reshape( - -1, pipeline_model_parallel_size).unbind(0) + _TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="tp", + ) + + assert _DCP is None, "decode context model parallel group is already initialized" + group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0) + group_ranks = [x.tolist() for x in group_ranks] + DCP_group_ranks = group_ranks + _DCP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="dcp", + ) + + assert _PCP is None, "prefill context parallel group is already initialized" + group_ranks = ( + all_ranks.transpose(3, 5) + .reshape(-1, prefill_context_model_parallel_size) + .unbind(0) + ) + group_ranks = [x.tolist() for x in group_ranks] + PCP_group_ranks = group_ranks + _PCP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="pcp", + ) + + assert _PP is None, "pipeline model parallel group is already initialized" + group_ranks = ( + all_ranks.transpose(2, 5) + .reshape(-1, pipeline_model_parallel_size) + .unbind(0) + ) group_ranks = [x.tolist() for x in group_ranks] PP_group_ranks = group_ranks - _PP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="pp") - - assert _DP is None, ("data parallel group is already initialized") - group_ranks = all_ranks.transpose(1, - 4).reshape(-1, - data_parallel_size).unbind(0) + _PP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="pp", + ) + + assert _DP is None, "data parallel group is already initialized" + group_ranks = ( + all_ranks.transpose(1, 5) + .reshape(-1, data_parallel_size) + .unbind(0) + ) group_ranks = [x.tolist() for x in group_ranks] DP_group_ranks = group_ranks - _DP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="dp") - - assert _EP is None, ("expert parallel group is already initialized") - group_ranks = all_ranks.transpose(1, 3).reshape( - -1, data_parallel_size * tensor_model_parallel_size).unbind(0) + _DP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="dp", + ) + + assert _EP is None, "expert parallel group is already initialized" + group_ranks = ( + all_ranks.permute(0, 4, 2, 1, 3, 5) # ExternalDP, SP, PP, DP, PCP, TP + .reshape(-1, data_parallel_size * prefill_context_model_parallel_size * tensor_model_parallel_size) + .unbind(0) + ) group_ranks = [x.tolist() for x in group_ranks] EP_group_ranks = group_ranks - _EP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="ep") - - # Build the sequence parallel groups. - assert parallel_state._SP is None, ( - "sequence parallel group is already initialized") - group_ranks = all_ranks.transpose(3, 4).reshape( - -1, sequence_parallel_size).unbind(0) + _EP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="ep", + ) + + assert parallel_state._SP is None, "sequence parallel group is already initialized" + group_ranks = ( + all_ranks.transpose(4, 5) + .reshape(-1, sequence_parallel_size) + .unbind(0) + ) group_ranks = [x.tolist() for x in group_ranks] SP_group_ranks = group_ranks - _SP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="sp") - - # Build full-TP groups for ShiftParallel - shift_parallel_size = (tensor_model_parallel_size * - sequence_parallel_size) - assert parallel_state._SP_TP is None, ( - "full-TP group is already initialized") - # transpose(3, 4) for obtaining the correct attn head order - group_ranks = all_ranks.transpose(3, 4).reshape( - -1, shift_parallel_size).unbind(0) + _SP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="sp", + ) + + shift_parallel_size = tensor_model_parallel_size * sequence_parallel_size + assert parallel_state._SP_TP is None, "full-TP group is already initialized" + group_ranks = ( + all_ranks.transpose(4, 5) # keep same head-order trick as your old transpose(3,4) + .reshape(-1, shift_parallel_size) + .unbind(0) + ) group_ranks = [x.tolist() for x in group_ranks] SP_TP_group_ranks = group_ranks - _SP_TP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="sp_tp") + _SP_TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="sp_tp", + ) parallel_state.logger.info( "rank %s in world size %s is assigned as DP rank %s, PP rank %s, " - "TP rank %s, EP rank %s, SP rank %s, SP_TP rank %s", rank, - world_size, _DP.rank_in_group, _PP.rank_in_group, - _TP.rank_in_group, _EP.rank_in_group, _SP.rank_in_group, - _SP_TP.rank_in_group) + "PCP rank %s, TP rank %s, DCP rank %s, EP rank %s, SP rank %s, SP_TP rank %s", + rank, + world_size, + _DP.rank_in_group, + _PP.rank_in_group, + _PCP.rank_in_group, + _TP.rank_in_group, + _DCP.rank_in_group, + _EP.rank_in_group, + _SP.rank_in_group, + _SP_TP.rank_in_group, + ) parallel_state._TP = _TP + parallel_state._DCP = _DCP + parallel_state._PCP = _PCP parallel_state._PP = _PP - parallel_state._SP = _SP - parallel_state._SP_TP = _SP_TP parallel_state._DP = _DP parallel_state._EP = _EP - - # check if SP requires kv replication - num_kv_heads = config.model_config._orig_get_num_kv_heads(config.parallel_config) + parallel_state._SP = _SP + parallel_state._SP_TP = _SP_TP if get_world_group().local_rank == 0: parallel_state.logger.info( - f"UlyssesParallelState initialized:\n" - f" PP {_PP.world_size} ranks {PP_group_ranks}\n" - f" TP {_TP.world_size} ranks {TP_group_ranks}\n" - f" SP {_SP.world_size} ranks {SP_group_ranks}\n" - f" DP {_DP.world_size} ranks {DP_group_ranks}\n" - f" EP {_EP.world_size} ranks {EP_group_ranks}\n" - f" SP_TP {_SP_TP.world_size} ranks {SP_TP_group_ranks}") - if num_kv_heads < sequence_parallel_size: - parallel_state.logger.info( - f" KV cache is replicated by factor {sequence_parallel_size // num_kv_heads}\n") + "UlyssesParallelState initialized:\n" + f" PP {_PP.world_size} ranks {PP_group_ranks}\n" + f" TP {_TP.world_size} ranks {TP_group_ranks}\n" + f" DCP {_DCP.world_size} ranks {DCP_group_ranks}\n" + f" PCP {_PCP.world_size} ranks {PCP_group_ranks}\n" + f" SP {_SP.world_size} ranks {SP_group_ranks}\n" + f" DP {_DP.world_size} ranks {DP_group_ranks}\n" + f" EP {_EP.world_size} ranks {EP_group_ranks}\n" + f" SP_TP {_SP_TP.world_size} ranks {SP_TP_group_ranks}" + ) + + num_kv_heads = config.model_config._orig_get_num_kv_heads(config.parallel_config) + if get_world_group().local_rank == 0 and num_kv_heads < sequence_parallel_size: + parallel_state.logger.info( + f"KV cache is replicated by factor {sequence_parallel_size // num_kv_heads}" + ) @contextmanager def graph_capture(device: torch.device): @@ -276,66 +354,109 @@ def _init_executor(self) -> None: self._finalizer = weakref.finalize(self, self.shutdown) self.is_failed = False self.shutdown_event = threading.Event() - self.failure_callback: Optional[FailureCallback] = None - self.io_thread_pool: Optional[ThreadPoolExecutor] = None + self.failure_callback: FailureCallback | None = None self.world_size = self.parallel_config.world_size - tensor_parallel_size = self.parallel_config.tensor_parallel_size - pp_parallel_size = self.parallel_config.pipeline_parallel_size - sp_parallel_size = self.parallel_config.ulysses_sequence_parallel_size - assert (self.world_size == - tensor_parallel_size * pp_parallel_size * sp_parallel_size), ( + assert self.world_size % self.parallel_config.nnodes_within_dp == 0, ( + f"global world_size ({self.parallel_config.world_size}) must be " + f"divisible by nnodes_within_dp " + f"({self.parallel_config.nnodes_within_dp}). " + ) + self.local_world_size = self.parallel_config.local_world_size + tp_size = self.parallel_config.tensor_parallel_size + pp_size = self.parallel_config.pipeline_parallel_size + pcp_size = self.parallel_config.prefill_context_parallel_size + sp_size = self.parallel_config.ulysses_sequence_parallel_size + + assert self.world_size == tp_size * pp_size * pcp_size * sp_size, ( f"world_size ({self.world_size}) must be equal to the " - f"tensor_parallel_size ({tensor_parallel_size}) x pipeline" - f"_parallel_size ({pp_parallel_size}) x ulysses_sequence_parallel" - f"_size ({sp_parallel_size}).") + f"tensor_parallel_size ({tp_size}) x pipeline" + f"_parallel_size ({pp_size}) x prefill_context" + f"_parallel_size ({pcp_size}) x ulysses_sequence_parallel" + f"_size ({sp_size})." + ) - # Set multiprocessing envs that are common to V0 and V1 + # Set multiprocessing envs set_multiprocessing_worker_envs() - # Multiprocessing-based executor does not support multi-node setting. - # Since it only works for single node, we can use the loopback address - # get_loopback_ip() for communication. + # use the loopback address get_loopback_ip() for communication. distributed_init_method = get_distributed_init_method( - get_loopback_ip(), get_open_port()) - + get_loopback_ip(), get_open_port() + ) + self.rpc_broadcast_mq: MessageQueue | None = None + scheduler_output_handle: Handle | None = None # Initialize worker and set up message queues for SchedulerOutputs # and ModelRunnerOutputs - max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024 - self.rpc_broadcast_mq = MessageQueue(self.world_size, - self.world_size, - max_chunk_bytes=max_chunk_bytes) - scheduler_output_handle = self.rpc_broadcast_mq.export_handle() - + if self.parallel_config.node_rank_within_dp == 0: + # For leader node within each dp rank, + # each dp will have its own leader multiproc executor. + max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024 + self.rpc_broadcast_mq = MessageQueue( + self.world_size, + self.local_world_size, + max_chunk_bytes=max_chunk_bytes, + connect_ip=self.parallel_config.master_addr, + ) + scheduler_output_handle = self.rpc_broadcast_mq.export_handle() + # Create workers - unready_workers: list[UnreadyWorkerProcHandle] = [] - from vllm.utils import get_mp_context + # FIX: Removed duplicate initialization and local import that caused UnboundLocalError context = get_mp_context() shared_worker_lock = context.Lock() + unready_workers: list[UnreadyWorkerProcHandle] = [] + success = False try: - for rank in range(self.world_size): + global_start_rank = ( + self.local_world_size * self.parallel_config.node_rank_within_dp + ) + for local_rank in range(self.local_world_size): + global_rank = global_start_rank + local_rank unready_workers.append( WorkerProc.make_worker_process( vllm_config=self.vllm_config, - local_rank=rank, - rank=rank, + local_rank=local_rank, + rank=global_rank, distributed_init_method=distributed_init_method, input_shm_handle=scheduler_output_handle, shared_worker_lock=shared_worker_lock, - )) + ) + ) # Workers must be created before wait_for_ready to avoid # deadlock, since worker.init_device() does a device sync. + + # Wait for all local workers to be ready. self.workers = WorkerProc.wait_for_ready(unready_workers) + # Start background thread to monitor worker health if not in headless mode. + if self.monitor_workers: + self.start_worker_monitor() + + self.response_mqs = [] + # Only leader node have remote response mqs + if self.parallel_config.node_rank_within_dp == 0: + for rank in range(self.world_size): + if rank < self.local_world_size: + local_message_queue = self.workers[rank].worker_response_mq + assert local_message_queue is not None + self.response_mqs.append(local_message_queue) + else: + remote_message_queue = self.workers[0].peer_worker_response_mqs[ + rank + ] + assert remote_message_queue is not None + self.response_mqs.append(remote_message_queue) + # Ensure message queues are ready. Will deadlock if re-ordered # Must be kept consistent with the WorkerProc. - self.rpc_broadcast_mq.wait_until_ready() - for w in self.workers: - w.worker_response_mq.wait_until_ready() - self.start_worker_monitor() + # Wait for all input mqs to be ready. + if self.rpc_broadcast_mq is not None: + self.rpc_broadcast_mq.wait_until_ready() + # Wait for all remote response mqs to be ready. + for response_mq in self.response_mqs: + response_mq.wait_until_ready() success = True finally: if not success: @@ -344,22 +465,11 @@ def _init_executor(self) -> None: for uw in unready_workers: if uw.death_writer is not None: uw.death_writer.close() - self._ensure_worker_termination( - [uw.proc for uw in unready_workers]) - - # For pipeline parallel, we use a thread pool for asynchronous - # execute_model. - if self.max_concurrent_batches > 1: - # Note: must use only 1 IO thread to keep dequeue sequence - # from the response queue - # _async_aggregate_workers_output also assumes a single IO thread - self.io_thread_pool = ThreadPoolExecutor( - max_workers=1, thread_name_prefix="mp_exec_io") + self._ensure_worker_termination([uw.proc for uw in unready_workers]) + + self.futures_queue = deque[tuple[FutureWrapper, Callable]]() self.output_rank = self._get_output_rank() - self.has_connector = self.vllm_config.kv_transfer_config is not None - self.kv_output_aggregator = KVOutputAggregator( - self.parallel_config.world_size) class UlyssesAttention(ArcticPatch[Attention]): @@ -431,31 +541,282 @@ class UlyssesCudagraphDispatcher(ArcticPatch[CudagraphDispatcher]): def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode, uniform_decode_query_len: int): - self._orig_initialize_cudagraph_keys(cudagraph_mode, uniform_decode_query_len) - # Ulysses specific keys for mixed prefill/decode mode - if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE: - sp_size = parallel_state._SP.world_size - for bs in self.compilation_config.cudagraph_capture_sizes: - self.add_cudagraph_key( - cudagraph_mode.mixed_mode(), - BatchDescriptor(num_tokens=bs * sp_size, uniform_decode=False)) - - # Ulyssses specific keys for full decode mode - if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL \ - and cudagraph_mode.separate_routine(): - max_num_tokens = uniform_decode_query_len * \ - self.vllm_config.scheduler_config.max_num_seqs - cudagraph_capture_sizes_for_decode = [ - x for x in self.compilation_config.cudagraph_capture_sizes - if x <= max_num_tokens and x >= uniform_decode_query_len + # sp_group = getattr(parallel_state, "_SP", None) + # sp_size = sp_group.world_size if sp_group is not None else 1 + # if sp_size <= 1: + # return + + # if self.vllm_config.lora_config: + # if self.compilation_config.cudagraph_specialize_lora: + # lora_cases = [True, False] + # else: + # lora_cases = [True] + # else: + # lora_cases = [False] + + # if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE: + # for bs, has_lora in product( + # self.compilation_config.cudagraph_capture_sizes, lora_cases + # ): + # bd = self._create_padded_batch_descriptor( + # num_tokens=bs, # * sp_size, + # uniform_decode=False, + # has_lora=has_lora, + # ).relax_for_mixed_batch_cudagraphs() + + # self.add_cudagraph_key(cudagraph_mode.mixed_mode(), bd) + + # if (cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + # and cudagraph_mode.separate_routine()): + # max_num_tokens = ( + # uniform_decode_query_len + # * self.vllm_config.scheduler_config.max_num_seqs + # ) + # cudagraph_capture_sizes_for_decode = [ + # x for x in self.compilation_config.cudagraph_capture_sizes + # if uniform_decode_query_len <= x <= max_num_tokens + # ] + # for bs, has_lora in product(cudagraph_capture_sizes_for_decode, lora_cases): + # bd = self._create_padded_batch_descriptor( + # num_tokens=bs, # * sp_size, + # uniform_decode=True, + # has_lora=has_lora, + # ) + # self.add_cudagraph_key(CUDAGraphMode.FULL, bd) + + +class UlyssesCompilationConfig(ArcticPatch[CompilationConfig]): + + _orig_post_init_cudagraph_sizes = CompilationConfig.post_init_cudagraph_sizes + + def post_init_cudagraph_sizes(self) -> None: + +# # print(f"Before post_init_cudagraph_sizes: max_cudagraph_capture_size={self.max_cudagraph_capture_size}, cudagraph_capture_sizes={self.cudagraph_capture_sizes}") + +# # Access the module-level variable set during engine config creation +# sp_size = _ulysses_sp_size + +# # scale sizes by Ulysses sequence parallel size +# self.max_cudagraph_capture_size *= sp_size +# self.cudagraph_capture_sizes = [size * sp_size for size in self.cudagraph_capture_sizes] + +# # print(f"After scaling for SP size {sp_size}: max_cudagraph_capture_size={self.max_cudagraph_capture_size}, cudagraph_capture_sizes={self.cudagraph_capture_sizes}") + + self._orig_post_init_cudagraph_sizes() + +# # revert back to original shapes +# self.max_cudagraph_capture_size //= sp_size +# self.cudagraph_capture_sizes = [size // sp_size for size in self.cudagraph_capture_sizes] + +# # print(f"self.bs_to_padded_graph_size {self.bs_to_padded_graph_size}") + +# # import traceback +# # traceback.print_stack() + +class UlyssesVllmConfig(ArcticPatch[VllmConfig]): + + _orig_set_cudagraph_sizes = VllmConfig._set_cudagraph_sizes + + @staticmethod + def _generate_capture_sizes(max_size: int) -> list[int]: + sizes = [i for i in [1, 2, 4] if i <= max_size] + if max_size >= 8: + sizes += list(range(8, min(max_size + 1, 256), 8)) + if max_size >= 256: + sizes += list(range(256, min(max_size + 1, 512), 16)) + if max_size >= 512: + sizes += list(range(512, max_size + 1, 32)) + return sizes + + @staticmethod + def _build_bs_to_padded(capture_sizes: list[int], + max_capture_size: int) -> list[int]: + table = [0] * (max_capture_size + 1) + for end, start in zip( + capture_sizes + [max_capture_size + 1], + [0] + capture_sizes, + ): + for bs in range(start, end): + table[bs] = start if bs == start else end + return table + + def _set_cudagraph_sizes(self): + sp_size = _ulysses_sp_size + + max_cudagraph_capture_size = self.compilation_config.max_cudagraph_capture_size + cudagraph_capture_sizes = self.compilation_config.cudagraph_capture_sizes + + if cudagraph_capture_sizes is None: + if max_cudagraph_capture_size is None: + max_cudagraph_capture_size = 512 + # Canonical (unscaled) sizes: [1, 2, 4, 8, ..., 512] + canonical_sizes = self._generate_capture_sizes( + max_cudagraph_capture_size) + + # Base model (Ulysses): scale by sp_size + self.compilation_config.cudagraph_capture_sizes = [ + s * sp_size for s in canonical_sizes ] - for bs in cudagraph_capture_sizes_for_decode: - self.add_cudagraph_key( - CUDAGraphMode.FULL, - BatchDescriptor(num_tokens=bs * sp_size, uniform_decode=True)) - self.keys_initialized = True + self.compilation_config.max_cudagraph_capture_size = ( + max_cudagraph_capture_size * sp_size + ) + + # Shift model: scale by 1 (use canonical sizes as-is) + shift_sizes = list(canonical_sizes) + shift_max = max_cudagraph_capture_size + else: + shift_sizes = list(cudagraph_capture_sizes) + shift_max = max(shift_sizes) if shift_sizes else 0 + self._shift_cudagraph_capture_sizes = shift_sizes + self._shift_max_cudagraph_capture_size = shift_max + self._shift_bs_to_padded_graph_size = self._build_bs_to_padded( + shift_sizes, shift_max) if shift_sizes else [] + print( + f"UlyssesVllmConfig: base max_cudagraph_capture_size=" + f"{self.compilation_config.max_cudagraph_capture_size}, " + f"base sizes={self.compilation_config.cudagraph_capture_sizes}, " + f"shift max={shift_max}, shift sizes={shift_sizes}" + ) + self._orig_set_cudagraph_sizes() + + def pad_for_cudagraph(self, batch_size: int) -> int: + from .model_runner import is_shift_parallel_mode + if is_shift_parallel_mode() and self._shift_bs_to_padded_graph_size: + return self._shift_bs_to_padded_graph_size[batch_size] + return self.compilation_config.bs_to_padded_graph_size[batch_size] + + +class UlyssesEngineCore(ArcticPatch[EngineCore]): + + iteration = 0 + + def step_with_batch_queue( + self, + ) -> tuple[dict[int, EngineCoreOutputs] | None, bool]: + """Schedule and execute batches with the batch queue. + Note that if nothing to output in this step, None is returned. + + The execution flow is as follows: + 1. Try to schedule a new batch if the batch queue is not full. + If a new batch is scheduled, directly return an empty engine core + output. In other words, fulfilling the batch queue has a higher priority + than getting model outputs. + 2. If there is no new scheduled batch, meaning that the batch queue + is full or no other requests can be scheduled, we block until the first + batch in the job queue is finished. + 3. Update the scheduler from the output. + """ + batch_queue = self.batch_queue + assert batch_queue is not None + + # Try to schedule a new batch if the batch queue is not full, but + # the scheduler may return an empty batch if all requests are scheduled. + # Note that this is not blocking. + assert len(batch_queue) < self.batch_queue_size + + step_start_time = time.monotonic() + + model_executed = False + deferred_scheduler_output = None + if self.scheduler.has_requests(): + scheduler_output = self.scheduler.schedule() + exec_future = self.model_executor.execute_model( + scheduler_output, non_block=True + ) + if not self.is_ec_producer: + model_executed = scheduler_output.total_num_scheduled_tokens > 0 + + if self.is_pooling_model or not model_executed: + # No sampling required (no requests scheduled). + future = cast(Future[ModelRunnerOutput], exec_future) + else: + if not scheduler_output.pending_structured_output_tokens: + # We aren't waiting for any tokens, get any grammar output + # and sample immediately. + grammar_output = self.scheduler.get_grammar_bitmask( + scheduler_output + ) + future = self.model_executor.sample_tokens( + grammar_output, non_block=True + ) + else: + # We need to defer sampling until we have processed the model output + # from the prior step. + deferred_scheduler_output = scheduler_output + + if not deferred_scheduler_output: + # Add this step's future to the queue. + batch_queue.appendleft((future, scheduler_output, exec_future)) + if ( + model_executed + and len(batch_queue) < self.batch_queue_size + and not batch_queue[-1][0].done() + ): + # Don't block on next worker response unless the queue is full + # or there are no more requests to schedule. + return None, True + + elif not batch_queue: + # Queue is empty. We should not reach here since this method should + # only be called when the scheduler contains requests or the queue + # is non-empty. + return None, False + + # Block until the next result is available. + future, scheduler_output, exec_model_fut = batch_queue.pop() + with ( + self.log_error_detail(scheduler_output), + self.log_iteration_details(scheduler_output), + ): + model_output = future.result() + if model_output is None: + # None from sample_tokens() implies that the original execute_model() + # call failed - raise that exception. + exec_model_fut.result() + raise RuntimeError("unexpected error") + + # Before processing the model output, process any aborts that happened + # during the model execution. + self._process_aborts_queue() + engine_core_outputs = self.scheduler.update_from_output( + scheduler_output, model_output + ) + + # NOTE(nick): We can either handle the deferred tasks here or save + # in a field and do it immediately once step_with_batch_queue is + # re-called. The latter slightly favors TTFT over TPOT/throughput. + if deferred_scheduler_output: + # If we are doing speculative decoding with structured output, + # we need to get the draft token ids from the prior step before + # we can compute the grammar bitmask for the deferred request. + if self.use_spec_decode: + draft_token_ids = self.model_executor.take_draft_token_ids() + assert draft_token_ids is not None + # Update the draft token ids in the scheduler output to + # filter out the invalid spec tokens, which will be padded + # with -1 and skipped by the grammar bitmask computation. + self.scheduler.update_draft_token_ids_in_output( + draft_token_ids, deferred_scheduler_output + ) + # We now have the tokens needed to compute the bitmask for the + # deferred request. Get the bitmask and call sample tokens. + grammar_output = self.scheduler.get_grammar_bitmask( + deferred_scheduler_output + ) + future = self.model_executor.sample_tokens(grammar_output, non_block=True) + batch_queue.appendleft((future, deferred_scheduler_output, exec_future)) + + total_time_ms = (time.monotonic() - step_start_time) * 1000 + + running, waiting = self.scheduler.get_request_counts() + scheduled_tokens = scheduler_output.total_num_scheduled_tokens + concurrency = len(scheduler_output.num_scheduled_tokens.keys()) + # print(f"iteration {self.iteration}, running: {running}, waiting: {waiting}, scheduled tokens: {scheduled_tokens}, concurrency: {concurrency}, total_time_ms: {total_time_ms:.2f}") + self.iteration += 1 + + return engine_core_outputs, model_executed \ No newline at end of file diff --git a/benchmark/rollout/README.md b/benchmark/rollout/README.md new file mode 100644 index 000000000..96fc91d47 --- /dev/null +++ b/benchmark/rollout/README.md @@ -0,0 +1,51 @@ +## Rollout Replay Patch for v0.14.1 + +This patch extencds SamplingParams to specify the length of each sequence when n > 1. + +The patch is applied as `source patch_sampling.sh`. + +As a result, you can specify `max_tokens_n` as a list in sampling params and set `ignore_eos` so that each sequence generates exactly specified number of tokens. + +``` +# Sample prompts. +prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + # "The future of AI is", +] + +sampling_params = [SamplingParams(n=2, + temperature=0.8, + top_p=1.0, + max_tokens_n=[25, 50], + ignore_eos=True, + ), + SamplingParams(n=3, + temperature=0.8, + top_p=1.0, + max_tokens_n=[5, 10, 15], + ignore_eos=True, + ), + SamplingParams(n=1, + temperature=0.8, + top_p=1.0, + max_tokens=100, + ignore_eos=True, + # max_tokens_n=[100], this will be ineffective since n = 1 + ), + ] + +outputs = llm.generate(prompts, sampling_params=sampling_params) +``` + +The number of resulting input and output tokens per sequence: +``` +prompt 0 seq 0: input 5 output 25 +prompt 0 seq 1: input 5 output 50 +prompt 1 seq 0: input 7 output 5 +prompt 1 seq 1: input 7 output 10 +prompt 1 seq 2: input 7 output 15 +prompt 2 seq 0: input 5 output 100 +``` + diff --git a/benchmark/rollout/parallel_sampling.patch b/benchmark/rollout/parallel_sampling.patch new file mode 100644 index 000000000..8c7bbb7a5 --- /dev/null +++ b/benchmark/rollout/parallel_sampling.patch @@ -0,0 +1,18 @@ +--- /home/yak/myenv_v14/lib/python3.12/site-packages/vllm/v1/engine/parallel_sampling.py 2026-01-31 22:11:45.340329657 +0000 ++++ parallel_sampling.py 2026-01-31 22:08:27.000000000 +0000 +@@ -66,11 +66,13 @@ + Child `sampling_params` instance. + """ + seed = self.sampling_params.seed +- if self.cached_child_sampling_params: ++ # if self.cached_child_sampling_params: + # Reuse child sampling_params data structure +- return self.cached_child_sampling_params ++ # return self.cached_child_sampling_params + # Build child sampling_params + child_sampling_params = copy(self.sampling_params) ++ print(f"child index: {index}, max_tokens_n: {child_sampling_params.max_tokens_n}") ++ child_sampling_params.max_tokens = child_sampling_params.max_tokens_n[index] + child_sampling_params.n = 1 + if seed is None: + # Cache child sampling_params for later reuse diff --git a/benchmark/rollout/patch_sampling.sh b/benchmark/rollout/patch_sampling.sh new file mode 100755 index 000000000..531bed562 --- /dev/null +++ b/benchmark/rollout/patch_sampling.sh @@ -0,0 +1,12 @@ + +VLLM_PATH="$(pip show vllm | awk '/^Location: /{print $2}')" + +if [ -z "$VLLM_PATH" ]; then + echo "Error: could not find VLLM in current env" + exit 1 +else + echo "VLLM path is: $VLLM_PATH" +fi + +patch $VLLM_PATH/vllm/sampling_params.py < sampling_params.patch +patch $VLLM_PATH/vllm/v1/engine/parallel_sampling.py < parallel_sampling.patch diff --git a/benchmark/rollout/sampling_params.patch b/benchmark/rollout/sampling_params.patch new file mode 100644 index 000000000..7647b8b6d --- /dev/null +++ b/benchmark/rollout/sampling_params.patch @@ -0,0 +1,26 @@ +--- /home/yak/myenv_v14/lib/python3.12/site-packages/vllm/sampling_params.py 2026-01-31 22:11:45.180329435 +0000 ++++ sampling_params.py 2026-01-31 22:16:23.000000000 +0000 +@@ -124,6 +124,7 @@ + """ + + n: int = 1 ++ max_tokens_n: list[int] | None = None + """Number of outputs to return for the given prompt request. + + NOTE: +@@ -250,6 +251,7 @@ + @staticmethod + def from_optional( + n: int | None = 1, ++ max_tokens_n: list[int] | None = None, + presence_penalty: float | None = 0.0, + frequency_penalty: float | None = 0.0, + repetition_penalty: float | None = 1.0, +@@ -289,6 +291,7 @@ + + return SamplingParams( + n=1 if n is None else n, ++ max_tokens_n=max_tokens_n, + presence_penalty=0.0 if presence_penalty is None else presence_penalty, + frequency_penalty=0.0 if frequency_penalty is None else frequency_penalty, + repetition_penalty=1.0 diff --git a/projects/spec_dec/offline_inference_spec_dec.py b/projects/spec_dec/offline_inference_spec_dec.py index 53604410a..06ed3e397 100644 --- a/projects/spec_dec/offline_inference_spec_dec.py +++ b/projects/spec_dec/offline_inference_spec_dec.py @@ -17,7 +17,8 @@ from vllm import LLM, SamplingParams import os -os.environ["VLLM_USE_V1"] = "1" +os.environ["ARCTIC_INFERENCE_ENABLED"] = "1" +os.environ["CUDA_VISIBLE_DEVICES"] = "4,5" vllm.plugins.load_general_plugins() @@ -29,9 +30,11 @@ "method": "arctic", "model": "Snowflake/Arctic-LSTM-Speculator-Llama-3.1-70B-Instruct", "num_speculative_tokens": 3, - "enable_suffix_decoding": True, + "enable_suffix_decoding": False, "disable_by_batch_size": 64, }, + enforce_eager=True, + async_scheduling=True, seed=0, ) diff --git a/projects/swiftkv/offline_inference_swiftkv.py b/projects/swiftkv/offline_inference_swiftkv.py index b3f86de34..e79853cd3 100644 --- a/projects/swiftkv/offline_inference_swiftkv.py +++ b/projects/swiftkv/offline_inference_swiftkv.py @@ -18,7 +18,7 @@ vllm.plugins.load_general_plugins() -llm = LLM(model="Snowflake/Llama-3.1-SwiftKV-8B-Instruct") +llm = LLM(model="Snowflake/Llama-3.1-SwiftKV-8B-Instruct", tensor_parallel_size=2) print("=" * 80) diff --git a/pyproject.toml b/pyproject.toml index 95c730f45..cb9329395 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ requires = [ # uv pip install vllm==0.9.1 --dry-run 2>&1 | grep protobuf | sed 's/^ + //' "protobuf==5.29.5", "grpcio-tools", - "torch==2.8.0", # need to be same as user env. + "torch==2.9.1", # need to be same as user env. ] build-backend = "setuptools.build_meta" @@ -43,7 +43,7 @@ packages = [ [project.optional-dependencies] vllm = [ - 'vllm==0.11.0', + 'vllm==0.14.1', ] docs = [ diff --git a/tests/benchmarks/test_benchmarks.py b/tests/benchmarks/test_benchmarks.py index e2ea1c7fa..891e83d49 100644 --- a/tests/benchmarks/test_benchmarks.py +++ b/tests/benchmarks/test_benchmarks.py @@ -9,7 +9,7 @@ import uvloop from vllm.entrypoints.openai.api_server import ( make_arg_parser, run_server, validate_parsed_serve_args) -from vllm.utils import FlexibleArgumentParser +from vllm.utils.argparse_utils import FlexibleArgumentParser from .benchmark_utils import (ACCURACY_TASKS, PERFORMANCE_TASKS, VLLM_CONFIGS, JSON_MODE_TASKS, update_benchmark_summary) diff --git a/tests/unit_tests/test_arctic_spec_max_len.py b/tests/unit_tests/test_arctic_spec_max_len.py index 43491f86b..591f18aba 100644 --- a/tests/unit_tests/test_arctic_spec_max_len.py +++ b/tests/unit_tests/test_arctic_spec_max_len.py @@ -49,13 +49,13 @@ def sampling_configs(): @pytest.fixture def model_name(): - return "Snowflake/Llama-3.1-SwiftKV-8B-Instruct" + return "meta-llama/Llama-3.3-70B-Instruct" # Define the speculative configurations that will be tested ARCTIC_SPEC_CONFIG = { "method": "arctic", - "model": "Snowflake/Arctic-LSTM-Speculator-Llama-3.1-8B-Instruct", + "model": "Snowflake/Arctic-LSTM-Speculator-Llama-3.3-70B-Instruct", "num_speculative_tokens": 3, "disable_by_batch_size": 64, "enable_suffix_decoding": True, @@ -85,6 +85,7 @@ def test_speculative_decoding( This test is parameterized to cover 'arctic' and 'suffix' methods. ''' with monkeypatch.context() as m: + m.setenv("ARCTIC_INFERENCE_ENABLED", "1") m.setenv("VLLM_PLUGINS", "arctic_inference") m.setenv("VLLM_USE_V1", "1") @@ -92,11 +93,12 @@ def test_speculative_decoding( spec_llm = LLM( model=model_name, - tensor_parallel_size=1, + tensor_parallel_size=2, quantization="fp8", speculative_config=spec_config, max_model_len=MAX_MODEL_LEN, enforce_eager=True, + trust_remote_code=True, ) for sampling_config in sampling_configs: