diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index 715477a..8120cb4 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -1,35 +1,37 @@ """FP8 training: scaling recipes, per-tensor state, and aten::linear dispatch. -Layered (see also ``ops/fp8.py`` for the CUDA interface adapter): - -1. Kernel interface: ``ops.fp8`` — the only module touching the pybind. -2. Training state (this module): scaling *recipes* (TE-style delayed scaling - or dynamic current-amax scaling), per-tensor scales + amax history, and - the ``fp8_autocast`` context (like ``torch.autocast``). -3. aten::linear integration (this module): registers the CUDA impl and the - dtype guard. +Layered (see ``ops/fp8.py`` for the CUDA interface adapter): +1. ``ops.fp8`` — the only module touching the pybind. +2. This module (strategy layer): scaling *recipes* (TE-style delayed scaling + or dynamic current-amax scaling), per-tensor scales + amax history, and the + ``fp8_autocast`` context manager (like ``torch.autocast``). +3. aten::linear integration: registers the CUDA + AutogradCUDA impls. Usage:: from astrai.extension.fp8 import fp8_autocast - with fp8_autocast(enabled=True, fp8_format="hybrid"): logits = model(input_ids) - loss.backward() # fp8 backward runs wherever it is called: the - # forward captures the fmt/recipe/meta on the autograd node + loss.backward() # fp8 backward runs anywhere; fwd captured state on the node -Importing this module registers the aten::linear CUDA and AutogradCUDA -implementations. +Format defaults follow the ecosystem consensus: E4M3 forward / E5M2 backward +("hybrid"); every operand's scale is a quantization step derived from its amax +history by the active recipe. -Format defaults follow the ecosystem consensus: E4M3 for the forward pass, -E5M2 for the backward (gradient) pass ("hybrid"); every operand's scale is a -quantization step derived from its amax history by the active recipe. +The context mirrors ``torch.autocast`` (``autocast_mode.py``): the active +``(enabled, recipe, fp8_format)`` triple is thread-local (a ``contextvars`` +``ContextVar``, absent outside any region), and the manager is class-based and +reentrant with nested ``enabled=False`` disabling dispatch inside it. The module +targets *training*: every step quantizes x/w/g fresh (no weight-cast cache — the +optimizer bumps the weight version each step, so a torch-style cached_cast would +miss anyway), and the per-operand scales come from the delayed/dynamic recipe. """ -from contextlib import contextmanager +import functools +from contextvars import ContextVar, Token from dataclasses import dataclass from enum import Enum -from typing import Optional +from typing import Dict, List, Optional import torch from torch.library import Library @@ -58,62 +60,46 @@ class FP8Format(str, Enum): class FP8Recipe: - """Scale-from-amax policy; the scale computation is the injection point. + """Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``. - ``scale_from_history`` receives the amax tensor for this operand (a ring - window for delayed scaling, the current amax for dynamic scaling) and - returns the quantization step: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``. + ``scale_from_history`` receives the operand's amax tensor (a ring window for + delayed scaling, the current amax for dynamic scaling) and returns the + quantization step. Subclasses set ``history_len`` / ``margin``. """ history_len: int = 16 margin: int = 0 def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor: - raise NotImplementedError + 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. - - The scale is computed from amax measured in *previous* steps (delayed one - step); the window length trades responsiveness against stability. - """ + """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 - 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 DynamicScaling(FP8Recipe): - """Current-amax scaling (torchao DYNAMIC): measure, then quantize. - - No history — the scale is derived from the amax of the tensor being - quantized in the same step, at the cost of an extra reduction pass. - """ + """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 - 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) - class _ScaleRing: - """One operand's delayed-scaling state, packed for in-kernel finalization. - - ``state`` is a single float32 CUDA buffer ``[hist[n] | scale | counter]`` - (``hist`` / ``scale`` are views). The quantize kernel's last-finishing - block records the freshly measured amax into ``hist[idx]``, reduces the - window and publishes the next step's scale entirely on device — the - Python-side hist-write / max / scale-write chain is gone. The counter - slot stays int32-zero (float bits) between launches. ``idx`` advances - host-side each step; ``margin`` is fixed by the recipe. + """One operand's delayed-scaling state: a float32 buffer + ``[hist[n] | scale | counter]`` (views). The quantize kernel's last-finishing + block records the measured amax into ``hist[idx]``, reduces the window and + publishes the next scale entirely on device — the Python-side write/max/write + chain is gone. The counter slot stays int32-zero (float bits) between + launches; ``idx`` advances host-side each step. """ __slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized") @@ -121,7 +107,6 @@ class _ScaleRing: def __init__(self, device: torch.device, recipe: FP8Recipe): self.recipe = recipe n = recipe.history_len - # [hist | scale | counter]; the counter slot must start at int 0. self.state = torch.zeros(n + 2, device=device, dtype=torch.float32) self.hist = self.state[:n] self.scale = self.state[n : n + 1] @@ -140,12 +125,10 @@ class _ScaleRing: class FP8TensorMeta: - """Per-weight delayed-scaling state: one ring per operand role. - - Holds the ``w`` / ``x`` / ``g`` rings; fused kernels record the amax - while quantizing, so the scale used at step N reflects amax from steps - < N. DynamicScaling never allocates a meta — it measures the current - amax inline (``_dynamic_scale``), so it needs no history storage. + """Per-weight delayed-scaling state: one ring per operand role (``w``/``x``/ + ``g``). Fused kernels record amax while quantizing, so the scale used at step + N reflects amax from steps < N. DynamicScaling never allocates a meta — it + measures the current amax inline. """ __slots__ = ("w", "x", "g") @@ -156,14 +139,70 @@ class FP8TensorMeta: self.g = _ScaleRing(device, recipe) +@dataclass(frozen=True) +class _ActiveConfig: + """The immutable (enabled, recipe, format) triple of one open region.""" + + enabled: bool + recipe: FP8Recipe + fp8_format: FP8Format + + +# Thread-local active configuration (torch's autocast TLS analog): set by +# fp8_autocast on __enter__, absent outside any region. Autograd engine +# threads run backwards with their own empty context — fine, since backward +# only reads state captured on ctx at forward time. +_active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar( + "astrai_fp8_active_config", default=None +) + + class FP8State: - """Global fp8 training state: active recipe + per-tensor metas.""" + """Global fp8 training state: per-tensor metas + out-of-region defaults. + + The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set + by ``fp8_autocast``. The properties below read that active config when a + region is open and the global defaults otherwise; the setters (and + ``fp8_linear_enable``) write the global defaults — the persistent switch + applying outside any region. The metas registry is shared across threads + (GIL-protected); fp8 backward runs on autograd engine threads and only + touches metas captured on ``ctx`` at forward time. + """ def __init__(self): - self.enabled = False - self.recipe: FP8Recipe = DelayedScaling() - self.fp8_format: FP8Format = FP8Format.HYBRID - self._metas: dict[tuple, FP8TensorMeta] = {} + self.default_enabled = False + self.default_recipe: FP8Recipe = DelayedScaling() + self.default_format: FP8Format = FP8Format.HYBRID + self._metas: Dict[tuple, FP8TensorMeta] = {} + + # Active-config views (region config if open, else the defaults). + @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) @@ -174,13 +213,11 @@ class FP8State: return meta def reset(self) -> None: - self.enabled = False + self.default_enabled = False self._metas.clear() -# Global singleton: autograd backward runs on the engine worker threads, so -# thread-local state would lose the fp8 flag during loss.backward(). The GIL -# protects Python-side mutation; the CUDA kernels take their own mutex. +# Process-wide singleton; per-thread/per-region state lives in _active_config. _state = FP8State() @@ -188,43 +225,76 @@ def fp8_state() -> FP8State: return _state -@contextmanager -def fp8_autocast( - enabled: bool = True, - update_interval: int = 16, - recipe: Optional[FP8Recipe] = None, - fp8_format: str = "hybrid", - margin: int = 0, -): +def _active() -> Optional[_ActiveConfig]: + """The active config when fp8 dispatch is on, else ``None`` (fast guard). + + A region config wins (honoring nested ``enabled=False`` regions); with no + region open this falls back to the persistent global switch + (``fp8_linear_enable``), so that flag still routes aten::linear to fp8. + """ + cfg = _active_config.get() + if cfg is not None: + return cfg if cfg.enabled else None + if _state.default_enabled: + return _ActiveConfig(True, _state.default_recipe, _state.default_format) + return None + + +def _current_config() -> _ActiveConfig: + """Like ``_active()`` but always returns a config (disabled regions and + out-of-region direct calls resolve to the global defaults).""" + cfg = _active_config.get() + if cfg is not None: + return cfg + return _ActiveConfig( + _state.default_enabled, _state.default_recipe, _state.default_format + ) + + +class fp8_autocast: """Autocast-style context: fp8 linear dispatch on this thread. - Usage:: + Mirrors ``torch.autocast`` — a class-based, reentrant, nestable context + over thread-local state:: with fp8_autocast(enabled=True, fp8_format="hybrid"): logits = model(input_ids) # aten::linear -> fp8 path loss.backward() # fp8 backward; state was captured at forward time - Args: - enabled: toggle fp8 dispatch for aten::linear. - update_interval: legacy alias for the delayed-scaling history window - (used only when ``recipe`` is not given). - recipe: scaling policy; defaults to ``DelayedScaling(update_interval)``. - fp8_format: ``"e4m3"`` / ``"e5m2"`` / ``"hybrid"`` (default) — hybrid - means E4M3 forward, E5M2 backward. - margin: scale headroom (``scale = (amax / FP8_MAX) / 2^margin``) used - with the default delayed recipe. + Nesting follows torch: each ``__enter__`` pushes the new active config, each + ``__exit__`` restores the previous one, and a nested ``enabled=False`` region + simply disables dispatch inside it. The instance doubles as a decorator. """ - state = fp8_state() - prev = (state.enabled, state.recipe, state.fp8_format) - if recipe is None: - recipe = DelayedScaling(history_len=update_interval, margin=margin) - state.enabled = enabled - state.recipe = recipe - state.fp8_format = FP8Format(fp8_format) - try: - yield - finally: - state.enabled, state.recipe, state.fp8_format = prev + + def __init__( + self, + enabled: bool = True, + update_interval: int = 16, + recipe: Optional[FP8Recipe] = None, + fp8_format: str = "hybrid", + margin: int = 0, + ): + if recipe is None: + recipe = DelayedScaling(history_len=update_interval, margin=margin) + self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format)) + self._tokens: List[Token] = [] + + def __enter__(self) -> "fp8_autocast": + self._tokens.append(_active_config.set(self._config)) + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> bool: + token = self._tokens.pop() + _active_config.reset(token) + return False + + def __call__(self, func): + @functools.wraps(func) + def decorate(*args, **kwargs): + with self: + return func(*args, **kwargs) + + return decorate # --------------------------------------------------------------------------- @@ -237,72 +307,96 @@ def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor return recipe.scale_from_history(amax, fmt) -def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): +_zero_bias: Dict[Optional[int], torch.Tensor] = {} + + +def _empty_bias(x: torch.Tensor) -> torch.Tensor: + """Per-device cached 0-element bf16 bias (the binding only checks numel — + never mutated), saving a CUDA allocation per bias-less linear.""" + key = x.device.index + t = _zero_bias.get(key) + if t is None: + t = torch.empty(0, device=x.device, dtype=torch.bfloat16) + _zero_bias[key] = t + return t + + +def fp8_linear_forward( + x: torch.Tensor, w: torch.Tensor, bias=None, cfg: Optional[_ActiveConfig] = None +): """Scaled fp8 linear forward (called from the aten::linear impl). - Pure FP8 path for both recipes: quantize x/w with the active scales, run - the pre-quantized GEMM. With delayed scaling the rings finalize inside - the quantize kernels (amax folded into the window, next step's scale - published on device); dynamic scaling measures the current amax itself. + Pure FP8 path for both recipes: quantize x/w with the active scales, run the + pre-quantized GEMM. Delayed scaling finalizes the rings inside the quantize + kernels (amax folded into the window, next scale published on device); + dynamic scaling measures the current amax itself. Training quantizes the + weight every step (the optimizer bumps its version, so there is no cast + cache, matching ``cached_cast``-less behavior). """ - if bias is None: - bias = torch.empty(0, device=x.device, dtype=x.dtype) state = fp8_state() - fmt = state.fp8_format.fwd() - 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) - out, amax_x, amax_w = linear_forward_fp8(x, w, bias, sx, sw, fmt) + if cfg is None: + cfg = _current_config() + fmt = cfg.fp8_format.fwd() + margin = cfg.recipe.margin + if bias is None: + bias = _empty_bias(x) + if isinstance(cfg.recipe, DynamicScaling): # measure-then-quantize, no state + sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt) + sw = _dynamic_scale(w, cfg.recipe, fmt) + out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt) + return out + + meta = state.get_weight_meta(w) + if not meta.w.initialized: + meta.w.seed(w, fmt) + if not meta.x.initialized: + meta.x.seed(x, fmt) + # The kernel finalizes each ring in-kernel (overwriting the scale slot), so + # the w/x scales are taken from the ring before the quantize. + if w.dtype is not torch.bfloat16: # static pre-quantized weight + w_arg, sw_arg, w_ring = w, meta.w.scale, None else: - meta = state.get_weight_meta(w) - if not meta.w.initialized: - meta.w.seed(w, fmt) - if not meta.x.initialized: - meta.x.seed(x, fmt) - # In-kernel ring finalization: the kernels write hist[idx] and the - # next scale; idx rotates host-side (the device counter self-rearms). - w_is_fp8 = w.dtype != torch.bfloat16 - out, amax_x, amax_w = linear_forward_fp8( - x, - w, - bias, - meta.x.scale, - meta.w.scale, - fmt, - None, - meta.x.state, - meta.x.idx, - state.recipe.margin, - None if w_is_fp8 else meta.w.state, - meta.w.idx, - state.recipe.margin, - ) - meta.x.advance() - if not w_is_fp8: - meta.w.advance() + w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state + out, _x8, _w8, _ax, _aw = linear_forward_fp8( + x, + w_arg, + bias, + meta.x.scale, + sw_arg, + fmt, + None, + meta.x.state, + meta.x.idx, + margin, + w_ring, + meta.w.idx, + margin, + ) + meta.x.advance() + if w_ring is not None: + meta.w.advance() return out class _LinearFp8(torch.autograd.Function): """The fp8 linear forward/backward pair (standard Function style). - The forward runs inside the ``fp8_autocast`` region and captures the - active fmt/recipe/meta on ``ctx``; the backward reads only the captured - state, so ``loss.backward()`` may run after the context exits. The - gradient is quantized once (E5M2 in hybrid mode) and the dX / dW GEMMs - share that quantization; output masks come from ``needs_input_grad``. + The forward runs inside ``fp8_autocast`` and captures the active + fmt/recipe/meta on ``ctx``; the backward reads only that captured state, so + ``loss.backward()`` may run after the context exits. The gradient is + quantized once (E5M2 in hybrid) and both dX/dW GEMMs share it; the output + masks come from ``needs_input_grad``. """ @staticmethod def forward(ctx, x, w, bias): - out = fp8_linear_forward(x, w, bias) - state = fp8_state() + cfg = _current_config() + out = fp8_linear_forward(x, w, bias, cfg) ctx.save_for_backward(x, w) - ctx.fmt_bwd = state.fp8_format.bwd() - ctx.recipe = state.recipe - ctx.is_dynamic = isinstance(state.recipe, DynamicScaling) - ctx.meta = None if ctx.is_dynamic else state.get_weight_meta(w) + ctx.fmt_bwd = cfg.fp8_format.bwd() + ctx.recipe = cfg.recipe + ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling) + ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w) return out @staticmethod @@ -310,32 +404,33 @@ class _LinearFp8(torch.autograd.Function): def backward(ctx, g): x, w = ctx.saved_tensors fmt = ctx.fmt_bwd + # Per-recipe scale/ring selection; both branches share one call below. if ctx.is_dynamic: sg = _dynamic_scale(g, ctx.recipe, fmt) sw = _dynamic_scale(w, ctx.recipe, fmt) sx = _dynamic_scale(x, ctx.recipe, fmt) - grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( - g, x, w, list(ctx.needs_input_grad), sg, sw, sx, fmt - ) + ring, idx = None, 0 else: meta = ctx.meta if not meta.g.initialized: meta.g.seed(g, fmt) - # The g quantize kernel finalizes the gradient's ring in-kernel. - grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( - g, - x, - w, - list(ctx.needs_input_grad), - meta.g.scale, - meta.w.scale, - meta.x.scale, - fmt, - meta.g.state, - meta.g.idx, - ctx.recipe.margin, - ) - meta.g.advance() + sg, ring, idx = meta.g.scale, meta.g.state, meta.g.idx + sw, sx = meta.w.scale, meta.x.scale + grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8( + g, + x, + w, + list(ctx.needs_input_grad), + sg, + sw, + sx, + fmt, + ring, + idx, + ctx.recipe.margin, + ) + if not ctx.is_dynamic: + meta.g.advance() # the g quantize kernel finalized the ring in-kernel return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None @@ -345,31 +440,29 @@ class _LinearFp8(torch.autograd.Function): def fp8_linear_enable(enabled: bool = True) -> None: - """Toggle fp8 dispatch for aten::linear (global; backward runs on engine - worker threads, so a thread-local flag would be lost during backward).""" - fp8_state().enabled = enabled + """Toggle fp8 dispatch for aten::linear globally (the out-of-region default; + ``fp8_autocast`` regions override it thread-locally).""" + fp8_state().default_enabled = enabled def fp8_linear_enabled() -> bool: - return fp8_state().enabled + """Whether fp8 dispatch is active right now (region config or global).""" + return _active() is not None def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool: - """Shape guard for the fp8 linear path. - - Unlike a strict 16-alignment requirement, the fp8 kernels handle unaligned - M/N via boundary checks (slower but correct) — so no whole-call bf16 - fallback for small decode batches. Only the K-dimension contraction must - match, and the weight must be 2D. - """ + """Shape guard for the fp8 path. Unlike a strict 16-alignment requirement, + the kernels handle unaligned M/N via boundary checks (slower but correct) — + so no whole-call bf16 fallback for small decode batches. Only the K-dimension + contraction must match and the weight must be 2D.""" return x.dim() >= 2 and w.dim() == 2 and x.size(-1) == w.size(1) def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): if ( - fp8_linear_enabled() - and x.dtype == torch.bfloat16 - and w.dtype == torch.bfloat16 + _active() is not None + and x.dtype is torch.bfloat16 + and w.dtype is torch.bfloat16 and _fp8_supported(x, w) ): return _LinearFp8.apply(x, w, bias) @@ -384,9 +477,8 @@ def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): _lib = Library("aten", "IMPL", "CUDA") _lib.impl("linear", _linear_cuda_impl) # Also replace torch's generated linear autograd formula (which would call -# aten::linear_backward after the fp8_autocast region exits). The fp8 -# backward is owned by _LinearFp8 with its state captured at forward time, -# so loss.backward() works wherever it is called; the same CUDA registration -# still covers inference_mode, where autograd keys are skipped entirely. +# aten::linear_backward after the fp8_autocast region exits). The fp8 backward +# is owned by _LinearFp8 with state captured at forward time, so loss.backward() +# works wherever it is called; the CUDA registration still covers inference_mode. _lib_autograd = Library("aten", "IMPL", "AutogradCUDA") _lib_autograd.impl("linear", _linear_cuda_impl) diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index 9966525..c399d2e 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -4,7 +4,7 @@ Isolates the ``fp8_ops`` CUDA extension behind stable Python primitives: - ``quantize_bf16(x, scale, fmt) -> (x8, amax)`` — BF16 → FP8 with fused amax - ``mm_fp8(a8, b8, sa, sb) -> out`` — pre-quantized FP8 GEMM (BF16 output) -- ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, amax_x, amax_w)`` +- ``linear_forward_fp8(x, w, bias, sx, sw) -> (out, x8, w8, amax_x, amax_w)`` - ``linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt) -> (gx, gw, gb, amax_g)`` Scale semantics: scales are *quantization steps* — the value divided out when @@ -16,6 +16,8 @@ Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``; this module is stateless. """ +from typing import List, Optional, Tuple + import torch from torch.library import custom_op @@ -39,7 +41,7 @@ def _fmt_dtype(fmt: str) -> torch.dtype: @custom_op("custom::fp8_quantize", mutates_args=()) def fp8_quantize( x: torch.Tensor, scale: torch.Tensor, fmt: int -) -> tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor]: """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``.""" @@ -73,7 +75,7 @@ def fp8_gemm( sa: torch.Tensor, sb: torch.Tensor, out_dtype: int = 0, - out_scale: torch.Tensor | None = None, + out_scale: Optional[torch.Tensor] = None, ) -> torch.Tensor: """FP8 GEMM: ``a @ b * (sa * sb)`` with FP32 accumulation. @@ -106,7 +108,9 @@ def _fp8_gemm_cpu(a, b, sa, sb, out_dtype=0, out_scale=None): return acc.to(torch.bfloat16) -def quantize_bf16(x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3"): +def quantize_bf16( + x: torch.Tensor, scale: torch.Tensor, fmt: str = "e4m3" +) -> Tuple[torch.Tensor, torch.Tensor]: """BF16 -> FP8 quantize with fused amax; returns ``(x8, amax)``. ``scale`` is the quantization step (device scalar); ``fmt`` selects @@ -122,7 +126,7 @@ def mm_fp8( sa: torch.Tensor, sb: torch.Tensor, out_dtype: str = "bf16", - out_scale: torch.Tensor | None = None, + out_scale: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Pre-quantized FP8 GEMM: ``a @ b * (sa * sb)``. @@ -139,24 +143,28 @@ def mm_fp8( def linear_forward_fp8( - x, - w, - bias, - sx, - sw, + x: torch.Tensor, + w: torch.Tensor, + bias: Optional[torch.Tensor], + sx: torch.Tensor, + sw: torch.Tensor, fmt: str = "e4m3", - bias_scale=None, - x_ring=None, + bias_scale: Optional[torch.Tensor] = None, + x_ring: Optional[torch.Tensor] = None, x_ring_idx: int = 0, x_ring_margin: int = 0, - w_ring=None, + w_ring: Optional[torch.Tensor] = None, w_ring_idx: int = 0, w_ring_margin: int = 0, -): +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Pure FP8 linear forward: quantize x/w to ``fmt``, pre-quantized GEMM. - Returns ``(out, amax_x, amax_w)``. ``bias`` may be ``None``. For static - fp8 inference, ``w`` and ``bias`` may arrive pre-quantized to ``fmt`` + Returns ``(out, x8, w8, amax_x, amax_w)`` — the quantized operands are + handed back so the policy layer can cache the weight quantization while + the weight tensor is unchanged (torch autocast's cached_cast analog). + ``x8`` is ``[M, K]`` and ``w8`` is ``[N, K]`` (the passed-in ``w`` itself + on the pre-quantized path). ``bias`` may be ``None``. For static fp8 + inference, ``w`` and ``bias`` may arrive pre-quantized to ``fmt`` (produced by :func:`quantize_bf16` with their scales as ``sw`` / ``bias_scale``); a pre-quantized ``bias`` requires ``bias_scale``, and its ``amax_w`` comes back 0. The bias is fused into the GEMM epilogue. @@ -190,18 +198,18 @@ def linear_forward_fp8( def linear_backward_fp8( - g, - x, - w, - masks, - sg, - sw, - sx, + g: torch.Tensor, + x: torch.Tensor, + w: torch.Tensor, + masks: List[bool], + sg: torch.Tensor, + sw: torch.Tensor, + sx: torch.Tensor, fmt: str = "e5m2", - g_ring=None, + g_ring: Optional[torch.Tensor] = None, g_ring_idx: int = 0, g_ring_margin: int = 0, -): +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``. The gradient (and the transposed w/x operands) are quantized to ``fmt`` diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index e2b02a1..1cb5d83 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -61,6 +61,12 @@ struct Fp8GemmTraits { static constexpr __nv_fp8_interpretation_t kNvFormat = kIsE5M2 ? __NV_E5M2 : __NV_E4M3; static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f; + + // Derived launch geometry: 64x32 warp tiles give the CTA thread count. + // The shared-memory budget is layout-aware (crosswise operands add K- + // major staging + a canonical buffer), so it lives in Fp8GemmSmem in + // gemm.cuh together with the resident-CTA hint for __launch_bounds__. + static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32; }; // Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index 657d243..9e43bd6 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -191,20 +191,11 @@ __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) { ((((col >> 4) ^ ((row >> kShift) & (kChunks - 1))) << 4) + (col & 15)); } -// Stage-load one GEMM operand into the canonical flat [rows * K] shared tile -// (addressing via tile_at, so stores land in the swizzled layout). The -// transpose is folded into the staging step via a CUTLASS-style crosswise -// layout: RowMajor (stored [rows][contract]) copies 16-byte K-contiguous runs -// with cp.async, while ColMajor (stored [contract][rows]) reads 16-byte runs -// along the operand's contiguous non-contract dim and scatters them across -// the tile's rows. Crosswise runs cannot use cp.async (the 16 destination -// bytes land on 16 different rows), so their global loads are plain LDGs — -// issued as one batch per row group before the first scatter so their -// latencies overlap instead of serializing behind the shared stores. -// RowsTile is the tile's row capacity (kBlockM / kBlockN) and kThreads the -// CTA size; the runtime `rows` bound may be smaller (tail predication). -// `block_row` is this block's origin in the operand's row dim. -template +// Stage-load a CONGRUOUS operand (stored [rows][contract], contract- +// contiguous — the only cp.async-able shape for the canonical tile) into the +// flat [rows * K] shared tile via tile_at's swizzle. Crosswise operands go +// through stage_crosswise_tile + transpose_crosswise_tile instead. +template __device__ __forceinline__ void load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows, int64_t contract, int64_t ld, int tid, int64_t k_base, @@ -213,101 +204,29 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows, static_assert(RowsTile * kChunks % kThreads == 0, "tile chunks must divide evenly across threads"); constexpr int kCpt = RowsTile * kChunks / kThreads; // chunks per thread - if constexpr (std::is_same_v) { - // Operand stored [contract][rows]: contiguous along the non-contract - // dim. Each thread scatters one 16-byte run per K/32 pass; when the - // tile has more 16-row groups than warps (RowsTile > kThreads/2), - // each thread covers several groups. - constexpr int kWarpsTile = kThreads / 32; - constexpr int kGroups = RowsTile / 16; - constexpr int kPasses = K / 32; - static_assert(kGroups % kWarpsTile == 0, - "row groups must divide evenly across warps"); - const int kl = tid & 31; // byte column within a 32B pass - // r0 is always a multiple of 16 (block_row is a multiple of RowsTile - // and each group covers 16 rows), so every run shares the base+ld - // alignment: one uniform check instead of one per pass. - const bool run_aligned = - ((reinterpret_cast(operand) | ld) & 15) == 0; + // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of + // one row (K=64: a contiguous 32B pair; K=32: a single chunk). + constexpr int kCpr = kChunks / kCpt; // chunks per row slice + const int r = tid / kCpr; + const int c0 = (tid % kCpr) * kCpt * 16; + const int64_t row = block_row + r; + const bool row_ok = row < rows; + // k_base and every c are multiples of 16, so the per-chunk sources + // share the row base's alignment. + const auto* src = operand + row * ld + k_base; + const bool chunk_aligned = (reinterpret_cast(src) & 15) == 0; #pragma unroll - for (int g = 0; g < kGroups / kWarpsTile; ++g) { - const int rg = (tid >> 5) + g * kWarpsTile; - const int64_t r0 = block_row + rg * 16; - const bool rows_full = r0 + 15 < rows; // pass-invariant - // Batch every 16B run load of this row group before the first - // scatter: the LDGs are independent, and the byte-granular - // shared stores would otherwise serialize behind each one. - uint4 v[kPasses]; - bool fast[kPasses]; + for (int j = 0; j < kCpt; ++j) { + const int c = c0 + j * 16; + T8* dst = tile_at(tile, r, c); + if (row_ok && chunk_aligned && k_base + c + 15 < contract) { + astrai::cp_async_16(dst, src + c, true); + } else { + // Tail chunk (or misaligned base): predicated scalar fill. #pragma unroll - for (int pass = 0; pass < kPasses; ++pass) { - const int64_t k_idx = k_base + kl + pass * 32; - fast[pass] = rows_full && run_aligned && k_idx < contract; - if (fast[pass]) - v[pass] = - *reinterpret_cast(operand + k_idx * ld + r0); - } -#pragma unroll - for (int pass = 0; pass < kPasses; ++pass) { - const int col = kl + pass * 32; - if (fast[pass]) { - const auto* bytes = reinterpret_cast(&v[pass]); - // Scatter 16 bytes along the tile rows through tile_at's - // swizzle. Rows sharing a physical chunk form groups of - // (8 / kChunks) consecutive rows (see tile_at), so each - // group is one tile_at address plus a K-byte row stride. - constexpr int kGrp = 8 / kChunks; -#pragma unroll - for (int j = 0; j < 16 / kGrp; ++j) { - T8* p = tile_at(tile, rg * 16 + j * kGrp, col); -#pragma unroll - for (int i = 0; i < kGrp; ++i) - p[i * K] = bytes[j * kGrp + i]; - } - } else if (k_base + col < contract) { - // Row-tail or misaligned run: byte-granular gather with - // per-row predication (the k column itself is in range). -#pragma unroll - for (int i = 0; i < 16; ++i) { - const int64_t r_idx = r0 + i; - *tile_at(tile, rg * 16 + i, col) = - r_idx < rows ? operand[(k_base + col) * ld + r_idx] - : T8(0.0f); - } - } else { - // Contract tail: straight zero-fill, no global traffic. -#pragma unroll - for (int i = 0; i < 16; ++i) - *tile_at(tile, rg * 16 + i, col) = T8(0.0f); - } - } - } - } else { - // Operand stored [rows][contract]: contiguous along the contract dim. - // Linear chunk mapping: thread covers kCpt consecutive 16B chunks of - // one row (K=64: a contiguous 32B pair; K=32: a single chunk). - constexpr int kCpr = kChunks / kCpt; // chunks per row slice - const int r = tid / kCpr; - const int c0 = (tid % kCpr) * kCpt * 16; - const int64_t row = block_row + r; - const bool row_ok = row < rows; - // k_base and every c are multiples of 16, so the per-chunk sources - // share the row base's alignment. - const auto* src = operand + row * ld + k_base; - const bool chunk_aligned = (reinterpret_cast(src) & 15) == 0; -#pragma unroll - for (int j = 0; j < kCpt; ++j) { - const int c = c0 + j * 16; - T8* dst = tile_at(tile, r, c); - if (row_ok && chunk_aligned && k_base + c + 15 < contract) { - astrai::cp_async_16(dst, src + c, true); - } else { - // Tail chunk (or misaligned base): predicated scalar fill. -#pragma unroll - for (int i = 0; i < 16; ++i) - dst[i] = - row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f); - } + for (int i = 0; i < 16; ++i) + dst[i] = + row_ok && k_base + c + i < contract ? src[c + i] : T8(0.0f); } } } @@ -320,15 +239,195 @@ load_operand_tile(T8* tile, const T8* __restrict__ operand, int64_t rows, // --------------------------------------------------------------------------- // Swizzled 16B-chunk address (tile_at's layout) as a raw shared-memory -// pointer for ldmatrix. Valid for kK in {32, 64} (the swizzle itself lives -// only in tile_at; this wrapper just converts the element address). +// pointer for ldmatrix. Valid for kK in {32, 64, 128} (the swizzle itself +// lives only in tile_at; this wrapper just converts the element address). template __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk) { - static_assert(kK == 32 || kK == 64, - "fragment swizzle offsets assume kK in {32, 64}"); + static_assert(kK == 32 || kK == 64 || kK == 128, + "fragment swizzle offsets assume kK in {32, 64, 128}"); return __cvta_generic_to_shared(tile_at(tile, row, chunk << 4)); } +// Crosswise operands (stored [contract][rows], rows-contiguous) cannot be +// cp.async'd into the canonical [rows][contract] tile — a 16B global run +// holds one contract byte for each of 16 rows. They stage K-major instead +// (byte (p, r) at p*RowsTile + r), where the very same runs land contiguously +// and cp.async applies unchanged; a per-tile smem->smem transpose (below) +// then produces the canonical swizzled tile the MMA fragments read. This +// keeps the whole global→shared path asynchronous — the synchronous +// LDG+byte-scatter staging this replaces left the kernel long-scoreboard +// bound (ncu: 4.6 stalled loads per issue vs 0.4 on the congruous path). +template +__device__ __forceinline__ void +stage_crosswise_tile(T8* staging, const T8* __restrict__ operand, int64_t rows, + int64_t contract, int64_t ld, int tid, int64_t k_base, + int64_t block_row) { + constexpr int kRuns = K * RowsTile / 16; // 16B runs per tile + // r0 is a multiple of 16 and p*ld keeps 16B alignment whenever ld has it, + // so one uniform verdict covers every run. + const bool run_aligned = + ((reinterpret_cast(operand) | ld) & 15) == 0; + for (int run = tid; run < kRuns; run += kThreads) { + const int pl = run % K; // local contract byte (column of the run) + const int rg = run / K; // 16-row group + const int64_t r0 = block_row + (int64_t)rg * 16; + T8* dst = staging + pl * RowsTile + rg * 16; + if (run_aligned && r0 + 15 < rows && k_base + pl < contract) + astrai::cp_async_16(dst, operand + (k_base + pl) * ld + r0, true); + else { + // Row tail, contract tail or misaligned base: predicated fill. +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int64_t r = r0 + i; + dst[i] = r < rows && k_base + pl < contract + ? operand[(k_base + pl) * ld + r] + : T8(0.0f); + } + } + } +} + +// Direct (synchronous) crosswise load into a canonical rotating stage: +// LDG.128 x4 (4 consecutive contract bytes x 16 rows) + in-register PRMT +// transpose + 16 STS.32. Used for crosswise operands whose global data is +// typically L2-resident (the A side of dW): the staging detour's extra +// shared-memory round trip costs more than the latency it hides there, +// while crosswise B operands (DRAM-streamed weights of dX) take the +// asynchronous stage_crosswise_tile path instead. +template +__device__ __forceinline__ void +load_crosswise_direct(T8* tile, const T8* __restrict__ operand, int64_t rows, + int64_t contract, int64_t ld, int tid, int64_t k_base, + int64_t block_row) { + constexpr int kQuads = K / 4; // 4-byte contract quads per tile + constexpr int kGroups = RowsTile / 16; + constexpr int kTChunks = kQuads * kGroups; // 64B chunks per tile + // r0 is always a multiple of 16 (block_row is a multiple of RowsTile and + // each group covers 16 rows), and p*ld keeps the base 16B-aligned + // whenever ld is, so every run of a chunk shares one alignment verdict. + const bool run_aligned = + ((reinterpret_cast(operand) | ld) & 15) == 0; + for (int chunk = tid; chunk < kTChunks; chunk += kThreads) { + const int quad = chunk / kGroups; + const int rg = chunk % kGroups; + const int64_t r0 = block_row + rg * 16; + const bool rows_full = r0 + 15 < rows; + if (rows_full && run_aligned) { + const int64_t p0 = k_base + quad * 4; + uint4 v[4]; +#pragma unroll + for (int s = 0; s < 4; ++s) { + // Contract tail: a run past k carries zero bytes; they flow + // through the PRMT transpose like any other value. + if (p0 + s < contract) + v[s] = *reinterpret_cast( + operand + (p0 + s) * ld + r0); + else + v[s] = make_uint4(0u, 0u, 0u, 0u); + } + const unsigned* bytes = reinterpret_cast(v); +#pragma unroll + for (int i = 0; i < 16; ++i) { + // word i = row r0+i's quad: byte i of each of the four runs + // [v0.b(i), v1.b(i), v2.b(i), v3.b(i)]. Byte i of a uint4 + // lives in its (i>>2)-th 32-bit register. + const unsigned nib = i & 3; + const unsigned sel = nib | ((nib + 4) << 4); + const unsigned w01 = + __byte_perm(bytes[0 + (i >> 2)], bytes[4 + (i >> 2)], sel); + const unsigned w23 = + __byte_perm(bytes[8 + (i >> 2)], bytes[12 + (i >> 2)], sel); + *reinterpret_cast(tile_at(tile, rg * 16 + i, + quad * 4)) = + __byte_perm(w01, w23, 0x5410u); + } + } else { + // Row-tail or misaligned chunk: byte-granular gather with + // per-row predication; contract-tail columns zero-fill. +#pragma unroll + for (int s = 0; s < 4; ++s) { + const int col = quad * 4 + s; + if (k_base + col >= contract) { +#pragma unroll + for (int i = 0; i < 16; ++i) + *tile_at(tile, rg * 16 + i, col) = T8(0.0f); + continue; + } +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int64_t r_idx = r0 + i; + *tile_at(tile, rg * 16 + i, col) = + r_idx < rows + ? operand[(k_base + col) * ld + r_idx] + : T8(0.0f); + } + } + } + } +} + +// K-major staging -> canonical [rows][kK] swizzled tile, one chunk at a time. +// Each chunk (indexed within a k_seg region of `quads_per_seg` quads) covers +// 4 consecutive contract bytes x 16 rows: four LDS.128 grab the staging runs, +// PRMT byte selects transpose them in registers, and sixteen STS.32 land the +// row quads through tile_at's swizzle — 4x fewer store instructions than a +// byte-granular scatter. Chunk-at-a-time lets the caller pool work across +// operands; the region restriction lets the main loop overlap one region's +// transpose with another region's MMAs (a whole-tile serial transpose put +// the crosswise GEMMs at 25% tensor utilization). +template +__device__ __forceinline__ void +transpose_crosswise_region(T8* tile, const T8* staging, int idx, int quad0) { + constexpr int kGroups = RowsTile / 16; + const int quad = quad0 + idx / kGroups; + const int rg = idx % kGroups; + // The four runs sit RowsTile bytes apart (one per contract byte of the + // quad); each run is 16 contiguous staging bytes = 16 rows. + const char* run0 = reinterpret_cast( + staging + quad * 4 * RowsTile + rg * 16); + uint4 v[4]; +#pragma unroll + for (int s = 0; s < 4; ++s) + v[s] = *reinterpret_cast(run0 + s * RowsTile); + const unsigned* bytes = reinterpret_cast(v); +#pragma unroll + for (int i = 0; i < 16; ++i) { + // word i = row r0+i's quad: byte i of each of the four runs + // [v0.b(i), v1.b(i), v2.b(i), v3.b(i)]. Byte i of a uint4 lives in + // its (i>>2)-th 32-bit register. + const unsigned nib = i & 3; + const unsigned sel = nib | ((nib + 4) << 4); + const unsigned w01 = + __byte_perm(bytes[0 + (i >> 2)], bytes[4 + (i >> 2)], sel); + const unsigned w23 = + __byte_perm(bytes[8 + (i >> 2)], bytes[12 + (i >> 2)], sel); + *reinterpret_cast(tile_at(tile, rg * 16 + i, quad * 4)) = + __byte_perm(w01, w23, 0x5410u); + } +} + +// Layout-aware shared-memory budget and occupancy hint. A congruous or +// direct-crosswise operand needs its kStages rotating canonical buffers; a +// staged-crosswise operand (crosswise B with kBStaged) needs kStages K-major +// staging buffers plus ONE canonical buffer (rewritten every tile by the +// in-kernel transpose). The 48KB static-smem watermark picks the resident-CTA +// hint for __launch_bounds__ (sm_89: 100KB smem per SM, so two CTAs fit while +// each stays within the static budget). +template +struct Fp8GemmSmem { + // Crosswise = the stage-load's view: A's tag directly, B's transposed. + // A-crosswise always loads direct (L2-typical activations); B-crosswise + // stages only when its contract dim is long enough to stream DRAM. + static constexpr bool kCrossA = std::is_same_v; + static constexpr bool kCrossB = std::is_same_v; + static constexpr bool kBStagePath = kCrossB && StagedB; + static constexpr int kBytes = + Traits::kStages * Traits::kBlockM * Traits::kK + + (kBStagePath ? Traits::kStages + 1 : Traits::kStages) * + Traits::kBlockN * Traits::kK; + static constexpr int kMinCtas = kBytes <= 48 * 1024 ? 2 : 1; +}; + // LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h): // A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] / // ColMajor = [N][K]. The kernel always computes @@ -342,25 +441,43 @@ __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk // (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and // exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the // launcher dispatches to it there (see launch_fp8_gemm). -template -__global__ void - __launch_bounds__((Traits::kBlockM / 64) * (Traits::kBlockN / 32) * 32, 2) +template +__global__ void __launch_bounds__(Traits::kCtaThreads, + Fp8GemmSmem::kMinCtas) fp8_gemm_kernel(FP8Params p) { using T8 = std::conditional_t; constexpr int kBlockM = Traits::kBlockM; constexpr int kBlockN = Traits::kBlockN; constexpr int kK = Traits::kK; constexpr int kStages = Traits::kStages; - constexpr int kCtaThreads = (kBlockM / 64) * (kBlockN / 32) * 32; + constexpr int kCtaThreads = Traits::kCtaThreads; + constexpr bool kCrossA = Fp8GemmSmem::kCrossA; + constexpr bool kCrossB = Fp8GemmSmem::kCrossB; + constexpr bool kBStagePath = + Fp8GemmSmem::kBStagePath; static_assert(kStages >= 1 && kStages <= 8, - "FP8 GEMM stages must be in the range [1, 8]"); + "FP8 GEMM stages must be in [1, 8]"); // Tiles are flat [rows * kK] with a 16B-chunk XOR swizzle (tile_at): // ldmatrix reads whole 16B chunks through the same mapping the staging // writes, and the swizzle removes the bank conflict the unswizzled - // 8-word row stride caused (see tile_at). - __shared__ __align__(16) T8 a_smem[kStages][kBlockM * kK]; - __shared__ __align__(16) T8 b_smem[kStages][kBlockN * kK]; + // 8-word row stride caused (see tile_at). The stages live in dynamic + // shared memory so deep pipelines (kStages * (kBlockM + kBlockN) * kK > + // 48KB static limit) opt in via cudaFuncSetAttribute in the launcher. + extern __shared__ __align__(16) char fp8_gemm_smem[]; + // Per operand: congruous or direct-crosswise = kStages rotating canonical + // buffers; staged-crosswise (B) = kStages K-major staging buffers (filled + // by cp.async, one per tile in flight) followed by one canonical buffer + // the per-tile transpose rewrites. + constexpr int kAStageBytes = kBlockM * kK; + constexpr int kBStageBytes = kBlockN * kK; + constexpr int kStB = kStages; // B staging ring size (see above) + T8* const a_base = reinterpret_cast(fp8_gemm_smem); + T8* const b_base = + reinterpret_cast(fp8_gemm_smem + kStages * kAStageBytes); + T8* const b_canon = b_base + kStages * kBStageBytes; // crosswise B only const auto* a = reinterpret_cast(p.a_ptr); const auto* b = reinterpret_cast(p.b_ptr); @@ -379,12 +496,14 @@ __global__ void // before advancing along N. All CTAs of one group share the same B column // stripe, so B tiles stay hot in L2 across the wave (the default // N-fastest order makes each wave touch every B tile instead). - // Measured win for the A-crosswise layouts (10-21% at K>=2048) and loss - // for A-congruous (-17..20%, A's cp.async stream prefers the N-fastest - // order) — so the branch follows LayoutA. + // kGroupRaster is a template knob (the launcher defaults it to the + // measured best per layout: grouped for A-crosswise (dW) and for the + // congruous NT forward — whose big B operand gains the most from the + // shared stripe — plain for dX's crosswise-B layouts, where it measured + // neutral). constexpr int kGroupM = 8; int block_m, block_n; - if constexpr (std::is_same_v) { + if constexpr (kGroupRaster) { const int blocks_m = gridDim.y; const int bid = blockIdx.y * gridDim.x + blockIdx.x; const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM; @@ -409,18 +528,48 @@ __global__ void const float sb = *p.scale_b; float acc[4][4][4] = {}; // [nt][mt][acc] - // Both operands are staged into the canonical [M][kK] / [N][kK] shared - // tiles regardless of their global layout (see load_operand_tile), so the - // MMA fragment reads below stay unchanged across the four layout - // combinations. A's tag already names the operand view ([M][K] = - // [rows][contract]); B's tag is relative to the canonical [K][N], so the - // stage-load sees its transpose (transpose_layout_t, see common.h). - auto load_tile = [&](int stage, int64_t k_base) { - load_operand_tile( - a_smem[stage], a, m, k, a_ld, tid, k_base, (int64_t)block_m * kBlockM); - load_operand_tile, kBlockN, - kCtaThreads>(b_smem[stage], b, n, k, b_ld, tid, k_base, - (int64_t)block_n * kBlockN); + // Both operands end up in the canonical [M][kK] / [N][kK] shared tiles + // the MMA fragments read, regardless of their global layout. A's tag + // already names the operand view ([M][K] = [rows][contract]); B's tag is + // relative to the canonical [K][N], so the stage-load sees its transpose + // (transpose_layout_t, see common.h). Congruous operands cp.async + // straight into their rotating canonical buffers; crosswise operands + // cp.async into K-major staging (zero transformation) and get a per-tile + // smem->smem transpose below. + auto load_tile = [&](int tile, int stage, int64_t k_base) { + if constexpr (kCrossA) { + load_crosswise_direct( + a_base + stage * kAStageBytes, a, m, k, a_ld, tid, k_base, + (int64_t)block_m * kBlockM); + } else { + load_operand_tile( + a_base + stage * kAStageBytes, a, m, k, a_ld, tid, k_base, + (int64_t)block_m * kBlockM); + } + if constexpr (kBStagePath) { + stage_crosswise_tile( + b_base + (tile % kStB) * kBStageBytes, b, n, k, b_ld, tid, + k_base, (int64_t)block_n * kBlockN); + } else if constexpr (kCrossB) { + load_crosswise_direct( + b_base + stage * kBStageBytes, b, n, k, b_ld, tid, k_base, + (int64_t)block_n * kBlockN); + } else { + load_operand_tile( + b_base + stage * kBStageBytes, b, n, k, b_ld, tid, k_base, + (int64_t)block_n * kBlockN); + } + }; + // smem->smem transpose of one k_seg region (kSegQuads contract quads) of + // this tile's staged-crosswise B into its single canonical buffer. + auto transpose_tile = [&](int tile, int seg) { + if constexpr (!kBStagePath) return; + constexpr int kSegQuads = kK / 4 / (kK / kMmaK); // quads per k_seg + constexpr int kBRegion = kSegQuads * (kBlockN / 16); + const T8* b_stg = b_base + (tile % kStB) * kBStageBytes; + for (int idx = tid; idx < kBRegion; idx += kCtaThreads) + transpose_crosswise_region(b_canon, b_stg, idx, + seg * kSegQuads); }; const int64_t tile_count = (k + kK - 1) / kK; @@ -445,7 +594,7 @@ __global__ void #pragma unroll for (int stage = 0; stage < kStages; ++stage) { if (stage < tile_count) { - load_tile(stage, static_cast(stage) * kK); + load_tile(stage, stage, static_cast(stage) * kK); astrai::cp_async_commit_group(); } } @@ -463,48 +612,77 @@ __global__ void // before any thread reads tiles written by other threads. __syncthreads(); + // Staged-crosswise B: produce the canonical tile one k_seg region at + // a time so each region's transpose overlaps the previous region's + // MMA sequence (the transposes are pure shared-memory traffic — B's + // global path stayed fully asynchronous above). + constexpr int kSegs = kK / kMmaK; + if constexpr (kBStagePath) { + transpose_tile(tile_index, 0); + // Barrier 2: region 0 visible to every thread before its + // fragment loads. (Compiled out for congruous/direct layouts.) + __syncthreads(); + } + + const T8* a_tile = a_base + (size_t)stage * kAStageBytes; + const T8* b_tile = + kBStagePath ? b_canon : b_base + (size_t)stage * kBStageBytes; + // 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg — // 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in // the 128x64-tile version (the kernel was LSU-issue-bound there). - constexpr int kSegs = kK / kMmaK; - // B fragments double-buffered across k_segs: the next k_seg's B load - // is issued before the current k_seg's MMA sequence, so its LDS - // latency hides behind the A pipeline + tensor-pipe work (same trick - // as the A mt+1 prefetch below; costs kSegs x 8 registers). + // B fragments double-buffer across k_segs while B is congruous (no + // region writes in flight); a crosswise B reloads per k_seg after + // the region's transpose became visible. unsigned b_frag[2][4][2]; + if constexpr (!kBStagePath) { #pragma unroll - for (int nt = 0; nt < 4; ++nt) { - const int row = b_row0 + nt * 8 + r7; - astrai::ldmatrix_x2_lane(b_frag[0][nt], - frag_addr(b_smem[stage], row, rh8)); + for (int nt = 0; nt < 4; ++nt) { + const int row = b_row0 + nt * 8 + r7; + astrai::ldmatrix_x2_lane(b_frag[0][nt], + frag_addr(b_tile, row, rh8)); + } } #pragma unroll for (int k_seg = 0; k_seg < kSegs; ++k_seg) { const int bcur = k_seg & 1, bnext = bcur ^ 1; - if (k_seg + 1 < kSegs) { + if constexpr (kBStagePath) { +#pragma unroll + for (int nt = 0; nt < 4; ++nt) { + const int row = b_row0 + nt * 8 + r7; + astrai::ldmatrix_x2_lane( + b_frag[bcur][nt], + frag_addr(b_tile, row, k_seg * 2 + rh8)); + } + } else if (k_seg + 1 < kSegs) { #pragma unroll for (int nt = 0; nt < 4; ++nt) { const int row = b_row0 + nt * 8 + r7; astrai::ldmatrix_x2_lane( b_frag[bnext][nt], - frag_addr(b_smem[stage], row, - (k_seg + 1) * 2 + rh8)); + frag_addr(b_tile, row, (k_seg + 1) * 2 + rh8)); } } + // Region k_seg+1's transpose overlaps this region's MMA work + // (disjoint canonical regions, no race). + if constexpr (kBStagePath) { + if (k_seg + 1 < kSegs) + transpose_tile(tile_index, k_seg + 1); + } // Software-pipelined A fragments: the ldmatrix.x4 for row mt+1 // is issued before the MMAs consuming row mt, so the LDS fixed // latency hides behind tensor-pipe work (cuts the `wait` stall, // ~2.3 cycles/issue before this). Costs 4 extra registers. unsigned a_frag[5][4]; astrai::ldmatrix_x4_lane( - a_frag[0], frag_addr(a_smem[stage], a_row0 + rh8 * 8 + r7, + a_frag[0], frag_addr(a_tile, a_row0 + rh8 * 8 + r7, k_seg * 2 + rh16)); #pragma unroll for (int mt = 0; mt < 4; ++mt) { if (mt < 3) astrai::ldmatrix_x4_lane( a_frag[mt + 1], - frag_addr(a_smem[stage], + frag_addr(a_tile, a_row0 + (mt + 1) * 16 + rh8 * 8 + r7, k_seg * 2 + rh16)); #pragma unroll @@ -512,12 +690,19 @@ __global__ void astrai::mma_sync(acc[nt][mt], a_frag[mt], b_frag[bcur][nt], acc[nt][mt]); } + // Barrier 3: region k_seg+1's transposes complete and become + // visible before the next k_seg reads them. + if constexpr (kBStagePath) { + if (k_seg + 1 < kSegs) __syncthreads(); + } } - // Barrier 2: every thread finished reading this stage's tiles before - // the prefetch for the (i+kStages)-th tile overwrites them. + // Barrier 4: every thread finished reading this stage's tiles before + // the prefetch for the (i+kStages)-th tile overwrites them (and the + // next iteration's transposes rewrite the canonical buffer). __syncthreads(); if (tile_index + kStages < tile_count) { - load_tile(stage, (tile_index + kStages) * kK); + load_tile(tile_index + kStages, stage, + (tile_index + kStages) * kK); astrai::cp_async_commit_group(); } } @@ -593,26 +778,74 @@ void launch_fp8_quantize(const FP8QuantizeParams& p, cudaStream_t stream) { fp8_quantize_kernel<<>>(p); } +// Launch one kernel instantiation with its shared-memory budget: stages live +// in dynamic smem, so budgets beyond the 48KB static limit opt in once per +// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared +// memory"). Templated on the kernel *value* (auto NTTP) so every +// instantiation owns its own armed flag — same-signature kernels must not +// share it (the attribute is per-function). +template +void launch_with_smem(int smem_bytes, dim3 grid, dim3 block, + cudaStream_t stream, Args... args) { + if (smem_bytes > 48 * 1024) { + static bool armed = false; // per instantiation + if (!armed) { + cudaFuncSetAttribute(Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes); + armed = true; + } + } + Kernel<<>>(args...); +} + // Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles). -// kK selects the K tile (32 or 64; 64 halves the __syncthreads count per K -// and doubles the MMA work per stage, at 2x the smem per stage — measured -// 10-35% across shapes, so 64 is the default). Stages=2 with kK=64 keeps the -// pipeline at 32KB smem; deeper pipelines only win on K >= 4096 squares and -// lose elsewhere. LayoutA/LayoutB mirror the kernel template (defaults keep -// the NN layout: out = a @ b). m <= 64 dispatches to the 64x128 CTA — a -// 128-row CTA would waste half its MMA work on predicated-off rows. +// kK selects the K tile (32 / 64 / 128; larger kK halves the __syncthreads +// count per K and doubles the MMA work per stage at more smem per stage). +// Stages is the cp.async pipeline depth (smem = Stages * (BM + BN) * kK +// bytes for congruous layouts; deep pipelines are dynamic-smem backed, 1 +// CTA/SM past 48KB). GroupRaster defaults to the historically-measured best +// per LayoutA (grouped for A-crosswise, plain for A-congruous). m <= 64 +// dispatches to the 64x128 CTA — a 128-row CTA would waste half its MMA work +// on predicated-off rows. +// Crosswise B takes the asynchronous staging+transpose pipeline only when +// the contract dim is long enough that B streams from DRAM (dX-class GEMMs, +// k = N_ffn); short-K crosswise GEMMs (dW: k = M tokens) read L2-resident +// operands, where the staging round trip costs more shared-memory traffic +// than the latency it hides (measured: dW ~37 TF direct vs ~29 TF staged, +// dX ~39 TF staged vs ~38 direct). +constexpr int64_t kCrossStageMinK = 8192; + template + typename LayoutB = RowMajor, int kK = 64, int Stages = 2, + bool GroupRaster = std::is_same_v || std::is_same_v> void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { dim3 grid((p.n + 127) / 128, (p.m + 127) / 128); + const bool b_staged = p.k >= kCrossStageMinK; if (p.m <= 64) { using Traits = Fp8GemmTraits; - fp8_gemm_kernel - <<>>(p); + if (b_staged) + launch_with_smem>( + Fp8GemmSmem::kBytes, grid, + dim3(Traits::kCtaThreads), stream, p); + else + launch_with_smem>( + Fp8GemmSmem::kBytes, grid, + dim3(Traits::kCtaThreads), stream, p); } else { using Traits = Fp8GemmTraits; - fp8_gemm_kernel - <<>>(p); + if (b_staged) + launch_with_smem>( + Fp8GemmSmem::kBytes, grid, + dim3(Traits::kCtaThreads), stream, p); + else + launch_with_smem>( + Fp8GemmSmem::kBytes, grid, + dim3(Traits::kCtaThreads), stream, p); } } diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index cdb8b2c..790f220 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -157,7 +157,10 @@ std::tuple quantize_bf16(torch::Tensor x, auto x8 = torch::empty_like( x_c, x_c.options().dtype(fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn)); - auto amax = torch::zeros({1}, x_c.options().dtype(torch::kFloat32)); + // amax feeds the cross-block atomic_max; zero it on the stream (empty + + // memset, not torch::zeros — the latter routes through a fill_ dispatcher). + auto amax = torch::empty({1}, x_c.options().dtype(torch::kFloat32)); + cudaMemsetAsync(amax.data_ptr(), 0, sizeof(float), stream.stream()); FP8QuantizeParams p; pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax, nullptr, 0, 0, x_c.numel()); @@ -230,18 +233,25 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor sa, return out; } -std::tuple linear_forward_fp8( - torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx, - torch::Tensor sw, int64_t fmt, c10::optional bias_scale, - c10::optional x_ring, int64_t x_ring_idx, - int64_t x_ring_margin, c10::optional w_ring, - int64_t w_ring_idx, int64_t w_ring_margin) { +std::tuple +linear_forward_fp8(torch::Tensor x, torch::Tensor w, torch::Tensor bias, + torch::Tensor sx, torch::Tensor sw, int64_t fmt, + c10::optional bias_scale, + c10::optional x_ring, int64_t x_ring_idx, + int64_t x_ring_margin, c10::optional w_ring, + int64_t w_ring_idx, int64_t w_ring_margin) { // Pure FP8 forward: quantize x/w (fmt: 0 = E4M3, 1 = E5M2), then the // pre-quantized GEMM; the dequantized BF16 output gets the bias added. - // amax_x / amax_w come from the quantize kernels (zero-initialized here; - // a pre-quantized w reports amax_w = 0 — nothing to feed a delayed ring). - // w may itself be pre-quantized fp8 storage matching fmt (static - // inference weights): the weight quantize is skipped, amax_w stays 0. + // Returns (out, x8, w8, amax_x, amax_w): the quantized operands are + // handed back so the policy layer can cache the weight quantization + // (torch autocast's cached_cast analog — w8 is reused while the weight + // tensor is unchanged, and the backward can share x8/w8 when the fwd/bwd + // formats match). amax_x / amax_w come from the quantize kernels + // (zero-initialized here; a pre-quantized w reports amax_w = 0 — nothing + // to feed a delayed ring). w may itself be pre-quantized fp8 storage + // matching fmt (static inference weights): the weight quantize is + // skipped, amax_w stays 0, and w8 returns the passed-in w. // When x_ring / w_ring are given (delayed scaling), the quantize kernels // finalize them in-kernel: the returned amax is already folded into the // ring window and the next step's scale is published on device. @@ -276,8 +286,16 @@ std::tuple linear_forward_fp8( if (b_prequant) check_scale(*bias_scale, x, "bias_scale"); } auto x8 = torch::empty({m, k}, x_c.options().dtype(f8opt)); - auto amax_x = torch::zeros({1}, x.options().dtype(torch::kFloat32)); - auto amax_w = torch::zeros({1}, x.options().dtype(torch::kFloat32)); + // Each amax slot feeds a cross-block atomic_max, so it must start at 0. + // torch::zeros would route through a fill_ dispatcher (~50us CPU per call + // in the profile); a caching-allocator empty + cudaMemsetAsync is ~2us. + // Zero both up front: the pre-quantized-w path never quantizes w, so its + // amax_w is never atomically written and must not carry stale bytes. The + // returned values are the freshly measured (or 0) amax either way. + auto amax_x = torch::empty({1}, x.options().dtype(torch::kFloat32)); + auto amax_w = torch::empty({1}, x.options().dtype(torch::kFloat32)); + cudaMemsetAsync(amax_x.data_ptr(), 0, sizeof(float), stream.stream()); + cudaMemsetAsync(amax_w.data_ptr(), 0, sizeof(float), stream.stream()); auto out = torch::empty({m, n}, x_c.options()); auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst, @@ -322,7 +340,7 @@ std::tuple linear_forward_fp8( std::vector shape(x.sizes().begin(), x.sizes().end() - 1); shape.push_back(n); - return {out.reshape(shape), amax_x, amax_w}; + return {out.reshape(shape), x8, w8, amax_x, amax_w}; } std::tuple @@ -360,7 +378,13 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w, auto grad_input = torch::empty_like(x); auto grad_weight = torch::empty_like(w); auto grad_bias = torch::empty({0}, g.options()); - auto amax_g = torch::zeros({1}, g.options().dtype(torch::kFloat32)); + // amax_g feeds a cross-block atomic_max in the g quantize kernel; zero it + // on the stream (empty + memset, not torch::zeros — see the forward). + // Only needed when a g quantize runs (mask[0]||mask[1]); the bias-only + // fallback below overwrites it via .copy_, so a wasted memset elsewhere + // is harmless. + auto amax_g = torch::empty({1}, g.options().dtype(torch::kFloat32)); + cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), stream.stream()); auto f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2) : g.options().dtype(torch::kFloat8_e4m3fn); @@ -453,7 +477,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "matching fmt (static inference path; fp8 bias requires bias_scale);" " x_ring/w_ring optionally finalize a delayed-scaling ring " "([hist | scale | counter] float32 buffer) in-kernel; returns " - "(out, amax_x, amax_w)"); + "(out, x8, w8, amax_x, amax_w) — x8 is [M,K], w8 is [N,K] (the " + "passed-in w on the pre-quantized path)"); m.def("linear_backward_fp8", &linear_backward_fp8, py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"), py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("fmt"), diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index 1437bff..dc75799 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -6,15 +6,22 @@ policy-level tests (recipes, autocast context, per-tensor meta, CPU fallbacks of the custom ops) run without a GPU. """ +import threading + import pytest import torch +import torch.nn.functional as F +import astrai.extension.fp8 as f8mod from astrai.extension.fp8 import ( DelayedScaling, DynamicScaling, FP8Format, FP8TensorMeta, + _ScaleRing, fp8_autocast, + fp8_linear_enable, + fp8_linear_enabled, fp8_state, ) from astrai.extension.ops.fp8 import ( @@ -91,8 +98,6 @@ def test_quantize_ring_in_kernel_finalize(): """The quantize kernel finalizes the delayed-scaling ring in-kernel: the measured amax lands in hist[idx], the window reduces to the next step's scale on device, and the counter re-arms for the next launch.""" - from astrai.extension.fp8 import _ScaleRing - torch.manual_seed(21) dev = torch.device("cuda") ring = _ScaleRing(dev, DelayedScaling(history_len=4, margin=0)) @@ -141,7 +146,7 @@ def test_fp8_linear_forward_and_backward(): bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) scale_x, scale_w, scale_g = _scale(x), _scale(weight), _scale(grad) - out, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w) + out, x8, w8, amax_x, amax_w = linear_forward_fp8(x, weight, bias, scale_x, scale_w) grad_x, grad_w, grad_b, amax_g = linear_backward_fp8( grad, x, weight, [1, 1, 1], scale_g, scale_w, scale_x, "e4m3" ) @@ -205,7 +210,7 @@ def test_fp8_linear_static_fp8_weight_and_bias(): w8, _ = quantize_bf16(weight, sw, "e4m3") b8, _ = quantize_bf16(bias, sb, "e4m3") - out, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb) + out, x8, w8_back, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb) qx = _quantize(x, sx) qw = _quantize(weight, sw) @@ -214,9 +219,10 @@ def test_fp8_linear_static_fp8_weight_and_bias(): torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1)) assert amax_w.item() == 0.0 # nothing measured on the static path + assert w8_back is w8 # pre-quantized w handed straight back # bf16 bias stays bf16 on the same fused-epilogue path - out_bf16bias, _, _ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3") + out_bf16bias, *_ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3") expected_b = (qx @ qw.t() * sx * sw + bias.float()).to(torch.bfloat16) torch.testing.assert_close(out_bf16bias, expected_b, atol=0.125, rtol=0.01) @@ -226,10 +232,6 @@ def test_fp8_linear_backward_outside_autocast(): """aten::linear records an fp8 autograd node inside fp8_autocast; the backward runs fp8 kernels even after the context exits (loss.backward() placement is free), instead of falling back to bf16 mm.""" - import torch.nn.functional as F - - import astrai.extension.fp8 as f8mod - torch.manual_seed(5) x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) weight = torch.randn( @@ -424,3 +426,63 @@ def test_mm_fp8_fp8_output_cpu(): assert out8.dtype == torch.float8_e4m3fn ref = (a8.float() @ b8.float() * 2.0 * 0.5 * 0.25).to(torch.float8_e4m3fn) assert torch.equal(out8, ref) + + +# -------------------------------------------------------------------------- +# torch-autocast parity: context semantics (nesting, thread locality, switch) +# -------------------------------------------------------------------------- + + +def _linear(): + """Shared helper: a small bf16 linear operand set on CUDA (grad-tracking + so aten::linear records an autograd node).""" + torch.manual_seed(31) + x = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + w = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + return x, w + + +@skip_no_fp8 +def test_nested_disabled_region_redispatches_bf16(): + """A nested fp8_autocast(enabled=False) region temporarily restores the + bf16 aten::linear path (torch's nested-disable semantics), and fp8 + resumes when it exits.""" + x, w = _linear() + with fp8_autocast(enabled=True): + out_fp8 = F.linear(x, w) + with fp8_autocast(enabled=False): + out_bf16 = F.linear(x, w) + assert type(out_bf16.grad_fn).__name__ != "_LinearFp8Backward" + assert out_bf16.dtype == torch.bfloat16 + out_again = F.linear(x, w) + assert type(out_again.grad_fn).__name__ == "_LinearFp8Backward" + + +@skip_no_fp8 +def test_global_switch_routes_without_region(): + """fp8_linear_enable(True) routes aten::linear to fp8 outside any region + (the persistent default); disabling restores bf16.""" + x, w = _linear() + state = fp8_state() + try: + fp8_linear_enable(True) + out = F.linear(x, w) + assert type(out.grad_fn).__name__ == "_LinearFp8Backward" + fp8_linear_enable(False) + out = F.linear(x, w) + assert type(out.grad_fn).__name__ != "_LinearFp8Backward" + finally: + state.reset() + + +def test_autocast_state_is_thread_local(): + """torch parity: the active config is thread-local — another thread does + not see an open region (CPU-only check of the flag, no kernels).""" + seen = {} + with fp8_autocast(enabled=True): + assert fp8_linear_enabled() + t = threading.Thread(target=lambda: seen.update(enabled=fp8_linear_enabled())) + t.start() + t.join() + assert seen["enabled"] is False + assert not fp8_linear_enabled()