Verified structures: catalog, gates, transactional attach, and provider export#150
Open
LiangSu8899 wants to merge 39 commits into
Open
Verified structures: catalog, gates, transactional attach, and provider export#150LiangSu8899 wants to merge 39 commits into
LiangSu8899 wants to merge 39 commits into
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
1 participant
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.
Implements the structural layer tracked in #149: structures as specs with references and gates, qualified implementations over Hub kernels, and an export path that packages a qualified host as a standard model runtime.
What's here
flash_rt/structures/decoder_ffn,vision_ffn): boundary contract, weight slots, variants, and a ~20-line plain-torch reference per structure (the ground truth the gate scores against).fp8_staticimplementations of both structures over the Hub FP8 FFN kernels:bind()packs weights and calibrates static scales from real-distribution data, in full-boundary and host-MLP-seam forms.detach()restores the original modules exactly.frt_model_runtime_v1, the same runtime face the native pipelines export — so a torch host with structures attached becomes one more producer of the standard serving contract.Structures vs standalone kernel swaps
The load-bearing A/B, measured on a real fine-tuned Pi0.5 checkpoint with real observation data — same Hub kernels, same calibrated scales; the only variable is the composition:
torch.compilemax-autotuneThe refusal rows are the point of the gate design: workload sensitivity is a guardrail, not a caveat.
Cross-model reuse
The same two structure definitions, zero spec changes, across model families and host forms — each new host costs a ~20-line binding. Layer gates at 0.999 on the structure boundary (residual included), then a family-level end-to-end gate (accuracy budget and net win):
Three qualitatively different gate outcomes, all correct: full activation where the workload wins, family-selective activation where amplification kills one family (Qwen3-VL's deepstack injects vision features into text layers — layer gates alone under-predict e2e there, which is exactly why the family-level gate exists), and whole-host refusal where small dims make FP8 boundary costs unprofitable. Nothing activates at a loss, ever.
The GR00T row is also the strongest reuse statement: the same
vision_ffndefinition qualifies both a SigLIP encoder tower and a DiT diffusion action head — two different host forms, one spec, and a third independent host stack (Isaac-GR00T native, neither transformers nor lerobot) adapted with three ~20-line bindings and no spec or implementation changes.Adoption cost vs standalone
What a standalone kernel swap asks of the integrator, per model and per site (this is what in-host integrations hand-roll today, typically ~100+ lines):
The same adoption through structures:
Packing, calibration, layout, parity, and rollback are the mechanism's job, not the integrator's. A new host is a ~20-line binding (module paths + weight layout); every binding in the table above cost exactly that. Next step on this axis: bindings become directly consumable, so adoption collapses further into a lookup + one call while keeping the full gate discipline.
Validation status
All gates run on real checkpoints with real-distribution data (calibration and evaluation sets kept separate): a real fine-tuned Pi0.5 checkpoint with real robot observations, stock transformers hosts with real long-document corpora, real images for the VLM host.
provider.pyand passes the same serving adoption/tick acceptance as the native pipelines (per-tick hot input with no recapture, deterministic replay, boundary snapshot/restore).Boundaries