flashrt-structures 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (122) hide show
  1. flashrt_structures/__init__.py +174 -0
  2. flashrt_structures/adapters/__init__.py +55 -0
  3. flashrt_structures/adapters/diffusers_attention.py +237 -0
  4. flashrt_structures/adapters/diffusers_rotary_attention.py +239 -0
  5. flashrt_structures/adapters/factored_qk_norm_rope.py +252 -0
  6. flashrt_structures/adapters/factored_two_way_attention.py +99 -0
  7. flashrt_structures/adapters/gemma_attention.py +226 -0
  8. flashrt_structures/adapters/packed_qkv_rope.py +342 -0
  9. flashrt_structures/adapters/packed_stream_qk_norm_rope.py +376 -0
  10. flashrt_structures/adapters/qwen_per_head_qk_norm_rope.py +231 -0
  11. flashrt_structures/adapters/sglang_engine.py +207 -0
  12. flashrt_structures/adapters/transformers_attention_interface.py +73 -0
  13. flashrt_structures/adapters/transformers_gated_delta.py +152 -0
  14. flashrt_structures/adapters/transformers_gated_delta_fused.py +96 -0
  15. flashrt_structures/adapters/vllm_engine.py +424 -0
  16. flashrt_structures/adjudicate.py +85 -0
  17. flashrt_structures/aot.py +191 -0
  18. flashrt_structures/autobuild.py +2052 -0
  19. flashrt_structures/beta/__init__.py +43 -0
  20. flashrt_structures/beta/conform.py +94 -0
  21. flashrt_structures/beta/joins.py +113 -0
  22. flashrt_structures/beta/negotiate.py +84 -0
  23. flashrt_structures/beta/ports.py +140 -0
  24. flashrt_structures/decisions.py +80 -0
  25. flashrt_structures/discover.py +623 -0
  26. flashrt_structures/explain.py +70 -0
  27. flashrt_structures/frontdoor.py +572 -0
  28. flashrt_structures/gates.py +465 -0
  29. flashrt_structures/guard.py +421 -0
  30. flashrt_structures/handle.py +189 -0
  31. flashrt_structures/impls/__init__.py +219 -0
  32. flashrt_structures/impls/adaln_producer/__init__.py +8 -0
  33. flashrt_structures/impls/adaln_producer/broker.py +116 -0
  34. flashrt_structures/impls/adaln_producer/fused.py +388 -0
  35. flashrt_structures/impls/adarms_stack/__init__.py +8 -0
  36. flashrt_structures/impls/adarms_stack/fp8_chain.py +832 -0
  37. flashrt_structures/impls/adarms_stack/region.py +102 -0
  38. flashrt_structures/impls/attention_core/__init__.py +132 -0
  39. flashrt_structures/impls/attention_core/fa2_seqused.py +458 -0
  40. flashrt_structures/impls/attention_core/fa4_cute.py +154 -0
  41. flashrt_structures/impls/attention_core/fa4_fp8.py +178 -0
  42. flashrt_structures/impls/attention_core/masked_mha.py +158 -0
  43. flashrt_structures/impls/attention_core/two_way_fa2.py +220 -0
  44. flashrt_structures/impls/cadence_static/__init__.py +18 -0
  45. flashrt_structures/impls/cadence_static/buffers.py +122 -0
  46. flashrt_structures/impls/cadence_static/cross_attention.py +187 -0
  47. flashrt_structures/impls/chain_elements.py +89 -0
  48. flashrt_structures/impls/decode_loop/__init__.py +0 -0
  49. flashrt_structures/impls/decode_loop/fp8_kv.py +206 -0
  50. flashrt_structures/impls/decode_loop/mtp_speculative.py +245 -0
  51. flashrt_structures/impls/decode_loop/whole_step.py +852 -0
  52. flashrt_structures/impls/decoder_block/__init__.py +6 -0
  53. flashrt_structures/impls/decoder_block/attn_sublayer.py +110 -0
  54. flashrt_structures/impls/decoder_block/fused.py +167 -0
  55. flashrt_structures/impls/decoder_ffn/__init__.py +0 -0
  56. flashrt_structures/impls/decoder_ffn/fp8_static.py +310 -0
  57. flashrt_structures/impls/decoder_ffn/fp8_static.yaml +22 -0
  58. flashrt_structures/impls/decoder_ffn/w4a16_static.py +221 -0
  59. flashrt_structures/impls/decoder_ffn/w8a16_static.py +183 -0
  60. flashrt_structures/impls/dit_stack/__init__.py +11 -0
  61. flashrt_structures/impls/dit_stack/fp4_chain.py +417 -0
  62. flashrt_structures/impls/dit_stack/region.py +86 -0
  63. flashrt_structures/impls/fixed_iter/__init__.py +29 -0
  64. flashrt_structures/impls/fixed_iter/openpi.py +264 -0
  65. flashrt_structures/impls/fixed_iter/protocol.py +94 -0
  66. flashrt_structures/impls/gated_delta_core/__init__.py +3 -0
  67. flashrt_structures/impls/gated_delta_core/fused_layer.py +545 -0
  68. flashrt_structures/impls/gated_delta_core/hub_v3.py +152 -0
  69. flashrt_structures/impls/graph_lowering/__init__.py +27 -0
  70. flashrt_structures/impls/graph_lowering/pi052_denoise.py +179 -0
  71. flashrt_structures/impls/graph_lowering/protocol.py +76 -0
  72. flashrt_structures/impls/graph_lowering/qwen3_vl.py +364 -0
  73. flashrt_structures/impls/linear_proj/__init__.py +0 -0
  74. flashrt_structures/impls/linear_proj/fp8_static.py +270 -0
  75. flashrt_structures/impls/linear_proj/nvfp4_balance.py +131 -0
  76. flashrt_structures/impls/linear_proj/nvfp4_dynamic.py +182 -0
  77. flashrt_structures/impls/linear_proj/w8a16_static.py +230 -0
  78. flashrt_structures/impls/modnorm_qkv_chain/__init__.py +0 -0
  79. flashrt_structures/impls/modnorm_qkv_chain/fp8_ptok_table.py +291 -0
  80. flashrt_structures/impls/moe_experts/__init__.py +9 -0
  81. flashrt_structures/impls/moe_experts/nvfp4_dynamic.py +208 -0
  82. flashrt_structures/impls/moe_experts/nvfp4_w4a16.py +129 -0
  83. flashrt_structures/impls/norm_fused/__init__.py +3 -0
  84. flashrt_structures/impls/norm_fused/bf16.py +94 -0
  85. flashrt_structures/impls/norm_fused/fp8_producer.py +84 -0
  86. flashrt_structures/impls/patch_projection/__init__.py +3 -0
  87. flashrt_structures/impls/patch_projection/bf16_flat.py +139 -0
  88. flashrt_structures/impls/prefill_tower/__init__.py +10 -0
  89. flashrt_structures/impls/prefill_tower/fp8_chain.py +955 -0
  90. flashrt_structures/impls/prefill_tower/region.py +99 -0
  91. flashrt_structures/impls/qk_norm_rope/__init__.py +12 -0
  92. flashrt_structures/impls/qk_norm_rope/per_head_gqa.py +199 -0
  93. flashrt_structures/impls/qk_norm_rope/projection_bf16.py +165 -0
  94. flashrt_structures/impls/qkv_pack/__init__.py +5 -0
  95. flashrt_structures/impls/qkv_pack/bf16.py +110 -0
  96. flashrt_structures/impls/qkv_pack/fp8_static.py +435 -0
  97. flashrt_structures/impls/qkv_pack/nvfp4_balance.py +218 -0
  98. flashrt_structures/impls/qkv_rope/__init__.py +3 -0
  99. flashrt_structures/impls/qkv_rope/packed_bias_bf16.py +143 -0
  100. flashrt_structures/impls/step_table.py +113 -0
  101. flashrt_structures/impls/vision_ffn/__init__.py +0 -0
  102. flashrt_structures/impls/vision_ffn/fp8_static.py +261 -0
  103. flashrt_structures/impls/vision_ffn/nvfp4_balance.py +211 -0
  104. flashrt_structures/impls/vision_tower/__init__.py +7 -0
  105. flashrt_structures/impls/vision_tower/fp8_chain.py +533 -0
  106. flashrt_structures/impls/vision_tower/region.py +91 -0
  107. flashrt_structures/matrix.py +126 -0
  108. flashrt_structures/points.py +368 -0
  109. flashrt_structures/prequantized.py +131 -0
  110. flashrt_structures/quantize_on_adopt.py +94 -0
  111. flashrt_structures/recipe.py +438 -0
  112. flashrt_structures/regions.py +208 -0
  113. flashrt_structures/schemes.py +490 -0
  114. flashrt_structures/stages.py +298 -0
  115. flashrt_structures/storage.py +255 -0
  116. flashrt_structures/swap.py +398 -0
  117. flashrt_structures/workspace.py +110 -0
  118. flashrt_structures-0.2.0.dist-info/METADATA +140 -0
  119. flashrt_structures-0.2.0.dist-info/RECORD +122 -0
  120. flashrt_structures-0.2.0.dist-info/WHEEL +5 -0
  121. flashrt_structures-0.2.0.dist-info/licenses/LICENSE +202 -0
  122. flashrt_structures-0.2.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,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.