Skip to content

groot_n17: make the RTX FP8 pipeline usable in a control loop - #182

Open
LiangSu8899 wants to merge 2 commits into
mainfrom
feat/groot-n17-rtx-control-loop
Open

groot_n17: make the RTX FP8 pipeline usable in a control loop#182
LiangSu8899 wants to merge 2 commits into
mainfrom
feat/groot-n17-rtx-control-loop

Conversation

@LiangSu8899

Copy link
Copy Markdown
Member

What this is

set_prompt builds a prompt-shaped pipeline for one observation, so three
things 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_features left the K/V standing: a 7.58 change in backbone
features moved the action by exactly zero, and the policy ran on proprioception
alone. refresh_observation() now recomputes the VL-self-attention output and
republishes 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 in
the 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 fresh pixel_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:

arm ms
stock host 41.75
host + structures attached 18.56
this frontend 15.85

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 the
backbone-graph work landed on main. It does not merge cleanly today: main
has since added an infer(aux=...) path with its own captured backbone graph
and a graph-contract check, and rewrote _run_kernel_backbone_fp8 around 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 main independently of this branch: the
cross-attention K/V slots described in (1) are filled when the attention
backend is built. run_backbone_graph updates _backbone_features; if nothing
republishes 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.

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.
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.

1 participant