flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/modnorm_qkv_chain/fp8_ptok_table.py

@@ -0,0 +1,291 @@
+"""Per-token-table implementation of the ``modnorm_qkv_chain`` structure.
+
+The video-DiT block form: modulation parameters live in the block's own
+``[1, chunks, D]`` table combined per token with a ``[B, M, chunks, D]``
+timestep embedding, inline in the block's forward. Only a block owner
+can reroute that math, so this impl binds the whole block:
+
+- both producer sites run the ``adaptive-layernorm-producers`` per-token
+  table entry (table add + chunk selection + no-affine LayerNorm +
+  modulation + static FP8 quantize, one pass — the per-block six-chunk
+  materialization never exists);
+- the self-attention Q/K/V consume the shared FP8 wire through wire
+  projections: the block hands them the quantized activation explicitly
+  before calling the host attention, so rotary/SDPA internals stay the
+  host's. The handoff is an explicit attribute set per call — never an
+  identity-keyed cache (a pointer-keyed bank measurably cross-fed CFG
+  branches on this host);
+- the FFN runs the fused FP8 GELU MLP straight from the second producer
+  site's wire;
+- the output projection and the whole cross-attention are *not* owned:
+  the forward calls whatever is attached there, so their individual
+  seams keep composing.
+
+Bind needs the four static activation scales the composition consumes
+(``attn_in``/``o_in``/``ffn_in``/``ffn_hid``), measured at the block's
+own sublayer inputs by the ordinary calibration pass.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+PRODUCER_DEP = {
+    "provider": "hf",
+    "repo": "flashrt/adaptive-layernorm-producers",
+    "version": ">=1",
+}
+FFN_DEP = {
+    "provider": "hf",
+    "repo": "flashrt/flashrt-fp8-ffn",
+    "version": ">=1",
+}
+
+_FP8 = torch.float8_e4m3fn
+_FP8_MAX = 448.0
+
+#: chunk indices in the block table: (shift, scale, gate) for the
+#: attention site and the FFN site, in table order
+_ATTN_CHUNKS = (0, 1, 2)
+_FFN_CHUNKS = (3, 4, 5)
+
+SUPPORT = {
+    "chunks": 6,
+    "D": {"min": 512, "max": 16384, "multiple_of": 2},
+}
+
+
+@lru_cache(maxsize=1)
+def _producer():
+    from flashrt_structures.impls import hub_kernel
+
+    pkg = hub_kernel(PRODUCER_DEP["repo"], PRODUCER_DEP["version"])
+    if not hasattr(pkg, "ada_layer_norm_quant_fp8_ptok_table_bf16"):
+        raise ValueError(
+            "refused: the installed adaptive-layernorm-producers build "
+            "predates the per-token table entry; a package release with "
+            "ada_layer_norm_quant_fp8_ptok_table_bf16 is required")
+    return pkg
+
+
+@lru_cache(maxsize=1)
+def _ffn_kernel():
+    from flashrt_structures.impls import hub_kernel
+
+    return hub_kernel(FFN_DEP["repo"], FFN_DEP["version"])
+
+
+def _q8(w: torch.Tensor):
+    s = (w.float().abs().amax() / _FP8_MAX).clamp_min(1e-8)
+    packed = (w.float() / s).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8) \
+        .contiguous()
+    return packed, s.reshape(1)
+
+
+class WireProj(torch.nn.Module):
+    """Projection that consumes the chain's FP8 wire.
+
+    The owning block sets ``take(x8, scale)`` immediately before the
+    host attention runs and the projection consumes it exactly once.
+    Called without a wire armed (someone invoking the projection outside
+    the chain), it falls back to the retained host projection — counted,
+    like every dispatch.
+    """
+
+    def __init__(self, lin, gemm):
+        super().__init__()
+        w8, ws = _q8(lin.weight.detach())
+        self.register_buffer("_w8", w8)
+        self.register_buffer("_ws", ws)
+        self._bias = (None if lin.bias is None
+                      else lin.bias.detach().to(torch.bfloat16))
+        self._gemm = gemm
+        self.host_linear = lin
+        self._wire = None
+        self.off_wire_calls = 0
+
+    def take(self, x8, scale):
+        self._wire = (x8, scale)
+
+    def forward(self, x):
+        wire = self._wire
+        if wire is None:
+            if not torch.compiler.is_compiling():
+                self.off_wire_calls += 1
+            return self.host_linear(x)
+        x8, scale = wire
+        y = self._gemm(x8, self._w8, scale, self._ws)
+        if self._bias is not None:
+            y = y + self._bias
+        return y.reshape(*x.shape[:-1], self._w8.shape[0]).type_as(x)
+
+
+class PerTokenModChainBlock(GuardedSeam, torch.nn.Module):
+    """Drop-in replacement for one per-token-table DiT block."""
+
+    _frt_host_attr = "host_block"
+    _frt_can_fallback = True
+
+    def __init__(self, block, wires, scales, ffn_state, producer_fn,
+                 ffn_fn, table, eps):
+        super().__init__()
+        self.host_block = block
+        self._wires = wires                  # (q, k, v) WireProj modules
+        self._scales = scales                # dict of [1] fp32 tensors
+        self._ffn = ffn_state
+        self._producer = producer_fn
+        self._ffn_fn = ffn_fn
+        self.register_buffer("_table", table)
+        self._eps = eps
+        guard = self._frt_arm(dtypes=CAST_OK, device=table.device,
+                              k=int(table.shape[1]))
+        guard.notes["host_form_calls"] = 0
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name == "host_block":
+                raise
+            return getattr(super().__getattr__("host_block"), name)
+
+    def _host_form(self, *args, **kwargs):
+        guard = self._frt_guard
+        if guard is not None and not torch.compiler.is_compiling():
+            guard.notes["host_form_calls"] += 1
+        return self.host_block(*args, **kwargs)
+
+    def forward(self, hidden_states, encoder_hidden_states, temb,
+                rotary_emb, *args, **kwargs):
+        admitted = self._frt_admit(hidden_states)
+        if admitted is not PROCEED:
+            return admitted
+        if temb.dim() != 4 or temb.shape[2] != self._table.shape[0]:
+            # the broadcast (per-sample) form is the host's own path
+            return self._host_form(hidden_states, encoder_hidden_states,
+                                   temb, rotary_emb, *args, **kwargs)
+        block = self.host_block
+        x = hidden_states.contiguous()
+        bsz, seq, dim = x.shape
+        tb = getattr(temb, "_frt_bf16", None)
+        if tb is None:
+            # one cast per transformer call, shared by every block: the
+            # attribute dies with the tensor, so there is no cross-call
+            # identity to poison
+            tb = temb.reshape(-1, temb.shape[2], dim) \
+                .to(torch.bfloat16).contiguous()
+            temb._frt_bf16 = tb
+        s_idx, c_idx, g_idx = _ATTN_CHUNKS
+        gate_msa = (self._table[g_idx]
+                    + temb[0, :, g_idx, :].float()).unsqueeze(0)
+        x2d = x.view(-1, dim)
+        x8 = self._producer(x2d, tb, self._table,
+                            self._scales["attn_in"], s_idx, c_idx,
+                            self._eps)
+        for wire in self._wires:
+            wire.take(x8, self._scales["attn_in"])
+        try:
+            # x is passed for its shape only: the wire projections
+            # consume the quantized activation, not this tensor's values
+            attn = block.attn1(x, None, None, rotary_emb)
+        finally:
+            for wire in self._wires:
+                wire._wire = None
+        x = (x.float() + attn * gate_msa).type_as(x)
+        n2 = block.norm2(x.float()).type_as(x)
+        x = x + block.attn2(n2, encoder_hidden_states, None, None)
+        fs_idx, fc_idx, fg_idx = _FFN_CHUNKS
+        c_gate = (self._table[fg_idx]
+                  + temb[0, :, fg_idx, :].float()).unsqueeze(0)
+        x8f = self._producer(x.contiguous().view(-1, dim), tb,
+                             self._table, self._scales["ffn_in"],
+                             fs_idx, fc_idx, self._eps)
+        st = self._ffn
+        ff = self._ffn_fn(x8f, st["up_w8"], st["up_b"], st["dn_w8"],
+                          st["dn_b"], self._scales["ffn_in"],
+                          st["up_ws"], self._scales["ffn_hid"],
+                          st["dn_ws"])
+        ff = ff.reshape(bsz, seq, dim)
+        return (x.float() + ff.float() * c_gate).type_as(x)
+
+
+@torch.no_grad()
+def bind_block_seam(model, seam, *, points):
+    """Bind one per-token-table block; returns the swap dict.
+
+    The dict carries the block wrapper plus the three wire projections
+    under the host attention, so attach/detach treats the whole
+    composition as one transaction.
+    """
+    from flashrt_structures.discover import _resolve
+
+    block = _resolve(model, seam.path)
+    table_param = block.scale_shift_table.detach()
+    chunks = int(table_param.shape[1])
+    dim = int(table_param.shape[2])
+    if chunks != SUPPORT["chunks"]:
+        raise ValueError(
+            f"refused: {chunks}-chunk table; this impl serves the "
+            f"6-chunk (dual-site) layout")
+    bounds = SUPPORT["D"]
+    if not bounds["min"] <= dim <= bounds["max"] or dim % 2:
+        raise ValueError(f"D={dim} outside support envelope")
+
+    # the collector keys each point by its own placement path (the
+    # block's sublayer input), exactly where points.resolve put it
+    sites = {"attn_in": ".attn1.to_q", "o_in": ".attn1.to_out.0",
+             "ffn_in": ".ffn", "ffn_hid": ".ffn.net.2"}
+    scales = {}
+    for name, child in sites.items():
+        value = (points.scale(seam.path + child, name)
+                 if points is not None else None)
+        if value is None:
+            raise ValueError(
+                f"refused: calibration point {name!r} was not measured "
+                "for this block")
+        scales[name] = torch.tensor([float(value)], device="cuda",
+                                    dtype=torch.float32)
+
+    producer_pkg = _producer()
+    ffn_pkg = _ffn_kernel()
+    gemm = ffn_pkg.fp8_gemm_bf16
+
+    wires = tuple(WireProj(getattr(block.attn1, a), gemm)
+                  for a in ("to_q", "to_k", "to_v"))
+    up, dn = block.ffn.net[0].proj, block.ffn.net[2]
+    ffn_state = {}
+    ffn_state["up_w8"], ffn_state["up_ws"] = _q8(up.weight.detach())
+    ffn_state["dn_w8"], ffn_state["dn_ws"] = _q8(dn.weight.detach())
+    ffn_state["up_b"] = up.bias.detach().to(torch.bfloat16).contiguous()
+    ffn_state["dn_b"] = dn.bias.detach().to(torch.bfloat16).contiguous()
+    for key in ("up_w8", "up_ws", "dn_w8", "dn_ws"):
+        ffn_state[key] = ffn_state[key].to("cuda")
+
+    eps = float(getattr(block.norm1, "eps", 1e-6))
+    table = table_param.reshape(chunks, dim).float().contiguous() \
+        .to("cuda")
+    wrapper = PerTokenModChainBlock(
+        block, wires, scales, ffn_state,
+        producer_pkg.ada_layer_norm_quant_fp8_ptok_table_bf16,
+        (getattr(ffn_pkg, "fp8_gelu_mlp_v2_bf16", None)
+         or ffn_pkg.fp8_gelu_mlp_bf16), table, eps)
+
+    # bind-time smoke: both producer sites launch once on zeros before
+    # the seam is handed out
+    z = torch.zeros(4, dim, device="cuda", dtype=torch.bfloat16)
+    zt = torch.zeros(4, chunks, dim, device="cuda", dtype=torch.bfloat16)
+    for s_idx, c_idx in (_ATTN_CHUNKS[:2], _FFN_CHUNKS[:2]):
+        probe = producer_pkg.ada_layer_norm_quant_fp8_ptok_table_bf16(
+            z, zt, table, scales["attn_in"], s_idx, c_idx, eps)
+        if probe.shape != (4, dim):
+            raise ValueError("refused: producer bind smoke shape "
+                             f"{tuple(probe.shape)}")
+
+    swaps = {seam.path: wrapper}
+    for attr, wire in zip(("to_q", "to_k", "to_v"), wires):
+        swaps[f"{seam.path}.attn1.{attr}"] = wire
+    return swaps

flashrt_structures/impls/moe_experts/__init__.py

@@ -0,0 +1,9 @@
+"""``moe_experts`` structure family.
+
+The structure is the expert bank of a sparse-MoE block: one module
+holding every expert's projection weights as stacked 3D tensors, called
+with the token batch plus the router's top-k assignment. On the hosts
+this family serves, that bank is where nearly all of the checkpoint's
+weight mass lives — which is exactly why it is the seam worth owning
+when the dense checkpoint does not fit the card.
+"""

flashrt_structures/impls/moe_experts/nvfp4_dynamic.py

@@ -0,0 +1,208 @@
+"""NVFP4 (W4A4, dynamic activation scales) ``moe_experts`` implementation.
+
+The expert bank of a sparse-MoE block stores every expert's projections
+as stacked 3D tensors (``gate_up_proj [E, 2I, H]``, ``down_proj
+[E, H, I]``). Each expert's matrices pack once, at bind time, into the
+grouped kernel's stacked NVFP4 layout, and the forward runs the bank as
+two grouped launches per call: one for every routed gate_up slot
+(``[T, top_k]``), one for every down slot (flattened to ``[T*top_k, 1]``
+because each routed pair carries its own intermediate activation). The
+routing tensor stays on the device end to end — no host sync, fixed
+shapes for a given ``T`` — so the step is legal inside a compiled
+region or a captured graph, and the *same* kernels serve the M=1 decode
+row and the M=K+1 verify pass: one numeric family across both, which
+is what token-identity between a spec verify and the plain step needs.
+
+Contributions accumulate in FP32 over the fixed top-k axis before the
+single cast back to the host dtype.
+
+Known ceiling, recorded not hidden: weight traffic is per routed slot.
+A long prefill (hundreds of tokens and up) re-reads shared expert
+weights once per slot where a per-expert grouping would read them once;
+until a grouped entry with per-expert accumulation ships, long prompts
+through this bank pay slot-linear traffic.
+
+There is no host fallback: binding exists to retire the dense weights
+whose footprint keeps the checkpoint off the card, so the guard refuses
+out-of-form calls instead of falling back.
+"""
+
+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/grouped-moe-gemv",
+    "version": ">=2",
+}
+
+#: mirrors the kernel's own checks (K divisible by 16, N by 8) — both
+#: contraction dims of an expert bank are K once: H for gate_up, I for
+#: down; both output dims are N once: 2I and H
+SUPPORT = {
+    "K": {"min": 16, "multiple_of": 16},
+    "N": {"min": 8, "multiple_of": 8},
+    "E": {"min": 1},
+}
+
+#: experts are streamed to the GPU in slabs of this many during bind so
+#: the transient footprint stays at slab size, not the full bank
+_BIND_SLAB = 32
+
+
+@lru_cache(maxsize=1)
+def _kernel():
+    from flashrt_structures.impls import hub_kernel
+
+    return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+
+
+def _sf_bytes(rows: int, dim: int) -> int:
+    """The kernel's swizzled scale-factor buffer size for one [rows, dim]."""
+    return ((rows + 127) // 128) * (((dim // 16) + 3) // 4) * 512
+
+
+def check_experts(weights: Mapping[str, torch.Tensor]) -> tuple[int, int, int]:
+    """Validate an expert bank's shapes; returns ``(E, H, I)``."""
+    gu, dn = weights["gate_up_proj"], weights["down_proj"]
+    if gu.dim() != 3 or dn.dim() != 3:
+        raise ValueError(
+            f"expert bank must be 3D stacks, got gate_up "
+            f"{tuple(gu.shape)}, down {tuple(dn.shape)}")
+    e, two_i, h = gu.shape
+    e2, h2, i = dn.shape
+    if e != e2 or h != h2 or two_i != 2 * i:
+        raise ValueError(
+            f"inconsistent expert bank: gate_up {tuple(gu.shape)} vs "
+            f"down {tuple(dn.shape)}")
+    if e < SUPPORT["E"]["min"]:
+        raise ValueError(f"E={e} outside support envelope")
+    for name, dim in (("H", h), ("I", i)):
+        if dim < SUPPORT["K"]["min"] or dim % SUPPORT["K"]["multiple_of"]:
+            raise ValueError(
+                f"{name}={dim} must be a positive multiple of "
+                f"{SUPPORT['K']['multiple_of']} (it is a contraction dim)")
+        if dim % SUPPORT["N"]["multiple_of"]:
+            raise ValueError(
+                f"{name}={dim} must be a multiple of "
+                f"{SUPPORT['N']['multiple_of']} (it is an output dim)")
+    return e, h, i
+
+
+class MoeExpertsNvfp4Dynamic(GuardedSeam, torch.nn.Module):
+    """Packed expert bank: two grouped FP4 launches behind the host
+    contract, for any short token batch."""
+
+    _frt_can_fallback = False
+
+    def __init__(self, gu_packed, gu_sfb, dn_packed, dn_sfb, act_fn,
+                 num_experts, hidden, inter):
+        super().__init__()
+        self.register_buffer("_gu_packed", gu_packed)
+        self.register_buffer("_gu_sfb", gu_sfb)
+        self.register_buffer("_dn_packed", dn_packed)
+        self.register_buffer("_dn_sfb", dn_sfb)
+        self.register_buffer("_alpha", torch.ones(
+            num_experts, device=gu_packed.device, dtype=torch.float32))
+        self._act = act_fn
+        self._e = num_experts
+        self._h = hidden
+        self._i = inter
+        self._grouped = _kernel().grouped_w4a4_gemv_from_bf16
+        self._frt_arm(dtypes=CAST_OK, device=gu_packed.device, k=hidden)
+
+    def forward(self, hidden_states: torch.Tensor,
+                top_k_index: torch.Tensor,
+                top_k_weights: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(hidden_states)
+        if admitted is not PROCEED:
+            return admitted
+        t = hidden_states.shape[0]
+        k = top_k_index.shape[1]
+        idx = top_k_index.to(torch.int32)
+        y = self._grouped(hidden_states.contiguous(), self._gu_packed,
+                          self._gu_sfb, self._alpha, idx)
+        gate, up = y.chunk(2, dim=-1)
+        inter = self._act(gate) * up                    # [T, k, I], fresh
+        d = self._grouped(inter.reshape(t * k, self._i), self._dn_packed,
+                          self._dn_sfb, self._alpha,
+                          idx.reshape(t * k, 1))
+        out = (d.view(t, k, self._h).float()
+               * top_k_weights[..., None].float()).sum(dim=1)
+        return out.to(hidden_states.dtype)
+
+
+@torch.no_grad()
+def _pack_bank(kern, bank: torch.Tensor, alpha: torch.Tensor,
+               probe_gen: torch.Generator) -> tuple[
+                   torch.Tensor, torch.Tensor, float]:
+    """Pack one 3D stack ``[E, N, K]``; returns worst probe-row relL2.
+
+    The kernel ships no dequantizer, so conversion is accounted at the
+    output: one random row through the grouped entry against the BF16
+    slab's own matmul, per slab, worst kept for the receipt.
+    """
+    e, n, kdim = bank.shape
+    packed = torch.empty(e, n, kdim // 2, device="cuda", dtype=torch.uint8)
+    sfb = torch.empty(e, _sf_bytes(n, kdim), device="cuda",
+                      dtype=torch.uint8)
+    worst = 0.0
+    grouped = kern.grouped_w4a4_gemv_from_bf16
+    for lo in range(0, e, _BIND_SLAB):
+        slab = bank[lo:lo + _BIND_SLAB].to("cuda", torch.bfloat16)
+        for j in range(slab.shape[0]):
+            kern.quantize_weights_nvfp4_bf16(
+                slab[j].contiguous(), packed=packed[lo + j],
+                sfb=sfb[lo + j])
+        x = (torch.randn(1, kdim, device="cuda", generator=probe_gen,
+                         dtype=torch.float32) * 0.05).to(torch.bfloat16)
+        got = grouped(x, packed, sfb, alpha,
+                      torch.tensor([[lo]], device="cuda",
+                                   dtype=torch.int32))[0, 0]
+        ref = x[0].float() @ slab[0].float().t()
+        rel = float((got.float() - ref).norm() / ref.norm().clamp_min(1e-12))
+        worst = max(worst, rel)
+        del slab
+    return packed, sfb, worst
+
+
+@torch.no_grad()
+def bind_experts_seam(
+    weights: Mapping[str, torch.Tensor], act_fn,
+) -> tuple[MoeExpertsNvfp4Dynamic, dict[str, float]]:
+    """Bind one expert bank from its dense 3D stacks.
+
+    Weights stream to the GPU in expert slabs and pack there; the
+    returned dict carries the worst probe-row relative L2 per stack,
+    for the adoption receipt. The bound module holds only the packed
+    layout — retiring the dense bank is the caller's move (and the
+    point).
+    """
+    e, h, i = check_experts(weights)
+    kern = _kernel()
+    alpha = torch.ones(e, device="cuda", dtype=torch.float32)
+    gen = torch.Generator(device="cuda")
+    gen.manual_seed(0)
+    gu_packed, gu_sfb, gu_rel = _pack_bank(
+        kern, weights["gate_up_proj"], alpha, gen)
+    dn_packed, dn_sfb, dn_rel = _pack_bank(
+        kern, weights["down_proj"], alpha, gen)
+    bound = MoeExpertsNvfp4Dynamic(gu_packed, gu_sfb, dn_packed, dn_sfb,
+                                   act_fn, e, h, i)
+    # bind-time smoke: one decode-shaped call through the real entries
+    probe = bound(
+        torch.zeros(1, h, device=gu_packed.device, dtype=torch.bfloat16),
+        torch.zeros(1, 1, device=gu_packed.device, dtype=torch.long),
+        torch.ones(1, 1, device=gu_packed.device, dtype=torch.bfloat16))
+    if probe.shape != (1, h) or not torch.isfinite(probe).all():
+        raise ValueError(
+            f"refused: moe_experts nvfp4 bind smoke produced shape "
+            f"{tuple(probe.shape)}, "
+            f"finite={bool(torch.isfinite(probe).all())}")
+    return bound, {"gate_up_proj": gu_rel, "down_proj": dn_rel}

flashrt_structures/impls/moe_experts/nvfp4_w4a16.py

@@ -0,0 +1,129 @@
+"""W4A16 implementation of the ``moe_experts`` structure.
+
+The SM110 twin of :mod:`.nvfp4_dynamic`: same packed NVFP4 expert bank,
+same external routing contract, but the grouped launch keeps activations
+in BF16 — the ``grouped_w4a4_*`` entries require SM120/SM121
+block-scaled MMA and refuse on Thor, while ``grouped_w4a16_gemv_bf16``
+serves one routed slot per activation row on every arch the package
+ships. The call convention therefore differs: rows are expanded to one
+per routed slot ([T*k, K]) instead of the W4A4 entry's [T, k] batch.
+
+Packing is byte-identical to the W4A4 impl (one ``quantize_weights_
+nvfp4_bf16`` layout serves both entries); only the bind-time probe and
+the forward launches change.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+from .nvfp4_dynamic import _BIND_SLAB, _kernel, _sf_bytes, check_experts
+
+__all__ = ["MoeExpertsW4A16", "bind_experts_seam"]
+
+
+class MoeExpertsW4A16(GuardedSeam, torch.nn.Module):
+    """Packed expert bank behind the host contract, BF16 activations."""
+
+    _frt_can_fallback = False
+
+    def __init__(self, gu_packed, gu_sfb, dn_packed, dn_sfb, act_fn,
+                 num_experts, hidden, inter):
+        super().__init__()
+        self.register_buffer("_gu_packed", gu_packed)
+        self.register_buffer("_gu_sfb", gu_sfb)
+        self.register_buffer("_dn_packed", dn_packed)
+        self.register_buffer("_dn_sfb", dn_sfb)
+        self.register_buffer("_alpha", torch.ones(
+            num_experts, device=gu_packed.device, dtype=torch.float32))
+        self._act = act_fn
+        self._e = num_experts
+        self._h = hidden
+        self._i = inter
+        self._grouped = _kernel().grouped_w4a16_gemv_bf16
+        self._frt_arm(dtypes=CAST_OK, device=gu_packed.device, k=hidden)
+
+    def _launch(self, x, packed, sfb, ids, n):
+        return self._grouped(
+            x, packed, sfb, self._alpha, ids, n=n,
+            w_stride=packed.shape[1] * packed.shape[2],
+            sfb_stride=sfb.shape[1])
+
+    def forward(self, hidden_states: torch.Tensor,
+                top_k_index: torch.Tensor,
+                top_k_weights: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(hidden_states)
+        if admitted is not PROCEED:
+            return admitted
+        t = hidden_states.shape[0]
+        k = top_k_index.shape[1]
+        ids = top_k_index.reshape(-1).to(torch.int32)
+        x = hidden_states.to(torch.bfloat16).contiguous()
+        xr = x.repeat_interleave(k, dim=0)              # [T*k, H]
+        y = self._launch(xr, self._gu_packed, self._gu_sfb, ids,
+                         2 * self._i)                    # [T*k, 2I]
+        gate, up = y.chunk(2, dim=-1)
+        inter = (self._act(gate) * up).contiguous()      # [T*k, I]
+        d = self._launch(inter, self._dn_packed, self._dn_sfb, ids,
+                         self._h)                        # [T*k, H]
+        out = (d.view(t, k, self._h).float()
+               * top_k_weights[..., None].float()).sum(dim=1)
+        return out.to(hidden_states.dtype)
+
+
+@torch.no_grad()
+def _pack_bank(kern, bank: torch.Tensor, alpha: torch.Tensor,
+               probe_gen: torch.Generator) -> tuple[
+                   torch.Tensor, torch.Tensor, float]:
+    """Pack one 3D stack ``[E, N, K]``; probe through the W4A16 entry."""
+    e, n, kdim = bank.shape
+    packed = torch.empty(e, n, kdim // 2, device="cuda", dtype=torch.uint8)
+    sfb = torch.empty(e, _sf_bytes(n, kdim), device="cuda",
+                      dtype=torch.uint8)
+    worst = 0.0
+    for lo in range(0, e, _BIND_SLAB):
+        slab = bank[lo:lo + _BIND_SLAB].to("cuda", torch.bfloat16)
+        for j in range(slab.shape[0]):
+            kern.quantize_weights_nvfp4_bf16(
+                slab[j].contiguous(), packed=packed[lo + j],
+                sfb=sfb[lo + j])
+        x = (torch.randn(1, kdim, device="cuda", generator=probe_gen,
+                         dtype=torch.float32) * 0.05).to(torch.bfloat16)
+        got = kern.grouped_w4a16_gemv_bf16(
+            x, packed, sfb, alpha,
+            torch.tensor([lo], device="cuda", dtype=torch.int32),
+            n=n, w_stride=n * kdim // 2, sfb_stride=sfb.shape[1])[0]
+        ref = x[0].float() @ slab[0].float().t()
+        rel = float((got.float() - ref).norm() / ref.norm().clamp_min(1e-12))
+        worst = max(worst, rel)
+        del slab
+    return packed, sfb, worst
+
+
+@torch.no_grad()
+def bind_experts_seam(
+    weights: Mapping[str, torch.Tensor], act_fn,
+) -> tuple[MoeExpertsW4A16, dict[str, float]]:
+    """Bind one expert bank from its dense 3D stacks (W4A16 launches)."""
+    e, h, i = check_experts(weights)
+    kern = _kernel()
+    alpha = torch.ones(e, device="cuda", dtype=torch.float32)
+    gen = torch.Generator(device="cuda")
+    gen.manual_seed(0)
+    gu_packed, gu_sfb, gu_rel = _pack_bank(
+        kern, weights["gate_up_proj"], alpha, gen)
+    dn_packed, dn_sfb, dn_rel = _pack_bank(
+        kern, weights["down_proj"], alpha, gen)
+    bound = MoeExpertsW4A16(gu_packed, gu_sfb, dn_packed, dn_sfb,
+                            act_fn, e, h, i)
+    probe = bound(torch.zeros(1, h, device="cuda", dtype=torch.bfloat16),
+                  torch.zeros(1, 1, device="cuda", dtype=torch.int64),
+                  torch.ones(1, 1, device="cuda", dtype=torch.float32))
+    if probe.shape != (1, h) or not torch.isfinite(probe).all():
+        raise ValueError(
+            f"refused: w4a16 experts bind smoke produced shape "
+            f"{tuple(probe.shape)}, finite={bool(torch.isfinite(probe).all())}")
+    return bound, {"gate_up_relL2": gu_rel, "down_relL2": dn_rel}

flashrt_structures/impls/norm_fused/__init__.py

@@ -0,0 +1,3 @@
+from .bf16 import FusedNorm, bind_norm_fused
+
+__all__ = ["FusedNorm", "bind_norm_fused"]