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