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