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,2052 @@
1
+ """Auto-assembly: discover seams, calibrate them in one pass, bind them.
2
+
3
+ This is the distribution layer. Given a host model and a way to run it,
4
+ it finds every structure seam (:mod:`.discover`), captures exactly the
5
+ calibration each one needs in a single forward pass, and binds each
6
+ through its library impl. The caller gets a ``path -> module`` swap map
7
+ and any outside-cadence update functions — the same thing the hand
8
+ recipes produced, derived from the model object rather than written by
9
+ hand. A host integrates by importing and calling; it writes no
10
+ per-seam scaffolding.
11
+
12
+ The calibration each structure needs, captured structure-aware:
13
+ linear_proj / qkv_pack : the shared input the projection(s) see, and
14
+ its per-tensor amax (the static act scale)
15
+ adaln_producer : the (cond, style) pairs the conditioning
16
+ projection emits across the tick, for the
17
+ step table and its fingerprint locator
18
+ decoder_ffn / vision_ffn : the normed input the MLP sees
19
+
20
+ Seam negotiation is resolved here: when an adaln_producer feeds a
21
+ sibling qkv_pack under the same parent, the producer emits fp8 and the
22
+ pack takes the shared act scale, skipping its own input quantization.
23
+ """
24
+
25
+ from __future__ import annotations
26
+
27
+ import itertools
28
+ from dataclasses import dataclass, field
29
+ from typing import Any, Callable, Iterable, Sequence
30
+
31
+ import torch
32
+
33
+ from .discover import (Seam, _resolve, discover, group_families,
34
+ seam_weights)
35
+ from .points import Collector, Point, resolve as resolve_points
36
+
37
+ _FP8 = torch.float8_e4m3fn
38
+ _FP8_CHAIN_MAX_ROWS = 256 # fp8 producer chain qualifies at denoise
39
+ # M (bandwidth-bound); large-M prefill skips
40
+
41
+ # Host-family attention adapters. Attention seams are not a static
42
+ # module pattern — where the attention math actually runs is
43
+ # host-specific (a function in one host, a processor in another), so
44
+ # auto-discovery of the attention_core structure is delegated to
45
+ # registered adapters. Each adapter, given the model and a way to run
46
+ # it, returns (swaps, update) or None (this host is not its family).
47
+ _ATTENTION_ADAPTERS: list = []
48
+ _QK_NORM_ROPE_ADAPTERS: list = []
49
+ _QKV_ROPE_ADAPTERS: list = []
50
+ _GATED_DELTA_ADAPTERS: list = []
51
+
52
+
53
+ def register_attention_adapter(adapter) -> None:
54
+ """Register a host-family attention adapter (callable)."""
55
+ _ATTENTION_ADAPTERS.append(adapter)
56
+
57
+
58
+ def register_qk_norm_rope_adapter(adapter) -> None:
59
+ """Register a host-family adapter for a Q/K norm + RoPE boundary."""
60
+ _QK_NORM_ROPE_ADAPTERS.append(adapter)
61
+
62
+
63
+ def register_qkv_rope_adapter(adapter) -> None:
64
+ """Register a host-family adapter for packed QKV + bias + RoPE."""
65
+ _QKV_ROPE_ADAPTERS.append(adapter)
66
+
67
+
68
+ def register_gated_delta_adapter(adapter) -> None:
69
+ """Register a host-family Gated Delta callable adapter."""
70
+ _GATED_DELTA_ADAPTERS.append(adapter)
71
+
72
+
73
+ # Per-structure binders registered from impls, consulted before the
74
+ # built-in routing in :func:`_bind_auto`. New structures land by
75
+ # registering here from their own module instead of editing the routing
76
+ # function — parallel additions then touch disjoint files. A binder is
77
+ # ``f(model, seam, cap, *, points, fmt) -> module | None`` with the same
78
+ # refusal contract as ``_bind_auto`` (raise ``ValueError`` with the
79
+ # reason; return ``None`` for "host keeps its path").
80
+ _STRUCTURE_BINDERS: dict[str, Any] = {}
81
+
82
+
83
+ def register_structure_binder(structure: str, binder) -> None:
84
+ """Route ``structure`` seams to ``binder`` (last write wins).
85
+
86
+ A binder is called as ``binder(model, seam, cap, *, points, fmt,
87
+ fmt_params)`` — ``points`` is the reduced collector, ``fmt`` the
88
+ scheme's format routing for this seam (or ``None`` for the default),
89
+ ``fmt_params`` the decision's recipe payload for that format (or
90
+ ``None``).
91
+ """
92
+ _STRUCTURE_BINDERS[structure] = binder
93
+
94
+
95
+ @dataclass
96
+ class AutoPlan:
97
+ """Discovered + calibrated swaps, ready to stage."""
98
+
99
+ swaps: dict[str, torch.nn.Module] = field(default_factory=dict)
100
+ updates: list[Callable[[], None]] = field(default_factory=list)
101
+ seams: list[Seam] = field(default_factory=list)
102
+ notes: dict[str, Any] = field(default_factory=dict)
103
+ #: modules that carry a guard but are not swapped at a path — an
104
+ #: adapter's routed seam. Reported by the attachment's ledger, never
105
+ #: installed by it, so a seam that cannot be swapped can still be
106
+ #: counted instead of being invisible.
107
+ observed: dict[str, torch.nn.Module] = field(default_factory=dict)
108
+ #: undo callables for host mutations that had to happen while the plan
109
+ #: was being built rather than when it was attached (a patched
110
+ #: module-level function). Handed to ``attach`` so ``detach`` really
111
+ #: does give back the model that came in.
112
+ revert: list[Callable[[], None]] = field(default_factory=list)
113
+ #: ``(enable, disable)`` pairs for those same non-module seams, so a
114
+ #: gate can put the host back for the baseline arm without unbinding
115
+ #: anything. A seam that cannot be turned off cannot be measured.
116
+ toggles: list[tuple[Callable[[], None], Callable[[], None]]] = field(
117
+ default_factory=list)
118
+ #: ``flash_rt.core.precision_spec.ModelPrecisionSpec`` for the scales
119
+ #: this plan baked in — the repo's introspection format, not a private
120
+ #: one, so ``plan.precision_spec`` reads like ``rt.precision_spec``
121
+ precision_spec: Any = None
122
+
123
+ def enable_routed(self) -> None:
124
+ for on, _ in self.toggles:
125
+ on()
126
+
127
+ def disable_routed(self) -> None:
128
+ for _, off in self.toggles:
129
+ off()
130
+
131
+ def abort(self) -> None:
132
+ """Roll back everything this plan touched without an attach.
133
+
134
+ The plan mutates the host as it builds (adapter routes enable,
135
+ interface switches land, streamed originals leave the device).
136
+ ``attach`` is the commit point; a caller that decides not to
137
+ commit — or a bind that dies midway — calls this instead, and
138
+ the model returns to the loaded form.
139
+ """
140
+ self.disable_routed()
141
+ self.revert_all()
142
+ store = self.notes.get("stream_store")
143
+ if store is not None:
144
+ store.restore_all()
145
+
146
+ def revert_all(self) -> None:
147
+ """Undo the plan-time host mutations. Idempotent per adapter."""
148
+ for undo in reversed(self.revert):
149
+ undo()
150
+ self.revert.clear()
151
+ self.observed.clear()
152
+
153
+
154
+ def _spec_points(seam) -> tuple[str, ...]:
155
+ """The point names this seam's structure spec declares."""
156
+ from flash_rt.catalog.registry import load
157
+
158
+ try:
159
+ calibration = load(seam.structure).calibration
160
+ if (seam.structure == "modnorm_qkv_chain"
161
+ and seam.variant.get("modulation") == "per_token_table"):
162
+ # the table form owns its sublayers' seams, so it carries
163
+ # their static scales itself; the spec names them separately
164
+ # because the scale_shift form keeps no points at all
165
+ return tuple(calibration.get("per_token_table_points", ()))
166
+ return tuple(calibration.get("points", ()))
167
+ except Exception: # noqa: BLE001
168
+ return ()
169
+
170
+
171
+ def _consumer_point(seam) -> tuple[str, str]:
172
+ """Where a negotiated consumer's input is observed.
173
+
174
+ The producer's output *is* this tensor, so one amax serves both sides
175
+ of the chain — which is why the pair can share a static scale at all.
176
+ """
177
+ if seam.structure == "qkv_pack":
178
+ return (seam.path + "." + (seam.pack_attrs or ("q_proj",))[0], "x")
179
+ if seam.structure == "decoder_ffn":
180
+ return (seam.path, "x_after_norm")
181
+ return (seam.path, "x")
182
+
183
+
184
+ def _seam_key(seam: Seam) -> str:
185
+ """Unique receipt/decision key for a discovered structure instance.
186
+
187
+ Most structures own one module path. A dual-path attention module may
188
+ expose two independent sibling QKV groups under the same parent; the
189
+ first projection is the stable, real host path that distinguishes them.
190
+ """
191
+ if seam.structure == "qkv_pack" and seam.pack_attrs:
192
+ return seam.path + "." + seam.pack_attrs[0]
193
+ if seam.structure == "modnorm_qkv_chain":
194
+ # This catalog structure describes the composition at the same host
195
+ # path as a possible block seam; keep both receipt identities.
196
+ return seam.path + "::modnorm_qkv_chain"
197
+ return seam.path
198
+
199
+
200
+ def _race_ms(fn, *, warmup: int = 3, iters: int = 10) -> float:
201
+ with torch.no_grad():
202
+ for _ in range(warmup):
203
+ fn()
204
+ torch.cuda.synchronize()
205
+ start = torch.cuda.Event(True)
206
+ end = torch.cuda.Event(True)
207
+ start.record()
208
+ for _ in range(iters):
209
+ fn()
210
+ end.record()
211
+ torch.cuda.synchronize()
212
+ return start.elapsed_time(end) / iters
213
+
214
+
215
+ def _pair_vision_norm_fp8(model, plan, points, stream,
216
+ probe=None) -> None:
217
+ """Seat an FP8-emitting norm producer where its FFN seat consumes.
218
+
219
+ Seat-to-seat only — the measured lesson stands: a *host* handed
220
+ FP8 keeps going and the output is silent garbage, so the producer
221
+ is created here, paired with the FFN seat as its direct consumer,
222
+ and nowhere else. Both forms quantize at the same static scale, so
223
+ the pair's smoke against the bf16 form isolates wiring breakage
224
+ rather than calibration; the flip itself is decided by a bind-time
225
+ measurement (the flip-only-if-faster house rule) and recorded next
226
+ to the other format races.
227
+
228
+ The premise itself — "the norm's sole consumer is the FFN" — is a
229
+ runtime fact, not a shape: one probe forward verifies per pair
230
+ that the norm's output tensor *is* the seat's input tensor. A
231
+ host that modulates between them, or recomputes the norm inside a
232
+ fused path, fails identity here and never flips (a paired seat on
233
+ such a host measured 0.97 teacher-forced and thirty dead seats).
234
+ No probe, no fact, no flip — the check fails closed.
235
+ """
236
+ import dataclasses
237
+ from .impls import KernelUnavailable
238
+ from .impls.vision_ffn import fp8_static as vis_impl
239
+
240
+ direct_feed: dict[str, bool] = {}
241
+ candidates = []
242
+ for seam in plan.seams:
243
+ if seam.structure != "vision_ffn" \
244
+ or not getattr(seam, "norm_attr", None):
245
+ continue
246
+ seat = plan.swaps.get(seam.path)
247
+ if not isinstance(seat, vis_impl.FusedGeluMlp):
248
+ continue
249
+ if seat._bound.in_dtype != "bf16":
250
+ continue
251
+ norm_path = (f"{seam.parent_path}.{seam.norm_attr}"
252
+ if seam.parent_path else seam.norm_attr)
253
+ try:
254
+ candidates.append((norm_path,
255
+ model.get_submodule(norm_path),
256
+ model.get_submodule(seam.path)))
257
+ except AttributeError:
258
+ continue
259
+ if candidates and probe is not None:
260
+ last_out: dict[str, int] = {}
261
+ hooks = []
262
+ for norm_path, host_norm, host_mlp in candidates:
263
+ direct_feed[norm_path] = True
264
+
265
+ def _out(_m, _a, out, _p=norm_path):
266
+ if torch.is_tensor(out):
267
+ last_out[_p] = out.data_ptr()
268
+
269
+ def _inp(_m, args, _p=norm_path):
270
+ if args and torch.is_tensor(args[0]) \
271
+ and last_out.get(_p) != args[0].data_ptr():
272
+ direct_feed[_p] = False
273
+
274
+ hooks.append(host_norm.register_forward_hook(_out))
275
+ hooks.append(host_mlp.register_forward_pre_hook(_inp))
276
+ try:
277
+ probe()
278
+ except Exception: # noqa: BLE001 — no fact, no flip
279
+ direct_feed.clear()
280
+ finally:
281
+ for hook in hooks:
282
+ hook.remove()
283
+
284
+ for seam in plan.seams:
285
+ if seam.structure != "vision_ffn" \
286
+ or not getattr(seam, "norm_attr", None):
287
+ continue
288
+ seat = plan.swaps.get(seam.path)
289
+ if not isinstance(seat, vis_impl.FusedGeluMlp):
290
+ continue
291
+ bound = seat._bound
292
+ if bound.in_dtype != "bf16" or not bound.fc1_fp8.is_cuda:
293
+ continue
294
+ norm_path = (f"{seam.parent_path}.{seam.norm_attr}"
295
+ if seam.parent_path else seam.norm_attr)
296
+ try:
297
+ host_norm = model.get_submodule(norm_path)
298
+ except AttributeError:
299
+ continue
300
+ if not direct_feed.get(norm_path):
301
+ plan.notes.setdefault("refused", []).append(
302
+ (f"{norm_path}::norm_fp8_producer",
303
+ "pair premise unverified: the norm's output is not "
304
+ "the seat's input on this host's runtime path"))
305
+ continue
306
+ try:
307
+ from .impls.norm_fused.fp8_producer import (
308
+ bind_norm_fp8_producer)
309
+ producer = bind_norm_fp8_producer(host_norm,
310
+ bound.input_scale)
311
+ kern = vis_impl._kernel()
312
+ twin = vis_impl.FusedGeluMlp(
313
+ dataclasses.replace(
314
+ bound, in_dtype="fp8_static",
315
+ fused_mlp=(getattr(kern, "fp8_gelu_mlp_v2_bf16",
316
+ None)
317
+ or kern.fp8_gelu_mlp_bf16)),
318
+ original=seat._frt_host())
319
+ except (KernelUnavailable, ValueError,
320
+ RuntimeError) as refusal:
321
+ plan.notes.setdefault("refused", []).append(
322
+ (f"{norm_path}::norm_fp8_producer",
323
+ str(refusal)[:200]))
324
+ continue
325
+ rows_seen = points.row_profile(seam.path, "x_after_norm")
326
+ rows = rows_seen[len(rows_seen) // 2] if rows_seen else 128
327
+ dim = int(host_norm.weight.shape[0])
328
+ dev = producer.w.device
329
+ x = torch.randn(rows, dim, device=dev, dtype=torch.bfloat16)
330
+ current_norm = plan.swaps.get(norm_path)
331
+ norm_a = current_norm if current_norm is not None else host_norm
332
+ # the host norm keeps whatever dtype its fidelity policy chose
333
+ # (fp32 norms under selective-bf16 hosts) — probe it in its own
334
+ # dtype, and a probe that still refuses records a refusal
335
+ # rather than killing the scan
336
+ w_dt = getattr(getattr(norm_a, "weight", None), "dtype",
337
+ torch.bfloat16)
338
+ try:
339
+ with torch.no_grad():
340
+ ref = seat(norm_a(x.to(w_dt)).to(torch.bfloat16))
341
+ got = twin(producer(x))
342
+ cos = torch.nn.functional.cosine_similarity(
343
+ got.float().flatten(), ref.float().flatten(),
344
+ dim=0)
345
+ except RuntimeError as refusal:
346
+ plan.notes.setdefault("refused", []).append(
347
+ (f"{norm_path}::norm_fp8_producer",
348
+ f"pair probe: {refusal}"))
349
+ continue
350
+ if float(cos) < 0.98:
351
+ plan.notes.setdefault("refused", []).append(
352
+ (f"{norm_path}::norm_fp8_producer",
353
+ f"pair smoke cos {float(cos):.6f} < 0.98"))
354
+ continue
355
+ ms_a = _race_ms(lambda: seat(norm_a(x.to(w_dt))
356
+ .to(torch.bfloat16)))
357
+ ms_b = _race_ms(lambda: twin(producer(x)))
358
+ plan.notes.setdefault("format_race", []).append(
359
+ {"layer": norm_path, "rows": rows, "dim": dim,
360
+ "bf16_chain_ms": round(ms_a, 4),
361
+ "fp8_norm_chain_ms": round(ms_b, 4),
362
+ "smoke_cos": round(float(cos), 6),
363
+ "winner": ("fp8_norm_chain" if ms_b < ms_a
364
+ else "bf16_chain")})
365
+ # the bind-time race is a qualification signal, not the
366
+ # activation: seat-level micro-timing was refuted in both
367
+ # directions by production-form measurement (28/28 race wins
368
+ # here measured +0.6ms end-to-end on one device). Activation
369
+ # needs a production-form receipt for this box.
370
+ from . import decisions as _dec
371
+ if ms_b >= ms_a or _dec.lookup(
372
+ "vision_norm_fp8") != "fp8_norm_chain":
373
+ continue
374
+ plan.swaps[seam.path] = twin
375
+ plan.swaps[norm_path] = producer
376
+ # the pair's guards know each other: one out-of-contract call
377
+ # at runtime demotes both seats as a unit (all-or-nothing, the
378
+ # same atomicity the gate group gives their bind-time verdict)
379
+ if (twin._frt_guard is not None
380
+ and producer._frt_guard is not None):
381
+ twin._frt_guard.pair = producer._frt_guard
382
+ producer._frt_guard.pair = twin._frt_guard
383
+ stream(norm_path)
384
+
385
+
386
+ def _bind_regions(model, seams, *, probe, say):
387
+ """Resolve and bind every registered region family on this host.
388
+
389
+ Returns ``None`` when no family identifies a region here. The
390
+ stability contract: resolution reads receipts (pin > cache >
391
+ seated) and a winning candidate binds *before* any seat, so its
392
+ failure — refused symbols, a smoke miss, an exception — leaves
393
+ every seam in place and lands on the trail. Only a successful
394
+ bind claims the seams under its root; the seated floor is what
395
+ remains everywhere else, always.
396
+ """
397
+ from . import regions
398
+ from .impls.dit_stack import region as _dit_region
399
+ from .impls.adarms_stack import region as _adarms_region
400
+ from .impls.prefill_tower import region as _prefill_region
401
+ from .impls.vision_tower import region as _vision_region
402
+ _dit_region.register()
403
+ _adarms_region.register()
404
+ _prefill_region.register()
405
+ _vision_region.register()
406
+
407
+ notes: dict = {}
408
+ extras = {"seams": list(seams), "observed": {}, "revert": [],
409
+ "toggles": [], "notes": notes}
410
+ engaged = False
411
+ for fam in regions.registered():
412
+ try:
413
+ roots = list(fam.identify(model))
414
+ except Exception: # noqa: BLE001 — identify never kills
415
+ continue
416
+ if not roots:
417
+ continue
418
+ engaged = True
419
+ try:
420
+ first = (model.get_submodule(roots[0]) if roots[0]
421
+ else model)
422
+ host_sig = regions.structural_signature(first)
423
+ except Exception: # noqa: BLE001 — scoping never kills a bind
424
+ host_sig = None
425
+ winner, source = regions.resolve(fam.family, host_sig=host_sig,
426
+ notes=notes)
427
+ if winner == regions.SEATED:
428
+ say(f"region {fam.family}: seated ({source})")
429
+ continue
430
+ cand = fam.candidate(winner)
431
+ for root in roots:
432
+ label = f"{root}::{winner}"
433
+ if probe is None:
434
+ notes.setdefault("regions_refused", []).append(
435
+ (label, "no probe callable"))
436
+ continue
437
+ try:
438
+ result = cand.bind(model, root, probe)
439
+ except Exception as exc: # noqa: BLE001 — never kills
440
+ result = {"refused": f"{type(exc).__name__}: {exc}"}
441
+ if not result or result.get("refused"):
442
+ reason = (result or {}).get(
443
+ "refused", "bind returned nothing")
444
+ notes.setdefault("regions_refused", []).append(
445
+ (label, str(reason)[:200]))
446
+ say(f"region {fam.family}@{root}: {winner} refused "
447
+ "-> seated")
448
+ continue
449
+ prefix = root + "." if root else ""
450
+ keep, claimed = [], 0
451
+ for s in extras["seams"]:
452
+ if s.path == root or s.path.startswith(prefix):
453
+ claimed += 1
454
+ else:
455
+ keep.append(s)
456
+ extras["seams"] = keep
457
+ extras["observed"].update(result.get("observed", {}))
458
+ extras["revert"].extend(result.get("revert", ()))
459
+ if result.get("toggle") is not None:
460
+ extras["toggles"].append(result["toggle"])
461
+ notes.setdefault("regions_bound", []).append(
462
+ {"family": fam.family, "root": root, "winner": winner,
463
+ "source": source, "claimed_seams": claimed,
464
+ "smoke_cos": result.get("smoke_cos")})
465
+ say(f"region {fam.family}@{root}: {winner} bound "
466
+ f"({claimed} seam(s) claimed, source={source})")
467
+ if engaged:
468
+ _wire_region_products(extras, notes, say)
469
+ return extras if engaged else None
470
+
471
+
472
+ def _wire_region_products(extras, notes, say) -> None:
473
+ """Cross-region wires: a bound producer feeds a bound consumer's
474
+ chain buffers directly, and the consumer drops its own in-graph
475
+ restaging. Structure-level, receipt-visible, and only ever armed
476
+ when both ends are bound with agreeing facts — a missing or
477
+ refused end simply leaves both chains in their standalone form.
478
+
479
+ First wire: a prefill tower's per-layer K/V (already in the shared
480
+ adjacent-pair chain layout — both chains assemble the same split
481
+ kernel) into a conditioned decoder stack's cache prefix rows. The
482
+ decoder's per-step prefix gather is the identity of this write, so
483
+ it leaves the graph.
484
+ """
485
+ towers, stacks = [], []
486
+ for label, bound in extras["observed"].items():
487
+ if "::prefill_" in label and getattr(bound, "prefix_sinks",
488
+ "no") != "no":
489
+ towers.append((label, bound))
490
+ if "::adarms_" in label and getattr(bound, "prefix_wired",
491
+ "no") != "no":
492
+ stacks.append((label, bound))
493
+ if len(towers) != 1 or len(stacks) != 1:
494
+ return
495
+ (t_label, tower), (s_label, stack) = towers[0], stacks[0]
496
+ tk = {k: tower.dims.get(k) for k in ("seq", "hd", "layers")}
497
+ if (tk["seq"] != stack.p_used or tk["hd"] != stack.dims.get("hd")
498
+ or tk["layers"] != stack.dims.get("layers")
499
+ or not stack.buf.get("kc")):
500
+ notes.setdefault("region_wires", []).append(
501
+ {"wire": "prefix_kv", "armed": False,
502
+ "reason": f"facts disagree: tower {tk} vs stack "
503
+ f"P={stack.p_used}"})
504
+ return
505
+ P = stack.p_used
506
+ tower.prefix_sinks = [
507
+ (stack.buf["kc"][l][0, :P, 0], stack.buf["vc"][l][0, :P, 0])
508
+ for l in range(tk["layers"])]
509
+ stack.prefix_wired = True
510
+
511
+ def unwire() -> None:
512
+ tower.prefix_sinks = None
513
+ stack.prefix_wired = False
514
+
515
+ extras["revert"].append(unwire)
516
+ notes.setdefault("region_wires", []).append(
517
+ {"wire": "prefix_kv", "armed": True, "producer": t_label,
518
+ "consumer": s_label, "rows": P})
519
+ say(f"region wire prefix_kv: {t_label} -> {s_label} ({P} rows)")
520
+
521
+
522
+ def _merge_region_extras(plan, extras) -> None:
523
+ plan.notes.update(extras["notes"])
524
+ plan.observed.update(extras["observed"])
525
+ plan.revert.extend(extras["revert"])
526
+ plan.toggles.extend(extras["toggles"])
527
+
528
+
529
+ def auto_swaps(
530
+ model: torch.nn.Module,
531
+ forward: Callable[..., Any] | Sequence[Callable[[], Any]],
532
+ *,
533
+ structures: tuple[str, ...] = ("decoder_ffn", "vision_ffn",
534
+ "qkv_pack", "adaln_producer",
535
+ "linear_proj", "patch_projection",
536
+ "norm_fused",
537
+ "attention_core", "decoder_block",
538
+ "modnorm_qkv_chain", "qk_norm_rope",
539
+ "qkv_rope",
540
+ "gated_delta_core"),
541
+ negotiate_fp8: bool = True,
542
+ prefix_cadence: bool = False,
543
+ observations: Iterable[Any] | None = None,
544
+ percentile: float = 99.9,
545
+ max_samples: int | None = None,
546
+ scheme: str | Any = "auto",
547
+ verbose: bool = False,
548
+ stream_store: Any = None,
549
+ ) -> AutoPlan:
550
+ """Discover, calibrate in one pass, and bind every applicable seam.
551
+
552
+ The calibration arguments are the repo's, not this layer's: the names,
553
+ the defaults and the meaning of ``percentile`` / ``max_samples`` /
554
+ ``observations`` are ``flash_rt.api.FlashRT.calibrate``'s, and the
555
+ reduction is ``flash_rt.core.calibration.accumulate_amax``. A second
556
+ vocabulary for the same thing is the one thing this layer must not
557
+ add.
558
+
559
+ auto_swaps(model, forward) # one sample
560
+ auto_swaps(model, [f0, f1, f2]) # one thunk each
561
+ auto_swaps(model, feed, observations=dataset) # feed(obs) per obs
562
+ auto_swaps(model, feed, observations=ds, percentile=95.0)
563
+
564
+ ``scheme`` names a registered quantisation scheme (:mod:`.schemes`):
565
+ what statistic each point needs, and per seam whether to bind or keep
566
+ the host at host precision. The default ``"auto"`` resolves to the
567
+ highest-performing profile the device can execute — on FP8-capable
568
+ hardware that is ``fp8_static``, bit-identical to the behaviour this
569
+ layer shipped with; elsewhere it is ``none`` (fusion structures
570
+ attach, quantised seams stay at host precision). Explicit selection
571
+ (``scheme="none"``, ``scheme="w4a16_decode"``, ...) overrides. It
572
+ adds no calibration entry.
573
+
574
+ ``forward`` is always "run the host once"; with ``observations`` it
575
+ takes one observation. That indirection is this layer's only
576
+ difference from a frontend's ``calibrate``, and it exists because a
577
+ host here is an arbitrary ``nn.Module`` with no common observation
578
+ contract — not because the calibration standard differs.
579
+
580
+ ``prefix_cadence`` declares that the caller will run ``plan.updates``
581
+ whenever the observation changes. Structures that hold per-observation
582
+ host state — the attention core keeps the prefix keys and values — are
583
+ only offered when it is set, because without the refresh they attend to
584
+ whatever the calibration saw. Leaving it off is the accurate default.
585
+
586
+ On ``percentile``: it reduces *across* samples. Within one sample the
587
+ reduction is a max over every call, which is required rather than
588
+ chosen — see ``docs/calibration.md`` §4.2. And note §10's own caveat
589
+ that at small N a 99.9 percentile barely clips at all (it interpolates
590
+ between the top two ranks); with N ≤ 64 and suspected outliers, pass a
591
+ lower one.
592
+ """
593
+
594
+ def say(msg: str) -> None:
595
+ if verbose:
596
+ print(f"[autobuild] {msg}", flush=True)
597
+
598
+ thunks, source = _calibration_thunks(forward, None, observations)
599
+ if max_samples is not None and len(thunks) > max_samples:
600
+ thunks = thunks[:max_samples]
601
+ plan_notes_calibration: dict[str, Any] = {}
602
+ plan_refusals: list[tuple[str, str]] = []
603
+
604
+ # A schedule adapter changes only the Python spelling of a qualified
605
+ # fixed loop. Calibration must observe that same canonical execution:
606
+ # otherwise a tensor-controlled host ``while`` re-enters a compiled
607
+ # graph break on every denoise step before the region hooks even run.
608
+ from .impls.fixed_iter import (
609
+ FixedIterationRefused,
610
+ normalize_fixed_iteration,
611
+ )
612
+
613
+ schedule_notes: list[dict[str, Any]] = []
614
+ normalized_thunks = []
615
+ for thunk in thunks:
616
+ try:
617
+ schedule = normalize_fixed_iteration(thunk, model)
618
+ except FixedIterationRefused as exc:
619
+ raise ValueError(str(exc)) from exc
620
+ if schedule is None:
621
+ normalized_thunks.append(thunk)
622
+ continue
623
+ normalized_thunks.append(schedule.forward)
624
+ schedule_notes.append({
625
+ "family": schedule.family,
626
+ "steps": schedule.steps,
627
+ "exact": schedule.exact,
628
+ **dict(schedule.details),
629
+ })
630
+ thunks = normalized_thunks
631
+
632
+ seams = discover(model, structures, refused=plan_refusals)
633
+ say(f"discovered {len(seams)} seam(s)")
634
+
635
+ # ---- region adjudication: the structure-level winner is a receipt.
636
+ # Resolution reads author pin > decision cache > seated and never
637
+ # experiments here; a winning candidate binds before any seat does
638
+ # (its failure leaves every seam in place — seated is always the
639
+ # floor), and only a *successful* bind claims the seams it absorbs.
640
+ # region binds observe every calibration sample: the probe runs
641
+ # all thunks and exposes the sample boundaries so a chain can keep
642
+ # per-sample statistics and reduce them with the house percentile
643
+ region_probe = thunks[0] if thunks else None
644
+ if thunks and len(thunks) > 1:
645
+ def _region_probe():
646
+ for t in thunks:
647
+ t()
648
+ _region_probe.samples = tuple(thunks)
649
+ region_probe = _region_probe
650
+ region_extras = _bind_regions(
651
+ model, seams, probe=region_probe, say=say)
652
+ if region_extras is not None:
653
+ seams = region_extras["seams"]
654
+
655
+ adapter_only = not seams and bool(
656
+ {"attention_core", "gated_delta_core"}.intersection(structures))
657
+ if not seams and not adapter_only:
658
+ plan = AutoPlan()
659
+ if region_extras is not None:
660
+ _merge_region_extras(plan, region_extras)
661
+ return plan
662
+
663
+ # ---- one calibration pass, structure-aware capture ----
664
+ # Activation scales go through the house two-level statistic: a max
665
+ # over every call inside one sample (docs/calibration.md §4.2 — a
666
+ # flow-matching host runs every step inside one forward and per-step
667
+ # scales crashed the compiler), then accumulate_amax's percentile
668
+ # across samples. Nothing holds an activation tensor for this: each
669
+ # point is one float, measured where the spec says it is
670
+ # (:mod:`.points`).
671
+ caps: dict[str, dict[str, Any]] = {}
672
+ hooks = []
673
+ all_points: list[Point] = []
674
+ seam_points: dict[str, list[Point]] = {}
675
+ for seam in seams:
676
+ try:
677
+ pts = resolve_points(seam, _spec_points(seam))
678
+ except ValueError as refusal:
679
+ plan_refusals.append((seam.path, str(refusal)[:120]))
680
+ continue
681
+ seam_points[_seam_key(seam)] = pts
682
+ all_points.extend(pts)
683
+
684
+ from . import schemes as _schemes
685
+ auto_resolved = isinstance(scheme, str) and scheme == "auto"
686
+ if auto_resolved:
687
+ scheme = _schemes.resolve_auto()
688
+ say(f"scheme auto -> {scheme}")
689
+ scheme_obj = (_schemes.get(scheme) if isinstance(scheme, str)
690
+ else scheme)
691
+ # loud wall before any calibration work: a scheme asking for a
692
+ # granularity the collector cannot measure must not silently get
693
+ # per-tensor numbers of the wrong shape
694
+ stat_request = scheme_obj.statistics(tuple(all_points))
695
+ _schemes.validate_request(stat_request)
696
+ collector = Collector(points=all_points, request=dict(stat_request))
697
+
698
+ # observed call order across the whole calibration pass. Anything
699
+ # that has to know which seam runs first (a stream-scoped buffer
700
+ # needs a writer, and the writer has to be the one the host calls
701
+ # first) reads it from here rather than assuming the module tree's
702
+ # order matches the forward's.
703
+ call_order = itertools.count()
704
+
705
+ def cap_cond(path):
706
+ def hook(module, args, out):
707
+ cap = caps[path]
708
+ if "order" not in cap:
709
+ cap["order"] = next(call_order)
710
+ cap.setdefault("pairs", []).append(
711
+ (args[0].detach().clone(), out.detach().clone()))
712
+ return None
713
+ return hook
714
+
715
+ def cap_shape(path):
716
+ # a block seam needs no tensors of its own, only the host's
717
+ # return convention (bare tensor or 1-tuple)
718
+ def hook(module, args, kwargs, out):
719
+ caps[path]["returns_tuple"] = isinstance(out, tuple)
720
+ return None
721
+ return hook
722
+
723
+ def cap_pack_input(path, attr):
724
+ """Record the executable shared-input property of a pack sibling.
725
+
726
+ Equal K dimensions only prove that projections *could* share an
727
+ input. Cross attention is the counterexample: Q consumes the live
728
+ latent while K/V consume encoder state. The leaf implementation
729
+ runs the first projection and turns later siblings into stash reads,
730
+ so exact storage identity and call order are part of its contract.
731
+ """
732
+ def hook(module, args):
733
+ x = args[0] if args else None
734
+ cap = caps[path]
735
+ cap.setdefault("pack_events", []).append(attr)
736
+ if not torch.is_tensor(x):
737
+ cap.setdefault("pack_inputs", {}).setdefault(attr, []).append(
738
+ None)
739
+ return None
740
+ signature = (
741
+ int(x.data_ptr()), int(x.storage_offset()), tuple(x.shape),
742
+ tuple(x.stride()), x.dtype, x.device,
743
+ )
744
+ cap.setdefault("pack_inputs", {}).setdefault(attr, []).append(
745
+ signature)
746
+ return None
747
+ return hook
748
+
749
+ def cap_patch_input(path, width):
750
+ """Record that the host really supplies complete flattened patches.
751
+
752
+ Matching Conv3d slots is not enough: a module with the same kernel
753
+ may consume an ordinary 5-D volume. The lowering is legal only when
754
+ the calibrated host input exposes K as its final dimension, exactly
755
+ as the processor-preflattened dataflow declares.
756
+ """
757
+ def hook(module, args):
758
+ x = args[0] if args else None
759
+ form = None
760
+ if torch.is_tensor(x) and x.numel() % width == 0:
761
+ form = (int(x.shape[-1]), int(x.numel() // width))
762
+ caps[path].setdefault("patch_inputs", []).append(form)
763
+ return None
764
+ return hook
765
+
766
+ # the amax points are hooked by the collector; only the two
767
+ # content/observation captures need their own hooks here, and neither
768
+ # is a statistic: a step table is memoised host output, a return
769
+ # convention is one boolean
770
+ for seam in seams:
771
+ key = _seam_key(seam)
772
+ caps[key] = {}
773
+ target = _resolve(model, seam.path)
774
+ if seam.structure == "decoder_block":
775
+ hooks.append(target.register_forward_hook(
776
+ cap_shape(key), with_kwargs=True))
777
+ elif seam.structure == "adaln_producer":
778
+ hooks.append(getattr(target, seam.cond_attr)
779
+ .register_forward_hook(cap_cond(key)))
780
+ elif seam.structure == "qkv_pack":
781
+ for attr in seam.pack_attrs or ():
782
+ hooks.append(getattr(target, attr).register_forward_pre_hook(
783
+ cap_pack_input(key, attr)))
784
+ elif seam.structure == "patch_projection":
785
+ hooks.append(target.register_forward_pre_hook(
786
+ cap_patch_input(key, seam.dims["K"])))
787
+ hooks.extend(collector.hooks(lambda path: _resolve(model, path)))
788
+
789
+ if hooks:
790
+ # the calibration pass is a transaction over the host: if a thunk
791
+ # raises, the hooks come off and no plan is returned. Removing
792
+ # them only on the success path leaves a failed calibration's
793
+ # hooks on the model, where they keep firing into a dict nobody
794
+ # reads and slow down every later forward for reasons that are
795
+ # nowhere in sight.
796
+ try:
797
+ with torch.no_grad():
798
+ for thunk in thunks:
799
+ thunk()
800
+ # one vector per sample, so the percentile across them
801
+ # is possible at all
802
+ collector.end_sample()
803
+ finally:
804
+ for h in hooks:
805
+ h.remove()
806
+ plan_notes_calibration = dict(
807
+ collector.reduce(percentile, verbose=verbose,
808
+ label=f"structures_N{len(thunks)}"),
809
+ source=source)
810
+ say(f"calibration pass done ({len(thunks)} sample(s) from "
811
+ f"{source}, {plan_notes_calibration['points']} point(s), "
812
+ f"{plan_notes_calibration['method']}"
813
+ + (f" p={percentile}" if len(thunks) > 1 else "") + ")")
814
+
815
+ # Adapters bind after the precision scheme may remove host-precision
816
+ # seams from ``plan.seams``. Preserve only the observed row capacity in
817
+ # the transient cap map so a structural adapter can preallocate without
818
+ # retaining activations or depending on a quantized sibling being bound.
819
+ for seam in seams:
820
+ key = _seam_key(seam)
821
+ rows = [
822
+ row
823
+ for point in seam_points.get(key, ())
824
+ for row in collector.row_profile(point.path, point.name)
825
+ ]
826
+ if rows:
827
+ caps[key]["rows"] = max(rows)
828
+
829
+ # A pack owns its sibling projections only after the calibration pass
830
+ # proves the property its execution relies on: identical input storage
831
+ # and q/k/v call order for every invocation. If it does not, retain the
832
+ # independent linear projections; they are valid structures at the
833
+ # narrower boundary. This is deliberately a data-flow check rather than
834
+ # a class/path allow-list, so self- and cross-attention implementations
835
+ # using the same module type are classified by what they actually do.
836
+ qualified_packs = []
837
+ for seam in (s for s in seams if s.structure == "qkv_pack"):
838
+ attrs = tuple(seam.pack_attrs or ())
839
+ cap = caps.get(_seam_key(seam), {})
840
+ inputs = cap.get("pack_inputs", {})
841
+ columns = [inputs.get(attr, []) for attr in attrs]
842
+ count_ok = bool(columns) and len({len(col) for col in columns}) == 1
843
+ calls = len(columns[0]) if count_ok else 0
844
+ shared = count_ok and calls > 0 and all(
845
+ all(col[i] is not None and col[i] == columns[0][i]
846
+ for col in columns[1:])
847
+ for i in range(calls)
848
+ )
849
+ ordered = cap.get("pack_events", []) == list(attrs) * calls
850
+ if shared and ordered:
851
+ qualified_packs.append(seam)
852
+ continue
853
+ plan_refusals.append((
854
+ _seam_key(seam),
855
+ "qkv_pack refused: sibling projections did not consume the "
856
+ "same tensor in fixed order during calibration",
857
+ ))
858
+ qualified_ids = {id(seam) for seam in qualified_packs}
859
+ seams = [s for s in seams
860
+ if s.structure != "qkv_pack" or id(s) in qualified_ids]
861
+ packed = {s.path + "." + a for s in qualified_packs
862
+ for a in (s.pack_attrs or ())}
863
+ seams = [s for s in seams
864
+ if not (s.structure == "linear_proj" and s.path in packed)]
865
+
866
+ # ---- per-token-table chains own their block's producer-fed members:
867
+ # the self-attention pack (the chain quantizes once for all three)
868
+ # and the FFN (the chain's second producer site feeds it fused).
869
+ # Everything else under the block — the output projection, the whole
870
+ # cross-attention — stays individually bindable, and the chain's
871
+ # forward calls whatever is attached there. ----
872
+ chain_blocks = {
873
+ s.path for s in seams
874
+ if s.structure == "modnorm_qkv_chain"
875
+ and s.variant.get("modulation") == "per_token_table"}
876
+ if chain_blocks:
877
+ def _chain_owns(seam):
878
+ for block in chain_blocks:
879
+ if (seam.structure == "qkv_pack"
880
+ and seam.path == block + ".attn1"):
881
+ return True
882
+ if (seam.structure == "vision_ffn"
883
+ and seam.path == block + ".ffn"):
884
+ return True
885
+ if (seam.structure == "linear_proj"
886
+ and seam.path.startswith(block + ".attn1.")
887
+ and seam.path.rsplit(".", 1)[1] in
888
+ ("to_q", "to_k", "to_v")):
889
+ return True
890
+ return False
891
+
892
+ seams = [s for s in seams if not _chain_owns(s)]
893
+
894
+ # ---- the scheme turns statistics into decisions. Keep-host is a
895
+ # first-class outcome recorded in the receipt, not a refusal: the
896
+ # seam is healthy, the scheme chose host precision for it. ----
897
+ class _SeamStats(dict):
898
+ def __init__(self, *args, structure: str, **kwargs):
899
+ super().__init__(*args, **kwargs)
900
+ self.structure = structure
901
+
902
+ seam_by_key = {_seam_key(seam): seam for seam in seams}
903
+ # Host-precision structural lowerings do not belong to a quantisation
904
+ # scheme decision. They still use the collector for real row/dtype
905
+ # qualification, but scheme="none" must not remove them.
906
+ scheme_independent = {"patch_projection"}
907
+ scheme_report = {
908
+ path: _SeamStats(
909
+ {f"{pt.path}|{pt.name}": collector.amax(pt.path, pt.name)
910
+ for pt in pts}, structure=seam_by_key[path].structure)
911
+ for path, pts in seam_points.items()
912
+ if path in seam_by_key
913
+ and seam_by_key[path].structure not in scheme_independent}
914
+ decision = scheme_obj.decide(scheme_report)
915
+ scheme_note: dict[str, Any] = {
916
+ "name": getattr(scheme_obj, "name", type(scheme_obj).__name__)}
917
+ if auto_resolved:
918
+ scheme_note["auto"] = True
919
+ if decision.keep_host:
920
+ kept = set(decision.keep_host)
921
+ seams = [s for s in seams if _seam_key(s) not in kept]
922
+ scheme_note["keep_host"] = {
923
+ p: decision.reasons.get(p, "") for p in sorted(kept)}
924
+ say(f"scheme {scheme_note['name']}: {len(kept)} seam(s) kept at "
925
+ f"host precision")
926
+ formats: dict[str, str] = dict(decision.formats or {})
927
+ fmt_params: dict[str, Any] = dict(getattr(decision, "params", None) or {})
928
+ if formats:
929
+ scheme_note["formats"] = dict(sorted(formats.items()))
930
+ if fmt_params:
931
+ scheme_note["params"] = {p: dict(v) for p, v
932
+ in sorted(fmt_params.items())}
933
+ say(f"scheme {scheme_note['name']}: {len(formats)} seam(s) "
934
+ f"routed to a non-default format")
935
+
936
+ # ---- fp8 seam negotiation: the load-bearing structure combination.
937
+ # A single kernel need not win alone (fp8 qkv at M=50 is marginal,
938
+ # fa2 in a bf16 stack loses); the *chain* wins — an adaln producer
939
+ # that emits fp8 lets the qkv pack skip its own input quantization
940
+ # and hands a clean fp8 seam down to the attention core. Bind the
941
+ # producer→pack pair together with one shared act scale wherever a
942
+ # producer feeds a pack under the same parent layer. ----
943
+ act_scales: dict[str, torch.Tensor] = {}
944
+ chain_rows: dict[str, int] = {}
945
+ negotiated: dict[str, dict[str, Seam]] = {}
946
+ if negotiate_fp8:
947
+ by_parent: dict[str, dict[str, Seam]] = {}
948
+ for seam in seams:
949
+ if formats.get(_seam_key(seam)):
950
+ # a chain shares one scale and one wire dtype; a member
951
+ # routed to another format has neither, so it binds
952
+ # standalone through its own impl instead
953
+ continue
954
+ layer = _layer_of(seam.path)
955
+ if seam.structure == "adaln_producer":
956
+ # a layer has two producer→consumer seams: the norm
957
+ # before attention feeds the projections, the norm after
958
+ # it feeds the MLP. Both can hand fp8 downstream.
959
+ slot = ("producer" if _feeds_attention(seam.path)
960
+ else "producer_ffn")
961
+ by_parent.setdefault(layer, {})[slot] = seam
962
+ elif seam.structure == "qkv_pack":
963
+ by_parent.setdefault(layer, {})["pack"] = seam
964
+ elif (seam.structure == "linear_proj"
965
+ and seam.proj_attr in ("q_proj", "to_q",
966
+ "add_q_proj")):
967
+ by_parent.setdefault(layer, {})["query"] = seam
968
+ elif seam.structure == "decoder_ffn":
969
+ by_parent.setdefault(layer, {})["ffn"] = seam
970
+ # the chain wins at small M (denoise): fp8 is bandwidth-bound and
971
+ # pays there, while a large-M prefill GEMM is compute-bound and
972
+ # fp8 buys little — and an fp8 producer feeding a big compiled
973
+ # prefill region is where the triton fp8 codegen chokes. Qualify
974
+ # on the calibrated row count, not on host names.
975
+ dev = next(model.parameters()).device
976
+ blocks = {s.path for s in seams if s.structure in (
977
+ "decoder_block", "modnorm_qkv_chain")}
978
+ for lay, g in by_parent.items():
979
+ # the attention pack is always negotiated. The FFN chain is
980
+ # negotiated only where a decoder_block owns the boundary,
981
+ # and the reason is the boundary rather than the kernel: at
982
+ # the norm seam the fused producer costs a kernel
983
+ # (gate_residual, +180 launches) plus its style
984
+ # materialization (+180) to save the FFN's own input
985
+ # quantize (-180) — measured net +0.17ms, so it is refused
986
+ # there. Inside a block the same kernel *replaces* the
987
+ # host's gated residual add instead of adding to it, which
988
+ # is the whole point of owning the block.
989
+ # Both chains need the block boundary, and for the same
990
+ # reason: a negotiated producer emits FP8, and only a caller
991
+ # that owns the block consumes it. Bound at the norm boundary
992
+ # the *host* is the consumer, and the host expects its norm to
993
+ # return a compute dtype — handed FP8 it keeps going and the
994
+ # output is garbage (measured 0.24 output match, and NaN on a
995
+ # neighbouring configuration) with nothing to see, because
996
+ # every shape and dtype is inside its contract. The FFN chain
997
+ # was already gated this way; the attention chain was not.
998
+ if lay not in blocks:
999
+ continue
1000
+ pairs = [("producer", "pack"), ("producer", "query"),
1001
+ ("producer_ffn", "ffn")]
1002
+ keep = {}
1003
+ for p_slot, c_slot in pairs:
1004
+ if p_slot not in g or c_slot not in g:
1005
+ continue
1006
+ c_path, c_name = _consumer_point(g[c_slot])
1007
+ amax = collector.amax(c_path, c_name)
1008
+ rows_seen = collector.row_profile(c_path, c_name)
1009
+ rows = rows_seen[len(rows_seen) // 2] if rows_seen else 1 << 30
1010
+ if amax is None or rows > _FP8_CHAIN_MAX_ROWS:
1011
+ continue
1012
+ # the consumer's input == the producer's output; its amax
1013
+ # is the one static scale both sides share
1014
+ keep[p_slot], keep[c_slot] = g[p_slot], g[c_slot]
1015
+ act_scales[f"{lay}|{c_slot}"] = torch.tensor(
1016
+ [max(amax / 448.0, 1e-8)], device=dev)
1017
+ chain_rows[f"{lay}|{c_slot}"] = rows
1018
+ if keep:
1019
+ negotiated[lay] = keep
1020
+
1021
+ # ---- the negotiated chain binds as one unit ----
1022
+ # producer and consumer must agree on the seam dtype: a pack bound
1023
+ # for fp8 input whose producer failed to bind would be handed BF16,
1024
+ # and the host would silently grow a quantize fused into whatever
1025
+ # produced it. Bind the pair together, or leave both on BF16.
1026
+ plan = AutoPlan(seams=seams)
1027
+ plan._requested_structures = frozenset(structures)
1028
+ if region_extras is not None:
1029
+ _merge_region_extras(plan, region_extras)
1030
+
1031
+ # ---- per-seat streaming consumption (the bind-peak lever) ----
1032
+ # With a weight store handed in, every leaf placement immediately
1033
+ # moves the replaced original's truth off the device: the originals
1034
+ # shrink as the quantized copies grow, and the bind peak stays near
1035
+ # one model instead of two. Regions whose later stages still read
1036
+ # host originals are excluded up front — the decoder_block seams
1037
+ # compose from their children, and the cross-attention K/V
1038
+ # projections serve the cadence banks — those wait for the
1039
+ # attachment's own consume().
1040
+ _stream_exclude: set[str] = set()
1041
+ if stream_store is not None:
1042
+ for _s in seams:
1043
+ if _s.structure == "decoder_block":
1044
+ _stream_exclude.add(_s.path)
1045
+ from .impls.cadence_static.cross_attention import (
1046
+ discover_cross_attention_kv as _discover_ckv)
1047
+ for _cand in _discover_ckv(model):
1048
+ _stream_exclude.add(_cand.path)
1049
+
1050
+ def _stream(placed) -> None:
1051
+ if stream_store is None:
1052
+ return
1053
+ paths = placed if isinstance(placed, (list, tuple, set)) \
1054
+ else [placed]
1055
+ from .swap import resolve_parent as _rp, _get as _sg
1056
+ for p in paths:
1057
+ if any(p == e or p.startswith(e + ".")
1058
+ for e in _stream_exclude):
1059
+ continue
1060
+ try:
1061
+ parent, attr = _rp(model, p)
1062
+ stream_store.stash_module(p, _sg(parent, attr))
1063
+ except (AttributeError, TypeError, ValueError):
1064
+ continue
1065
+ plan.notes["streamed_bytes"] = stream_store.stats["freed_bytes"]
1066
+ # ---- backbone attention interface: measured band decision. This
1067
+ # must precede every adapter that resolves the host's attention
1068
+ # registry into a closure (the rope routes, the vision pin), or a
1069
+ # switched interface arrives after the traffic already left ----
1070
+ if "attention_core" in structures:
1071
+ from .adapters.transformers_attention_interface import (
1072
+ TransformersAttentionInterfaceAdapter)
1073
+ try:
1074
+ iface = TransformersAttentionInterfaceAdapter()(model, plan)
1075
+ except (ValueError, RuntimeError) as refusal:
1076
+ plan.notes.setdefault("refused", []).append(
1077
+ ("backbone_attn", str(refusal)[:200]))
1078
+ iface = None
1079
+ if iface:
1080
+ if iface.get("refused"):
1081
+ plan.notes.setdefault("refused", []).extend(
1082
+ iface["refused"])
1083
+ plan.revert.extend(iface.get("revert", ()))
1084
+ plan.notes.update(iface.get("notes", {}))
1085
+ plan.notes["scheme"] = scheme_note
1086
+ if schedule_notes:
1087
+ plan.notes["schedules"] = schedule_notes
1088
+ handled: set[str] = set()
1089
+ for lay, g in negotiated.items():
1090
+ for p_slot, c_slot in (("producer", "pack"),
1091
+ ("producer", "query"),
1092
+ ("producer_ffn", "ffn")):
1093
+ if p_slot not in g or c_slot not in g:
1094
+ continue
1095
+ p_seam, c_seam = g[p_slot], g[c_slot]
1096
+ p_cap = caps.get(_seam_key(p_seam), {})
1097
+ if not p_cap.get("pairs"):
1098
+ continue
1099
+ try:
1100
+ pair = _bind_negotiated(
1101
+ model, p_seam, c_seam, p_cap, collector,
1102
+ act_scales[f"{lay}|{c_slot}"],
1103
+ chain_rows[f"{lay}|{c_slot}"], plan)
1104
+ except (ValueError, RuntimeError) as refusal:
1105
+ plan.notes.setdefault("refused", []).append(
1106
+ (f"{lay} [{c_slot} chain]", str(refusal)[:200]))
1107
+ continue
1108
+ plan.swaps.update(pair)
1109
+ _stream(list(pair))
1110
+ handled.update({_seam_key(p_seam), _seam_key(c_seam)})
1111
+ for chain in (
1112
+ seam for seam in seams
1113
+ if seam.structure == "modnorm_qkv_chain"
1114
+ and seam.path == lay
1115
+ ):
1116
+ handled.add(_seam_key(chain))
1117
+ plan.notes.setdefault("composed_structures", []).append(
1118
+ _seam_key(chain))
1119
+ plan.notes["negotiated_layers"] = sorted(
1120
+ lay for lay, g in negotiated.items()
1121
+ if any(_seam_key(sm) in handled for sm in g.values()))
1122
+
1123
+ # ---- bind the remaining seams individually ----
1124
+ for name, members in group_families(seams).items():
1125
+ for seam in members:
1126
+ key = _seam_key(seam)
1127
+ if key in handled:
1128
+ continue
1129
+ cap = caps.get(key, {})
1130
+ try:
1131
+ bound = _bind_auto(model, seam, cap, plan, act_scales,
1132
+ negotiate_fp8, points=collector,
1133
+ fmt=formats.get(key),
1134
+ fmt_params=fmt_params.get(key))
1135
+ except (ValueError, RuntimeError) as refusal:
1136
+ plan.notes.setdefault("refused", []).append(
1137
+ (key, str(refusal)[:200]))
1138
+ continue
1139
+ if bound is None:
1140
+ continue
1141
+ if isinstance(bound, dict):
1142
+ plan.swaps.update(bound)
1143
+ _stream(list(bound))
1144
+ else:
1145
+ plan.swaps[seam.path] = bound
1146
+ _stream(seam.path)
1147
+ # ---- pre-FFN norm → FP8 producer pairs (seat-to-seat only) ----
1148
+ # The norm's sole consumer is the FFN seat, so the pair moves the
1149
+ # FFN's input quantize into the norm kernel. Both ends are seats
1150
+ # (the FP8_ONLY guard refuses anything else between them), the
1151
+ # smoke compares the pair against the bf16 form it would replace,
1152
+ # and the flip happens only when the measured chain is faster.
1153
+ if negotiate_fp8:
1154
+ _pair_vision_norm_fp8(model, plan, collector, _stream,
1155
+ probe=region_probe)
1156
+
1157
+ # ---- streaming window: the later stages record through the model,
1158
+ # and the streamed originals are meta now — so the bound seats stand
1159
+ # in for them for the duration. Temporarily attached with plain
1160
+ # setattr (no guards armed), undone before returning: the caller's
1161
+ # attach() must still find the originals at every path.
1162
+ _stream_temp = []
1163
+ if stream_store is not None:
1164
+ from .swap import resolve_parent as _srp, _get as _sgt, _set as _sst
1165
+ for _p, _m in list(plan.swaps.items()):
1166
+ try:
1167
+ _par, _at = _srp(model, _p)
1168
+ _stream_temp.append((_par, _at, _sgt(_par, _at)))
1169
+ _sst(_par, _at, _m)
1170
+ except (AttributeError, TypeError, ValueError):
1171
+ continue
1172
+
1173
+ if stream_store is not None:
1174
+ plan.notes["stream_store"] = stream_store
1175
+
1176
+ def _stream_window_undo():
1177
+ if not _stream_temp:
1178
+ return
1179
+ from .swap import _set as _sst2
1180
+ for _par, _at, _orig in reversed(_stream_temp):
1181
+ _sst2(_par, _at, _orig)
1182
+ _stream_temp.clear()
1183
+
1184
+ # every later stage may record through the seated model; whatever
1185
+ # they do the window closes, and a stage that dies rolls the whole
1186
+ # plan back — a half-routed model is worse than no plan
1187
+ try:
1188
+ # ---- qk_norm_rope: compose a packed QKV seam with host attention ----
1189
+ if "qk_norm_rope" in structures:
1190
+ from . import adapters as _adapters # noqa: F401 (registers)
1191
+ import inspect as _inspect
1192
+ _probe0 = (forward if callable(forward)
1193
+ else (forward[0] if forward else None))
1194
+ for adapter in _QK_NORM_ROPE_ADAPTERS:
1195
+ try:
1196
+ if "probe" in _inspect.signature(
1197
+ adapter.__call__).parameters:
1198
+ result = adapter(model, plan, probe=_probe0)
1199
+ else:
1200
+ result = adapter(model, plan)
1201
+ except (ValueError, RuntimeError) as refusal:
1202
+ plan.notes.setdefault("refused", []).append(
1203
+ ("qk_norm_rope", str(refusal)[:200]))
1204
+ continue
1205
+ if result is None:
1206
+ continue
1207
+ extras = result
1208
+ if extras.get("refused"):
1209
+ plan.notes.setdefault("refused", []).extend(
1210
+ extras["refused"])
1211
+ engaged = bool(
1212
+ extras.get("observed")
1213
+ or extras.get("revert")
1214
+ or extras.get("toggle")
1215
+ )
1216
+ if not engaged:
1217
+ continue
1218
+ plan.observed.update(extras.get("observed", {}))
1219
+ plan.revert.extend(extras.get("revert", ()))
1220
+ if extras.get("toggle") is not None:
1221
+ plan.notes["qk_norm_rope_toggle_index"] = len(plan.toggles)
1222
+ plan.toggles.append(extras["toggle"])
1223
+ plan.notes["qk_norm_rope_adapter"] = (
1224
+ type(adapter).__name__
1225
+ if hasattr(adapter, "__name__")
1226
+ else str(adapter)
1227
+ )
1228
+ break
1229
+ # ---- qkv_rope: packed biased QKV plus rotate-half RoPE ----
1230
+ if "qkv_rope" in structures:
1231
+ from . import adapters as _adapters # noqa: F401 (registers)
1232
+ for adapter in _QKV_ROPE_ADAPTERS:
1233
+ try:
1234
+ result = adapter(model, plan, caps)
1235
+ except (ValueError, RuntimeError) as refusal:
1236
+ plan.notes.setdefault("refused", []).append(
1237
+ ("qkv_rope", str(refusal)[:100]))
1238
+ continue
1239
+ if result is None:
1240
+ continue
1241
+ if result.get("refused"):
1242
+ plan.notes.setdefault("refused", []).extend(result["refused"])
1243
+ engaged = bool(
1244
+ result.get("observed")
1245
+ or result.get("revert")
1246
+ or result.get("toggle")
1247
+ )
1248
+ if not engaged:
1249
+ continue
1250
+ plan.observed.update(result.get("observed", {}))
1251
+ plan.revert.extend(result.get("revert", ()))
1252
+ if result.get("toggle") is not None:
1253
+ plan.toggles.append(result["toggle"])
1254
+ plan.notes["qkv_rope_adapter"] = type(adapter).__name__
1255
+ break
1256
+ # ---- attention_core: host-family adapters (fa2 seam) ----
1257
+ if "attention_core" in structures:
1258
+ from . import adapters as _adapters # noqa: F401 (registers)
1259
+ for adapter in _ATTENTION_ADAPTERS:
1260
+ try:
1261
+ # the adapter needs "run the host once", which is what a
1262
+ # thunk is. Handing it the caller's callable breaks the
1263
+ # sample entry, where that callable takes a sample —
1264
+ # the whole point of normalising the three ways in was
1265
+ # that nothing downstream should see the difference
1266
+ result = adapter(model, thunks[0],
1267
+ prefix_cadence=prefix_cadence)
1268
+ except (ValueError, RuntimeError) as refusal:
1269
+ plan.notes.setdefault("refused", []).append(
1270
+ ("attention_core", str(refusal)[:200]))
1271
+ continue
1272
+ if result is None:
1273
+ continue
1274
+ # an adapter may hand back a third element for the parts of
1275
+ # its seam that are not modules at paths: how to undo them,
1276
+ # and what to report
1277
+ att_swaps, update = result[0], result[1]
1278
+ extras = result[2] if len(result) > 2 else {}
1279
+ if extras.get("refused"):
1280
+ plan.notes.setdefault("refused", []).extend(
1281
+ extras["refused"])
1282
+ engaged = bool(
1283
+ att_swaps or update
1284
+ or extras.get("observed")
1285
+ or extras.get("revert")
1286
+ or extras.get("toggle")
1287
+ )
1288
+ if not engaged:
1289
+ continue
1290
+ plan.swaps.update(att_swaps)
1291
+ plan.observed.update(extras.get("observed", {}))
1292
+ if extras.get("attention_variants"):
1293
+ plan.notes.setdefault(
1294
+ "attention_core_variants", {}).update(
1295
+ extras["attention_variants"])
1296
+ plan.revert.extend(extras.get("revert", ()))
1297
+ if extras.get("toggle") is not None:
1298
+ plan.toggles.append(extras["toggle"])
1299
+ if update is not None:
1300
+ plan.updates.append(update)
1301
+ plan.notes["attention_adapter"] = type(adapter).__name__ \
1302
+ if hasattr(adapter, "__name__") else str(adapter)
1303
+ break
1304
+
1305
+
1306
+ # ---- gated_delta_core: stateful host callable adapters ----
1307
+ if "gated_delta_core" in structures:
1308
+ from . import adapters as _adapters # noqa: F401 (registers)
1309
+ for adapter in _GATED_DELTA_ADAPTERS:
1310
+ try:
1311
+ # adapters that declare scheme awareness receive the
1312
+ # active scheme; the rest keep the two-argument call
1313
+ if getattr(adapter, "scheme_aware", False):
1314
+ result = adapter(model, thunks[0], scheme=scheme_obj)
1315
+ else:
1316
+ result = adapter(model, thunks[0])
1317
+ except (ValueError, RuntimeError) as refusal:
1318
+ plan.notes.setdefault("refused", []).append(
1319
+ ("gated_delta_core", str(refusal)[:120]))
1320
+ continue
1321
+ if result is None:
1322
+ continue
1323
+ plan.observed.update(result.get("observed", {}))
1324
+ plan.revert.extend(result.get("revert", ()))
1325
+ if result.get("toggle") is not None:
1326
+ plan.toggles.append(result["toggle"])
1327
+ plan.notes["gated_delta_adapter"] = type(adapter).__name__ \
1328
+ if hasattr(adapter, "__name__") else str(adapter)
1329
+ break
1330
+
1331
+ # ---- one step-scoped style materialisation per conditioning stream
1332
+ # Every adaptive-norm producer on one stream resolves the same step,
1333
+ # so the whole stream's styles are fixed for the step's duration.
1334
+ # Materialising them once beats materialising them per call by the
1335
+ # launch count, which is what that work actually costs. Runs before
1336
+ # the block assembly: a block holds its producers directly and drops
1337
+ # them from the swap map, so afterwards they are no longer findable
1338
+ # here.
1339
+ _attach_brokers(caps, plan, say)
1340
+
1341
+ # ---- decoder_block: compose the bound sublayers into one block ----
1342
+ # last, because it is assembled from what the region structures
1343
+ # produced. The swaps it absorbs are dropped from the plan: the
1344
+ # block holds those modules directly, and a swap that also targeted
1345
+ # the host child would leave two live copies of the same seam.
1346
+ for seam in (s for s in seams if s.structure == "decoder_block"):
1347
+ try:
1348
+ block = _bind_block(
1349
+ model, seam, caps.get(_seam_key(seam), {}), plan)
1350
+ except (ValueError, RuntimeError) as refusal:
1351
+ plan.notes.setdefault("refused", []).append(
1352
+ (seam.path + " [block]", str(refusal)[:200]))
1353
+ continue
1354
+ if block is None:
1355
+ continue
1356
+ for child in _BLOCK_OWNED:
1357
+ plan.swaps.pop(seam.path + "." + child, None)
1358
+ plan.swaps[seam.path] = block
1359
+
1360
+ except BaseException:
1361
+ _stream_window_undo()
1362
+ plan.abort()
1363
+ raise
1364
+ finally:
1365
+ _stream_window_undo()
1366
+
1367
+ # what discovery took on trust, for the seams that actually bound. An
1368
+ # assumption that reaches the model without reaching the receipt is
1369
+ # indistinguishable from something that was checked.
1370
+ assumed = [(s.path, note) for s in seams if s.assumptions
1371
+ and s.path in plan.swaps for note in s.assumptions]
1372
+ if assumed:
1373
+ plan.notes["assumed"] = assumed
1374
+ say(f"{len(assumed)} seam(s) carry an assumption the parity gate "
1375
+ f"has to check (see notes['assumed'])")
1376
+
1377
+ if plan_notes_calibration:
1378
+ # the calibration method is part of the result, not a detail of
1379
+ # how it was produced: a parity band means something different
1380
+ # depending on how much of the distribution it was scaled from
1381
+ plan.notes["calibration"] = plan_notes_calibration
1382
+ # and the receipt itself is the repo's, so a structures attachment
1383
+ # answers ``precision_spec`` the same way a frontend does
1384
+ from .points import precision_spec as _spec
1385
+ plan.precision_spec = _spec(collector, plan_notes_calibration)
1386
+ if plan_refusals:
1387
+ plan.notes.setdefault("refused", []).extend(plan_refusals)
1388
+ # Skipping a package this host cannot supply is what lets one plan
1389
+ # build everywhere, but it must never be silent: a package that is
1390
+ # broken here and one that was never shipped here both come out as
1391
+ # "skipped", and only the first is a defect. Carry the raw failures
1392
+ # into the plan so they reach the receipt.
1393
+ from .impls import unavailable_report
1394
+ unavailable = unavailable_report()
1395
+ if unavailable:
1396
+ plan.notes["kernel_unavailable"] = unavailable
1397
+ say(f"{len(unavailable)} kernel package(s) unavailable here: "
1398
+ + ", ".join(f"{row['repo']} ({row['error']})"
1399
+ for row in unavailable))
1400
+ say(f"bound {len(plan.swaps)} seam(s), "
1401
+ f"{len(plan.notes.get('refused', []))} refused")
1402
+ return plan
1403
+
1404
+
1405
+ def _attach_brokers(caps, plan, say) -> None:
1406
+ from .impls.adaln_producer import AdaLNProducer, bind_style_broker
1407
+
1408
+ groups: dict[tuple, list] = {}
1409
+ for path, module in plan.swaps.items():
1410
+ if not isinstance(module, AdaLNProducer):
1411
+ continue
1412
+ cap = caps.get(path, {})
1413
+ order = cap.get("order")
1414
+ if order is None or not cap.get("pairs"):
1415
+ continue
1416
+ # one broker per (stream, style width, row count): producers
1417
+ # that differ in any of those cannot share a buffer
1418
+ key = (_stream_key(cap["pairs"]), int(module.styles.shape[-1]),
1419
+ int(module.resid.shape[0]))
1420
+ groups.setdefault(key, []).append((order, path, module))
1421
+
1422
+ for key, members in groups.items():
1423
+ # the writer is the producer the host calls first, taken from the
1424
+ # observed order of the calibration pass — not from the module
1425
+ # tree's order, which need not match the forward's
1426
+ members.sort(key=lambda entry: entry[0])
1427
+ try:
1428
+ broker = bind_style_broker([m for _, _, m in members], key[2])
1429
+ except (ValueError, RuntimeError) as refusal:
1430
+ plan.notes.setdefault("refused", []).append(
1431
+ (f"style_broker[{key[1]}x{key[2]}]", str(refusal)[:200]))
1432
+ continue
1433
+ if broker is None:
1434
+ continue
1435
+ plan.notes.setdefault("brokers", []).append(
1436
+ {"slots": broker.slots, "rows": key[2], "width": key[1],
1437
+ "writer": members[0][1]})
1438
+ say(f"style broker: {broker.slots} producer(s) share one "
1439
+ f"step-scoped materialisation (writer {members[0][1]})")
1440
+
1441
+
1442
+ _BLOCK_OWNED = ("input_layernorm", "post_attention_layernorm", "mlp")
1443
+
1444
+
1445
+ def _cond_kw(host) -> str:
1446
+ """The keyword the host threads its conditioning through."""
1447
+ import inspect
1448
+ try:
1449
+ params = list(inspect.signature(host.forward).parameters)
1450
+ except (TypeError, ValueError):
1451
+ params = []
1452
+ for name in ("adarms_cond", "cond", "temb", "emb"):
1453
+ if name in params:
1454
+ return name
1455
+ return "adarms_cond"
1456
+
1457
+
1458
+ def _bind_block(model, seam, cap, plan):
1459
+ """Assemble one decoder_block from its already-bound sublayers."""
1460
+ from .impls.decoder_block import bind_decoder_block
1461
+
1462
+ prod_in = plan.swaps.get(seam.path + ".input_layernorm")
1463
+ prod_out = plan.swaps.get(seam.path + ".post_attention_layernorm")
1464
+ ffn = plan.swaps.get(seam.path + ".mlp")
1465
+ if prod_in is None or prod_out is None or ffn is None:
1466
+ # a sublayer that did not bind leaves the host block intact:
1467
+ # the block structure adds composition, it does not substitute
1468
+ # for the region seams it is made of
1469
+ return None
1470
+ host = _resolve(model, seam.path)
1471
+ # the attention sublayer is family-specific (where the attention runs
1472
+ # and which rotary form it uses), so it comes from the same adapters
1473
+ # that bound the attention core. None keeps the host's attention
1474
+ # module, which is the pre-block behaviour.
1475
+ attn = None
1476
+ for adapter in _ATTENTION_ADAPTERS:
1477
+ builder = getattr(adapter, "sublayer", None)
1478
+ if builder is None:
1479
+ continue
1480
+ attn = builder(host)
1481
+ if attn is not None:
1482
+ break
1483
+ if attn is not None:
1484
+ _alias_kv_region(plan, seam.path, attn)
1485
+ return bind_decoder_block(
1486
+ host, prod_in, prod_out, ffn, cond_kw=_cond_kw(host),
1487
+ returns_tuple=bool(cap.get("returns_tuple")), attn=attn)
1488
+
1489
+
1490
+ def _alias_kv_region(plan, path: str, sublayer) -> None:
1491
+ """Let the packed projections write into the core's packed KV region.
1492
+
1493
+ Both sides can express this (see ``beta.joins``); the qualification
1494
+ is that nothing transforms the tensor in between. Value goes straight
1495
+ from the projection to the kernel and qualifies. Key does not on this
1496
+ family: a rotary embedding runs after the projection, so aliasing it
1497
+ would leave untransformed keys in the packed region — writing the
1498
+ transformed ones back is the copy this was meant to remove. Hosts
1499
+ without a rotary step qualify for both; the attribute is general and
1500
+ the qualification is per join.
1501
+ """
1502
+ from .impls.qkv_pack import PackedLinear
1503
+
1504
+ head = plan.swaps.get(path + ".self_attn.q_proj")
1505
+ core = getattr(sublayer, "core", None)
1506
+ if not isinstance(head, PackedLinear) or core is None:
1507
+ return
1508
+ if not hasattr(core, "alias_suffix"):
1509
+ return
1510
+ _, v_region = core.alias_suffix(key=False, value=True)
1511
+ if v_region is None:
1512
+ return
1513
+ try:
1514
+ head.alias_stash(2, v_region) # sibling order q, k, v
1515
+ except (ValueError, RuntimeError) as refusal:
1516
+ core._alias_v = False
1517
+ plan.notes.setdefault("refused", []).append(
1518
+ (path + " [kv alias]", str(refusal)[:200]))
1519
+ return
1520
+ plan.notes.setdefault("aliased_kv", []).append(path)
1521
+
1522
+ # The joint q|k view is deliberately not enabled here. q and k are
1523
+ # one contiguous run of the packed output and take the same rotary
1524
+ # arithmetic, so one pass over the pair should replace two — and
1525
+ # measured, it replaces nothing: the rotary kernels keep their exact
1526
+ # launch counts (180/162/63) because the compiler splits the merged
1527
+ # pass back apart, fusing it into each consumer (q is made
1528
+ # contiguous for the kernel, k is copied into the packed region).
1529
+ # Paired timing: -0.014 ms on 23.1, which is below the margin this
1530
+ # stack ships at. Expressing "do this once" in tensor ops does not
1531
+ # survive a compiler that re-derives its fusion boundaries from the
1532
+ # consumers; the style broker only survived by being opaque, and an
1533
+ # opaque wrapper here would be worse, since the rotary would then run
1534
+ # as several eager ops instead of one fused kernel. Merging these
1535
+ # needs a rotary kernel, not a rearrangement. The capability stays
1536
+ # on the impl for a host where the arithmetic is not launch-bound.
1537
+
1538
+
1539
+ def _bind_auto(model, seam, cap, plan, act_scales, negotiate_fp8,
1540
+ points=None, fmt=None, fmt_params=None):
1541
+ """Route one seam to its impl with the calibrated scales.
1542
+
1543
+ ``points`` is the reduced collector: every scale an impl needs is one
1544
+ float looked up by (path, spec point name). No activation tensors are
1545
+ threaded through here, because none are needed — the two scales that
1546
+ used to be recomputed from held inputs are measured at the GEMM whose
1547
+ input they are (:mod:`.points`).
1548
+
1549
+ ``fmt`` is the scheme's per-seam format routing. ``None`` binds the
1550
+ structure's default impl; a named format binds that variant instead,
1551
+ and a name with no variant for this structure fails loudly — the
1552
+ scheme author's error surfaces at bind time, not as accuracy.
1553
+ ``fmt_params`` is the decision's recipe payload for that format
1554
+ (algorithm parameters, never bytes), handed to the variant's binder.
1555
+ """
1556
+ dev0 = None
1557
+ if model is not None:
1558
+ dev0 = next(model.parameters(), torch.empty(0)).device
1559
+ if dev0 is not None and dev0.type == "cuda":
1560
+ free, _total = torch.cuda.mem_get_info(dev0)
1561
+ if free < (512 << 20):
1562
+ # binding is a transaction against a budget: below the
1563
+ # headroom every further seat is refused with the number,
1564
+ # instead of eating the remainder and failing later as an
1565
+ # unattributable OOM in the first treated forward
1566
+ raise ValueError(
1567
+ f"insufficient_vram(free={free >> 20}MiB, "
1568
+ "headroom=512MiB) — host keeps this seam")
1569
+
1570
+ from .impls.decoder_ffn import fp8_static as ffn_impl
1571
+ from .impls.vision_ffn import fp8_static as vis_impl
1572
+
1573
+ def scale(name, path=None):
1574
+ return None if points is None else points.scale(path or seam.path,
1575
+ name)
1576
+
1577
+ custom = _STRUCTURE_BINDERS.get(seam.structure)
1578
+ if custom is not None:
1579
+ return custom(model, seam, cap, points=points, fmt=fmt,
1580
+ fmt_params=fmt_params)
1581
+
1582
+ if fmt and not (seam.structure == "qkv_pack"
1583
+ and fmt in ("bf16_pack", "nvfp4_balance")) \
1584
+ and not (seam.structure == "vision_ffn"
1585
+ and fmt == "nvfp4_balance") \
1586
+ and seam.structure not in ("decoder_ffn", "linear_proj"):
1587
+ raise ValueError(f"scheme routed {seam.structure} to format "
1588
+ f"{fmt!r}, which has no impl variant here")
1589
+
1590
+ if seam.structure == "decoder_ffn":
1591
+ if fmt in ("w8a16_static", "w4a16_static"):
1592
+ if fmt == "w8a16_static":
1593
+ from .impls.decoder_ffn import w8a16_static as wq_impl
1594
+ else:
1595
+ from .impls.decoder_ffn import w4a16_static as wq_impl
1596
+
1597
+ # two callers, two layout conventions: ``seam_weights``
1598
+ # serves the fp8 impl transposed ([D, F]); these binders are
1599
+ # checkpoint-native ([F, D]) and their dim check passes with
1600
+ # the names swapped, so handing them the transposed dict
1601
+ # binds a guard with k = F and every call falls back.
1602
+ # Transpose back here, at the seam between the conventions.
1603
+ w = seam_weights(model, seam)
1604
+ w = dict(w,
1605
+ w_gate=w["w_gate"].t().contiguous(),
1606
+ w_up=w["w_up"].t().contiguous(),
1607
+ w_down=w["w_down"].t().contiguous())
1608
+ return wq_impl.bind_mlp_seam(
1609
+ w, variant=seam.variant,
1610
+ original=_resolve(model, seam.path))
1611
+ if fmt not in (None, "fp8_static"):
1612
+ raise ValueError(f"scheme routed decoder_ffn to format "
1613
+ f"{fmt!r}, which has no impl variant here")
1614
+ in_s = scale("x_after_norm")
1615
+ hid_s = scale("act_after_mul", seam.path + ".down_proj")
1616
+ if in_s is None or hid_s is None:
1617
+ return None
1618
+ return ffn_impl.bind_mlp_seam(
1619
+ seam_weights(model, seam), variant=seam.variant,
1620
+ input_scale=in_s, hidden_scale=hid_s,
1621
+ original=_resolve(model, seam.path))
1622
+
1623
+ if seam.structure == "vision_ffn":
1624
+ fc2 = (seam.fc_attrs or ("fc1", "fc2"))[1]
1625
+ if fmt == "nvfp4_balance":
1626
+ from .impls.vision_ffn import nvfp4_balance as vis_w4
1627
+ chan_in = points.channel_amax(seam.path, "x_after_norm")
1628
+ chan_hid = points.channel_amax(
1629
+ seam.path + "." + fc2, "hidden_after_act")
1630
+ if chan_in is None or chan_hid is None:
1631
+ return None
1632
+ return vis_w4.bind_mlp_seam(
1633
+ seam_weights(model, seam), channel_in=chan_in,
1634
+ channel_hidden=chan_hid,
1635
+ original=_resolve(model, seam.path),
1636
+ **dict(fmt_params or {}))
1637
+ in_s = scale("x_after_norm")
1638
+ hid_s = scale("hidden_after_act", seam.path + "." + fc2)
1639
+ if in_s is None or hid_s is None:
1640
+ return None
1641
+ rows_seen = points.row_profile(seam.path, "x_after_norm")
1642
+ rows_med = rows_seen[len(rows_seen) // 2] if rows_seen else 1 << 30
1643
+ if rows_med <= 64:
1644
+ # the small-M denoise band: the measured band decision for
1645
+ # this box (recorded by a band run, cached per device)
1646
+ # routes these seats — never a shape rule alone, never a
1647
+ # device name
1648
+ from .decisions import lookup as _band_lookup
1649
+ if _band_lookup("groot_dit", default="fp8") == "fp4":
1650
+ from .impls.vision_ffn import nvfp4_balance as vis_w4
1651
+ w = seam_weights(model, seam)
1652
+ dev = w["w_fc1"].device
1653
+ try:
1654
+ # flat channel vectors: no balance folded — the W4
1655
+ # quantizer alone, judged by the parity gate
1656
+ return vis_w4.bind_mlp_seam(
1657
+ w,
1658
+ channel_in=torch.ones(
1659
+ w["w_fc1"].shape[1], device=dev),
1660
+ channel_hidden=torch.ones(
1661
+ w["w_fc2"].shape[1], device=dev),
1662
+ original=_resolve(model, seam.path),
1663
+ fuse_wire=True)
1664
+ except (ValueError, RuntimeError):
1665
+ pass
1666
+ return vis_impl.bind_mlp_seam(
1667
+ seam_weights(model, seam), input_scale=in_s,
1668
+ hidden_scale=hid_s, original=_resolve(model, seam.path))
1669
+
1670
+ if seam.structure == "modnorm_qkv_chain":
1671
+ if seam.variant.get("modulation") == "per_token_table":
1672
+ from .impls.modnorm_qkv_chain import fp8_ptok_table as chain
1673
+ return chain.bind_block_seam(model, seam, points=points)
1674
+ # the scale_shift form composes through producer negotiation
1675
+ return None
1676
+
1677
+ if seam.structure == "norm_fused":
1678
+ from .impls.norm_fused import bind_norm_fused
1679
+ return bind_norm_fused(
1680
+ _resolve(model, seam.path),
1681
+ host_dtypes=(None if points is None
1682
+ else points.seen_dtypes(seam.path, "x")))
1683
+
1684
+ if seam.structure == "linear_proj":
1685
+ if fmt == "nvfp4_balance":
1686
+ from .impls.linear_proj import nvfp4_balance as proj_w4
1687
+ chan = points.channel_amax(seam.path, "x")
1688
+ if chan is None:
1689
+ return None
1690
+ return proj_w4.bind_proj_seam(
1691
+ seam_weights(model, seam), channel_amax=chan,
1692
+ original=_resolve(model, seam.path),
1693
+ **dict(fmt_params or {}))
1694
+ if fmt == "w8a16_static":
1695
+ # weight-only decode band: no calibration scale to look up,
1696
+ # and the weight dict is already the kernel's [N, K] layout
1697
+ from .impls.linear_proj import w8a16_static as proj_w8
1698
+ return proj_w8.bind_proj_seam(
1699
+ seam_weights(model, seam),
1700
+ original=_resolve(model, seam.path))
1701
+ if fmt not in (None, "fp8_static"):
1702
+ raise ValueError(f"scheme routed linear_proj to format "
1703
+ f"{fmt!r}, which has no impl variant here")
1704
+ in_s = scale("x")
1705
+ if in_s is None:
1706
+ return None
1707
+ from .impls.linear_proj import fp8_static as proj_impl
1708
+ return proj_impl.bind_proj_seam(
1709
+ seam_weights(model, seam), input_scale=in_s,
1710
+ row_profile=points.row_profile(seam.path, "x"),
1711
+ original=_resolve(model, seam.path))
1712
+
1713
+ if seam.structure == "patch_projection":
1714
+ from .impls.patch_projection import bind_flat_patch_projection
1715
+
1716
+ forms = tuple(cap.get("patch_inputs", ()))
1717
+ if not forms or any(
1718
+ form is None or form[0] != seam.dims["K"] for form in forms
1719
+ ):
1720
+ raise ValueError(
1721
+ "patch_projection: calibrated host input is not "
1722
+ f"preflattened full-patch rows with K={seam.dims['K']}"
1723
+ )
1724
+ rows = points.row_profile(seam.path, "x") if points else ()
1725
+ dtypes = points.seen_dtypes(seam.path, "x") if points else ()
1726
+ return bind_flat_patch_projection(
1727
+ seam_weights(model, seam),
1728
+ row_profile=rows,
1729
+ host_dtypes=dtypes,
1730
+ original=_resolve(model, seam.path),
1731
+ )
1732
+
1733
+ if seam.structure == "qkv_pack":
1734
+ if fmt == "bf16_pack":
1735
+ if seam.variant.get("bind") == "module":
1736
+ raise ValueError(
1737
+ "qkv_pack bf16_pack v1 supports leaf binding only")
1738
+ first = (seam.pack_attrs or ("q_proj",))[0]
1739
+ rows_seen = points.row_profile(
1740
+ seam.path + "." + first, "x")
1741
+ if not rows_seen:
1742
+ return None
1743
+ from .impls.qkv_pack import bf16 as pack_impl
1744
+ block = _resolve(model, seam.path)
1745
+ mods = [getattr(block, attr) for attr in seam.pack_attrs]
1746
+ parts = pack_impl.bind_qkv_pack(mods, rows=max(rows_seen))
1747
+ return {seam.path + "." + attr: mod
1748
+ for attr, mod in zip(seam.pack_attrs, parts)}
1749
+ if fmt == "nvfp4_balance":
1750
+ if seam.variant.get("bind") == "module":
1751
+ raise ValueError(
1752
+ "qkv_pack nvfp4_balance supports leaf binding only")
1753
+ from .impls.qkv_pack import nvfp4_balance as pack_w4
1754
+ first = (seam.pack_attrs or ("q_proj",))[0]
1755
+ chan = points.channel_amax(seam.path + "." + first, "x")
1756
+ rows_seen = points.row_profile(seam.path + "." + first, "x")
1757
+ if chan is None or not rows_seen:
1758
+ return None
1759
+ block = _resolve(model, seam.path)
1760
+ mods = [getattr(block, a) for a in seam.pack_attrs]
1761
+ parts = pack_w4.bind_qkv_pack(
1762
+ mods, channel_amax=chan, rows=max(rows_seen),
1763
+ **dict(fmt_params or {}))
1764
+ return {seam.path + "." + a: m
1765
+ for a, m in zip(seam.pack_attrs, parts)}
1766
+ from .impls.qkv_pack import bind_attn_block, bind_qkv_pack
1767
+ first = (seam.pack_attrs or ("q_proj",))[0]
1768
+ amax = None if points is None else points.amax(
1769
+ seam.path + "." + first, "x")
1770
+ if amax is None:
1771
+ return None
1772
+ block = _resolve(model, seam.path)
1773
+ rows = points.row_profile(seam.path + "." + first, "x")
1774
+ # The packed implementation preallocates scratch/stash storage but
1775
+ # the Hub entry accepts any logical M covered by that storage. Use
1776
+ # the largest calibrated observation as capacity; choosing the
1777
+ # median here turns a valid variable-row call into a guard fallback
1778
+ # and can also under-allocate when calibration itself has buckets.
1779
+ cap = dict(cap or {}, rows=max(rows) if rows else 1)
1780
+ act_scale = torch.tensor(
1781
+ [max(amax / 448.0, 1e-8)],
1782
+ device=getattr(block, first).weight.device)
1783
+ if seam.variant.get("bind") == "module":
1784
+ # the whole block: packed projections *and* the attention
1785
+ # compute dtype (hosts that run SDPA in fp32 pay for it)
1786
+ return {seam.path: bind_attn_block(
1787
+ block, act_scale, rows=cap["rows"],
1788
+ sdpa_dtype=torch.bfloat16)}
1789
+ mods = [getattr(block, a) for a in seam.pack_attrs]
1790
+ parts = bind_qkv_pack(mods, act_scale, rows=cap["rows"],
1791
+ in_dtype="bf16_fused_quant")
1792
+ return {seam.path + "." + a: m
1793
+ for a, m in zip(seam.pack_attrs, parts)}
1794
+
1795
+ if seam.structure == "adaln_producer":
1796
+ from .impls.adaln_producer import (bind_adaln_producer,
1797
+ bind_style_table)
1798
+ if not cap.get("pairs"):
1799
+ return None
1800
+ norm = _resolve(model, seam.path)
1801
+ proj = getattr(norm, seam.cond_attr)
1802
+ key = _stream_key(cap["pairs"])
1803
+ loc = plan.notes.setdefault("_locators", {}).get(key)
1804
+ table = bind_style_table(proj, cap["pairs"], locator=loc)
1805
+ plan.notes["_locators"][key] = table.locator
1806
+ return {seam.path + "." + seam.cond_attr: table}
1807
+
1808
+ return None
1809
+
1810
+
1811
+ class _Eager(torch.nn.Module):
1812
+ """Wrap a module so its forward runs outside the compiled region.
1813
+
1814
+ An fp8-emitting seam's arithmetic, if traced by inductor, gets fused
1815
+ into fp8 math (illegal on sm120 triton) — and the quantize even
1816
+ reaches back across the boundary, so inductor casts the host's own
1817
+ gated residual to fp8 to feed it. The hand recipes never hit this
1818
+ because the whole denoise block froze to eager. A swapped-in module
1819
+ does not inherit that freezing, so fp8 seams declare it. Overriding
1820
+ the instance ``forward`` is not enough (dynamo inlines the class
1821
+ forward); the disable must sit on a class method, which is what this
1822
+ wrapper provides. The kernels are opaque either way, so eager here
1823
+ is a graph break, not real work.
1824
+ """
1825
+
1826
+ def __init__(self, inner: torch.nn.Module):
1827
+ super().__init__()
1828
+ self.inner = inner
1829
+
1830
+ @torch._dynamo.disable
1831
+ def forward(self, *args, **kwargs):
1832
+ return self.inner(*args, **kwargs)
1833
+
1834
+ def __getattr__(self, name):
1835
+ try:
1836
+ return super().__getattr__(name)
1837
+ except AttributeError:
1838
+ return getattr(super().__getattr__("inner"), name)
1839
+
1840
+
1841
+ def _eager(module):
1842
+ return _Eager(module)
1843
+
1844
+
1845
+ def _bind_negotiated(model, p_seam, k_seam, p_cap, points, scale, rows,
1846
+ plan):
1847
+ """Bind an fp8 producer and the pack it feeds as one chain.
1848
+
1849
+ This is the combination the structure layer exists for: neither half
1850
+ is worth much alone (a small-M fp8 projection barely beats BF16, a
1851
+ producer that only reshapes styles saves nothing), but together the
1852
+ producer's fused quantize removes the consumer's input quantization
1853
+ entirely and hands a clean fp8 seam downstream.
1854
+ """
1855
+ from .impls.adaln_producer import bind_adaln_producer
1856
+ from .impls.qkv_pack import bind_qkv_pack
1857
+
1858
+ norm = _resolve(model, p_seam.path)
1859
+ consumer = _resolve(model, k_seam.path)
1860
+ key = _stream_key(p_cap["pairs"])
1861
+ loc = plan.notes.setdefault("_locators", {}).get(key)
1862
+ dim, form = _adaln_form(p_cap, points, p_seam)
1863
+ prod = bind_adaln_producer(
1864
+ norm, p_cap["pairs"], act_scale=scale, rows=rows,
1865
+ dim=dim, locator=loc, norm=form)
1866
+ plan.notes["_locators"][key] = prod.locator
1867
+
1868
+ swaps = {p_seam.path: prod}
1869
+ if k_seam.structure == "decoder_ffn":
1870
+ from .impls.decoder_ffn import fp8_static as ffn_impl
1871
+ # the input scale is the one the producer upstream will quantize
1872
+ # with — the same number, because the producer's output is this
1873
+ # consumer's input; the hidden scale is measured at the down
1874
+ # projection whose input it is
1875
+ w = seam_weights(model, k_seam)
1876
+ bound = ffn_impl.bind_mlp_seam(
1877
+ w, variant={**k_seam.variant, "in_dtype": "fp8_static"},
1878
+ input_scale=float(scale.item()),
1879
+ hidden_scale=points.scale(k_seam.path + ".down_proj",
1880
+ "act_after_mul"),
1881
+ original=consumer)
1882
+ swaps[k_seam.path] = bound
1883
+ return swaps
1884
+ if k_seam.structure == "linear_proj":
1885
+ from .impls.linear_proj import fp8_static as proj_impl
1886
+ swaps[k_seam.path] = proj_impl.bind_proj_seam(
1887
+ seam_weights(model, k_seam),
1888
+ input_scale=float(scale.item()),
1889
+ row_profile=points.row_profile(k_seam.path, "x"),
1890
+ original=consumer,
1891
+ in_dtype="fp8_static",
1892
+ )
1893
+ return swaps
1894
+ mods = [getattr(consumer, a) for a in k_seam.pack_attrs]
1895
+ parts = bind_qkv_pack(mods, scale, rows=rows,
1896
+ in_dtype="fp8_static")
1897
+ # ---- the FP4 wire is a second candidate for the same seats: a
1898
+ # producer that norms straight into packed NVFP4 + swizzled scale
1899
+ # factors, and a pack that consumes the wire with no quantize of
1900
+ # its own. Which chain is faster is a property of this device's
1901
+ # GEMM bands at this row count — so it is measured here, on the
1902
+ # calibrated shape with the real conditioning, and the winner takes
1903
+ # the seats. A candidate that cannot build or run loses by default.
1904
+ import os
1905
+ if dim >= 512 and not os.environ.get("FRT_DISABLE_FP4_RACE"):
1906
+ try:
1907
+ prod4 = bind_adaln_producer(
1908
+ norm, p_cap["pairs"], act_scale=None, rows=rows,
1909
+ dim=dim, locator=prod.locator, norm=form,
1910
+ out_format="nvfp4")
1911
+ from .impls.qkv_pack import nvfp4_balance as pack_w4
1912
+ parts4 = pack_w4.bind_qkv_pack(
1913
+ mods, channel_amax=None, rows=rows, wire=True)
1914
+ parts4[0].accept_wire(prod4.wire_sfa)
1915
+ cond0 = p_cap["pairs"][0][0].detach()
1916
+ dev = prod4.wire_sfa.device
1917
+ x_bench = torch.randn(rows, dim, device=dev,
1918
+ dtype=torch.bfloat16)
1919
+
1920
+ def _chain_ms(producer, head, iters=30):
1921
+ def once():
1922
+ y = producer(x_bench, cond0)
1923
+ head(y)
1924
+ for _ in range(5):
1925
+ once()
1926
+ torch.cuda.synchronize()
1927
+ start = torch.cuda.Event(True)
1928
+ end = torch.cuda.Event(True)
1929
+ start.record()
1930
+ for _ in range(iters):
1931
+ once()
1932
+ end.record()
1933
+ torch.cuda.synchronize()
1934
+ return start.elapsed_time(end) / iters
1935
+
1936
+ with torch.no_grad():
1937
+ a_ms = _chain_ms(prod, parts[0])
1938
+ b_ms = _chain_ms(prod4, parts4[0])
1939
+ race = {"layer": p_seam.path, "rows": int(rows),
1940
+ "dim": int(dim), "fp8_chain_ms": round(a_ms, 4),
1941
+ "nvfp4_wire_ms": round(b_ms, 4),
1942
+ "winner": "nvfp4_wire" if b_ms < a_ms else
1943
+ "fp8_chain"}
1944
+ plan.notes.setdefault("format_race", []).append(race)
1945
+ if b_ms < a_ms:
1946
+ prod, parts = prod4, parts4
1947
+ swaps[p_seam.path] = prod4
1948
+ except (ValueError, RuntimeError, KeyError, OSError) as lost:
1949
+ plan.notes.setdefault("format_race", []).append(
1950
+ {"layer": p_seam.path,
1951
+ "winner": "fp8_chain",
1952
+ "nvfp4_wire": f"refused: {str(lost)[:200]}"})
1953
+ swaps.update({k_seam.path + "." + a: m
1954
+ for a, m in zip(k_seam.pack_attrs, parts)})
1955
+ return swaps
1956
+
1957
+
1958
+ def _calibration_thunks(forward, frames, samples):
1959
+ """Turn the three ways of asking for calibration into one list.
1960
+
1961
+ They differ only in where a frame's input comes from, so they end as
1962
+ the same thing: a list of callables, each of which runs the host
1963
+ once. Keeping them one axis is what stops "how much calibration" and
1964
+ "how to run the host" from becoming two interfaces that can disagree.
1965
+ """
1966
+ if samples is not None:
1967
+ if not callable(forward):
1968
+ raise ValueError(
1969
+ "auto_swaps: with samples=, forward takes one sample")
1970
+ taken = list(samples) if frames is None else [
1971
+ s for _, s in zip(range(max(1, frames)), samples)]
1972
+ if not taken:
1973
+ raise ValueError("auto_swaps: samples is empty")
1974
+ return [(lambda s=s: forward(s)) for s in taken], "samples"
1975
+ if isinstance(forward, (list, tuple)):
1976
+ if not forward:
1977
+ raise ValueError("auto_swaps: no forward thunks given")
1978
+ if frames is not None and frames != len(forward):
1979
+ raise ValueError(
1980
+ f"auto_swaps: {len(forward)} thunk(s) given but "
1981
+ "observations= and a thunk list are alternatives, "
1982
+ "not a pair; the thunks decide")
1983
+ return list(forward), "thunks"
1984
+ if not callable(forward):
1985
+ raise ValueError("auto_swaps: forward must be callable")
1986
+ return [forward] * max(1, frames or 1), "forward"
1987
+
1988
+
1989
+ def _adaln_form(cap, points, seam) -> tuple[int, str]:
1990
+ """Read the producer's width and form off the calibration.
1991
+
1992
+ Both were assumed before: the form was hard-coded to rms and the
1993
+ width taken as ``style_width // 3``. That holds only where the style
1994
+ carries three parts. A host whose adaptive norm emits (scale, shift)
1995
+ — the layer form — got bound as rms at two thirds of its real width,
1996
+ and the plan built cleanly and then could not run. It took a second
1997
+ host and an actual forward to see it, because nothing on the way
1998
+ there had to disagree.
1999
+
2000
+ The norm's own input says how wide it is, and the ratio to the style
2001
+ says which form it is. Neither is a guess.
2002
+ """
2003
+ dim = points.width(seam.path, "x")
2004
+ if dim is None:
2005
+ raise ValueError(
2006
+ "adaln_producer: the norm's own input was never observed, so "
2007
+ "the form cannot be told from the style width alone")
2008
+ style_width = int(cap["pairs"][0][1].shape[-1])
2009
+ if style_width == 3 * dim:
2010
+ return dim, "rms" # scale, shift, gate
2011
+ if style_width == 2 * dim:
2012
+ return dim, "layer" # scale, shift
2013
+ raise ValueError(
2014
+ f"adaln_producer: style width {style_width} is neither two nor "
2015
+ f"three times the norm width {dim} — the modulation is a shape "
2016
+ "this structure does not model")
2017
+
2018
+
2019
+ def _stream_key(pairs) -> str:
2020
+ """Identify the conditioning stream a producer was calibrated on.
2021
+
2022
+ Locators were keyed by seam family, which gives every family its own
2023
+ lookup even when they all read the same conditioning — the two norms
2024
+ of one block among them. Keying by the observed conditioning instead
2025
+ shares one locator across the whole stream. It is safe by
2026
+ construction rather than by convention: the key is a digest of the
2027
+ conditioning rows themselves, so two seams share a locator only when
2028
+ they saw byte-identical inputs, and identical inputs resolve to
2029
+ identical indices whichever seam built the table.
2030
+ """
2031
+ import hashlib
2032
+
2033
+ digest = hashlib.blake2b(digest_size=16)
2034
+ for cond, _ in pairs:
2035
+ c = cond.detach().reshape(-1, cond.shape[-1]).to(torch.float32)
2036
+ digest.update(c.cpu().numpy().tobytes())
2037
+ return digest.hexdigest()
2038
+
2039
+
2040
+ def _layer_of(path: str) -> str:
2041
+ """The parent layer key: a.layers.7.self_attn -> a.layers.7."""
2042
+ import re
2043
+ m = re.search(
2044
+ r"((?:.*\.)?(?:layers|transformer_blocks)\.\d+)\.", path)
2045
+ return m.group(1) if m else path.rsplit(".", 1)[0]
2046
+
2047
+
2048
+ def _feeds_attention(path: str) -> bool:
2049
+ """An adaln producer that feeds attention (input_layernorm) rather
2050
+ than the MLP (post_attention_layernorm)."""
2051
+ leaf = path.rsplit(".", 1)[-1]
2052
+ return "input" in leaf or leaf in ("norm1", "ln1")