diff --git a/flash_rt/structures/__init__.py b/flash_rt/structures/__init__.py new file mode 100644 index 00000000..3c0bc7da --- /dev/null +++ b/flash_rt/structures/__init__.py @@ -0,0 +1,70 @@ +"""FlashRT structures — verified, composable model sub-blocks. + +A structure is a versioned specification of one model region: boundary +tensors, framework-neutral weight slots, a plain reference implementation +used as ground truth, and qualification gates. This package hosts the +structure catalog and its registry. Implementations, host adapters, and +the qualification harness build on top of these specifications. +""" + +from flash_rt.structures.registry import StructureSpec, list_structures, load + + +def get(name): + """Explicit door: pull one structure and bind it yourself. + + Mirrors ``kernels.get_kernel``: ``get("decoder_ffn").bind(module, + calibration=[...])`` returns a gated drop-in replacement you plug in + where you choose. See :mod:`flash_rt.structures.handle`. + """ + from flash_rt.structures.handle import get as _get + + return _get(name) + + +def capture(fn, **kwargs): + """Capture door: graph a hot stage with declared swap windows. + + See :func:`flash_rt.structures.stages.capture`. + """ + from flash_rt.structures.stages import capture as _capture + + return _capture(fn, **kwargs) + + +def auto_swaps(model, forward, **kwargs): + """Distribution layer: discover, calibrate, bind — one pass, no + per-seam scaffolding. Returns an :class:`AutoPlan` of swaps. + + See :func:`flash_rt.structures.autobuild.auto_swaps`. + """ + from flash_rt.structures.autobuild import auto_swaps as _auto + + return _auto(model, forward, **kwargs) + + +def run_recipe(recipe, model, ctx=None, **kwargs): + """Recipe door: assemble declared levers, audit same-process on the + graph, certify or refuse — one call, one receipt. + + See :mod:`flash_rt.structures.recipe` for ``Recipe``/``Lever``/ + ``Gates`` and the switch lifecycle. + """ + from flash_rt.structures.recipe import run_recipe as _run + + return _run(recipe, model, ctx, **kwargs) + + +def attach(model, forward, **kwargs): + """One-call front door: discover, calibrate, gate, activate. + + See :func:`flash_rt.structures.frontdoor.attach`. Imported lazily so + that spec-only consumers do not pay for torch-side machinery. + """ + from flash_rt.structures.frontdoor import attach as _attach + + return _attach(model, forward, **kwargs) + + +__all__ = ["StructureSpec", "attach", "capture", "get", "list_structures", + "load", "run_recipe"] diff --git a/flash_rt/structures/adapters/__init__.py b/flash_rt/structures/adapters/__init__.py new file mode 100644 index 00000000..f5941854 --- /dev/null +++ b/flash_rt/structures/adapters/__init__.py @@ -0,0 +1,13 @@ +"""Host-family adapters — where a structure's seam is host-specific. + +Importing this package registers the built-in adapters with autobuild. +Attention seams (attention_core) live here because where the attention +math runs differs by host family; a static module pattern cannot find +them, so each family gets a small adapter. +""" +from ..autobuild import register_attention_adapter +from .gemma_attention import GemmaAttentionAdapter + +register_attention_adapter(GemmaAttentionAdapter()) + +__all__ = ["GemmaAttentionAdapter"] diff --git a/flash_rt/structures/adapters/gemma_attention.py b/flash_rt/structures/adapters/gemma_attention.py new file mode 100644 index 00000000..29e0325a --- /dev/null +++ b/flash_rt/structures/adapters/gemma_attention.py @@ -0,0 +1,113 @@ +"""Attention adapter for Gemma-family denoise hosts (pi05 / pi_gemma). + +Where the attention math runs is host-specific. In this family the +transformer's own forward calls ``modeling_gemma.eager_attention_forward`` +directly, bypassing the config/interface dispatch entirely, so the seam +is that function, not a module. This adapter locates it by capturing one +denoise pass, binds an :mod:`..impls.attention_core` per layer from the +captured shapes and masks, and installs a function-level patch that +routes the fixed denoise shape to the packed-KV kernel while leaving +prefill and any other shape on the host path. + +Registering this adapter lets ``autobuild`` pick up the attention_core +structure for this host family with no per-host scaffolding at the call +site — the host still just calls ``auto_swaps``. +""" + +from __future__ import annotations + +from ..impls.attention_core import bind_attention_core + + +class GemmaAttentionAdapter: + """Recognise a Gemma-family denoise host and wire its fa2 seam.""" + + __name__ = "gemma_attention" + + def __call__(self, model, forward): + try: + import transformers.models.gemma.modeling_gemma as mg + except ImportError: + return None + orig = mg.eager_attention_forward + + recs = {"q": None, "masks": [], "keys": [], "values": []} + + def record(module, query, key, value, attention_mask, **kw): + if query.shape[2] < 128: # denoise (short) vs prefill + recs["q"] = query.detach() + recs["masks"].append( + attention_mask.detach() + if attention_mask is not None else None) + recs["keys"].append(key.detach().clone()) + recs["values"].append(value.detach().clone()) + return orig(module, query, key, value, attention_mask, **kw) + + mg.eager_attention_forward = record + try: + with __import__("torch").no_grad(): + forward() + finally: + mg.eager_attention_forward = orig + if recs["q"] is None: + return None # host never called this seam — not our family + + n_layers = _infer_layers(model) + if n_layers == 0 or len(recs["keys"]) % n_layers != 0: + return None + steps = len(recs["keys"]) // n_layers + captures = [{ + "q": recs["q"], + "keys": [recs["keys"][i + s * n_layers] for s in range(steps)], + "values": [recs["values"][i + s * n_layers] + for s in range(steps)], + "mask": recs["masks"][i], + } for i in range(n_layers)] + + bound = bind_attention_core(captures) + if bound is None: + return None # head_dim unsupported → host keeps its path + cores, prefix_update = bound + seq_q = recs["q"].shape[2] + expert = _expert_layers(model) + for i, layer in enumerate(expert): + layer.self_attn._fa2_core = cores[i] + + fired = {"n": 0} + + def fa2_fn(module, query, key, value, attention_mask, **kw): + if query.shape[2] != seq_q or not hasattr(module, "_fa2_core"): + return orig(module, query, key, value, attention_mask, **kw) + fired["n"] += 1 + return module._fa2_core(query, key, value, + scale=kw.get("scaling")), None + + mg.eager_attention_forward = fa2_fn + # self-verify the patch actually takes over the seam on the same + # forward that will be captured; if it never fires (the host's + # captured entry uses a different attention shape than the one + # we calibrated on) leave the host untouched rather than install + # a patch that only adds a Python check and blocks fusion + with __import__("torch").no_grad(): + forward() + if fired["n"] == 0: + mg.eager_attention_forward = orig + for layer in expert: + if hasattr(layer.self_attn, "_fa2_core"): + del layer.self_attn._fa2_core + return None + # the swap map is empty (the seam is a function, not a module); + # the patch and the per-layer core buffers are the swap + return {}, None + + +def _infer_layers(model) -> int: + layers = _expert_layers(model) + return len(layers) if layers is not None else 0 + + +def _expert_layers(model): + try: + return model.paligemma_with_expert.gemma_expert.model.layers + except AttributeError: + return None diff --git a/flash_rt/structures/autobuild.py b/flash_rt/structures/autobuild.py new file mode 100644 index 00000000..88bcd590 --- /dev/null +++ b/flash_rt/structures/autobuild.py @@ -0,0 +1,223 @@ +"""Auto-assembly: discover seams, calibrate them in one pass, bind them. + +This is the distribution layer. Given a host model and a way to run it, +it finds every structure seam (:mod:`.discover`), captures exactly the +calibration each one needs in a single forward pass, and binds each +through its library impl. The caller gets a ``path -> module`` swap map +and any outside-cadence update functions — the same thing the hand +recipes produced, derived from the model object rather than written by +hand. A host integrates by importing and calling; it writes no +per-seam scaffolding. + +The calibration each structure needs, captured structure-aware: + linear_proj / qkv_pack : the shared input the projection(s) see, and + its per-tensor amax (the static act scale) + adaln_producer : the (cond, style) pairs the conditioning + projection emits across the tick, for the + step table and its fingerprint locator + decoder_ffn / vision_ffn : the normed input the MLP sees + +Seam negotiation is resolved here: when an adaln_producer feeds a +sibling qkv_pack under the same parent, the producer emits fp8 and the +pack takes the shared act scale, skipping its own input quantization. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable + +import torch + +from .discover import (Seam, _resolve, discover, group_families, + seam_weights) + +_FP8 = torch.float8_e4m3fn + +# Host-family attention adapters. Attention seams are not a static +# module pattern — where the attention math actually runs is +# host-specific (a function in one host, a processor in another), so +# auto-discovery of the attention_core structure is delegated to +# registered adapters. Each adapter, given the model and a way to run +# it, returns (swaps, update) or None (this host is not its family). +_ATTENTION_ADAPTERS: list = [] + + +def register_attention_adapter(adapter) -> None: + """Register a host-family attention adapter (callable).""" + _ATTENTION_ADAPTERS.append(adapter) + + +@dataclass +class AutoPlan: + """Discovered + calibrated swaps, ready to stage.""" + + swaps: dict[str, torch.nn.Module] = field(default_factory=dict) + updates: list[Callable[[], None]] = field(default_factory=list) + seams: list[Seam] = field(default_factory=list) + notes: dict[str, Any] = field(default_factory=dict) + + +def _amax(tensors) -> float: + return max(t.float().abs().max().item() for t in tensors) + + +def _rows(t: torch.Tensor) -> int: + return int(t.reshape(-1, t.shape[-1]).shape[0]) + + +def auto_swaps( + model: torch.nn.Module, + forward: Callable[[], Any], + *, + structures: tuple[str, ...] = ("decoder_ffn", "vision_ffn", + "qkv_pack", "adaln_producer"), + negotiate_fp8: bool = True, + frames: int = 1, + verbose: bool = False, +) -> AutoPlan: + """Discover, calibrate in one pass, and bind every applicable seam.""" + + def say(msg: str) -> None: + if verbose: + print(f"[autobuild] {msg}", flush=True) + + seams = discover(model, structures) + say(f"discovered {len(seams)} seam(s)") + if not seams: + return AutoPlan() + + # ---- one calibration pass, structure-aware capture ---- + caps: dict[str, dict[str, Any]] = {} + hooks = [] + + def cap_input(path): + def hook(module, args, kwargs, out): + x = args[0] if args else next(iter(kwargs.values())) + caps[path].setdefault("x", []).append(x.detach()) + caps[path]["rows"] = _rows(x) + return None + return hook + + def cap_cond(path): + def hook(module, args, out): + caps[path].setdefault("pairs", []).append( + (args[0].detach().clone(), out.detach().clone())) + return None + return hook + + for seam in seams: + caps[seam.path] = {} + target = _resolve(model, seam.path) + if seam.structure == "adaln_producer": + # the norm's own input gives rows; the cond projection gives + # the (cond, style) pairs the table is built from + hooks.append(target.register_forward_hook( + cap_input(seam.path), with_kwargs=True)) + hooks.append(getattr(target, seam.cond_attr) + .register_forward_hook(cap_cond(seam.path))) + elif seam.structure == "qkv_pack": + first = getattr(target, seam.pack_attrs[0]) + hooks.append(first.register_forward_hook( + cap_input(seam.path), with_kwargs=True)) + else: + hooks.append(target.register_forward_hook( + cap_input(seam.path), with_kwargs=True)) + + with torch.no_grad(): + for _ in range(max(1, frames)): + forward() + for h in hooks: + h.remove() + say("calibration pass done") + + # ---- bind each seam via its impl (act_scales carries the shared + # producer/consumer scale for the fp8 seam negotiation) ---- + act_scales: dict[str, torch.Tensor] = {} + plan = AutoPlan(seams=seams) + for name, members in group_families(seams).items(): + for seam in members: + cap = caps.get(seam.path, {}) + try: + bound = _bind_auto(model, seam, cap, plan, act_scales, + negotiate_fp8) + except (ValueError, RuntimeError) as refusal: + plan.notes.setdefault("refused", []).append( + (seam.path, str(refusal)[:80])) + continue + if bound is None: + continue + if isinstance(bound, dict): + plan.swaps.update(bound) + else: + plan.swaps[seam.path] = bound + # ---- attention_core: host-family adapters (fa2 seam) ---- + if "attention_core" in structures: + from . import adapters as _adapters # noqa: F401 (registers) + for adapter in _ATTENTION_ADAPTERS: + try: + result = adapter(model, forward) + except (ValueError, RuntimeError) as refusal: + plan.notes.setdefault("refused", []).append( + ("attention_core", str(refusal)[:80])) + continue + if result is None: + continue + att_swaps, update = result + plan.swaps.update(att_swaps) + if update is not None: + plan.updates.append(update) + plan.notes["attention_adapter"] = type(adapter).__name__ \ + if hasattr(adapter, "__name__") else str(adapter) + break + + say(f"bound {len(plan.swaps)} seam(s), " + f"{len(plan.notes.get('refused', []))} refused") + return plan + + +def _bind_auto(model, seam, cap, plan, act_scales, negotiate_fp8): + """Route one seam to its impl with the captured calibration.""" + from .handle import get + + if seam.structure in ("decoder_ffn", "vision_ffn"): + if not cap.get("x"): + return None + h = get(seam.structure) + return h.bind(_resolve(model, seam.path), cap["x"], gate_cos=0.0) + + if seam.structure == "linear_proj": + if not cap.get("x"): + return None + from .impls.linear_proj import fp8_static as proj_impl + return proj_impl.bind_proj_seam( + seam_weights(model, seam), calibration=cap["x"], + original=_resolve(model, seam.path)) + + if seam.structure == "qkv_pack": + from .impls.qkv_pack import bind_qkv_pack + if not cap.get("x"): + return None + parent = _resolve(model, seam.path) + mods = [getattr(parent, a) for a in seam.pack_attrs] + act_scale = torch.tensor( + [max(_amax(cap["x"]) / 448.0, 1e-8)], + device=mods[0].weight.device) + act_scales[seam.path] = act_scale + parts = bind_qkv_pack(mods, act_scale, rows=cap["rows"], + in_dtype="bf16_fused_quant") + return {seam.path + "." + a: m + for a, m in zip(seam.pack_attrs, parts)} + + if seam.structure == "adaln_producer": + from .impls.adaln_producer import bind_style_table + if not cap.get("pairs"): + return None + norm = _resolve(model, seam.path) + proj = getattr(norm, seam.cond_attr) + loc = plan.notes.setdefault("_locators", {}).get(seam.family) + table = bind_style_table(proj, cap["pairs"], locator=loc) + plan.notes["_locators"][seam.family] = table.locator + return {seam.path + "." + seam.cond_attr: table} + + return None diff --git a/flash_rt/structures/bindings/groot_n16_dit.yaml b/flash_rt/structures/bindings/groot_n16_dit.yaml new file mode 100644 index 00000000..f6536212 --- /dev/null +++ b/flash_rt/structures/bindings/groot_n16_dit.yaml @@ -0,0 +1,23 @@ +binding: groot_n16_dit +structure: vision_ffn +dims: {D: 1536, F: 6144, layers: 32} +variant: {activation: gelu} +boundary_dtype: bf16 + +m_profile: + denoise: {m_class: small, phase: denoise} # M = action tokens per step + +weights_map: + isaac_gr00t: + layout: out_in + w_fc1: "action_head.model.transformer_blocks.{i}.ff.net.0.proj.weight" + b_fc1: "action_head.model.transformer_blocks.{i}.ff.net.0.proj.bias" + w_fc2: "action_head.model.transformer_blocks.{i}.ff.net.2.weight" + b_fc2: "action_head.model.transformer_blocks.{i}.ff.net.2.bias" + w_norm: {const: ones} # DiT norm3 seam; affine may be absent + b_norm: {const: zeros} + +hosts: + isaac_gr00t: + module_path: "action_head.model.transformer_blocks.{i}" + versions: "n1.6" diff --git a/flash_rt/structures/bindings/groot_n16_llm.yaml b/flash_rt/structures/bindings/groot_n16_llm.yaml new file mode 100644 index 00000000..5e74980e --- /dev/null +++ b/flash_rt/structures/bindings/groot_n16_llm.yaml @@ -0,0 +1,21 @@ +binding: groot_n16_llm +structure: decoder_ffn +dims: {D: 2048, F: 6144, layers: 16} +variant: {activation: silu, norm_weight_mode: direct, conditioning: none} +boundary_dtype: bf16 + +m_profile: + prefix: {m_class: medium, phase: prefill} # M = vision+language prefix + +weights_map: + isaac_gr00t: + layout: out_in + w_gate: "backbone.model.language_model.model.layers.{i}.mlp.gate_proj.weight" + w_up: "backbone.model.language_model.model.layers.{i}.mlp.up_proj.weight" + w_down: "backbone.model.language_model.model.layers.{i}.mlp.down_proj.weight" + w_norm: "backbone.model.language_model.model.layers.{i}.post_attention_layernorm.weight" + +hosts: + isaac_gr00t: + module_path: "backbone.model.language_model.model.layers.{i}" + versions: "n1.6" diff --git a/flash_rt/structures/bindings/groot_n16_tick.yaml b/flash_rt/structures/bindings/groot_n16_tick.yaml new file mode 100644 index 00000000..d640910a --- /dev/null +++ b/flash_rt/structures/bindings/groot_n16_tick.yaml @@ -0,0 +1,26 @@ +binding: groot_n16_tick +structure: vla_tick_pipeline + +stages: + obs_encode: + seam: "model.backbone" + capture: eager # SigLIP2 window reordering uses CPU indexing + outputs: + cond_features: "backbone_features" + mutation_guard: shallow_copy # action_head reassigns backbone_features + # (vlln write-back); loop must receive a + # fresh shallow-copy wrapper per call + action_denoise: + seam: "model.action_head.get_action" + loop_steps: 4 # num_inference_timesteps + noise_window: single_randn_site # hoisted; recorded at first eager run + +cadences: + observation: obs_encode + tick: action_denoise + replan: action_denoise # dit-only replay while observation unchanged + +hosts: + isaac_gr00t: + module_path: "policy.model" + versions: "n1.6" diff --git a/flash_rt/structures/bindings/groot_n16_vision.yaml b/flash_rt/structures/bindings/groot_n16_vision.yaml new file mode 100644 index 00000000..beaafc07 --- /dev/null +++ b/flash_rt/structures/bindings/groot_n16_vision.yaml @@ -0,0 +1,23 @@ +binding: groot_n16_vision +structure: vision_ffn +dims: {D: 1152, F: 4304, layers: 27} +variant: {activation: gelu} +boundary_dtype: bf16 + +m_profile: + encode: {m_class: medium, phase: vision_encode} # M = patch tokens/view + +weights_map: + isaac_gr00t: + layout: out_in + w_fc1: "backbone.model.vision_model.vision_model.encoder.layers.{i}.mlp.fc1.weight" + b_fc1: "backbone.model.vision_model.vision_model.encoder.layers.{i}.mlp.fc1.bias" + w_fc2: "backbone.model.vision_model.vision_model.encoder.layers.{i}.mlp.fc2.weight" + b_fc2: "backbone.model.vision_model.vision_model.encoder.layers.{i}.mlp.fc2.bias" + w_norm: "backbone.model.vision_model.vision_model.encoder.layers.{i}.layer_norm2.weight" + b_norm: "backbone.model.vision_model.vision_model.encoder.layers.{i}.layer_norm2.bias" + +hosts: + isaac_gr00t: + module_path: "backbone.model.vision_model.vision_model.encoder.layers.{i}" + versions: "n1.6" diff --git a/flash_rt/structures/bindings/pi05.yaml b/flash_rt/structures/bindings/pi05.yaml new file mode 100644 index 00000000..ada4f0d7 --- /dev/null +++ b/flash_rt/structures/bindings/pi05.yaml @@ -0,0 +1,29 @@ +binding: pi05 +structure: decoder_ffn +dims: {D: 1024, F: 4096, layers: 18} +variant: {activation: gelu, norm_weight_mode: offset, conditioning: ada_ln} +boundary_dtype: bf16 + +m_profile: + denoise: {m_class: small, phase: denoise} # M = host action chunk size + +weights_map: + transformers: + layout: out_in # nn.Linear stores [out, in]; slots are [in, out] + w_gate: "paligemma_with_expert.gemma_expert.model.layers.{i}.mlp.gate_proj.weight" + w_up: "paligemma_with_expert.gemma_expert.model.layers.{i}.mlp.up_proj.weight" + w_down: "paligemma_with_expert.gemma_expert.model.layers.{i}.mlp.down_proj.weight" + w_norm: {const: zeros} # AdaRMS conditional branch has no learned norm weight + +conditioning: + source: time_embedding + projection: + weight: "paligemma_with_expert.gemma_expert.model.layers.{i}.post_attention_layernorm.dense.weight" + bias: "paligemma_with_expert.gemma_expert.model.layers.{i}.post_attention_layernorm.dense.bias" + chunk_order: [scale, shift, gate] + residual_gate: true # output residual is x + ffn_out * gate + +hosts: + lerobot: + module_path: "model.paligemma_with_expert.gemma_expert.model.layers.{i}" + versions: ">=0.5,<0.6" diff --git a/flash_rt/structures/bindings/pi05_prefix.yaml b/flash_rt/structures/bindings/pi05_prefix.yaml new file mode 100644 index 00000000..03d26e28 --- /dev/null +++ b/flash_rt/structures/bindings/pi05_prefix.yaml @@ -0,0 +1,21 @@ +binding: pi05_prefix +structure: decoder_ffn +dims: {D: 2048, F: 16384, layers: 18} +variant: {activation: gelu, norm_weight_mode: offset, conditioning: none} +boundary_dtype: bf16 + +m_profile: + prefill: {m_class: medium, phase: vision_text_prefix} # M = prefix tokens + +weights_map: + transformers: + layout: out_in + w_gate: "paligemma_with_expert.paligemma.model.language_model.layers.{i}.mlp.gate_proj.weight" + w_up: "paligemma_with_expert.paligemma.model.language_model.layers.{i}.mlp.up_proj.weight" + w_down: "paligemma_with_expert.paligemma.model.language_model.layers.{i}.mlp.down_proj.weight" + w_norm: "paligemma_with_expert.paligemma.model.language_model.layers.{i}.post_attention_layernorm.weight" + +hosts: + lerobot: + module_path: "model.paligemma_with_expert.paligemma.model.language_model.layers.{i}" + versions: ">=0.5,<0.6" diff --git a/flash_rt/structures/bindings/pi05_vision.yaml b/flash_rt/structures/bindings/pi05_vision.yaml new file mode 100644 index 00000000..05cafbbb --- /dev/null +++ b/flash_rt/structures/bindings/pi05_vision.yaml @@ -0,0 +1,23 @@ +binding: pi05_vision +structure: vision_ffn +dims: {D: 1152, F: 4304, layers: 27} +variant: {activation: gelu} +boundary_dtype: float32 # lerobot selective-bf16 recipe keeps vision fp32 + +m_profile: + vision: {m_class: medium, phase: vision} # M = views x 256 tokens + +weights_map: + transformers: + layout: out_in_native # fc weights consumed in checkpoint layout + w_fc1: "paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}.mlp.fc1.weight" + b_fc1: "paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}.mlp.fc1.bias" + w_fc2: "paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}.mlp.fc2.weight" + b_fc2: "paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}.mlp.fc2.bias" + w_norm: "paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}.layer_norm2.weight" + b_norm: "paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}.layer_norm2.bias" + +hosts: + lerobot: + module_path: "model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers.{i}" + versions: ">=0.5,<0.6" diff --git a/flash_rt/structures/bindings/qwen25_15b.yaml b/flash_rt/structures/bindings/qwen25_15b.yaml new file mode 100644 index 00000000..791cb508 --- /dev/null +++ b/flash_rt/structures/bindings/qwen25_15b.yaml @@ -0,0 +1,22 @@ +binding: qwen25_15b +structure: decoder_ffn +dims: {D: 1536, F: 8960, layers: 28} +variant: {activation: silu, norm_weight_mode: direct, conditioning: none} +boundary_dtype: bf16 + +m_profile: + prefill: {m_class: medium, phase: prefill} + decode: {m_class: micro, phase: decode} + +weights_map: + transformers: + layout: out_in + w_gate: "model.layers.{i}.mlp.gate_proj.weight" + w_up: "model.layers.{i}.mlp.up_proj.weight" + w_down: "model.layers.{i}.mlp.down_proj.weight" + w_norm: "model.layers.{i}.post_attention_layernorm.weight" + +hosts: + transformers: + module_path: "model.layers.{i}" + versions: ">=4.50,<6" diff --git a/flash_rt/structures/bindings/qwen3_8b.yaml b/flash_rt/structures/bindings/qwen3_8b.yaml new file mode 100644 index 00000000..771b1ffd --- /dev/null +++ b/flash_rt/structures/bindings/qwen3_8b.yaml @@ -0,0 +1,22 @@ +binding: qwen3_8b +structure: decoder_ffn +dims: {D: 4096, F: 12288, layers: 36} +variant: {activation: silu, norm_weight_mode: direct, conditioning: none} +boundary_dtype: bf16 + +m_profile: + prefill: {m_class: medium, phase: prefill} + decode: {m_class: micro, phase: decode} + +weights_map: + transformers: + layout: out_in + w_gate: "model.layers.{i}.mlp.gate_proj.weight" + w_up: "model.layers.{i}.mlp.up_proj.weight" + w_down: "model.layers.{i}.mlp.down_proj.weight" + w_norm: "model.layers.{i}.post_attention_layernorm.weight" + +hosts: + transformers: + module_path: "model.layers.{i}" + versions: ">=4.51,<6" diff --git a/flash_rt/structures/bindings/qwen3_vl_8b_text.yaml b/flash_rt/structures/bindings/qwen3_vl_8b_text.yaml new file mode 100644 index 00000000..b7337e9b --- /dev/null +++ b/flash_rt/structures/bindings/qwen3_vl_8b_text.yaml @@ -0,0 +1,22 @@ +binding: qwen3_vl_8b_text +structure: decoder_ffn +dims: {D: 4096, F: 12288, layers: 36} +variant: {activation: silu, norm_weight_mode: direct, conditioning: none} +boundary_dtype: bf16 + +m_profile: + prefill: {m_class: medium, phase: prefill} # M includes vision tokens + decode: {m_class: micro, phase: decode} + +weights_map: + transformers: + layout: out_in + w_gate: "model.language_model.layers.{i}.mlp.gate_proj.weight" + w_up: "model.language_model.layers.{i}.mlp.up_proj.weight" + w_down: "model.language_model.layers.{i}.mlp.down_proj.weight" + w_norm: "model.language_model.layers.{i}.post_attention_layernorm.weight" + +hosts: + transformers: + module_path: "model.language_model.layers.{i}" + versions: ">=5.0,<6" diff --git a/flash_rt/structures/bindings/qwen3_vl_8b_vision.yaml b/flash_rt/structures/bindings/qwen3_vl_8b_vision.yaml new file mode 100644 index 00000000..4bdecbda --- /dev/null +++ b/flash_rt/structures/bindings/qwen3_vl_8b_vision.yaml @@ -0,0 +1,23 @@ +binding: qwen3_vl_8b_vision +structure: vision_ffn +dims: {D: 1152, F: 4304, layers: 27} +variant: {activation: gelu} +boundary_dtype: bf16 + +m_profile: + encode: {m_class: medium, phase: vision_encode} # M = vision patch tokens + +weights_map: + transformers: + layout: out_in + w_fc1: "model.visual.blocks.{i}.mlp.linear_fc1.weight" + b_fc1: "model.visual.blocks.{i}.mlp.linear_fc1.bias" + w_fc2: "model.visual.blocks.{i}.mlp.linear_fc2.weight" + b_fc2: "model.visual.blocks.{i}.mlp.linear_fc2.bias" + w_norm: "model.visual.blocks.{i}.norm2.weight" + b_norm: "model.visual.blocks.{i}.norm2.bias" + +hosts: + transformers: + module_path: "model.visual.blocks.{i}" + versions: ">=5.0,<6" diff --git a/flash_rt/structures/bindings/smolvla_base.yaml b/flash_rt/structures/bindings/smolvla_base.yaml new file mode 100644 index 00000000..880e67aa --- /dev/null +++ b/flash_rt/structures/bindings/smolvla_base.yaml @@ -0,0 +1,21 @@ +binding: smolvla_base +structure: decoder_ffn +dims: {D: 960, F: 2560, layers: 16} +variant: {activation: silu, norm_weight_mode: direct, conditioning: none} +boundary_dtype: f32 + +m_profile: + prefix: {m_class: medium, phase: prefill} # M = vision+language prefix + +weights_map: + lerobot: + layout: out_in + w_gate: "model.vlm_with_expert.vlm.model.text_model.layers.{i}.mlp.gate_proj.weight" + w_up: "model.vlm_with_expert.vlm.model.text_model.layers.{i}.mlp.up_proj.weight" + w_down: "model.vlm_with_expert.vlm.model.text_model.layers.{i}.mlp.down_proj.weight" + w_norm: "model.vlm_with_expert.vlm.model.text_model.layers.{i}.post_attention_layernorm.weight" + +hosts: + lerobot: + module_path: "model.vlm_with_expert.vlm.model.text_model.layers.{i}" + versions: ">=0.5,<0.6" diff --git a/flash_rt/structures/bindings/smolvla_expert.yaml b/flash_rt/structures/bindings/smolvla_expert.yaml new file mode 100644 index 00000000..6924c257 --- /dev/null +++ b/flash_rt/structures/bindings/smolvla_expert.yaml @@ -0,0 +1,21 @@ +binding: smolvla_expert +structure: decoder_ffn +dims: {D: 720, F: 2048, layers: 16} +variant: {activation: silu, norm_weight_mode: direct, conditioning: none} +boundary_dtype: f32 + +m_profile: + denoise: {m_class: small, phase: denoise} # M = action chunk size + +weights_map: + lerobot: + layout: out_in + w_gate: "model.vlm_with_expert.lm_expert.layers.{i}.mlp.gate_proj.weight" + w_up: "model.vlm_with_expert.lm_expert.layers.{i}.mlp.up_proj.weight" + w_down: "model.vlm_with_expert.lm_expert.layers.{i}.mlp.down_proj.weight" + w_norm: "model.vlm_with_expert.lm_expert.layers.{i}.post_attention_layernorm.weight" + +hosts: + lerobot: + module_path: "model.vlm_with_expert.lm_expert.layers.{i}" + versions: ">=0.5,<0.6" diff --git a/flash_rt/structures/bindings/smolvla_tick.yaml b/flash_rt/structures/bindings/smolvla_tick.yaml new file mode 100644 index 00000000..1959c9b6 --- /dev/null +++ b/flash_rt/structures/bindings/smolvla_tick.yaml @@ -0,0 +1,27 @@ +binding: smolvla_tick +structure: vla_tick_pipeline + +stages: + obs_encode: + seam: "model.embed_prefix + model.vlm_with_expert(fill_kv_cache=True)" + capture: eager # SmolVLM vision embeddings use boolean-mask + # scatter (position_ids[mask]=...) which is + # illegal on a capturing stream — same class + # as the GR00T SigLIP2 CPU-indexing ruling + outputs: + cond_features: "past_key_values" # prefix KV, copied into static + # buffers, read-only in the loop + action_denoise: + seam: "model.denoise_step" + loop_steps: 10 # config.num_steps + noise_window: injectable # sample_actions(noise=...) native window + +cadences: + observation: obs_encode + tick: [obs_encode, action_denoise] # eager prefill + denoise graph replay + replan: action_denoise # denoise-only replay, static prefix KV + +hosts: + lerobot: + module_path: "policy.model" + versions: "smolvla_base" diff --git a/flash_rt/structures/bindings/smolvla_vision.yaml b/flash_rt/structures/bindings/smolvla_vision.yaml new file mode 100644 index 00000000..fc1a8e36 --- /dev/null +++ b/flash_rt/structures/bindings/smolvla_vision.yaml @@ -0,0 +1,23 @@ +binding: smolvla_vision +structure: vision_ffn +dims: {D: 768, F: 3072, layers: 12} +variant: {activation: gelu} +boundary_dtype: f32 + +m_profile: + encode: {m_class: medium, phase: vision_encode} # M = patch tokens/view + +weights_map: + lerobot: + layout: out_in + w_fc1: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}.mlp.fc1.weight" + b_fc1: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}.mlp.fc1.bias" + w_fc2: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}.mlp.fc2.weight" + b_fc2: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}.mlp.fc2.bias" + w_norm: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}.layer_norm2.weight" + b_norm: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}.layer_norm2.bias" + +hosts: + lerobot: + module_path: "model.vlm_with_expert.vlm.model.vision_model.encoder.layers.{i}" + versions: ">=0.5,<0.6" diff --git a/flash_rt/structures/catalog/__init__.py b/flash_rt/structures/catalog/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/adaln_producer/__init__.py b/flash_rt/structures/catalog/adaln_producer/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/adaln_producer/reference.py b/flash_rt/structures/catalog/adaln_producer/reference.py new file mode 100644 index 00000000..97e160bf --- /dev/null +++ b/flash_rt/structures/catalog/adaln_producer/reference.py @@ -0,0 +1,40 @@ +"""Ground-truth reference for the ``adaln_producer`` structure.""" + +from __future__ import annotations + +import torch + + +def adaln_producer_ref( + x: torch.Tensor, + cond: torch.Tensor, + style_w: torch.Tensor, + style_b: torch.Tensor | None = None, + norm_w: torch.Tensor | None = None, + *, + variant: dict[str, str] | None = None, + eps: float = 1e-6, +): + """Plain-torch reference: norm(x) * (1 + scale) + shift, and gate. + + The style projection is computed here rather than looked up — the + step table is an execution decision, and at the declared boundary + the structure is the conditioning projection plus the modulated + norm. ``style_w`` uses the declared [C, 3*D] slot layout, its three + chunks being scale, shift and gate. + """ + variant = variant or {} + style = cond.to(torch.float32) @ style_w.to(torch.float32) + if style_b is not None: + style = style + style_b.to(torch.float32) + scale, shift, gate = style.chunk(3, dim=-1) + xf = x.to(torch.float32) + if variant.get("norm", "rms") == "rms": + normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) + else: + normed = (xf - xf.mean(-1, keepdim=True)) * torch.rsqrt( + xf.var(-1, keepdim=True, unbiased=False) + eps) + if norm_w is not None: + normed = normed * norm_w.to(torch.float32) + y = normed * (1.0 + scale) + shift + return y.to(x.dtype), gate.to(x.dtype) diff --git a/flash_rt/structures/catalog/adaln_producer/structure.yaml b/flash_rt/structures/catalog/adaln_producer/structure.yaml new file mode 100644 index 00000000..00895c44 --- /dev/null +++ b/flash_rt/structures/catalog/adaln_producer/structure.yaml @@ -0,0 +1,72 @@ +structure: adaln_producer +version: 1 +description: > + An adaptive norm driven by a step-quantized conditioning vector: + norm, modulation by a conditioning-derived style, and an optional + static FP8 quantization of the output, resolved as one region. The + conditioning takes one of a small fixed set of values over a tick, + so the style projection is memoized at bind time and the step is + recovered in-graph by a fingerprint lookup. The FP8 output form is + the upstream half of a producer-consumer seam: the shared activation + scale lets the consuming projection skip its own input quantization. + +reference: + module: adaln_producer.reference + entrypoint: adaln_producer_ref + +boundary: + symbolic_dims: [M, D] + inputs: + - {name: x, dims: [M, D], dtype: "@binding"} + - {name: cond, dims: [1, C], dtype: "@binding"} + outputs: + - {name: y, dims: [M, D], dtype: "@binding"} + - {name: gate, dims: [1, D], dtype: "@binding", optional: true} + +weights: + - {slot: style_w, dims: [C, "3*D"]} + - {slot: style_b, dims: ["3*D"], optional: true} + - {slot: norm_w, dims: [D], optional: true} + +variants: + bind: [table_only, fused_norm] # table_only replaces just the + # conditioning projection; fused_norm + # replaces the whole adaptive norm + norm: [rms, layer] + out_dtype: [bf16, fp8_static] + epilogue: [none, gate_residual] + +calibration: + points: [cond, x] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, vendor] + rule: net_positive_including_boundary + per_shape: true + qualification: + max_steps: 64 + fingerprint_rel_margin: 0.001 + +qualification: + - "the conditioning must be step-quantized: more distinct vectors + than max_steps means it depends on more than the step and a table + would alias different inputs onto one row" + - "the fingerprint coordinates must separate the stored vectors by a + real margin, otherwise the in-graph lookup is not decidable" + - "the step lookup is pure tensor work (no host-side state), so it + stays valid under both dynamo tracing and graph capture" + - "bind scope is a measurement, not a default: the fused_norm form + pays a kernel the host may already fuse, so it earns its place + only where it also feeds a downstream seam (pi05 r15: fused_norm + at the 19 bf16 sites cost 1.9% against table_only)" + +evidence: + - "pi05 expert AdaRMS — r4 step table (fp32 GEMV 3.67ms bucket + removed), r11 fingerprint locator (cuBLAS GEMV 370x removed)" + - "pi05 expert fused fp8 producer — r7..r14, feeds the packed + projection seam" + - "GROOT DiT AdaLayerNorm — round3 fp8 producer, 32 blocks" diff --git a/flash_rt/structures/catalog/attention_core/__init__.py b/flash_rt/structures/catalog/attention_core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/attention_core/reference.py b/flash_rt/structures/catalog/attention_core/reference.py new file mode 100644 index 00000000..777a5d32 --- /dev/null +++ b/flash_rt/structures/catalog/attention_core/reference.py @@ -0,0 +1,36 @@ +"""Ground-truth reference for the ``attention_core`` structure.""" + +from __future__ import annotations + +import torch + + +def attention_core_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None = None, + *, + scale: float | None = None, + variant: dict[str, str] | None = None, +) -> torch.Tensor: + """Plain-torch reference: masked softmax attention, host layout in, + (B, Sq, H, D) out. + + Packing and kernel choice are execution decisions; at the declared + boundary this is ordinary attention against the host's own mask, + which is what the packing must reproduce. + """ + del variant + scale = scale if scale is not None else q.shape[-1] ** -0.5 + heads, kv_heads = q.shape[1], k.shape[1] + if kv_heads != heads: + k = k.expand(-1, heads, -1, -1) if kv_heads == 1 else \ + k.repeat_interleave(heads // kv_heads, dim=1) + v = v.expand(-1, heads, -1, -1) if kv_heads == 1 else \ + v.repeat_interleave(heads // kv_heads, dim=1) + scores = (q.float() @ k.float().transpose(-1, -2)) * scale + if mask is not None: + scores = scores + mask.float() + out = torch.softmax(scores, dim=-1) @ v.float() + return out.transpose(1, 2).to(q.dtype) diff --git a/flash_rt/structures/catalog/attention_core/structure.yaml b/flash_rt/structures/catalog/attention_core/structure.yaml new file mode 100644 index 00000000..383104b4 --- /dev/null +++ b/flash_rt/structures/catalog/attention_core/structure.yaml @@ -0,0 +1,65 @@ +structure: attention_core +version: 1 +description: > + Attention run on a fused kernel, with the host's additive mask + resolved into packed keys and values at bind time. A mask that blocks + one contiguous run of positions has no form the kernel accepts, but + the surviving positions do: packing them leaves a dense attention. + The packed prefix is encoder output for the current observation and + therefore belongs to a slower cadence — it is filled once and + refreshed by an update function, while the suffix is written per + call. + +reference: + module: attention_core.reference + entrypoint: attention_core_ref + +boundary: + symbolic_dims: [B, H, Sq, Skv, D] + inputs: + - {name: q, dims: [B, H, Sq, D], dtype: "@binding"} + - {name: k, dims: [B, Hkv, Skv, D], dtype: "@binding"} + - {name: v, dims: [B, Hkv, Skv, D], dtype: "@binding"} + - {name: mask, dims: [B, 1, Sq, Skv], dtype: "@binding", optional: true} + outputs: + - {name: out, dims: [B, Sq, H, D], dtype: "@binding"} + +weights: [] + +variants: + kernel: [fa2_seqused, host_fallback] + kv_cadence: [static_prefix, per_call] + +calibration: + points: [q, k, v, mask] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, vendor] + rule: net_positive_including_boundary + per_shape: true + qualification: + head_dims: [64, 96, 128, 256] + mask_runs_max: 1 + prefix_static_rtol: 0.001 + +qualification: + - "head_dim must be one the kernel supports; binding returns nothing + rather than failing, so the host keeps a fallback attention" + - "every query row must see the same mask and the blocked positions + must form a single contiguous run — anything else is not a packed + dense attention" + - "the prefix keys must not move across the hot loop; a moving + prefix means the cadence split is wrong, and it is rejected here + where the reason is still legible rather than downstream as a + parity failure" + +evidence: + - "standalone at pi05 denoise shapes — cos 0.9999979 vs eager+mask, + 2.04x (11.7 vs 23.9us), split-KV workspace selected" + - "pi05 structures stack — round13, 2.246 -> 2.308x" + - "pi05 hub-kernel demo e2e — own FA2 21.42ms vs community 21.61ms, + action cosine 0.99987 against the native FlashRT pipeline" diff --git a/flash_rt/structures/catalog/cadence_static/__init__.py b/flash_rt/structures/catalog/cadence_static/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/cadence_static/reference.py b/flash_rt/structures/catalog/cadence_static/reference.py new file mode 100644 index 00000000..b71ae992 --- /dev/null +++ b/flash_rt/structures/catalog/cadence_static/reference.py @@ -0,0 +1,24 @@ +"""Ground-truth reference for the ``cadence_static`` structure.""" + +from __future__ import annotations + +from typing import Callable + +import torch + + +def cadence_static_ref( + x: torch.Tensor, + host_fn: Callable[[torch.Tensor], torch.Tensor], + *, + variant: dict[str, str] | None = None, +) -> torch.Tensor: + """Plain-torch reference: simply run the host's own computation. + + Holding the result in a buffer is an execution decision. At the + declared boundary this structure is the identity on the host's + module, which is exactly why binding must first prove the output + does not move within the hot loop. + """ + del variant + return host_fn(x) diff --git a/flash_rt/structures/catalog/cadence_static/structure.yaml b/flash_rt/structures/catalog/cadence_static/structure.yaml new file mode 100644 index 00000000..7e1fdd4e --- /dev/null +++ b/flash_rt/structures/catalog/cadence_static/structure.yaml @@ -0,0 +1,56 @@ +structure: cadence_static +version: 1 +description: > + Work that changes on a slower cadence than the hot loop, held in + static buffers the captured graph reads and refreshed by an explicit + update function outside it. Cross-attention key/value projections + over a fixed prompt, encoder features over one observation, and + prefix KV in an action-chunk decoder are all instances: the hot loop + re-executes them every tick although their inputs only move when a + new observation arrives. + +reference: + module: cadence_static.reference + entrypoint: cadence_static_ref + +boundary: + symbolic_dims: [M, N] + inputs: + - {name: x, dims: [M, "*"], dtype: "@binding"} + outputs: + - {name: y, dims: [M, N], dtype: "@binding"} + +weights: [] + +variants: + cadence: [observation, prompt] + update: [host_recompute, external_write] + +calibration: + points: [y] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, vendor] + rule: net_positive_including_boundary + per_shape: true + qualification: + invariance_rtol: 0.001 + min_hot_loop_captures: 2 + +qualification: + - "the wrapped output must be constant across the fast loop, proven + from captures of several hot-loop iterations — a per-step output + is not a cadence substructure and freezing it would silently + change the model" + - "the update function is returned to the caller, not hidden inside + the replacement: the recipe registers it as an outside update so + the slower cadence stays visible and its cost stays measured" + +evidence: + - "GROOT cross-attention KV — round2, 16 blocks, L3 14.32 -> 12.71ms" + - "pi05 prefix KV in the packed attention core — round13, the + obs-cadence mismatch was caught by the parity gate (0.984)" diff --git a/flash_rt/structures/catalog/decoder_ffn/__init__.py b/flash_rt/structures/catalog/decoder_ffn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/decoder_ffn/reference.py b/flash_rt/structures/catalog/decoder_ffn/reference.py new file mode 100644 index 00000000..8ed5b57f --- /dev/null +++ b/flash_rt/structures/catalog/decoder_ffn/reference.py @@ -0,0 +1,64 @@ +"""Reference implementation for the ``decoder_ffn`` structure. + +Ground truth for the qualification gates: the plainest possible PyTorch, +never executed on a serving hot path. Norm statistics are computed in +float32 and cast back to the boundary dtype, matching the numerical +convention of the model families this structure binds to. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +_ACTIVATIONS = { + "gelu": lambda t: F.gelu(t, approximate="tanh"), + "silu": F.silu, +} + + +def decoder_ffn_ref( + x: torch.Tensor, + w_norm: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + *, + activation: str = "gelu", + norm_weight_mode: str = "offset", + cond_scale: torch.Tensor | None = None, + cond_shift: torch.Tensor | None = None, + cond_gate: torch.Tensor | None = None, + eps: float = 1e-6, +) -> torch.Tensor: + """x -> RMSNorm -> gate/up GEMM -> act(gate) * up -> down GEMM -> + x. + + ``norm_weight_mode``: "offset" multiplies by ``1 + w_norm`` (Gemma + convention), "direct" multiplies by ``w_norm`` (Qwen convention). + ``cond_scale``/``cond_shift``/``cond_gate`` apply AdaLN-style + modulation under the ``ada_ln`` conditioning variant: the norm output + becomes ``normed * (1 + scale) + shift`` and the residual becomes + ``x + ffn_out * gate``. A binding whose AdaLN branch has no learned + norm weight maps ``w_norm`` to zeros under "offset" mode. + """ + act = _ACTIVATIONS[activation] + + h = x.float() + h = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + eps) + if norm_weight_mode == "offset": + h = h * (1.0 + w_norm.float()) + elif norm_weight_mode == "direct": + h = h * w_norm.float() + else: + raise ValueError(f"unknown norm_weight_mode: {norm_weight_mode!r}") + if cond_scale is not None: + h = h * (1.0 + cond_scale.float()) + if cond_shift is not None: + h = h + cond_shift.float() + h = h.to(x.dtype) + + hidden = act(h @ w_gate) * (h @ w_up) + out = hidden @ w_down + if cond_gate is not None: + out = out * cond_gate + return x + out diff --git a/flash_rt/structures/catalog/decoder_ffn/structure.yaml b/flash_rt/structures/catalog/decoder_ffn/structure.yaml new file mode 100644 index 00000000..2f418192 --- /dev/null +++ b/flash_rt/structures/catalog/decoder_ffn/structure.yaml @@ -0,0 +1,44 @@ +structure: decoder_ffn +version: 1 +description: > + Post-attention feed-forward block of a transformer decoder layer: + x -> RMSNorm (optionally AdaLN-modulated) -> gate/up projections -> + activation(gate) * up -> down projection -> residual add (optionally + gated by an AdaLN gate term). + +reference: + module: decoder_ffn.reference + entrypoint: decoder_ffn_ref + +boundary: + symbolic_dims: [M, D, F] + inputs: + - {name: x, dims: [M, D], dtype: "@binding"} + - {name: cond_scale, dims: [M, D], dtype: "@binding", optional: true} + - {name: cond_shift, dims: [M, D], dtype: "@binding", optional: true} + - {name: cond_gate, dims: [M, D], dtype: "@binding", optional: true} + outputs: + - {name: x, dims: [M, D], dtype: "@binding"} + +weights: + - {slot: w_norm, dims: [D]} + - {slot: w_gate, dims: [D, F]} + - {slot: w_up, dims: [D, F]} + - {slot: w_down, dims: [F, D]} + +variants: + activation: [gelu, silu] + norm_weight_mode: [offset, direct] + conditioning: [none, ada_ln] + +calibration: + points: [x_after_norm, act_after_mul] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, unfused_chain, vendor] + rule: net_positive_including_boundary + per_shape: true diff --git a/flash_rt/structures/catalog/linear_proj/__init__.py b/flash_rt/structures/catalog/linear_proj/__init__.py new file mode 100644 index 00000000..ddccb7a4 --- /dev/null +++ b/flash_rt/structures/catalog/linear_proj/__init__.py @@ -0,0 +1,3 @@ +"""linear_proj: the finest generally-useful region — a calibrated +projection with epilogue variants. See structure.yaml. +""" diff --git a/flash_rt/structures/catalog/linear_proj/reference.py b/flash_rt/structures/catalog/linear_proj/reference.py new file mode 100644 index 00000000..5dc7c1a3 --- /dev/null +++ b/flash_rt/structures/catalog/linear_proj/reference.py @@ -0,0 +1,35 @@ +"""Ground-truth reference for the ``linear_proj`` structure.""" + +from __future__ import annotations + +import torch + + +def linear_proj_ref( + x: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + *, + variant: dict[str, str] | None = None, +) -> torch.Tensor: + """Plain-torch reference: y = x @ w (+ b) (+ epilogue). + + ``w`` uses the declared [K, N] slot layout. Epilogue variants: + ``none`` returns the projection; ``residual_add`` adds the residual + input; ``gelu_quant_fp8`` applies tanh-GELU (the fp8 quantization of + the epilogue output is an implementation detail of fused impls — at + the declared boundary the reference stays in the compute dtype). + """ + variant = variant or {} + y = x.to(torch.float32) @ w.to(torch.float32) + if b is not None and variant.get("bias", "add") != "none": + y = y + b.to(torch.float32) + epilogue = variant.get("epilogue", "none") + if epilogue == "gelu_quant_fp8": + y = torch.nn.functional.gelu(y, approximate="tanh") + elif epilogue == "residual_add": + if residual is None: + raise ValueError("residual_add epilogue requires residual input") + y = y + residual.to(torch.float32) + return y.to(x.dtype) diff --git a/flash_rt/structures/catalog/linear_proj/structure.yaml b/flash_rt/structures/catalog/linear_proj/structure.yaml new file mode 100644 index 00000000..404f451a --- /dev/null +++ b/flash_rt/structures/catalog/linear_proj/structure.yaml @@ -0,0 +1,51 @@ +structure: linear_proj +version: 1 +description: > + A single projection with calibration semantics: x[M,K] -> y[M,N] + with optional bias and a fused epilogue. The finest generally-useful + region boundary — attention Q/K/V/O projections, DiT projections and + non-gated MLP halves are all instances. Elementwise work (residual, + gate, activation+quant) exists only as epilogue variants of this + boundary, never as standalone regions. + +reference: + module: linear_proj.reference + entrypoint: linear_proj_ref + +boundary: + symbolic_dims: [M, K, N] + inputs: + - {name: x, dims: [M, K], dtype: "@binding"} + - {name: residual, dims: [M, N], dtype: "@binding", optional: true} + outputs: + - {name: y, dims: [M, N], dtype: "@binding"} + +weights: + - {slot: w, dims: [K, N]} + - {slot: b, dims: [N], optional: true} + +variants: + bias: [none, add] + epilogue: [none, gelu_quant_fp8, residual_add] + in_dtype: [bf16, fp8_static] # fp8_static = producer-negotiated seam: + # upstream norm/adaln producer emits + # fp8+scale, this region skips its own + # input quantization (plan-level contract, + # re-certified at the composed boundary) + +calibration: + points: [x] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, vendor] + rule: net_positive_including_boundary + per_shape: true + qualification: + weight_elements_min: 1048576 # do not candidate tiny projections; + # discovery groups sibling q/k/v/o + # into one family, never blind-scans + # every nn.Linear diff --git a/flash_rt/structures/catalog/qkv_pack/__init__.py b/flash_rt/structures/catalog/qkv_pack/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/qkv_pack/reference.py b/flash_rt/structures/catalog/qkv_pack/reference.py new file mode 100644 index 00000000..d2b8baf6 --- /dev/null +++ b/flash_rt/structures/catalog/qkv_pack/reference.py @@ -0,0 +1,33 @@ +"""Ground-truth reference for the ``qkv_pack`` structure.""" + +from __future__ import annotations + +from typing import Sequence + +import torch + + +def qkv_pack_ref( + x: torch.Tensor, + weights: Sequence[torch.Tensor], + biases: Sequence[torch.Tensor] | None = None, + *, + variant: dict[str, str] | None = None, +) -> list[torch.Tensor]: + """Plain-torch reference: the sibling projections, computed apart. + + Packing is an execution decision — at the declared boundary the + structure is exactly the group of independent projections, in + sibling order. ``weights`` use the declared [K, N_i] slot layout. + The module bind form is judged at the same boundary: its attention + and output projection belong to the host, and only the packed + projections are this structure's contract. + """ + del variant + outs = [] + for i, w in enumerate(weights): + y = x.to(torch.float32) @ w.to(torch.float32) + if biases is not None and biases[i] is not None: + y = y + biases[i].to(torch.float32) + outs.append(y.to(x.dtype)) + return outs diff --git a/flash_rt/structures/catalog/qkv_pack/structure.yaml b/flash_rt/structures/catalog/qkv_pack/structure.yaml new file mode 100644 index 00000000..6f84316c --- /dev/null +++ b/flash_rt/structures/catalog/qkv_pack/structure.yaml @@ -0,0 +1,65 @@ +structure: qkv_pack +version: 1 +description: > + Pack sibling linears that share one input and a fixed consumption + order (attention q/k/v, gated-MLP gate/up) into a single GEMM. + Small-M sibling GEMMs sit on the launch/latency floor, so the win is + structural (three launches become one), not per-kernel. Two bind + forms: "leaf" replaces the first sibling slot with the packed GEMM + (later slots read preallocated stash buffers; the host's own call + order is the in-graph data dependency) and "module" replaces a whole + q/k/v/out attention block with packed GEMM + SDPA at a declared + compute dtype + the original out projection. + +reference: + module: qkv_pack.reference + entrypoint: qkv_pack_ref + +boundary: + symbolic_dims: [M, K, N_i] + inputs: + - {name: x, dims: [M, K], dtype: "@binding"} + outputs: + - {name: y_i, dims: [M, N_i], dtype: "@binding", repeated: true} + +weights: + - {slot: w_i, dims: [K, N_i], repeated: true} + - {slot: b_i, dims: [N_i], optional: true, repeated: true} + +variants: + bind: [leaf, module] + in_dtype: [bf16_fused_quant, fp8_static] + sdpa_dtype: [none, bf16] # module form only + +qualification: + - "siblings share the input tensor and K dim; call order is fixed and + observable (the leaf form relies on it for stash ordering)" + - "joint per-tensor weight scale — the rounding difference against + per-sibling scales is adjudicated by the layer parity gate" + - "module form requires q_proj/k_proj/v_proj/out_proj plus head_dim + and scale attributes, with the standard block dataflow" + +calibration: + points: [x] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, vendor] + rule: net_positive_including_boundary + per_shape: true + qualification: + siblings_min: 2 + weight_elements_min: 1048576 + +discovery: + families: > + attention sibling groups already produced by the linear_proj + discovery patterns (q/k/v[/out] under one parent) + +evidence: + - "pi05 expert leaf form — r8, +0.087 ratio (launch 3->1, M=50)" + - "pi05 vision module form — r9, +0.060 ratio (SigLIP, 27 blocks)" + - "pi05 trunk leaf bf16 entry — r12 (M=712 prefill)" diff --git a/flash_rt/structures/catalog/vision_ffn/__init__.py b/flash_rt/structures/catalog/vision_ffn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/catalog/vision_ffn/reference.py b/flash_rt/structures/catalog/vision_ffn/reference.py new file mode 100644 index 00000000..485b45a4 --- /dev/null +++ b/flash_rt/structures/catalog/vision_ffn/reference.py @@ -0,0 +1,33 @@ +"""Reference implementation for the ``vision_ffn`` structure. + +Ground truth for the qualification gates: the plainest possible PyTorch, +never executed on a serving hot path. LayerNorm statistics follow the +boundary dtype's promotion rules via float32, matching the vision-tower +convention of the model families this structure binds to. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def vision_ffn_ref( + x: torch.Tensor, + w_norm: torch.Tensor, + b_norm: torch.Tensor, + w_fc1: torch.Tensor, + b_fc1: torch.Tensor, + w_fc2: torch.Tensor, + b_fc2: torch.Tensor, + *, + activation: str = "gelu", + eps: float = 1e-6, +) -> torch.Tensor: + """x -> LayerNorm -> fc1 -> GELU(tanh) -> fc2 -> + x.""" + if activation != "gelu": + raise ValueError(f"unsupported activation: {activation!r}") + h = F.layer_norm(x.float(), (x.shape[-1],), + w_norm.float(), b_norm.float(), eps).to(x.dtype) + hidden = F.gelu(h @ w_fc1.t() + b_fc1, approximate="tanh") + return x + (hidden @ w_fc2.t() + b_fc2) diff --git a/flash_rt/structures/catalog/vision_ffn/structure.yaml b/flash_rt/structures/catalog/vision_ffn/structure.yaml new file mode 100644 index 00000000..d6721afa --- /dev/null +++ b/flash_rt/structures/catalog/vision_ffn/structure.yaml @@ -0,0 +1,41 @@ +structure: vision_ffn +version: 1 +description: > + Feed-forward block of a vision-transformer encoder layer: + x -> LayerNorm (affine + bias) -> fc1 (+bias) -> GELU -> + fc2 (+bias) -> residual add. Non-gated MLP, as used by the SigLIP + towers shared across the supported VLA families. + +reference: + module: vision_ffn.reference + entrypoint: vision_ffn_ref + +boundary: + symbolic_dims: [M, D, F] + inputs: + - {name: x, dims: [M, D], dtype: "@binding"} + outputs: + - {name: x, dims: [M, D], dtype: "@binding"} + +weights: + - {slot: w_norm, dims: [D]} + - {slot: b_norm, dims: [D]} + - {slot: w_fc1, dims: [F, D]} + - {slot: b_fc1, dims: [F]} + - {slot: w_fc2, dims: [D, F]} + - {slot: b_fc2, dims: [D]} + +variants: + activation: [gelu] + +calibration: + points: [x_after_norm, hidden_after_act] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + latency: + baselines: [torch_compile, unfused_chain, vendor] + rule: net_positive_including_boundary + per_shape: true diff --git a/flash_rt/structures/catalog/vla_tick_pipeline/__init__.py b/flash_rt/structures/catalog/vla_tick_pipeline/__init__.py new file mode 100644 index 00000000..de84ec36 --- /dev/null +++ b/flash_rt/structures/catalog/vla_tick_pipeline/__init__.py @@ -0,0 +1,7 @@ +"""vla_tick_pipeline: schedule-layer stage structure (cond_iter_pipeline +family, tick specialization). + +Unlike region structures there is no standalone torch reference module: +the parity reference is the host's own eager path under the same noise +window (see structure.yaml gates.parity.reference). +""" diff --git a/flash_rt/structures/catalog/vla_tick_pipeline/structure.yaml b/flash_rt/structures/catalog/vla_tick_pipeline/structure.yaml new file mode 100644 index 00000000..ddae10fa --- /dev/null +++ b/flash_rt/structures/catalog/vla_tick_pipeline/structure.yaml @@ -0,0 +1,69 @@ +structure: vla_tick_pipeline +kind: stage_pipeline +family: cond_iter_pipeline +version: 1 +description: > + Schedule-layer structure for tick-driven VLA inference: a condition + encoding stage (vision tower + LLM + adapters) that runs at observation + cadence, feeding a fixed-shape iterative denoise hot loop (diffusion or + flow action head, K steps) that runs at tick cadence and is captured as + a single graph. Noise is an explicit SWAP window; condition features are + static read-only buffers across the loop. + +family_note: > + cond_iter_pipeline family = cond_encode once + fixed-shape iterative hot + loop + explicit noise/init window + static condition buffers. Sibling + specializations: one-shot image/video generation (cadence axis degenerates + to once per generation; capture yield ~0 in the compute-bound large-step + regime, net-win gate routes to in-step region structures) and streaming + frame loops (TTS frames, speculative-decode draft loops). This + specialization adds tick/cadence semantics: observation, replan and + prompt cadences plus action-chunk readback. + +stages: + - name: obs_encode + cadence: observation + inputs: + - {name: images, window: swap} + - {name: language, cadence: prompt} + - {name: state, window: swap} + outputs: + - {name: cond_features, buffer: static} + capture: host_dependent # CPU-side indexing in vision towers is common; + # compile-then-capture or eager+copy both conform + - name: action_denoise + cadence: tick + loop: {steps: "@binding", fixed_shape: true} + inputs: + - {name: cond_features, buffer: static, access: read_only} + - {name: state, window: swap, optional: true} + - {name: noise, window: swap} + outputs: + - {name: action_chunk, window: swap, readback: true} + capture: required_for_yield + +embedded_regions: [decoder_ffn, vision_ffn] + +cadence_substructures: + - {name: cross_kv_precompute, cadence: prompt} + - {name: adaln_step_table, cadence: static_per_plan} + +conformance: + # audited per host before binding; failures route to sibling + # specializations or refusal + - stage_seam_exists + - hot_loop_fixed_shape_fixed_steps + - noise_single_hoistable_site + - condition_read_only_in_loop + - action_chunk_cadence + - capture_blockers_enumerated + +gates: + parity: + reference: host_eager_same_seed + boundary: action_chunk + rule: noise_must_be_explicit_window # in-graph RNG never matches eager + metrics: [cosine, max_abs] + net_win: + rule: tick_latency_ab_vs_host_eager + report: per_stage_times # cadence amortization is part of the account diff --git a/flash_rt/structures/discover.py b/flash_rt/structures/discover.py new file mode 100644 index 00000000..0a7de260 --- /dev/null +++ b/flash_rt/structures/discover.py @@ -0,0 +1,261 @@ +"""Structure discovery — find catalog structures inside a host model. + +Walks the module tree and matches region-structure seams by shape, not +by model name: a gated gate/up/down MLP is a ``decoder_ffn`` seam, a +fc1/fc2 MLP with a sibling LayerNorm is a ``vision_ffn`` seam. The +result is the same information a hand-written binding file carries +(paths, dims, variant), derived from the model object itself; bindings +become generated receipts instead of required inputs. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +import torch +from torch import nn + +_DECODER_PROJ = ("gate_proj", "up_proj", "down_proj") +_VISION_PROJ = (("fc1", "fc2"), ("linear_fc1", "linear_fc2")) +_NORM_ATTRS = ("post_attention_layernorm", "layer_norm2", "norm2") +_ATTN_PROJ = (("q_proj", "k_proj", "v_proj", "o_proj"), + ("q_proj", "k_proj", "v_proj", "out_proj")) +# sibling groups that qkv_pack packs into one GEMM: same input, fixed +# consumption order. The trailing o_proj/out_proj is not part of the +# pack (it consumes the attention output, not the shared input). +_QKV_PACK = (("q_proj", "k_proj", "v_proj"), + ("to_q", "to_k", "to_v")) +# adaptive-norm modules: a norm that also projects a conditioning +# vector. The child that produces the modulation is the tell. +_COND_PROJ_ATTRS = ("dense", "linear", "adaLN_modulation", "modulation") +_PROJ_WEIGHT_FLOOR = 262144 # candidacy filter only; impls add their own + # work-based qualification and gates decide + + +def _has_cond_forward(module: nn.Module) -> bool: + """A norm takes a conditioning argument (adaptive norm) if its + forward accepts a second positional / a ``cond``/``temb`` keyword.""" + import inspect + try: + params = list(inspect.signature(module.forward).parameters) + except (TypeError, ValueError): + return False + if any(p in params for p in ("cond", "temb", "emb", "c")): + return True + # (self is bound out of module.forward already) x + one more positional + positional = [p for p in params if p not in ("args", "kwargs")] + return len(positional) >= 2 + + +@dataclass +class Seam: + """One replaceable site: a structure instance found in the host.""" + + structure: str + path: str # dotted path of the swappable module + parent_path: str + norm_attr: str | None + dims: dict[str, int] + variant: dict[str, str] + fc_attrs: tuple[str, str] | None = None + proj_attr: str | None = None # linear_proj: attr name in parent + pack_attrs: tuple[str, ...] | None = None # qkv_pack: sibling attrs + cond_attr: str | None = None # adaln_producer: cond-proj child + family: str = "" + layer_index: int = -1 + m_profile: list[int] = field(default_factory=list) + + +def _activation_of(module: nn.Module) -> str | None: + fn = None + for attr in ("act_fn", "activation_fn", "act", "activation"): + fn = getattr(module, attr, None) + if fn is not None: + break + if fn is None: + return None + label = " ".join( + [getattr(fn, "__name__", ""), type(fn).__name__, repr(fn)]).lower() + if "silu" in label or "swish" in label: + return "silu" + if "gelu" in label: + return "gelu" + return None + + +def _resolve(root: nn.Module, path: str) -> nn.Module: + node = root + for part in path.split("."): + if part: + node = node[int(part)] if part.isdigit() else getattr(node, part) + return node + + +def _family_key(path: str) -> tuple[str, int]: + """Template the trailing layer index: a.layers.12.mlp -> a.layers.{i}.mlp.""" + matches = list(re.finditer(r"\.(\d+)\.", "." + path + ".")) + if not matches: + return path, -1 + m = matches[-1] + start, end = m.start(1) - 1, m.end(1) - 1 # offsets in original path + return path[:start] + "{i}" + path[end:], int(m.group(1)) + + +def _norm_attr_of(root: nn.Module, parent_path: str) -> str | None: + try: + parent = _resolve(root, parent_path) + except (AttributeError, IndexError, KeyError): + return None + for attr in _NORM_ATTRS: + if isinstance(getattr(parent, attr, None), nn.Module): + return attr + return None + + +def discover( + model: nn.Module, + structures: tuple[str, ...] = ("decoder_ffn", "vision_ffn"), +) -> list[Seam]: + """Find every region-structure seam in ``model``.""" + seams: list[Seam] = [] + for path, module in model.named_modules(): + if not path: + continue + parent_path = path.rsplit(".", 1)[0] if "." in path else "" + if "decoder_ffn" in structures and all( + isinstance(getattr(module, a, None), nn.Linear) + for a in _DECODER_PROJ + ): + gate = module.gate_proj + act = _activation_of(module) or "silu" + family, idx = _family_key(path) + seams.append(Seam( + structure="decoder_ffn", path=path, parent_path=parent_path, + norm_attr=_norm_attr_of(model, parent_path), + dims={"D": gate.in_features, "F": gate.out_features}, + variant={"activation": act, "norm_weight_mode": "direct"}, + family=family, layer_index=idx)) + continue + if "qkv_pack" in structures: + for group in _QKV_PACK: + projs = [getattr(module, a, None) for a in group] + if not all(isinstance(p, nn.Linear) for p in projs): + continue + if len({p.in_features for p in projs}) != 1: + continue # siblings must share the input dim + if projs[0].weight.numel() < _PROJ_WEIGHT_FLOOR: + continue + family, idx = _family_key(path) + seams.append(Seam( + structure="qkv_pack", path=path, parent_path=parent_path, + norm_attr=None, pack_attrs=group, + dims={"K": projs[0].in_features, + "N": sum(p.out_features for p in projs)}, + variant={"bind": "leaf", "in_dtype": "bf16_fused_quant"}, + family=family, layer_index=idx)) + break + if "adaln_producer" in structures: + cond_attr = next( + (a for a in _COND_PROJ_ATTRS + if isinstance(getattr(module, a, None), nn.Linear)), None) + if (cond_attr is not None and _has_cond_forward(module) + and getattr(module, cond_attr).out_features + % 2 == 0): + cond_proj = getattr(module, cond_attr) + family, idx = _family_key(path) + # style width is a multiple of the model dim: 3x (scale, + # shift, gate) for RMS AdaLN, 2x (scale, shift) for LN + seams.append(Seam( + structure="adaln_producer", path=path, + parent_path=parent_path, norm_attr=None, + cond_attr=cond_attr, + dims={"C": cond_proj.in_features, + "S": cond_proj.out_features}, + variant={"bind": "table_only", "out_dtype": "bf16"}, + family=family, layer_index=idx)) + if "linear_proj" in structures: + for group in _ATTN_PROJ: + projs = [getattr(module, a, None) for a in group] + if not all(isinstance(p, nn.Linear) for p in projs): + continue + for attr, proj in zip(group, projs): + if proj.weight.numel() < _PROJ_WEIGHT_FLOOR: + continue + family, idx = _family_key(path) + seams.append(Seam( + structure="linear_proj", + path=path + "." + attr, parent_path=path, + norm_attr=None, proj_attr=attr, + dims={"K": proj.in_features, + "N": proj.out_features}, + variant={"bias": ("add" if proj.bias is not None + else "none"), + "epilogue": "none", "in_dtype": "bf16"}, + family=family + "." + attr, layer_index=idx)) + break + if "vision_ffn" in structures: + for fc1_attr, fc2_attr in _VISION_PROJ: + fc1 = getattr(module, fc1_attr, None) + fc2 = getattr(module, fc2_attr, None) + if not (isinstance(fc1, nn.Linear) + and isinstance(fc2, nn.Linear)): + continue + if (fc1.out_features != fc2.in_features + or fc1.in_features != fc2.out_features + or fc1.bias is None or fc2.bias is None): + continue + norm_attr = _norm_attr_of(model, parent_path) + if norm_attr is None: + continue # vision_ffn boundary includes a LayerNorm + act = _activation_of(module) or "gelu" + family, idx = _family_key(path) + seams.append(Seam( + structure="vision_ffn", path=path, + parent_path=parent_path, norm_attr=norm_attr, + dims={"D": fc1.in_features, "F": fc1.out_features}, + variant={"activation": act}, fc_attrs=(fc1_attr, fc2_attr), + family=family, layer_index=idx)) + break + return seams + + +def group_families(seams: list[Seam]) -> dict[str, list[Seam]]: + """Group seams into families (same template path), index-sorted.""" + families: dict[str, list[Seam]] = {} + for seam in seams: + families.setdefault(seam.family, []).append(seam) + for members in families.values(): + members.sort(key=lambda s: s.layer_index) + return families + + +def seam_weights(model: nn.Module, seam: Seam) -> dict[str, torch.Tensor]: + """Extract the impl-facing weight dict for one seam.""" + module = _resolve(model, seam.path) + norm = (_resolve(model, seam.parent_path + "." + seam.norm_attr) + if seam.norm_attr else None) + if seam.structure == "linear_proj": + return {"w": module.weight.detach(), + "b": (module.bias.detach() + if module.bias is not None else None)} + if seam.structure == "decoder_ffn": + w_norm = (norm.weight.detach() if norm is not None + and getattr(norm, "weight", None) is not None + else torch.ones(seam.dims["D"])) + return { + "w_norm": w_norm, + "w_gate": module.gate_proj.weight.detach().t().contiguous(), + "w_up": module.up_proj.weight.detach().t().contiguous(), + "w_down": module.down_proj.weight.detach().t().contiguous(), + } + fc1_attr, fc2_attr = seam.fc_attrs + fc1, fc2 = getattr(module, fc1_attr), getattr(module, fc2_attr) + return { + "w_norm": norm.weight.detach(), + "b_norm": norm.bias.detach(), + "w_fc1": fc1.weight.detach(), + "b_fc1": fc1.bias.detach(), + "w_fc2": fc2.weight.detach(), + "b_fc2": fc2.bias.detach(), + } diff --git a/flash_rt/structures/frontdoor.py b/flash_rt/structures/frontdoor.py new file mode 100644 index 00000000..cdbf53fc --- /dev/null +++ b/flash_rt/structures/frontdoor.py @@ -0,0 +1,327 @@ +"""One-call front door: ``structures.attach(model, forward, ...)``. + +The consumption contract mirrors ``kernels.get_kernel``: one import, one +call. Everything the structure layer needs — seam discovery, real +distribution calibration, per-layer parity gates, family-level accuracy +and net-win gates, transactional swap, receipt — runs inside the call. +An attachment that does not both stay accurate and win latency is +refused; the model is left untouched and the refusal is reported, never +silently absorbed. + +Calibration entry points (``calibration=``): + "auto" capture activations by running ``forward`` under hooks + ".pt" load a previously captured calibration file if it + exists, otherwise capture then save there +Pass ``forward`` as one callable (optionally repeated with ``frames=N``, +e.g. a closure advancing a dataloader) or a sequence of callables, one +per calibration frame. +""" + +from __future__ import annotations + +import hashlib +import json +import pathlib +import time +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, Sequence + +import torch + +from .swap import AttachHandle as _AttachHandle, attach as _swap_attach +from .discover import Seam, discover, group_families, seam_weights +from .gates import parity_metrics + + +def _bind_seam(seam: Seam, weights, calib, original): + if seam.structure == "decoder_ffn": + from .impls.decoder_ffn import fp8_static as impl + elif seam.structure == "vision_ffn": + from .impls.vision_ffn import fp8_static as impl + elif seam.structure == "linear_proj": + from .impls.linear_proj import fp8_static as proj_impl + return proj_impl.bind_proj_seam(weights, calibration=calib, + original=original) + else: + raise ValueError(f"no auto impl for structure {seam.structure!r}") + return impl.bind_mlp_seam(weights, variant=seam.variant, + calibration_normed=calib, original=original) + + +def _cuda_time_ms(fn: Callable[[], Any], warmup: int = 5, + iters: int = 20) -> float: + for _ in range(warmup): + fn() + if not torch.cuda.is_available(): + t0 = time.perf_counter() + for _ in range(iters): + fn() + return (time.perf_counter() - t0) * 1e3 / iters + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if torch.is_tensor(value): + return value + if isinstance(value, (tuple, list)): + for item in value: + found = _first_tensor(item) + if found is not None: + return found + if isinstance(value, Mapping): + for item in value.values(): + found = _first_tensor(item) + if found is not None: + return found + return None + + +@dataclass +class Plan: + """Result of one ``attach`` call: what was activated, why, evidence.""" + + activated: dict[str, torch.nn.Module] + families: dict[str, dict[str, Any]] + receipt: dict[str, Any] + _handle: _AttachHandle | None = None + active: bool = field(init=False) + + def __post_init__(self) -> None: + self.active = self._handle is not None + + def report(self) -> str: + lines = [f"structures.attach: {len(self.activated)} seam(s) active"] + for name, stat in self.families.items(): + lines.append( + f" {name}: {stat['layer_pass']}/{stat['layers']} layer-pass" + f" (worst cos {stat['gate_worst_cos']:.6f}) -> " + f"{stat['outcome']}" + + (f" [{stat.get('reason', '')}]" + if stat["outcome"] == "refused" else "")) + e2e = self.receipt.get("e2e") + if e2e: + lines.append( + f" e2e: {e2e['base_ms']:.2f} -> {e2e['ms']:.2f} ms " + f"({e2e['speedup']:.3f}x), cos={e2e.get('cos', float('nan')):.6f}") + return "\n".join(lines) + + def detach(self) -> None: + if self._handle is not None: + self._handle.detach() + self.active = False + + def save_receipt(self, directory: str | pathlib.Path) -> pathlib.Path: + directory = pathlib.Path(directory) + directory.mkdir(parents=True, exist_ok=True) + path = directory / f"attach_{self.receipt['digest'][:12]}.json" + path.write_text(json.dumps(self.receipt, indent=2, default=str)) + return path + + +def attach( + model: torch.nn.Module, + forward: Callable[[], Any] | Sequence[Callable[[], Any]], + *, + calibration: str = "auto", + frames: int = 1, + structures: tuple[str, ...] = ("decoder_ffn", "vision_ffn"), + gate_cos: float = 0.999, + e2e_budget: float = 0.995, + min_speedup: float = 1.02, + output: Callable[[], torch.Tensor] | None = None, + verbose: bool = True, +) -> Plan: + """Discover, calibrate, gate and activate structures in one call.""" + + def say(msg: str) -> None: + if verbose: + print(f"[structures] {msg}", flush=True) + + thunks = (list(forward) if isinstance(forward, (list, tuple)) + else [forward] * max(1, frames)) + seams = discover(model, structures) + families = group_families(seams) + say(f"discovered {len(seams)} seam(s) in {len(families)} family(ies)") + if not seams: + return Plan({}, {}, {"digest": "none", "seams": 0}) + + # ---- calibration: load cached captures or run forward under hooks ---- + cache = (pathlib.Path(calibration) + if calibration not in ("auto",) else None) + caps: dict[str, dict[str, list[torch.Tensor]]] + if cache is not None and cache.is_file(): + caps = torch.load(cache, weights_only=True) + say(f"calibration loaded from {cache} " + f"({len(next(iter(caps.values()))['in'])} frame(s))") + else: + caps = {s.path: {"x": [], "in": [], "out": []} for s in seams} + hooks = [] + + def mlp_hook(path): + def hook(module, args, out): + caps[path]["in"].append(args[0].detach().to("cpu")) + caps[path]["out"].append(out.detach().to("cpu")) + return hook + + def norm_hook(path): + def hook(module, args, out): + caps[path]["x"].append(args[0].detach().to("cpu")) + return hook + + from .discover import _resolve + for seam in seams: + hooks.append(_resolve(model, seam.path).register_forward_hook( + mlp_hook(seam.path))) + if seam.norm_attr: + hooks.append(_resolve( + model, seam.parent_path + "." + seam.norm_attr + ).register_forward_hook(norm_hook(seam.path))) + with torch.no_grad(): + for thunk in thunks: + thunk() + for hook in hooks: + hook.remove() + say(f"calibrated on {len(thunks)} frame(s)") + if cache is not None: + torch.save(caps, cache) + say(f"calibration saved to {cache}") + + device = next(model.parameters()).device + + # ---- bind + per-layer parity gate (structure boundary incl. residual) -- + fam_stats: dict[str, dict[str, Any]] = {} + fam_swaps: dict[str, dict[str, torch.nn.Module]] = {} + for name, members in families.items(): + swaps, refused, worst = {}, [], 1.0 + for seam in members: + cap = caps.get(seam.path, {"x": [], "in": [], "out": []}) + if not cap["in"]: + refused.append((seam.layer_index, "no calibration hits")) + continue + seam.m_profile = sorted( + {int(t.reshape(-1, t.shape[-1]).shape[0]) + for t in cap["in"]}) + from .discover import _resolve + try: + bound = _bind_seam(seam, seam_weights(model, seam), + [t for t in cap["in"]], + original=_resolve(model, seam.path)) + except ValueError as refusal: + refused.append((seam.layer_index, str(refusal)[:80])) + continue + layer_worst = 1.0 + with torch.no_grad(): + for k in range(len(cap["in"])): + xin = cap["in"][k].to(device) + got = bound(xin) + want = cap["out"][k].to(device) + if cap["x"]: + got = cap["x"][k].to(device) + got + want = cap["x"][k].to(device) + want + layer_worst = min( + layer_worst, parity_metrics(got, want)["cosine"]) + worst = min(worst, layer_worst) + if layer_worst >= gate_cos: + swaps[seam.path] = bound + else: + refused.append((seam.layer_index, round(layer_worst, 6))) + fam_swaps[name] = swaps + fam_stats[name] = { + "structure": members[0].structure, "dims": members[0].dims, + "variant": members[0].variant, "layers": len(members), + "layer_pass": len(swaps), "refused_layers": refused, + "gate_worst_cos": round(worst, 6), + "m_profile": sorted({m for s in members for m in s.m_profile}), + "outcome": "pending"} + say(f"{name}: {len(swaps)}/{len(members)} layer-pass " + f"@{gate_cos} (worst {worst:.6f})") + + # ---- family-level e2e gate: accuracy budget AND net win ---- + eval_thunk = thunks[-1] + get_out = output or (lambda: _first_tensor(eval_thunk())) + + def run_out() -> torch.Tensor | None: + with torch.no_grad(): + out = get_out() + return out.detach().float().cpu() if out is not None else None + + base_out = run_out() + base_ms = _cuda_time_ms(lambda: eval_thunk()) + say(f"baseline {base_ms:.2f} ms") + + active: dict[str, torch.nn.Module] = {} + for name, swaps in fam_swaps.items(): + stat = fam_stats[name] + if not swaps: + stat["outcome"], stat["reason"] = "refused", "no passing layers" + continue + handle = _swap_attach(model, swaps) + cos = (parity_metrics(run_out(), base_out)["cosine"] + if base_out is not None else None) + ms = _cuda_time_ms(lambda: eval_thunk()) + handle.detach() + wins = ((cos is None or cos >= e2e_budget) + and base_ms / ms >= min_speedup) + stat["e2e"] = {"cos": cos, "ms": round(ms, 3), + "speedup": round(base_ms / ms, 4)} + if wins: + stat["outcome"] = "activated" + active.update(swaps) + else: + stat["outcome"] = "refused" + stat["reason"] = (f"e2e cos {cos:.6f} < {e2e_budget}" + if cos is not None and cos < e2e_budget + else f"no net win ({base_ms / ms:.3f}x)") + say(f"{name}: e2e cos=" + f"{'n/a' if cos is None else f'{cos:.6f}'} " + f"{ms:.2f}ms ({base_ms / ms:.3f}x) -> {stat['outcome']}") + + # ---- union re-check, then commit ---- + e2e_final: dict[str, Any] | None = None + handle = None + if active: + handle = _swap_attach(model, active) + cos = (parity_metrics(run_out(), base_out)["cosine"] + if base_out is not None else None) + ms = _cuda_time_ms(lambda: eval_thunk()) + if ((cos is not None and cos < e2e_budget) + or base_ms / ms < min_speedup): + handle.detach() + handle, active = None, {} + for stat in fam_stats.values(): + if stat["outcome"] == "activated": + stat["outcome"] = "refused" + stat["reason"] = "union fails combined budget" + say("union of activated families fails combined budget -> " + "refuse all") + else: + e2e_final = {"base_ms": round(base_ms, 3), "ms": round(ms, 3), + "speedup": round(base_ms / ms, 4), + "cos": None if cos is None else round(cos, 7)} + say(f"active: {len(active)} seam(s), " + f"{base_ms:.2f} -> {ms:.2f} ms ({base_ms / ms:.3f}x)") + if not active: + say("outcome: whole-host refusal — model left untouched") + + receipt = { + "model": type(model).__name__, + "frames": len(thunks), + "calibration": calibration, + "eval_note": ("eval frame == calibration frame" + if len(thunks) == 1 else "last frame held for eval"), + "gate_cos": gate_cos, "e2e_budget": e2e_budget, + "min_speedup": min_speedup, + "families": fam_stats, "e2e": e2e_final, + "seams_active": sorted(active), + } + receipt["digest"] = hashlib.sha256( + json.dumps(receipt, sort_keys=True, default=str).encode() + ).hexdigest() + return Plan(active, fam_stats, receipt, _handle=handle) diff --git a/flash_rt/structures/gates.py b/flash_rt/structures/gates.py new file mode 100644 index 00000000..baae8c4d --- /dev/null +++ b/flash_rt/structures/gates.py @@ -0,0 +1,204 @@ +"""Qualification gates — parity judgment against a structure's reference. + +The harness is structure-agnostic. Implementations follow the structure +calling convention: required boundary inputs in declared order, then +weight tensors in slot order, then variant selections and any optional +boundary inputs as keyword arguments — the same signature the reference +implementation exposes. + +A qualification produces a machine-readable record whose ``plan_digest`` +binds the spec content, variant, resolved workload dims, thresholds, +implementation identity, and environment. A record certifies exactly one +execution plan; change any component and the record no longer applies. +""" + +from __future__ import annotations + +import hashlib +import json +import pathlib +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping + +import torch + +from flash_rt.structures.registry import StructureSpec, _CATALOG_DIR + + +@dataclass(frozen=True) +class QualificationCase: + """One workload to qualify: boundary inputs, weights, and variant.""" + + inputs: Mapping[str, torch.Tensor] + weights: Mapping[str, torch.Tensor] + variant: Mapping[str, str] = field(default_factory=dict) + + +def solve_dims( + spec: StructureSpec, + inputs: Mapping[str, torch.Tensor], + weights: Mapping[str, torch.Tensor], +) -> dict[str, int]: + """Resolve symbolic dims from actual tensors, rejecting inconsistency.""" + dims: dict[str, int] = {} + + def bind(declared: list[str], tensor: torch.Tensor, what: str) -> None: + if tensor.ndim != len(declared): + raise ValueError( + f"{what}: expected rank {len(declared)} {declared}, " + f"got shape {tuple(tensor.shape)}" + ) + for name, size in zip(declared, tensor.shape): + if dims.setdefault(name, int(size)) != int(size): + raise ValueError( + f"{what}: dim {name}={int(size)} conflicts with " + f"{name}={dims[name]} resolved earlier" + ) + + for entry in spec.boundary["inputs"]: + tensor = inputs.get(entry["name"]) + if tensor is None: + if entry.get("optional", False): + continue + raise ValueError(f"missing required input: {entry['name']!r}") + bind(entry["dims"], tensor, f"input {entry['name']!r}") + for entry in spec.weights: + slot = entry["slot"] + if slot not in weights: + raise ValueError(f"missing weight slot: {slot!r}") + bind(entry["dims"], weights[slot], f"weight {slot!r}") + return dims + + +def _call(spec: StructureSpec, fn: Callable[..., Any], + case: QualificationCase, *, bound: bool = False) -> torch.Tensor: + """Invoke ``fn`` per the structure calling convention. + + ``bound=False`` targets full-signature callables (the reference): + required inputs, then weight slots, with variants and optional inputs + as keywords. ``bound=True`` targets bound implementations whose + weights and variant were baked in at bind time: inputs only. + """ + args: list[torch.Tensor] = [] + kwargs: dict[str, Any] = {} if bound else dict(case.variant) + for entry in spec.boundary["inputs"]: + tensor = case.inputs.get(entry["name"]) + if entry.get("optional", False): + if tensor is not None: + kwargs[entry["name"]] = tensor + else: + args.append(tensor) + if not bound: + args.extend(case.weights[slot] for slot in spec.weight_slots) + return fn(*args, **kwargs) + + +def parity_metrics(got: torch.Tensor, want: torch.Tensor) -> dict[str, float]: + """Cosine / max-abs / p99-abs between an implementation and a truth.""" + return _parity_metrics(got, want) + + +def _parity_metrics(got: torch.Tensor, want: torch.Tensor) -> dict[str, float]: + if got.shape != want.shape: + raise ValueError( + f"output shape mismatch: impl {tuple(got.shape)} vs " + f"reference {tuple(want.shape)}" + ) + diff = (got.double() - want.double()).abs().flatten() + cosine = torch.nn.functional.cosine_similarity( + got.double().flatten(), want.double().flatten(), dim=0 + ) + # kthvalue instead of quantile: exact and free of quantile's input + # size limit (qualification outputs can exceed it, e.g. LLM logits) + k = max(1, int(0.99 * diff.numel())) + return { + "cosine": float(cosine), + "max_abs": float(diff.max()), + "p99_abs": float(diff.kthvalue(k).values), + } + + +def _spec_digest(spec: StructureSpec) -> str: + path = _CATALOG_DIR / spec.name / "structure.yaml" + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _environment() -> dict[str, str]: + env = {"torch": torch.__version__} + if torch.cuda.is_available(): + env["device"] = torch.cuda.get_device_name(0) + major, minor = torch.cuda.get_device_capability(0) + env["arch"] = f"sm_{major}{minor}" + else: + env["device"] = "cpu" + return env + + +def qualify_parity( + spec: StructureSpec, + impl: Callable[..., Any], + case: QualificationCase, + *, + impl_id: str, + thresholds: Mapping[str, float], + bound: bool = False, +) -> dict[str, Any]: + """Judge ``impl`` against the structure reference on one workload. + + ``thresholds`` maps metric name to its passing bound (``cosine`` is a + floor, absolute-error metrics are ceilings). Every thresholded metric + must pass for a PASS verdict; metrics without thresholds are recorded + as evidence only. Set ``bound=True`` when ``impl`` was produced by an + implementation's ``bind`` and takes boundary inputs only. + """ + for key in case.variant: + if key not in spec.variants: + raise ValueError(f"unknown variant key: {key!r}") + workload = solve_dims(spec, case.inputs, case.weights) + reference = spec.reference() + + want = _call(spec, reference, case) + got = _call(spec, impl, case, bound=bound) + metrics = _parity_metrics(got, want) + + passed = True + for name, bound in thresholds.items(): + if name not in metrics: + raise ValueError(f"threshold on unknown metric: {name!r}") + ok = metrics[name] >= bound if name == "cosine" else metrics[name] <= bound + passed = passed and ok + + x_dtype = next(iter(case.inputs.values())).dtype + record = { + "structure": f"{spec.name}@{spec.version}", + "spec_digest": _spec_digest(spec), + "impl": impl_id, + "variant": dict(case.variant), + "workload": {**workload, "dtype": str(x_dtype)}, + "env": _environment(), + "gate": "parity", + "metrics": metrics, + "thresholds": dict(thresholds), + "verdict": "PASS" if passed else "FAIL", + } + record["plan_digest"] = "sha256:" + hashlib.sha256( + json.dumps( + {k: record[k] for k in + ("structure", "spec_digest", "impl", "variant", "workload", + "env", "thresholds")}, + sort_keys=True, + ).encode("utf-8") + ).hexdigest() + return record + + +def save_record(record: Mapping[str, Any], directory: str | pathlib.Path) -> pathlib.Path: + """Write one qualification record as JSON, named by its plan digest.""" + directory = pathlib.Path(directory) + directory.mkdir(parents=True, exist_ok=True) + digest = record["plan_digest"].split(":", 1)[1][:16] + path = directory / f"{record['gate']}_{digest}.json" + with open(path, "w", encoding="utf-8") as handle: + json.dump(record, handle, indent=2, sort_keys=True) + handle.write("\n") + return path diff --git a/flash_rt/structures/handle.py b/flash_rt/structures/handle.py new file mode 100644 index 00000000..471d1ae7 --- /dev/null +++ b/flash_rt/structures/handle.py @@ -0,0 +1,160 @@ +"""Explicit front door: ``structures.get(name)`` — pull one structure, +bind it to one host module, plug it in yourself. + +This is the single-site, fully visible counterpart of ``attach``. The +consumption mirrors ``kernels.get_kernel``: + + ffn = structures.get("decoder_ffn") + new_mlp = ffn.bind(model.model.layers[3].mlp, calibration=[x1, x2]) + model.model.layers[3].mlp = new_mlp # your swap, your call + +``bind`` extracts the weights from the module you hand it, calibrates on +the samples you provide, builds the fused replacement, and self-checks +it against the original module on those same samples. A replacement +that misses the parity gate raises ``GateRefused`` instead of returning +— the explicit path fails loudly, it never hands back a bad part. +``attach`` is this same operation, batched over every discovered site +with the net-win A/B added. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from .discover import _VISION_PROJ, _activation_of +from .gates import parity_metrics +from .registry import StructureSpec, load + + +class GateRefused(RuntimeError): + """The bound replacement did not meet the parity gate.""" + + +class StructureHandle: + """One pulled structure, ready to bind to host modules.""" + + def __init__(self, spec: StructureSpec): + self.spec = spec + self.name = spec.name + + def __repr__(self) -> str: + return f"StructureHandle({self.name!r}, v{self.spec.version})" + + def bind( + self, + module: torch.nn.Module, + calibration: torch.Tensor | Sequence[torch.Tensor], + *, + variant: dict[str, str] | None = None, + norm: torch.nn.Module | None = None, + residual: torch.Tensor | Sequence[torch.Tensor] | None = None, + gate_cos: float = 0.999, + ) -> torch.nn.Module: + """Build a gated drop-in replacement for ``module``. + + ``calibration``: one or more real inputs of ``module`` (the + normed hidden states it actually sees). ``residual``: the + matching pre-norm hidden states — the structure gate is declared + at the boundary *including* the residual add, so provide these + to gate at the declared boundary; without them the self-check + runs on the bare seam output, which is strictly harsher. + ``norm``: optional sibling norm module, recorded for + full-structure use; the MLP seam itself leaves the norm in the + host. ``gate_cos=0`` skips the self-check (not recommended). + """ + samples = ([calibration] if torch.is_tensor(calibration) + else list(calibration)) + if not samples: + raise ValueError("calibration must contain at least one sample") + residuals = (None if residual is None else + [residual] if torch.is_tensor(residual) + else list(residual)) + if residuals is not None and len(residuals) != len(samples): + raise ValueError("residual must align 1:1 with calibration") + + if self.name == "decoder_ffn": + from .impls.decoder_ffn import fp8_static as impl + gate = module.gate_proj + device = gate.weight.device + dims = {"D": gate.in_features, "F": gate.out_features} + var = {"activation": _activation_of(module) or "silu", + "norm_weight_mode": "direct"} + weights = { + "w_norm": (norm.weight.detach() if norm is not None + and getattr(norm, "weight", None) is not None + else torch.ones(dims["D"], device=device)), + "w_gate": gate.weight.detach().t().contiguous(), + "w_up": module.up_proj.weight.detach().t().contiguous(), + "w_down": module.down_proj.weight.detach().t().contiguous(), + } + elif self.name == "vision_ffn": + from .impls.vision_ffn import fp8_static as impl + fc_attrs = next( + (pair for pair in _VISION_PROJ + if all(isinstance(getattr(module, a, None), torch.nn.Linear) + for a in pair)), None) + if fc_attrs is None: + raise ValueError( + "module has no fc1/fc2 (or linear_fc1/linear_fc2) pair") + fc1 = getattr(module, fc_attrs[0]) + fc2 = getattr(module, fc_attrs[1]) + device = fc1.weight.device + dims = {"D": fc1.in_features, "F": fc1.out_features} + var = {"activation": _activation_of(module) or "gelu"} + weights = { + "w_norm": (norm.weight.detach() if norm is not None + else torch.ones(dims["D"], device=device)), + "b_norm": (norm.bias.detach() if norm is not None + else torch.zeros(dims["D"], device=device)), + "w_fc1": fc1.weight.detach(), "b_fc1": fc1.bias.detach(), + "w_fc2": fc2.weight.detach(), "b_fc2": fc2.bias.detach(), + } + else: + raise ValueError( + f"structure {self.name!r} has no module-seam bind; " + "stage pipelines go through the capture door") + var.update(variant or {}) + + bound = impl.bind_mlp_seam(weights, variant=var, + calibration_normed=samples, + original=module) + + worst = 1.0 + with torch.no_grad(): + for k, sample in enumerate(samples): + x = sample.to(device) + got, want = bound(x), module(x) + if residuals is not None: + res = residuals[k].to(device) + got, want = res + got, res + want + worst = min(worst, parity_metrics(got, want)["cosine"]) + boundary = ("structure (incl. residual)" if residuals is not None + else "bare seam") + if gate_cos and worst < gate_cos: + raise GateRefused( + f"{self.name} replacement worst cos {worst:.6f} < " + f"{gate_cos} at {boundary} boundary on {len(samples)} " + "calibration sample(s) — not handing back a part that " + "fails its own gate" + + ("" if residuals is not None else + "; pass residual= to gate at the declared boundary")) + bound.certification = { + "structure": self.name, "version": self.spec.version, + "dims": dims, "variant": var, "worst_cos": round(worst, 7), + "gate_boundary": boundary, "samples": len(samples), + "m_profile": sorted({ + int(s.reshape(-1, s.shape[-1]).shape[0]) for s in samples}), + } + return bound + + +def get(name: str) -> StructureHandle: + """Pull one structure from the catalog (get_kernel-style).""" + spec = load(name) + if spec.kind != "region": + raise ValueError( + f"{name!r} is a {spec.kind} structure — module-level bind " + "does not apply; use the capture/provider route") + return StructureHandle(spec) diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py new file mode 100644 index 00000000..da9d2ec8 --- /dev/null +++ b/flash_rt/structures/impls/__init__.py @@ -0,0 +1,16 @@ +"""Structure implementations. + +``hub_kernel`` is the shared, process-wide hub loader: two impls that +depend on the same kernel repo must share one loaded module — a second +``kernels.get_kernel`` import of the same repo re-registers its fake +ops and torch.library raises. +""" + +from functools import lru_cache + + +@lru_cache(maxsize=None) +def hub_kernel(repo: str, version: str): + from kernels import get_kernel + + return get_kernel(repo, version=version) diff --git a/flash_rt/structures/impls/adaln_producer/__init__.py b/flash_rt/structures/impls/adaln_producer/__init__.py new file mode 100644 index 00000000..7a9a30d6 --- /dev/null +++ b/flash_rt/structures/impls/adaln_producer/__init__.py @@ -0,0 +1,7 @@ +from .fused import (AdaLNProducer, StepLocator, StyleTable, + bind_adaln_producer, bind_step_locator, + bind_style_table) + +__all__ = ["AdaLNProducer", "StepLocator", "StyleTable", + "bind_adaln_producer", "bind_step_locator", + "bind_style_table"] diff --git a/flash_rt/structures/impls/adaln_producer/fused.py b/flash_rt/structures/impls/adaln_producer/fused.py new file mode 100644 index 00000000..afbaea63 --- /dev/null +++ b/flash_rt/structures/impls/adaln_producer/fused.py @@ -0,0 +1,224 @@ +"""adaln_producer — conditioning-driven norm, resolved per step. + +Diffusion hosts modulate every layer with a projection of the current +timestep embedding. Two facts make this a structure rather than a plain +GEMM: the conditioning vector takes one of a small fixed set of values +over a tick (it is a function of the step), and the norm that consumes +it can fuse modulation and output quantization into one kernel. + +This implementation splits those two concerns: + +- :class:`StepLocator` resolves "which step is this" from the + conditioning tensor using a few high-separation coordinates — a + fingerprint — instead of a full-width matmul against every stored + vector. It is pure tensor work (index_select, squared distance, + argmax), so it traces under dynamo and captures into a graph without + host-side state. Sibling producers fed by the same conditioning + stream share one locator, letting the compiler fold the repeated + lookups. +- :class:`AdaLNProducer` replaces the host's adaptive norm: it looks up + the precomputed style row for the current step and runs the fused + norm+modulate(+static FP8 quantize) kernel, emitting either BF16 or + FP8 plus the host's gate. The FP8 form is the upstream half of a + producer→consumer seam: the shared ``act_scale`` lets a packed + projection skip its own input quantization. + +Qualification: the conditioning must actually be step-quantized (more +distinct vectors than ``max_steps`` means it depends on more than the +step, and a table would alias different inputs onto one row), and the +chosen fingerprint coordinates must separate the stored vectors by a +real margin. Either failure raises ``ValueError`` — the host keeps its +own producer. +""" + +from __future__ import annotations + +import torch + +from .. import hub_kernel + + +def _dedup(pairs, max_steps, rtol): + conds, outs = [], [] + for cond, out in pairs: + c = cond.detach().reshape(-1, cond.shape[-1]) + o = out.detach().reshape(-1, out.shape[-1]) + for row in range(c.shape[0]): + cr = c[row] + if any(torch.allclose(cr, seen, rtol=rtol, + atol=1e-6 * cr.abs().max().item() + 1e-12) + for seen in conds): + continue + conds.append(cr.clone()) + outs.append(o[row].clone()) + if len(conds) > max_steps: + raise ValueError( + f"adaln_producer: >{max_steps} distinct conditioning " + "vectors — not step-quantized, keeping the host path") + if not conds: + raise ValueError("adaln_producer: no calibration pairs") + return torch.stack(conds), torch.stack(outs) + + +class StepLocator(torch.nn.Module): + """Resolve the current step index from the conditioning tensor.""" + + def __init__(self, conds: torch.Tensor, n_dims: int = 8, + rel_margin: float = 1e-3): + super().__init__() + c = conds.float() + steps = c.shape[0] + if steps == 1: + dims = torch.zeros(1, dtype=torch.long, device=c.device) + else: + diffs = (c.unsqueeze(0) - c.unsqueeze(1)).abs() + eye = torch.eye(steps, device=c.device, dtype=torch.bool) + diffs = diffs.masked_fill(eye.unsqueeze(-1), float("inf")) + minsep = diffs.amin(dim=(0, 1)) + k = min(n_dims, c.shape[1]) + dims = minsep.topk(k).indices.sort().values + margin = minsep[dims].min().item() + if margin < rel_margin * c.abs().mean().item(): + raise ValueError( + "adaln_producer: conditioning vectors are not " + "separable on any coordinate subset") + self.register_buffer("fp_dims", dims) + self.register_buffer("fp_conds", c.index_select( + 1, dims).contiguous()) + + def forward(self, cond: torch.Tensor) -> torch.Tensor: + flat = cond.reshape(-1, cond.shape[-1]).float() + c = flat.index_select(1, self.fp_dims) + scores = -((c.unsqueeze(1) - self.fp_conds) ** 2).sum(-1) + return scores.argmax(-1) + + +class StyleTable(torch.nn.Module): + """Replace only the conditioning projection with its step table. + + The narrower of the two bind forms: the host keeps its own norm + (often already fused by the compiler) and only the per-step style + projection is memoized. Prefer this wherever the norm itself is not + being upgraded — measurement decides, and the fused form is worth + its kernel only when it also serves a downstream seam. + """ + + def __init__(self, host_proj: torch.nn.Module, styles: torch.Tensor, + locator: StepLocator): + super().__init__() + self.host_linear = host_proj + self.locator = locator + self.register_buffer("table", styles.contiguous()) + + def forward(self, cond: torch.Tensor) -> torch.Tensor: + out = self.table.index_select(0, self.locator(cond)) + return out.reshape(*cond.shape[:-1], out.shape[-1]) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_linear"), name) + + +class AdaLNProducer(torch.nn.Module): + """Adaptive norm replacement: step lookup + fused norm/quantize.""" + + def __init__(self, host_norm: torch.nn.Module, + styles: torch.Tensor, locator: StepLocator, + act_scale: torch.Tensor | None, rows: int, dim: int, + norm: str = "rms"): + super().__init__() + self.host_norm = host_norm + self.locator = locator + self.norm = norm + self.register_buffer("styles", + styles.to(torch.bfloat16).contiguous()) + self.out_fp8 = act_scale is not None + dev = styles.device + if norm == "layer": + # LayerNorm hosts (DiT AdaLayerNorm): style is (scale, + # shift), no gate, and the fused kernel takes the raw + # scale — it applies the (1 + scale) itself. + if not self.out_fp8: + raise ValueError("adaln_producer: layer norm form " + "currently requires fp8 output") + kq = hub_kernel("flashrt/adaptive-layernorm-producers", ">=1") + self._fn = kq.ada_layer_norm_quant_fp8_bf16 + self.register_buffer("act_scale", act_scale) + else: + ka = hub_kernel("flashrt/flashrt-adaptive-norms", ">=1") + if self.out_fp8: + self._fn = ka.gate_residual_ada_norm_fp8_static_bf16 + self.register_buffer("act_scale", act_scale) + else: + self._fn = ka.ada_rms_norm_style_bf16 + # residual=0 / gate=1 turn the gated-residual kernel into a + # plain modulated norm; both are preallocated for graph replay + self.register_buffer("w_ones", torch.ones( + dim, device=dev, dtype=torch.bfloat16)) + self.register_buffer("resid", torch.zeros( + rows, dim, device=dev, dtype=torch.bfloat16)) + self.register_buffer("gate_ones", torch.ones( + rows, dim, device=dev, dtype=torch.bfloat16)) + + def forward(self, x: torch.Tensor, cond: torch.Tensor | None = None): + idx = self.locator(cond) + style = self.styles.index_select(0, idx) + if self.norm == "layer": + scale, shift = style[0].chunk(2, dim=-1) + y = self._fn( + x.reshape(-1, x.shape[-1]).to(torch.bfloat16) + .contiguous(), + scale.contiguous(), shift.contiguous(), + self.act_scale) + return y.reshape(x.shape) + rows = self.resid.shape[0] + style2d = style.expand(rows, -1).contiguous() + x2d = x.reshape(-1, x.shape[-1]) + if self.out_fp8: + _, y, gate = self._fn(self.resid, x2d, self.gate_ones, + self.w_ones, style2d, self.act_scale) + else: + y, gate = self._fn(x2d, self.w_ones, style2d) + return (y.reshape(x.shape), + gate[:1].reshape(1, 1, gate.shape[-1]).to(x.dtype)) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_norm"), name) + + +def bind_step_locator(pairs, *, max_steps: int = 64, + dedup_rtol: float = 1e-5, n_dims: int = 8): + """Build a locator and the step table from ``(cond, out)`` pairs.""" + conds, outs = _dedup(pairs, max_steps, dedup_rtol) + return StepLocator(conds, n_dims=n_dims), outs + + +def bind_style_table(host_proj: torch.nn.Module, pairs, *, + locator: StepLocator | None = None, + max_steps: int = 64) -> StyleTable: + """Bind the table-only form onto the conditioning projection.""" + built, styles = bind_step_locator(pairs, max_steps=max_steps) + return StyleTable(host_proj, styles, locator or built) + + +def bind_adaln_producer(host_norm: torch.nn.Module, pairs, *, + act_scale: torch.Tensor | None = None, + rows: int, dim: int, + locator: StepLocator | None = None, + max_steps: int = 64, norm: str = "rms"): + """Bind an adaptive-norm producer from real ``(cond, style)`` pairs. + + ``pairs`` come from hooking the host's own conditioning projection + over at least one full tick, so the stored style rows are exactly + what the host computed. Pass ``act_scale`` to emit FP8 for a + downstream packed projection; pass ``locator`` to share the step + lookup with sibling producers on the same conditioning stream. + """ + built, styles = bind_step_locator(pairs, max_steps=max_steps) + return AdaLNProducer(host_norm, styles, locator or built, + act_scale, rows, dim, norm=norm) diff --git a/flash_rt/structures/impls/attention_core/__init__.py b/flash_rt/structures/impls/attention_core/__init__.py new file mode 100644 index 00000000..dd3cb0c5 --- /dev/null +++ b/flash_rt/structures/impls/attention_core/__init__.py @@ -0,0 +1,5 @@ +from .fa2_seqused import (PackedKVAttention, SUPPORTED_HEAD_DIMS, + bind_attention_core, plan_packed_kv) + +__all__ = ["PackedKVAttention", "SUPPORTED_HEAD_DIMS", + "bind_attention_core", "plan_packed_kv"] diff --git a/flash_rt/structures/impls/attention_core/fa2_seqused.py b/flash_rt/structures/impls/attention_core/fa2_seqused.py new file mode 100644 index 00000000..3391c886 --- /dev/null +++ b/flash_rt/structures/impls/attention_core/fa2_seqused.py @@ -0,0 +1,193 @@ +"""attention_core — run attention on the FlashRT FA2 kernel. + +A fused attention kernel wants contiguous keys and values; hosts hand +it an additive mask instead, and a mask that blocks a run of positions +in the middle of the sequence has no equivalent the kernel accepts. The +structure resolves that at bind time rather than per call: it reads the +host's own mask, and when the blocked positions form one contiguous run +it packs the surviving keys and values into a compact buffer. What is +left is a dense attention the kernel runs directly. + +The packed prefix (everything before the blocked run) belongs to a +slower cadence — it is the encoder's output for the current +observation, unchanged across the denoise loop — so it is filled once +at bind time and refreshed through an update function, exactly as +:mod:`..cadence_static` does for whole modules. The suffix is written +per call. + +Three qualifications, all decided from real captures: + +- ``head_dim`` must be one the kernel supports; otherwise the caller + keeps its own path (``bind_attention_core`` returns ``None`` so the + host can fall back to a community kernel rather than fail), +- every query row must see the same mask, and the blocked positions + must form a single contiguous run — anything else is not expressible + as a packed dense attention, +- the prefix keys must not move across the loop; a moving prefix means + the split is wrong and the parity gate would catch it downstream, so + it is rejected here where the reason is still legible. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from .. import hub_kernel + +SUPPORTED_HEAD_DIMS = (64, 96, 128, 256) + + +@dataclass +class _Scratch: + """Output/LSE/workspace shared by same-shaped attention sites.""" + + out: torch.Tensor + lse: torch.Tensor + workspace: object + + +@dataclass +class PackedKVPlan: + """How the host's masked attention maps onto a dense one.""" + + prefix: int + suffix_start: int + suffix_len: int + seq_kv: int + + @property + def packed(self) -> int: + return self.prefix + self.suffix_len + + +def plan_packed_kv(mask: torch.Tensor | None, seq_kv: int) -> PackedKVPlan: + """Derive the packing plan from one captured attention mask.""" + if mask is None: + return PackedKVPlan(seq_kv, seq_kv, 0, seq_kv) + if mask.dim() < 3: + raise ValueError("attention_core: unexpected mask rank") + rows = mask.reshape(-1, mask.shape[-2], mask.shape[-1])[0] + if not bool((rows == rows[0]).all()): + raise ValueError("attention_core: mask differs per query row") + row = rows[0].float() + blocked = ((row < -1e5) | row.isneginf()).nonzero().flatten() + if blocked.numel() == 0: + return PackedKVPlan(seq_kv, seq_kv, 0, seq_kv) + runs, start, prev = [], None, None + for i in blocked.tolist(): + if start is None: + start = prev = i + elif i == prev + 1: + prev = i + else: + runs.append((start, prev)) + start = prev = i + runs.append((start, prev)) + if len(runs) != 1: + raise ValueError( + f"attention_core: mask blocks {len(runs)} separate runs — " + "not expressible as one packed dense attention") + lo, hi = runs[0] + return PackedKVPlan(lo, hi + 1, seq_kv - (hi + 1), seq_kv) + + +class PackedKVAttention(torch.nn.Module): + """Attention over packed keys/values, on the FlashRT FA2 kernel. + + Holds one host module's packed buffers. Call it in place of the + host's attention body; the prefix half is refreshed by the update + function returned from :func:`bind_attention_core`. + """ + + def __init__(self, plan: PackedKVPlan, q_shape, kv_heads: int, + dtype: torch.dtype, device, prefix_kv=None, + scratch: "_Scratch | None" = None): + super().__init__() + self.plan = plan + b, heads, seq_q, head_dim = q_shape + self.seq_q = seq_q + self._kfa = hub_kernel("flashrt/fa2-seqused-runtime", ">=1") + self.register_buffer("packed_k", torch.zeros( + b, plan.packed, kv_heads, head_dim, device=device, + dtype=dtype)) + self.register_buffer("packed_v", torch.zeros_like(self.packed_k)) + if prefix_kv is not None: + k0, v0 = prefix_kv + self.packed_k[:, :plan.prefix] = k0 + self.packed_v[:, :plan.prefix] = v0 + # Output, LSE and split-KV workspace are scratch: each site + # consumes its result before the next one runs, so sites with + # the same shapes share one set. Per-site copies cost real + # latency (pi05 r16: 18 private scratches ran 1.5% slower than + # one shared set) and buy nothing. + if scratch is None: + q_sample = torch.empty(b, seq_q, heads, head_dim, + device=device, dtype=dtype) + out, lse = self._kfa.allocate_outputs(q_sample) + scratch = _Scratch(out, lse, self._kfa.allocate_workspace( + q_sample, self.packed_k)) + self._scratch = scratch + + def forward(self, query, key, value, *, scale=None): + """``query``/``key``/``value`` in the host's (B, H, S, D).""" + plan = self.plan + q = query.transpose(1, 2).contiguous() + if plan.suffix_len: + self.packed_k[:, plan.prefix:].copy_( + key.transpose(1, 2)[:, plan.suffix_start:]) + self.packed_v[:, plan.prefix:].copy_( + value.transpose(1, 2)[:, plan.suffix_start:]) + sc = self._scratch + return self._kfa.forward_static( + q, self.packed_k, self.packed_v, out=sc.out, + softmax_lse=sc.lse, workspace=sc.workspace, + softmax_scale=scale) + + +def bind_attention_core(captures, *, prefix_static_rtol: float = 1e-3): + """Bind one packed-KV attention per site from real captures. + + ``captures`` is a sequence of per-site dicts holding ``q`` (one + captured query, host layout), ``keys``/``values`` (the tensors that + site produced across the hot loop, oldest first) and ``mask``. + Returns ``(modules, update)``, or ``None`` when the head dim is + unsupported so the caller can keep a fallback path. + """ + if not captures: + raise ValueError("attention_core: no captures") + head_dim = captures[0]["q"].shape[-1] + if head_dim not in SUPPORTED_HEAD_DIMS: + return None + + modules, scratch = [], None + for site, cap in enumerate(captures): + keys = cap["keys"] + plan = plan_packed_kv(cap.get("mask"), keys[0].shape[2]) + first = keys[0][:, :, :plan.prefix] + for other in keys[1:]: + if not torch.allclose(first, other[:, :, :plan.prefix], + rtol=prefix_static_rtol, + atol=prefix_static_rtol): + raise ValueError( + f"attention_core: site {site} prefix keys move " + "across the loop — the cadence split is wrong") + q = cap["q"] + core = PackedKVAttention( + plan, q.shape, keys[0].shape[1], q.dtype, q.device, + prefix_kv=(keys[0].transpose(1, 2)[:, :plan.prefix], + cap["values"][0].transpose(1, 2)[:, :plan.prefix]), + scratch=scratch) + scratch = scratch or core._scratch + modules.append(core) + + def update(fresh_kv) -> None: + """Refresh every site's prefix from freshly computed K/V.""" + with torch.no_grad(): + for mod, (k, v) in zip(modules, fresh_kv): + p = mod.plan.prefix + mod.packed_k[:, :p].copy_(k.transpose(1, 2)[:, :p]) + mod.packed_v[:, :p].copy_(v.transpose(1, 2)[:, :p]) + + return modules, update diff --git a/flash_rt/structures/impls/cadence_static/__init__.py b/flash_rt/structures/impls/cadence_static/__init__.py new file mode 100644 index 00000000..b44f496c --- /dev/null +++ b/flash_rt/structures/impls/cadence_static/__init__.py @@ -0,0 +1,3 @@ +from .buffers import StaticOutput, bind_cadence_static + +__all__ = ["StaticOutput", "bind_cadence_static"] diff --git a/flash_rt/structures/impls/cadence_static/buffers.py b/flash_rt/structures/impls/cadence_static/buffers.py new file mode 100644 index 00000000..c628c096 --- /dev/null +++ b/flash_rt/structures/impls/cadence_static/buffers.py @@ -0,0 +1,91 @@ +"""cadence_static — hold work that changes slower than the hot loop. + +A tick pipeline usually runs several cadences at once: a denoise loop +that repeats every tick, and encoder-side work that only changes when a +new observation or prompt arrives. Modules on the slower cadence still +sit inside the hot path, so a graph captures them and pays for them +every tick even though their inputs are unchanged. + +This structure moves such a module out of the loop: its output becomes +a static buffer the captured graph reads, and the real computation runs +in an update function the host calls at the module's own cadence. The +resulting split is explicit — the update callable is returned to the +caller so it can be registered as a recipe ``outside_update`` rather +than hidden inside the replacement. + +Qualification is empirical: the wrapped output must actually be +constant across the fast loop. Callers pass calibration captures from +several iterations of the hot loop and binding refuses when they +disagree — a module whose output moves per step is not a cadence +substructure, and freezing it would silently change the model. +""" + +from __future__ import annotations + +from typing import Callable, Sequence + +import torch + + +class StaticOutput(torch.nn.Module): + """Return a buffer the host refreshes at the slower cadence.""" + + def __init__(self, original: torch.nn.Module, value: torch.Tensor): + super().__init__() + self.host_module = original + self.register_buffer("buffer", value.contiguous().clone()) + + def forward(self, *args, **kwargs): + return self.buffer + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_module"), name) + + +def bind_cadence_static( + modules: Sequence[torch.nn.Module], + captures: Sequence[Sequence[torch.Tensor]], + *, + recompute: Callable[[], Sequence[torch.Tensor]] | None = None, + rtol: float = 1e-3, + atol: float = 1e-3, +): + """Freeze module outputs into buffers plus one update function. + + ``captures[i]`` holds the outputs module ``i`` produced across + several iterations of the hot loop; they must agree, or the module + is not on a slower cadence and binding refuses. ``recompute`` + returns fresh outputs for the same modules — typically by running + the host's own encoder path — and the returned update callable + copies them into the buffers. Register that callable as the + recipe's outside update so the split cadence stays explicit and + the slower work is still timed. + """ + if len(modules) != len(captures): + raise ValueError("cadence_static: modules/captures mismatch") + statics = [] + for i, (mod, caps) in enumerate(zip(modules, captures)): + if not caps: + raise ValueError(f"cadence_static: module {i} has no captures") + first = caps[0] + for other in caps[1:]: + if not torch.allclose(first, other, rtol=rtol, atol=atol): + raise ValueError( + f"cadence_static: module {i} output varies within " + "the hot loop — not a cadence substructure") + statics.append(StaticOutput(mod, first)) + + def update() -> None: + if recompute is None: + raise RuntimeError( + "cadence_static: no recompute function was supplied, so " + "the buffers cannot be refreshed for a new observation") + with torch.no_grad(): + fresh = recompute() + for static, value in zip(statics, fresh): + static.buffer.copy_(value) + + return statics, update diff --git a/flash_rt/structures/impls/decoder_ffn/__init__.py b/flash_rt/structures/impls/decoder_ffn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/impls/decoder_ffn/fp8_static.py b/flash_rt/structures/impls/decoder_ffn/fp8_static.py new file mode 100644 index 00000000..61ed06da --- /dev/null +++ b/flash_rt/structures/impls/decoder_ffn/fp8_static.py @@ -0,0 +1,265 @@ +"""FP8-static implementation of the ``decoder_ffn`` structure. + +Composes the fused FP8 gate/up -> activation -> down block from the +``flashrt/flashrt-fp8-swiglu-ffn`` Hub kernel behind the structure +boundary. Two bind entrypoints share the packing and calibration code: +``bind`` covers the full structure (norm and AdaLN modulation run in +torch ahead of the fused block); ``bind_mlp_seam`` covers the +normed-input -> ffn-output slice for hosts whose replaceable module +boundary is the MLP. Activation scales are static per-tensor, +calibrated from caller-provided representative inputs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, Mapping, Sequence + +import torch + +KERNEL_DEP = { + "provider": "hf", + "repo": "flashrt/flashrt-fp8-swiglu-ffn", + "version": ">=1", +} + +_FP8 = torch.float8_e4m3fn +_FP8_MAX = 448.0 +_ENTRYPOINTS = {"gelu": "bf16_fp8_geglu_mlp_bf16", + "silu": "bf16_fp8_swiglu_mlp_bf16"} + +SUPPORT = { + "D": {"min": 512, "max": 16384}, + "F": {"min": 1024, "max": 16384}, + "m_classes": ("micro", "small", "medium"), +} + + +@lru_cache(maxsize=1) +def _kernel(): + from flash_rt.structures.impls import hub_kernel + + return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) + + +def _activation(variant: Mapping[str, str]) -> tuple[str, Callable]: + name = variant.get("activation", "gelu") + if name not in _ENTRYPOINTS: + raise ValueError(f"unsupported activation: {name!r}") + if name == "gelu": + return name, lambda t: torch.nn.functional.gelu(t, approximate="tanh") + return name, torch.nn.functional.silu + + +def _amax_scale(tensor: torch.Tensor) -> torch.Tensor: + return (tensor.float().abs().max() / _FP8_MAX).clamp(min=1e-8) + + +def _quantize(tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return (tensor.float() / scale).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8) + + +def _normalize( + x: torch.Tensor, + w_norm: torch.Tensor, + mode: str, + cond_scale: torch.Tensor | None, + cond_shift: torch.Tensor | None, + eps: float, +) -> torch.Tensor: + h = x.float() + h = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + eps) + if mode == "offset": + h = h * (1.0 + w_norm.float()) + elif mode == "direct": + h = h * w_norm.float() + else: + raise ValueError(f"unknown norm_weight_mode: {mode!r}") + if cond_scale is not None: + h = h * (1.0 + cond_scale.float()) + if cond_shift is not None: + h = h + cond_shift.float() + return h.to(torch.bfloat16) + + +def _check_and_pack(weights: Mapping[str, torch.Tensor]): + """Validate dims against the support envelope; pack FP8 weights.""" + w_gate, w_up, w_down = weights["w_gate"], weights["w_up"], weights["w_down"] + dim_d, dim_f = w_gate.shape + if w_up.shape != (dim_d, dim_f) or w_down.shape != (dim_f, dim_d): + raise ValueError( + f"inconsistent weight dims: gate {tuple(w_gate.shape)}, " + f"up {tuple(w_up.shape)}, down {tuple(w_down.shape)}" + ) + for name, dim in (("D", dim_d), ("F", dim_f)): + bounds = SUPPORT[name] + if not bounds["min"] <= dim <= bounds["max"]: + raise ValueError( + f"{name}={dim} outside support envelope " + f"[{bounds['min']}, {bounds['max']}]" + ) + if not (w_gate.is_cuda and w_up.is_cuda and w_down.is_cuda): + raise ValueError("fp8_static requires CUDA-resident weights") + gate_up = torch.cat([w_gate.t(), w_up.t()], dim=0).contiguous() + down = w_down.t().contiguous() + return gate_up, down, _amax_scale(gate_up), _amax_scale(down) + + +def _calibrate_scales( + normed_samples: Sequence[torch.Tensor], + w_gate: torch.Tensor, + w_up: torch.Tensor, + act: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Static per-tensor input/hidden scales from normed activations.""" + if not normed_samples: + raise ValueError("calibration samples must be non-empty") + device = w_gate.device + input_amax = torch.zeros((), device=device) + hidden_amax = torch.zeros((), device=device) + with torch.no_grad(): + for h in normed_samples: + flat = h.reshape(-1, h.shape[-1]).float().to(device) + hidden = act(flat @ w_gate.float()) * (flat @ w_up.float()) + input_amax = torch.maximum(input_amax, flat.abs().max()) + hidden_amax = torch.maximum(hidden_amax, hidden.abs().max()) + return ((input_amax / _FP8_MAX).clamp(min=1e-8), + (hidden_amax / _FP8_MAX).clamp(min=1e-8)) + + +@dataclass(frozen=True) +class BoundDecoderFfnFp8: + """Bound callable for the full structure boundary.""" + + fused_mlp: Callable[..., torch.Tensor] + w_norm: torch.Tensor + gate_up_fp8: torch.Tensor + down_fp8: torch.Tensor + input_scale: torch.Tensor + gate_up_scale: torch.Tensor + hidden_scale: torch.Tensor + down_scale: torch.Tensor + norm_weight_mode: str + eps: float + + def ffn(self, normed: torch.Tensor) -> torch.Tensor: + """The normed-input -> ffn-output slice (no norm, no residual). + + Input quantization is fused inside the kernel's BF16 entry.""" + shape = normed.shape + out = self.fused_mlp( + normed.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous(), + self.gate_up_fp8, + self.down_fp8, + self.input_scale.view(1), + self.gate_up_scale.view(1), + self.hidden_scale.view(1), + self.down_scale.view(1), + ) + return out.reshape(shape).to(normed.dtype) + + def __call__( + self, + x: torch.Tensor, + *, + cond_scale: torch.Tensor | None = None, + cond_shift: torch.Tensor | None = None, + cond_gate: torch.Tensor | None = None, + ) -> torch.Tensor: + h = _normalize(x, self.w_norm, self.norm_weight_mode, + cond_scale, cond_shift, self.eps) + out = self.ffn(h) + if cond_gate is not None: + out = out * cond_gate + return x + out.to(x.dtype) + + +class FusedGeGluMlp(torch.nn.Module): + """MLP-seam module for hosts whose replaceable boundary is the MLP. + + The host keeps its own norm, AdaLN gate, and residual. ``original`` + is retained whole (host MLP naming varies across model families), and + attribute lookups fall through to it so hosts that introspect the + projection attributes of the module they call keep working. + """ + + def __init__(self, bound: BoundDecoderFfnFp8, + original: torch.nn.Module | None = None): + super().__init__() + self._bound = bound + if original is not None: + self.host_mlp = original + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_mlp": + raise + return getattr(super().__getattr__("host_mlp"), name) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self._bound.ffn(hidden) + + +def _build(weights, variant, input_scale, hidden_scale, eps): + name, _ = _activation(variant) + gate_up, down, gate_up_scale, down_scale = _check_and_pack(weights) + return BoundDecoderFfnFp8( + fused_mlp=getattr(_kernel(), _ENTRYPOINTS[name]), + w_norm=weights["w_norm"], + gate_up_fp8=_quantize(gate_up, gate_up_scale), + down_fp8=_quantize(down, down_scale), + input_scale=input_scale, + gate_up_scale=gate_up_scale, + hidden_scale=hidden_scale, + down_scale=down_scale, + norm_weight_mode=variant.get("norm_weight_mode", "offset"), + eps=eps, + ) + + +@torch.no_grad() +def bind( + weights: Mapping[str, torch.Tensor], + *, + variant: Mapping[str, str], + calibration_inputs: Sequence[Mapping[str, torch.Tensor]], + eps: float = 1e-6, +) -> BoundDecoderFfnFp8: + """Bind the full structure: calibration inputs are boundary inputs. + + ``calibration_inputs`` must be drawn from the real input distribution + of the target binding; static FP8 scales are only as trustworthy as + the data they were measured on. + """ + if not calibration_inputs: + raise ValueError("calibration_inputs must be non-empty") + _, act = _activation(variant) + mode = variant.get("norm_weight_mode", "offset") + normed = [ + _normalize(sample["x"], weights["w_norm"], mode, + sample.get("cond_scale"), sample.get("cond_shift"), eps) + for sample in calibration_inputs + ] + input_scale, hidden_scale = _calibrate_scales( + normed, weights["w_gate"], weights["w_up"], act) + return _build(weights, variant, input_scale, hidden_scale, eps) + + +@torch.no_grad() +def bind_mlp_seam( + weights: Mapping[str, torch.Tensor], + *, + variant: Mapping[str, str], + calibration_normed: Sequence[torch.Tensor], + original: torch.nn.Module | None = None, + eps: float = 1e-6, +) -> FusedGeGluMlp: + """Bind the MLP-seam slice: calibration inputs are normed activations.""" + _, act = _activation(variant) + input_scale, hidden_scale = _calibrate_scales( + calibration_normed, weights["w_gate"], weights["w_up"], act) + bound = _build(weights, variant, input_scale, hidden_scale, eps) + return FusedGeGluMlp(bound, original=original) diff --git a/flash_rt/structures/impls/decoder_ffn/fp8_static.yaml b/flash_rt/structures/impls/decoder_ffn/fp8_static.yaml new file mode 100644 index 00000000..ad9bc08a --- /dev/null +++ b/flash_rt/structures/impls/decoder_ffn/fp8_static.yaml @@ -0,0 +1,22 @@ +impl: fp8_static +structure: decoder_ffn +version: 1 +recipe: fp8_static +backends: [cuda, rocm] +form: hub +module: fp8_static +entrypoint: bind + +kernels: + - {provider: hf, repo: flashrt/flashrt-fp8-swiglu-ffn, version: ">=1"} + +envelope: + support: + D: {min: 512, max: 16384} + F: {min: 1024, max: 16384} + align: {} + m_classes: [micro, small, medium] + notes: > + Support bounds mirror the checks enforced in bind(); a kernel-side + can_implement probe should replace them once exposed. Activation + scales are static per-tensor and require calibration inputs at bind. diff --git a/flash_rt/structures/impls/decoder_ffn/w8a16_static.py b/flash_rt/structures/impls/decoder_ffn/w8a16_static.py new file mode 100644 index 00000000..6df62fe8 --- /dev/null +++ b/flash_rt/structures/impls/decoder_ffn/w8a16_static.py @@ -0,0 +1,168 @@ +"""Weight-only INT8 implementation of the ``decoder_ffn`` structure. + +Composes the fused W8A16 gate/up -> activation -> down block from the +``flashrt/weight-only-ffn`` Hub kernel behind the structure boundary. +Activations stay BF16, so binding needs no calibration data: packing is +a pure weight transform, and qualification still runs the parity gate on +real host activations like every other implementation. + +The kernel's optimized dispatch covers the decode band (M in [1, 8]); +larger M is outside the support envelope and is refused at call time +rather than routed to a slow path. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from functools import lru_cache + +import torch + +KERNEL_DEP = { + "provider": "huggingface_kernels", + "repo": "flashrt/weight-only-ffn", + "version": ">=1", +} + +_ENTRYPOINTS = {"gelu": "w8a16_geglu_ffn_bf16", "silu": "w8a16_swiglu_ffn_bf16"} + +SUPPORT = { + "D": {"min": 512, "max": 16384, "multiple_of": 64}, + "F": {"min": 1024, "max": 16384, "multiple_of": 64}, + "M": {"min": 1, "max": 8}, + "m_classes": ("micro",), +} + + +@lru_cache(maxsize=1) +def _kernel(): + from flash_rt.structures.impls import hub_kernel + + return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) + + +def _entrypoint(variant: Mapping[str, str]): + name = variant.get("activation", "gelu") + if name not in _ENTRYPOINTS: + raise ValueError(f"unsupported activation: {name!r}") + return getattr(_kernel(), _ENTRYPOINTS[name]) + + +def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]: + w_gate, w_up, w_down = (weights["w_gate"], weights["w_up"], + weights["w_down"]) + dim_f, dim_d = w_gate.shape + if w_up.shape != (dim_f, dim_d) or w_down.shape != (dim_d, dim_f): + raise ValueError( + f"inconsistent weight dims: gate {tuple(w_gate.shape)}, " + f"up {tuple(w_up.shape)}, down {tuple(w_down.shape)}") + for name, dim in (("D", dim_d), ("F", dim_f)): + bounds = SUPPORT[name] + if not bounds["min"] <= dim <= bounds["max"]: + raise ValueError( + f"{name}={dim} outside support envelope " + f"[{bounds['min']}, {bounds['max']}]") + if dim % bounds["multiple_of"]: + raise ValueError( + f"{name}={dim} must be a multiple of " + f"{bounds['multiple_of']}") + return dim_d, dim_f + + +class BoundDecoderFfnW8A16: + """MLP-seam callable: normed activations in, FFN output out (BF16).""" + + def __init__(self, ffn_fn, gate_up_q, gate_up_scale, down_q, down_scale, + dim_d): + self._ffn = ffn_fn + self._gate_up_q = gate_up_q + self._gate_up_scale = gate_up_scale + self._down_q = down_q + self._down_scale = down_scale + self._dim_d = dim_d + + def ffn(self, normed: torch.Tensor) -> torch.Tensor: + shape = normed.shape + x = normed.reshape(-1, shape[-1]) + m = x.shape[0] + m_max = SUPPORT["M"]["max"] + if m > m_max: + raise ValueError( + f"M={m} outside the weight-only decode envelope " + f"[1, {m_max}]") + variant = 0 if m <= 4 else 3 + out = self._ffn(x.to(torch.bfloat16).contiguous(), + self._gate_up_q, self._gate_up_scale, + self._down_q, self._down_scale, variant=variant) + return out.reshape(shape).to(normed.dtype) + + __call__ = ffn + + +class FusedGluMlpW8A16(torch.nn.Module): + """MLP-seam module with declared M-dispatch. + + The weight-only kernel covers the decode band (M in [1, 8]); calls + with larger M 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 qualification + record states which band the kernel serves. + + ``original`` is retained whole (host MLP naming varies across model + families), and attribute lookups fall through to it so hosts that + introspect the module they call keep working. + """ + + def __init__(self, bound: BoundDecoderFfnW8A16, + original: torch.nn.Module | None = None): + super().__init__() + self._bound = bound + if original is not None: + self.host_mlp = original + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_mlp": + raise + return getattr(super().__getattr__("host_mlp"), name) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + m = hidden.numel() // hidden.shape[-1] + if m > SUPPORT["M"]["max"]: + try: + host = super().__getattr__("host_mlp") + except AttributeError: + pass + else: + return host(hidden) + return self._bound.ffn(hidden) + + +@torch.no_grad() +def bind_mlp_seam( + weights: Mapping[str, torch.Tensor], + *, + variant: Mapping[str, str], + original: torch.nn.Module | None = None, +): + """Bind the MLP-seam slice of ``decoder_ffn`` with weight-only INT8. + + ``weights`` uses checkpoint-native ``[out, in]`` projection layout + (``w_gate``/``w_up``: ``[F, D]``, ``w_down``: ``[D, F]``). No + calibration data is required: quantization is per-output-channel on + weights only. + """ + dim_d, _ = _check(weights) + k = _kernel() + ffn_fn = _entrypoint(variant) + gate_up = torch.cat( + [weights["w_gate"].to("cuda", torch.bfloat16), + weights["w_up"].to("cuda", torch.bfloat16)], dim=0).contiguous() + down = weights["w_down"].to("cuda", torch.bfloat16).contiguous() + gate_up_q, gate_up_scale = k.quantize_w8_weight_bf16(gate_up) + down_q, down_scale = k.quantize_w8_weight_bf16(down) + bound = BoundDecoderFfnW8A16( + ffn_fn, gate_up_q, gate_up_scale, down_q, down_scale, dim_d) + return FusedGluMlpW8A16(bound, original=original) diff --git a/flash_rt/structures/impls/linear_proj/__init__.py b/flash_rt/structures/impls/linear_proj/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/impls/linear_proj/fp8_static.py b/flash_rt/structures/impls/linear_proj/fp8_static.py new file mode 100644 index 00000000..a20dccc6 --- /dev/null +++ b/flash_rt/structures/impls/linear_proj/fp8_static.py @@ -0,0 +1,151 @@ +"""FP8-static implementation of the ``linear_proj`` structure. + +Wraps the fused BF16-entry FP8 projection from ``flashrt/flashrt-fp8-ffn`` +(``bf16_fp8_linear_bias_bf16``: fused input quantization, FP8 weights, +BF16 bias/output) behind the linear_proj boundary. Activation scale is +static per-tensor, calibrated from caller-provided representative +inputs. + +Qualification is work-based, derived from standalone preflight on the +target shapes: the fused entry pays a fixed input-quantization cost, so +projections whose GEMM is too small to amortize it (M*N*K below the +floor) are refused at bind time — the caller records the refusal, the +host keeps its own Linear. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Mapping, Sequence + +import torch + +KERNEL_DEP = { + "provider": "hf", + "repo": "flashrt/flashrt-fp8-ffn", + "version": ">=1", +} + +_FP8 = torch.float8_e4m3fn +_FP8_MAX = 448.0 + +SUPPORT = { + # measured on RTX 5090: wins at M=712 K=2048 (1.2-1.9x), loses at + # M=51 K=1024 (quant cost > GEMM). Floor sits between those bands. + "flops_min": 2.0e8, + "K": {"min": 512, "max": 16384}, + "N": {"min": 128, "max": 16384}, +} + + +@lru_cache(maxsize=1) +def _kernel(): + from flash_rt.structures.impls import hub_kernel + + return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) + + +def _amax_scale(t: torch.Tensor) -> torch.Tensor: + return (t.float().abs().max() / _FP8_MAX).clamp(min=1e-8) + + +class FusedLinearProj(torch.nn.Module): + """Drop-in replacement for one nn.Linear projection. + + ``original`` is retained whole and attribute lookups fall through to + it, so host code that introspects ``weight``/``bias``/``in_features`` + keeps working. + """ + + def __init__(self, w_fp8, bias, input_scale, weight_scale, + original: torch.nn.Module | None = None): + super().__init__() + self._w_fp8 = w_fp8 + self._bias = bias + self._input_scale = input_scale + self._weight_scale = weight_scale + self._bufs: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + # resolve the op at bind time: calling the hub loader inside + # forward makes dynamo trace through kernels.get_kernel's + # version resolution (network + inspect.Signature) — 26 graph + # breaks that fragment the surrounding compiled region + self._fn = _kernel().bf16_fp8_linear_bias_bf16 + if original is not None: + self.host_linear = original + + 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: + shape = x.shape + flat = x.reshape(-1, shape[-1]) + m = flat.shape[0] + bufs = self._bufs.get(m) + if bufs is None: + bufs = (torch.empty_like(flat, dtype=_FP8), + torch.empty(m, self._w_fp8.shape[0], device=x.device, + dtype=torch.bfloat16)) + self._bufs[m] = bufs + x_fp8, out = bufs + y = self._fn( + flat.to(torch.bfloat16).contiguous(), self._w_fp8, self._bias, + self._input_scale, self._weight_scale, + input_fp8=x_fp8, out=out) + return y.reshape(*shape[:-1], y.shape[-1]).to(x.dtype) + + +@torch.no_grad() +def bind_proj_seam( + weights: Mapping[str, torch.Tensor], + *, + calibration: Sequence[torch.Tensor], + original: torch.nn.Module | None = None, +) -> FusedLinearProj: + """Bind one projection: ``weights['w']`` is checkpoint-layout [N, K]. + + ``calibration``: real inputs of the projection. Work-based + qualification uses the median calibration M. + """ + if not calibration: + raise ValueError("calibration must be non-empty") + w = weights["w"] + n, k = w.shape + for name, dim in (("K", k), ("N", n)): + lo, hi = SUPPORT[name]["min"], SUPPORT[name]["max"] + if not lo <= dim <= hi: + raise ValueError(f"{name}={dim} outside support envelope") + ms = sorted(int(t.reshape(-1, t.shape[-1]).shape[0]) + for t in calibration) + m_med = ms[len(ms) // 2] + if m_med * n * k < SUPPORT["flops_min"]: + raise ValueError( + f"projection work {m_med}x{n}x{k} below amortization floor " + f"({SUPPORT['flops_min']:.0e}) — fused quant cost would not " + "pay for itself; host keeps its Linear") + if not w.is_cuda: + raise ValueError("fp8_static requires CUDA-resident weights") + + device = w.device + w_scale = _amax_scale(w) + w_fp8 = (w.float() / w_scale).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8) + amax = torch.zeros((), device=device) + for t in calibration: + amax = torch.maximum(amax, t.float().abs().max().to(device)) + input_scale = (amax / _FP8_MAX).clamp(min=1e-8) + bias = weights.get("b") + if bias is None: + bias = torch.zeros(n, device=device, dtype=torch.bfloat16) + else: + bias = bias.detach().to(torch.bfloat16) + bound = FusedLinearProj(w_fp8, bias, input_scale.view(1), + w_scale.view(1), original=original) + for m in set(ms): # pre-allocate per calibrated M: keeps the hot + bound._bufs[m] = ( # path allocation-free (graph/compile safe) + torch.empty(m, k, device=device, dtype=_FP8), + torch.empty(m, n, device=device, dtype=torch.bfloat16)) + return bound diff --git a/flash_rt/structures/impls/qkv_pack/__init__.py b/flash_rt/structures/impls/qkv_pack/__init__.py new file mode 100644 index 00000000..bbe73e9b --- /dev/null +++ b/flash_rt/structures/impls/qkv_pack/__init__.py @@ -0,0 +1,5 @@ +from .fp8_static import (AttnBlockPacked, PackedLinear, StashReader, + bind_attn_block, bind_qkv_pack) + +__all__ = ["AttnBlockPacked", "PackedLinear", "StashReader", + "bind_attn_block", "bind_qkv_pack"] diff --git a/flash_rt/structures/impls/qkv_pack/fp8_static.py b/flash_rt/structures/impls/qkv_pack/fp8_static.py new file mode 100644 index 00000000..6b860633 --- /dev/null +++ b/flash_rt/structures/impls/qkv_pack/fp8_static.py @@ -0,0 +1,225 @@ +"""qkv_pack — pack sibling linears that share one input into one GEMM. + +Sibling projections consumed in a fixed call order (q/k/v of an +attention block, gate/up of an MLP) each pay a small-M GEMM whose cost +is launch/latency floor, not bandwidth. Packing their weights into one +``[sum(N_i), K]`` matrix turns the group into a single GEMM; the later +siblings become buffer reads. Two bind forms cover the hosts seen so +far: + +- **leaf**: the host calls the sibling modules separately and there is + no enclosing attention-module boundary. The first sibling's slot gets + a :class:`PackedLinear` (runs the packed GEMM, writes the other + outputs into preallocated buffers); the later slots get + :class:`StashReader` (return the buffer). The host's own call order + is the data dependency — functionalization keeps copy/read ordered + inside compiled and captured graphs. +- **module**: the host has an attention module with + ``q_proj/k_proj/v_proj/out_proj`` attributes and a standard + projections → attention → out_proj forward. :class:`AttnBlockPacked` + replaces the whole module: packed GEMM, SDPA at a declared compute + dtype, original ``out_proj``. + +Both forms quantize the packed weight to FP8 with one joint per-tensor +scale (the joint-scale rounding difference is covered by the parity +gate). Inputs enter either as FP8 (a producer seam supplies the shared +``act_scale``) or as BF16 through the fused-quantize entry with a +calibrated ``act_scale``. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from .. import hub_kernel + +_FP8 = torch.float8_e4m3fn + + +def _pack_weights(mods: Sequence[torch.nn.Module]): + ws = [m.weight.detach() for m in mods] + w = torch.cat(ws, 0) + scale = (w.float().abs().max() / 448.0).clamp(min=1e-8).view(1) + w8 = (w.float() / scale).clamp(-448, 448).to(_FP8) + splits = [wi.shape[0] for wi in ws] + biases = [] + for m in mods: + b = getattr(m, "bias", None) + biases.append(b.detach().to(torch.bfloat16) if b is not None + else torch.zeros(m.weight.shape[0], + device=w.device, + dtype=torch.bfloat16)) + return w8, scale, torch.cat(biases), splits + + +class PackedLinear(torch.nn.Module): + """Leaf-form head: one packed GEMM, later siblings stashed.""" + + def __init__(self, mods: Sequence[torch.nn.Module], + act_scale: torch.Tensor, rows: int, + in_dtype: str = "fp8_static"): + super().__init__() + self.host_linear = mods[0] + kf = hub_kernel("flashrt/flashrt-fp8-ffn", ">=1") + self.in_dtype = in_dtype + w8, w_scale, bias, splits = _pack_weights(mods) + self.splits = splits + self.register_buffer("w8", w8) + self.register_buffer("w_scale", w_scale) + self.register_buffer("bias_cat", bias) + self.register_buffer("act_scale", act_scale) + dev = w8.device + for i, n in enumerate(splits[1:], 1): + self.register_buffer(f"stash{i}", torch.zeros( + rows, n, device=dev, dtype=torch.bfloat16)) + if in_dtype == "fp8_static": + self._fn = kf.fp8_linear_bias_bf16 + else: + self._fn = kf.bf16_fp8_linear_bias_bf16 + k = mods[0].weight.shape[1] + self.register_buffer("x8_buf", torch.empty( + rows, k, device=dev, dtype=_FP8)) + self.register_buffer("y_buf", torch.empty( + rows, sum(splits), device=dev, dtype=torch.bfloat16)) + + def forward(self, x): + flat = x.reshape(-1, x.shape[-1]) + if self.in_dtype == "fp8_static": + y = self._fn(flat, self.w8, self.bias_cat, self.act_scale, + self.w_scale) + else: + y = self._fn(flat.to(torch.bfloat16).contiguous(), + self.w8, self.bias_cat, self.act_scale, + self.w_scale, input_fp8=self.x8_buf, + out=self.y_buf) + off = self.splits[0] + for i, n in enumerate(self.splits[1:], 1): + getattr(self, f"stash{i}").copy_(y[:, off:off + n]) + off += n + out = y[:, :self.splits[0]].contiguous() + out = out.reshape(*x.shape[:-1], self.splits[0]) + # the kernel's output dtype is BF16 by contract; only cast back + # when the host boundary itself is a compute dtype. On the + # fp8_static entry the input is FP8 (a producer seam supplies + # it) and casting to it would hand FP8 activations to the + # host's next op. + return out if x.dtype is _FP8 else out.to(x.dtype) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_linear"), name) + + +class StashReader(torch.nn.Module): + """Leaf-form tail: return the packed head's stashed output.""" + + def __init__(self, orig: torch.nn.Module, packed: PackedLinear, + index: int): + super().__init__() + self.host_linear = orig + self._packed = (packed,) + self.index = index + + def forward(self, x): + buf = getattr(self._packed[0], f"stash{self.index}") + out = buf.reshape(*x.shape[:-1], buf.shape[-1]) + return out if x.dtype is _FP8 else out.to(x.dtype) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_linear"), name) + + +def bind_qkv_pack(mods: Sequence[torch.nn.Module], + act_scale: torch.Tensor, rows: int, + in_dtype: str = "fp8_static"): + """Bind a sibling group; returns replacements in sibling order.""" + if len(mods) < 2: + raise ValueError("qkv_pack: need at least two siblings") + kdims = {m.weight.shape[1] for m in mods} + if len(kdims) != 1: + raise ValueError(f"qkv_pack: sibling K dims differ {kdims}") + packed = PackedLinear(mods, act_scale, rows, in_dtype=in_dtype) + out = [packed] + for i, m in enumerate(mods[1:], 1): + out.append(StashReader(m, packed, i)) + return out + + +class AttnBlockPacked(torch.nn.Module): + """Module-form: packed QKV + SDPA at a declared dtype + out_proj. + + Fits attention modules exposing ``q_proj/k_proj/v_proj/out_proj``, + ``head_dim`` and ``scale`` with the standard block forward + (SigLIP/CLIP-family vision towers and friends). + """ + + def __init__(self, orig: torch.nn.Module, act_scale: torch.Tensor, + rows: int, sdpa_dtype: torch.dtype = torch.bfloat16): + super().__init__() + self.host_attn = orig + kf = hub_kernel("flashrt/flashrt-fp8-ffn", ">=1") + self._fn = kf.bf16_fp8_linear_bias_bf16 + w8, w_scale, bias, splits = _pack_weights( + [orig.q_proj, orig.k_proj, orig.v_proj]) + if len(set(splits)) != 1: + raise ValueError("attn_block: q/k/v widths differ") + self.e = splits[0] + self.register_buffer("w8", w8) + self.register_buffer("w_scale", w_scale) + self.register_buffer("bias_cat", bias) + self.register_buffer("in_scale", act_scale) + k = orig.q_proj.weight.shape[1] + dev = w8.device + self.register_buffer("x8_buf", torch.empty( + rows, k, device=dev, dtype=_FP8)) + self.register_buffer("y_buf", torch.empty( + rows, 3 * self.e, device=dev, dtype=torch.bfloat16)) + self.sdpa_dtype = sdpa_dtype + + def forward(self, hidden_states, attention_mask=None, **kw): + a = self.host_attn + bsz, seq, dim = hidden_states.shape + flat = hidden_states.reshape(-1, dim).to( + torch.bfloat16).contiguous() + y = self._fn(flat, self.w8, self.bias_cat, self.in_scale, + self.w_scale, input_fp8=self.x8_buf, + out=self.y_buf) + hd = a.head_dim + + def split(t): + return t.contiguous().view(bsz, seq, -1, hd).transpose( + 1, 2).to(self.sdpa_dtype) + + e = self.e + mask = (attention_mask.to(self.sdpa_dtype) + if attention_mask is not None else None) + o = torch.nn.functional.scaled_dot_product_attention( + split(y[:, :e]), split(y[:, e:2 * e]), split(y[:, 2 * e:]), + attn_mask=mask, scale=a.scale) + o = o.to(hidden_states.dtype).transpose(1, 2).reshape( + bsz, seq, dim).contiguous() + return a.out_proj(o), None + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_attn"), name) + + +def bind_attn_block(orig: torch.nn.Module, act_scale: torch.Tensor, + rows: int, + sdpa_dtype: torch.dtype = torch.bfloat16 + ) -> AttnBlockPacked: + for attr in ("q_proj", "k_proj", "v_proj", "out_proj", "head_dim", + "scale"): + if not hasattr(orig, attr): + raise ValueError(f"attn_block: host lacks {attr!r}") + return AttnBlockPacked(orig, act_scale, rows, sdpa_dtype=sdpa_dtype) diff --git a/flash_rt/structures/impls/step_table.py b/flash_rt/structures/impls/step_table.py new file mode 100644 index 00000000..5e014057 --- /dev/null +++ b/flash_rt/structures/impls/step_table.py @@ -0,0 +1,113 @@ +"""Step-table memoization for step-constant conditioning producers. + +Diffusion-style hosts recompute ``dense(cond)`` in every layer at every +denoise step, yet ``cond`` depends only on the timestep: over a tick the +producer emits a small fixed set of vectors. This implementation +replaces such a producer with a calibrated table — the distinct +conditioning vectors seen during calibration and the outputs the host's +own producer computed for them. At runtime the module locates the +current step by nearest-neighbour match against the stored vectors +(pure tensor ops: safe under both compile tracing and graph capture, +no Python state) and gathers the stored row instead of running the +GEMV. Outputs are bit-identical to calibration by construction; the +match itself is arbitrated by the caller's parity gate. + +Qualification refuses hosts whose conditioning is not actually +step-quantized: if calibration sees more distinct vectors than +``max_steps``, the producer depends on more than the step and a table +would silently mis-hit — that host keeps its GEMV. +""" + +from __future__ import annotations + +import torch + + +class StepTableLinear(torch.nn.Module): + """Drop-in for a ``nn.Linear`` whose input is step-constant.""" + + def __init__(self, original: torch.nn.Module, conds: torch.Tensor, + table: torch.Tensor, + locator: "StepTableLinear | None" = None): + super().__init__() + self.host_linear = original + # match score: argmax(2 c·k - |k|^2) == nearest neighbour. + # Sibling tables fed by the same conditioning stream share the + # locator buffers (same tensor objects), so a compiling host + # sees one common subexpression per step instead of one locate + # per table — the redundant matches fold away. + if locator is not None: + self.register_buffer("conds_t", locator.conds_t) + self.register_buffer("cond_sq", locator.cond_sq) + else: + self.register_buffer("conds_t", conds.float().t().contiguous()) + self.register_buffer("cond_sq", + (conds.float() ** 2).sum(-1).contiguous()) + self.register_buffer("table", table.contiguous()) + + def forward(self, cond: torch.Tensor) -> torch.Tensor: + flat = cond.reshape(-1, cond.shape[-1]).float() + scores = 2.0 * (flat @ self.conds_t) - self.cond_sq + idx = scores.argmax(dim=-1) + out = self.table.index_select(0, idx) + return out.reshape(*cond.shape[:-1], out.shape[-1]) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("host_linear"), name) + + +def bind_step_table(original: torch.nn.Module, + calibration: list[tuple[torch.Tensor, torch.Tensor]], + *, max_steps: int = 64, + dedup_rtol: float = 1e-5, + share_locator_with: StepTableLinear | None = None + ) -> StepTableLinear: + """Build a step table from real ``(cond, out)`` calibration pairs. + + Pairs come from hooking the host's own producer over at least one + full tick, so the table rows are exactly what the host computed. + Refuses (``ValueError``) when the distinct-vector count exceeds + ``max_steps`` — the producer is then not step-constant and a table + would alias different inputs onto one row. + + ``share_locator_with``: a sibling table bound from the same + conditioning stream; when its stored vectors match this + calibration exactly (same set, same order), the new table reuses + the sibling's locator buffers so redundant per-table step matches + can fold into one. On any mismatch the table keeps its own + locator — sharing is an optimization, never an assumption. + """ + if not calibration: + raise ValueError("step_table: no calibration pairs captured") + conds: list[torch.Tensor] = [] + outs: list[torch.Tensor] = [] + for cond, out in calibration: + c = cond.detach().reshape(-1, cond.shape[-1]) + o = out.detach().reshape(-1, out.shape[-1]) + for row in range(c.shape[0]): + cr = c[row] + if any(torch.allclose(cr, seen, rtol=dedup_rtol, + atol=1e-6 * cr.abs().max().item() + 1e-12) + for seen in conds): + continue + conds.append(cr.clone()) + outs.append(o[row].clone()) + if len(conds) > max_steps: + raise ValueError( + f"step_table: >{max_steps} distinct conditioning " + "vectors — producer is not step-constant, keeping " + "the host GEMV") + stacked = torch.stack(conds) + locator = None + if (share_locator_with is not None + and share_locator_with.conds_t.shape[1] == stacked.shape[0] + and torch.allclose(share_locator_with.conds_t.t(), + stacked.float().to( + share_locator_with.conds_t.device), + rtol=dedup_rtol, atol=1e-6)): + locator = share_locator_with + return StepTableLinear(original, stacked, torch.stack(outs), + locator=locator) diff --git a/flash_rt/structures/impls/vision_ffn/__init__.py b/flash_rt/structures/impls/vision_ffn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/structures/impls/vision_ffn/fp8_static.py b/flash_rt/structures/impls/vision_ffn/fp8_static.py new file mode 100644 index 00000000..9134377f --- /dev/null +++ b/flash_rt/structures/impls/vision_ffn/fp8_static.py @@ -0,0 +1,214 @@ +"""FP8-static implementation of the ``vision_ffn`` structure. + +Composes the fused FP8 fc1 -> GELU -> fc2 block (biases included) from +the ``flashrt/flashrt-fp8-ffn`` Hub kernel. ``bind`` covers the full +structure boundary; ``bind_mlp_seam`` covers the normed-input -> +ffn-output slice for hosts whose replaceable module boundary is the MLP. +Weights use the checkpoint-native (out, in) layout directly. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, Mapping, Sequence + +import torch + +KERNEL_DEP = { + "provider": "hf", + "repo": "flashrt/flashrt-fp8-ffn", + "version": ">=1", +} + +_FP8 = torch.float8_e4m3fn +_FP8_MAX = 448.0 + +SUPPORT = { + "D": {"min": 512, "max": 16384}, + "F": {"min": 1024, "max": 16384}, + "m_classes": ("small", "medium", "large"), +} + + +@lru_cache(maxsize=1) +def _kernel(): + from flash_rt.structures.impls import hub_kernel + + return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) + + +def _amax_scale(tensor: torch.Tensor) -> torch.Tensor: + return (tensor.float().abs().max() / _FP8_MAX).clamp(min=1e-8) + + +def _quantize(tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return (tensor.float() / scale).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8) + + +@dataclass(frozen=True) +class BoundVisionFfnFp8: + """Bound callable for the full structure boundary.""" + + fused_mlp: Callable[..., torch.Tensor] + w_norm: torch.Tensor + b_norm: torch.Tensor + fc1_fp8: torch.Tensor + b_fc1: torch.Tensor + fc2_fp8: torch.Tensor + b_fc2: torch.Tensor + input_scale: torch.Tensor + fc1_scale: torch.Tensor + hidden_scale: torch.Tensor + fc2_scale: torch.Tensor + eps: float + + def ffn(self, normed: torch.Tensor) -> torch.Tensor: + """The normed-input -> ffn-output slice (no norm, no residual). + + Input quantization is fused inside the kernel's BF16 entry.""" + shape = normed.shape + out = self.fused_mlp( + normed.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous(), + self.fc1_fp8, + self.b_fc1, + self.fc2_fp8, + self.b_fc2, + self.input_scale.view(1), + self.fc1_scale.view(1), + self.hidden_scale.view(1), + self.fc2_scale.view(1), + ) + return out.reshape(shape).to(normed.dtype) + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + h = torch.nn.functional.layer_norm( + x.float(), (x.shape[-1],), + self.w_norm.float(), self.b_norm.float(), self.eps).to(x.dtype) + return x + self.ffn(h).to(x.dtype) + + +class FusedGeluMlp(torch.nn.Module): + """MLP-seam module: the host keeps its own norm and residual. + + ``original`` is retained whole (host MLP naming varies across model + families), and attribute lookups fall through to it so hosts that + introspect the module they call keep working. + """ + + def __init__(self, bound: BoundVisionFfnFp8, + original: torch.nn.Module | None = None): + super().__init__() + self._bound = bound + if original is not None: + self.host_mlp = original + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_mlp": + raise + return getattr(super().__getattr__("host_mlp"), name) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self._bound.ffn(hidden) + + +def _calibrate(normed_samples, w_fc1, b_fc1): + if not normed_samples: + raise ValueError("calibration samples must be non-empty") + device = w_fc1.device + input_amax = torch.zeros((), device=device) + hidden_amax = torch.zeros((), device=device) + for h in normed_samples: + flat = h.reshape(-1, h.shape[-1]).float().to(device) + hidden = torch.nn.functional.gelu( + flat @ w_fc1.float().t() + b_fc1.float(), approximate="tanh") + input_amax = torch.maximum(input_amax, flat.abs().max()) + hidden_amax = torch.maximum(hidden_amax, hidden.abs().max()) + return ((input_amax / _FP8_MAX).clamp(min=1e-8), + (hidden_amax / _FP8_MAX).clamp(min=1e-8)) + + +def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]: + w_fc1, w_fc2 = weights["w_fc1"], weights["w_fc2"] + dim_f, dim_d = w_fc1.shape + if w_fc2.shape != (dim_d, dim_f): + raise ValueError( + f"inconsistent weight dims: fc1 {tuple(w_fc1.shape)}, " + f"fc2 {tuple(w_fc2.shape)}" + ) + for name, dim in (("D", dim_d), ("F", dim_f)): + bounds = SUPPORT[name] + if not bounds["min"] <= dim <= bounds["max"]: + raise ValueError( + f"{name}={dim} outside support envelope " + f"[{bounds['min']}, {bounds['max']}]" + ) + if not (w_fc1.is_cuda and w_fc2.is_cuda): + raise ValueError("fp8_static requires CUDA-resident weights") + return dim_d, dim_f + + +def _build(weights, input_scale, hidden_scale, eps): + _check(weights) + fc1_scale = _amax_scale(weights["w_fc1"]) + fc2_scale = _amax_scale(weights["w_fc2"]) + to_bf16 = lambda t: t.to(torch.bfloat16) + return BoundVisionFfnFp8( + fused_mlp=_kernel().bf16_fp8_gelu_mlp_bf16, + w_norm=weights["w_norm"], + b_norm=weights["b_norm"], + fc1_fp8=_quantize(weights["w_fc1"], fc1_scale), + b_fc1=to_bf16(weights["b_fc1"]), + fc2_fp8=_quantize(weights["w_fc2"], fc2_scale), + b_fc2=to_bf16(weights["b_fc2"]), + input_scale=input_scale, + fc1_scale=fc1_scale, + hidden_scale=hidden_scale, + fc2_scale=fc2_scale, + eps=eps, + ) + + +@torch.no_grad() +def bind( + weights: Mapping[str, torch.Tensor], + *, + variant: Mapping[str, str], + calibration_inputs: Sequence[Mapping[str, torch.Tensor]], + eps: float = 1e-6, +) -> BoundVisionFfnFp8: + """Bind the full structure: calibration inputs are boundary inputs.""" + if variant.get("activation", "gelu") != "gelu": + raise ValueError("vision_ffn fp8_static supports gelu only") + if not calibration_inputs: + raise ValueError("calibration_inputs must be non-empty") + normed = [ + torch.nn.functional.layer_norm( + s["x"].float(), (s["x"].shape[-1],), + weights["w_norm"].float(), weights["b_norm"].float(), eps) + for s in calibration_inputs + ] + input_scale, hidden_scale = _calibrate( + normed, weights["w_fc1"], weights["b_fc1"]) + return _build(weights, input_scale, hidden_scale, eps) + + +@torch.no_grad() +def bind_mlp_seam( + weights: Mapping[str, torch.Tensor], + *, + variant: Mapping[str, str], + calibration_normed: Sequence[torch.Tensor], + original: torch.nn.Module | None = None, + eps: float = 1e-6, +) -> FusedGeluMlp: + """Bind the MLP-seam slice: calibration inputs are normed activations.""" + if variant.get("activation", "gelu") != "gelu": + raise ValueError("vision_ffn fp8_static supports gelu only") + input_scale, hidden_scale = _calibrate( + calibration_normed, weights["w_fc1"], weights["b_fc1"]) + bound = _build(weights, input_scale, hidden_scale, eps) + return FusedGeluMlp(bound, original=original) diff --git a/flash_rt/structures/provider.py b/flash_rt/structures/provider.py new file mode 100644 index 00000000..505b5715 --- /dev/null +++ b/flash_rt/structures/provider.py @@ -0,0 +1,113 @@ +"""Package a host-captured CUDA graph as an ``frt_model_runtime_v1``. + +The provider route for external hosts (torch or any framework that can +capture its hot path): the host warms up and captures on a dedicated +stream, keeps every boundary tensor at a fixed device address, and hands +the raw graph exec plus those boundary windows here. The result is the +same runtime face the native pipelines export from +``flash_rt.models.*.runtime_export``, so any ``cap_model_runtime`` +consumer adopts it without knowing which producer built it. + +The caller keeps ownership of the captured graph and the window memory; +the returned runtime anchors the exec context and the wrapped buffers, +and must not outlive the host objects backing them. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +from flash_rt.runtime import export as _rt +from flash_rt.runtime.exec import _import_native + +_DIR_ROLE = {"in": ("input",), "out": ("output",)} + + +def export_captured_runtime( + stream_handle: int, + graphs: Sequence[tuple[str, int]], + windows: Mapping[str, tuple[int, int]], + ports: Sequence[Mapping], + regions: Sequence[tuple[str, str]] = (), + stages: Sequence[str] = (), + roles: Mapping[str, object] | None = None, + identity: Mapping[str, object] | None = None, + manifest_extra: Mapping[str, object] | None = None, +): + """Build an ``frt_model_runtime_v1`` over host-captured graphs. + + - ``stream_handle``: native stream the graphs were captured on (for + torch, ``torch.cuda.Stream().cuda_stream``). + - ``graphs``: ``(name, raw_graph_exec)`` pairs; for torch, + ``torch.cuda.CUDAGraph().raw_cuda_graph_exec()``. + - ``windows``: ``name -> (device_ptr, nbytes)`` boundary windows the + captured graphs read/write in place (SWAP semantics). + - ``ports``: dicts with :class:`flash_rt.runtime.export.PortSpec` + fields, plus ``window`` naming the backing window. An output port + may share its window with an input port (in-place convergence). + - ``regions``: ``(region_name, window_name)`` restorable-state + regions for capsule snapshot/restore. + - ``stages``: graph names in sequential step order; defaults to the + declared graph order. + - ``roles``: optional ``window_name -> role`` overrides; by default + a window's role is the union of its ports' directions. + """ + if not graphs: + raise ValueError("export_captured_runtime needs at least one graph") + c = _import_native() + ctx = c.Ctx() + sid = ctx.wrap_stream(int(stream_handle)) + + graph_specs = [] + for name, graph_exec in graphs: + g = ctx.graph(name) + g.adopt(0, int(graph_exec)) + graph_specs.append(_rt.GraphSpec(name, g, 0, (0,), stream="main")) + + wraps = {name: ctx.wrap(name, int(ptr), int(nbytes)) + for name, (ptr, nbytes) in windows.items()} + + port_specs = [] + port_roles: dict[str, set] = {name: set() for name in wraps} + for p in ports: + p = dict(p) + window = p.pop("window") + if window not in wraps: + raise ValueError(f"port {p.get('name')!r} names unknown window " + f"{window!r}") + port_roles[window].update(_DIR_ROLE.get(p.get("direction", "in"), ())) + port_specs.append(_rt.PortSpec(buffer=wraps[window], **p)) + + buffer_specs = [] + for name, buf in wraps.items(): + role = (roles or {}).get(name) + if role is None: + role = tuple(sorted(port_roles[name])) or ("state",) + buffer_specs.append(_rt.BufferSpec(name, buf, role)) + + region_specs = [] + for rname, window in regions: + if window not in wraps: + raise ValueError(f"region {rname!r} names unknown window " + f"{window!r}") + region_specs.append(_rt.RegionSpec(rname, wraps[window])) + + stage_names = list(stages) or [name for name, _ in graphs] + known = {g.name for g in graph_specs} + unknown = [s for s in stage_names if s not in known] + if unknown: + raise ValueError(f"stages name unknown graphs: {unknown}") + + return _rt.build_model_runtime( + ctx, + streams=[_rt.StreamSpec("main", sid, + native_handle=int(stream_handle))], + graphs=graph_specs, + buffers=buffer_specs, + regions=region_specs, + ports=port_specs, + stages=[_rt.StageSpec(name) for name in stage_names], + identity={str(k): str(v) for k, v in (identity or {}).items()}, + manifest_extra=dict(manifest_extra or {}), + owner=(ctx, wraps), + ) diff --git a/flash_rt/structures/recipe.py b/flash_rt/structures/recipe.py new file mode 100644 index 00000000..e02b858c --- /dev/null +++ b/flash_rt/structures/recipe.py @@ -0,0 +1,438 @@ +"""Declarative e2e recipes: assemble levers, audit on the graph, certify. + +A recipe declares how a host pipeline eats structure yield end to end: +which levers engage (region families, cadence pieces, seam negotiations, +dtype changes), how the hot stage is built, and under which explicit +gates the result is judged. ``run_recipe`` assembles the levers +transactionally, audits baseline vs treated **in the same process** +(cross-run anchors carry sub-millisecond drift, so sub-2% verdicts are +unreliable across runs), and emits a receipt recording every switch +state, every gate number, and every refusal reason. + +Switch discipline (the lever lifecycle): + + off declared but not engaged; recorded, never built + candidate engaged and audited this run; a refused lever stays a + candidate — refusal is an outcome, not an error + certified won a same-process audit under the current plan digest; + any digest change (shapes, weights, environment, gates) + demotes it back to candidate for re-audit — certification + never outlives the plan it was earned on + +Gate thresholds are switches too: they are folded into the plan digest +and written into the receipt, so a relaxed gate is always visible and +invalidates prior certifications instead of silently inheriting them. +""" + +from __future__ import annotations + +import hashlib +import json +import pathlib +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping + +import torch + +from .gates import parity_metrics +from .swap import AttachHandle, attach as _swap_attach + +_LEVER_KINDS = ("regions", "cadence", "seam_negotiation", "dtype", + "capture_form") + + +@dataclass +class Gates: + """Explicit gate thresholds. Part of the plan digest: changing one + is a plan change, not a tweak. + + ``parity_cos`` is the hard floor — below it the recipe refuses. + ``parity_warn`` marks the comfort line: results in + ``[parity_cos, parity_warn)`` pass but the receipt records + ``parity_band: "warn"`` with a note that calibration owes the + remaining accuracy. During the performance-assembly phase the floor + is intentionally loose (collapse-only); tighter banded gates arrive + with the calibration/fallback design, not before. + """ + + parity_cos: float = 0.99 + parity_warn: float = 0.999 + min_speedup: float = 1.02 + drift_budget: float = 0.02 + + def as_dict(self) -> dict[str, float]: + return {"parity_cos": self.parity_cos, + "parity_warn": self.parity_warn, + "min_speedup": self.min_speedup, + "drift_budget": self.drift_budget} + + +@dataclass +class Lever: + """One named yield switch: a group of swaps engaged and judged + together. + + ``build(model, ctx)`` returns either a mapping ``path -> module`` of + swaps, or a ``(swaps, outside_update)`` tuple where + ``outside_update`` is a callable the host must run at the lever's + cadence outside the captured stage (e.g. refreshing static KV + buffers on observation ticks). Outside updates are collected into + ``ctx["outside_updates"]`` before ``build_stage`` runs. + """ + + name: str + kind: str + build: Callable[[torch.nn.Module, dict], Any] | None = None + state: str = "candidate" + notes: str = "" + + def __post_init__(self) -> None: + if self.kind not in _LEVER_KINDS: + raise ValueError(f"lever kind {self.kind!r} not in " + f"{_LEVER_KINDS}") + if self.state not in ("off", "candidate", "certified"): + raise ValueError(f"lever state {self.state!r}") + + +@dataclass +class Arm: + """One executable form of the pipeline under audit. + + ``refs`` must pin every object whose device memory the arm's graph + reads (input buffers, static intermediates, closures holding them). + A captured graph keeps no Python references of its own: if a tensor + it reads is garbage-collected, the next arm's allocations reuse that + memory and this arm replays over foreign data — the same-process + baseline retime is exactly where that corruption surfaces. + """ + + tick: Callable[[], Any] + output: Callable[[], torch.Tensor] + teardown: Callable[[], None] | None = None + refs: Any = None + stage: Any = None # the CapturedStage, if the arm graphed one — + # lets the winning arm feed stage.export() + + +@dataclass +class Recipe: + """The declaration ``run_recipe`` executes. + + ``build_stage(model, ctx)`` constructs the hot stage for whatever is + currently attached to the model (it is called once for the untouched + baseline and once with the levers engaged) and returns an + :class:`Arm`. ``reference(model, ctx)`` produces the stock eager + output that anchors parity. ``between_arms`` runs after the baseline + arm is built and before levers attach — compile hosts pass + ``torch._dynamo.reset`` here so the treated arm gets a fresh + compile budget instead of inheriting cache-limit fallout. + """ + + name: str + levers: list[Lever] + build_stage: Callable[[torch.nn.Module, dict], Arm] + reference: Callable[[torch.nn.Module, dict], torch.Tensor] + gates: Gates = field(default_factory=Gates) + between_arms: Callable[[], None] | None = None + + +@dataclass +class RecipeRun: + """Outcome of one audit: verdict, evidence, and the live stage.""" + + verdict: str + receipt: dict[str, Any] + arm: Arm | None = None + _handle: AttachHandle | None = None + + def report(self) -> str: + r = self.receipt + lines = [f"recipe {r['recipe']}: {self.verdict}" + + (f" [{r['reason']}]" if r.get("reason") else "")] + for name, stat in r["levers"].items(): + demote = (f" (demoted: {stat['demoted']})" + if stat.get("demoted") else "") + lines.append(f" {name} [{stat['kind']}]: " + f"{stat['state_in']} -> {stat['state_out']}" + f", {stat.get('seams', 0)} seam(s){demote}") + base, treat = r.get("baseline"), r.get("treated") + if base: + lines.append(f" baseline {base['ms']:.2f} ms " + f"(retime {base.get('retime_ms')}, " + f"drift {base.get('drift')})") + if treat: + lines.append(f" treated {treat['ms']:.2f} ms " + f"({treat['speedup']:.3f}x), parity " + f"{treat['parity_vs_reference']:.6f}") + return "\n".join(lines) + + def detach(self) -> None: + if self.arm is not None and self.arm.teardown is not None: + self.arm.teardown() + self.arm = None + if self._handle is not None: + self._handle.detach() + self._handle = None + + def save_receipt(self, directory) -> pathlib.Path: + directory = pathlib.Path(directory) + directory.mkdir(parents=True, exist_ok=True) + path = directory / f"recipe_{self.receipt['recipe']}.json" + path.write_text(json.dumps(self.receipt, indent=2, default=str)) + return path + + +def _plan_digest(recipe: Recipe, model: torch.nn.Module, + ctx: Mapping[str, Any]) -> str: + ident = { + "recipe": recipe.name, + "levers": sorted((lv.name, lv.kind) for lv in recipe.levers), + "gates": recipe.gates.as_dict(), + "torch": torch.__version__, + "device": (torch.cuda.get_device_name() + if torch.cuda.is_available() else "cpu"), + "model": type(model).__name__, + "shape_sig": str(ctx.get("shape_sig", "")), + } + return hashlib.sha256(json.dumps( + ident, sort_keys=True, default=str).encode()).hexdigest() + + +def _time_ms(fn: Callable[[], Any], warmup: int, iters: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def run_recipe( + recipe: Recipe, + model: torch.nn.Module, + ctx: dict[str, Any] | None = None, + *, + receipts_dir: str | pathlib.Path | None = None, + reaudit: str = "always", + warmup: int = 5, + iters: int = 30, + verbose: bool = True, +) -> RecipeRun: + """Assemble, audit, and certify a recipe in one call. + + ``reaudit="always"`` runs the full same-process A/B. + ``reaudit="on_change"`` skips the timing audit when every engaged + lever is already certified under the current plan digest in the + stored receipt — parity is still re-checked, timing numbers carry + over marked ``cached``. Any digest mismatch forces the full audit. + """ + + def say(msg: str) -> None: + if verbose: + print(f"[recipe] {msg}", flush=True) + + ctx = ctx if ctx is not None else {} + gates = recipe.gates + digest = _plan_digest(recipe, model, ctx) + say(f"{recipe.name}: digest {digest[:12]}") + + prior = None + if receipts_dir is not None: + prior_path = (pathlib.Path(receipts_dir) + / f"recipe_{recipe.name}.json") + if prior_path.is_file(): + prior = json.loads(prior_path.read_text()) + + # ---- lever lifecycle: demote certifications the digest no longer + # covers, adopt ones the stored receipt still backs ---- + lever_stats: dict[str, dict[str, Any]] = {} + engaged: list[Lever] = [] + for lever in recipe.levers: + stat = {"kind": lever.kind, "state_in": lever.state, + "notes": lever.notes} + state = lever.state + prior_rec = (prior or {}).get("levers", {}).get(lever.name) + prior_certified = (prior is not None + and prior.get("digest") == digest + and prior_rec is not None + and prior_rec.get("state_out") == "certified") + if state == "certified" and not prior_certified: + state = "candidate" + stat["demoted"] = ("plan digest changed" + if prior is not None else "no stored receipt") + elif state == "candidate" and prior_certified: + state = "certified" + stat["state"] = state + lever_stats[lever.name] = stat + if state != "off": + engaged.append(lever) + if not engaged: + say("no engaged levers — nothing to audit") + receipt = {"recipe": recipe.name, "digest": digest, + "levers": {}, "verdict": "empty"} + return RecipeRun("empty", receipt) + + cached_ok = (reaudit == "on_change" and prior is not None + and prior.get("digest") == digest + and prior.get("verdict") == "win" + and all(lever_stats[lv.name]["state"] == "certified" + for lv in engaged)) + + # ---- stock reference + baseline arm ---- + with torch.no_grad(): + reference_out = recipe.reference(model, ctx).detach().float().cpu() + + base_ms = base_retime = drift = None + base_arm: Arm | None = None + if not cached_ok: + base_arm = recipe.build_stage(model, ctx) + base_arm.tick() # outputs are only defined after a full tick + base_out = base_arm.output().detach().float().cpu() + base_parity = parity_metrics(base_out, reference_out)["cosine"] + if base_parity < gates.parity_cos: + if base_arm.teardown is not None: + base_arm.teardown() + say(f"baseline arm parity {base_parity:.6f} < " + f"{gates.parity_cos} vs stock eager — audit invalid") + for stat in lever_stats.values(): + stat["state_out"] = stat.pop("state") + receipt = {"recipe": recipe.name, "digest": digest, + "levers": lever_stats, + "baseline": {"parity_vs_reference": base_parity}, + "verdict": "invalid_baseline", + "reason": "baseline arm does not match stock eager"} + return RecipeRun("invalid_baseline", receipt) + base_ms = _time_ms(base_arm.tick, warmup, iters) + say(f"baseline arm {base_ms:.2f} ms " + f"(parity {base_parity:.6f})") + + if recipe.between_arms is not None: + recipe.between_arms() + + # ---- engage levers (one transaction) ---- + swaps: dict[str, torch.nn.Module] = {} + updates: list[Callable[[], None]] = [] + refused_build: list[Lever] = [] + for lever in engaged: + try: + built = lever.build(model, ctx) if lever.build else {} + except ValueError as refusal: + # a lever may disqualify itself at build time (calibration + # shows its precondition does not hold on this host) — + # that is an outcome to record, not a reason to abort the + # other levers + lever_stats[lever.name]["state_out"] = "refused" + lever_stats[lever.name]["reason"] = str(refusal)[:120] + lever_stats[lever.name]["seams"] = 0 + refused_build.append(lever) + say(f"{lever.name}: refused at build " + f"[{str(refusal)[:80]}]") + continue + if isinstance(built, tuple): + lever_swaps, update = built + else: + lever_swaps, update = built, None + overlap = set(lever_swaps) & set(swaps) + if overlap: + raise ValueError(f"lever {lever.name!r} overlaps prior " + f"levers at {sorted(overlap)[:3]}") + swaps.update(lever_swaps) + if update is not None: + updates.append(update) + lever_stats[lever.name]["seams"] = len(lever_swaps) + ctx["outside_updates"] = updates + handle = _swap_attach(model, swaps) if swaps else None + say(f"engaged {len(engaged)} lever(s), {len(swaps)} seam(s)") + + # ---- treated arm: parity gate, then net-win gate ---- + arm = recipe.build_stage(model, ctx) + arm.tick() + treated_out = arm.output().detach().float().cpu() + parity = parity_metrics(treated_out, reference_out)["cosine"] + parity_ok = parity >= gates.parity_cos + parity_band = ("ok" if parity >= gates.parity_warn + else "warn" if parity_ok else "fail") + say(f"treated parity vs stock eager: {parity:.6f}" + + (" [WARN: below comfort line — calibration owes the rest]" + if parity_band == "warn" else "")) + + treated_ms = speedup = None + if parity_ok and cached_ok: + base_ms = prior["baseline"]["ms"] + treated_ms = prior["treated"]["ms"] + speedup = prior["treated"]["speedup"] + win = True + say(f"digest hit — timings carried from stored receipt " + f"({treated_ms:.2f} ms)") + elif parity_ok: + treated_ms = _time_ms(arm.tick, warmup, iters) + # baseline retime after the treated arm: same-process drift + # bound; the win must clear the *faster* of the two baselines + base_retime = _time_ms(base_arm.tick, warmup, iters) + drift = round(abs(base_retime - base_ms) / base_ms, 4) + base_floor = min(base_ms, base_retime) + speedup = round(base_floor / treated_ms, 4) + win = speedup >= gates.min_speedup + say(f"treated {treated_ms:.2f} ms vs baseline floor " + f"{base_floor:.2f} ({speedup:.3f}x, drift {drift})") + else: + win = False + + if base_arm is not None and base_arm.teardown is not None: + base_arm.teardown() + + verdict = "win" if win else "refused" + reason = None + if not parity_ok: + reason = f"parity {parity:.6f} < {gates.parity_cos}" + elif not win: + reason = f"no net win ({speedup}x < {gates.min_speedup})" + for lever in engaged: + if lever in refused_build: + continue + lever_stats[lever.name]["state_out"] = ( + "certified" if win else "candidate") + for name, stat in lever_stats.items(): + stat.setdefault("state_out", stat["state"]) + stat.pop("state", None) + + if not win: + if arm.teardown is not None: + arm.teardown() + arm = None + if handle is not None: + handle.detach() + handle = None + say(f"refused [{reason}] — model restored untouched") + else: + say(f"win: {len(engaged)} lever(s) certified under digest " + f"{digest[:12]}") + + receipt = { + "recipe": recipe.name, "digest": digest, + "model": type(model).__name__, + "device": (torch.cuda.get_device_name() + if torch.cuda.is_available() else "cpu"), + "torch": torch.__version__, + "gates": gates.as_dict(), + "reaudit": reaudit, "cached": bool(cached_ok), + "levers": lever_stats, + "baseline": {"ms": None if base_ms is None else round(base_ms, 3), + "retime_ms": (None if base_retime is None + else round(base_retime, 3)), + "drift": drift}, + "treated": (None if treated_ms is None else + {"ms": round(treated_ms, 3), "speedup": speedup, + "parity_vs_reference": round(parity, 7), + "parity_band": parity_band}), + "verdict": verdict, "reason": reason, + } + run = RecipeRun(verdict, receipt, arm=arm, _handle=handle) + if receipts_dir is not None: + path = run.save_receipt(receipts_dir) + say(f"receipt -> {path}") + return run diff --git a/flash_rt/structures/registry.py b/flash_rt/structures/registry.py new file mode 100644 index 00000000..15fe324f --- /dev/null +++ b/flash_rt/structures/registry.py @@ -0,0 +1,105 @@ +"""Structure catalog registry — pure lookup, no execution logic. + +Loads structure specifications from the on-disk catalog and resolves +their reference implementations. Dispatch, tuning, calibration, and +activation live in separate layers; the registry only answers "what is +structure X and where is its ground truth". +""" + +from __future__ import annotations + +import importlib +import pathlib +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, Sequence + +import yaml + +_CATALOG_DIR = pathlib.Path(__file__).resolve().parent / "catalog" + + +@dataclass(frozen=True) +class StructureSpec: + """One catalog entry, as declared in ``catalog//structure.yaml``. + + Two kinds exist. ``region`` structures declare a tensor boundary, + weight slots and a torch reference. ``stage_pipeline`` structures + declare a stage graph with cadence attributes; their parity reference + is the host's own eager path, so ``boundary``/``weights``/``reference`` + are empty for them and ``stages`` is populated instead. + """ + + name: str + version: int + description: str + boundary: Mapping[str, Any] + weights: Sequence[Mapping[str, Any]] + variants: Mapping[str, Sequence[str]] + calibration: Mapping[str, Any] + gates: Mapping[str, Any] + _reference: Mapping[str, str] = field(repr=False) + kind: str = "region" + family: str = "" + stages: Sequence[Mapping[str, Any]] = () + conformance: Sequence[str] = () + + @property + def symbolic_dims(self) -> Sequence[str]: + return tuple(self.boundary.get("symbolic_dims", ())) + + @property + def weight_slots(self) -> Sequence[str]: + return tuple(entry["slot"] for entry in self.weights) + + def reference(self) -> Callable[..., Any]: + """Resolve the reference entrypoint (ground truth for gates).""" + if not self._reference: + raise LookupError( + f"structure {self.name!r} (kind={self.kind!r}) has no " + "standalone reference; its parity ground truth is the " + "host's own eager path under the same noise window" + ) + module = importlib.import_module( + f"{__package__}.catalog.{self._reference['module']}" + ) + return getattr(module, self._reference["entrypoint"]) + + +def list_structures() -> list[str]: + """Names of all structures present in the catalog.""" + return sorted( + path.parent.name for path in _CATALOG_DIR.glob("*/structure.yaml") + ) + + +def load(name: str) -> StructureSpec: + """Load one structure specification by catalog name.""" + path = _CATALOG_DIR / name / "structure.yaml" + if not path.is_file(): + raise KeyError(f"unknown structure: {name!r}") + with open(path, "r", encoding="utf-8") as handle: + data = yaml.safe_load(handle) + kind = data.get("kind", "region") + spec = StructureSpec( + name=data["structure"], + version=int(data["version"]), + description=str(data.get("description", "")).strip(), + boundary=data.get("boundary", {}) if kind != "region" + else data["boundary"], + weights=data.get("weights", ()) if kind != "region" + else data["weights"], + variants=data.get("variants", {}), + calibration=data.get("calibration", {}), + gates=data["gates"], + _reference=data.get("reference", {}) if kind != "region" + else data["reference"], + kind=kind, + family=str(data.get("family", "")), + stages=data.get("stages", ()), + conformance=data.get("conformance", ()), + ) + if spec.name != name: + raise ValueError( + f"catalog directory {name!r} declares structure {spec.name!r}" + ) + return spec diff --git a/flash_rt/structures/stages.py b/flash_rt/structures/stages.py new file mode 100644 index 00000000..1296c1da --- /dev/null +++ b/flash_rt/structures/stages.py @@ -0,0 +1,177 @@ +"""Capture door: ``structures.capture(fn, ...)`` — graph a hot stage. + +The schedule-structure counterpart of the region doors. ``fn`` is the +hot stage of a cond_iter pipeline (a denoise loop, a decode step chain); +``capture`` warms it, records it into a CUDA graph on a side stream, and +returns a runtime whose ``replay()`` re-executes the whole stage as one +launch. Declared ``windows`` are the tensors you may rewrite between +replays (noise, observations, condition buffers) — replay reads their +current contents in place, which is what makes the graph a reusable +stage rather than a frozen trace. + +Gates are built in: ``reference`` (an eager thunk producing the same +output under the same window contents) certifies parity, and the replay +is timed against eager ``fn``; a capture that is not both accurate and +net-faster raises ``CaptureRefused`` — the host keeps its eager path, +never a silently wrong or slower graph. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping + +import torch + +from .gates import parity_metrics + + +class CaptureRefused(RuntimeError): + """The captured stage failed its parity or net-win gate.""" + + +def _time_ms(fn: Callable[[], Any], warmup: int = 5, iters: int = 20) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +@dataclass +class CapturedStage: + """A replayable hot stage with declared swap windows.""" + + graph: torch.cuda.CUDAGraph + stream: torch.cuda.Stream + output: Any + windows: Mapping[str, torch.Tensor] + certification: dict[str, Any] = field(default_factory=dict) + + def replay(self, sync: bool = True) -> Any: + with torch.cuda.stream(self.stream): + self.graph.replay() + if sync: + torch.cuda.synchronize() + return self.output + + def write(self, name: str, value: torch.Tensor) -> None: + """Rewrite one declared window in place.""" + self.windows[name].copy_(value) + + def export(self, *, ports, regions=(), stages=None, roles=None, + identity=None, manifest_extra=None): + """One call from this captured stage to ``frt_model_runtime_v1``. + + The declared windows become the contract's boundary windows + (name -> device pointer + bytes); ``ports`` reference them by + name exactly as in + :func:`flash_rt.structures.provider.export_captured_runtime`. + This is the absorption edge: a stage captured out of any torch + host becomes a runtime the FlashRT serving mechanisms (Nexus + tick, capsule snapshot/restore) consume like a whitebox one. + """ + from .provider import export_captured_runtime + + window_map = { + name: (t.data_ptr(), t.numel() * t.element_size()) + for name, t in self.windows.items()} + return export_captured_runtime( + stream_handle=self.stream.cuda_stream, + graphs=[("infer", self.graph.raw_cuda_graph_exec())], + windows=window_map, + ports=ports, regions=regions, + stages=tuple(stages) if stages else ("infer",), + roles=roles, identity=identity, + manifest_extra=manifest_extra) + + +def capture( + fn: Callable[[], Any], + *, + windows: Mapping[str, torch.Tensor] | None = None, + reference: Callable[[], torch.Tensor] | None = None, + output_of: Callable[[Any], torch.Tensor] | None = None, + warmup: int = 3, + gate_cos: float = 0.999, + min_speedup: float = 1.02, + verbose: bool = True, +) -> CapturedStage: + """Capture ``fn`` into a replayable, gated stage. + + ``fn`` must be graph-safe: fixed shapes, no host-side branching on + tensor values, all varying inputs read from the declared ``windows`` + buffers. ``reference`` runs the host's own eager path under the same + window contents; parity is judged between its output and the + replayed stage output. Pass ``gate_cos=0`` / ``min_speedup=0`` to + skip a gate explicitly (recorded in the certification). + """ + + def say(msg: str) -> None: + if verbose: + print(f"[structures] {msg}", flush=True) + + windows = dict(windows or {}) + stream = torch.cuda.Stream() + with torch.no_grad(), torch.cuda.stream(stream): + for _ in range(max(1, warmup)): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = fn() + torch.cuda.synchronize() + say(f"stage captured ({len(windows)} window(s))") + + stage = CapturedStage(graph=graph, stream=stream, output=output, + windows=windows) + pick = output_of or (lambda out: out) + + cert: dict[str, Any] = {"windows": sorted(windows), + "gate_cos": gate_cos, + "min_speedup": min_speedup} + if reference is not None and gate_cos: + with torch.no_grad(): + want = reference().detach().float().cpu() + got = pick(stage.replay()).detach().float().cpu() + cos = parity_metrics(got, want)["cosine"] + cert["parity_cos"] = round(cos, 7) + if cos < gate_cos: + raise CaptureRefused( + f"captured stage parity cos {cos:.6f} < {gate_cos} vs " + "eager reference — check that every varying input is a " + "declared window (in-graph RNG never matches eager)") + say(f"parity vs eager: cos={cos:.6f}") + if min_speedup: + eager_ms = _time_ms(lambda: fn()) + # replays run on the side stream: the timing events must be + # recorded on that same stream or they only measure enqueue time + for _ in range(5): + stage.replay(sync=False) + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + with torch.cuda.stream(stream): + start.record() + for _ in range(20): + stage.replay(sync=False) + with torch.cuda.stream(stream): + end.record() + torch.cuda.synchronize() + replay_ms = start.elapsed_time(end) / 20 + cert["eager_ms"] = round(eager_ms, 3) + cert["replay_ms"] = round(replay_ms, 3) + cert["speedup"] = round(eager_ms / replay_ms, 4) + if eager_ms / replay_ms < min_speedup: + raise CaptureRefused( + f"captured stage is not a net win " + f"({eager_ms / replay_ms:.3f}x vs eager, margin " + f"{min_speedup}) — host keeps its eager path") + say(f"net win: {eager_ms:.2f} -> {replay_ms:.2f} ms " + f"({eager_ms / replay_ms:.3f}x)") + stage.certification = cert + return stage diff --git a/flash_rt/structures/swap.py b/flash_rt/structures/swap.py new file mode 100644 index 00000000..ddac5dfa --- /dev/null +++ b/flash_rt/structures/swap.py @@ -0,0 +1,77 @@ +"""Transactional module attachment — mechanism only, no policy. + +Swaps host modules for bound structure implementations atomically: every +staged swap is applied only if all of them can be applied, and a handle +restores the originals exactly. Which modules to swap, and whether an +implementation passed its gates, is decided by the caller before +staging; this layer refuses partial application by construction. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping + +import torch + + +def resolve_parent(root: torch.nn.Module, path: str) -> tuple[torch.nn.Module, str]: + """Resolve the parent module and attribute name for a dotted path.""" + parts = path.split(".") + parent = root + for part in parts[:-1]: + parent = parent[int(part)] if part.isdigit() else getattr(parent, part) + attr = parts[-1] + leaf = parent[int(attr)] if attr.isdigit() else getattr(parent, attr) + if not isinstance(leaf, torch.nn.Module): + raise TypeError(f"{path!r} does not resolve to a torch.nn.Module") + return parent, attr + + +@dataclass +class AttachHandle: + """Restores the originals of one committed attachment.""" + + _entries: list[tuple[torch.nn.Module, str, torch.nn.Module]] + records: Mapping[str, Any] = field(default_factory=dict) + active: bool = True + + def detach(self) -> None: + """Restore every swapped module. Idempotent.""" + if not self.active: + return + for parent, attr, original in reversed(self._entries): + setattr(parent, attr, original) + self.active = False + + +def attach( + root: torch.nn.Module, + swaps: Mapping[str, torch.nn.Module], + *, + records: Mapping[str, Any] | None = None, +) -> AttachHandle: + """Atomically replace the modules at ``swaps`` paths. + + All paths are resolved and validated before the first swap is + applied; any resolution failure leaves the model untouched. Callers + pass the qualification ``records`` that justified the attachment so + the handle carries its own evidence. + """ + if not swaps: + raise ValueError("no swaps staged") + staged: list[tuple[torch.nn.Module, str, torch.nn.Module, torch.nn.Module]] = [] + for path, replacement in swaps.items(): + parent, attr = resolve_parent(root, path) + staged.append((parent, attr, getattr(parent, attr), replacement)) + + entries = [] + try: + for parent, attr, original, replacement in staged: + setattr(parent, attr, replacement) + entries.append((parent, attr, original)) + except Exception: + for parent, attr, original in reversed(entries): + setattr(parent, attr, original) + raise + return AttachHandle(_entries=entries, records=dict(records or {}))