| """Request-owned CUDA graph for ABC and semantic top-k/top-p sampling.""" |
| import torch |
|
|
| from .protocol import MUSIC_END, CODEC_OFFSET, CODEC_SIZE, ABC_END, EOD |
|
|
|
|
| class CapturedSampler: |
| """Return one vocabulary ID in shared output storage, valid until next call. |
| |
| The caller owns the RNG and already validates the generated history. Capture |
| and warmup restore its state; only real replays consume random numbers. |
| """ |
|
|
| @torch.inference_mode() |
| def __init__(self, logits, sampling, generator, *, phase="semantic"): |
| if phase not in {"abc", "semantic"}: |
| raise ValueError("Captured sampling phase must be abc or semantic") |
| if logits.device.type != "cuda" or logits.ndim != 2 or logits.shape[0] != 1: |
| raise ValueError("Captured sampling requires one CUDA logit row after CFG") |
| required = ABC_END + 1 if phase == "abc" else CODEC_OFFSET + CODEC_SIZE |
| if logits.shape[-1] < required or sampling.temperature <= 0: |
| raise ValueError("Captured sampling requires the phase vocabulary and positive temperature") |
| |
| if (generator.device.type != "cuda" or |
| generator.device.index is not None and generator.device.index != logits.device.index): |
| raise ValueError("Sampling generator and logits must share a CUDA device") |
| self.sampling = sampling |
| self.phase = phase |
| self.offset = 0 if phase == "abc" else MUSIC_END |
| self.eos_index = EOD if phase == "abc" else 0 |
| self.generator = generator |
| self.graph = self.output = None |
| self.closed = False |
| size = EOD + 1 if phase == "abc" else CODEC_OFFSET + CODEC_SIZE - MUSIC_END |
| self.logits = torch.zeros((1, size), |
| device=logits.device, dtype=torch.float32) |
| self.recent = torch.full((1, sampling.penalty_window), -1, |
| device=logits.device, dtype=torch.long) |
| self.allow_eos = torch.ones((), device=logits.device, dtype=torch.bool) |
| saved = generator.get_state() |
| try: |
| with torch.cuda.device(logits.device): |
| current = torch.cuda.current_stream() |
| warmup = torch.cuda.Stream() |
| warmup.wait_stream(current) |
| with torch.cuda.stream(warmup): |
| for _ in range(3): |
| self._compute() |
| current.wait_stream(warmup) |
| torch.cuda.synchronize() |
| self.graph = torch.cuda.CUDAGraph() |
| self.graph.register_generator_state(generator) |
| with torch.cuda.graph(self.graph): |
| self.output = self._compute() |
| finally: |
| generator.set_state(saved) |
|
|
| def _compute(self): |
| sampling = self.sampling |
| scores = self.logits.clone() |
| scores[..., self.eos_index] = torch.where(self.allow_eos, scores[..., self.eos_index], -torch.inf) |
| if sampling.repetition_penalty != 1: |
| frequency = torch.zeros_like(scores) |
| frequency.scatter_add_(-1, self.recent.clamp_min(0), (self.recent >= 0).float()) |
| alpha = sampling.repetition_penalty ** frequency |
| scores = torch.where(scores < 0, scores * alpha, scores / alpha) |
| if sampling.temperature != 1: |
| scores /= sampling.temperature |
| values, indices = scores.topk(min(sampling.top_k, scores.shape[-1])) |
| if sampling.top_p < 1: |
| probabilities = values.softmax(-1) |
| removed = probabilities.cumsum(-1) - probabilities > sampling.top_p |
| removed[..., :1] = False |
| values = values.masked_fill(removed, -torch.inf) |
| choice = torch.multinomial(values.softmax(-1), 1, generator=self.generator) |
| selected = indices.gather(-1, choice) |
| if self.phase == "abc": |
| return torch.where(selected == EOD, ABC_END, selected) |
| return selected + MUSIC_END |
|
|
| @torch.inference_mode() |
| def __call__(self, logits, history, step): |
| if self.closed: |
| raise RuntimeError("CapturedSampler is closed") |
| recent = [token - self.offset for token in history[-self.sampling.penalty_window:]] |
| recent += [-1] * (self.sampling.penalty_window - len(recent)) |
| self.recent.copy_(torch.tensor([recent], dtype=torch.long, device="cpu")) |
| if self.phase == "abc": |
| |
| self.logits[..., :EOD].copy_(logits[..., :EOD]) |
| self.logits[..., EOD].copy_(logits[..., ABC_END]) |
| else: |
| self.logits.copy_(logits[..., MUSIC_END:CODEC_OFFSET + CODEC_SIZE]) |
| self.allow_eos.fill_(step >= self.sampling.min_tokens) |
| self.graph.replay() |
| return self.output |
|
|
| def close(self): |
| self.graph = self.output = self.generator = None |
| self.logits = self.recent = self.allow_eos = None |
| self.closed = True |
|
|
