Skip to content

vllm.v1.worker.gpu.spec_decode.utils

Classes:

Functions:

DraftTokensHandler

Methods:

  • get_draft_tokens –

    Return the latest drafts needed for the current step's deferred

Source code in vllm/v1/worker/gpu/spec_decode/utils.py 
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class DraftTokensHandler:
    def __init__(self, device: torch.device | None = None):
        self.device = device
        self.copy_stream = torch.cuda.Stream(device)
        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
        self.copy_event = torch.cuda.Event(blocking=True)
        self._copy_pending = False

        # Latest drafts per live request: a row view into an async-copied
        # pinned array (synced in get_draft_tokens), or just the draft width
        # when the scheduler doesn't need the token ids.
        self._draft_tokens_by_req: dict[str, np.ndarray | int] = {}
        # Requests whose drafts the scheduler will fetch for the latest step's
        # deferred grammar bitmask; None means all requests.
        self._deferred_req_ids: Iterable[str] | None = None

    def begin_step(self, scheduler_output: SchedulerOutput) -> None:
        # get_draft_tokens is only called for the most recently executed
        # scheduler output.
        self._deferred_req_ids = (
            scheduler_output.scheduled_spec_decode_tokens.keys()
            if scheduler_output.pending_structured_output_tokens
            else None
        )

    def set_draft_tokens(
        self, input_batch: InputBatch, draft_tokens: torch.Tensor
    ) -> None:
        if not input_batch.has_structured_output_reqs:
            # No draft token validation needs to be performed by
            # the scheduler for this batch. Record draft token widths only.
            num_draft_tokens = draft_tokens.shape[1]
            for req_id in input_batch.req_ids:
                self._draft_tokens_by_req[req_id] = num_draft_tokens
            return

        # For spec decoding + structured outputs, we must transfer the
        # draft tokens back to the scheduler for grammar validation.
        current_stream = torch.cuda.current_stream(self.device)
        self.copy_stream.wait_stream(current_stream)
        with torch.cuda.stream(self.copy_stream):
            draft_tokens_np = async_copy_to_np(draft_tokens)
            # draft_tokens is a temporary allocation on the main stream and read here on
            # copy_stream; without record_stream, the caching allocator may reuse its
            # memory before the async copy executes.
            draft_tokens.record_stream(self.copy_stream)
            self.copy_event.record()
            self._copy_pending = True

        for req_id, row in zip(input_batch.req_ids, draft_tokens_np):
            self._draft_tokens_by_req[req_id] = row

    def discard(self, req_id: str) -> None:
        self._draft_tokens_by_req.pop(req_id, None)

    def get_draft_tokens(self) -> DraftTokenIds | None:
        """Return the latest drafts needed for the current step's deferred
        grammar bitmask, or for all live requests if sampling isn't deferred.
        Requests without retained drafts are omitted."""
        if self._copy_pending:
            self.copy_event.synchronize()
            self._copy_pending = False

        drafts_by_req = self._draft_tokens_by_req
        req_ids = self._deferred_req_ids
        if req_ids is None:
            out_req_ids = list(drafts_by_req)
        else:
            out_req_ids = [req_id for req_id in req_ids if req_id in drafts_by_req]
        draft_token_ids: list[list[int]] = []
        for req_id in out_req_ids:
            draft_tokens = drafts_by_req[req_id]
            if isinstance(draft_tokens, int):
                # Construct placeholders from draft token width
                draft_tokens = [-1] * draft_tokens
            else:
                # Convert numpy view to list
                draft_tokens = draft_tokens.tolist()
            draft_token_ids.append(draft_tokens)

        return DraftTokenIds(out_req_ids, draft_token_ids)

get_draft_tokens()

Return the latest drafts needed for the current step's deferred grammar bitmask, or for all live requests if sampling isn't deferred. Requests without retained drafts are omitted.

Source code in vllm/v1/worker/gpu/spec_decode/utils.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def get_draft_tokens(self) -> DraftTokenIds | None:
    """Return the latest drafts needed for the current step's deferred
    grammar bitmask, or for all live requests if sampling isn't deferred.
    Requests without retained drafts are omitted."""
    if self._copy_pending:
        self.copy_event.synchronize()
        self._copy_pending = False

    drafts_by_req = self._draft_tokens_by_req
    req_ids = self._deferred_req_ids
    if req_ids is None:
        out_req_ids = list(drafts_by_req)
    else:
        out_req_ids = [req_id for req_id in req_ids if req_id in drafts_by_req]
    draft_token_ids: list[list[int]] = []
    for req_id in out_req_ids:
        draft_tokens = drafts_by_req[req_id]
        if isinstance(draft_tokens, int):
            # Construct placeholders from draft token width
            draft_tokens = [-1] * draft_tokens
        else:
            # Convert numpy view to list
            draft_tokens = draft_tokens.tolist()
        draft_token_ids.append(draft_tokens)

    return DraftTokenIds(out_req_ids, draft_token_ids)

get_drafter_hidden_states(model, hidden_states)

Hidden states fed to the drafter.

Targets such as DeepSeek V4 expose the pre-hc_head residual through get_mtp_target_hidden_states(). The buffer is sized at max_num_batched_tokens and only allocated for drafters that consume target hidden states, so None means "use the regular hidden states".

Source code in vllm/v1/worker/gpu/spec_decode/utils.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def get_drafter_hidden_states(
    model: torch.nn.Module, hidden_states: torch.Tensor
) -> torch.Tensor:
    """Hidden states fed to the drafter.

    Targets such as DeepSeek V4 expose the pre-hc_head residual through
    get_mtp_target_hidden_states(). The buffer is sized at
    max_num_batched_tokens and only allocated for drafters that consume
    target hidden states, so None means "use the regular hidden states".
    """
    get_target_hidden_states = getattr(model, "get_mtp_target_hidden_states", None)
    if get_target_hidden_states is None:
        return hidden_states
    target_hidden_states = get_target_hidden_states()
    if target_hidden_states is None:
        return hidden_states
    return target_hidden_states[: hidden_states.shape[0]]

get_parallel_drafting_token_id(hf_config)

Resolve the mask token id used for parallel drafting slots.

Checks (in order): dflash_config.mask_token_id, top-level mask_token_id, dspark_noise_token_id, pard_token, ptd_token_id. Raises ValueError if none are present.

Source code in vllm/v1/worker/gpu/spec_decode/utils.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def get_parallel_drafting_token_id(hf_config) -> int:
    """Resolve the mask token id used for parallel drafting slots.

    Checks (in order): `dflash_config.mask_token_id`, top-level `mask_token_id`,
    `dspark_noise_token_id`, `pard_token`, `ptd_token_id`. Raises ValueError if
    none are present.
    """
    dflash_config = getattr(hf_config, "dflash_config", None) or {}
    if "mask_token_id" in dflash_config:
        return int(dflash_config["mask_token_id"])
    if getattr(hf_config, "mask_token_id", None) is not None:
        return int(hf_config.mask_token_id)
    if hasattr(hf_config, "dspark_noise_token_id"):
        return int(hf_config.dspark_noise_token_id)
    if hasattr(hf_config, "pard_token"):
        return int(hf_config.pard_token)
    if hasattr(hf_config, "ptd_token_id"):
        return int(hf_config.ptd_token_id)
    raise ValueError(
        "Model config must specify `dflash_config.mask_token_id`,"
        " `mask_token_id`, `dspark_noise_token_id`, `pard_token`, or"
        " `ptd_token_id` for parallel drafting."
    )

get_pp_safe_draft_load_config(load_config)

Avoid collectives that include PP ranks without a draft model.

Source code in vllm/v1/worker/gpu/spec_decode/utils.py 
38
39
40
41
42
43
44
45
46
47
48
def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig:
    """Avoid collectives that include PP ranks without a draft model."""
    if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors":
        logger.warning_once(
            "fastsafetensors cannot load a draft model instantiated on only "
            "one pipeline stage; falling back to the standard safetensors "
            "loader for the draft model. The target model still uses "
            "fastsafetensors."
        )
        return replace(load_config, load_format="auto")
    return load_config