| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | |
| 5 | import numpy as np |
| 6 | import torch |
| 7 | |
| 8 | |
| 9 | @dataclass |
| 10 | class SpecDecodeMetadata: |
| 11 | # [num_tokens] |
| 12 | draft_token_ids: torch.Tensor |
| 13 | # [batch_size] |
| 14 | num_draft_tokens: list[int] |
| 15 | # [batch_size] |
| 16 | cu_num_draft_tokens: torch.Tensor |
| 17 | # [batch_size] |
| 18 | cu_num_sampled_tokens: torch.Tensor |
| 19 | # [num_tokens] |
| 20 | target_logits_indices: torch.Tensor |
| 21 | # [batch_size] |
| 22 | bonus_logits_indices: torch.Tensor |
| 23 | # [num_tokens + batch_size] |
| 24 | logits_indices: torch.Tensor |
| 25 | |
| 26 | def __post_init__(self): |
| 27 | self.max_spec_len = max(self.num_draft_tokens) |
| 28 | |
| 29 | @classmethod |
| 30 | def make_dummy( |
| 31 | cls, |
| 32 | draft_token_ids: list[list[int]], |
| 33 | device: torch.device, |
| 34 | ) -> "SpecDecodeMetadata": |
| 35 | batch_size = len(draft_token_ids) |
| 36 | num_draft_tokens = [len(ids) for ids in draft_token_ids] |
| 37 | num_sampled_tokens = [len(ids) + 1 for ids in draft_token_ids] |
| 38 | flattened_draft_token_ids = sum(draft_token_ids, []) |
| 39 | num_tokens = len(flattened_draft_token_ids) |
| 40 | |
| 41 | draft_token_ids_tensor = torch.tensor( |
| 42 | flattened_draft_token_ids, dtype=torch.int32, device=device |
| 43 | ) |
| 44 | cu_num_draft_tokens = np.cumsum(num_draft_tokens, dtype=np.int32) |
| 45 | cu_num_draft_tokens_tensor = torch.from_numpy(cu_num_draft_tokens).to(device) |
| 46 | cu_num_sampled_tokens = np.cumsum(num_sampled_tokens, dtype=np.int32) |
| 47 | cu_num_sampled_tokens_tensor = torch.from_numpy(cu_num_sampled_tokens).to( |
| 48 | device |
| 49 | ) |
| 50 | |
| 51 | target_logits_indices = torch.zeros( |
| 52 | num_tokens, dtype=torch.int32, device=device |
| 53 | ) |
| 54 | bonus_logits_indices = torch.zeros(batch_size, dtype=torch.int32, device=device) |
| 55 | logits_indices = torch.zeros( |
| 56 | num_tokens + batch_size, dtype=torch.int32, device=device |
| 57 | ) |
| 58 | return cls( |
| 59 | draft_token_ids=draft_token_ids_tensor, |
| 60 | num_draft_tokens=num_draft_tokens, |
| 61 | cu_num_draft_tokens=cu_num_draft_tokens_tensor, |
| 62 | cu_num_sampled_tokens=cu_num_sampled_tokens_tensor, |
| 63 | target_logits_indices=target_logits_indices, |
| 64 | bonus_logits_indices=bonus_logits_indices, |
| 65 | logits_indices=logits_indices, |
| 66 | ) |