refactor: collapse fp8 recipe hierarchy and state property layers

- Merge DelayedScaling/DynamicScaling and the abstract FP8Recipe base into one FP8Recipe dataclass with a dynamic flag; dispatch now reads cfg.recipe.dynamic instead of isinstance checks
- Drop the _ActiveOrDefault descriptor and the FP8State property views; the persistent defaults are plain default_* attributes and get_weight_meta takes the active recipe explicitly
- Convert FP8TensorMeta to a NamedTuple of the three per-operand rings
- Update tests to the new API; the autocast context test now asserts _active_config push/restore directly

astrai/extension/fp8.py

@@ -31,7 +31,7 @@ import functools
 from contextvars import ContextVar, Token
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Dict, List, NamedTuple, Optional
 
 import torch
 from torch.library import Library
@@ -56,40 +56,27 @@ class FP8Format(str, Enum):
         return "e5m2" if self is FP8Format.HYBRID else self.value
 
 
+@dataclass
 class FP8Recipe:
     """Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
 
-    ``scale_from_history`` receives the operand's amax tensor (a ring window for
+    ``dynamic=False`` (default) is TE-style delayed scaling: max over the
-    delayed scaling, the current amax for dynamic scaling) and returns the
+    amax history window (amax from *previous* steps; the window trades
-    quantization step. Subclasses set ``history_len`` / ``margin``.
+    responsiveness against stability). ``dynamic=True`` is current-amax
+    scaling (torchao DYNAMIC): measure, then quantize — no history, at an
+    extra pass. ``scale_from_history`` receives the operand's amax tensor
+    (a ring window / the current amax) and returns the quantization step.
     """
 
     history_len: int = 16
     margin: int = 0
+    dynamic: bool = False
 
     def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
         peak = amax.max()
         return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
 
 
-@dataclass
-class DelayedScaling(FP8Recipe):
-    """TE-style delayed scaling: max over the amax history window (amax from
-    *previous* steps; the window trades responsiveness against stability)."""
-
-    history_len: int = 16
-    margin: int = 0
-
-
-@dataclass
-class DynamicScaling(FP8Recipe):
-    """Current-amax scaling (torchao DYNAMIC): measure, then quantize. No
-    history — the scale is derived from the same-step amax, at an extra pass."""
-
-    history_len: int = 1
-    margin: int = 0
-
-
 class _ScaleRing:
     """One operand's delayed-scaling state: a float32 buffer
     ``[hist[n] | scale | legacy | amax | done]`` (views). The quantize
@@ -131,18 +118,15 @@ class _ScaleRing:
         }
 
 
-class FP8TensorMeta:
+class FP8TensorMeta(NamedTuple):
-    """Per-weight delayed-scaling state for ``w``, ``x`` and ``g``.
+    """Per-weight delayed-scaling rings for ``w``, ``x`` and ``g``.
 
-    DynamicScaling never allocates a meta; it measures the current amax inline.
+    Dynamic scaling never allocates a meta; it measures the current amax inline.
     """
 
-    __slots__ = ("w", "x", "g")
+    w: _ScaleRing
-
+    x: _ScaleRing
-    def __init__(self, device: torch.device, recipe: FP8Recipe):
+    g: _ScaleRing
-        self.w = _ScaleRing(device, recipe)
-        self.x = _ScaleRing(device, recipe)
-        self.g = _ScaleRing(device, recipe)
 
 
 @dataclass(frozen=True)
@@ -166,55 +150,29 @@ _active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
 class FP8State:
     """Global fp8 training state: per-tensor metas + out-of-region defaults.
 
-    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set
+    The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar``
-    by ``fp8_autocast``. The properties below read that active config when a
+    set by ``fp8_autocast`` (see ``_active``/``_current_config``); these plain
-    region is open and the global defaults otherwise; the setters (and
+    attributes are the persistent defaults applied outside any region —
-    ``fp8_linear_enable``) write the global defaults — the persistent switch
+    ``fp8_linear_enable`` writes ``default_enabled``. The metas registry is
-    applying outside any region. The metas registry is shared across threads
+    shared across threads (GIL-protected); fp8 backward runs on autograd
-    (GIL-protected); fp8 backward runs on autograd engine threads and only
+    engine threads and only touches metas captured on ``ctx`` at forward time.
-    touches metas captured on ``ctx`` at forward time.
     """
 
     def __init__(self):
         self.default_enabled = False
-        self.default_recipe: FP8Recipe = DelayedScaling()
+        self.default_recipe: FP8Recipe = FP8Recipe()
         self.default_format: FP8Format = FP8Format.HYBRID
         self._metas: Dict[tuple, FP8TensorMeta] = {}
 
-    # Active-config views (region config if open, else the defaults).
+    def get_weight_meta(self, w: torch.Tensor, recipe: FP8Recipe) -> FP8TensorMeta:
-    @property
-    def enabled(self) -> bool:
-        cfg = _active_config.get()
-        return cfg.enabled if cfg is not None else self.default_enabled
-
-    @property
-    def recipe(self) -> FP8Recipe:
-        cfg = _active_config.get()
-        return cfg.recipe if cfg is not None else self.default_recipe
-
-    @property
-    def fp8_format(self) -> FP8Format:
-        cfg = _active_config.get()
-        return cfg.fp8_format if cfg is not None else self.default_format
-
-    # Persistent (out-of-region) defaults.
-    @enabled.setter
-    def enabled(self, value: bool) -> None:
-        self.default_enabled = bool(value)
-
-    @recipe.setter
-    def recipe(self, value: FP8Recipe) -> None:
-        self.default_recipe = value
-
-    @fp8_format.setter
-    def fp8_format(self, value: FP8Format) -> None:
-        self.default_format = FP8Format(value)
-
-    def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
         key = (w.data_ptr(), w.shape, w.dtype)
         meta = self._metas.get(key)
         if meta is None:
-            meta = FP8TensorMeta(w.device, self.recipe)
+            meta = FP8TensorMeta(
+                _ScaleRing(w.device, recipe),
+                _ScaleRing(w.device, recipe),
+                _ScaleRing(w.device, recipe),
+            )
             self._metas[key] = meta
         return meta
 
@@ -222,7 +180,7 @@ class FP8State:
         """Restore construction defaults (switch, recipe, format) and drop all
         per-weight metas — a full state reset for tests / reconfiguration."""
         self.default_enabled = False
-        self.default_recipe = DelayedScaling()
+        self.default_recipe = FP8Recipe()
         self.default_format = FP8Format.HYBRID
         self._metas.clear()
 
@@ -285,7 +243,7 @@ class fp8_autocast:
         margin: int = 0,
     ):
         if recipe is None:
-            recipe = DelayedScaling(history_len=update_interval, margin=margin)
+            recipe = FP8Recipe(history_len=update_interval, margin=margin)
         self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
         self._tokens: List[Token] = []
 
@@ -339,7 +297,7 @@ def fp8_linear_forward(
     if cfg is None:
         cfg = _current_config()
     fmt = cfg.fp8_format.fwd()
-    if isinstance(cfg.recipe, DynamicScaling):
+    if cfg.recipe.dynamic:
         sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
         sw = _dynamic_scale(w, cfg.recipe, fmt)
         x8, _ = quantize(x, sx.reciprocal(), fmt)
@@ -352,7 +310,7 @@ def fp8_linear_forward(
         ).reshape(*x.shape[:-1], w.size(0))
         return out, sx, sw
 
-    meta = state.get_weight_meta(w)
+    meta = state.get_weight_meta(w, cfg.recipe)
     if not meta.w.initialized:
         meta.w.seed(w, fmt)
     if not meta.x.initialized:
@@ -392,8 +350,8 @@ class _LinearFp8(torch.autograd.Function):
         ctx.save_for_backward(x, w, sx, sw)
         ctx.fmt_bwd = cfg.fp8_format.bwd()
         ctx.recipe = cfg.recipe
-        ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
+        ctx.is_dynamic = cfg.recipe.dynamic
-        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w)
+        ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w, cfg.recipe)
         return out
 
     @staticmethod

tests/extension/test_fp8_mma.py

@@ -15,9 +15,8 @@ import torch.nn.functional as F
 
 import astrai.extension.fp8 as f8mod
 from astrai.extension.fp8 import (
-    DelayedScaling,
-    DynamicScaling,
     FP8Format,
+    FP8Recipe,
     FP8TensorMeta,
     _ScaleRing,
     fp8_autocast,
@@ -234,7 +233,7 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
     dev = torch.device("cuda")
     state = f8mod.fp8_state()
     state.reset()
-    state.default_recipe = DelayedScaling(history_len=1, margin=0)
+    state.default_recipe = FP8Recipe(history_len=1, margin=0)
     state.default_format = FP8Format.E4M3
     try:
         m, n, k = 32, 16, 64
@@ -273,7 +272,7 @@ def test_fp8_linear_forward_and_backward():
 
     state = f8mod.fp8_state()
     state.reset()
-    state.default_recipe = DynamicScaling()
+    state.default_recipe = FP8Recipe(dynamic=True)
     try:
         out, _, _ = f8mod.fp8_linear_forward(x, weight, bias)
 
@@ -400,13 +399,13 @@ def test_mm_fp8_matches_scaled_mm():
 def test_recipe_scale_from_history():
     """Delayed: max over the window + margin; dynamic: current amax."""
     hist = torch.tensor([1.0, 2.0, 0.5])
-    d = DelayedScaling(history_len=3, margin=0)
+    d = FP8Recipe(history_len=3, margin=0)
     assert torch.allclose(d.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0))
-    d_m = DelayedScaling(history_len=3, margin=2)
+    d_m = FP8Recipe(history_len=3, margin=2)
     assert torch.allclose(
         d_m.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0 / 4.0)
     )
-    dyn = DynamicScaling()
+    dyn = FP8Recipe(dynamic=True)
     amax = torch.tensor([0.25])
     assert torch.allclose(
         dyn.scale_from_history(amax, "e4m3"), torch.tensor(0.25 / 448.0)
@@ -424,27 +423,37 @@ def test_fp8_format_enum():
 
 
 def test_fp8_autocast_context():
-    """fp8_autocast sets and restores recipe + format on the global state."""
+    """fp8_autocast pushes and restores the thread-local active config."""
     state = fp8_state()
-    prev = (state.enabled, state.recipe, state.fp8_format)
+    state.reset()
     try:
         with fp8_autocast(enabled=True, fp8_format="hybrid", update_interval=8):
-            assert state.enabled
+            cfg = f8mod._active_config.get()
-            assert isinstance(state.recipe, DelayedScaling)
+            assert cfg is not None and cfg.enabled
-            assert state.recipe.history_len == 8
+            assert not cfg.recipe.dynamic
-            assert state.fp8_format is FP8Format.HYBRID
+            assert cfg.recipe.history_len == 8
-            with fp8_autocast(enabled=True, recipe=DynamicScaling(), fp8_format="e4m3"):
+            assert cfg.fp8_format is FP8Format.HYBRID
-                assert isinstance(state.recipe, DynamicScaling)
+            with fp8_autocast(
-                assert state.fp8_format is FP8Format.E4M3
+                enabled=True, recipe=FP8Recipe(dynamic=True), fp8_format="e4m3"
-            assert state.fp8_format is FP8Format.HYBRID  # restored on exit
+            ):
-        assert not state.enabled
+                inner = f8mod._active_config.get()
+                assert inner.recipe.dynamic
+                assert inner.fp8_format is FP8Format.E4M3
+            assert f8mod._active_config.get() is cfg  # restored on exit
+        assert f8mod._active_config.get() is None
+        assert not fp8_linear_enabled()
     finally:
-        state.enabled, state.recipe, state.fp8_format = prev
+        state.reset()
 
 
 def test_fp8_tensor_meta_delayed_update():
     """Meta seeds from data; hist/scale are packed views of one state buffer."""
-    meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0))
+    recipe = FP8Recipe(history_len=4, margin=0)
+    meta = FP8TensorMeta(
+        _ScaleRing(torch.device("cpu"), recipe),
+        _ScaleRing(torch.device("cpu"), recipe),
+        _ScaleRing(torch.device("cpu"), recipe),
+    )
     w = torch.randn(8, 8)
     meta.w.seed(w, "e4m3")
     assert meta.w.initialized