pi05 thor update: NVFP4 encoder/decoder tier + FA4 attention - #164
Merged
Merged
Conversation
The scalar qkv_split_rope_kvcache_fp16 and quantize_fp4_dynamic_sfa_fp16 kernels issue one 2-byte load per thread, which starves memory-level parallelism: the encoder-sized split-RoPE runs at ~50 GB/s and costs 85 us per layer. Add bit-exact vectorized variants that move 16 bytes per thread (int4 loads, 8-byte packed stores) and wire them into the Pi0.5 FP4 encoder and decoder paths. Both variants validate alignment and return nonzero without launching so callers can keep the scalar kernels on unaligned shapes; the fixed-shape/devpos path is unchanged. Measured on Thor SM110 (3-view FP4+FP4, CUDA-graph pipeline, nsys): encoder split-RoPE 84.5 -> 21.8 us per layer, decoder split-RoPE 2.3 -> 1.6 us, decoder O-input quantize 3.6 -> 3.3 us; total GPU time per inference 40.5 -> 39.4 ms. Kernel outputs are bit-identical to the scalar implementations on all production shapes.
Same 16-byte-load / 8-byte-store treatment as the split-RoPE and activation-quantize kernels, applied to the decoder GeGLU + NVFP4/SFA kernel. Arithmetic and evaluation order are unchanged, so outputs are bit-identical to gate_geglu_fp4_sfa_v2_fp16; unaligned buffers return nonzero without launching. Measured on Thor SM110 in the 3-view FP4+FP4 pipeline: 4.19 -> 3.46 us per launch across 180 launches per inference.
Quantize the 17 encoder attention output projections to NVFP4 (per-block MSE weight scales, dynamic FP4 activation quantize) behind the explicit use_fp4_encoder_attn opt-in; the strict E2E harness enables it by default for the FP4 child. QKV projections deliberately stay FP8: 4-bit Q/K weights shift the attention logits enough to break the raw-cosine acceptance gate (worst sample 0.97 versus the required 0.995), while O-only passes every gate with margin (matched-noise 3-view raw cosine 0.9987 / worst 0.9963, final-action 0.99951 / worst 0.99873 against the FP8 reference). The last encoder layer runs no attention, so it keeps no O weight, and unsupported configurations raise instead of silently selecting FP8.
Formal 2/3-view acceptance at 1ee2d37: FP4+FP4 30.277/35.877 ms p50 against same-run FP8 38.584/46.372 ms, all gates passed. Documents the sustained-load clock-regime drift that moves absolute milliseconds between runs and the like-for-like comparison rule it imposes.
Run the SigLIP vision-tower FFN in NVFP4 behind the use_fp4_siglip_ffn opt-in (harness default for the FP4 child). Three new pieces: - layer_norm_fp4_sfa_fp16: fused gamma/beta LayerNorm + NVFP4/SFA quantize, bit-identical to a fp16 LayerNorm followed by the vectorized quantize kernel. - cutlass_fp4_gemm_bias_gelu_fp4out: block-scaled NVFP4 GEMM with a fused per-channel bias + tanh-GELU + fp4/SFA epilogue, so the Up output feeds the Down GEMM without an fp16 round-trip or a separate quantize launch. - cutlass_fp4_gemm_bias_res_fp16: block-scaled NVFP4 GEMM with fused per-channel bias + residual accumulate (C may alias D). The FFN hidden dimension (4304) is zero-padded to 4320 for the 32-element fp4 TMA alignment; pad rows and columns carry zero weights and biases so padding is mathematically inert. Both GEMMs match a chain built from the proven fp16-out NVFP4 GEMM plus reference activation code to cosine 0.99993 on random inputs. Quantizing all 27 layers pushes the worst-sample raw cosine below the 0.995 gate (0.9889, and 0.9936-0.9943 at 20/23 layers), so the validated preset keeps the first 16 layers in NVFP4 and the rest in FP8; layers absent from the weight set keep the FP8 FFN path. Matched-noise 3-view fidelity for the preset against the FP8 reference: raw cosine 0.9987 / worst 0.9969, final action 0.99967 / worst 0.99923.
Formal 2/3-view acceptance at 1e07ea2: 30.076/37.442 ms p50 against same-run FP8 38.616/49.401 ms, speedups 1.2839/1.3194, all gates passed.
Expose per-projection decoder GEMM variants and switch qkv, o, and down from v7 to the v10 tile (128x64x256), matching gate_up. At M=10 these GEMMs are launch-wave-bound: the narrow-N tile doubles the CTA count (o and down run 8 CTAs under v7). Per-kernel nsys times inside the CUDA-graph pipeline, cross-checked against the unchanged gate_up kernel as a clock-regime canary: qkv 10.6 -> 10.1 us, o 9.4 -> 9.0 us, down 14.3 -> 13.5 us, about 0.3 ms per 3-view inference in total.
2-view 29.782 ms p50 (same-regime speedup 1.2953, matching the kernel-level prediction); 3-view 34.922 ms p50 in the fast clock regime with the cross-regime speedup caveat spelled out. All gates passed.
Add activation-aware weight quantization for the SigLIP FFN Up projections. During multi-sample calibration an eager FP8 SigLIP pass snapshots the FFN LayerNorm output at every NVFP4 layer; the per-channel amax is percentile-reduced across samples, the Up weights are rescaled and requantized in place, and the fused LayerNorm kernel gains a per-channel inverse-scale multiply (identity until calibration) so the GEMM still computes the unscaled product. The Down projections keep plain MSE quantization: the GELU nonlinearity blocks folding an activation scale through the Up output. Without AWQ, quantizing all 27 layers left the worst-sample raw cosine at 0.9889 against the 0.995 gate (0.9936-0.9943 at 20/23 layers); with it the full 27-layer preset measures raw 0.99906 / worst 0.99833 and final action 0.99963 / worst 0.99900 on the matched-noise 3-view comparison, so the default preset moves from 16 to all 27 layers.
2-view 29.587 ms p50 (speedup 1.3057) and 3-view 34.864 ms p50 (speedup 1.3675) against same-run FP8, both children regime-stable across all 100 samples, fidelity better than the 16-layer preset on every gate.
…s inert The tile-interleaved SF layouts round K up to 64-element atoms while the quantize kernels and fp4out epilogues only write entries for real (row, block) coordinates. With K not a multiple of 64 (SigLIP FFN H_pad=4320) the padding entries kept allocation garbage, which can decode as UE4M3 NaN and poison the block-scaled GEMM accumulator. Whether this fired depended on allocator history: reusing freed weight-staging memory produced NaN vision embeddings end to end. Zero scales make the padding contribute exactly zero. Init-time cost only; healthy runs are bit-identical.
Quantize the encoder attention QKV weights (input_layernorm folded, GQA interleave, all 18 layers) to NVFP4 and replace the FP8 RMSNorm+GEMM pair with a fused weightless-RMSNorm x AWQ-inverse-scale -> NVFP4 quantize kernel feeding a block-scaled GEMM. Plain 4-bit Q/K weights shift the attention logits past the raw-cosine gate, so the QKV path relies on activation-aware requant: per-channel amax is collected at the QKV input during multi-sample calibration and the weights are requantized in place at stable addresses, keeping the captured graph valid. Full-pipeline precision with SigLIP FFN NVFP4 and all prior FP4 stages enabled: raw cos 0.9987 (worst sample 0.9966), action cos 0.9995 (worst 0.9989) against the FP8 reference.
Back-to-back formal 3-view runs show the QKV NVFP4 path passes every gate (raw min 0.9966, action min 0.9989) but costs ~0.36 ms: at encoder sequence length the QKV GEMM is compute-bound, so NVFP4 only matches FP8 while the fused quantize kernel adds ~18 us per layer. Default the bench flag to 0 and document the negative result.
… descriptor
The SM110 tcgen05 block-scaled MMA decodes its 3-bit element-format
descriptor field at run time; value 0 selects a sign-magnitude uniform
INT4 grid (E0M3, magnitudes 0..7), validated element-level with distinct
payloads including negative codes. The uniform grid removes E2M1's
non-uniform rounding bins, lowering weight quantization error.
Adds, all additive and parameter-gated (default path unchanged):
- quantize_e0m3_dynamic_sfa_fp16: per-16 UE4M3-scaled E0M3 quantizer,
bit-exact against a host reference (IEEE division intrinsics so
--use_fast_math cannot perturb fp8 rounding ties)
- cutlass_fp4_gemm_e0m3w: runtime-datatype GEMM issuing E2M1 activations
against E0M3 weights on the production decoder tile (128x64x256);
cosine vs fp32 reference 0.999998+ on all four decoder shapes
- decoder_weight_format frontend kwarg and --decoder-weight-format bench
flag ('nvfp4' default | 'e0m3'); the decoder pipeline routes all four
projections through the E0M3 runner when selected
…tation Extends the uniform-INT4 decoder path to the activation side. New additive kernels mirror the four production activation quantize exits (AdaRMS entry, gated-residual AdaRMS, attention-context vec, GeGLU vec) with the E0M3 encoder, and optionally rotate each 16-value block by the orthonormal 16x16 Hadamard matrix before quantization; the same rotation applied to the weight blocks offline leaves the GEMM mathematically unchanged (validated end to end: rotated and unrotated W4A4 products match the fp16 reference equally). The lane-parallel kernels implement the rotation as a 4-stage shfl_xor butterfly, the vectorized kernels as an in-register FWHT. The e0m3w GEMM gains an a_format argument to issue E0M3 activations against E0M3 weights. Parameter-gated end to end: decoder_act_format and decoder_rht frontend kwargs, --decoder-act-format / --decoder-rht bench flags; defaults leave every existing path untouched.
The gated-residual AdaRMS that quantizes the next layer's QKV input was still emitting E2M1 codes when decoder_act_format='e0m3', so the E0M3 GEMM misdecoded one of five activation exits (and with RHT the unrotated activations met rotated weights). All five exits now follow the activation-format dispatch.
The rotated W4A4 INT4 decoder beats the NVFP4 baseline on every cosine metric (raw min 0.99904 vs 0.99833) at ~0.35 ms.
Add a forked SM100 block-scale row-store visitor whose visit() computes gelu(gate)*up on adjacent accumulator column pairs of a pairwise interleaved gate/up weight, and a GEMM runner wired to it. One interleaved GEMM replaces the separate gate/up fp4out GEMMs plus the GeGLU combiner kernel; the down projection consumes the full-width duplicated output through a K-expanded weight (odd columns zero). The down-projection AWQ inv_s is folded into the up weight rows at quantization time (algebraically exact), so the epilogue needs no per-column vector. Parameter-isolated behind encoder_p1_combiner='epilogue'; default paths unchanged.
Pipeline A/B on Thor inverts the isolated-benchmark ranking for the K=16384 Down: cluster variants regress heavily under real cache pressure (2x1 +2.2 ms e2e, 2x2/2x4 +11-14 ms) while the plain 128x256x128 tile runs the expanded Down at baseline-Down cost, turning the fused-GeGLU path into a net win.
Add a second store-node variant that quantizes gelu(gate)*up at compact granularity (16 unique values per scale block, combiner-equivalent) and writes the packed FP4 + SFA buffers directly from the visitor, so the down projection keeps its original K and weight; the collective's D path lands in one small reusable dummy buffer. The thread's global base coordinate is recovered as problem extent minus the per-thread residue (the coordinate tensor itself is thread-relative), and quantization uses the hardware e2m1 converter (the branch-ladder version cost 2.9x kernel time). Parameter-isolated behind encoder_p1_combiner='epilogue_hw'; defaults unchanged.
Formal same-batch A/B (identical FP8 references): 3v 32.25 vs 34.30 ms (speedup 1.4355 vs 1.3478), 2v 27.74 vs 29.13 ms, with better raw and action cosines on both view counts (single quantization). lut_native remains available for comparison.
Two decoder-side launch/kernel eliminations, parameter-isolated and default-off: - attention_qkv_fp16_seqused_v2 folds the seqused -inf mask into the softmax kernel (positions beyond the valid length become exact zero probabilities, matching the reference mask+softmax result), removing one kernel per attention call. Selected via the attention backend's use_fused_softmax attribute (decoder_fused_attn). - A skinny-M instantiation of the compact GeGLU store epilogue on the decoder GEMM tile replaces the gate_up GEMM + GeGLU-quantize kernel pair with one interleaved GEMM writing the down-projection input directly (decoder_fused_geglu, nvfp4 weights only; the interleaved weight is MSE-quantized like the other decoder projections).
Formal same-batch A/B (matched FP8 references): 3v 32.26 -> 31.64 ms (speedup 1.4417 -> 1.4726), 2v 27.70 -> 27.14 ms (1.3917 -> 1.4169), with equal-or-better cosine floors on both view counts.
… fold The fold used a scalar register-to-column mapping while the reference softmax kernel pairs columns through __half2, so the warp reduction summed in a different order and the attention output drifted by one to two fp16 ulp. Mirroring the reference mapping makes the fold's output bit-identical to the mask + softmax chain, which is the whole point of the change: it is a launch-count optimization, not a numerics change.
Contract tests for the kernels this branch adds: the seqused softmax fold must reproduce the mask + softmax chain exactly, the fused GeGLU epilogue must be at least as accurate as the gate/up GEMM + combiner chain it replaces, and the vectorized SigLIP LayerNorms must agree with the reference norm + quantize pair.
Only the fixed-shape state-prompt path routes decoder attention through the seqused kernels; the strict FP4 E2E suite passes no state, so it takes the plain attention path where no separate mask kernel exists and the fold has nothing to remove. A single-frame kernel trace confirms it never runs there. The fold stays available for the fixed-shape path, but it ships off since no exercised configuration measures a gain, and the end-to-end improvement is attributed to the fused GeGLU FFN alone.
The ported series renamed every sm_101a mention to sm_110a, which was right when it was written but no longer matches the loader: it now picks the chip string from the installed nvidia-cutlass-dsl (sm_101a on 4.5+, sm_110a on 4.4.x). Restore main's LingBot docs, example, comments and pyproject note so they do not assert a fixed target the code does not use. No functional change.
Two issues the GROOT N1.7 review surfaced in the same kernel class, found here by inspection: - The seqused softmax fold capped at SMV2_MAX_COLS columns with no bound check, so an S_kv_max above 1024 left the tail of every logits row unnormalized while the PV GEMM still consumed it — a silent wrong answer on a public binding. Rows past the register tile now use a multi-pass kernel that holds no per-column registers, so any S_kv_max is correct; pinned at 1024 / 1025 / 2048 against torch SDPA. The binding docstring no longer advertises a limit that no longer exists. - attention_seqused_fused.cu and rope_vec.cu were compiled into the shared flash_rt_kernels module on every architecture. They join the existing SM100-class group (FLASHRT_HAVE_THOR_VLA_KERNELS) alongside the N1.7 helpers, and their bindings are guarded by the same define so source and binding drop together on SM8x.
The ported Thor NVFP4 work added use_fp4_decoder, use_fa4, encoder_p1_combiner, encoder_down_variant and decoder_gate_up_variant to load_model, and changed awq_alpha's default from 0.5 to None (which resolves to the per-stage production value). docs/stable_api.md now carries all of them, as the contributing guide requires for public API changes.
The ported series carried a developer's checkpoint, fixture, interpreter and output directories as literal defaults in tests/bench_pi05_decoder_fp4_e2e.py and in the reproduction command of docs/pi05_thor_decoder_fp4_e2e.md. The harness now takes the checkpoint and fixture directory from --checkpoint / --fixture or the PI05_CHECKPOINT / PI05_FIXTURE_DIR environment variables and fails with a clear message when neither is given; the doc shows a placeholder command and states the FA4 runtime requirements in prose instead of a machine-specific path.
…mbers Same-session A/B against the FP8 path on Jetson AGX Thor (locked clocks, MAXN, LIBERO fixtures, 10 denoise steps, 100 iterations after 20 warmup), via tests/bench_pi05_decoder_fp4_e2e.py: 1 view FP8 32.92 -> NVFP4 23.01 ms (1.43x) 2 views FP8 38.70 -> NVFP4 27.17 ms (1.42x) 3 views FP8 49.02 -> NVFP4 31.74 ms (1.54x) The 2- and 3-view runs clear every precision gate (action cosine 0.99972 and 0.99974, worst sample 0.99916 and 0.99944). The 1-view row is marked: its worst-sample action cosine is 0.971, which the docs already characterize as a flow-matching bifurcation rather than a kernel defect, with a configuration that does clear the gates recorded alongside.
Member
Author
|
@heiheiha798 Nice tile push |
LiangSu8899
force-pushed
the
pi05-thor-update
branch
from
August 5, 2026 22:33
7c6ef59 to
5993751
Compare
The measured configuration was only reachable by constructing the frontend directly: use_fp4_encoder_attn and use_fp4_siglip_ffn were not exposed at all, and encoder_p1_combiner defaulted to lut_native while the benchmark ran the fused GeGLU epilogue. load_model(use_fp4=True, use_fp4_decoder=True) now resolves the full tier the latency table is measured with, the same way awq_alpha already keyed off use_fp4_decoder, and every sub-flag remains individually overridable. The benchmark harness gains --construct load_model (the new default), which builds both processes through the public API and refuses to run when any sweep knob deviates from the published preset, so a published number cannot come from a configuration load_model does not produce. The recorded result.json carries the construction mode and the exact call. use_fa4 stays independent of use_fp4: both the FP8 frontend and the NVFP4 subclass accept it, and the FP8 baseline behind the published speedups runs FA4 as well, so the comparison isolates NVFP4 rather than the attention backend. The stable-API doc claimed it required use_fp4; that was wrong and is corrected. Its config/framework/hardware validation also moves ahead of pipeline-class resolution so an unsupported combination fails before a frontend import. Also: check the NVFP4 GEMM return codes in the FP4 encoder FFN path, which report can_implement/initialize/allocate/run failures through a status code rather than an exception; drop the layer_norm_fp4_sfa_fp16 binding, which had no runtime caller and duplicates layer_norm_mul_fp4_sfa_fp16 with a null inverse scale; and record the actual SigLIP layer count in the benchmark result instead of a stale 16-layer label.
The headline table keeps the measured 23.01 / 27.17 / 31.74 ms. A second session re-ran all three view counts through the harness default --construct load_model, which builds both children with the public load_model(); every cosine came out identical at every view count, so the two construction paths demonstrably produce the same configuration, and the FP4 latencies land within 0.15 ms (23.07 / 27.29 / 31.89 ms). The FP8 baseline drifts a few ms between sessions on this hardware, which moves a ratio without moving the FP4 result, so the doc states that only within-row ratios are meaningful. The 1-view alternative configuration and the encoder-QKV NVFP4 flag deviate from the published preset, so their commands now carry --construct frontend.
LiangSu8899
force-pushed
the
pi05-thor-update
branch
from
August 5, 2026 22:37
5993751 to
ab45ac1
Compare
This was referenced Aug 18, 2026
DXICM
added a commit
to DXICM/FlashRT
that referenced
this pull request
Aug 18, 2026
cutlass-dsl caches the device arch at import time. The previous code imported cutlass to check its version, then set CUTE_DSL_ARCH=sm_101a — too late; NVVM already cached sm_110a and ICEs on the hd256 2CTA kernel (introduced in flashrt-project#164, commit 7fd75d2). Fix: set CUTE_DSL_ARCH=sm_101a unconditionally before any cutlass import. Also revert the hd256 2CTA dispatch to SM100-only (the dedicated kernel was never validated on SM110) and restore the _fa4_trimmed lazy loader for BlackwellFusedMultiHeadAttentionForward. Verified: all-tier E2E on Thor — median 27.7 ms, p95 28.5 ms, actions finite, cos 0.999933 vs HF eager.
LiangSu8899
added a commit
that referenced
this pull request
Aug 24, 2026
…177) * fix(groot): 12 HF-alignment bugs for N1.6 Thor (SM110) frontend Root-cause and fix 12 real bugs where the upstream N1.6 frontend inherited openpi-family (Pi0/Pi0.5) vision/kernel assumptions that do not hold for GR00T N1.6's HF behaviour: 1. Tokenization: reproduce Eagle chat template (system/user headers, formalize, per-view image blocks) instead of bare encode() 2. Resolution: HF eval chain outputs 252x252, not 224 3. SigLIP attention scope: HF(sdpa) does cross-view full attention on the packed 648-token sequence, not per-view 4. Patch flatten order: HF NaFlex uses (ph,pw,C), not (C,ph,pw) 5. Strided FMHA divergence on non-power-of-2 seq with real data: parity mode routes SigLIP attention through torch sdpa 6. CKernelQwen3 diverges from HF on real sequences: parity mode runs HF-native Qwen3Model (bf16, sdpa, graph-captured) 7. Wild pointer after re-capture: Qwen3 graph-captured LN referenced local tensors; promote to persistent attributes + finiteness guard 8. adaLN chunk order reversed: HF proj_out_1 is (shift, scale) 9. Single-frame FP8 calibration too narrow: multi-frame calibrate (current + 7 synthetic frames, percentile=99.9) 10. Prompt switch rejected after graph bake: detect change, reset graph runtime, re-set prompt, re-capture 11. Idle-first-frame garbage: Thor GPU idle reset invalidates captured graphs; add replay finiteness self-check + re-capture retry 12. Prompt-switch re-capture device-side assert: stale DiT static buffers/indices not rebuilt; add to stale list Precision vs HF eager: cos 0.999933 / maxd 0.059 (denormalized action). No inference hyperparameters changed (4-step, 252x252, T=50, bf16). Also adds tools/convert_groot_n16_hf_checkpoint.py for HF safetensors to FlashRT layout conversion (Qwen3 16-layer truncation, DiT repack, SigLIP mlp1 layout). * perf(groot): FA4 + NVFP4 full-kernelization (130 -> 28.5 ms) New CUDA kernels for the N1.6 Thor NVFP4 pipeline: - fused_fp4/silu_mul_fp4_sfa_bf16: SiLU(gate)*up (bf16) direct to NVFP4+SFA, bit-exact vs torch two-step chain - fused_fp4/dit_norm_fp4_sfa: AdaLN / no-affine LN / weighted RMSNorm direct to NVFP4+SFA (bf16 input variants) - gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100: bias / bias+residual / bias+tanh-GELU+fp4out epilogue variants - quantize/quantize_fp4_sfa_bf16: vectorized bf16 dynamic quantize - kernels/qk_norm_rope_rotate_half_bf16: fused per-head RMSNorm + rotate-half RoPE (bf16, in-place, one launch per Q/K) Performance rounds (no hyperparameter changes): - DiT NVFP4 fused epilogue: 36.6 -> 15.7 ms (8 kernels/layer) - Qwen3 fused norm/rope/GQA: 12.7 -> 5.0 ms (cos 0.999986) - SigLIP FA4 + fp4 encoder: 10.3 -> 6.9 ms (cos 0.999988) - SigLIP embeddings in-graph: 34 -> 28.5 ms (bit-exact) - E2E total: 130 -> 28.5 ms (4-step, 2-camera, 252x252, T=50) Bandwidth ceiling: Thor measured 252-255 GB/s (~93% of 273 spec); DiT 15.2 ms is weight-bandwidth-bound floor for this config. Tier switches (all default ON, independently fall back): FLASHRT_N16_DIT_FP4, FLASHRT_N16_QWEN3_FP4, FLASHRT_N16_SIGLIP_FP4, FLASHRT_N16_FA4 * docs(groot): N1.6 authoritative adaptation doc + companion docs - docs/groot_n16_thor_sm110.md: single authoritative document covering architecture facts, 12-bug root-cause table, falsified hypotheses, full optimization record (130 -> 28.5 ms), roofline/bandwidth ceiling analysis (252-255 GB/s, ~93% of spec), precision tier switches, and verification methodology. - docs/groot_transformers5_weight_corruption.md: transformers>=5 silent weight corruption via _initialize_missing_keys re-randomizing SigLIP2 vision tower (282 tensors). One-line fix + integrity guard. - docs/thor_gpu_idle_reset_workaround.md: Thor GPU idle reset defect and three-layer CUDA Graph protection (keepalive, idle reinit, finiteness). * fix(fa4): set CUTE_DSL_ARCH before cutlass-dsl import (SM110 NVVM ICE) cutlass-dsl caches the device arch at import time. The previous code imported cutlass to check its version, then set CUTE_DSL_ARCH=sm_101a — too late; NVVM already cached sm_110a and ICEs on the hd256 2CTA kernel (introduced in #164, commit 7fd75d2). Fix: set CUTE_DSL_ARCH=sm_101a unconditionally before any cutlass import. Also revert the hd256 2CTA dispatch to SM100-only (the dedicated kernel was never validated on SM110) and restore the _fa4_trimmed lazy loader for BlackwellFusedMultiHeadAttentionForward. Verified: all-tier E2E on Thor — median 27.7 ms, p95 28.5 ms, actions finite, cos 0.999933 vs HF eager. * fix(groot): close PR 177 audit gaps * fix(groot): keep parity routing coherent --------- Co-authored-by: LiangSu8899 <7thuniversels@gmail.com>
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
2 participants
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.
Summary
Brings the Pi0.5 Thor NVFP4 work up to the current production configuration. The FP8 path is unchanged; everything new sits behind
use_fp4=Trueand its sub-flags.Same-session A/B against the FP8 path on Jetson AGX Thor (locked clocks, MAXN, LIBERO fixtures, 10 denoise steps, medians over 100 iterations after 20 warmup), measured with the committed harness:
Each row is one same-session A/B; rows are separate sessions, and the FP8 baseline drifts a few ms between sessions on this hardware, so only within-row ratios are meaningful. Both children run FA4, so the speedup isolates NVFP4 rather than the attention backend.
A second session re-ran all three view counts through the harness default
--construct load_model, which builds both children with the publicload_model(). Every cosine came out identical at every view count and the FP4 latencies landed within 0.15 ms (23.07 / 27.29 / 31.89 ms), so the two construction paths demonstrably produce the same configuration. Both tables and the artifact hashes are indocs/pi05_thor_decoder_fp4_e2e.md.The 2- and 3-view runs clear every gate the harness enforces (raw and action cosine, per-sample floors, and the latency targets). † At 1 view the per-sample gate does not pass;
docs/pi05_thor_decoder_fp4_e2e.mdcharacterizes this as a flow-matching bifurcation rather than a kernel defect — with single-view input the velocity field itself is near a decision boundary on some samples — and records an encoder-FP8 + rotated-INT4-decoder configuration that does clear the gates at 1 view. The README row is marked accordingly.What is in the port
csrc/fp4_utils): the tile-interleaved SF layout rounds K up to 64-element atoms, and the quantize kernels only write real coordinates, so allocation garbage in the padding could decode as UE4M3 NaN and poison a block-scaled accumulator. This one is not Pi0.5-specific — any NVFP4 caller with a non-multiple-of-64 K was exposed.Public API
use_fp4_decoder=Trueselects the complete tier the table above is measured with, resolving the sub-flags left atNonethe same wayawq_alphaalready keyed off it:That resolves to
use_fp4_encoder_attn=True,use_fp4_siglip_ffn=True,encoder_p1_combiner="epilogue_hw"andawq_alpha=0.8.use_fp4=Truealone keeps the earlier encoder-only preset (False,False,"lut_native",0.5). Any sub-flag passed explicitly still overrides the preset.use_fa4is an attention-backend choice and is deliberately independent ofuse_fp4: both the FP8 frontend and the NVFP4 subclass accept it, and the FP8 baseline above runs it too.Behavior changes to call out
use_fp4=Truenow resolves its encoder preset totuple(range(17))instead ofrange(18), andawq_alphadefaults toNone(0.8 with the FP4 decoder, 0.5 otherwise) instead of a flat 0.5.use_fp4_decoder=Trueadditionally resolvesencoder_p1_combinerto"epilogue_hw"(was"lut_native") and turns on the NVFP4 encoder attention output and SigLIP FFN, which is what makes the published latency reachable fromload_model().load_modeloptions:use_fp4_decoder,use_fp4_encoder_attn,use_fp4_siglip_ffn,use_fa4,encoder_p1_combiner,encoder_down_variant,decoder_gate_up_variant, and a read-onlyVLAModel.pipelineaccessor. Documented indocs/stable_api.md.Fixes made while porting
S_kv_max > 1024was silently wrong in the seqused softmax fold: the register-tiled kernel caps at 1024 columns with no bound check, so wider logits rows kept an unnormalized tail that the PV GEMM still consumed, while the public binding accepted any length and its docstring advertised the limit. Rows past the tile now use a multi-pass kernel that holds no per-column registers, so anyS_kv_maxis correct; pinned at 1024 / 1025 / 2048 against torch SDPA.attention_seqused_fused.cuandrope_vec.cuwere compiled into the shared module on every architecture. They join the existing SM100-class group (FLASHRT_HAVE_THOR_VLA_KERNELS), and their bindings are guarded by the same define so source and binding drop together on SM8x.--checkpoint/--fixtureor$PI05_CHECKPOINT/$PI05_FIXTURE_DIR, with a clear error when neither is set), and the doc's reproduction command and two "handoff reference in private issue" mentions are gone.sm_101amention tosm_110aacross LingBot docs, example, comments andpyproject.toml. That is no longer accurate — the FA4 loader picks the chip string from the installednvidia-cutlass-dsl— so those files are restored to main.Review follow-ups
load_model().use_fp4_encoder_attnanduse_fp4_siglip_ffnwere not exposed at all, andencoder_p1_combinerdefaulted tolut_nativewhile the harness ran the fused GeGLU epilogue. Both are now public options anduse_fp4_decoder=Trueresolves the measured preset. The harness gained--construct load_model(the new default), which builds both children through the public API and refuses to run when any sweep knob deviates from the preset;--construct frontendremains for the exploratory knobs and is labelled as not producing public-API numbers.result.jsonrecords the mode and the exact call. Re-measured in that mode, the action and raw cosines came out identical at every view count to the direct-construction runs the table above reports, which is the evidence the two build the same thing.use_fa4did not match its documented contract. The doc claimed it requireduse_fp4=True; the code allowed FP8 + FA4. The code was right — that combination is the measured baseline — so the doc is corrected and the independence is now stated indocs/stable_api.md, theload_modeldocstring and a comment at the check. Its config/framework/hardware validation also moved ahead of pipeline-class resolution, so an unsupported combination fails before a frontend import instead of after.RuntimeErrornaming the kernel, layer and M/N/K. (The fused elementwise bindings in the same block validate in C++ and return nothing; those are deliberately not wrapped, and the helper's docstring says so.)nvfp4_16_layerswhile the SigLIP default covers all 27 — now recorded as the actual count; thelayer_norm_fp4_sfa_fp16binding had no runtime caller and duplicatedlayer_norm_mul_fp4_sfa_fp16with a null inverse scale, so it is removed; the title now names the scope.Tests
tests/test_pi05_fp4_fusion_kernels.pyandtests/test_pi05_decoder_fp4_kernels.pypin the FP4 kernels against the unfused chains they replace, including the new wide-logits boundary cases.tests/test_pi05_thor_fp4_routing.py(new, 24 cases, no GPU needed) covers theload_modelcontract: FA4 reaching both frontends and being rejected elsewhere, each NVFP4 sub-flag requiringuse_fp4and Pi0.5-torch-Thor, explicit values overriding the preset, and — the important one — the resolved constructor arguments asserted against the harness's own preset table, so the published configuration and the public API cannot drift apart. Two cases also assert that the knobsload_modeldoes not forward still default to the published values in the frontend signature.Full-suite regression on Thor: 747 passed, 198 skipped. The 25 failures in that run are pre-existing and unrelated (qwen36 GDN numerics,
test_groot_n17_e2eneeding thegr00tpackage, and test-ordering artifacts inweight_loader/minimax/qwen3_vlthat pass in isolation); the identical set fails with this PR's new test file deselected.Reproduce
The harness runs the FP8 and NVFP4 children in one session, writes both action tensors, and fails on any gate. It requires a clean tracked worktree so the recorded commit identifies the build.
Base
Rebased onto current
main(after #163). The Pi0.5 series and the GROOT N1.7 tier share thecsrc/gemm/fp4andcsrc/fused_fp4trees and the SM100-class kernel group; the rebase conflicts were all additive list entries plus the FA4 loader's chip-string selection, where main's version-aware implementation was kept.