GROOT N1.7 update: Thor NVFP4 + FA4 performance tier

[LiangSu8899]

Summary

Adds an opt-in NVFP4 + FA4 performance tier for GROOT N1.7 on Jetson AGX Thor (SM110). The default FP8 path is unchanged.

Measured on the GR00T-N1.7 LIBERO fine-tune (libero_10), one camera, 4 denoising steps, batch 1, medians over 20 iterations after 5 warmup iterations, on a fixture captured from a real libero_10 observation through the official preprocessing path (JetPack 7.2, MAXN):

Tier	Backbone	Action head	E2E	Frequency
FP8 (default, unchanged)	11.0 ms	25.8 ms	36.8 ms	27.2 Hz
NVFP4 + FA4 (use_fp4=True)	8.4 ms	15.2 ms	23.7 ms	42.3 Hz

Run as one same-session A/B/A: FP8 36.76 ms → NVFP4 23.65 ms → FP8 36.85 ms. Action cosine between the tiers on that fixture is 1.000000.

On the base GR00T-N1.7-3B checkpoint with a 2-view fixture (T=40, same protocol): FP8 50.2 ms (20 Hz) → NVFP4 29.9 ms (33 Hz), also A/B/A (51.6 / 29.9 / 50.2), action cosine 0.99994.

CUDA-graph replays on the new tier are bit-identical (the masked-softmax attention also removes a long-standing source of replay nondeterminism in the FP8 tier's logits padding). On an 8-sample aux set captured from the base checkpoint, 7 of 8 samples hold action cosine >= 0.9993 against the FP8 tier with a worst sample of 0.9976 — deployments that are accuracy-critical should validate on the target task rather than on cosine alone, or stay on the FP8 tier.

An earlier revision of this branch carried a dit_fp8_layers precision ladder. It was removed before review: with every DiT layer marked for FP8 the captured graph still produced the all-NVFP4 result, so the option did not do what it documented. The last commit removes the knob, its per-layer gate, and the benchmark flag rather than shipping unverified behavior.

What is in the tier

• NVFP4 DiT action head — every DiT GEMM (fused self-attn QKV, cross-attn Q, attention output, both FFN projections) runs as a block-scaled NVFP4 GEMM with fused bf16 bias / bias+residual / bias+tanh-GELU+FP4-output epilogues. At M=41 the DiT is weight-bandwidth-bound, so halving weight bytes is the dominant win. Activation quantization is dynamic per-16-element (no calibration pass).
• NVFP4 cross-KV projections — the 32 per-frame K/V GEMMs quarter their weight traffic.
• Fused norm→FP4 / norm→FP8 front-ends — AdaLN and the pre-FFN LayerNorm emit FP4 directly (bit-identical to the two-step chain); the ViT / VL-self-attn LayerNorms emit FP8 directly, dropping the fp16 intermediate's write and read-back.
• Vectorized backbone small kernels — 16-byte-load rewrites of the per-head RMSNorms, LayerNorms, static FP8 quantize, split-half RoPE, GQA head expand, and residual add. The scalar originals issue 2–4 B accesses and sit far under the LPDDR5X floor. Element math matches; RoPE / quantize / head-expand / residual are bit-exact.
• Masked-softmax MHA — the softmax touches only the valid S_kv columns, so the per-layer full-buffer -inf logits pre-fill disappears; the DiT self-attention also reads the fused QKV GEMM output in place (token stride 3D), dropping three split copies per layer.
• FA4 attention — the three backbone MHA sites (ViT multi-view, LLM causal GQA with pack_gqa, VL-self-attn) run the vendored FlashAttention-4 (CuTe-DSL) forward, falling back to the existing fmha/cuBLAS chain when the thor-fa4 runtime deps are absent.

Shared-code changes

Only two files outside groot_n17 carry behavior:

• flash_rt/api.py — use_fp4=True with config="groot_n17" on Thor routes to the new frontend and falls back to the FP8 tier when flash_rt_fp4 is unavailable. Pi0.5's use_fp4 routing is untouched.
• flash_rt/hardware/thor/fa4_backend.py — the Thor chip string now follows the installed nvidia-cutlass-dsl (4.4.x only accepts sm_110a; 4.5+ needs the sm_101a alias). Previously the loader always requested sm_101a, so on 4.4.x FA4 failed to compile and every caller silently fell back to fmha. Call sites: LingBot (flash_rt/models/lingbot/kernel_ops.py) and Cosmos3-Edge (flash_rt/models/cosmos3_edge/layer_ref.py) — both keep their existing behavior on 4.5+ and now get the working FA4 path on 4.4.x.

The shared build surface also changes, in gated form: two new compilation units (csrc/kernels/attention_mha_masked.cu, csrc/kernels/vec_fp16_backbone.cu) join flash_rt_kernels, but only on SM100-class builds — CMake adds them under ENABLE_SM100_CUTLASS and defines FLASHRT_HAVE_THOR_VLA_KERNELS, and their bindings in csrc/bindings.cpp are guarded by the same define, so source and binding drop together on SM8x. The NVFP4 frontend probes for those bindings up front and fails naming the build flag rather than dying inside graph capture.

Everything else is additive: new csrc kernels with new symbols, a new frontend subclass, and gated branches in the N1.7 pipeline.

Tests

tests/test_groot_n17_thor_fp4_kernels.py — 16 contracts covering the NVFP4 GEMM epilogues against a torch reference, the bf16 activation quantizer, the fused DiT norms (bit-exact vs the two-step chain), the vectorized backbone helpers, and masked-softmax MHA against the pre-filled variant with a poisoned logits scratch. The masked softmax is additionally pinned at S_kv = 1024 / 1025 / 2048 in both dtypes against torch SDPA, since 1024 is where its register-tiled path hands over to the multi-pass one. Skips cleanly without CUDA or the optional flash_rt_fp4 extension.

tests/test_groot_n17_thor_fp4_routing.py — 5 contracts on the public load_model behavior the tier adds: use_fp4=True selects the NVFP4 frontend, a missing or NVFP4-less flash_rt_fp4 falls back to FP8, use_fp4 with use_fp16 is rejected, and the default route is unchanged. Frontends are stubbed, so these need no GPU or checkpoint.

Regression run on Thor (tests/test_groot_n17_*, test_thor_groot_attn_backend, test_build_inventory, test_load_model_use_fp8_kwarg): 156 passed, 36 skipped, 1 pre-existing environment failure (test_e2e_action_cosines_with_hf_noise needs a newer transformers for denormalize_action; it fails the same way on main).

Reproduce

The A/B/A above is three runs of the same command with the tier flipped;
--ref-out / --ref-in carry the FP8 output forward so the FP4 run reports
the cross-tier action cosine, and --min-cosine turns it into a gate:

# A: FP8 tier, save its actions as the precision reference
python benchmarks/groot_n17_thor_latency.py \
    --ckpt <n17-libero-checkpoint-dir> --aux <1-camera-aux-fixture.pt> \
    --tier fp8 --views 1 --embodiment libero_sim \
    --warmup 5 --iters 20 --ref-out fp8_actions.pt

# B: NVFP4 tier, compared against it (exits non-zero below the threshold)
python benchmarks/groot_n17_thor_latency.py \
    --ckpt <n17-libero-checkpoint-dir> --aux <1-camera-aux-fixture.pt> \
    --tier fp4 --views 1 --embodiment libero_sim \
    --warmup 5 --iters 20 --ref-in fp8_actions.pt --min-cosine 0.999

# A': FP8 again, to show the pair is not regime drift
python benchmarks/groot_n17_thor_latency.py \
    --ckpt <n17-libero-checkpoint-dir> --aux <1-camera-aux-fixture.pt> \
    --tier fp8 --views 1 --embodiment libero_sim --warmup 5 --iters 20

tests/_helpers/groot_n17/capture_aux_multi.py --views N captures a fixture with a restricted camera count, which is how the one-camera fixture above was produced.

Base

Rebased onto current main (after #159 / #160 / #161). The overlap with those is docs and registration only — README.md, docs/stable_api.md, CMakeLists.txt, flash_rt/api.py — and the rebase was clean. #160's RTX backbone-graph work is untouched; this PR only refreshes the GROOT N1.7 benchmark table's RTX 5090 row to the 16.61 ms full-graph number that #160 produced, since that table still carried the pre-graph 22 ms.

Docs

README.md, USAGE.md, docs/stable_api.md, and docs/benchmark_comparison.md record the tier, its flag, the precision ladder, and the measured numbers with the harness definition. docs/benchmark_comparison.md also carries NVIDIA's published TensorRT reference table for the same model family on the same board, alongside the camera-count-matched FlashRT rows and the 2-view base-checkpoint rows, each with its harness stated.