feat(cosmos3-edge): add Thor AV denoise engine and TeaCache

CMakeLists.txt

@@ -1289,6 +1289,7 @@ pybind11_add_module(flash_rt_kernels
   csrc/kernels/attention_mha.cu
   csrc/kernels/attention_mha_causal.cu
   csrc/kernels/rope_qwen3.cu
+  csrc/kernels/cosmos3_edge_misc.cu
   csrc/kernels/decoder_fused.cu
   csrc/kernels/dit_bf16.cu
   csrc/kernels/attention_dit_bf16.cu
@@ -1562,6 +1563,17 @@ if(GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121" OR
     ENABLE_QWEN36_FLASHINFER_XQA=1)
 endif()
 
+if(GPU_ARCH STREQUAL "110")
+  target_sources(flash_rt_kernels PRIVATE
+    csrc/conv/fp8_conv3d_sm120_v17.cu
+    csrc/conv/fp8_conv3d_sm120_v18.cu
+    csrc/conv/bf16_conv3d_v0_sm110.cu)
+  target_compile_definitions(flash_rt_kernels PRIVATE
+    ENABLE_FP8_CONV3D_V17=1
+    ENABLE_FP8_CONV3D_V18=1
+    ENABLE_BF16_CONV3D_V0=1)
+endif()
+
 if(FLASHRT_ENABLE_MOTUS AND GPU_ARCH STREQUAL "120")
   target_sources(flash_rt_kernels PRIVATE
     $<TARGET_OBJECTS:motus_aux_sm120_obj>
@@ -1630,6 +1642,7 @@ if(ENABLE_SM100_CUTLASS)
     csrc/fused_fp4/norm_silu_fp4_sfa.cu
     csrc/fused_fp4/silu_mul_fp4_sfa_v2.cu
     csrc/fused_fp4/res_rms_fp4_sfa_v2.cu
+    csrc/fused_fp4/cosmos3_edge_fp4.cu
     csrc/fused_fp4/per_channel_mul_fp16.cu
     csrc/fused_fp4/res_rms_mul_fp4_sfa.cu
     csrc/fused_fp4/silu_mul_mul_fp4_sfa_v2.cu

README.md

@@ -158,6 +158,20 @@ DGX Spark / GB10:
 | RTX 5090 | 720p, 121f, 20 steps | **178.6 s** | [Wan2.2 benchmarks](docs/wan22_usage.md#benchmarks) |
 | RTX 5090 | TeaCache 0.3 | **114.2 s** | [Wan2.2 benchmarks](docs/wan22_usage.md#benchmarks) |
 
+#### Cosmos3-Edge AV inverse dynamics
+
+Jetson AGX Thor, 30-step denoise boundary, 15 hot back-to-back iterations:
+
+| Mode | Denoise P50 | Speedup | Final-action check | Source |
+|---|---:|---:|---:|---|
+| Official eager | 33.14 s | 1.00x | reference | [Cosmos3-Edge usage](docs/cosmos3_edge_thor.md#performance) |
+| FlashRT FP8, no step cache | 5.792 s | **5.72x** | cos 0.999983 | [Cosmos3-Edge usage](docs/cosmos3_edge_thor.md#performance) |
+| FlashRT FP8 + NVFP4 FFN, no step cache | 5.456 s | **6.07x** | cos 0.999707 | [Cosmos3-Edge usage](docs/cosmos3_edge_thor.md#performance) |
+| FlashRT FP8 + TeaCache, 3 computes | 0.576 s | **57.56x** | cos 0.999986 | [Cosmos3-Edge usage](docs/cosmos3_edge_thor.md#performance) |
+
+TeaCache results are workload-specific and require checkpoint- and task-level
+validation. They are not a general policy-quality guarantee.
+
 ## Getting Started
 
 - [Install FlashRT](#build--install)
@@ -205,6 +219,7 @@ First call: ~3 s (calibration + CUDA Graph capture). Every subsequent call: 44 m
 | **Run Higgs Audio v3 TTS** | [`docs/higgs_audio_v3.md`](docs/higgs_audio_v3.md) — usage + performance · [`serving/higgs_audio_agent/`](serving/higgs_audio_agent/README.md) — HTTP serving |
 | **Run Motus RTX beta, TeaCache, or legacy async chunk runner** | [`docs/motus_usage_beta.md`](docs/motus_usage_beta.md) · [`docs/rtc_lite_design.md`](docs/rtc_lite_design.md) |
 | **Run Wan2.2 TI2V-5B official-pipeline baseline** | [`docs/wan22_usage.md`](docs/wan22_usage.md) |
+| **Run Cosmos3-Edge AV inverse dynamics on Jetson AGX Thor** | [`docs/cosmos3_edge_thor.md`](docs/cosmos3_edge_thor.md) — official baseline, optimized denoise, TeaCache qualification, and benchmarks |
 | **Use FlashRT kernels through Hugging Face Kernel Hub** | [`LiangSu8899/FlashRT-HF-kernels`](https://github.com/LiangSu8899/FlashRT-HF-kernels) · [`huggingface.co/flashrt`](https://huggingface.co/flashrt) |
 | **Run serving hosts** | [`serving/README.md`](serving/README.md) — scenario hosts · [`docs/serving_design.md`](docs/serving_design.md) — capsules and roadmap · [`docs/serving_production.md`](docs/serving_production.md) — production notes |
 | **Look up the stable Python API surface** | [`docs/stable_api.md`](docs/stable_api.md) |
@@ -972,6 +987,7 @@ examples/
 - **LingBot-VLA** — [LingBot usage](docs/lingbot_usage.md), [Thor latency](docs/lingbot_usage.md#5-accuracy--latency-thor-sm_110-cuda-graph-replay)
 - **Motus Stage3 RTX beta** (`config="motus"`) — [Motus usage](docs/motus_usage_beta.md), [legacy async chunk runner](docs/rtc_lite_design.md)
 - **Wan2.2 TI2V-5B** (`config="wan22_ti2v_5b"`) — [Wan2.2 usage](docs/wan22_usage.md)
+- **Cosmos3-Edge AV inverse dynamics** (`config="cosmos3_edge"`) — Jetson AGX Thor official baseline and 6x-class FlashRT denoise; [usage and performance](docs/cosmos3_edge_thor.md)
 - **Higgs Audio v3 TTS-4B** — [Higgs usage](docs/higgs_audio_v3.md#3-quickstart), [Higgs serving](serving/higgs_audio_agent/README.md)
 - **Qwen3.6-27B NVFP4** — [Qwen3.6 usage](docs/qwen36_nvfp4.md), [parameter reference](docs/qwen36_usage.md), [serving](serving/qwen36_agent/README.md)
 - **Qwen3-8B NVFP4** — [Qwen3-8B usage](docs/qwen3_8b_nvfp4.md), [OpenAI server example](examples/qwen3_openai_server.py)
@@ -988,7 +1004,7 @@ artifacts and dispatch map are.
 
 | Hardware | SM | Status | Validated paths / notes |
 |---|---:|---|---|
-| Jetson AGX Thor | SM110 | Production target | Pi0, Pi0.5, GROOT N1.6, Pi0-FAST, Qwen3.6 Thor path, Lingbot; CUTLASS FMHA / Thor attention paths; Pi0.5 FP8 and NVFP4 validation live in [examples/thor](examples/thor/README.md#thor-vla-performance). |
+| Jetson AGX Thor | SM110 | Production target | Pi0, Pi0.5, GROOT N1.6, Pi0-FAST, Qwen3.6 Thor path, Lingbot, and Cosmos3-Edge AV denoise; CUTLASS FMHA / Thor attention paths; Pi0.5 FP8 and NVFP4 validation live in [examples/thor](examples/thor/README.md#thor-vla-performance). |
 | RTX 5090 | SM120 | Production target | Pi0/Pi0.5/GROOT/Pi0-FAST RTX paths, Qwen3.6, Qwen3-8B, Higgs Audio v3 FP8, Motus, Wan2.2, HF Kernel Hub package validation; see [RTX 5090 latency](examples/blackwell/README.md#vla-latency-rtx-5090). |
 | RTX 4090 | SM89 | Validated / supported target | RTX VLA build path and deployment recipe; Higgs BF16 path compiles/configures. See [deployment_rtx4090.md](docs/deployment_rtx4090.md). |
 | RTX 5060 Ti | SM120 | Community validated | Pi0.5 FP8 and LIBERO Spatial submission; see [Community benchmarks](#community-benchmarks). |

USAGE.md

@@ -213,7 +213,7 @@ model = flash_rt.load_model(
 | `autotune` | `int\|bool` | `3` | CUDA Graph autotune intensity. See [Autotune](#autotune). |
 | `recalibrate` | `bool` | `False` | Force fresh FP8 calibration (and weight cache for JAX), ignoring cache. See [Calibration](#calibration). |
 | `weight_cache` | `bool` | `True` | Cache FP8-quantized weights to disk. **JAX only** — reduces cold start from ~42s to ~6s. Torch loads in ~3s and ignores this. See [Weight Cache](#weight-cache-jax-only). |
-| `config` | `str` | `"pi05"` | Model architecture config: `"pi05"`, `"pi0"`, `"groot"`, `"groot_n17"`, `"pi0fast"`, `"motus"`, `"wan22_ti2v_5b"`, `"cosmos3_video"`. `"cosmos3_video"` is a non-VLA text2video denoise model — drive it with `set_prompt(ref=...)` + `infer(...)`, not `predict()`. |
+| `config` | `str` | `"pi05"` | Model architecture config: `"pi05"`, `"pi0"`, `"groot"`, `"groot_n17"`, `"pi0fast"`, `"motus"`, `"wan22_ti2v_5b"`, `"cosmos3_video"`, `"cosmos3_edge"`. Cosmos video/Edge routes use `set_prompt(...)` + `infer(...)`, not the VLA `predict()` convenience API. |
 | `decode_cuda_graph` | `bool` | `False` | **Pi0-FAST only.** Capture action-phase decode as CUDA Graph. Trades startup time for per-token speed. See [Pi0-FAST](#pi0-fast). |
 | `decode_graph_steps` | `int` | `80` | **Pi0-FAST only.** Number of action tokens to capture in the decode graph. Should cover your longest expected action sequence. |
 | `use_fp4` | `bool` | `False` | **Pi0.5 torch + jax on Thor.** Enable NVFP4 quantization on the encoder FFN stack. When `True`, resolves to the production preset (`fp4_layers=tuple(range(18))` + `use_awq=True` + `use_p1_split_gu=True`). Requires SM100+ GPU. Other configs emit a warning and fall back to FP8. See [NVFP4](#nvfp4-pi05-only). |
@@ -677,6 +677,37 @@ The model-local kernel is built once on the target GPU (`cd
 flash_rt/models/cosmos3_video/kernels && python3 setup.py build_ext --inplace`).
 See [`docs/cosmos3_video_usage.md`](docs/cosmos3_video_usage.md).
 
+### Cosmos3-Edge AV inverse dynamics
+
+Cosmos3-Edge is registered for PyTorch on Jetson AGX Thor. Run the official
+action-only path first to establish the end-to-end output and baseline:
+
+```python
+import flash_rt
+
+model = flash_rt.load_model(
+    "/path/to/Cosmos3-Edge",
+    framework="torch",
+    config="cosmos3_edge",
+    hardware="thor",
+)
+model.set_prompt(input_json="/path/to/av_inverse_input.json")
+result = model.infer(
+    output_dir="/path/to/output",
+    backend="official_action_only",
+    cosmos_root="/path/to/cosmos-framework",
+    vae_path="/path/to/Wan2.2_VAE.pth",
+    benchmark=True,
+)
+```
+
+The FlashRT denoise engine executes all 30 steps in one CUDA Graph and reaches
+5.72x with FP8 or 6.07x with FP8 plus NVFP4 FFN on the measured Thor workload.
+Optional TeaCache schedules can be much faster, but must be qualified for each
+task distribution and fine-tuned checkpoint. See
+[`docs/cosmos3_edge_thor.md`](docs/cosmos3_edge_thor.md) for the complete table,
+benchmark commands, accuracy gates, and limitations.
+
 ### `model.predict()`
 
 ```python