refactor: dedupe fp8 meta state into per-operand rings

- collapse FP8TensorMeta's 12 slots + 6 copy-paste methods into three _ScaleRing objects (hist/idx/scale/initialized + update/seed)
- skip meta allocation entirely on the DynamicScaling path (zero rings, scales measured inline)
- drop write-only FP8State._last_device and unused E4M3_MAX alias

astrai/extension/fp8.py

@@ -41,7 +41,6 @@ from astrai.extension.ops.fp8 import (
 
 # Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
 FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
-E4M3_MAX = FP8_MAX["e4m3"]  # legacy alias
 
 
 class FP8Format(str, Enum):
@@ -105,86 +104,50 @@ class DynamicScaling(FP8Recipe):
         return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
 
 
-class FP8TensorMeta:
+class _ScaleRing:
-    """Per-tensor scaling state: amax history rings + derived scales.
+    """One operand's delayed-scaling state: amax history ring + derived scale.
 
-    One ring per operand (weight / activation / gradient). Scales are derived
+    The ring captures its recipe at construction; ``update`` records a fresh
-    from the ring by the recipe; fused kernels record the amax while
+    amax and refreshes the scale for the *next* step (delayed one step).
-    quantizing, so the scale used at step N reflects amax from steps < N
-    (delayed one step).
     """
 
-    __slots__ = (
+    __slots__ = ("recipe", "hist", "idx", "scale", "initialized")
-        "recipe",
-        "w_hist",
-        "x_hist",
-        "g_hist",
-        "w_idx",
-        "x_idx",
-        "g_idx",
-        "w_scale",
-        "x_scale",
-        "g_scale",
-        "w_init",
-        "x_init",
-        "g_init",
-    )
 
     def __init__(self, device: torch.device, recipe: FP8Recipe):
         self.recipe = recipe
         n = recipe.history_len
-        self.w_hist = torch.ones(n, device=device, dtype=torch.float32)
+        self.hist = torch.ones(n, device=device, dtype=torch.float32)
-        self.x_hist = torch.ones(n, device=device, dtype=torch.float32)
+        self.idx = 0
-        self.g_hist = torch.ones(n, device=device, dtype=torch.float32)
+        self.scale = torch.ones(1, device=device, dtype=torch.float32)
-        self.w_idx = self.x_idx = self.g_idx = 0
+        self.initialized = False
-        self.w_scale = torch.ones(1, device=device, dtype=torch.float32)
-        self.x_scale = torch.ones(1, device=device, dtype=torch.float32)
-        self.g_scale = torch.ones(1, device=device, dtype=torch.float32)
-        self.w_init = self.x_init = self.g_init = False
 
-    # -- ring helpers -------------------------------------------------------
+    def update(self, amax: torch.Tensor, fmt: str) -> None:
+        self.hist[self.idx] = amax.reshape(())
+        self.idx = (self.idx + 1) % self.hist.numel()
+        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
 
-    def _record(self, hist: torch.Tensor, idx: int, amax: torch.Tensor) -> int:
+    def seed(self, t: torch.Tensor, fmt: str) -> None:
-        hist[idx] = amax.reshape(())
-        return (idx + 1) % hist.numel()
-
-    def _refresh(self, hist: torch.Tensor, scale: torch.Tensor, fmt: str) -> None:
-        scale.copy_(self.recipe.scale_from_history(hist, fmt))
-
-    def _seed(
-        self, hist: torch.Tensor, scale: torch.Tensor, t: torch.Tensor, fmt: str
-    ) -> None:
         amax = t.abs().amax().to(torch.float32).clamp_min(1e-12)
-        hist.fill_(amax)
+        self.hist.fill_(amax)
-        scale.copy_(self.recipe.scale_from_history(hist, fmt))
+        self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
+        self.initialized = True
 
-    # -- per-operand updates (delayed: record now, refresh for next step) ---
 
-    def update_w(self, amax: torch.Tensor, fmt: str) -> None:
+class FP8TensorMeta:
-        self.w_idx = self._record(self.w_hist, self.w_idx, amax)
+    """Per-weight delayed-scaling state: one ring per operand role.
-        self._refresh(self.w_hist, self.w_scale, fmt)
 
-    def update_x(self, amax: torch.Tensor, fmt: str) -> None:
+    Holds the ``w`` / ``x`` / ``g`` rings; fused kernels record the amax
-        self.x_idx = self._record(self.x_hist, self.x_idx, amax)
+    while quantizing, so the scale used at step N reflects amax from steps
-        self._refresh(self.x_hist, self.x_scale, fmt)
+    < N. DynamicScaling never allocates a meta — it measures the current
+    amax inline (``_dynamic_scale``), so it needs no history storage.
+    """
 
-    def update_g(self, amax: torch.Tensor, fmt: str) -> None:
+    __slots__ = ("w", "x", "g")
-        self.g_idx = self._record(self.g_hist, self.g_idx, amax)
-        self._refresh(self.g_hist, self.g_scale, fmt)
 
-    # -- first-use seeding --------------------------------------------------
+    def __init__(self, device: torch.device, recipe: FP8Recipe):
-
+        self.w = _ScaleRing(device, recipe)
-    def init_w(self, w: torch.Tensor, fmt: str) -> None:
+        self.x = _ScaleRing(device, recipe)
-        self._seed(self.w_hist, self.w_scale, w, fmt)
+        self.g = _ScaleRing(device, recipe)
-        self.w_init = True
-
-    def init_x(self, x: torch.Tensor, fmt: str) -> None:
-        self._seed(self.x_hist, self.x_scale, x, fmt)
-        self.x_init = True
-
-    def init_g(self, g: torch.Tensor, fmt: str) -> None:
-        self._seed(self.g_hist, self.g_scale, g, fmt)
-        self.g_init = True
 
 
 class FP8State:
@@ -195,14 +158,11 @@ class FP8State:
         self.recipe: FP8Recipe = DelayedScaling()
         self.fp8_format: FP8Format = FP8Format.HYBRID
         self._metas: dict[tuple, FP8TensorMeta] = {}
-        self._last_device: Optional[torch.device] = None
 
     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:
-            if self._last_device is None:
-                self._last_device = w.device
             meta = FP8TensorMeta(w.device, self.recipe)
             self._metas[key] = meta
         return meta
@@ -210,7 +170,6 @@ class FP8State:
     def reset(self) -> None:
         self.enabled = False
         self._metas.clear()
-        self._last_device = None
 
 
 # Global singleton: autograd backward runs on the engine worker threads, so
@@ -283,20 +242,21 @@ def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
         bias = torch.empty(0, device=x.device, dtype=x.dtype)
     state = fp8_state()
     fmt = state.fp8_format.fwd()
-    meta = state.get_weight_meta(w)
-    if not meta.w_init:
-        meta.init_w(w, fmt)
     if isinstance(state.recipe, DynamicScaling):
+        meta = None
         sx = _dynamic_scale(x.reshape(-1, w.size(1)), state.recipe, fmt)
         sw = _dynamic_scale(w, state.recipe, fmt)
     else:
-        if not meta.x_init:
+        meta = state.get_weight_meta(w)
-            meta.init_x(x, fmt)
+        if not meta.w.initialized:
-        sx, sw = meta.x_scale, meta.w_scale
+            meta.w.seed(w, fmt)
+        if not meta.x.initialized:
+            meta.x.seed(x, fmt)
+        sx, sw = meta.x.scale, meta.w.scale
     out, amax_x, amax_w = linear_forward_fp8(x, w, bias, sx, sw, fmt)
-    if not isinstance(state.recipe, DynamicScaling):
+    if meta is not None:
-        meta.update_x(amax_x, fmt)
+        meta.x.update(amax_x, fmt)
-        meta.update_w(amax_w, fmt)
+        meta.w.update(amax_w, fmt)
     return out
 
 
@@ -317,8 +277,8 @@ class _LinearFp8(torch.autograd.Function):
         ctx.save_for_backward(x, w)
         ctx.fmt_bwd = state.fp8_format.bwd()
         ctx.recipe = state.recipe
-        ctx.meta = state.get_weight_meta(w)
         ctx.is_dynamic = isinstance(state.recipe, DynamicScaling)
+        ctx.meta = None if ctx.is_dynamic else state.get_weight_meta(w)
         return out
 
     @staticmethod
@@ -332,15 +292,15 @@ class _LinearFp8(torch.autograd.Function):
             sx = _dynamic_scale(x, ctx.recipe, fmt)
         else:
             meta = ctx.meta
-            if not meta.g_init:
+            if not meta.g.initialized:
-                meta.init_g(g, fmt)
+                meta.g.seed(g, fmt)
-            sg, sw, sx = meta.g_scale, meta.w_scale, meta.x_scale
+            sg, sw, sx = meta.g.scale, meta.w.scale, meta.x.scale
         masks = list(ctx.needs_input_grad)
         grad_x, grad_w, grad_b, amax_g = linear_backward_fp8(
             g, x, w, masks, sg, sw, sx, fmt
         )
         if not ctx.is_dynamic:
-            ctx.meta.update_g(amax_g, fmt)
+            ctx.meta.g.update(amax_g, fmt)
         return grad_x, grad_w, grad_b if masks[2] else None
 
 

tests/extension/test_fp8_mma.py

@@ -334,11 +334,11 @@ def test_fp8_tensor_meta_delayed_update():
     """Meta seeds from data and refreshes the scale from the amax ring."""
     meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0))
     w = torch.randn(8, 8)
-    meta.init_w(w, "e4m3")
+    meta.w.seed(w, "e4m3")
-    assert meta.w_init
+    assert meta.w.initialized
-    torch.testing.assert_close(meta.w_scale, (w.abs().amax() / 448.0).reshape(1))
+    torch.testing.assert_close(meta.w.scale, (w.abs().amax() / 448.0).reshape(1))
-    meta.update_w(torch.tensor([4.0]), "e4m3")
+    meta.w.update(torch.tensor([4.0]), "e4m3")
-    torch.testing.assert_close(meta.w_scale, torch.tensor(4.0 / 448.0).reshape(1))
+    torch.testing.assert_close(meta.w.scale, torch.tensor(4.0 / 448.0).reshape(1))
 
 
 def test_quantize_bf16_cpu_fallback():