Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Verified structures: catalog, gates, transactional attach, and provider export#150

Open
LiangSu8899 wants to merge 48 commits into
mainflashrt-project/FlashRT:mainfrom
feat/hf-kernels-structuresflashrt-project/FlashRT:feat/hf-kernels-structuresCopy head branch name to clipboard
Open

Verified structures: catalog, gates, transactional attach, and provider export#150
LiangSu8899 wants to merge 48 commits into
mainflashrt-project/FlashRT:mainfrom
feat/hf-kernels-structuresflashrt-project/FlashRT:feat/hf-kernels-structuresCopy head branch name to clipboard

Conversation

@LiangSu8899

@LiangSu8899 LiangSu8899 commented Jul 17, 2026

Copy link
Copy Markdown
Member

Structures as specs with references and gates, qualified implementations over Hub kernels, one-call assembly onto a host, and an export path that packages a qualified host as a standard model runtime. Tracks #149.

Usage

One call

from flash_rt import structures
from flash_rt.structures.swap import attach

plan = structures.auto_swaps(model, forward)   # discover, calibrate, bind
handle = attach(model, plan.swaps)             # transactional swap-in
...
handle.detach()                                # exact restore

forward is any callable that runs the host once. It is used to capture calibration; nothing else is asked of the host — no module paths, no hooks, no scale plumbing.

plan.swaps is {module_path: replacement}. plan.notes carries the receipt: what was refused and why, which seams were negotiated as a chain, how calibration was obtained.

auto_swaps builds a plan; it does not judge one. Binding-time qualification (work bands, unsupported head dims, non-step-quantised conditioning) is applied and recorded in plan.notes["refused"], but the accuracy and net-win verdict comes from the bench below, or from the gated front door:

plan = structures.attach(model, forward)       # discover, calibrate, gate, activate
plan.activated                                 # {} when the gates refused

structures.attach runs per-layer parity, a family-level accuracy budget and a net-win A/B before touching the model, and covers decoder_ffn / vision_ffn by default.

Calibration

Three ways in, one axis. forward is always "run the host once".

structures.auto_swaps(model, forward)                    # one frame (default)
structures.auto_swaps(model, forward, frames=8)          # closure advances its own data
structures.auto_swaps(model, [f0, f1, f2])               # one thunk per frame
structures.auto_swaps(model, feed, samples=dataset)      # feed(sample) per sample

samples is any iterable; with it, forward takes one sample. frames defaults to one, and to the whole of samples when a sample source is given.

Static scales accumulate over every call a seam sees, so they are exact regardless of how many activations are held. keep_samples bounds what is held for the parity self-check; it defaults to unbounded at one frame and to a cap above it, so more frames cost time and not memory (measured: 14.4 GiB at one frame, 18.6 GiB at four).

plan.notes["calibration"] records {frames, source, stat, keep_samples}.

Single site, visible

ffn = structures.get("decoder_ffn")
new_mlp = ffn.bind(model.model.layers[3].mlp, calibration=[x1, x2],
                   residual=[r1, r2])          # gate at the declared boundary
model.model.layers[3].mlp = new_mlp            # your swap, your call

bind extracts weights, calibrates, builds the replacement and self-checks it against the original on those samples. Missing the parity gate raises GateRefused rather than returning a bad part.

Graph boundary

stage = structures.capture(hot, windows={"noise": noise_buf},
                           reference=eager_reference)
stage.replay()
runtime = stage.export(ports=...)               # frt_model_runtime_v1

Every varying input must be a declared window; in-graph RNG never matches eager, so noise is a window rather than a call.

Development bench

run = structures.run_recipe(recipe, model, shape_sig)
run.verdict          # win | refused
run.receipt          # per-lever state, parity band, drift, seam counts

Same-process A/B with a baseline re-time. A lever that does not both stay accurate and win latency is refused and recorded.

Catalog

Structure Boundary
decoder_ffn gated MLP with its norm
vision_ffn fc1/fc2 MLP with its LayerNorm
linear_proj one projection; forms bias / no_bias / fp8_in, each with its own measured work band
qkv_pack sibling projections sharing one input, packed into one GEMM; leaf and module forms
adaln_producer conditioning-driven norm; step table, fingerprint locator, fused fp8 producer
norm_fused an affine norm the host runs in fp32, collapsed to a fused bf16 kernel
attention_core attention over packed keys/values, mask resolved at bind time
decoder_block the pre-norm block as one boundary: folds the pending gated residual into the producer's kernel and removes the cancelling q/k/v transposes
cadence_static work that changes slower than the hot loop
vla_tick_pipeline schedule level: observation-cadence encode feeding a fixed-shape captured hot loop

Each entry is catalog/<name>/structure.yaml (boundary, weight slots, variants, qualification, gates, evidence) plus a plain-torch reference used as the gate's ground truth.

Code

Path Role
catalog/ structure specs and references
registry.py spec lookup; pure metadata
discover.py finds seams by shape and slot, not by module name
autobuild.py one-call assembly: discover, calibrate, negotiate, bind
impls/ implementations over Hub kernels
adapters/ host-family adapters for seams that are not a static module pattern
swap.py transactional attach/detach
handle.py get(name).bind(...) explicit door
stages.py capture(...) graph door
recipe.py development-side A/B bench and receipts
provider.py export a captured host as frt_model_runtime_v1
bindings/ per-host receipts (module paths, weight layouts)
beta/ join declarations between adjacent structures; opt-in, not consumed by assembly yet

Hosts

Four independent host stacks, no shared integration path between them:

Host Weights Host stack
Pi0.5 fine-tuned checkpoint lerobot policy
SmolVLA lerobot/smolvla_base lerobot policy
GR00T N1.6-3B nvidia/GR00T-N1.6-3B Isaac-GR00T native — neither transformers nor lerobot
Qwen2.5-1.5B / Qwen3-8B / Qwen3-VL-8B stock Hub weights stock transformers

Results

Every parity figure is same input, compared output, against that host's own unmodified form.

Host Baseline With structures Speedup Parity
Pi0.5 — tick 59.0 ms · torch.compile + CUDA graph 24.46 ms 2.41× 0.99990 action
GR00T N1.6-3B — tick 57.4 ms · official pipeline, eager 11.26 ms 5.10× 0.999994 action
SmolVLA — tick 87.6 ms · official pipeline, eager 29.3 ms 2.99× 0.999995 action
SmolVLA — replan 87.6 ms · official pipeline, eager 12.6 ms 6.95× 0.999995 action
Qwen3-8B — prefill, 360 tokens 51.2 ms · eager 24.8 ms 2.06× top-1 99.4%, last-position cos 0.9995

The two VLA tick figures combine both levers. On GR00T they decompose cleanly: the stage split carries 57.4 → 15.42 ms (3.72×), and the region structures take 15.46 → 11.26 ms (1.37×) on top of it.

Earlier matrix, decoder_ffn / vision_ffn only, eager baselines:

Host Structures Layer gates Speedup Parity
Qwen2.5-1.5B decoder_ffn ×28 28/28 1.39× prefill logits 0.9982
Qwen3-8B decoder_ffn ×36 36/36 1.51× prefill (70.7 → 46.9 ms) logits 0.9975
Qwen3-VL-8B decoder_ffn + vision_ffn text 35/36, vision 27/27 1.43× text (150.7 → 105.1 ms) text 0.9972

Qwen3-VL's vision family is refused by the family gate at 0.9906 while its text family activates: deepstack injects vision features into text layers, so layer gates alone under-predict end to end there.

The GR00T row is the reuse statement — the same vision_ffn definition qualifies both a SigLIP encoder tower and a DiT diffusion action head, on a host stack that is neither transformers nor lerobot.

Structure vs standalone kernel swap

Same Hub kernels, same calibrated scales; the only variable is the composition.

Comparison Result
Per-op standalone swap vs eager host net negative at every tested M — boundary quant/dequant eats the kernel win
Structure (composed region) vs per-op standalone 1.65–2.1× across the same M sweep
Structures vs torch.compile max-autotune 1.35× survives on top of the strongest compile

Choosing the parity metric for a host

Whole-tensor cosine over logits is the wrong measure for a language host, and measuring it exposed why. Same bindings, same weights, two prompt lengths:

Structures bound Tokens Cosine over all logits Top-1 agreement Last-position cosine KL per token
decoder_ffn ×36 15 0.9991 93.3% 0.99973 0.011 nats
decoder_ffn ×36 360 0.9450 99.4% 0.99975 0.007 nats
full set (180 seams) 15 0.9990 93.3% 0.99935 0.011 nats
full set (180 seams) 360 0.9241 99.2% 0.99951 0.007 nats

The two metrics move in opposite directions with length: the aggregate cosine falls while token agreement rises. Aggregated over every position it is dominated by positions that never drive a decision, so it tracks sequence length more than it tracks output fidelity. The generation-relevant quantities — the last position, top-1 agreement, per-token KL — are stable across both lengths and across structure sets.

This applies to the older LLM rows in the table above, which were scored the same way and are therefore length-dependent; they are kept as recorded rather than restated.

For hosts whose output is a distribution rather than a value, the gate metric belongs to the host: token agreement and last-position fidelity for a language model, action cosine for a policy. That selection is not yet part of the gate and is the next thing owed here.

Parity is per workload

Static per-tensor scales calibrated from a single frame accumulate across depth, so a parity figure belongs to a workload rather than to a host. The receipt carries the calibration method next to the band, and the multi-frame and dataset entries exist for this.

Validation

  • Real checkpoints and real-distribution data; calibration and evaluation sets kept separate.
  • References match live model outputs at the declared structure boundary.
  • Every host is taken through discover → bind → attach → real forward → parity. A plan that builds is not evidence that it runs.
  • attach/detach is transactional: failed resolution leaves the model untouched, detach restores bit-exactly.
  • A captured pipeline exports through provider.py and passes the same serving adoption/tick acceptance as the native pipelines.

Boundaries

  • Additive only: no changes to existing kernels, runtime bindings, or pipelines.
  • Hosts are never modified on disk; attach is an in-memory, reversible module swap.
  • Implementations stay native to each host; nothing is generated or translated.

Introduce flash_rt.structures: a registry over versioned structure
specifications (boundary tensors, framework-neutral weight slots,
reference implementation, qualification gates). First catalog entry is
decoder_ffn, covering the gelu/silu activation, offset/direct norm
weight, and none/ada_ln conditioning variants shared across the
supported decoder families.
Structure-agnostic harness: derives the calling convention from the
specification, resolves symbolic dims from actual tensors with
consistency checks, judges an implementation against the reference
(cosine / max-abs / p99 metrics vs explicit thresholds), and emits a
machine-readable record keyed by a plan digest over spec content,
variant, workload, implementation identity, and environment.
Bind-time weight packing and static-scale calibration feeding the fused
FP8 gate/up -> activation -> down block from
flashrt/flashrt-fp8-swiglu-ffn (gelu and silu entrypoints). Enforces
the support envelope and non-empty calibration inputs at bind. The
parity gate gains a bound-callable mode for implementations whose
weights and variant are baked at bind time.
Maps the transformers-convention checkpoint layout (out-in linear
weights, AdaRMS conditional norm with no learned weight, time-embedding
modulation projected to scale/shift/gate, gated residual) onto the
decoder_ffn slots and conditioning inputs.
Add the optional cond_gate boundary input: under the ada_ln
conditioning variant the output residual becomes x + ffn_out * gate,
matching the AdaRMS decoder families whose norm projection emits
scale/shift/gate. The reference and the fp8_static implementation both
honor it.
attach() swaps host modules atomically: all staged paths are resolved
and validated before the first swap, any failure rolls back completely,
and the returned handle restores the exact originals idempotently.
fp8_static gains bind_mlp_seam for hosts whose replaceable boundary is
the MLP module: shared packing and calibration with the full-structure
bind, the host keeping its own norm, AdaLN gate, and residual. Parity
metrics are exported for host-side acceptance checks.
Weight packing and calibration are setup-time operations; without
no_grad the quantization chain recorded autograd history on the packed
FP8 tensors, keeping every fp32 intermediate alive for the lifetime of
the bound implementation (~1.7 GB per trunk-sized layer).
Second binding of the same structure: the PaliGemma trunk MLPs
(standard learned-weight RMSNorm, no conditioning, plain residual) at
prefix token counts. Together with the pi05 action-expert binding this
covers both M regimes of the model with one structure definition.
Third binding of the same structure, first outside the pi05 family:
Qwen2.5-1.5B-Instruct under a plain transformers host (silu activation,
direct norm-weight convention, no conditioning).
torch.quantile rejects large inputs (LLM-logits-sized qualification
outputs exceed its limit); kthvalue is exact and unbounded.
Second catalog structure: the non-gated LayerNorm + fc1/GELU/fc2 vision
FFN block shared by the SigLIP towers of the supported VLA families,
implemented over the fused FP8 GELU MLP Hub kernel (checkpoint-native
weight layout, biases included).
Introduce the first schedule-layer structure. stage_pipeline specs declare
a stage graph with cadence attributes instead of a tensor boundary; their
parity ground truth is the host's own eager path under an explicit noise
window, so the registry now carries kind/family/stages/conformance fields
and region entries are unchanged.

vla_tick_pipeline (cond_iter_pipeline family, tick specialization):
obs_encode at observation cadence feeding a fixed-shape K-step denoise
loop at tick cadence, noise as a SWAP window, condition buffers read-only.
Includes the GR00T N1.6 binding with its capture rulings (eager backbone,
shallow-copy mutation guard, hoisted noise site).
Same structure declaration as the GR00T N1.6 binding. obs_encode ruled
eager: SmolVLM vision embeddings use a boolean-mask scatter that is
illegal on a capturing stream (same blocker class as SigLIP2 CPU
indexing), so the tick form is eager prefill into static KV buffers plus
a captured denoise loop; noise uses the host's native injectable window.
Consumption now mirrors kernels.get_kernel: a single call discovers
structure seams in the host (pattern-matched from the module tree, no
hand-written binding required), calibrates on the caller's real forward
passes, runs per-layer parity gates and family-level accuracy plus
net-win gates, transactionally activates only what earns, and returns a
Plan with a receipt and exact detach.

Calibration entries: calibration="auto" captures during the provided
forward(s); a file path loads a prior capture or saves this one, so
dataset-scale calibration is a path string. Multi-frame via a sequence
of forwards or frames=N. The net-win gate requires a margin
(min_speedup, default 1.02x) so measurement noise can never activate a
family.

The swap machinery module is renamed attach.py -> swap.py to free the
attach name for the front door; no behavior change.
The single-site counterpart of attach, with get_kernel ergonomics: pull
one structure, bind it to the module you point at, plug it in yourself.
bind extracts weights from the module, calibrates on caller-provided
real inputs, and self-checks the replacement against the original
before handing it back; a part that misses the parity gate raises
GateRefused instead of being returned. Pass residual= to gate at the
declared structure boundary (residual included) — the same measurement
attach uses; without it the check runs on the bare seam, which is
strictly harsher. The returned module carries a certification record
(worst cos, gate boundary, dims, m profile, variant).
structures.capture(fn, windows=..., reference=...) is the schedule-
structure counterpart of the region doors: it warms and records the hot
stage into a CUDA graph on a side stream and returns a replayable
CapturedStage. Declared windows are the tensors the caller may rewrite
between replays (noise, observations, condition buffers); replay reads
their current contents in place. Parity against the host's eager path
under the same window contents and a margin-gated net-win check run
inside the call; a stage that fails either raises CaptureRefused.

Replay timing records its events on the replay stream — default-stream
events only measure enqueue time and overstate the win by orders of
magnitude.
A calibrated single projection (x[M,K] -> y[M,N], optional bias) with
epilogue variants (gelu+quant, residual add) and a negotiable input
dtype: fp8_static marks a producer-negotiated seam where an upstream
norm/adaln producer emits fp8+scale and this region skips its own input
quantization — decided at plan level and re-certified at the composed
boundary. Elementwise work exists only as epilogue variants here, never
as standalone regions. Discovery is qualified (weight-size floor,
sibling q/k/v/o grouped into one family), not a blind scan of every
nn.Linear.
linear_proj gains its first implementation: the fused BF16-entry FP8
projection (FP8 weights, static per-tensor activation scale, fused
input quantization), with work-based qualification measured from
standalone preflight — projections whose GEMM cannot amortize the fixed
quantization cost are refused at bind time and the host keeps its
Linear. Per-calibrated-M buffers are pre-allocated so the hot path is
allocation-free under compile and graph capture.

Discovery matches sibling q/k/v/o(out) projections as one family per
attention block behind a weight-size candidacy floor — never a blind
scan of every nn.Linear. The front door records bind-time refusals
per layer, and impls now share one process-wide hub loader: importing
the same kernel repo twice re-registers its fake ops and raises.
A captured stage's declared windows become frt_model_runtime_v1
boundary windows directly; ports reference them by name. This closes
the absorption edge: attach + capture + export takes a torch host to a
runtime the FlashRT serving mechanisms (Nexus tick, capsule
snapshot/restore) consume exactly like a whitebox-produced one.
Calling the hub loader from forward makes dynamo trace through
kernels.get_kernel's version resolution (network calls,
inspect.Signature) — 26 graph breaks that fragment the surrounding
compiled region, and a fragment boundary can drop host-side constant
construction into eager code that is illegal under CUDA graph capture.
Binding the op function once at bind time keeps the module a single
traceable custom-op call, which is why the FFN implementations (which
already did this) never fragmented.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Tracking] Verified structures — composing HF Kernels into net-win building blocks

1 participant

Morty Proxy This is a proxified and sanitized view of the page, visit original site.