flashrt-structures 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (122) hide show
  1. flashrt_structures/__init__.py +174 -0
  2. flashrt_structures/adapters/__init__.py +55 -0
  3. flashrt_structures/adapters/diffusers_attention.py +237 -0
  4. flashrt_structures/adapters/diffusers_rotary_attention.py +239 -0
  5. flashrt_structures/adapters/factored_qk_norm_rope.py +252 -0
  6. flashrt_structures/adapters/factored_two_way_attention.py +99 -0
  7. flashrt_structures/adapters/gemma_attention.py +226 -0
  8. flashrt_structures/adapters/packed_qkv_rope.py +342 -0
  9. flashrt_structures/adapters/packed_stream_qk_norm_rope.py +376 -0
  10. flashrt_structures/adapters/qwen_per_head_qk_norm_rope.py +231 -0
  11. flashrt_structures/adapters/sglang_engine.py +207 -0
  12. flashrt_structures/adapters/transformers_attention_interface.py +73 -0
  13. flashrt_structures/adapters/transformers_gated_delta.py +152 -0
  14. flashrt_structures/adapters/transformers_gated_delta_fused.py +96 -0
  15. flashrt_structures/adapters/vllm_engine.py +424 -0
  16. flashrt_structures/adjudicate.py +85 -0
  17. flashrt_structures/aot.py +191 -0
  18. flashrt_structures/autobuild.py +2052 -0
  19. flashrt_structures/beta/__init__.py +43 -0
  20. flashrt_structures/beta/conform.py +94 -0
  21. flashrt_structures/beta/joins.py +113 -0
  22. flashrt_structures/beta/negotiate.py +84 -0
  23. flashrt_structures/beta/ports.py +140 -0
  24. flashrt_structures/decisions.py +80 -0
  25. flashrt_structures/discover.py +623 -0
  26. flashrt_structures/explain.py +70 -0
  27. flashrt_structures/frontdoor.py +572 -0
  28. flashrt_structures/gates.py +465 -0
  29. flashrt_structures/guard.py +421 -0
  30. flashrt_structures/handle.py +189 -0
  31. flashrt_structures/impls/__init__.py +219 -0
  32. flashrt_structures/impls/adaln_producer/__init__.py +8 -0
  33. flashrt_structures/impls/adaln_producer/broker.py +116 -0
  34. flashrt_structures/impls/adaln_producer/fused.py +388 -0
  35. flashrt_structures/impls/adarms_stack/__init__.py +8 -0
  36. flashrt_structures/impls/adarms_stack/fp8_chain.py +832 -0
  37. flashrt_structures/impls/adarms_stack/region.py +102 -0
  38. flashrt_structures/impls/attention_core/__init__.py +132 -0
  39. flashrt_structures/impls/attention_core/fa2_seqused.py +458 -0
  40. flashrt_structures/impls/attention_core/fa4_cute.py +154 -0
  41. flashrt_structures/impls/attention_core/fa4_fp8.py +178 -0
  42. flashrt_structures/impls/attention_core/masked_mha.py +158 -0
  43. flashrt_structures/impls/attention_core/two_way_fa2.py +220 -0
  44. flashrt_structures/impls/cadence_static/__init__.py +18 -0
  45. flashrt_structures/impls/cadence_static/buffers.py +122 -0
  46. flashrt_structures/impls/cadence_static/cross_attention.py +187 -0
  47. flashrt_structures/impls/chain_elements.py +89 -0
  48. flashrt_structures/impls/decode_loop/__init__.py +0 -0
  49. flashrt_structures/impls/decode_loop/fp8_kv.py +206 -0
  50. flashrt_structures/impls/decode_loop/mtp_speculative.py +245 -0
  51. flashrt_structures/impls/decode_loop/whole_step.py +852 -0
  52. flashrt_structures/impls/decoder_block/__init__.py +6 -0
  53. flashrt_structures/impls/decoder_block/attn_sublayer.py +110 -0
  54. flashrt_structures/impls/decoder_block/fused.py +167 -0
  55. flashrt_structures/impls/decoder_ffn/__init__.py +0 -0
  56. flashrt_structures/impls/decoder_ffn/fp8_static.py +310 -0
  57. flashrt_structures/impls/decoder_ffn/fp8_static.yaml +22 -0
  58. flashrt_structures/impls/decoder_ffn/w4a16_static.py +221 -0
  59. flashrt_structures/impls/decoder_ffn/w8a16_static.py +183 -0
  60. flashrt_structures/impls/dit_stack/__init__.py +11 -0
  61. flashrt_structures/impls/dit_stack/fp4_chain.py +417 -0
  62. flashrt_structures/impls/dit_stack/region.py +86 -0
  63. flashrt_structures/impls/fixed_iter/__init__.py +29 -0
  64. flashrt_structures/impls/fixed_iter/openpi.py +264 -0
  65. flashrt_structures/impls/fixed_iter/protocol.py +94 -0
  66. flashrt_structures/impls/gated_delta_core/__init__.py +3 -0
  67. flashrt_structures/impls/gated_delta_core/fused_layer.py +545 -0
  68. flashrt_structures/impls/gated_delta_core/hub_v3.py +152 -0
  69. flashrt_structures/impls/graph_lowering/__init__.py +27 -0
  70. flashrt_structures/impls/graph_lowering/pi052_denoise.py +179 -0
  71. flashrt_structures/impls/graph_lowering/protocol.py +76 -0
  72. flashrt_structures/impls/graph_lowering/qwen3_vl.py +364 -0
  73. flashrt_structures/impls/linear_proj/__init__.py +0 -0
  74. flashrt_structures/impls/linear_proj/fp8_static.py +270 -0
  75. flashrt_structures/impls/linear_proj/nvfp4_balance.py +131 -0
  76. flashrt_structures/impls/linear_proj/nvfp4_dynamic.py +182 -0
  77. flashrt_structures/impls/linear_proj/w8a16_static.py +230 -0
  78. flashrt_structures/impls/modnorm_qkv_chain/__init__.py +0 -0
  79. flashrt_structures/impls/modnorm_qkv_chain/fp8_ptok_table.py +291 -0
  80. flashrt_structures/impls/moe_experts/__init__.py +9 -0
  81. flashrt_structures/impls/moe_experts/nvfp4_dynamic.py +208 -0
  82. flashrt_structures/impls/moe_experts/nvfp4_w4a16.py +129 -0
  83. flashrt_structures/impls/norm_fused/__init__.py +3 -0
  84. flashrt_structures/impls/norm_fused/bf16.py +94 -0
  85. flashrt_structures/impls/norm_fused/fp8_producer.py +84 -0
  86. flashrt_structures/impls/patch_projection/__init__.py +3 -0
  87. flashrt_structures/impls/patch_projection/bf16_flat.py +139 -0
  88. flashrt_structures/impls/prefill_tower/__init__.py +10 -0
  89. flashrt_structures/impls/prefill_tower/fp8_chain.py +955 -0
  90. flashrt_structures/impls/prefill_tower/region.py +99 -0
  91. flashrt_structures/impls/qk_norm_rope/__init__.py +12 -0
  92. flashrt_structures/impls/qk_norm_rope/per_head_gqa.py +199 -0
  93. flashrt_structures/impls/qk_norm_rope/projection_bf16.py +165 -0
  94. flashrt_structures/impls/qkv_pack/__init__.py +5 -0
  95. flashrt_structures/impls/qkv_pack/bf16.py +110 -0
  96. flashrt_structures/impls/qkv_pack/fp8_static.py +435 -0
  97. flashrt_structures/impls/qkv_pack/nvfp4_balance.py +218 -0
  98. flashrt_structures/impls/qkv_rope/__init__.py +3 -0
  99. flashrt_structures/impls/qkv_rope/packed_bias_bf16.py +143 -0
  100. flashrt_structures/impls/step_table.py +113 -0
  101. flashrt_structures/impls/vision_ffn/__init__.py +0 -0
  102. flashrt_structures/impls/vision_ffn/fp8_static.py +261 -0
  103. flashrt_structures/impls/vision_ffn/nvfp4_balance.py +211 -0
  104. flashrt_structures/impls/vision_tower/__init__.py +7 -0
  105. flashrt_structures/impls/vision_tower/fp8_chain.py +533 -0
  106. flashrt_structures/impls/vision_tower/region.py +91 -0
  107. flashrt_structures/matrix.py +126 -0
  108. flashrt_structures/points.py +368 -0
  109. flashrt_structures/prequantized.py +131 -0
  110. flashrt_structures/quantize_on_adopt.py +94 -0
  111. flashrt_structures/recipe.py +438 -0
  112. flashrt_structures/regions.py +208 -0
  113. flashrt_structures/schemes.py +490 -0
  114. flashrt_structures/stages.py +298 -0
  115. flashrt_structures/storage.py +255 -0
  116. flashrt_structures/swap.py +398 -0
  117. flashrt_structures/workspace.py +110 -0
  118. flashrt_structures-0.2.0.dist-info/METADATA +140 -0
  119. flashrt_structures-0.2.0.dist-info/RECORD +122 -0
  120. flashrt_structures-0.2.0.dist-info/WHEEL +5 -0
  121. flashrt_structures-0.2.0.dist-info/licenses/LICENSE +202 -0
  122. flashrt_structures-0.2.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,70 @@
1
+ """structures.explain — the coverage table, from the plan's own notes.
2
+
3
+ Adoption lives or dies on one question: *what did the system actually
4
+ do to my model, and why not more?* The plan already knows — discovered
5
+ seams, bound swaps, seams kept at host precision with the scheme's
6
+ reasons, refusals with theirs, adapter routes. This renders that
7
+ knowledge as one table instead of leaving it in ``plan.notes``.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from typing import Any
13
+
14
+
15
+ def explain(plan: Any) -> str:
16
+ """Render one plan's coverage as human-readable text."""
17
+ lines: list[str] = []
18
+ notes = getattr(plan, "notes", {}) or {}
19
+ scheme = notes.get("scheme", {}) or {}
20
+ name = scheme.get("name") or notes.get("scheme_name") or "?"
21
+ lines.append(f"scheme: {name}"
22
+ + (" (auto)" if scheme.get("auto") else ""))
23
+
24
+ swaps = getattr(plan, "swaps", {}) or {}
25
+ observed = getattr(plan, "observed", {}) or {}
26
+ lines.append(f"bound: {len(swaps)} swapped seam(s), "
27
+ f"{len(observed)} adapter-routed seam(s)")
28
+ by_kind: dict[str, int] = {}
29
+ for path in swaps:
30
+ seam = None
31
+ for s in getattr(plan, "seams", []) or []:
32
+ if str(getattr(s, "path", "")) == str(path):
33
+ seam = s
34
+ break
35
+ kind = getattr(seam, "structure", None) or "seam"
36
+ by_kind[kind] = by_kind.get(kind, 0) + 1
37
+ for kind in sorted(by_kind):
38
+ lines.append(f" {kind}: {by_kind[kind]}")
39
+ adapter = notes.get("gated_delta_adapter")
40
+ if adapter:
41
+ lines.append(f" gated-delta adapter: {adapter}")
42
+
43
+ routed = scheme.get("formats", {}) or {}
44
+ if routed:
45
+ lines.append(f"routed to non-default formats: {len(routed)}")
46
+ for path, fmt in list(sorted(routed.items()))[:8]:
47
+ lines.append(f" {path} -> {fmt}")
48
+ if len(routed) > 8:
49
+ lines.append(f" ... {len(routed) - 8} more")
50
+
51
+ kept = scheme.get("keep_host", {}) or {}
52
+ if kept:
53
+ lines.append(f"kept at host precision: {len(kept)}")
54
+ for path, why in list(sorted(kept.items()))[:8]:
55
+ lines.append(f" {path}: {why or 'scheme decision'}")
56
+ if len(kept) > 8:
57
+ lines.append(f" ... {len(kept) - 8} more")
58
+
59
+ refused = notes.get("refused", []) or []
60
+ if refused:
61
+ from collections import Counter
62
+
63
+ lines.append(f"refused: {len(refused)}")
64
+ reasons = Counter(str(r[1] if isinstance(r, (tuple, list))
65
+ else r)[:88] for r in refused)
66
+ for why, cnt in reasons.most_common(8):
67
+ lines.append(f" x{cnt}: {why}")
68
+ if not refused and not kept:
69
+ lines.append("refused: 0")
70
+ return "\n".join(lines)
@@ -0,0 +1,572 @@
1
+ """One-call front door: ``structures.attach(model, forward, ...)``.
2
+
3
+ The consumption contract mirrors ``kernels.get_kernel``: one import, one
4
+ call. Everything the structure layer needs — seam discovery, real
5
+ distribution calibration, accuracy and net-win gates, transactional swap,
6
+ receipt — runs inside the call. An attachment that does not both stay
7
+ accurate and win latency is refused; the model is left untouched and the
8
+ refusal is reported, never silently absorbed.
9
+
10
+ Three things about this gate are deliberate, and each replaced something
11
+ that used to be assumed:
12
+
13
+ **It gates what ``auto_swaps`` binds, not a subset of it.** Binding lives
14
+ in one place (:mod:`.autobuild`) and this module only judges. The earlier
15
+ arrangement had its own binding path covering three structures, so the
16
+ one call a caller was told to make was the one call that judged the least.
17
+
18
+ **The accuracy metric follows the host's output type.** A cosine over a
19
+ whole logits tensor falls as the sequence grows while token agreement
20
+ rises (see :mod:`.gates`); scoring a language host that way measures
21
+ sequence length. Distribution outputs are judged on top-1 agreement and
22
+ the last position, value outputs on cosine.
23
+
24
+ **Latency is judged by paired alternating timing.** Timing one arm and
25
+ then the other attributes machine drift to whichever arm it landed on;
26
+ this stack has produced an 8% spread between runs of the same
27
+ configuration, which is four times the margin the gate is asked to
28
+ resolve. Both arms are timed in every round and the decision is the
29
+ median of the per-round ratios.
30
+
31
+ And one thing it checks that nothing used to: after each scoring forward
32
+ it reads the attachment's ledger. A family whose seams fell back to the
33
+ host module did not run, however good its parity looks — that parity is
34
+ the host's own.
35
+ """
36
+
37
+ from __future__ import annotations
38
+
39
+ import hashlib
40
+ import json
41
+ import pathlib
42
+ import statistics
43
+ import time
44
+ import warnings
45
+ from dataclasses import dataclass, field
46
+ from typing import Any, Callable, Iterable, Mapping, Sequence
47
+
48
+ import torch
49
+
50
+ from .autobuild import AutoPlan, _layer_of, auto_swaps
51
+ from .gates import (DEFAULT_FLOORS, band_note, band_of, infer_output_kind,
52
+ metrics_for, passes)
53
+ from .guard import GuardRefused
54
+ from .swap import AttachHandle as _AttachHandle, attach as _swap_attach
55
+
56
+ #: the full catalog, which is what "one call" has to mean
57
+ ALL_STRUCTURES = ("decoder_ffn", "vision_ffn", "qkv_pack", "adaln_producer",
58
+ "linear_proj", "patch_projection", "norm_fused",
59
+ "attention_core",
60
+ "decoder_block", "modnorm_qkv_chain", "qk_norm_rope",
61
+ "qkv_rope", "gated_delta_core")
62
+
63
+ #: structure name per implementation class, for swaps whose path is not
64
+ #: itself a discovered seam (a pack's sibling readers, a composed block's
65
+ #: absorbed children). The type of the module that was bound is a local
66
+ #: fact; inferring it from the path would be a naming guess.
67
+ _STRUCTURE_BY_IMPL = {
68
+ "FusedGeGluMlp": "decoder_ffn", "FusedGluMlpW8A16": "decoder_ffn",
69
+ "FusedGeluMlp": "vision_ffn", "FusedLinearProj": "linear_proj",
70
+ "FlatPatchProjection": "patch_projection",
71
+ "PackedLinear": "qkv_pack", "StashReader": "qkv_pack",
72
+ "AttnBlockPacked": "qkv_pack", "AdaLNProducer": "adaln_producer",
73
+ # the vision pre-FFN norm producer is one gate unit with its FFN
74
+ # consumer: judged apart, an on-producer/off-consumer arm hands the
75
+ # host MLP an FP8 tensor (the exact failure the note below names)
76
+ "FusedNormFp8Producer": "vision_ffn",
77
+ "StyleTable": "adaln_producer", "FusedNorm": "norm_fused",
78
+ "FusedDecoderBlock": "decoder_block", "StaticOutput": "cadence_static",
79
+ }
80
+
81
+ #: a negotiated fp8 chain is one gate unit. The producer emits fp8 under a
82
+ #: scale the consumer was bound for, so attaching one without the other
83
+ #: hands the consumer a dtype it refuses — which the ledger would report
84
+ #: as a family that fell back, from a split this gate created itself.
85
+ _CHAIN = "negotiated_fp8_chain"
86
+ _ROUTED = "attention_core_routed"
87
+
88
+
89
+ def _cuda_time_ms(fn: Callable[[], Any], warmup: int = 3,
90
+ iters: int = 10) -> float:
91
+ with torch.no_grad():
92
+ for _ in range(warmup):
93
+ fn()
94
+ if not torch.cuda.is_available():
95
+ t0 = time.perf_counter()
96
+ for _ in range(iters):
97
+ fn()
98
+ return (time.perf_counter() - t0) * 1e3 / iters
99
+ torch.cuda.synchronize()
100
+ start, end = torch.cuda.Event(True), torch.cuda.Event(True)
101
+ start.record()
102
+ for _ in range(iters):
103
+ fn()
104
+ end.record()
105
+ torch.cuda.synchronize()
106
+ return start.elapsed_time(end) / iters
107
+
108
+
109
+ def _paired_ab(thunk: Callable[[], Any], on: Callable[[], Any],
110
+ off: Callable[[], Any], *, rounds: int = 5,
111
+ iters: int = 10) -> dict[str, float]:
112
+ """Time both arms in every round; decide on the median paired ratio.
113
+
114
+ Drift on this machine lands on whichever arm happens to be running,
115
+ so an A-then-B measurement attributes it to that arm. Alternating and
116
+ pairing cancels it: every ratio comes from two measurements taken
117
+ seconds apart, and the spread across rounds says how much to trust
118
+ the answer rather than leaving it to be assumed.
119
+ """
120
+ rows: list[tuple[float, float]] = []
121
+ for _ in range(max(1, rounds)):
122
+ on()
123
+ treated = _cuda_time_ms(thunk, warmup=1, iters=iters)
124
+ off()
125
+ base = _cuda_time_ms(thunk, warmup=1, iters=iters)
126
+ rows.append((treated, base))
127
+ ratios = sorted(b / t for t, b in rows)
128
+ return {
129
+ "ms": round(statistics.median(t for t, _ in rows), 3),
130
+ "base_ms": round(statistics.median(b for _, b in rows), 3),
131
+ "speedup": round(statistics.median(ratios), 4),
132
+ "speedup_min": round(ratios[0], 4),
133
+ "speedup_max": round(ratios[-1], 4),
134
+ "spread": round(ratios[-1] - ratios[0], 4),
135
+ "rounds": len(rows),
136
+ }
137
+
138
+
139
+ def _score_tensor(value: Any) -> torch.Tensor | None:
140
+ """The tensor a host's output should be judged on.
141
+
142
+ Logits when the host produces them, otherwise the first tensor in a
143
+ deterministic walk. Mappings are walked in sorted key order so the two
144
+ arms of a comparison never pick different leaves.
145
+ """
146
+ if value is None:
147
+ return None
148
+ if torch.is_tensor(value):
149
+ return value
150
+ logits = getattr(value, "logits", None)
151
+ if torch.is_tensor(logits):
152
+ return logits
153
+ if isinstance(value, Mapping):
154
+ if torch.is_tensor(value.get("logits")):
155
+ return value["logits"]
156
+ for key in sorted(value):
157
+ found = _score_tensor(value[key])
158
+ if found is not None:
159
+ return found
160
+ return None
161
+ if isinstance(value, (tuple, list)):
162
+ for item in value:
163
+ found = _score_tensor(item)
164
+ if found is not None:
165
+ return found
166
+ return None
167
+
168
+
169
+ def _gate_groups(plan: AutoPlan) -> dict[str, dict[str, torch.nn.Module]]:
170
+ """Split the plan into units that can be judged independently."""
171
+ structure_of = {s.path: s.structure for s in plan.seams}
172
+ negotiated = set(plan.notes.get("negotiated_layers", ()))
173
+ groups: dict[str, dict[str, torch.nn.Module]] = {}
174
+ for path, module in plan.swaps.items():
175
+ if negotiated and _layer_of(path) in negotiated:
176
+ key = _CHAIN
177
+ else:
178
+ key = (structure_of.get(path)
179
+ or _STRUCTURE_BY_IMPL.get(type(module).__name__)
180
+ or "other")
181
+ groups.setdefault(key, {})[path] = module
182
+ return groups
183
+
184
+
185
+ class _Arm:
186
+ """One gate unit, switchable on and off for the paired timing loop.
187
+
188
+ Holds a single handle rather than one per switch: attaching twice over
189
+ the same paths would record the first replacement as the "original",
190
+ and detaching would then restore a structure instead of the host.
191
+ Idempotent on both sides so the timing loop can call them freely.
192
+ """
193
+
194
+ def __init__(self, model: torch.nn.Module,
195
+ swaps: Mapping[str, torch.nn.Module], plan: AutoPlan,
196
+ mode: str, *, routed: bool = False,
197
+ revert: Any = None) -> None:
198
+ self.model = model
199
+ self.swaps = dict(swaps)
200
+ self.plan = plan
201
+ self.mode = mode
202
+ self.routed = routed
203
+ self.revert = revert
204
+ self.handle: _AttachHandle | None = None
205
+
206
+ def on(self) -> None:
207
+ if self.handle is None:
208
+ self.handle = _swap_attach(
209
+ self.model, self.swaps,
210
+ observe=self.plan.observed if self.routed else None,
211
+ on_guard_fail=self.mode, revert=self.revert)
212
+ if self.routed:
213
+ self.plan.enable_routed()
214
+
215
+ def off(self) -> None:
216
+ if self.handle is not None:
217
+ self.handle.detach()
218
+ self.handle = None
219
+ if self.routed:
220
+ self.plan.disable_routed()
221
+
222
+
223
+ @dataclass
224
+ class Plan:
225
+ """Result of one ``attach`` call: what was activated, why, evidence."""
226
+
227
+ activated: dict[str, torch.nn.Module]
228
+ families: dict[str, dict[str, Any]]
229
+ receipt: dict[str, Any]
230
+ _handle: _AttachHandle | None = None
231
+ _plan: AutoPlan | None = None
232
+ active: bool = field(init=False)
233
+
234
+ def __post_init__(self) -> None:
235
+ self.active = self._handle is not None
236
+
237
+ def report(self) -> str:
238
+ lines = [f"structures.attach: {len(self.activated)} seam(s) active"]
239
+ for name, stat in self.families.items():
240
+ line = (f" {name}: {stat['seams']} seam(s) -> "
241
+ f"{stat['outcome']}")
242
+ if stat.get("metrics"):
243
+ line += f" [{stat.get('band', '?')}]"
244
+ if stat["outcome"] == "refused":
245
+ line += f" ({stat.get('reason', '')})"
246
+ lines.append(line)
247
+ e2e = self.receipt.get("e2e")
248
+ if e2e:
249
+ lines.append(
250
+ f" e2e: {e2e['base_ms']:.2f} -> {e2e['ms']:.2f} ms "
251
+ f"({e2e['speedup']:.3f}x, spread {e2e['spread']:.3f})")
252
+ led = self.receipt.get("ledger")
253
+ if led:
254
+ lines.append(f" ledger: {led['fallbacks']} fallback(s) over "
255
+ f"{led['guarded_calls']} guarded call(s)"
256
+ + ("" if led["clean"] else
257
+ f" — {led['seams_fell_back']}"))
258
+ return "\n".join(lines)
259
+
260
+ def ledger(self) -> dict[str, Any]:
261
+ """The live per-seam ledger of the committed attachment."""
262
+ return {} if self._handle is None else self._handle.report()
263
+
264
+ def detach(self) -> None:
265
+ """Restore the host, including seams that are not modules."""
266
+ if self._handle is not None:
267
+ self._handle.detach()
268
+ elif self._plan is not None:
269
+ self._plan.revert_all()
270
+ self.active = False
271
+
272
+ def save_receipt(self, directory: str | pathlib.Path) -> pathlib.Path:
273
+ directory = pathlib.Path(directory)
274
+ directory.mkdir(parents=True, exist_ok=True)
275
+ path = directory / f"attach_{self.receipt['digest'][:12]}.json"
276
+ path.write_text(json.dumps(self.receipt, indent=2, default=str))
277
+ return path
278
+
279
+
280
+ def attach(
281
+ model: torch.nn.Module,
282
+ forward: Callable[[], Any] | Sequence[Callable[[], Any]],
283
+ *,
284
+ structures: tuple[str, ...] = ALL_STRUCTURES,
285
+ observations: Iterable[Any] | None = None,
286
+ prefix_cadence: bool = False,
287
+ percentile: float = 99.9,
288
+ max_samples: int | None = None,
289
+ output: Callable[[], Any] | None = None,
290
+ output_kind: str = "auto",
291
+ floors: Mapping[str, float] | None = None,
292
+ min_speedup: float = 1.02,
293
+ rounds: int = 5,
294
+ iters: int = 10,
295
+ on_guard_fail: str = "fallback",
296
+ scheme: str | Any = "auto",
297
+ negotiate_fp8: bool = True,
298
+ verbose: bool = True,
299
+ ) -> Plan:
300
+ """Discover, calibrate, gate and activate structures in one call.
301
+
302
+ ``forward`` runs the host once; ``observations`` / ``percentile`` /
303
+ ``max_samples`` are the repo's calibration arguments and mean exactly
304
+ what they mean in ``flash_rt.api.FlashRT.calibrate``.
305
+
306
+ ``scheme`` selects the precision profile by registered name
307
+ (:mod:`.schemes`); the default ``"auto"`` resolves to the fastest
308
+ profile the device can execute, and ``"none"`` is the explicit
309
+ quantisation off-switch.
310
+
311
+ ``output_kind`` picks how accuracy is measured — ``"values"`` for a
312
+ host whose output is the answer, ``"distribution"`` for one whose
313
+ output scores a vocabulary, ``"auto"`` to read it off what the host
314
+ returns. ``floors`` overrides the per-kind parity floors.
315
+
316
+ Refusal is per gate unit and per form: a unit that is refused is
317
+ refused *in the form it was measured, at the shape it was measured
318
+ at*, and the receipt records both.
319
+ """
320
+
321
+ def say(msg: str) -> None:
322
+ if verbose:
323
+ print(f"[structures] {msg}", flush=True)
324
+
325
+ thunks = (list(forward) if isinstance(forward, (list, tuple))
326
+ else [forward])
327
+ eval_thunk = thunks[-1]
328
+
329
+ # ---- reference: the host as it shipped, before anything is built ---
330
+ with torch.no_grad():
331
+ base_out = eval_thunk()
332
+ kind = (infer_output_kind(base_out) if output_kind == "auto"
333
+ else output_kind)
334
+ get_out = output or eval_thunk
335
+ want = _score_tensor(base_out)
336
+ want = None if want is None else want.detach().float().cpu()
337
+ # Some generation-style model outputs retain a full KV cache alongside
338
+ # logits. Only the scored CPU tensor is needed after this point; keeping
339
+ # the host output alive through candidate binding can consume the memory
340
+ # needed by a reversible FP8 arm on near-capacity models.
341
+ del base_out
342
+ say(f"host output scored as {kind!r}"
343
+ + ("" if want is not None else " (no tensor found — accuracy gate "
344
+ "cannot run, latency gate only)"))
345
+ band_floors = dict(floors or DEFAULT_FLOORS[kind])
346
+
347
+ # ---- bind: one path, the same one the plain call uses --------------
348
+ plan = auto_swaps(model, forward, structures=structures,
349
+ observations=observations, percentile=percentile,
350
+ max_samples=max_samples, prefix_cadence=prefix_cadence,
351
+ scheme=scheme, negotiate_fp8=negotiate_fp8,
352
+ verbose=verbose)
353
+ if not plan.swaps and not plan.toggles:
354
+ plan.revert_all()
355
+ return Plan({}, {}, {"digest": "none", "seams": 0,
356
+ "output_kind": kind}, _plan=plan)
357
+
358
+ calibration = plan.notes.get("calibration") or {}
359
+ # Adapters build their routed seam enabled so plain ``auto_swaps`` can
360
+ # be consumed directly. The front door must start its A/B gate from the
361
+ # untouched host and judge that routed seam as its own unit.
362
+ plan.disable_routed()
363
+ groups = _gate_groups(plan)
364
+ if plan.toggles:
365
+ groups[_ROUTED] = {}
366
+ bound_count = len(plan.swaps) + len(plan.observed)
367
+ say(f"{bound_count} bound seam(s) in {len(groups)} gate unit(s): "
368
+ + ", ".join(
369
+ f"{k}×{len(plan.observed) if k == _ROUTED else len(v)}"
370
+ for k, v in sorted(groups.items())))
371
+
372
+ def scored() -> torch.Tensor | None:
373
+ with torch.no_grad():
374
+ got = _score_tensor(get_out())
375
+ return None if got is None else got.detach().float().cpu()
376
+
377
+ # ---- per unit: accuracy, then that it ran, then net win -----------
378
+ stats: dict[str, dict[str, Any]] = {}
379
+ winners: dict[str, torch.nn.Module] = {}
380
+ routed_winner = False
381
+ for name, swaps in sorted(groups.items()):
382
+ routed = name == _ROUTED
383
+ paths = plan.observed if routed else swaps
384
+ stat: dict[str, Any] = {"seams": len(paths),
385
+ "paths": sorted(paths)[:4],
386
+ "outcome": "pending"}
387
+ stats[name] = stat
388
+ arm = _Arm(model, swaps, plan, on_guard_fail, routed=routed)
389
+ try:
390
+ arm.on()
391
+ if want is not None:
392
+ metrics = metrics_for(kind, scored(), want)
393
+ stat["metrics"] = _round(metrics)
394
+ stat["band"] = band_of(metrics, kind)
395
+ stat["band_note"] = band_note(metrics, kind, calibration)
396
+ _say_band(name, stat["band"], stat["band_note"], say)
397
+ ok, why = passes(metrics, band_floors)
398
+ if not ok:
399
+ stat["outcome"] = "refused"
400
+ stat["reason"] = f"{why} (caller floor)"
401
+ continue
402
+ # read before the timing loop: it is this unit's own scoring
403
+ # forward that the accuracy number came from
404
+ led = arm.handle.summary()
405
+ stat["ledger"] = led
406
+ if not led["clean"]:
407
+ # the parity above looked fine because the host computed it
408
+ stat["outcome"] = "refused"
409
+ stat["reason"] = (
410
+ f"{len(led['seams_fell_back'])} seam(s) fell back to "
411
+ f"the host module, so this unit did not run: "
412
+ f"{led['seams_fell_back'][:3]}")
413
+ continue
414
+ timing = _paired_ab(eval_thunk, arm.on, arm.off,
415
+ rounds=rounds, iters=iters)
416
+ stat["e2e"] = timing
417
+ if timing["speedup"] < min_speedup:
418
+ stat["outcome"] = "refused"
419
+ stat["reason"] = (
420
+ f"no net win ({timing['speedup']:.3f}x, spread "
421
+ f"{timing['spread']:.3f}) at {_shape_note(plan)}")
422
+ continue
423
+ stat["outcome"] = "activated"
424
+ if routed:
425
+ routed_winner = True
426
+ else:
427
+ winners.update(swaps)
428
+ except GuardRefused as refusal:
429
+ stat["outcome"] = "refused"
430
+ stat["reason"] = f"runtime form refused: {refusal}"
431
+ finally:
432
+ arm.off()
433
+ say(f"{name}: {stat['outcome']}"
434
+ + (f" ({stat.get('reason')})" if stat.get("reason") else
435
+ f" {stat['e2e']['speedup']:.3f}x, band {stat.get('band')}"))
436
+
437
+ # ---- union re-check, then commit ---------------------------------
438
+ e2e_final: dict[str, Any] | None = None
439
+ ledger_final: dict[str, Any] | None = None
440
+ handle = None
441
+ if winners or routed_winner:
442
+ arm = _Arm(model, winners, plan, on_guard_fail,
443
+ routed=routed_winner, revert=plan.revert)
444
+ reason = ""
445
+ try:
446
+ arm.on()
447
+ metrics = (metrics_for(kind, scored(), want)
448
+ if want is not None else {})
449
+ ok, why = (
450
+ passes(metrics, band_floors) if metrics else (True, ""))
451
+ ledger_final = arm.handle.summary()
452
+ timing = _paired_ab(eval_thunk, arm.on, arm.off,
453
+ rounds=rounds, iters=iters)
454
+ arm.on() # the loop ends on the off arm
455
+ handle = arm.handle
456
+ if not ok or not ledger_final["clean"] \
457
+ or timing["speedup"] < min_speedup:
458
+ reason = (why if not ok
459
+ else "seams fell back" if not ledger_final["clean"]
460
+ else f"no net win ({timing['speedup']:.3f}x)")
461
+ except GuardRefused as refusal:
462
+ reason = f"runtime form refused: {refusal}"
463
+ if reason:
464
+ arm.off() # also reverts the routed seams
465
+ handle, winners, routed_winner = None, {}, False
466
+ for stat in stats.values():
467
+ if stat["outcome"] == "activated":
468
+ stat["outcome"] = "refused"
469
+ stat["reason"] = f"union: {reason}"
470
+ say(f"union of activated units refused — {reason}; host "
471
+ "restored, including the routed seams")
472
+ else:
473
+ e2e_final = timing
474
+ e2e_final["metrics"] = _round(metrics)
475
+ e2e_final["band"] = band_of(metrics, kind) if metrics else "n/a"
476
+ if metrics:
477
+ e2e_final["band_note"] = band_note(metrics, kind, calibration)
478
+ _say_band("e2e", e2e_final["band"],
479
+ e2e_final["band_note"], say)
480
+ say(f"active: {len(winners)} seam(s), "
481
+ f"{timing['base_ms']:.2f} -> {timing['ms']:.2f} ms "
482
+ f"({timing['speedup']:.3f}x, spread {timing['spread']:.3f})"
483
+ f", band {e2e_final['band']}")
484
+ if not winners and not routed_winner:
485
+ plan.revert_all()
486
+ say("outcome: whole-host refusal — model left untouched")
487
+
488
+ activated = dict(winners)
489
+ if routed_winner:
490
+ activated.update(plan.observed)
491
+ receipt = {
492
+ "schema_version": 1,
493
+ "environment": _environment(),
494
+ "model": type(model).__name__,
495
+ "output_kind": kind,
496
+ "scheme": plan.notes.get("scheme"),
497
+ "floors": band_floors,
498
+ "min_speedup": min_speedup,
499
+ "timing": {"method": "paired alternating", "rounds": rounds,
500
+ "iters": iters},
501
+ "calibration": plan.notes.get("calibration"),
502
+ "assumed": plan.notes.get("assumed", []),
503
+ "refused_at_bind": plan.notes.get("refused", []),
504
+ "negotiated_layers": plan.notes.get("negotiated_layers", []),
505
+ "units": stats,
506
+ "e2e": e2e_final,
507
+ "ledger": ledger_final,
508
+ "seams_active": sorted(activated),
509
+ }
510
+ receipt["digest"] = hashlib.sha256(
511
+ json.dumps(receipt, sort_keys=True, default=str).encode()
512
+ ).hexdigest()
513
+ return Plan(activated, stats, receipt, _handle=handle, _plan=plan)
514
+
515
+
516
+ def _environment() -> dict[str, str]:
517
+ """Version fingerprint baked into every receipt.
518
+
519
+ A parity or latency figure without the environment it was measured
520
+ in is not comparable to anything. The concrete case: a host
521
+ modelling contract that moved between two library versions changed
522
+ every activation scale while both arms of an A/B stayed mutually
523
+ consistent — two receipts differing only in this block is exactly
524
+ how that shows up.
525
+ """
526
+ import platform
527
+
528
+ env = {"python": platform.python_version(),
529
+ "torch": torch.__version__}
530
+ if torch.cuda.is_available():
531
+ env["cuda"] = str(torch.version.cuda)
532
+ env["device"] = torch.cuda.get_device_name(0)
533
+ cap = torch.cuda.get_device_capability(0)
534
+ env["sm"] = f"sm{cap[0]}{cap[1]}"
535
+ try:
536
+ import transformers
537
+ env["transformers"] = transformers.__version__
538
+ except ImportError:
539
+ pass
540
+ return env
541
+
542
+
543
+ def _say_band(where: str, band: str, note: str, say) -> None:
544
+ """Report the band; say it out loud when it is the one to look at.
545
+
546
+ A ``low`` band is not a refusal — see :mod:`.gates`. It is the caller's
547
+ call, so it has to reach the caller rather than sit in a receipt nobody
548
+ opens.
549
+ """
550
+ say(f"{where}: {note}")
551
+ if band == "low":
552
+ warnings.warn(
553
+ f"structures: {where} is in the low accuracy band — {note}. "
554
+ "This is reported, not refused: whether it is acceptable "
555
+ "depends on the deployment. Pass floors={...} to make it a "
556
+ "hard requirement, or widen the calibration set.",
557
+ RuntimeWarning, stacklevel=3)
558
+
559
+
560
+ def _round(metrics: Mapping[str, Any]) -> dict[str, Any]:
561
+ return {k: (round(v, 7) if isinstance(v, float) else v)
562
+ for k, v in metrics.items()}
563
+
564
+
565
+ def _shape_note(plan: AutoPlan) -> str:
566
+ """The workload a refusal was measured at, for the receipt.
567
+
568
+ A refusal with no shape attached turns into "that structure does not
569
+ work here"; the shape is what makes it "not at this size".
570
+ """
571
+ rows = sorted({m for s in plan.seams for m in (s.m_profile or ())})
572
+ return f"rows={rows}" if rows else "rows unrecorded"