Skip to content

vllm.v1.worker.gpu.spec_decode.draft_logits

Classes:

  • DraftLogitsCache –

    Pre-temperature draft logits, cached by a draft step for the next

Functions:

  • draft_logits_row –

    Row of a request in this rank's draft logits cache, or -1 if no row.

DraftLogitsCache

Pre-temperature draft logits, cached by a draft step for the next verification.

A replicated cache holds every request on every TP rank, at row req_state_idx.

A sharded cache (batch-sharded sampling) holds on each TP rank only the requests that rank verifies. Every draft step gives the request at batch position i to rank i % tp, at row i // tp of its cache, so no rank caches more than cdiv(num_reqs, tp) requests however the request slots are spread.

Rows are reused every step. A request that the scheduler skips for a step keeps its drafts, but its row may since have been given to another request. inverse_row_mapping detects this, and stale_drafts marks those drafts invalid, so verification rejects them and samples from the target.

Methods:

  • assign –

    Give a sharded cache's rows out by batch position. Call after

  • owners_of –

    The TP rank that verifies each of a batch's requests: the rank that

  • stale_drafts –

    Mark the drafts whose draft logits were overwritten since drafting.

Source code in vllm/v1/worker/gpu/spec_decode/draft_logits.py 
 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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
class DraftLogitsCache:
    """Pre-temperature draft logits, cached by a draft step for the next
    verification.

    A replicated cache holds every request on every TP rank, at row
    req_state_idx.

    A sharded cache (batch-sharded sampling) holds on each TP rank only the
    requests that rank verifies. Every draft step gives the request at batch
    position i to rank i % tp, at row i // tp of its cache, so no rank caches
    more than cdiv(num_reqs, tp) requests however the request slots are spread.

    Rows are reused every step. A request that the scheduler skips for a step
    keeps its drafts, but its row may since have been given to another
    request. `inverse_row_mapping` detects this, and `stale_drafts` marks those
    drafts invalid, so verification rejects them and samples from the target.
    """

    def __init__(
        self,
        max_num_reqs: int,
        num_steps: int,
        vocab_size: int,
        *,
        dtype: torch.dtype,
        fill: float,
        device: torch.device | str,
        sharded: bool = False,
    ):
        num_rows = max_num_reqs
        # [max_num_reqs] Each request's row in values, -1 if another rank
        # caches it. None for a replicated cache.
        self.row_mapping: torch.Tensor | None = None
        if sharded:
            tp_group = get_tp_group()
            self._num_shards = tp_group.world_size
            self._shard_idx = tp_group.rank_in_group
            num_rows = cdiv(max_num_reqs, self._num_shards)
            # Requests never drafted keep a slot-based owner. They have no
            # drafts, so any owner works as long as every rank agrees.
            slots = np.arange(max_num_reqs, dtype=np.int32)
            # [max_num_reqs] The TP rank that caches, and verifies, each request.
            self.shard_mapping_np = slots % self._num_shards
            self.shard_mapping = torch.from_numpy(self.shard_mapping_np).to(device)
            rows_np = np.where(
                self.shard_mapping_np == self._shard_idx,
                shard_rows(max_num_reqs, self._num_shards),
                -1,
            )
            self.row_mapping = torch.from_numpy(rows_np.astype(np.int32)).to(device)
            # [num_rows] The request last drafted into each row of this rank.
            self.inverse_row_mapping = torch.full(
                (num_rows,), -1, dtype=torch.int32, device=device
            )
        # [num_rows, num_steps, vocab_size]
        self.values = torch.full(
            (num_rows, num_steps, vocab_size), fill, dtype=dtype, device=device
        )

    def assign(self, idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray) -> None:
        """Give a sharded cache's rows out by batch position. Call after
        verification and before the draft step writes the draft logits."""
        num_reqs = idx_mapping_np.shape[0]
        if self.row_mapping is None or num_reqs == 0:
            return
        self.shard_mapping_np[idx_mapping_np] = shard_owners(num_reqs, self._num_shards)
        block_size = 1024
        _assign_rows_kernel[(triton.cdiv(num_reqs, block_size),)](
            idx_mapping,
            self.shard_mapping,
            self.row_mapping,
            self.inverse_row_mapping,
            num_reqs,
            NUM_SHARDS=self._num_shards,
            SHARD_IDX=self._shard_idx,
            BLOCK_SIZE=block_size,
        )

    def owners_of(
        self, idx_mapping_np: np.ndarray
    ) -> tuple[np.ndarray, torch.Tensor | None]:
        """The TP rank that verifies each of a batch's requests: the rank that
        cached its draft logits, or for a replicated cache, round-robin over
        batch positions.

        Returns:
            [num_reqs] each request's owner, and the [max_num_reqs] device owner
            table, or None for a replicated cache.

        """
        if self.row_mapping is None:
            num_reqs = idx_mapping_np.shape[0]
            return shard_owners(num_reqs, get_tp_group().world_size), None
        return self.shard_mapping_np[idx_mapping_np], self.shard_mapping

    def stale_drafts(self, input_batch: InputBatch) -> torch.Tensor | None:
        """Mark the drafts whose draft logits were overwritten since drafting.
        Only valid for requests this rank owns.

        Returns:
            [num_logits] True for each stale draft, or None for a replicated
            cache, whose rows are never reused.

        """
        if self.row_mapping is None:
            return None
        expanded_idx_mapping = input_batch.expanded_idx_mapping
        num_logits = expanded_idx_mapping.shape[0]
        stale = expanded_idx_mapping.new_empty(num_logits, dtype=torch.bool)
        if num_logits == 0:
            return stale
        block_size = 1024
        _stale_drafts_kernel[(triton.cdiv(num_logits, block_size),)](
            stale,
            expanded_idx_mapping,
            input_batch.expanded_local_pos,
            self.row_mapping,
            self.inverse_row_mapping,
            num_logits,
            BLOCK_SIZE=block_size,
        )
        return stale

assign(idx_mapping, idx_mapping_np)

Give a sharded cache's rows out by batch position. Call after verification and before the draft step writes the draft logits.

Source code in vllm/v1/worker/gpu/spec_decode/draft_logits.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def assign(self, idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray) -> None:
    """Give a sharded cache's rows out by batch position. Call after
    verification and before the draft step writes the draft logits."""
    num_reqs = idx_mapping_np.shape[0]
    if self.row_mapping is None or num_reqs == 0:
        return
    self.shard_mapping_np[idx_mapping_np] = shard_owners(num_reqs, self._num_shards)
    block_size = 1024
    _assign_rows_kernel[(triton.cdiv(num_reqs, block_size),)](
        idx_mapping,
        self.shard_mapping,
        self.row_mapping,
        self.inverse_row_mapping,
        num_reqs,
        NUM_SHARDS=self._num_shards,
        SHARD_IDX=self._shard_idx,
        BLOCK_SIZE=block_size,
    )

owners_of(idx_mapping_np)

The TP rank that verifies each of a batch's requests: the rank that cached its draft logits, or for a replicated cache, round-robin over batch positions.

Returns:

  • ndarray –

    [num_reqs] each request's owner, and the [max_num_reqs] device owner

  • Tensor | None –

    table, or None for a replicated cache.

Source code in vllm/v1/worker/gpu/spec_decode/draft_logits.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def owners_of(
    self, idx_mapping_np: np.ndarray
) -> tuple[np.ndarray, torch.Tensor | None]:
    """The TP rank that verifies each of a batch's requests: the rank that
    cached its draft logits, or for a replicated cache, round-robin over
    batch positions.

    Returns:
        [num_reqs] each request's owner, and the [max_num_reqs] device owner
        table, or None for a replicated cache.

    """
    if self.row_mapping is None:
        num_reqs = idx_mapping_np.shape[0]
        return shard_owners(num_reqs, get_tp_group().world_size), None
    return self.shard_mapping_np[idx_mapping_np], self.shard_mapping

stale_drafts(input_batch)

Mark the drafts whose draft logits were overwritten since drafting. Only valid for requests this rank owns.

Returns:

  • Tensor | None –

    [num_logits] True for each stale draft, or None for a replicated

  • Tensor | None –

    cache, whose rows are never reused.

Source code in vllm/v1/worker/gpu/spec_decode/draft_logits.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def stale_drafts(self, input_batch: InputBatch) -> torch.Tensor | None:
    """Mark the drafts whose draft logits were overwritten since drafting.
    Only valid for requests this rank owns.

    Returns:
        [num_logits] True for each stale draft, or None for a replicated
        cache, whose rows are never reused.

    """
    if self.row_mapping is None:
        return None
    expanded_idx_mapping = input_batch.expanded_idx_mapping
    num_logits = expanded_idx_mapping.shape[0]
    stale = expanded_idx_mapping.new_empty(num_logits, dtype=torch.bool)
    if num_logits == 0:
        return stale
    block_size = 1024
    _stale_drafts_kernel[(triton.cdiv(num_logits, block_size),)](
        stale,
        expanded_idx_mapping,
        input_batch.expanded_local_pos,
        self.row_mapping,
        self.inverse_row_mapping,
        num_logits,
        BLOCK_SIZE=block_size,
    )
    return stale

draft_logits_row(req_state_idx, row_mapping_ptr)

Row of a request in this rank's draft logits cache, or -1 if no row.

row_mapping_ptr is None for a replicated cache, which is indexed by request state. Otherwise it is DraftLogitsCache.row_mapping.

Source code in vllm/v1/worker/gpu/spec_decode/draft_logits.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
@triton.jit
def draft_logits_row(req_state_idx, row_mapping_ptr):
    """Row of a request in this rank's draft logits cache, or -1 if no row.

    ``row_mapping_ptr`` is None for a replicated cache, which is indexed by request
    state. Otherwise it is `DraftLogitsCache.row_mapping`.
    """
    row = req_state_idx
    if row_mapping_ptr is not None:
        row = tl.load(
            row_mapping_ptr + req_state_idx, mask=req_state_idx >= 0, other=-1
        )
    return row.to(tl.int64)