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.
- flashrt_structures/__init__.py +174 -0
- flashrt_structures/adapters/__init__.py +55 -0
- flashrt_structures/adapters/diffusers_attention.py +237 -0
- flashrt_structures/adapters/diffusers_rotary_attention.py +239 -0
- flashrt_structures/adapters/factored_qk_norm_rope.py +252 -0
- flashrt_structures/adapters/factored_two_way_attention.py +99 -0
- flashrt_structures/adapters/gemma_attention.py +226 -0
- flashrt_structures/adapters/packed_qkv_rope.py +342 -0
- flashrt_structures/adapters/packed_stream_qk_norm_rope.py +376 -0
- flashrt_structures/adapters/qwen_per_head_qk_norm_rope.py +231 -0
- flashrt_structures/adapters/sglang_engine.py +207 -0
- flashrt_structures/adapters/transformers_attention_interface.py +73 -0
- flashrt_structures/adapters/transformers_gated_delta.py +152 -0
- flashrt_structures/adapters/transformers_gated_delta_fused.py +96 -0
- flashrt_structures/adapters/vllm_engine.py +424 -0
- flashrt_structures/adjudicate.py +85 -0
- flashrt_structures/aot.py +191 -0
- flashrt_structures/autobuild.py +2052 -0
- flashrt_structures/beta/__init__.py +43 -0
- flashrt_structures/beta/conform.py +94 -0
- flashrt_structures/beta/joins.py +113 -0
- flashrt_structures/beta/negotiate.py +84 -0
- flashrt_structures/beta/ports.py +140 -0
- flashrt_structures/decisions.py +80 -0
- flashrt_structures/discover.py +623 -0
- flashrt_structures/explain.py +70 -0
- flashrt_structures/frontdoor.py +572 -0
- flashrt_structures/gates.py +465 -0
- flashrt_structures/guard.py +421 -0
- flashrt_structures/handle.py +189 -0
- flashrt_structures/impls/__init__.py +219 -0
- flashrt_structures/impls/adaln_producer/__init__.py +8 -0
- flashrt_structures/impls/adaln_producer/broker.py +116 -0
- flashrt_structures/impls/adaln_producer/fused.py +388 -0
- flashrt_structures/impls/adarms_stack/__init__.py +8 -0
- flashrt_structures/impls/adarms_stack/fp8_chain.py +832 -0
- flashrt_structures/impls/adarms_stack/region.py +102 -0
- flashrt_structures/impls/attention_core/__init__.py +132 -0
- flashrt_structures/impls/attention_core/fa2_seqused.py +458 -0
- flashrt_structures/impls/attention_core/fa4_cute.py +154 -0
- flashrt_structures/impls/attention_core/fa4_fp8.py +178 -0
- flashrt_structures/impls/attention_core/masked_mha.py +158 -0
- flashrt_structures/impls/attention_core/two_way_fa2.py +220 -0
- flashrt_structures/impls/cadence_static/__init__.py +18 -0
- flashrt_structures/impls/cadence_static/buffers.py +122 -0
- flashrt_structures/impls/cadence_static/cross_attention.py +187 -0
- flashrt_structures/impls/chain_elements.py +89 -0
- flashrt_structures/impls/decode_loop/__init__.py +0 -0
- flashrt_structures/impls/decode_loop/fp8_kv.py +206 -0
- flashrt_structures/impls/decode_loop/mtp_speculative.py +245 -0
- flashrt_structures/impls/decode_loop/whole_step.py +852 -0
- flashrt_structures/impls/decoder_block/__init__.py +6 -0
- flashrt_structures/impls/decoder_block/attn_sublayer.py +110 -0
- flashrt_structures/impls/decoder_block/fused.py +167 -0
- flashrt_structures/impls/decoder_ffn/__init__.py +0 -0
- flashrt_structures/impls/decoder_ffn/fp8_static.py +310 -0
- flashrt_structures/impls/decoder_ffn/fp8_static.yaml +22 -0
- flashrt_structures/impls/decoder_ffn/w4a16_static.py +221 -0
- flashrt_structures/impls/decoder_ffn/w8a16_static.py +183 -0
- flashrt_structures/impls/dit_stack/__init__.py +11 -0
- flashrt_structures/impls/dit_stack/fp4_chain.py +417 -0
- flashrt_structures/impls/dit_stack/region.py +86 -0
- flashrt_structures/impls/fixed_iter/__init__.py +29 -0
- flashrt_structures/impls/fixed_iter/openpi.py +264 -0
- flashrt_structures/impls/fixed_iter/protocol.py +94 -0
- flashrt_structures/impls/gated_delta_core/__init__.py +3 -0
- flashrt_structures/impls/gated_delta_core/fused_layer.py +545 -0
- flashrt_structures/impls/gated_delta_core/hub_v3.py +152 -0
- flashrt_structures/impls/graph_lowering/__init__.py +27 -0
- flashrt_structures/impls/graph_lowering/pi052_denoise.py +179 -0
- flashrt_structures/impls/graph_lowering/protocol.py +76 -0
- flashrt_structures/impls/graph_lowering/qwen3_vl.py +364 -0
- flashrt_structures/impls/linear_proj/__init__.py +0 -0
- flashrt_structures/impls/linear_proj/fp8_static.py +270 -0
- flashrt_structures/impls/linear_proj/nvfp4_balance.py +131 -0
- flashrt_structures/impls/linear_proj/nvfp4_dynamic.py +182 -0
- flashrt_structures/impls/linear_proj/w8a16_static.py +230 -0
- flashrt_structures/impls/modnorm_qkv_chain/__init__.py +0 -0
- flashrt_structures/impls/modnorm_qkv_chain/fp8_ptok_table.py +291 -0
- flashrt_structures/impls/moe_experts/__init__.py +9 -0
- flashrt_structures/impls/moe_experts/nvfp4_dynamic.py +208 -0
- flashrt_structures/impls/moe_experts/nvfp4_w4a16.py +129 -0
- flashrt_structures/impls/norm_fused/__init__.py +3 -0
- flashrt_structures/impls/norm_fused/bf16.py +94 -0
- flashrt_structures/impls/norm_fused/fp8_producer.py +84 -0
- flashrt_structures/impls/patch_projection/__init__.py +3 -0
- flashrt_structures/impls/patch_projection/bf16_flat.py +139 -0
- flashrt_structures/impls/prefill_tower/__init__.py +10 -0
- flashrt_structures/impls/prefill_tower/fp8_chain.py +955 -0
- flashrt_structures/impls/prefill_tower/region.py +99 -0
- flashrt_structures/impls/qk_norm_rope/__init__.py +12 -0
- flashrt_structures/impls/qk_norm_rope/per_head_gqa.py +199 -0
- flashrt_structures/impls/qk_norm_rope/projection_bf16.py +165 -0
- flashrt_structures/impls/qkv_pack/__init__.py +5 -0
- flashrt_structures/impls/qkv_pack/bf16.py +110 -0
- flashrt_structures/impls/qkv_pack/fp8_static.py +435 -0
- flashrt_structures/impls/qkv_pack/nvfp4_balance.py +218 -0
- flashrt_structures/impls/qkv_rope/__init__.py +3 -0
- flashrt_structures/impls/qkv_rope/packed_bias_bf16.py +143 -0
- flashrt_structures/impls/step_table.py +113 -0
- flashrt_structures/impls/vision_ffn/__init__.py +0 -0
- flashrt_structures/impls/vision_ffn/fp8_static.py +261 -0
- flashrt_structures/impls/vision_ffn/nvfp4_balance.py +211 -0
- flashrt_structures/impls/vision_tower/__init__.py +7 -0
- flashrt_structures/impls/vision_tower/fp8_chain.py +533 -0
- flashrt_structures/impls/vision_tower/region.py +91 -0
- flashrt_structures/matrix.py +126 -0
- flashrt_structures/points.py +368 -0
- flashrt_structures/prequantized.py +131 -0
- flashrt_structures/quantize_on_adopt.py +94 -0
- flashrt_structures/recipe.py +438 -0
- flashrt_structures/regions.py +208 -0
- flashrt_structures/schemes.py +490 -0
- flashrt_structures/stages.py +298 -0
- flashrt_structures/storage.py +255 -0
- flashrt_structures/swap.py +398 -0
- flashrt_structures/workspace.py +110 -0
- flashrt_structures-0.2.0.dist-info/METADATA +140 -0
- flashrt_structures-0.2.0.dist-info/RECORD +122 -0
- flashrt_structures-0.2.0.dist-info/WHEEL +5 -0
- flashrt_structures-0.2.0.dist-info/licenses/LICENSE +202 -0
- 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")
|