| """Hashed n-gram addressing shared by Senryu-2's model, its host table store, and serving. |
| |
| Each table row is addressed with integer ops that give the same result on CPU |
| and CUDA: canonical token ids (NFKC-normalized, lowercased, stripped surface |
| forms) are multiplied by odd 31-bit constants, XOR-combined, mixed, and reduced |
| modulo a distinct prime per table. Token ids stay below 2**16, so every |
| product stays below 2**47 and int64 never overflows. |
| |
| Tables live in one concatenated ``[total_rows, embedding_dim]`` tensor; table |
| ``t`` owns rows ``[offset_t, offset_t + rows_t)``. The prime row counts, the |
| multipliers, and the canonical map are saved beside every table file and |
| checked on load, so a table can never be read with a different address map. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import unicodedata |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| from tokenizers import Tokenizer |
|
|
| from tiny_gdn.config import TinyGDNConfig |
|
|
| NGRAM_TABLES_FILE = "ngram_tables.safetensors" |
| NGRAM_TABLES_EMA_FILE = "ngram_tables_ema.safetensors" |
| NGRAM_OPTIMIZER_FILE = "ngram_optimizer.safetensors" |
|
|
| PRIME_STRIDE = 4099 |
| HASH_SEED = 0x6B696D6F |
| MIX_SHIFT = 29 |
| LOAD_CHUNK_ROWS = 1 << 20 |
| _MASK64 = (1 << 64) - 1 |
| _MILLER_RABIN_BASES = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37) |
|
|
|
|
| def _splitmix64(state: int) -> tuple[int, int]: |
| state = (state + 0x9E3779B97F4A7C15) & _MASK64 |
| mixed = ((state ^ (state >> 30)) * 0xBF58476D1CE4E5B9) & _MASK64 |
| mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & _MASK64 |
| return state, mixed ^ (mixed >> 31) |
|
|
|
|
| def _is_prime(value: int) -> bool: |
| if value < 2: |
| return False |
| for base in _MILLER_RABIN_BASES: |
| if value % base == 0: |
| return value == base |
| odd, twos = value - 1, 0 |
| while odd % 2 == 0: |
| odd //= 2 |
| twos += 1 |
| for base in _MILLER_RABIN_BASES: |
| witness = pow(base, odd, value) |
| if witness in (1, value - 1): |
| continue |
| for _ in range(twos - 1): |
| witness = witness * witness % value |
| if witness == value - 1: |
| break |
| else: |
| return False |
| return True |
|
|
|
|
| def _next_prime(value: int) -> int: |
| while not _is_prime(value): |
| value += 1 |
| return value |
|
|
|
|
| @dataclass(frozen=True) |
| class NGramLayout: |
| """Where every n-gram of every order lands in the concatenated tables.""" |
|
|
| orders: tuple[int, ...] |
| heads_per_order: int |
| table_rows: tuple[int, ...] |
| table_offsets: tuple[int, ...] |
| multipliers: tuple[tuple[int, ...], ...] |
| embedding_dim: int |
| pad_token_id: int |
|
|
| @classmethod |
| def from_config(cls, config: TinyGDNConfig) -> NGramLayout: |
| if not config.ngram_orders: |
| raise ValueError("config has no n-gram memory") |
| rows: list[int] = [] |
| candidate = config.ngram_table_rows |
| for _ in range(config.ngram_table_count): |
| candidate = _next_prime(candidate + PRIME_STRIDE) |
| rows.append(candidate) |
| offsets = tuple(int(value) for value in np.cumsum([0, *rows[:-1]])) |
| state = HASH_SEED |
| multipliers: list[tuple[int, ...]] = [] |
| max_order = max(config.ngram_orders) |
| for _ in range(config.ngram_table_count): |
| table: list[int] = [] |
| for _ in range(max_order): |
| state, value = _splitmix64(state) |
| table.append((value >> 33) | (1 << 30) | 1) |
| multipliers.append(tuple(table)) |
| return cls( |
| orders=tuple(config.ngram_orders), |
| heads_per_order=config.ngram_heads_per_order, |
| table_rows=tuple(rows), |
| table_offsets=offsets, |
| multipliers=tuple(multipliers), |
| embedding_dim=config.ngram_embedding_dim, |
| pad_token_id=config.bos_token_id, |
| ) |
|
|
| @property |
| def table_count(self) -> int: |
| return len(self.table_rows) |
|
|
| @property |
| def total_rows(self) -> int: |
| return self.table_offsets[-1] + self.table_rows[-1] |
|
|
| @property |
| def max_order(self) -> int: |
| return max(self.orders) |
|
|
| @property |
| def memory_dim(self) -> int: |
| return self.table_count * self.embedding_dim |
|
|
| @property |
| def table_orders(self) -> tuple[int, ...]: |
| return tuple(order for order in self.orders for _ in range(self.heads_per_order)) |
|
|
| def metadata_tensors(self) -> dict[str, torch.Tensor]: |
| return { |
| "table_rows": torch.tensor(self.table_rows, dtype=torch.int64), |
| "hash_multipliers": torch.tensor(self.multipliers, dtype=torch.int64), |
| } |
|
|
| def check_metadata(self, tensors: dict[str, torch.Tensor], source: Path) -> None: |
| for name, expected in self.metadata_tensors().items(): |
| found = tensors.get(name) |
| if found is None or not torch.equal(found.cpu(), expected): |
| raise ValueError(f"{source} was written for a different n-gram layout ({name})") |
|
|
|
|
| def canonical_token_ids(tokenizer: Tokenizer, vocab_size: int) -> torch.Tensor: |
| """Map every token id to the smallest id sharing its normalized surface form. |
| |
| Normalization is NFKC, then strip, then lowercase. Special tokens keep |
| their own ids, whitespace-only tokens are only NFKC-merged, and tokens |
| that decode to partial UTF-8 stay distinct. |
| """ |
|
|
| special = set(tokenizer.get_added_tokens_decoder()) |
| canonical = np.arange(vocab_size, dtype=np.int64) |
| first_by_key: dict[tuple[str, str], int] = {} |
| for token_id in range(min(vocab_size, tokenizer.get_vocab_size())): |
| if token_id in special: |
| continue |
| text = tokenizer.decode([token_id], skip_special_tokens=False) |
| if "\ufffd" in text: |
| key = ("bytes", tokenizer.id_to_token(token_id) or str(token_id)) |
| else: |
| normalized = unicodedata.normalize("NFKC", text) |
| folded = normalized.strip().lower() |
| key = ("text", folded) if folded else ("space", normalized) |
| canonical[token_id] = first_by_key.setdefault(key, token_id) |
| return torch.from_numpy(canonical) |
|
|
|
|
| def document_distance(tokens: torch.Tensor, bos_token_id: int, row_starts: torch.Tensor | None = None) -> torch.Tensor: |
| """Offset of every token in a flat ``[N]`` stream from its document's first token. |
| |
| Documents start at position 0, at every BOS token, and at ``row_starts``. |
| """ |
|
|
| starts = tokens == bos_token_id |
| starts[0] = True |
| if row_starts is not None: |
| starts[row_starts] = True |
| positions = torch.arange(tokens.numel(), device=tokens.device) |
| first = torch.where(starts, positions, torch.zeros_like(positions)).cummax(dim=0).values |
| return positions - first |
|
|
|
|
| def ngram_row_indices( |
| canonical: torch.Tensor, |
| layout: NGramLayout, |
| *, |
| distance: torch.Tensor | None = None, |
| history: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Global table rows ``[..., L, tables]`` for the n-grams ending at each token. |
| |
| ``canonical`` holds int64 canonical ids ``[..., L]``. Tokens before a |
| document start read as the pad id: with ``distance`` (same shape) a |
| document starts wherever the distance is zero; with ``history`` (the |
| ``max_order - 1`` canonical ids before position 0) the sequence continues |
| a cached one; with neither, position 0 starts the only document. |
| """ |
|
|
| span = layout.max_order |
| length = canonical.shape[-1] |
| pad = layout.pad_token_id |
| if history is not None: |
| if history.shape[-1] != span - 1: |
| raise ValueError(f"history must hold {span - 1} tokens") |
| extended = torch.cat([history, canonical], dim=-1) |
| shifted = [extended[..., span - 1 - back : span - 1 - back + length] for back in range(span)] |
| else: |
| shifted = [canonical] |
| for back in range(1, span): |
| if back >= length: |
| previous = torch.full_like(canonical, pad) |
| else: |
| padding = canonical.new_full((*canonical.shape[:-1], back), pad) |
| previous = torch.cat([padding, canonical[..., :-back]], dim=-1) |
| if distance is not None: |
| previous = torch.where(distance >= back, previous, pad) |
| shifted.append(previous) |
| columns = [] |
| for table, order in enumerate(layout.table_orders): |
| multipliers = layout.multipliers[table] |
| mixed = shifted[0] * multipliers[0] |
| for back in range(1, order): |
| mixed = mixed ^ (shifted[back] * multipliers[back]) |
| mixed = mixed ^ (mixed >> MIX_SHIFT) |
| columns.append(torch.remainder(mixed, layout.table_rows[table]) + layout.table_offsets[table]) |
| return torch.stack(columns, dim=-1) |
|
|
|
|
| def save_table_file(path: Path, tables: torch.Tensor, layout: NGramLayout, canonical: torch.Tensor) -> None: |
| if tables.shape != (layout.total_rows, layout.embedding_dim): |
| raise ValueError(f"tables have shape {tuple(tables.shape)}, layout needs {(layout.total_rows, layout.embedding_dim)}") |
| save_file( |
| { |
| "tables": tables.contiguous(), |
| "canonical_ids": canonical.to(dtype=torch.int64, device="cpu").contiguous(), |
| **layout.metadata_tensors(), |
| }, |
| str(path), |
| ) |
|
|
|
|
| def read_table_metadata(path: Path, layout: NGramLayout) -> torch.Tensor: |
| """Check ``path`` against ``layout`` and return its canonical token map.""" |
|
|
| with safe_open(str(path), framework="pt", device="cpu") as handle: |
| metadata = {name: handle.get_tensor(name) for name in ("table_rows", "hash_multipliers")} |
| layout.check_metadata(metadata, path) |
| |
| |
| return handle.get_tensor("canonical_ids").clone() |
|
|
|
|
| def copy_rows_from_file(path: Path, name: str, destination: torch.Tensor) -> None: |
| """Copy tensor ``name`` of a safetensors file into ``destination`` in row chunks.""" |
|
|
| with safe_open(str(path), framework="pt", device="cpu") as handle: |
| source = handle.get_slice(name) |
| shape = tuple(source.get_shape()) |
| if shape != tuple(destination.shape): |
| raise ValueError(f"{path} holds {name} of shape {shape}, expected {tuple(destination.shape)}") |
| for start in range(0, shape[0], LOAD_CHUNK_ROWS): |
| stop = min(start + LOAD_CHUNK_ROWS, shape[0]) |
| destination[start:stop].copy_(source[start:stop]) |
|
|
|
|
| def load_table_file( |
| path: Path, |
| layout: NGramLayout, |
| *, |
| device: torch.device | str = "cpu", |
| dtype: torch.dtype = torch.float32, |
| out: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """``(tables, canonical_ids)`` from ``path``, copied in row chunks to bound host memory.""" |
|
|
| canonical = read_table_metadata(path, layout) |
| tables = ( |
| torch.empty((layout.total_rows, layout.embedding_dim), dtype=dtype, device=device) |
| if out is None |
| else out |
| ) |
| copy_rows_from_file(path, "tables", tables) |
| return tables, canonical |
|
|
