feat(groot): GR00T N1.6 Thor parity fixes and opt-in 28.5 ms profile

[DXICM]

Summary

GR00T N1.6-3B support for Jetson AGX Thor (SM110), including HF-alignment fixes, CUDA Graph lifecycle handling, and an explicit opt-in FA4/NVFP4 performance profile.

The existing public route remains the default: use_fp8=True, image_size=224, and parity=False. The new HF-native parity path and approximate FP4/FA4 tiers do not activate unless requested.

The reported opt-in profile measured 28.5 ms E2E on Thor (median 27.7 ms, p95 28.5 ms) with denormalized action cosine 0.999933 versus HF eager. These are the contributor's original Thor measurements; the audit follow-up below validated build and maintenance contracts but could not rerun Thor E2E while that machine was offline.

Audit follow-up

Commits eb01fa8d and 4c830e89 close the merge blockers found during review:

• Gates the qk_norm_rope_rotate_half_bf16 header and pybind symbol with the same FLASHRT_HAVE_THOR_VLA_KERNELS definition as its SM100/SM110 source.
• Keeps CUTE_DSL_ARCH=sm_101a as a CuTe-DSL compilation alias without overriding FLASH_ATTENTION_ARCH; runtime FA4 dispatch now sees the physical SM110 capability and does not enter the SM100-only hd256 kernel.
• Restores the pre-existing N1.6 defaults and makes FLASHRT_N16_DIT_FP4, FLASHRT_N16_QWEN3_FP4, FLASHRT_N16_SIGLIP_FP4, and FLASHRT_N16_FA4 opt-in.
• Routes the HF-native backbone and DiT only when parity=True, and rejects the unsupported FP8-backbone/Torch-DiT hybrid configuration.
• Implements prompt switching by resetting graph-owned state and recapturing on the next inference. FP4 patch state is reset with the graph runtime.
• Uses row-wise concat2_bf16 in DiT calibration, matching the production action/time embedding layout.
• Resolves Eagle remote code only from the checkpoint or an explicitly pinned FLASHRT_N16_EAGLE_DIR; no arbitrary Hugging Face cache revision is selected.
• Adds focused maintenance coverage for compile gates, FA4 runtime dispatch, defaults, prompt lifecycle, calibration layout, and remote-code resolution.

Opt-in performance profile

Construct the parity frontend explicitly with the model's 252x252 evaluation resolution, then enable the measured tiers:

export FLASHRT_N16_DIT_FP4=1
export FLASHRT_N16_QWEN3_FP4=1
export FLASHRT_N16_SIGLIP_FP4=1
export FLASHRT_N16_FA4=1

frontend = GrootTorchFrontendThor(
    checkpoint,
    use_fp8=False,
    image_size=252,
    parity=True,
)

All four tiers independently fall back to BF16/Torch when disabled or unavailable. FLASHRT_N16_DIT_STEPS remains 4; reducing diffusion steps is experimental and is not part of the reported profile.

Validation

• 24 passed across the N1.6 maintenance, prompt lifecycle, attention backend, calibration, and precision contract subset.
• Python compileall: passed.
• git diff --check: passed.
• SM89 flash_rt_kernels: full compile, link, and import passed; Thor-only binding absent.
• SM120 flash_rt_kernels: full compile, link, and import passed; Thor-only binding absent.
• SM110 flash_rt_kernels and flash_rt_fp4: full cross-compile, link, and import passed; Thor-only binding present.
• CMake inventory confirms the qk-norm source and binding macro are present only for SM100/SM110.

Thor hardware was offline during this audit follow-up, so the updated commit still requires one final Thor E2E rerun covering default fallback, the opt-in profile, prompt switching, finite outputs, and HF numerical comparison before merge.

Scope

The runtime changes are limited to the GR00T N1.6 Thor frontend and its isolated Thor FA4 backend. No new N1.6 path is enabled for other models or architectures.

[LiangSu8899]

Audit follow-up pushed in eb01fa8d and 4c830e89.

Resolved:

• SM100/SM110 qk-norm source, header, and pybind now share one compile gate.
• FA4 keeps the sm_101a DSL compilation alias but uses the physical GPU capability for runtime dispatch, so Thor hd256 cannot enter the SM100-only stub.
• Existing N1.6 behavior is restored as the public default; parity and all FP4/FA4 tiers are explicit opt-ins.
• HF-native capture is now controlled exclusively by parity=True; the unsupported FP8-backbone/Torch-DiT hybrid is rejected.
• Prompt changes reset and recapture graph-owned state, including FP4 module patches.
• DiT calibration now uses the same row-wise action/time concat as production.
• Eagle remote code must be checkpoint-local or selected through an explicit pinned directory.

Validation completed:

• 24 related contract tests passed.
• compileall and git diff --check passed.
• Full SM89 and SM120 flash_rt_kernels compile/link/import passed; the Thor-only binding is absent.
• Full SM110 flash_rt_kernels + flash_rt_fp4 cross-compile/link/import passed; the Thor-only binding is present.
• Merge simulation against current main is clean.

The Thor host was offline during this follow-up. I have therefore not represented the cross-build as a Thor runtime result; one final Thor E2E/default-fallback/prompt-switch numerical run remains before merge. The PR body now states this distinction and documents the opt-in 28.5 ms profile accurately.

[LiangSu8899]

Great Work!! @DXICM Could you help test whether this Thor behaves as expected?
Then I think we ready to merge!
Thx!!!

[DXICM]

Final Thor E2E verification (Jetson AGX Thor, SM110) — PR head 4c830e89

Fresh build of this PR head in a clean worktree on Thor hardware; all five
items from the pre-merge checklist covered.

Environment

• NVIDIA Thor, SM110 (CC 11.0), CUDA 13.0, torch 2.9.0a0+nv25.09, Python 3.12.3
• cmake -B build -DGPU_ARCH=110: SM100 CUTLASS FP8/FP16 → sm_110, Thor VLA
  helpers ENABLED, flash_rt_fp4 → sm_110a; CUTLASS v4.4.2.
• Thor-only symbol qk_norm_rope_rotate_half_bf16 present and importable.
• GPU idle-verified before timed runs; clocks at max throughout
  (GPC 1.575 GHz / NVD 1.692 GHz, BSP governor already pinned).
• Runtime isolation via PYTHONPATH only; system install untouched.

Contract tests

test_groot_n16_pr177_maintenance.py + test_prompt_runtime_lifecycle.py +
test_thor_groot_attn_backend.py → 24 passed, matching the PR report.

Results

Config	E2E p50	p95	Notes
Default FP8 224 (parity=False)	54.79 ms	55.95	vs unmodified main same harness: 54.72 ms → no regression (the USAGE.md 44 ms reference was not reproduced by either arm in this environment)
Parity BF16 252, all tiers off	63.77 ms	64.25	negative control
Opt-in FP4+FA4 252	28.39 ms	28.67	matches the reported 28.5 ms profile

The five checklist items

1. Default fallback — PASS. Finite outputs, sanity PASS, prompt reuse OK;
   opt-in env vars verified not to affect the default route; latency identical
   to the unmodified-main baseline in a paired comparison.

2. Opt-in profile — PASS. 28.39 ms p50 / 28.67 p95 (reported: median 27.7 /
   p95 28.5). Rejections verified: use_fp8=True, parity=True and
   parity=True, image_size≠252 both raise ValueError.

3. Prompt switching — PASS. A→B→A on one instance, no error; switch frame
   performs one recapture (~2.3 s), subsequent frames back to ~28.4 ms; all
   outputs finite.

4. Finite outputs — PASS. 60-frame stress with three 4 s idle gaps:
   0/60 non-finite; idle guard re-captures cleanly and latency returns to the
   28–29 ms class immediately.

5. Numerical comparison on real data — PASS. Real frames from the
   fine-tuned SO-101 dataset (lerobot_so101_sim_v1_gr00t_n1d6_sim_fruits_cubes_10w,
   new_embodiment slot), identical pixels/state/prompt and matched
   denoise-noise seed on both arms:

   frame	opt-in vs parity-BF16 cos	maxd
   50	0.999873	0.111
   150	0.999893	0.089
   300	0.999823	0.122

   Null checks: each arm rerun against itself → cos 1.000000, so the protocol
   and noise matching are exact. Consistent with the reported 0.999933 vs
   HF eager. One methodology note: precision comparisons must be run on an
   in-distribution checkpoint/inputs — out-of-distribution slots amplify
   kernel-level numerical differences through the 4-step flow-matching and
   grossly overstate quantization error.

Path-execution evidence

(fallback is numerically exact, so identical output alone proves nothing)

• Negative control: same parity frontend with all tiers off → 63.8 ms vs
  28.4 ms opt-in (2.2×) — the FP4/FA4 tiers are doing the work.
• flash_rt.hardware.thor.fa4_backend.status() → active; CuTe-DSL
  compilation exercised during the opt-in run.
• Per-tier isolation runs each engage their tier independently and stay at
  cos ≥0.9998 vs parity on real data.

One environment note for users

The Isaac n1.6.1-release Eagle-Block2A-2B-v2/modeling_siglip2.py raises
NameError: Siglip2Model under transformers ≥5 (_init_weights references a
removed class). A patched copy (branches removed) works and was used here via
an explicitly pinned FLASHRT_N16_EAGLE_DIR. Worth a note in the doc or a
version pin for users on transformers 5.

Cleanup

No system changes left behind (PYTHONPATH-only isolation, editable install
untouched); worktrees kept at /mnt/FlashRT-pr177 and /mnt/FlashRT-main-bcec
for reference.

Verdict: all five pre-merge items pass on Thor hardware; ready to merge.

[LiangSu8899]

Thx!!