| """Minimum-window suffix budgets; pure contract helpers do not import Torch. |
| |
| This changes the distribution of retained lengths, not attention visibility or |
| inference caps. A sampled K always includes the current position within K. |
| """ |
| from __future__ import annotations |
|
|
| import math |
|
|
|
|
| def validate_minimum(minimum): |
| if type(minimum) is not int or minimum < 1 or minimum & (minimum - 1): |
| raise ValueError('Minimum retained history must be a positive power of two') |
|
|
|
|
| def retained_candidates(available, minimum): |
| """Unique ascending candidates, including full available history once.""" |
| validate_minimum(minimum) |
| if type(available) is not int or available < 1: |
| raise ValueError('Available history must be a positive integer') |
| if available < minimum: |
| return (available,) |
| values = [] |
| value = minimum |
| while value <= available: |
| values.append(value) |
| value *= 2 |
| if values[-1] != available: |
| values.append(available) |
| return tuple(values) |
|
|
|
|
| def retained_from_uniform(available, minimum, uniform): |
| """Scalar reference used to verify equal-probability sampling intervals.""" |
| if not math.isfinite(uniform) or not 0 <= uniform < 1: |
| raise ValueError('Uniform variate must lie in [0, 1)') |
| choices = retained_candidates(available, minimum) |
| return choices[int(uniform * len(choices))] |
|
|
|
|
| def minimum_budget_policy(canonical_policy, minimum): |
| """Bind the new distribution while preserving legacy min1 serialization.""" |
| validate_minimum(minimum) |
| if minimum == 1: |
| return canonical_policy |
| if canonical_policy is None: |
| raise ValueError('Minimum-window training requires an explicit canonical budget policy') |
| return { |
| **canonical_policy, |
| 'retained_length_distribution': { |
| 'schema': 1, |
| 'name': 'uniform_unique_dyadic_minimum_plus_full', |
| 'minimum_history': minimum, |
| 'short_history': 'dense', |
| 'full_history_included_once': True, |
| 'scope': 'independent_sequence_query_shared_heads_and_middle_layers', |
| }, |
| } |
|
|
|
|
| def sample_suffix_lengths(batch, length, *, minimum, generator, device): |
| """Consume one original-shaped uniform draw; dense short queries still draw. |
| |
| CanonicalBudgetPlan calls this separately for each original 16x2048 draw, |
| preserving per-example quantiles and generator advancement on world2/4. |
| The original min1 trainer continues calling its unchanged sample_prefixes. |
| """ |
| import torch |
|
|
| validate_minimum(minimum) |
| if batch < 1 or length < 1: |
| raise ValueError('Sampler dimensions must be positive') |
| n = torch.arange(1, length + 1, device=device, dtype=torch.int64) |
| powers = torch.tensor( |
| [minimum << p for p in range(max(1, length // minimum).bit_length())], |
| device=device, |
| ) |
| power_count = torch.bucketize(n, powers, right=True) |
| extra_full = ((n & (n - 1)) != 0).long() |
| count = torch.where(n < minimum, 1, power_count + extra_full) |
| choice = (torch.rand((batch, length), device=device, generator=generator) * count).long() |
| retained = torch.where(choice < power_count, minimum << choice, n) |
| return torch.where(n < minimum, n, retained) |
|
|
