vllm.v1.spec_decode.draft_model ¶

class DraftModelProposer(SpecDecodeBaseProposer):
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        runner=None,
    ):
        super().__init__(
            vllm_config=vllm_config,
            device=device,
            pass_hidden_states_to_model=False,
            runner=runner,
        )
        self._raise_if_multimodal()
        self._raise_if_mrope()
        self._raise_if_padded_drafter_batch_disabled()
        self._raise_if_vocab_size_mismatch()
        self._raise_if_draft_tp_mismatch()

    def _block_size(self) -> int:
        builder = self._get_attention_metadata_builder()
        return builder.kv_cache_spec.block_size

    def _raise_if_multimodal(self):
        if self.supports_mm_inputs:
            raise NotImplementedError(
                "Speculative Decoding with draft models "
                "does not support multimodal models yet"
            )

    def _raise_if_mrope(self):
        if self.draft_model_config.uses_mrope:
            raise NotImplementedError(
                "Speculative Decoding with draft models does not support M-RoPE yet"
            )

    def _raise_if_padded_drafter_batch_disabled(self):
        if self.vllm_config.speculative_config.disable_padded_drafter_batch:
            raise NotImplementedError(
                "Speculative Decoding with draft models only supports "
                "padded drafter batch. Please don't pass --disable-padded-drafter-batch"
                " in the speculative_config."
            )

    def _raise_if_vocab_size_mismatch(self):
        self.vllm_config.speculative_config.verify_equal_vocab_size_if_draft_model()

    def _raise_if_draft_tp_mismatch(self):
        # Note(Tomas Ruiz) If we run the target model with TP > 1 and
        # the draft model with TP = 1, then the different TP ranks collide.
        # Specifically when all ranks compile the draft model on rank 0
        # (because TP=1), then the torch compile cache is overwritten and corrupted.
        # We need a mechanism like this: https://github.com/vllm-project/vllm/pull/5414
        # To prevent this error, we assert that both TP sizes must be the same.
        spec_cfg: SpeculativeConfig = self.vllm_config.speculative_config
        tgt_tp = spec_cfg.target_parallel_config.tensor_parallel_size
        draft_tp = spec_cfg.draft_parallel_config.tensor_parallel_size
        if draft_tp != tgt_tp:
            raise ValueError(
                f"Currently, 'draft_tensor_parallel_size' and 'tensor_parallel_size' "
                f"must be the same. Got {draft_tp} and {tgt_tp}. "
                "Please pass 'draft_tensor_parallel_size' in the speculative_config."
            )

    def set_inputs_first_pass(
        self,
        target_token_ids: torch.Tensor,
        next_token_ids: torch.Tensor,
        target_positions: torch.Tensor,
        last_token_indices: torch.Tensor | None,
        cad: CommonAttentionMetadata,
        num_rejected_tokens_gpu: torch.Tensor | None,
    ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
        batch_size = cad.batch_size()
        grid = (batch_size,)
        start_locs = cad.query_start_loc[:-1]
        end_locs = cad.query_start_loc[1:] - 1
        if num_rejected_tokens_gpu is not None:
            end_locs -= num_rejected_tokens_gpu

        num_tokens = target_token_ids.shape[0] + batch_size
        is_rejected_tok = torch.empty(
            (num_tokens,), device=self.input_ids.device, dtype=torch.bool
        )
        merge_toks_kernel[grid](
            target_toks_ptr=target_token_ids,
            next_toks_ptr=next_token_ids,
            query_start_locs_ptr=start_locs,
            query_end_locs_ptr=end_locs,
            out_ptr_merged_toks=self.input_ids,
            out_ptr_is_rejected_tok=is_rejected_tok,
            target_toks_size=target_token_ids.shape[0],
            # passing a negative rejected_tok_fill value will raise an error
            # when the value is used to index into embeddings.
            # Therefore, we pass a valid integer, e.g. 0.
            rejected_tok_fill=0,
        )
        merge_toks_kernel[grid](
            target_toks_ptr=target_positions,
            next_toks_ptr=target_positions[end_locs] + 1,
            query_start_locs_ptr=start_locs,
            query_end_locs_ptr=end_locs,
            out_ptr_merged_toks=self.positions,
            out_ptr_is_rejected_tok=is_rejected_tok,
            target_toks_size=target_positions.shape[0],
            rejected_tok_fill=0,
        )

        # recompute slot mapping
        new_slot_mapping = compute_new_slot_mapping(
            cad=cad,
            new_positions=self.positions[:num_tokens],
            is_rejected_token_mask=is_rejected_tok,
            block_size=self._block_size(),
            max_model_len=self.max_model_len,
        )
        # update common_attn_metadata
        new_cad: CommonAttentionMetadata = extend_all_queries_by_1(
            cad,
            arange=self.arange,
            new_slot_mapping=new_slot_mapping,
        )

        new_last_token_indices = new_cad.query_start_loc[1:] - 1
        if num_rejected_tokens_gpu is not None:
            new_last_token_indices -= num_rejected_tokens_gpu

        return num_tokens, new_last_token_indices, new_cad

    def load_model(self, target_model: Any) -> None:
        """Takes target_model to satisfy the type checker."""

        # This must be computed before loading the draft model
        # because that mutates the forward_context of the vllm_config
        target_attn_layer_names = set(
            get_layers_from_vllm_config(self.vllm_config, Attention).keys()
        )

        from vllm.compilation.backends import set_model_tag

        draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(
            target_model_vllm_config=self.vllm_config
        )
        logger.info(
            "Starting to load draft model %s. TP=%d, rank=%d",
            draft_vllm_config.model_config.model,
            draft_vllm_config.parallel_config.tensor_parallel_size,
            draft_vllm_config.parallel_config.rank,
        )
        with set_model_tag("draft_model"):
            self.model = get_model(vllm_config=draft_vllm_config, prefix="draft_model")

        # This must be computed after loading the draft model
        # because that mutates the forward_context of the vllm_config
        draft_attn_layer_names = (
            get_layers_from_vllm_config(self.vllm_config, Attention).keys()
            - target_attn_layer_names
        )
        self.attn_layer_names = list(draft_attn_layer_names)