# TRUE TERNARY REFACTOR 15

Date: 2026-05-20

## Goal

Fix the two regressions reported after the platform restructure:

- model capacity had fallen back to about 1.9B instead of the 3B target
- training was hitting OOM from fp16/fp32 state leaking into ternary paths

## Changes

### 1. Restored the 3B target shape

`arbitor/config.py` now restores the large VQ targets while keeping the motif width small enough to stay near 3B:

- `CODEBOOK_DIM = 64`
- `SHARED_VQ_SIZE = 10_000_000`
- `KGVQ_CODEBOOK_SIZE = 5_000_000`
- `KGVQ_CODEBOOK_DIM = 64`

A no-allocation constructor trace of the assembled default model reports:

```text
dummy logical ternary total: 3,011,944,672
```

This keeps the requested 10M shared VQ and 5M KG VQ without the accidental 1024-wide VQ explosion.

### 2. Removed MoEGraph fp16 edge EMA

`MoEGraph` no longer allocates dense `codebook_size * 10` graph edges for large VQ graphs and no longer registers `edge_ema` as `float16`.

Large graphs now use bounded active edge state:

- `active_edge_src`: int32
- `active_edge_dst`: int32
- `active_edge_attr`: int8 ternary edge sign
- `active_edge_score`: int8 residual score
- `edge_index`: empty compatibility buffer for large active mode

Small graph tests still use dense edges, but the score path is now int8 `edge_score`, not fp16 EMA.

### 3. Removed float KG VQ buffers

The old `KGVQCodebook` kept float32 `embed` and `embed_avg` buffers. It is now a compatibility wrapper around `TernaryVQCodebook`, so the KG/composite VQ uses packed ternary rows, int8 scales, int8 accumulators, and int16 usage counts.

### 4. Large VQ initialization is now packed-first

`TernaryEmbeddingTable` now detects million-entry tables and initializes directly into:

- packed `uint8` trits
- int8 `E`
- int8 `E_accum`
- int8 `T_accum`

This avoids building temporary multi-GB float tensors for the 10M shared VQ and 5M KG VQ.

### 5. Removed persistent fp32 Triton training hooks

The Triton ternary backward path now stores `_hook_grad_T_sign` as int8 instead of keeping `_hook_grad_2d` and `_hook_x_2d` fp32 activation/gradient views on each ternary module after backward.

The direct fp32 hook fallback remains only for non-Triton compatibility paths, and the tests now assert that the CUDA Triton path does not retain fp32 grad/x hooks.

### 6. Ternary MoE centroids

MoEGraph routing centroids are now a `TernaryEmbeddingTable` instead of a float `nn.Parameter`.

## Validation

Passed:

```bash
python -m compileall -q arbitor training testing
python -m pytest -q testing/kg/test_kg_edges.py testing/kg/test_composite_head.py testing/test_gradient_capture.py testing/test_tilelang_training.py
python -m pytest -q testing/test_tscale.py::test_cuda_triton_tscale_path
python -m pytest -q --import-mode=importlib testing/model/test_tscale.py::test_cuda_triton_tscale_path
```

Additional targeted checks passed:

- large active MoEGraph with `codebook_size=10_000_000` has `edge_index.shape == (2, 0)` and no float edge buffers
- 1M-entry `TernaryVQCodebook` has no float buffers and trains through sparse forward/backward/update
- small active MoEGraph forward/backward remains finite with ternary centroids

## Remaining Risk

The kernels still use fp32 accumulators internally for numeric accumulation and losses still produce floating scalar loss values. This pass removes persistent fp16/fp32 ternary state and retained fp32 training hooks, which were the memory leak/OOM concern. A fully integer activation/loss path would be a separate kernel-level redesign.