flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/linear_proj/nvfp4_balance.py

@@ -0,0 +1,131 @@
+"""NVFP4 (W4A4) ``linear_proj`` with per-input-channel balance.
+
+The W4 recipe as a formal impl variant: the weight is folded with the
+activation-only channel balance fitted on calibrated per-channel amax
+(``fit_input_channel_balance`` — the production formula), packed to
+NVFP4 (E2M1 data plus per-16-element-block scale factors), and the
+activation is quantized dynamically per call after the inverse fold.
+The fold is exact (``x' @ W'.T == x @ W.T``) and only then is either
+side quantized, so the balance costs no arithmetic identity — it moves
+quantization error out of the hot channels.
+
+Activations carry dynamic per-block scales computed per call, so there
+is no static activation scale to drift across a denoise schedule; the
+per-channel amax calibration feeds the *balance*, not a scale.
+
+Unlike :mod:`.nvfp4_dynamic` (the adoption path for checkpoints that
+are already packed and have no host form to return to), this variant
+retains the host module whole: a call outside the calibrated form runs
+it, and detach restores it bit-exact.
+"""
+
+from __future__ import annotations
+
+from typing import Mapping
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+from .nvfp4_dynamic import KERNEL_DEP, _check, _kernel  # noqa: F401
+
+_VARIANT = 2  # the qualified GEMM dispatch across the served shapes
+
+
+class LinearProjNvfp4Balance(GuardedSeam, torch.nn.Module):
+    """Balanced-fold projection: dynamic FP4 quantize + FP4 GEMM."""
+
+    _frt_host_attr = "host_linear"
+    _frt_can_fallback = True
+
+    def __init__(self, w_packed, w_sfb, inv_s, bias, n, k,
+                 original: torch.nn.Module | None = None):
+        super().__init__()
+        self.register_buffer("_w_packed", w_packed)
+        self.register_buffer("_w_sfb", w_sfb)
+        self.register_buffer("_inv_s", inv_s)
+        if bias is not None:
+            self.register_buffer("_bias", bias)
+        else:
+            self._bias = None
+        self._n = n
+        kern = _kernel()
+        self._kern = kern
+        self._gemm = kern.fp4_w4a16_linear_bf16
+        # capability probe: the fused-bias epilogue entry, where the
+        # installed package variant ships it — absence is the ordinary
+        # two-launch path, never a refusal
+        self._gemm_bias = getattr(kern, "nvfp4_gemm_bias_bf16", None)
+        if original is not None:
+            self.host_linear = original
+        self._frt_arm(dtypes=CAST_OK, device=w_packed.device, k=int(k))
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name == "host_linear":
+                raise
+            return getattr(super().__getattr__("host_linear"), name)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        shape = x.shape
+        flat = (x.reshape(-1, shape[-1]).to(torch.float16)
+                * self._inv_s).contiguous()
+        a_packed, a_sfa = self._kern.quantize_fp4_sfa_fp16(flat)
+        if self._bias is not None and self._gemm_bias is not None:
+            y = self._gemm_bias(a_packed, self._w_packed, a_sfa,
+                                self._w_sfb, self._bias)
+        else:
+            y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb,
+                           variant=_VARIANT)
+            if self._bias is not None:
+                y = y + self._bias
+        return y.reshape(*shape[:-1], self._n).type_as(x)
+
+
+@torch.no_grad()
+def bind_proj_seam(
+    weights: Mapping[str, torch.Tensor],
+    *,
+    channel_amax,
+    original: torch.nn.Module | None = None,
+    alpha: float = 0.5,
+    clamp=(0.25, 4.0),
+) -> LinearProjNvfp4Balance:
+    """Bind one projection from a dense ``[N, K]`` weight.
+
+    ``channel_amax`` is the calibrated per-input-channel amax vector at
+    this projection's input (``[K]``) — it parameterises the balance
+    fold, never a scale. ``alpha``/``clamp`` are the decision's recipe
+    payload.
+    """
+    from flash_rt.core.quantization import fit_input_channel_balance
+
+    n, k = _check(weights)
+    kern = _kernel()
+    w = weights["w"].detach()
+    amax = torch.as_tensor(channel_amax, device=w.device,
+                           dtype=torch.float32)
+    w_bal, inv_s = fit_input_channel_balance(
+        w.float(), amax, alpha=alpha,
+        clamp=(float(clamp[0]), float(clamp[1])),
+        out_dtype=torch.float32)
+    w_packed, w_sfb = kern.quantize_fp4_sfa_fp16(
+        w_bal.to("cuda", torch.float16).contiguous(), is_sfb=True)
+    bias = weights.get("b")
+    if bias is not None:
+        bias = bias.detach().to("cuda", torch.bfloat16)
+    bound = LinearProjNvfp4Balance(
+        w_packed, w_sfb, inv_s.to("cuda", torch.float16), bias, n, k,
+        original=original)
+    probe = bound(torch.zeros(1, k, device=w_packed.device,
+                              dtype=torch.bfloat16))
+    if probe.shape != (1, n) or not torch.isfinite(probe).all():
+        raise ValueError(
+            f"refused: nvfp4_balance bind smoke produced shape "
+            f"{tuple(probe.shape)}, "
+            f"finite={bool(torch.isfinite(probe).all())}")
+    return bound

flashrt_structures/impls/linear_proj/nvfp4_dynamic.py

@@ -0,0 +1,182 @@
+"""NVFP4 (W4A4, dynamic activation scales) ``linear_proj`` implementation.
+
+Weights are packed to NVFP4 (E2M1 data plus per-16-element-block scale
+factors) at bind time; activations are quantized to the same format at
+runtime, per call, with dynamically computed block scales — no
+calibration data at either end. This is the execution form behind the
+27B enablement line: checkpoints whose upstream loader decompresses
+4-bit weights to BF16 inside ``forward`` (and therefore cannot fit the
+card) run on the same card once their projections consume the packed
+layout directly.
+
+The ``flashrt/fp4-gemm`` entry point ``fp4_w4a16_linear_bf16`` takes the
+pre-quantized activation tensor plus its scale factors — despite the
+``a16`` in its historical name, the GEMM it runs is W4A4. ``variant=2``
+is the qualified dispatch across the decode and short-prefill shapes
+this impl serves (same-token 1.0000 against an exact reference on the
+27B host, decode and prefill both through this path).
+
+There is no host fallback: the module this replaces holds packed
+weights the host cannot execute. The guard therefore refuses instead of
+falling back, and adoption of a whole checkpoint is a load-time
+transform, not a reversible attachment.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from functools import lru_cache
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+KERNEL_DEP = {
+    "provider": "huggingface_kernels",
+    "repo": "flashrt/fp4-gemm",
+    "version": ">=1",
+}
+
+#: mirrors the kernel's own shape checks (``torch_binding.cpp``: K
+#: divisible by 16 for the per-block scale factors, positive dims) —
+#: no invented size walls: the adoption path serves whatever the
+#: checkpoint author packed, and the 27B host's 17408-wide FFN is a
+#: qualified shape, not an edge case
+SUPPORT = {
+    "K": {"min": 16, "multiple_of": 16},
+    "N": {"min": 1},
+}
+
+
+@lru_cache(maxsize=1)
+def _kernel():
+    from flashrt_structures.impls import hub_kernel
+
+    return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+
+
+def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]:
+    w = weights["w"]
+    if w.dim() != 2:
+        raise ValueError(f"w must be [N, K], got {tuple(w.shape)}")
+    n, k = w.shape
+    for name, dim in (("K", k), ("N", n)):
+        bounds = SUPPORT[name]
+        if dim < bounds["min"]:
+            raise ValueError(
+                f"{name}={dim} outside support envelope "
+                f"(min {bounds['min']})")
+        if bounds.get("multiple_of") and dim % bounds["multiple_of"]:
+            raise ValueError(
+                f"{name}={dim} must be a multiple of "
+                f"{bounds['multiple_of']}")
+    b = weights.get("b")
+    if b is not None and tuple(b.shape) != (n,):
+        raise ValueError(
+            f"bias shape {tuple(b.shape)} does not match N={n}")
+    return n, k
+
+
+def _quantize_activation(kern, flat: torch.Tensor):
+    """Use the direct BF16 producer when the installed artifact carries it."""
+    if flat.dtype is torch.bfloat16:
+        direct = getattr(kern, "quantize_fp4_sfa_bf16", None)
+        if direct is not None:
+            return direct(flat.contiguous())
+    return kern.quantize_fp4_sfa_fp16(
+        flat.to(torch.float16).contiguous())
+
+
+class LinearProjNvfp4Dynamic(GuardedSeam, torch.nn.Module):
+    """Packed-weight projection: FP4 GEMM with runtime activation scales."""
+
+    _frt_can_fallback = False
+
+    def __init__(self, w_packed, w_sfb, bias, n, k):
+        super().__init__()
+        self.register_buffer("_w_packed", w_packed)
+        self.register_buffer("_w_sfb", w_sfb)
+        self._bias = bias
+        self._n = n
+        self._k = k
+        kern = _kernel()
+        self._kern = kern
+        self._gemm = kern.fp4_w4a16_linear_bf16
+        # M=1 decode rows route to the warp-split GEMV where the build
+        # carries it and the shape qualifies (its own contract: N%8,
+        # K a multiple of 64*warps). Absence is not a refusal - the
+        # tiled GEMM serves every shape correctly, the GEMV just fills
+        # the SMs it underfills at long-K decode shapes.
+        # The entry's presence in the build is not its qualification to
+        # run: the aarch64 package carries it and the kernel refuses at
+        # call time on anything below SM120, which surfaces as a runtime
+        # error on the first M=1 row rather than as a choice made here.
+        # The engine adapters already withhold it off SM120; deciding it
+        # once, where the arm is selected, means an impl bound directly -
+        # through structures.get(), or any hand assembly - behaves the
+        # same as one bound through a door.
+        gemv = getattr(kern, "fp4_w4a4_gemv_warpsplit_bf16", None)
+        cc = (torch.cuda.get_device_capability(w_packed.device)
+              if w_packed.is_cuda else (0, 0))
+        self._gemv = (gemv if gemv is not None and cc >= (12, 0)
+                      and n % 8 == 0 and k % (64 * 4) == 0 else None)
+        self._frt_arm(dtypes=CAST_OK, device=w_packed.device, k=int(k))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        shape = x.shape
+        flat = x.reshape(-1, shape[-1])
+        a_packed, a_sfa = _quantize_activation(self._kern, flat)
+        if flat.shape[0] == 1 and self._gemv is not None:
+            y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb)
+        else:
+            y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb,
+                           variant=2)
+        if self._bias is not None:
+            y = y + self._bias
+        return y.reshape(*shape[:-1], self._n).type_as(x)
+
+
+@torch.no_grad()
+def bind_proj_seam(
+    weights: Mapping[str, torch.Tensor],
+) -> tuple[LinearProjNvfp4Dynamic, float]:
+    """Bind one projection from a dense ``[N, K]`` weight.
+
+    The weight is packed to the Hub kernel's NVFP4 layout on the GPU;
+    the returned float is the pack-and-unpack relative L2 against the
+    input weight — the *conversion* cost of regridding into this layout,
+    reported so a caller adopting a whole checkpoint can put it in the
+    receipt instead of losing it.
+    """
+    n, k = _check(weights)
+    kern = _kernel()
+    w = weights["w"].to("cuda", torch.float16).contiguous()
+    w_packed, w_sfb = kern.quantize_fp4_sfa_fp16(w, is_sfb=True)
+    # the conversion check accumulates in row slabs: a whole-tensor
+    # FP32 dequant doubles the bind's transient footprint, and on
+    # head-class weights that spike is what fails under a tight budget
+    deq = kern.dequantize_fp4_sfa_fp16(w_packed, w_sfb)
+    num_sq = den_sq = 0.0
+    for i in range(0, n, 4096):
+        diff = deq[i:i + 4096].float() - w[i:i + 4096].float()
+        num_sq += float(diff.square().sum())
+        den_sq += float(w[i:i + 4096].float().square().sum())
+    del deq
+    rel = (num_sq ** 0.5) / max(den_sq ** 0.5, 1e-12)
+    bias = weights.get("b")
+    if bias is not None:
+        bias = bias.detach().to("cuda", torch.bfloat16)
+    bound = LinearProjNvfp4Dynamic(w_packed, w_sfb, bias, n, k)
+    # bind-time smoke: one M=1 launch through the real entry point before
+    # the seam is handed out — a stale build or missing symbol surfaces
+    # as a clean bind refusal, not later inside the host's forward
+    probe = bound(torch.zeros(1, k, device=w_packed.device,
+                              dtype=torch.bfloat16))
+    if probe.shape != (1, n) or not torch.isfinite(probe).all():
+        raise ValueError(
+            f"refused: nvfp4 bind smoke produced shape "
+            f"{tuple(probe.shape)}, finite={bool(torch.isfinite(probe).all())}")
+    return bound, rel

flashrt_structures/impls/linear_proj/w8a16_static.py

@@ -0,0 +1,230 @@
+"""Weight-only INT8 implementation of the ``linear_proj`` structure.
+
+The decode-band twin of the FP8 projection impl: weights are quantized
+per output channel to INT8 at bind time, activations stay BF16, so the
+seam needs no calibration data at all. This is the projection-shaped
+slice of the recipe already shipped for ``decoder_ffn`` —
+``w8a16_static`` there covers the gated MLP, this file covers the
+attention Q/K/V/O family and any other single projection the discovery
+qualifies.
+
+The ``flashrt/weight-only-ffn`` package's linear entry point qualifies
+its auto dispatch narrowly, and this impl mirrors that table exactly
+rather than stretching it (``torch_binding.cpp``: ``check_variant`` and
+``w8_auto_linear_supported``):
+
+- M in [1, 4] — the decode band; and
+- K <= 1024 always qualifies; K <= 4096 needs N >= 1024; larger K needs
+  N >= 1024 for M <= 2 and N >= 2048 for M in {3, 4}.
+
+Calls outside the band are dispatched to the retained host module by
+declared plan, counted in the ledger — prefill runs the host GEMM,
+decode runs the kernel, and the qualification record states which band
+the kernel serves.
+
+The linear entry point carries no bias operand. A projection with a
+bias gets it added in BF16 after the GEMM — one [M<=4, N] elementwise
+add inside the decode band, where the weight read dominates end to end.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from functools import lru_cache
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+KERNEL_DEP = {
+    "provider": "huggingface_kernels",
+    "repo": "flashrt/weight-only-ffn",
+    "version": ">=1",
+}
+
+SUPPORT = {
+    "K": {"min": 512, "max": 16384},
+    "N": {"min": 128, "max": 262144},
+    "M": {"min": 1, "max": 4},
+    "m_classes": ("micro",),
+}
+
+
+def _qualified(m: int, n: int, k: int) -> bool:
+    """The kernel's own auto-dispatch qualification, mirrored.
+
+    Copied from ``w8_auto_linear_supported`` plus the ``variant=0``
+    M-bound in the package's ``torch_binding.cpp`` — the kernel raises
+    outside this table, so the band dispatch must agree with it, not
+    rediscover it as runtime errors.
+    """
+    if not SUPPORT["M"]["min"] <= m <= SUPPORT["M"]["max"]:
+        return False
+    if k <= 1024:
+        return True
+    if k <= 4096:
+        return n >= 1024
+    return n >= (1024 if m <= 2 else 2048)
+
+
+@lru_cache(maxsize=1)
+def _kernel():
+    from flashrt_structures.impls import hub_kernel
+
+    return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+
+
+def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]:
+    w = weights["w"]
+    if w.dim() != 2:
+        raise ValueError(f"w must be [N, K], got {tuple(w.shape)}")
+    n, k = w.shape
+    for name, dim in (("K", k), ("N", n)):
+        bounds = SUPPORT[name]
+        if not bounds["min"] <= dim <= bounds["max"]:
+            raise ValueError(
+                f"{name}={dim} outside support envelope "
+                f"[{bounds['min']}, {bounds['max']}]")
+    b = weights.get("b")
+    if b is not None and tuple(b.shape) != (n,):
+        raise ValueError(
+            f"bias shape {tuple(b.shape)} does not match N={n}")
+    return n, k
+
+
+class BoundLinearProjW8A16:
+    """Projection callable: x[M, K] in, y[M, N] out (BF16)."""
+
+    def __init__(self, linear_fn, w_q, w_scale, bias, n, k):
+        self._linear = linear_fn
+        self._w_q = w_q
+        self._w_scale = w_scale
+        self._bias = bias
+        self._n = n
+        self._k = k
+
+    def project(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        flat = x.reshape(-1, shape[-1])
+        m = flat.shape[0]
+        if not _qualified(m, self._n, self._k):
+            raise ValueError(
+                f"M={m} outside the W8A16 auto-dispatch qualification "
+                f"for N={self._n}, K={self._k} (decode band M in "
+                f"[1, {SUPPORT['M']['max']}])")
+        y = self._linear(flat.to(torch.bfloat16).contiguous(),
+                         self._w_q, self._w_scale)
+        if self._bias is not None:
+            y = y + self._bias
+        return y.reshape(*shape[:-1], self._n).to(x.dtype)
+
+    __call__ = project
+
+
+class LinearProjW8A16(GuardedSeam, torch.nn.Module):
+    """Drop-in projection module with declared M-dispatch.
+
+    The weight-only kernel covers the decode band; calls with larger M
+    (prefill) are dispatched to the retained host module. This is part
+    of the declared plan — per-M dispatch on the real workload — not a
+    fallback: both paths are first-class, and the ledger counts the
+    dispatch so neither path's share of the calls is ever unknown.
+
+    ``original`` is retained whole and attribute lookups fall through to
+    it, so host code that introspects ``weight``/``in_features`` keeps
+    working.
+    """
+
+    _frt_host_attr = "host_linear"
+    _frt_can_fallback = True
+
+    def __init__(self, bound: BoundLinearProjW8A16,
+                 original: torch.nn.Module | None = None):
+        super().__init__()
+        self._bound = bound
+        # the same tensors, reachable through *module* attributes: an
+        # exporter attributes a tensor by its access path, and a tensor
+        # reached only through a plain object gets lifted as an
+        # anonymous immutable constant — unnameable in a
+        # weights-external package. Identity is unchanged.
+        self.register_buffer("_frt_w_q", bound._w_q)
+        self.register_buffer("_frt_w_scale", bound._w_scale)
+        if bound._bias is not None:
+            self.register_buffer("_frt_bias", bound._bias)
+        else:
+            self._frt_bias = None
+        if original is not None:
+            self.host_linear = original
+        guard = self._frt_arm(dtypes=CAST_OK,
+                              device=bound._w_q.device,
+                              k=int(bound._k))
+        guard.notes["dispatched_by_band"] = 0
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name == "host_linear":
+                raise
+            return getattr(super().__getattr__("host_linear"), name)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        m = x.numel() // x.shape[-1]
+        if not _qualified(m, self._bound._n, self._bound._k):
+            host = self._frt_host()
+            if host is not None:
+                guard = self._frt_guard
+                if guard is not None and not torch.compiler.is_compiling():
+                    guard.notes["dispatched_by_band"] += 1
+                return host(x)
+            return self._bound.project(x)   # states the refusal
+        shape = x.shape
+        flat = x.reshape(-1, shape[-1])
+        y = self._bound._linear(flat.to(torch.bfloat16).contiguous(),
+                                self._frt_w_q, self._frt_w_scale)
+        if self._frt_bias is not None:
+            y = y + self._frt_bias
+        return y.reshape(*shape[:-1], self._bound._n).to(x.dtype)
+
+
+@torch.no_grad()
+def bind_proj_seam(
+    weights: Mapping[str, torch.Tensor],
+    *,
+    original: torch.nn.Module | None = None,
+) -> LinearProjW8A16:
+    """Bind one projection with weight-only INT8.
+
+    ``weights['w']`` is checkpoint-layout ``[N, K]``, exactly what the
+    kernel consumes — no transpose. No calibration data is required:
+    quantization is per-output-channel on weights only, and the optional
+    ``weights['b']`` is kept in BF16.
+    """
+    n, k = _check(weights)
+    if not _qualified(1, n, k):
+        raise ValueError(
+            f"refused: N={n}, K={k} has no qualified fast path even at "
+            f"M=1; the W8A16 projection cannot serve this seam at any M")
+    kern = _kernel()
+    w = weights["w"].to("cuda", torch.bfloat16).contiguous()
+    w_q, w_scale = kern.quantize_w8_weight_bf16(w)
+    bias = weights.get("b")
+    if bias is not None:
+        bias = bias.detach().to("cuda", torch.bfloat16)
+    bound = BoundLinearProjW8A16(kern.w8a16_linear_bf16, w_q, w_scale,
+                                 bias, n, k)
+    # bind-time smoke: one M=1 launch through the real entry point before
+    # the seam is handed out. A stale build or missing symbol must
+    # surface here as a clean bind refusal, not later inside the host's
+    # forward — identical output cannot catch it there, because the
+    # fallback path is numerically exact.
+    probe = bound.project(torch.zeros(1, k, device=w_q.device,
+                                      dtype=torch.bfloat16))
+    if probe.shape != (1, n) or not torch.isfinite(probe).all():
+        raise ValueError(
+            f"refused: w8a16 bind smoke produced shape "
+            f"{tuple(probe.shape)}, finite={bool(torch.isfinite(probe).all())}")
+    return LinearProjW8A16(bound, original=original)