groot_n17: make the RTX FP8 pipeline usable in a control loop - #182
Open
LiangSu8899 wants to merge 2 commits into
Open
groot_n17: make the RTX FP8 pipeline usable in a control loop#182LiangSu8899 wants to merge 2 commits into
LiangSu8899 wants to merge 2 commits into
Conversation
The RTX FP8 frontend bakes the observation into set_prompt and refuses a second call, so a control loop cannot hand it new camera frames. The prompt-shaped pipeline, the activation scales and the captured DiT graphs depend on the token layout rather than pixel values, so a fresh frame only needs the feature pass: refresh_observation() recomputes the backbone features in place and leaves the rest of the pipeline standing. It refuses a bundle whose token count differs from the one the pipeline was built for.
set_prompt built a prompt-shaped pipeline for a single observation, so
three things only ever had to be right once. Driving it frame by frame
exposes all three.
* The action head reads cross-attention K/V, not the backbone, and on
RTX those live in the attention backend's own padded slots, filled
when the backend is built. refresh_observation replaced the backbone
features but nothing republished the K/V, so the head kept answering
about the first frame: a 7.58 change in backbone features moved the
action by exactly zero, and the policy ran on proprioception alone.
* The per-frame work rebuilt the whole plan -- buffers, per-layer
weight pointer tables, weight-scale scalars -- although all of it is
a function of the token layout. Record it on the first pass and
replay it: 9.16 -> 6.23 ms per frame, bitwise identical output.
* adopt_visual_merge() lets the frontend run the checkpoint's final
patch merger itself (cos 0.99999 against the shipped contract), so
the bundle carries text rows only, constant for a fixed prompt, and
the host's vision tower is not re-run per frame. Activation scales
are calibrated on the merged rows, not the text-only ones.
Measured on LIBERO-10 with 2 views: 15.85 ms per decision end to end,
against 41.75 ms for the stock host and 18.56 ms attached, and the first
frame's action agrees with the host's to within its own sampling spread.
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.
What this is
set_promptbuilds a prompt-shaped pipeline for one observation, so threethings in the RTX FP8 GR00T N1.7 frontend only ever had to be right once.
Driving it frame by frame in a control loop exposes all three.
The three
1. The action head kept answering about the first frame.
The head reads cross-attention K/V, not the backbone, and on RTX those live in
the attention backend's own padded slots, filled when the backend is built.
Replacing
_backbone_featuresleft the K/V standing: a 7.58 change in backbonefeatures moved the action by exactly zero, and the policy ran on proprioception
alone.
refresh_observation()now recomputes the VL-self-attention output andrepublishes the cross K/V into those slots.
2. Every frame rebuilt the whole plan.
Buffers, per-layer weight pointer tables and weight-scale scalars are all a
function of the token layout, not of pixel values. The first pass now records
them and later frames replay: 9.16 → 6.23 ms per frame, bitwise identical
output.
3. The host's vision tower was re-run per frame just to fill the visual rows.
adopt_visual_merge()(opt-in) runs the checkpoint's own final patch merger inthe frontend — it ships in the same checkpoint and is already loaded, and this
pipeline runs the full 24-layer ViT anyway to tap DeepStack. Cosine 0.99999
against the shipped contract. After the call the bundle carries text-only
llm_input_embeds, constant for a fixed prompt, plus freshpixel_features.Activation scales are calibrated on the merged rows, not the text-only ones —
calibrating on text-only rows would under-range every LLM GEMM.
It is off by default because it changes what the bundle is expected to contain.
Measured
LIBERO-10, 2 views, per decision end to end:
The first frame's action agrees with the host's to within its own sampling
spread.
Base and conflicts
This branch was developed and measured on the tree at
b5a97d07, before thebackbone-graph work landed on
main. It does not merge cleanly today:mainhas since added an
infer(aux=...)path with its own captured backbone graphand a graph-contract check, and rewrote
_run_kernel_backbone_fp8around it.The two approaches overlap — both re-run the backbone for a fresh observation,
one through a captured graph and one through the kernel path with a replayed
plan.
Rebasing is not mechanical and the result has to be re-measured on the target
card, so I am opening this against the measured tree rather than pushing an
unverified merge. Happy to redo it either way once we agree which of the two
paths should own the control loop.
One thing worth checking on
mainindependently of this branch: thecross-attention K/V slots described in (1) are filled when the attention
backend is built.
run_backbone_graphupdates_backbone_features; if nothingrepublishes the cross K/V after it, the
infer(aux=...)path has the same stale-observation behaviour this PR fixes. I have not measured that path, so this is
a question rather than a claim.