From 1bc2926983ed4917c42d6896621fc2e4241b8ccc Mon Sep 17 00:00:00 2001
From: DXICM <10598463@qq.com>
Date: Mon, 17 Aug 2026 08:26:20 +0000
Subject: [PATCH 1/6] 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).
---
flash_rt/frontends/torch/groot_thor.py | 1532 +++++++++++++++++-
flash_rt/hardware/thor/attn_backend_groot.py | 43 +-
flash_rt/models/groot/pipeline_thor.py | 79 +-
tools/convert_groot_n16_hf_checkpoint.py | 217 +++
4 files changed, 1762 insertions(+), 109 deletions(-)
create mode 100644 tools/convert_groot_n16_hf_checkpoint.py
diff --git a/flash_rt/frontends/torch/groot_thor.py b/flash_rt/frontends/torch/groot_thor.py
index 2d5cde0c..394c5b0c 100644
--- a/flash_rt/frontends/torch/groot_thor.py
+++ b/flash_rt/frontends/torch/groot_thor.py
@@ -13,6 +13,7 @@
import json
import logging
import math
+import os
import pathlib
import time
from typing import Optional
@@ -80,7 +81,8 @@ class GrootTorchFrontendThor:
"""GROOT N1.6 inference pipeline on Thor SM110."""
def __init__(self, checkpoint, num_views=2, autotune=3,
- embodiment_tag="new_embodiment", use_fp8=True):
+ embodiment_tag="new_embodiment", use_fp8=False,
+ image_size=252, parity=True):
"""Initialize GROOT pipeline.
Args:
@@ -88,6 +90,10 @@ def __init__(self, checkpoint, num_views=2, autotune=3,
num_views: camera views (default 2)
autotune: CUDA Graph autotune intensity (0=off, 3=default)
embodiment_tag: target embodiment for per-embodiment MLPs
+ image_size: SigLIP input edge (must be divisible by 14).
+ 224 = 16x16 patches (legacy); 252 = 18x18 patches, the
+ GR00T N1.6 training/eval resolution (HF processor chain
+ LetterBoxPad -> 256 -> 0.95 crop -> 252).
"""
if embodiment_tag not in EMBODIMENT_TAG_TO_INDEX:
raise ValueError(
@@ -98,8 +104,20 @@ def __init__(self, checkpoint, num_views=2, autotune=3,
self._num_views = num_views
self._autotune = autotune
self.use_fp8 = bool(use_fp8)
+ # Working dtype of the vision/LLM feature path: FP8 mode keeps the
+ # legacy fp16 buffers; parity mode stays bf16 end-to-end to match HF.
+ self._bd = (torch.float16 if (self.use_fp8 or not parity)
+ else torch.bfloat16)
+ # parity=True: HF-native torch modules (exact, ~110ms).
+ # parity=False: FlashRT kernel fast path (fp16/fp8 GEMM+FMHA, ~30ms).
+ self.parity = bool(parity)
self._embodiment_tag = embodiment_tag
self._embodiment_id = EMBODIMENT_TAG_TO_INDEX[embodiment_tag]
+ if image_size % 14 != 0 or (image_size // 14) % 2 != 0:
+ raise ValueError(
+ f"image_size must be divisible by 28 (got {image_size}); "
+ "the patch grid must be even for pixel_unshuffle(2).")
+ self.image_size = image_size
self._real_data_calibrated = False
self.calibrated = False
@@ -133,8 +151,8 @@ def __init__(self, checkpoint, num_views=2, autotune=3,
self.HD_sig = 72
self.H_sig = 4304
self.L_sig = 27
- self.spv_raw = 256 # raw patches per view (before pixel unshuffle)
- self.spv = 64 # tokens per view after 2x2 pixel_unshuffle (C4)
+ self.spv_raw = (image_size // 14) ** 2 # raw patches per view
+ self.spv = (image_size // 28) ** 2 # tokens per view after 2x2 pixel_unshuffle (C4)
self.mlp1_in = 4608 # 1152 * 4 after pixel unshuffle (C5)
# Qwen3 (★16 layers★, select_layer=16, checkpoint truncated) (C1)
@@ -158,7 +176,7 @@ def __init__(self, checkpoint, num_views=2, autotune=3,
self.action_dim = 128 # ★ NOT 29 — padded max, per-embodiment actual varies ★
self.state_dim = 128 # ★ NOT 29 ★
self.action_horizon = 50 # ★ NOT 16 — padded max ★
- self.num_steps = 4 # flow-matching Euler steps
+ self.num_steps = max(1, int(os.environ.get("FLASHRT_N16_DIT_STEPS", "4")))
self.Sa = self.action_horizon + 1 # 51 = 1 state + 50 actions
# ── Load checkpoint (keep full sd; SigLIP init deferred to _capture_all_graphs) ──
@@ -227,8 +245,18 @@ def _load_siglip2_weights(self, sd):
self._sig_patch_w = sd[f"{VIS_PREFIX}.embeddings.patch_embedding.weight"].to(fp16) # [1152, 588]
self._sig_patch_b = sd[f"{VIS_PREFIX}.embeddings.patch_embedding.bias"].to(fp16) # [1152]
- # Position embedding: [256, 1152] — 256 patches, NO CLS (C15)
- self._sig_pos_embed = sd[f"{VIS_PREFIX}.embeddings.position_embedding.weight"].to(fp16) # [256, 1152]
+ # Position embedding: [256, 1152] — 256 patches, NO CLS (C15).
+ # Resize 16x16 -> grid x grid to match HF resize_positional_embeddings
+ # (bilinear, align_corners=False, antialias=True, fp32).
+ pos = sd[f"{VIS_PREFIX}.embeddings.position_embedding.weight"].float()
+ grid = int(math.sqrt(self.spv_raw))
+ if grid != 16:
+ pos = (torch.nn.functional.interpolate(
+ pos.reshape(1, 16, 16, -1).permute(0, 3, 1, 2),
+ size=(grid, grid), mode="bilinear",
+ align_corners=False, antialias=True)
+ .permute(0, 2, 3, 1).reshape(-1, self.D_sig))
+ self._sig_pos_embed = pos.to(fp16) # [spv_raw, 1152]
# mlp1: LN(4608) → Linear(4608,2048) → GELU → Linear(2048,2048) (C5)
self._mlp1_ln_w = sd[f"{MLP1_PREFIX}.0.weight"].to(fp16) # [4608]
@@ -598,7 +626,7 @@ def _precompute_timesteps(self):
# flip_sin_to_cos=True, downscale_freq_shift=1
half_dim = 128 # 256 / 2
exponent = -torch.arange(half_dim, dtype=torch.float32, device='cuda') * \
- (math.log(10000.0) / half_dim)
+ (math.log(10000.0) / (half_dim - 1))
emb_freqs = exponent.exp() # [128]
t_values = [0, 250, 500, 750]
@@ -734,7 +762,7 @@ def _timestep_encode(self, t_disc):
"""Timesteps(256) → TimestepEmbedding → temb [1, 1536]."""
half_dim = 128
exp = -torch.arange(half_dim, dtype=torch.float32, device='cuda') * \
- (math.log(10000.0) / half_dim)
+ (math.log(10000.0) / (half_dim - 1))
t_tensor = torch.tensor([t_disc], dtype=torch.float32, device='cuda')
args = t_tensor[:, None] * exp.exp()
sincos = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(fp16)
@@ -752,6 +780,931 @@ def _state_encode(self, state):
h = self._fp16_gemm(h, self._state_enc_w2, 1, self.D_dit, 1024) + self._state_enc_b2
return h.unsqueeze(0) # [1, 1, D_dit]
+ def _fill_dit_kv(self):
+ """Compact valid backbone rows into the DiT cross-attention KV buffers.
+
+ HF excludes the other modality's backbone tokens from each DiT
+ cross-attention layer via an attention mask. Writing zero-masked
+ full-length rows instead would leave bias-only K/V rows that still
+ receive attention mass and corrupt the velocity field (the
+ flow-matching noise then fails to collapse). Gather only the valid
+ rows and report their counts so the kernels attend the right length.
+ """
+ txt = self._g_vlln_buf[self._g_non_img[0]] # (n_text, D) fp16
+ img = self._g_vlln_buf[self._g_img_m[0]] # (n_img, D) fp16
+ self._g_dit.b_kv_text[:txt.shape[0]].copy_(txt)
+ self._g_dit.b_kv_img[:img.shape[0]].copy_(img)
+ self._g_dit.set_kv_counts(txt.shape[0], img.shape[0])
+
+ def _setup_torch_dit(self):
+ """Extract action-head weights from ``_full_sd`` for the
+ HF-faithful torch DiT path. Call while ``_full_sd`` is alive.
+ Parity mode runs the DiT in fp32 (bf16 denoising amplifies
+ rounding into visible chunk wobble vs HF)."""
+ sd = self._full_sd
+ AH = "action_head"
+ eid = self._embodiment_id
+ # HF runs the DiT in bf16; parity mode matches it exactly.
+ dt = torch.bfloat16
+ self._dit_dt = dt
+ w = {}
+ ts = f"{AH}.model.timestep_encoder.timestep_embedder"
+ w['ts_l1_w'] = sd[f"{ts}.linear_1.weight"].to(dt)
+ w['ts_l1_b'] = sd[f"{ts}.linear_1.bias"].to(dt)
+ w['ts_l2_w'] = sd[f"{ts}.linear_2.weight"].to(dt)
+ w['ts_l2_b'] = sd[f"{ts}.linear_2.bias"].to(dt)
+ w['proj1_w'] = sd[f"{AH}.model.proj_out_1.weight"].to(dt)
+ w['proj1_b'] = sd[f"{AH}.model.proj_out_1.bias"].to(dt)
+ w['proj2_w'] = sd[f"{AH}.model.proj_out_2.weight"].to(dt)
+ w['proj2_b'] = sd[f"{AH}.model.proj_out_2.bias"].to(dt)
+ blocks = []
+ for l in range(32):
+ p = f"{AH}.model.transformer_blocks.{l}"
+ blocks.append({k: sd[f"{p}.{n}"].to(dt) for k, n in (
+ ('n1w', 'norm1.linear.weight'), ('n1b', 'norm1.linear.bias'),
+ ('qw', 'attn1.to_q.weight'), ('qb', 'attn1.to_q.bias'),
+ ('kw', 'attn1.to_k.weight'), ('kb', 'attn1.to_k.bias'),
+ ('vw', 'attn1.to_v.weight'), ('vb', 'attn1.to_v.bias'),
+ ('ow', 'attn1.to_out.0.weight'), ('ob', 'attn1.to_out.0.bias'),
+ ('fw', 'ff.net.0.proj.weight'), ('fb', 'ff.net.0.proj.bias'),
+ ('dw', 'ff.net.2.weight'), ('db', 'ff.net.2.bias'))})
+ w['blocks'] = blocks
+ # Per-embodiment MLPs (CategorySpecificLinear: x @ W[eid] + b[eid])
+ w['ae_w1'] = sd[f"{AH}.action_encoder.W1.W"][eid].to(dt)
+ w['ae_b1'] = sd[f"{AH}.action_encoder.W1.b"][eid].to(dt)
+ w['ae_w2'] = sd[f"{AH}.action_encoder.W2.W"][eid].to(dt)
+ w['ae_b2'] = sd[f"{AH}.action_encoder.W2.b"][eid].to(dt)
+ w['ae_w3'] = sd[f"{AH}.action_encoder.W3.W"][eid].to(dt)
+ w['ae_b3'] = sd[f"{AH}.action_encoder.W3.b"][eid].to(dt)
+ w['se_w1'] = sd[f"{AH}.state_encoder.layer1.W"][eid].to(dt)
+ w['se_b1'] = sd[f"{AH}.state_encoder.layer1.b"][eid].to(dt)
+ w['se_w2'] = sd[f"{AH}.state_encoder.layer2.W"][eid].to(dt)
+ w['se_b2'] = sd[f"{AH}.state_encoder.layer2.b"][eid].to(dt)
+ w['ad_w1'] = sd[f"{AH}.action_decoder.layer1.W"][eid].to(dt)
+ w['ad_b1'] = sd[f"{AH}.action_decoder.layer1.b"][eid].to(dt)
+ w['ad_w2'] = sd[f"{AH}.action_decoder.layer2.W"][eid].to(dt)
+ w['ad_b2'] = sd[f"{AH}.action_decoder.layer2.b"][eid].to(dt)
+ w['posemb'] = sd[f"{AH}.position_embedding.weight"].to(dt)
+ # diffusers Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=1)
+ half = 128
+ self._dit_ts_freqs = torch.exp(
+ -torch.arange(half, dtype=torch.float32, device='cuda')
+ * (math.log(10000.0) / (half - 1)))
+ # SinusoidalPositionalEncoding(1536): sin-first, exact half_dim denom
+ hd = self.D_dit // 2
+ self._dit_tau_freqs = torch.exp(
+ -torch.arange(hd, dtype=torch.float32, device='cuda')
+ * (math.log(10000.0) / hd))
+ self._dit_tw = w
+
+ @staticmethod
+ def _resolve_eagle_dir() -> pathlib.Path:
+ """Locate the Eagle-Block2A-2B-v2 remote-code directory."""
+ override = os.environ.get("FLASHRT_N16_EAGLE_DIR")
+ if override:
+ p = pathlib.Path(override)
+ if (p / "config.json").exists():
+ return p
+ raise RuntimeError(
+ f"FLASHRT_N16_EAGLE_DIR={override} has no config.json")
+ cache = pathlib.Path.home() / ".cache/huggingface/modules/transformers_modules"
+ for cand in sorted(cache.glob(
+ "Eagle_hyphen_Block2A_hyphen_2B_hyphen_v2/*/config.json")):
+ return cand.parent
+ raise RuntimeError(
+ "Eagle-Block2A-2B-v2 remote code not found. Load the GR00T N1.6 "
+ "model once with transformers (AutoModel.from_pretrained) to "
+ "populate the HF cache, or set FLASHRT_N16_EAGLE_DIR.")
+
+ def _setup_torch_siglip(self):
+ """Parity mode: HF-native Siglip2VisionModel (bf16) from the Eagle
+ remote code, weights from the checkpoint state dict."""
+ import glob as _glob
+ import importlib.util as _ilu
+ import json as _json
+ eagle_dir = self._resolve_eagle_dir()
+ cfg = _json.load(open(eagle_dir / "config.json"))
+ mod_path = None
+ for cand in _glob.glob(str(eagle_dir / "modeling_siglip2.py")) + \
+ _glob.glob(str(pathlib.Path.home() /
+ ".cache/huggingface/modules/transformers_modules/"
+ "Eagle_hyphen_Block2A_hyphen_2B_hyphen_v2/*/"
+ "modeling_siglip2.py")):
+ mod_path = cand
+ break
+ if mod_path is None:
+ raise RuntimeError("modeling_siglip2.py not found for parity SigLIP")
+ spec = _ilu.spec_from_file_location("eagle_siglip2_parity", mod_path)
+ mod = _ilu.module_from_spec(spec)
+ spec.loader.exec_module(mod)
+ vc = dict(cfg["vision_config"])
+ vc.pop("_attn_implementation_autoset", None)
+ model = mod.Siglip2VisionModel(mod.Siglip2VisionConfig(**vc))
+ model = model.to(torch.bfloat16).cuda().eval()
+ prefix = "backbone.model.vision_model."
+ sub = {k[len(prefix):]: v.to(torch.bfloat16)
+ for k, v in self._full_sd.items() if k.startswith(prefix)}
+ missing, unexpected = model.load_state_dict(sub, strict=False)
+ if missing or unexpected:
+ raise RuntimeError(
+ f"torch SigLIP weight mismatch: missing={missing[:5]} "
+ f"unexpected={unexpected[:5]}")
+ self._torch_siglip = model
+ self._patch_siglip_fa4(mod)
+ self._setup_siglip_fp4()
+ logger.info("Torch SigLIP (parity mode) loaded: 27L bf16")
+
+ def _patch_siglip_fa4(self, mod):
+ """Route Siglip2Attention through FA4 (flash_rt.hardware.thor.fa4_backend).
+
+ The parity encoder does cross-view FULL attention over the packed
+ 648-token sequence (sdpa silently ignores the NaFlex window
+ segmentation — see docs/groot_n16_thor_sm110.md bug #3), so FA4's
+ causal=False path is an exact drop-in for the sdpa call. Measured
+ encoder-graph 10.3 -> 9.0 ms, actions cos 1.000000 vs sdpa
+ (2026-08-14). FLASHRT_N16_FA4=0 forces sdpa; a missing FA4 runtime
+ falls back silently. Must run before the encoder graph is captured.
+ """
+ if os.environ.get("FLASHRT_N16_FA4", "1") == "0":
+ return
+ try:
+ from flash_rt.hardware.thor import fa4_backend
+ except Exception: # noqa: BLE001
+ return
+ if not fa4_backend.is_available():
+ logger.info("SigLIP FA4 skipped: %s", fa4_backend.status())
+ return
+ fa4 = fa4_backend.fa4_func()
+ attn_cls = type(self._torch_siglip.vision_model.encoder.layers[0].self_attn)
+ if getattr(attn_cls, "_flashrt_fa4_patched", False):
+ return
+ apply_rope = attn_cls.forward.__globals__.get("apply_rope")
+
+ def fa4_forward(self, hidden_states, output_attentions=False,
+ rope_freqs_cis=None, win_meta_list=None,
+ windows_attn=False):
+ B, S, E = hidden_states.shape
+ q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+ k = self.k_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+ v = self.v_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+ if self.use_rope and apply_rope is not None:
+ q, k = apply_rope(q, k, rope_freqs_cis)
+ o = fa4(q.contiguous(), k.contiguous(), v.contiguous(),
+ causal=False, softmax_scale=self.scale)
+ if isinstance(o, tuple):
+ o = o[0]
+ return self.out_proj(o.view(B, S, E)), None
+
+ attn_cls.forward = fa4_forward
+ attn_cls._flashrt_fa4_patched = True
+ logger.info("SigLIP attention routed through FA4 (%s)", fa4_backend.status())
+
+ def _setup_torch_qwen3(self):
+ """Parity mode: load HF-native Qwen3Model (bf16, sdpa) from the
+ checkpoint state dict. Same math as the eager HF baseline."""
+ import json as _json
+ from transformers import Qwen3Config, Qwen3Model
+ eagle_dir = self._resolve_eagle_dir()
+ cfg = _json.load(open(eagle_dir / "config.json"))
+ tc = dict(cfg.get("text_config", cfg))
+ tc["num_hidden_layers"] = 16 # checkpoint is truncated
+ tc["_attn_implementation"] = "sdpa"
+ config = Qwen3Config(**tc)
+ model = Qwen3Model(config).to(torch.bfloat16).cuda().eval()
+ prefix = "backbone.model.language_model.model."
+ sub = {k[len(prefix):]: v.to(torch.bfloat16)
+ for k, v in self._full_sd.items()
+ if k.startswith(prefix)}
+ missing, unexpected = model.load_state_dict(sub, strict=False)
+ missing = [m for m in missing if "embed_tokens" not in m]
+ if missing or unexpected:
+ raise RuntimeError(
+ f"torch Qwen3 weight mismatch: missing={missing[:5]} "
+ f"unexpected={unexpected[:5]}")
+ self._torch_qwen3 = model
+ logger.info("Torch Qwen3 (parity mode) loaded: 16L bf16 sdpa")
+
+ def _setup_qwen3_fp4(self):
+ """FLASHRT_N16_QWEN3_FP4=1: NVFP4 fused-epilogue GEMMs for all 16
+ Qwen3 layers (same kernel family as the DiT fused chain, upstream
+ #163). RMSNorm / q-k norms / RoPE / sdpa / SiLU-mul stay bf16 torch;
+ every projection GEMM runs W4A4 with fused bias (+residual) epilogue
+ and the residual stream is updated in place. Weight traffic drops
+ 2.85 GB -> ~0.8 GB per inference.
+ """
+ if os.environ.get("FLASHRT_N16_QWEN3_FP4", "1") != "1":
+ return
+ if getattr(self, "_qwen3_fp4_done", False):
+ return
+ try:
+ import flash_rt.flash_rt_fp4 as _f4
+ except ImportError:
+ logger.warning("FLASHRT_N16_QWEN3_FP4=1 but flash_rt_fp4 is "
+ "missing; Qwen3 stays bf16")
+ return
+ if not hasattr(_f4, "cutlass_fp4_gemm_bias_res_bf16"):
+ logger.warning("Qwen3 FP4 needs the fused-epilogue kernels "
+ "(rebuild flash_rt_fp4); staying bf16")
+ return
+ from transformers.models.qwen3.modeling_qwen3 import apply_rotary_pos_emb
+ dt = torch.bfloat16
+ Se = self._Se
+ D = self.D_llm
+ KV = self.NHKV * self.HD_llm
+ FF = self.H_llm
+ NH, NHKV, HD = self.NHQ, self.NHKV, self.HD_llm
+
+ def _qw(w_bf):
+ w16 = w_bf.to(torch.float16).contiguous()
+ N, K = w16.shape
+ packed = torch.empty(N, K // 2, dtype=torch.uint8, device='cuda')
+ sfb = torch.zeros(_f4.sfa_size_bytes(N, K, True),
+ dtype=torch.uint8, device='cuda')
+ rc = _f4.quantize_fp4_dynamic_sfa_fp16(
+ w16.data_ptr(), packed.data_ptr(), sfb.data_ptr(), N, K, True, 0)
+ if rc != 0:
+ raise RuntimeError(f"Qwen3 fp4 weight quantize rc={rc}")
+ return packed, sfb
+
+ def _sfa(cols):
+ return torch.zeros(_f4.sfa_size_bytes(Se, cols, False),
+ dtype=torch.uint8, device='cuda')
+
+ def _ck(rc, what, li):
+ if rc != 0:
+ raise RuntimeError(f"Qwen3 fp4 {what} layer {li} rc={rc}")
+
+ def _make_fwd(tab, bufs, ln1_w, ln2_w, ln_eps, qn, kn, ob, gb, ub,
+ db, li):
+ def fwd(hidden_states, position_embeddings, attention_mask=None,
+ **kwargs):
+ s = torch.cuda.current_stream().cuda_stream
+ h = hidden_states.view(Se, D)
+ hp = h.data_ptr()
+ xn_p, xn_s = bufs['xn_p'].data_ptr(), bufs['xn_s'].data_ptr()
+ # ── attn block ──
+ _ck(_f4.rms_norm_weight_fp4_sfa_bf16(
+ hp, ln1_w.data_ptr(), xn_p, xn_s, Se, D, ln_eps, s),
+ "q-ln", li)
+ packed, sfb = tab['q']
+ _ck(_f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ attn_q_b[li].data_ptr(), bufs['q'].data_ptr(),
+ Se, D, D, s), "q", li)
+ packed, sfb = tab['k']
+ _ck(_f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ attn_k_b[li].data_ptr(), bufs['k'].data_ptr(),
+ Se, KV, D, s), "k", li)
+ packed, sfb = tab['v']
+ _ck(_f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ attn_v_b[li].data_ptr(), bufs['v'].data_ptr(),
+ Se, KV, D, s), "v", li)
+ cos, sin = position_embeddings
+ _ck(fvk.qk_norm_rope_rotate_half_bf16(
+ bufs['q'].data_ptr(), qn.weight.data_ptr(),
+ cos.data_ptr(), sin.data_ptr(),
+ Se, NH, HD, qn.variance_epsilon, s), "q-norm-rope", li)
+ _ck(fvk.qk_norm_rope_rotate_half_bf16(
+ bufs['k'].data_ptr(), kn.weight.data_ptr(),
+ cos.data_ptr(), sin.data_ptr(),
+ Se, NHKV, HD, kn.variance_epsilon, s), "k-norm-rope", li)
+ q = bufs['q'].view(1, Se, NH, HD).transpose(1, 2)
+ k = bufs['k'].view(1, Se, NHKV, HD).transpose(1, 2)
+ v = bufs['v'].view(1, Se, NHKV, HD).transpose(1, 2)
+ # Torch's fused attention backends reject mismatched Q/KV
+ # head counts (enable_gqa falls back to unfused math: 2 fp16
+ # GEMMs + materialized softmax per layer). Expanding KV
+ # enables the fused path (measured Qwen3 7.0 -> 5.0 ms).
+ if NHKV != NH:
+ k = k.repeat_interleave(NH // NHKV, dim=1)
+ v = v.repeat_interleave(NH // NHKV, dim=1)
+ if attention_mask is None:
+ o = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, is_causal=q.shape[2] > 1)
+ else:
+ o = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attention_mask)
+ o_flat = o.transpose(1, 2).reshape(Se, D)
+ o_p, o_s = bufs['o_p'].data_ptr(), bufs['o_s'].data_ptr()
+ _ck(_f4.quantize_fp4_dynamic_sfa_bf16_vec(
+ o_flat.data_ptr(), o_p, o_s, Se, D, False, s), "o-q", li)
+ packed, sfb = tab['o']
+ _ck(_f4.cutlass_fp4_gemm_bias_res_bf16(
+ o_p, o_s, packed.data_ptr(), sfb.data_ptr(),
+ ob.data_ptr(), hp, hp, Se, D, D, s), "o", li)
+ # ── FFN block ──
+ _ck(_f4.rms_norm_weight_fp4_sfa_bf16(
+ hp, ln2_w.data_ptr(), xn_p, xn_s, Se, D, ln_eps, s),
+ "f-ln", li)
+ packed, sfb = tab['gate']
+ _ck(_f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ gb.data_ptr(), bufs['gate'].data_ptr(),
+ Se, FF, D, s), "gate", li)
+ packed, sfb = tab['up']
+ _ck(_f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ ub.data_ptr(), bufs['up'].data_ptr(),
+ Se, FF, D, s), "up", li)
+ gu_p, gu_s = bufs['gu_p'].data_ptr(), bufs['gu_s'].data_ptr()
+ _ck(_f4.silu_mul_fp4_sfa_bf16(
+ bufs['gate'].data_ptr(), bufs['up'].data_ptr(),
+ gu_p, gu_s, Se, FF, False, s), "gu-q", li)
+ packed, sfb = tab['dn']
+ _ck(_f4.cutlass_fp4_gemm_bias_res_bf16(
+ gu_p, gu_s, packed.data_ptr(), sfb.data_ptr(),
+ db.data_ptr(), hp, hp, Se, D, FF, s), "dn", li)
+ return h.view(1, Se, D)
+ return fwd
+
+ zero_bias = torch.zeros(FF, dtype=dt, device='cuda')
+
+ def _bias(t):
+ return t if t is not None else zero_bias
+
+ attn_q_b, attn_k_b, attn_v_b = [], [], []
+ for li, layer in enumerate(self._torch_qwen3.layers):
+ attn, mlp = layer.self_attn, layer.mlp
+ tab = {n: _qw(getattr(attn, n + "_proj").weight)
+ for n in ('q', 'k', 'v', 'o')}
+ tab['gate'] = _qw(mlp.gate_proj.weight)
+ tab['up'] = _qw(mlp.up_proj.weight)
+ tab['dn'] = _qw(mlp.down_proj.weight)
+ attn_q_b.append(_bias(attn.q_proj.bias))
+ attn_k_b.append(_bias(attn.k_proj.bias))
+ attn_v_b.append(_bias(attn.v_proj.bias))
+ bufs = {
+ 'xn_p': torch.empty(Se, D // 2, dtype=torch.uint8, device='cuda'),
+ 'xn_s': _sfa(D),
+ 'q': torch.empty(Se, D, dtype=dt, device='cuda'),
+ 'k': torch.empty(Se, KV, dtype=dt, device='cuda'),
+ 'v': torch.empty(Se, KV, dtype=dt, device='cuda'),
+ 'o_p': torch.empty(Se, D // 2, dtype=torch.uint8, device='cuda'),
+ 'o_s': _sfa(D),
+ 'gate': torch.empty(Se, FF, dtype=dt, device='cuda'),
+ 'up': torch.empty(Se, FF, dtype=dt, device='cuda'),
+ 'gu_p': torch.empty(Se, FF // 2, dtype=torch.uint8, device='cuda'),
+ 'gu_s': _sfa(FF),
+ }
+ layer.forward = _make_fwd(
+ tab, bufs, layer.input_layernorm.weight,
+ layer.post_attention_layernorm.weight,
+ layer.input_layernorm.variance_epsilon,
+ attn.q_norm, attn.k_norm,
+ _bias(attn.o_proj.bias), _bias(mlp.gate_proj.bias),
+ _bias(mlp.up_proj.bias), _bias(mlp.down_proj.bias), li)
+ self._qwen3_fp4_done = True
+ logger.info("Qwen3 NVFP4 fused-epilogue tier enabled (16 layers)")
+
+ def _setup_siglip_fp4(self):
+ """FLASHRT_N16_SIGLIP_FP4 (default on): NVFP4 fused-epilogue encoder.
+
+ Every layer runs the kernel chain: LN->fp4 producer (affine LayerNorm
+ expressed as AdaLN with scale=w-1, shift=b), q/k/v bias GEMMs into
+ contiguous buffers, FA4 full attention, o bias+residual GEMM, LN->fp4,
+ fc1 bias+tanh-GELU+fp4out (N padded 4304->4352 so both fp4 dims are
+ mult-64; the pad columns stay zero end-to-end), fc2 bias+residual.
+ Measured encoder 9.0 -> 7.6 ms; actions vs HF cos 0.99997 / maxd
+ 0.038 (vs 0.023 bf16) — the vision-quality trade is simulation-gated;
+ set FLASHRT_N16_SIGLIP_FP4=0 to revert to bf16.
+ """
+ if os.environ.get("FLASHRT_N16_SIGLIP_FP4", "1") != "1":
+ return
+ if getattr(self, "_siglip_fp4_done", False):
+ return
+ try:
+ import flash_rt.flash_rt_fp4 as _f4
+ except ImportError:
+ logger.warning("FLASHRT_N16_SIGLIP_FP4 on but flash_rt_fp4 "
+ "missing; SigLIP stays bf16")
+ return
+ if not hasattr(_f4, "cutlass_fp4_gemm_bias_gelu_fp4out_bf16"):
+ logger.warning("SigLIP FP4 needs the fused-epilogue kernels "
+ "(rebuild flash_rt_fp4); staying bf16")
+ return
+ from flash_rt.hardware.thor import fa4_backend
+ if not fa4_backend.is_available():
+ logger.warning("SigLIP FP4 tier requires FA4; staying bf16")
+ return
+ fa4 = fa4_backend.fa4_func()
+ dt = torch.bfloat16
+ enc = self._torch_siglip.vision_model.encoder
+ S = self._num_views * self.spv_raw
+ head0 = enc.layers[0]
+ attn0, D = head0.self_attn, head0.self_attn.embed_dim
+ NH, HD = attn0.num_heads, attn0.head_dim
+ FF = head0.mlp.fc1.out_features
+ FFP = (FF + 63) // 64 * 64
+ eps = head0.layer_norm1.eps
+ gmod = type(attn0).forward.__globals__
+ apply_rope = gmod.get("apply_rope")
+
+ def _qw(w_bf, pad_rows=0, pad_cols=0):
+ w16 = w_bf.to(torch.float16)
+ if pad_cols:
+ w16 = F.pad(w16, (0, pad_cols))
+ if pad_rows:
+ w16 = F.pad(w16, (0, 0, 0, pad_rows))
+ w16 = w16.contiguous()
+ N, K = w16.shape
+ packed = torch.empty(N, K // 2, dtype=torch.uint8, device='cuda')
+ sfb = torch.zeros(_f4.sfa_size_bytes(N, K, True),
+ dtype=torch.uint8, device='cuda')
+ rc = _f4.quantize_fp4_dynamic_sfa_fp16(
+ w16.data_ptr(), packed.data_ptr(), sfb.data_ptr(),
+ N, K, True, 0)
+ if rc != 0:
+ raise RuntimeError(f"SigLIP fp4 weight quantize rc={rc}")
+ return packed, sfb
+
+ def _ck(rc, what, l):
+ if rc != 0:
+ raise RuntimeError(f"SigLIP fp4 {what} layer {l} rc={rc}")
+
+ def _make_fwd(tab, bufs, use_rope, scale, li):
+ def fwd(hidden_states, output_attentions=False,
+ rope_freqs_cis=None, win_meta_list=None,
+ windows_attn=False):
+ s = torch.cuda.current_stream().cuda_stream
+ h = hidden_states.view(S, D)
+ hp = h.data_ptr()
+ xn_p, xn_s = bufs['xn_p'].data_ptr(), bufs['xn_s'].data_ptr()
+ _ck(_f4.ada_layer_norm_fp4_sfa_bf16(
+ hp, tab['ln1_sc'].data_ptr(), tab['ln1_sh'].data_ptr(),
+ xn_p, xn_s, S, D, eps, s), "ln1", li)
+ for n in ('q', 'k', 'v'):
+ packed, sfb = tab[n]
+ _ck(_f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ tab[n + '_b'].data_ptr(), bufs[n].data_ptr(),
+ S, D, D, s), n, li)
+ q = bufs['q'].view(1, S, NH, HD)
+ k = bufs['k'].view(1, S, NH, HD)
+ v = bufs['v'].view(1, S, NH, HD)
+ if use_rope and apply_rope is not None:
+ q, k = apply_rope(q, k, rope_freqs_cis)
+ q = q.contiguous()
+ k = k.contiguous()
+ o = fa4(q, k, v, causal=False, softmax_scale=scale)
+ if isinstance(o, tuple):
+ o = o[0]
+ o_p, o_s = bufs['o_p'].data_ptr(), bufs['o_s'].data_ptr()
+ _ck(_f4.quantize_fp4_dynamic_sfa_bf16_vec(
+ o.data_ptr(), o_p, o_s, S, D, False, s), "o-q", li)
+ packed, sfb = tab['o']
+ _ck(_f4.cutlass_fp4_gemm_bias_res_bf16(
+ o_p, o_s, packed.data_ptr(), sfb.data_ptr(),
+ tab['o_b'].data_ptr(), hp, hp, S, D, D, s), "o", li)
+ _ck(_f4.ada_layer_norm_fp4_sfa_bf16(
+ hp, tab['ln2_sc'].data_ptr(), tab['ln2_sh'].data_ptr(),
+ xn_p, xn_s, S, D, eps, s), "ln2", li)
+ packed, sfb = tab['fc1']
+ hid_p, hid_s = bufs['hid_p'].data_ptr(), bufs['hid_s'].data_ptr()
+ _ck(_f4.cutlass_fp4_gemm_bias_gelu_fp4out_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ tab['fc1_b'].data_ptr(), hid_p, hid_s,
+ S, FFP, D, s), "fc1", li)
+ packed, sfb = tab['fc2']
+ _ck(_f4.cutlass_fp4_gemm_bias_res_bf16(
+ hid_p, hid_s, packed.data_ptr(), sfb.data_ptr(),
+ tab['fc2_b'].data_ptr(), hp, hp, S, D, FFP, s), "fc2", li)
+ return (h.view(1, S, D),)
+ return fwd
+
+ for li, layer in enumerate(enc.layers):
+ attn, mlp = layer.self_attn, layer.mlp
+ ln1, ln2 = layer.layer_norm1, layer.layer_norm2
+ tab = {
+ 'q': _qw(attn.q_proj.weight), 'q_b': attn.q_proj.bias,
+ 'k': _qw(attn.k_proj.weight), 'k_b': attn.k_proj.bias,
+ 'v': _qw(attn.v_proj.weight), 'v_b': attn.v_proj.bias,
+ 'o': _qw(attn.out_proj.weight), 'o_b': attn.out_proj.bias,
+ 'fc1': _qw(mlp.fc1.weight, pad_rows=FFP - FF),
+ 'fc1_b': F.pad(mlp.fc1.bias, (0, FFP - FF)),
+ 'fc2': _qw(mlp.fc2.weight, pad_cols=FFP - FF),
+ 'fc2_b': mlp.fc2.bias,
+ 'ln1_sc': (ln1.weight - 1).contiguous(),
+ 'ln1_sh': ln1.bias.contiguous(),
+ 'ln2_sc': (ln2.weight - 1).contiguous(),
+ 'ln2_sh': ln2.bias.contiguous(),
+ }
+ bufs = {
+ 'xn_p': torch.empty(S, D // 2, dtype=torch.uint8, device='cuda'),
+ 'xn_s': torch.zeros(_f4.sfa_size_bytes(S, D, False),
+ dtype=torch.uint8, device='cuda'),
+ 'q': torch.empty(S, D, dtype=dt, device='cuda'),
+ 'k': torch.empty(S, D, dtype=dt, device='cuda'),
+ 'v': torch.empty(S, D, dtype=dt, device='cuda'),
+ 'o_p': torch.empty(S, D // 2, dtype=torch.uint8, device='cuda'),
+ 'o_s': torch.zeros(_f4.sfa_size_bytes(S, D, False),
+ dtype=torch.uint8, device='cuda'),
+ 'hid_p': torch.empty(S, FFP // 2, dtype=torch.uint8, device='cuda'),
+ 'hid_s': torch.zeros(_f4.sfa_size_bytes(S, FFP, False),
+ dtype=torch.uint8, device='cuda'),
+ }
+ layer.forward = _make_fwd(tab, bufs, attn.use_rope, attn.scale, li)
+ self._siglip_fp4_done = True
+ logger.info("SigLIP NVFP4 fused-epilogue encoder enabled "
+ "(27 layers, S=%d)", S)
+
+ def _run_torch_dit(self, state):
+ """4-step flow-matching DiT in bf16 torch — HF-faithful.
+
+ The CUDA-kernel DiT cannot reproduce HF's numerics on this model
+ (see docs/groot_n16_dit_kernel_nonconvergence.md), so the action head
+ runs as bf16 torch. To avoid the 32-layer x 4-step eager kernel-launch
+ overhead, the step-invariant parts are precomputed and the per-frame
+ compute is captured as a CUDA graph over static buffers (set
+ FLASHRT_N16_DIT_GRAPH=0 to force the eager fallback).
+ """
+ if not hasattr(self, '_dit_in_state'):
+ self._dit_setup_graph_buffers()
+ dt = self._dit_dt
+ # Fill static input buffers from the current frame.
+ st = torch.as_tensor(np.asarray(state), dtype=torch.float32).cuda().reshape(1, -1).to(dt)
+ self._dit_in_state.copy_(st)
+ ehs = self._g_vlln_buf
+ self._dit_in_kvt.copy_(ehs.index_select(0, self._dit_txt_idx).to(dt))
+ self._dit_in_kvi.copy_(ehs.index_select(0, self._dit_img_idx).to(dt))
+ self._dit_in_noise.normal_()
+
+ if self._dit_torch_graph is not None:
+ self._dit_torch_graph.replay()
+ return self._dit_out
+ # Eager fallback.
+ self._dit_body()
+ return self._dit_out
+
+ def _dit_setup_graph_buffers(self):
+ """Allocate static DiT I/O buffers, precompute step-invariant tensors,
+ and capture the 4-step DiT as a CUDA graph."""
+ w = self._dit_tw
+ dt = self._dit_dt
+ T, D = self.action_horizon, self.D_dit
+ N, NH, HD = self.num_steps, self.NH_dit, self.HD_dit
+
+ self._dit_txt_idx = self._g_non_img[0].nonzero(as_tuple=True)[0].long()
+ self._dit_img_idx = self._g_img_m[0].nonzero(as_tuple=True)[0].long()
+ n_txt, n_img = self._dit_txt_idx.numel(), self._dit_img_idx.numel()
+ self._dit_in_kvt = torch.empty(n_txt, 2048, dtype=dt, device='cuda')
+ self._dit_in_kvi = torch.empty(n_img, 2048, dtype=dt, device='cuda')
+ self._dit_in_state = torch.empty(1, self.state_dim, dtype=dt, device='cuda')
+ self._dit_in_noise = torch.empty(1, T, self.action_dim, dtype=dt, device='cuda')
+ self._dit_out = torch.empty(1, T, self.action_dim, dtype=dt, device='cuda')
+
+ # Step-invariant embeddings (depend only on weights + step).
+ # Recomputed from self.num_steps so step-count experiments keep the
+ # uniform schedule t_disc = step/N*1000 (HF FlowMatchEuler).
+ self._dit_posemb = w['posemb'][:T].unsqueeze(0)
+ self._dit_silu_tembs, self._dit_taus = [], []
+ for step in range(self.num_steps):
+ t_disc = int(step / float(self.num_steps) * 1000)
+ args = t_disc * self._dit_ts_freqs
+ sincos = torch.cat([torch.cos(args), torch.sin(args)], -1).to(dt).unsqueeze(0)
+ temb = F.linear(F.silu(F.linear(sincos, w['ts_l1_w'], w['ts_l1_b'])),
+ w['ts_l2_w'], w['ts_l2_b'])
+ self._dit_silu_tembs.append(F.silu(temb))
+ freqs = t_disc * self._dit_tau_freqs
+ self._dit_taus.append(torch.cat([torch.sin(freqs), torch.cos(freqs)], -1)
+ .to(dt).unsqueeze(0).expand(1, T, -1).contiguous())
+
+ # Precompute the per-block adaLN conditioning embeddings for every
+ # (step, layer): they depend only on weights + step, yet the HF loop
+ # re-runs the M=1 F.linear inside the denoising loop, streaming
+ # 4*32*9.4MB of n1w weights per inference. Hoisting them out is
+ # numerically identical (same bf16 kernel, same inputs).
+ self._dit_ada_embs = [
+ [F.linear(temb, blk['n1w'], blk['n1b']) for blk in w['blocks']]
+ for temb in self._dit_silu_tembs]
+
+ # Optional NVFP4 tier (FLASHRT_N16_DIT_FP4=1): every block GEMM runs
+ # as a block-scaled W4A4 CUTLASS GEMM (port of the N1.7 FP4 tier,
+ # #163). At M=51 the DiT is weight-bandwidth-bound; fp4 quarters the
+ # weight bytes. Measured DiT 36.6 -> 29.4 ms, actions vs HF
+ # cos 0.999994 / maxd 0.012 (2026-08-14). Weight tables are static
+ # and survive graph re-captures.
+ self._dit_use_fp4 = False
+ if os.environ.get("FLASHRT_N16_DIT_FP4", "1") == "1":
+ try:
+ import flash_rt.flash_rt_fp4 as _f4
+ except ImportError:
+ _f4 = None
+ logger.warning("FLASHRT_N16_DIT_FP4=1 but flash_rt_fp4 is "
+ "missing; DiT stays bf16")
+ if _f4 is not None:
+ if not hasattr(self, "_dit_fq"):
+ def _qw(wt):
+ w16 = wt.to(torch.float16).contiguous()
+ N, K = w16.shape
+ packed = torch.empty(N, K // 2, dtype=torch.uint8,
+ device='cuda')
+ # Zero-init: the tile-interleaved SF layout rounds up
+ # to atom size; unwritten padding must not decode as
+ # garbage scales (upstream #163 pitfall).
+ sfb = torch.zeros(_f4.sfa_size_bytes(N, K, True),
+ dtype=torch.uint8, device='cuda')
+ rc = _f4.quantize_fp4_dynamic_sfa_fp16(
+ w16.data_ptr(), packed.data_ptr(), sfb.data_ptr(),
+ N, K, True, 0)
+ if rc != 0:
+ raise RuntimeError(f"DiT fp4 quantize rc={rc}")
+ return packed, sfb
+ self._dit_fq = [
+ {n: _qw(blk[n]) for n in ('qw', 'kw', 'vw', 'ow',
+ 'fw', 'dw')}
+ for blk in w['blocks']]
+ self._dit_fp4_mod = _f4
+ self._dit_fp4_bufs = {}
+ self._dit_use_fp4 = True
+ logger.info("DiT NVFP4 tier enabled (block GEMMs W4A4)")
+ # Fused-epilogue chain (upstream N1.7 #163 port): norm->fp4
+ # producers + bias/residual/GELU inside the GEMM epilogues.
+ if hasattr(_f4, "ada_layer_norm_fp4_sfa_bf16"):
+ self._dit_fp4_fused = True
+ if not hasattr(self, "_dit_fqf"):
+ Sa, FF = T + 1, self.H_dit
+ fqf = []
+ for l, blk in enumerate(w['blocks']):
+ e = {}
+ if l % 2 == 1: # self-attn: fused QKV
+ qkvw = torch.cat([blk['qw'], blk['kw'], blk['vw']], 0)
+ e['qkv'] = _qw(qkvw)
+ e['qkv_b'] = torch.cat(
+ [blk['qb'], blk['kb'], blk['vb']]).contiguous()
+ else: # cross-attn: Q only (K/V precomputed)
+ e['q'] = _qw(blk['qw'])
+ e['q_b'] = blk['qb'].contiguous()
+ e['o'] = _qw(blk['ow'])
+ e['up'] = _qw(blk['fw'])
+ e['dn'] = _qw(blk['dw'])
+ fqf.append(e)
+ self._dit_fqf = fqf
+ # Per-step AdaLN modulators (scale, shift) and the final
+ # proj1 (shift, scale) as pointer-stable tensors.
+ ada_sc = torch.stack(
+ [torch.stack([emb[0, :D] for emb in embs])
+ for embs in self._dit_ada_embs]).contiguous()
+ ada_sh = torch.stack(
+ [torch.stack([emb[0, D:] for emb in embs])
+ for embs in self._dit_ada_embs]).contiguous()
+ self._dit_ada_scale = ada_sc
+ self._dit_ada_shift = ada_sh
+ self._dit_proj1 = torch.stack([
+ F.linear(temb, w['proj1_w'], w['proj1_b'])[0]
+ for temb in self._dit_silu_tembs]).contiguous()
+ # Static scratch (graph-capture stable).
+ self._dit_h = torch.empty(Sa, D, dtype=dt, device='cuda')
+ self._dit_qkv_buf = torch.empty(Sa, 3 * D, dtype=dt, device='cuda')
+ self._dit_q_buf = torch.empty(Sa, D, dtype=dt, device='cuda')
+ self._dit_xn_p = torch.empty(Sa, D // 2, dtype=torch.uint8, device='cuda')
+ self._dit_xn_s = torch.zeros(_f4.sfa_size_bytes(Sa, D, False),
+ dtype=torch.uint8, device='cuda')
+ self._dit_o_p = torch.empty(Sa, D // 2, dtype=torch.uint8, device='cuda')
+ self._dit_o_s = torch.zeros(_f4.sfa_size_bytes(Sa, D, False),
+ dtype=torch.uint8, device='cuda')
+ self._dit_hid_p = torch.empty(Sa, FF // 2, dtype=torch.uint8, device='cuda')
+ self._dit_hid_s = torch.zeros(_f4.sfa_size_bytes(Sa, FF, False),
+ dtype=torch.uint8, device='cuda')
+ logger.info("DiT NVFP4 fused-epilogue chain enabled")
+ else:
+ self._dit_fp4_fused = False
+
+ # Seed the inputs so warmup/capture run on valid data.
+ self._dit_in_state.zero_()
+ self._dit_in_kvt.zero_(); self._dit_in_kvi.zero_()
+ self._dit_in_noise.normal_()
+
+ self._dit_torch_graph = None
+ # Graph the torch DiT by default (~15ms vs 40ms eager); reset paths
+ # delete and re-capture it safely. Set FLASHRT_N16_DIT_GRAPH=0 to
+ # force the eager fallback.
+ if os.environ.get("FLASHRT_N16_DIT_GRAPH", "1") != "0":
+ try:
+ s = torch.cuda.Stream()
+ s.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(s):
+ for _ in range(2):
+ self._dit_body()
+ torch.cuda.current_stream().wait_stream(s)
+ g = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(g):
+ self._dit_body()
+ self._dit_torch_graph = g
+ logger.info("torch DiT captured as CUDA graph (n_txt=%d, n_img=%d)",
+ n_txt, n_img)
+ except Exception as e: # noqa: BLE001
+ logger.warning("torch DiT graph capture failed (%r); using eager", e)
+ self._dit_torch_graph = None
+
+ def _dit_fp4_linear(self, x_bf, fq, N, K, bias):
+ """NVFP4 W4A4 GEMM replacement for F.linear inside the DiT graph.
+
+ Static per-shape buffers (graph-safe): bf16->fp16 cast, dynamic
+ per-16 activation quant, cutlass_fp4_sq_fp16, bf16 bias add.
+ """
+ f4 = self._dit_fp4_mod
+ M = x_bf.shape[-2]
+ key = (M, K, N)
+ bufs = self._dit_fp4_bufs
+ if key not in bufs:
+ bufs[key] = (torch.empty(M, K, dtype=torch.float16, device='cuda'),
+ torch.empty(M, K // 2, dtype=torch.uint8, device='cuda'),
+ torch.zeros(f4.sfa_size_bytes(M, K, False),
+ dtype=torch.uint8, device='cuda'),
+ torch.empty(M, N, dtype=torch.float16, device='cuda'))
+ x16b, xq, xs, out = bufs[key]
+ packed, sfb = fq
+ x16b.copy_(x_bf.reshape(M, K).to(torch.float16))
+ s = torch.cuda.current_stream().cuda_stream
+ f4.quantize_fp4_dynamic_sfa_fp16(x16b.data_ptr(), xq.data_ptr(),
+ xs.data_ptr(), M, K, False, s)
+ f4.cutlass_fp4_sq_fp16(xq.data_ptr(), xs.data_ptr(), packed.data_ptr(),
+ sfb.data_ptr(), out.data_ptr(), M, N, K,
+ 1.0, 0.0, s)
+ y = out.to(torch.bfloat16)
+ if bias is not None:
+ y = y + bias
+ return y.reshape(*x_bf.shape[:-1], N)
+
+ def _dit_body_fp4_fused(self):
+ """Fused-epilogue NVFP4 DiT chain (upstream N1.7 #163 port).
+
+ Per block: adaLN-modulated norm emits fp4 directly; QKV/cross-Q/O/
+ FFN run as block-scaled NVFP4 GEMMs with bias (and residual /
+ tanh-GELU+fp4out) fused into the epilogues — 8 kernels per layer,
+ no elementwise traffic between them. Attention stays bf16 sdpa.
+ """
+ w = self._dit_tw
+ dt = self._dit_dt
+ f4 = self._dit_fp4_mod
+ T, D = self.action_horizon, self.D_dit
+ N, NH, HD = self.num_steps, self.NH_dit, self.HD_dit
+ FF = self.H_dit
+ Sa = T + 1
+
+ h = self._dit_h
+ qkv_buf, q_buf = self._dit_qkv_buf, self._dit_q_buf
+ xn_p, xn_s = self._dit_xn_p.data_ptr(), self._dit_xn_s.data_ptr()
+ o_p, o_s = self._dit_o_p.data_ptr(), self._dit_o_s.data_ptr()
+ hid_p, hid_s = self._dit_hid_p.data_ptr(), self._dit_hid_s.data_ptr()
+ fqf = self._dit_fqf
+
+ state = self._dit_in_state
+ st = F.relu(state @ w['se_w1'] + w['se_b1'])
+ state_feat = (st @ w['se_w2'] + w['se_b2']) # (1, D)
+ kv_text, kv_img = self._dit_in_kvt, self._dit_in_kvi
+
+ # Cross-attn K/V in bf16, once per inference (backbone constant).
+ # fp4 here was measured to degrade actions 5x (cos 0.9992 / maxd
+ # 0.105) for ~0 ms gain — cross features feed every cross layer.
+ cross_kv = {}
+ for l, blk in enumerate(w['blocks']):
+ if l % 2 == 0:
+ kv = kv_text if l % 4 == 0 else kv_img
+ cross_kv[l] = (
+ F.linear(kv, blk['kw'], blk['kb']).view(1, -1, NH, HD).transpose(1, 2),
+ F.linear(kv, blk['vw'], blk['vb']).view(1, -1, NH, HD).transpose(1, 2))
+
+ def _ck(rc, what, l):
+ if rc != 0:
+ raise RuntimeError(f"DiT fp4 fused {what} layer {l} rc={rc}")
+
+ actions = self._dit_in_noise
+ for step in range(N):
+ s = torch.cuda.current_stream().cuda_stream
+ a_emb = actions @ w['ae_w1'] + w['ae_b1']
+ x = torch.cat([a_emb, self._dit_taus[step]], -1)
+ af = (F.silu(x @ w['ae_w2'] + w['ae_b2']) @ w['ae_w3']
+ + w['ae_b3']) + self._dit_posemb
+ h[:1].copy_(state_feat)
+ h[1:].copy_(af.view(T, D))
+ h_ptr = h.data_ptr()
+ for l, blk in enumerate(w['blocks']):
+ e = fqf[l]
+ _ck(f4.ada_layer_norm_fp4_sfa_bf16(
+ h_ptr, self._dit_ada_scale[step, l].data_ptr(),
+ self._dit_ada_shift[step, l].data_ptr(), xn_p, xn_s,
+ Sa, D, 1e-5, s), "adaln", l)
+ if l % 2 == 1: # self-attn
+ packed, sfb = e['qkv']
+ _ck(f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ e['qkv_b'].data_ptr(), qkv_buf.data_ptr(),
+ Sa, 3 * D, D, s), "qkv", l)
+ q = qkv_buf[:, :D].view(1, Sa, NH, HD).transpose(1, 2)
+ k = qkv_buf[:, D:2 * D].view(1, Sa, NH, HD).transpose(1, 2)
+ v = qkv_buf[:, 2 * D:].view(1, Sa, NH, HD).transpose(1, 2)
+ else: # cross-attn
+ packed, sfb = e['q']
+ _ck(f4.cutlass_fp4_gemm_bias_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ e['q_b'].data_ptr(), q_buf.data_ptr(),
+ Sa, D, D, s), "q", l)
+ q = q_buf.view(1, Sa, NH, HD).transpose(1, 2)
+ k, v = cross_kv[l]
+ o = F.scaled_dot_product_attention(q, k, v)
+ o_flat = o.transpose(1, 2).reshape(Sa, D)
+ _ck(f4.quantize_fp4_dynamic_sfa_bf16_vec(
+ o_flat.data_ptr(), o_p, o_s, Sa, D, False, s), "o-quant", l)
+ packed, sfb = e['o']
+ _ck(f4.cutlass_fp4_gemm_bias_res_bf16(
+ o_p, o_s, packed.data_ptr(), sfb.data_ptr(),
+ blk['ob'].data_ptr(), h_ptr, h_ptr,
+ Sa, D, D, s), "o", l)
+ _ck(f4.layer_norm_no_affine_fp4_sfa_bf16(
+ h_ptr, xn_p, xn_s, Sa, D, 1e-5, s), "ffn-ln", l)
+ packed, sfb = e['up']
+ _ck(f4.cutlass_fp4_gemm_bias_gelu_fp4out_bf16(
+ xn_p, xn_s, packed.data_ptr(), sfb.data_ptr(),
+ blk['fb'].data_ptr(), hid_p, hid_s,
+ Sa, FF, D, s), "ffn-up", l)
+ packed, sfb = e['dn']
+ _ck(f4.cutlass_fp4_gemm_bias_res_bf16(
+ hid_p, hid_s, packed.data_ptr(), sfb.data_ptr(),
+ blk['db'].data_ptr(), h_ptr, h_ptr,
+ Sa, D, FF, s), "ffn-down", l)
+ p1 = self._dit_proj1[step]
+ shift, scale = p1.chunk(2, dim=0)
+ hn = F.layer_norm(h, (D,), None, None, 1e-6) * (1 + scale) + shift
+ out = F.linear(hn, w['proj2_w'], w['proj2_b'])
+ dec = F.relu(out @ w['ad_w1'] + w['ad_b1'])
+ pred = dec @ w['ad_w2'] + w['ad_b2']
+ actions = actions + (1.0 / N) * pred[-T:]
+ self._dit_out.copy_(actions)
+
+ def _dit_body(self):
+ """The 4-step DiT over static input buffers; writes self._dit_out."""
+ if getattr(self, "_dit_fp4_fused", False):
+ return self._dit_body_fp4_fused()
+ w = self._dit_tw
+ dt = self._dit_dt
+ T, D = self.action_horizon, self.D_dit
+ N, NH, HD = self.num_steps, self.NH_dit, self.HD_dit
+ use_fp4 = getattr(self, "_dit_use_fp4", False)
+ H = self.H_dit
+
+ def lin(x, l, name, Nout, Kin, bias):
+ if use_fp4:
+ return self._dit_fp4_linear(x, self._dit_fq[l][name],
+ Nout, Kin, bias)
+ return F.linear(x, w['blocks'][l][name], bias)
+
+ state = self._dit_in_state
+ st = F.relu(state @ w['se_w1'] + w['se_b1'])
+ state_feat = (st @ w['se_w2'] + w['se_b2']).unsqueeze(0) # (1,1,D)
+ kv_text, kv_img = self._dit_in_kvt, self._dit_in_kvi
+
+ # Cross-attn K/V: backbone is constant across the 4 steps.
+ cross_kv = {}
+ for l, blk in enumerate(w['blocks']):
+ if l % 2 == 0:
+ kv = kv_text if l % 4 == 0 else kv_img
+ Kkv = blk['kw'].shape[1]
+ cross_kv[l] = (
+ lin(kv, l, 'kw', D, Kkv, blk['kb']).view(1, -1, NH, HD).transpose(1, 2),
+ lin(kv, l, 'vw', D, Kkv, blk['vb']).view(1, -1, NH, HD).transpose(1, 2))
+
+ actions = self._dit_in_noise
+ for step in range(N):
+ silu_temb = self._dit_silu_tembs[step]
+ a_emb = actions @ w['ae_w1'] + w['ae_b1']
+ x = torch.cat([a_emb, self._dit_taus[step]], -1)
+ af = (F.silu(x @ w['ae_w2'] + w['ae_b2']) @ w['ae_w3']
+ + w['ae_b3']) + self._dit_posemb
+ hidden = torch.cat([state_feat, af], 1) # (1,1+T,D)
+ for l, blk in enumerate(w['blocks']):
+ emb = self._dit_ada_embs[step][l]
+ scale, shift = emb.chunk(2, dim=1)
+ hnorm = (F.layer_norm(hidden, (D,), None, None, 1e-5)
+ * (1 + scale[:, None]) + shift[:, None])
+ q = lin(hnorm, l, 'qw', D, D, blk['qb'])
+ S_q = q.shape[1]
+ q = q.view(1, S_q, NH, HD).transpose(1, 2)
+ if l % 2 == 1: # self-attn
+ k = lin(hnorm, l, 'kw', D, D, blk['kb']).view(1, S_q, NH, HD).transpose(1, 2)
+ v = lin(hnorm, l, 'vw', D, D, blk['vb']).view(1, S_q, NH, HD).transpose(1, 2)
+ else: # cross-attn
+ k, v = cross_kv[l]
+ o = F.scaled_dot_product_attention(q, k, v)
+ o = o.transpose(1, 2).reshape(1, S_q, D)
+ hidden = hidden + lin(o, l, 'ow', D, D, blk['ob'])
+ hnorm = F.layer_norm(hidden, (D,), None, None, 1e-5)
+ ff = F.gelu(lin(hnorm, l, 'fw', H, D, blk['fb']),
+ approximate='tanh')
+ hidden = hidden + lin(ff, l, 'dw', D, H, blk['db'])
+ emb = F.linear(silu_temb, w['proj1_w'], w['proj1_b'])
+ # HF DiT proj_out_1 chunk order is (shift, scale) — not (scale, shift).
+ shift, scale = emb.chunk(2, dim=1)
+ hidden = (F.layer_norm(hidden, (D,), None, None, 1e-6)
+ * (1 + scale[:, None]) + shift[:, None])
+ out = F.linear(hidden, w['proj2_w'], w['proj2_b'])
+ dec = F.relu(out @ w['ad_w1'] + w['ad_b1'])
+ pred = dec @ w['ad_w2'] + w['ad_b2']
+ actions = actions + (1.0 / N) * pred[:, -T:]
+ self._dit_out.copy_(actions)
+
def _copy_state_feature_to_dit(self, state):
"""Encode current robot state into the captured DiT input buffer.
@@ -896,48 +1849,109 @@ def _dit_forward(self, sa_embs, backbone_features, image_mask, backbone_mask, te
output = self._fp16_gemm(hidden.squeeze(0), self._dit_proj_out_2_w, Sa, self.output_dim, D)
return output.unsqueeze(0) # [1, Sa, output_dim]
+ # ─────────────────────────────────────────────────────────────
+ # CUDA-graph idle guard
+ #
+ # Jetson Thor silently invalidates captured graphs after a few seconds
+ # of GPU idle (replays then run fast but emit garbage); see
+ # docs/thor_gpu_idle_reset_workaround.md. infer() re-captures instead
+ # of replaying once the GPU has been idle longer than
+ # FLASHRT_GRAPH_IDLE_REINIT_S.
+ # ─────────────────────────────────────────────────────────────
+
+ @property
+ def _graph_idle_limit_s(self) -> float:
+ return float(os.environ.get("FLASHRT_GRAPH_IDLE_REINIT_S", "2"))
+
+ def _graph_idle_stale(self) -> bool:
+ return (time.monotonic() - getattr(
+ self, "_last_graph_use", time.monotonic())) > self._graph_idle_limit_s
+
+ def invalidate_graphs(self) -> None:
+ """Force the captured CUDA graphs to be re-captured on next use."""
+ self._last_graph_use = 0.0
+
+ def reset_graph_runtime(self) -> None:
+ """Drop every capture-time artifact so the next ``infer`` re-captures.
+
+ Prompt state (``set_prompt`` outputs) is kept; ``_full_sd`` was
+ released after the first capture and is reloaded by ``infer`` before
+ re-capture. The calibration disk cache makes re-capture cheap.
+ """
+ stale_prefixes = ("_sig_", "_g_", "_mlp1_")
+ stale_exact = ("_siglip_graph", "_qwen3_graph", "_dit_graph",
+ "_attn", "_vision_features", "_unit_scale",
+ "_qwen3_torch_graph", "_siglip_torch_graph",
+ # DiT static buffers/indexes depend on Se/masks and
+ # must be rebuilt on prompt-switch re-capture.
+ "_dit_in_state", "_dit_in_kvt", "_dit_in_kvi",
+ "_dit_in_noise", "_dit_out", "_dit_txt_idx",
+ "_dit_img_idx")
+ for name in list(vars(self)):
+ if name.startswith(stale_prefixes) or name in stale_exact:
+ delattr(self, name)
+ self._graphs_built = False
+ torch.cuda.empty_cache()
+
# ─────────────────────────────────────────────────────────────
# Public API
# ─────────────────────────────────────────────────────────────
- def set_prompt(self, prompt):
- """Tokenize prompt and prepare text embeddings for Qwen3 backbone."""
+ def set_prompt(self, prompt, input_ids=None):
+ """Tokenize prompt and prepare text embeddings for Qwen3 backbone.
+
+ Args:
+ prompt: language instruction (used for logging / change-detect).
+ input_ids: optional pre-computed HF token sequence (list[int]).
+ HF wraps the instruction in a chat template (system + user
+ headers) and expands each view into a per-view
+ ``![]()
…IMG_CONTEXT…`` block; a naive ``encode(prompt)``
+ yields a DIFFERENT sequence whose backbone features are
+ uncorrelated with HF's (cos~0), breaking DiT conditioning.
+ When supplied (e.g. from the serving aux builder, which
+ reproduces HF exactly), these ids are used verbatim.
+ """
if getattr(self, '_graphs_built', False):
raise RuntimeError(
"set_prompt() after the pipeline is built is not supported; "
"construct a new GrootTorchFrontendThor instance for a new prompt")
- from transformers import AutoTokenizer
-
- if not hasattr(self, '_tokenizer'):
- eagle_dir = pathlib.Path(__file__).parent.parent.parent.parent / "configs"
- # Try multiple tokenizer locations
- for tok_path in [
- str(self._checkpoint_path), # checkpoint dir may have tokenizer
- str(self._checkpoint_path / "tokenizer"), # subfolder
- # Local GROOT code Eagle dir
- str(pathlib.Path(__file__).parent.parent.parent.parent.parent /
- "GR00T" / "Isaac-GR00T" / "gr00t" / "model" / "modules" / "nvidia" / "Eagle-Block2A-2B-v2"),
- "nvidia/Eagle-Block2A-2B-v2", # HF hub (fallback)
- ]:
- try:
- self._tokenizer = AutoTokenizer.from_pretrained(tok_path, trust_remote_code=True)
- break
- except Exception:
- continue
- if not hasattr(self, '_tokenizer'):
- raise RuntimeError("Cannot load Qwen3 tokenizer")
- self._img_token_id = 151669 #
- self._img_start_id = 151670 # ![]()
- self._img_end_id = 151671 #
+ # Image special-token ids (fixed for the Eagle vocab).
+ self._img_token_id = 151669 #
+ self._img_start_id = 151670 # ![]()
+ self._img_end_id = 151671 #
+
+ if input_ids is not None:
+ full_ids = list(input_ids)
+ text_count = sum(1 for t in full_ids if t != self._img_token_id)
+ else:
+ from transformers import AutoTokenizer
- S_img = self._num_views * self.spv # image tokens after pixel unshuffle
- text_ids = self._tokenizer.encode(prompt, add_special_tokens=False)
- # Build: text + ![]()
+ *S_img +
- full_ids = text_ids + [self._img_start_id] + [self._img_token_id] * S_img + [self._img_end_id]
+ if not hasattr(self, '_tokenizer'):
+ for tok_path in [
+ str(self._checkpoint_path), # checkpoint dir may have tokenizer
+ str(self._checkpoint_path / "tokenizer"), # subfolder
+ # Local GROOT code Eagle dir
+ str(pathlib.Path(__file__).parent.parent.parent.parent.parent /
+ "GR00T" / "Isaac-GR00T" / "gr00t" / "model" / "modules" / "nvidia" / "Eagle-Block2A-2B-v2"),
+ "nvidia/Eagle-Block2A-2B-v2", # HF hub (fallback)
+ ]:
+ try:
+ self._tokenizer = AutoTokenizer.from_pretrained(tok_path, trust_remote_code=True)
+ break
+ except Exception:
+ continue
+ if not hasattr(self, '_tokenizer'):
+ raise RuntimeError("Cannot load Qwen3 tokenizer")
+
+ S_img = self._num_views * self.spv # image tokens after pixel unshuffle
+ text_ids = self._tokenizer.encode(prompt, add_special_tokens=False)
+ # Build: text + ![]()
+ *S_img +
+ full_ids = text_ids + [self._img_start_id] + [self._img_token_id] * S_img + [self._img_end_id]
+ text_count = len(text_ids)
self._input_ids = torch.tensor([full_ids], dtype=torch.long, device='cuda')
- self._text_len = len(text_ids)
+ self._text_len = text_count
self._Se = len(full_ids)
self._prompt_text = prompt
@@ -948,8 +1962,10 @@ def set_prompt(self, prompt):
self._image_mask = (self._input_ids == self._img_token_id) # [1, Se]
self._backbone_mask = torch.ones(1, self._Se, dtype=torch.bool, device='cuda')
- logger.info("Prompt set: '%s' (%d text + %d img = %d total tokens)",
- prompt[:50], self._text_len, S_img, self._Se)
+ n_img_tok = int(self._image_mask.sum())
+ logger.info("Prompt set: '%s' (%d non-img + %d IMG_CONTEXT = %d total tokens%s)",
+ prompt[:50], self._text_len, n_img_tok, self._Se,
+ ", HF input_ids" if input_ids is not None else "")
def infer_action_head(self, backbone_features, image_mask, backbone_mask,
state, action_horizon=None, noise_seed=None):
@@ -1207,7 +2223,11 @@ def _patch_embed_image(self, img_np):
pixel = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(fp16).cuda()
patches = pixel.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size)
- patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(S, -1) # [S, 588]
+ # HF NaFlex convert_images_to_patches flattens each patch as
+ # (ph, pw, C) — channel-LAST. The openpi-lineage (C, ph, pw) order
+ # feeds the HF-trained patch_embedding the wrong column order
+ # (embed cos 0.966 vs HF).
+ patches = patches.permute(0, 2, 3, 4, 5, 1).reshape(S, -1) # [S, 588]
# Linear patch embed + position embed → write into SigLIP buffer
self._sig_x[:S].copy_(
@@ -1445,45 +2465,111 @@ def infer(self, obs):
if not hasattr(self, '_input_ids'):
raise RuntimeError("Call set_prompt() before infer()")
- # ── Lazy graph capture on first call ──
+ # ── Lazy graph capture on first call; idle re-capture guard ──
+ # Thor invalidates captured graphs after a few seconds of GPU idle
+ # (replays then run fast but emit garbage) — re-capture instead.
+ if getattr(self, '_graphs_built', False) and self._graph_idle_stale():
+ logger.info("CUDA graphs idle-stale (>%.1fs) — re-capturing",
+ self._graph_idle_limit_s)
+ self.reset_graph_runtime()
if not getattr(self, '_graphs_built', False):
+ if not hasattr(self, '_full_sd'):
+ self._load_checkpoint()
self._capture_all_graphs(obs)
# ── 1. SigLIP: patch embed + graph replay + mlp1 ──
views = [obs['image']]
if 'wrist_image' in obs and self._num_views >= 2:
views.append(obs['wrist_image'])
- self._patch_embed_2views(views)
- self._siglip_graph.replay()
- torch.cuda.synchronize() # pixel_unshuffle needs SigLIP output
- self._run_pixel_unshuffle_mlp1()
+ if getattr(self, "_torch_siglip", None) is not None:
+ # Parity mode: HF-native bf16 SigLIP. Embeddings (patch embed +
+ # pos + NaFlex window split) run eager; encoder runs as a graph.
+ self._run_torch_siglip(views)
+ else:
+ self._patch_embed_2views(views)
+ self._siglip_graph.replay()
+ torch.cuda.synchronize() # pixel_unshuffle needs SigLIP output
+ if getattr(self, "_torch_siglip", None) is not None:
+ self._run_pixel_unshuffle_mlp1_bd()
+ else:
+ self._run_pixel_unshuffle_mlp1()
# ── 2. Build input embeddings (pre-allocated buffer) ──
- self._g_ie_buf.copy_(self._text_embeds.squeeze(0).to(fp16))
- self._g_ie_buf[self._image_mask[0]] = self._g_vision_out.to(fp16)
- fvk.gpu_copy(self._g_qwen3.b_x.data_ptr(), self._g_ie_buf.data_ptr(),
- self._Se * self.D_llm * 2, 0)
+ self._g_ie_buf.copy_(self._text_embeds.squeeze(0).to(self._bd))
+ self._g_ie_buf[self._image_mask[0]] = self._g_vision_out.to(self._bd)
# ── 3. Qwen3 graph replay (+ vlln) ──
- self._qwen3_graph.replay()
- torch.cuda.synchronize() # KV update needs Qwen3 output
-
- # ── 4. Update DiT KV + init noise ──
- bb = self._g_vlln_buf.unsqueeze(0)
- self._g_dit.b_kv_text.copy_((bb * self._g_non_img.unsqueeze(-1).to(fp16)).squeeze(0))
- self._g_dit.b_kv_img.copy_((bb * self._g_img_m.unsqueeze(-1).to(fp16)).squeeze(0))
-
- # Precompute cross-attention K/V projections (runs once, reused across 4 steps)
- self._g_dit.precompute_cross_kv()
- self._copy_state_feature_to_dit(
- obs.get('state', np.zeros(self.state_dim, dtype=np.float32)))
- self._g_dit.b_actions.normal_()
-
- # ── 5. DiT graph replay ──
- self._dit_graph.replay()
+ if getattr(self, "_torch_qwen3", None) is not None:
+ # Parity mode (--no-fp8): HF-native Qwen3Model (bf16, sdpa)
+ # captured as a CUDA graph — identical math to the HF eager
+ # baseline backbone at graph-replay latency.
+ self._tq_in.copy_(self._g_ie_buf.to(torch.bfloat16))
+ self._qwen3_torch_graph.replay()
+ torch.cuda.synchronize()
+ else:
+ fvk.gpu_copy(self._g_qwen3.b_x.data_ptr(),
+ self._g_ie_buf.data_ptr(),
+ self._Se * self.D_llm * 2, 0)
+ self._qwen3_graph.replay()
+ torch.cuda.synchronize() # KV update needs Qwen3 output
+ # Thor can silently invalidate captured graphs (replays stay fast
+ # but emit NaN/garbage — see docs/thor_gpu_idle_reset_workaround.md).
+ # The vlln LN output is the first cheap full-width signal after the
+ # backbone replay; if it is non-finite, re-capture once and retry
+ # this observation instead of serving garbage.
+ if not bool(torch.isfinite(self._g_vlln_buf).all()):
+ logger.error(
+ "Non-finite backbone output detected after graph replay — "
+ "graphs invalidated; re-capturing and retrying this frame")
+ self.reset_graph_runtime()
+ if not hasattr(self, '_full_sd'):
+ self._load_checkpoint()
+ self._capture_all_graphs(obs)
+ if getattr(self, "_torch_siglip", None) is not None:
+ self._run_torch_siglip(views)
+ self._run_pixel_unshuffle_mlp1_bd()
+ else:
+ self._patch_embed_2views(views)
+ self._siglip_graph.replay()
+ torch.cuda.synchronize()
+ self._run_pixel_unshuffle_mlp1()
+ torch.cuda.synchronize()
+ self._g_ie_buf.copy_(self._text_embeds.squeeze(0).to(self._bd))
+ self._g_ie_buf[self._image_mask[0]] = self._g_vision_out.to(self._bd)
+ if getattr(self, "_torch_qwen3", None) is not None:
+ self._tq_in.copy_(self._g_ie_buf.to(torch.bfloat16))
+ self._qwen3_torch_graph.replay()
+ torch.cuda.synchronize()
+ else:
+ fvk.gpu_copy(self._g_qwen3.b_x.data_ptr(),
+ self._g_ie_buf.data_ptr(),
+ self._Se * self.D_llm * 2, 0)
+ self._qwen3_graph.replay()
+ torch.cuda.synchronize()
+ if not bool(torch.isfinite(self._g_vlln_buf).all()):
+ raise RuntimeError(
+ "backbone graph output still non-finite after re-capture")
+
+ # ── 4. DiT ──
+ # parity mode: HF-faithful fp32 torch DiT (max accuracy).
+ # fast mode (--no-parity): kernel DiT (concat bug fixed; ~4ms).
+ if not self.parity:
+ kd = self._g_dit
+ self._fill_dit_kv()
+ kd.precompute_cross_kv()
+ self._copy_state_feature_to_dit(
+ obs.get('state', np.zeros(self.state_dim, dtype=np.float32)))
+ kd.b_actions.normal_()
+ self._dit_graph.replay()
+ torch.cuda.synchronize()
+ actions = kd.b_actions
+ else:
+ actions = self._run_torch_dit(
+ obs.get('state', np.zeros(self.state_dim, dtype=np.float32)))
torch.cuda.synchronize()
- return {'actions': self._g_dit.b_actions.squeeze(0).cpu().numpy()}
+ self._last_graph_use = time.monotonic()
+ return {'actions': actions.squeeze(0).float().cpu().numpy()}
# ─────────────────────────────────────────────────────────────
# FP8 Activation Calibration
@@ -1789,6 +2875,8 @@ def _calibrate_dit(self):
def _calibrate_single_frame(self, obs):
"""N=1 path: build pipeline (which calibrates) and snapshot the spec."""
if not getattr(self, "_graphs_built", False):
+ if not hasattr(self, "_full_sd"):
+ self._load_checkpoint()
self._capture_all_graphs(obs)
self._calibrated = True
self._precision_spec = self._snapshot_precision_spec(
@@ -1822,6 +2910,8 @@ def _calibrate_multi_frame(
n, percentile)
if not getattr(self, "_graphs_built", False):
+ if not hasattr(self, "_full_sd"):
+ self._load_checkpoint()
self._capture_all_graphs(obs_list[0], release_full_sd=False)
# ── Pass A: collect per-sample Qwen3 amax (uses _full_sd) ──
@@ -1898,9 +2988,7 @@ def _collect_dit_amax_for_obs(self, obs):
vlln_b.data_ptr(), self._g_vlln_buf.data_ptr(),
Se, self.D_llm, 1e-5, 0)
torch.cuda.synchronize()
- bb = self._g_vlln_buf.unsqueeze(0)
- self._g_dit.b_kv_text.copy_((bb * self._g_non_img.unsqueeze(-1).to(fp16)).squeeze(0))
- self._g_dit.b_kv_img.copy_((bb * self._g_img_m.unsqueeze(-1).to(fp16)).squeeze(0))
+ self._fill_dit_kv()
self._g_dit.precompute_cross_kv()
torch.cuda.synchronize()
@@ -1992,12 +3080,90 @@ def _patch_embed_2views(self, views):
arr = (arr - 0.5) / 0.5
pixel = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(fp16).cuda()
patches = pixel.unfold(2, 14, 14).unfold(3, 14, 14)
- patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(self.spv_raw, -1)
+ # HF NaFlex patch order is channel-LAST (ph, pw, C); see
+ # _patch_embed_image for why the openpi (C, ph, pw) order is wrong.
+ patches = patches.permute(0, 2, 3, 4, 5, 1).reshape(self.spv_raw, -1)
embedded = (F.linear(patches.float(), self._sig_patch_w.float(),
self._sig_patch_b.float())
+ self._sig_pos_embed[:self.spv_raw].float()).to(fp16)
self._sig_x[idx * self.spv_raw:(idx + 1) * self.spv_raw].copy_(embedded)
+ def _run_torch_siglip(self, views):
+ """Parity mode: HF-native bf16 SigLIP (embeddings eager + encoder graph)."""
+ if getattr(self, "_tsig_u8", None) is None:
+ self._tsig_u8 = torch.empty(
+ self._num_views, self.image_size, self.image_size, 3,
+ dtype=torch.uint8, device='cuda')
+ for i, v in enumerate(views):
+ arr = np.array(v, dtype=np.uint8, copy=True)
+ if not arr.flags['C_CONTIGUOUS']:
+ arr = np.ascontiguousarray(arr)
+ self._tsig_u8[i].copy_(torch.from_numpy(arr))
+ x = self._tsig_u8[i].to(torch.float32).div_(255.0).sub_(0.5).div_(0.5)
+ self._tsig_in[i].copy_(x.permute(2, 0, 1).to(torch.bfloat16))
+ if not getattr(self, "_sig_embed_in_graph", False):
+ if getattr(self, "_tsig_gidx", None) is not None:
+ with torch.no_grad():
+ pv = self._tsig_in
+ nv, _, nph, ps, npw = (pv.shape[0], pv.shape[1],
+ self._tsig_nph, self._tsig_ps,
+ self._tsig_npw)
+ patched = (pv.reshape(nv, 3, nph, ps, npw, ps)
+ .permute(0, 2, 4, 3, 5, 1)
+ .reshape(nv * nph * npw, -1))
+ emb_mod = self._torch_siglip.vision_model.embeddings
+ pe = emb_mod.patch_embedding(patched)
+ self._tw_in.copy_(
+ (pe + self._tsig_pos).index_select(
+ 0, self._tsig_gidx).unsqueeze(0))
+ else:
+ with torch.no_grad():
+ _wt, _, _, _ = self._torch_siglip.vision_model.embeddings(
+ [self._tsig_in[i:i + 1] for i in range(len(views))])
+ self._tw_in.copy_(_wt)
+ if getattr(self, "_siglip_torch_graph", None) is not None:
+ self._siglip_torch_graph.replay()
+ else:
+ # Capture-time fallback: encoder runs eager.
+ with torch.no_grad():
+ _enc = self._torch_siglip.vision_model.encoder(
+ inputs_embeds=self._tw_in,
+ win_meta_list=self._tw_meta,
+ spatial_shapes=self._tw_shapes)
+ _lh = self._torch_siglip.vision_model.post_layernorm(
+ _enc.last_hidden_state)
+ self._g_sig_postln.copy_(
+ _lh[:, self._tw_rm].to(self._bd).squeeze(0))
+ torch.cuda.synchronize()
+
+ def _run_pixel_unshuffle_mlp1_bd(self):
+ """Parity mode: pixel unshuffle + mlp1 in bf16 torch (HF-exact math)."""
+ if getattr(self, "_mlp1_fc1_w_lin", None) is None:
+ # Cache F.linear-layout bf16 weights once (the fp16 sources are
+ # transposed for the kernel NN convention); avoids per-frame
+ # .T.contiguous()/.float() re-materialization.
+ self._mlp1_ln_w_f32 = self._mlp1_ln_w.to(torch.float32)
+ self._mlp1_ln_b_f32 = self._mlp1_ln_b.to(torch.float32)
+ self._mlp1_fc1_w_lin = self._mlp1_fc1_w.t().contiguous().to(self._bd)
+ self._mlp1_fc1_b_bd = self._mlp1_fc1_b.to(self._bd)
+ self._mlp1_fc2_w_lin = self._mlp1_fc2_w.t().contiguous().to(self._bd)
+ self._mlp1_fc2_b_bd = self._mlp1_fc2_b.to(self._bd)
+ nH = int(math.sqrt(self.spv_raw))
+ D, S_img = self.D_sig, self._num_views * self.spv
+ all_flat = []
+ for v in range(self._num_views):
+ view_out = self._g_sig_postln[v * self.spv_raw:(v + 1) * self.spv_raw]
+ spatial = view_out.view(1, nH, nH, D).permute(0, 3, 1, 2)
+ flat = F.pixel_unshuffle(spatial, 2).view(self.mlp1_in, -1).T.contiguous()
+ all_flat.append(flat)
+ combined = torch.cat(all_flat, dim=0)
+ h = F.layer_norm(combined.float(), (self.mlp1_in,),
+ self._mlp1_ln_w_f32, self._mlp1_ln_b_f32,
+ 1e-5).to(self._bd)
+ h = F.gelu(F.linear(h, self._mlp1_fc1_w_lin, self._mlp1_fc1_b_bd))
+ h = F.linear(h, self._mlp1_fc2_w_lin, self._mlp1_fc2_b_bd)
+ self._g_vision_out.copy_(h)
+
def _run_pixel_unshuffle_mlp1(self):
"""Pixel unshuffle + mlp1 (runs outside graph, ~0.3ms)."""
nH = int(math.sqrt(self.spv_raw))
@@ -2037,6 +3203,29 @@ def _capture_all_graphs(self, obs, release_full_sd: bool = True):
from flash_rt.models.groot.pipeline_thor import CKernelQwen3, CKernelDiTHead
logger.info("Capturing CUDA Graphs for E2E pipeline...")
+ # Extract bf16 action-head weights for the HF-faithful torch DiT
+ # path (must happen while _full_sd is alive).
+ self._setup_torch_dit()
+ if not self.use_fp8 and not hasattr(self, "_torch_qwen3"):
+ self._setup_torch_qwen3()
+ if not self.use_fp8:
+ self._setup_qwen3_fp4()
+ if not self.use_fp8 and not hasattr(self, "_torch_siglip"):
+ self._setup_torch_siglip()
+ # Static NaFlex window meta (shapes fixed per image_size):
+ # needed by the capture-time eager encoder run before 6c.
+ self._tsig_in = torch.empty(
+ self._num_views, 3, self.image_size, self.image_size,
+ dtype=torch.bfloat16, device='cuda')
+ with torch.no_grad():
+ _wt0, self._tw_meta, self._tw_shapes, self._tw_rm = \
+ self._torch_siglip.vision_model.embeddings(
+ [torch.zeros(1, 3, self.image_size, self.image_size,
+ dtype=torch.bfloat16, device='cuda')
+ for _ in range(self._num_views)])
+ self._tw_rm = self._tw_rm.to(torch.long).cuda()
+ self._tw_in = torch.empty_like(_wt0)
+
Se = self._Se
S_sig = self._num_views * self.spv_raw
S_img = self._num_views * self.spv
@@ -2044,15 +3233,22 @@ def _capture_all_graphs(self, obs, release_full_sd: bool = True):
# ── 1. Build CKernel objects FIRST (from sd dict, before SigLIP FP8 init) ──
# This ensures cuBLASLt workspace is not polluted by SigLIP FP8 quantization
self._g_qwen3 = CKernelQwen3(self._full_sd, Se, use_fp8=self.use_fp8)
- vlln_w = self._vlln_w.to(fp16)
- vlln_b = self._vlln_b.to(fp16)
+ # Persistent (NOT local): these pointers are baked into the Qwen3
+ # graph's vlln LayerNorm. Local tensors would be freed when this
+ # function returns and their blocks reused — replays would then
+ # dereference dangling memory and emit NaN.
+ self._g_vlln_w = self._vlln_w.to(fp16).contiguous()
+ self._g_vlln_b = self._vlln_b.to(fp16).contiguous()
+ vlln_w = self._g_vlln_w
+ vlln_b = self._g_vlln_b
self._g_non_img = (~self._image_mask) & self._backbone_mask
self._g_img_m = self._image_mask & self._backbone_mask
T = self.action_horizon
self._g_dit = CKernelDiTHead(self._full_sd, self._embodiment_id, T,
- (1, Se, self.D_llm), use_fp8=self.use_fp8)
+ (1, Se, self.D_llm), use_fp8=self.use_fp8,
+ num_steps=self.num_steps)
# ── 2. SigLIP init (after CKernel objects, so FP8 quant doesn't pollute cuBLAS state) ──
self._load_siglip2_weights(self._full_sd)
@@ -2073,11 +3269,20 @@ def _capture_all_graphs(self, obs, release_full_sd: bool = True):
qwen3_seq_max=Se,
sa=self._g_dit.Sa,
s_kv=self._g_dit.S_kv,
+ # Live HF N1.6 runs SigLIP2 under sdpa, which ignores the
+ # NaFlex per-view seq_len_list segmentation -> all packed
+ # views attend jointly (cross-view full attention).
+ siglip_cross_view=True,
),
siglip_slots={
"qkv": self._sig_qkv.data_ptr(),
"O": self._sig_attn.data_ptr(),
"D": self.D_sig,
+ # tensor views let the backend run exact torch sdpa for
+ # SigLIP (strided FMHA kernel unreliable at non-power-of-2
+ # seqs like 648)
+ "qkv_t": self._sig_qkv,
+ "O_t": self._sig_attn,
},
qwen3_slots={
"ctx": self._g_qwen3.ctx, # Qwen3's cuBLAS handle
@@ -2112,16 +3317,19 @@ def _capture_all_graphs(self, obs, release_full_sd: bool = True):
self._g_qwen3.attn = self._attn
self._g_dit.attn = self._attn
+ # HF N1.6 (sdpa) does cross-view full attention over all packed
+ # views, so run SigLIP as ONE batch-1 sequence of S_sig tokens
+ # instead of num_views independent 256-token views.
sig_dims = {'S': S_sig, 'D': self.D_sig, 'H': self.H_sig,
'NH': self.NH_sig, 'HD': self.HD_sig, 'L': self.L_sig,
- 'num_views': self._num_views, 'seq_per_view': self.spv_raw}
+ 'num_views': 1, 'seq_per_view': S_sig}
# Allocate graph-persistent buffers
- self._g_sig_postln = torch.empty(S_sig, self.D_sig, dtype=fp16, device='cuda')
- self._g_mlp1_ln = torch.empty(S_img, self.mlp1_in, dtype=fp16, device='cuda')
- self._g_mlp1_fc1 = torch.empty(S_img, self.D_llm, dtype=fp16, device='cuda')
- self._g_vision_out = torch.empty(S_img, self.D_llm, dtype=fp16, device='cuda')
- self._g_vlln_buf = torch.empty(Se, self.D_llm, dtype=fp16, device='cuda')
+ self._g_sig_postln = torch.empty(S_sig, self.D_sig, dtype=self._bd, device='cuda')
+ self._g_mlp1_ln = torch.empty(S_img, self.mlp1_in, dtype=self._bd, device='cuda')
+ self._g_mlp1_fc1 = torch.empty(S_img, self.D_llm, dtype=self._bd, device='cuda')
+ self._g_vision_out = torch.empty(S_img, self.D_llm, dtype=self._bd, device='cuda')
+ self._g_vlln_buf = torch.empty(Se, self.D_llm, dtype=self._bd, device='cuda')
# ── 3. SigLIP graph ──
stream = torch.cuda.Stream()
@@ -2148,15 +3356,19 @@ def _run_sig(si):
views = [obs['image']]
if 'wrist_image' in obs and self._num_views >= 2:
views.append(obs['wrist_image'])
- self._patch_embed_2views(views)
- self._siglip_graph.replay()
- torch.cuda.synchronize()
- self._run_pixel_unshuffle_mlp1()
- torch.cuda.synchronize()
+ if getattr(self, "_torch_siglip", None) is not None:
+ self._run_torch_siglip(views)
+ self._run_pixel_unshuffle_mlp1_bd()
+ else:
+ self._patch_embed_2views(views)
+ self._siglip_graph.replay()
+ torch.cuda.synchronize()
+ self._run_pixel_unshuffle_mlp1()
+ torch.cuda.synchronize()
input_embeds = self._text_embeds.clone()
input_embeds[0, self._image_mask[0]] = self._g_vision_out.to(input_embeds.dtype)
- ie_fp16 = input_embeds.squeeze(0).to(fp16).contiguous()
+ ie_fp16 = input_embeds.squeeze(0).to(self._bd).contiguous()
# ── 4. Calibrate + run Qwen3 ──
self._calibrate_qwen3(ie_fp16)
@@ -2169,9 +3381,7 @@ def _run_sig(si):
Se, self.D_llm, 1e-5, 0)
torch.cuda.synchronize()
- bb = self._g_vlln_buf.unsqueeze(0)
- self._g_dit.b_kv_text.copy_((bb * self._g_non_img.unsqueeze(-1).to(fp16)).squeeze(0))
- self._g_dit.b_kv_img.copy_((bb * self._g_img_m.unsqueeze(-1).to(fp16)).squeeze(0))
+ self._fill_dit_kv()
# Precompute cross-attention K/V projections (constant across steps)
self._g_dit.precompute_cross_kv()
@@ -2223,6 +3433,150 @@ def _run_sig(si):
torch.cuda.synchronize()
logger.info(" Qwen3 graph captured (Se=%d)", Se)
+ # ── 6b. Torch Qwen3 graph (parity mode) ──
+ if getattr(self, "_torch_qwen3", None) is not None:
+ self._tq_in = torch.empty(
+ 1, Se, self.D_llm, dtype=torch.bfloat16, device='cuda')
+
+ def _run_tq():
+ with torch.no_grad():
+ _out = self._torch_qwen3(
+ inputs_embeds=self._tq_in).last_hidden_state.squeeze(0)
+ _v = torch.nn.functional.layer_norm(
+ _out.float(), (self.D_llm,),
+ self._g_vlln_w.float(), self._g_vlln_b.float(),
+ 1e-5).to(self._bd)
+ self._g_vlln_buf.copy_(_v)
+
+ stream_tq = torch.cuda.Stream()
+ with torch.cuda.stream(stream_tq):
+ for _ in range(2):
+ _run_tq()
+ torch.cuda.synchronize()
+ self._qwen3_torch_graph = torch.cuda.CUDAGraph()
+ with torch.cuda.stream(stream_tq):
+ self._qwen3_torch_graph.capture_begin()
+ _run_tq()
+ self._qwen3_torch_graph.capture_end()
+ torch.cuda.synchronize()
+ logger.info(" Torch Qwen3 graph captured (Se=%d)", Se)
+
+ # ── 6c. Torch SigLIP graph (parity mode) ──
+ # The NaFlex window split contains a non-capturable GPU index op, so
+ # embeddings run eager per frame; only the 27-layer encoder (+post-LN
+ # + reverse mapping) is captured, over a static windows buffer.
+ if getattr(self, "_torch_siglip", None) is not None:
+ nv = self._num_views
+ if not hasattr(self, "_tsig_in"):
+ self._tsig_in = torch.empty(
+ nv, 3, self.image_size, self.image_size,
+ dtype=torch.bfloat16, device='cuda')
+ if not hasattr(self, "_tw_in"):
+ with torch.no_grad():
+ _wt, self._tw_meta, self._tw_shapes, self._tw_rm = \
+ self._torch_siglip.vision_model.embeddings(
+ [self._tsig_in[i:i + 1] for i in range(nv)])
+ self._tw_rm = self._tw_rm.to(torch.long).cuda()
+ self._tw_in = torch.empty_like(_wt)
+
+ # Fast embeddings: the NaFlex window split is a static gather for a
+ # fixed image size, and the antialias positional-embedding resize
+ # depends only on static shapes. Precompute both once so the
+ # per-frame path is patchify + patch_embedding + pos_add + gather
+ # (all capture-safe), skipping the unfold/im2col + antialias
+ # interpolate (~0.4 ms/frame). Verified bit-exact vs the HF
+ # embeddings forward (max diff 0.0).
+ if not hasattr(self, "_tsig_gidx"):
+ emb = self._torch_siglip.vision_model.embeddings
+ ps = emb.patch_size
+ nph = self.image_size // ps
+ npw = self.image_size // ps
+ self._tsig_ps, self._tsig_nph, self._tsig_npw = ps, nph, npw
+ bchw = [torch.Size((1, 3, self.image_size, self.image_size))
+ for _ in range(nv)]
+ ss = emb.get_spatial_shapes(bchw)
+ fake = torch.arange(nv * nph * npw, dtype=torch.float32,
+ device='cuda').view(1, nv * nph * npw, 1)
+ wt, _, _ = emb.split_patch_embeddings_to_windows_with_meta(
+ fake, ss, emb.window_size)
+ self._tsig_gidx = wt.view(-1).long().cuda()
+ pos = emb.position_embedding.weight.reshape(
+ emb.position_embedding_size, emb.position_embedding_size, -1)
+ self._tsig_pos = emb.resize_positional_embeddings(pos, ss)[0]
+
+ def _sig_embed_to_twin():
+ # Capture-safe fast embeddings: patchify + patch_embedding +
+ # precomputed pos add + static gather, _tsig_in -> _tw_in.
+ pv = self._tsig_in
+ nv, nph, ps, npw = (pv.shape[0], self._tsig_nph,
+ self._tsig_ps, self._tsig_npw)
+ patched = (pv.reshape(nv, 3, nph, ps, npw, ps)
+ .permute(0, 2, 4, 3, 5, 1)
+ .reshape(nv * nph * npw, -1))
+ pe = self._torch_siglip.vision_model.embeddings.patch_embedding(
+ patched)
+ self._tw_in.copy_(
+ (pe + self._tsig_pos).index_select(
+ 0, self._tsig_gidx).unsqueeze(0))
+
+ def _run_ts():
+ with torch.no_grad():
+ _sig_embed_to_twin()
+ _enc = self._torch_siglip.vision_model.encoder(
+ inputs_embeds=self._tw_in,
+ win_meta_list=self._tw_meta,
+ spatial_shapes=self._tw_shapes)
+ _lh = self._torch_siglip.vision_model.post_layernorm(
+ _enc.last_hidden_state)
+ self._g_sig_postln.copy_(
+ _lh[:, self._tw_rm].to(self._bd).squeeze(0))
+
+ stream_ts = torch.cuda.Stream()
+ with torch.cuda.stream(stream_ts):
+ for _ in range(2):
+ _run_ts()
+ torch.cuda.synchronize()
+ # Try capturing embeddings+encoder together. If any NaFlex op is
+ # not capture-safe, fall back to an encoder-only graph and run the
+ # (fast) embeddings eager per frame.
+ self._sig_embed_in_graph = False
+ try:
+ g = torch.cuda.CUDAGraph()
+ with torch.cuda.stream(stream_ts):
+ g.capture_begin()
+ _run_ts()
+ g.capture_end()
+ torch.cuda.synchronize()
+ self._siglip_torch_graph = g
+ self._sig_embed_in_graph = True
+ logger.info(" Torch SigLIP graph captured embeddings+encoder "
+ "(S=%d)", nv * self.spv_raw)
+ except Exception as e: # noqa: BLE001
+ logger.warning("SigLIP embeddings not capture-safe (%r); "
+ "encoder-only graph + eager embeddings", e)
+ torch.cuda.synchronize()
+ g = torch.cuda.CUDAGraph()
+ def _enc_only():
+ with torch.no_grad():
+ _enc = self._torch_siglip.vision_model.encoder(
+ inputs_embeds=self._tw_in,
+ win_meta_list=self._tw_meta,
+ spatial_shapes=self._tw_shapes)
+ _lh = self._torch_siglip.vision_model.post_layernorm(
+ _enc.last_hidden_state)
+ self._g_sig_postln.copy_(
+ _lh[:, self._tw_rm].to(self._bd).squeeze(0))
+ with torch.cuda.stream(stream_ts):
+ for _ in range(2):
+ _enc_only()
+ torch.cuda.synchronize()
+ with torch.cuda.stream(stream_ts):
+ g.capture_begin()
+ _enc_only()
+ g.capture_end()
+ torch.cuda.synchronize()
+ self._siglip_torch_graph = g
+
# Reuse the exact buffer the Qwen3 graph baked into its captured
# ``gpu_copy(b_x, ie_fp16)`` so per-frame ``_g_ie_buf`` writes land at
# the pointer the graph reads on replay. (A fresh allocation here only
diff --git a/flash_rt/hardware/thor/attn_backend_groot.py b/flash_rt/hardware/thor/attn_backend_groot.py
index 08fe2ae6..3547065a 100644
--- a/flash_rt/hardware/thor/attn_backend_groot.py
+++ b/flash_rt/hardware/thor/attn_backend_groot.py
@@ -211,6 +211,21 @@ def run(self, site: str, layer_idx: int, q_seq: int,
HD = site_spec.head_dim
if kv_seq is None:
kv_seq = q_seq
+ # The strided FMHA kernel is only validated for power-of-two
+ # sequences on real data (648-seq cross-view SigLIP at the 252
+ # training resolution diverges silently). Use torch sdpa when
+ # tensor views are provided (graph-capturable, exact).
+ if "qkv_t" in s:
+ import torch.nn.functional as F
+ qkv_t = s["qkv_t"]
+ O_t = s["O_t"]
+ seq = int(q_seq)
+ q = qkv_t[:seq, :D].view(seq, NH, HD).transpose(0, 1)
+ k = qkv_t[:seq, D:2 * D].view(seq, NH, HD).transpose(0, 1)
+ v = qkv_t[:seq, 2 * D:].view(seq, NH, HD).transpose(0, 1)
+ o = F.scaled_dot_product_attention(q, k, v)
+ O_t[:seq].copy_(o.transpose(0, 1).reshape(seq, D))
+ return int(s["O"])
stride = 3 * D
Q = int(s["qkv"])
K = Q + D * 2
@@ -261,7 +276,8 @@ def run(self, site: str, layer_idx: int, q_seq: int,
def make_groot_attention_spec(*, num_views: int, qwen3_seq_max: int,
- sa: int, s_kv: int) -> AttentionSpec:
+ sa: int, s_kv: int,
+ siglip_cross_view: bool = False) -> AttentionSpec:
"""Build the GROOT AttentionSpec (4 sites).
Args:
@@ -269,13 +285,28 @@ def make_groot_attention_spec(*, num_views: int, qwen3_seq_max: int,
qwen3_seq_max: max Qwen3 sequence length (prompt + vision tokens).
sa: DiT action sequence length (hidden tokens = 1 state + T actions).
s_kv: DiT cross-attention KV length (non_img + img backbone features).
+ siglip_cross_view: N1.6 HF semantics — the live HF baseline runs
+ SigLIP2 under sdpa, which ignores the NaFlex per-view
+ ``seq_len_list`` segmentation, so all packed views attend
+ jointly. Model the SigLIP site as one batch-1 sequence of
+ ``num_views * 256`` tokens instead of ``num_views``
+ independent 256-token views. Default False keeps the legacy
+ per-view behaviour.
"""
spec = AttentionSpec()
- spec.add_site(
- "siglip",
- num_layers=27, num_q_heads=16, num_kv_heads=16, head_dim=72,
- max_q_seq=256, max_kv_seq=256, batch_axis=int(num_views),
- )
+ if siglip_cross_view:
+ spec.add_site(
+ "siglip",
+ num_layers=27, num_q_heads=16, num_kv_heads=16, head_dim=72,
+ max_q_seq=256 * int(num_views), max_kv_seq=256 * int(num_views),
+ batch_axis=1,
+ )
+ else:
+ spec.add_site(
+ "siglip",
+ num_layers=27, num_q_heads=16, num_kv_heads=16, head_dim=72,
+ max_q_seq=256, max_kv_seq=256, batch_axis=int(num_views),
+ )
spec.add_site(
"qwen3",
num_layers=16, num_q_heads=16, num_kv_heads=16, head_dim=128,
diff --git a/flash_rt/models/groot/pipeline_thor.py b/flash_rt/models/groot/pipeline_thor.py
index 00ccc2ce..d6d45a81 100644
--- a/flash_rt/models/groot/pipeline_thor.py
+++ b/flash_rt/models/groot/pipeline_thor.py
@@ -597,6 +597,9 @@ def forward(self, x_in, s=0):
fvk.gpu_strided_copy_fp16(self.b_qkv.data_ptr(), self.b_attn.data_ptr(), Se, NHKV*HD, self.QKV, NHQ*HD+NHKV*HD, s)
fvk.gpu_repeat_interleave_heads(self.b_k.data_ptr(), self.b_k_exp.data_ptr(), Se, NHKV, HD, NHQ//NHKV, s)
fvk.gpu_repeat_interleave_heads(self.b_attn.data_ptr(), self.b_v_exp.data_ptr(), Se, NHKV, HD, NHQ//NHKV, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('k', self.b_k.clone()))
+ self._dbg.append(('v', self.b_attn.clone()))
fvk.gpu_fill_neginf_fp16(self.b_logits.data_ptr(), self.b_logits.nelement(), s)
if self.attn is not None:
self.attn.run("qwen3", i, q_seq=Se, stream=s)
@@ -604,6 +607,8 @@ def forward(self, x_in, s=0):
fvk.attention_mha_fp16(self.ctx, self.b_q.data_ptr(), self.b_k_exp.data_ptr(), self.b_v_exp.data_ptr(),
self.b_logits.data_ptr(), self.b_o.data_ptr(), Se, Se, NHQ, HD, 1.0/math.sqrt(HD), s)
self.gemm.fp16_nn(self.b_o.data_ptr(), w['o_w'].data_ptr(), self.b_xn.data_ptr(), Se, D, D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('o', self.b_xn.clone()))
fvk.residual_add_fp16(self.b_x.data_ptr(), self.b_xn.data_ptr(), Se * D, s)
fvk.rms_norm_fp16(self.b_x.data_ptr(), w['ln2_w'].data_ptr(), self.b_xn.data_ptr(), Se, D, 1e-6, s)
if self.use_fp8:
@@ -628,6 +633,9 @@ def forward(self, x_in, s=0):
self.gemm.fp16_nn(self.b_gu.data_ptr(), w['down_fp16'].data_ptr(),
self.b_down.data_ptr(), Se, D, H, s)
fvk.residual_add_fp16(self.b_x.data_ptr(), self.b_down.data_ptr(), Se * D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('ff', self.b_down.clone()))
+ self._dbg.append(('x', self.b_x.clone()))
fvk.rms_norm_fp16(self.b_x.data_ptr(), self.final_norm_w.data_ptr(), self.b_xn.data_ptr(), Se, D, 1e-6, s)
return self.b_xn
@@ -640,14 +648,21 @@ class CKernelDiTHead:
"""
def __init__(self, sd_or_path, embodiment_id, action_horizon, backbone_shape,
- use_fp8=True):
+ use_fp8=True, num_steps=4):
self.gemm = fvk.GemmRunner()
self.ctx = fvk.FvkContext()
self.use_fp8 = bool(use_fp8)
self.D = 1536; self.H = 6144; self.NH = 32; self.HD = 48
- self.L = 32; self.action_dim = 128; self.num_steps = 4
+ self.L = 32; self.action_dim = 128; self.num_steps = int(num_steps)
self.T = action_horizon; self.Sa = 1 + action_horizon
self.S_kv = backbone_shape[1]; self.D_kv = backbone_shape[2]
+ # Compact-KV counts: number of VALID backbone rows written into
+ # b_kv_text / b_kv_img by the frontend. HF excludes the other
+ # modality's tokens from each cross-attention layer via an
+ # attention mask; we instead attend only to the valid rows, so the
+ # per-layer kv length must come from these counts (not S_kv).
+ self.n_kv_text = self.S_kv
+ self.n_kv_img = self.S_kv
# Optional AttentionBackend (ThorGrootAttnBackend); set post-construct
# by the frontend. When None, _run_step uses the direct
# fvk.attention_mha_fp16 calls below (bit-identical legacy path).
@@ -721,7 +736,9 @@ def _precompute(self, sd):
proj_out_1_w = sd["action_head.model.proj_out_1.weight"].T.contiguous().to(fp16)
proj_out_1_b = sd["action_head.model.proj_out_1.bias"].to(fp16)
half_dim = 128
- exp = -torch.arange(half_dim, dtype=torch.float32, device='cuda') * (math.log(10000.0) / half_dim)
+ # diffusers Timesteps(downscale_freq_shift=1): denominator is
+ # (half_dim - 1), NOT half_dim — HF's exact convention.
+ exp = -torch.arange(half_dim, dtype=torch.float32, device='cuda') * (math.log(10000.0) / (half_dim - 1))
half_d = D // 2
exp_d = (-torch.arange(half_d, dtype=torch.float, device='cuda') * (math.log(10000.0) / half_d)).exp()
self.ada_scales = []; self.ada_shifts = []
@@ -755,6 +772,11 @@ def _precompute(self, sd):
self.out_scales = torch.stack(self.out_scales); self.out_shifts = torch.stack(self.out_shifts)
self.action_time_embeds = torch.stack(self.action_time_embeds)
+ def set_kv_counts(self, n_text, n_img):
+ """Report the number of valid compact rows in b_kv_text / b_kv_img."""
+ self.n_kv_text = int(n_text)
+ self.n_kv_img = int(n_img)
+
def precompute_cross_kv(self, s=0):
"""Precompute K/V projections for all cross-attention blocks.
@@ -762,22 +784,28 @@ def precompute_cross_kv(self, s=0):
through per-block K/V weight matrices. The backbone features are CONSTANT
across diffusion steps, so these projections only need to run once.
+ Only the first n_kv_text / n_kv_img rows are valid (compact layout);
+ the other modality's tokens are excluded, matching HF's attention mask.
+
Without this: 16 blocks × 2 projections × 4 steps = 128 GEMMs
With this: 16 blocks × 2 projections × 1 time = 32 GEMMs (save 96)
"""
- D, S_kv = self.D, self.S_kv
+ D = self.D
for block_idx in range(self.L // 2):
l = block_idx * 2 # cross-attention layers: 0, 2, 4, ..., 30
w = self.dit[l]
- kv_src = self.b_kv_text if l % 4 == 0 else self.b_kv_img
+ if l % 4 == 0:
+ kv_src, n_kv = self.b_kv_text, self.n_kv_text
+ else:
+ kv_src, n_kv = self.b_kv_img, self.n_kv_img
self._fp16_gemm(kv_src.data_ptr(), w['k_w'].data_ptr(),
- self._precomp_k[block_idx].data_ptr(), S_kv, D, self.D_kv, s)
+ self._precomp_k[block_idx].data_ptr(), n_kv, D, self.D_kv, s)
fvk.add_bias_fp16(self._precomp_k[block_idx].data_ptr(),
- w['k_b'].data_ptr(), S_kv, D, s)
+ w['k_b'].data_ptr(), n_kv, D, s)
self._fp16_gemm(kv_src.data_ptr(), w['v_w'].data_ptr(),
- self._precomp_v[block_idx].data_ptr(), S_kv, D, self.D_kv, s)
+ self._precomp_v[block_idx].data_ptr(), n_kv, D, self.D_kv, s)
fvk.add_bias_fp16(self._precomp_v[block_idx].data_ptr(),
- w['v_b'].data_ptr(), S_kv, D, s)
+ w['v_b'].data_ptr(), n_kv, D, s)
def _alloc_buffers(self):
D, H, T, Sa = self.D, self.H, self.T, self.Sa
@@ -871,20 +899,36 @@ def _run_step(self, step, s=0):
fvk.gpu_cast_fp32_to_fp16(self.b_actions.data_ptr(), self.b_actions_fp16.data_ptr(), T*self.action_dim, s)
self._fp16_gemm(self.b_actions_fp16.data_ptr(), self.ae_w1.data_ptr(), self.b_a_emb.data_ptr(), T, D, self.action_dim, s)
fvk.add_bias_fp16(self.b_a_emb.data_ptr(), self.ae_b1.data_ptr(), T, D, s)
- fvk.gpu_copy(self.b_concat.data_ptr(), self.b_a_emb.data_ptr(), T*D*2, s)
- fvk.gpu_copy(self.b_concat.data_ptr()+T*D*2, self.action_time_embeds[step].data_ptr(), T*D*2, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('ae1', self.b_a_emb.clone()))
+ # gpu_copy is a raw memcpy: it cannot place (T,D) rows into the first
+ # D columns of a (T,2D) buffer (row stride differs). Use concat2_bf16.
+ fvk.concat2_bf16(self.b_a_emb.data_ptr(),
+ self.action_time_embeds[step].data_ptr(),
+ self.b_concat.data_ptr(), T, D, D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('concat', self.b_concat.clone()))
self._fp16_gemm(self.b_concat.data_ptr(), self.ae_w2.data_ptr(), self.b_enc_h.data_ptr(), T, D, 2*D, s)
fvk.add_bias_fp16(self.b_enc_h.data_ptr(), self.ae_b2.data_ptr(), T, D, s)
fvk.silu_inplace_fp16(self.b_enc_h.data_ptr(), T*D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('ae2', self.b_enc_h.clone()))
self._fp16_gemm(self.b_enc_h.data_ptr(), self.ae_w3.data_ptr(), self.b_a_emb.data_ptr(), T, D, D, s)
fvk.add_bias_fp16(self.b_a_emb.data_ptr(), self.ae_b3.data_ptr(), T, D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('ae3', self.b_a_emb.clone()))
+ self._dbg.append(('tau', self.action_time_embeds[step].clone()))
fvk.residual_add_fp16(self.b_a_emb.data_ptr(), self.pos_emb[:T].data_ptr(), T*D, s)
fvk.gpu_copy(self.b_hidden.data_ptr(), self.b_state_feat.data_ptr(), D*2, s)
fvk.gpu_copy(self.b_hidden.data_ptr()+D*2, self.b_a_emb.data_ptr(), T*D*2, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('h0', self.b_hidden.clone()))
for l in range(self.L):
is_self = (l % 2 == 1); w = self.dit[l]
fvk.ada_layer_norm_fp16(self.b_hidden.data_ptr(), self.ada_scales[step,l].data_ptr(),
self.ada_shifts[step,l].data_ptr(), self.b_h_norm.data_ptr(), Sa, D, 1e-5, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('ada', self.b_h_norm.clone()))
as_qkv_ptr = self._dit_act_scales_dev.data_ptr() + (l * 3 + 0) * 4
as_up_ptr = self._dit_act_scales_dev.data_ptr() + (l * 3 + 1) * 4
as_dn_ptr = self._dit_act_scales_dev.data_ptr() + (l * 3 + 2) * 4
@@ -909,8 +953,11 @@ def _run_step(self, step, s=0):
else:
self._fp16_gemm(self.b_h_norm.data_ptr(), w['q_w'].data_ptr(), self.b_q_cross.data_ptr(), Sa, D, D, s)
fvk.add_bias_fp16(self.b_q_cross.data_ptr(), w['q_b'].data_ptr(), Sa, D, s)
- # Use precomputed K/V projections (computed once before step loop)
+ # Use precomputed K/V projections (computed once before step loop).
+ # kv length = compact valid-row count for this layer's modality
+ # (text layers see n_kv_text rows, image layers n_kv_img).
cross_idx = l // 2
+ kv_seq = self.n_kv_text if l % 4 == 0 else self.n_kv_img
k_ptr = self._precomp_k[cross_idx].data_ptr()
v_ptr = self._precomp_v[cross_idx].data_ptr()
fvk.gpu_fill_neginf_fp16(self.b_attn_logits_cross.data_ptr(), self.b_attn_logits_cross.nelement(), s)
@@ -918,12 +965,14 @@ def _run_step(self, step, s=0):
# dit_cross site: cross_idx = l // 2 indexes the 16 cross layers
# (even DiT layers 0, 2, ..., 30).
self.attn.run("dit_cross", cross_idx,
- q_seq=Sa, kv_seq=S_kv, stream=s)
+ q_seq=Sa, kv_seq=kv_seq, stream=s)
else:
fvk.attention_mha_fp16(self.ctx, self.b_q_cross.data_ptr(), k_ptr, v_ptr,
- self.b_attn_logits_cross.data_ptr(), self.b_attn_out.data_ptr(), Sa, S_kv, NH, HD, 1.0/math.sqrt(HD), s)
+ self.b_attn_logits_cross.data_ptr(), self.b_attn_out.data_ptr(), Sa, kv_seq, NH, HD, 1.0/math.sqrt(HD), s)
self._fp16_gemm(self.b_attn_out.data_ptr(), w['o_w'].data_ptr(), self.b_o.data_ptr(), Sa, D, D, s)
fvk.add_bias_fp16(self.b_o.data_ptr(), w['o_b'].data_ptr(), Sa, D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(('attn_o', self.b_o.clone()))
fvk.residual_add_fp16(self.b_hidden.data_ptr(), self.b_o.data_ptr(), Sa*D, s)
fvk.layer_norm_no_affine_fp16(self.b_hidden.data_ptr(), self.b_h_norm.data_ptr(), Sa, D, 1e-5, s)
if self.use_fp8:
@@ -942,6 +991,8 @@ def _run_step(self, step, s=0):
self._fp16_gemm(self.b_ff_h.data_ptr(), w['ff_dn_fp16'].data_ptr(), self.b_ff_out.data_ptr(), Sa, D, H, s)
fvk.add_bias_fp16(self.b_ff_out.data_ptr(), w['ff_dn_b'].data_ptr(), Sa, D, s)
fvk.residual_add_fp16(self.b_hidden.data_ptr(), self.b_ff_out.data_ptr(), Sa*D, s)
+ if getattr(self, '_dbg', None) is not None:
+ self._dbg.append(self.b_hidden.clone())
fvk.ada_layer_norm_fp16(self.b_hidden.data_ptr(), self.out_scales[step].data_ptr(),
self.out_shifts[step].data_ptr(), self.b_h_norm.data_ptr(), Sa, D, 1e-6, s)
self._fp16_gemm(self.b_h_norm.data_ptr(), self.proj_out_2_w.data_ptr(), self.b_model_out.data_ptr(), Sa, 1024, D, s)
diff --git a/tools/convert_groot_n16_hf_checkpoint.py b/tools/convert_groot_n16_hf_checkpoint.py
new file mode 100644
index 00000000..0a9244c6
--- /dev/null
+++ b/tools/convert_groot_n16_hf_checkpoint.py
@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""Offline converter: HF GR00T N1.6 checkpoint -> FlashRT weight layout.
+
+Phase-1b of the N1.6 parity plan. Converts a fine-tuned (or base) HF GR00T
+N1.6 checkpoint into the exact tensor layout the FlashRT N1.6 kernels
+consume, emitting BOTH the converted ``.safetensors`` and a JSON manifest
+recording, for every tensor: HF key, HF shape, transform, FlashRT key,
+FlashRT shape. This makes every weight mapping explicit and auditable, so a
+backbone/weight-layout mismatch can be localized to a specific rule instead
+of guessed from the final action.
+
+Layout rules (one explicit rule per weight family):
+
+ A. SigLIP2 (``backbone.model.vision_model.vision_model.*``)
+ - attention q/k/v/o and FFN fc1/fc2: HF ``[out,in]`` -> ``[in,out]``
+ (FlashRT GEMMs take ``[in,out]``); QKV kept separate, order Q,K,V.
+ - layernorm / position_embedding / patch_embedding bias: passthrough.
+ - (the double ``vision_model`` prefix is part of the HF key and kept).
+
+ B. Qwen3 (``backbone.model.language_model.model.layers.*``)
+ - q ``[2048,2048]``, k ``[1024,2048]``, v ``[1024,2048]`` are fused as
+ ``cat([q,k,v], dim=0).T.contiguous()`` -> ``[2048, 4096]``
+ (Q first, then K, then V; NO interleaving).
+ - FFN ``cat([gate_proj, up_proj], dim=0).T.contiguous()`` with order
+ ``gate | up`` (must not be swapped); down_proj transposed.
+ - layernorm / q_norm / k_norm: passthrough.
+
+ C. DiT (``action_head.model.transformer_blocks.{l}.*``)
+ - even block = cross-attention, odd block = self-attention.
+ - self-attn (odd): QKV fused ``cat([q,k,v],0).T`` -> ``[1536, 4608]``
+ (K/V input dim 1536).
+ - cross-attn (even): q transposed; k/v ``[1536,2048]`` transposed,
+ kept separate (NOT fused).
+ - FFN is GELU (not GEGLU): net.0.proj and net.2 transposed.
+ - norm1.linear (AdaLN) transposed; norm1.norm / norm3 / norm_out have
+ NO affine parameters (absent from the checkpoint); output
+ conditioning chunk order is (shift, scale).
+ - proj_out_1 / proj_out_2 / timestep_encoder linears transposed.
+
+ D. Embodiment (``action_head.{action_encoder,state_encoder,action_decoder}``)
+ - CategorySpecificLinear ``W`` is already ``[num_categories,in,out]``;
+ after selecting ``W[eid]`` it is ``[in,out]`` already -> NO extra
+ transpose. Biases passthrough.
+
+Usage:
+ python tools/convert_groot_n16_hf_checkpoint.py \
+ --src /mnt/lerobot_so101_sim_v1_gr00t_n1d6_sim_fruits_cubes_10w \
+ --dst /mnt/.../n1d6_flashrt_layout \
+ [--dtype fp16]
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+def load_src(src: Path) -> dict:
+ sd = {}
+ for f in sorted(src.glob("*.safetensors")):
+ with safe_open(str(f), framework="pt") as sf:
+ for k in sf.keys():
+ sd[k] = sf.get_tensor(k)
+ return sd
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--src", required=True)
+ ap.add_argument("--dst", required=True)
+ ap.add_argument("--dtype", default="keep", choices=["keep", "fp16", "bf16"])
+ args = ap.parse_args()
+
+ src = Path(args.src)
+ dst = Path(args.dst)
+ dst.mkdir(parents=True, exist_ok=True)
+ sd = load_src(src)
+
+ cast = {"keep": None, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype]
+ def maybe_cast(t: torch.Tensor) -> torch.Tensor:
+ return t.to(cast) if (cast is not None and t.is_floating_point()) else t
+
+ out: dict[str, torch.Tensor] = {}
+ manifest: list[dict] = []
+
+ def emit(hf_key: str, fr_key: str, tensor: torch.Tensor, transform: str) -> None:
+ hf_shape = list(sd[hf_key].shape) if hf_key in sd else None
+ tensor = maybe_cast(tensor)
+ out[fr_key] = tensor
+ manifest.append({
+ "hf_key": hf_key, "hf_shape": hf_shape, "transform": transform,
+ "fr_key": fr_key, "fr_shape": list(tensor.shape),
+ })
+
+ T = lambda t: t.T.contiguous() # [out,in] -> [in,out]
+
+ for k, v in sd.items():
+ # ── A. SigLIP2 ──
+ if k.startswith("backbone.model.vision_model.vision_model."):
+ if any(k.endswith(s) for s in (".to_q.weight", ".to_k.weight",
+ ".to_v.weight", ".to_out.0.weight",
+ ".mlp.fc1.weight", ".mlp.fc2.weight")):
+ emit(k, k, T(v), "transpose[out,in]->[in,out]")
+ else:
+ emit(k, k, v, "passthrough")
+
+ # ── B. Qwen3 ──
+ elif k.startswith("backbone.model.language_model.model.layers."):
+ if k.endswith(".self_attn.q_proj.weight"):
+ # fuse when the sibling k/v are present (handled at q key)
+ pre = k[: -len(".self_attn.q_proj.weight")]
+ q = sd[f"{pre}.self_attn.q_proj.weight"]
+ kk = sd[f"{pre}.self_attn.k_proj.weight"]
+ vv = sd[f"{pre}.self_attn.v_proj.weight"]
+ fused = torch.cat([q, kk, vv], dim=0).T.contiguous()
+ emit(k, f"{pre}.self_attn.qkv_fused", fused,
+ "cat([q,k,v],0).T -> [in,out], order Q,K,V")
+ elif k.endswith((".self_attn.k_proj.weight",
+ ".self_attn.v_proj.weight")):
+ continue # already fused into qkv_fused
+ elif k.endswith(".self_attn.q_proj.bias"):
+ pre = k[: -len(".self_attn.q_proj.bias")]
+ q = sd.get(f"{pre}.self_attn.q_proj.bias")
+ if q is not None:
+ kk = sd[f"{pre}.self_attn.k_proj.bias"]; vv = sd[f"{pre}.self_attn.v_proj.bias"]
+ emit(k, f"{pre}.self_attn.qkv_bias", torch.cat([q, kk, vv], 0),
+ "cat([qb,kb,vb],0)")
+ else:
+ emit(k, k, v, "passthrough")
+ elif k.endswith((".self_attn.k_proj.bias", ".self_attn.v_proj.bias")):
+ pre = k[: -len(".self_attn.k_proj.bias")]
+ if f"{pre}.self_attn.q_proj.bias" in sd:
+ continue # fused into qkv_bias
+ emit(k, k, v, "passthrough")
+ elif k.endswith(".mlp.gate_proj.weight"):
+ pre = k[: -len(".mlp.gate_proj.weight")]
+ g = sd[f"{pre}.mlp.gate_proj.weight"]
+ u = sd[f"{pre}.mlp.up_proj.weight"]
+ emit(k, f"{pre}.mlp.gate_up_fused",
+ torch.cat([g, u], dim=0).T.contiguous(),
+ "cat([gate,up],0).T -> [in,out], order gate|up")
+ elif k.endswith(".mlp.up_proj.weight"):
+ continue # fused into gate_up_fused
+ elif k.endswith(".mlp.down_proj.weight"):
+ emit(k, k, T(v), "transpose[out,in]->[in,out]")
+ else:
+ emit(k, k, v, "passthrough")
+
+ # ── C. DiT ──
+ elif k.startswith("action_head.model.transformer_blocks."):
+ parts = k.split(".")
+ l = int(parts[3])
+ is_self = (l % 2 == 1)
+ if k.endswith(".attn1.to_q.weight"):
+ pre = k[: -len(".attn1.to_q.weight")]
+ q = sd[f"{pre}.attn1.to_q.weight"]
+ if is_self:
+ kk = sd[f"{pre}.attn1.to_k.weight"]
+ vv = sd[f"{pre}.attn1.to_v.weight"]
+ emit(k, f"{pre}.attn1.qkv_fused",
+ torch.cat([q, kk, vv], dim=0).T.contiguous(),
+ "self-attn cat([q,k,v],0).T -> [in,out]")
+ else:
+ emit(k, k, T(q), "cross-attn q transpose")
+ elif k.endswith(".attn1.to_k.weight") or k.endswith(".attn1.to_v.weight"):
+ if is_self:
+ continue # fused
+ emit(k, k, T(v), "cross-attn k/v transpose [1536,2048]->[2048,1536]")
+ elif k.endswith((".attn1.to_q.bias",)):
+ pre = k[: -len(".attn1.to_q.bias")]
+ q = sd[f"{pre}.attn1.to_q.bias"]; kk = sd[f"{pre}.attn1.to_k.bias"]; vv = sd[f"{pre}.attn1.to_v.bias"]
+ if is_self:
+ emit(k, f"{pre}.qkv_bias", torch.cat([q, kk, vv], 0), "cat([qb,kb,vb],0)")
+ else:
+ emit(k, k, v, "passthrough")
+ elif k.endswith((".attn1.to_k.bias", ".attn1.to_v.bias")):
+ if is_self:
+ continue
+ emit(k, k, v, "passthrough")
+ elif k.endswith(".ff.net.0.proj.weight") or k.endswith(".ff.net.2.weight") \
+ or k.endswith(".norm1.linear.weight") or k.endswith(".attn1.to_out.0.weight"):
+ emit(k, k, T(v), "transpose[out,in]->[in,out] (GELU FFN, not GEGLU)")
+ else:
+ emit(k, k, v, "passthrough")
+
+ # ── DiT top-level linears ──
+ elif k.startswith("action_head.model.") and k.endswith(".weight") \
+ and any(s in k for s in ("proj_out_1", "proj_out_2", "timestep_embedder")):
+ emit(k, k, T(v), "transpose[out,in]->[in,out]")
+
+ # ── D. Embodiment CategorySpecificLinear (NO transpose) ──
+ elif k.startswith("action_head.") and any(
+ s in k for s in ("action_encoder", "state_encoder", "action_decoder")):
+ emit(k, k, v, "passthrough (W[eid] already [in,out]; no transpose)")
+
+ # ── everything else (embeddings, vlln, mlp1, norms, etc.) ──
+ else:
+ emit(k, k, v, "passthrough")
+
+ save_file(out, str(dst / "model_flashrt.safetensors"))
+ with open(dst / "flashrt_layout_manifest.json", "w") as f:
+ json.dump(manifest, f, indent=1)
+ print(f"wrote {dst/'model_flashrt.safetensors'} ({len(out)} tensors)")
+ print(f"wrote {dst/'flashrt_layout_manifest.json'} ({len(manifest)} rules)")
+ # summary of transforms
+ from collections import Counter
+ c = Counter(m["transform"] for m in manifest)
+ for t, n in c.most_common():
+ print(f" {n:5d} {t}")
+
+
+if __name__ == "__main__":
+ main()
From e89ecde20b846e181efb8dde1d3b0efde68463e9 Mon Sep 17 00:00:00 2001
From: DXICM <10598463@qq.com>
Date: Mon, 17 Aug 2026 08:26:47 +0000
Subject: [PATCH 2/6] 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
---
CMakeLists.txt | 4 +-
csrc/bindings.cpp | 11 ++
csrc/fp4_bindings.cpp | 32 ++++
csrc/fused_fp4/dit_norm_fp4_sfa.cu | 101 ++++++++++++
csrc/fused_fp4/dit_norm_fp4_sfa.cuh | 6 +
csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu | 152 ++++++++++++++++++
csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cuh | 26 +++
csrc/kernels/qk_norm_rope_rotate_half_bf16.cu | 91 +++++++++++
.../kernels/qk_norm_rope_rotate_half_bf16.cuh | 15 ++
flash_rt/models/groot/pipeline_thor.py | 25 ---
10 files changed, 437 insertions(+), 26 deletions(-)
create mode 100644 csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu
create mode 100644 csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cuh
create mode 100644 csrc/kernels/qk_norm_rope_rotate_half_bf16.cu
create mode 100644 csrc/kernels/qk_norm_rope_rotate_half_bf16.cuh
diff --git a/CMakeLists.txt b/CMakeLists.txt
index a1356a2c..99fec573 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1588,7 +1588,8 @@ if(ENABLE_SM100_CUTLASS)
csrc/kernels/attention_mha_masked.cu
csrc/kernels/vec_fp16_backbone.cu
csrc/kernels/attention_seqused_fused.cu
- csrc/kernels/rope_vec.cu)
+ csrc/kernels/rope_vec.cu
+ csrc/kernels/qk_norm_rope_rotate_half_bf16.cu)
target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_THOR_VLA_KERNELS=1)
message(STATUS "Thor VLA helper kernels: ENABLED (sm_${GPU_ARCH})")
else()
@@ -2056,6 +2057,7 @@ if(ENABLE_SM100_CUTLASS)
csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cu
csrc/quantize/quantize_fp4_sfa_bf16.cu
csrc/fused_fp4/dit_norm_fp4_sfa.cu
+ csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu
csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cu
csrc/quantize/quantize_e0m3_sfa.cu
csrc/gemm/fp4/cutlass_fp4_gemm_e0m3w_sm100.cu
diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp
index 3a1e0bca..0b401947 100644
--- a/csrc/bindings.cpp
+++ b/csrc/bindings.cpp
@@ -312,6 +312,7 @@ extern "C" void flash_rt_awq_quant_fp8_static_fp16(
#include "quantize/bf16_quant_fp8_ncdhw_to_ndhwc.cuh"
#endif
#include "quantize/qkv_split_norm_rope_bf16.cuh"
+#include "kernels/qk_norm_rope_rotate_half_bf16.cuh"
#include "attention/fmha_dispatch.h"
#ifdef ENABLE_MOTUS_SAGE2_RAW
#include "attention/sage2/sage2_attn_raw.cuh"
@@ -1800,6 +1801,16 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
py::arg("eps") = 1e-5f, py::arg("stream") = 0);
#endif // FLASHRT_ENABLE_CHAMELEON
+ m.def("qk_norm_rope_rotate_half_bf16",
+ [](uintptr_t x, uintptr_t w, uintptr_t cos_t, uintptr_t sin_t,
+ int S, int NH, int HD, float eps, uintptr_t stream) -> int {
+ return flash_rt::kernels::qk_norm_rope_rotate_half_bf16(
+ to_ptr(x), to_ptr(w), to_ptr(cos_t), to_ptr(sin_t),
+ S, NH, HD, eps, to_stream(stream));
+ }, py::arg("x"), py::arg("w"), py::arg("cos_table"), py::arg("sin_table"),
+ py::arg("S"), py::arg("NH"), py::arg("HD"), py::arg("eps") = 1e-6f,
+ py::arg("stream") = 0);
+
m.def("gate_mul_residual_fp16",
[](uintptr_t residual, uintptr_t x, uintptr_t gate,
int n, uintptr_t stream) {
diff --git a/csrc/fp4_bindings.cpp b/csrc/fp4_bindings.cpp
index 78dcca56..6a0194f7 100644
--- a/csrc/fp4_bindings.cpp
+++ b/csrc/fp4_bindings.cpp
@@ -41,6 +41,7 @@
#include "gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cuh"
#include "quantize/quantize_fp4_sfa_bf16.cuh"
#include "fused_fp4/dit_norm_fp4_sfa.cuh"
+#include "fused_fp4/silu_mul_fp4_sfa_bf16.cuh"
#include "fused_fp4/layer_norm_fp4_sfa.cuh"
#include "gemm/fp4/cutlass_fp4_gemm_siglip_ffn_sm100.cuh"
@@ -1228,6 +1229,37 @@ same contract as cutlass_fp4_gemm_geglu_il_hw.
py::arg("stream") = 0,
"Fused no-affine LayerNorm (bf16) -> NVFP4 packed + SFA.");
+ m.def("rms_norm_weight_fp4_sfa_bf16",
+ [](uintptr_t x, uintptr_t weight, uintptr_t packed, uintptr_t sfa,
+ int seq_len, int dim, float eps, uintptr_t stream) -> int {
+ return flash_rt::fused_fp4::rms_norm_weight_fp4_sfa_bf16(
+ reinterpret_cast(x),
+ reinterpret_cast(weight),
+ reinterpret_cast(packed),
+ reinterpret_cast(sfa),
+ seq_len, dim, eps,
+ reinterpret_cast(stream));
+ },
+ py::arg("x"), py::arg("weight"), py::arg("packed"), py::arg("sfa"),
+ py::arg("seq_len"), py::arg("dim"), py::arg("eps") = 1e-5f,
+ py::arg("stream") = 0,
+ "Fused weighted RMSNorm (bf16) -> NVFP4 packed + SFA.");
+
+ m.def("silu_mul_fp4_sfa_bf16",
+ [](uintptr_t gate, uintptr_t up, uintptr_t packed, uintptr_t sfa,
+ int N, int D, bool is_sfb, uintptr_t stream) -> int {
+ return flash_rt::fused_fp4::silu_mul_fp4_sfa_bf16(
+ reinterpret_cast(gate),
+ reinterpret_cast(up),
+ reinterpret_cast(packed),
+ reinterpret_cast(sfa),
+ N, D, is_sfb,
+ reinterpret_cast(stream));
+ },
+ py::arg("gate"), py::arg("up"), py::arg("packed"), py::arg("sfa"),
+ py::arg("N"), py::arg("D"), py::arg("is_sfb"), py::arg("stream") = 0,
+ "Fused SiLU(gate)*up (bf16) -> NVFP4 packed + SFA.");
+
m.attr("__version__") = "0.1.0-dev";
m.attr("layout_note") = "scales are linear [N, D/16]; Phase 4 adds tile-interleave conversion";
}
diff --git a/csrc/fused_fp4/dit_norm_fp4_sfa.cu b/csrc/fused_fp4/dit_norm_fp4_sfa.cu
index bde3c2ff..15e0420a 100644
--- a/csrc/fused_fp4/dit_norm_fp4_sfa.cu
+++ b/csrc/fused_fp4/dit_norm_fp4_sfa.cu
@@ -206,3 +206,104 @@ int layer_norm_no_affine_fp4_sfa_bf16(
} // namespace fused_fp4
} // namespace flash_rt
+
+// ----------------------------------------------------------------------------
+// Weighted RMSNorm (bf16) -> NVFP4 quantize + SFA (GR00T N1.6 Qwen3 tier).
+// y = x * rsqrt(mean(x^2) + eps) * weight, rounded through bf16 before
+// quantization (matches the torch RMSNorm(bf16) + quantize chain).
+// ----------------------------------------------------------------------------
+namespace flash_rt {
+namespace fused_fp4 {
+
+#if FV_HAVE_CUTLASS
+
+namespace {
+
+template
+__global__ void rms_norm_w_fp4_sfa_kernel(
+ const __nv_bfloat16* __restrict__ x,
+ const __nv_bfloat16* __restrict__ weight,
+ uint2* __restrict__ packed,
+ uint8_t* __restrict__ dst_sfa,
+ LayoutSF layout,
+ int D, float eps) {
+ const int r = blockIdx.x;
+ const __nv_bfloat162* row2 =
+ reinterpret_cast(x + static_cast(r) * D);
+ const int D2 = D >> 1;
+ __shared__ float sh[32];
+
+ float ssq = 0.f;
+ for (int i = threadIdx.x; i < D2; i += blockDim.x) {
+ const __nv_bfloat162 v = row2[i];
+ const float a = __bfloat162float(v.x);
+ const float b = __bfloat162float(v.y);
+ ssq += a * a + b * b;
+ }
+ const float rstd = rsqrtf(block_sum_dn(ssq, sh) / D + eps);
+
+ const int n_blocks = D >> 4;
+ const int4* x4 = reinterpret_cast(x + static_cast(r) * D);
+ const int4* w4 = reinterpret_cast(weight);
+ for (int blk = threadIdx.x; blk < n_blocks; blk += blockDim.x) {
+ const int4 xr[2] = {x4[2 * blk], x4[2 * blk + 1]};
+ const int4 wr[2] = {w4[2 * blk], w4[2 * blk + 1]};
+ const __nv_bfloat16* xh = reinterpret_cast(xr);
+ const __nv_bfloat16* wh = reinterpret_cast(wr);
+ float vals[16];
+ float amax = 0.f;
+ #pragma unroll
+ for (int i = 0; i < 16; ++i) {
+ const float normed =
+ __bfloat162float(xh[i]) * rstd * __bfloat162float(wh[i]);
+ vals[i] = __bfloat162float(__float2bfloat16(normed));
+ const float a = fabsf(vals[i]);
+ if (a > amax) amax = a;
+ }
+ float desired = amax / 6.f;
+ if (desired < 1e-12f) desired = 1e-12f;
+ __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
+ const float bs_dq = static_cast(bs_q);
+ dst_sfa[layout(r, blk * 16, 0)] = *reinterpret_cast(&bs_q);
+ const float inv_bs = 1.f / bs_dq;
+ uint2 out;
+ uint8_t* ob = reinterpret_cast(&out);
+ #pragma unroll
+ for (int p = 0; p < 8; ++p) {
+ const uint8_t lo = fp32_to_e2m1_dn(vals[2 * p] * inv_bs);
+ const uint8_t hi = fp32_to_e2m1_dn(vals[2 * p + 1] * inv_bs);
+ ob[p] = static_cast(lo | (hi << 4));
+ }
+ packed[static_cast(r) * n_blocks + blk] = out;
+ }
+}
+
+} // namespace
+
+#endif // FV_HAVE_CUTLASS
+
+int rms_norm_weight_fp4_sfa_bf16(
+ const void* x, const void* weight, void* packed, void* sfa,
+ int seq_len, int dim, float eps, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+ if (check_dn_args(x, packed, dim) != 0) return -1;
+ if (reinterpret_cast(weight) & 15) return -1;
+ auto shape = cute::make_shape(seq_len, 1, dim, 1);
+ auto layout = CfgDN::tile_atom_to_shape_SFA(shape);
+ rms_norm_w_fp4_sfa_kernel<<>>(
+ reinterpret_cast(x),
+ reinterpret_cast(weight),
+ reinterpret_cast(packed),
+ reinterpret_cast(sfa),
+ layout, dim, eps);
+ const cudaError_t e = cudaGetLastError();
+ return (e == cudaSuccess) ? 0 : -static_cast(e);
+#else
+ (void)x; (void)weight; (void)packed; (void)sfa;
+ (void)seq_len; (void)dim; (void)eps; (void)stream;
+ return -2;
+#endif
+}
+
+} // namespace fused_fp4
+} // namespace flash_rt
diff --git a/csrc/fused_fp4/dit_norm_fp4_sfa.cuh b/csrc/fused_fp4/dit_norm_fp4_sfa.cuh
index de19c3fd..11a369ea 100644
--- a/csrc/fused_fp4/dit_norm_fp4_sfa.cuh
+++ b/csrc/fused_fp4/dit_norm_fp4_sfa.cuh
@@ -29,5 +29,11 @@ int layer_norm_no_affine_fp4_sfa_bf16(
const void* x, void* packed, void* sfa,
int seq_len, int dim, float eps, cudaStream_t stream);
+// packed/sfa = quantize(RMSNorm(x[row]) * weight) — no mean removal.
+// x: bf16 [S, D]; weight: bf16 [D]. Qwen3 pre-attn / pre-FF norms.
+int rms_norm_weight_fp4_sfa_bf16(
+ const void* x, const void* weight, void* packed, void* sfa,
+ int seq_len, int dim, float eps, cudaStream_t stream);
+
} // namespace fused_fp4
} // namespace flash_rt
diff --git a/csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu b/csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu
new file mode 100644
index 00000000..1cc46515
--- /dev/null
+++ b/csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu
@@ -0,0 +1,152 @@
+// ============================================================================
+// Fused SiLU(gate) * up (bf16) + NVFP4 quantize + SFA write.
+// One thread per 16-element block: two int4 loads per operand, silu in
+// fp32 rounded to bf16, bf16 multiply, then the standard per-block scale
+// selection + e2m1 rounding + tile-interleaved SFA byte.
+// ============================================================================
+#include "silu_mul_fp4_sfa_bf16.cuh"
+
+#include
+#include
+
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
+# include "cutlass/cutlass.h"
+# include "cutlass/detail/sm100_blockscaled_layout.hpp"
+# include "cute/tensor.hpp"
+# define FV_HAVE_CUTLASS 1
+#else
+# define FV_HAVE_CUTLASS 0
+#endif
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+#if FV_HAVE_CUTLASS
+
+namespace {
+
+using CfgSM = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+__device__ __forceinline__ uint8_t fp32_to_e2m1_sm(float x) {
+ uint8_t sign = (x < 0.f) ? 0x8u : 0x0u;
+ float ax = fabsf(x);
+ uint8_t mant;
+ if (ax <= 0.25f) mant = 0u;
+ else if (ax <= 0.75f) mant = 1u;
+ else if (ax <= 1.25f) mant = 2u;
+ else if (ax <= 1.75f) mant = 3u;
+ else if (ax <= 2.5f) mant = 4u;
+ else if (ax <= 3.5f) mant = 5u;
+ else if (ax <= 5.0f) mant = 6u;
+ else mant = 7u;
+ return sign | mant;
+}
+
+template
+__global__ void kernel_silu_mul_fp4_sfa_bf16(
+ const int4* __restrict__ gate,
+ const int4* __restrict__ up,
+ uint2* __restrict__ dst_packed,
+ uint8_t* __restrict__ dst_sfa,
+ LayoutSF layout,
+ int N, int D8) {
+ const int block_idx = blockIdx.x * blockDim.x + threadIdx.x;
+ const int row = blockIdx.y;
+ const int n_blocks = D8 >> 1;
+ if (row >= N || block_idx >= n_blocks) return;
+
+ const int4 g0 = gate[row * D8 + 2 * block_idx];
+ const int4 g1 = gate[row * D8 + 2 * block_idx + 1];
+ const int4 u0 = up[row * D8 + 2 * block_idx];
+ const int4 u1 = up[row * D8 + 2 * block_idx + 1];
+ const __nv_bfloat16* gh0 = reinterpret_cast(&g0);
+ const __nv_bfloat16* gh1 = reinterpret_cast(&g1);
+ const __nv_bfloat16* uh0 = reinterpret_cast(&u0);
+ const __nv_bfloat16* uh1 = reinterpret_cast(&u1);
+
+ float vals[16];
+ float amax = 0.f;
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) {
+ const float g[2] = {__bfloat162float(gh0[i]), __bfloat162float(gh1[i])};
+ const __nv_bfloat16 u[2] = {uh0[i], uh1[i]};
+ #pragma unroll
+ for (int h = 0; h < 2; ++h) {
+ const float s = g[h] / (1.f + expf(-g[h])); // silu fp32
+ const __nv_bfloat16 sb = __float2bfloat16(s); // round like torch
+ const __nv_bfloat16 prod = __hmul(sb, u[h]); // bf16 multiply
+ vals[h * 8 + i] = __bfloat162float(prod);
+ const float a = fabsf(vals[h * 8 + i]);
+ if (a > amax) amax = a;
+ }
+ }
+
+ float desired = amax / 6.f;
+ if (desired < 1e-12f) desired = 1e-12f;
+ __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
+ const float bs_dq = static_cast(bs_q);
+
+ dst_sfa[layout(row, block_idx * 16, 0)] =
+ *reinterpret_cast(&bs_q);
+
+ const float inv_bs = 1.f / bs_dq;
+ uint2 out;
+ uint8_t* ob = reinterpret_cast(&out);
+ #pragma unroll
+ for (int p = 0; p < 8; ++p) {
+ const uint8_t lo = fp32_to_e2m1_sm(vals[2 * p] * inv_bs);
+ const uint8_t hi = fp32_to_e2m1_sm(vals[2 * p + 1] * inv_bs);
+ ob[p] = static_cast(lo | (hi << 4));
+ }
+ dst_packed[row * n_blocks + block_idx] = out;
+}
+
+} // namespace
+
+#endif // FV_HAVE_CUTLASS
+
+int silu_mul_fp4_sfa_bf16(
+ const void* gate, const void* up, void* packed, void* sfa,
+ int N, int D, bool is_sfb, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+ if (D % 16 != 0) return -1;
+ if ((reinterpret_cast(gate) & 15) ||
+ (reinterpret_cast(up) & 15) ||
+ (reinterpret_cast(packed) & 7)) return -1;
+ const int n_blocks = D / 16;
+ const int threads = 128;
+ dim3 grid((n_blocks + threads - 1) / threads, N);
+
+ auto shape = cute::make_shape(
+ is_sfb ? 1 : N,
+ is_sfb ? N : 1,
+ D, 1);
+
+ if (is_sfb) {
+ auto layout = CfgSM::tile_atom_to_shape_SFB(shape);
+ kernel_silu_mul_fp4_sfa_bf16<<>>(
+ reinterpret_cast(gate),
+ reinterpret_cast(up),
+ reinterpret_cast(packed),
+ reinterpret_cast(sfa),
+ layout, N, D >> 3);
+ } else {
+ auto layout = CfgSM::tile_atom_to_shape_SFA(shape);
+ kernel_silu_mul_fp4_sfa_bf16<<>>(
+ reinterpret_cast(gate),
+ reinterpret_cast(up),
+ reinterpret_cast(packed),
+ reinterpret_cast(sfa),
+ layout, N, D >> 3);
+ }
+ const cudaError_t e = cudaGetLastError();
+ return (e == cudaSuccess) ? 0 : -static_cast(e);
+#else
+ (void)gate; (void)up; (void)packed; (void)sfa;
+ (void)N; (void)D; (void)is_sfb; (void)stream;
+ return -2;
+#endif
+}
+
+} // namespace fused_fp4
+} // namespace flash_rt
diff --git a/csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cuh b/csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cuh
new file mode 100644
index 00000000..af130d56
--- /dev/null
+++ b/csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cuh
@@ -0,0 +1,26 @@
+// ============================================================================
+// FlashRT — fused SiLU(gate) * up (bf16 in) + NVFP4 quantize + SFA write.
+//
+// GR00T N1.6 Qwen3 FFN hand-off for the NVFP4 tier: replaces the torch
+// silu + mul + separate quantize chain (three elementwise passes over
+// [Se, Dff]) with one kernel emitting packed e2m1 + tile-interleaved UE4M3
+// scales. Value path mirrors torch: silu in fp32, rounded to bf16, then
+// bf16 multiply.
+//
+// Additive: new symbols only.
+// ============================================================================
+#pragma once
+
+#include
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+// packed/sfa = quantize(bf16(silu(gate)) * up)
+// gate/up: bf16 [N, D] row-major. is_sfb selects SFB (weights) layout.
+int silu_mul_fp4_sfa_bf16(
+ const void* gate, const void* up, void* packed, void* sfa,
+ int N, int D, bool is_sfb, cudaStream_t stream);
+
+} // namespace fused_fp4
+} // namespace flash_rt
diff --git a/csrc/kernels/qk_norm_rope_rotate_half_bf16.cu b/csrc/kernels/qk_norm_rope_rotate_half_bf16.cu
new file mode 100644
index 00000000..bb5d1fdc
--- /dev/null
+++ b/csrc/kernels/qk_norm_rope_rotate_half_bf16.cu
@@ -0,0 +1,91 @@
+// ============================================================================
+// FlashRT — fused per-head RMSNorm + rotate-half RoPE (bf16, in-place).
+//
+// GR00T N1.6 Qwen3 tier: replaces the torch chain per layer
+// q = q_norm(q_proj_out.view(1,S,NH,HD)) (RMSNorm over HD)
+// q, k = apply_rotary_pos_emb(q, k, cos, sin) (cat/mul/rotate_half)
+// with one warp-per-(row, head) kernel. Launched once for Q (NHQ heads)
+// and once for K (NHKV heads), so GQA is handled by the caller.
+//
+// x : [S, NH*HD] bf16 row-major, modified in place
+// w : [HD] bf16 RMSNorm weight
+// cos/sin: [S, HD] bf16 (HF rotary_emb output; duplicated halves)
+// Rounding follows the HF chain closely: norm output rounded to bf16
+// before the RoPE multiplies; RoPE accumulated in fp32, rounded once.
+// ============================================================================
+#include "qk_norm_rope_rotate_half_bf16.cuh"
+
+#include
+
+namespace flash_rt {
+namespace kernels {
+
+namespace {
+
+__global__ void qk_norm_rope_rotate_half_bf16_kernel(
+ __nv_bfloat16* __restrict__ x,
+ const __nv_bfloat16* __restrict__ w,
+ const __nv_bfloat16* __restrict__ cos_t,
+ const __nv_bfloat16* __restrict__ sin_t,
+ int S, int NH, int HD, float eps) {
+ const int half = HD >> 1; // pairs per head (HD==128 -> 64)
+ const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
+ const int lane = threadIdx.x & 31;
+ if (warp_id >= S * NH) return;
+ const int n = warp_id % NH;
+ const int s = warp_id / NH;
+ const int base = s * NH * HD + n * HD;
+
+ // Lane l handles pairs p = l and p = l + 32 (half == 64).
+ float vals[4];
+ int p[2] = {lane, lane + 32};
+ float ssq = 0.f;
+ #pragma unroll
+ for (int j = 0; j < 2; ++j) {
+ const int d = p[j];
+ const float a = __bfloat162float(x[base + d]);
+ const float b = __bfloat162float(x[base + d + half]);
+ vals[2 * j] = a; vals[2 * j + 1] = b;
+ ssq += a * a + b * b;
+ }
+ #pragma unroll
+ for (int o = 16; o > 0; o >>= 1)
+ ssq += __shfl_xor_sync(0xffffffffu, ssq, o);
+ const float inv = rsqrtf(ssq / HD + eps);
+
+ #pragma unroll
+ for (int j = 0; j < 2; ++j) {
+ const int d = p[j];
+ // norm, rounded to bf16 like the HF RMSNorm output
+ const float n_lo = __bfloat162float(__float2bfloat16(
+ vals[2 * j] * inv * __bfloat162float(w[d])));
+ const float n_hi = __bfloat162float(__float2bfloat16(
+ vals[2 * j + 1] * inv * __bfloat162float(w[d + half])));
+ const float c = __bfloat162float(cos_t[s * HD + d]);
+ const float si = __bfloat162float(sin_t[s * HD + d]);
+ x[base + d] = __float2bfloat16(n_lo * c - n_hi * si);
+ x[base + d + half] = __float2bfloat16(n_hi * c + n_lo * si);
+ }
+}
+
+} // namespace
+
+int qk_norm_rope_rotate_half_bf16(
+ void* x, const void* w, const void* cos_t, const void* sin_t,
+ int S, int NH, int HD, float eps, cudaStream_t stream) {
+ if (HD != 128) return -1;
+ const int total_warps = S * NH;
+ const int threads = 128; // 4 warps per block
+ const int blocks = (total_warps * 32 + threads - 1) / threads;
+ qk_norm_rope_rotate_half_bf16_kernel<<>>(
+ reinterpret_cast<__nv_bfloat16*>(x),
+ reinterpret_cast(w),
+ reinterpret_cast(cos_t),
+ reinterpret_cast(sin_t),
+ S, NH, HD, eps);
+ const cudaError_t e = cudaGetLastError();
+ return (e == cudaSuccess) ? 0 : -static_cast(e);
+}
+
+} // namespace kernels
+} // namespace flash_rt
diff --git a/csrc/kernels/qk_norm_rope_rotate_half_bf16.cuh b/csrc/kernels/qk_norm_rope_rotate_half_bf16.cuh
new file mode 100644
index 00000000..214d5a1a
--- /dev/null
+++ b/csrc/kernels/qk_norm_rope_rotate_half_bf16.cuh
@@ -0,0 +1,15 @@
+#pragma once
+#include
+
+namespace flash_rt {
+namespace kernels {
+
+// Fused per-head RMSNorm + rotate-half RoPE (bf16, in-place). See .cu.
+// Launched once for Q (NHQ heads) and once for K (NHKV heads) for GQA.
+// Returns 0 on success, -1 unsupported HD, -cudaError otherwise.
+int qk_norm_rope_rotate_half_bf16(
+ void* x, const void* w, const void* cos_t, const void* sin_t,
+ int S, int NH, int HD, float eps, cudaStream_t stream);
+
+} // namespace kernels
+} // namespace flash_rt
diff --git a/flash_rt/models/groot/pipeline_thor.py b/flash_rt/models/groot/pipeline_thor.py
index d6d45a81..27b7c24d 100644
--- a/flash_rt/models/groot/pipeline_thor.py
+++ b/flash_rt/models/groot/pipeline_thor.py
@@ -597,9 +597,6 @@ def forward(self, x_in, s=0):
fvk.gpu_strided_copy_fp16(self.b_qkv.data_ptr(), self.b_attn.data_ptr(), Se, NHKV*HD, self.QKV, NHQ*HD+NHKV*HD, s)
fvk.gpu_repeat_interleave_heads(self.b_k.data_ptr(), self.b_k_exp.data_ptr(), Se, NHKV, HD, NHQ//NHKV, s)
fvk.gpu_repeat_interleave_heads(self.b_attn.data_ptr(), self.b_v_exp.data_ptr(), Se, NHKV, HD, NHQ//NHKV, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('k', self.b_k.clone()))
- self._dbg.append(('v', self.b_attn.clone()))
fvk.gpu_fill_neginf_fp16(self.b_logits.data_ptr(), self.b_logits.nelement(), s)
if self.attn is not None:
self.attn.run("qwen3", i, q_seq=Se, stream=s)
@@ -607,8 +604,6 @@ def forward(self, x_in, s=0):
fvk.attention_mha_fp16(self.ctx, self.b_q.data_ptr(), self.b_k_exp.data_ptr(), self.b_v_exp.data_ptr(),
self.b_logits.data_ptr(), self.b_o.data_ptr(), Se, Se, NHQ, HD, 1.0/math.sqrt(HD), s)
self.gemm.fp16_nn(self.b_o.data_ptr(), w['o_w'].data_ptr(), self.b_xn.data_ptr(), Se, D, D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('o', self.b_xn.clone()))
fvk.residual_add_fp16(self.b_x.data_ptr(), self.b_xn.data_ptr(), Se * D, s)
fvk.rms_norm_fp16(self.b_x.data_ptr(), w['ln2_w'].data_ptr(), self.b_xn.data_ptr(), Se, D, 1e-6, s)
if self.use_fp8:
@@ -633,9 +628,6 @@ def forward(self, x_in, s=0):
self.gemm.fp16_nn(self.b_gu.data_ptr(), w['down_fp16'].data_ptr(),
self.b_down.data_ptr(), Se, D, H, s)
fvk.residual_add_fp16(self.b_x.data_ptr(), self.b_down.data_ptr(), Se * D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('ff', self.b_down.clone()))
- self._dbg.append(('x', self.b_x.clone()))
fvk.rms_norm_fp16(self.b_x.data_ptr(), self.final_norm_w.data_ptr(), self.b_xn.data_ptr(), Se, D, 1e-6, s)
return self.b_xn
@@ -899,36 +891,23 @@ def _run_step(self, step, s=0):
fvk.gpu_cast_fp32_to_fp16(self.b_actions.data_ptr(), self.b_actions_fp16.data_ptr(), T*self.action_dim, s)
self._fp16_gemm(self.b_actions_fp16.data_ptr(), self.ae_w1.data_ptr(), self.b_a_emb.data_ptr(), T, D, self.action_dim, s)
fvk.add_bias_fp16(self.b_a_emb.data_ptr(), self.ae_b1.data_ptr(), T, D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('ae1', self.b_a_emb.clone()))
# gpu_copy is a raw memcpy: it cannot place (T,D) rows into the first
# D columns of a (T,2D) buffer (row stride differs). Use concat2_bf16.
fvk.concat2_bf16(self.b_a_emb.data_ptr(),
self.action_time_embeds[step].data_ptr(),
self.b_concat.data_ptr(), T, D, D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('concat', self.b_concat.clone()))
self._fp16_gemm(self.b_concat.data_ptr(), self.ae_w2.data_ptr(), self.b_enc_h.data_ptr(), T, D, 2*D, s)
fvk.add_bias_fp16(self.b_enc_h.data_ptr(), self.ae_b2.data_ptr(), T, D, s)
fvk.silu_inplace_fp16(self.b_enc_h.data_ptr(), T*D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('ae2', self.b_enc_h.clone()))
self._fp16_gemm(self.b_enc_h.data_ptr(), self.ae_w3.data_ptr(), self.b_a_emb.data_ptr(), T, D, D, s)
fvk.add_bias_fp16(self.b_a_emb.data_ptr(), self.ae_b3.data_ptr(), T, D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('ae3', self.b_a_emb.clone()))
- self._dbg.append(('tau', self.action_time_embeds[step].clone()))
fvk.residual_add_fp16(self.b_a_emb.data_ptr(), self.pos_emb[:T].data_ptr(), T*D, s)
fvk.gpu_copy(self.b_hidden.data_ptr(), self.b_state_feat.data_ptr(), D*2, s)
fvk.gpu_copy(self.b_hidden.data_ptr()+D*2, self.b_a_emb.data_ptr(), T*D*2, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('h0', self.b_hidden.clone()))
for l in range(self.L):
is_self = (l % 2 == 1); w = self.dit[l]
fvk.ada_layer_norm_fp16(self.b_hidden.data_ptr(), self.ada_scales[step,l].data_ptr(),
self.ada_shifts[step,l].data_ptr(), self.b_h_norm.data_ptr(), Sa, D, 1e-5, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('ada', self.b_h_norm.clone()))
as_qkv_ptr = self._dit_act_scales_dev.data_ptr() + (l * 3 + 0) * 4
as_up_ptr = self._dit_act_scales_dev.data_ptr() + (l * 3 + 1) * 4
as_dn_ptr = self._dit_act_scales_dev.data_ptr() + (l * 3 + 2) * 4
@@ -971,8 +950,6 @@ def _run_step(self, step, s=0):
self.b_attn_logits_cross.data_ptr(), self.b_attn_out.data_ptr(), Sa, kv_seq, NH, HD, 1.0/math.sqrt(HD), s)
self._fp16_gemm(self.b_attn_out.data_ptr(), w['o_w'].data_ptr(), self.b_o.data_ptr(), Sa, D, D, s)
fvk.add_bias_fp16(self.b_o.data_ptr(), w['o_b'].data_ptr(), Sa, D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(('attn_o', self.b_o.clone()))
fvk.residual_add_fp16(self.b_hidden.data_ptr(), self.b_o.data_ptr(), Sa*D, s)
fvk.layer_norm_no_affine_fp16(self.b_hidden.data_ptr(), self.b_h_norm.data_ptr(), Sa, D, 1e-5, s)
if self.use_fp8:
@@ -991,8 +968,6 @@ def _run_step(self, step, s=0):
self._fp16_gemm(self.b_ff_h.data_ptr(), w['ff_dn_fp16'].data_ptr(), self.b_ff_out.data_ptr(), Sa, D, H, s)
fvk.add_bias_fp16(self.b_ff_out.data_ptr(), w['ff_dn_b'].data_ptr(), Sa, D, s)
fvk.residual_add_fp16(self.b_hidden.data_ptr(), self.b_ff_out.data_ptr(), Sa*D, s)
- if getattr(self, '_dbg', None) is not None:
- self._dbg.append(self.b_hidden.clone())
fvk.ada_layer_norm_fp16(self.b_hidden.data_ptr(), self.out_scales[step].data_ptr(),
self.out_shifts[step].data_ptr(), self.b_h_norm.data_ptr(), Sa, D, 1e-6, s)
self._fp16_gemm(self.b_h_norm.data_ptr(), self.proj_out_2_w.data_ptr(), self.b_model_out.data_ptr(), Sa, 1024, D, s)
From a5e33078202fd4874aa7c7fb4877cabcc5ca1c14 Mon Sep 17 00:00:00 2001
From: DXICM <10598463@qq.com>
Date: Mon, 17 Aug 2026 08:27:01 +0000
Subject: [PATCH 3/6] 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).
---
docs/groot_n16_thor_sm110.md | 336 ++++++++++++++++++
docs/groot_transformers5_weight_corruption.md | 107 ++++++
docs/thor_gpu_idle_reset_workaround.md | 161 +++++++++
3 files changed, 604 insertions(+)
create mode 100644 docs/groot_n16_thor_sm110.md
create mode 100644 docs/groot_transformers5_weight_corruption.md
create mode 100644 docs/thor_gpu_idle_reset_workaround.md
diff --git a/docs/groot_n16_thor_sm110.md b/docs/groot_n16_thor_sm110.md
new file mode 100644
index 00000000..0200b1dc
--- /dev/null
+++ b/docs/groot_n16_thor_sm110.md
@@ -0,0 +1,336 @@
+# GR00T N1.6 × FlashRT（Jetson Thor SM110）— 权威适配与优化文档
+
+> 自权重对齐开始的完整记录：HF 数值对齐 → 12 个真实 bug 修复 → parity 通路 →
+> 逐轮性能优化（130 ms → 28.5 ms）。配套两份独立主题文档见 §10。
+>
+> 本文是 N1.6 的唯一权威文档（one authoritative doc per model/platform）。
+> 提交 PR 时以本文 + §11 文件清单为准。
+
+---
+
+## 0. 状态总览
+
+**完成并通过仿真验证。** FlashRT 与 HF eager 基线数值和动作行为
+一致，仿真可完成任务。
+
+| 指标 | 数值 | 备注 |
+|---|---|---|
+| GPU 推理（e2e，中位数） | **~28.5 ms** | 4 步 DiT、2 相机、252×252、T=50 |
+| 服务侧 total | ~32 ms | 含预处理 2.8 ms + 解码 0.3 ms |
+| HTTP 往返 / RCP 侧 | ~38 ms | 20 请求 median 39 / p95 40 / 全 finite |
+| 精度 vs HF eager | cos **0.999933** / maxd **0.059** | 去归一化动作空间 |
+| 起点（HF eager） | ~110–150 ms | NVIDIA 官方 eager ~126 ms |
+| Thor 实测带宽 | **252–255 GB/s**（官方 ~273 GB/s，~93%） | 见 §8 |
+
+推理参数与 HF 完全一致（4 步 flow-matching、252×252、T=50、bf16 数学），
+**未为提速改动任何推理超参**；加速全部来自 kernel 化 / 量化 / 图融合。
+
+---
+
+## 1. 背景与模型架构
+
+GR00T N1.6-3B = Eagle-Block2A-2B-v2 backbone + AlternateVLDiT action head：
+
+| 组件 | 结构 | 说明 |
+|---|---|---|
+| 视觉塔 | SigLIP2 NaFlex，27 层，D=1152，HD=72，16 头 | 252×252 → 18×18=324 patch/视图，2 视图打包成 648 token 单序列 |
+| LLM | Qwen3，16 层，D=2048，GQA 16Q/8KV，HD=128 | ★ checkpoint 只截断到 16 层（非 28） |
+| mlp1 | pixel-unshuffle + LN + 2×Linear | 视觉 → LLM 维度投影 |
+| Action head | AlternateVLDiT，32 层，D=1536，NH=32，HD=48，FF=6144 | Sa = T+1 = 51；4 步 flow-matching；奇数层 self-attn、偶数层 cross-attn（交替 text/image KV） |
+| 动作 | action_dim=128（padding），so101 实际仅前 ~6 维有效 | T=50（padded max） |
+
+关键架构事实（对性能分析至关重要）：
+
+- **SigLIP 是 cross-view full attention**：HF(sdpa) 对打包的 648 token 做全注意力，
+ **不是** per-view。NaFlex 代码里按 img_idx 的 `seq_len_list` 分段**只有
+ flash_attention_2 后端消费，sdpa 静默忽略**；Eagle `extra_kwargs` 里的
+ `attn_implementation="flash_attention_2"` 是死代码（未传入加载）。训练/部署实际都是 sdpa。
+- **DiT 是权重带宽主导**：32 层 × 4 步，每步重读 ~415 MB fp4 权重；M=Sa=51 极小，
+ GEMM 完全 weight-bandwidth-bound。
+- **HF eval 图像链输出 252×252**（256→0.95 crop→Eagle smart_resize 到 14 的倍数），
+ 不是 224。
+
+---
+
+## 2. 适配路线总览（自权重对齐开始）
+
+```
+阶段一 权重对齐 & HF 基线 → 发现 transformers>=5 静默权重损坏，修复；
+ 建立 HF eager 数值基线
+阶段二 12 个真实 bug 定位修复 → tokenization / 分辨率 / cross-view / patch 序 /
+ FMHA 发散 / Qwen3 发散 / adaLN chunk 序 /
+ 图捕获野指针 / 校准 / prompt 切换 / 空闲重置 …
+阶段三 parity 通路落地 → SigLIP/Qwen3/DiT 全 HF 原生 bf16 + CUDA graph，
+ 与 HF eager cos≈1.0（~111 ms）
+阶段四 服务侧优化 → 预处理 11→2.5 ms（apply_state 直连 / 线程池 /
+ GPU 归一化 / JPEG→cv2）
+阶段五 FA4 + NVFP4 + 全 kernel 化 → 130 ms → 28.5 ms（§6 逐轮）
+```
+
+---
+
+## 3. 阶段一：权重对齐与 HF 基线
+
+### 3.1 transformers>=5 静默权重损坏（第一个核心问题）
+
+详见配套文档 `docs/groot_transformers5_weight_corruption.md`。摘要：
+
+- `from_pretrained` 加载本身正确（1106/1106，0 missing），但收尾阶段
+ `_initialize_missing_keys` 因 `_auto_class` 为空 → `is_remote_code()==False`
+ → 无视 `_is_hf_initialized` → 用 `_init_weights` **重新随机化整个 SigLIP2
+ 视觉塔（282 张量）+ mlp1.1/3**。
+- 症状：策略"不看图"，只跟随 state；黑图 vs 真实图 Δ 被噪声淹没。
+- 修复（一行）：`type(self)._auto_class = "AutoModel"`（`Gr00tN1d6.__init__`，
+ `post_init()` 之前）。
+- 预防：权重完整性校验（`verify_weight_integrity()`） 抽样 ~12 张量 live vs safetensors
+ 比对，不一致拒绝启动（日志 `[weight-check]`）。
+
+### 3.2 HF eager 基线
+
+HF eager 基线（独立部署仓库）：`Gr00tPolicy`（AutoModel/AutoProcessor）纯 eager，
+作为数值/行为的 ground truth。所有 FlashRT 精度都对它度量。
+
+---
+
+## 4. 阶段二：真实 Bug 清单（12 项，已修）
+
+上游 N1.6 前端把 **openpi 系（Pi0/Pi0.5）的视觉/核假设**直接套到 N1.6，而 N1.6 的
+HF 实际行为不同；叠加若干实现 bug。全部修复集中在
+`flash_rt/frontends/torch/groot_thor.py`。
+
+| # | 问题 | 根因 | 修复 | 验证 |
+|---|------|------|------|------|
+| 1 | backbone 特征与 HF 正交 | **tokenization**：前端裸 `encode(prompt)`；HF 用 Eagle chat template（system/user 头 + formalize + 每视图图像块，168→196 tokens） | 前端 `build_input_ids` 逐 token 复现 HF | `torch.equal` |
+| 2 | 视觉输入分布错 | **分辨率**：HF eval 链输出 252×252，前端固定 224 | `image_size=252`；aux 在 252 用 processor 默认链 + 复刻 smart_resize（PIL bicubic） | 像素 maxΔ 0.002 |
+| 3 | SigLIP 深层发散（L26 cos 0.86） | **注意力范围**：HF(sdpa) 做 cross-view full attention，前端做 per-view（见 §1） | SigLIP 改 cross-view（batch=1 单序列） | 27 层逐层 cos≥0.999 |
+| 4 | patch-embed cos 0.966 | **patch 展平序**：openpi 系 `(C,ph,pw)`；HF NaFlex `convert_images_to_patches` 是 `(ph,pw,C)` | permute 改 `(0,2,3,4,5,1)` | patch-embed cos 1.0 |
+| 5 | 648-seq 注意力静默错 | **strided FMHA kernel 在非 2 幂 seq + 真实数据下发散**（256/512 与随机数据单测均正常，极隐蔽） | parity 模式 SigLIP attention 走 torch sdpa | kernel 单测 vs sdpa |
+| 6 | Qwen3 输出与 HF 正交 | **CKernelQwen3 与 HF 在真实序列上发散**（同输入对拍 block 级逐步衰减） | parity 模式改跑 HF 原生 `Qwen3Model`（bf16、sdpa、图捕获） | vlln cos 0.9998 |
+| 7 | 重捕获后 replay 出 NaN | **野指针**：Qwen3 图捕获时 vlln LN 引用的 `vlln_w/vlln_b` 是 `_capture_all_graphs` 局部张量，函数返回即被 allocator 复用 | 改持久属性 `_g_vlln_w/_g_vlln_b`；加 `_g_vlln_buf` 有限性自检 + 重捕获兜底 | 重捕获回归 |
+| 8 | DiT chunk 震荡、动作失真 | **最终 adaLN chunk 顺序反**：HF `proj_out_1` 为 `(shift, scale)`，前端写成 `(scale, shift)` | 交换 | chunk 单调平滑，maxΔ 0.018 |
+| 9 | 首帧后动作饱和/NaN | **单帧 FP8 校准过窄**：捕获由首帧触发，scale 只测该帧 | 服务侧「当前帧 + 7 合成帧」`calibrate(percentile=99.9)` 多帧校准 | 饱和消失 |
+| 10 | prompt 切换报错/卡死 | prompt 烘焙进图，建图后 `set_prompt` 被拒 | 检测 prompt 变化 → `reset_graph_runtime()` + `set_prompt` + 重捕获 | 多 prompt A/B |
+| 11 | 空闲后首帧垃圾 | Thor 空闲重置使已捕获图失效（见配套文档） | replay 后 `_g_vlln_buf` 有限性自检，非有限即重捕获重试 | 压测 |
+| 12 | prompt 切换重捕获 device-side assert | `reset_graph_runtime` 未删 DiT 静态缓冲/索引（`_dit_txt_idx` 等），重捕获后按旧 Se 越界 | 加入 stale 列表，重捕获时重建 | 切换+空闲 live 全 finite |
+
+---
+
+## 5. 证伪/排除项（**非** bug，勿再追）
+
+- **FP8/FP16 精度**：FP32/BF16 参考给出相同分叉形态，精度不是上述任何一项的根因
+ （parity DiT 用 bf16，与 HF eager 完全一致）。
+- **GELU 近似（exact/tanh）**：两种近似结果相同。
+- **position_embed resize**：16×16→16×16 为恒等；252 的 18×18 resize 与 HF 同参一致。
+- **每块 sinusoidal position embed「缺失」**：**假 bug**。本模型
+ `diffusion_model_cfg.positional_embeddings=None`，每块本就不加 pe；按 diffusers
+ 默认补加反而引入大偏差（曾误修，已撤）。
+- **「kernel DiT 独立不收敛」**：**归因错误**。噪声不收敛由 backbone 失真（#3/#4）+
+ DiT 实现 bug（#8）共同造成；无独立 kernel DiT 问题。旧文档
+ `groot_n16_dit_kernel_nonconvergence.md` 结论作废，已删除。
+- **训练/部署注意力不一致（NVIDIA 侧）**：不存在。flash_attention_2 为死代码，
+ 训练与部署均为 sdpa cross-view。
+- **DiT cross-KV fp4**：实测动作精度恶化 5×（cos 0.9992 / maxd 0.105）且无时延收益，
+ cross 特征喂所有 cross 层 —— 已排除（commit b5afeecc）。
+
+---
+
+## 6. 阶段四/五：性能优化全记录（130 → 28.5 ms）
+
+### 6.0 优化前服务侧预处理优化（阶段四）
+
+- `processor()` 全调用（5.7 ms，含用不到的 VLM/tokenizer 路径）替换为
+ `state_action_processor.apply_state` + 零填充 —— 与 processor 的 state 输出
+ **bit 一致**（28 组真实/合成 state 验证）。
+- 双视图 albumentations + PIL bicubic 变换改线程池并行（4.2→2.1 ms）。
+- SigLIP 像素归一化移 GPU（uint8 H2D + fp32 除法，数学不变）；mlp1 权重布局缓存
+ （去掉每帧 `.T.contiguous()`）。
+- JPEG 解码 PIL→cv2（`pb_utils._jpeg_to_rgb`，带 PIL 回退）：2.13→0.94 ms/图，
+ 解码像素 bit 一致（同为 libjpeg）。
+- 结果：预处理 11→2.5 ms，parity 整链 ~111→~73 ms（4 步）。
+
+### 6.1 Roofline / 天花板分析（阶段五的指导）
+
+- Thor 实测带宽 **252–255 GB/s**（官方 ~273 GB/s，~93%）；GPU GPC 1575 MHz /
+ NVD 1692 MHz 均满载，无节流。
+- DiT 32 层权重 fp4 ~415 MB/步 × 4 步 = 1.66 GB → 纯权重地板 ~6.6 ms；加注意力/
+ norm/quant/激活，DiT ~15 ms 为该配置接近地板的水平。
+- 结论：**DiT/Qwen3 fp4 GEMM 已带宽受限**（L2-cold 实测 170–239 GB/s），tile 重调
+ 无空间；减步数伤行为（见 §6.4）；**28.5 ms 已接近 2 相机 252px T=50 的实际下限**。
+
+### 6.2 FA4 + NVFP4 移植（参考 Chameleon/HyVLA/N1.7 #163）
+
+- **FA4（FlashAttention-4 CuTe-DSL，`flash_rt/hardware/thor/fa4_backend.py`）用于
+ SigLIP**：SigLIP 本就是 cross-view full attention（648 token 单序列），FA4
+ causal=False 为 sdpa 精确替换；encoder graph 10.3→9.0 ms，动作 cos 1.000000。
+ 开关 `FLASHRT_N16_FA4`（默认开，缺失自动回退）。
+- **FA4 不用于 DiT/Qwen3**：小 seq（51/208）下 FA4 0.042–0.394 ms ≫ sdpa 0.013–0.015 ms。
+- **NVFP4（W4A4 CUTLASS）用于 DiT**：block GEMM 走
+ `quantize_fp4_dynamic_sfa_fp16` + `cutlass_fp4_sq_fp16`（per-16 block scale、动态激活
+ 量化、无校准）。关键事实：**M=51 冷 L2 流式下 fp4 GEMM 比 bf16 cuBLAS 快 2–5×**
+ （此前 L2-hot 微基准误判为无收益）。开关 `FLASHRT_N16_DIT_FP4`。
+
+### 6.3 融合 epilogue kernel 移植（上游 N1.7 #163，本地重编译）
+
+从上游移植（`csrc/`，加入 `fp4_kernels_obj` + `fp4_bindings`）：
+
+- `gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cu` — bias / bias+residual /
+ bias+tanh-GELU+fp4out 三种 epilogue。
+- `fused_fp4/dit_norm_fp4_sfa.cu` — AdaLN norm 与 no-affine LN 直接输出 fp4+SFA。
+- `quantize/quantize_fp4_sfa_bf16.cu` — bf16 向量化量化。
+
+DiT 链每层仅 8 kernel（无层间逐元素流量）：**DiT 36.6→15.7 ms**，
+精度不变（vs HF cos 0.999995）。
+
+### 6.4 DiT 减步数实验（**已回退，勿用于生产**）
+
+| 步数 | e2e | 精度/行为 |
+|---|---|---|
+| 4（生产） | ~28.5 ms | 基线 |
+| 2（`FLASHRT_N16_DIT_STEPS=2`） | ~46 ms* | ⚠️ 仿真反馈动作犹豫、夹爪收起慢（Euler 大步长对速度场积分欠冲，夹取等速度突变阶段显现；离线录制帧诊断不可见）→ **回退** |
+| 1（`FLASHRT_N16_DIT_STEPS=1`） | ~36 ms* | maxΔ 0.055，更激进 |
+
+\* 减步数同时减少权重读，但**行为退化**，生产保持 4 步。该开关保留作实验用。
+
+### 6.5 全 kernel 化三轮（37.5 → 28.5 ms）
+
+- **第一轮 — Qwen3 融合**（`FLASHRT_N16_QWEN3_FP4`，生产已启用）：16 层全部
+ q/k/v/o/gate/up/down 走 W4A4 融合 GEMM（residual 原地更新）。
+ - 新写两个 bf16 输入 producer kernel（`csrc/fused_fp4/`，与 torch 两步链 **bit 一致**）：
+ `rms_norm_weight_fp4_sfa_bf16`（加权 RMSNorm→fp4 直出）、
+ `silu_mul_fp4_sfa_bf16`（silu·mul→fp4 直出）。
+ - torch LN/quant/silu/mul 五趟 → 两个融合 kernel：10.9→9.2 ms。
+ - **fused per-head RMSNorm + rotate-half RoPE**（`qk_norm_rope_rotate_half_bf16`）：
+ −2.2 ms。
+ - **fused GQA attention**：torch `enable_gqa` 会退回非融合 math 路径（每层 2 个
+ fp16 GEMM + 显式 softmax）；改 `repeat_interleave` 展开 KV 走融合后端：−2.0 ms。
+ - Qwen3 12.7→5.0 ms；vs HF cos 0.999986。
+- **第二轮 — SigLIP encoder fp4**（`FLASHRT_N16_SIGLIP_FP4`，默认开）：LN producer
+ 用 scale=w−1/shift=b 表达 affine；q/k/v 分离 GEMM 写连续缓冲直喂 FA4；o/ffn 残差进
+ epilogue；fc1 N 与 fc2 K 4304→4352 pad（pad 维全零端到端无贡献）。encoder 9.0→6.9 ms。
+ 全开时动作精度不降（vs HF cos 0.999988）。
+- **第三轮 — SigLIP embeddings 入图**：推理 34→**28.5 ms**。NaFlex window split 对
+ 固定 image_size 是静态 gather、antialias 位置编码 resize 只依赖静态形状；setup 时
+ 预计算 gather 索引 + resize 后位置编码，per-frame 只做
+ patchify+patch_embedding+pos_add+gather 并入 SigLIP 图。对 HF embeddings 前向
+ **bit 一致（max diff 0.0）**，省去每帧 antialias interpolate + unfold/im2col。
+ 图捕获失败时自动回退「encoder-only 图 + eager embeddings」。
+
+### 6.6 已排除的优化路径（dead-end，勿重复尝试）
+
+- **torch.compile（DiT/Qwen3，max-autotune）**：无收益（43→41 ms），且曾出现图捕获
+ 数值异常，勿在生产启用。
+- **torch fp8（`torch._scaled_mm`）**：Thor 仅支持 tensorwise scale + bf16 输出；
+ DiT 权重 per-channel 量化精度好（cos 0.99993），但图捕获后量化/反量化开销吃掉权重
+ 减半收益（29.6 > bf16 21.5 ms），不采用。
+- **legacy fp8 kernel 快路径（`--fp8`）**：实测 siglip 32.8 / qwen3 12.8 / dit 41.5
+ = 91 ms，全面慢于 torch bf16，且含 #5/#6 数值问题。**N1.6 生产请用默认 `--no-fp8`**。
+- **SigLIP FFN torch 级 fp4**：量化开销 47→53 ms，弃用；SigLIP down-proj K=4304 需 pad。
+- **NVFP4 扩展到 SigLIP encoder（早期 torch 级）**：162 Linear fp4 后 9.0→20.7 ms
+ （M=648 激活流量 12.6× 于 DiT，无融合 kernel 时量化开销主导）且 postln cos 降至 0.94
+ → torch 级 fp4 只对 M 小、权重流量主导的 DiT 有效。（后被 §6.5 的融合 kernel 路线取代。）
+- **DiT attention head_dim 48→64 padding**：注意力仅 1.28 ms，pad 浪费 ~0.3 ms，
+ 需自写 head_dim-48 kernel，性价比低，不做。
+- **DiT GEMM tile 重调**：已带宽受限（§6.1），无空间。
+
+---
+
+## 7. 最终架构与配置
+
+### 7.1 数据流（parity + 三层 NVFP4 + FA4，生产默认）
+
+```
+obs ──> preprocess（apply_state 直连 + 线程池图像变换） ~2.8 ms
+ ──> SigLIP 图（embeddings 入图 + 27 层 fp4 encoder + FA4） ~6.5 ms
+ ──> mlp1 / pixel-unshuffle（bf16 torch） ~0.4 ms
+ ──> Qwen3 图（16 层 fp4 + fused norm/rope/GQA） ~5.0 ms
+ ──> DiT 图（4 步 × 32 层 fp4 fused epilogue） ~15.2 ms
+ ──> denormalize_actions（复刻 HF decode_action） ~0.3 ms
+```
+
+全部 CUDA Graph 捕获；三层各自独立 tier 开关，可单独回退。
+
+### 7.2 精度 tier 与开关
+
+| 开关 | 默认 | 作用 |
+|---|---|---|
+| `FLASHRT_N16_DIT_FP4` | 生产开 | DiT NVFP4 融合链 |
+| `FLASHRT_N16_QWEN3_FP4` | 生产开 | Qwen3 NVFP4 融合层（含 fused norm/rope/GQA） |
+| `FLASHRT_N16_SIGLIP_FP4` | 开 | SigLIP encoder fp4 层 |
+| `FLASHRT_N16_FA4` | 开 | SigLIP FA4 注意力（缺失自动回退） |
+| `FLASHRT_N16_DIT_STEPS` | 4 | flow-matching 步数（**生产勿改**，减步伤行为） |
+| `--no-fp8` / `parity` | 生产开 | HF 原生 parity 通路（非 legacy kernel 快路径） |
+
+精度汇总（去归一化动作 vs HF eager）：
+
+| 配置 | cos | maxd |
+|---|---|---|
+| 全开（生产） | 0.999933 | 0.059 |
+| Qwen3 fp4 单独 | 0.999986 | 0.023 |
+| DiT fp4 单独 | 0.999995 | 0.012 |
+
+### 7.3 生产稳定性（继承 N1.7 三层防护）
+
+见 `docs/thor_gpu_idle_reset_workaround.md`：GPU 心跳保活
+（`FLASHRT_GPU_KEEPALIVE=0.15`）+ 空闲重捕获守卫（`FLASHRT_GRAPH_IDLE_REINIT_S=2`）
++ 有限性自检。对抗 Thor GPU 空闲 ~200–300 ms 后驱动重置破坏已捕获 graph 的缺陷。
+
+---
+
+## 8. 带宽与天花板分析
+
+| 项 | 值 |
+|---|---|
+| Thor 实测带宽 | 252–255 GB/s（读/拷贝、fp16/fp32、512/1024MB 一致） |
+| 官方标称 | ~273 GB/s（LPDDR5X），实测 ~93% |
+| GPU 时钟 | GPC 1575 MHz / NVD 1692 MHz（满载，无节流） |
+| DiT fp4 权重地板 | 1.66 GB / 253 GB/s ≈ 6.6 ms（4 步） |
+
+> 注：早期手测曾报 63 GB/s，系把 float32 **元素数当成字节数**少算 4× 的统计错误；
+> 修正后与 roofline.py 一致（~254 GB/s）。roofline 脚本：
+> `.qoder/skills/flashrt-model-adaptation/scripts/roofline.py --measure-bw`。
+
+**天花板结论**：DiT 15.2 ms（4 步权重重读，带宽主导）是剩余大头，无法在不改推理
+超参（减步数）的前提下进一步压缩。**28.5 ms 为该任务配置的实际下限附近。**
+
+---
+
+## 9. 验证方法（可复现）
+
+- **数值一致性**：离线 `cos / maxd` vs HF eager（`/tmp/hf_act_ref.npy` 参考），
+ 以及 live A/B（FlashRT vs HF eager 逐关节 Δ）。
+- **bit 一致验证**：fast embeddings、fused norm/rope/silu-mul producer 均对
+ torch 两步链 / HF 前向 bit 一致（max diff 0.0）。
+- **稳定性**：20 请求 median/p95/全 finite；stress_gaps（间隙压测）+ selfcheck。
+- **带宽**：`roofline.py --measure-bw`。
+- 时延分解：`FLASHRT_GROOT_TIMING=1` 输出 `[timing] preprocess/infer/total`。
+
+---
+
+## 10. 相关文件与配套文档
+
+**本 PR 改动的核心文件**（`git diff main..HEAD`）：
+
+- 前端：`flash_rt/frontends/torch/groot_thor.py`（核心，+1526 行）
+- attention backend：`flash_rt/hardware/thor/attn_backend_groot.py`
+- 模型管线：`flash_rt/models/groot/pipeline_thor.py`
+- 新 kernel：`csrc/fused_fp4/{dit_norm_fp4_sfa,silu_mul_fp4_sfa_bf16}.{cu,cuh}`、
+ `csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.{cu,cuh}`、
+ `csrc/quantize/quantize_fp4_sfa_bf16.{cu,cuh}`、`csrc/fp4_bindings.cpp`
+- 权重转换：`tools/convert_groot_n16_hf_checkpoint.py`
+- 文档：本文 + `docs/groot_transformers5_weight_corruption.md`
+
+**配套文档（独立主题）**：
+
+- `docs/thor_gpu_idle_reset_workaround.md` — Thor GPU 空闲重置缺陷与 CUDA Graph 防护。
+- `docs/groot_transformers5_weight_corruption.md` — transformers>=5 权重静默损坏。
+
+---
+
+## 11. 遗留工作（猜想/未测）
+
+- [ ] **N1.7 同清单巡检**：N1.7 视觉同为 Eagle/SigLIP2 NaFlex，猜想存在同款
+ cross-view/patch 序问题；N1.7 仿真"可用"但未与其 HF 基线数值对拍。
+- [ ] **`FLASHRT_N16_DIT_STEPS=2/1` 仿真验证**：离线数值已测，闭环任务成功率未测
+ （N=2 已见夹爪犹豫，默认不用）。
+- [ ] **parity 模式小时级长稳压测**：prompt 切换 + 空闲自愈已测，长稳未测。
+- [ ] **SigLIP embeddings 之外的进一步 kernel 化**：已接近地板，收益有限。
diff --git a/docs/groot_transformers5_weight_corruption.md b/docs/groot_transformers5_weight_corruption.md
new file mode 100644
index 00000000..9945262a
--- /dev/null
+++ b/docs/groot_transformers5_weight_corruption.md
@@ -0,0 +1,107 @@
+# GR00T N1.6/N1.7: transformers>=5 静默权重损坏与修复
+
+## 状态
+
+已修复（2026-08-08）。影响所有通过 `AutoModel.from_pretrained` 加载 Gr00tN1d6 / Gr00tN1d7
+的 transformers **5.x** 环境；transformers 4.51.3（训练同版本）不受影响。
+
+## 症状
+
+- 服务侧（`serving/groot_n16/eager_server.py`，HF eager 服务端）表现为"不使用实时图像"：
+ 换图 / 换 prompt 时动作输出只变化 ~0.05，策略几乎只跟随 state。
+- 加载日志完全正常：`Loading weights: 100% 1106/1106`，missing / unexpected keys 均为 0。
+- checkpoint 本身完好（其他设备、transformers 4.51.3 下推理正常）。
+
+## 根因
+
+**加载本身是正确的，损坏发生在 `from_pretrained` 的收尾阶段。**
+
+1. `from_pretrained` 把全部 1106 个张量正确写入模型（加载结束瞬间逐张量与磁盘
+ safetensors 比对，maxΔ = 0）。
+2. 随后 `_finalize_model_loading` → `_initialize_missing_keys` → `initialize_weights()`
+ 遍历所有模块调用各自的 `_init_weights`。跳过已加载参数的唯一依据是参数上的
+ `_is_hf_initialized` 标记，而 `PreTrainedModel._initialize_weights` 只在
+ `is_remote_code() == True` 时才检查该标记：
+
+ ```python
+ # transformers/modeling_utils.py (5.10)
+ if getattr(module, "_is_hf_initialized", False):
+ return
+ if is_remote_code and all(getattr(p, "_is_hf_initialized", False) for p in module.parameters(recurse=False)) ...:
+ return
+ self._init_weights(module) # <-- 重新随机初始化
+ ```
+
+3. Gr00tN1d6 / Gr00tN1d7 虽然经 `trust_remote_code` 动态加载，但类的 `_auto_class`
+ 属性为空 → `is_remote_code()` 返回 `False` → 参数级标记被无视 →
+ Siglip2 视觉塔的 `_init_weights` 对已加载的 `nn.Linear` / `nn.Embedding`
+ 执行重新随机初始化。
+
+4. 受损范围：整个 Siglip2 视觉塔 282 个张量（`vision_model.vision_model.*`，
+ 含 `position_embedding`、全部 encoder layers、head）+ 投影层 `mlp1.1/3`。
+ Qwen3 LLM 与 action head 因各自 `_init_weights` 的行为未被波及。
+
+5. 后果：视觉特征是随机噪声 → DiT 交叉注意拿不到有效图像信息 → 策略退化为
+ 纯 state 条件模型。
+
+实测特征（可用于快速诊断）：受损张量的 live 值 std ≈ 0.0294（随机初始化分布），
+与磁盘值的余弦相似度 ≈ 0；例如
+`backbone.model.vision_model.vision_model.encoder.layers.0.self_attn.q_proj.weight`
+磁盘 std 0.02049，加载后变为 0.02936。
+
+## 修复
+
+一行核心改动，在 `Gr00tN1d6.__init__` / `Gr00tN1d7.__init__` 的 `post_init()` 之前：
+
+```python
+# transformers>=5: 必须标记为 auto/remote-code 类，否则加载收尾阶段的
+# `_initialize_missing_keys` 会无视 `_is_hf_initialized` 标记，
+# 用子模型的 `_init_weights` 重新随机化已加载的权重。
+type(self)._auto_class = "AutoModel"
+```
+
+涉及文件：
+
+- `/gr00t/model/gr00t_n1d6/gr00t_n1d6.py`
+- `/gr00t/model/gr00t_n1d7/gr00t_n1d7.py`
+
+## 预防措施
+
+权重完整性自检收敛为共享函数
+`serving/groot_n17/aux_builder.py:verify_weight_integrity()`:对 checkpoint
+safetensors 均匀抽样 ~12 个张量与 live 参数数值比对(容差取 bf16 舍入量级),
+不一致则抛错、拒绝启动。两条服务链路都在构建 `Gr00tPolicy` 后立即执行:
+
+- FlashRT 服务:`Gr00tN17AuxBuilder.__init__`(HF 仅作预处理器,权重同样经
+ transformers 5 加载,必须检查);
+- HF 基线服务:`serving/groot_n16/eager_server.py` /
+ `serving/groot_n17/eager_server.py`(后者已复用共享函数)。
+
+日志标记 `[weight-check]`。
+
+## 验证数据（2026-08-08）
+
+| 指标（HF eager 服务端） | 修复前 | 修复后 |
+| --- | --- | --- |
+| 权重与磁盘不一致张量 | 282（整个视觉塔） | 0 / 1106 |
+| 相同请求重复噪声底 | 0.65 ~ 0.71 | 0.011 ~ 0.016 |
+| 黑图 vs 真实图 Δ（mean） | 被噪声淹没 | 0.135 |
+| 随机图 vs 真实图 Δ | — | 0.137 |
+| prompt 切换 Δ | — | 0.036 |
+| state +0.3 Δ | — | 0.307 |
+
+离线 `Gr00tPolicy` 同条件复现一致（seeded 噪声底 0，黑图 Δ ≈ 0.13）。
+
+N1.7（5558）权重 1030/1030 核对一致；其图像敏感度偏弱（反色图 Δ ≈ 0.03 <
+噪声底 0.07，state Δ 0.22 / prompt Δ 0.07 正常），属该 epoch5 checkpoint 自身
+特性，非加载问题。
+
+## 教训
+
+- 跨 transformers 大版本迁移 `trust_remote_code` 模型时，
+ `from_pretrained` 成功（0 missing / 0 unexpected）**不等于**权重正确，
+ 必须做数值级核对（抽样 live vs safetensors）。
+- 复合模型（PreTrainedModel 嵌套 PreTrainedModel）在 5.x 下的收尾初始化
+ 依赖 `is_remote_code()`；自定义顶层模型类务必显式设置 `_auto_class`。
+- 诊断此类问题的快捷手段：比较 live 权重 std 与随机初始化分布；
+ 对 `_load_pretrained_model` / `_initialize_missing_keys` 打桩二分定位。
diff --git a/docs/thor_gpu_idle_reset_workaround.md b/docs/thor_gpu_idle_reset_workaround.md
new file mode 100644
index 00000000..5d2dec02
--- /dev/null
+++ b/docs/thor_gpu_idle_reset_workaround.md
@@ -0,0 +1,161 @@
+# Thor 部署稳定性:GPU 空闲重置缺陷与 CUDA Graph 防护体系
+
+## 状态
+
+已修复并验证(2026-08-08/09)。本文记录 GR00T N1.7 FlashRT 服务
+(`serving/groot_n17/`,端口 5558)排查"推理结果周期性震荡 + 偶发 4s 卡顿"
+全过程的结论与修复方案,作为后续对 **所有 Thor(SM110)serving 代码**
+做系统化修复的依据。
+
+配套文档:
+- `docs/groot_transformers5_weight_corruption.md`(transformers>=5 权重静默损坏,
+ 本轮修复的第一个核心问题,已单独记录)。
+- `docs/groot_n16_thor_sm110.md`(N1.6×FlashRT Thor 权威文档,含 GPU 空闲重置
+ 防护在本服务上的适用说明)。
+
+## 一、问题现象
+
+1. **震荡**:服务稳定推理 30–50 帧后,action 输出突然剧烈震动;错误帧的
+ 耗时反而比正常帧短(50ms vs 正常 70ms)。
+2. **卡顿**:偶发单帧推理卡住 3.5–4.5s(约每 20–30s 一次),输出本身正常。
+3. 两者相关:卡顿帧之后若 graph 已失效,后续帧持续输出垃圾。
+
+## 二、根因(两个独立但同源的缺陷)
+
+### 根因 A:Jetson Thor CUDA 驱动空闲重置缺陷(平台级,无法从应用层修复)
+
+- **GPU 连续空闲超过约 200–300ms** 后,驱动进入某种低功耗/重置状态;
+ 下一次任意 CUDA 调用(哪怕 2KB 的 H2D 拷贝或一个 kernel launch)
+ 会在**内核态自旋约 3.5 秒**(实测 3.3–3.9s,`stime` 主导、wchan=0,
+ 周期 22–30s 随空闲节奏出现)。
+- 该重置过程**同时破坏已捕获的 CUDA graph**:之后的 `graph.replay()`
+ 跑得很快但输出垃圾——这就是现象 1 的根源。
+
+证据链(可复现,探针脚本见"复现方法"):
+
+| 实验 | 结果 |
+|---|---|
+| 纯 2KB H2D 拷贝 @1Hz,无任何模型 | 每 22–27s 卡一次,每次 ~3.49s |
+| 纯 kernel launch + sync | 同样卡顿 → 与拷贝/分配无关 |
+| `gc.disable()` | 无效 → 排除 Python GC |
+| CPU 调频(schedutil,实测锁 2.6GHz)/内存压力 | 无异常 → 排除 |
+| GPU 持续繁忙(matmul 循环)60s | **0 卡顿** |
+| 心跳周期 100/200ms | 0 卡顿 |
+| 心跳周期 300/500/1000ms | 复现卡顿 |
+
+结论:**空闲阈值在 200–300ms 之间;保活算子大小无关,只要 GPU 不空闲
+超过约 200ms 即可**。机器上无 `jetson_clocks`/`nvpmodel`,无法用官方
+工具锁定电源状态;只能靠保活规避。
+
+### 根因 B:CUDA graph 无空闲防护(代码级)
+
+前端原有逻辑:graph 一旦捕获就永远 replay。在根因 A 触发后,
+replay 的是已被驱动重置破坏的 graph → 垃圾输出。需要在 replay 前
+判断"距上次使用是否超过安全窗口",超时则丢弃重捕获。
+
+## 三、修复方案(三层防护,已在 N1.7 落地)
+
+### 防护 1:GPU 心跳保活(主修复,规避根因 A)
+
+位置:`serving/groot_n17/run_http_policy.py` lifespan。
+
+- 每 **150ms**(环境变量 `FLASHRT_GPU_KEEPALIVE`,秒,默认 `0.15`)
+ 向 GPU 提交一个微小算子(`torch.ones(16)` 加法 + `synchronize`)。
+- 心跳通过与推理相同的 `ThreadPoolExecutor` 提交,推理忙时自动排队让路。
+- **关键约束:心跳绝不能更新前端的 `_last_graph_use` 时间戳**,
+ 否则会掩盖防护 2 的空闲判断(早期踩过这个坑:心跳更新时间戳后,
+ 40s 空闲后直接输出垃圾)。
+- 设为 `0` 可关闭(不建议:机器人按 action chunk 消费时请求间隙
+ 常达 0.5–2s,必然触发)。
+
+### 防护 2:空闲重捕获守卫(兜底,对抗根因 B)
+
+空闲判断收敛在基类 `GrootN17TorchFrontendThor` 的共享助手:
+`_graph_idle_limit_s`(读 `FLASHRT_GRAPH_IDLE_REINIT_S`)、
+`_graph_idle_stale()`、`invalidate_graphs()`(外部强制重捕获入口)。
+两处对称使用:
+
+1. **backbone graph** — `flash_rt/frontends/torch/groot_n17_thor_fp8.py`
+ `set_prompt()`:比较 `time.monotonic() - _last_graph_use` 与
+ `FLASHRT_GRAPH_IDLE_REINIT_S`(默认 **2s**);超时或形状变化则
+ `reset_prompt_runtime()` 重新捕获,否则走快路径(copy_ + replay)。
+2. **DiT graph** — `flash_rt/frontends/torch/groot_n17_thor.py`
+ `infer()`:`use_dit_graph` 路径同样检查空闲时间,超时则
+ `del _k_dit_graph` 后重新 `_capture_kernel_dit_graphs()`。
+
+阈值演进:10s → 5s(9s 间隙实测损坏)→ 2s(4–5s 间隙实测损坏,
+且 5.0s 整的间隙因严格 `>` 比较漏网)。有心跳后正常不会触发,
+只在心跳失效或极端长间隙时兜底;触发时代价 ~500ms/帧。
+
+### 防护 3:graph vs eager 自检(保险,保证输出正确性)
+
+`serving/groot_n17/run_http_policy.py`,`FLASHRT_SELFCHECK=N`(每 N 帧):
+
+- 用**固定噪声**分别跑 graph 路径和 eager 路径(`use_dit_graph=False`),
+ 比较输出 maxΔ:健康基线 0.015–0.05;graph 损坏时 3.3–3.8。
+- 阈值 0.5 超限 → `fe.invalidate_graphs()` 强制重捕获 → 重跑该帧,
+ 客户端拿到的仍是正确结果。
+- 注意开销:每帧自检会让平均耗时从 ~70ms 升到 ~170ms,
+ 生产用 `N=10` 或更大。
+
+## 四、验证结果(5558,2026-08-08)
+
+- `stress_gaps`(60 帧 @1Hz + 每 10 帧 4.5s 间隙):**0 卡顿、0 坏帧**;
+ 普通帧 ~65ms,gap 后重捕获帧 ~500ms,selfcheck maxΔ=0.03。
+- 修复前同样测试:必现 1 次 3.5–4.5s 卡顿,4–5s 间隙后 4 次 graph
+ 损坏(d=3.29–3.84)。
+
+## 五、系统化修复清单(待执行)
+
+上述三层防护目前只落在 GR00T N1.7 链路上,需推广:
+
+- [ ] **N1.6 serving**(`serving/groot_n16/`):同样使用 captured graph,
+ 需移植空闲守卫 + 心跳 + 自检(先做离线数值一致性验证,任务 #10)。
+- [ ] **其他 Thor 前端**:Chameleon-7B、HyVLA-0.5 等 SM110 前端的
+ serving 入口,凡使用 CUDA graph 的都要加:
+ 1. 心跳保活(默认 0.15s);
+ 2. graph replay 前的空闲检查 + 重捕获;
+ 3. 可选自检。
+- [ ] **守卫参数统一**:把 `FLASHRT_GRAPH_IDLE_REINIT_S`(默认 2s)、
+ `FLASHRT_GPU_KEEPALIVE`(默认 0.15s)、`FLASHRT_SELFCHECK`
+ 三个环境变量的语义写进 `docs/serving_production.md`,
+ 避免各 serving 脚本各自实现走样。
+- [x] **代码清理**(2026-08-09 已完成,dev→生产整理):
+ `gc.freeze()` 已删(与本问题无关);debug 开关
+ `FLASHRT_REPROMPT_DEBUG`/`FLASHRT_FORCE_RECAPTURE`/`FLASHRT_DUMP_OBS`/
+ `FLASHRT_BUILD_STAGE_TIMING` 及 `[act]` 数组打印已删;
+ `aux_builder.py` 的 `[builder-slow]`/`[preprocess-slow]` 慢帧日志保留
+ (仅 >100ms 触发,零常态开销)。
+- [ ] **提交**:改动文件(均未提交)——
+ `flash_rt/frontends/torch/groot_n17_thor_fp8.py`、
+ `flash_rt/frontends/torch/groot_n17_thor.py`、
+ `serving/groot_n17/run_http_policy.py`、
+ `serving/groot_n17/aux_builder.py`、
+ `serving/groot_n17/eager_server.py`。建议 conventional commits 分拆。
+- [ ] **AGENTS.md**:在架构规则处加一条 Thor 空闲重置缺陷提示,
+ 指向本文档。
+
+## 六、复现与诊断方法(备查)
+
+- 最小复现:`/tmp/h2d_probe.py`(2KB H2D @1Hz)、
+ `/tmp/cuda_op_probe.py`(kernel/alloc/pinned 分模式)、
+ `/tmp/keepalive_probe.py`(`BEAT_MS`/`BEAT_TINY` 扫周期)。
+ 典型输出:`wall=3494ms cpu=3485ms STALL`(内核态自旋特征)。
+- 服务侧诊断:`FLASHRT_GROOT_TIMING=1` 输出每帧
+ `[timing] preprocess+backbone=… set_prompt=… infer=… total=…`。
+- 判别口诀:
+ - 错误帧**更快**(50ms)+ 间隙后出现 → graph 损坏,查防护 2;
+ - 单帧**更慢** 3.5s+ 且输出正常 → 驱动空闲重置,查防护 1;
+ - 平均耗时整体抬升 → 检查自检频率(`FLASHRT_SELFCHECK`)。
+
+## 七、关键事实速查
+
+| 项 | 值 |
+|---|---|
+| 空闲触发阈值 | 200–300ms(200 安全,300 复现) |
+| 驱动停顿时长 | ~3.5s(内核态自旋) |
+| graph 损坏表现 | replay 快(50ms)但输出垃圾,maxΔ≈3.3–3.8 |
+| 健康自检 maxΔ | 0.015–0.05,阈值 0.5 |
+| 心跳默认周期 | 150ms(`FLASHRT_GPU_KEEPALIVE=0.15`) |
+| 空闲守卫阈值 | 2s(`FLASHRT_GRAPH_IDLE_REINIT_S=2`) |
+| 正常帧耗时 | ~65–70ms;重捕获帧 ~500ms;自检帧 +100ms |
From 4cae6c0eb27ce47c3bd80c775fb8acf320478037 Mon Sep 17 00:00:00 2001
From: DXICM <10598463@qq.com>
Date: Mon, 17 Aug 2026 09:24:39 +0000
Subject: [PATCH 4/6] fix(fa4): set CUTE_DSL_ARCH before cutlass-dsl import
(SM110 NVVM ICE)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
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 7fd75d20).
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.
---
.../flashrt_fa4/cute/interface_fwd_sm100.py | 12 +-
docs/groot_n16_thor_sm110.md | 2 +-
flash_rt/frontends/torch/groot_thor.py | 58 +++--
flash_rt/hardware/thor/fa4_backend.py | 21 +-
tools/convert_groot_n16_hf_checkpoint.py | 217 ------------------
5 files changed, 48 insertions(+), 262 deletions(-)
delete mode 100644 tools/convert_groot_n16_hf_checkpoint.py
diff --git a/csrc/attention/flash_attn_4_src/flashrt_fa4/cute/interface_fwd_sm100.py b/csrc/attention/flash_attn_4_src/flashrt_fa4/cute/interface_fwd_sm100.py
index 3eb0e664..ef1250fb 100644
--- a/csrc/attention/flash_attn_4_src/flashrt_fa4/cute/interface_fwd_sm100.py
+++ b/csrc/attention/flash_attn_4_src/flashrt_fa4/cute/interface_fwd_sm100.py
@@ -56,9 +56,7 @@ def _raise(*_a, **_k):
FlashAttentionForwardSm90 = _fa4_trimmed("FlashAttentionForwardSm90")
FlashAttentionForwardSm120 = _fa4_trimmed("FlashAttentionForwardSm120")
FlashAttentionMLAForwardSm100 = _fa4_trimmed("FlashAttentionMLAForwardSm100")
-from flashrt_fa4.cute.sm100_hd256_2cta_fmha_forward import (
- BlackwellFusedMultiHeadAttentionForward,
-)
+BlackwellFusedMultiHeadAttentionForward = _fa4_trimmed("BlackwellFusedMultiHeadAttentionForward")
from flashrt_fa4.cute.block_sparsity import (
BlockSparseTensorsTorch,
@@ -580,12 +578,8 @@ def _flash_attn_fwd(
and (tile_m % qhead_per_kvhead == 0 or not pack_gqa)
)
- # hd=256 2CTA forward uses the dedicated kernel on both SM100 and SM110.
- use_dedicated_hd256_kernel = (
- arch // 10 in [10, 11]
- and head_dim == 256
- and head_dim_v == 256
- )
+ # hd=256 2CTA forward uses dedicated kernel (SM100 only; SM110 NVVM ICE)
+ use_dedicated_hd256_kernel = arch // 10 == 10 and head_dim == 256 and head_dim_v == 256
use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel
if softcap is not None:
diff --git a/docs/groot_n16_thor_sm110.md b/docs/groot_n16_thor_sm110.md
index 0200b1dc..d98019ac 100644
--- a/docs/groot_n16_thor_sm110.md
+++ b/docs/groot_n16_thor_sm110.md
@@ -316,7 +316,7 @@ obs ──> preprocess（apply_state 直连 + 线程池图像变换）
- 新 kernel：`csrc/fused_fp4/{dit_norm_fp4_sfa,silu_mul_fp4_sfa_bf16}.{cu,cuh}`、
`csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.{cu,cuh}`、
`csrc/quantize/quantize_fp4_sfa_bf16.{cu,cuh}`、`csrc/fp4_bindings.cpp`
-- 权重转换：`tools/convert_groot_n16_hf_checkpoint.py`
+- 权重转换（开发调试工具，不含于本 PR）：前端已内置完整的 HF→FlashRT layout 变换（transpose、QKV fuse），直接加载 HF 原始 safetensors 即可推理。离线审计脚本 `convert_groot_n16_hf_checkpoint.py` 仅用于 parity 对拍时定位 weight mapping 问题。
- 文档：本文 + `docs/groot_transformers5_weight_corruption.md`
**配套文档（独立主题）**：
diff --git a/flash_rt/frontends/torch/groot_thor.py b/flash_rt/frontends/torch/groot_thor.py
index 394c5b0c..21ed7962 100644
--- a/flash_rt/frontends/torch/groot_thor.py
+++ b/flash_rt/frontends/torch/groot_thor.py
@@ -859,17 +859,22 @@ def _setup_torch_dit(self):
@staticmethod
def _resolve_eagle_dir() -> pathlib.Path:
- """Locate the Eagle-Block2A-2B-v2 remote-code directory."""
+ """Locate the Eagle-Block2A-2B-v2 remote-code directory.
+
+ The HF modules cache stores only .py files (no config.json), so we
+ use modeling_siglip2.py as the marker.
+ """
override = os.environ.get("FLASHRT_N16_EAGLE_DIR")
if override:
p = pathlib.Path(override)
- if (p / "config.json").exists():
+ if (p / "modeling_siglip2.py").exists() or (p / "config.json").exists():
return p
raise RuntimeError(
- f"FLASHRT_N16_EAGLE_DIR={override} has no config.json")
+ f"FLASHRT_N16_EAGLE_DIR={override} has neither "
+ "modeling_siglip2.py nor config.json")
cache = pathlib.Path.home() / ".cache/huggingface/modules/transformers_modules"
for cand in sorted(cache.glob(
- "Eagle_hyphen_Block2A_hyphen_2B_hyphen_v2/*/config.json")):
+ "Eagle_hyphen_Block2A_hyphen_2B_hyphen_v2/*/modeling_siglip2.py")):
return cand.parent
raise RuntimeError(
"Eagle-Block2A-2B-v2 remote code not found. Load the GR00T N1.6 "
@@ -879,25 +884,26 @@ def _resolve_eagle_dir() -> pathlib.Path:
def _setup_torch_siglip(self):
"""Parity mode: HF-native Siglip2VisionModel (bf16) from the Eagle
remote code, weights from the checkpoint state dict."""
- import glob as _glob
import importlib.util as _ilu
import json as _json
eagle_dir = self._resolve_eagle_dir()
- cfg = _json.load(open(eagle_dir / "config.json"))
- mod_path = None
- for cand in _glob.glob(str(eagle_dir / "modeling_siglip2.py")) + \
- _glob.glob(str(pathlib.Path.home() /
- ".cache/huggingface/modules/transformers_modules/"
- "Eagle_hyphen_Block2A_hyphen_2B_hyphen_v2/*/"
- "modeling_siglip2.py")):
- mod_path = cand
- break
- if mod_path is None:
- raise RuntimeError("modeling_siglip2.py not found for parity SigLIP")
- spec = _ilu.spec_from_file_location("eagle_siglip2_parity", mod_path)
+ mod_path = eagle_dir / "modeling_siglip2.py"
+ if not mod_path.exists():
+ raise RuntimeError(
+ f"modeling_siglip2.py not found in {eagle_dir}")
+ spec = _ilu.spec_from_file_location("eagle_siglip2_parity", str(mod_path))
mod = _ilu.module_from_spec(spec)
spec.loader.exec_module(mod)
- vc = dict(cfg["vision_config"])
+ cfg_json = eagle_dir / "config.json"
+ if cfg_json.exists():
+ vc = dict(_json.load(open(cfg_json))["vision_config"])
+ else:
+ vc = dict(
+ hidden_size=self.D_sig, num_attention_heads=self.NH_sig,
+ num_hidden_layers=self.L_sig, intermediate_size=self.H_sig,
+ image_size=252, patch_size=14, hidden_act="gelu_pytorch_tanh",
+ layer_norm_eps=1e-6, attention_dropout=0.0, dropout=0.0,
+ )
vc.pop("_attn_implementation_autoset", None)
model = mod.Siglip2VisionModel(mod.Siglip2VisionConfig(**vc))
model = model.to(torch.bfloat16).cuda().eval()
@@ -965,8 +971,20 @@ def _setup_torch_qwen3(self):
import json as _json
from transformers import Qwen3Config, Qwen3Model
eagle_dir = self._resolve_eagle_dir()
- cfg = _json.load(open(eagle_dir / "config.json"))
- tc = dict(cfg.get("text_config", cfg))
+ cfg_json = eagle_dir / "config.json"
+ if cfg_json.exists():
+ tc = dict(_json.load(open(cfg_json)).get("text_config", {}))
+ else:
+ tc = {}
+ if not tc:
+ tc = dict(
+ hidden_size=self.D_llm, num_attention_heads=self.NHQ,
+ num_key_value_heads=self.NHKV, head_dim=self.HD_llm,
+ intermediate_size=self.H_llm, hidden_act="silu",
+ max_position_embeddings=4096, rms_norm_eps=1e-6,
+ rope_theta=1000000.0, use_sliding_window=False,
+ vocab_size=151680,
+ )
tc["num_hidden_layers"] = 16 # checkpoint is truncated
tc["_attn_implementation"] = "sdpa"
config = Qwen3Config(**tc)
diff --git a/flash_rt/hardware/thor/fa4_backend.py b/flash_rt/hardware/thor/fa4_backend.py
index a8df381d..9b292913 100644
--- a/flash_rt/hardware/thor/fa4_backend.py
+++ b/flash_rt/hardware/thor/fa4_backend.py
@@ -53,21 +53,12 @@ def _load() -> None:
_REASON = "disabled (FLASHRT_THOR_FA4=0)"
return
- # Thor compilation target. The accepted chip string depends on the
- # installed nvidia-cutlass-dsl: 4.5+ needs ``sm_101a`` (the legacy
- # sm_110 alias; its ``sm_110a`` path hits a chip-string bug), while
- # 4.4.x only knows ``sm_110a``. Either way the vendored SM100-compatible
- # forward is what runs — this tree ships no separate SM110 kernel — so
- # FLASH_ATTENTION_ARCH stays sm_100a below. Don't clobber a user value.
- try:
- import cutlass as _ctl # nvidia-cutlass-dsl
- _dsl_ver = tuple(
- int(x) for x in str(getattr(_ctl, "__version__", "0.0"))
- .split(".")[:2])
- except Exception: # noqa: BLE001 — any failure: assume the newer alias
- _dsl_ver = (4, 5)
- os.environ.setdefault(
- "CUTE_DSL_ARCH", "sm_101a" if _dsl_ver >= (4, 5) else "sm_110a")
+ # Thor compilation target: sm_101a (the legacy alias of sm_110).
+ # Must be set BEFORE importing cutlass-dsl, which caches the device
+ # arch at import time; setting it after is too late (NVVM ICE on
+ # sm_110a). Requires nvidia-cutlass-dsl >= 4.5 (which accepts the
+ # sm_101a alias); 4.4.x only knows sm_110a and is NOT supported.
+ os.environ.setdefault("CUTE_DSL_ARCH", "sm_101a")
os.environ.setdefault("FLASH_ATTENTION_ARCH", "sm_100a")
# Allow an override dir (e.g. for local FA4 development), else the vendor.
diff --git a/tools/convert_groot_n16_hf_checkpoint.py b/tools/convert_groot_n16_hf_checkpoint.py
deleted file mode 100644
index 0a9244c6..00000000
--- a/tools/convert_groot_n16_hf_checkpoint.py
+++ /dev/null
@@ -1,217 +0,0 @@
-#!/usr/bin/env python3
-"""Offline converter: HF GR00T N1.6 checkpoint -> FlashRT weight layout.
-
-Phase-1b of the N1.6 parity plan. Converts a fine-tuned (or base) HF GR00T
-N1.6 checkpoint into the exact tensor layout the FlashRT N1.6 kernels
-consume, emitting BOTH the converted ``.safetensors`` and a JSON manifest
-recording, for every tensor: HF key, HF shape, transform, FlashRT key,
-FlashRT shape. This makes every weight mapping explicit and auditable, so a
-backbone/weight-layout mismatch can be localized to a specific rule instead
-of guessed from the final action.
-
-Layout rules (one explicit rule per weight family):
-
- A. SigLIP2 (``backbone.model.vision_model.vision_model.*``)
- - attention q/k/v/o and FFN fc1/fc2: HF ``[out,in]`` -> ``[in,out]``
- (FlashRT GEMMs take ``[in,out]``); QKV kept separate, order Q,K,V.
- - layernorm / position_embedding / patch_embedding bias: passthrough.
- - (the double ``vision_model`` prefix is part of the HF key and kept).
-
- B. Qwen3 (``backbone.model.language_model.model.layers.*``)
- - q ``[2048,2048]``, k ``[1024,2048]``, v ``[1024,2048]`` are fused as
- ``cat([q,k,v], dim=0).T.contiguous()`` -> ``[2048, 4096]``
- (Q first, then K, then V; NO interleaving).
- - FFN ``cat([gate_proj, up_proj], dim=0).T.contiguous()`` with order
- ``gate | up`` (must not be swapped); down_proj transposed.
- - layernorm / q_norm / k_norm: passthrough.
-
- C. DiT (``action_head.model.transformer_blocks.{l}.*``)
- - even block = cross-attention, odd block = self-attention.
- - self-attn (odd): QKV fused ``cat([q,k,v],0).T`` -> ``[1536, 4608]``
- (K/V input dim 1536).
- - cross-attn (even): q transposed; k/v ``[1536,2048]`` transposed,
- kept separate (NOT fused).
- - FFN is GELU (not GEGLU): net.0.proj and net.2 transposed.
- - norm1.linear (AdaLN) transposed; norm1.norm / norm3 / norm_out have
- NO affine parameters (absent from the checkpoint); output
- conditioning chunk order is (shift, scale).
- - proj_out_1 / proj_out_2 / timestep_encoder linears transposed.
-
- D. Embodiment (``action_head.{action_encoder,state_encoder,action_decoder}``)
- - CategorySpecificLinear ``W`` is already ``[num_categories,in,out]``;
- after selecting ``W[eid]`` it is ``[in,out]`` already -> NO extra
- transpose. Biases passthrough.
-
-Usage:
- python tools/convert_groot_n16_hf_checkpoint.py \
- --src /mnt/lerobot_so101_sim_v1_gr00t_n1d6_sim_fruits_cubes_10w \
- --dst /mnt/.../n1d6_flashrt_layout \
- [--dtype fp16]
-"""
-from __future__ import annotations
-
-import argparse
-import json
-from pathlib import Path
-
-import torch
-from safetensors import safe_open
-from safetensors.torch import save_file
-
-
-def load_src(src: Path) -> dict:
- sd = {}
- for f in sorted(src.glob("*.safetensors")):
- with safe_open(str(f), framework="pt") as sf:
- for k in sf.keys():
- sd[k] = sf.get_tensor(k)
- return sd
-
-
-def main() -> None:
- ap = argparse.ArgumentParser()
- ap.add_argument("--src", required=True)
- ap.add_argument("--dst", required=True)
- ap.add_argument("--dtype", default="keep", choices=["keep", "fp16", "bf16"])
- args = ap.parse_args()
-
- src = Path(args.src)
- dst = Path(args.dst)
- dst.mkdir(parents=True, exist_ok=True)
- sd = load_src(src)
-
- cast = {"keep": None, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype]
- def maybe_cast(t: torch.Tensor) -> torch.Tensor:
- return t.to(cast) if (cast is not None and t.is_floating_point()) else t
-
- out: dict[str, torch.Tensor] = {}
- manifest: list[dict] = []
-
- def emit(hf_key: str, fr_key: str, tensor: torch.Tensor, transform: str) -> None:
- hf_shape = list(sd[hf_key].shape) if hf_key in sd else None
- tensor = maybe_cast(tensor)
- out[fr_key] = tensor
- manifest.append({
- "hf_key": hf_key, "hf_shape": hf_shape, "transform": transform,
- "fr_key": fr_key, "fr_shape": list(tensor.shape),
- })
-
- T = lambda t: t.T.contiguous() # [out,in] -> [in,out]
-
- for k, v in sd.items():
- # ── A. SigLIP2 ──
- if k.startswith("backbone.model.vision_model.vision_model."):
- if any(k.endswith(s) for s in (".to_q.weight", ".to_k.weight",
- ".to_v.weight", ".to_out.0.weight",
- ".mlp.fc1.weight", ".mlp.fc2.weight")):
- emit(k, k, T(v), "transpose[out,in]->[in,out]")
- else:
- emit(k, k, v, "passthrough")
-
- # ── B. Qwen3 ──
- elif k.startswith("backbone.model.language_model.model.layers."):
- if k.endswith(".self_attn.q_proj.weight"):
- # fuse when the sibling k/v are present (handled at q key)
- pre = k[: -len(".self_attn.q_proj.weight")]
- q = sd[f"{pre}.self_attn.q_proj.weight"]
- kk = sd[f"{pre}.self_attn.k_proj.weight"]
- vv = sd[f"{pre}.self_attn.v_proj.weight"]
- fused = torch.cat([q, kk, vv], dim=0).T.contiguous()
- emit(k, f"{pre}.self_attn.qkv_fused", fused,
- "cat([q,k,v],0).T -> [in,out], order Q,K,V")
- elif k.endswith((".self_attn.k_proj.weight",
- ".self_attn.v_proj.weight")):
- continue # already fused into qkv_fused
- elif k.endswith(".self_attn.q_proj.bias"):
- pre = k[: -len(".self_attn.q_proj.bias")]
- q = sd.get(f"{pre}.self_attn.q_proj.bias")
- if q is not None:
- kk = sd[f"{pre}.self_attn.k_proj.bias"]; vv = sd[f"{pre}.self_attn.v_proj.bias"]
- emit(k, f"{pre}.self_attn.qkv_bias", torch.cat([q, kk, vv], 0),
- "cat([qb,kb,vb],0)")
- else:
- emit(k, k, v, "passthrough")
- elif k.endswith((".self_attn.k_proj.bias", ".self_attn.v_proj.bias")):
- pre = k[: -len(".self_attn.k_proj.bias")]
- if f"{pre}.self_attn.q_proj.bias" in sd:
- continue # fused into qkv_bias
- emit(k, k, v, "passthrough")
- elif k.endswith(".mlp.gate_proj.weight"):
- pre = k[: -len(".mlp.gate_proj.weight")]
- g = sd[f"{pre}.mlp.gate_proj.weight"]
- u = sd[f"{pre}.mlp.up_proj.weight"]
- emit(k, f"{pre}.mlp.gate_up_fused",
- torch.cat([g, u], dim=0).T.contiguous(),
- "cat([gate,up],0).T -> [in,out], order gate|up")
- elif k.endswith(".mlp.up_proj.weight"):
- continue # fused into gate_up_fused
- elif k.endswith(".mlp.down_proj.weight"):
- emit(k, k, T(v), "transpose[out,in]->[in,out]")
- else:
- emit(k, k, v, "passthrough")
-
- # ── C. DiT ──
- elif k.startswith("action_head.model.transformer_blocks."):
- parts = k.split(".")
- l = int(parts[3])
- is_self = (l % 2 == 1)
- if k.endswith(".attn1.to_q.weight"):
- pre = k[: -len(".attn1.to_q.weight")]
- q = sd[f"{pre}.attn1.to_q.weight"]
- if is_self:
- kk = sd[f"{pre}.attn1.to_k.weight"]
- vv = sd[f"{pre}.attn1.to_v.weight"]
- emit(k, f"{pre}.attn1.qkv_fused",
- torch.cat([q, kk, vv], dim=0).T.contiguous(),
- "self-attn cat([q,k,v],0).T -> [in,out]")
- else:
- emit(k, k, T(q), "cross-attn q transpose")
- elif k.endswith(".attn1.to_k.weight") or k.endswith(".attn1.to_v.weight"):
- if is_self:
- continue # fused
- emit(k, k, T(v), "cross-attn k/v transpose [1536,2048]->[2048,1536]")
- elif k.endswith((".attn1.to_q.bias",)):
- pre = k[: -len(".attn1.to_q.bias")]
- q = sd[f"{pre}.attn1.to_q.bias"]; kk = sd[f"{pre}.attn1.to_k.bias"]; vv = sd[f"{pre}.attn1.to_v.bias"]
- if is_self:
- emit(k, f"{pre}.qkv_bias", torch.cat([q, kk, vv], 0), "cat([qb,kb,vb],0)")
- else:
- emit(k, k, v, "passthrough")
- elif k.endswith((".attn1.to_k.bias", ".attn1.to_v.bias")):
- if is_self:
- continue
- emit(k, k, v, "passthrough")
- elif k.endswith(".ff.net.0.proj.weight") or k.endswith(".ff.net.2.weight") \
- or k.endswith(".norm1.linear.weight") or k.endswith(".attn1.to_out.0.weight"):
- emit(k, k, T(v), "transpose[out,in]->[in,out] (GELU FFN, not GEGLU)")
- else:
- emit(k, k, v, "passthrough")
-
- # ── DiT top-level linears ──
- elif k.startswith("action_head.model.") and k.endswith(".weight") \
- and any(s in k for s in ("proj_out_1", "proj_out_2", "timestep_embedder")):
- emit(k, k, T(v), "transpose[out,in]->[in,out]")
-
- # ── D. Embodiment CategorySpecificLinear (NO transpose) ──
- elif k.startswith("action_head.") and any(
- s in k for s in ("action_encoder", "state_encoder", "action_decoder")):
- emit(k, k, v, "passthrough (W[eid] already [in,out]; no transpose)")
-
- # ── everything else (embeddings, vlln, mlp1, norms, etc.) ──
- else:
- emit(k, k, v, "passthrough")
-
- save_file(out, str(dst / "model_flashrt.safetensors"))
- with open(dst / "flashrt_layout_manifest.json", "w") as f:
- json.dump(manifest, f, indent=1)
- print(f"wrote {dst/'model_flashrt.safetensors'} ({len(out)} tensors)")
- print(f"wrote {dst/'flashrt_layout_manifest.json'} ({len(manifest)} rules)")
- # summary of transforms
- from collections import Counter
- c = Counter(m["transform"] for m in manifest)
- for t, n in c.most_common():
- print(f" {n:5d} {t}")
-
-
-if __name__ == "__main__":
- main()
From eb01fa8d470df024d4245d64a1ffa774f1f4f94f Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Sun, 23 Aug 2026 13:03:51 -0400
Subject: [PATCH 5/6] fix(groot): close PR 177 audit gaps
---
csrc/bindings.cpp | 4 ++
docs/groot_n16_thor_sm110.md | 20 +++---
flash_rt/frontends/torch/groot_thor.py | 65 ++++++++++---------
flash_rt/hardware/thor/fa4_backend.py | 8 +--
tests/test_groot_n16_pr177_maintenance.py | 79 +++++++++++++++++++++++
tests/test_prompt_runtime_lifecycle.py | 13 ++--
6 files changed, 138 insertions(+), 51 deletions(-)
create mode 100644 tests/test_groot_n16_pr177_maintenance.py
diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp
index 0b401947..2f7f7e35 100644
--- a/csrc/bindings.cpp
+++ b/csrc/bindings.cpp
@@ -312,7 +312,9 @@ extern "C" void flash_rt_awq_quant_fp8_static_fp16(
#include "quantize/bf16_quant_fp8_ncdhw_to_ndhwc.cuh"
#endif
#include "quantize/qkv_split_norm_rope_bf16.cuh"
+#ifdef FLASHRT_HAVE_THOR_VLA_KERNELS
#include "kernels/qk_norm_rope_rotate_half_bf16.cuh"
+#endif
#include "attention/fmha_dispatch.h"
#ifdef ENABLE_MOTUS_SAGE2_RAW
#include "attention/sage2/sage2_attn_raw.cuh"
@@ -1801,6 +1803,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
py::arg("eps") = 1e-5f, py::arg("stream") = 0);
#endif // FLASHRT_ENABLE_CHAMELEON
+#ifdef FLASHRT_HAVE_THOR_VLA_KERNELS
m.def("qk_norm_rope_rotate_half_bf16",
[](uintptr_t x, uintptr_t w, uintptr_t cos_t, uintptr_t sin_t,
int S, int NH, int HD, float eps, uintptr_t stream) -> int {
@@ -1810,6 +1813,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
}, py::arg("x"), py::arg("w"), py::arg("cos_table"), py::arg("sin_table"),
py::arg("S"), py::arg("NH"), py::arg("HD"), py::arg("eps") = 1e-6f,
py::arg("stream") = 0);
+#endif // FLASHRT_HAVE_THOR_VLA_KERNELS
m.def("gate_mul_residual_fp16",
[](uintptr_t residual, uintptr_t x, uintptr_t gate,
diff --git a/docs/groot_n16_thor_sm110.md b/docs/groot_n16_thor_sm110.md
index d98019ac..f594442c 100644
--- a/docs/groot_n16_thor_sm110.md
+++ b/docs/groot_n16_thor_sm110.md
@@ -162,7 +162,7 @@ HF 实际行为不同；叠加若干实现 bug。全部修复集中在
- **FA4（FlashAttention-4 CuTe-DSL，`flash_rt/hardware/thor/fa4_backend.py`）用于
SigLIP**：SigLIP 本就是 cross-view full attention（648 token 单序列），FA4
causal=False 为 sdpa 精确替换；encoder graph 10.3→9.0 ms，动作 cos 1.000000。
- 开关 `FLASHRT_N16_FA4`（默认开，缺失自动回退）。
+ 开关 `FLASHRT_N16_FA4=1`（显式开启，缺失自动回退）。
- **FA4 不用于 DiT/Qwen3**：小 seq（51/208）下 FA4 0.042–0.394 ms ≫ sdpa 0.013–0.015 ms。
- **NVFP4（W4A4 CUTLASS）用于 DiT**：block GEMM 走
`quantize_fp4_dynamic_sfa_fp16` + `cutlass_fp4_sq_fp16`（per-16 block scale、动态激活
@@ -193,7 +193,7 @@ DiT 链每层仅 8 kernel（无层间逐元素流量）：**DiT 36.6→15.7 ms**
### 6.5 全 kernel 化三轮（37.5 → 28.5 ms）
-- **第一轮 — Qwen3 融合**（`FLASHRT_N16_QWEN3_FP4`，生产已启用）：16 层全部
+- **第一轮 — Qwen3 融合**（`FLASHRT_N16_QWEN3_FP4=1`，显式 fast profile）：16 层全部
q/k/v/o/gate/up/down 走 W4A4 融合 GEMM（residual 原地更新）。
- 新写两个 bf16 输入 producer kernel（`csrc/fused_fp4/`，与 torch 两步链 **bit 一致**）：
`rms_norm_weight_fp4_sfa_bf16`（加权 RMSNorm→fp4 直出）、
@@ -204,7 +204,7 @@ DiT 链每层仅 8 kernel（无层间逐元素流量）：**DiT 36.6→15.7 ms**
- **fused GQA attention**：torch `enable_gqa` 会退回非融合 math 路径（每层 2 个
fp16 GEMM + 显式 softmax）；改 `repeat_interleave` 展开 KV 走融合后端：−2.0 ms。
- Qwen3 12.7→5.0 ms；vs HF cos 0.999986。
-- **第二轮 — SigLIP encoder fp4**（`FLASHRT_N16_SIGLIP_FP4`，默认开）：LN producer
+- **第二轮 — SigLIP encoder fp4**（`FLASHRT_N16_SIGLIP_FP4=1`，显式 fast profile）：LN producer
用 scale=w−1/shift=b 表达 affine；q/k/v 分离 GEMM 写连续缓冲直喂 FA4；o/ffn 残差进
epilogue；fc1 N 与 fc2 K 4304→4352 pad（pad 维全零端到端无贡献）。encoder 9.0→6.9 ms。
全开时动作精度不降（vs HF cos 0.999988）。
@@ -236,7 +236,7 @@ DiT 链每层仅 8 kernel（无层间逐元素流量）：**DiT 36.6→15.7 ms**
## 7. 最终架构与配置
-### 7.1 数据流（parity + 三层 NVFP4 + FA4，生产默认）
+### 7.1 数据流（parity + 三层 NVFP4 + FA4，显式 fast profile）
```
obs ──> preprocess（apply_state 直连 + 线程池图像变换） ~2.8 ms
@@ -253,18 +253,18 @@ obs ──> preprocess（apply_state 直连 + 线程池图像变换）
| 开关 | 默认 | 作用 |
|---|---|---|
-| `FLASHRT_N16_DIT_FP4` | 生产开 | DiT NVFP4 融合链 |
-| `FLASHRT_N16_QWEN3_FP4` | 生产开 | Qwen3 NVFP4 融合层（含 fused norm/rope/GQA） |
-| `FLASHRT_N16_SIGLIP_FP4` | 开 | SigLIP encoder fp4 层 |
-| `FLASHRT_N16_FA4` | 开 | SigLIP FA4 注意力（缺失自动回退） |
+| `FLASHRT_N16_DIT_FP4` | 关 | DiT NVFP4 融合链；设为 `1` 开启 |
+| `FLASHRT_N16_QWEN3_FP4` | 关 | Qwen3 NVFP4 融合层（含 fused norm/rope/GQA）；设为 `1` 开启 |
+| `FLASHRT_N16_SIGLIP_FP4` | 关 | SigLIP encoder fp4 层；设为 `1` 开启 |
+| `FLASHRT_N16_FA4` | 关 | SigLIP FA4 注意力；设为 `1` 开启，缺失时回退 |
| `FLASHRT_N16_DIT_STEPS` | 4 | flow-matching 步数（**生产勿改**，减步伤行为） |
-| `--no-fp8` / `parity` | 生产开 | HF 原生 parity 通路（非 legacy kernel 快路径） |
+| `parity` | 关 | HF 原生 parity 通路（非 legacy kernel 快路径）；直接构造 frontend 时显式开启 |
精度汇总（去归一化动作 vs HF eager）：
| 配置 | cos | maxd |
|---|---|---|
-| 全开（生产） | 0.999933 | 0.059 |
+| 全开（显式 fast profile） | 0.999933 | 0.059 |
| Qwen3 fp4 单独 | 0.999986 | 0.023 |
| DiT fp4 单独 | 0.999995 | 0.012 |
diff --git a/flash_rt/frontends/torch/groot_thor.py b/flash_rt/frontends/torch/groot_thor.py
index 21ed7962..b580401e 100644
--- a/flash_rt/frontends/torch/groot_thor.py
+++ b/flash_rt/frontends/torch/groot_thor.py
@@ -81,8 +81,8 @@ class GrootTorchFrontendThor:
"""GROOT N1.6 inference pipeline on Thor SM110."""
def __init__(self, checkpoint, num_views=2, autotune=3,
- embodiment_tag="new_embodiment", use_fp8=False,
- image_size=252, parity=True):
+ embodiment_tag="new_embodiment", use_fp8=True,
+ image_size=224, parity=False):
"""Initialize GROOT pipeline.
Args:
@@ -94,6 +94,8 @@ def __init__(self, checkpoint, num_views=2, autotune=3,
224 = 16x16 patches (legacy); 252 = 18x18 patches, the
GR00T N1.6 training/eval resolution (HF processor chain
LetterBoxPad -> 256 -> 0.95 crop -> 252).
+ parity: opt into the HF-native BF16 backbone and DiT. The default
+ keeps the pre-existing fully-kernelized execution path.
"""
if embodiment_tag not in EMBODIMENT_TAG_TO_INDEX:
raise ValueError(
@@ -857,29 +859,29 @@ def _setup_torch_dit(self):
* (math.log(10000.0) / hd))
self._dit_tw = w
- @staticmethod
- def _resolve_eagle_dir() -> pathlib.Path:
+ def _resolve_eagle_dir(self) -> pathlib.Path:
"""Locate the Eagle-Block2A-2B-v2 remote-code directory.
- The HF modules cache stores only .py files (no config.json), so we
- use modeling_siglip2.py as the marker.
+ Only checkpoint-local code or an explicitly selected checkout is
+ accepted. Selecting a cache directory by sort order is not
+ reproducible when several revisions are installed.
"""
override = os.environ.get("FLASHRT_N16_EAGLE_DIR")
if override:
p = pathlib.Path(override)
- if (p / "modeling_siglip2.py").exists() or (p / "config.json").exists():
+ if (p / "modeling_siglip2.py").is_file():
return p
raise RuntimeError(
- f"FLASHRT_N16_EAGLE_DIR={override} has neither "
- "modeling_siglip2.py nor config.json")
- cache = pathlib.Path.home() / ".cache/huggingface/modules/transformers_modules"
- for cand in sorted(cache.glob(
- "Eagle_hyphen_Block2A_hyphen_2B_hyphen_v2/*/modeling_siglip2.py")):
- return cand.parent
+ f"FLASHRT_N16_EAGLE_DIR={override} has no modeling_siglip2.py")
+ for candidate in (
+ self._checkpoint_path,
+ self._checkpoint_path / "Eagle-Block2A-2B-v2"):
+ if (candidate / "modeling_siglip2.py").is_file():
+ return candidate
raise RuntimeError(
- "Eagle-Block2A-2B-v2 remote code not found. Load the GR00T N1.6 "
- "model once with transformers (AutoModel.from_pretrained) to "
- "populate the HF cache, or set FLASHRT_N16_EAGLE_DIR.")
+ "Eagle-Block2A-2B-v2 remote code was not found in the checkpoint. "
+ "Set FLASHRT_N16_EAGLE_DIR to an explicitly pinned checkout; "
+ "FlashRT does not select an arbitrary Hugging Face cache revision.")
def _setup_torch_siglip(self):
"""Parity mode: HF-native Siglip2VisionModel (bf16) from the Eagle
@@ -931,7 +933,7 @@ def _patch_siglip_fa4(self, mod):
(2026-08-14). FLASHRT_N16_FA4=0 forces sdpa; a missing FA4 runtime
falls back silently. Must run before the encoder graph is captured.
"""
- if os.environ.get("FLASHRT_N16_FA4", "1") == "0":
+ if os.environ.get("FLASHRT_N16_FA4", "0") != "1":
return
try:
from flash_rt.hardware.thor import fa4_backend
@@ -1010,7 +1012,7 @@ def _setup_qwen3_fp4(self):
and the residual stream is updated in place. Weight traffic drops
2.85 GB -> ~0.8 GB per inference.
"""
- if os.environ.get("FLASHRT_N16_QWEN3_FP4", "1") != "1":
+ if os.environ.get("FLASHRT_N16_QWEN3_FP4", "0") != "1":
return
if getattr(self, "_qwen3_fp4_done", False):
return
@@ -1177,7 +1179,7 @@ def _bias(t):
logger.info("Qwen3 NVFP4 fused-epilogue tier enabled (16 layers)")
def _setup_siglip_fp4(self):
- """FLASHRT_N16_SIGLIP_FP4 (default on): NVFP4 fused-epilogue encoder.
+ """FLASHRT_N16_SIGLIP_FP4=1: NVFP4 fused-epilogue encoder.
Every layer runs the kernel chain: LN->fp4 producer (affine LayerNorm
expressed as AdaLN with scale=w-1, shift=b), q/k/v bias GEMMs into
@@ -1188,7 +1190,7 @@ def _setup_siglip_fp4(self):
0.038 (vs 0.023 bf16) — the vision-quality trade is simulation-gated;
set FLASHRT_N16_SIGLIP_FP4=0 to revert to bf16.
"""
- if os.environ.get("FLASHRT_N16_SIGLIP_FP4", "1") != "1":
+ if os.environ.get("FLASHRT_N16_SIGLIP_FP4", "0") != "1":
return
if getattr(self, "_siglip_fp4_done", False):
return
@@ -1404,7 +1406,7 @@ def _dit_setup_graph_buffers(self):
# cos 0.999994 / maxd 0.012 (2026-08-14). Weight tables are static
# and survive graph re-captures.
self._dit_use_fp4 = False
- if os.environ.get("FLASHRT_N16_DIT_FP4", "1") == "1":
+ if os.environ.get("FLASHRT_N16_DIT_FP4", "0") == "1":
try:
import flash_rt.flash_rt_fp4 as _f4
except ImportError:
@@ -1900,6 +1902,7 @@ def reset_graph_runtime(self) -> None:
stale_exact = ("_siglip_graph", "_qwen3_graph", "_dit_graph",
"_attn", "_vision_features", "_unit_scale",
"_qwen3_torch_graph", "_siglip_torch_graph",
+ "_qwen3_fp4_done", "_siglip_fp4_done",
# DiT static buffers/indexes depend on Se/masks and
# must be rebuilt on prompt-switch re-capture.
"_dit_in_state", "_dit_in_kvt", "_dit_in_kvi",
@@ -1929,11 +1932,6 @@ def set_prompt(self, prompt, input_ids=None):
When supplied (e.g. from the serving aux builder, which
reproduces HF exactly), these ids are used verbatim.
"""
- if getattr(self, '_graphs_built', False):
- raise RuntimeError(
- "set_prompt() after the pipeline is built is not supported; "
- "construct a new GrootTorchFrontendThor instance for a new prompt")
-
# Image special-token ids (fixed for the Eagle vocab).
self._img_token_id = 151669 #
self._img_start_id = 151670 # ![]()
@@ -1952,7 +1950,6 @@ def set_prompt(self, prompt, input_ids=None):
# Local GROOT code Eagle dir
str(pathlib.Path(__file__).parent.parent.parent.parent.parent /
"GR00T" / "Isaac-GR00T" / "gr00t" / "model" / "modules" / "nvidia" / "Eagle-Block2A-2B-v2"),
- "nvidia/Eagle-Block2A-2B-v2", # HF hub (fallback)
]:
try:
self._tokenizer = AutoTokenizer.from_pretrained(tok_path, trust_remote_code=True)
@@ -1968,7 +1965,16 @@ def set_prompt(self, prompt, input_ids=None):
full_ids = text_ids + [self._img_start_id] + [self._img_token_id] * S_img + [self._img_end_id]
text_count = len(text_ids)
+ input_id_values = tuple(full_ids)
+ if getattr(self, '_graphs_built', False):
+ if input_id_values == getattr(self, '_input_id_values', None):
+ self._prompt_text = prompt
+ return
+ logger.info("Prompt changed after graph capture; resetting graph runtime")
+ self.reset_graph_runtime()
+
self._input_ids = torch.tensor([full_ids], dtype=torch.long, device='cuda')
+ self._input_id_values = input_id_values
self._text_len = text_count
self._Se = len(full_ids)
self._prompt_text = prompt
@@ -2761,8 +2767,9 @@ def _collect_dit_amax(self):
ae_concat = torch.empty(T, 2*D, dtype=fp16, device='cuda')
self._gemm.fp16_nn(actions_fp16.data_ptr(), dit.ae_w1.data_ptr(), a_emb_out.data_ptr(), T, D, dit.action_dim, 0)
fvk.add_bias_fp16(a_emb_out.data_ptr(), dit.ae_b1.data_ptr(), T, D, 0)
- fvk.gpu_copy(ae_concat.data_ptr(), a_emb_out.data_ptr(), T*D*2, 0)
- fvk.gpu_copy(ae_concat.data_ptr()+T*D*2, dit.action_time_embeds[0].data_ptr(), T*D*2, 0)
+ fvk.concat2_bf16(a_emb_out.data_ptr(),
+ dit.action_time_embeds[0].data_ptr(),
+ ae_concat.data_ptr(), T, D, D, 0)
enc_h = torch.empty(T, D, dtype=fp16, device='cuda')
self._gemm.fp16_nn(ae_concat.data_ptr(), dit.ae_w2.data_ptr(), enc_h.data_ptr(), T, D, 2*D, 0)
fvk.add_bias_fp16(enc_h.data_ptr(), dit.ae_b2.data_ptr(), T, D, 0)
diff --git a/flash_rt/hardware/thor/fa4_backend.py b/flash_rt/hardware/thor/fa4_backend.py
index 9b292913..d7dd0fde 100644
--- a/flash_rt/hardware/thor/fa4_backend.py
+++ b/flash_rt/hardware/thor/fa4_backend.py
@@ -1,9 +1,8 @@
"""Isolated FlashAttention-4 (CuTe-DSL) backend for Thor (sm_110).
-FA4's SM100 Blackwell forward kernel runs on Thor when compiled for ``sm_110a``.
-The vendored tree contains the SM100-compatible kernel only, so the loader
-keeps ``FLASH_ATTENTION_ARCH=sm_100a`` as the runtime dispatch key while
-``CUTE_DSL_ARCH=sm_110a`` selects the Thor compilation target.
+FA4's SM100-compatible forward kernel runs on Thor through the ``sm_101a``
+CuTe-DSL compilation alias. Runtime dispatch still uses the physical device
+capability (SM110); the compilation alias must never select SM100-only kernels.
At the LingBot denoise shape (Sq=51, Skv~891, GQA 16/2, HD=128) with
``pack_gqa`` it is ~17% faster than the vendored fmha kernel, cos=1.0, and is
CUDA-graph capture-safe.
@@ -59,7 +58,6 @@ def _load() -> None:
# sm_110a). Requires nvidia-cutlass-dsl >= 4.5 (which accepts the
# sm_101a alias); 4.4.x only knows sm_110a and is NOT supported.
os.environ.setdefault("CUTE_DSL_ARCH", "sm_101a")
- os.environ.setdefault("FLASH_ATTENTION_ARCH", "sm_100a")
# Allow an override dir (e.g. for local FA4 development), else the vendor.
src = os.environ.get("LINGBOT_FA4_SRC") or (
diff --git a/tests/test_groot_n16_pr177_maintenance.py b/tests/test_groot_n16_pr177_maintenance.py
new file mode 100644
index 00000000..aecabe3a
--- /dev/null
+++ b/tests/test_groot_n16_pr177_maintenance.py
@@ -0,0 +1,79 @@
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _read(path: str) -> str:
+ return (ROOT / path).read_text()
+
+
+def test_thor_qk_norm_source_and_binding_share_one_gate():
+ cmake = _read("CMakeLists.txt")
+ bindings = _read("csrc/bindings.cpp")
+
+ gate = cmake.index("if(ENABLE_SM100_CUTLASS)", cmake.index("Thor-class VLA"))
+ source = cmake.index("csrc/kernels/qk_norm_rope_rotate_half_bf16.cu", gate)
+ gate_end = cmake.index("endif()", gate)
+ assert gate < source < gate_end
+ assert '#ifdef FLASHRT_HAVE_THOR_VLA_KERNELS\n#include "kernels/qk_norm_rope_rotate_half_bf16.cuh"' in bindings
+ assert '#ifdef FLASHRT_HAVE_THOR_VLA_KERNELS\n m.def("qk_norm_rope_rotate_half_bf16"' in bindings
+
+
+def test_fa4_compilation_alias_does_not_override_runtime_arch():
+ backend = _read("flash_rt/hardware/thor/fa4_backend.py")
+ interface = _read(
+ "csrc/attention/flash_attn_4_src/flashrt_fa4/cute/"
+ "interface_fwd_sm100.py"
+ )
+
+ assert 'setdefault("CUTE_DSL_ARCH", "sm_101a")' in backend
+ assert 'setdefault("FLASH_ATTENTION_ARCH"' not in backend
+ assert "torch.cuda.get_device_capability()" in interface
+ assert "use_dedicated_hd256_kernel = arch // 10 == 10" in interface
+
+
+def test_n16_new_precision_tiers_are_opt_in():
+ source = _read("flash_rt/frontends/torch/groot_thor.py")
+
+ assert 'embodiment_tag="new_embodiment", use_fp8=True,' in source
+ assert "image_size=224, parity=False" in source
+ for name in (
+ "FLASHRT_N16_FA4",
+ "FLASHRT_N16_QWEN3_FP4",
+ "FLASHRT_N16_SIGLIP_FP4",
+ "FLASHRT_N16_DIT_FP4",
+ ):
+ assert f'os.environ.get("{name}", "0")' in source
+
+
+def test_prompt_switch_and_calibration_follow_runtime_layout_contracts():
+ source = _read("flash_rt/frontends/torch/groot_thor.py")
+ prompt_start = source.index(" def set_prompt(")
+ prompt = source[prompt_start:source.index(" def infer_action_head(", prompt_start)]
+
+ assert "self.reset_graph_runtime()" in prompt
+ assert prompt.index("self.reset_graph_runtime()") < prompt.index(
+ "self._input_ids = torch.tensor"
+ )
+ reset_start = source.index(" def reset_graph_runtime(")
+ reset = source[reset_start:source.index(" def set_prompt(", reset_start)]
+ assert '"_qwen3_fp4_done", "_siglip_fp4_done"' in reset
+
+ collect_start = source.index(" def _collect_dit_amax(")
+ collect = source[collect_start:source.index(" def _calibrate_dit(", collect_start)]
+ assert "fvk.concat2_bf16(a_emb_out.data_ptr()," in collect
+ assert "ae_concat.data_ptr()+T*D*2" not in collect
+
+
+def test_eagle_remote_code_requires_checkpoint_local_or_explicit_revision():
+ source = _read("flash_rt/frontends/torch/groot_thor.py")
+ resolver_start = source.index(" def _resolve_eagle_dir(")
+ resolver = source[resolver_start:source.index(
+ " def _setup_torch_siglip(", resolver_start
+ )]
+
+ assert "FLASHRT_N16_EAGLE_DIR" in resolver
+ assert "self._checkpoint_path" in resolver
+ assert "transformers_modules" not in resolver
+ assert ".glob(" not in resolver
diff --git a/tests/test_prompt_runtime_lifecycle.py b/tests/test_prompt_runtime_lifecycle.py
index f32a0841..239d2ede 100644
--- a/tests/test_prompt_runtime_lifecycle.py
+++ b/tests/test_prompt_runtime_lifecycle.py
@@ -219,15 +219,14 @@ def infer(self, obs):
assert BucketFrontend.calibrations == 2
-def test_groot_thor_rejects_prompt_changes_after_graph_build():
+def test_groot_thor_resets_graphs_before_accepting_prompt_changes():
source = Path("flash_rt/frontends/torch/groot_thor.py").read_text()
- set_prompt_pos = source.index(" def set_prompt(self, prompt):")
- guard_pos = source.index("getattr(self, '_graphs_built', False)",
- set_prompt_pos)
- tokenizer_pos = source.index("from transformers import AutoTokenizer",
- set_prompt_pos)
+ set_prompt_pos = source.index(" def set_prompt(self, prompt, input_ids=None):")
+ reset_pos = source.index("self.reset_graph_runtime()", set_prompt_pos)
+ assign_pos = source.index("self._input_ids = torch.tensor", set_prompt_pos)
- assert guard_pos < tokenizer_pos
+ assert reset_pos < assign_pos
+ assert "set_prompt() after the pipeline is built is not supported" not in source
def test_groot_thor_infer_refreshes_state_before_dit_graph_replay():
From 4c830e898482cd48960a4401e732216aeb0edcf2 Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Sun, 23 Aug 2026 13:06:43 -0400
Subject: [PATCH 6/6] fix(groot): keep parity routing coherent
---
flash_rt/frontends/torch/groot_thor.py | 17 +++++++++++++----
tests/test_groot_n16_pr177_maintenance.py | 8 ++++++++
2 files changed, 21 insertions(+), 4 deletions(-)
diff --git a/flash_rt/frontends/torch/groot_thor.py b/flash_rt/frontends/torch/groot_thor.py
index b580401e..1d2f4937 100644
--- a/flash_rt/frontends/torch/groot_thor.py
+++ b/flash_rt/frontends/torch/groot_thor.py
@@ -106,6 +106,14 @@ def __init__(self, checkpoint, num_views=2, autotune=3,
self._num_views = num_views
self._autotune = autotune
self.use_fp8 = bool(use_fp8)
+ if parity and self.use_fp8:
+ raise ValueError(
+ "parity=True requires use_fp8=False; mixing the FP8 backbone "
+ "with the HF-native DiT is not a supported precision profile")
+ if parity and image_size != 252:
+ raise ValueError(
+ "parity=True requires image_size=252 to match the GR00T "
+ "N1.6 evaluation preprocessing contract")
# Working dtype of the vision/LLM feature path: FP8 mode keeps the
# legacy fp16 buffers; parity mode stays bf16 end-to-end to match HF.
self._bd = (torch.float16 if (self.use_fp8 or not parity)
@@ -3230,12 +3238,13 @@ def _capture_all_graphs(self, obs, release_full_sd: bool = True):
# Extract bf16 action-head weights for the HF-faithful torch DiT
# path (must happen while _full_sd is alive).
- self._setup_torch_dit()
- if not self.use_fp8 and not hasattr(self, "_torch_qwen3"):
+ if self.parity:
+ self._setup_torch_dit()
+ if self.parity and not hasattr(self, "_torch_qwen3"):
self._setup_torch_qwen3()
- if not self.use_fp8:
+ if self.parity:
self._setup_qwen3_fp4()
- if not self.use_fp8 and not hasattr(self, "_torch_siglip"):
+ if self.parity and not hasattr(self, "_torch_siglip"):
self._setup_torch_siglip()
# Static NaFlex window meta (shapes fixed per image_size):
# needed by the capture-time eager encoder run before 6c.
diff --git a/tests/test_groot_n16_pr177_maintenance.py b/tests/test_groot_n16_pr177_maintenance.py
index aecabe3a..fcc9fbf8 100644
--- a/tests/test_groot_n16_pr177_maintenance.py
+++ b/tests/test_groot_n16_pr177_maintenance.py
@@ -38,6 +38,8 @@ def test_n16_new_precision_tiers_are_opt_in():
assert 'embodiment_tag="new_embodiment", use_fp8=True,' in source
assert "image_size=224, parity=False" in source
+ assert "parity=True requires use_fp8=False" in source
+ assert "parity=True requires image_size=252" in source
for name in (
"FLASHRT_N16_FA4",
"FLASHRT_N16_QWEN3_FP4",
@@ -46,6 +48,12 @@ def test_n16_new_precision_tiers_are_opt_in():
):
assert f'os.environ.get("{name}", "0")' in source
+ capture_start = source.index(" def _capture_all_graphs(")
+ capture = source[capture_start:]
+ assert "if self.parity:\n self._setup_torch_dit()" in capture
+ assert 'if self.parity and not hasattr(self, "_torch_qwen3")' in capture
+ assert 'if self.parity and not hasattr(self, "_torch_siglip")' in capture
+
def test_prompt_switch_and_calibration_follow_runtime_layout_contracts():
source = _read("flash_rt/frontends/torch/groot_thor.py")