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,398 @@
|
|
|
1
|
+
"""Transactional module attachment — mechanism only, no policy.
|
|
2
|
+
|
|
3
|
+
Swaps host modules for bound structure implementations atomically: every
|
|
4
|
+
staged swap is applied only if all of them can be applied, and a handle
|
|
5
|
+
restores the originals exactly. Which modules to swap, and whether an
|
|
6
|
+
implementation passed its gates, is decided by the caller before
|
|
7
|
+
staging; this layer refuses partial application by construction.
|
|
8
|
+
|
|
9
|
+
The handle is also where the attachment answers for itself at runtime.
|
|
10
|
+
Every swapped-in structure carries a guard recording the form it was
|
|
11
|
+
calibrated for (:mod:`flashrt_structures.guard`); attaching gives each
|
|
12
|
+
guard its path and a way to restore the host module, and
|
|
13
|
+
``handle.report()`` reads them back. An attachment whose report shows
|
|
14
|
+
fallbacks is one that is not running the structures it claims to — that
|
|
15
|
+
has to be visible here, because nothing downstream can tell the
|
|
16
|
+
difference.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
from dataclasses import dataclass, field
|
|
22
|
+
from typing import Any, Mapping
|
|
23
|
+
|
|
24
|
+
import torch
|
|
25
|
+
|
|
26
|
+
from .guard import GUARD_ATTR, SeamGuard, collect as _collect_guards
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def resolve_parent(root: torch.nn.Module, path: str) -> tuple[torch.nn.Module, str]:
|
|
30
|
+
"""Resolve the parent module and attribute name for a dotted path."""
|
|
31
|
+
parts = path.split(".")
|
|
32
|
+
parent = root
|
|
33
|
+
for part in parts[:-1]:
|
|
34
|
+
parent = parent[int(part)] if part.isdigit() else getattr(parent, part)
|
|
35
|
+
attr = parts[-1]
|
|
36
|
+
leaf = parent[int(attr)] if attr.isdigit() else getattr(parent, attr)
|
|
37
|
+
if not isinstance(leaf, torch.nn.Module):
|
|
38
|
+
raise TypeError(f"{path!r} does not resolve to a torch.nn.Module")
|
|
39
|
+
return parent, attr
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def _get(parent: torch.nn.Module, attr: str) -> torch.nn.Module:
|
|
43
|
+
return parent[int(attr)] if attr.isdigit() else getattr(parent, attr)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _set(parent: torch.nn.Module, attr: str,
|
|
47
|
+
module: torch.nn.Module) -> None:
|
|
48
|
+
if attr.isdigit():
|
|
49
|
+
parent[int(attr)] = module
|
|
50
|
+
else:
|
|
51
|
+
setattr(parent, attr, module)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclass
|
|
55
|
+
class AttachHandle:
|
|
56
|
+
"""Restores the originals of one committed attachment."""
|
|
57
|
+
|
|
58
|
+
_entries: list[tuple[torch.nn.Module, str, torch.nn.Module]]
|
|
59
|
+
records: Mapping[str, Any] = field(default_factory=dict)
|
|
60
|
+
active: bool = True
|
|
61
|
+
_guards: dict[str, SeamGuard] = field(default_factory=dict)
|
|
62
|
+
_revert: list[Any] = field(default_factory=list)
|
|
63
|
+
_store: Any = None
|
|
64
|
+
_paths: list[str] = field(default_factory=list)
|
|
65
|
+
|
|
66
|
+
def consume(self, store=None) -> dict:
|
|
67
|
+
"""Move every replaced module's weights off the device.
|
|
68
|
+
|
|
69
|
+
The resident tier is gone: an attachment no longer keeps a
|
|
70
|
+
second copy of the model in device memory for the sake of an
|
|
71
|
+
instant fallback. Each original's truth moves to the weight
|
|
72
|
+
store — the checkpoint file when provenance verifies, pinned
|
|
73
|
+
host memory otherwise — and its device storage is freed.
|
|
74
|
+
Fallback and ``detach`` both survive as restore-from-store:
|
|
75
|
+
slower, never wrong. Reversible until :meth:`finalize`.
|
|
76
|
+
"""
|
|
77
|
+
from .storage import WeightStore
|
|
78
|
+
|
|
79
|
+
if store is None:
|
|
80
|
+
store = self._store or WeightStore(
|
|
81
|
+
checkpoint=self.records.get("checkpoint"))
|
|
82
|
+
self._store = store
|
|
83
|
+
freed = 0
|
|
84
|
+
serving = 0
|
|
85
|
+
paths = self._paths or [""] * len(self._entries)
|
|
86
|
+
for path, (parent, attr, original) in zip(paths, self._entries):
|
|
87
|
+
# a seat that actively calls its retained host (a cadence
|
|
88
|
+
# bank refreshing through the host projection) declares it:
|
|
89
|
+
# that original is serving, not merely held for fallback,
|
|
90
|
+
# and consuming it would corrupt the live path
|
|
91
|
+
current = _get(parent, attr)
|
|
92
|
+
if getattr(current, "_frt_host_serving", False):
|
|
93
|
+
serving += 1
|
|
94
|
+
continue
|
|
95
|
+
freed += store.stash_module(path, original)
|
|
96
|
+
# the caching allocator keeps the bind-era blocks reserved;
|
|
97
|
+
# hand them back so co-tenant CUDA consumers see the space —
|
|
98
|
+
# only unused cached blocks are released, capture pools are
|
|
99
|
+
# private and untouched
|
|
100
|
+
returned = 0
|
|
101
|
+
if torch.cuda.is_available():
|
|
102
|
+
before = torch.cuda.memory_reserved()
|
|
103
|
+
torch.cuda.empty_cache()
|
|
104
|
+
returned = before - torch.cuda.memory_reserved()
|
|
105
|
+
self.records = dict(self.records, consumed=dict(
|
|
106
|
+
store.stats, freed_bytes=freed, kept_serving=serving,
|
|
107
|
+
cache_returned_bytes=returned))
|
|
108
|
+
return {"consumed": True, "freed_bytes": freed,
|
|
109
|
+
"kept_serving": serving,
|
|
110
|
+
"cache_returned_bytes": returned,
|
|
111
|
+
"tiers": {"disk": store.stats["disk"],
|
|
112
|
+
"ram": store.stats["ram"]}}
|
|
113
|
+
|
|
114
|
+
def finalize(self) -> dict:
|
|
115
|
+
"""Make the consumption permanent.
|
|
116
|
+
|
|
117
|
+
Drops the restore tickets (including any host-RAM spill), flips
|
|
118
|
+
fallback off — a contract miss must refuse now, not restore —
|
|
119
|
+
and forbids ``detach``. Consumes first when the caller has not.
|
|
120
|
+
Irreversible, and says so in the receipt it returns.
|
|
121
|
+
"""
|
|
122
|
+
if self._store is None:
|
|
123
|
+
self.consume()
|
|
124
|
+
freed = self.records.get("consumed", {}).get("freed_bytes", 0)
|
|
125
|
+
paths = set(self._paths)
|
|
126
|
+
for parent, attr, original in self._entries:
|
|
127
|
+
self._store.drop_module(original)
|
|
128
|
+
current = _get(parent, attr)
|
|
129
|
+
if hasattr(current, "_frt_can_fallback"):
|
|
130
|
+
current._frt_can_fallback = False
|
|
131
|
+
# the guard is the one that answers at call time: flip its own
|
|
132
|
+
# fallback bit too, or a contract miss would run an emptied host
|
|
133
|
+
for site, guard in self._guards.items():
|
|
134
|
+
root_site = site.split("::", 1)[0]
|
|
135
|
+
if root_site in paths:
|
|
136
|
+
guard.can_fallback = False
|
|
137
|
+
self.records = dict(self.records,
|
|
138
|
+
finalized={"freed_bytes": freed})
|
|
139
|
+
self._finalized = True
|
|
140
|
+
return {"finalized": True, "freed_bytes": freed}
|
|
141
|
+
|
|
142
|
+
def detach(self) -> None:
|
|
143
|
+
"""Restore every swapped module. Idempotent.
|
|
144
|
+
|
|
145
|
+
Consumed weights come back from the store first — the promise
|
|
146
|
+
is the same bit-exact host, backed by the checkpoint file or
|
|
147
|
+
the host-RAM spill instead of a resident device copy.
|
|
148
|
+
|
|
149
|
+
Also runs any ``revert`` callables the caller handed over. Some
|
|
150
|
+
seams are not modules at paths — an adapter that patches a
|
|
151
|
+
library-level function — and restoring only what ``setattr`` can
|
|
152
|
+
reach would leave those live while reporting the model as restored.
|
|
153
|
+
"""
|
|
154
|
+
if getattr(self, "_finalized", False):
|
|
155
|
+
raise RuntimeError(
|
|
156
|
+
"attachment was finalized: the host originals were "
|
|
157
|
+
"released and detach is impossible")
|
|
158
|
+
if not self.active:
|
|
159
|
+
return
|
|
160
|
+
for parent, attr, original in reversed(self._entries):
|
|
161
|
+
if self._store is not None:
|
|
162
|
+
self._store.restore_module(original)
|
|
163
|
+
_set(parent, attr, original)
|
|
164
|
+
for undo in reversed(self._revert):
|
|
165
|
+
undo()
|
|
166
|
+
self._revert.clear()
|
|
167
|
+
for guard in self._guards.values():
|
|
168
|
+
guard.release_site()
|
|
169
|
+
self.active = False
|
|
170
|
+
|
|
171
|
+
# ---- what actually ran -----------------------------------------
|
|
172
|
+
|
|
173
|
+
def report(self) -> dict[str, dict[str, Any]]:
|
|
174
|
+
"""Per-seam ledger: calls, fallbacks, and why.
|
|
175
|
+
|
|
176
|
+
The unit of truth for "did this attachment run what it says". A
|
|
177
|
+
seam with ``fallbacks`` above zero ran the host module for that
|
|
178
|
+
many calls; one with ``detached`` set gave up and put the host
|
|
179
|
+
module back. Tests assert on this rather than assuming it.
|
|
180
|
+
|
|
181
|
+
Counts are eager-only, and necessarily so: inside a compiled or
|
|
182
|
+
captured region the kernel runs without re-entering Python, so
|
|
183
|
+
there is nothing left to count. A captured host therefore reports
|
|
184
|
+
zero calls for the seams inside its graph — read that as "this
|
|
185
|
+
ledger does not cover the graph", not as "the seam did not run".
|
|
186
|
+
The graph's own parity check is what covers it.
|
|
187
|
+
"""
|
|
188
|
+
return {site: guard.entry()
|
|
189
|
+
for site, guard in sorted(self._guards.items())}
|
|
190
|
+
|
|
191
|
+
def summary(self) -> dict[str, Any]:
|
|
192
|
+
"""The report reduced to what a caller has to act on.
|
|
193
|
+
|
|
194
|
+
``seams_never_called`` means "no eager call reached this seam",
|
|
195
|
+
which for a captured host is every seam in the graph. It is a
|
|
196
|
+
finding on an eager host and noise on a captured one.
|
|
197
|
+
"""
|
|
198
|
+
entries = self.report()
|
|
199
|
+
fell_back = sorted(site for site, e in entries.items()
|
|
200
|
+
if e["fallbacks"])
|
|
201
|
+
return {
|
|
202
|
+
"seams": len(entries),
|
|
203
|
+
"guarded_calls": sum(e["calls"] for e in entries.values()),
|
|
204
|
+
"fallbacks": sum(e["fallbacks"] for e in entries.values()),
|
|
205
|
+
"seams_fell_back": fell_back,
|
|
206
|
+
"seams_self_detached": sorted(
|
|
207
|
+
site for site, e in entries.items() if e["detached"]),
|
|
208
|
+
"seams_never_called": sorted(
|
|
209
|
+
site for site, e in entries.items() if not e["calls"]),
|
|
210
|
+
"clean": not fell_back,
|
|
211
|
+
}
|
|
212
|
+
|
|
213
|
+
def manifest(self) -> dict[str, Any]:
|
|
214
|
+
"""One document answering "why does this box run this form".
|
|
215
|
+
|
|
216
|
+
The receipts exist — band decisions, variant trails, format
|
|
217
|
+
races, the workspace ledger, the weight-residency receipt — but
|
|
218
|
+
scattered, each with its own schema. Cross-box drift diagnosis
|
|
219
|
+
and decision transport both need the single view: device
|
|
220
|
+
fingerprint, every seam with its kind and its story, the
|
|
221
|
+
decisions consumed, the memory plan, and what happened to the
|
|
222
|
+
weights. Read-only; safe to serialize.
|
|
223
|
+
"""
|
|
224
|
+
import json as _json
|
|
225
|
+
|
|
226
|
+
device = (torch.cuda.get_device_name(0)
|
|
227
|
+
if torch.cuda.is_available() else "cpu")
|
|
228
|
+
decisions: dict[str, Any] = {}
|
|
229
|
+
try:
|
|
230
|
+
from .decisions import _cache_path
|
|
231
|
+
decisions = _json.loads(_cache_path().read_text())
|
|
232
|
+
except Exception:
|
|
233
|
+
pass
|
|
234
|
+
workspace_report: Any = None
|
|
235
|
+
try:
|
|
236
|
+
from .workspace import report as _ws_report
|
|
237
|
+
workspace_report = _ws_report()
|
|
238
|
+
except Exception:
|
|
239
|
+
pass
|
|
240
|
+
seams = {}
|
|
241
|
+
for site, guard in sorted(self._guards.items()):
|
|
242
|
+
entry = guard.entry()
|
|
243
|
+
seams[site] = {
|
|
244
|
+
"kind": entry.get("kind"),
|
|
245
|
+
"calls": entry.get("calls"),
|
|
246
|
+
"fallbacks": entry.get("fallbacks"),
|
|
247
|
+
"notes": entry.get("notes"),
|
|
248
|
+
}
|
|
249
|
+
return {
|
|
250
|
+
"device": device,
|
|
251
|
+
"seams": seams,
|
|
252
|
+
"summary": self.summary(),
|
|
253
|
+
"records": dict(self.records),
|
|
254
|
+
"decisions": {k: v for k, v in decisions.items()
|
|
255
|
+
if k.startswith(device + "|")},
|
|
256
|
+
"workspace": workspace_report,
|
|
257
|
+
}
|
|
258
|
+
|
|
259
|
+
def raise_on_fallback(self) -> None:
|
|
260
|
+
"""Assertion helper: fail loudly if any seam fell back.
|
|
261
|
+
|
|
262
|
+
For tests, and for callers who would rather not ship a run whose
|
|
263
|
+
seams quietly reverted. Reads the ledger, so it costs nothing in
|
|
264
|
+
the hot path.
|
|
265
|
+
"""
|
|
266
|
+
summary = self.summary()
|
|
267
|
+
if summary["clean"]:
|
|
268
|
+
return
|
|
269
|
+
detail = {site: entry["last_reason"]
|
|
270
|
+
for site, entry in self.report().items()
|
|
271
|
+
if entry["fallbacks"]}
|
|
272
|
+
raise RuntimeError(
|
|
273
|
+
f"{len(detail)} seam(s) fell back to the host module during "
|
|
274
|
+
f"this attachment ({summary['fallbacks']} call(s) total): "
|
|
275
|
+
f"{detail}")
|
|
276
|
+
|
|
277
|
+
|
|
278
|
+
def attach(
|
|
279
|
+
root: torch.nn.Module,
|
|
280
|
+
swaps: Mapping[str, torch.nn.Module],
|
|
281
|
+
*,
|
|
282
|
+
records: Mapping[str, Any] | None = None,
|
|
283
|
+
on_guard_fail: str = "fallback",
|
|
284
|
+
allow_training: bool = False,
|
|
285
|
+
observe: Mapping[str, torch.nn.Module] | None = None,
|
|
286
|
+
revert: Any = None,
|
|
287
|
+
store: Any = None,
|
|
288
|
+
consume: bool = False,
|
|
289
|
+
) -> AttachHandle:
|
|
290
|
+
"""Atomically replace the modules at ``swaps`` paths.
|
|
291
|
+
|
|
292
|
+
All paths are resolved and validated before the first swap is
|
|
293
|
+
applied; any resolution failure leaves the model untouched. Callers
|
|
294
|
+
pass the qualification ``records`` that justified the attachment so
|
|
295
|
+
the handle carries its own evidence.
|
|
296
|
+
|
|
297
|
+
``on_guard_fail`` selects what a seam does when it is called outside
|
|
298
|
+
the form it was bound for: ``"fallback"`` runs the host module and
|
|
299
|
+
counts it in the ledger, ``"raise"`` refuses immediately. Development
|
|
300
|
+
and CI should use ``"raise"`` — a fallback a test tolerates is a
|
|
301
|
+
fallback nobody reads.
|
|
302
|
+
|
|
303
|
+
Attaching to a model in training mode is refused: these
|
|
304
|
+
implementations pack and quantise the host weights once at bind time,
|
|
305
|
+
so a gradient step would update the host copy and leave the kernel's
|
|
306
|
+
stale. Pass ``allow_training=True`` for a forward-only pass that
|
|
307
|
+
happens to sit inside a training script.
|
|
308
|
+
|
|
309
|
+
``observe`` names modules that carry a guard but are not swapped here
|
|
310
|
+
— an adapter's routed seam. They are reported in the ledger and never
|
|
311
|
+
installed. ``revert`` collects undo callables for host mutations made
|
|
312
|
+
before this call, so ``detach`` restores those too.
|
|
313
|
+
|
|
314
|
+
Weight residency is a lifecycle, not a mode: attach → validate →
|
|
315
|
+
``handle.consume()`` → optionally ``handle.finalize()``. Consuming
|
|
316
|
+
moves each replaced original's truth to the weight ``store`` (the
|
|
317
|
+
checkpoint file when provenance verifies, pinned host memory
|
|
318
|
+
otherwise) and frees its device storage — there is no resident
|
|
319
|
+
tier to keep. Fallback and ``detach`` restore from the store;
|
|
320
|
+
seats that declare their retained host as actively serving keep it
|
|
321
|
+
whole. Consumption comes after the caller's validation pass
|
|
322
|
+
because the attached model still owes the host schema (state_dict,
|
|
323
|
+
A/B reference arms, captures that alias host weights) until then;
|
|
324
|
+
``consume=True`` collapses the steps for callers with no such
|
|
325
|
+
pass.
|
|
326
|
+
"""
|
|
327
|
+
if not swaps and not observe:
|
|
328
|
+
raise ValueError("no swaps or routed seams staged")
|
|
329
|
+
if on_guard_fail not in ("fallback", "raise"):
|
|
330
|
+
raise ValueError(
|
|
331
|
+
f"on_guard_fail must be 'fallback' or 'raise', "
|
|
332
|
+
f"got {on_guard_fail!r}")
|
|
333
|
+
if root.training and not allow_training:
|
|
334
|
+
raise ValueError(
|
|
335
|
+
"structures.attach: the model is in training mode. These "
|
|
336
|
+
"implementations pack and quantise the host weights once at "
|
|
337
|
+
"bind time, so training through them would update the host "
|
|
338
|
+
"copy and leave the kernel's stale. Call model.eval() first, "
|
|
339
|
+
"or pass allow_training=True for a forward-only pass.")
|
|
340
|
+
|
|
341
|
+
staged: list[tuple[torch.nn.Module, str, torch.nn.Module,
|
|
342
|
+
torch.nn.Module]] = []
|
|
343
|
+
for path, replacement in swaps.items():
|
|
344
|
+
parent, attr = resolve_parent(root, path)
|
|
345
|
+
staged.append((parent, attr, _get(parent, attr), replacement))
|
|
346
|
+
|
|
347
|
+
entries: list[tuple[torch.nn.Module, str, torch.nn.Module]] = []
|
|
348
|
+
try:
|
|
349
|
+
for parent, attr, original, replacement in staged:
|
|
350
|
+
_set(parent, attr, replacement)
|
|
351
|
+
entries.append((parent, attr, original))
|
|
352
|
+
except Exception:
|
|
353
|
+
for parent, attr, original in reversed(entries):
|
|
354
|
+
_set(parent, attr, original)
|
|
355
|
+
raise
|
|
356
|
+
|
|
357
|
+
# ---- give every guard its site and its own way out ----
|
|
358
|
+
# keyed by guard identity as well as by site: one module can be
|
|
359
|
+
# reachable by two routes (a composed block holds the host block, so a
|
|
360
|
+
# core hanging off the host's attention is found twice), and counting
|
|
361
|
+
# it twice would overstate how many seams an attachment has
|
|
362
|
+
guards: dict[str, SeamGuard] = {}
|
|
363
|
+
seen: set[int] = set()
|
|
364
|
+
# named seams first: an adapter knows what its routed seam should be
|
|
365
|
+
# called, and the same object found by walking a composed structure
|
|
366
|
+
# would otherwise claim it under an incidental path
|
|
367
|
+
for site, module in (observe or {}).items():
|
|
368
|
+
for child, guard in _collect_guards(module):
|
|
369
|
+
if id(guard) in seen:
|
|
370
|
+
continue
|
|
371
|
+
seen.add(id(guard))
|
|
372
|
+
key = site if not child else f"{site}::{child}"
|
|
373
|
+
guard.bind_site(key, restore=None, mode=on_guard_fail)
|
|
374
|
+
guards[key] = guard
|
|
375
|
+
for (parent, attr, original), (path, replacement) in zip(
|
|
376
|
+
entries, swaps.items()):
|
|
377
|
+
for child, guard in _collect_guards(replacement):
|
|
378
|
+
if id(guard) in seen:
|
|
379
|
+
continue
|
|
380
|
+
seen.add(id(guard))
|
|
381
|
+
site = path if not child else f"{path}::{child}"
|
|
382
|
+
# only the module actually swapped in at a path can restore
|
|
383
|
+
# itself; a guard held inside a composed structure reports but
|
|
384
|
+
# cannot exit on its own
|
|
385
|
+
restore = ((lambda p=parent, a=attr, o=original: _set(p, a, o))
|
|
386
|
+
if not child else None)
|
|
387
|
+
guard.bind_site(site, restore=restore, mode=on_guard_fail)
|
|
388
|
+
guards[site] = guard
|
|
389
|
+
|
|
390
|
+
handle = AttachHandle(_entries=entries, records=dict(records or {}),
|
|
391
|
+
_guards=guards, _revert=list(revert or ()),
|
|
392
|
+
_store=store, _paths=list(swaps.keys()))
|
|
393
|
+
if consume:
|
|
394
|
+
handle.consume(store)
|
|
395
|
+
return handle
|
|
396
|
+
|
|
397
|
+
|
|
398
|
+
__all__ = ["AttachHandle", "attach", "resolve_parent", "GUARD_ATTR"]
|
|
@@ -0,0 +1,110 @@
|
|
|
1
|
+
"""Shared workspaces: seat scratch lives in a pool, not in every seat.
|
|
2
|
+
|
|
3
|
+
A pack's sibling stash, a producer's residual scratch, a wire's packed
|
|
4
|
+
buffer — their lifetimes all sit inside one layer's forward. Sequential
|
|
5
|
+
layers therefore never need their own copies: every seat that asks for
|
|
6
|
+
the same (shape, dtype, device, tag) receives the *same* tensor, and
|
|
7
|
+
the pool's footprint is one layer's worth instead of layers x tokens.
|
|
8
|
+
Capture-compatible by the same argument that makes memory pools
|
|
9
|
+
capture-compatible: the graph records fixed pointers, and same-stream
|
|
10
|
+
sequential lifetimes never overlap.
|
|
11
|
+
|
|
12
|
+
Two lease kinds:
|
|
13
|
+
- ``scratch``: contents are call-transient; the seat must write before
|
|
14
|
+
it reads (packs and producers already do).
|
|
15
|
+
- ``ones``: constant-filled; shared freely and never written.
|
|
16
|
+
|
|
17
|
+
The pool is also the accounting surface: :func:`report` returns bytes
|
|
18
|
+
held and the reuse count per tag — the memory column of the receipt.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from __future__ import annotations
|
|
22
|
+
|
|
23
|
+
import weakref
|
|
24
|
+
|
|
25
|
+
import torch
|
|
26
|
+
|
|
27
|
+
_POOL: dict[tuple, torch.Tensor] = {}
|
|
28
|
+
_LEASES: dict[str, int] = {}
|
|
29
|
+
# exclusive leases are owned by their seats, not the pool; the receipt
|
|
30
|
+
# tracks them through weak references so bytes leave the ledger the
|
|
31
|
+
# moment a detached seat releases its buffer
|
|
32
|
+
_EXCLUSIVE: dict[str, list] = {}
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def lease(shape, dtype, device, *, tag: str,
|
|
36
|
+
fill: str = "scratch",
|
|
37
|
+
exclusive: bool = False) -> torch.Tensor:
|
|
38
|
+
"""One shared tensor for every seat asking this (shape, tag).
|
|
39
|
+
|
|
40
|
+
The pool's safety argument is single-stream sequential execution:
|
|
41
|
+
layer i's scratch is dead before layer i+1 writes it *because they
|
|
42
|
+
share one stream*. A caller leasing from a non-default stream is
|
|
43
|
+
outside that argument, and the pool refuses rather than corrupt
|
|
44
|
+
silently — multi-stream serving needs per-lane isolation first.
|
|
45
|
+
|
|
46
|
+
``exclusive=True`` allocates a private tensor instead of joining
|
|
47
|
+
the shared slab, still counted in the pool's receipt. It exists
|
|
48
|
+
for *state*, not scratch: a buffer whose consumer is the host may
|
|
49
|
+
be retained past the tick (a KV cache holding a reader's view was
|
|
50
|
+
the measured failure — every layer's write clobbered every other
|
|
51
|
+
layer's cached slice, with every ledger clean). Sharing such a
|
|
52
|
+
buffer requires the immediacy of consumption as a verified fact;
|
|
53
|
+
absent that fact, state is exclusive.
|
|
54
|
+
"""
|
|
55
|
+
if (torch.cuda.is_available()
|
|
56
|
+
and str(device).startswith("cuda")
|
|
57
|
+
and torch.cuda.current_stream()
|
|
58
|
+
!= torch.cuda.default_stream()):
|
|
59
|
+
raise RuntimeError(
|
|
60
|
+
"workspace: leasing from a non-default CUDA stream — the "
|
|
61
|
+
"shared pool's lifetime argument only covers single-stream "
|
|
62
|
+
"sequential execution; isolate per-stream lanes first")
|
|
63
|
+
if exclusive:
|
|
64
|
+
buf = torch.zeros(*shape, dtype=dtype, device=device)
|
|
65
|
+
if fill == "ones":
|
|
66
|
+
buf.fill_(1)
|
|
67
|
+
_LEASES[tag] = _LEASES.get(tag, 0) + 1
|
|
68
|
+
_EXCLUSIVE.setdefault(tag, []).append(weakref.ref(buf))
|
|
69
|
+
return buf
|
|
70
|
+
key = (tuple(shape), dtype, str(device), tag, fill)
|
|
71
|
+
buf = _POOL.get(key)
|
|
72
|
+
if buf is None:
|
|
73
|
+
buf = torch.zeros(*shape, dtype=dtype, device=device)
|
|
74
|
+
if fill == "ones":
|
|
75
|
+
buf.fill_(1)
|
|
76
|
+
_POOL[key] = buf
|
|
77
|
+
_LEASES[tag] = _LEASES.get(tag, 0) + 1
|
|
78
|
+
return buf
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def report() -> dict:
|
|
82
|
+
"""Bytes held and lease counts — the receipt's memory column.
|
|
83
|
+
|
|
84
|
+
``held_bytes`` covers both the shared pool and every live exclusive
|
|
85
|
+
lease; ``exclusive_by_tag`` breaks the latter out so the receipt
|
|
86
|
+
shows what stayed private and what the pool actually deduplicated.
|
|
87
|
+
"""
|
|
88
|
+
by_tag: dict[str, int] = {}
|
|
89
|
+
for (shape, dtype, _dev, tag, _fill), buf in _POOL.items():
|
|
90
|
+
by_tag[tag] = by_tag.get(tag, 0) + buf.numel() * buf.element_size()
|
|
91
|
+
exclusive_by_tag: dict[str, int] = {}
|
|
92
|
+
for tag, refs in _EXCLUSIVE.items():
|
|
93
|
+
live = [ref() for ref in refs]
|
|
94
|
+
refs[:] = [ref for ref, buf in zip(refs, live) if buf is not None]
|
|
95
|
+
held = sum(buf.numel() * buf.element_size()
|
|
96
|
+
for buf in live if buf is not None)
|
|
97
|
+
if held:
|
|
98
|
+
exclusive_by_tag[tag] = held
|
|
99
|
+
return {"held_bytes": sum(by_tag.values())
|
|
100
|
+
+ sum(exclusive_by_tag.values()),
|
|
101
|
+
"by_tag": by_tag,
|
|
102
|
+
"exclusive_by_tag": exclusive_by_tag,
|
|
103
|
+
"leases": dict(_LEASES)}
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def clear() -> None:
|
|
107
|
+
"""Drop every pooled buffer (between hosts, or in tests)."""
|
|
108
|
+
_POOL.clear()
|
|
109
|
+
_LEASES.clear()
|
|
110
|
+
_EXCLUSIVE.clear()
|
|
@@ -0,0 +1,140 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flashrt-structures
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Attach FlashRT's verified acceleration structures to an unmodified PyTorch host (lerobot, Isaac-GR00T, openpi, transformers, diffusers, vLLM, SGLang) with no fork and no edit to the host.
|
|
5
|
+
License: Apache-2.0
|
|
6
|
+
Project-URL: Homepage, https://github.com/flashrt-project/FlashRT-Structures
|
|
7
|
+
Project-URL: Repository, https://github.com/flashrt-project/FlashRT-Structures
|
|
8
|
+
Project-URL: Issues, https://github.com/flashrt-project/FlashRT-Structures/issues
|
|
9
|
+
Project-URL: Engine, https://github.com/flashrt-project/FlashRT
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Requires-Python: >=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: flash-rt>=0.1.0
|
|
24
|
+
Requires-Dist: numpy
|
|
25
|
+
Requires-Dist: pyyaml
|
|
26
|
+
Provides-Extra: hub
|
|
27
|
+
Requires-Dist: kernels>=0.12; extra == "hub"
|
|
28
|
+
Provides-Extra: torch
|
|
29
|
+
Requires-Dist: torch; extra == "torch"
|
|
30
|
+
Requires-Dist: safetensors; extra == "torch"
|
|
31
|
+
Dynamic: license-file
|
|
32
|
+
|
|
33
|
+
# FlashRT Structures
|
|
34
|
+
|
|
35
|
+
Attach [FlashRT](https://github.com/flashrt-project/FlashRT)'s verified
|
|
36
|
+
acceleration structures to an **unmodified** PyTorch host: `lerobot`,
|
|
37
|
+
Isaac-GR00T, `openpi`, `transformers`, `diffusers`, and inside vLLM /
|
|
38
|
+
SGLang. No fork, no edit to the host's source. Kernels arrive from the
|
|
39
|
+
Hugging Face kernel hub at bind time.
|
|
40
|
+
|
|
41
|
+
```python
|
|
42
|
+
import flashrt_structures as structures
|
|
43
|
+
|
|
44
|
+
plan = structures.attach(model, forward) # discover → calibrate → gate → activate
|
|
45
|
+
print(structures.explain(plan)) # bound / routed / kept-at-host / refused, with reasons
|
|
46
|
+
|
|
47
|
+
loop = structures.decode_loop(model, max_len=4096) # serving door
|
|
48
|
+
out = loop.generate(input_ids, max_new_tokens=256)
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
Measured on VLA, VLM, LLM and video models across RTX 5090 and Jetson AGX
|
|
52
|
+
Thor; the films are in the
|
|
53
|
+
[walkthrough](https://huggingface.co/spaces/liangsu9988/fast-kernels-are-not-fast-pipelines).
|
|
54
|
+
|
|
55
|
+
## What this package is, and what it is not
|
|
56
|
+
|
|
57
|
+
A *structure* is a versioned specification of one model region: boundary
|
|
58
|
+
tensors, framework-neutral weight slots, calibration points, gates, and a
|
|
59
|
+
plain-torch reference that is the gate's ground truth. Those
|
|
60
|
+
specifications, and the per-host *bindings* that say where the positions
|
|
61
|
+
sit on a concrete host, live in the FlashRT repository as
|
|
62
|
+
`flash_rt.catalog` and ship in the pure-Python `flash-rt` wheel.
|
|
63
|
+
|
|
64
|
+
This package is everything that acts on them for a PyTorch host:
|
|
65
|
+
|
|
66
|
+
| layer | here |
|
|
67
|
+
|---|---|
|
|
68
|
+
| `impls/` | executable forms per structure family: FP8 / NVFP4 / weight-only projections and FFNs, fused norm producers, attention cores, fused decoder and vision towers, the whole-step decode loop, graph-lowering pins |
|
|
69
|
+
| `adapters/` | host adapters: `transformers`, `diffusers`, Gemma / Qwen attention, gated-delta, and the vLLM / SGLang engine hooks |
|
|
70
|
+
| `discover.py`, `autobuild.py`, `points.py`, `schemes.py`, `gates.py` | structural discovery, the one-pass discover → calibrate → bind assembly, calibration collection, quantisation schemes, accuracy judgment |
|
|
71
|
+
| `guard.py`, `swap.py`, `frontdoor.py`, `stages.py`, `recipe.py` | the runtime contract and ledger, attach / detach, the one-call door, graph capture with declared swap windows, recipes |
|
|
72
|
+
| `examples/` | the explicit pipeline: seat tables, calibration hooks and binder calls written out by hand for GR00T N1.7 and π0.5 |
|
|
73
|
+
|
|
74
|
+
Structure specs, references and bindings are **not** here. A new structure
|
|
75
|
+
or a new host binding is a change to `flash_rt/catalog/` in FlashRT; a new
|
|
76
|
+
executable form, adapter or door is a change here.
|
|
77
|
+
|
|
78
|
+
## Install
|
|
79
|
+
|
|
80
|
+
```bash
|
|
81
|
+
pip install flashrt-structures # pulls flash-rt (pure Python) for the catalog
|
|
82
|
+
pip install "flashrt-structures[hub]" # + the kernel hub client, needed to bind
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
`flash-rt` must be a version that ships the catalog as `flash_rt.catalog`:
|
|
86
|
+
FlashRT `main` after the structures split, or a wheel newer than 0.1.0
|
|
87
|
+
(the 0.1.0 wheel on PyPI predates the split and still carries the layer
|
|
88
|
+
inside `flash_rt.structures`). An older flash-rt makes `import
|
|
89
|
+
flashrt_structures` fail with a message that says exactly this.
|
|
90
|
+
|
|
91
|
+
Bring your own torch. Which kernels exist is decided by the torch version:
|
|
92
|
+
the hub's published face is thickest at `torch 2.11 / cu128` on x86-64 and
|
|
93
|
+
much thinner at the newest release. If binds refuse on a fresh install,
|
|
94
|
+
check the torch version first.
|
|
95
|
+
|
|
96
|
+
Three boundaries produce refusals that look like bugs and are not:
|
|
97
|
+
|
|
98
|
+
- **Kernel availability follows the hub's build matrix, not this
|
|
99
|
+
package's.** The wheel installs on any torch; the kernels do not exist
|
|
100
|
+
for every torch.
|
|
101
|
+
- **`HF_HUB_OFFLINE=1` makes every kernel unavailable, even with a fully
|
|
102
|
+
warm cache**, because a version specifier has to resolve refs online.
|
|
103
|
+
Air-gapped deployments should stage packages and point at them with
|
|
104
|
+
`LOCAL_KERNELS=<repo>=<path>` rather than switching the hub offline.
|
|
105
|
+
- **aarch64 (Jetson Thor, sm_110) is not covered by the current
|
|
106
|
+
qualification pass.** Nothing in the wheel is architecture-bound, but
|
|
107
|
+
the engine adapters were last verified against vLLM 0.26 on Thor in
|
|
108
|
+
August 2026, not in the release run.
|
|
109
|
+
|
|
110
|
+
## Where to read next
|
|
111
|
+
|
|
112
|
+
- [`docs/hosts.md`](docs/hosts.md) — which door your host takes, how to
|
|
113
|
+
tell a seated run from a refused one from a door that never fired, how
|
|
114
|
+
to read a refusal, and what has been measured where.
|
|
115
|
+
- [`docs/structures.md`](docs/structures.md) — the norm: what a structure
|
|
116
|
+
is, the three-layer split, calibration reuse, accuracy bands, the
|
|
117
|
+
runtime contract and ledger, and the norms that came from being wrong.
|
|
118
|
+
- [`docs/serving_engines.md`](docs/serving_engines.md) — attaching inside
|
|
119
|
+
vLLM or SGLang without forking either.
|
|
120
|
+
- [`docs/adopt_in_20_lines.md`](docs/adopt_in_20_lines.md) — a measured
|
|
121
|
+
Qwen3-VL-8B adoption, twenty lines end to end.
|
|
122
|
+
- [`examples/README.md`](examples/README.md) — the explicit pipeline
|
|
123
|
+
against the automatic one, measured on two GR00T N1.7 hosts.
|
|
124
|
+
- [`docs/structure_release_qualification.md`](docs/structure_release_qualification.md)
|
|
125
|
+
— how a structure moves from a reusable boundary to a released hardware
|
|
126
|
+
route.
|
|
127
|
+
- [`AGENTS.md`](AGENTS.md) — the operating procedure for producing
|
|
128
|
+
structures; [`docs/structure_contributing.md`](docs/structure_contributing.md)
|
|
129
|
+
— the contribution boundary and PR self-review checklist.
|
|
130
|
+
|
|
131
|
+
## Tests
|
|
132
|
+
|
|
133
|
+
```bash
|
|
134
|
+
pip install -e . pytest
|
|
135
|
+
pytest tests # the CPU set; two *_gpu tests need a CUDA device and hub access
|
|
136
|
+
```
|
|
137
|
+
|
|
138
|
+
## License
|
|
139
|
+
|
|
140
|
+
Apache-2.0, same as FlashRT.
|