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#150LiangSu8899 wants to merge 48 commits intomainflashrt-project/FlashRT:mainfrom feat/hf-kernels-structuresflashrt-project/FlashRT:feat/hf-kernels-structuresCopy head branch name to clipboard
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
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.
…ication lifecycle
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
forwardis 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.swapsis{module_path: replacement}.plan.notescarries the receipt: what was refused and why, which seams were negotiated as a chain, how calibration was obtained.auto_swapsbuilds a plan; it does not judge one. Binding-time qualification (work bands, unsupported head dims, non-step-quantised conditioning) is applied and recorded inplan.notes["refused"], but the accuracy and net-win verdict comes from the bench below, or from the gated front door:structures.attachruns per-layer parity, a family-level accuracy budget and a net-win A/B before touching the model, and coversdecoder_ffn/vision_ffnby default.Calibration
Three ways in, one axis.
forwardis always "run the host once".samplesis any iterable; with it,forwardtakes one sample.framesdefaults to one, and to the whole ofsampleswhen 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_samplesbounds 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
bindextracts weights, calibrates, builds the replacement and self-checks it against the original on those samples. Missing the parity gate raisesGateRefusedrather than returning a bad part.Graph boundary
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
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
decoder_ffnvision_ffnlinear_projbias/no_bias/fp8_in, each with its own measured work bandqkv_packadaln_producernorm_fusedattention_coredecoder_blockcadence_staticvla_tick_pipelineEach 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
catalog/registry.pydiscover.pyautobuild.pyimpls/adapters/swap.pyhandle.pyget(name).bind(...)explicit doorstages.pycapture(...)graph doorrecipe.pyprovider.pyfrt_model_runtime_v1bindings/beta/Hosts
Four independent host stacks, no shared integration path between them:
lerobot/smolvla_basenvidia/GR00T-N1.6-3BtransformersResults
Every parity figure is same input, compared output, against that host's own unmodified form.
torch.compile+ CUDA graphThe 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_ffnonly, eager baselines:decoder_ffn×28decoder_ffn×36decoder_ffn+vision_ffnQwen3-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_ffndefinition 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.
torch.compilemax-autotuneChoosing 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:
decoder_ffn×36decoder_ffn×36The 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
attach/detachis transactional: failed resolution leaves the model untouched,detachrestores bit-exactly.provider.pyand passes the same serving adoption/tick acceptance as the native pipelines.Boundaries