File size: 1,304 Bytes
e1acc15 | 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 | """Sparse transport -> FP32 full-vocabulary probabilities for standard rejection.
The target vocabulary is never truncated to the sparse proposal support. State
is indexed by engine request slots, not current batch order. Failed replies have
valid=False; the probability-aware runner must invalidate every draft for them.
"""
from __future__ import annotations
import torch
class ProbabilityBuffers:
def __init__(self, slots: int, depth: int, vocab: int, device: torch.device):
self.probs = torch.zeros(slots, depth, vocab, dtype=torch.float32, device=device)
self.valid = torch.zeros(slots, dtype=torch.bool, device=device)
self.batch = torch.empty_like(self.probs)
def update(self, mapping: torch.Tensor, support: torch.Tensor, q: torch.Tensor,
ok: torch.Tensor) -> None:
n = mapping.numel()
if (support.shape != (n, self.probs.shape[1], 32) or q.shape != support.shape
or q.dtype != torch.float32 or ok.shape != (n,)):
raise ValueError('sparse probability GPU buffer ABI mismatch')
dense = self.batch[:n]
dense.zero_()
dense.scatter_(2, support.long(), q)
idx = mapping.long()
self.probs.index_copy_(0, idx, dense)
self.valid.index_copy_(0, idx, ok.bool())
|