Make HistoryEncoder a pluggable registry, like Router/Objective (gitea #35)

Stage2Autoregressive.init_history_cache and .history_step both
isinstance-checked self.history_encoder against AttentionHistory to decide
whether to use its real incremental-cache methods or a no-op fallback, so a
third history type couldn't be added without editing Stage2Autoregressive
itself. The two-value "markov"/"attention" enum was also independently
hardcoded in three places (Stage2Autoregressive's own validation,
config.py's validate_config, and AutoregressiveConfig.from_dict's default).

Mirrors the Router (giant/model/routers.py) and Objective
(giant/model/objectives.py, gitea #32) pattern: HistoryEncoder now declares
working O(1) init_cache/step defaults (init_cache -> None, step -> one
forward() call), so every registered history type satisfies the incremental
interface without opting in; AttentionHistory overrides both with its real
KV-cache versions since its forward() needs the full prefix. Added
HISTORY_REGISTRY/register_history/build_history, registered "markov" and
"attention", and deleted both isinstance checks in models.py.

Per user decision during planning, config.py's validate_config now imports
HISTORY_REGISTRY and checks membership dynamically instead of keeping its own
hardcoded tuple, making the registry the single source of truth end to end
(verified no import cycle: config.py had no prior dependency on giant.model,
and giant.model.history has none on giant.config).

No config-schema change and no checkpoint impact — this is a pure
internal-interface refactor.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

giant/config.py

@@ -15,6 +15,7 @@ import numpy as np
 import torch
 
 from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
+from giant.model.history import HISTORY_REGISTRY
 
 
 class Conditioning(str, Enum):
@@ -1341,8 +1342,10 @@ def validate_config(cfg: dict) -> None:
                 "AutoregressiveConfig.order's docstring)"
             )
         history = _get_path(cfg, "stage2_model.autoregressive.history")
-        if history not in ("markov", "attention"):
-            raise ValueError(f"stage2_model.autoregressive.history = {history!r} — must be 'markov' or 'attention'")
+        if history not in HISTORY_REGISTRY:
+            raise ValueError(
+                f"stage2_model.autoregressive.history = {history!r} — must be one of {sorted(HISTORY_REGISTRY)}"
+            )
         teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
         if teacher_forcing not in ("always", "scheduled", "never"):
             raise ValueError(

giant/model/history.py

@@ -1,5 +1,9 @@
 """History encoders — stage-2 autoregressive only. Self-contained, no
-dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+dependency on any other `giant.model` submodule (issues.md Issue 8), except
+for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
+`giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
+
+import inspect
 
 import torch
 import torch.nn as nn
@@ -9,19 +13,57 @@ class HistoryEncoder(nn.Module):
     """Interface for stage-2 autoregressive per-token history summaries:
     `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
     a full (teacher-forced) token sequence — used by training. `MarkovHistory`
-    and `AttentionHistory` are the two implementations. Inference
-    (`giant/sample.py`) generates one token at a
-    time and cannot afford `forward`'s per-step cost to be O(K) (attention
-    would then be O(K^2) over a rollout's k_max loop); encoders that need
-    incremental state for that path additionally implement `init_cache`/
-    `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since
-    its per-step cost is already O(1) (it only ever looks at the previous
-    token, not the full prefix)."""
+    and `AttentionHistory` are the two registered implementations (see
+    `HISTORY_REGISTRY`/`build_history`). Inference (`giant/sample.py`)
+    generates one token at a time and cannot afford `forward`'s per-step cost
+    to be O(K) (attention would then be O(K^2) over a rollout's k_max loop),
+    so this interface also declares `init_cache`/`step` for that incremental
+    path, with working O(1) defaults here (`init_cache` -> `None`, `step` ->
+    one `forward` call ignoring `cache`) — correct for any encoder whose
+    per-step cost is already O(1) (i.e. it only ever looks at the previous
+    token, not the full prefix), which is what `MarkovHistory` relies on.
+    `AttentionHistory` overrides both with real incremental-cache versions,
+    since its `forward` genuinely needs the full prefix."""
 
     def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError
 
+    def init_cache(self) -> object:
+        return None
 
+    def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
+        return self.forward(feat, has_prev), cache
+
+
+HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
+
+
+def register_history(name: str):
+    def decorator(cls: type[HistoryEncoder]) -> type[HistoryEncoder]:
+        HISTORY_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEncoder:
+    """Factory: look up a `HistoryEncoder` subclass by name from the registry.
+
+    Every registered history type is fed the same `stage2_model.autoregressive`
+    kwargs; kwargs not declared by that type's constructor are silently
+    dropped, so per-type hyperparameters (e.g. `AttentionHistory`'s
+    `n_heads`/`n_layers`) can coexist in one config without special-casing —
+    mirrors `giant.model.routers.build_router`.
+    """
+    if name not in HISTORY_REGISTRY:
+        raise ValueError(f"unknown history type {name!r}; available: {sorted(HISTORY_REGISTRY)}")
+    cls = HISTORY_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "in_dim", "out_dim"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(in_dim, out_dim, **filtered)
+
+
+@register_history("markov")
 class MarkovHistory(HistoryEncoder):
     """Summarizes the previous secondary's own `(energy_fraction, direction,
     type_representation)` through one small MLP — the "markov" history:
@@ -90,6 +132,7 @@ class _CausalAttnBlock(nn.Module):
         return x, kv
 
 
+@register_history("attention")
 class AttentionHistory(HistoryEncoder):
     """Causal self-attention over the emitted-token prefix — the more
     expressive alternative to `MarkovHistory`'s fixed previous-token-only
@@ -137,16 +180,19 @@ class AttentionHistory(HistoryEncoder):
 
     def step(
         self,
-        token_feat: torch.Tensor,
+        feat: torch.Tensor,
         has_prev: torch.Tensor,
-        cache: list[torch.Tensor | None],
-    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
-        """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest
-        token's own features (what would be `feat[:, k]` in `forward`).
+        cache: object,
+    ) -> tuple[torch.Tensor, object]:
+        """`feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest token's
+        own features (what would be `feat[:, k]` in `forward`). `cache`: the
+        `list[Tensor | None]` from `init_cache`/a previous `step` call (typed
+        `object` here to match `HistoryEncoder.step`'s base signature).
         Advances every block's cache by this position and returns this
         position's output (`(B, 1, out_dim)`, the correct history summary for
         the NEXT slot) plus the updated cache."""
-        x = self._embed(token_feat, has_prev)
+        assert isinstance(cache, list)
+        x = self._embed(feat, has_prev)
         new_cache: list[torch.Tensor | None] = []
         for block, kv in zip(self.blocks, cache):
             x, kv_new = block.step(x, kv)

giant/model/models.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 
 from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
 from giant.model.encoders import ConditionEncoder
-from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory
+from giant.model.history import HistoryEncoder, build_history
 from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
 from giant.model.objectives import build_objective
 from giant.model.routers import Router
@@ -358,8 +358,6 @@ class Stage2Autoregressive(nn.Module):
         cond_enc: ConditionEncoder | None = None,
     ) -> None:
         super().__init__()
-        if history not in ("markov", "attention"):
-            raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
         self.history_kind = history
         self.generator_kind = generator
         self.noise_dim = noise_dim
@@ -384,10 +382,8 @@ class Stage2Autoregressive(nn.Module):
         # this, a reasonable default rather than a design-doc-specified value.
         history_dim = cond_out_dim
         hist_in_dim = CONT_SLOT_DIM + self.type_dim
-        self.history_encoder: HistoryEncoder = (
-            AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
-            if history == "attention"
-            else MarkovHistory(hist_in_dim, history_dim)
+        self.history_encoder: HistoryEncoder = build_history(
+            history, hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
         )
         token_fuse_in = cond_out_dim + context_dim + history_dim + 2  # +2: remaining_frac, slot_idx
         self.token_fuse = nn.Sequential(
@@ -462,12 +458,12 @@ class Stage2Autoregressive(nn.Module):
 
     def init_history_cache(self):
         """Inference-only incremental-decoding state for `self.history_encoder`
-        (`giant/sample.py`'s AR loop): `None` under `history="markov"` (its
-        per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
-        `AttentionHistory.init_cache()` under `history="attention"`."""
-        if isinstance(self.history_encoder, AttentionHistory):
-            return self.history_encoder.init_cache()
-        return None
+        (`giant/sample.py`'s AR loop) — whatever `self.history_encoder.init_cache()`
+        returns for the configured `history` type: `None` under `history="markov"`
+        (its per-step cost is already O(1) — see `HistoryEncoder`'s docstring),
+        or `AttentionHistory.init_cache()`'s real per-block KV cache under
+        `history="attention"`."""
+        return self.history_encoder.init_cache()
 
     def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
         """One inference slot's worth of history encoding: advances `cache`
@@ -479,9 +475,7 @@ class Stage2Autoregressive(nn.Module):
         every model call made for this slot), `new_cache` is what to pass into
         the *next* slot's `history_step`. Must be called exactly once per
         slot — see `AttentionHistory.step`'s docstring."""
-        if isinstance(self.history_encoder, AttentionHistory):
-            return self.history_encoder.step(token_feat, has_prev, cache)
-        return self.history_encoder(token_feat, has_prev), cache
+        return self.history_encoder.step(token_feat, has_prev, cache)
 
     def forward(
         self,

giant/model/network.py

@@ -10,7 +10,15 @@ import X` call site keeps working unchanged.
 from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_state_dict
 from giant.model.builders import build_critics, build_models
 from giant.model.encoders import ConditionEncoder, cat_col_layout
-from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock
+from giant.model.history import (
+    HISTORY_REGISTRY,
+    AttentionHistory,
+    HistoryEncoder,
+    MarkovHistory,
+    _CausalAttnBlock,
+    build_history,
+    register_history,
+)
 from giant.model.layers import (
     BLOCK_REGISTRY,
     AdaLNResBlock,
@@ -78,6 +86,7 @@ __all__ = [
     "ExpertTrunk",
     "FilmResBlock",
     "FlowObjective",
+    "HISTORY_REGISTRY",
     "HistoryEncoder",
     "MarkovHistory",
     "OBJECTIVE_REGISTRY",
@@ -106,6 +115,7 @@ __all__ = [
     "build_composed_router",
     "build_critics",
     "build_expert_body",
+    "build_history",
     "build_models",
     "build_objective",
     "build_router",
@@ -113,6 +123,7 @@ __all__ = [
     "cat_col_layout",
     "migrate_legacy_state_dict",
     "register_block",
+    "register_history",
     "register_objective",
     "register_router",
     "register_trunk",

tests/test_network.py

@@ -5,14 +5,17 @@ import torch
 from giant import config as gconfig
 from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
+    HISTORY_REGISTRY,
     AttentionHistory,
     ConditionEncoder,
+    HistoryEncoder,
     MarkovHistory,
     SinusoidalEmbedding,
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
     build_critics,
+    build_history,
     build_models,
     cat_col_layout,
     stage2_trunk_sec_dim,
@@ -407,6 +410,48 @@ def test_attention_history_step_matches_forward():
     assert torch.allclose(stepped, expected, atol=1e-5)
 
 
+# --- HISTORY_REGISTRY / build_history (gitea #35) ----------------------------
+
+
+def test_history_registry_has_exactly_the_two_known_histories():
+    assert set(HISTORY_REGISTRY) == {"markov", "attention"}
+
+
+def test_build_history_returns_correct_concrete_type():
+    assert isinstance(build_history("markov", 4, 6), MarkovHistory)
+    assert isinstance(build_history("attention", 4, 8), AttentionHistory)
+
+
+def test_build_history_unknown_name_raises():
+    with pytest.raises(ValueError):
+        build_history("bogus", 4, 6)
+
+
+def test_build_history_filters_kwargs_by_signature():
+    """Attention-only kwargs (n_heads/n_layers) must be silently dropped when
+    building a MarkovHistory, matching build_router's documented behavior for
+    per-type hyperparameters coexisting in one config."""
+    hist = build_history("markov", 4, 6, n_heads=2, n_layers=1)
+    assert isinstance(hist, MarkovHistory)
+
+
+def test_history_encoder_base_default_init_cache_and_step():
+    """A HistoryEncoder subclass implementing only forward() must still get
+    working O(1) init_cache/step defaults from the base class."""
+
+    class _StubHistory(HistoryEncoder):
+        def forward(self, feat, has_prev):
+            return feat * 2
+
+    hist = _StubHistory()
+    assert hist.init_cache() is None
+    feat = torch.randn(2, 1, 4)
+    has_prev = torch.ones(2, 1, dtype=torch.bool)
+    out, cache = hist.step(feat, has_prev, "unused-cache")
+    assert torch.equal(out, hist.forward(feat, has_prev))
+    assert cache == "unused-cache"
+
+
 # --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------