File size: 2,214 Bytes
2e0601e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Shared evidence prefix and isolated candidate branch layout."""
from dataclasses import dataclass
import torch
@dataclass(frozen=True)
class CandidateTree:
prefix_length: int
branch_lengths: tuple[int, ...]
def __post_init__(self):
if self.prefix_length <= 0 or not self.branch_lengths or min(self.branch_lengths) <= 0:
raise ValueError("A nonempty prefix and nonempty branches are required")
@property
def length(self):
return self.prefix_length + sum(self.branch_lengths)
def allowed(self, *, device=None):
groups = [-1] * self.prefix_length
for index, length in enumerate(self.branch_lengths):
groups.extend([index] * length)
groups = torch.tensor(groups, device=device)
physical = torch.arange(self.length, device=device)
causal = physical[None, :] <= physical[:, None]
return causal & ((groups[None, :] == -1) | (groups[:, None] == groups[None, :]))
def additive_mask(self, *, dtype=torch.float32, device=None):
allowed = self.allowed(device=device)
return torch.zeros_like(allowed, dtype=dtype).masked_fill(~allowed, float("-inf"))[None, None]
def position_ids(self, prefix_positions):
if prefix_positions.ndim != 3 or prefix_positions.shape[-1] != self.prefix_length:
raise ValueError("Expected native prefix positions with shape (axes, batch, prefix)")
if prefix_positions.shape[1] != 1:
raise ValueError("This reference packs one request")
origin = int(prefix_positions.max()) + 1
parts = [prefix_positions]
for length in self.branch_lengths:
branch = torch.arange(origin, origin + length, device=prefix_positions.device)
parts.append(branch.view(1, 1, -1).expand(prefix_positions.shape[0], 1, -1))
return torch.cat(parts, dim=-1)
def attention_pairs(self):
prefix = self.prefix_length
return prefix * (prefix + 1) // 2 + sum(
prefix * length + length * (length + 1) // 2 for length in self.branch_lengths
)
|