From dfa63529a960e64ad834988a23718a7bf77071ed Mon Sep 17 00:00:00 2001 From: LiangSu8899 Date: Tue, 21 Jul 2026 07:38:44 -0400 Subject: [PATCH 1/3] feat(cosmos3-edge): Thor AV inverse-dynamics denoise with FP8/FP4 whole-graph engine Add config="cosmos3_edge" for Cosmos3-Edge (4B, two-tower MoT, 28L/2048, GQA 16/8, non-gated relu2 FFN) AV inverse dynamics on Thor SM110: - Official Cosmos Framework baseline runner plus action-only, live-handoff, boundary/prepare dump-replay, and VAE probe paths behind one frontend. - Torch/FVK bring-up references, boundary dumps, static und-K/V engine, and a native one-launch UniPC step for the fixed 30-step AV schedule. - models/cosmos3_edge/pipeline_thor.py: the optimized denoise engine. All six gen-tower projections in per-tensor FP8 E4M3 with calibrated static activation scales; fused quant chain (residual_add_rms_norm_fp8, quantize_fp8_static, relu2->fp8); warp-per-head fused qk-norm+RoPE; FA4 attention over per-layer joint und+gen K/V; M=60 slim last layer; the whole 30-step denoise replayed as a single CUDA graph. - Optional NVFP4 W4A4 FFN via flash_rt_fp4 with fused bf16 res+rms->fp4 and fp16 relu2->fp4 quantizers. - benchmarks/cosmos3_edge_thor_denoise.py with latency and final-action gates. Thor, 15 back-to-back iterations, official eager denoise 33.14s, gates cos>=0.999 / rel_l2<3% (both engines also pass on a second-seed dump): --engine bf16-eager 11.998s (2.76x) cos 0.9999959 rel_l2 0.29% --engine fp8 5.792s (5.72x) cos 0.9999834 rel_l2 0.59% --engine fp8 --ffn-fp4 5.456s (6.07x) cos 0.9997067 rel_l2 2.42% --- CMakeLists.txt | 13 + benchmarks/cosmos3_edge_thor_denoise.py | 359 ++ csrc/bindings.cpp | 243 ++ csrc/conv/bf16_conv3d_v0_sm110.cu | 291 ++ csrc/fp4_bindings.cpp | 34 + csrc/fused_fp4/cosmos3_edge_fp4.cu | 212 + csrc/fused_fp4/cosmos3_edge_fp4.cuh | 22 + csrc/kernels/activation.cu | 27 + csrc/kernels/activation.cuh | 1 + csrc/kernels/cosmos3_edge_misc.cu | 695 ++++ csrc/kernels/cosmos3_edge_misc.cuh | 143 + docs/cosmos3_edge_thor.md | 1370 +++++++ docs/stable_api.md | 7 +- examples/cosmos3_edge_thor_baseline.py | 281 ++ flash_rt/api.py | 15 +- flash_rt/configs/cosmos3_edge.yaml | 21 + .../torch/_cosmos3_edge_thor_spec.py | 177 + flash_rt/frontends/torch/cosmos3_edge_thor.py | 571 +++ flash_rt/hardware/__init__.py | 4 + flash_rt/models/cosmos3_edge/__init__.py | 37 + .../cosmos3_edge/action_only_official.py | 3447 +++++++++++++++++ flash_rt/models/cosmos3_edge/boundary_dump.py | 151 + flash_rt/models/cosmos3_edge/denoise_ref.py | 322 ++ flash_rt/models/cosmos3_edge/dump_replay.py | 195 + flash_rt/models/cosmos3_edge/layer_ref.py | 1035 +++++ flash_rt/models/cosmos3_edge/pipeline_thor.py | 505 +++ flash_rt/models/cosmos3_edge/static_engine.py | 346 ++ flash_rt/models/cosmos3_edge/static_unipc.py | 226 ++ flash_rt/models/cosmos3_edge/vae_native.py | 304 ++ flash_rt/models/cosmos3_edge/weights.py | 91 + tests/test_cosmos3_edge_boundary_dump.py | 220 ++ tests/test_cosmos3_edge_dump_replay.py | 172 + tests/test_cosmos3_edge_kernel_bindings.py | 596 +++ tests/test_cosmos3_edge_layer_ref.py | 122 + tests/test_cosmos3_edge_step0_reference.py | 175 + tests/test_cosmos3_edge_thor_spec.py | 32 + tests/test_cosmos3_edge_weights.py | 49 + tests/test_load_model_use_fp8_kwarg.py | 1136 ++++++ 38 files changed, 13643 insertions(+), 4 deletions(-) create mode 100644 benchmarks/cosmos3_edge_thor_denoise.py create mode 100644 csrc/conv/bf16_conv3d_v0_sm110.cu create mode 100644 csrc/fused_fp4/cosmos3_edge_fp4.cu create mode 100644 csrc/fused_fp4/cosmos3_edge_fp4.cuh create mode 100644 csrc/kernels/cosmos3_edge_misc.cu create mode 100644 csrc/kernels/cosmos3_edge_misc.cuh create mode 100644 docs/cosmos3_edge_thor.md create mode 100644 examples/cosmos3_edge_thor_baseline.py create mode 100644 flash_rt/configs/cosmos3_edge.yaml create mode 100644 flash_rt/frontends/torch/_cosmos3_edge_thor_spec.py create mode 100644 flash_rt/frontends/torch/cosmos3_edge_thor.py create mode 100644 flash_rt/models/cosmos3_edge/__init__.py create mode 100644 flash_rt/models/cosmos3_edge/action_only_official.py create mode 100644 flash_rt/models/cosmos3_edge/boundary_dump.py create mode 100644 flash_rt/models/cosmos3_edge/denoise_ref.py create mode 100644 flash_rt/models/cosmos3_edge/dump_replay.py create mode 100644 flash_rt/models/cosmos3_edge/layer_ref.py create mode 100644 flash_rt/models/cosmos3_edge/pipeline_thor.py create mode 100644 flash_rt/models/cosmos3_edge/static_engine.py create mode 100644 flash_rt/models/cosmos3_edge/static_unipc.py create mode 100644 flash_rt/models/cosmos3_edge/vae_native.py create mode 100644 flash_rt/models/cosmos3_edge/weights.py create mode 100644 tests/test_cosmos3_edge_boundary_dump.py create mode 100644 tests/test_cosmos3_edge_dump_replay.py create mode 100644 tests/test_cosmos3_edge_kernel_bindings.py create mode 100644 tests/test_cosmos3_edge_layer_ref.py create mode 100644 tests/test_cosmos3_edge_step0_reference.py create mode 100644 tests/test_cosmos3_edge_thor_spec.py create mode 100644 tests/test_cosmos3_edge_weights.py diff --git a/CMakeLists.txt b/CMakeLists.txt index f4474fb4..2b9877f7 100644 --- a/CMakeLists.txt +++ b/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 $ @@ -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 diff --git a/benchmarks/cosmos3_edge_thor_denoise.py b/benchmarks/cosmos3_edge_thor_denoise.py new file mode 100644 index 00000000..cc46b05c --- /dev/null +++ b/benchmarks/cosmos3_edge_thor_denoise.py @@ -0,0 +1,359 @@ +#!/usr/bin/env python +"""Cosmos3-Edge Thor denoise latency + accuracy benchmark. + +This benchmark measures the current FlashRT denoise-only path against the P0 +official dump. It intentionally does not run Cosmos Framework; pass the +official ``benchmark.json`` only to print the recorded eager baseline next to +the FlashRT numbers. + +Run inside the Thor container and isolated venv: + + cd /work/official/flashrt-public + source /work/.venv_cosmos_thor/bin/activate + python benchmarks/cosmos3_edge_thor_denoise.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --reference-dump dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors \ + --boundary-dump dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors \ + --official-benchmark dev_scratch_cosmos3_thor/edge_av_inverse_0/official_outputs/benchmark.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +import warnings +from pathlib import Path +from typing import Any + + +def _percentile(sorted_values: list[float], q: float) -> float: + if not sorted_values: + raise ValueError("cannot compute percentile over empty list") + if len(sorted_values) == 1: + return sorted_values[0] + pos = (len(sorted_values) - 1) * q + lo = int(pos) + hi = min(lo + 1, len(sorted_values) - 1) + frac = pos - lo + return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac + + +def _load_official_benchmark(path: Path | None) -> dict[str, Any] | None: + if path is None: + return None + if not path.exists(): + raise FileNotFoundError(f"official benchmark does not exist: {path}") + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def _official_denoise_s(bench: dict[str, Any] | None) -> float | None: + return _official_average(bench, "OmniMoTModel.generate_samples_from_batch") + + +def _official_average(bench: dict[str, Any] | None, key: str) -> float | None: + if not bench: + return None + try: + return float(bench["average"][key]) + except (KeyError, TypeError, ValueError): + return None + + +def _summarize(values: list[float]) -> dict[str, float]: + ordered = sorted(values) + return { + "count": float(len(values)), + "total_s": float(sum(values)), + "p50_s": float(statistics.median(ordered)), + "p90_s": float(_percentile(ordered, 0.90)), + "min_s": float(min(ordered)), + "max_s": float(max(ordered)), + } + + +def _profile_loop_breakdown(runner: Any, *, device: str) -> dict[str, Any]: + import torch + + if device != "cuda": + raise RuntimeError("--profile-loop-breakdown requires CUDA") + + latent = runner.denoise_dump.step_noise(0).to(device=runner.device, dtype=torch.float32) + engine = getattr(runner, "engine", None) + if engine is not None: + engine.latent.copy_(latent) + latent = engine.latent + if runner.static_scheduler is not None and runner.static_scheduler.prev_m1 is None: + # Never re-allocate scheduler state that a captured graph already binds. + runner.static_scheduler.reset(latent) + timesteps = torch.tensor(runner.denoise_dump.timesteps(), device=runner.device, dtype=torch.int64) + + events: list[tuple[Any, Any, Any]] = [] + for step, timestep in enumerate(timesteps[: runner.denoise_dump.num_steps]): + start = torch.cuda.Event(enable_timing=True) + mid = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + velocity = runner.velocity_for_step(latent, timestep.reshape(1, 1), step_index=step) + mid.record() + if runner.static_scheduler is None: + raise RuntimeError("--profile-loop-breakdown requires the native static scheduler") + latent = runner.static_scheduler.step(latent, velocity, step) + end.record() + events.append((start, mid, end)) + + torch.cuda.synchronize() + velocity_s = [start.elapsed_time(mid) / 1000.0 for start, mid, _end in events] + scheduler_s = [mid.elapsed_time(end) / 1000.0 for _start, mid, end in events] + step_s = [start.elapsed_time(end) / 1000.0 for start, _mid, end in events] + return { + "note": "CUDA event timings inside one profiled run; per-step sync is not inserted.", + "steps": len(events), + "velocity": _summarize(velocity_s), + "scheduler": _summarize(scheduler_s), + "step_total": _summarize(step_s), + "native_scheduler": bool(getattr(runner, "native_scheduler_available", False)), + "use_cuda_graphs": bool(getattr(runner, "use_cuda_graphs", False)), + } + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", default="/work/models/Cosmos3-Edge") + ap.add_argument( + "--reference-dump", + default="dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors", + ) + ap.add_argument( + "--boundary-dump", + default="dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors", + ) + ap.add_argument( + "--official-benchmark", + default="dev_scratch_cosmos3_thor/edge_av_inverse_0/official_outputs/benchmark.json", + ) + ap.add_argument("--device", default="cuda") + ap.add_argument("--warmup-steps", type=int, default=1) + ap.add_argument("--iters", type=int, default=3) + ap.add_argument("--use-cuda-graphs", action="store_true") + ap.add_argument( + "--engine", + choices=("bf16-eager", "fp8", "quant-bf16"), + default="bf16-eager", + help=( + "bf16-eager: original per-op FlashRT engine. " + "fp8: quantized static engine (whole-denoise CUDA graph unless --no-quant-graph). " + "quant-bf16: static engine with all projections in bf16 (quant-path parity check)." + ), + ) + ap.add_argument("--no-quant-graph", action="store_true") + ap.add_argument("--bf16-projs", default="", help="comma-separated projections kept in bf16 (fp8 engine)") + ap.add_argument("--ffn-fp4", action="store_true", help="NVFP4 W4A4 FFN (up/down) via flash_rt_fp4") + ap.add_argument("--no-slim-last", action="store_true", help="disable the M=60 slim last layer") + ap.add_argument("--json-out", default=None) + ap.add_argument( + "--profile-loop-breakdown", + action="store_true", + help="Add CUDA event timings for velocity vs scheduler inside one 30-step run.", + ) + ap.add_argument("--enforce-gates", action="store_true") + ap.add_argument("--min-speedup", type=float, default=2.5) + ap.add_argument("--min-final-action-cos", type=float, default=0.999) + ap.add_argument("--max-final-action-rel-l2", type=float, default=0.03) + ap.add_argument( + "--skip-speedup-gate", + action="store_true", + help="Do not fail --enforce-gates on speedup. Useful for additional-seed accuracy checks.", + ) + args = ap.parse_args() + + if args.iters < 1: + raise ValueError("--iters must be >= 1") + if args.warmup_steps < 0: + raise ValueError("--warmup-steps must be >= 0") + + warnings.filterwarnings("ignore", category=DeprecationWarning) + + import torch + + from flash_rt.hardware.thor import fa4_backend + from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump + from flash_rt.models.cosmos3_edge.denoise_ref import EdgeDenoiseFlashRT, EdgeDenoiseFlashRTQuant + from flash_rt.models.cosmos3_edge.dump_replay import EdgeDenoiseDump + from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights + + if args.device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA device is not available") + + checkpoint = Path(args.checkpoint) + reference_dump = Path(args.reference_dump) + boundary_dump = Path(args.boundary_dump) + official_benchmark_path = Path(args.official_benchmark) if args.official_benchmark else None + for label, path in ( + ("checkpoint", checkpoint), + ("reference dump", reference_dump), + ("boundary dump", boundary_dump), + ): + if not path.exists(): + raise FileNotFoundError(f"{label} does not exist: {path}") + + official_benchmark = _load_official_benchmark(official_benchmark_path) + official_denoise = _official_denoise_s(official_benchmark) + official_e2e = _official_average(official_benchmark, "OmniInference.generate_batch") + official_decode = _official_average(official_benchmark, "OmniMoTModel.decode") + + if args.device == "cuda": + torch.cuda.reset_peak_memory_stats() + t0 = time.perf_counter() + if args.engine == "bf16-eager": + runner = EdgeDenoiseFlashRT( + EdgeDenoiseDump(reference_dump), + EdgeBoundaryDump(boundary_dump), + EdgeTransformerWeights(checkpoint), + device=args.device, + use_cuda_graphs=args.use_cuda_graphs, + ) + else: + runner = EdgeDenoiseFlashRTQuant( + EdgeDenoiseDump(reference_dump), + EdgeBoundaryDump(boundary_dump), + EdgeTransformerWeights(checkpoint), + device=args.device, + quant="bf16" if args.engine == "quant-bf16" else "fp8", + bf16_projs=tuple(p for p in args.bf16_projs.split(",") if p), + ffn_fp4=args.ffn_fp4, + slim_last=not args.no_slim_last, + use_cuda_graphs=not args.no_quant_graph, + ) + if args.device == "cuda": + torch.cuda.synchronize() + init_s = time.perf_counter() - t0 + + warmup_result = None + if args.warmup_steps: + with torch.no_grad(): + warmup_result = runner.run(max_steps=args.warmup_steps) + if args.device == "cuda": + torch.cuda.synchronize() + + times: list[float] = [] + last = None + for _ in range(args.iters): + if args.device == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + with torch.no_grad(): + last = runner.run(max_steps=runner.denoise_dump.num_steps) + if args.device == "cuda": + torch.cuda.synchronize() + times.append(time.perf_counter() - t0) + + assert last is not None + sorted_times = sorted(times) + official_action = runner.denoise_dump.final_action.float() + diff = last.final_action.float() - official_action + rel_l2 = float((torch.linalg.vector_norm(diff) / torch.linalg.vector_norm(official_action)).item()) + cos = float( + torch.nn.functional.cosine_similarity( + last.final_action.flatten(), + official_action.flatten(), + dim=0, + ).item() + ) + max_abs = float(diff.abs().max().item()) + p50_s = statistics.median(sorted_times) + official_non_denoise_non_decode = ( + official_e2e - official_denoise - official_decode + if official_e2e is not None and official_denoise is not None and official_decode is not None + else None + ) + estimated_action_only_e2e = ( + official_non_denoise_non_decode + p50_s + if official_non_denoise_non_decode is not None + else None + ) + native_attention = bool(getattr(runner, "native_attention_available", False)) + graph_attention = bool(getattr(runner, "graph_attention_available", False)) + native_scheduler = bool(getattr(runner, "native_scheduler_available", False)) + speedup = official_denoise / p50_s if official_denoise is not None else None + gate_checks = { + "native_attention": native_attention, + "native_scheduler": native_scheduler, + "speedup_vs_official_denoise_p50": ( + True if args.skip_speedup_gate else speedup is not None and speedup >= args.min_speedup + ), + "final_action_cos": cos >= args.min_final_action_cos, + "final_action_rel_l2": rel_l2 < args.max_final_action_rel_l2, + } + result: dict[str, Any] = { + "backend": "flashrt", + "engine": args.engine, + "bf16_projs": args.bf16_projs, + "quant_graph": bool(args.engine != "bf16-eager" and not args.no_quant_graph), + "checkpoint": str(checkpoint), + "reference_dump": str(reference_dump), + "boundary_dump": str(boundary_dump), + "device": args.device, + "use_cuda_graphs": bool(args.use_cuda_graphs), + "fa4_status": fa4_backend.status(), + "native_attention": native_attention, + "graph_attention": graph_attention, + "native_scheduler": native_scheduler, + "init_s": init_s, + "warmup_steps": args.warmup_steps, + "warmup_max_velocity_abs_diff": ( + warmup_result.max_velocity_abs_diff if warmup_result is not None else None + ), + "iters": args.iters, + "times_s": times, + "p50_s": p50_s, + "p90_s": _percentile(sorted_times, 0.90), + "min_s": min(sorted_times), + "max_s": max(sorted_times), + "steps": last.steps_run, + "max_input_abs_diff": last.max_input_abs_diff, + "max_velocity_abs_diff": last.max_velocity_abs_diff, + "final_action_cos": cos, + "final_action_rel_l2": rel_l2, + "final_action_max_abs": max_abs, + "official_denoise_s": official_denoise, + "official_e2e_s": official_e2e, + "official_decode_s": official_decode, + "official_non_denoise_non_decode_s": official_non_denoise_non_decode, + "estimated_action_only_e2e_s": estimated_action_only_e2e, + "speedup_vs_official_denoise_p50": speedup, + "estimated_action_only_e2e_speedup": ( + official_e2e / estimated_action_only_e2e + if official_e2e is not None and estimated_action_only_e2e is not None + else None + ), + "peak_mem_gib": ( + torch.cuda.max_memory_allocated() / (1024**3) if args.device == "cuda" else None + ), + "gate_thresholds": { + "min_speedup": args.min_speedup, + "min_final_action_cos": args.min_final_action_cos, + "max_final_action_rel_l2": args.max_final_action_rel_l2, + "skip_speedup_gate": args.skip_speedup_gate, + }, + "gate_checks": gate_checks, + "gates_passed": all(gate_checks.values()), + } + if args.profile_loop_breakdown: + result["loop_breakdown"] = _profile_loop_breakdown(runner, device=args.device) + + text = json.dumps(result, indent=2, sort_keys=True) + print(text) + if args.json_out: + out = Path(args.json_out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(text + "\n", encoding="utf-8") + if args.enforce_gates and not result["gates_passed"]: + failed = ", ".join(name for name, passed in gate_checks.items() if not passed) + raise SystemExit(f"Cosmos3-Edge FlashRT benchmark gates failed: {failed}") + + +if __name__ == "__main__": + main() diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 5afa35ed..c2a04116 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -72,6 +72,14 @@ extern "C" int fp8_conv3d_v18_ncdhw_res_bf16out( int N, int T_cache, int T_new, int H, int W, int Ci, int Co, float alpha, cudaStream_t stream); #endif +#ifdef ENABLE_BF16_CONV3D_V0 +extern "C" int bf16_conv3d_v0_ndhwc_bf16out( + const void* cache_x_bf16, const void* new_x_bf16, + const void* w_bf16, void* y_bf16, + const void* bias_bf16, + int N, int T_cache, int T_new, int H, int W, int Ci, int Co, + float alpha, cudaStream_t stream); +#endif #ifdef ENABLE_FP8_CONV2D_3X3_V1 extern "C" int fp8_conv2d_3x3_v1_nhwc_bf16out( const void* x_fp8, const void* w_fp8, void* y_bf16, @@ -155,6 +163,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #ifdef FLASHRT_HAVE_AUDIO_CODEBOOK #include "kernels/delayed_codebook_kernels.cuh" #endif +#include "kernels/cosmos3_edge_misc.cuh" #ifdef FLASHRT_HAVE_QWEN36_KERNELS #include "kernels/qwen36_misc.cuh" #endif @@ -2286,6 +2295,11 @@ PYBIND11_MODULE(flash_rt_kernels, m) { relu_inplace_bf16(reinterpret_cast<__nv_bfloat16*>(x), n, to_stream(stream)); }, py::arg("x"), py::arg("n"), py::arg("stream") = 0); + m.def("relu2_inplace_bf16", [](uintptr_t x, int n, uintptr_t stream) { + extern void relu2_inplace_bf16(__nv_bfloat16*, int, cudaStream_t); + relu2_inplace_bf16(reinterpret_cast<__nv_bfloat16*>(x), n, to_stream(stream)); + }, py::arg("x"), py::arg("n"), py::arg("stream") = 0); + // GQA KV repeat interleave (for Qwen3 8→16 heads) m.def("gpu_repeat_interleave_heads", [](uintptr_t src, uintptr_t dst, int S, int NH_src, int HD, int repeat, uintptr_t stream) { @@ -4557,6 +4571,213 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("step"), py::arg("stream") = 0); #endif + m.def("cosmos3_edge_qk_norm_rope_bf16", + [](uintptr_t q_in, uintptr_t k_in, uintptr_t q_weight, uintptr_t k_weight, + uintptr_t cos, uintptr_t sin, uintptr_t q_out, uintptr_t k_out, + int rows, int q_heads, int k_heads, int head_dim, int rope_dim, + float eps, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_qk_norm_rope_bf16( + reinterpret_cast(q_in), + reinterpret_cast(k_in), + reinterpret_cast(q_weight), + reinterpret_cast(k_weight), + reinterpret_cast(cos), + reinterpret_cast(sin), + reinterpret_cast<__nv_bfloat16*>(q_out), + reinterpret_cast<__nv_bfloat16*>(k_out), + rows, q_heads, k_heads, head_dim, rope_dim, eps, + to_stream(stream)); + }, + py::arg("q_in"), py::arg("k_in"), py::arg("q_weight"), + py::arg("k_weight"), py::arg("cos"), py::arg("sin"), + py::arg("q_out"), py::arg("k_out"), py::arg("rows"), + py::arg("q_heads"), py::arg("k_heads"), py::arg("head_dim"), + py::arg("rope_dim"), py::arg("eps"), py::arg("stream") = 0); + + m.def("cosmos3_edge_fill_flat_velocity_bf16", + [](uintptr_t action, uintptr_t velocity, int flat_dim, int action_numel, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_fill_flat_velocity_bf16( + reinterpret_cast(action), + reinterpret_cast<__nv_bfloat16*>(velocity), + flat_dim, + action_numel, + to_stream(stream)); + }, + py::arg("action"), py::arg("velocity"), py::arg("flat_dim"), + py::arg("action_numel"), py::arg("stream") = 0); + + m.def("cosmos3_edge_add_bias_zero_action_tail_bf16", + [](uintptr_t action, uintptr_t bias, int rows, int cols, int valid_cols, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_add_bias_zero_action_tail_bf16( + reinterpret_cast<__nv_bfloat16*>(action), + reinterpret_cast(bias), + rows, + cols, + valid_cols, + to_stream(stream)); + }, + py::arg("action"), py::arg("bias"), py::arg("rows"), py::arg("cols"), + py::arg("valid_cols"), py::arg("stream") = 0); + + m.def("cosmos3_edge_scatter_rows_bf16", + [](uintptr_t src, uintptr_t dst, uintptr_t row_indices, int rows, int hidden, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_scatter_rows_bf16( + reinterpret_cast(src), + reinterpret_cast<__nv_bfloat16*>(dst), + reinterpret_cast(row_indices), + rows, + hidden, + to_stream(stream)); + }, + py::arg("src"), py::arg("dst"), py::arg("row_indices"), + py::arg("rows"), py::arg("hidden"), py::arg("stream") = 0); + + m.def("cosmos3_edge_gather_rows_bf16", + [](uintptr_t src, uintptr_t dst, uintptr_t row_indices, int rows, int hidden, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_gather_rows_bf16( + reinterpret_cast(src), + reinterpret_cast<__nv_bfloat16*>(dst), + reinterpret_cast(row_indices), + rows, + hidden, + to_stream(stream)); + }, + py::arg("src"), py::arg("dst"), py::arg("row_indices"), + py::arg("rows"), py::arg("hidden"), py::arg("stream") = 0); + + m.def("cosmos3_edge_qk_norm_rope_strided_bf16", + [](uintptr_t q_in, uintptr_t k_in, uintptr_t q_weight, uintptr_t k_weight, + uintptr_t cos, uintptr_t sin, uintptr_t q_out, uintptr_t k_out, + int rows, int q_heads, int k_heads, int q_in_row_stride, int k_in_row_stride, + float eps, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_qk_norm_rope_strided_bf16( + reinterpret_cast(q_in), + reinterpret_cast(k_in), + reinterpret_cast(q_weight), + reinterpret_cast(k_weight), + reinterpret_cast(cos), + reinterpret_cast(sin), + reinterpret_cast<__nv_bfloat16*>(q_out), + reinterpret_cast<__nv_bfloat16*>(k_out), + rows, q_heads, k_heads, q_in_row_stride, k_in_row_stride, eps, + to_stream(stream)); + }, + py::arg("q_in"), py::arg("k_in"), py::arg("q_weight"), py::arg("k_weight"), + py::arg("cos"), py::arg("sin"), py::arg("q_out"), py::arg("k_out"), + py::arg("rows"), py::arg("q_heads"), py::arg("k_heads"), + py::arg("q_in_row_stride"), py::arg("k_in_row_stride"), + py::arg("eps"), py::arg("stream") = 0); + + m.def("cosmos3_edge_relu2_to_fp8_static_bf16", + [](uintptr_t x, uintptr_t out, uintptr_t d_scale, int numel, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_relu2_to_fp8_static_bf16( + reinterpret_cast(x), + reinterpret_cast<__nv_fp8_e4m3*>(out), + reinterpret_cast(d_scale), + numel, + to_stream(stream)); + }, + py::arg("x"), py::arg("out"), py::arg("d_scale"), + py::arg("numel"), py::arg("stream") = 0); + + m.def("cosmos3_edge_copy_action_tail_f32_to_bf16", + [](uintptr_t flat_noise, uintptr_t action, int flat_dim, int action_numel, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_copy_action_tail_f32_to_bf16( + reinterpret_cast(flat_noise), + reinterpret_cast<__nv_bfloat16*>(action), + flat_dim, + action_numel, + to_stream(stream)); + }, + py::arg("flat_noise"), py::arg("action"), py::arg("flat_dim"), + py::arg("action_numel"), py::arg("stream") = 0); + + m.def("cosmos3_edge_add_action_bias_timestep_bf16", + [](uintptr_t x, uintptr_t static_bias, uintptr_t timestep, int rows, int hidden, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_add_action_bias_timestep_bf16( + reinterpret_cast<__nv_bfloat16*>(x), + reinterpret_cast(static_bias), + reinterpret_cast(timestep), + rows, + hidden, + to_stream(stream)); + }, + py::arg("x"), py::arg("static_bias"), py::arg("timestep"), + py::arg("rows"), py::arg("hidden"), py::arg("stream") = 0); + + m.def("cosmos3_edge_add_bf16", + [](uintptr_t a, uintptr_t b, uintptr_t out, int numel, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_add_bf16( + reinterpret_cast(a), + reinterpret_cast(b), + reinterpret_cast<__nv_bfloat16*>(out), + numel, + to_stream(stream)); + }, + py::arg("a"), py::arg("b"), py::arg("out"), py::arg("numel"), + py::arg("stream") = 0); + + m.def("cosmos3_edge_avgdown3d_bf16", + [](uintptr_t x, uintptr_t out, int b, int c, int t, int h, int w, + int out_c, int factor_t, int factor_s, int group_size, uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_avgdown3d_bf16( + reinterpret_cast(x), + reinterpret_cast<__nv_bfloat16*>(out), + b, + c, + t, + h, + w, + out_c, + factor_t, + factor_s, + group_size, + to_stream(stream)); + }, + py::arg("x"), py::arg("out"), py::arg("b"), py::arg("c"), + py::arg("t"), py::arg("h"), py::arg("w"), py::arg("out_c"), + py::arg("factor_t"), py::arg("factor_s"), py::arg("group_size"), + py::arg("stream") = 0); + + m.def("cosmos3_edge_unipc_step_f32_bf16", + [](uintptr_t sample, uintptr_t velocity, uintptr_t prev_m1, + uintptr_t prev_m2, uintptr_t prev_last_sample, uintptr_t next_sample, + uintptr_t current_m, uintptr_t current_last_sample, int numel, + float sigma, int corrector_order, int predictor_order, + float c_sample, float c_last, float c_prev_m1, float c_prev_m2, + float c_curr_m, float p_sample, float p_curr_m, float p_prev_m1, + uintptr_t stream) { + flash_rt::kernels::cosmos3_edge_unipc_step_f32_bf16( + reinterpret_cast(sample), + reinterpret_cast(velocity), + reinterpret_cast(prev_m1), + reinterpret_cast(prev_m2), + reinterpret_cast(prev_last_sample), + reinterpret_cast(next_sample), + reinterpret_cast(current_m), + reinterpret_cast(current_last_sample), + numel, + sigma, + corrector_order, + predictor_order, + c_sample, + c_last, + c_prev_m1, + c_prev_m2, + c_curr_m, + p_sample, + p_curr_m, + p_prev_m1, + to_stream(stream)); + }, + py::arg("sample"), py::arg("velocity"), py::arg("prev_m1"), + py::arg("prev_m2"), py::arg("prev_last_sample"), py::arg("next_sample"), + py::arg("current_m"), py::arg("current_last_sample"), py::arg("numel"), + py::arg("sigma"), py::arg("corrector_order"), py::arg("predictor_order"), + py::arg("c_sample"), py::arg("c_last"), py::arg("c_prev_m1"), + py::arg("c_prev_m2"), py::arg("c_curr_m"), py::arg("p_sample"), + py::arg("p_curr_m"), py::arg("p_prev_m1"), py::arg("stream") = 0); + #ifdef FLASHRT_HAVE_QWEN36_KERNELS m.def("qwen36_embedding_lookup_bf16", [](uintptr_t token_ids, uintptr_t embed, uintptr_t out, @@ -6069,6 +6290,28 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("Ci"), py::arg("Co"), py::arg("alpha") = 1.0f, py::arg("stream") = 0); #endif +#ifdef ENABLE_BF16_CONV3D_V0 + m.def("bf16_conv3d_v0_ndhwc_bf16out", + [](uintptr_t cache_x_bf16, uintptr_t new_x_bf16, + uintptr_t w_bf16, uintptr_t y_bf16, + uintptr_t bias_bf16, + int N, int T_cache, int T_new, int H, int W, int Ci, int Co, + float alpha, uintptr_t stream) { + return ::bf16_conv3d_v0_ndhwc_bf16out( + to_ptr(cache_x_bf16), to_ptr(new_x_bf16), + to_ptr(w_bf16), to_ptr(y_bf16), + to_ptr(bias_bf16), + N, T_cache, T_new, H, W, Ci, Co, alpha, + to_stream(stream)); + }, + py::arg("cache_x_bf16"), py::arg("new_x_bf16"), + py::arg("w_bf16"), py::arg("y_bf16"), + py::arg("bias_bf16") = 0, + py::arg("N"), py::arg("T_cache"), py::arg("T_new"), + py::arg("H"), py::arg("W"), + py::arg("Ci"), py::arg("Co"), + py::arg("alpha") = 1.0f, py::arg("stream") = 0); +#endif #ifdef ENABLE_FP8_CONV2D_3X3_V1 m.def("fp8_conv2d_3x3_v1_nhwc_bf16out", [](uintptr_t x_fp8, uintptr_t w_fp8, uintptr_t y_bf16, diff --git a/csrc/conv/bf16_conv3d_v0_sm110.cu b/csrc/conv/bf16_conv3d_v0_sm110.cu new file mode 100644 index 00000000..a1ed50cc --- /dev/null +++ b/csrc/conv/bf16_conv3d_v0_sm110.cu @@ -0,0 +1,291 @@ +// ================================================================ +// flash_rt — BF16 Conv3d fprop, SM110 bring-up v0 +// +// Cosmos3-Edge Wan2.2 VAE steady CausalConv3d probe: +// cache_x [B, T_cache=2, H, W, Ci] BF16 +// new_x [B, T_new, H, W, Ci] BF16 +// w [Co, 3, 3, 3, Ci] BF16 +// y [B, T_new, H, W, Co] BF16 +// +// This mirrors fp8_conv3d_v17's virtual cache concat and direct causal +// output, but uses BF16 tensor cores. It is a probe kernel, not a default +// production path. +// ================================================================ + +#include +#include +#include +#include + +namespace flash_rt { +namespace conv { + +constexpr int B3D_BLOCK_M = 128; +constexpr int B3D_BLOCK_N = 128; +constexpr int B3D_BLOCK_K = 16; +constexpr int B3D_N_ATOMS = B3D_BLOCK_N / 8; +constexpr int B3D_NUM_WARPS = 8; +constexpr int B3D_THREADS = B3D_NUM_WARPS * 32; +constexpr int B3D_STAGES = 2; +constexpr int B3D_SMEM_K_STRIDE_BYTES = 48; + +__device__ __forceinline__ +void b3d_mma_m16n8k16_bf16( + float &d0, float &d1, float &d2, float &d3, + uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, + uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ +const uint8_t* b3d_x_byte_ptr(const __nv_bfloat16* cache_x, + const __nv_bfloat16* new_x, + int m_global, int k_global, + int N, int T_cache, int T_new, + int H, int W, int Ci) { + int K_total = 27 * Ci; + int M_total = N * T_new * H * W; + if (k_global >= K_total || m_global >= M_total) return nullptr; + int spatial = T_new * H * W; + int b_idx = m_global / spatial; + int rem = m_global - b_idx * spatial; + int t_out = rem / (H * W); + rem -= t_out * (H * W); + int h_out = rem / W; + int w_out = rem - h_out * W; + int q = k_global / Ci; + int ci0 = k_global - q * Ci; + int ks = q % 3; q /= 3; + int kr = q % 3; + int kt = q / 3; + int d_in = t_out + kt; + int h_in = h_out + kr - 1; + int w_in = w_out + ks - 1; + if (h_in < 0 || h_in >= H || w_in < 0 || w_in >= W) return nullptr; + if (d_in < T_cache) { + int idx = (((b_idx * T_cache + d_in) * H + h_in) * W + w_in) * Ci + ci0; + return reinterpret_cast(&cache_x[idx]); + } + int d_new = d_in - T_cache; + int idx = (((b_idx * T_new + d_new) * H + h_in) * W + w_in) * Ci + ci0; + return reinterpret_cast(&new_x[idx]); +} + +__device__ __forceinline__ +const uint8_t* b3d_w_byte_ptr(const __nv_bfloat16* w, + int co, int k_global, int Co, int Ci) { + int K_total = 27 * Ci; + if (co >= Co || k_global >= K_total) return nullptr; + int q = k_global / Ci; + int ci0 = k_global - q * Ci; + int ks = q % 3; q /= 3; + int kr = q % 3; + int kt = q / 3; + int idx = (((co * 3 + kt) * 3 + kr) * 3 + ks) * Ci + ci0; + return reinterpret_cast(&w[idx]); +} + +__device__ __forceinline__ +void b3d_cp_async_16(uint32_t smem_int_ptr, const uint8_t* src) { + int src_bytes = (src == nullptr) ? 0 : 16; + asm volatile( + "cp.async.ca.shared.global [%0], [%1], 16, %2;\n" + :: "r"(smem_int_ptr), "l"(src), "r"(src_bytes)); +} + +__device__ __forceinline__ +uint32_t b3d_to_smem_int(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +__global__ void __launch_bounds__(B3D_THREADS, 2) +bf16_conv3d_v0_kernel( + const __nv_bfloat16* __restrict__ cache_x, + const __nv_bfloat16* __restrict__ new_x, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, + const __nv_bfloat16* __restrict__ bias, + int N, int T_cache, int T_new, int H, int W, int Ci, int Co, + float alpha, + int M_tiles, int N_tiles) +{ + __shared__ __align__(16) uint8_t A_smem[B3D_STAGES][B3D_BLOCK_M * B3D_SMEM_K_STRIDE_BYTES]; + __shared__ __align__(16) uint8_t B_smem[B3D_STAGES][B3D_BLOCK_N * B3D_SMEM_K_STRIDE_BYTES]; + + const int t = threadIdx.x; + const int warp_id = t / 32; + const int lane = t % 32; + const int l = lane % 4; + const int h = lane / 4; + + const int M_total = N * T_new * H * W; + const int K_total = 27 * Ci; + + const int ld_row_a = t / 2; + const int ld_k_elem_a = (t & 1) * 8; + const int ld_row_b = t / 2; + const int ld_k_elem_b = (t & 1) * 8; + + int tile_idx = blockIdx.x; + int m_idx = tile_idx / N_tiles; + int n_idx = tile_idx - m_idx * N_tiles; + int m_base = m_idx * B3D_BLOCK_M; + int co_base = n_idx * B3D_BLOCK_N; + if (m_base >= M_total || co_base >= Co) return; + + float dA[B3D_N_ATOMS] = {0}; + float dB[B3D_N_ATOMS] = {0}; + float dC[B3D_N_ATOMS] = {0}; + float dD[B3D_N_ATOMS] = {0}; + + auto issue_load = [&](int stage, int k_base) { + { + const uint8_t* src = b3d_x_byte_ptr(cache_x, new_x, + m_base + ld_row_a, + k_base + ld_k_elem_a, + N, T_cache, T_new, H, W, Ci); + uint32_t smem_int = b3d_to_smem_int( + &A_smem[stage][ld_row_a * B3D_SMEM_K_STRIDE_BYTES + ld_k_elem_a * 2]); + b3d_cp_async_16(smem_int, src); + } + { + const uint8_t* src = b3d_w_byte_ptr(w, co_base + ld_row_b, + k_base + ld_k_elem_b, + Co, Ci); + uint32_t smem_int = b3d_to_smem_int( + &B_smem[stage][ld_row_b * B3D_SMEM_K_STRIDE_BYTES + ld_k_elem_b * 2]); + b3d_cp_async_16(smem_int, src); + } + }; + + issue_load(0, 0); + asm volatile("cp.async.commit_group;\n" ::); + + int compute_stage = 0; + for (int k_base = 0; k_base < K_total; k_base += B3D_BLOCK_K) { + int next_stage = compute_stage ^ 1; + int k_next = k_base + B3D_BLOCK_K; + if (k_next < K_total) { + issue_load(next_stage, k_next); + } + asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group 1;\n" ::); + __syncthreads(); + + const int warp_M_off = warp_id * 16; + const int kA0 = 4 * l; + const int kA2 = 4 * l + 16; + int rA0 = warp_M_off + h; + int rA1 = warp_M_off + h + 8; + uint32_t A0 = *reinterpret_cast( + &A_smem[compute_stage][rA0 * B3D_SMEM_K_STRIDE_BYTES + kA0]); + uint32_t A1 = *reinterpret_cast( + &A_smem[compute_stage][rA1 * B3D_SMEM_K_STRIDE_BYTES + kA0]); + uint32_t A2 = *reinterpret_cast( + &A_smem[compute_stage][rA0 * B3D_SMEM_K_STRIDE_BYTES + kA2]); + uint32_t A3 = *reinterpret_cast( + &A_smem[compute_stage][rA1 * B3D_SMEM_K_STRIDE_BYTES + kA2]); + + #pragma unroll + for (int n_atom = 0; n_atom < B3D_N_ATOMS; ++n_atom) { + int co_n = n_atom * 8 + h; + uint32_t B0 = *reinterpret_cast( + &B_smem[compute_stage][co_n * B3D_SMEM_K_STRIDE_BYTES + kA0]); + uint32_t B1 = *reinterpret_cast( + &B_smem[compute_stage][co_n * B3D_SMEM_K_STRIDE_BYTES + kA2]); + b3d_mma_m16n8k16_bf16( + dA[n_atom], dB[n_atom], dC[n_atom], dD[n_atom], + A0, A1, A2, A3, B0, B1); + } + + compute_stage = next_stage; + } + + asm volatile("cp.async.wait_all;\n" ::); + + const int warp_M_off = warp_id * 16; + #pragma unroll + for (int n_atom = 0; n_atom < B3D_N_ATOMS; ++n_atom) { + int co_pair = co_base + n_atom * 8 + 2 * l; + int row0 = m_base + warp_M_off + h; + int row1 = m_base + warp_M_off + h + 8; + float b0 = 0.f, b1 = 0.f; + if (bias != nullptr && co_pair < Co) { + b0 = __bfloat162float(bias[co_pair]); + if (co_pair + 1 < Co) b1 = __bfloat162float(bias[co_pair + 1]); + } + if (co_pair + 1 < Co) { + __nv_bfloat162 packAB; + packAB.x = __float2bfloat16(dA[n_atom] * alpha + b0); + packAB.y = __float2bfloat16(dB[n_atom] * alpha + b1); + __nv_bfloat162 packCD; + packCD.x = __float2bfloat16(dC[n_atom] * alpha + b0); + packCD.y = __float2bfloat16(dD[n_atom] * alpha + b1); + if (row0 < M_total) { + *reinterpret_cast<__nv_bfloat162*>(&y[row0 * Co + co_pair]) = packAB; + } + if (row1 < M_total) { + *reinterpret_cast<__nv_bfloat162*>(&y[row1 * Co + co_pair]) = packCD; + } + } else { + auto store = [&](int row, int co, float v, float bv) { + if (row < M_total && co < Co) { + y[row * Co + co] = __float2bfloat16(v * alpha + bv); + } + }; + store(row0, co_pair, dA[n_atom], b0); + store(row0, co_pair + 1, dB[n_atom], b1); + store(row1, co_pair, dC[n_atom], b0); + store(row1, co_pair + 1, dD[n_atom], b1); + } + } +} + +extern "C" int bf16_conv3d_v0_ndhwc_bf16out( + const void* cache_x_bf16, const void* new_x_bf16, + const void* w_bf16, void* y_bf16, + const void* bias_bf16, + int N, int T_cache, int T_new, int H, int W, int Ci, int Co, + float alpha, cudaStream_t stream) +{ + if (Ci % B3D_BLOCK_K != 0 || Co % 8 != 0) { + std::fprintf(stderr, + "[bf16_conv3d_v0] Ci%%%d (got %d) or Co%%8 (got %d) bad\n", + B3D_BLOCK_K, Ci, Co); + return -1; + } + if (T_cache != 2) { + std::fprintf(stderr, "[bf16_conv3d_v0] T_cache must be 2 (got %d)\n", T_cache); + return -3; + } + int M = N * T_new * H * W; + int M_tiles = (M + B3D_BLOCK_M - 1) / B3D_BLOCK_M; + int N_tiles = (Co + B3D_BLOCK_N - 1) / B3D_BLOCK_N; + int total_tiles = M_tiles * N_tiles; + + bf16_conv3d_v0_kernel<<>>( + reinterpret_cast(cache_x_bf16), + reinterpret_cast(new_x_bf16), + reinterpret_cast(w_bf16), + reinterpret_cast<__nv_bfloat16*>(y_bf16), + reinterpret_cast(bias_bf16), + N, T_cache, T_new, H, W, Ci, Co, alpha, + M_tiles, N_tiles); + cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { + std::fprintf(stderr, "[bf16_conv3d_v0] launch err: %s\n", cudaGetErrorString(e)); + return -2; + } + return 0; +} + +} // namespace conv +} // namespace flash_rt diff --git a/csrc/fp4_bindings.cpp b/csrc/fp4_bindings.cpp index ea7f36e1..ead34813 100644 --- a/csrc/fp4_bindings.cpp +++ b/csrc/fp4_bindings.cpp @@ -20,6 +20,7 @@ #include "quantize/quantize_fp4_sfa.cuh" #include "quantize/reshape_scales_sfa.cuh" #include "fused_fp16/rms_norm_noweight_fp16.cuh" +#include "fused_fp4/cosmos3_edge_fp4.cuh" #include "fused_fp4/norm_silu_fp4_sfa.cuh" #include "fused_fp4/silu_mul_two_fp4_to_fp4.cuh" @@ -235,6 +236,39 @@ reshape_linear_scales_to_sfa, in a single kernel launch. py::arg("seq_len"), py::arg("dim"), py::arg("stream") = 0, "F3: fused residual+rms_norm + fp4_quant + SFA write."); + // Cosmos3-Edge model-specific fused FP4 quant kernels (bf16 residual chain). + m.def("cosmos3_edge_res_rms_fp4_sfa_bf16", + [](uintptr_t residual, uintptr_t x, uintptr_t weight, + uintptr_t packed, uintptr_t sfa, + int seq_len, int dim, float eps, uintptr_t stream) { + flash_rt::fused_fp4::cosmos3_edge_res_rms_fp4_sfa_bf16( + reinterpret_cast<__nv_bfloat16*>(residual), + reinterpret_cast(x), + reinterpret_cast(weight), + reinterpret_cast(packed), + reinterpret_cast(sfa), + seq_len, dim, eps, + reinterpret_cast(stream)); + }, + py::arg("residual"), py::arg("x"), py::arg("weight"), + py::arg("packed"), py::arg("sfa"), + py::arg("seq_len"), py::arg("dim"), py::arg("eps"), py::arg("stream") = 0, + "Cosmos3-Edge: bf16 residual += x; weighted RMSNorm; NVFP4 quant + SFA."); + + m.def("cosmos3_edge_relu2_fp4_sfa_fp16", + [](uintptr_t x, uintptr_t packed, uintptr_t sfa, + int seq_len, int dim, uintptr_t stream) { + flash_rt::fused_fp4::cosmos3_edge_relu2_fp4_sfa_fp16( + reinterpret_cast(x), + reinterpret_cast(packed), + reinterpret_cast(sfa), + seq_len, dim, + reinterpret_cast(stream)); + }, + py::arg("x"), py::arg("packed"), py::arg("sfa"), + py::arg("seq_len"), py::arg("dim"), py::arg("stream") = 0, + "Cosmos3-Edge: relu(x)^2 (fp16 in) -> NVFP4 quant + SFA."); + // GEGLU (tanh-approx GELU(gate) * up) fused FP4 kernels. m.def("gate_geglu_fp4_sfa_fp16", [](uintptr_t merged, uintptr_t packed, uintptr_t sfa, diff --git a/csrc/fused_fp4/cosmos3_edge_fp4.cu b/csrc/fused_fp4/cosmos3_edge_fp4.cu new file mode 100644 index 00000000..88ae978e --- /dev/null +++ b/csrc/fused_fp4/cosmos3_edge_fp4.cu @@ -0,0 +1,212 @@ +// ============================================================================ +// Cosmos3-Edge model-specific NVFP4 fused quant kernels (additive). +// +// 1. cosmos3_edge_res_rms_fp4_sfa_bf16: +// residual(bf16, updated in place) += x(bf16); RMSNorm with gamma weight +// (bf16, eps parameterized); NVFP4 quantize + CUTLASS tile-interleaved SFA. +// Register-only F3-v2 design: 1 thread per NVFP4 block, D = 16*blockDim. +// +// 2. cosmos3_edge_relu2_fp4_sfa_fp16: +// x(fp16, up-proj GEMM output) -> relu(x)^2 -> NVFP4 quantize + SFA. +// 1 thread per NVFP4 block, F4-v2 design. +// +// Both feed cutlass_fp4_sq_fp16 (A-side packed + SFA). +// ============================================================================ +#include +#include +#include +#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 + +using CfgEdge = cutlass::detail::Sm1xxBlockScaledConfig<16>; + +__device__ __forceinline__ uint8_t fp32_to_e2m1_edge(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; +} + +__device__ __forceinline__ void quant_pack_sfa_edge( + float (&vals)[16], float amax, int row, int col_base, int block_idx, + uint8_t* __restrict__ packed_row, uint8_t* __restrict__ dst_sfa, int sfa_off) { + 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 inv_bs = 1.f / static_cast(bs_q); + dst_sfa[sfa_off] = *reinterpret_cast(&bs_q); + #pragma unroll + for (int p = 0; p < 8; ++p) { + uint8_t lo = fp32_to_e2m1_edge(vals[2 * p ] * inv_bs); + uint8_t hi = fp32_to_e2m1_edge(vals[2 * p + 1] * inv_bs); + packed_row[block_idx * 8 + p] = lo | (hi << 4); + } +} + +template +__global__ void edge_res_rms_fp4_sfa_bf16_kernel( + __nv_bfloat16* __restrict__ residual, + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ weight, + uint8_t* __restrict__ packed, + uint8_t* __restrict__ dst_sfa, + LayoutSF layout, + int D, + float eps) { + const int r = blockIdx.x; + __nv_bfloat16* res_row = residual + r * D; + const __nv_bfloat16* x_row = x + r * D; + uint8_t* packed_row = packed + r * (D / 2); + + const int col_base = threadIdx.x * 16; + if (col_base >= D) return; + + __nv_bfloat162* res2 = reinterpret_cast<__nv_bfloat162*>(res_row + col_base); + const __nv_bfloat162* x2 = reinterpret_cast(x_row + col_base); + const __nv_bfloat162* w2 = reinterpret_cast(weight + col_base); + + float vals[16]; + float local_ssq = 0.f; + #pragma unroll + for (int p = 0; p < 8; ++p) { + __nv_bfloat162 rv = res2[p]; + __nv_bfloat162 xv = x2[p]; + float a = __bfloat162float(rv.x) + __bfloat162float(xv.x); + float b = __bfloat162float(rv.y) + __bfloat162float(xv.y); + vals[2 * p] = a; + vals[2 * p + 1] = b; + res2[p] = __floats2bfloat162_rn(a, b); + local_ssq += a * a + b * b; + } + + __shared__ float sh[16]; + int lane = threadIdx.x % 32, wid = threadIdx.x / 32; + #pragma unroll + for (int o = 16; o > 0; o >>= 1) local_ssq += __shfl_xor_sync(0xffffffff, local_ssq, o); + if (!lane) sh[wid] = local_ssq; + __syncthreads(); + float ssq; + if (!wid) { + ssq = (lane < (blockDim.x / 32)) ? sh[lane] : 0.f; + for (int o = 16; o > 0; o >>= 1) ssq += __shfl_xor_sync(0xffffffff, ssq, o); + } + __syncthreads(); + if (!threadIdx.x) sh[0] = ssq; + __syncthreads(); + + const float rms = __frsqrt_rn(sh[0] / D + eps); + + float amax = 0.f; + #pragma unroll + for (int p = 0; p < 8; ++p) { + __nv_bfloat162 wv = w2[p]; + float v0 = vals[2 * p] * rms * __bfloat162float(wv.x); + float v1 = vals[2 * p + 1] * rms * __bfloat162float(wv.y); + vals[2 * p] = v0; + vals[2 * p + 1] = v1; + float a0 = fabsf(v0), a1 = fabsf(v1); + if (a0 > amax) amax = a0; + if (a1 > amax) amax = a1; + } + + const int block_idx = threadIdx.x; + quant_pack_sfa_edge(vals, amax, r, col_base, block_idx, packed_row, dst_sfa, + layout(r, col_base, 0)); +} + +template +__global__ void edge_relu2_fp4_sfa_fp16_kernel( + const __half* __restrict__ x, // [S, H] + uint8_t* __restrict__ packed, // [S, H/2] + uint8_t* __restrict__ dst_sfa, + LayoutSF layout, + int H) { + const int block_idx = blockIdx.y * blockDim.x + threadIdx.x; + const int row = blockIdx.x; + const int n_blocks = H / 16; + if (block_idx >= n_blocks) return; + + const int col_base = block_idx * 16; + const __half2* x2 = reinterpret_cast(x + row * H + col_base); + + float vals[16]; + float amax = 0.f; + #pragma unroll + for (int i = 0; i < 8; ++i) { + __half2 v2 = x2[i]; + float a = fmaxf(__half2float(v2.x), 0.f); + float b = fmaxf(__half2float(v2.y), 0.f); + a = a * a; + b = b * b; + vals[2 * i] = a; + vals[2 * i + 1] = b; + if (a > amax) amax = a; + if (b > amax) amax = b; + } + + uint8_t* packed_row = packed + row * (H / 2); + quant_pack_sfa_edge(vals, amax, row, col_base, block_idx, packed_row, dst_sfa, + layout(row, col_base, 0)); +} + +#endif // FV_HAVE_CUTLASS + +void cosmos3_edge_res_rms_fp4_sfa_bf16( + __nv_bfloat16* residual, const __nv_bfloat16* x, const __nv_bfloat16* weight, + uint8_t* packed, uint8_t* sfa, + int seq_len, int dim, float eps, cudaStream_t stream) { +#if FV_HAVE_CUTLASS + int threads = dim / 16; + if (threads <= 0 || threads > 1024 || (dim % 16) != 0) return; + auto shape = cute::make_shape(seq_len, 1, dim, 1); + auto layout = CfgEdge::tile_atom_to_shape_SFA(shape); + edge_res_rms_fp4_sfa_bf16_kernel<<>>( + residual, x, weight, packed, sfa, layout, dim, eps); +#else + (void)residual; (void)x; (void)weight; (void)packed; (void)sfa; + (void)seq_len; (void)dim; (void)eps; (void)stream; +#endif +} + +void cosmos3_edge_relu2_fp4_sfa_fp16( + const __half* x, uint8_t* packed, uint8_t* sfa, + int seq_len, int dim, cudaStream_t stream) { +#if FV_HAVE_CUTLASS + if ((dim % 16) != 0) return; + auto shape = cute::make_shape(seq_len, 1, dim, 1); + auto layout = CfgEdge::tile_atom_to_shape_SFA(shape); + const int n_blocks = dim / 16; + const int threads = 256; + const int y_groups = (n_blocks + threads - 1) / threads; + dim3 grid(seq_len, y_groups); + edge_relu2_fp4_sfa_fp16_kernel<<>>( + x, packed, sfa, layout, dim); +#else + (void)x; (void)packed; (void)sfa; (void)seq_len; (void)dim; (void)stream; +#endif +} + +} // namespace fused_fp4 +} // namespace flash_rt diff --git a/csrc/fused_fp4/cosmos3_edge_fp4.cuh b/csrc/fused_fp4/cosmos3_edge_fp4.cuh new file mode 100644 index 00000000..287a9694 --- /dev/null +++ b/csrc/fused_fp4/cosmos3_edge_fp4.cuh @@ -0,0 +1,22 @@ +// Cosmos3-Edge model-specific fused NVFP4 quant kernels (additive). +#pragma once + +#include +#include +#include +#include + +namespace flash_rt { +namespace fused_fp4 { + +void cosmos3_edge_res_rms_fp4_sfa_bf16( + __nv_bfloat16* residual, const __nv_bfloat16* x, const __nv_bfloat16* weight, + uint8_t* packed, uint8_t* sfa, + int seq_len, int dim, float eps, cudaStream_t stream); + +void cosmos3_edge_relu2_fp4_sfa_fp16( + const __half* x, uint8_t* packed, uint8_t* sfa, + int seq_len, int dim, cudaStream_t stream); + +} // namespace fused_fp4 +} // namespace flash_rt diff --git a/csrc/kernels/activation.cu b/csrc/kernels/activation.cu index 6e827bdc..b5813793 100644 --- a/csrc/kernels/activation.cu +++ b/csrc/kernels/activation.cu @@ -397,3 +397,30 @@ void relu_inplace_bf16(__nv_bfloat16* x, int n, cudaStream_t stream) { int n2 = n >> 1; relu_inplace_kernel<__nv_bfloat16><<<(n2 + 255) / 256, 256, 0, stream>>>(x, n); } + +// ── ReLU² in-place: x = max(x, 0)^2, used by Cosmos3-Edge FFN ── +template +__global__ void relu2_inplace_kernel(T* __restrict__ x, int n) { + using T2 = typename packed2::type; + T2* x2 = reinterpret_cast(x); + int n2 = n >> 1; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n2) { + T2 val = x2[idx]; + float v0 = fmaxf(to_f32(val.x), 0.0f); + float v1 = fmaxf(to_f32(val.y), 0.0f); + x2[idx] = make_packed2(from_f32(v0 * v0), from_f32(v1 * v1)); + } + int tail = n2 * 2; + if (idx == 0 && tail < n) { + float v = fmaxf(to_f32(x[tail]), 0.0f); + x[tail] = from_f32(v * v); + } +} + +template __global__ void relu2_inplace_kernel<__nv_bfloat16>(__nv_bfloat16*, int); + +void relu2_inplace_bf16(__nv_bfloat16* x, int n, cudaStream_t stream) { + int n2 = n >> 1; + relu2_inplace_kernel<__nv_bfloat16><<<(n2 + 255) / 256, 256, 0, stream>>>(x, n); +} diff --git a/csrc/kernels/activation.cuh b/csrc/kernels/activation.cuh index 93284d96..c176337b 100644 --- a/csrc/kernels/activation.cuh +++ b/csrc/kernels/activation.cuh @@ -16,6 +16,7 @@ void gate_silu_mul(const __nv_bfloat16* gate, const __nv_bfloat16* up, __nv_bfloat16* out, int n, cudaStream_t stream = 0); void gelu_inplace(__nv_bfloat16* x, int n, cudaStream_t stream = 0); +void relu2_inplace_bf16(__nv_bfloat16* x, int n, cudaStream_t stream = 0); // G7.11 — fused (bias add + GELU(tanh)) in-place on bf16 tensor. // x: (M, N) bf16; bias: (N,) bf16 broadcast over rows. Replaces diff --git a/csrc/kernels/cosmos3_edge_misc.cu b/csrc/kernels/cosmos3_edge_misc.cu new file mode 100644 index 00000000..c1d1b849 --- /dev/null +++ b/csrc/kernels/cosmos3_edge_misc.cu @@ -0,0 +1,695 @@ +#include "cosmos3_edge_misc.cuh" + +#include "common.cuh" + +#include + +namespace flash_rt::kernels { +namespace { + +__device__ __forceinline__ __nv_bfloat16 norm_value_bf16( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ weight, + int base, + int col, + float rms) +{ + return __float2bfloat16(static_cast(x[base + col]) * rms * static_cast(weight[col])); +} + +__device__ __forceinline__ __nv_bfloat16 norm_rope_value_bf16( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ weight, + const __nv_bfloat16* __restrict__ cos, + const __nv_bfloat16* __restrict__ sin, + int row, + int base, + int col, + int rope_dim, + float rms) +{ + const __nv_bfloat16 xv_bf = norm_value_bf16(x, weight, base, col, rms); + if (col >= rope_dim) { + return xv_bf; + } + + const int half = rope_dim >> 1; + const int rot_col = (col < half) ? (col + half) : (col - half); + float rot = static_cast(norm_value_bf16(x, weight, base, rot_col, rms)); + if (col < half) { + rot = -rot; + } + const float xv = static_cast(xv_bf); + const float cv = static_cast(cos[row * rope_dim + col]); + const float sv = static_cast(sin[row * rope_dim + col]); + const float rot_sin_bf = static_cast(__float2bfloat16(rot * sv)); + return __float2bfloat16(rot_sin_bf + xv * cv); +} + +__global__ void qk_norm_rope_bf16_kernel( + const __nv_bfloat16* __restrict__ q_in, + const __nv_bfloat16* __restrict__ k_in, + const __nv_bfloat16* __restrict__ q_weight, + const __nv_bfloat16* __restrict__ k_weight, + const __nv_bfloat16* __restrict__ cos, + const __nv_bfloat16* __restrict__ sin, + __nv_bfloat16* __restrict__ q_out, + __nv_bfloat16* __restrict__ k_out, + int rows, + int q_heads, + int k_heads, + int head_dim, + int rope_dim, + float eps) +{ + const int q_vecs = rows * q_heads; + const int vec = blockIdx.x; + const bool is_q = vec < q_vecs; + const int heads = is_q ? q_heads : k_heads; + const int local_vec = is_q ? vec : vec - q_vecs; + const int row = local_vec / heads; + const int head = local_vec - row * heads; + const __nv_bfloat16* x = is_q ? q_in : k_in; + const __nv_bfloat16* weight = is_q ? q_weight : k_weight; + __nv_bfloat16* out = is_q ? q_out : k_out; + const int base = (row * heads + head) * head_dim; + + extern __shared__ float shared[]; + const int dim2 = head_dim >> 1; + float local_sum = 0.0f; + if (threadIdx.x < dim2) { + const int c0 = threadIdx.x << 1; + const int c1 = c0 + 1; + const float v0 = static_cast(x[base + c0]); + const float v1 = static_cast(x[base + c1]); + local_sum = v0 * v0 + v1 * v1; + } + const float rms = rsqrtf(block_reduce_sum(local_sum, shared) / head_dim + eps); + + if (threadIdx.x < dim2) { + const int c0 = threadIdx.x << 1; + const int c1 = c0 + 1; + out[base + c0] = norm_rope_value_bf16(x, weight, cos, sin, row, base, c0, rope_dim, rms); + out[base + c1] = norm_rope_value_bf16(x, weight, cos, sin, row, base, c1, rope_dim, rms); + } +} + +// Fast path for head_dim == rope_dim == 128: one warp per (row, head) vector. +// Per-head RMSNorm via register warp-reduce (no shared memory, no barriers); +// the rotate-half partner values live 16 lanes away, exchanged with +// __shfl_xor_sync instead of re-reading global memory. All loads/stores are +// __nv_bfloat162-vectorized. +__global__ void qk_norm_rope_bf16_warp128_kernel( + const __nv_bfloat16* __restrict__ q_in, + const __nv_bfloat16* __restrict__ k_in, + const __nv_bfloat16* __restrict__ q_weight, + const __nv_bfloat16* __restrict__ k_weight, + const __nv_bfloat16* __restrict__ cos, + const __nv_bfloat16* __restrict__ sin, + __nv_bfloat16* __restrict__ q_out, + __nv_bfloat16* __restrict__ k_out, + int rows, + int q_heads, + int k_heads, + float eps, + int total_vecs, + int q_in_row_stride, + int k_in_row_stride) +{ + const int warp = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + if (warp >= total_vecs) { + return; + } + const int lane = threadIdx.x & 31; + const int q_vecs = rows * q_heads; + const bool is_q = warp < q_vecs; + const int heads = is_q ? q_heads : k_heads; + const int local = is_q ? warp : warp - q_vecs; + const int row = local / heads; + const int head = local - row * heads; + const __nv_bfloat16* x = is_q ? q_in : k_in; + const __nv_bfloat16* w = is_q ? q_weight : k_weight; + __nv_bfloat16* out = is_q ? q_out : k_out; + const int in_stride = is_q ? q_in_row_stride : k_in_row_stride; + const int in_base = row * in_stride + (head << 7); + const int base = (row * heads + head) << 7; + + const __nv_bfloat162* x2 = reinterpret_cast(x + in_base); + const int e0 = lane << 1; + const __nv_bfloat162 a = x2[e0]; + const __nv_bfloat162 b = x2[e0 + 1]; + const float v0 = __bfloat162float(a.x), v1 = __bfloat162float(a.y); + const float v2 = __bfloat162float(b.x), v3 = __bfloat162float(b.y); + float ss = v0 * v0 + v1 * v1 + v2 * v2 + v3 * v3; +#pragma unroll + for (int off = 16; off; off >>= 1) { + ss += __shfl_xor_sync(0xffffffffu, ss, off); + } + const float rms = rsqrtf(ss * (1.0f / 128.0f) + eps); + + const __nv_bfloat162* w2 = reinterpret_cast(w); + const __nv_bfloat162 wa = w2[e0], wb = w2[e0 + 1]; + const float n0 = v0 * rms * __bfloat162float(wa.x); + const float n1 = v1 * rms * __bfloat162float(wa.y); + const float n2 = v2 * rms * __bfloat162float(wb.x); + const float n3 = v3 * rms * __bfloat162float(wb.y); + + const float p0 = __shfl_xor_sync(0xffffffffu, n0, 16); + const float p1 = __shfl_xor_sync(0xffffffffu, n1, 16); + const float p2 = __shfl_xor_sync(0xffffffffu, n2, 16); + const float p3 = __shfl_xor_sync(0xffffffffu, n3, 16); + const float sgn = (lane < 16) ? -1.0f : 1.0f; + + const __nv_bfloat162* c2 = reinterpret_cast(cos + (row << 7)); + const __nv_bfloat162* s2 = reinterpret_cast(sin + (row << 7)); + const __nv_bfloat162 ca = c2[e0], cb = c2[e0 + 1]; + const __nv_bfloat162 sa = s2[e0], sb = s2[e0 + 1]; + const float o0 = n0 * __bfloat162float(ca.x) + sgn * p0 * __bfloat162float(sa.x); + const float o1 = n1 * __bfloat162float(ca.y) + sgn * p1 * __bfloat162float(sa.y); + const float o2 = n2 * __bfloat162float(cb.x) + sgn * p2 * __bfloat162float(sb.x); + const float o3 = n3 * __bfloat162float(cb.y) + sgn * p3 * __bfloat162float(sb.y); + + __nv_bfloat162* out2 = reinterpret_cast<__nv_bfloat162*>(out + base); + out2[e0] = __floats2bfloat162_rn(o0, o1); + out2[e0 + 1] = __floats2bfloat162_rn(o2, o3); +} + +__global__ void fill_flat_velocity_bf16_kernel( + const __nv_bfloat16* __restrict__ action, + __nv_bfloat16* __restrict__ velocity, + int flat_dim, + int action_numel) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= flat_dim) { + return; + } + const int action_start = flat_dim - action_numel; + velocity[idx] = (idx < action_start) ? __float2bfloat16(0.0f) : action[idx - action_start]; +} + +__global__ void add_bias_zero_action_tail_bf16_kernel( + __nv_bfloat16* __restrict__ action, + const __nv_bfloat16* __restrict__ bias, + int total, + int cols, + int valid_cols) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int col = idx % cols; + if (col >= valid_cols) { + action[idx] = __float2bfloat16(0.0f); + return; + } + action[idx] = __float2bfloat16(static_cast(action[idx]) + static_cast(bias[col])); +} + +__global__ void scatter_rows_bf16_kernel( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ dst, + const int64_t* __restrict__ row_indices, + int total, + int hidden) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int row = idx / hidden; + const int col = idx - row * hidden; + const int64_t dst_row = row_indices[row]; + dst[dst_row * hidden + col] = src[idx]; +} + +__global__ void gather_rows_bf16_kernel( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ dst, + const int64_t* __restrict__ row_indices, + int total, + int hidden) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int row = idx / hidden; + const int col = idx - row * hidden; + const int64_t src_row = row_indices[row]; + dst[idx] = src[src_row * hidden + col]; +} + +__global__ void copy_action_tail_f32_to_bf16_kernel( + const float* __restrict__ flat_noise, + __nv_bfloat16* __restrict__ action, + int flat_dim, + int action_numel) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= action_numel) { + return; + } + action[idx] = __float2bfloat16(flat_noise[flat_dim - action_numel + idx]); +} + +__global__ void add_action_bias_timestep_bf16_kernel( + __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ static_bias, + const __nv_bfloat16* __restrict__ timestep, + int total, + int hidden) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int col = idx % hidden; + const __nv_bfloat16 with_static = __float2bfloat16( + static_cast(x[idx]) + static_cast(static_bias[col])); + x[idx] = __float2bfloat16(static_cast(with_static) + static_cast(timestep[col])); +} + +__global__ void add_bf16_kernel( + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ out, + int numel) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= numel) { + return; + } + out[idx] = __float2bfloat16(static_cast(a[idx]) + static_cast(b[idx])); +} + +__global__ void avgdown3d_bf16_kernel( + const __nv_bfloat16* __restrict__ x, + __nv_bfloat16* __restrict__ out, + int total, + int c, + int t, + int h, + int w, + int out_c, + int out_t, + int out_h, + int out_w, + int factor_t, + int factor_s, + int group_size, + int pad_t) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + + int rem = idx; + const int ow = rem % out_w; + rem /= out_w; + const int oh = rem % out_h; + rem /= out_h; + const int ot = rem % out_t; + rem /= out_t; + const int oc = rem % out_c; + const int b = rem / out_c; + + const int factor = factor_t * factor_s * factor_s; + float acc = 0.0f; + #pragma unroll + for (int g = 0; g < 8; ++g) { + if (g >= group_size) { + break; + } + int cf = oc * group_size + g; + const int ws = cf % factor_s; + cf /= factor_s; + const int hs = cf % factor_s; + cf /= factor_s; + const int tt = cf % factor_t; + const int ic = cf / factor_t; + if (ic >= c || oc * group_size + g >= c * factor) { + continue; + } + const int padded_t = ot * factor_t + tt; + if (padded_t < pad_t) { + continue; + } + const int it = padded_t - pad_t; + if (it >= t) { + continue; + } + const int ih = oh * factor_s + hs; + const int iw = ow * factor_s + ws; + const int64_t in_idx = + (((static_cast(b) * c + ic) * t + it) * h + ih) * w + iw; + acc += static_cast(x[in_idx]); + } + out[idx] = __float2bfloat16(acc / static_cast(group_size)); +} + +__global__ void unipc_step_f32_bf16_kernel( + const float* __restrict__ sample, + const __nv_bfloat16* __restrict__ velocity, + const float* __restrict__ prev_m1, + const float* __restrict__ prev_m2, + const float* __restrict__ prev_last_sample, + float* __restrict__ next_sample, + float* __restrict__ current_m, + float* __restrict__ current_last_sample, + int numel, + float sigma, + int corrector_order, + int predictor_order, + float c_sample, + float c_last, + float c_prev_m1, + float c_prev_m2, + float c_curr_m, + float p_sample, + float p_curr_m, + float p_prev_m1) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= numel) { + return; + } + + const float sample_v = sample[idx]; + const float sigma_velocity = static_cast(__float2bfloat16(sigma * static_cast(velocity[idx]))); + const float curr_m = sample_v - sigma_velocity; + + float corrected = c_sample * sample_v + c_curr_m * curr_m; + if (corrector_order >= 1) { + corrected += c_last * prev_last_sample[idx] + c_prev_m1 * prev_m1[idx]; + } + if (corrector_order >= 2) { + corrected += c_prev_m2 * prev_m2[idx]; + } + + float next = p_sample * corrected + p_curr_m * curr_m; + if (predictor_order >= 2) { + next += p_prev_m1 * prev_m1[idx]; + } + + current_m[idx] = curr_m; + current_last_sample[idx] = corrected; + next_sample[idx] = next; +} + +} // namespace + +void cosmos3_edge_qk_norm_rope_bf16( + const __nv_bfloat16* q_in, + const __nv_bfloat16* k_in, + const __nv_bfloat16* q_weight, + const __nv_bfloat16* k_weight, + const __nv_bfloat16* cos, + const __nv_bfloat16* sin, + __nv_bfloat16* q_out, + __nv_bfloat16* k_out, + int rows, + int q_heads, + int k_heads, + int head_dim, + int rope_dim, + float eps, + cudaStream_t stream) +{ + if (rows <= 0 || q_heads <= 0 || k_heads <= 0 || head_dim <= 0 || rope_dim <= 0) { + return; + } + const int vecs = rows * (q_heads + k_heads); + if (head_dim == 128 && rope_dim == 128) { + constexpr int kWarpsPerBlock = 8; + const int blocks = (vecs + kWarpsPerBlock - 1) / kWarpsPerBlock; + qk_norm_rope_bf16_warp128_kernel<<>>( + q_in, k_in, q_weight, k_weight, cos, sin, q_out, k_out, + rows, q_heads, k_heads, eps, vecs, + q_heads * head_dim, k_heads * head_dim); + return; + } + qk_norm_rope_bf16_kernel<<>>( + q_in, k_in, q_weight, k_weight, cos, sin, q_out, k_out, + rows, q_heads, k_heads, head_dim, rope_dim, eps); +} + +void cosmos3_edge_qk_norm_rope_strided_bf16( + const __nv_bfloat16* q_in, + const __nv_bfloat16* k_in, + const __nv_bfloat16* q_weight, + const __nv_bfloat16* k_weight, + const __nv_bfloat16* cos, + const __nv_bfloat16* sin, + __nv_bfloat16* q_out, + __nv_bfloat16* k_out, + int rows, + int q_heads, + int k_heads, + int q_in_row_stride, + int k_in_row_stride, + float eps, + cudaStream_t stream) +{ + if (rows <= 0 || q_heads <= 0 || k_heads <= 0) { + return; + } + const int vecs = rows * (q_heads + k_heads); + constexpr int kWarpsPerBlock = 8; + const int blocks = (vecs + kWarpsPerBlock - 1) / kWarpsPerBlock; + qk_norm_rope_bf16_warp128_kernel<<>>( + q_in, k_in, q_weight, k_weight, cos, sin, q_out, k_out, + rows, q_heads, k_heads, eps, vecs, q_in_row_stride, k_in_row_stride); +} + +void cosmos3_edge_fill_flat_velocity_bf16( + const __nv_bfloat16* action, + __nv_bfloat16* velocity, + int flat_dim, + int action_numel, + cudaStream_t stream) +{ + if (action == nullptr || velocity == nullptr || flat_dim <= 0 || action_numel <= 0 || action_numel > flat_dim) { + return; + } + constexpr int kBlock = 256; + fill_flat_velocity_bf16_kernel<<<(flat_dim + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + action, velocity, flat_dim, action_numel); +} + +void cosmos3_edge_add_bias_zero_action_tail_bf16( + __nv_bfloat16* action, + const __nv_bfloat16* bias, + int rows, + int cols, + int valid_cols, + cudaStream_t stream) +{ + if (action == nullptr || bias == nullptr || rows <= 0 || cols <= 0 || valid_cols < 0 || valid_cols > cols) { + return; + } + constexpr int kBlock = 256; + const int total = rows * cols; + add_bias_zero_action_tail_bf16_kernel<<<(total + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + action, bias, total, cols, valid_cols); +} + +void cosmos3_edge_scatter_rows_bf16( + const __nv_bfloat16* src, + __nv_bfloat16* dst, + const int64_t* row_indices, + int rows, + int hidden, + cudaStream_t stream) +{ + if (src == nullptr || dst == nullptr || row_indices == nullptr || rows <= 0 || hidden <= 0) { + return; + } + constexpr int kBlock = 256; + const int total = rows * hidden; + scatter_rows_bf16_kernel<<<(total + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + src, dst, row_indices, total, hidden); +} + +void cosmos3_edge_gather_rows_bf16( + const __nv_bfloat16* src, + __nv_bfloat16* dst, + const int64_t* row_indices, + int rows, + int hidden, + cudaStream_t stream) +{ + if (src == nullptr || dst == nullptr || row_indices == nullptr || rows <= 0 || hidden <= 0) { + return; + } + constexpr int kBlock = 256; + const int total = rows * hidden; + gather_rows_bf16_kernel<<<(total + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + src, dst, row_indices, total, hidden); +} + +void cosmos3_edge_copy_action_tail_f32_to_bf16( + const float* flat_noise, + __nv_bfloat16* action, + int flat_dim, + int action_numel, + cudaStream_t stream) +{ + if (flat_noise == nullptr || action == nullptr || flat_dim <= 0 || action_numel <= 0 || action_numel > flat_dim) { + return; + } + constexpr int kBlock = 256; + copy_action_tail_f32_to_bf16_kernel<<<(action_numel + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + flat_noise, action, flat_dim, action_numel); +} + +void cosmos3_edge_add_action_bias_timestep_bf16( + __nv_bfloat16* x, + const __nv_bfloat16* static_bias, + const __nv_bfloat16* timestep, + int rows, + int hidden, + cudaStream_t stream) +{ + if (x == nullptr || static_bias == nullptr || timestep == nullptr || rows <= 0 || hidden <= 0) { + return; + } + constexpr int kBlock = 256; + const int total = rows * hidden; + add_action_bias_timestep_bf16_kernel<<<(total + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + x, static_bias, timestep, total, hidden); +} + +void cosmos3_edge_add_bf16( + const __nv_bfloat16* a, + const __nv_bfloat16* b, + __nv_bfloat16* out, + int numel, + cudaStream_t stream) +{ + if (a == nullptr || b == nullptr || out == nullptr || numel <= 0) { + return; + } + constexpr int kBlock = 256; + add_bf16_kernel<<<(numel + kBlock - 1) / kBlock, kBlock, 0, stream>>>(a, b, out, numel); +} + +void cosmos3_edge_avgdown3d_bf16( + const __nv_bfloat16* x, + __nv_bfloat16* out, + int b, + int c, + int t, + int h, + int w, + int out_c, + int factor_t, + int factor_s, + int group_size, + cudaStream_t stream) +{ + if (x == nullptr || out == nullptr || b <= 0 || c <= 0 || t <= 0 || h <= 0 || w <= 0 || + out_c <= 0 || factor_t <= 0 || factor_s <= 0 || group_size <= 0) { + return; + } + const int factor = factor_t * factor_s * factor_s; + if (c * factor != out_c * group_size || h % factor_s != 0 || w % factor_s != 0) { + return; + } + const int pad_t = (factor_t - (t % factor_t)) % factor_t; + const int out_t = (t + pad_t) / factor_t; + const int out_h = h / factor_s; + const int out_w = w / factor_s; + const int64_t total64 = static_cast(b) * out_c * out_t * out_h * out_w; + if (total64 <= 0 || total64 > static_cast(INT_MAX)) { + return; + } + constexpr int kBlock = 256; + const int total = static_cast(total64); + avgdown3d_bf16_kernel<<<(total + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + x, out, total, c, t, h, w, out_c, out_t, out_h, out_w, + factor_t, factor_s, group_size, pad_t); +} + +void cosmos3_edge_unipc_step_f32_bf16( + const float* sample, + const __nv_bfloat16* velocity, + const float* prev_m1, + const float* prev_m2, + const float* prev_last_sample, + float* next_sample, + float* current_m, + float* current_last_sample, + int numel, + float sigma, + int corrector_order, + int predictor_order, + float c_sample, + float c_last, + float c_prev_m1, + float c_prev_m2, + float c_curr_m, + float p_sample, + float p_curr_m, + float p_prev_m1, + cudaStream_t stream) +{ + if (sample == nullptr || velocity == nullptr || next_sample == nullptr || + current_m == nullptr || current_last_sample == nullptr || numel <= 0) { + return; + } + if ((corrector_order >= 1 && (prev_m1 == nullptr || prev_last_sample == nullptr)) || + (corrector_order >= 2 && prev_m2 == nullptr) || + (predictor_order >= 2 && prev_m1 == nullptr)) { + return; + } + constexpr int kBlock = 256; + unipc_step_f32_bf16_kernel<<<(numel + kBlock - 1) / kBlock, kBlock, 0, stream>>>( + sample, velocity, prev_m1, prev_m2, prev_last_sample, next_sample, + current_m, current_last_sample, numel, sigma, corrector_order, + predictor_order, c_sample, c_last, c_prev_m1, c_prev_m2, c_curr_m, + p_sample, p_curr_m, p_prev_m1); +} + +namespace { + +__global__ void relu2_to_fp8_static_kernel( + const __nv_bfloat16* __restrict__ x, + __nv_fp8_e4m3* __restrict__ out, + const float* __restrict__ d_scale, + int n2) +{ + const float inv_scale = 1.0f / (*d_scale); + const __nv_bfloat162* x2 = reinterpret_cast(x); + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n2; i += gridDim.x * blockDim.x) { + const __nv_bfloat162 v = x2[i]; + float a = fmaxf(__bfloat162float(v.x), 0.0f); + float b = fmaxf(__bfloat162float(v.y), 0.0f); + a = fminf(a * a * inv_scale, 448.0f); + b = fminf(b * b * inv_scale, 448.0f); + out[2 * i] = __nv_fp8_e4m3(a); + out[2 * i + 1] = __nv_fp8_e4m3(b); + } +} + +} // namespace + +void cosmos3_edge_relu2_to_fp8_static_bf16( + const __nv_bfloat16* x, + __nv_fp8_e4m3* out, + const float* d_scale, + int numel, + cudaStream_t stream) +{ + if (x == nullptr || out == nullptr || d_scale == nullptr || numel <= 0 || (numel & 1) != 0) { + return; + } + constexpr int kBlock = 256; + const int n2 = numel >> 1; + const int blocks = min((n2 + kBlock - 1) / kBlock, 65535); + relu2_to_fp8_static_kernel<<>>(x, out, d_scale, n2); +} + +} // namespace flash_rt::kernels diff --git a/csrc/kernels/cosmos3_edge_misc.cuh b/csrc/kernels/cosmos3_edge_misc.cuh new file mode 100644 index 00000000..2316e52f --- /dev/null +++ b/csrc/kernels/cosmos3_edge_misc.cuh @@ -0,0 +1,143 @@ +// Cosmos3-Edge hot-path helpers. + +#pragma once + +#include +#include +#include +#include + +namespace flash_rt::kernels { + +void cosmos3_edge_qk_norm_rope_bf16( + const __nv_bfloat16* q_in, + const __nv_bfloat16* k_in, + const __nv_bfloat16* q_weight, + const __nv_bfloat16* k_weight, + const __nv_bfloat16* cos, + const __nv_bfloat16* sin, + __nv_bfloat16* q_out, + __nv_bfloat16* k_out, + int rows, + int q_heads, + int k_heads, + int head_dim, + int rope_dim, + float eps, + cudaStream_t stream); + +void cosmos3_edge_fill_flat_velocity_bf16( + const __nv_bfloat16* action, + __nv_bfloat16* velocity, + int flat_dim, + int action_numel, + cudaStream_t stream); + +void cosmos3_edge_add_bias_zero_action_tail_bf16( + __nv_bfloat16* action, + const __nv_bfloat16* bias, + int rows, + int cols, + int valid_cols, + cudaStream_t stream); + +void cosmos3_edge_scatter_rows_bf16( + const __nv_bfloat16* src, + __nv_bfloat16* dst, + const int64_t* row_indices, + int rows, + int hidden, + cudaStream_t stream); + +void cosmos3_edge_gather_rows_bf16( + const __nv_bfloat16* src, + __nv_bfloat16* dst, + const int64_t* row_indices, + int rows, + int hidden, + cudaStream_t stream); + +void cosmos3_edge_copy_action_tail_f32_to_bf16( + const float* flat_noise, + __nv_bfloat16* action, + int flat_dim, + int action_numel, + cudaStream_t stream); + +void cosmos3_edge_add_action_bias_timestep_bf16( + __nv_bfloat16* x, + const __nv_bfloat16* static_bias, + const __nv_bfloat16* timestep, + int rows, + int hidden, + cudaStream_t stream); + +void cosmos3_edge_add_bf16( + const __nv_bfloat16* a, + const __nv_bfloat16* b, + __nv_bfloat16* out, + int numel, + cudaStream_t stream); + +void cosmos3_edge_avgdown3d_bf16( + const __nv_bfloat16* x, + __nv_bfloat16* out, + int b, + int c, + int t, + int h, + int w, + int out_c, + int factor_t, + int factor_s, + int group_size, + cudaStream_t stream); + +void cosmos3_edge_unipc_step_f32_bf16( + const float* sample, + const __nv_bfloat16* velocity, + const float* prev_m1, + const float* prev_m2, + const float* prev_last_sample, + float* next_sample, + float* current_m, + float* current_last_sample, + int numel, + float sigma, + int corrector_order, + int predictor_order, + float c_sample, + float c_last, + float c_prev_m1, + float c_prev_m2, + float c_curr_m, + float p_sample, + float p_curr_m, + float p_prev_m1, + cudaStream_t stream); + +void cosmos3_edge_qk_norm_rope_strided_bf16( + const __nv_bfloat16* q_in, + const __nv_bfloat16* k_in, + const __nv_bfloat16* q_weight, + const __nv_bfloat16* k_weight, + const __nv_bfloat16* cos, + const __nv_bfloat16* sin, + __nv_bfloat16* q_out, + __nv_bfloat16* k_out, + int rows, + int q_heads, + int k_heads, + int q_in_row_stride, + int k_in_row_stride, + float eps, + cudaStream_t stream); + +void cosmos3_edge_relu2_to_fp8_static_bf16( + const __nv_bfloat16* x, + __nv_fp8_e4m3* out, + const float* d_scale, + int numel, + cudaStream_t stream); + +} // namespace flash_rt::kernels diff --git a/docs/cosmos3_edge_thor.md b/docs/cosmos3_edge_thor.md new file mode 100644 index 00000000..9856b2c5 --- /dev/null +++ b/docs/cosmos3_edge_thor.md @@ -0,0 +1,1370 @@ +# Cosmos3-Edge Thor Baseline + +`config="cosmos3_edge"` is the Thor baseline runner for NVIDIA's official +Cosmos Framework inference path. It is intentionally upstream-first: run this +baseline, capture latency/outputs, then port the FlashRT optimized denoise/action +path behind the same config. + +Current status: + +- Official Cosmos Framework baseline is runnable from FlashRT. +- P0 denoise dump is captured for `av_inverse_0`. +- P1 scheduler/latent replay scaffold is wired as `backend="replay"`. +- P1 Torch reference validates the native transformer math path through all 28 + layers for step 0 action velocity, and through the full 30-step UniPC denoise + loop as `backend="torch_ref"`. +- `backend="flashrt"` is wired to the current optimized eager denoise engine: + static AV buffers, static und K/V cache, BF16 FlashRT GEMMs/RMSNorm/relu2, + fused qk-norm+RoPE, cached timestep embeddings, native action input tail copy, + native flat velocity fill, native action bias+timestep add, native action + bias+tail clear, native action row scatter/gather, and FA4 native gen + attention. +- `backend="official_action_only"` runs the official no-dump Cosmos Framework + path, but skips inverse-dynamics vision decode/save and rewrites + `sample_outputs.json` to action-only. +- `official_action_only` also has an opt-in `live_dump_out` capture mode that + records a denoise safetensors dump from the live official run without editing + the external Cosmos Framework checkout. +- `official_action_only` has an opt-in live FlashRT handoff prototype: step 0 + runs official `_get_velocity` to capture the live boundary, then later + denoise steps short-circuit `_get_velocity` through FlashRT velocity. +- `--live-prelayer-bootstrap` is a stronger handoff mode: it captures the live + boundary before official decoder layers, aborts the official transformer, and + returns FlashRT velocity for all 30 denoise steps. +- `--live-warm-request` is the measured request path for the live handoff: it + enables pre-layer bootstrap and one official warmup request so cold FA4/engine + setup is paid before the timed batch. +- In live handoff mode, the Cosmos UniPC sampler is patched in-process to use + FlashRT's native `cosmos3_edge_unipc_step_f32_bf16` path for the fixed + single-sample Edge AV schedule when available. Set + `FLASHRT_COSMOS3_EDGE_LIVE_NATIVE_UNIPC=0` to fall back to the official + Python scheduler. + +## Quantized whole-graph denoise engine (`models/cosmos3_edge/pipeline_thor.py`) + +`CosmosEdgeThor` is the optimized denoise engine: pre-quantized gen-tower +weights, static buffers, per-layer joint und+gen K/V with the static und prefix +pre-installed, FA4 attention, native UniPC, and the entire 30-step denoise +captured in **one CUDA graph**. It reads no environment variables; precision is +selected by the caller. + +- `quant="fp8"` (default): per-tensor FP8 E4M3 weights for all six gen-tower + projections, calibrated static activation scales (one dynamic-scale denoise + pass records per-site amax ceilings via `fp8_accumulate_scale_max`, margin + 1.25), fused `residual_add_rms_norm_fp8` / `quantize_fp8_static` / + `cosmos3_edge_relu2_to_fp8_static_bf16` quant chain. +- `ffn_fp4=True` (opt-in): NVFP4 W4A4 up/down FFN GEMMs via `flash_rt_fp4` + (`cutlass_fp4_sq_fp16` + fused `cosmos3_edge_res_rms_fp4_sfa_bf16` and + `cosmos3_edge_relu2_fp4_sfa_fp16` quantizers, dynamic per-16-block scales). +- The gen-stream fused qk-norm+RoPE kernel is the warp-per-head register + variant (one warp per `(row, head)` vector, `__shfl_xor` rope-partner + exchange, no shared memory or barriers). +- The last layer runs slim at M=60: only the action rows feed the head, so + final-layer Q/attention/o-proj/FFN shrink to the 60 action tokens while K/V + still cover the full sequence (math-identical; `--no-slim-last` to disable). + +Measured on Thor (15 back-to-back iterations, hot regime, `av_inverse_0`, +official eager denoise baseline 33.14s; gate cos >= 0.999, rel_l2 < 3%, +both engines re-validated on the second-seed dump): + +| engine | denoise P50 | speedup | final action cos | rel_l2 | +|---|---|---|---|---| +| `--engine bf16-eager` (per-op path) | 11.998s | 2.76x | 0.9999959 | 0.29% | +| `--engine fp8` | 5.792s | 5.72x | 0.9999834 | 0.59% | +| `--engine fp8 --ffn-fp4` | 5.456s | 6.07x | 0.9997067 | 2.42% | + +```bash +python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 --iters 15 \ + --warmup-steps 30 --enforce-gates # FP8 default +python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 --ffn-fp4 \ + --iters 15 --warmup-steps 30 --enforce-gates # max-quantization variant +``` + +FP8 is the recommended default (large precision margin); `--ffn-fp4` trades +~6% latency for a still-passing but thinner gate margin. Measured negatives, +kept default-off or not landed: a fused QKV wide GEMM (cuBLASLt N=4096 tactic +regression, behind `qkv_fused=`) and SageAttention2 int8 gen attention (the +SM80-era int8 mma core is both numerically wrong and slower than FA4 on +SM110). The largest remaining denoise levers are a relu2 epilogue fused into +the up GEMM and a Thor-native int8/fp8 attention kernel. + +## Quickstart + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_flashrt \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework +``` + +The example must run inside the same container/venv used for the official +baseline, with `cosmos-framework` importable. `--vae-path` is optional only when +the container can download `Wan2.2_VAE.pth` itself. + +To validate the no-dump action-only official path without decoded vision output: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_action_only_official \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only +``` + +To capture a live official denoise dump while keeping the action-only output: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_dump \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-dump-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_denoise.safetensors +``` + +This capture mode is for boundary validation, not baseline timing: enabling the +official velocity postprocess hook makes Cosmos run an extra uncond velocity +branch even when guidance is 1.0. + +To exercise the live FlashRT handoff prototype: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_handoff \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-flashrt-handoff \ + --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_handoff_boundary.safetensors +``` + +This is a correctness/architecture prototype. It still pays official step 0, +FlashRT engine initialization, and cold FA4 setup inside +`generate_samples_from_batch`; the current archived run reports 26.10s +`generate_batch`. `--live-boundary-out` is only for debug artifact archival; the +runtime can construct the boundary in memory, which measured 26.20s and the same +action metrics. + +To skip the official step-0 transformer as well: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_trace.json +``` + +The archived pre-layer run proves the official decoder was aborted +(`prelayer_aborted=true`) and returns action within the gate, but one-shot +latency is still dominated by FA4 first-use compilation: first FlashRT velocity +8.50s, steady FlashRT velocity 0.398s/step. + +To archive the smaller native-denoise/VLM input contract at the pre-layer +boundary: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_boundary \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary_upstream_trace.json +``` + +Archived `official_action_only_live_prelayer_boundary.safetensors` metrics: + +- 60.36 MiB file, 31 tensors / 60.36 MiB tensor footprint +- contains `lm_in/full_only_seq`, layer-0 input, RoPE, VFM tokens, action + tokens, timestep, and `noise_x` +- does not store layer-0 output or step-0 velocity; the older step-0 handoff + boundary was 87.77 MiB / 36 tensors because it archived those outputs +- run action gate: cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 +- `generate_batch`: 26.70s; native UniPC scheduler: 30 steps + +This is the input contract for a future native VLM/prefill-to-denoise engine. +It still relies on official upstream code to synthesize `lm_in/full_only_seq`, +RoPE, VFM tokens, and action tokens from the raw request. + +To replay that pre-layer boundary as the live denoise engine input: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_boundary_in \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-boundary-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_upstream_trace.json +``` + +Archived `official_action_only_live_boundary_in_outputs` metrics: + +- `OmniInference.generate_batch`: 26.03s +- official velocity calls: 0; FlashRT velocity calls: 30 +- boundary load/validate: 0.8ms; engine construct: 1.29s +- native UniPC scheduler: 1 run / 30 steps, scheduler step total 0.0071s +- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 + +This is still artifact-backed. Its purpose is to make the native prefill target +executable: once native VAE/SigLIP/VLM code can synthesize the same boundary +tensors live, the denoise engine no longer needs official transformer capture. + +The boundary-in path can also be combined with the current warm repeated-request +prepare cache: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_boundary_in_warm_prepare_cache \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-boundary-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors \ + --live-warm-request \ + --cache-warmup-prepare \ + --prepare-slim-no-raw-state-vision \ + --prepare-slim-derive-condition-reference \ + --prepare-slim-derive-initial-noise \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_warm_request_prepare_cache_slim_derive_noise_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_warm_request_prepare_cache_slim_derive_noise_upstream_trace.json +``` + +Archived metrics for this service-style path: + +- measured `OmniInference.generate_batch`: 12.08s; denoise: 12.07s +- warmup `OmniInference.generate_batch`: 26.08s +- official velocity calls: 0; FlashRT velocity calls: 60 +- measured `_prepare_inference_data`: 0.042s; prepare hit/store: 1/1 +- measured VAE encode calls: 0 +- native UniPC scheduler: 2 runs / 60 steps, scheduler step total 0.0134s +- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 + +The helper `_derive_step0_vfm_boundary_from_prepare_payload` formalizes the +part of the pre-layer boundary that is already implied by the slim prepare +contract. From `official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt` +and seed 0 plus `embed_tokens.weight` it derives 28 step-0 +VFM/noise/pack/position/RoPE/text-causal tensors exactly: + +- `steps/00/noise_x` +- VFM vision tokens, condition mask, sequence indexes, and empty no-loss/noisy + metadata +- VFM action noisy tokens, condition mask, sequence/timestep/loss/noisy indexes, + domain id, and raw action dim +- fixed packed metadata: causal/full indices, offsets, sample offsets, and + `position_ids` +- layer-0 RoPE cos/sin tensors for causal and full-only sequences +- causal LM hidden state from `cond_text_tokens + eos + <|vision_start|>` and + `embed_tokens.weight` + +Verifier metrics: 28 tensors / 10.65 MiB derived exactly. The layer-0 input +tensors are checked as aliases of `lm_in` (2 tensors / 25.10 MiB). The remaining +1 tensor / 24.61 MiB is the full-only LM packed hidden state. The verifier now +probes that tensor too and requires all 6300 rows to be bit-exact; the action +rows require the reference action projection/timestep path to run with CUDA +matmul TF32 disabled. This narrows the remaining native VLM/prefill task to +synthesizing the same prepare/prefill state from live raw request data rather +than resolving an unresolved packed-hidden-state formula gap. + +To run the live denoise engine from the 9.18 MiB slim prepare artifact instead +of the 60.36 MiB pre-layer boundary artifact: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_prepare_boundary_in \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --prepare-replay-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ + --live-boundary-prepare-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ + --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_boundary_in_derived_boundary.safetensors \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_boundary_in_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_boundary_in_upstream_trace.json +``` + +Archived `official_action_only_live_prepare_boundary_in_outputs` metrics: + +- input artifact: 9.18 MiB slim-derived-noise prepare payload +- derived debug boundary: 62.65 MiB / 31 tensors; larger than the 60.36 MiB + pre-layer artifact only because the debug save clones shared tensor views +- boundary derive: 0.258s; engine construct: 1.074s +- `OmniInference.generate_batch`: 22.24s; denoise: 22.22s; decode: 0.0084s +- official velocity calls: 0; FlashRT velocity calls: 30 +- measured `_prepare_inference_data`: 0.101s; measured VAE encode calls: 0 +- native UniPC scheduler: 1 run / 30 steps +- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 + +This removes the 60.36 MiB executable boundary artifact from the runtime input +contract for replay-style testing. It is still artifact-backed prepare replay: +raw-video no-artifact live inference still needs native VAE/SigLIP/VLM prefill. + +To run the same executable-boundary synthesis from this request's in-memory +prepared state, without reading a prepare artifact: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-boundary-prepare-live \ + --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_derived_boundary.safetensors \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_upstream_trace.json \ + --prepare-slim-no-raw-state-vision \ + --prepare-slim-derive-condition-reference \ + --prepare-slim-derive-initial-noise +``` + +Archived `official_action_only_live_prepare_live_outputs` metrics: + +- runtime source: in-memory `_prepare_inference_data` payload, not a prepare or + pre-layer artifact +- derived debug boundary: 62.65 MiB / 31 tensors; debug save clones shared views +- boundary derive: 0.143s; engine construct: 1.267s +- `OmniInference.generate_batch`: 25.93s; denoise: 25.91s; decode: 0.0088s +- official velocity calls: 0; FlashRT velocity calls: 30 +- measured `_prepare_inference_data`: 3.88s; measured VAE encode: 3.66s; text tokenization: 30.8ms; pack: 2.65ms +- native UniPC scheduler: 1 run / 30 steps +- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 + +This is the first no-prepare-artifact live denoise handoff. It still uses the +official raw-video prepare path, so strict native E2E still requires native +VAE/SigLIP/VLM prefill to produce the same prepared state. + +The same bridge can be combined with upstream Wan2.2 AOT encode using +`--vae-compile-encode`. This is not a FlashRT native VAE implementation, but it +is the current no-prepare-artifact service-warmup target line for native VAE +work: + +- `official_action_only_live_prepare_live_vae_compile_encode_outputs` +- AOT setup compile/load: 149.39s; loaded AOT functions: 5 +- measured `OmniInference.generate_batch`: 24.33s; denoise: 24.32s; decode: + 0.0085s +- measured `_prepare_inference_data`: 2.65s; measured VAE encode: 2.43s +- boundary derive: 0.146s; engine construct: 1.714s +- official velocity calls: 0; FlashRT velocity calls: 30; native UniPC 30 steps +- final action cos 0.9999948 / rel_l2 0.322% / max_abs 0.00802 + +For the current no-prepare/pre-layer-artifact service floor, combine the live +prepare bridge with one warmup request plus the VAE and prepare caches, and keep +`--live-boundary-out` unset: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_warm_cache_nosave_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --warmup 1 \ + --cache-warmup-vae \ + --cache-warmup-prepare \ + --live-boundary-prepare-live \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_warm_cache_nosave_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_warm_cache_nosave_upstream_trace.json \ + --prepare-slim-no-raw-state-vision \ + --prepare-slim-derive-condition-reference \ + --prepare-slim-derive-initial-noise +``` + +Archived `official_action_only_live_prepare_live_warm_cache_nosave_outputs` +metrics: + +- measured `OmniInference.generate_batch`: 13.33s; denoise: 13.32s; decode: + 0.00020s +- warmup `OmniInference.generate_batch`: 25.77s +- measured `_prepare_inference_data`: 0.075s; measured VAE encode calls: 0 +- official velocity calls: 0; FlashRT velocity calls: 60; native UniPC 60 steps +- measured native UniPC run: 13.23s +- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 + +A debug-save variant with `--live-boundary-out` measured 22.04s because the +archival path perturbs the service run, so production/service floor measurements +should leave that flag off. + +The `--live-warm-request` path moves cold FA4/engine setup and repeat VAE encode +out of the measured batch for the same sample by enabling pre-layer bootstrap, +one official warmup request, and a FlashRT warmup VAE latent cache. FlashRT +patches warmup to clone `data_batch` first, because upstream mutates video +tensors in-place during preprocessing. + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_warm_request_vae_cache \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-warm-request \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_vae_cache_handoff_trace.json +``` + +Archived `official_action_only_live_warm_request_vae_cache_outputs` metrics: + +- `OmniInference.generate_batch`: 12.12s +- `OmniMoTModel.generate_samples_from_batch`: 12.11s +- `OmniMoTModel.decode`: 0.00012s +- live native UniPC scheduler: 2 runs / 60 steps, native scheduler step total + 0.0128s +- final action cos 0.9999958 / rel_l2 0.291% / max_abs 0.00687 + +The warmup batch still pays cold cost: `[warmup] OmniInference.generate_batch` +25.95s. The VAE cache key includes shape, dtype, device, element count, and a +small sampled content fingerprint so same-shape different-video tensors do not +reuse the warmup latent silently. + +For P4 upstream breakdown, add `--upstream-trace-out`: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_warm_request_vae_cache \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-warm-request \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_vae_cache_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_vae_cache_upstream_trace.json +``` + +Archived uncached upstream trace metrics: + +- measured `OmniMoTModel.encode`: 3.41s +- measured `OmniMoTModel.get_data_and_condition`: 3.49s +- measured `OmniMoTModel._prepare_inference_data`: 3.52s +- input batch preparation: 0.57s total +- warmup official `Cosmos3VFMNetwork.forward`: 9.79s + +With warmup VAE cache enabled by `--live-warm-request`, the measured upstream +trace changes to: + +- measured `OmniMoTModel.encode`: 0.00051s +- measured `OmniMoTModel._prepare_inference_data`: 0.1068s +- measured cache hits: 1, warmup cache stores: 1, measured cache misses: 0 +- warmup `OmniMoTModel.encode`: 3.65s +- warmup prelayer handoff trace events are marked as expected/ok + +For the no-warmup VAE replacement boundary, dump the measured VAE encode latent +and replay it: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_vae_boundary \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --vae-encode-dump-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_encode.safetensors \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_boundary_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_boundary_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_vae_latent_replay \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --vae-latent-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_encode.safetensors \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_latent_replay_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_latent_replay_upstream_trace.json +``` + +Archived VAE boundary/replay metrics: + +- boundary dump: `official_action_only_live_prelayer_vae_encode.safetensors`, + 4.57 MiB, tensor `vae_encode/output` shape `[1,48,16,30,52]`, plus input + signature metadata +- boundary run measured `OmniMoTModel.encode`: 3.72s; + `_prepare_inference_data`: 3.86s; `generate_batch`: 26.22s +- latent replay measured `OmniMoTModel.encode`: 0.0108s; + `vae_latent_replay_hit`: 0.0050s; `_prepare_inference_data`: 0.141s; + `generate_batch`: 22.55s +- latent replay final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 + +This is not the final no-dump VAE path. It is the correctness and latency +boundary a native Wan2.2 VAE encoder must match to remove the remaining +no-warmup official encode. + +For a wider no-warmup prepare/prefill replacement boundary, dump and replay the +full `_prepare_inference_data` return tuple: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_boundary_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare.pt \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_boundary_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_boundary_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_replay_v2_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare.pt \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_replay_v2_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_replay_v2_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory_replay_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory.json \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_inventory.json \ + --prepare-slim-no-raw-state-vision \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_inventory.json \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_inventory.json \ + --prepare-slim-no-raw-state-vision \ + --prepare-slim-derive-condition-reference \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_inventory.json \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_inventory.json \ + --prepare-slim-no-raw-state-vision \ + --prepare-slim-derive-condition-reference \ + --prepare-slim-derive-initial-noise \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_upstream_trace.json + +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_outputs \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ + --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_inventory.json \ + --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_handoff_trace.json \ + --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_upstream_trace.json +``` + +Archived prepare boundary/replay metrics: + +- prepare artifact: `official_action_only_live_prelayer_prepare.pt`, 297 MiB, + with input signature and the 7-tuple consumed by `generate_samples_from_batch` +- dump run measured `_prepare_inference_data`: 4.40s; `prepare_boundary_dump` + event: 3.87s +- replay measured `_prepare_inference_data`: 0.232s; `prepare_replay_hit`: + 0.182s; `generate_batch`: 22.66s +- replay final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 +- inventory replay writes + `official_action_only_live_prelayer_prepare_inventory.json` using schema + `flashrt_cosmos3_edge_prepare_inventory_v1`: 11 tensors, 297.13 MiB total. + The largest tensor is `gen_data_clean.raw_state_vision[0]` + `[1,3,61,480,832]` at 278.79 MiB; the denoise seed/reference/mask vectors + are `initial_noise[0]`, `condition_reference[0]`, and `condition_mask[0]`, + each `[1201920]` float32 at 4.58 MiB. This is the current machine-readable + native prefill contract. +- inventory replay measured `_prepare_inference_data`: 0.207s; + `prepare_replay_hit`: 0.159s; `generate_batch`: 22.75s; final action + cos 0.9999958 / rel_l2 0.291% / max_abs 0.00687. +- `--prepare-slim-no-raw-state-vision` proves the post-prepare denoise path + does not need the 278.79 MiB RGB `raw_state_vision` tensor once + `temporal_positions_vision` and `x0_tokens_vision` are present. The archived + slim dump is 18.35 MiB, 10 tensors / 18.34 MiB total, with `x0_tokens_vision`, + action state, and the three 4.58 MiB noise/reference/mask vectors intact. + The slim dump run keeps action cos 0.9999958 / rel_l2 0.291%; replaying the + slim `.pt` gives `prepare_replay_hit` 0.026s, `_prepare_inference_data` + 0.075s, `generate_batch` 22.62s, and unchanged action metrics. +- `--prepare-slim-derive-condition-reference` proves `condition_reference` is + exactly derivable from `gen_data_clean.x0_tokens_vision/action` and + `raw_action_dim` after raw video has been dropped. The derived slim artifact + stores `condition_reference=None`, is 13.76 MiB, and has 9 tensors / + 13.75 MiB. Replay restores `condition_reference` before entering the official + sampler, so the runtime payload inventory returns to 10 tensors / 18.34 MiB. + The archived replay has `prepare_replay_hit` 0.026s, + `_prepare_inference_data` 0.074s, `generate_batch` 22.56s, final action cos + 0.9999959 / rel_l2 0.291% / max_abs 0.00687, and native UniPC 1 run / + 30 steps. +- `--prepare-slim-derive-initial-noise` further proves `initial_noise` is + bit-exactly derivable from the request seed, `x0_tokens_vision/action`, + `condition_mask`, and `raw_action_dim`. The slim-derived-noise artifact stores + `raw_state_vision=None`, `condition_reference=None`, and `initial_noise=None`; + it is 9.18 MiB with 8 tensors / 9.17 MiB, consisting of the prepared latent + state plus `condition_mask`. Replay restores both `initial_noise` and + `condition_reference`, so runtime inventory returns to 10 tensors / + 18.34 MiB. The archived replay has `prepare_replay_hit` 0.051s, + `_prepare_inference_data` 0.100s, `generate_batch` 22.59s, final action cos + 0.9999959 / rel_l2 0.291% / max_abs 0.00687, and native UniPC 1 run / + 30 steps. + +This is a full prepare/prefill correctness boundary for future native +VAE/SigLIP/VLM-prefill replacement. It is not a production path: the artifact is +input-specific and still loaded via `torch.save`/`torch.load`. The derived slim +result shows the native denoise contract can be reduced to latent/action state, +temporal positions, condition mask, request seed/signature metadata, and enough +metadata to reconstruct condition references and initial noise. + +For the first native Wan2.2 VAE encode subgraph, enable native BF16 +`RMS_norm -> SiLU` inside encode `ResidualBlock`s: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_vae_native_rms_silu \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-prelayer-bootstrap \ + --vae-native-rms-silu \ + --vae-encode-profile-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_native_rms_silu_profile.json \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_native_rms_silu_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_native_rms_silu_upstream_trace.json +``` + +Archived native RMS+SiLU metrics: + +- profile inventory artifact: + `official_action_only_live_prelayer_vae_encode_profile.json` shows one + `WanVAE_.encode` call from BF16 `[1,3,61,480,832]` to BF16 + `[1,48,16,30,52]`, 27 `CausalConv3d`, 22 `RMS_norm`, 21 `SiLU`, and + top time in the early high-resolution 160-channel residual convs. +- baseline profile-only `WanVAE_.encode`: 3.778s; native RMS+SiLU profile + `WanVAE_.encode`: 3.252s. The refreshed profile artifacts now include + `shape_summary` and `native_candidate_summary`; the top native candidate is + `encoder.downsamples.0.downsamples.0.residual.2` as + `steady_cached_causal_conv3d` on `[1,160,24,240,416]`. +- upstream measured `OmniMoTModel.encode`: 4.197s -> 3.295s; measured + `_prepare_inference_data`: 4.348s -> 3.436s. +- native RMS+SiLU run `generate_batch`: 25.63s; decode 0.0088s; live native + UniPC scheduler 1 run / 30 steps; final action cos 0.9999955 / + rel_l2 0.299% / max_abs 0.00929. + +The remaining VAE bottleneck is still the early CausalConv3d chain: +`encoder.downsamples.0` and `encoder.downsamples.1` dominate the profile. The +next native VAE step should fuse/quantize the `RMS_norm -> SiLU -> CausalConv3d` +triplet for those fixed shapes, reusing the Motus VAE FP8/FP4 conv swap +structure but recalibrated against the Cosmos3-Edge action gate. + +An additional no-go probe rewrites no-cache single-frame Wan2.2 +`CausalConv3d` calls to mathematically equivalent `Conv2d` via +`--vae-t1-conv2d`. This is correct but slower on Thor when combined with native +RMS+SiLU: + +- `WanVAE_.encode`: 3.291s -> 3.507s +- measured `OmniMoTModel.encode`: 3.295s -> 3.511s +- `OmniInference.generate_batch`: 25.63s -> 25.92s +- final action cos 0.9999960 / rel_l2 0.282% / max_abs 0.00799 + +The T=1 Conv2d rewrite is therefore kept default-off as an opt-in probe. The +steady `T=24` cached chunks are still the important target. + +Another opt-in native VAE probe, `--vae-native-avgdown3d`, replaces Wan2.2 +`AvgDown3D` shortcut pooling with `cosmos3_edge_avgdown3d_bf16`. The kernel +canary is bit-exact against the upstream reshape/permute/group-mean semantics, +including the temporal front-pad case. In the real no-warmup pre-layer run, +local `AvgDown3D` profile time moves from about 77ms to 67ms, but the complete +path is not better when combined with native RMS+SiLU: + +- `WanVAE_.encode`: 3.252s -> 3.504s +- measured `OmniMoTModel.encode`: 3.295s -> 3.508s +- `OmniInference.generate_batch`: 25.63s -> 26.03s +- final action cos 0.9999955 / rel_l2 0.299% / max_abs 0.00929 + +This is kept as an opt-in correctness scaffold and default-off/no-go. The +worthwhile native VAE target remains the high-resolution CausalConv3d chain. + +The steady cached `T=24` Conv3d-to-Conv2d-slices decomposition has also been +probed via `dev_log_cosmos3_thor/probe_vae_conv2d_decompose.py`. It is local +math-equivalent enough for the action gate, but not a useful default path: + +- stage0 160-channel `[1,160,24,240,416]`: cos 0.999995 / + rel_l2 0.303%, but 0.185s vs BF16 cuDNN Conv3d 0.112s. +- stage1 320-channel `[1,320,24,120,208]`: cos 0.999995 / + rel_l2 0.311%, but 0.099s vs BF16 cuDNN Conv3d 0.086s. +- stage2 640-channel `[1,640,12,60,104]`: cos 0.999996 / + rel_l2 0.293%, 0.0294s vs 0.0296s, only ~1.01x and not the dominant site. + +This keeps the decomposition default-off too. The remaining viable path is a +single Thor-specific CausalConv3d kernel for the early steady chunks, not a +multi-call Torch/cuDNN decomposition. + +A first Thor-specific BF16 MMA CausalConv3d bring-up kernel is also wired as +`bf16_conv3d_v0_ndhwc_bf16out` for `sm_110a`. It mirrors the Motus v17 virtual +cache concat/direct causal-output structure, but uses BF16 tensor cores and +stays probe-only. The small-shape canary passes, and the real VAE shape probe +is accurate enough for the local gate, but still slower than cuDNN: + +- stage0 160-channel `[1,160,24,240,416]`: cos 0.999996 / + rel_l2 0.282%, kernel-only 0.275s vs BF16 cuDNN 0.112s. +- stage1 320-channel `[1,320,24,120,208]`: cos 0.999996 / + rel_l2 0.284%, kernel-only 0.214s vs BF16 cuDNN 0.087s. +- stage2 640-channel `[1,640,12,60,104]`: cos 0.999996 / + rel_l2 0.301%, kernel-only 0.089s vs BF16 cuDNN 0.029s. + +The v0 kernel is therefore default-off/no-go. It is useful as a native +correctness scaffold, but the next attempt needs better tiling/persistence or +triplet fusion, not this one-kernel v0 swap. + +The first steady cached chunk also has a cache-shape corner case: upstream +passes a one-frame cache into `CausalConv3d.forward`, which then prepends one +temporal zero via `F.pad`. A cache2 normalization probe confirms that replacing +that with an explicit `[zero, cache]` two-frame cache is bit-exact, but not a +useful cuDNN optimization by itself: + +- stage0 160-channel: prebuilt cache2 Conv3d 0.11231s vs 0.11247s reference; + building cache2 in Python plus Conv3d is 0.11308s. +- stage1 320-channel: prebuilt cache2 Conv3d 0.08665s vs 0.08669s reference; + building cache2 in Python plus Conv3d is 0.08697s. +- stage2 640-channel: prebuilt cache2 Conv3d 0.02971s vs 0.02960s reference; + building cache2 in Python plus Conv3d is 0.02976s. +- native `update_cache2_ncdhw_bf16` matches Python cache2 exactly, but still + adds 0.30-1.19ms per call at these shapes. + +This remains a native-kernel interface fact, not a standalone default path. +Cache2 is only worth standardizing if a future fused CausalConv3d/triplet +kernel consumes it directly and wins the full-site A/B. + +The `vae_channels_last3d_probe.json` artifact tests cuDNN BF16 Conv3d with +`channels_last_3d` tensors at the same three steady sites. It is bit-exact and +shows a useful preconverted-layout signal, but not a default-ready per-conv +rewrite: + +- stage0 160-channel: total conversion path 0.119s vs 0.112s reference + (0.95x), preconverted-layout conv 0.073s (1.55x). +- stage1 320-channel: total conversion path 0.076s vs 0.086s reference + (1.13x), preconverted-layout conv 0.053s (1.62x). +- stage2 640-channel: total conversion path 0.031s vs 0.030s reference + (0.95x), preconverted-layout conv 0.026s (1.16x). + +Conclusion: `channels_last_3d` is worth pursuing only if the VAE encode segment +can keep activations and weights in that layout across adjacent residual blocks +or if a fused native triplet consumes the layout directly. Per-conv conversion +stays default-off. + +An opt-in full-run patch, `--vae-channels-last3d-conv320`, applies this only to +steady cached 320-channel CausalConv3d sites and returns contiguous tensors to +preserve the upstream cache contract. Combined with `--vae-native-rms-silu`, it +is correct but slower in the real no-warmup pre-layer run: measured +`OmniMoTModel.encode` 4.04s vs 3.29s for native RMS+SiLU alone, +`_prepare_inference_data` 4.18s, `generate_batch` 26.41s, final action cos +0.9999955 / rel_l2 0.299% / max_abs 0.00929. The profile shows the 320-channel +cached sites become the top regression at 0.13-0.24s each, so the patch remains +default-off/no-go. + +The follow-up `vae_resblock_channels_last3d_probe.json` artifact tests a +block-local layout propagation variant across three real steady +`ResidualBlock`s. It keeps the second half of each block in `channels_last_3d` +instead of returning to contiguous after every conv. The probe remains correct, +but all three blocks are slower than the official block: + +- stage0 160-channel block: 0.441s vs 0.323s reference (0.73x). +- stage1 320-channel block: 0.259s vs 0.222s reference (0.85x). +- stage2 640-channel block: 0.077s vs 0.072s reference (0.93x). + +This closes the small-step `channels_last_3d` path: block-local propagation is +not enough. Future VAE work should use a fused RMS/SiLU/CausalConv3d triplet or +a stronger Thor-specific steady CausalConv3d kernel that consumes a stable +layout directly. + +For a non-native service-warmup comparison, `--vae-compile-encode` enables the +upstream Wan2.2 chunk-level AOT `compile_encode` path for 480p 16:9. The +archived run compiled 5 loadable AOT functions in 133.7s during setup, then +reduced measured no-warmup pre-layer VAE encode: + +- measured `OmniMoTModel.encode`: 4.197s eager baseline -> 2.428s AOT +- measured `_prepare_inference_data`: 2.554s +- `OmniInference.generate_batch`: 24.43s, decode 0.0085s +- final action cos 0.9999958 / rel_l2 0.293% / max_abs 0.00553 + +This is useful as an opt-in long-lived-service baseline and a latency target for +native VAE work. It is not the final FlashRT native implementation because it +pays large startup compile cost and still relies on Torch Inductor/AOT. + +The direct Motus FP8 conv3d reuse path has been probed on Thor: + +```bash +python dev_log_cosmos3_thor/probe_vae_triplet_fp8.py \ + --json-out dev_scratch_cosmos3_thor/edge_av_inverse_0/vae_triplet_fp8_probe_site0.json +``` + +The SM120-origin v17/v18 FP8 conv3d kernels compile for `sm_110a`; CMake now +exposes `fp8_conv3d_v17_ndhwc_bf16out`, +`fp8_conv3d_v17_anyco_ndhwc_bf16out`, and +`fp8_conv3d_v18_ncdhw_res_bf16out` on Thor. The small-shape kernel canary +matches torch exactly, but the real Cosmos VAE triplet A/B is negative: + +- stage0 160-channel `[1,160,24,240,416]`: FP8 v18 cos 0.999868 / + rel_l2 1.69%, but 0.245s vs BF16 cuDNN 0.113s. +- stage1 320-channel `[1,320,24,120,208]`: FP8 v18 cos 0.999816 / + rel_l2 1.92%, but 0.182s vs BF16 cuDNN 0.086s. +- stage2 640-channel `[1,640,12,60,104]`: FP8 v18 cos 0.999931 / + rel_l2 1.38%, but 0.082s vs BF16 cuDNN 0.030s. + +Therefore the Motus v17/v18 FP8 conv3d path is kept default-off for +Cosmos3-Edge. The first Thor-specific BF16 v0 kernel above is also slower, so +the next VAE conv step needs a better Thor-specific BF16/FP8 CausalConv3d +tiling/fusion strategy, not a direct Motus conv swap or the v0 BF16 probe. + +To also reuse the whole prepared inference state from the warmup request: + +```bash +python examples/cosmos3_edge_thor_baseline.py \ + --checkpoint /work/models/Cosmos3-Edge \ + --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ + --output-dir /work/.tmp_cosmos_edge_outputs_live_warm_request_prepare_cache \ + --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ + --cosmos-root /work/external/cosmos-framework \ + --backend official_action_only \ + --live-warm-request \ + --cache-warmup-prepare \ + --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_prepare_cache_handoff_trace.json \ + --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_prepare_cache_upstream_trace.json +``` + +Archived `official_action_only_live_warm_request_prepare_cache_outputs` metrics: + +- `OmniInference.generate_batch`: 11.98s +- `OmniMoTModel.generate_samples_from_batch`: 11.97s +- `OmniMoTModel.decode`: 0.00015s +- measured `OmniMoTModel._prepare_inference_data`: 0.0185s +- measured prepare cache hits: 1, stores: 1, misses: 0 +- measured VAE encode calls: 0 +- live native UniPC scheduler: 2 runs / 60 steps, native scheduler step total + 0.0133s +- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 +- Adding `--prepare-slim-no-raw-state-vision` slims the cached prepare payload + by dropping raw RGB state after warmup prepare. Archived slim cache metrics: + `generate_batch` 12.08s, denoise 12.07s, measured prepare 0.0146s, prepare + hit/store 1/1 with `slim_no_raw_state_vision=true`, measured VAE encode calls + 0, native UniPC 2 runs / 60 steps, final action cos 0.9999958 / rel_l2 0.291%. +- Adding `--prepare-slim-derive-condition-reference` on top stores the warmup + prepare payload without raw RGB and without `condition_reference`; the + measured cache hit derives `condition_reference` before the official sampler. + Archived derived slim cache metrics: `generate_batch` 12.00s, denoise 11.99s, + measured prepare 0.0149s, prepare hit/store 1/1 with both slim flags recorded, + measured VAE encode calls 0, native UniPC 2 runs / 60 steps, final action cos + 0.9999959 / rel_l2 0.291%. +- Adding `--prepare-slim-derive-initial-noise` on top stores the repeated-request + prepare payload without raw RGB, stored `condition_reference`, or stored + `initial_noise`; the measured cache hit derives noise from seed/x0/mask and + then derives the reference. Archived slim-derived-noise cache metrics: + `generate_batch` 12.05s, denoise 12.04s, measured prepare 0.0382s, prepare + hit/store 1/1 with all three slim flags recorded, measured VAE encode calls 0, + native UniPC 2 runs / 60 steps, final action cos 0.9999959 / rel_l2 0.291%. + +## P1 Replay Scaffold + +The replay scaffold validates the denoise boundary without rerunning the +transformer. It reads the P0 dump, replays the recorded velocity tensors through +FlashRT's local UniPC scheduler, and splits the final flat latent into: + +- vision latent: `[1,48,16,30,52]` +- action model latent: `[60,64]` + +Run it inside the container/venv: + +```bash +python - <<'PY' +from flash_rt.frontends.torch.cosmos3_edge_thor import Cosmos3EdgeTorchFrontendThor + +pipe = Cosmos3EdgeTorchFrontendThor("/work/models/Cosmos3-Edge", hardware="thor") +result = pipe.infer( + backend="replay", + output_dir="/work/.tmp_cosmos_edge_replay", + reference_dump=( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/tensors.safetensors" + ), + replay_device="cuda", +) +print(result) +PY +``` + +Expected validation facts for the current dump: + +- `num_steps`: 30 +- `flat_dim`: 1201920 +- `timesteps`: `999 ... 256` +- CUDA replay `max_input_abs_diff`: 0.0 + +## P1 Torch Reference + +The native bring-up path now has two local correctness gates: + +- `EdgeLayer0TorchReference`: layer 0 causal/full outputs match the official + boundary dump exactly after correcting the public checkpoint tower mapping. +- `EdgeTransformerTorchReference`: all 28 layers plus the action head match the + step-0 official action velocity. +- `EdgeDenoiseTorchReference`: the full 30-step UniPC loop computes action + velocity live from the Torch reference instead of replaying recorded velocity. +- `EdgeDenoiseFlashRT`: the public optimized eager engine used by + `backend="flashrt"`. The velocity path is the static FlashRT engine rather + than the two-tower reference replay path, and the fixed UniPC latent update + uses the native scheduler step when available. The frontend writes an + action-only `sample_outputs.json` for this backend and intentionally does not + emit decoded vision files. +- `EdgeStaticBufferEngine`: owns the static AV geometry, static vision tokens, + static und K/V cache, action slots, and the flat velocity contract. Its action + in/out projections use `GemmRunner.bf16_nn` when `flash_rt_kernels` is + available. The action input bias and modality embedding are precombined once + during static engine initialization. +- `EdgeTransformerFvkLinearReference`: routes cached-path BF16 linear projections + through `GemmRunner.bf16_nn` and cached-path RMSNorm through `fvk.rms_norm`; + cached-path main q/k RoPE through `qwen36_partial_rope_qk_bf16`, gen q/k + norm+RoPE through `cosmos3_edge_qk_norm_rope_bf16`, and FFN activation through + `relu2_inplace_bf16`. On Thor, non-causal gen attention uses the vendored FA4 + CuTe path when the `thor-fa4` extra is installed. `flash_rt_fa2` is not + expected on Thor. +- Hot-path weights/transposes are resident in the static engine cache. The + current checkpoint keeps 532 bf16 tensors, 4 fp32 timestep tensors, 4 boundary + tensors, and 336 transposed GEMM weights resident. +- `EdgeStaticBufferEngine.flat_velocity_for_step` uses + `cosmos3_edge_fill_flat_velocity_bf16` when the binding is available. This + replaces the Torch `zero_()` plus action-tail assignment for the flat velocity + output buffer. +- `EdgeStaticBufferEngine._decode_action_velocity` uses + `cosmos3_edge_add_bias_zero_action_tail_bf16` when the binding is available. + This fuses action output bias add and invalid action-dimension clearing. +- `EdgeStaticBufferEngine` uses `cosmos3_edge_scatter_rows_bf16` and + `cosmos3_edge_gather_rows_bf16` when available for action token row movement + between the compact action buffer and full sequence buffer. +- `EdgeStaticBufferEngine.full_sequence_for_step` uses + `cosmos3_edge_copy_action_tail_f32_to_bf16` when available for CUDA float32 + latents, avoiding the Torch action-tail slice and BF16 cast. +- `EdgeStaticBufferEngine._encode_action_tokens` uses + `cosmos3_edge_add_action_bias_timestep_bf16` when available. The timestep + embedding is computed once for the scalar timestep and broadcast natively + across the 60 action rows. +- `EdgeDenoiseFlashRT` precomputes all dump timesteps into + `EdgeStaticBufferEngine.timestep_embed_cache`, so measured denoise steps reuse + resident BF16 timestep embeddings instead of recomputing the Torch + cos/sin/MLP path. + +Install the FA4 dependency set in the isolated container venv before measuring +the optimized eager path: + +```bash +cd /work/official/flashrt-public +source /work/.venv_cosmos_thor/bin/activate +python -m pip install -e ".[thor-fa4]" +``` + +Run the full reference denoise check inside the container/venv: + +```bash +python - <<'PY' +from flash_rt.frontends.torch.cosmos3_edge_thor import Cosmos3EdgeTorchFrontendThor + +pipe = Cosmos3EdgeTorchFrontendThor("/work/models/Cosmos3-Edge", hardware="thor") +result = pipe.infer( + backend="flashrt", + output_dir="/work/.tmp_cosmos_edge_flashrt", + reference_dump=( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/tensors.safetensors" + ), + boundary_dump=( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + replay_device="cuda", +) +print(result) +PY +``` + +Current 30-step optimized eager result with FA4 attention: + +- und-cache precompute/init: 1.44s +- 30-step run wall time after one warm step: P50 11.998s over 10 iters +- speedup vs official `OmniMoTModel.generate_samples_from_batch`: 2.762x +- `max_input_abs_diff`: 0.00711 +- `max_velocity_abs_diff`: 0.046875 +- final action vs official: cos 0.9999959, rel_l2 0.286%, max_abs 0.00734 +- peak allocated VRAM during warm run: 11.02 GiB +- estimated action-only E2E with official upstream, skipped vision decode, and + FlashRT denoise: 14.52s, or 3.26x vs official 47.32s E2E +- hard gate run: `flashrt_benchmark_gate_10iter.json` with `gates_passed=true` +- native flat velocity fill canary: + `flashrt_benchmark_fill_velocity_native_3iter.json`, P50 11.97s, speedup + 2.77x, with the same action accuracy gate. +- native action bias+tail-clear canary: + `flashrt_benchmark_action_tail_native_3iter.json`, P50 12.02s, speedup + 2.76x, with the same action accuracy gate. +- native action row scatter/gather canary: + `flashrt_benchmark_action_rows_native_3iter.json`, P50 12.03s, speedup + 2.75x, with the same action accuracy gate. +- native action input tail-copy canary: + `flashrt_benchmark_action_input_native_3iter.json`, P50 12.02s, speedup + 2.76x, with the same action accuracy gate. +- native action bias+timestep canary: + `flashrt_benchmark_action_bias_timestep_native_3iter.json`, P50 12.03s, + speedup 2.75x, with the same action accuracy gate. +- timestep cache canary: + `flashrt_benchmark_timestep_cache_3iter.json`, P50 12.04s, speedup 2.75x, + with the same action accuracy gate. +- native UniPC scheduler canary: + `flashrt_benchmark_native_unipc_3iter.json`, P50 12.00s, speedup 2.76x, + final action cos 0.9999959, rel_l2 0.286%. This is enabled by default when + the binding is available. +- opt-in native residual add canary: + `flashrt_benchmark_residual_add_native_3iter.json`, P50 12.24s, speedup + 2.71x, with the same action accuracy gate. This kernel is correct but is not + default because the extra launches are slower than Torch add in the current + eager stack. +- opt-in CUDA graph run: + `flashrt_benchmark_graph_10iter.json` with `use_cuda_graphs=true`, + `graph_attention=true`, `native_scheduler=true`, P50 12.01s, speedup 2.76x, + final action cos 0.9999962, rel_l2 0.279%. Graph replay is verified, but + eager remains the default because it is marginally faster on this Thor run. +- graph loop breakdown: + `flashrt_benchmark_graph_loop_profile_3iter.json` adds CUDA event timing for + one profiled 30-step graph run. Velocity graph replay accounts for 11.966s of + the 11.970s step total; native UniPC scheduler accounts for only 0.0041s + total, with p50 0.141ms/step. Capturing the Python 30-step loop as one larger + graph is therefore not the next useful latency target; the remaining time is + inside per-step velocity compute. +- additional seed check: seed=1 dump + `dev_scratch_cosmos3_thor/edge_av_inverse_1/tensors.safetensors`, final + action cos 0.9999956 / rel_l2 0.300% / max_abs 0.00923. This run skips the + speedup gate because the official seed=1 denoise baseline was a single faster + run at 29.16s; it is used as the anti-overfit accuracy check. + +Reproduce the number: + +```bash +python benchmarks/cosmos3_edge_thor_denoise.py \ + --iters 10 \ + --warmup-steps 1 \ + --enforce-gates \ + --json-out dev_scratch_cosmos3_thor/edge_av_inverse_0/flashrt_benchmark_gate_10iter.json +``` + +Reproduce the opt-in CUDA graph number: + +```bash +FLASHRT_COSMOS3_EDGE_FA4_FWD=1 \ +python benchmarks/cosmos3_edge_thor_denoise.py \ + --iters 10 \ + --warmup-steps 1 \ + --use-cuda-graphs \ + --enforce-gates \ + --json-out dev_scratch_cosmos3_thor/edge_av_inverse_0/flashrt_benchmark_graph_10iter.json +``` + +Verify the existing delivery artifacts without rerunning GPU benchmarks: + +```bash +python dev_log_cosmos3_thor/verify_delivery.py +``` + +Attention notes: + +- FA4 (`flash_rt.hardware.thor.fa4_backend`) is the default native eager + attention path for gen/full GQA. It is accurate at the Cosmos3 shape and avoids + the long-context failure seen in the cuBLAS MHA fallback. FA4 accepts BF16 + Q/K/V directly, so the hot path avoids per-layer BF16→FP16 casts. +- `fmha_strided_full` remains behind `FLASHRT_COSMOS3_EDGE_FMHA=1`; it is + accurate on some large random shapes but increased full step-0 velocity error + in the Cosmos3 stack, so it is not the default. +- `attention_mha_bf16` remains behind `FLASHRT_COSMOS3_EDGE_MHA=1`; its + `softmax_bf16` path is not reliable for Cosmos3's long `S_kv=6425`. +- Velocity-only CUDA graph capture is available behind + `FLASHRT_COSMOS3_EDGE_FA4_FWD=1` plus `use_cuda_graphs=True`. The failure was + narrowed to the final gen norm bypassing the weight cache; after routing + `norm_moe_gen.weight` through the resident cache, the 28-layer + `flat_velocity` probe and the 10-iter denoise graph benchmark both pass. Graph + replay is opt-in rather than default because the measured P50 is 12.01s vs + 12.00s for eager in the current artifact set. +- The graph loop profile shows native scheduler and Python-loop glue are no + longer material at the 12s denoise scale: scheduler event time is 4.1ms across + all 30 steps. This closes the immediate "capture the outer loop" question as + a no-go for latency; further denoise work must reduce the per-layer velocity + kernels themselves. +- A forced graph experiment with SDPA attention also exited with code 139, so + there is no SDPA graph fallback enabled in this Python eager stack. +- `FLASHRT_COSMOS3_EDGE_FA4_FWD=1` switches eager attention to the lower-level + FA4 `_flash_attn_fwd` entry with preallocated output/LSE buffers. It is also + required for the opt-in graph path. As an eager-only experiment, seed0 3-iter + P50 was 12.008s with the same action accuracy, which does not beat the + default gate result. +- A/B attempts to replace per-layer `torch.cat` with manual static K/V copies + and precompute action timestep embeddings did not improve P50 latency on Thor, + so those overrides are not in the default path. The earlier two-op + `relu_inplace_bf16 + square_` experiment was also not useful, but the fused + `relu2_inplace_bf16` kernel is now enabled by default. Seed0 3-iter A/B: + native relu2 P50 12.274s vs Torch relu2 P50 12.926s, with identical final + action metrics. +- `cosmos3_edge_qk_norm_rope_bf16` fuses gen q/k RMSNorm and RoPE and is enabled + by default. Seed0 3-iter A/B: fused P50 11.938s vs two-step native path P50 + 12.255s, with identical final action metrics. +- `cosmos3_edge_fill_flat_velocity_bf16` is enabled by default when + `flash_rt_kernels` exposes the binding. It fills the `[1201920]` flat velocity + buffer in one native launch, zeroing the vision segment and copying the + `[60,64]` action tail. The CUDA canary in + `tests/test_cosmos3_edge_kernel_bindings.py` checks exact equality with the + Torch construction. +- `cosmos3_edge_add_bias_zero_action_tail_bf16` is enabled by default when + available. It adds the action output bias for the valid action dimensions and + clears columns `raw_action_dim:64` in one native launch. The CUDA canary checks + exact equality with the prior Torch `add_` plus slice-clear construction. +- `cosmos3_edge_scatter_rows_bf16` and `cosmos3_edge_gather_rows_bf16` are + enabled by default when available. They replace the action-row advanced + indexing operations used to install encoded action tokens into the full + sequence and gather decoded action hidden states back out. The CUDA canary + checks exact equality with PyTorch indexed row copy. +- `cosmos3_edge_copy_action_tail_f32_to_bf16` is enabled by default when + available and the latent is CUDA float32. It copies the action tail from the + flat denoise latent into the BF16 action input buffer in one native launch. + The CUDA canary checks exact equality with PyTorch tail slicing plus BF16 + conversion. +- `cosmos3_edge_add_action_bias_timestep_bf16` is enabled by default when + available. It fuses the action static bias add and timestep embedding add, + preserving the prior two-step BF16 rounding order. The timestep embedding is + computed as one row instead of 60 duplicate rows; this changes final action + metrics only in the last few decimals and remains well inside the reference + gate. +- `EdgeDenoiseFlashRT` now precomputes the 30 fixed timestep embeddings from + the dump schedule. The measured denoise loop indexes this cache by step, so + the per-step action encode path no longer calls the Torch timestep + `cos/sin/linear/silu/linear` sequence. +- `cosmos3_edge_unipc_step_f32_bf16` is enabled by default when available. It + folds the fixed 30-step UniPC x0 conversion, corrector, and predictor update + into one native launch per denoise step. The kernel preserves the prior + BF16 `sigma * velocity` rounding order before updating the float32 latent. +- `cosmos3_edge_add_bf16` is available as an opt-in native residual add + experiment with `FLASHRT_COSMOS3_EDGE_NATIVE_ADD=1`. The CUDA canary checks + exact equality with Torch BF16 add. It remains disabled by default because the + measured 3-iter canary regressed P50 to 12.24s versus the 12.00s default. + +P2 NVFP4 canary: + +- SM110 exposes the NVFP4-named bindings in `flash_rt_kernels`, and minimal + `quantize_bf16_to_nvfp4_swizzled` plus `fp4_w4a16_gemm_sm120_bf16out` + canaries execute. +- A real Cosmos3 layer-0 gen `add_q_proj` canary with correct `[N,K]` weight + layout and `bf16_weight_to_nvfp4_swizzled` alpha produced cos 0.9974, + rel_l2 10.3%, and was slower than BF16 GEMM when activation quantization was + included. This does not pass the P2 accuracy/perf gate, so NVFP4 is not wired + into the default backend. + +P4 E2E note: + +- Official benchmark for `av_inverse_0` reports `OmniMoTModel.decode` at + 11.65s. Since inverse dynamics only emits action in the final JSON, this is a + strong candidate for an E2E skip in a future no-dump path. +- The external Cosmos Framework path confirms the skip opportunity: + `inference.py` unconditionally does `outputs.pop("vision") -> decode_vision` + and saves `vision.mp4` before writing `content["action"]`; the action payload + itself is taken from `outputs["action"]` and does not depend on decoded vision. + The external checkout remains unmodified. +- FlashRT-owned `backend="official_action_only"` validates that skip on the + real no-dump official path: `OmniMoTModel.decode` is 0.00025s, generated + action is bit-exact to the official baseline JSON, `files=[]`, and no + `vision*` artifact remains. This still uses the official eager denoise path; + the remaining P4 gap is wiring the FlashRT denoise engine to live upstream + conditioning instead of captured dumps. +- The opt-in live capture path produced + `official_action_only_live_denoise.safetensors` from the real official run: + 30 steps, valid `EdgeDenoiseDump` geometry, action-only JSON with `files=[]`, + decode 0.00048s, and selected tensors match the original reference dump with + max_abs 0.0. This proves a FlashRT-owned live capture boundary, but not yet + live FlashRT denoise handoff. +- The live FlashRT handoff prototype produced + `official_action_only_live_handoff_boundary.safetensors` and an action-only + run with `OmniInference.generate_batch` 26.10s, decode 0.0083s, files `[]`, + final action cos 0.9999959 / rel_l2 0.292% / max_abs 0.00729 vs official + baseline. This removes captured denoise dumps from the online path after + step 0; remaining work is removing the official step-0 bootstrap and reducing + cold handoff overhead. +- The pre-layer bootstrap variant removes the official step-0 decoder forward: + `prelayer_aborted=true`, 30 FlashRT velocity calls, `generate_batch` 26.24s, + decode 0.0090s, final action cos 0.9999955 / rel_l2 0.308% / max_abs 0.00780. + It confirms the remaining one-shot latency is cold FA4/engine setup rather + than official decoder compute. +- The archived pre-layer boundary contract + `official_action_only_live_prelayer_boundary.safetensors` is the smaller + native-denoise/VLM input boundary: 60.36 MiB, 31 tensors, no layer-0 output, + and no step-0 velocity. The older step-0 handoff boundary was 87.77 MiB / + 36 tensors because it also stored layer-0 output and velocity. The run keeps + the same action gate (cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687) and + native UniPC covers all 30 steps. This boundary defines what native prefill + must synthesize live from raw input before the denoise engine can be fully + detached from official upstream code. +- `--live-boundary-in` can now use that 60.36 MiB boundary artifact to + initialize the live FlashRT denoise engine directly. The archived replay run + has `official_velocity_call_count=0`, 30 FlashRT velocities, native UniPC 30 + steps, `generate_batch` 26.03s, and the same action gate. It is a runnable + native prefill target contract, not the final no-artifact live path. +- Combined with `--live-warm-request --cache-warmup-prepare` and all slim derive + flags, boundary-in gives the current artifact/cache service path: + measured `generate_batch` 12.08s, official velocity calls 0, FlashRT velocity + calls 60, measured VAE encode calls 0, measured prepare 0.042s, and the same + action gate. +- The prepare-to-boundary derive helper proves 28 VFM/noise/pack/position/RoPE/text + tensors / 10.65 MiB are already implied by the 9.18 MiB slim prepare contract + plus `embed_tokens.weight`, + and verifies 2 layer-0 input tensors / 25.10 MiB as `lm_in` aliases. The + remaining pre-layer boundary gap is 1 tensor / 24.61 MiB: + `lm_in/full_only_seq`; the verifier now requires the derived full 6300-row + tensor to be bit-exact, including the 60 action rows. +- `--live-boundary-prepare-in` now derives the executable 31-tensor boundary + in-process from the same 9.18 MiB slim prepare artifact plus checkpoint + weights, then drives the live FlashRT denoise engine with 0 official velocity + calls. The archived run measured `generate_batch` 22.24s, denoise 22.22s, + boundary derive 0.258s, engine construct 1.074s, native UniPC 30 steps, and + the same action gate. The optional `--live-boundary-out` debug artifact is + 62.65 MiB because shared views are cloned for safetensors archival; it is not + the runtime input. +- With `--live-warm-request`, the real measured pre-layer batch drops to + `generate_batch` 12.12s / denoise 12.11s while preserving the action gate; + the trace shows 60 FlashRT velocity calls and native UniPC scheduler coverage + for both warmup and measured sampler runs. Only the first warmup velocity call + pays cold compile, and a measured VAE cache hit removes the previous 3.41s + encode from the real request. The warmup VAE cache key includes a sampled + content fingerprint in addition to tensor metadata. +- The opt-in `--cache-warmup-prepare` path reuses the full prepared inference + state from warmup for a repeated request: `generate_batch` 11.98s / denoise + 11.97s, measured `_prepare_inference_data` 0.0185s, prepare hit/store 1/1, + measured misses 0, and measured VAE encode calls 0. The live sampler now uses + the native FlashRT UniPC step for this fixed Edge schedule: 2 native scheduler + runs / 60 steps, 0.0133s total scheduler step time. With + `--prepare-slim-no-raw-state-vision`, the same cache stores a slim post-prepare + payload without raw RGB frames: measured `generate_batch` 12.08s, prepare + 0.0146s, and trace hit/store events both mark + `slim_no_raw_state_vision=true`. With + `--prepare-slim-derive-condition-reference`, the cached payload also omits + `condition_reference` and restores it from `x0_tokens_vision/action` on hit: + measured `generate_batch` 12.00s, prepare 0.0149s, and both cache events mark + `slim_derive_condition_reference=true`. With + `--prepare-slim-derive-initial-noise`, the cache also omits `initial_noise` + and restores it from seed/x0/mask metadata on hit: measured `generate_batch` + 12.05s, prepare 0.0382s, and cache events mark all three slim flags. This is + a service-style repeated-request prefill cache, not the final no-warmup live + VAE/SigLIP/VLM replacement. +- The opt-in upstream trace without VAE cache identified measured VAE + encode/get-data as the next P4 target; the current warmup cache moves that + repeated work out of the measured request for repeated-sample/service-style + warm paths. Native live VAE/SigLIP/VLM prefill remains the broader no-warmup + P4 target. + +Important mapping correction from the layer-0 gate: + +- `to_q/to_k/to_v/to_out` = und/causal tower +- `add_q_proj/add_k_proj/add_v_proj/to_add_out` = gen/full tower +- `norm_added_q/norm_added_k` belong to gen/full +- `k_norm_und_for_gen` is required for gen attending to und keys + +## Optimization Port + +The Cosmos3-Nano FlashRT structure is reusable, but not drop-in compatible: + +- Reuse: hardware dispatch, frontend/pipeline split, UniPC loop shape, static + conditioning cache, CUDA graph capture, TeaCache step reuse, and model-local + kernels. +- Change: Edge uses hidden size 2048, 28 layers, 16/8 GQA heads, `relu2`, and + Nemotron Dense VL reasoner/processor instead of Nano's 4096/36-layer Qwen path. +- Thor work: add a model-local `pipeline_thor.py`, SM110 joint-attention/mrope + kernels, per-layer BF16 reference taps, then FP8/NVFP4 calibration. + +Keep this official runner as the accuracy and benchmark gate while the optimized +Thor pipeline is filled in. diff --git a/docs/stable_api.md b/docs/stable_api.md index ab1f9821..ca6d67fa 100644 --- a/docs/stable_api.md +++ b/docs/stable_api.md @@ -29,7 +29,7 @@ def load_model( autotune: int = 3, # 0=off, 3=default, 5+=thorough recalibrate: bool = False, weight_cache: bool = True, # JAX only - config: str = "pi05", # "pi05" | "pi0" | "groot" | "groot_n17" | "pi0fast" | "motus" | "wan22_ti2v_5b" | "cosmos3_video" + config: str = "pi05", # "pi05" | "pi0" | "groot" | "groot_n17" | "pi0fast" | "motus" | "wan22_ti2v_5b" | "cosmos3_video" | "cosmos3_edge" device=None, # reserved # Pi0-FAST-specific: decode_cuda_graph: bool = False, @@ -113,6 +113,11 @@ Returns a `VLAModel` wrapping the appropriate frontend for the detected compare_ref=..., return_metadata=...)`, returning the denoised vision latent; `predict()` is not part of this API. Precision is selected with `load_model(..., use_fp8=True|False)`. See `docs/cosmos3_video_usage.md`. +- `config="cosmos3_edge"` is a Thor official-baseline runner for NVIDIA + Cosmos Framework. It exposes `set_prompt(input_json=...)` or + `set_prompt(sample=...)`, then `infer(output_dir=..., vae_path=..., + cosmos_root=...)`. It is the accuracy/latency gate before the optimized Thor + Edge pipeline replaces the reference execution. See `docs/cosmos3_edge_thor.md`. - `config="groot_n17"` is registered for `framework="torch"` on `hardware in {"thor", "rtx_sm120", "rtx_sm89"}`. On RTX, `rtx_sm120` resolves through the historical shared RTX registration and diff --git a/examples/cosmos3_edge_thor_baseline.py b/examples/cosmos3_edge_thor_baseline.py new file mode 100644 index 00000000..b6185339 --- /dev/null +++ b/examples/cosmos3_edge_thor_baseline.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""Run the Cosmos3-Edge official Thor baseline through FlashRT dispatch.""" + +from __future__ import annotations + +import argparse +import pathlib +import sys + +ROOT = pathlib.Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +import flash_rt + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True, + help="Local Cosmos3-Edge directory or registered name") + parser.add_argument("--input-json", required=True, + help="Official Cosmos Framework sample JSON") + parser.add_argument("--output-dir", required=True) + parser.add_argument("--vae-path", + help="Local Wan2.2_VAE.pth; avoids HF download inside the container") + parser.add_argument("--cosmos-root", + help="Path to a Cosmos Framework checkout") + parser.add_argument("--config-file", + help="Cosmos3-Edge YAML; defaults under --cosmos-root") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--live-dump-out", + help="Optional safetensors path for an official live denoise dump; requires official_action_only", + ) + parser.add_argument( + "--live-flashrt-handoff", + action="store_true", + help="Use official step 0 to capture live boundary, then replace later denoise velocities with FlashRT", + ) + parser.add_argument( + "--live-prelayer-bootstrap", + action="store_true", + help="Build the live FlashRT boundary before official decoder layers and skip the official step-0 transformer", + ) + parser.add_argument( + "--live-warm-request", + action="store_true", + help="Run the pre-layer FlashRT handoff with one official warmup request before the measured batch", + ) + parser.add_argument( + "--warmup", + type=int, + default=0, + help="Number of official warmup requests to run before the measured request", + ) + parser.add_argument( + "--cache-warmup-vae", + action="store_true", + help="Reuse VAE latents encoded during official warmup for the measured request", + ) + parser.add_argument( + "--cache-warmup-prepare", + action="store_true", + help="Reuse prepared inference state from official warmup for the measured request", + ) + parser.add_argument( + "--live-boundary-out", + help="Optional safetensors path for the step-0 live boundary used by --live-flashrt-handoff", + ) + parser.add_argument( + "--live-boundary-in", + help="Optional safetensors pre-layer boundary whose tensors initialize the live FlashRT denoise engine", + ) + parser.add_argument( + "--live-boundary-prepare-in", + help="Optional torch.save slim prepare artifact used to derive the live FlashRT pre-layer boundary", + ) + parser.add_argument( + "--live-boundary-prepare-live", + action="store_true", + help="Derive the live FlashRT pre-layer boundary from this request's in-memory prepared state", + ) + parser.add_argument( + "--live-handoff-trace-out", + help="Optional JSON path for per-stage live FlashRT handoff timings", + ) + parser.add_argument( + "--upstream-trace-out", + help="Optional JSON path for opt-in official upstream stage timings", + ) + parser.add_argument( + "--vae-encode-dump-out", + help="Optional safetensors path for the measured request's VAE encode latent boundary", + ) + parser.add_argument( + "--vae-latent-in", + help="Optional safetensors path whose vae_encode/output tensor replaces measured VAE encode", + ) + parser.add_argument( + "--vae-encode-dump-input", + action="store_true", + help="Also include the full VAE input video tensor in --vae-encode-dump-out", + ) + parser.add_argument( + "--vae-encode-profile-out", + help="Optional JSON path for Wan2.2 VAE encode module/shape/timing profile", + ) + parser.add_argument( + "--vae-native-rms-silu", + action="store_true", + help="Enable FlashRT native BF16 RMS+SiLU inside Wan2.2 VAE encode ResidualBlocks", + ) + parser.add_argument( + "--vae-t1-conv2d", + action="store_true", + help="Run no-cache single-frame Wan2.2 CausalConv3d sites as equivalent Conv2d", + ) + parser.add_argument( + "--vae-native-avgdown3d", + action="store_true", + help="Enable FlashRT native BF16 AvgDown3D shortcut pooling inside Wan2.2 VAE encode", + ) + parser.add_argument( + "--vae-channels-last3d-conv320", + action="store_true", + help="Run steady 320-channel Wan2.2 CausalConv3d sites with channels_last_3d cuDNN layout", + ) + parser.add_argument( + "--vae-compile-encode", + action="store_true", + help="Enable upstream Wan2.2 VAE chunk AOT compile_encode for 480p 16:9", + ) + parser.add_argument( + "--vae-compile-trace-out", + help="Optional JSON path for VAE compile_encode setup trace", + ) + parser.add_argument( + "--prepare-dump-out", + help="Optional torch.save path for the measured request's prepared inference state", + ) + parser.add_argument( + "--prepare-replay-in", + help="Optional torch.save path whose prepared inference state replaces measured prepare", + ) + parser.add_argument( + "--prepare-inventory-out", + help="Optional JSON path describing the measured/replayed prepared inference state tensor footprint", + ) + parser.add_argument( + "--prepare-slim-no-raw-state-vision", + action="store_true", + help="Drop gen_data_clean.raw_state_vision from the prepared state; keeps latent/noise/mask state", + ) + parser.add_argument( + "--prepare-slim-derive-condition-reference", + action="store_true", + help="Do not store condition_reference in prepared state; derive it from x0_tokens on replay/cache hit", + ) + parser.add_argument( + "--prepare-slim-derive-initial-noise", + action="store_true", + help="Do not store initial_noise in prepared state; derive it from seed/x0_tokens/mask on replay/cache hit", + ) + parser.add_argument( + "--backend", + choices=("official", "official_action_only"), + default="official", + help="official runs the upstream Cosmos output path; official_action_only skips inverse-dynamics vision decode/save", + ) + args = parser.parse_args() + + model = flash_rt.load_model( + args.checkpoint, + framework="torch", + config="cosmos3_edge", + hardware="thor", + ) + model.set_prompt(input_json=args.input_json) + out = model.infer( + output_dir=args.output_dir, + seed=args.seed, + backend=args.backend, + vae_path=args.vae_path, + cosmos_root=args.cosmos_root, + config_file=args.config_file, + live_dump_out=args.live_dump_out, + live_flashrt_handoff=args.live_flashrt_handoff, + live_boundary_out=args.live_boundary_out, + live_boundary_in=args.live_boundary_in, + live_boundary_prepare_in=args.live_boundary_prepare_in, + live_boundary_prepare_live=args.live_boundary_prepare_live, + live_handoff_trace_out=args.live_handoff_trace_out, + upstream_trace_out=args.upstream_trace_out, + vae_encode_dump_out=args.vae_encode_dump_out, + vae_latent_in=args.vae_latent_in, + vae_encode_dump_input=args.vae_encode_dump_input, + vae_encode_profile_out=args.vae_encode_profile_out, + vae_native_rms_silu=args.vae_native_rms_silu, + vae_t1_conv2d=args.vae_t1_conv2d, + vae_native_avgdown3d=args.vae_native_avgdown3d, + vae_channels_last3d_conv320=args.vae_channels_last3d_conv320, + vae_compile_encode=args.vae_compile_encode, + vae_compile_trace_out=args.vae_compile_trace_out, + prepare_dump_out=args.prepare_dump_out, + prepare_replay_in=args.prepare_replay_in, + prepare_inventory_out=args.prepare_inventory_out, + prepare_slim_no_raw_state_vision=args.prepare_slim_no_raw_state_vision, + prepare_slim_derive_condition_reference=args.prepare_slim_derive_condition_reference, + prepare_slim_derive_initial_noise=args.prepare_slim_derive_initial_noise, + live_prelayer_bootstrap=args.live_prelayer_bootstrap, + live_warm_request=args.live_warm_request, + cache_warmup_vae=args.cache_warmup_vae, + cache_warmup_prepare=args.cache_warmup_prepare, + warmup=args.warmup, + ) + print(f"[cosmos3_edge] {args.backend} wall={out['wall_s']:.2f}s") + print(f"[cosmos3_edge] output_dir={out['output_dir']}") + for path in out.get("sample_outputs", []): + print(f"[cosmos3_edge] sample_outputs={path}") + if "live_dump_out" in out: + print(f"[cosmos3_edge] live_dump_out={out['live_dump_out']}") + if out.get("live_flashrt_handoff"): + print("[cosmos3_edge] live_flashrt_handoff=true") + if out.get("live_prelayer_bootstrap"): + print("[cosmos3_edge] live_prelayer_bootstrap=true") + if out.get("live_warm_request"): + print("[cosmos3_edge] live_warm_request=true") + if out.get("cache_warmup_vae"): + print("[cosmos3_edge] cache_warmup_vae=true") + if out.get("cache_warmup_prepare"): + print("[cosmos3_edge] cache_warmup_prepare=true") + if "warmup" in out: + print(f"[cosmos3_edge] warmup={out['warmup']}") + if "live_boundary_out" in out: + print(f"[cosmos3_edge] live_boundary_out={out['live_boundary_out']}") + if "live_boundary_in" in out: + print(f"[cosmos3_edge] live_boundary_in={out['live_boundary_in']}") + if "live_boundary_prepare_in" in out: + print(f"[cosmos3_edge] live_boundary_prepare_in={out['live_boundary_prepare_in']}") + if out.get("live_boundary_prepare_live"): + print("[cosmos3_edge] live_boundary_prepare_live=true") + if "live_handoff_trace_out" in out: + print(f"[cosmos3_edge] live_handoff_trace_out={out['live_handoff_trace_out']}") + if "upstream_trace_out" in out: + print(f"[cosmos3_edge] upstream_trace_out={out['upstream_trace_out']}") + if "vae_encode_dump_out" in out: + print(f"[cosmos3_edge] vae_encode_dump_out={out['vae_encode_dump_out']}") + if "vae_latent_in" in out: + print(f"[cosmos3_edge] vae_latent_in={out['vae_latent_in']}") + if out.get("vae_encode_dump_input"): + print("[cosmos3_edge] vae_encode_dump_input=true") + if "vae_encode_profile_out" in out: + print(f"[cosmos3_edge] vae_encode_profile_out={out['vae_encode_profile_out']}") + if out.get("vae_native_rms_silu"): + print("[cosmos3_edge] vae_native_rms_silu=true") + if out.get("vae_t1_conv2d"): + print("[cosmos3_edge] vae_t1_conv2d=true") + if out.get("vae_native_avgdown3d"): + print("[cosmos3_edge] vae_native_avgdown3d=true") + if out.get("vae_channels_last3d_conv320"): + print("[cosmos3_edge] vae_channels_last3d_conv320=true") + if out.get("vae_compile_encode"): + print("[cosmos3_edge] vae_compile_encode=true") + if "vae_compile_trace_out" in out: + print(f"[cosmos3_edge] vae_compile_trace_out={out['vae_compile_trace_out']}") + if "prepare_dump_out" in out: + print(f"[cosmos3_edge] prepare_dump_out={out['prepare_dump_out']}") + if "prepare_replay_in" in out: + print(f"[cosmos3_edge] prepare_replay_in={out['prepare_replay_in']}") + if "prepare_inventory_out" in out: + print(f"[cosmos3_edge] prepare_inventory_out={out['prepare_inventory_out']}") + if out.get("prepare_slim_no_raw_state_vision"): + print("[cosmos3_edge] prepare_slim_no_raw_state_vision=true") + if out.get("prepare_slim_derive_condition_reference"): + print("[cosmos3_edge] prepare_slim_derive_condition_reference=true") + if out.get("prepare_slim_derive_initial_noise"): + print("[cosmos3_edge] prepare_slim_derive_initial_noise=true") + + +if __name__ == "__main__": + main() diff --git a/flash_rt/api.py b/flash_rt/api.py index 1ebe3875..d6fcf0c5 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -318,9 +318,13 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, weight_cache: if True (default), cache FP8-quantized weights to disk after first load. Only affects JAX. config: model config name: "pi05", "pi0", "groot", "groot_n17", - "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video". + "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video", + "cosmos3_edge". "cosmos3_video" is a non-VLA text2video denoise model: drive it with set_prompt(ref=) + infer(...), not predict(). + "cosmos3_edge" is the official Cosmos Framework Thor baseline + runner: drive it with set_prompt(sample=... or input_json=...) + + infer(output_dir=..., vae_path=...). device: ignored (auto-detects GPU). Reserved for future multi-GPU. decode_cuda_graph: Pi0-FAST only. Capture action-phase decode as CUDA Graph for max throughput (trades startup time for per-token speed). @@ -433,11 +437,12 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, "See docs/qwen3_vl_fp8_sm89.md and docs/qwen3_vl_nvfp4.md.") if config not in ("pi05", "groot", "groot_n17", "pi0", "pi0fast", - "motus", "wan22_ti2v_5b", "cosmos3_video", "nexn2"): + "motus", "wan22_ti2v_5b", "cosmos3_video", + "cosmos3_edge", "nexn2"): raise ValueError( f"Unknown config: {config}. " f"Supported: pi05, groot, groot_n17, pi0, pi0fast, motus, " - f"wan22_ti2v_5b, cosmos3_video, nexn2") + f"wan22_ti2v_5b, cosmos3_video, cosmos3_edge, nexn2") if framework not in ("torch", "jax"): raise ValueError( f"Unknown framework: {framework}. Supported: torch, jax") @@ -656,6 +661,10 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, elif config == "wan22_ti2v_5b": if "autotune" in sig.parameters: kwargs["autotune"] = autotune + elif config == "cosmos3_edge": + # Official Cosmos Framework baseline runner. Runtime knobs such as + # output_dir, seed, benchmark, and local Wan VAE path are infer() args. + pass else: # pi05, pi0 — both Thor and rtx variants take (checkpoint, num_views, autotune) # or (checkpoint, num_views). Feature-detect. diff --git a/flash_rt/configs/cosmos3_edge.yaml b/flash_rt/configs/cosmos3_edge.yaml new file mode 100644 index 00000000..f05a0631 --- /dev/null +++ b/flash_rt/configs/cosmos3_edge.yaml @@ -0,0 +1,21 @@ +# Cosmos3-Edge official Thor baseline. +# Metadata only - the reference execution is delegated to NVIDIA Cosmos +# Framework so the first Thor baseline remains directly comparable upstream. +name: cosmos3_edge +description: Cosmos3-Edge official Cosmos Framework baseline runner for Thor +arch: cosmos3_edge_omni +hidden_dim: 2048 +ffn_hidden_dim: 9216 +num_layers: 28 +num_heads: 16 +num_kv_heads: 8 +head_dim: 128 +patch_latent_dim: 192 +latent_channels: 48 +max_action_dim: 64 +action_gen: true +vision_gen: true +sound_gen: false +precision: bf16 +backend: cosmos_framework +hardware: thor diff --git a/flash_rt/frontends/torch/_cosmos3_edge_thor_spec.py b/flash_rt/frontends/torch/_cosmos3_edge_thor_spec.py new file mode 100644 index 00000000..6024221b --- /dev/null +++ b/flash_rt/frontends/torch/_cosmos3_edge_thor_spec.py @@ -0,0 +1,177 @@ +"""Cosmos3-Edge Thor weight specification. + +The public Edge checkpoint is a diffusers-style directory. This file is the +single source of truth for mapping those tensor names onto the FlashRT Thor +denoise engine. It is intentionally explicit: in the public Edge checkpoint as +loaded by Cosmos Framework, ``to_*`` maps to the understanding/causal tower and +``add_*`` maps to the generation/full tower. The latter also owns +``norm_added_{q,k}``. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + + +@dataclass(frozen=True) +class Cosmos3EdgeThorSpec: + num_layers: int = 28 + hidden_size: int = 2048 + ffn_size: int = 9216 + num_heads: int = 16 + num_kv_heads: int = 8 + head_dim: int = 128 + patch_latent_dim: int = 192 + latent_channels: int = 48 + action_dim: int = 64 + num_action_domains: int = 32 + vocab_size: int = 131072 + rms_eps: float = 1e-5 + + @property + def q_width(self) -> int: + return self.num_heads * self.head_dim + + @property + def kv_width(self) -> int: + return self.num_kv_heads * self.head_dim + + +SPEC = Cosmos3EdgeThorSpec() + + +GLOBAL_SHAPES: dict[str, tuple[int, ...]] = { + "action_modality_embed": (SPEC.hidden_size,), + "action_proj_in.bias.weight": (SPEC.num_action_domains, SPEC.hidden_size), + "action_proj_in.fc.weight": ( + SPEC.num_action_domains, + SPEC.action_dim * SPEC.hidden_size, + ), + "action_proj_out.bias.weight": (SPEC.num_action_domains, SPEC.action_dim), + "action_proj_out.fc.weight": ( + SPEC.num_action_domains, + SPEC.action_dim * SPEC.hidden_size, + ), + "embed_tokens.weight": (SPEC.vocab_size, SPEC.hidden_size), + "lm_head.weight": (SPEC.vocab_size, SPEC.hidden_size), + "norm.weight": (SPEC.hidden_size,), + "norm_moe_gen.weight": (SPEC.hidden_size,), + "proj_in.bias": (SPEC.hidden_size,), + "proj_in.weight": (SPEC.hidden_size, SPEC.patch_latent_dim), + "proj_out.bias": (SPEC.patch_latent_dim,), + "proj_out.weight": (SPEC.patch_latent_dim, SPEC.hidden_size), + "time_embedder.linear_1.bias": (SPEC.hidden_size,), + "time_embedder.linear_1.weight": (SPEC.hidden_size, 256), + "time_embedder.linear_2.bias": (SPEC.hidden_size,), + "time_embedder.linear_2.weight": (SPEC.hidden_size, SPEC.hidden_size), +} + + +LAYER_SHAPES: dict[str, tuple[int, ...]] = { + # Und tower norms/MLP. + "input_layernorm.weight": (SPEC.hidden_size,), + "post_attention_layernorm.weight": (SPEC.hidden_size,), + "mlp.up_proj.weight": (SPEC.ffn_size, SPEC.hidden_size), + "mlp.down_proj.weight": (SPEC.hidden_size, SPEC.ffn_size), + # Gen tower norms/MLP. + "input_layernorm_moe_gen.weight": (SPEC.hidden_size,), + "post_attention_layernorm_moe_gen.weight": (SPEC.hidden_size,), + "mlp_moe_gen.up_proj.weight": (SPEC.ffn_size, SPEC.hidden_size), + "mlp_moe_gen.down_proj.weight": (SPEC.hidden_size, SPEC.ffn_size), + # Und tower attention: diffusers ``to_*``. + "self_attn.to_q.weight": (SPEC.q_width, SPEC.hidden_size), + "self_attn.to_k.weight": (SPEC.kv_width, SPEC.hidden_size), + "self_attn.to_v.weight": (SPEC.kv_width, SPEC.hidden_size), + "self_attn.to_out.weight": (SPEC.hidden_size, SPEC.hidden_size), + # Gen tower attention: diffusers ``add_*``. + "self_attn.add_q_proj.weight": (SPEC.q_width, SPEC.hidden_size), + "self_attn.add_k_proj.weight": (SPEC.kv_width, SPEC.hidden_size), + "self_attn.add_v_proj.weight": (SPEC.kv_width, SPEC.hidden_size), + "self_attn.to_add_out.weight": (SPEC.hidden_size, SPEC.hidden_size), + "self_attn.norm_added_q.weight": (SPEC.head_dim,), + "self_attn.norm_added_k.weight": (SPEC.head_dim,), + # Edge-only: gen tower consumes und K through this extra normalization. + "self_attn.k_norm_und_for_gen.weight": (SPEC.head_dim,), +} + + +GEN_ATTENTION_KEYS = ( + "self_attn.add_q_proj.weight", + "self_attn.add_k_proj.weight", + "self_attn.add_v_proj.weight", + "self_attn.to_add_out.weight", +) + +UND_ATTENTION_KEYS = ( + "self_attn.to_q.weight", + "self_attn.to_k.weight", + "self_attn.to_v.weight", + "self_attn.to_out.weight", +) + +GEN_MLP_KEYS = ( + "mlp_moe_gen.up_proj.weight", + "mlp_moe_gen.down_proj.weight", +) + +UND_MLP_KEYS = ( + "mlp.up_proj.weight", + "mlp.down_proj.weight", +) + + +def layer_key(layer: int, suffix: str) -> str: + if layer < 0 or layer >= SPEC.num_layers: + raise ValueError(f"layer index out of range: {layer}") + return f"layers.{layer}.{suffix}" + + +def iter_expected_shapes() -> Iterable[tuple[str, tuple[int, ...]]]: + yield from GLOBAL_SHAPES.items() + for layer in range(SPEC.num_layers): + for suffix, shape in LAYER_SHAPES.items(): + yield layer_key(layer, suffix), shape + + +def transformer_index_path(checkpoint: str | Path) -> Path: + root = Path(checkpoint) + return root / "transformer" / "diffusion_pytorch_model.safetensors.index.json" + + +def load_transformer_weight_map(checkpoint: str | Path) -> dict[str, str]: + path = transformer_index_path(checkpoint) + data = json.loads(path.read_text(encoding="utf-8")) + weight_map = data.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"{path} does not contain a weight_map.") + return {str(k): str(v) for k, v in weight_map.items()} + + +def validate_transformer_index(checkpoint: str | Path) -> None: + weight_map = load_transformer_weight_map(checkpoint) + expected = dict(iter_expected_shapes()) + missing = sorted(set(expected) - set(weight_map)) + extra = sorted(set(weight_map) - set(expected)) + if missing or extra: + raise ValueError( + "Cosmos3-Edge transformer index mismatch: " + f"missing={missing[:10]} extra={extra[:10]}" + ) + + +def validate_transformer_shapes(checkpoint: str | Path) -> None: + from safetensors import safe_open + + root = Path(checkpoint) + weight_map = load_transformer_weight_map(root) + for key, expected_shape in iter_expected_shapes(): + shard = root / "transformer" / weight_map[key] + with safe_open(str(shard), framework="pt", device="cpu") as f: + actual_shape = tuple(f.get_slice(key).get_shape()) + if actual_shape != expected_shape: + raise ValueError( + f"{key} shape mismatch: expected {expected_shape}, got {actual_shape}" + ) diff --git a/flash_rt/frontends/torch/cosmos3_edge_thor.py b/flash_rt/frontends/torch/cosmos3_edge_thor.py new file mode 100644 index 00000000..701febfe --- /dev/null +++ b/flash_rt/frontends/torch/cosmos3_edge_thor.py @@ -0,0 +1,571 @@ +"""Cosmos3-Edge Thor official-baseline frontend. + +This frontend deliberately runs NVIDIA's Cosmos Framework inference entrypoint +first. It gives FlashRT a reproducible Thor baseline before the optimized +model-local Thor pipeline is ported from the Cosmos3-Nano FlashRT work. +""" + +from __future__ import annotations + +import json +import os +import pathlib +import subprocess +import sys +import tempfile +import time +from typing import Any + + +class Cosmos3EdgeTorchFrontendThor: + """Thin wrapper around ``cosmos_framework.scripts.inference``. + + ``checkpoint`` is a local Cosmos3-Edge Hugging Face/diffusers directory or + a Cosmos Framework registered name. Use ``set_prompt(sample=...)`` or + ``set_prompt(input_json=...)`` to provide official sample JSON, then + ``infer(output_dir=..., vae_path=...)`` to run the baseline. + """ + + def __init__(self, checkpoint: str, num_views: int = 1, + hardware: str | None = None): + del num_views + if hardware not in (None, "thor"): + raise ValueError( + "Cosmos3-Edge baseline frontend is registered for Thor only; " + f"got hardware={hardware!r}.") + self.checkpoint = str(checkpoint) + self.hardware = "thor" + self._input_json: str | None = None + self._sample: dict[str, Any] | None = None + + def set_prompt(self, *, input_json: str | None = None, + sample: dict[str, Any] | None = None, + **sample_fields: Any) -> None: + """Set the official Cosmos Framework sample input. + + Exactly one source should be provided: ``input_json`` for an existing + JSON file, ``sample`` for a dict, or keyword fields for a single sample. + """ + provided = sum(x is not None for x in (input_json, sample)) + if provided + bool(sample_fields) != 1: + raise ValueError( + "Provide exactly one of input_json=..., sample=..., or " + "sample keyword fields.") + self._input_json = str(input_json) if input_json is not None else None + if sample is not None: + self._sample = dict(sample) + elif sample_fields: + self._sample = dict(sample_fields) + else: + self._sample = None + + @staticmethod + def _default_cosmos_root() -> pathlib.Path: + here = pathlib.Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "external" / "cosmos-framework" + if candidate.exists(): + return candidate + return pathlib.Path.cwd() + + @staticmethod + def _default_config_file(cosmos_root: pathlib.Path) -> pathlib.Path: + return (cosmos_root / "cosmos_framework" / "inference" / "configs" / + "model" / "Cosmos3-Edge.yaml") + + @staticmethod + def _write_sample_json(sample: dict[str, Any], work_dir: pathlib.Path) -> pathlib.Path: + work_dir.mkdir(parents=True, exist_ok=True) + fd, path = tempfile.mkstemp( + prefix="cosmos3_edge_sample_", suffix=".json", dir=str(work_dir)) + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(sample, f, indent=2) + f.write("\n") + return pathlib.Path(path) + + @staticmethod + def _write_action_only_output(action: Any, output_dir: pathlib.Path) -> pathlib.Path: + """Write a Cosmos-style action-only output without decoded vision files.""" + output_dir.mkdir(parents=True, exist_ok=True) + payload = { + "status": "success", + "args": {}, + "outputs": [ + { + "content": {"action": action}, + "files": [], + } + ], + } + output_path = output_dir / "sample_outputs.json" + output_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + return output_path + + @staticmethod + def _absolute_output_path(path: str) -> str: + out = pathlib.Path(path).expanduser() + if not out.is_absolute(): + out = pathlib.Path.cwd() / out + return str(out.resolve()) + + def infer(self, *, output_dir: str, seed: int = 0, + backend: str = "official", + benchmark: bool = True, + parallelism_preset: str = "latency", + guardrails: bool = False, + use_torch_compile: bool = False, + use_cuda_graphs: bool = False, + cosmos_root: str | None = None, + config_file: str | None = None, + python_executable: str | None = None, + vae_path: str | None = None, + reference_dump: str | None = None, + boundary_dump: str | None = None, + live_dump_out: str | None = None, + live_flashrt_handoff: bool = False, + live_boundary_out: str | None = None, + live_boundary_in: str | None = None, + live_boundary_prepare_in: str | None = None, + live_boundary_prepare_live: bool = False, + live_handoff_trace_out: str | None = None, + upstream_trace_out: str | None = None, + vae_encode_dump_out: str | None = None, + vae_latent_in: str | None = None, + vae_encode_dump_input: bool = False, + vae_encode_profile_out: str | None = None, + vae_native_rms_silu: bool = False, + vae_t1_conv2d: bool = False, + vae_native_avgdown3d: bool = False, + vae_channels_last3d_conv320: bool = False, + vae_compile_encode: bool = False, + vae_compile_trace_out: str | None = None, + prepare_dump_out: str | None = None, + prepare_replay_in: str | None = None, + prepare_inventory_out: str | None = None, + prepare_slim_no_raw_state_vision: bool = False, + prepare_slim_derive_condition_reference: bool = False, + prepare_slim_derive_initial_noise: bool = False, + live_prelayer_bootstrap: bool = False, + live_warm_request: bool = False, + cache_warmup_vae: bool = False, + cache_warmup_prepare: bool = False, + warmup: int = 0, + replay_device: str = "cpu", + extra_args: list[str] | None = None, + env: dict[str, str] | None = None, + timeout: float | None = None) -> dict[str, Any]: + """Run the official baseline or P1 denoise replay scaffold. + + ``vae_path`` should point at the official ``Wan2.2_VAE.pth`` when the + container cannot download it. It is passed as a config override with an + empty bucket name so Cosmos Framework treats it as a local file. + """ + if backend not in {"official", "official_action_only", "replay", "torch_ref", "flashrt"}: + raise ValueError( + "Cosmos3-Edge backend must be one of 'official', " + "'official_action_only', 'replay', 'torch_ref', or " + f"'flashrt'; got {backend!r}.") + if warmup < 0: + raise ValueError("warmup must be non-negative") + if live_warm_request: + if backend != "official_action_only": + raise ValueError("live_warm_request is only supported with backend='official_action_only'") + live_prelayer_bootstrap = True + warmup = max(int(warmup), 1) + cache_warmup_vae = True + if live_dump_out is not None and backend != "official_action_only": + raise ValueError("live_dump_out is only supported with backend='official_action_only'") + if upstream_trace_out is not None and backend != "official_action_only": + raise ValueError("upstream_trace_out is only supported with backend='official_action_only'") + if ( + vae_encode_dump_out is not None + or vae_latent_in is not None + or vae_encode_dump_input + or vae_encode_profile_out is not None + or vae_native_rms_silu + or vae_t1_conv2d + or vae_native_avgdown3d + or vae_channels_last3d_conv320 + or vae_compile_encode + or vae_compile_trace_out is not None + or prepare_dump_out is not None + or prepare_replay_in is not None + or prepare_inventory_out is not None + or prepare_slim_no_raw_state_vision + or prepare_slim_derive_condition_reference + or prepare_slim_derive_initial_noise + ) and backend != "official_action_only": + raise ValueError( + "VAE/prepare boundary/profile/inventory dump/replay is only supported with backend='official_action_only'" + ) + if vae_compile_encode and (vae_native_rms_silu or vae_native_avgdown3d): + raise ValueError("vae_compile_encode cannot be combined with native VAE monkeypatches") + if cache_warmup_vae and (vae_encode_dump_out is not None or vae_latent_in is not None): + raise ValueError("cache_warmup_vae cannot be combined with VAE encode boundary dump/replay") + if cache_warmup_prepare and ( + prepare_dump_out is not None + or prepare_replay_in is not None + or prepare_inventory_out is not None + ): + raise ValueError("cache_warmup_prepare cannot be combined with prepare boundary dump/replay/inventory") + if ( + live_flashrt_handoff + or live_boundary_out is not None + or live_boundary_in is not None + or live_boundary_prepare_in is not None + or live_boundary_prepare_live + or live_handoff_trace_out is not None + or live_prelayer_bootstrap + or live_warm_request + ) and backend != "official_action_only": + raise ValueError("live FlashRT handoff is only supported with backend='official_action_only'") + if live_dump_out is not None and live_flashrt_handoff: + raise ValueError("live_dump_out cannot be combined with live_flashrt_handoff") + if live_boundary_in is not None and live_boundary_prepare_in is not None: + raise ValueError("live_boundary_in cannot be combined with live_boundary_prepare_in") + if live_boundary_prepare_live and (live_boundary_in is not None or live_boundary_prepare_in is not None): + raise ValueError("live_boundary_prepare_live cannot be combined with live boundary input artifacts") + if backend in {"replay", "torch_ref", "flashrt"}: + if reference_dump is None: + raise ValueError(f"backend={backend!r} requires reference_dump=...") + out_dir = pathlib.Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + from flash_rt.models.cosmos3_edge import ( + EdgeBoundaryDump, + EdgeDenoiseDump, + EdgeDenoiseFlashRT, + EdgeDenoiseReplay, + EdgeDenoiseTorchReference, + EdgeTransformerWeights, + ) + + t0 = time.perf_counter() + dump = EdgeDenoiseDump(reference_dump) + if backend == "replay": + runner = EdgeDenoiseReplay(dump, device=replay_device) + replay_result = runner.replay() + wall_s = time.perf_counter() - t0 + return { + "backend": "replay", + "reference_dump": str(reference_dump), + "output_dir": str(out_dir), + "wall_s": wall_s, + "num_steps": len(replay_result.timesteps), + "timesteps": list(replay_result.timesteps), + "flat_dim": int(replay_result.final_flat.numel()), + "vision_shape": list(replay_result.final_parts.vision.shape), + "action_model_shape": list(replay_result.final_parts.action_model.shape), + "max_input_abs_diff": replay_result.max_input_abs_diff, + } + + if boundary_dump is None: + raise ValueError(f"backend={backend!r} requires boundary_dump=...") + runner_cls = EdgeDenoiseFlashRT if backend == "flashrt" else EdgeDenoiseTorchReference + runner = runner_cls( + dump, + EdgeBoundaryDump(boundary_dump), + EdgeTransformerWeights(self.checkpoint), + device=replay_device, + **({"use_cuda_graphs": use_cuda_graphs} if backend == "flashrt" else {}), + ) + replay_result = runner.run() + wall_s = time.perf_counter() - t0 + result = { + "backend": backend, + "reference_dump": str(reference_dump), + "boundary_dump": str(boundary_dump), + "output_dir": str(out_dir), + "wall_s": wall_s, + "num_steps": replay_result.steps_run, + "timesteps": list(replay_result.timesteps), + "flat_dim": int(replay_result.final_flat.numel()), + "vision_shape": list(replay_result.final_parts.vision.shape), + "action_model_shape": list(replay_result.final_parts.action_model.shape), + "max_input_abs_diff": replay_result.max_input_abs_diff, + "max_velocity_abs_diff": replay_result.max_velocity_abs_diff, + } + if backend == "flashrt": + result["use_cuda_graphs"] = bool(use_cuda_graphs) + if runner.static_engine is not None: + result["native_attention"] = bool( + getattr(runner.static_engine.reference, "native_attention_available", False) + ) + result["graph_attention"] = bool( + getattr(runner.static_engine.reference, "graph_attention_available", False) + ) + official_action = dump.final_action.float() + if tuple(official_action.shape) == tuple(replay_result.final_action.shape): + import torch + + diff = replay_result.final_action.float() - official_action + result["final_action_cos"] = float( + torch.nn.functional.cosine_similarity( + replay_result.final_action.flatten(), + official_action.flatten(), + dim=0, + ).item() + ) + result["final_action_rel_l2"] = float( + (torch.linalg.vector_norm(diff) / torch.linalg.vector_norm(official_action)).item() + ) + result["final_action_max_abs"] = float(diff.abs().max().item()) + if backend == "flashrt": + action_json = self._write_action_only_output( + replay_result.final_action.detach().cpu().tolist(), + out_dir, + ) + result["sample_outputs"] = str(action_json) + result["action_only_output"] = True + return result + + if self._input_json is None and self._sample is None: + raise ValueError( + "Call set_prompt(input_json=...) or set_prompt(sample=...) " + "before infer().") + + out_dir = pathlib.Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + if self._input_json is not None: + input_json = pathlib.Path(self._input_json) + else: + input_json = self._write_sample_json(self._sample or {}, out_dir) + + cosmos_root_path = pathlib.Path(cosmos_root) if cosmos_root else self._default_cosmos_root() + config_path = pathlib.Path(config_file) if config_file else self._default_config_file(cosmos_root_path) + py = python_executable or sys.executable + live_dump_path = self._absolute_output_path(live_dump_out) if live_dump_out is not None else None + live_boundary_path = self._absolute_output_path(live_boundary_out) if live_boundary_out is not None else None + live_boundary_in_path = self._absolute_output_path(live_boundary_in) if live_boundary_in is not None else None + live_boundary_prepare_in_path = ( + self._absolute_output_path(live_boundary_prepare_in) + if live_boundary_prepare_in is not None + else None + ) + live_trace_path = ( + self._absolute_output_path(live_handoff_trace_out) + if live_handoff_trace_out is not None + else None + ) + upstream_trace_path = ( + self._absolute_output_path(upstream_trace_out) + if upstream_trace_out is not None + else None + ) + vae_encode_dump_path = ( + self._absolute_output_path(vae_encode_dump_out) + if vae_encode_dump_out is not None + else None + ) + vae_latent_path = ( + self._absolute_output_path(vae_latent_in) + if vae_latent_in is not None + else None + ) + vae_encode_profile_path = ( + self._absolute_output_path(vae_encode_profile_out) + if vae_encode_profile_out is not None + else None + ) + vae_compile_trace_path = ( + self._absolute_output_path(vae_compile_trace_out) + if vae_compile_trace_out is not None + else None + ) + prepare_dump_path = ( + self._absolute_output_path(prepare_dump_out) + if prepare_dump_out is not None + else None + ) + prepare_replay_path = ( + self._absolute_output_path(prepare_replay_in) + if prepare_replay_in is not None + else None + ) + prepare_inventory_path = ( + self._absolute_output_path(prepare_inventory_out) + if prepare_inventory_out is not None + else None + ) + + module = ( + "flash_rt.models.cosmos3_edge.action_only_official" + if backend == "official_action_only" + else "cosmos_framework.scripts.inference" + ) + cmd = [ + py, "-m", module, + "--parallelism-preset", parallelism_preset, + "-i", str(input_json), + "-o", str(out_dir), + "--checkpoint-path", self.checkpoint, + "--config-file", str(config_path), + "--seed", str(seed), + ] + cmd.append("--benchmark" if benchmark else "--no-benchmark") + cmd.append("--guardrails" if guardrails else "--no-guardrails") + cmd.append("--use-torch-compile" if use_torch_compile else "--no-use-torch-compile") + cmd.append("--use-cuda-graphs" if use_cuda_graphs else "--no-use-cuda-graphs") + if warmup: + cmd.extend(["--warmup", str(int(warmup))]) + if vae_path: + cmd.extend([ + "--experiment-overrides", + "model.config.tokenizer.bucket_name=", + f"model.config.tokenizer.vae_path={vae_path}", + ]) + if live_dump_path is not None: + cmd.extend(["--flashrt-live-dump-out", live_dump_path]) + if live_flashrt_handoff: + cmd.append("--flashrt-live-flashrt-handoff") + if live_prelayer_bootstrap: + cmd.append("--flashrt-live-prelayer-bootstrap") + if live_boundary_path is not None: + cmd.extend(["--flashrt-live-boundary-out", live_boundary_path]) + if live_boundary_in_path is not None: + cmd.extend(["--flashrt-live-boundary-in", live_boundary_in_path]) + if live_boundary_prepare_in_path is not None: + cmd.extend(["--flashrt-live-boundary-prepare-in", live_boundary_prepare_in_path]) + if live_boundary_prepare_live: + cmd.append("--flashrt-live-boundary-prepare-live") + if live_trace_path is not None: + cmd.extend(["--flashrt-live-handoff-trace-out", live_trace_path]) + if upstream_trace_path is not None: + cmd.extend(["--flashrt-upstream-trace-out", upstream_trace_path]) + if vae_encode_dump_path is not None: + cmd.extend(["--flashrt-vae-encode-dump-out", vae_encode_dump_path]) + if vae_latent_path is not None: + cmd.extend(["--flashrt-vae-latent-in", vae_latent_path]) + if vae_encode_dump_input: + cmd.append("--flashrt-vae-encode-dump-input") + if vae_encode_profile_path is not None: + cmd.extend(["--flashrt-vae-encode-profile-out", vae_encode_profile_path]) + if vae_native_rms_silu: + cmd.append("--flashrt-vae-native-rms-silu") + if vae_t1_conv2d: + cmd.append("--flashrt-vae-t1-conv2d") + if vae_native_avgdown3d: + cmd.append("--flashrt-vae-native-avgdown3d") + if vae_channels_last3d_conv320: + cmd.append("--flashrt-vae-channels-last3d-conv320") + if vae_compile_encode: + cmd.append("--flashrt-vae-compile-encode") + if vae_compile_trace_path is not None: + cmd.extend(["--flashrt-vae-compile-trace-out", vae_compile_trace_path]) + if prepare_dump_path is not None: + cmd.extend(["--flashrt-prepare-dump-out", prepare_dump_path]) + if prepare_replay_path is not None: + cmd.extend(["--flashrt-prepare-replay-in", prepare_replay_path]) + if prepare_inventory_path is not None: + cmd.extend(["--flashrt-prepare-inventory-out", prepare_inventory_path]) + if prepare_slim_no_raw_state_vision: + cmd.append("--flashrt-prepare-slim-no-raw-state-vision") + if prepare_slim_derive_condition_reference: + cmd.append("--flashrt-prepare-slim-derive-condition-reference") + if prepare_slim_derive_initial_noise: + cmd.append("--flashrt-prepare-slim-derive-initial-noise") + if cache_warmup_vae: + cmd.append("--flashrt-cache-warmup-vae") + if cache_warmup_prepare: + cmd.append("--flashrt-cache-warmup-prepare") + if extra_args: + cmd.extend(extra_args) + + run_env = os.environ.copy() + if env: + run_env.update(env) + + t0 = time.perf_counter() + proc = subprocess.run( + cmd, + cwd=str(cosmos_root_path), + env=run_env, + text=True, + capture_output=True, + timeout=timeout, + ) + wall_s = time.perf_counter() - t0 + result: dict[str, Any] = { + "backend": backend, + "returncode": proc.returncode, + "cmd": cmd, + "stdout": proc.stdout, + "stderr": proc.stderr, + "output_dir": str(out_dir), + "wall_s": wall_s, + } + bench = out_dir / "benchmark.json" + if bench.exists(): + try: + result["benchmark"] = json.loads(bench.read_text(encoding="utf-8")) + except json.JSONDecodeError: + result["benchmark_path"] = str(bench) + sample_outputs = list(out_dir.glob("*/sample_outputs.json")) + if sample_outputs: + result["sample_outputs"] = [str(p) for p in sample_outputs] + if backend == "official_action_only": + result["action_only_output"] = True + if live_dump_path is not None: + result["live_dump_out"] = live_dump_path + if live_flashrt_handoff: + result["live_flashrt_handoff"] = True + if live_prelayer_bootstrap: + result["live_prelayer_bootstrap"] = True + if live_warm_request: + result["live_warm_request"] = True + if cache_warmup_vae: + result["cache_warmup_vae"] = True + if cache_warmup_prepare: + result["cache_warmup_prepare"] = True + if warmup: + result["warmup"] = int(warmup) + if live_boundary_path is not None: + result["live_boundary_out"] = live_boundary_path + if live_boundary_in_path is not None: + result["live_boundary_in"] = live_boundary_in_path + if live_boundary_prepare_in_path is not None: + result["live_boundary_prepare_in"] = live_boundary_prepare_in_path + if live_boundary_prepare_live: + result["live_boundary_prepare_live"] = True + if live_trace_path is not None: + result["live_handoff_trace_out"] = live_trace_path + if upstream_trace_path is not None: + result["upstream_trace_out"] = upstream_trace_path + if vae_encode_dump_path is not None: + result["vae_encode_dump_out"] = vae_encode_dump_path + if vae_latent_path is not None: + result["vae_latent_in"] = vae_latent_path + if vae_encode_dump_input: + result["vae_encode_dump_input"] = True + if vae_encode_profile_path is not None: + result["vae_encode_profile_out"] = vae_encode_profile_path + if vae_native_rms_silu: + result["vae_native_rms_silu"] = True + if vae_t1_conv2d: + result["vae_t1_conv2d"] = True + if vae_native_avgdown3d: + result["vae_native_avgdown3d"] = True + if vae_channels_last3d_conv320: + result["vae_channels_last3d_conv320"] = True + if vae_compile_encode: + result["vae_compile_encode"] = True + if vae_compile_trace_path is not None: + result["vae_compile_trace_out"] = vae_compile_trace_path + if prepare_dump_path is not None: + result["prepare_dump_out"] = prepare_dump_path + if prepare_replay_path is not None: + result["prepare_replay_in"] = prepare_replay_path + if prepare_inventory_path is not None: + result["prepare_inventory_out"] = prepare_inventory_path + if prepare_slim_no_raw_state_vision: + result["prepare_slim_no_raw_state_vision"] = True + if prepare_slim_derive_condition_reference: + result["prepare_slim_derive_condition_reference"] = True + if prepare_slim_derive_initial_noise: + result["prepare_slim_derive_initial_noise"] = True + + if proc.returncode != 0: + raise RuntimeError( + "Cosmos3-Edge official baseline failed with return code " + f"{proc.returncode}.\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}" + ) + return result diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 60b74ee6..4ee54985 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -146,6 +146,10 @@ def detect_arch() -> str: ("flash_rt.frontends.torch.qwen3_vl_fp8_sm89_multimodal", "Qwen3VlFp8Sm89Frontend"), + # Cosmos3-Edge official Thor baseline. + ("cosmos3_edge", "torch", "thor"): + ("flash_rt.frontends.torch.cosmos3_edge_thor", "Cosmos3EdgeTorchFrontendThor"), + # ── Nex-N2-mini / Qwen3.6-35B-A3B (qwen3_5_moe) ── # Text LLM, not a VLA: GDN linear-attn + full-attn-every-4th + 256-expert # NVFP4 MoE. RTX 5090 (SM120) only, and requires the gated kernel build diff --git a/flash_rt/models/cosmos3_edge/__init__.py b/flash_rt/models/cosmos3_edge/__init__.py new file mode 100644 index 00000000..8cc188da --- /dev/null +++ b/flash_rt/models/cosmos3_edge/__init__.py @@ -0,0 +1,37 @@ +"""Cosmos3-Edge Thor denoise scaffolding.""" + +from .boundary_dump import EdgeBoundaryDump, EdgeBoundaryShapes +from .denoise_ref import EdgeDenoiseFlashRT, EdgeDenoiseTorchReference, ReferenceDenoiseResult +from .dump_replay import ( + EDGE_ACTION_MODEL_SHAPE, + EDGE_FLAT_DIM, + EDGE_VISION_SHAPE, + EdgeDenoiseDump, + EdgeDenoiseReplay, + EdgeLatentParts, + ReplayResult, +) +from .layer_ref import EdgeLayer0TorchReference, EdgeTransformerFvkLinearReference, EdgeTransformerTorchReference +from .static_engine import EdgeStaticBufferEngine +from .weights import EdgeTransformerWeights, WeightRef + +__all__ = [ + "EDGE_ACTION_MODEL_SHAPE", + "EDGE_FLAT_DIM", + "EDGE_VISION_SHAPE", + "EdgeDenoiseDump", + "EdgeDenoiseReplay", + "EdgeDenoiseFlashRT", + "EdgeDenoiseTorchReference", + "EdgeLatentParts", + "ReferenceDenoiseResult", + "ReplayResult", + "EdgeTransformerWeights", + "WeightRef", + "EdgeBoundaryDump", + "EdgeBoundaryShapes", + "EdgeLayer0TorchReference", + "EdgeTransformerFvkLinearReference", + "EdgeTransformerTorchReference", + "EdgeStaticBufferEngine", +] diff --git a/flash_rt/models/cosmos3_edge/action_only_official.py b/flash_rt/models/cosmos3_edge/action_only_official.py new file mode 100644 index 00000000..2b0a9fbe --- /dev/null +++ b/flash_rt/models/cosmos3_edge/action_only_official.py @@ -0,0 +1,3447 @@ +"""FlashRT-owned Cosmos3-Edge official action-only runner.""" + +from __future__ import annotations + +import json +import os +import sys +import tempfile +import time +from contextlib import contextmanager +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any, Callable, Sequence + + +def _extract_output_dir(argv: Sequence[str]) -> Path: + for idx, arg in enumerate(argv): + if arg in {"-o", "--output-dir"} and idx + 1 < len(argv): + return Path(argv[idx + 1]) + for prefix in ("--output-dir=", "--setup.output-dir="): + if arg.startswith(prefix): + return Path(arg.split("=", 1)[1]) + raise SystemExit("official action-only runner requires -o/--output-dir") + + +def _extract_live_dump_out(argv: Sequence[str]) -> tuple[list[str], Path | None]: + clean: list[str] = [] + live_dump: Path | None = None + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-live-dump-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-live-dump-out requires a path") + live_dump = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-live-dump-out="): + live_dump = Path(arg.split("=", 1)[1]) + idx += 1 + continue + clean.append(arg) + idx += 1 + + env_live_dump = os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_DUMP_OUT") + if live_dump is None and env_live_dump: + live_dump = Path(env_live_dump) + return clean, live_dump + + +def _extract_checkpoint_path(argv: Sequence[str]) -> Path: + for idx, arg in enumerate(argv): + if arg == "--checkpoint-path" and idx + 1 < len(argv): + return Path(argv[idx + 1]) + for prefix in ("--checkpoint-path=", "--setup.checkpoint-path="): + if arg.startswith(prefix): + return Path(arg.split("=", 1)[1]) + raise SystemExit("FlashRT live handoff requires --checkpoint-path") + + +def _extract_live_handoff_args( + argv: Sequence[str], +) -> tuple[list[str], bool, Path | None, Path | None, Path | None, bool, Path | None, bool]: + clean: list[str] = [] + enabled = False + boundary_out: Path | None = None + boundary_in: Path | None = None + boundary_prepare_in: Path | None = None + boundary_prepare_live = False + trace_out: Path | None = None + prelayer_bootstrap = False + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-live-flashrt-handoff": + enabled = True + idx += 1 + continue + if arg == "--flashrt-live-prelayer-bootstrap": + enabled = True + prelayer_bootstrap = True + idx += 1 + continue + if arg == "--flashrt-live-boundary-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-live-boundary-out requires a path") + boundary_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-live-boundary-out="): + boundary_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-live-boundary-in": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-live-boundary-in requires a path") + enabled = True + boundary_in = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-live-boundary-in="): + enabled = True + boundary_in = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-live-boundary-prepare-in": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-live-boundary-prepare-in requires a path") + enabled = True + boundary_prepare_in = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-live-boundary-prepare-in="): + enabled = True + boundary_prepare_in = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-live-boundary-prepare-live": + enabled = True + boundary_prepare_live = True + idx += 1 + continue + if arg == "--flashrt-live-handoff-trace-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-live-handoff-trace-out requires a path") + trace_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-live-handoff-trace-out="): + trace_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + clean.append(arg) + idx += 1 + + if os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_FLASHRT_HANDOFF"): + enabled = True + if os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_PRELAYER_BOOTSTRAP"): + enabled = True + prelayer_bootstrap = True + env_boundary = os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_OUT") + if boundary_out is None and env_boundary: + boundary_out = Path(env_boundary) + env_boundary_in = os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_IN") + if boundary_in is None and env_boundary_in: + enabled = True + boundary_in = Path(env_boundary_in) + env_boundary_prepare = os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_PREPARE_IN") + if boundary_prepare_in is None and env_boundary_prepare: + enabled = True + boundary_prepare_in = Path(env_boundary_prepare) + if os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_PREPARE_LIVE"): + enabled = True + boundary_prepare_live = True + env_trace = os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_HANDOFF_TRACE_OUT") + if trace_out is None and env_trace: + trace_out = Path(env_trace) + return clean, enabled, boundary_out, boundary_in, boundary_prepare_in, boundary_prepare_live, trace_out, prelayer_bootstrap + + +def _extract_upstream_trace_out(argv: Sequence[str]) -> tuple[list[str], Path | None]: + clean: list[str] = [] + trace_out: Path | None = None + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-upstream-trace-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-upstream-trace-out requires a path") + trace_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-upstream-trace-out="): + trace_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + clean.append(arg) + idx += 1 + + env_trace = os.environ.get("FLASHRT_COSMOS3_EDGE_UPSTREAM_TRACE_OUT") + if trace_out is None and env_trace: + trace_out = Path(env_trace) + return clean, trace_out + + +def _extract_warmup_vae_cache(argv: Sequence[str]) -> tuple[list[str], bool]: + clean: list[str] = [] + enabled = False + for arg in argv: + if arg == "--flashrt-cache-warmup-vae": + enabled = True + continue + clean.append(arg) + if os.environ.get("FLASHRT_COSMOS3_EDGE_CACHE_WARMUP_VAE"): + enabled = True + return clean, enabled + + +def _extract_warmup_prepare_cache(argv: Sequence[str]) -> tuple[list[str], bool]: + clean: list[str] = [] + enabled = False + for arg in argv: + if arg == "--flashrt-cache-warmup-prepare": + enabled = True + continue + clean.append(arg) + if os.environ.get("FLASHRT_COSMOS3_EDGE_CACHE_WARMUP_PREPARE"): + enabled = True + return clean, enabled + + +def _extract_vae_encode_boundary_args(argv: Sequence[str]) -> tuple[list[str], Path | None, Path | None, bool]: + clean: list[str] = [] + dump_out: Path | None = None + latent_in: Path | None = None + dump_input = False + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-vae-encode-dump-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-vae-encode-dump-out requires a path") + dump_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-vae-encode-dump-out="): + dump_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-vae-latent-in": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-vae-latent-in requires a path") + latent_in = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-vae-latent-in="): + latent_in = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-vae-encode-dump-input": + dump_input = True + idx += 1 + continue + clean.append(arg) + idx += 1 + + env_dump = os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_ENCODE_DUMP_OUT") + if dump_out is None and env_dump: + dump_out = Path(env_dump) + env_latent = os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_LATENT_IN") + if latent_in is None and env_latent: + latent_in = Path(env_latent) + if os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_ENCODE_DUMP_INPUT"): + dump_input = True + return clean, dump_out, latent_in, dump_input + + +def _extract_vae_encode_profile_out(argv: Sequence[str]) -> tuple[list[str], Path | None]: + clean: list[str] = [] + profile_out: Path | None = None + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-vae-encode-profile-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-vae-encode-profile-out requires a path") + profile_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-vae-encode-profile-out="): + profile_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + clean.append(arg) + idx += 1 + + env_profile = os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_ENCODE_PROFILE_OUT") + if profile_out is None and env_profile: + profile_out = Path(env_profile) + return clean, profile_out + + +def _extract_vae_native_rms_silu(argv: Sequence[str]) -> tuple[list[str], bool]: + clean: list[str] = [] + enabled = False + for arg in argv: + if arg == "--flashrt-vae-native-rms-silu": + enabled = True + continue + clean.append(arg) + if os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_NATIVE_RMS_SILU"): + enabled = True + return clean, enabled + + +def _extract_vae_t1_conv2d(argv: Sequence[str]) -> tuple[list[str], bool]: + clean: list[str] = [] + enabled = False + for arg in argv: + if arg == "--flashrt-vae-t1-conv2d": + enabled = True + continue + clean.append(arg) + if os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_T1_CONV2D"): + enabled = True + return clean, enabled + + +def _extract_vae_native_avgdown3d(argv: Sequence[str]) -> tuple[list[str], bool]: + clean: list[str] = [] + enabled = False + for arg in argv: + if arg == "--flashrt-vae-native-avgdown3d": + enabled = True + continue + clean.append(arg) + if os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_NATIVE_AVGDOWN3D"): + enabled = True + return clean, enabled + + +def _extract_vae_channels_last3d_conv320(argv: Sequence[str]) -> tuple[list[str], bool]: + clean: list[str] = [] + enabled = False + for arg in argv: + if arg == "--flashrt-vae-channels-last3d-conv320": + enabled = True + continue + clean.append(arg) + if os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_CHANNELS_LAST3D_CONV320"): + enabled = True + return clean, enabled + + +def _extract_vae_compile_encode(argv: Sequence[str]) -> tuple[list[str], bool, Path | None]: + clean: list[str] = [] + enabled = False + trace_out: Path | None = None + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-vae-compile-encode": + enabled = True + idx += 1 + continue + if arg == "--flashrt-vae-compile-trace-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-vae-compile-trace-out requires a path") + trace_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-vae-compile-trace-out="): + trace_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + clean.append(arg) + idx += 1 + if os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_COMPILE_ENCODE"): + enabled = True + env_trace = os.environ.get("FLASHRT_COSMOS3_EDGE_VAE_COMPILE_TRACE_OUT") + if trace_out is None and env_trace: + trace_out = Path(env_trace) + return clean, enabled, trace_out + + +def _extract_prepare_boundary_args( + argv: Sequence[str], +) -> tuple[list[str], Path | None, Path | None, Path | None, bool, bool, bool]: + clean: list[str] = [] + dump_out: Path | None = None + replay_in: Path | None = None + inventory_out: Path | None = None + slim_no_raw_state_vision = False + slim_derive_condition_reference = False + slim_derive_initial_noise = False + idx = 0 + while idx < len(argv): + arg = argv[idx] + if arg == "--flashrt-prepare-dump-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-prepare-dump-out requires a path") + dump_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-prepare-dump-out="): + dump_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-prepare-replay-in": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-prepare-replay-in requires a path") + replay_in = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-prepare-replay-in="): + replay_in = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-prepare-inventory-out": + if idx + 1 >= len(argv): + raise SystemExit("--flashrt-prepare-inventory-out requires a path") + inventory_out = Path(argv[idx + 1]) + idx += 2 + continue + if arg.startswith("--flashrt-prepare-inventory-out="): + inventory_out = Path(arg.split("=", 1)[1]) + idx += 1 + continue + if arg == "--flashrt-prepare-slim-no-raw-state-vision": + slim_no_raw_state_vision = True + idx += 1 + continue + if arg == "--flashrt-prepare-slim-derive-condition-reference": + slim_derive_condition_reference = True + idx += 1 + continue + if arg == "--flashrt-prepare-slim-derive-initial-noise": + slim_derive_initial_noise = True + idx += 1 + continue + clean.append(arg) + idx += 1 + env_dump = os.environ.get("FLASHRT_COSMOS3_EDGE_PREPARE_DUMP_OUT") + if dump_out is None and env_dump: + dump_out = Path(env_dump) + env_replay = os.environ.get("FLASHRT_COSMOS3_EDGE_PREPARE_REPLAY_IN") + if replay_in is None and env_replay: + replay_in = Path(env_replay) + env_inventory = os.environ.get("FLASHRT_COSMOS3_EDGE_PREPARE_INVENTORY_OUT") + if inventory_out is None and env_inventory: + inventory_out = Path(env_inventory) + if os.environ.get("FLASHRT_COSMOS3_EDGE_PREPARE_SLIM_NO_RAW_STATE_VISION"): + slim_no_raw_state_vision = True + if os.environ.get("FLASHRT_COSMOS3_EDGE_PREPARE_SLIM_DERIVE_CONDITION_REFERENCE"): + slim_derive_condition_reference = True + if os.environ.get("FLASHRT_COSMOS3_EDGE_PREPARE_SLIM_DERIVE_INITIAL_NOISE"): + slim_derive_initial_noise = True + return ( + clean, + dump_out, + replay_in, + inventory_out, + slim_no_raw_state_vision, + slim_derive_condition_reference, + slim_derive_initial_noise, + ) + + +class _VAECompileEncodePatch: + def __init__(self, trace_out: Path | None) -> None: + self.trace_out = trace_out + self.events: list[dict[str, Any]] = [] + + def wrap_setup(self, original: Callable[..., Any]) -> Callable[..., Any]: + def setup_with_compile(model: Any, *args: Any, **kwargs: Any) -> Any: + out = original(model, *args, **kwargs) + tokenizer = getattr(model, "tokenizer_vision_gen", None) + if tokenizer is None or not hasattr(tokenizer, "compile_encode"): + self.events.append({"name": "vae_compile_encode", "ok": False, "reason": "missing_tokenizer"}) + return out + t0 = time.perf_counter() + with tempfile.TemporaryDirectory(prefix="cosmos3_edge_vae_aot_") as tmp: + tokenizer.compile_encode(["480"], tmp, aspect_ratio="16,9") + loaded = getattr(getattr(tokenizer, "model", None), "model", None) + loaded_fns = getattr(loaded, "_aot_chunk_fns", {}) + self.events.append( + { + "name": "vae_compile_encode", + "ok": True, + "compile_s": time.perf_counter() - t0, + "resolution": "480", + "aspect_ratio": "16,9", + "loaded_aot_functions": len(loaded_fns) if isinstance(loaded_fns, dict) else 0, + } + ) + return out + + return setup_with_compile + + def save(self) -> None: + if self.trace_out is None: + return + self.trace_out.parent.mkdir(parents=True, exist_ok=True) + self.trace_out.write_text(json.dumps({"version": 1, "events": self.events}, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _install_vae_compile_encode_patch(trace_out: Path | None) -> _VAECompileEncodePatch: + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + patch = _VAECompileEncodePatch(trace_out) + OmniMoTModel.set_up_tokenizers = patch.wrap_setup(OmniMoTModel.set_up_tokenizers) + return patch + + +def _fake_decode(self: Any, vision_latent: Any) -> Any: + tokenizer = self.tokenizer_vision_gen + frames = tokenizer.get_pixel_num_frames(int(vision_latent.shape[2])) + height = int(vision_latent.shape[3]) * int(tokenizer.spatial_compression_factor) + width = int(vision_latent.shape[4]) * int(tokenizer.spatial_compression_factor) + return vision_latent.new_zeros((vision_latent.shape[0], 3, frames, height, width)) + + +def _fake_save_img_or_video(vision: Any, path_without_suffix: str, *, fps: int = 24, quality: int = 7) -> None: + del fps, quality + path = Path(path_without_suffix) + frames = int(vision.shape[1]) if hasattr(vision, "shape") and len(vision.shape) > 1 else 2 + output = path.with_suffix(".png" if frames == 1 else ".mp4") + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(b"") + + +def _rewrite_sample_outputs_action_only(output_dir: Path) -> int: + rewritten = 0 + for sample_outputs in output_dir.glob("*/sample_outputs.json"): + payload = json.loads(sample_outputs.read_text(encoding="utf-8")) + for output in payload.get("outputs", []): + content = output.get("content") + if not isinstance(content, dict) or "action" not in content: + continue + for rel in output.get("files", []): + try: + path = sample_outputs.parent / rel + if path.exists(): + path.unlink() + except OSError: + pass + output["files"] = [] + rewritten += 1 + sample_outputs.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return rewritten + + +def _install_action_only_patches() -> None: + import cosmos_framework.inference.inference as inference_mod + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + inference_mod.save_img_or_video = _fake_save_img_or_video + OmniMoTModel.decode = _fake_decode + + +def _clone_for_warmup(value: Any) -> Any: + try: + import torch + except Exception: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + return value.clone() + if is_dataclass(value) and not isinstance(value, type): + return type(value)(**{field.name: _clone_for_warmup(getattr(value, field.name)) for field in fields(value)}) + if isinstance(value, dict): + return {key: _clone_for_warmup(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_for_warmup(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_for_warmup(item) for item in value) + return value + + +def _install_warmup_clone_patch() -> None: + from cosmos_framework.inference.inference import OmniInference + + original = OmniInference.generate_batch + + def generate_batch_clone_warmup(self: Any, sample_args_list: Any, data_batch: dict[str, Any], *, warmup: bool = False): + if warmup: + data_batch = _clone_for_warmup(data_batch) + return original(self, sample_args_list, data_batch, warmup=warmup) + + OmniInference.generate_batch = generate_batch_clone_warmup + + +class _WarmupVAECache: + def __init__(self) -> None: + self.phase = "setup" + self.cached: list[tuple[tuple[Any, ...], Any]] = [] + self.events: list[dict[str, Any]] = [] + + @staticmethod + def _signature(state: Any) -> tuple[Any, ...] | None: + if not hasattr(state, "shape") or not hasattr(state, "dtype") or not hasattr(state, "device"): + return None + signature: list[Any] = [ + tuple(int(dim) for dim in state.shape), + str(state.dtype), + str(state.device), + int(state.numel()) if hasattr(state, "numel") else None, + ] + # Guard against same-shape different-video cache hits. A full content + # hash would erase too much of the saved VAE time, so sample stable + # positions across the flattened tensor. + try: + import torch + + if isinstance(state, torch.Tensor) and state.numel() > 0: + flat = state.detach().reshape(-1) + sample_count = min(32, int(flat.numel())) + if sample_count == 1: + index_values = [0] + else: + last = int(flat.numel()) - 1 + index_values = [ + (idx * last + (sample_count - 1) // 2) // (sample_count - 1) + for idx in range(sample_count) + ] + indexes = torch.tensor(index_values, device=flat.device, dtype=torch.long) + sample = flat.index_select(0, indexes).to(device="cpu", dtype=torch.float32) + signature.append(tuple(float(value) for value in sample.tolist())) + except Exception: + signature.append("content_signature_unavailable") + return tuple(signature) + + def wrap_generate_batch(self, original: Callable[..., Any]) -> Callable[..., Any]: + def generate_batch_with_phase(self_obj: Any, sample_args_list: Any, data_batch: dict[str, Any], *, warmup: bool = False): + previous = self.phase + self.phase = "warmup" if warmup else "measured" + try: + return original(self_obj, sample_args_list, data_batch, warmup=warmup) + finally: + self.phase = previous + + return generate_batch_with_phase + + def wrap_encode(self, original: Callable[..., Any]) -> Callable[..., Any]: + def encode_with_cache(model: Any, state: Any) -> Any: + signature = self._signature(state) + if self.phase == "measured" and signature is not None and self.cached: + cached_signature, cached_value = self.cached.pop(0) + if cached_signature == signature: + self.events.append( + {"phase": self.phase, "name": "vae_cache_hit", "s": 0.0, "ok": True, "signature": str(signature)} + ) + return cached_value.clone() + self.events.append( + { + "phase": self.phase, + "name": "vae_cache_miss_signature", + "s": 0.0, + "ok": True, + "expected": str(cached_signature), + "actual": str(signature), + } + ) + encoded = original(model, state) + if self.phase == "warmup" and signature is not None: + self.cached.append((signature, encoded.detach().clone())) + self.events.append( + {"phase": self.phase, "name": "vae_cache_store", "s": 0.0, "ok": True, "signature": str(signature)} + ) + elif self.phase == "measured": + self.events.append( + { + "phase": self.phase, + "name": "vae_cache_miss_empty", + "s": 0.0, + "ok": True, + "signature": str(signature), + } + ) + return encoded + + return encode_with_cache + + +def _signature_for_value(value: Any) -> Any: + try: + import torch + except Exception: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + return _WarmupVAECache._signature(value) + if isinstance(value, dict): + return tuple((str(key), _signature_for_value(item)) for key, item in sorted(value.items(), key=lambda kv: str(kv[0]))) + if isinstance(value, (list, tuple)): + return tuple(_signature_for_value(item) for item in value) + if is_dataclass(value) and not isinstance(value, type): + return tuple((field.name, _signature_for_value(getattr(value, field.name))) for field in fields(value)) + if isinstance(value, (str, int, float, bool, type(None))): + return value + return repr(value) + + +class _WarmupPrepareCache: + def __init__( + self, + *, + slim_no_raw_state_vision: bool = False, + slim_derive_condition_reference: bool = False, + slim_derive_initial_noise: bool = False, + ) -> None: + self.phase = "setup" + self.cached: list[tuple[tuple[Any, ...], Any]] = [] + self.slim_no_raw_state_vision = slim_no_raw_state_vision + self.slim_derive_condition_reference = slim_derive_condition_reference + self.slim_derive_initial_noise = slim_derive_initial_noise + self.events: list[dict[str, Any]] = [] + + @staticmethod + def _signature(data_batch: Any, seed: Any, has_negative_prompt: Any) -> tuple[Any, ...]: + return ( + _signature_for_value(data_batch), + _signature_for_value(seed), + bool(has_negative_prompt), + ) + + def wrap_generate_batch(self, original: Callable[..., Any]) -> Callable[..., Any]: + def generate_batch_with_phase(self_obj: Any, sample_args_list: Any, data_batch: dict[str, Any], *, warmup: bool = False): + previous = self.phase + self.phase = "warmup" if warmup else "measured" + try: + return original(self_obj, sample_args_list, data_batch, warmup=warmup) + finally: + self.phase = previous + + return generate_batch_with_phase + + def wrap_prepare(self, original: Callable[..., Any]) -> Callable[..., Any]: + def prepare_with_cache(model: Any, *args: Any, **kwargs: Any) -> Any: + data_batch = args[0] if args else kwargs.get("data_batch") + seed = args[1] if len(args) > 1 else kwargs.get("seed") + has_negative_prompt = ( + args[2] if len(args) > 2 else kwargs.get("has_negative_prompt", False) + ) + signature = self._signature(data_batch, seed, has_negative_prompt) + if self.phase == "measured" and self.cached: + cached_signature, cached_value = self.cached.pop(0) + if cached_signature == signature: + self.events.append( + { + "phase": self.phase, + "name": "prepare_cache_hit", + "s": 0.0, + "ok": True, + "slim_no_raw_state_vision": self.slim_no_raw_state_vision, + "slim_derive_condition_reference": self.slim_derive_condition_reference, + "slim_derive_initial_noise": self.slim_derive_initial_noise, + } + ) + payload = _prepare_payload_with_derived_initial_noise(_clone_for_warmup(cached_value), seed) + return _prepare_payload_with_derived_condition_reference(payload) + self.events.append( + { + "phase": self.phase, + "name": "prepare_cache_miss_signature", + "s": 0.0, + "ok": True, + "expected": str(cached_signature), + "actual": str(signature), + } + ) + result = original(model, *args, **kwargs) + if self.phase == "warmup": + cached_value = _slim_prepare_payload( + result, + no_raw_state_vision=self.slim_no_raw_state_vision, + derive_condition_reference=self.slim_derive_condition_reference, + derive_initial_noise=self.slim_derive_initial_noise, + ) + self.cached.append((signature, cached_value)) + self.events.append( + { + "phase": self.phase, + "name": "prepare_cache_store", + "s": 0.0, + "ok": True, + "slim_no_raw_state_vision": self.slim_no_raw_state_vision, + "slim_derive_condition_reference": self.slim_derive_condition_reference, + "slim_derive_initial_noise": self.slim_derive_initial_noise, + } + ) + elif self.phase == "measured": + self.events.append( + {"phase": self.phase, "name": "prepare_cache_miss_empty", "s": 0.0, "ok": True} + ) + return result + + return prepare_with_cache + + +def _first_tensor_device(value: Any) -> str | None: + try: + import torch + except Exception: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + return str(value.device) + if isinstance(value, dict): + for item in value.values(): + device = _first_tensor_device(item) + if device is not None: + return device + if isinstance(value, (list, tuple)): + for item in value: + device = _first_tensor_device(item) + if device is not None: + return device + if is_dataclass(value) and not isinstance(value, type): + for field in fields(value): + device = _first_tensor_device(getattr(value, field.name)) + if device is not None: + return device + return None + + +def _prepare_model_device(model: Any, data_batch: Any) -> str | None: + tensor_kwargs = getattr(model, "tensor_kwargs", None) + if isinstance(tensor_kwargs, dict) and tensor_kwargs.get("device") is not None: + return str(tensor_kwargs["device"]) + tensor_kwargs_fp32 = getattr(model, "tensor_kwargs_fp32", None) + if isinstance(tensor_kwargs_fp32, dict) and tensor_kwargs_fp32.get("device") is not None: + return str(tensor_kwargs_fp32["device"]) + return _first_tensor_device(data_batch) + + +_PREPARE_PAYLOAD_FIELD_NAMES = ( + "sequence_plans", + "gen_data_clean", + "cond_text_tokens", + "uncond_text_tokens", + "initial_noise", + "condition_reference", + "condition_mask", +) + + +def _json_safe_scalar(value: Any) -> Any: + if isinstance(value, (str, int, float, bool, type(None))): + return value + return repr(value) + + +def _tensor_inventory(path: str, tensor: Any) -> dict[str, Any]: + return { + "path": path, + "shape": [int(dim) for dim in tensor.shape], + "dtype": str(tensor.dtype), + "device": str(tensor.device), + "stride": [int(dim) for dim in tensor.stride()], + "contiguous": bool(tensor.is_contiguous()), + "numel": int(tensor.numel()), + "element_size": int(tensor.element_size()), + "bytes": int(tensor.numel()) * int(tensor.element_size()), + } + + +def _prepare_inventory_value( + value: Any, + *, + path: str, + tensors: list[dict[str, Any]], + seen: set[int], + depth: int = 0, + max_depth: int = 6, + max_children: int = 16, +) -> Any: + try: + import torch + except Exception: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + info = _tensor_inventory(path, value) + tensors.append(info) + return {key: info[key] for key in ("shape", "dtype", "device", "stride", "contiguous", "numel", "bytes")} + if isinstance(value, (str, int, float, bool, type(None))): + return {"type": type(value).__name__, "value": _json_safe_scalar(value)} + if id(value) in seen: + return {"type": value.__class__.__name__, "cycle": True} + seen.add(id(value)) + if depth >= max_depth: + return {"type": value.__class__.__name__, "truncated": "max_depth"} + if is_dataclass(value) and not isinstance(value, type): + children: dict[str, Any] = {} + for field in fields(value): + field_value = getattr(value, field.name) + children[field.name] = _prepare_inventory_value( + field_value, + path=f"{path}.{field.name}", + tensors=tensors, + seen=seen, + depth=depth + 1, + max_depth=max_depth, + max_children=max_children, + ) + return { + "type": value.__class__.__name__, + "module": value.__class__.__module__, + "fields": children, + } + if isinstance(value, dict): + children = {} + for idx, (key, item) in enumerate(value.items()): + if idx >= max_children: + break + key_text = str(key) + children[key_text] = _prepare_inventory_value( + item, + path=f"{path}[{key_text!r}]", + tensors=tensors, + seen=seen, + depth=depth + 1, + max_depth=max_depth, + max_children=max_children, + ) + return { + "type": "dict", + "len": len(value), + "truncated_children": max(0, len(value) - len(children)), + "children": children, + } + if isinstance(value, (list, tuple)): + children = [] + for idx, item in enumerate(value[:max_children]): + children.append( + _prepare_inventory_value( + item, + path=f"{path}[{idx}]", + tensors=tensors, + seen=seen, + depth=depth + 1, + max_depth=max_depth, + max_children=max_children, + ) + ) + return { + "type": type(value).__name__, + "len": len(value), + "truncated_children": max(0, len(value) - len(children)), + "children": children, + } + return {"type": value.__class__.__name__, "repr": repr(value)[:256]} + + +def _prepare_field_name(index: int) -> str: + if index < len(_PREPARE_PAYLOAD_FIELD_NAMES): + return _PREPARE_PAYLOAD_FIELD_NAMES[index] + return f"field_{index}" + + +def _prepare_payload_inventory(payload: Any, *, signature: Any = None, source: str | None = None) -> dict[str, Any]: + tensors: list[dict[str, Any]] = [] + seen: set[int] = set() + if isinstance(payload, tuple): + root = { + _prepare_field_name(idx): _prepare_inventory_value( + item, + path=_prepare_field_name(idx), + tensors=tensors, + seen=seen, + ) + for idx, item in enumerate(payload) + } + else: + root = _prepare_inventory_value(payload, path="payload", tensors=tensors, seen=seen) + + by_top_level: dict[str, dict[str, Any]] = {} + by_dtype: dict[str, dict[str, Any]] = {} + for tensor in tensors: + top = str(tensor["path"]).split(".", 1)[0].split("[", 1)[0] + field_item = by_top_level.setdefault(top, {"field": top, "tensor_count": 0, "bytes": 0}) + field_item["tensor_count"] += 1 + field_item["bytes"] += int(tensor["bytes"]) + dtype = str(tensor["dtype"]) + dtype_item = by_dtype.setdefault(dtype, {"dtype": dtype, "tensor_count": 0, "bytes": 0}) + dtype_item["tensor_count"] += 1 + dtype_item["bytes"] += int(tensor["bytes"]) + for items in (by_top_level.values(), by_dtype.values()): + for item in items: + item["mib"] = float(item["bytes"]) / (1024.0 * 1024.0) + total_bytes = sum(int(tensor["bytes"]) for tensor in tensors) + return { + "version": 1, + "schema": "flashrt_cosmos3_edge_prepare_inventory_v1", + "source": source, + "signature": signature, + "field_names": list(_PREPARE_PAYLOAD_FIELD_NAMES), + "payload_type": payload.__class__.__name__, + "tensor_count": len(tensors), + "tensor_bytes": total_bytes, + "tensor_mib": float(total_bytes) / (1024.0 * 1024.0), + "tensor_bytes_by_dtype": sorted(by_dtype.values(), key=lambda item: -int(item["bytes"])), + "tensor_bytes_by_top_level": sorted(by_top_level.values(), key=lambda item: -int(item["bytes"])), + "largest_tensors": sorted(tensors, key=lambda item: -int(item["bytes"]))[:32], + "tree": root, + } + + +def _write_prepare_inventory(path: Path, payload: Any, *, signature: Any = None, source: str | None = None) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + _prepare_payload_inventory(payload, signature=signature, source=source), + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + + +def _as_python_int(value: Any) -> int: + if hasattr(value, "item"): + return int(value.item()) + return int(value) + + +def _indexed_optional_int(values: Any, index: int) -> int | None: + if values is None: + return None + value = values[index] if isinstance(values, (list, tuple)) else values + if value is None: + return None + return _as_python_int(value) + + +def _indexed_optional_float(values: Any, index: int) -> float | None: + if values is None: + return None + if isinstance(values, (list, tuple)): + value = values[index] + else: + try: + value = values[index] + except (IndexError, TypeError): + value = values + if value is None: + return None + if hasattr(value, "item"): + return float(value.item()) + return float(value) + + +def _num_vision_items_for_sample(values: Any, index: int) -> int: + if values is None: + return 1 + value = values[index] + if isinstance(value, (list, tuple)): + return len(value) + return _as_python_int(value) + + +def _derive_condition_reference_from_prepare_payload(payload: Any) -> list[Any]: + if not isinstance(payload, tuple) or len(payload) <= 5: + raise RuntimeError("prepare payload cannot derive condition_reference from a non-v1 tuple") + sequence_plans = payload[0] + gen_data_clean = payload[1] + initial_noise = payload[4] + if not isinstance(sequence_plans, list) or not isinstance(initial_noise, list): + raise RuntimeError("prepare payload cannot derive condition_reference without sequence_plans and initial_noise") + + x0_tokens_vision = getattr(gen_data_clean, "x0_tokens_vision", None) + x0_tokens_action = getattr(gen_data_clean, "x0_tokens_action", None) + x0_tokens_sound = getattr(gen_data_clean, "x0_tokens_sound", None) + raw_action_dim = getattr(gen_data_clean, "raw_action_dim", None) + num_vision_items_per_sample = getattr(gen_data_clean, "num_vision_items_per_sample", None) + + references: list[Any] = [] + idx_vision = 0 + idx_action = 0 + idx_sound = 0 + for sample_idx, plan in enumerate(sequence_plans): + parts = [] + sample_noise = initial_noise[sample_idx] + if getattr(plan, "has_vision", False): + if x0_tokens_vision is None: + raise RuntimeError("prepare payload cannot derive vision condition_reference without x0_tokens_vision") + for _ in range(_num_vision_items_for_sample(num_vision_items_per_sample, sample_idx)): + x0_vision = x0_tokens_vision[idx_vision].to(dtype=sample_noise.dtype, device=sample_noise.device) + parts.append(x0_vision.reshape(-1)) + idx_vision += 1 + if getattr(plan, "has_action", False): + if x0_tokens_action is None: + raise RuntimeError("prepare payload cannot derive action condition_reference without x0_tokens_action") + x0_action = x0_tokens_action[idx_action].to(dtype=sample_noise.dtype, device=sample_noise.device) + action_dim = _indexed_optional_int(raw_action_dim, idx_action) + if action_dim is not None: + x0_action = x0_action.clone() + x0_action[:, action_dim:] = 0 + parts.append(x0_action.reshape(-1)) + idx_action += 1 + if getattr(plan, "has_sound", False): + if x0_tokens_sound is None: + raise RuntimeError("prepare payload cannot derive sound condition_reference without x0_tokens_sound") + x0_sound = x0_tokens_sound[idx_sound].to(dtype=sample_noise.dtype, device=sample_noise.device) + parts.append(x0_sound.reshape(-1)) + idx_sound += 1 + if not parts: + raise RuntimeError("prepare payload cannot derive an empty condition_reference sample") + import torch + + references.append(torch.cat(parts, dim=0)) + return references + + +def _seed_list_for_prepare(seed: Any, n_sample: int) -> list[int]: + if seed is None: + raise RuntimeError("prepare payload cannot derive initial_noise without seed") + if isinstance(seed, (list, tuple)): + values = list(seed) + elif hasattr(seed, "tolist"): + raw = seed.tolist() + values = raw if isinstance(raw, list) else [raw] + else: + values = [seed] + if len(values) != n_sample: + raise RuntimeError(f"prepare initial_noise seed count mismatch: got {len(values)} expected {n_sample}") + return [_as_python_int(value) for value in values] + + +def _arch_invariant_rand_like(shape: Any, *, dtype: Any, device: Any, seed: int) -> Any: + import numpy as np + import torch + + random_array = np.random.RandomState(seed).standard_normal(tuple(int(dim) for dim in shape)).astype(np.float32) + return torch.from_numpy(random_array).to(dtype=dtype, device=device) + + +def _append_derived_noise_part( + parts: list[Any], + *, + x0: Any, + flat_mask: Any, + offset: int, + seed: int, + pure_noise_dtype: Any, + raw_action_dim: int | None = None, +) -> int: + numel = int(x0.numel()) + mask = flat_mask[offset : offset + numel].reshape(x0.shape).to(device=x0.device, dtype=flat_mask.dtype) + pure_noise = _arch_invariant_rand_like(x0.shape, dtype=pure_noise_dtype, device=x0.device, seed=seed) + x0_t = x0.to(device=x0.device, dtype=pure_noise_dtype) + noise = mask * x0_t + (1.0 - mask) * pure_noise + if raw_action_dim is not None: + noise = noise.clone() + noise[:, raw_action_dim:] = 0 + parts.append(noise.reshape(-1)) + return offset + numel + + +def _derive_initial_noise_from_prepare_payload(payload: Any, seed: Any) -> list[Any]: + if not isinstance(payload, tuple) or len(payload) <= 6: + raise RuntimeError("prepare payload cannot derive initial_noise from a non-v1 tuple") + sequence_plans = payload[0] + gen_data_clean = payload[1] + condition_mask = payload[6] + if not isinstance(sequence_plans, list) or not isinstance(condition_mask, list): + raise RuntimeError("prepare payload cannot derive initial_noise without sequence_plans and condition_mask") + + x0_tokens_vision = getattr(gen_data_clean, "x0_tokens_vision", None) + x0_tokens_action = getattr(gen_data_clean, "x0_tokens_action", None) + x0_tokens_sound = getattr(gen_data_clean, "x0_tokens_sound", None) + raw_action_dim = getattr(gen_data_clean, "raw_action_dim", None) + num_vision_items_per_sample = getattr(gen_data_clean, "num_vision_items_per_sample", None) + seeds = _seed_list_for_prepare(seed, len(sequence_plans)) + + noises: list[Any] = [] + idx_vision = 0 + idx_action = 0 + idx_sound = 0 + for sample_idx, plan in enumerate(sequence_plans): + flat_mask = condition_mask[sample_idx] + offset = 0 + parts: list[Any] = [] + sample_seed = seeds[sample_idx] + if getattr(plan, "has_vision", False): + if x0_tokens_vision is None: + raise RuntimeError("prepare payload cannot derive vision initial_noise without x0_tokens_vision") + for _ in range(_num_vision_items_for_sample(num_vision_items_per_sample, sample_idx)): + x0_vision = x0_tokens_vision[idx_vision] + offset = _append_derived_noise_part( + parts, + x0=x0_vision, + flat_mask=flat_mask, + offset=offset, + seed=sample_seed, + pure_noise_dtype=flat_mask.dtype, + ) + idx_vision += 1 + if getattr(plan, "has_action", False): + if x0_tokens_action is None: + raise RuntimeError("prepare payload cannot derive action initial_noise without x0_tokens_action") + x0_action = x0_tokens_action[idx_action] + offset = _append_derived_noise_part( + parts, + x0=x0_action, + flat_mask=flat_mask, + offset=offset, + seed=sample_seed, + pure_noise_dtype=x0_action.dtype, + raw_action_dim=_indexed_optional_int(raw_action_dim, idx_action), + ) + idx_action += 1 + if getattr(plan, "has_sound", False): + if x0_tokens_sound is None: + raise RuntimeError("prepare payload cannot derive sound initial_noise without x0_tokens_sound") + x0_sound = x0_tokens_sound[idx_sound] + offset = _append_derived_noise_part( + parts, + x0=x0_sound, + flat_mask=flat_mask, + offset=offset, + seed=sample_seed, + pure_noise_dtype=x0_sound.dtype, + ) + idx_sound += 1 + if not parts: + raise RuntimeError("prepare payload cannot derive an empty initial_noise sample") + if offset != int(flat_mask.numel()): + raise RuntimeError( + f"prepare initial_noise derivation consumed {offset} mask values, expected {int(flat_mask.numel())}" + ) + import torch + + noises.append(torch.cat(parts, dim=0).to(dtype=flat_mask.dtype, device=flat_mask.device)) + return noises + + +def _replace_prepare_payload_field(payload: Any, index: int, value: Any) -> Any: + if not isinstance(payload, tuple) or index >= len(payload): + raise RuntimeError("prepare payload cannot replace field in a non-v1 tuple") + items = list(payload) + items[index] = value + return tuple(items) + + +def _prepare_payload_with_derived_condition_reference(payload: Any) -> Any: + if isinstance(payload, tuple) and len(payload) > 5 and payload[5] is None: + return _replace_prepare_payload_field(payload, 5, _derive_condition_reference_from_prepare_payload(payload)) + return payload + + +def _prepare_payload_with_derived_initial_noise(payload: Any, seed: Any) -> Any: + if isinstance(payload, tuple) and len(payload) > 4 and payload[4] is None: + return _replace_prepare_payload_field(payload, 4, _derive_initial_noise_from_prepare_payload(payload, seed)) + return payload + + +def _flat_parts_for_edge_av(flat: Any) -> tuple[Any, Any]: + from flash_rt.models.cosmos3_edge.dump_replay import EDGE_ACTION_MODEL_SHAPE, EDGE_FLAT_DIM, EDGE_VISION_SHAPE + + if int(flat.numel()) != EDGE_FLAT_DIM: + raise RuntimeError(f"expected Edge flat latent dim {EDGE_FLAT_DIM}, got {int(flat.numel())}") + vision_dim = 1 + for dim in EDGE_VISION_SHAPE: + vision_dim *= int(dim) + flat = flat.reshape(-1) + return flat[:vision_dim].reshape(EDGE_VISION_SHAPE), flat[vision_dim:].reshape(EDGE_ACTION_MODEL_SHAPE) + + +def _edge_vae_mrope_ids( + *, + grid_t: int, + grid_h: int, + grid_w: int, + temporal_offset: int | float, + fps: float, + base_fps: float, + temporal_compression_factor: int, + base_temporal_compression_factor: int | None = None, + start_frame_offset: int = 0, + device: Any, +) -> tuple[Any, int | float]: + import math + import torch + + effective_base_tcf = ( + base_temporal_compression_factor + if base_temporal_compression_factor is not None + else temporal_compression_factor + ) + tps = fps / temporal_compression_factor + base_tps = base_fps / effective_base_tcf + frame_indices = torch.arange(grid_t, dtype=torch.float32, device=device) + scaled_t = (frame_indices + start_frame_offset) / tps * base_tps + temporal_offset + t_index = scaled_t.view(-1, 1).expand(-1, grid_h * grid_w).flatten() + h_index = torch.arange(grid_h, dtype=torch.long, device=device).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() + w_index = torch.arange(grid_w, dtype=torch.long, device=device).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten() + mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0) + return mrope_ids, math.ceil(mrope_ids.max().item()) + 1 + + +def _derive_step0_position_ids_from_prepare_payload(payload: Any) -> Any: + """Derive Edge AV step-0 3D MRoPE position ids from slim prepare state.""" + + import math + import torch + + if not isinstance(payload, tuple) or len(payload) <= 3: + raise RuntimeError("prepare payload cannot derive position_ids from a non-v1 tuple") + sequence_plans = payload[0] + gen_data_clean = payload[1] + cond_text_tokens = payload[2] + if not isinstance(sequence_plans, list) or len(sequence_plans) != 1: + raise RuntimeError("position_ids derivation currently supports batch size 1") + if not isinstance(cond_text_tokens, list) or len(cond_text_tokens) != 1: + raise RuntimeError("position_ids derivation requires one cond_text_tokens entry") + + plan = sequence_plans[0] + if not (getattr(plan, "has_text", False) and getattr(plan, "has_vision", False) and getattr(plan, "has_action", False)): + raise RuntimeError("position_ids derivation currently supports Edge AV text+vision+action packs") + if getattr(plan, "has_sound", False): + raise RuntimeError("position_ids derivation does not support sound packs") + if getattr(plan, "share_vision_temporal_positions", False): + raise RuntimeError("position_ids derivation does not support shared vision temporal positions") + + x0_tokens_vision = getattr(gen_data_clean, "x0_tokens_vision", None) + x0_tokens_action = getattr(gen_data_clean, "x0_tokens_action", None) + if not isinstance(x0_tokens_vision, list) or len(x0_tokens_vision) != 1: + raise RuntimeError("position_ids derivation requires one vision token tensor") + if not isinstance(x0_tokens_action, list) or len(x0_tokens_action) != 1: + raise RuntimeError("position_ids derivation requires one action token tensor") + + vision = x0_tokens_vision[0] + action = x0_tokens_action[0] + if len(tuple(vision.shape)) != 5: + raise RuntimeError(f"unexpected vision token shape for position_ids derivation: {tuple(vision.shape)}") + if len(tuple(action.shape)) != 2: + raise RuntimeError(f"unexpected action token shape for position_ids derivation: {tuple(action.shape)}") + + device = vision.device + latent_patch_size = 2 + temporal_compression_factor = 4 + base_fps = 24.0 + temporal_modality_margin = 15000 + text_len = len(cond_text_tokens[0]) + 2 # official Edge packs EOS and start-of-generation after prompt ids. + _, _, latent_t, latent_h, latent_w = vision.shape + patch_h = math.ceil(int(latent_h) / latent_patch_size) + patch_w = math.ceil(int(latent_w) / latent_patch_size) + action_len = int(action.shape[0]) + + vision_fps = _indexed_optional_float(getattr(gen_data_clean, "fps_vision", None), 0) + action_fps = _indexed_optional_float(getattr(gen_data_clean, "fps_action", None), 0) + if vision_fps is None or action_fps is None: + raise RuntimeError("position_ids derivation requires fps_vision and fps_action") + + text_ids = torch.arange(text_len, dtype=torch.float32, device=device).unsqueeze(0).expand(3, -1).contiguous() + vision_temporal_offset = text_len + temporal_modality_margin + vision_ids, _ = _edge_vae_mrope_ids( + grid_t=int(latent_t), + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=vision_temporal_offset, + fps=vision_fps, + base_fps=base_fps, + temporal_compression_factor=temporal_compression_factor, + device=device, + ) + action_ids, _ = _edge_vae_mrope_ids( + grid_t=action_len, + grid_h=1, + grid_w=1, + temporal_offset=vision_temporal_offset, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + base_temporal_compression_factor=temporal_compression_factor, + start_frame_offset=_as_python_int(getattr(plan, "action_start_frame_offset", 1)), + device=device, + ) + return torch.cat([text_ids, vision_ids, action_ids], dim=1).contiguous() + + +def _derive_step0_rope_from_position_ids(position_ids: Any, causal_indices: Any, full_indices: Any) -> dict[str, Any]: + import torch + + head_dim = 128 + rope_theta = 100000000.0 + mrope_section = [24, 20, 20] + device = position_ids.device + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64, device=device).to(dtype=torch.float32) / head_dim) + ) + expanded_position_ids = position_ids.unsqueeze(1) if position_ids.ndim == 2 else position_ids + inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, expanded_position_ids.shape[1], -1, 1) + position_ids_expanded = expanded_position_ids[:, :, None, :].float() + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs_t = freqs[0].clone() + for dim, offset in enumerate((1, 2), start=1): + length = mrope_section[dim] * 3 + freqs_t[..., slice(offset, length, 3)] = freqs[dim, ..., slice(offset, length, 3)] + emb = torch.cat((freqs_t, freqs_t), dim=-1) + cos = emb.cos().squeeze(0).to(dtype=torch.bfloat16).contiguous() + sin = emb.sin().squeeze(0).to(dtype=torch.bfloat16).contiguous() + causal_long = causal_indices.to(dtype=torch.long) + full_long = full_indices.to(dtype=torch.long) + return { + "s00/layers/00/rope/cos/causal_seq": cos.index_select(0, causal_long).contiguous(), + "s00/layers/00/rope/cos/full_only_seq": cos.index_select(0, full_long).contiguous(), + "s00/layers/00/rope/sin/causal_seq": sin.index_select(0, causal_long).contiguous(), + "s00/layers/00/rope/sin/full_only_seq": sin.index_select(0, full_long).contiguous(), + } + + +def _derive_step0_causal_seq_from_prepare_payload(payload: Any, text_embedding_weight: Any) -> Any: + import torch + + if not isinstance(payload, tuple) or len(payload) <= 2: + raise RuntimeError("prepare payload cannot derive causal_seq from a non-v1 tuple") + cond_text_tokens = payload[2] + if not isinstance(cond_text_tokens, list) or len(cond_text_tokens) != 1: + raise RuntimeError("causal_seq derivation requires one cond_text_tokens entry") + if len(tuple(text_embedding_weight.shape)) != 2 or int(text_embedding_weight.shape[1]) != 2048: + raise RuntimeError(f"unexpected text embedding shape: {tuple(text_embedding_weight.shape)}") + + eos_token_id = 11 + start_of_generation_token_id = 20 + token_ids = list(cond_text_tokens[0]) + [eos_token_id, start_of_generation_token_id] + if len(token_ids) != 125: + raise RuntimeError(f"unexpected Edge causal token count: {len(token_ids)}") + ids = torch.tensor(token_ids, dtype=torch.long, device=text_embedding_weight.device) + return text_embedding_weight.index_select(0, ids).contiguous() + + +def _derive_step0_vfm_boundary_from_prepare_payload( + payload: Any, + seed: Any, + text_embedding_weight: Any | None = None, +) -> dict[str, Any]: + """Derive the step-0 VFM/noise boundary tensors from a prepared Edge payload. + + The full-only LM packed hidden states still require native VLM/prefill work; + this helper isolates the tensors that already follow from the slim prepare + contract, model text embeddings, and static Edge AV pack/RoPE rules. + """ + + import torch + + payload = _prepare_payload_with_derived_initial_noise(payload, seed) + if not isinstance(payload, tuple) or len(payload) <= 6: + raise RuntimeError("prepare payload cannot derive step0 boundary from a non-v1 tuple") + sequence_plans = payload[0] + gen_data_clean = payload[1] + initial_noise = payload[4] + condition_mask = payload[6] + if not isinstance(sequence_plans, list) or len(sequence_plans) != 1: + raise RuntimeError("step0 boundary derivation currently supports batch size 1") + if not isinstance(initial_noise, list) or len(initial_noise) != 1: + raise RuntimeError("step0 boundary derivation requires one initial_noise tensor") + if not isinstance(condition_mask, list) or len(condition_mask) != 1: + raise RuntimeError("step0 boundary derivation requires one condition_mask tensor") + + x0_tokens_vision = getattr(gen_data_clean, "x0_tokens_vision", None) + action_domain_id = getattr(gen_data_clean, "action_domain_id", None) + raw_action_dim = getattr(gen_data_clean, "raw_action_dim", None) + if not isinstance(x0_tokens_vision, list) or len(x0_tokens_vision) != 1: + raise RuntimeError("step0 boundary derivation requires one vision token tensor") + + noise_vision, noise_action = _flat_parts_for_edge_av(initial_noise[0]) + mask_vision, mask_action = _flat_parts_for_edge_av(condition_mask[0]) + device = initial_noise[0].device + + vision_condition = mask_vision.to(dtype=torch.float32) + if tuple(vision_condition.shape) != tuple(noise_vision.shape): + raise RuntimeError("unexpected vision condition mask shape") + vision_condition = vision_condition.amax(dim=(0, 1, 3, 4)).reshape(-1, 1, 1).contiguous() + action_condition = mask_action.to(dtype=torch.float32).amax(dim=1, keepdim=True).contiguous() + raw_action_dim_value = _indexed_optional_int(raw_action_dim, 0) + if raw_action_dim_value is None: + raise RuntimeError("step0 boundary derivation requires raw_action_dim") + domain_id_value = _indexed_optional_int(action_domain_id, 0) + if domain_id_value is None: + raise RuntimeError("step0 boundary derivation requires action_domain_id") + + causal_indices = torch.arange(125, dtype=torch.int32, device=device) + full_indices = torch.arange(125, 6425, dtype=torch.int32, device=device) + position_ids = _derive_step0_position_ids_from_prepare_payload(payload).to(device=device) + derived = { + "once/num_velocity_calls": torch.tensor([1], dtype=torch.int64, device=device), + "s00/lm_in/_causal_indices": causal_indices, + "s00/lm_in/_causal_seq_offsets": torch.tensor([0, 125], dtype=torch.int32, device=device), + "s00/lm_in/_full_indices": full_indices, + "s00/lm_in/_full_only_seq_offsets": torch.tensor([0, 6300], dtype=torch.int32, device=device), + "s00/lm_in/position_ids": position_ids, + "s00/lm_in/sample_offsets": torch.tensor([0, 6425], dtype=torch.int32, device=device), + "steps/00/noise_x": initial_noise[0].contiguous(), + "steps/00/timestep": torch.tensor([[999]], dtype=torch.int64, device=device), + "s00/vfm_in/vision/tokens/0": noise_vision.contiguous(), + "s00/vfm_in/vision/condition_mask/0": vision_condition, + "s00/vfm_in/vision/sequence_indexes": torch.arange(125, 6365, dtype=torch.int64, device=device), + "s00/vfm_in/vision/timesteps": torch.empty((0,), dtype=torch.float32, device=device), + "s00/vfm_in/vision/mse_loss_indexes": torch.empty((0,), dtype=torch.int64, device=device), + "s00/vfm_in/vision/noisy_frame_indexes/0": torch.empty((0,), dtype=torch.int64, device=device), + "s00/vfm_in/action/tokens/0": noise_action.contiguous(), + "s00/vfm_in/action/condition_mask/0": action_condition, + "s00/vfm_in/action/sequence_indexes": torch.arange(6365, 6425, dtype=torch.int64, device=device), + "s00/vfm_in/action/timesteps": torch.full((60,), 999.0, dtype=torch.float32, device=device), + "s00/vfm_in/action/mse_loss_indexes": torch.arange(6365, 6425, dtype=torch.int64, device=device), + "s00/vfm_in/action/noisy_frame_indexes/0": torch.arange(60, dtype=torch.int64, device=device), + "s00/vfm_in/action/domain_id/0": torch.tensor(domain_id_value, dtype=torch.int64, device=device), + "s00/vfm_in/action/raw_action_dim/0": torch.tensor(raw_action_dim_value, dtype=torch.int64, device=device), + } + derived.update(_derive_step0_rope_from_position_ids(position_ids, causal_indices, full_indices)) + if text_embedding_weight is not None: + derived["s00/lm_in/causal_seq"] = _derive_step0_causal_seq_from_prepare_payload(payload, text_embedding_weight).to( + device=device + ) + return derived + + +def _prepare_artifact_seed(artifact: dict[str, Any]) -> Any | None: + signature = artifact.get("signature") + if isinstance(signature, (list, tuple)) and len(signature) > 1: + seed_sig = signature[1] + if isinstance(seed_sig, tuple): + return list(seed_sig) + return seed_sig + return None + + +def _load_text_embedding_weight_for_checkpoint(checkpoint: Path) -> Any: + from safetensors import safe_open + + transformer_dir = checkpoint / "transformer" + candidates = [ + transformer_dir / "diffusion_pytorch_model-00001-of-00002.safetensors", + *sorted(transformer_dir.glob("*.safetensors")), + ] + seen: set[Path] = set() + for path in candidates: + if path in seen or not path.exists(): + continue + seen.add(path) + with safe_open(str(path), framework="pt", device="cpu") as handle: + if "embed_tokens.weight" in handle.keys(): + return handle.get_tensor("embed_tokens.weight") + raise RuntimeError(f"{checkpoint} is missing transformer embed_tokens.weight") + + +def _derive_step0_executable_boundary_from_prepare_artifact( + prepare_path: Path, + checkpoint: Path, + *, + seed: Any | None = None, + device: Any = "cuda", +) -> dict[str, Any]: + """Build the executable pre-layer boundary from a slim prepare artifact. + + The returned tensor map is sufficient for ``EdgeStaticBufferEngine``. Most + tensors come from the bit-exact slim-prepare contract; ``full_only_seq`` is + synthesized through the same Torch reference used by the live engine bring-up. + """ + + import torch + + loaded = torch.load(str(prepare_path), map_location="cpu", weights_only=False) + if not isinstance(loaded, dict) or loaded.get("version") != 1 or "payload" not in loaded: + raise RuntimeError(f"{prepare_path} is not a FlashRT prepare boundary v1 artifact") + effective_seed = seed if seed is not None else _prepare_artifact_seed(loaded) + if effective_seed is None: + raise RuntimeError("prepare-derived live boundary requires a seed or a prepare artifact signature") + return _derive_step0_executable_boundary_from_prepare_payload( + loaded["payload"], + checkpoint, + seed=effective_seed, + device=device, + path=f"{prepare_path}:derived", + ) + + +def _derive_step0_executable_boundary_from_prepare_payload( + payload: Any, + checkpoint: Path, + *, + seed: Any, + device: Any = "cuda", + path: str = "live_prepare_result:derived", +) -> dict[str, Any]: + """Build the executable pre-layer boundary directly from a prepare payload.""" + + import torch + + from flash_rt.models.cosmos3_edge import EdgeBoundaryDump, EdgeTransformerWeights + from flash_rt.models.cosmos3_edge.layer_ref import EdgeTransformerTorchReference + + text_embedding = _load_text_embedding_weight_for_checkpoint(checkpoint) + derived = _derive_step0_vfm_boundary_from_prepare_payload( + payload, + seed=seed, + text_embedding_weight=text_embedding, + ) + tensors = {key: tensor.detach().cpu().contiguous() for key, tensor in derived.items()} + causal = tensors["s00/lm_in/causal_seq"] + full_placeholder = torch.zeros((6300, causal.shape[1]), dtype=causal.dtype) + tensors["s00/lm_in/full_only_seq"] = full_placeholder + tensors["s00/layers/00/input/causal_seq"] = causal.clone() + tensors["s00/layers/00/input/full_only_seq"] = full_placeholder.clone() + + boundary = EdgeBoundaryDump.from_tensors(tensors, path=path) + weights = EdgeTransformerWeights(checkpoint) + ref = EdgeTransformerTorchReference(weights, device=device, dtype=torch.bfloat16) + full = ref.full_sequence_for_step( + boundary, + tensors["steps/00/noise_x"].to(device=device), + tensors["steps/00/timestep"].to(device=device), + ).detach().cpu().contiguous() + tensors["s00/lm_in/full_only_seq"] = full + tensors["s00/layers/00/input/full_only_seq"] = full.clone() + EdgeBoundaryDump.from_tensors(tensors, path=path).validate_geometry() + return tensors + + +def _slim_prepare_payload( + payload: Any, + *, + no_raw_state_vision: bool = False, + derive_condition_reference: bool = False, + derive_initial_noise: bool = False, +) -> Any: + slim = _clone_for_warmup(payload) + if isinstance(slim, tuple) and len(slim) > 1: + gen_data_clean = slim[1] + if no_raw_state_vision and hasattr(gen_data_clean, "raw_state_vision"): + gen_data_clean.raw_state_vision = None + if derive_condition_reference: + slim = _replace_prepare_payload_field(slim, 5, None) + if derive_initial_noise: + slim = _replace_prepare_payload_field(slim, 4, None) + return slim + + +def _slim_prepare_payload_no_raw_state_vision(payload: Any) -> Any: + return _slim_prepare_payload(payload, no_raw_state_vision=True) + + +class _PrepareBoundary: + def __init__( + self, + dump_out: Path | None, + replay_in: Path | None, + inventory_out: Path | None, + *, + slim_no_raw_state_vision: bool = False, + slim_derive_condition_reference: bool = False, + slim_derive_initial_noise: bool = False, + ) -> None: + self.dump_out = dump_out + self.replay_in = replay_in + self.inventory_out = inventory_out + self.slim_no_raw_state_vision = slim_no_raw_state_vision + self.slim_derive_condition_reference = slim_derive_condition_reference + self.slim_derive_initial_noise = slim_derive_initial_noise + self.events: list[dict[str, Any]] = [] + self._loaded: dict[str, Any] | None = None + self._last_effective_slim_no_raw_state_vision = slim_no_raw_state_vision + self._last_effective_slim_derive_condition_reference = slim_derive_condition_reference + self._last_effective_slim_derive_initial_noise = slim_derive_initial_noise + self.latest_payload: Any | None = None + self.latest_signature: tuple[Any, ...] | None = None + self.latest_seed: Any | None = None + self.latest_source: str | None = None + + @staticmethod + def _signature(data_batch: Any, seed: Any, has_negative_prompt: Any) -> tuple[Any, ...]: + return _WarmupPrepareCache._signature(data_batch, seed, has_negative_prompt) + + @staticmethod + def _signature_json(signature: tuple[Any, ...]) -> str: + return json.dumps(signature, sort_keys=True) + + def _remember(self, payload: Any, *, signature: tuple[Any, ...], seed: Any, source: str) -> None: + self.latest_payload = _clone_for_warmup(payload) + self.latest_signature = signature + self.latest_seed = seed + self.latest_source = source + + def latest_effective_payload(self) -> Any: + if self.latest_payload is None: + raise RuntimeError("live prepare boundary requested before _prepare_inference_data produced a payload") + return self.latest_payload + + def _load(self, *, model: Any, data_batch: Any, seed: Any, has_negative_prompt: Any) -> Any: + if self.replay_in is None: + raise RuntimeError("prepare replay path is not configured") + import torch + + if self._loaded is None: + map_location = _prepare_model_device(model, data_batch) + load_kwargs: dict[str, Any] = {"map_location": map_location} if map_location is not None else {} + try: + loaded = torch.load(str(self.replay_in), weights_only=False, **load_kwargs) + except TypeError: + loaded = torch.load(str(self.replay_in), **load_kwargs) + if not isinstance(loaded, dict) or loaded.get("version") != 1 or "payload" not in loaded: + raise RuntimeError(f"{self.replay_in} is not a FlashRT prepare boundary v1 artifact") + self._loaded = loaded + signature = self._signature(data_batch, seed, has_negative_prompt) + loaded_signature = self._loaded.get("signature") + if loaded_signature is not None: + loaded_signature_json = self._signature_json(tuple(loaded_signature)) + current_signature_json = self._signature_json(signature) + if loaded_signature_json != current_signature_json: + raise RuntimeError("prepare replay signature mismatch; refusing to reuse a different prepared state") + effective_slim_no_raw_state_vision = self.slim_no_raw_state_vision or bool( + self._loaded.get("slim_no_raw_state_vision") + ) + effective_slim_derive_condition_reference = self.slim_derive_condition_reference or bool( + self._loaded.get("slim_derive_condition_reference") + ) + effective_slim_derive_initial_noise = self.slim_derive_initial_noise or bool( + self._loaded.get("slim_derive_initial_noise") + ) + self._last_effective_slim_no_raw_state_vision = effective_slim_no_raw_state_vision + self._last_effective_slim_derive_condition_reference = effective_slim_derive_condition_reference + self._last_effective_slim_derive_initial_noise = effective_slim_derive_initial_noise + payload = _slim_prepare_payload( + self._loaded["payload"], + no_raw_state_vision=effective_slim_no_raw_state_vision, + derive_condition_reference=effective_slim_derive_condition_reference, + derive_initial_noise=effective_slim_derive_initial_noise, + ) + payload = _prepare_payload_with_derived_initial_noise(payload, seed) + payload = _prepare_payload_with_derived_condition_reference(payload) + if self.inventory_out is not None: + _write_prepare_inventory( + self.inventory_out, + payload, + signature=signature, + source=str(self.replay_in), + ) + return payload + + def wrap_prepare(self, original: Callable[..., Any]) -> Callable[..., Any]: + def prepare_with_boundary(model: Any, *args: Any, **kwargs: Any) -> Any: + data_batch = args[0] if args else kwargs.get("data_batch") + seed = args[1] if len(args) > 1 else kwargs.get("seed") + has_negative_prompt = args[2] if len(args) > 2 else kwargs.get("has_negative_prompt", False) + signature = self._signature(data_batch, seed, has_negative_prompt) + if self.replay_in is not None: + t0 = time.perf_counter() + payload = self._load(model=model, data_batch=data_batch, seed=seed, has_negative_prompt=has_negative_prompt) + self._remember( + payload, + signature=signature, + seed=seed, + source=str(self.replay_in), + ) + self.events.append( + { + "name": "prepare_replay_hit", + "phase": "measured", + "s": time.perf_counter() - t0, + "ok": True, + "wrote_inventory": self.inventory_out is not None, + "slim_no_raw_state_vision": self._last_effective_slim_no_raw_state_vision, + "slim_derive_condition_reference": self._last_effective_slim_derive_condition_reference, + "slim_derive_initial_noise": self._last_effective_slim_derive_initial_noise, + } + ) + return payload + + t0 = time.perf_counter() + result = original(model, *args, **kwargs) + elapsed = time.perf_counter() - t0 + output_payload = _slim_prepare_payload( + result, + no_raw_state_vision=self.slim_no_raw_state_vision, + derive_condition_reference=self.slim_derive_condition_reference, + derive_initial_noise=self.slim_derive_initial_noise, + ) + if self.dump_out is not None: + import torch + + self.dump_out.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "version": 1, + "signature": signature, + "payload": _clone_for_warmup(output_payload), + "slim_no_raw_state_vision": self.slim_no_raw_state_vision, + "slim_derive_condition_reference": self.slim_derive_condition_reference, + "slim_derive_initial_noise": self.slim_derive_initial_noise, + }, + str(self.dump_out), + ) + if self.inventory_out is not None: + _write_prepare_inventory( + self.inventory_out, + output_payload, + signature=signature, + source=str(self.dump_out) if self.dump_out is not None else "live_prepare_result", + ) + self.events.append( + { + "name": "prepare_boundary_dump", + "phase": "measured", + "s": elapsed, + "ok": True, + "wrote_dump": self.dump_out is not None, + "wrote_inventory": self.inventory_out is not None, + "slim_no_raw_state_vision": self.slim_no_raw_state_vision, + "slim_derive_condition_reference": self.slim_derive_condition_reference, + "slim_derive_initial_noise": self.slim_derive_initial_noise, + } + ) + output_payload = _prepare_payload_with_derived_initial_noise(output_payload, seed) + output_payload = _prepare_payload_with_derived_condition_reference(output_payload) + self._remember( + output_payload, + signature=signature, + seed=seed, + source=str(self.dump_out) if self.dump_out is not None else "live_prepare_result", + ) + return output_payload + + return prepare_with_boundary + + +def _install_prepare_boundary_patch( + dump_out: Path | None, + replay_in: Path | None, + inventory_out: Path | None, + *, + slim_no_raw_state_vision: bool = False, + slim_derive_condition_reference: bool = False, + slim_derive_initial_noise: bool = False, +) -> _PrepareBoundary: + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + boundary = _PrepareBoundary( + dump_out, + replay_in, + inventory_out, + slim_no_raw_state_vision=slim_no_raw_state_vision, + slim_derive_condition_reference=slim_derive_condition_reference, + slim_derive_initial_noise=slim_derive_initial_noise, + ) + OmniMoTModel._prepare_inference_data = boundary.wrap_prepare(OmniMoTModel._prepare_inference_data) + return boundary + + +class _VAEEncodeBoundary: + def __init__( + self, + dump_out: Path | None, + latent_in: Path | None, + *, + dump_input: bool = False, + ) -> None: + self.dump_out = dump_out + self.latent_in = latent_in + self.dump_input = dump_input + self.events: list[dict[str, Any]] = [] + self._loaded_latent: Any = None + self._loaded_signature: tuple[Any, ...] | None = None + self._loaded_metadata: dict[str, str] | None = None + + @staticmethod + def _signature(state: Any) -> tuple[Any, ...] | None: + return _WarmupVAECache._signature(state) + + @staticmethod + def _signature_json(signature: tuple[Any, ...] | None) -> str: + return json.dumps(signature, sort_keys=True) + + def _load_latent(self, state: Any) -> Any: + if self.latent_in is None: + raise RuntimeError("latent input path is not configured") + if self._loaded_latent is None: + from safetensors.torch import load_file, safe_open + + device = str(state.device) if hasattr(state, "device") else "cpu" + tensors = load_file(str(self.latent_in), device=device) + if "vae_encode/output" not in tensors: + raise RuntimeError(f"{self.latent_in} is missing vae_encode/output") + with safe_open(str(self.latent_in), framework="pt", device="cpu") as f: + self._loaded_metadata = dict(f.metadata() or {}) + raw_signature = (self._loaded_metadata or {}).get("state_signature") + self._loaded_signature = tuple(json.loads(raw_signature)) if raw_signature else None + self._loaded_latent = tensors["vae_encode/output"] + signature = self._signature(state) + if self._loaded_signature is not None and signature is not None: + loaded_signature_json = self._signature_json(self._loaded_signature) + current_signature_json = self._signature_json(signature) + if loaded_signature_json != current_signature_json: + raise RuntimeError( + "VAE latent input signature mismatch; refusing to reuse a latent from a different input" + ) + return self._loaded_latent.clone() + + def wrap_encode(self, original: Callable[..., Any]) -> Callable[..., Any]: + def encode_with_boundary(model: Any, state: Any) -> Any: + signature = self._signature(state) + if self.latent_in is not None: + t0 = time.perf_counter() + latent = self._load_latent(state) + self.events.append( + { + "name": "vae_latent_replay_hit", + "phase": "measured", + "s": time.perf_counter() - t0, + "ok": True, + } + ) + return latent + + t0 = time.perf_counter() + encoded = original(model, state) + elapsed = time.perf_counter() - t0 + if self.dump_out is not None: + from safetensors.torch import save_file + + tensors = {"vae_encode/output": encoded.detach().cpu().contiguous()} + if self.dump_input: + tensors["vae_encode/input"] = state.detach().cpu().contiguous() + metadata = { + "state_signature": self._signature_json(signature), + "dump_input": "1" if self.dump_input else "0", + "version": "1", + } + self.dump_out.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(self.dump_out), metadata=metadata) + self.events.append( + { + "name": "vae_encode_boundary_dump", + "phase": "measured", + "s": elapsed, + "ok": True, + "wrote_dump": self.dump_out is not None, + "dump_input": self.dump_input, + } + ) + return encoded + + return encode_with_boundary + + +def _install_vae_encode_boundary_patch( + dump_out: Path | None, + latent_in: Path | None, + *, + dump_input: bool = False, +) -> _VAEEncodeBoundary: + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + boundary = _VAEEncodeBoundary(dump_out, latent_in, dump_input=dump_input) + OmniMoTModel.encode = boundary.wrap_encode(OmniMoTModel.encode) + return boundary + + +class _VAEEncodeProfiler: + def __init__(self, profile_out: Path) -> None: + self.profile_out = profile_out + self.events: list[dict[str, Any]] = [] + self.modules: list[dict[str, Any]] = [] + self.encode_calls: list[dict[str, Any]] = [] + self._installed_model_ids: set[int] = set() + self._handles: list[Any] = [] + self._call_index = 0 + self._active_call: int | None = None + self._torch: Any = None + + @staticmethod + def _tensor_sig(value: Any) -> Any: + try: + import torch + except Exception: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + return { + "shape": [int(dim) for dim in value.shape], + "dtype": str(value.dtype), + "device": str(value.device), + "stride": [int(dim) for dim in value.stride()], + "contiguous": bool(value.is_contiguous()), + } + if isinstance(value, (list, tuple)): + return [self_item for self_item in (_VAEEncodeProfiler._tensor_sig(item) for item in value)] + if isinstance(value, dict): + return {str(key): _VAEEncodeProfiler._tensor_sig(item) for key, item in value.items()} + return None + + @staticmethod + def _module_inventory(name: str, module: Any) -> dict[str, Any]: + params = [] + for param_name, param in module.named_parameters(recurse=False): + params.append( + { + "name": param_name, + "shape": [int(dim) for dim in param.shape], + "dtype": str(param.dtype), + "numel": int(param.numel()), + } + ) + attrs: dict[str, Any] = {} + for attr in ( + "in_channels", + "out_channels", + "kernel_size", + "stride", + "padding", + "in_dim", + "out_dim", + "dim", + "mode", + "factor_t", + "factor_s", + "group_size", + ): + if hasattr(module, attr): + value = getattr(module, attr) + if isinstance(value, tuple): + value = [int(item) if isinstance(item, int) else item for item in value] + elif isinstance(value, (int, float, str, bool)): + pass + else: + value = repr(value) + attrs[attr] = value + return { + "name": name, + "class": module.__class__.__name__, + "children": len(list(module.children())), + "parameters": params, + "parameter_numel": sum(int(item["numel"]) for item in params), + "attrs": attrs, + } + + @staticmethod + def _is_timed_leaf(module: Any) -> bool: + if len(list(module.children())) != 0: + return False + return module.__class__.__name__ in { + "CausalConv3d", + "RMS_norm", + "SiLU", + "Dropout", + "Identity", + "AttentionBlock", + "AvgDown3D", + } or module.__class__.__module__.startswith("torch.nn.modules") + + def _sync(self) -> None: + if self._torch is None: + try: + import torch + + self._torch = torch + except Exception: + self._torch = False + torch = self._torch + if ( + torch + and torch.cuda.is_available() + and torch.cuda.is_initialized() + ): + torch.cuda.synchronize() + + def _install_for_model(self, model: Any) -> None: + model_id = id(model) + if model_id in self._installed_model_ids: + return + self._installed_model_ids.add(model_id) + self.modules = [] + for name, module in model.named_modules(): + if name and not (name == "conv1" or name.startswith("encoder")): + continue + self.modules.append(self._module_inventory(name or "", module)) + if not self._is_timed_leaf(module): + continue + + label = name or "" + cls_name = module.__class__.__name__ + param_numel = sum(int(param.numel()) for param in module.parameters(recurse=False)) + + def pre_hook(_module: Any, inputs: Any, *, _label: str = label, _cls: str = cls_name): + if self._active_call is None: + return + self._sync() + setattr( + _module, + "_flashrt_vae_profile_pre", + { + "t": time.perf_counter(), + "input": self._tensor_sig(inputs), + "call": self._active_call, + "label": _label, + "class": _cls, + }, + ) + + def post_hook( + _module: Any, + _inputs: Any, + output: Any, + *, + _param_numel: int = param_numel, + ): + meta = getattr(_module, "_flashrt_vae_profile_pre", None) + if not isinstance(meta, dict): + return + self._sync() + self.events.append( + { + "encode_call": int(meta["call"]), + "name": meta["label"], + "class": meta["class"], + "input": meta["input"], + "output": self._tensor_sig(output), + "s": time.perf_counter() - float(meta["t"]), + "parameter_numel": int(_param_numel), + } + ) + setattr(_module, "_flashrt_vae_profile_pre", None) + + self._handles.append(module.register_forward_pre_hook(pre_hook)) + self._handles.append(module.register_forward_hook(post_hook)) + + def wrap_encode(self, original: Callable[..., Any]) -> Callable[..., Any]: + def encode_with_profile(model: Any, x: Any, scale: Any) -> Any: + self._install_for_model(model) + call_index = self._call_index + self._call_index += 1 + previous = self._active_call + self._active_call = call_index + input_sig = self._tensor_sig(x) + self._sync() + t0 = time.perf_counter() + try: + out = original(model, x, scale) + finally: + self._sync() + elapsed = time.perf_counter() - t0 + self._active_call = previous + self.encode_calls.append( + { + "encode_call": call_index, + "input": input_sig, + "output": self._tensor_sig(out), + "s": elapsed, + } + ) + return out + + return encode_with_profile + + def save(self) -> None: + by_name: dict[str, dict[str, Any]] = {} + for event in self.events: + key = str(event["name"]) + item = by_name.setdefault( + key, + { + "name": event["name"], + "class": event["class"], + "count": 0, + "total_s": 0.0, + "max_s": 0.0, + "parameter_numel": event.get("parameter_numel", 0), + "input": event.get("input"), + "output": event.get("output"), + }, + ) + seconds = float(event["s"]) + item["count"] += 1 + item["total_s"] += seconds + item["max_s"] = max(float(item["max_s"]), seconds) + item["avg_s"] = float(item["total_s"]) / max(1, int(item["count"])) + summary = sorted(by_name.values(), key=lambda item: -float(item["total_s"])) + + def shape_of(value: Any) -> list[int] | None: + if isinstance(value, dict) and isinstance(value.get("shape"), list): + return [int(dim) for dim in value["shape"]] + return None + + def input_shape_at(event: dict[str, Any], idx: int) -> list[int] | None: + inputs = event.get("input") + if not isinstance(inputs, list) or idx >= len(inputs): + return None + return shape_of(inputs[idx]) + + by_shape: dict[tuple[Any, ...], dict[str, Any]] = {} + candidate_by_shape: dict[tuple[Any, ...], dict[str, Any]] = {} + for event in self.events: + input0_shape = input_shape_at(event, 0) + cache_shape = input_shape_at(event, 1) + output_shape = shape_of(event.get("output")) + key = ( + event.get("class"), + event.get("name"), + tuple(input0_shape or ()), + tuple(cache_shape or ()), + tuple(output_shape or ()), + ) + item = by_shape.setdefault( + key, + { + "name": event.get("name"), + "class": event.get("class"), + "count": 0, + "total_s": 0.0, + "max_s": 0.0, + "parameter_numel": int(event.get("parameter_numel", 0)), + "input_shape": input0_shape, + "cache_shape": cache_shape, + "output_shape": output_shape, + }, + ) + seconds = float(event["s"]) + item["count"] += 1 + item["total_s"] += seconds + item["max_s"] = max(float(item["max_s"]), seconds) + item["avg_s"] = float(item["total_s"]) / max(1, int(item["count"])) + + if event.get("class") != "CausalConv3d": + continue + t_new = input0_shape[2] if input0_shape and len(input0_shape) == 5 else None + if cache_shape and t_new and t_new > 1: + candidate_type = "steady_cached_causal_conv3d" + elif cache_shape is None and t_new == 1: + candidate_type = "prime_t1_no_cache" + else: + candidate_type = "other_causal_conv3d" + candidate_key = ( + candidate_type, + event.get("name"), + tuple(input0_shape or ()), + tuple(cache_shape or ()), + tuple(output_shape or ()), + ) + candidate = candidate_by_shape.setdefault( + candidate_key, + { + "candidate_type": candidate_type, + "name": event.get("name"), + "count": 0, + "total_s": 0.0, + "max_s": 0.0, + "parameter_numel": int(event.get("parameter_numel", 0)), + "input_shape": input0_shape, + "cache_shape": cache_shape, + "output_shape": output_shape, + }, + ) + candidate["count"] += 1 + candidate["total_s"] += seconds + candidate["max_s"] = max(float(candidate["max_s"]), seconds) + candidate["avg_s"] = float(candidate["total_s"]) / max(1, int(candidate["count"])) + + shape_summary = sorted(by_shape.values(), key=lambda item: -float(item["total_s"])) + native_candidate_summary = sorted( + candidate_by_shape.values(), + key=lambda item: ( + str(item.get("candidate_type")) != "steady_cached_causal_conv3d", + -float(item["total_s"]), + ), + ) + module_class_counts: dict[str, int] = {} + for module in self.modules: + cls = str(module["class"]) + module_class_counts[cls] = module_class_counts.get(cls, 0) + 1 + self.profile_out.parent.mkdir(parents=True, exist_ok=True) + self.profile_out.write_text( + json.dumps( + { + "version": 1, + "module_class_counts": module_class_counts, + "encode_calls": self.encode_calls, + "summary": summary, + "shape_summary": shape_summary, + "native_candidate_summary": native_candidate_summary, + "modules": self.modules, + "events": self.events, + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + + +def _install_vae_encode_profile_patch(profile_out: Path) -> _VAEEncodeProfiler: + from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import WanVAE_ + + profiler = _VAEEncodeProfiler(profile_out) + WanVAE_.encode = profiler.wrap_encode(WanVAE_.encode) + return profiler + + +def _install_warmup_vae_cache_patch() -> _WarmupVAECache: + from cosmos_framework.inference.inference import OmniInference + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + cache = _WarmupVAECache() + OmniInference.generate_batch = cache.wrap_generate_batch(OmniInference.generate_batch) + OmniMoTModel.encode = cache.wrap_encode(OmniMoTModel.encode) + return cache + + +def _install_warmup_prepare_cache_patch( + *, + slim_no_raw_state_vision: bool = False, + slim_derive_condition_reference: bool = False, + slim_derive_initial_noise: bool = False, +) -> _WarmupPrepareCache: + from cosmos_framework.inference.inference import OmniInference + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + cache = _WarmupPrepareCache( + slim_no_raw_state_vision=slim_no_raw_state_vision, + slim_derive_condition_reference=slim_derive_condition_reference, + slim_derive_initial_noise=slim_derive_initial_noise, + ) + OmniInference.generate_batch = cache.wrap_generate_batch(OmniInference.generate_batch) + OmniMoTModel._prepare_inference_data = cache.wrap_prepare(OmniMoTModel._prepare_inference_data) + return cache + + +class _UpstreamTrace: + def __init__( + self, + trace_out: Path, + *, + vae_cache: _WarmupVAECache | None = None, + prepare_cache: _WarmupPrepareCache | None = None, + vae_boundary: _VAEEncodeBoundary | None = None, + prepare_boundary: _PrepareBoundary | None = None, + ): + self.trace_out = trace_out + self.phase = "setup" + self.vae_cache = vae_cache + self.prepare_cache = prepare_cache + self.vae_boundary = vae_boundary + self.prepare_boundary = prepare_boundary + self.events: list[dict[str, Any]] = [] + try: + import torch + except Exception: + torch = None + self.torch = torch + + def _sync(self) -> None: + if ( + self.torch is not None + and self.torch.cuda.is_available() + and self.torch.cuda.is_initialized() + ): + self.torch.cuda.synchronize() + + @contextmanager + def record( + self, + name: str, + *, + phase: str | None = None, + metadata: dict[str, Any] | None = None, + expected_exceptions: tuple[str, ...] = (), + ): + self._sync() + t0 = time.perf_counter() + ok = False + event_metadata = dict(metadata) if metadata else None + try: + yield + ok = True + except Exception as exc: + if exc.__class__.__name__ in expected_exceptions: + ok = True + if event_metadata is None: + event_metadata = {} + event_metadata["expected_exception"] = exc.__class__.__name__ + raise + finally: + self._sync() + event = { + "name": name, + "phase": phase if phase is not None else self.phase, + "s": time.perf_counter() - t0, + "ok": ok, + } + if event_metadata: + event["metadata"] = event_metadata + self.events.append(event) + + def save(self) -> None: + events = list(self.events) + if self.vae_cache is not None: + events.extend(self.vae_cache.events) + if self.prepare_cache is not None: + events.extend(self.prepare_cache.events) + if self.vae_boundary is not None: + events.extend(self.vae_boundary.events) + if self.prepare_boundary is not None: + events.extend(self.prepare_boundary.events) + by_key: dict[str, dict[str, Any]] = {} + for event in events: + key = f"{event['phase']}::{event['name']}" + item = by_key.setdefault( + key, + { + "name": event["name"], + "phase": event["phase"], + "count": 0, + "total_s": 0.0, + "max_s": 0.0, + "ok_count": 0, + }, + ) + seconds = float(event["s"]) + item["count"] += 1 + item["total_s"] += seconds + item["max_s"] = max(float(item["max_s"]), seconds) + if event.get("ok"): + item["ok_count"] += 1 + summary = [] + for item in by_key.values(): + item["avg_s"] = float(item["total_s"]) / max(1, int(item["count"])) + summary.append(item) + summary.sort(key=lambda item: (str(item["phase"]), -float(item["total_s"]), str(item["name"]))) + self.trace_out.parent.mkdir(parents=True, exist_ok=True) + self.trace_out.write_text( + json.dumps({"summary": summary, "events": events}, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _install_upstream_trace_patch( + trace_out: Path, + *, + vae_cache: _WarmupVAECache | None = None, + prepare_cache: _WarmupPrepareCache | None = None, + vae_boundary: _VAEEncodeBoundary | None = None, + prepare_boundary: _PrepareBoundary | None = None, +) -> _UpstreamTrace: + import cosmos_framework.inference.inference as inference_mod + from cosmos_framework.inference.inference import OmniInference + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork + from cosmos_framework.model.generator.mot.unified_mot import ( + Nemotron3DenseVLTextForCausalLM, + Nemotron3DenseVLTextModel, + Qwen3VLTextForCausalLM, + Qwen3VLTextModel, + Qwen3VLMoeTextForCausalLM, + Qwen3VLMoeTextModel, + ) + + trace = _UpstreamTrace( + trace_out, + vae_cache=vae_cache, + prepare_cache=prepare_cache, + vae_boundary=vae_boundary, + prepare_boundary=prepare_boundary, + ) + orig_create_batches = OmniInference.create_batches + orig_generate_batch = OmniInference.generate_batch + orig_finalize_data_batch = inference_mod._finalize_data_batch + orig_prepare_inference_data = OmniMoTModel._prepare_inference_data + orig_get_data_and_condition = OmniMoTModel.get_data_and_condition + orig_get_inference_text_tokens = OmniMoTModel._get_inference_text_tokens + orig_pack_input_sequence = OmniMoTModel._pack_input_sequence + orig_encode = OmniMoTModel.encode + orig_get_velocity = OmniMoTModel._get_velocity + orig_vfm_forward = Cosmos3VFMNetwork.forward + text_model_classes = (Nemotron3DenseVLTextModel, Qwen3VLTextModel, Qwen3VLMoeTextModel) + lm_classes = (Nemotron3DenseVLTextForCausalLM, Qwen3VLTextForCausalLM, Qwen3VLMoeTextForCausalLM) + orig_text_forwards = {cls: cls.forward for cls in (*text_model_classes, *lm_classes)} + + def create_batches_traced(self: Any, sample_args_list: Any): + idx = 0 + iterator = orig_create_batches(self, sample_args_list) + while True: + trace._sync() + t0 = time.perf_counter() + try: + item = next(iterator) + except StopIteration: + trace._sync() + return + trace._sync() + trace.events.append( + { + "name": "OmniInference.create_batches.next", + "phase": "input_prep", + "s": time.perf_counter() - t0, + "ok": True, + "metadata": {"batch_index": idx}, + } + ) + idx += 1 + yield item + + def generate_batch_traced(self: Any, sample_args_list: Any, data_batch: dict[str, Any], *, warmup: bool = False): + previous = trace.phase + trace.phase = "warmup" if warmup else "measured" + try: + with trace.record("OmniInference.generate_batch"): + return orig_generate_batch(self, sample_args_list, data_batch, warmup=warmup) + finally: + trace.phase = previous + + def finalize_data_batch_traced(*args: Any, **kwargs: Any): + with trace.record("inference._finalize_data_batch"): + return orig_finalize_data_batch(*args, **kwargs) + + def prepare_inference_data_traced(self: Any, *args: Any, **kwargs: Any): + with trace.record("OmniMoTModel._prepare_inference_data"): + return orig_prepare_inference_data(self, *args, **kwargs) + + def get_data_and_condition_traced(self: Any, *args: Any, **kwargs: Any): + with trace.record("OmniMoTModel.get_data_and_condition"): + return orig_get_data_and_condition(self, *args, **kwargs) + + def get_inference_text_tokens_traced(self: Any, *args: Any, **kwargs: Any): + with trace.record("OmniMoTModel._get_inference_text_tokens"): + return orig_get_inference_text_tokens(self, *args, **kwargs) + + def pack_input_sequence_traced(self: Any, *args: Any, **kwargs: Any): + metadata = {"skip_text_tokens": bool(kwargs.get("skip_text_tokens", False))} + with trace.record("OmniMoTModel._pack_input_sequence", metadata=metadata): + return orig_pack_input_sequence(self, *args, **kwargs) + + def encode_traced(self: Any, state: Any): + metadata = {"shape": list(state.shape)} if hasattr(state, "shape") else None + with trace.record("OmniMoTModel.encode", metadata=metadata): + return orig_encode(self, state) + + def get_velocity_traced(self: Any, *args: Any, **kwargs: Any): + with trace.record("OmniMoTModel._get_velocity", expected_exceptions=("_FlashRTVelocityReady",)): + return orig_get_velocity(self, *args, **kwargs) + + def vfm_forward_traced(self: Any, *args: Any, **kwargs: Any): + with trace.record("Cosmos3VFMNetwork.forward", expected_exceptions=("_FlashRTVelocityReady",)): + return orig_vfm_forward(self, *args, **kwargs) + + def make_text_forward_traced(cls: type[Any]) -> Callable[..., Any]: + orig = orig_text_forwards[cls] + + def text_forward_traced(self: Any, *args: Any, **kwargs: Any): + with trace.record(f"{cls.__name__}.forward", expected_exceptions=("_FlashRTVelocityReady",)): + return orig(self, *args, **kwargs) + + return text_forward_traced + + OmniInference.create_batches = create_batches_traced + OmniInference.generate_batch = generate_batch_traced + inference_mod._finalize_data_batch = finalize_data_batch_traced + OmniMoTModel._prepare_inference_data = prepare_inference_data_traced + OmniMoTModel.get_data_and_condition = get_data_and_condition_traced + OmniMoTModel._get_inference_text_tokens = get_inference_text_tokens_traced + OmniMoTModel._pack_input_sequence = pack_input_sequence_traced + OmniMoTModel.encode = encode_traced + OmniMoTModel._get_velocity = get_velocity_traced + Cosmos3VFMNetwork.forward = vfm_forward_traced + for cls in (*text_model_classes, *lm_classes): + cls.forward = make_text_forward_traced(cls) + return trace + + +def _dump_tensor(tensor: Any, *, flatten: bool = False) -> Any: + out = tensor.detach().cpu().contiguous() + if flatten: + out = out.reshape(-1).contiguous() + return out + + +def _normalise_final_action(action: Any) -> Any: + out = _dump_tensor(action) + if out.ndim == 3 and out.shape[0] == 1: + out = out.squeeze(0).contiguous() + return out + + +def _install_live_dump_patch(live_dump_out: Path) -> None: + import torch + from safetensors.torch import save_file + + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + + original = OmniMoTModel.generate_samples_from_batch + + def wrapped_generate_samples_from_batch(self: Any, *args: Any, **kwargs: Any) -> dict[str, list[Any]]: + steps: list[dict[str, Any]] = [] + user_builder: Callable[..., Any] | None = kwargs.get("velocity_postprocess_builder") + + def velocity_postprocess_builder(**builder_kwargs: Any) -> Callable[[list[Any], list[Any], Any], list[Any]]: + user_postprocess = user_builder(**builder_kwargs) if user_builder is not None else None + + def velocity_postprocess(cond_v_full: list[Any], noise_x: list[Any], timestep: Any) -> list[Any]: + post_v = ( + user_postprocess(cond_v_full, noise_x, timestep) + if user_postprocess is not None + else cond_v_full + ) + if len(noise_x) == 1 and len(post_v) == 1: + steps.append( + { + "noise_x": _dump_tensor(noise_x[0], flatten=True), + "velocity": _dump_tensor(post_v[0], flatten=True), + "timestep": _dump_tensor(timestep), + } + ) + return post_v + + return velocity_postprocess + + kwargs["velocity_postprocess_builder"] = velocity_postprocess_builder + outputs = original(self, *args, **kwargs) + + tensors: dict[str, torch.Tensor] = { + "once/num_velocity_calls": torch.tensor([len(steps)], dtype=torch.int64) + } + for idx, step in enumerate(steps): + tensors[f"steps/{idx:02d}/noise_x"] = step["noise_x"] + tensors[f"steps/{idx:02d}/velocity"] = step["velocity"] + tensors[f"steps/{idx:02d}/timestep"] = step["timestep"] + + if not outputs.get("vision"): + raise RuntimeError("live dump capture did not receive a final vision latent") + if not outputs.get("action"): + raise RuntimeError("live dump capture did not receive a final action tensor") + tensors["once/final_vision"] = _dump_tensor(outputs["vision"][0]) + tensors["once/final_action"] = _normalise_final_action(outputs["action"][0]) + + live_dump_out.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(live_dump_out)) + return outputs + + OmniMoTModel.generate_samples_from_batch = wrapped_generate_samples_from_batch + + +class _FlashRTVelocityReady(RuntimeError): + def __init__(self, velocity: list[Any]): + super().__init__("FlashRT live pre-layer handoff velocity is ready") + self.velocity = velocity + + +class _LiveFlashRTHandoff: + def __init__( + self, + checkpoint: Path, + boundary_out: Path | None, + boundary_in: Path | None, + boundary_prepare_in: Path | None, + boundary_prepare_live: _PrepareBoundary | None, + trace_out: Path | None, + *, + prelayer_bootstrap: bool = False, + ): + import torch + + self.checkpoint = checkpoint + self.boundary_out = boundary_out + self.boundary_in = boundary_in + self.boundary_prepare_in = boundary_prepare_in + self.boundary_prepare_live = boundary_prepare_live + self.trace_out = trace_out + self.prelayer_bootstrap = prelayer_bootstrap + self.torch = torch + self.tensors: dict[str, torch.Tensor] = {} + self.active_step: int | None = None + self.step = 0 + self.engine: Any = None + self.current_noise_x: list[Any] | None = None + self.current_timestep: Any | None = None + self.current_scheduler_seed: Any | None = None + self.current_skip_text_tokens = False + self._engine_prepare_signature: tuple[Any, ...] | None = None + self._captured_vfm = False + self._captured_lm = False + self._captured_layer0 = False + self.trace: dict[str, Any] = { + "official_velocity_calls": [], + "flashrt_velocity_calls": [], + "engine_init": {}, + "native_scheduler": { + "enabled": os.environ.get("FLASHRT_COSMOS3_EDGE_LIVE_NATIVE_UNIPC", "1") != "0", + "runs": [], + "fallbacks": [], + "failures": [], + }, + "prelayer_bootstrap": prelayer_bootstrap, + "boundary_in": str(boundary_in) if boundary_in is not None else None, + "boundary_prepare_in": str(boundary_prepare_in) if boundary_prepare_in is not None else None, + "boundary_prepare_live": boundary_prepare_live is not None, + } + + def _sync_for_trace(self) -> None: + if ( + self.trace_out is not None + and self.torch.cuda.is_available() + and self.torch.cuda.is_initialized() + ): + self.torch.cuda.synchronize() + + def _now_for_trace(self) -> float: + self._sync_for_trace() + return time.perf_counter() + + def _cpu(self, value: Any) -> Any: + return value.detach().cpu().contiguous() + + def _add_tensor(self, name: str, value: Any) -> None: + if isinstance(value, self.torch.Tensor): + self.tensors[name] = self._cpu(value) + + def _add_tensor_list(self, prefix: str, values: Any) -> None: + if isinstance(values, list): + for idx, value in enumerate(values): + self._add_tensor(f"{prefix}/{idx}", value) + + def _capture_modality(self, prefix: str, modality: Any) -> None: + if modality is None: + return + for field in ("sequence_indexes", "timesteps", "mse_loss_indexes"): + self._add_tensor(f"{prefix}/{field}", getattr(modality, field, None)) + for field in ("tokens", "condition_mask", "noisy_frame_indexes", "domain_id", "raw_action_dim"): + self._add_tensor_list(f"{prefix}/{field}", getattr(modality, field, None)) + + def capture_vfm_boundary(self, packed_seq: Any) -> None: + if self.active_step != 0 or self._captured_vfm: + return + self._captured_vfm = True + self._capture_modality("s00/vfm_in/vision", getattr(packed_seq, "vision", None)) + self._capture_modality("s00/vfm_in/action", getattr(packed_seq, "action", None)) + + def capture_lm_boundary(self, pack: Any, _attention_mask: Any, position_ids: Any) -> None: + if self.active_step != 0 or self._captured_lm: + return + self._captured_lm = True + if isinstance(pack, dict): + for key, value in pack.items(): + self._add_tensor(f"s00/lm_in/{key}", value) + self._add_tensor("s00/lm_in/position_ids", position_ids) + + def capture_layer0(self, module: Any, input_pack: Any, args: tuple[Any, ...], output: Any) -> None: + if self.active_step != 0 or self._captured_layer0: + return + layer_idx = getattr(getattr(module, "self_attn", None), "layer_idx", None) + if layer_idx != 0: + return + self._captured_layer0 = True + if isinstance(input_pack, dict): + for key in ("causal_seq", "full_only_seq", "_causal_indices", "_full_indices"): + self._add_tensor(f"s00/layers/00/input/{key}", input_pack.get(key)) + packed_position_embeddings = args[1] if len(args) > 1 else None + if ( + isinstance(packed_position_embeddings, tuple) + and len(packed_position_embeddings) == 2 + and all(isinstance(item, dict) for item in packed_position_embeddings) + ): + for name, pack in (("cos", packed_position_embeddings[0]), ("sin", packed_position_embeddings[1])): + for key in ("causal_seq", "full_only_seq"): + self._add_tensor(f"s00/layers/00/rope/{name}/{key}", pack.get(key)) + hidden = output[0] if isinstance(output, tuple) and output else None + if isinstance(hidden, dict): + for key in ("causal_seq", "full_only_seq"): + self._add_tensor(f"s00/layers/00/output/{key}", hidden.get(key)) + + def _initialise_engine(self, noise_x: list[Any], timestep: Any, velocity: list[Any] | None = None) -> None: + if len(noise_x) != 1 or (velocity is not None and len(velocity) != 1): + raise RuntimeError("FlashRT live handoff currently supports batch size 1") + from safetensors.torch import save_file + + from flash_rt.models.cosmos3_edge import EdgeBoundaryDump, EdgeStaticBufferEngine, EdgeTransformerWeights + + t0 = self._now_for_trace() + self.tensors["steps/00/noise_x"] = self._cpu(noise_x[0]) + if velocity is not None: + self.tensors["steps/00/velocity"] = self._cpu(velocity[0]) + self.tensors["steps/00/timestep"] = self._cpu(timestep) + self.tensors.setdefault("once/num_velocity_calls", self.torch.tensor([1], dtype=self.torch.int64)) + t_assembled = self._now_for_trace() + boundary_path: Path | str = "" + if self.boundary_out is not None: + self.boundary_out.parent.mkdir(parents=True, exist_ok=True) + save_file(self.tensors, str(self.boundary_out)) + boundary_path = self.boundary_out + t_saved = self._now_for_trace() + boundary = EdgeBoundaryDump.from_tensors(self.tensors, path=boundary_path) + boundary.validate_geometry() + t_validated = self._now_for_trace() + self.engine = EdgeStaticBufferEngine( + EdgeTransformerWeights(str(self.checkpoint)), + boundary, + device=noise_x[0].device, + ) + t_done = self._now_for_trace() + self.trace["engine_init"] = { + "mode": "captured_live", + "boundary_assembly_s": t_assembled - t0, + "boundary_save_s": t_saved - t_assembled, + "boundary_validate_s": t_validated - t_saved, + "engine_construct_s": t_done - t_validated, + "total_s": t_done - t0, + "wrote_boundary": self.boundary_out is not None, + } + + def _initialise_engine_from_boundary_in(self, device: Any) -> None: + if self.boundary_in is None: + raise RuntimeError("missing FlashRT live boundary input") + + from flash_rt.models.cosmos3_edge import EdgeBoundaryDump, EdgeStaticBufferEngine, EdgeTransformerWeights + + t0 = self._now_for_trace() + boundary = EdgeBoundaryDump(self.boundary_in) + t_loaded = self._now_for_trace() + boundary.validate_geometry() + t_validated = self._now_for_trace() + self.engine = EdgeStaticBufferEngine( + EdgeTransformerWeights(str(self.checkpoint)), + boundary, + device=device, + ) + self._engine_prepare_signature = self.boundary_prepare_live.latest_signature + t_done = self._now_for_trace() + self.trace["engine_init"] = { + "mode": "boundary_in", + "boundary_in": str(self.boundary_in), + "boundary_load_s": t_loaded - t0, + "boundary_validate_s": t_validated - t_loaded, + "engine_construct_s": t_done - t_validated, + "total_s": t_done - t0, + "wrote_boundary": False, + } + + def _initialise_engine_from_prepare_in(self, device: Any) -> None: + if self.boundary_prepare_in is None: + raise RuntimeError("missing FlashRT live prepare-boundary input") + from safetensors.torch import save_file + + from flash_rt.models.cosmos3_edge import EdgeBoundaryDump, EdgeStaticBufferEngine, EdgeTransformerWeights + + t0 = self._now_for_trace() + tensors = _derive_step0_executable_boundary_from_prepare_artifact( + self.boundary_prepare_in, + self.checkpoint, + seed=self.current_scheduler_seed, + device=device, + ) + t_derived = self._now_for_trace() + boundary_path: Path | str = f"{self.boundary_prepare_in}:derived" + if self.boundary_out is not None: + self.boundary_out.parent.mkdir(parents=True, exist_ok=True) + save_file({key: tensor.clone() for key, tensor in tensors.items()}, str(self.boundary_out)) + boundary_path = self.boundary_out + t_saved = self._now_for_trace() + boundary = EdgeBoundaryDump.from_tensors(tensors, path=boundary_path) + boundary.validate_geometry() + t_validated = self._now_for_trace() + self.engine = EdgeStaticBufferEngine( + EdgeTransformerWeights(str(self.checkpoint)), + boundary, + device=device, + ) + t_done = self._now_for_trace() + self.trace["engine_init"] = { + "mode": "prepare_in", + "boundary_prepare_in": str(self.boundary_prepare_in), + "prepare_artifact_mib": self.boundary_prepare_in.stat().st_size / (1024.0 * 1024.0), + "boundary_derive_s": t_derived - t0, + "boundary_save_s": t_saved - t_derived, + "boundary_validate_s": t_validated - t_saved, + "engine_construct_s": t_done - t_validated, + "total_s": t_done - t0, + "wrote_boundary": self.boundary_out is not None, + } + + def _initialise_engine_from_prepare_live(self, device: Any) -> None: + if self.boundary_prepare_live is None: + raise RuntimeError("missing FlashRT live prepare-boundary source") + from safetensors.torch import save_file + + from flash_rt.models.cosmos3_edge import EdgeBoundaryDump, EdgeStaticBufferEngine, EdgeTransformerWeights + + t0 = self._now_for_trace() + payload = self.boundary_prepare_live.latest_effective_payload() + seed = self.current_scheduler_seed + if seed is None: + seed = self.boundary_prepare_live.latest_seed + if seed is None: + raise RuntimeError("live prepare-derived boundary requires the sampler seed") + tensors = _derive_step0_executable_boundary_from_prepare_payload( + payload, + self.checkpoint, + seed=seed, + device=device, + path="live_prepare_result:derived", + ) + t_derived = self._now_for_trace() + boundary_path: Path | str = "live_prepare_result:derived" + if self.boundary_out is not None: + self.boundary_out.parent.mkdir(parents=True, exist_ok=True) + save_file({key: tensor.clone() for key, tensor in tensors.items()}, str(self.boundary_out)) + boundary_path = self.boundary_out + t_saved = self._now_for_trace() + boundary = EdgeBoundaryDump.from_tensors(tensors, path=boundary_path) + boundary.validate_geometry() + t_validated = self._now_for_trace() + self.engine = EdgeStaticBufferEngine( + EdgeTransformerWeights(str(self.checkpoint)), + boundary, + device=device, + ) + t_done = self._now_for_trace() + self.trace["engine_init"] = { + "mode": "prepare_live", + "prepare_source": self.boundary_prepare_live.latest_source, + "prepare_slim_no_raw_state_vision": self.boundary_prepare_live._last_effective_slim_no_raw_state_vision, + "prepare_slim_derive_condition_reference": self.boundary_prepare_live._last_effective_slim_derive_condition_reference, + "prepare_slim_derive_initial_noise": self.boundary_prepare_live._last_effective_slim_derive_initial_noise, + "boundary_derive_s": t_derived - t0, + "boundary_save_s": t_saved - t_derived, + "boundary_validate_s": t_validated - t_saved, + "engine_construct_s": t_done - t_validated, + "total_s": t_done - t0, + "wrote_boundary": self.boundary_out is not None, + } + + def should_capture_pre_layers(self) -> bool: + return ( + self.prelayer_bootstrap + and self.engine is None + and self.active_step == 0 + and not self.current_skip_text_tokens + and self.current_noise_x is not None + and self.current_timestep is not None + ) + + def capture_pre_layers_and_raise( + self, + text_model: Any, + pack: Any, + _attention_mask: Any, + position_ids: Any, + ) -> None: + if not self.should_capture_pre_layers(): + return + from cosmos_framework.data.generator.sequence_packing.runtime import from_all_seq, get_device_and_dtype + + if isinstance(pack, dict): + for key, value in pack.items(): + self._add_tensor(f"s00/lm_in/{key}", value) + if key in {"causal_seq", "full_only_seq"}: + self._add_tensor(f"s00/layers/00/input/{key}", value) + self._add_tensor("s00/lm_in/position_ids", position_ids) + + device, dtype = get_device_and_dtype(pack) + meta = self.torch.tensor([], dtype=dtype, device=device) + cos, sin = text_model.rotary_emb( + meta, + position_ids=position_ids.unsqueeze(0) if position_ids.ndim == 1 else position_ids.unsqueeze(1), + ) + cos_pack = from_all_seq(cos.squeeze(0), pack) + sin_pack = from_all_seq(sin.squeeze(0), pack) + if isinstance(cos_pack, dict) and isinstance(sin_pack, dict): + for key in ("causal_seq", "full_only_seq"): + self._add_tensor(f"s00/layers/00/rope/cos/{key}", cos_pack.get(key)) + self._add_tensor(f"s00/layers/00/rope/sin/{key}", sin_pack.get(key)) + + assert self.current_noise_x is not None + assert self.current_timestep is not None + self._initialise_engine(self.current_noise_x, self.current_timestep) + assert self.engine is not None + t0 = self._now_for_trace() + velocity = self.engine.flat_velocity_for_step(self.current_noise_x[0], self.current_timestep) + t1 = self._now_for_trace() + self.tensors["steps/00/velocity"] = self._cpu(velocity) + self.trace["flashrt_velocity_calls"].append({"step": self.step, "s": t1 - t0, "bootstrap": "prelayer"}) + raise _FlashRTVelocityReady([velocity]) + + def get_velocity(self, original: Callable[..., Any], model: Any, *args: Any, **kwargs: Any) -> list[Any]: + noise_x = kwargs.get("noise_x") + timestep = kwargs.get("timestep") + skip_text_tokens = bool(kwargs.get("skip_text_tokens", False)) + if self.engine is None and self.boundary_in is not None and not skip_text_tokens: + if not isinstance(noise_x, list) or len(noise_x) != 1: + raise RuntimeError("FlashRT live boundary input currently supports batch size 1") + self._initialise_engine_from_boundary_in(noise_x[0].device) + if self.engine is None and self.boundary_prepare_in is not None and not skip_text_tokens: + if not isinstance(noise_x, list) or len(noise_x) != 1: + raise RuntimeError("FlashRT live prepare-boundary input currently supports batch size 1") + self._initialise_engine_from_prepare_in(noise_x[0].device) + if self.engine is None and self.boundary_prepare_live is not None and not skip_text_tokens: + if not isinstance(noise_x, list) or len(noise_x) != 1: + raise RuntimeError("FlashRT live prepare-boundary source currently supports batch size 1") + self._initialise_engine_from_prepare_live(noise_x[0].device) + if self.engine is not None and not skip_text_tokens: + if not isinstance(noise_x, list) or len(noise_x) != 1: + raise RuntimeError("FlashRT live handoff currently supports batch size 1") + t0 = self._now_for_trace() + velocity = self.engine.flat_velocity_for_step(noise_x[0], timestep) + t1 = self._now_for_trace() + self.trace["flashrt_velocity_calls"].append({"step": self.step, "s": t1 - t0}) + self.step += 1 + return [velocity] + + self.active_step = self.step + self.current_noise_x = noise_x if isinstance(noise_x, list) else None + self.current_timestep = timestep + self.current_skip_text_tokens = skip_text_tokens + t0 = self._now_for_trace() + try: + velocity = original(model, *args, **kwargs) + except _FlashRTVelocityReady as ready: + t1 = self._now_for_trace() + self.trace["official_velocity_calls"].append( + { + "step": self.step, + "s": t1 - t0, + "skip_text_tokens": skip_text_tokens, + "prelayer_aborted": True, + } + ) + self.step += 1 + return ready.velocity + finally: + self.active_step = None + self.current_noise_x = None + self.current_timestep = None + self.current_skip_text_tokens = False + t1 = self._now_for_trace() + self.trace["official_velocity_calls"].append( + {"step": self.step, "s": t1 - t0, "skip_text_tokens": skip_text_tokens} + ) + if self.step == 0 and not skip_text_tokens: + if not isinstance(noise_x, list) or timestep is None: + raise RuntimeError("FlashRT live handoff could not see step-0 noise/timestep") + self._initialise_engine(noise_x, timestep, velocity) + self.step += 1 + return velocity + + def _record_native_scheduler_fallback(self, reason: str) -> None: + native = self.trace.get("native_scheduler") + if isinstance(native, dict): + native.setdefault("fallbacks", []).append({"reason": reason}) + + def _record_native_scheduler_failure(self, reason: str) -> None: + native = self.trace.get("native_scheduler") + if isinstance(native, dict): + native.setdefault("failures", []).append({"reason": reason}) + + @staticmethod + def _single_latent(noise: Any) -> tuple[Any | None, bool, str | None]: + if isinstance(noise, list): + if len(noise) != 1: + return None, True, "noise_list_len_not_1" + return noise[0], True, None + return noise, False, None + + def try_native_unipc_forward( + self, + velocity_fn: Callable[..., Any], + noise: Any, + *, + num_steps: int, + shift: float | None, + seed: Any, + ) -> tuple[bool, Any]: + self.current_scheduler_seed = seed + native_trace = self.trace.get("native_scheduler") + if not isinstance(native_trace, dict) or not native_trace.get("enabled", False): + self._record_native_scheduler_fallback("disabled") + return False, None + if seed is not None and isinstance(seed, list) and len(seed) != 1: + self._record_native_scheduler_fallback("seed_list_len_not_1") + return False, None + if shift is None: + self._record_native_scheduler_fallback("missing_shift") + return False, None + + from flash_rt.models.cosmos3_edge.dump_replay import EDGE_FLAT_DIM, EDGE_NUM_STEPS + + if int(num_steps) != EDGE_NUM_STEPS: + self._record_native_scheduler_fallback(f"num_steps_{num_steps}") + return False, None + + latent_in, return_list, reason = self._single_latent(noise) + if reason is not None: + self._record_native_scheduler_fallback(reason) + return False, None + if not isinstance(latent_in, self.torch.Tensor): + self._record_native_scheduler_fallback("noise_not_tensor") + return False, None + if latent_in.device.type != "cuda": + self._record_native_scheduler_fallback("noise_not_cuda") + return False, None + if latent_in.dtype != self.torch.float32: + self._record_native_scheduler_fallback(f"noise_dtype_{latent_in.dtype}") + return False, None + if int(latent_in.numel()) != EDGE_FLAT_DIM: + self._record_native_scheduler_fallback(f"noise_numel_{int(latent_in.numel())}") + return False, None + if ( + self.boundary_prepare_live is not None + and self.engine is not None + and self.boundary_prepare_live.latest_signature != self._engine_prepare_signature + ): + self.engine = None + self._engine_prepare_signature = None + + from cosmos_framework.utils.progress_bar import progress_bar + from flash_rt.models.cosmos3_edge.static_unipc import EdgeStaticUniPCScheduler + + scheduler = EdgeStaticUniPCScheduler(int(num_steps), device=latent_in.device, shift=float(shift)) + if not scheduler.native_available: + self._record_native_scheduler_fallback("native_binding_unavailable") + return False, None + + latent = latent_in.clone() + scheduler.reset(latent) + step_total_s = 0.0 + velocity_dtype: str | None = None + t_run0 = self._now_for_trace() + try: + for step_index, timestep in enumerate( + progress_bar(scheduler.timesteps, desc="Sampling", total=len(scheduler.timesteps)) + ): + velocity_pred = velocity_fn([latent] if return_list else latent, timestep.reshape(1, 1)) + velocity = velocity_pred[0] if return_list and isinstance(velocity_pred, list) else velocity_pred + if not isinstance(velocity, self.torch.Tensor): + raise TypeError(f"native UniPC expected tensor velocity, got {type(velocity).__name__}") + if velocity_dtype is None: + velocity_dtype = str(velocity.dtype) + if velocity.dtype != self.torch.bfloat16: + raise TypeError(f"native UniPC expected bf16 velocity, got {velocity.dtype}") + if tuple(velocity.shape) != tuple(latent.shape): + raise ValueError(f"native UniPC velocity shape {tuple(velocity.shape)} != latent {tuple(latent.shape)}") + t_step0 = self._now_for_trace() + scheduler.step(latent, velocity, step_index) + t_step1 = self._now_for_trace() + step_total_s += t_step1 - t_step0 + except Exception as exc: + self._record_native_scheduler_failure(type(exc).__name__) + raise + t_run1 = self._now_for_trace() + native_trace.setdefault("runs", []).append( + { + "num_steps": int(num_steps), + "shift": float(shift), + "total_s": t_run1 - t_run0, + "scheduler_step_total_s": step_total_s, + "scheduler_step_avg_s": step_total_s / int(num_steps), + "latent_dtype": str(latent_in.dtype), + "velocity_dtype": velocity_dtype, + "return_list": return_list, + } + ) + return True, [latent] if return_list else latent + + def save_trace(self) -> None: + if self.trace_out is None: + return + flashrt_calls = self.trace["flashrt_velocity_calls"] + official_calls = self.trace["official_velocity_calls"] + flashrt_total = sum(float(item["s"]) for item in flashrt_calls) + official_total = sum(float(item["s"]) for item in official_calls) + native_scheduler = self.trace.get("native_scheduler", {}) + native_runs = native_scheduler.get("runs", []) if isinstance(native_scheduler, dict) else [] + native_step_total = sum(float(item.get("scheduler_step_total_s", 0.0)) for item in native_runs) + summary = { + "flashrt_velocity_call_count": len(flashrt_calls), + "flashrt_velocity_total_s": flashrt_total, + "flashrt_velocity_avg_s": flashrt_total / len(flashrt_calls) if flashrt_calls else 0.0, + "official_velocity_call_count": len(official_calls), + "official_velocity_total_s": official_total, + "native_scheduler_enabled": bool(native_scheduler.get("enabled")) if isinstance(native_scheduler, dict) else False, + "native_scheduler_run_count": len(native_runs), + "native_scheduler_step_count": sum(int(item.get("num_steps", 0)) for item in native_runs), + "native_scheduler_step_total_s": native_step_total, + } + self.trace_out.parent.mkdir(parents=True, exist_ok=True) + self.trace_out.write_text( + json.dumps({"summary": summary, **self.trace}, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _install_live_flashrt_handoff( + checkpoint: Path, + boundary_out: Path | None, + boundary_in: Path | None, + boundary_prepare_in: Path | None, + boundary_prepare_live: _PrepareBoundary | None, + trace_out: Path | None, + *, + prelayer_bootstrap: bool = False, +) -> _LiveFlashRTHandoff: + from cosmos_framework.model.generator.diffusion.samplers.unipc import UniPCSampler + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel + from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork + from cosmos_framework.model.generator.mot.unified_mot import ( + MoTDecoderLayer, + Nemotron3DenseVLTextForCausalLM, + Qwen3VLTextForCausalLM, + Qwen3VLTextModel, + Qwen3VLMoeTextForCausalLM, + Qwen3VLMoeTextModel, + Nemotron3DenseVLTextModel, + ) + + handoff = _LiveFlashRTHandoff( + checkpoint, + boundary_out, + boundary_in, + boundary_prepare_in, + boundary_prepare_live, + trace_out, + prelayer_bootstrap=prelayer_bootstrap, + ) + orig_unipc_forward = UniPCSampler.forward + orig_get_velocity = OmniMoTModel._get_velocity + orig_vfm_forward = Cosmos3VFMNetwork.forward + orig_layer_forward = MoTDecoderLayer.forward + lm_classes = (Nemotron3DenseVLTextForCausalLM, Qwen3VLTextForCausalLM, Qwen3VLMoeTextForCausalLM) + orig_lm_forwards = {cls: cls.forward for cls in lm_classes} + text_model_classes = (Nemotron3DenseVLTextModel, Qwen3VLTextModel, Qwen3VLMoeTextModel) + orig_text_model_forwards = {cls: cls.forward for cls in text_model_classes} + + def _get_velocity_hook(self: Any, *args: Any, **kwargs: Any) -> list[Any]: + return handoff.get_velocity(orig_get_velocity, self, *args, **kwargs) + + def _unipc_forward_hook( + self: Any, + velocity_fn: Callable[..., Any], + noise: Any, + num_steps: int = 35, + shift: float | None = None, + seed: Any = None, + ) -> Any: + handled, output = handoff.try_native_unipc_forward( + velocity_fn, + noise, + num_steps=num_steps, + shift=self.cfg.shift if shift is None else shift, + seed=seed, + ) + if handled: + return output + return orig_unipc_forward( + self, + velocity_fn, + noise, + num_steps=num_steps, + shift=shift, + seed=seed, + ) + + def _vfm_forward_hook(self: Any, packed_seq: Any, *args: Any, **kwargs: Any) -> Any: + handoff.capture_vfm_boundary(packed_seq) + return orig_vfm_forward(self, packed_seq, *args, **kwargs) + + def _layer_forward_hook(self: Any, input_pack: Any, *args: Any, **kwargs: Any) -> Any: + output = orig_layer_forward(self, input_pack, *args, **kwargs) + handoff.capture_layer0(self, input_pack, args, output) + return output + + def _make_lm_forward_hook(cls: type[Any]) -> Callable[..., Any]: + orig = orig_lm_forwards[cls] + + def _lm_forward_hook(self: Any, pack: Any, attention_mask: Any, position_ids: Any, *args: Any, **kwargs: Any) -> Any: + handoff.capture_lm_boundary(pack, attention_mask, position_ids) + return orig(self, pack, attention_mask, position_ids, *args, **kwargs) + + return _lm_forward_hook + + def _make_text_model_forward_hook(cls: type[Any]) -> Callable[..., Any]: + orig = orig_text_model_forwards[cls] + + def _text_model_forward_hook( + self: Any, + pack: Any, + attention_mask: Any, + position_ids: Any, + *args: Any, + **kwargs: Any, + ) -> Any: + handoff.capture_pre_layers_and_raise(self, pack, attention_mask, position_ids) + return orig(self, pack, attention_mask, position_ids, *args, **kwargs) + + return _text_model_forward_hook + + UniPCSampler.forward = _unipc_forward_hook + OmniMoTModel._get_velocity = _get_velocity_hook + Cosmos3VFMNetwork.forward = _vfm_forward_hook + MoTDecoderLayer.forward = _layer_forward_hook + for cls in lm_classes: + cls.forward = _make_lm_forward_hook(cls) + for cls in text_model_classes: + cls.forward = _make_text_model_forward_hook(cls) + return handoff + + +def main(argv: Sequence[str] | None = None) -> None: + argv = list(sys.argv[1:] if argv is None else argv) + argv, live_dump_out = _extract_live_dump_out(argv) + argv, upstream_trace_out = _extract_upstream_trace_out(argv) + argv, cache_warmup_vae = _extract_warmup_vae_cache(argv) + argv, cache_warmup_prepare = _extract_warmup_prepare_cache(argv) + argv, vae_encode_dump_out, vae_latent_in, vae_encode_dump_input = _extract_vae_encode_boundary_args(argv) + argv, vae_encode_profile_out = _extract_vae_encode_profile_out(argv) + argv, vae_native_rms_silu = _extract_vae_native_rms_silu(argv) + argv, vae_t1_conv2d = _extract_vae_t1_conv2d(argv) + argv, vae_native_avgdown3d = _extract_vae_native_avgdown3d(argv) + argv, vae_channels_last3d_conv320 = _extract_vae_channels_last3d_conv320(argv) + argv, vae_compile_encode, vae_compile_trace_out = _extract_vae_compile_encode(argv) + ( + argv, + prepare_dump_out, + prepare_replay_in, + prepare_inventory_out, + prepare_slim_no_raw_state_vision, + prepare_slim_derive_condition_reference, + prepare_slim_derive_initial_noise, + ) = _extract_prepare_boundary_args(argv) + ( + argv, + live_handoff, + live_boundary_out, + live_boundary_in, + live_boundary_prepare_in, + live_boundary_prepare_live, + live_handoff_trace_out, + live_prelayer_bootstrap, + ) = _extract_live_handoff_args(argv) + if live_dump_out is not None and live_handoff: + raise SystemExit("--flashrt-live-dump-out cannot be combined with --flashrt-live-flashrt-handoff") + if live_boundary_in is not None and live_boundary_prepare_in is not None: + raise SystemExit("--flashrt-live-boundary-in cannot be combined with --flashrt-live-boundary-prepare-in") + if live_boundary_prepare_live and (live_boundary_in is not None or live_boundary_prepare_in is not None): + raise SystemExit( + "--flashrt-live-boundary-prepare-live cannot be combined with live boundary input artifacts" + ) + if cache_warmup_vae and (vae_encode_dump_out is not None or vae_latent_in is not None): + raise SystemExit("--flashrt-cache-warmup-vae cannot be combined with VAE encode boundary dump/replay") + if cache_warmup_prepare and ( + prepare_dump_out is not None + or prepare_replay_in is not None + or prepare_inventory_out is not None + ): + raise SystemExit( + "--flashrt-cache-warmup-prepare cannot be combined with prepare boundary dump/replay/inventory" + ) + if vae_compile_encode and (vae_native_rms_silu or vae_native_avgdown3d): + raise SystemExit( + "--flashrt-vae-compile-encode cannot be combined with native VAE monkeypatches" + ) + output_dir = _extract_output_dir(argv) + _install_action_only_patches() + if live_dump_out is not None: + _install_live_dump_patch(live_dump_out) + upstream_trace: _UpstreamTrace | None = None + handoff: _LiveFlashRTHandoff | None = None + vae_profiler: _VAEEncodeProfiler | None = None + vae_compile_patch: _VAECompileEncodePatch | None = None + if vae_compile_encode: + vae_compile_patch = _install_vae_compile_encode_patch(vae_compile_trace_out) + if vae_t1_conv2d: + from flash_rt.models.cosmos3_edge.vae_native import install_wan_vae_encode_t1_conv2d + + install_wan_vae_encode_t1_conv2d() + if vae_native_avgdown3d: + from flash_rt.models.cosmos3_edge.vae_native import install_wan_vae_encode_native_avgdown3d + + install_wan_vae_encode_native_avgdown3d() + if vae_channels_last3d_conv320: + from flash_rt.models.cosmos3_edge.vae_native import install_wan_vae_encode_channels_last3d_conv320 + + install_wan_vae_encode_channels_last3d_conv320() + if vae_native_rms_silu: + from flash_rt.models.cosmos3_edge.vae_native import install_wan_vae_encode_native_rms_silu + + install_wan_vae_encode_native_rms_silu() + if vae_encode_profile_out is not None: + vae_profiler = _install_vae_encode_profile_patch(vae_encode_profile_out) + if live_handoff: + _install_warmup_clone_patch() + vae_cache: _WarmupVAECache | None = None + if cache_warmup_vae: + vae_cache = _install_warmup_vae_cache_patch() + prepare_cache: _WarmupPrepareCache | None = None + if cache_warmup_prepare: + prepare_cache = _install_warmup_prepare_cache_patch( + slim_no_raw_state_vision=prepare_slim_no_raw_state_vision, + slim_derive_condition_reference=prepare_slim_derive_condition_reference, + slim_derive_initial_noise=prepare_slim_derive_initial_noise, + ) + prepare_boundary: _PrepareBoundary | None = None + if ( + prepare_dump_out is not None + or prepare_replay_in is not None + or prepare_inventory_out is not None + or prepare_slim_no_raw_state_vision + or prepare_slim_derive_condition_reference + or prepare_slim_derive_initial_noise + or live_boundary_prepare_live + ): + prepare_boundary = _install_prepare_boundary_patch( + prepare_dump_out, + prepare_replay_in, + prepare_inventory_out, + slim_no_raw_state_vision=prepare_slim_no_raw_state_vision, + slim_derive_condition_reference=prepare_slim_derive_condition_reference, + slim_derive_initial_noise=prepare_slim_derive_initial_noise, + ) + vae_boundary: _VAEEncodeBoundary | None = None + if vae_encode_dump_out is not None or vae_latent_in is not None: + vae_boundary = _install_vae_encode_boundary_patch( + vae_encode_dump_out, + vae_latent_in, + dump_input=vae_encode_dump_input, + ) + if upstream_trace_out is not None: + upstream_trace = _install_upstream_trace_patch( + upstream_trace_out, + vae_cache=vae_cache, + prepare_cache=prepare_cache, + vae_boundary=vae_boundary, + prepare_boundary=prepare_boundary, + ) + if live_handoff: + handoff = _install_live_flashrt_handoff( + _extract_checkpoint_path(argv), + live_boundary_out, + live_boundary_in, + live_boundary_prepare_in, + prepare_boundary if live_boundary_prepare_live else None, + live_handoff_trace_out, + prelayer_bootstrap=live_prelayer_bootstrap, + ) + + from cosmos_framework.scripts.inference import main as official_main + + old_argv = sys.argv + try: + sys.argv = [old_argv[0], *argv] + official_main() + finally: + sys.argv = old_argv + + rewritten = _rewrite_sample_outputs_action_only(output_dir) + if rewritten == 0: + raise SystemExit(f"no action outputs were rewritten under {output_dir}") + if handoff is not None: + handoff.save_trace() + if upstream_trace is not None: + upstream_trace.save() + if vae_compile_patch is not None: + vae_compile_patch.save() + if vae_profiler is not None: + vae_profiler.save() + + +if __name__ == "__main__": + main() diff --git a/flash_rt/models/cosmos3_edge/boundary_dump.py b/flash_rt/models/cosmos3_edge/boundary_dump.py new file mode 100644 index 00000000..dd25c311 --- /dev/null +++ b/flash_rt/models/cosmos3_edge/boundary_dump.py @@ -0,0 +1,151 @@ +"""Boundary dump loader for Cosmos3-Edge native forward bring-up.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import torch +from safetensors.torch import load_file + + +EDGE_BOUNDARY_UND_TOKENS = 125 +EDGE_BOUNDARY_GEN_TOKENS = 6300 +EDGE_BOUNDARY_VISION_TOKENS = 6240 +EDGE_BOUNDARY_ACTION_TOKENS = 60 +EDGE_HIDDEN_SIZE = 2048 +EDGE_HEAD_DIM = 128 +EDGE_PATCH_SPATIAL = 2 + + +@dataclass(frozen=True) +class EdgeBoundaryShapes: + und_tokens: int + gen_tokens: int + vision_tokens: int + action_tokens: int + hidden_size: int + head_dim: int + + +class EdgeBoundaryDump: + """Step-0 transformer boundary captured from the official Edge pipeline.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + self.tensors = load_file(str(self.path), device="cpu") + self._validate_required_tensors() + + @classmethod + def from_tensors(cls, tensors: dict[str, torch.Tensor], *, path: str | Path = "") -> "EdgeBoundaryDump": + obj = cls.__new__(cls) + obj.path = Path(path) + obj.tensors = tensors + obj._validate_required_tensors() + return obj + + def _validate_required_tensors(self) -> None: + required = { + "s00/lm_in/causal_seq", + "s00/lm_in/full_only_seq", + "s00/lm_in/position_ids", + "s00/layers/00/rope/cos/causal_seq", + "s00/layers/00/rope/cos/full_only_seq", + "s00/layers/00/rope/sin/causal_seq", + "s00/layers/00/rope/sin/full_only_seq", + "s00/layers/00/input/causal_seq", + "s00/layers/00/input/full_only_seq", + "s00/vfm_in/vision/tokens/0", + "s00/vfm_in/action/tokens/0", + "s00/vfm_in/action/domain_id/0", + "s00/vfm_in/action/raw_action_dim/0", + } + missing = sorted(required.difference(self.tensors)) + if missing: + raise ValueError(f"{self.path} is missing boundary tensors: {missing[:8]}") + + @property + def shapes(self) -> EdgeBoundaryShapes: + return EdgeBoundaryShapes( + und_tokens=int(self.causal_seq.shape[0]), + gen_tokens=int(self.full_only_seq.shape[0]), + vision_tokens=int( + self.vision_tokens.shape[2] + * (self.vision_tokens.shape[3] // EDGE_PATCH_SPATIAL) + * (self.vision_tokens.shape[4] // EDGE_PATCH_SPATIAL) + ), + action_tokens=int(self.action_tokens.shape[0]), + hidden_size=int(self.causal_seq.shape[1]), + head_dim=int(self.rope_cos_causal.shape[1]), + ) + + @property + def causal_seq(self) -> torch.Tensor: + return self.tensors["s00/lm_in/causal_seq"] + + @property + def full_only_seq(self) -> torch.Tensor: + return self.tensors["s00/lm_in/full_only_seq"] + + @property + def position_ids(self) -> torch.Tensor: + return self.tensors["s00/lm_in/position_ids"] + + @property + def rope_cos_causal(self) -> torch.Tensor: + return self.tensors["s00/layers/00/rope/cos/causal_seq"] + + @property + def rope_cos_full(self) -> torch.Tensor: + return self.tensors["s00/layers/00/rope/cos/full_only_seq"] + + @property + def rope_sin_causal(self) -> torch.Tensor: + return self.tensors["s00/layers/00/rope/sin/causal_seq"] + + @property + def rope_sin_full(self) -> torch.Tensor: + return self.tensors["s00/layers/00/rope/sin/full_only_seq"] + + @property + def layer0_input_causal(self) -> torch.Tensor: + return self.tensors["s00/layers/00/input/causal_seq"] + + @property + def layer0_input_full(self) -> torch.Tensor: + return self.tensors["s00/layers/00/input/full_only_seq"] + + @property + def layer0_output_causal(self) -> torch.Tensor: + return self.tensors["s00/layers/00/output/causal_seq"] + + @property + def layer0_output_full(self) -> torch.Tensor: + return self.tensors["s00/layers/00/output/full_only_seq"] + + @property + def vision_tokens(self) -> torch.Tensor: + return self.tensors["s00/vfm_in/vision/tokens/0"] + + @property + def action_tokens(self) -> torch.Tensor: + return self.tensors["s00/vfm_in/action/tokens/0"] + + def validate_geometry(self) -> None: + shapes = self.shapes + expected = EdgeBoundaryShapes( + und_tokens=EDGE_BOUNDARY_UND_TOKENS, + gen_tokens=EDGE_BOUNDARY_GEN_TOKENS, + vision_tokens=EDGE_BOUNDARY_VISION_TOKENS, + action_tokens=EDGE_BOUNDARY_ACTION_TOKENS, + hidden_size=EDGE_HIDDEN_SIZE, + head_dim=EDGE_HEAD_DIM, + ) + if shapes != expected: + raise ValueError(f"unexpected Edge boundary geometry: expected {expected}, got {shapes}") + if self.layer0_input_causal.shape != self.causal_seq.shape: + raise ValueError("layer0 causal input shape does not match lm_in causal_seq") + if self.layer0_input_full.shape != self.full_only_seq.shape: + raise ValueError("layer0 full input shape does not match lm_in full_only_seq") + if self.position_ids.shape != (3, EDGE_BOUNDARY_UND_TOKENS + EDGE_BOUNDARY_GEN_TOKENS): + raise ValueError(f"unexpected position_ids shape: {tuple(self.position_ids.shape)}") diff --git a/flash_rt/models/cosmos3_edge/denoise_ref.py b/flash_rt/models/cosmos3_edge/denoise_ref.py new file mode 100644 index 00000000..4f8b1825 --- /dev/null +++ b/flash_rt/models/cosmos3_edge/denoise_ref.py @@ -0,0 +1,322 @@ +"""Torch reference denoise loop for Cosmos3-Edge AV inverse dynamics.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump +from flash_rt.models.cosmos3_edge.dump_replay import ( + EDGE_ACTION_MODEL_SHAPE, + EDGE_NUM_TRAIN_TIMESTEPS, + EDGE_SHIFT, + EdgeDenoiseDump, + EdgeLatentParts, +) +from flash_rt.models.cosmos3_edge.layer_ref import EdgeTransformerTorchReference +from flash_rt.models.cosmos3_edge.static_engine import EdgeStaticBufferEngine +from flash_rt.models.cosmos3_edge.static_unipc import EdgeStaticUniPCScheduler +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights +from flash_rt.models.cosmos3_video.fm_solvers_unipc import FlowUniPCMultistepScheduler + + +@dataclass(frozen=True) +class ReferenceDenoiseResult: + final_flat: torch.Tensor + final_parts: EdgeLatentParts + final_action: torch.Tensor + timesteps: tuple[int, ...] + max_input_abs_diff: float + max_velocity_abs_diff: float + steps_run: int + + +class EdgeDenoiseTorchReference: + """Full denoise loop that computes action velocity with the Torch reference. + + This is still a correctness scaffold, not the optimized FlashRT engine. For + AV inverse dynamics the official dump predicts zero velocity for the vision + segment, so only the action tail is produced by the transformer reference. + """ + + def __init__( + self, + denoise_dump: EdgeDenoiseDump, + boundary_dump: EdgeBoundaryDump, + weights: EdgeTransformerWeights, + *, + device: str | torch.device = "cuda", + shift: float = EDGE_SHIFT, + use_static_und_cache: bool = True, + use_cuda_graphs: bool = False, + ): + self.denoise_dump = denoise_dump + self.boundary_dump = boundary_dump + self.device = torch.device(device) + self.shift = float(shift) + self.use_cuda_graphs = bool(use_cuda_graphs) + self.transformer = EdgeTransformerTorchReference(weights, device=self.device) + self.static_engine = None + self.denoise_dump.validate_geometry() + self.boundary_dump.validate_geometry() + self.und_cache = None + self.static_scheduler = None + if use_static_und_cache: + self.static_engine = EdgeStaticBufferEngine(weights, self.boundary_dump, device=self.device) + self.und_cache = self.static_engine.und_cache + + def _scheduler(self) -> FlowUniPCMultistepScheduler: + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=EDGE_NUM_TRAIN_TIMESTEPS, + shift=1.0, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(self.denoise_dump.num_steps, device=self.device, shift=self.shift) + expected = tuple(int(t.item()) for t in scheduler.timesteps.cpu()) + actual = self.denoise_dump.timesteps() + if expected != actual: + raise ValueError(f"dump timesteps {actual} do not match UniPC schedule {expected}") + return scheduler + + @staticmethod + def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max().item()) + + def velocity_for_step( + self, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + *, + step_index: int | None = None, + ) -> torch.Tensor: + latent = flat_noise.to(device=self.device) + if self.static_engine is not None: + if self.use_cuda_graphs: + return self.static_engine.replay_velocity_graph(latent, timestep) + return self.static_engine.flat_velocity_for_step(latent, timestep, timestep_index=step_index) + + velocity = torch.zeros_like(latent) + if self.und_cache is None: + action_velocity = self.transformer.action_velocity_for_step(self.boundary_dump, latent, timestep) + else: + action_velocity = self.transformer.action_velocity_for_step_with_und_cache( + self.boundary_dump, + latent, + timestep, + self.und_cache, + ) + action_flat = action_velocity.reshape(EDGE_ACTION_MODEL_SHAPE[0] * EDGE_ACTION_MODEL_SHAPE[1]) + velocity[-action_flat.numel() :] = action_flat.to(dtype=velocity.dtype) + return velocity + + def run(self, *, max_steps: int | None = None, check_against_dump: bool = True) -> ReferenceDenoiseResult: + scheduler = self._scheduler() + steps_run = self.denoise_dump.num_steps if max_steps is None else min(max_steps, self.denoise_dump.num_steps) + latent = self.denoise_dump.step_noise(0).to(device=self.device) + input_diffs: list[float] = [] + velocity_diffs: list[float] = [] + static_scheduler = self.static_scheduler + if static_scheduler is not None: + static_scheduler.reset(latent) + + for step, timestep in enumerate(scheduler.timesteps[:steps_run]): + if check_against_dump: + expected_noise = self.denoise_dump.step_noise(step).to(device=self.device) + input_diffs.append(self._max_abs_diff(latent, expected_noise)) + velocity = self.velocity_for_step(latent, timestep.reshape(1, 1), step_index=step) + if check_against_dump: + expected_velocity = self.denoise_dump.step_velocity(step).to(device=self.device) + velocity_diffs.append(self._max_abs_diff(velocity, expected_velocity)) + if static_scheduler is not None: + latent = static_scheduler.step(latent, velocity, step) + else: + latent = scheduler.step( + model_output=velocity, + timestep=timestep, + sample=latent.unsqueeze(0), + return_dict=False, + )[0].squeeze(0) + + final_flat = latent.detach().cpu().contiguous() + parts = self.denoise_dump.split_flat(final_flat) + raw_action_dim = int(self.boundary_dump.tensors["s00/vfm_in/action/raw_action_dim/0"].item()) + final_action = parts.action_model[:, :raw_action_dim].float().contiguous() + return ReferenceDenoiseResult( + final_flat=final_flat, + final_parts=parts, + final_action=final_action, + timesteps=tuple(int(t.item()) for t in scheduler.timesteps[:steps_run].cpu()), + max_input_abs_diff=max(input_diffs) if input_diffs else 0.0, + max_velocity_abs_diff=max(velocity_diffs) if velocity_diffs else 0.0, + steps_run=steps_run, + ) + + +class EdgeDenoiseFlashRTQuant: + """Quantized static-graph FlashRT denoise runner (Thor). + + Wraps ``CosmosEdgeThor``: pre-quantized gen-tower weights, static und K/V, + FA4 attention, native UniPC, and (by default) the entire 30-step denoise + captured in a single CUDA graph. The eager path is kept for warmup and + per-step dump comparisons. + """ + + def __init__( + self, + denoise_dump: EdgeDenoiseDump, + boundary_dump: EdgeBoundaryDump, + weights: EdgeTransformerWeights, + *, + device: str | torch.device = "cuda", + shift: float = EDGE_SHIFT, + quant: str = "fp8", + bf16_projs: tuple[str, ...] = (), + ffn_fp4: bool = False, + slim_last: bool = True, + use_cuda_graphs: bool = True, + ): + from flash_rt.models.cosmos3_edge.pipeline_thor import CosmosEdgeThor + + if torch.device(device).type != "cuda": + raise RuntimeError("EdgeDenoiseFlashRTQuant requires CUDA") + self.denoise_dump = denoise_dump + self.boundary_dump = boundary_dump + self.device = torch.device(device) + self.use_cuda_graphs = bool(use_cuda_graphs) + self.denoise_dump.validate_geometry() + self.boundary_dump.validate_geometry() + self.engine = CosmosEdgeThor( + weights, + boundary_dump, + self.denoise_dump.timesteps(), + quant=quant, + bf16_projs=bf16_projs, + ffn_fp4=ffn_fp4, + slim_last=slim_last, + shift=shift, + ) + self.static_engine = None + self.static_scheduler = self.engine.unipc + self.engine.calibrate(self.denoise_dump.step_noise(0)) + if self.use_cuda_graphs: + self.engine.capture(warmup_noise=self.denoise_dump.step_noise(0)) + + @property + def native_attention_available(self) -> bool: + return True + + @property + def graph_attention_available(self) -> bool: + return self.engine.graph is not None + + @property + def native_scheduler_available(self) -> bool: + return True + + def velocity_for_step( + self, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + *, + step_index: int | None = None, + ) -> torch.Tensor: + if step_index is None: + raise ValueError("quant engine requires an explicit step index") + del timestep # baked into the per-step embed table + if flat_noise.data_ptr() != self.engine.latent.data_ptr(): + self.engine.latent.copy_(flat_noise.to(device=self.device, dtype=torch.float32)) + return self.engine.forward_step(step_index, self.engine.latent) + + def _result(self, steps_run: int, input_diffs: list[float], velocity_diffs: list[float]) -> ReferenceDenoiseResult: + final_flat = self.engine.latent.detach().cpu().contiguous() + parts = self.denoise_dump.split_flat(final_flat) + final_action = parts.action_model[:, : self.engine.raw_action_dim].float().contiguous() + return ReferenceDenoiseResult( + final_flat=final_flat, + final_parts=parts, + final_action=final_action, + timesteps=self.denoise_dump.timesteps()[:steps_run], + max_input_abs_diff=max(input_diffs) if input_diffs else 0.0, + max_velocity_abs_diff=max(velocity_diffs) if velocity_diffs else 0.0, + steps_run=steps_run, + ) + + def run(self, *, max_steps: int | None = None, check_against_dump: bool = True) -> ReferenceDenoiseResult: + num_steps = self.denoise_dump.num_steps + steps_run = num_steps if max_steps is None else min(max_steps, num_steps) + noise0 = self.denoise_dump.step_noise(0) + if self.engine.graph is not None and steps_run == num_steps: + # Whole-denoise graph replay: per-step dump checks are not + # observable inside the graph; the final-action gate is authoritative. + self.engine.denoise(noise0) + return self._result(steps_run, [], []) + self.engine.latent.copy_(noise0.to(device=self.device, dtype=torch.float32)) + input_diffs: list[float] = [] + velocity_diffs: list[float] = [] + for step in range(steps_run): + if check_against_dump: + expected = self.denoise_dump.step_noise(step).to(device=self.device) + input_diffs.append(float((self.engine.latent.float() - expected.float()).abs().max().item())) + velocity = self.engine.forward_step(step, self.engine.latent) + if check_against_dump: + expected_v = self.denoise_dump.step_velocity(step).to(device=self.device) + velocity_diffs.append(float((velocity.float() - expected_v.float()).abs().max().item())) + self.engine.unipc.step(self.engine.latent, velocity, step) + return self._result(steps_run, input_diffs, velocity_diffs) + + +class EdgeDenoiseFlashRT(EdgeDenoiseTorchReference): + """Optimized eager FlashRT denoise engine for Cosmos3-Edge AV. + + The scheduler and correctness accounting intentionally stay shared with + ``EdgeDenoiseTorchReference``. The model velocity path is the FlashRT static + engine: static vision tokens, cached und K/V, FVK BF16 GEMMs/RMSNorm/RoPE, + and native Thor attention when available. + """ + + def __init__( + self, + denoise_dump: EdgeDenoiseDump, + boundary_dump: EdgeBoundaryDump, + weights: EdgeTransformerWeights, + *, + device: str | torch.device = "cuda", + shift: float = EDGE_SHIFT, + use_cuda_graphs: bool = False, + ): + super().__init__( + denoise_dump, + boundary_dump, + weights, + device=device, + shift=shift, + use_static_und_cache=True, + use_cuda_graphs=use_cuda_graphs, + ) + if self.static_engine is None: + raise RuntimeError("EdgeDenoiseFlashRT requires the static denoise engine") + self.static_engine.precompute_timestep_embeds(self.denoise_dump.timesteps()) + self.static_scheduler = EdgeStaticUniPCScheduler( + self.denoise_dump.num_steps, + device=self.device, + shift=self.shift, + ) + + @property + def native_attention_available(self) -> bool: + return bool( + self.static_engine is not None + and getattr(self.static_engine.reference, "native_attention_available", False) + ) + + @property + def graph_attention_available(self) -> bool: + return bool( + self.static_engine is not None + and getattr(self.static_engine.reference, "graph_attention_available", False) + ) + + @property + def native_scheduler_available(self) -> bool: + return bool(self.static_scheduler is not None and self.static_scheduler.native_available) diff --git a/flash_rt/models/cosmos3_edge/dump_replay.py b/flash_rt/models/cosmos3_edge/dump_replay.py new file mode 100644 index 00000000..2b47c454 --- /dev/null +++ b/flash_rt/models/cosmos3_edge/dump_replay.py @@ -0,0 +1,195 @@ +"""Replay utilities for Cosmos3-Edge AV denoise dumps. + +This module is a P1 scaffold: it validates the denoise boundary captured from +Cosmos Framework before the native FlashRT transformer forward is connected. +It deliberately reuses the local self-contained UniPC scheduler so the next +step can replace recorded velocities with FlashRT-computed velocities without +changing scheduler or latent splitting code. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +import torch +from safetensors.torch import load_file + +from flash_rt.models.cosmos3_video.fm_solvers_unipc import FlowUniPCMultistepScheduler + + +EDGE_VISION_SHAPE = (1, 48, 16, 30, 52) +EDGE_ACTION_MODEL_SHAPE = (60, 64) +EDGE_FLAT_DIM = 1_201_920 +EDGE_NUM_STEPS = 30 +EDGE_SHIFT = 10.0 +EDGE_NUM_TRAIN_TIMESTEPS = 1000 + + +@dataclass(frozen=True) +class EdgeLatentParts: + vision: torch.Tensor + action_model: torch.Tensor + + +@dataclass(frozen=True) +class ReplayResult: + final_flat: torch.Tensor + final_parts: EdgeLatentParts + timesteps: tuple[int, ...] + max_input_abs_diff: float + per_step_input_abs_diff: tuple[float, ...] + + +class EdgeDenoiseDump: + """Loaded P0 denoise boundary dump for one AV inverse-dynamics sample.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + self.tensors = load_file(str(self.path), device="cpu") + self._validate_required_tensors() + + def _validate_required_tensors(self) -> None: + required = {"once/final_action", "once/final_vision", "once/num_velocity_calls"} + for step in range(EDGE_NUM_STEPS): + required.add(f"steps/{step:02d}/noise_x") + required.add(f"steps/{step:02d}/velocity") + required.add(f"steps/{step:02d}/timestep") + missing = sorted(required.difference(self.tensors)) + if missing: + preview = ", ".join(missing[:8]) + raise ValueError(f"{self.path} is missing required tensors: {preview}") + + @property + def num_steps(self) -> int: + return int(self.tensors["once/num_velocity_calls"].reshape(-1)[0].item()) + + @property + def final_action(self) -> torch.Tensor: + return self.tensors["once/final_action"] + + @property + def final_vision(self) -> torch.Tensor: + return self.tensors["once/final_vision"] + + @property + def flat_dim(self) -> int: + return int(self.step_noise(0).numel()) + + @property + def vision_dim(self) -> int: + return int(torch.tensor(self.final_vision.shape).prod().item()) + + @property + def action_model_dim(self) -> int: + return self.flat_dim - self.vision_dim + + def step_noise(self, step: int) -> torch.Tensor: + return self.tensors[f"steps/{step:02d}/noise_x"] + + def step_velocity(self, step: int) -> torch.Tensor: + return self.tensors[f"steps/{step:02d}/velocity"] + + def step_timestep(self, step: int) -> torch.Tensor: + return self.tensors[f"steps/{step:02d}/timestep"] + + def timesteps(self) -> tuple[int, ...]: + return tuple(int(self.step_timestep(i).reshape(-1)[0].item()) for i in range(self.num_steps)) + + def validate_geometry(self) -> None: + if self.num_steps != EDGE_NUM_STEPS: + raise ValueError(f"expected {EDGE_NUM_STEPS} steps, got {self.num_steps}") + if tuple(self.final_vision.shape) != EDGE_VISION_SHAPE: + raise ValueError(f"unexpected final_vision shape: {tuple(self.final_vision.shape)}") + if tuple(self.final_action.shape) != (60, 9): + raise ValueError(f"unexpected final_action shape: {tuple(self.final_action.shape)}") + if self.flat_dim != EDGE_FLAT_DIM: + raise ValueError(f"unexpected flat latent dim: {self.flat_dim}") + if self.action_model_dim != EDGE_ACTION_MODEL_SHAPE[0] * EDGE_ACTION_MODEL_SHAPE[1]: + raise ValueError(f"unexpected action-model latent dim: {self.action_model_dim}") + for step in range(self.num_steps): + if tuple(self.step_noise(step).shape) != (EDGE_FLAT_DIM,): + raise ValueError(f"unexpected noise_x shape at step {step}") + if tuple(self.step_velocity(step).shape) != (EDGE_FLAT_DIM,): + raise ValueError(f"unexpected velocity shape at step {step}") + + def split_flat(self, flat: torch.Tensor) -> EdgeLatentParts: + if flat.numel() != EDGE_FLAT_DIM: + raise ValueError(f"expected flat latent dim {EDGE_FLAT_DIM}, got {flat.numel()}") + flat = flat.reshape(-1) + vision = flat[: self.vision_dim].reshape(self.final_vision.shape) + action = flat[self.vision_dim :].reshape(EDGE_ACTION_MODEL_SHAPE) + return EdgeLatentParts(vision=vision, action_model=action) + + +class EdgeDenoiseReplay: + """Replay a captured Edge denoise run with recorded velocity tensors.""" + + def __init__( + self, + dump: EdgeDenoiseDump, + *, + device: str | torch.device = "cpu", + shift: float = EDGE_SHIFT, + ): + self.dump = dump + self.device = torch.device(device) + self.shift = float(shift) + self.dump.validate_geometry() + + def _scheduler(self) -> FlowUniPCMultistepScheduler: + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=EDGE_NUM_TRAIN_TIMESTEPS, + shift=1.0, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(self.dump.num_steps, device=self.device, shift=self.shift) + expected = tuple(int(t.item()) for t in scheduler.timesteps.cpu()) + actual = self.dump.timesteps() + if expected != actual: + raise ValueError(f"dump timesteps {actual} do not match UniPC schedule {expected}") + return scheduler + + @staticmethod + def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max().item()) + + def replay(self, *, check_inputs: bool = True) -> ReplayResult: + scheduler = self._scheduler() + latent = self.dump.step_noise(0).to(device=self.device) + input_diffs: list[float] = [] + + for step, timestep in enumerate(scheduler.timesteps): + if check_inputs: + expected_noise = self.dump.step_noise(step).to(device=self.device) + input_diffs.append(self._max_abs_diff(latent, expected_noise)) + velocity = self.dump.step_velocity(step).to(device=self.device) + latent = scheduler.step( + model_output=velocity, + timestep=timestep, + sample=latent.unsqueeze(0), + return_dict=False, + )[0].squeeze(0) + + final_flat = latent.detach().cpu().contiguous() + parts = self.dump.split_flat(final_flat) + max_diff = max(input_diffs) if input_diffs else 0.0 + return ReplayResult( + final_flat=final_flat, + final_parts=parts, + timesteps=self.dump.timesteps(), + max_input_abs_diff=max_diff, + per_step_input_abs_diff=tuple(input_diffs), + ) + + +def cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: + a_f = a.reshape(-1).float() + b_f = b.reshape(-1).float() + return float(torch.nn.functional.cosine_similarity(a_f, b_f, dim=0).item()) + + +def max_abs(values: Iterable[float]) -> float: + values = tuple(values) + return max(values) if values else 0.0 diff --git a/flash_rt/models/cosmos3_edge/layer_ref.py b/flash_rt/models/cosmos3_edge/layer_ref.py new file mode 100644 index 00000000..2df9750f --- /dev/null +++ b/flash_rt/models/cosmos3_edge/layer_ref.py @@ -0,0 +1,1035 @@ +"""PyTorch reference pieces for Cosmos3-Edge layer bring-up.""" + +from __future__ import annotations + +import math +import os +from pathlib import Path + +import torch +import torch.nn.functional as F + +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import SPEC +from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights + + +EdgeUndKVCache = tuple[tuple[torch.Tensor, torch.Tensor], ...] + + +class _CudaMatmulTf32: + def __init__(self, enabled: bool): + self.enabled = enabled + self.previous_allow_tf32: bool | None = None + + def __enter__(self) -> None: + if not torch.cuda.is_available(): + return + self.previous_allow_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = self.enabled + + def __exit__(self, exc_type, exc, tb) -> None: + if self.previous_allow_tf32 is not None: + torch.backends.cuda.matmul.allow_tf32 = self.previous_allow_tf32 + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = SPEC.rms_eps) -> torch.Tensor: + dtype = x.dtype + y = x.float() + y = y * torch.rsqrt(y.pow(2).mean(dim=-1, keepdim=True) + eps) + return (y * weight.float()).to(dtype) + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rope_partial( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + rot_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rot_dim], q[..., rot_dim:] + k_rot, k_pass = k[..., :rot_dim], k[..., rot_dim:] + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + return torch.cat((q_embed, q_pass), dim=-1), torch.cat((k_embed, k_pass), dim=-1) + + +def _linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return F.linear(x, weight) + + +def _attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor: + groups = q.shape[1] // k.shape[1] + if groups != 1: + k = k.repeat_interleave(groups, dim=1) + v = v.repeat_interleave(groups, dim=1) + q_b = q.transpose(0, 1).unsqueeze(0) + k_b = k.transpose(0, 1).unsqueeze(0) + v_b = v.transpose(0, 1).unsqueeze(0) + out = F.scaled_dot_product_attention(q_b, k_b, v_b, is_causal=is_causal, scale=SPEC.head_dim ** -0.5) + return out.squeeze(0).transpose(0, 1).reshape(q.shape[0], SPEC.hidden_size) + + +def _apply_timestep_embeds_to_noisy_tokens( + packed_tokens: torch.Tensor, + packed_timestep_embeds: torch.Tensor, + noisy_frame_indexes: torch.Tensor, + token_shape: tuple[int, ...], +) -> torch.Tensor: + if noisy_frame_indexes.numel() == 0: + return packed_tokens + spatial_numel = math.prod(token_shape[1:]) + spatial_indexes = torch.arange(spatial_numel, device=packed_tokens.device) + token_indexes = (noisy_frame_indexes * spatial_numel).unsqueeze(-1).expand(-1, spatial_numel) + token_indexes = (token_indexes + spatial_indexes).flatten().to(torch.long) + token_indexes = token_indexes.unsqueeze(-1).expand(-1, packed_tokens.shape[1]) + return packed_tokens.scatter_add(dim=0, index=token_indexes, src=packed_timestep_embeds) + + +class EdgeLayer0TorchReference: + """Torch math reference for Edge transformer layer 0. + + This is intentionally not the optimized path. It exists to validate the + Edge tower mapping and formulas before replacing individual pieces with + FlashRT kernels. + """ + + def __init__( + self, + weights: EdgeTransformerWeights, + *, + device: str | torch.device = "cuda", + dtype: torch.dtype = torch.bfloat16, + ): + self.weights = weights + self.device = torch.device(device) + self.dtype = dtype + self._cache: dict[str, torch.Tensor] = {} + self._float_cache: dict[str, torch.Tensor] = {} + self._dump_cache: dict[str, torch.Tensor] = {} + + def w(self, key: str) -> torch.Tensor: + if key not in self._cache: + self._cache[key] = self.weights.load_tensor(key, device=self.device, dtype=self.dtype) + return self._cache[key] + + def wf(self, key: str) -> torch.Tensor: + if key not in self._float_cache: + self._float_cache[key] = self.weights.load_tensor(key, device=self.device, dtype=torch.float32) + return self._float_cache[key] + + def dump_tensor(self, dump: EdgeBoundaryDump, key: str, *, dtype: torch.dtype | None = None) -> torch.Tensor: + target_dtype = dtype or self.dtype + cache_key = f"{id(dump)}:{key}:{target_dtype}" + if cache_key not in self._dump_cache: + self._dump_cache[cache_key] = dump.tensors[key].to(device=self.device, dtype=target_dtype).contiguous() + return self._dump_cache[cache_key] + + def layer_w(self, suffix: str) -> torch.Tensor: + return self.w(f"layers.0.{suffix}") + + def linear(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return _linear(x, weight) + + def norm(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return rms_norm(x, weight) + + def relu2(self, x: torch.Tensor) -> torch.Tensor: + return F.relu(x).square() + + def add(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return a + b + + def concat_gen_kv( + self, + k_und_for_gen_rope: torch.Tensor, + k_gen_rope: torch.Tensor, + v_und: torch.Tensor, + v_gen: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.cat([k_und_for_gen_rope, k_gen_rope], dim=0), torch.cat([v_und, v_gen], dim=0) + + def apply_rope( + self, + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return apply_rope_partial(q, k, cos, sin) + + def norm_rope( + self, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.apply_rope( + self.norm(q, q_weight), + self.norm(k, k_weight), + cos, + sin, + ) + + def attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor: + return _attention(q, k, v, is_causal=is_causal) + + def clear_cache(self) -> None: + self._cache.clear() + self._float_cache.clear() + self._dump_cache.clear() + + def _tower_attention( + self, + causal: torch.Tensor, + full: torch.Tensor, + dump: EdgeBoundaryDump, + ) -> tuple[torch.Tensor, torch.Tensor]: + n_und = causal.shape[0] + n_gen = full.shape[0] + + und_norm = rms_norm(causal, self.layer_w("input_layernorm.weight")) + gen_norm = rms_norm(full, self.layer_w("input_layernorm_moe_gen.weight")) + + q_und = _linear(und_norm, self.layer_w("self_attn.to_q.weight")).view(n_und, SPEC.num_heads, SPEC.head_dim) + k_und = _linear(und_norm, self.layer_w("self_attn.to_k.weight")).view(n_und, SPEC.num_kv_heads, SPEC.head_dim) + v_und = _linear(und_norm, self.layer_w("self_attn.to_v.weight")).view(n_und, SPEC.num_kv_heads, SPEC.head_dim) + + q_gen = _linear(gen_norm, self.layer_w("self_attn.add_q_proj.weight")).view( + n_gen, SPEC.num_heads, SPEC.head_dim + ) + k_gen = _linear(gen_norm, self.layer_w("self_attn.add_k_proj.weight")).view( + n_gen, SPEC.num_kv_heads, SPEC.head_dim + ) + v_gen = _linear(gen_norm, self.layer_w("self_attn.add_v_proj.weight")).view( + n_gen, SPEC.num_kv_heads, SPEC.head_dim + ) + + cos_c = self.dump_tensor(dump, "s00/layers/00/rope/cos/causal_seq") + sin_c = self.dump_tensor(dump, "s00/layers/00/rope/sin/causal_seq") + cos_f = self.dump_tensor(dump, "s00/layers/00/rope/cos/full_only_seq") + sin_f = self.dump_tensor(dump, "s00/layers/00/rope/sin/full_only_seq") + q_und_rope, k_und_rope = self.apply_rope(q_und, k_und, cos_c, sin_c) + q_gen_rope, k_gen_rope = self.norm_rope( + q_gen, + k_gen, + self.layer_w("self_attn.norm_added_q.weight"), + self.layer_w("self_attn.norm_added_k.weight"), + cos_f, + sin_f, + ) + + k_und_for_gen = rms_norm(k_und, self.layer_w("self_attn.k_norm_und_for_gen.weight")) + _, k_und_for_gen_rope = apply_rope_partial(q_und_rope, k_und_for_gen, cos_c, sin_c) + + und_attn = self.attention(q_und_rope, k_und_rope, v_und, is_causal=True) + k_gen_attn, v_gen_attn = self.concat_gen_kv(k_und_for_gen_rope, k_gen_rope, v_und, v_gen) + gen_attn = self.attention(q_gen_rope, k_gen_attn, v_gen_attn, is_causal=False) + + und_out = _linear(und_attn, self.layer_w("self_attn.to_out.weight")) + gen_out = _linear(gen_attn, self.layer_w("self_attn.to_add_out.weight")) + return und_out, gen_out + + def forward(self, dump: EdgeBoundaryDump) -> tuple[torch.Tensor, torch.Tensor]: + causal = dump.layer0_input_causal.to(device=self.device, dtype=self.dtype) + full = dump.layer0_input_full.to(device=self.device, dtype=self.dtype) + + und_attn, gen_attn = self._tower_attention(causal, full, dump) + residual_und = causal + und_attn + residual_gen = full + gen_attn + + und_mlp_in = rms_norm(residual_und, self.layer_w("post_attention_layernorm.weight")) + gen_mlp_in = rms_norm(residual_gen, self.layer_w("post_attention_layernorm_moe_gen.weight")) + + und_mlp = _linear( + self.relu2(_linear(und_mlp_in, self.layer_w("mlp.up_proj.weight"))), + self.layer_w("mlp.down_proj.weight"), + ) + gen_mlp = _linear( + self.relu2(_linear(gen_mlp_in, self.layer_w("mlp_moe_gen.up_proj.weight"))), + self.layer_w("mlp_moe_gen.down_proj.weight"), + ) + + return residual_und + und_mlp, residual_gen + gen_mlp + + +class EdgeTransformerTorchReference(EdgeLayer0TorchReference): + """Full 28-layer Torch reference for step-0 velocity validation.""" + + def layer_w_at(self, layer: int, suffix: str) -> torch.Tensor: + return self.w(f"layers.{layer}.{suffix}") + + def _tower_attention_at( + self, + layer: int, + causal: torch.Tensor, + full: torch.Tensor, + dump: EdgeBoundaryDump, + ) -> tuple[torch.Tensor, torch.Tensor]: + n_und = causal.shape[0] + n_gen = full.shape[0] + lw = lambda suffix: self.layer_w_at(layer, suffix) + + und_norm = rms_norm(causal, lw("input_layernorm.weight")) + gen_norm = rms_norm(full, lw("input_layernorm_moe_gen.weight")) + + q_und = self.linear(und_norm, lw("self_attn.to_q.weight")).view(n_und, SPEC.num_heads, SPEC.head_dim) + k_und = self.linear(und_norm, lw("self_attn.to_k.weight")).view(n_und, SPEC.num_kv_heads, SPEC.head_dim) + v_und = self.linear(und_norm, lw("self_attn.to_v.weight")).view(n_und, SPEC.num_kv_heads, SPEC.head_dim) + + q_gen = _linear(gen_norm, lw("self_attn.add_q_proj.weight")).view(n_gen, SPEC.num_heads, SPEC.head_dim) + k_gen = _linear(gen_norm, lw("self_attn.add_k_proj.weight")).view(n_gen, SPEC.num_kv_heads, SPEC.head_dim) + v_gen = _linear(gen_norm, lw("self_attn.add_v_proj.weight")).view(n_gen, SPEC.num_kv_heads, SPEC.head_dim) + + cos_c = dump.rope_cos_causal.to(device=self.device, dtype=self.dtype) + sin_c = dump.rope_sin_causal.to(device=self.device, dtype=self.dtype) + cos_f = dump.rope_cos_full.to(device=self.device, dtype=self.dtype) + sin_f = dump.rope_sin_full.to(device=self.device, dtype=self.dtype) + q_und_rope, k_und_rope = apply_rope_partial(q_und, k_und, cos_c, sin_c) + q_gen_rope, k_gen_rope = self.norm_rope( + q_gen, + k_gen, + lw("self_attn.norm_added_q.weight"), + lw("self_attn.norm_added_k.weight"), + cos_f, + sin_f, + ) + + k_und_for_gen = rms_norm(k_und, lw("self_attn.k_norm_und_for_gen.weight")) + _, k_und_for_gen_rope = apply_rope_partial(q_und_rope, k_und_for_gen, cos_c, sin_c) + + und_attn = _attention(q_und_rope, k_und_rope, v_und, is_causal=True) + k_gen_attn, v_gen_attn = self.concat_gen_kv(k_und_for_gen_rope, k_gen_rope, v_und, v_gen) + gen_attn = _attention(q_gen_rope, k_gen_attn, v_gen_attn, is_causal=False) + + return _linear(und_attn, lw("self_attn.to_out.weight")), _linear( + gen_attn, lw("self_attn.to_add_out.weight") + ) + + def forward_layer( + self, + layer: int, + causal: torch.Tensor, + full: torch.Tensor, + dump: EdgeBoundaryDump, + ) -> tuple[torch.Tensor, torch.Tensor]: + lw = lambda suffix: self.layer_w_at(layer, suffix) + und_attn, gen_attn = self._tower_attention_at(layer, causal, full, dump) + residual_und = causal + und_attn + residual_gen = full + gen_attn + + und_mlp_in = rms_norm(residual_und, lw("post_attention_layernorm.weight")) + gen_mlp_in = rms_norm(residual_gen, lw("post_attention_layernorm_moe_gen.weight")) + + und_mlp = _linear(self.relu2(_linear(und_mlp_in, lw("mlp.up_proj.weight"))), lw("mlp.down_proj.weight")) + gen_mlp = _linear( + self.relu2(_linear(gen_mlp_in, lw("mlp_moe_gen.up_proj.weight"))), + lw("mlp_moe_gen.down_proj.weight"), + ) + return residual_und + und_mlp, residual_gen + gen_mlp + + def forward_transformer(self, dump: EdgeBoundaryDump) -> tuple[torch.Tensor, torch.Tensor]: + causal = dump.layer0_input_causal.to(device=self.device, dtype=self.dtype) + full = dump.layer0_input_full.to(device=self.device, dtype=self.dtype) + return self.forward_transformer_from_sequences(dump, causal, full) + + def forward_transformer_from_sequences( + self, + dump: EdgeBoundaryDump, + causal: torch.Tensor, + full: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + for layer in range(SPEC.num_layers): + causal, full = self.forward_layer(layer, causal, full, dump) + self.clear_cache() + causal = rms_norm(causal, self.w("norm.weight")) + full = rms_norm(full, self.w("norm_moe_gen.weight")) + return causal, full + + def _causal_layer_with_gen_kv( + self, + layer: int, + causal: torch.Tensor, + dump: EdgeBoundaryDump, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + n_und = causal.shape[0] + lw = lambda suffix: self.layer_w_at(layer, suffix) + und_norm = self.norm(causal, lw("input_layernorm.weight")) + + q_und = self.linear(und_norm, lw("self_attn.to_q.weight")).view(n_und, SPEC.num_heads, SPEC.head_dim) + k_und = self.linear(und_norm, lw("self_attn.to_k.weight")).view(n_und, SPEC.num_kv_heads, SPEC.head_dim) + v_und = self.linear(und_norm, lw("self_attn.to_v.weight")).view(n_und, SPEC.num_kv_heads, SPEC.head_dim) + + cos_c = self.dump_tensor(dump, "s00/layers/00/rope/cos/causal_seq") + sin_c = self.dump_tensor(dump, "s00/layers/00/rope/sin/causal_seq") + q_und_rope, k_und_rope = self.apply_rope(q_und, k_und, cos_c, sin_c) + + k_und_for_gen = self.norm(k_und, lw("self_attn.k_norm_und_for_gen.weight")) + _, k_und_for_gen_rope = apply_rope_partial(k_und_for_gen, k_und_for_gen, cos_c, sin_c) + + und_attn = self.attention(q_und_rope, k_und_rope, v_und, is_causal=True) + und_out = self.linear(und_attn, lw("self_attn.to_out.weight")) + residual_und = causal + und_out + und_mlp_in = self.norm(residual_und, lw("post_attention_layernorm.weight")) + und_mlp = self.linear( + self.relu2(self.linear(und_mlp_in, lw("mlp.up_proj.weight"))), + lw("mlp.down_proj.weight"), + ) + return residual_und + und_mlp, k_und_for_gen_rope.contiguous(), v_und.contiguous() + + def precompute_und_kv_cache(self, dump: EdgeBoundaryDump) -> EdgeUndKVCache: + causal = dump.layer0_input_causal.to(device=self.device, dtype=self.dtype) + cache: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer in range(SPEC.num_layers): + causal, k_und_for_gen, v_und = self._causal_layer_with_gen_kv(layer, causal, dump) + cache.append((k_und_for_gen.clone(), v_und.clone())) + self.clear_cache() + return tuple(cache) + + def _gen_layer_with_und_cache( + self, + layer: int, + full: torch.Tensor, + dump: EdgeBoundaryDump, + k_und_for_gen_rope: torch.Tensor, + v_und: torch.Tensor, + ) -> torch.Tensor: + n_gen = full.shape[0] + lw = lambda suffix: self.layer_w_at(layer, suffix) + gen_norm = self.norm(full, lw("input_layernorm_moe_gen.weight")) + + q_gen = self.linear(gen_norm, lw("self_attn.add_q_proj.weight")).view(n_gen, SPEC.num_heads, SPEC.head_dim) + k_gen = self.linear(gen_norm, lw("self_attn.add_k_proj.weight")).view(n_gen, SPEC.num_kv_heads, SPEC.head_dim) + v_gen = self.linear(gen_norm, lw("self_attn.add_v_proj.weight")).view(n_gen, SPEC.num_kv_heads, SPEC.head_dim) + + cos_f = self.dump_tensor(dump, "s00/layers/00/rope/cos/full_only_seq") + sin_f = self.dump_tensor(dump, "s00/layers/00/rope/sin/full_only_seq") + q_gen_rope, k_gen_rope = self.norm_rope( + q_gen, + k_gen, + lw("self_attn.norm_added_q.weight"), + lw("self_attn.norm_added_k.weight"), + cos_f, + sin_f, + ) + + k_gen_attn, v_gen_attn = self.concat_gen_kv(k_und_for_gen_rope, k_gen_rope, v_und, v_gen) + gen_attn = self.attention(q_gen_rope, k_gen_attn, v_gen_attn, is_causal=False) + gen_out = self.linear(gen_attn, lw("self_attn.to_add_out.weight")) + residual_gen = self.add(full, gen_out) + gen_mlp_in = self.norm(residual_gen, lw("post_attention_layernorm_moe_gen.weight")) + gen_mlp = self.linear( + self.relu2(self.linear(gen_mlp_in, lw("mlp_moe_gen.up_proj.weight"))), + lw("mlp_moe_gen.down_proj.weight"), + ) + return self.add(residual_gen, gen_mlp) + + def forward_gen_with_und_cache( + self, + dump: EdgeBoundaryDump, + full: torch.Tensor, + und_cache: EdgeUndKVCache, + ) -> torch.Tensor: + if len(und_cache) != SPEC.num_layers: + raise ValueError(f"expected {SPEC.num_layers} cached layers, got {len(und_cache)}") + for layer, (k_und_for_gen, v_und) in enumerate(und_cache): + full = self._gen_layer_with_und_cache(layer, full, dump, k_und_for_gen, v_und) + self.clear_cache() + return self.norm(full, self.w("norm_moe_gen.weight")) + + def timestep_embed(self, timestep: torch.Tensor) -> torch.Tensor: + t = timestep.to(device=self.device, dtype=torch.float32).reshape(-1) * 0.001 + half = 128 + freqs = torch.exp(-math.log(10000) * torch.arange(0, half, dtype=torch.float32, device=self.device) / half) + args = t[:, None] * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + h = F.linear( + emb, + self.wf("time_embedder.linear_1.weight"), + self.wf("time_embedder.linear_1.bias"), + ) + h = F.silu(h) + h = F.linear( + h, + self.wf("time_embedder.linear_2.weight"), + self.wf("time_embedder.linear_2.bias"), + ) + return h.to(self.dtype) + + def patchify_vision_tokens(self, vision: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int]]: + latent = vision.to(device=self.device, dtype=self.dtype).squeeze(0) + if latent.dim() != 4: + raise ValueError(f"expected vision latent [1,C,T,H,W] or [C,T,H,W], got {tuple(vision.shape)}") + channels, latent_t, latent_h, latent_w = latent.shape + if channels != SPEC.latent_channels: + raise ValueError(f"expected {SPEC.latent_channels} latent channels, got {channels}") + + patch = 2 + h_padded = ((latent_h + patch - 1) // patch) * patch + w_padded = ((latent_w + patch - 1) // patch) * patch + if h_padded != latent_h or w_padded != latent_w: + padded = torch.zeros( + (channels, latent_t, h_padded, w_padded), + device=self.device, + dtype=self.dtype, + ) + padded[:, :, :latent_h, :latent_w] = latent + latent = padded + + h_patches = h_padded // patch + w_patches = w_padded // patch + latent = latent.reshape(channels, latent_t, h_patches, patch, w_patches, patch) + packed = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, SPEC.patch_latent_dim) + return packed, (latent_t, h_patches, w_patches) + + def encode_vision_tokens( + self, + vision: torch.Tensor, + timestep: torch.Tensor, + noisy_frame_indexes: torch.Tensor | None = None, + ) -> torch.Tensor: + packed, token_shape = self.patchify_vision_tokens(vision) + encoded = F.linear( + packed, + self.weights.load_tensor("proj_in.weight", device=self.device, dtype=self.dtype), + self.weights.load_tensor("proj_in.bias", device=self.device, dtype=self.dtype), + ) + if noisy_frame_indexes is None or noisy_frame_indexes.numel() == 0: + return encoded + noisy_frame_indexes = noisy_frame_indexes.to(device=self.device) + noisy_patches = noisy_frame_indexes.numel() * math.prod(token_shape[1:]) + timestep_embeds = self.timestep_embed(timestep.reshape(-1)[:1].expand(noisy_patches)) + return _apply_timestep_embeds_to_noisy_tokens(encoded, timestep_embeds, noisy_frame_indexes, token_shape) + + def encode_action_tokens(self, action: torch.Tensor, timestep: torch.Tensor, domain_id: int) -> torch.Tensor: + action = action.to(device=self.device, dtype=self.dtype) + fc = self.weights.load_tensor("action_proj_in.fc.weight", device=self.device, dtype=self.dtype)[domain_id] + bias = self.weights.load_tensor("action_proj_in.bias.weight", device=self.device, dtype=self.dtype)[domain_id] + modality = self.weights.load_tensor("action_modality_embed", device=self.device, dtype=self.dtype) + with _CudaMatmulTf32(False): + encoded = action @ fc.view(SPEC.action_dim, SPEC.hidden_size) + encoded = encoded + bias + modality + encoded = encoded + self.timestep_embed(timestep.reshape(-1)[:1].expand(action.shape[0])) + return encoded + + def full_sequence_for_step(self, dump: EdgeBoundaryDump, flat_noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + domain_id = int(dump.tensors["s00/vfm_in/action/domain_id/0"].item()) + action = flat_noise[-60 * SPEC.action_dim :].reshape(60, SPEC.action_dim) + base = dump.layer0_input_causal.shape[0] + full = torch.zeros_like(dump.layer0_input_full, device=self.device, dtype=self.dtype) + + vision_indexes = dump.tensors["s00/vfm_in/vision/sequence_indexes"].to(device=self.device) - base + vision_noisy = dump.tensors["s00/vfm_in/vision/noisy_frame_indexes/0"].to(device=self.device) + full[vision_indexes] = self.encode_vision_tokens(dump.vision_tokens, timestep, vision_noisy) + + action_indexes = dump.tensors["s00/vfm_in/action/sequence_indexes"].to(device=self.device) - base + full[action_indexes] = self.encode_action_tokens(action, timestep, domain_id) + return full + + def step0_action_velocity(self, dump: EdgeBoundaryDump) -> torch.Tensor: + _, full = self.forward_transformer(dump) + return self.action_velocity_from_final_full(dump, full) + + def action_velocity_from_final_full(self, dump: EdgeBoundaryDump, full: torch.Tensor) -> torch.Tensor: + base = dump.layer0_input_causal.shape[0] + action_indexes = dump.tensors["s00/vfm_in/action/sequence_indexes"].to(device=self.device) - base + action_hidden = full[action_indexes] + domain_id = int(dump.tensors["s00/vfm_in/action/domain_id/0"].item()) + raw_action_dim = int(dump.tensors["s00/vfm_in/action/raw_action_dim/0"].item()) + fc = self.weights.load_tensor("action_proj_out.fc.weight", device=self.device, dtype=self.dtype)[domain_id] + bias = self.weights.load_tensor("action_proj_out.bias.weight", device=self.device, dtype=self.dtype)[domain_id] + out = action_hidden @ fc.view(SPEC.hidden_size, SPEC.action_dim) + bias + out[:, raw_action_dim:] = 0 + return out + + def action_velocity_for_step( + self, + dump: EdgeBoundaryDump, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + causal = dump.layer0_input_causal.to(device=self.device, dtype=self.dtype) + full = self.full_sequence_for_step(dump, flat_noise, timestep) + _, full_out = self.forward_transformer_from_sequences(dump, causal, full) + return self.action_velocity_from_final_full(dump, full_out) + + def action_velocity_for_step_with_und_cache( + self, + dump: EdgeBoundaryDump, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + und_cache: EdgeUndKVCache, + ) -> torch.Tensor: + full = self.full_sequence_for_step(dump, flat_noise, timestep) + full_out = self.forward_gen_with_und_cache(dump, full, und_cache) + return self.action_velocity_from_final_full(dump, full_out) + + +class EdgeTransformerFvkLinearReference(EdgeTransformerTorchReference): + """Torch reference with BF16 linear calls routed through ``GemmRunner``. + + This is an incremental P1 bridge: it validates FlashRT's BF16 GEMM calling + convention for Edge's static cached path while leaving norm, RoPE, attention, + and relu2 in the reference implementation. + """ + + def __init__( + self, + weights: EdgeTransformerWeights, + *, + device: str | torch.device = "cuda", + dtype: torch.dtype = torch.bfloat16, + ): + super().__init__(weights, device=device, dtype=dtype) + self._linear_t_cache: dict[int, torch.Tensor] = {} + self._fmha_buffers: dict[tuple[int, int], tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = {} + self._mha_buffers: dict[tuple[int, int], tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = {} + self.gemm = None + self.ctx = None + self.fa2 = None + self._fmha_loaded: bool | None = None + self._fa4_tried = False + self._fa4_func = None + self._fa4_fwd_tried = False + self._fa4_fwd = None + self._fa4_fwd_buffers: dict[tuple[int, int], tuple[torch.Tensor, torch.Tensor]] = {} + if self.device.type == "cuda" and self.dtype == torch.bfloat16: + try: + import flash_rt.flash_rt_kernels as fvk + + self.gemm = fvk.GemmRunner() + self.ctx = fvk.FvkContext() + except Exception: + self.gemm = None + self.ctx = None + try: + from flash_rt import flash_rt_fa2 as fa2 + + self.fa2 = fa2 + except ImportError: + self.fa2 = None + + @property + def native_attention_available(self) -> bool: + return ( + self.fa2 is not None + or self._get_fa4_func() is not None + or (os.environ.get("FLASHRT_COSMOS3_EDGE_FMHA", "0") == "1" and self._ensure_fmha_strided()) + or (os.environ.get("FLASHRT_COSMOS3_EDGE_MHA", "0") == "1" and self.ctx is not None) + ) + + @property + def graph_attention_available(self) -> bool: + return self.fa2 is not None or ( + os.environ.get("FLASHRT_COSMOS3_EDGE_FA4_FWD", "0") == "1" + and self._get_fa4_fwd() is not None + ) + + def _get_fa4_func(self): + if self._fa4_tried: + return self._fa4_func + self._fa4_tried = True + if self.gemm is None or self.device.type != "cuda": + return None + try: + from flash_rt.hardware.thor import fa4_backend + + self._fa4_func = fa4_backend.fa4_func() + except Exception: + self._fa4_func = None + return self._fa4_func + + def _get_fa4_fwd(self): + if self._fa4_fwd_tried: + return self._fa4_fwd + self._fa4_fwd_tried = True + if self.gemm is None or self.device.type != "cuda": + return None + try: + from flash_rt.hardware.thor import fa4_backend + + self._fa4_fwd = fa4_backend.fa4_fwd() + except Exception: + self._fa4_fwd = None + return self._fa4_fwd + + def _ensure_fmha_strided(self) -> bool: + if self._fmha_loaded is not None: + return self._fmha_loaded + self._fmha_loaded = False + if self.gemm is None or self.device.type != "cuda": + return False + try: + import flash_rt.flash_rt_kernels as fvk + + so_path = Path(fvk.__file__).with_name("libfmha_fp16_strided.so") + if not so_path.exists() or not hasattr(fvk, "fmha_strided_full"): + return False + fvk.load_fmha_strided_library(str(so_path)) + self._fmha_loaded = True + except Exception: + self._fmha_loaded = False + return self._fmha_loaded + + def clear_cache(self) -> None: + # The static P1 engine owns this subclass and reuses the same fixed + # geometry across all denoise steps. Keep weights and transposes resident; + # graph capture and latency measurements are meaningless if each layer + # reloads safetensors from disk. + return None + + def linear(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if self.gemm is None or x.dim() != 2 or weight.dim() != 2: + return super().linear(x, weight) + x_in = x.contiguous() + cache_key = id(weight) + if cache_key not in self._linear_t_cache: + self._linear_t_cache[cache_key] = weight.t().contiguous() + weight_t = self._linear_t_cache[cache_key] + out = torch.empty(x_in.shape[0], weight.shape[0], device=self.device, dtype=self.dtype) + self.gemm.bf16_nn( + x_in.data_ptr(), + weight_t.data_ptr(), + out.data_ptr(), + x_in.shape[0], + weight.shape[0], + x_in.shape[1], + torch.cuda.current_stream().cuda_stream, + ) + return out + + def relu2(self, x: torch.Tensor) -> torch.Tensor: + if ( + self.gemm is None + or self.device.type != "cuda" + or self.dtype != torch.bfloat16 + or os.environ.get("FLASHRT_COSMOS3_EDGE_NATIVE_RELU2", "1") == "0" + ): + return super().relu2(x) + import flash_rt.flash_rt_kernels as fvk + + if not hasattr(fvk, "relu2_inplace_bf16"): + return super().relu2(x) + out = x.contiguous() + fvk.relu2_inplace_bf16( + out.data_ptr(), + out.numel(), + torch.cuda.current_stream().cuda_stream, + ) + return out + + def add(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + if ( + self.gemm is None + or self.device.type != "cuda" + or self.dtype != torch.bfloat16 + or a.shape != b.shape + or os.environ.get("FLASHRT_COSMOS3_EDGE_NATIVE_ADD", "0") != "1" + ): + return super().add(a, b) + import flash_rt.flash_rt_kernels as fvk + + if not hasattr(fvk, "cosmos3_edge_add_bf16"): + return super().add(a, b) + a_in = a.contiguous() + b_in = b.contiguous() + out = torch.empty_like(a_in) + fvk.cosmos3_edge_add_bf16( + a_in.data_ptr(), + b_in.data_ptr(), + out.data_ptr(), + out.numel(), + torch.cuda.current_stream().cuda_stream, + ) + return out.view_as(a) + + def norm(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if self.gemm is None or self.device.type != "cuda" or self.dtype != torch.bfloat16: + return super().norm(x, weight) + import flash_rt.flash_rt_kernels as fvk + + x_in = x.contiguous() + out = torch.empty_like(x_in) + dim = x_in.shape[-1] + rows = x_in.numel() // dim + fvk.rms_norm( + x_in.data_ptr(), + weight.contiguous().data_ptr(), + out.data_ptr(), + rows, + dim, + SPEC.rms_eps, + torch.cuda.current_stream().cuda_stream, + ) + return out.view_as(x) + + def apply_rope( + self, + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.gemm is None or self.device.type != "cuda" or self.dtype != torch.bfloat16: + return super().apply_rope(q, k, cos, sin) + if q.dim() != 3 or k.dim() != 3 or q.shape[1:] != (SPEC.num_heads, SPEC.head_dim): + return super().apply_rope(q, k, cos, sin) + if k.shape[1:] != (SPEC.num_kv_heads, SPEC.head_dim): + return super().apply_rope(q, k, cos, sin) + import flash_rt.flash_rt_kernels as fvk + + q_in = q.contiguous() + k_in = k.contiguous() + cos_in = cos.contiguous() + sin_in = sin.contiguous() + q_out = torch.empty_like(q_in) + k_out = torch.empty_like(k_in) + fvk.qwen36_partial_rope_qk_bf16( + q_in.data_ptr(), + k_in.data_ptr(), + cos_in.data_ptr(), + sin_in.data_ptr(), + q_out.data_ptr(), + k_out.data_ptr(), + q_in.shape[0], + SPEC.num_heads, + SPEC.num_kv_heads, + SPEC.head_dim, + SPEC.head_dim, + torch.cuda.current_stream().cuda_stream, + ) + return q_out, k_out + + def norm_rope( + self, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if ( + self.gemm is None + or self.device.type != "cuda" + or self.dtype != torch.bfloat16 + or os.environ.get("FLASHRT_COSMOS3_EDGE_FUSED_QK_NORM_ROPE", "1") == "0" + ): + return super().norm_rope(q, k, q_weight, k_weight, cos, sin) + if q.dim() != 3 or k.dim() != 3 or q.shape[1:] != (SPEC.num_heads, SPEC.head_dim): + return super().norm_rope(q, k, q_weight, k_weight, cos, sin) + if k.shape[1:] != (SPEC.num_kv_heads, SPEC.head_dim) or q.shape[0] != k.shape[0]: + return super().norm_rope(q, k, q_weight, k_weight, cos, sin) + if cos.shape != (q.shape[0], SPEC.head_dim) or sin.shape != (q.shape[0], SPEC.head_dim): + return super().norm_rope(q, k, q_weight, k_weight, cos, sin) + import flash_rt.flash_rt_kernels as fvk + + if not hasattr(fvk, "cosmos3_edge_qk_norm_rope_bf16"): + return super().norm_rope(q, k, q_weight, k_weight, cos, sin) + q_in = q.contiguous() + k_in = k.contiguous() + q_out = torch.empty_like(q_in) + k_out = torch.empty_like(k_in) + fvk.cosmos3_edge_qk_norm_rope_bf16( + q_in.data_ptr(), + k_in.data_ptr(), + q_weight.contiguous().data_ptr(), + k_weight.contiguous().data_ptr(), + cos.contiguous().data_ptr(), + sin.contiguous().data_ptr(), + q_out.data_ptr(), + k_out.data_ptr(), + q_in.shape[0], + SPEC.num_heads, + SPEC.num_kv_heads, + SPEC.head_dim, + SPEC.head_dim, + SPEC.rms_eps, + torch.cuda.current_stream().cuda_stream, + ) + return q_out, k_out + + def attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor: + if self.gemm is None or self.device.type != "cuda" or self.dtype != torch.bfloat16: + return super().attention(q, k, v, is_causal=is_causal) + if q.dim() != 3 or k.dim() != 3 or v.dim() != 3: + return super().attention(q, k, v, is_causal=is_causal) + if q.shape[1:] != (SPEC.num_heads, SPEC.head_dim): + return super().attention(q, k, v, is_causal=is_causal) + if k.shape[1:] != (SPEC.num_kv_heads, SPEC.head_dim) or v.shape[1:] != (SPEC.num_kv_heads, SPEC.head_dim): + return super().attention(q, k, v, is_causal=is_causal) + if not is_causal and os.environ.get("FLASHRT_COSMOS3_EDGE_FA4_FWD", "0") == "1": + fa4_fwd = self._get_fa4_fwd() + if fa4_fwd is not None: + return self._attention_fa4_fwd(q, k, v, fa4_fwd) + fa4 = self._get_fa4_func() + if not is_causal and fa4 is not None: + return self._attention_fa4(q, k, v, fa4) + if ( + os.environ.get("FLASHRT_COSMOS3_EDGE_FMHA", "0") == "1" + and not is_causal + and q.shape[0] >= 128 + and k.shape[0] >= 128 + and self._ensure_fmha_strided() + ): + return self._attention_fmha_strided(q, k, v) + if os.environ.get("FLASHRT_COSMOS3_EDGE_MHA", "0") == "1" and not is_causal and self.ctx is not None: + return self._attention_mha_bf16(q, k, v) + if self.fa2 is None: + return super().attention(q, k, v, is_causal=is_causal) + + q_in = q.contiguous() + k_in = k.contiguous() + v_in = v.contiguous() + out = torch.empty_like(q_in) + lse = torch.empty(1, SPEC.num_heads, max(q_in.shape[0], k_in.shape[0]), dtype=torch.float32, device=self.device) + fwd = self.fa2.fwd_bf16_causal if is_causal else self.fa2.fwd_bf16 + q_strides = (q_in.shape[0] * SPEC.num_heads * SPEC.head_dim, SPEC.num_heads * SPEC.head_dim, SPEC.head_dim) + k_strides = (k_in.shape[0] * SPEC.num_kv_heads * SPEC.head_dim, SPEC.num_kv_heads * SPEC.head_dim, SPEC.head_dim) + fwd( + Q=q_in.data_ptr(), + K=k_in.data_ptr(), + V=v_in.data_ptr(), + O=out.data_ptr(), + softmax_lse=lse.data_ptr(), + softmax_lse_accum=0, + o_accum=0, + batch=1, + seqlen_q=q_in.shape[0], + seqlen_k=k_in.shape[0], + num_heads_q=SPEC.num_heads, + num_heads_kv=SPEC.num_kv_heads, + head_dim=SPEC.head_dim, + q_strides=q_strides, + k_strides=k_strides, + v_strides=k_strides, + o_strides=q_strides, + softmax_scale=SPEC.head_dim ** -0.5, + num_sms=torch.cuda.get_device_properties(0).multi_processor_count, + stream=torch.cuda.current_stream().cuda_stream, + ) + return out.reshape(q_in.shape[0], SPEC.hidden_size) + + def _attention_fa4(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, fa4) -> torch.Tensor: + q_in = q.contiguous().unsqueeze(0) + k_in = k.contiguous().unsqueeze(0) + v_in = v.contiguous().unsqueeze(0) + out = fa4(q_in, k_in, v_in, causal=False, pack_gqa=True) + if isinstance(out, tuple): + out = out[0] + return out.reshape(q.shape[0], SPEC.hidden_size).to(dtype=self.dtype) + + def _attention_fa4_fwd(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, fa4_fwd) -> torch.Tensor: + q_in = q.contiguous().unsqueeze(0) + k_in = k.contiguous().unsqueeze(0) + v_in = v.contiguous().unsqueeze(0) + shape_key = (q_in.shape[1], k_in.shape[1]) + if shape_key not in self._fa4_fwd_buffers: + sq, _ = shape_key + self._fa4_fwd_buffers[shape_key] = ( + torch.empty_like(q_in), + torch.empty(1, SPEC.num_heads, sq, dtype=torch.float32, device=self.device), + ) + out, lse = self._fa4_fwd_buffers[shape_key] + fa4_fwd( + q_in, + k_in, + v_in, + softmax_scale=SPEC.head_dim ** -0.5, + causal=False, + pack_gqa=True, + out=out, + lse=lse, + ) + return out.reshape(q.shape[0], SPEC.hidden_size).to(dtype=self.dtype) + + def _attention_fmha_strided(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + import flash_rt.flash_rt_kernels as fvk + + q_in = q.contiguous() + k_in = k.contiguous() + v_in = v.contiguous() + shape_key = (q_in.shape[0], k_in.shape[0]) + if shape_key not in self._fmha_buffers: + sq, skv = shape_key + self._fmha_buffers[shape_key] = ( + torch.empty(1, sq, SPEC.num_heads, SPEC.head_dim, dtype=torch.float16, device=self.device), + torch.empty(1, skv, SPEC.num_kv_heads, SPEC.head_dim, dtype=torch.float16, device=self.device), + torch.empty(1, skv, SPEC.num_kv_heads, SPEC.head_dim, dtype=torch.float16, device=self.device), + torch.empty(1, sq, SPEC.num_heads, SPEC.head_dim, dtype=torch.float16, device=self.device), + torch.empty(1, sq, SPEC.num_heads, SPEC.head_dim, dtype=self.dtype, device=self.device), + ) + q_fp16, k_fp16, v_fp16, out_fp16, out_bf16 = self._fmha_buffers[shape_key] + q_fp16.copy_(q_in.view_as(q_fp16)) + k_fp16.copy_(k_in.view_as(k_fp16)) + v_fp16.copy_(v_in.view_as(v_fp16)) + fvk.fmha_strided_full( + q_fp16.data_ptr(), + k_fp16.data_ptr(), + v_fp16.data_ptr(), + out_fp16.data_ptr(), + 1, + q_in.shape[0], + k_in.shape[0], + SPEC.num_heads, + SPEC.num_kv_heads, + SPEC.head_dim, + SPEC.num_heads * SPEC.head_dim, + SPEC.num_kv_heads * SPEC.head_dim, + torch.cuda.current_stream().cuda_stream, + ) + out_bf16.copy_(out_fp16) + return out_bf16.view(q_in.shape[0], SPEC.hidden_size) + + def _attention_mha_bf16(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + import flash_rt.flash_rt_kernels as fvk + + q_in = q.contiguous() + k_source = k.contiguous() + v_source = v.contiguous() + shape_key = (q_in.shape[0], k_source.shape[0]) + if shape_key not in self._mha_buffers: + sq, skv = shape_key + skv_stride = ((skv + 7) // 8) * 8 + self._mha_buffers[shape_key] = ( + torch.empty(SPEC.num_heads * sq, skv_stride, dtype=self.dtype, device=self.device), + torch.empty(sq, SPEC.num_heads, SPEC.head_dim, dtype=self.dtype, device=self.device), + torch.empty(skv, SPEC.num_heads, SPEC.head_dim, dtype=self.dtype, device=self.device), + torch.empty(skv, SPEC.num_heads, SPEC.head_dim, dtype=self.dtype, device=self.device), + ) + logits, out, k_expanded, v_expanded = self._mha_buffers[shape_key] + groups = SPEC.num_heads // SPEC.num_kv_heads + k_expanded.view(k_source.shape[0], SPEC.num_kv_heads, groups, SPEC.head_dim).copy_( + k_source.unsqueeze(2).expand(-1, -1, groups, -1) + ) + v_expanded.view(v_source.shape[0], SPEC.num_kv_heads, groups, SPEC.head_dim).copy_( + v_source.unsqueeze(2).expand(-1, -1, groups, -1) + ) + fvk.gpu_fill_neginf_bf16(logits.data_ptr(), logits.numel(), torch.cuda.current_stream().cuda_stream) + fvk.attention_mha_bf16( + self.ctx, + q_in.data_ptr(), + k_expanded.data_ptr(), + v_expanded.data_ptr(), + logits.data_ptr(), + out.data_ptr(), + q_in.shape[0], + k_source.shape[0], + SPEC.num_heads, + SPEC.head_dim, + SPEC.head_dim ** -0.5, + logits.shape[1], + torch.cuda.current_stream().cuda_stream, + ) + return out.view(q_in.shape[0], SPEC.hidden_size) diff --git a/flash_rt/models/cosmos3_edge/pipeline_thor.py b/flash_rt/models/cosmos3_edge/pipeline_thor.py new file mode 100644 index 00000000..ab156a5a --- /dev/null +++ b/flash_rt/models/cosmos3_edge/pipeline_thor.py @@ -0,0 +1,505 @@ +#!/usr/bin/env python3 +"""Cosmos3-Edge AV inverse-dynamics denoise compute path (Thor SM110). + +Per docs/adding_new_model.md §0 rule 1 this is the model's own +``models/cosmos3_edge/pipeline_thor.py``: the quantized two-tower MoT +kernel-call sequence, the static und K/V cache, and the whole-denoise +CUDA-graph capture live here, self-contained. + +The gen tower carries static vision conditioning plus 60 action tokens; the +head is the per-domain action projection. The und (text) tower is identical +across denoise steps, so its per-layer K/V (k_norm_und_for_gen + rope applied) +is computed once and cached; each step runs only the gen tower against the +cached und K/V through FA4 attention. + +Precision (selected by the caller, not the environment): + - quant="fp8" : w8a8 FP8 E4M3 GEMMs (fp8_gemm_descale_bf16out) with + per-tensor weight scales and device-side dynamic activation + scales. Production default. + - quant="bf16" : reference-accuracy path (bf16 GemmRunner GEMMs). +bf16_projs keeps named projections ("q", "k", "v", "o", "up", "down") in bf16. +This module reads no environment variables. +""" +from __future__ import annotations + +import torch + +import flash_rt.flash_rt_kernels as fvk +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import SPEC +from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump +from flash_rt.models.cosmos3_edge.dump_replay import EDGE_ACTION_MODEL_SHAPE, EDGE_FLAT_DIM +from flash_rt.models.cosmos3_edge.static_engine import EdgeStaticBufferEngine +from flash_rt.models.cosmos3_edge.static_unipc import EdgeStaticUniPCScheduler +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights + +DEV = "cuda" +BF = torch.bfloat16 +H, KVH, D, FF, HID, NL = ( + SPEC.num_heads, + SPEC.num_kv_heads, + SPEC.head_dim, + SPEC.ffn_size, + SPEC.hidden_size, + SPEC.num_layers, +) +EPS = SPEC.rms_eps +NA, AD = EDGE_ACTION_MODEL_SHAPE +ALL_PROJS = ("q", "k", "v", "o", "up", "down") + +_GEN_PROJ_KEYS = { + "q": "self_attn.add_q_proj.weight", + "k": "self_attn.add_k_proj.weight", + "v": "self_attn.add_v_proj.weight", + "o": "self_attn.to_add_out.weight", + "up": "mlp_moe_gen.up_proj.weight", + "down": "mlp_moe_gen.down_proj.weight", +} + + +class CosmosEdgeThor: + """Static-buffer quantized gen-tower denoise engine for Cosmos3-Edge AV.""" + + def __init__( + self, + weights: EdgeTransformerWeights, + boundary_dump: EdgeBoundaryDump, + timesteps: tuple[int, ...], + *, + quant: str = "fp8", + bf16_projs: tuple[str, ...] = (), + qkv_fused: bool = False, + ffn_fp4: bool = False, + slim_last: bool = True, + shift: float, + ): + self.slim_last = bool(slim_last) + if quant not in ("fp8", "bf16"): + raise ValueError(f"unsupported quant mode: {quant}") + self.quant = quant + self.fp4 = None + self.ffn_fp4 = False + if ffn_fp4 and quant == "fp8" and not ({"up", "down"} & set(bf16_projs)): + import flash_rt.flash_rt_fp4 as fp4mod + + if not fp4mod.has_nvfp4(): + raise RuntimeError("ffn_fp4 requested but flash_rt_fp4 lacks NVFP4 support") + self.fp4 = fp4mod + self.ffn_fp4 = True + self.bf16_projs = set(ALL_PROJS) if quant == "bf16" else {p for p in bf16_projs if p} + self.num_steps = len(timesteps) + + # Static conditioning, und K/V cache, action-head constants, and the + # timestep-embed table come from the validated bring-up engine; the + # denoise hot path below never calls back into it. + base = EdgeStaticBufferEngine(weights, boundary_dump, device=DEV, dtype=BF) + base.precompute_timestep_embeds(timesteps) + fa4_getter = getattr(base.reference, "_get_fa4_fwd", None) + self._fa4_fwd = fa4_getter() if fa4_getter is not None else None + if self._fa4_fwd is None: + raise RuntimeError("CosmosEdgeThor requires the FA4 forward entry on this build") + + self.NU = int(base.und_cache[0][0].shape[0]) + self.NG = int(base.full.shape[0]) + self.NJ = self.NU + self.NG + NU, NG, NJ = self.NU, self.NG, self.NJ + + self.gemm = fvk.GemmRunner() + z = lambda *s: torch.zeros(*s, device=DEV, dtype=BF) + + # --- weights: quantized gen-tower projections + bf16 norms --- + def qf8(w_nk: torch.Tensor): # per-tensor FP8 E4M3; GEMM wants B as [K,N] + w = w_nk.t().contiguous() + s = max(w.float().abs().max().item() / 448.0, 1e-12) + f8 = (w.float() / s).clamp(-448, 448).to(torch.float8_e4m3fn).contiguous() + return f8, torch.tensor([s], dtype=torch.float32, device=DEV) + + # QKV fused wide GEMM measured NEGATIVE on Thor (P50 6.056s vs 5.883s + # separate: cuBLASLt N=4096 tactic worse than 3×separate + strided V + # copy). Keep the path for re-evaluation; off by default. + self.qkv_fused = bool(qkv_fused) and not ({"q", "k", "v"} & self.bf16_projs) and quant == "fp8" + self.Wt = {} + self.Wf8 = {} + self.Wds = {} + self.Wf4 = {} + self.Wn = {} + for li in range(NL): + loaded = {} + for nm, key in _GEN_PROJ_KEYS.items(): + loaded[nm] = weights.load_tensor(f"layers.{li}.{key}", device=DEV, dtype=BF) + if self.qkv_fused: + merged = torch.cat([loaded.pop("q"), loaded.pop("k"), loaded.pop("v")], dim=0) + self.Wf8[(li, "qkv")], self.Wds[(li, "qkv")] = qf8(merged) + del merged + if self.ffn_fp4: + for nm in ("up", "down"): + w16 = loaded.pop(nm).to(torch.float16).contiguous() + n_out, k_in = w16.shape + packed = torch.empty(n_out, k_in // 2, dtype=torch.uint8, device=DEV) + sfb = torch.empty( + self.fp4.sfa_size_bytes(n_out, k_in, True), dtype=torch.uint8, device=DEV) + rc = self.fp4.quantize_fp4_dynamic_sfa_fp16( + w16.data_ptr(), packed.data_ptr(), sfb.data_ptr(), n_out, k_in, True, 0) + if rc != 0: + raise RuntimeError(f"fp4 weight quant failed rc={rc} layer {li} {nm}") + self.Wf4[(li, nm)] = (packed, sfb) + del w16 + for nm, w in loaded.items(): + if nm in self.bf16_projs: + self.Wt[(li, nm)] = w.t().contiguous() + else: + self.Wf8[(li, nm)], self.Wds[(li, nm)] = qf8(w) + for nm, key in ( + ("in_ln", "input_layernorm_moe_gen.weight"), + ("post_ln", "post_attention_layernorm_moe_gen.weight"), + ("q_norm", "self_attn.norm_added_q.weight"), + ("k_norm", "self_attn.norm_added_k.weight"), + ): + self.Wn[(li, nm)] = weights.load_tensor(f"layers.{li}.{key}", device=DEV, dtype=BF) + self.norm_g = weights.load_tensor("norm_moe_gen.weight", device=DEV, dtype=BF) + + # --- action head constants (already domain-selected by the base engine) --- + self.action_in_w = base.action_in_w + self.action_static_bias = base.action_static_bias + self.action_out_w = base.action_out_w + self.action_out_bias = base.action_out_bias + self.raw_action_dim = base.raw_action_dim + self.action_indexes = base.action_indexes.to(torch.int64).contiguous() + self.t_emb = base.timestep_embed_cache + if self.t_emb is None or self.t_emb.shape[0] != self.num_steps: + raise RuntimeError("timestep embed table missing or mismatched") + + # --- static activation buffers --- + self.gen_init = base.full.clone() # vision conditioning; action rows zero + self.Hb = z(NG, HID) + self.nrm = z(NG, HID) + self.nrm2 = z(NG, HID) + self.Qt = z(NG, H * D) + self.Kt = z(NG, KVH * D) + self.QKVt = z(NG, (H + 2 * KVH) * D) if self.qkv_fused else None + self.Qb = z(NG, H, D) + self.attn = z(1, NG, H, D) + self.lse = torch.empty(1, H, NG, dtype=torch.float32, device=DEV) + self.ob = z(NG, HID) + self.up = z(NG, FF) + self.dn = z(NG, HID) + self.action_input = z(NA, AD) + self.action_encoded = z(NA, HID) + self.action_hidden = z(NA, HID) + self.action_norm = z(NA, HID) + self.action_out = z(NA, AD) + # Slim last layer: only the NA action rows feed the head, so the final + # layer's Q/attention/o/FFN run at M=NA (K/V still need all rows). + self.Qs = z(NA, H, D) + self.attn_s = z(1, NA, H, D) + self.lse_s = torch.empty(1, H, NA, dtype=torch.float32, device=DEV) + self.velocity = z(EDGE_FLAT_DIM) + self.latent = torch.zeros(EDGE_FLAT_DIM, device=DEV, dtype=torch.float32) + if self.ffn_fp4: + self.a4h = torch.empty(NG, HID // 2, dtype=torch.uint8, device=DEV) + self.sfa_h = torch.empty(self.fp4.sfa_size_bytes(NG, HID, False), dtype=torch.uint8, device=DEV) + self.up16 = torch.empty(NG, FF, dtype=torch.float16, device=DEV) + self.a4f = torch.empty(NG, FF // 2, dtype=torch.uint8, device=DEV) + self.sfa_f = torch.empty(self.fp4.sfa_size_bytes(NG, FF, False), dtype=torch.uint8, device=DEV) + self.dn16 = torch.empty(NG, HID, dtype=torch.float16, device=DEV) + # Fastest verified GEMM variant at the Edge FFN shapes; -1 = default API. + self.fp4_variant = 3 + probe = torch.zeros(2, HID // 2, dtype=torch.uint8, device=DEV) + probe_sfa = torch.zeros(self.fp4.sfa_size_bytes(2, HID, False), dtype=torch.uint8, device=DEV) + probe_out = torch.empty(2, FF, dtype=torch.float16, device=DEV) + rc = self.fp4.cutlass_fp4_gemm_variant( + self.fp4_variant, probe.data_ptr(), probe_sfa.data_ptr(), + self.Wf4[(0, "up")][0].data_ptr(), self.Wf4[(0, "up")][1].data_ptr(), + probe_out.data_ptr(), 2, FF, HID, 1.0, 0.0, 0) + if rc != 0: + self.fp4_variant = -1 + del probe, probe_sfa, probe_out + self.a8h = self.a8f = self.asc = None + # Per-(layer, site) static activation scales for the fused quant chain. + # Sites: 0 = qkv input, 1 = attention output, 2 = FFN input, 3 = down input. + self.site_scale = torch.full((NL, 4), 1e-12, dtype=torch.float32, device=DEV) + self.calibrated = False + if quant == "fp8": + self.a8h = torch.empty(NG, HID, dtype=torch.float8_e4m3fn, device=DEV) + self.a8f = torch.empty(NG, FF, dtype=torch.float8_e4m3fn, device=DEV) + self.asc = torch.empty(1, dtype=torch.float32, device=DEV) + + # --- per-layer joint K/V with the static und prefix pre-installed --- + self.Kj = torch.zeros(NL, 1, NJ, KVH, D, device=DEV, dtype=BF) + self.Vj = torch.zeros(NL, 1, NJ, KVH, D, device=DEV, dtype=BF) + for li, (k_und, v_und) in enumerate(base.und_cache): + self.Kj[li, 0, :NU].copy_(k_und) + self.Vj[li, 0, :NU].copy_(v_und) + + # --- static rope tables (full_only sequence) --- + t = boundary_dump.tensors + self.cos_f = t["s00/layers/00/rope/cos/full_only_seq"].to(device=DEV, dtype=BF).contiguous() + self.sin_f = t["s00/layers/00/rope/sin/full_only_seq"].to(device=DEV, dtype=BF).contiguous() + + self.unipc = EdgeStaticUniPCScheduler(self.num_steps, device=torch.device(DEV), shift=shift) + if not self.unipc.native_available: + raise RuntimeError("native UniPC step binding is required") + self.graph = None + del base + torch.cuda.synchronize() + + # ---- kernel helpers ---- + def _s(self): + return torch.cuda.current_stream().cuda_stream + + def _site_ptr(self, li: int, site: int) -> int: + return self.site_scale.data_ptr() + 4 * (li * 4 + site) + + def _proj( + self, + a_bf16: torch.Tensor, + a_f8: torch.Tensor | None, + li: int, + nm: str, + out: torch.Tensor, + n_out: int, + act_scale_ptr: int | None = None, + ): + m, k = a_bf16.shape + if (li, nm) in self.Wf8: + assert a_f8 is not None + scale_ptr = self.asc.data_ptr() if act_scale_ptr is None else act_scale_ptr + fvk.fp8_gemm_descale_bf16out( + a_f8.data_ptr(), self.Wf8[(li, nm)].data_ptr(), out.data_ptr(), + m, n_out, k, scale_ptr, self.Wds[(li, nm)].data_ptr(), self._s(), + ) + else: + self.gemm.bf16_nn(a_bf16.data_ptr(), self.Wt[(li, nm)].data_ptr(), out.data_ptr(), m, n_out, k, self._s()) + + def _fp4_gemm(self, a_packed, a_sfa, li: int, nm: str, out16, m: int, n: int, k: int, s: int): + w_packed, w_sfb = self.Wf4[(li, nm)] + if self.fp4_variant >= 0: + rc = self.fp4.cutlass_fp4_gemm_variant( + self.fp4_variant, a_packed.data_ptr(), a_sfa.data_ptr(), + w_packed.data_ptr(), w_sfb.data_ptr(), out16.data_ptr(), m, n, k, 1.0, 0.0, s) + else: + rc = self.fp4.cutlass_fp4_sq_fp16( + a_packed.data_ptr(), a_sfa.data_ptr(), + w_packed.data_ptr(), w_sfb.data_ptr(), out16.data_ptr(), m, n, k, 1.0, 0.0, s) + if rc != 0: + raise RuntimeError(f"fp4 GEMM failed rc={rc} layer {li} {nm}") + + def _ffn_fp4_block(self, li: int, x_bf16: torch.Tensor, s: int): + """FP4 W4A4 FFN: fused bf16 res+rms -> fp4, up GEMM, relu2 -> fp4, down GEMM.""" + NG = self.NG + self.fp4.cosmos3_edge_res_rms_fp4_sfa_bf16( + self.Hb.data_ptr(), x_bf16.data_ptr(), self.Wn[(li, "post_ln")].data_ptr(), + self.a4h.data_ptr(), self.sfa_h.data_ptr(), NG, HID, EPS, s) + self._fp4_gemm(self.a4h, self.sfa_h, li, "up", self.up16, NG, FF, HID, s) + self.fp4.cosmos3_edge_relu2_fp4_sfa_fp16( + self.up16.data_ptr(), self.a4f.data_ptr(), self.sfa_f.data_ptr(), NG, FF, s) + self._fp4_gemm(self.a4f, self.sfa_f, li, "down", self.dn16, NG, HID, FF, s) + self.dn.copy_(self.dn16) + + def _quant(self, a_bf16: torch.Tensor, a_f8: torch.Tensor, needed: bool, li: int, site: int): + """Dynamic quantization pass; accumulates the per-site scale ceiling.""" + if needed and self.quant == "fp8": + fvk.quantize_fp8_device(a_bf16.data_ptr(), a_f8.data_ptr(), self.asc.data_ptr(), a_bf16.numel(), self._s()) + fvk.fp8_accumulate_scale_max(self.asc.data_ptr(), self._site_ptr(li, site), self._s()) + + def _layer_needs_fp8(self, li: int, names: tuple[str, ...]) -> bool: + return any((li, nm) in self.Wf8 for nm in names) + + @property + def _fused_static(self) -> bool: + return self.quant == "fp8" and self.calibrated and not self.bf16_projs + + # ---- per-step forward: latent(f32 flat) -> velocity(bf16 flat) ---- + def forward_step(self, step: int, latent: torch.Tensor) -> torch.Tensor: + s = self._s() + NG, NJ, NU = self.NG, self.NJ, self.NU + + fvk.cosmos3_edge_copy_action_tail_f32_to_bf16( + latent.data_ptr(), self.action_input.data_ptr(), latent.numel(), NA * AD, s) + self.gemm.bf16_nn( + self.action_input.data_ptr(), self.action_in_w.data_ptr(), + self.action_encoded.data_ptr(), NA, HID, AD, s) + fvk.cosmos3_edge_add_action_bias_timestep_bf16( + self.action_encoded.data_ptr(), self.action_static_bias.data_ptr(), + self.t_emb[step : step + 1].data_ptr(), NA, HID, s) + self.Hb.copy_(self.gen_init) + fvk.cosmos3_edge_scatter_rows_bf16( + self.action_encoded.data_ptr(), self.Hb.data_ptr(), self.action_indexes.data_ptr(), NA, HID, s) + + fused = self._fused_static + if fused: + fvk.rms_norm(self.Hb.data_ptr(), self.Wn[(0, "in_ln")].data_ptr(), self.nrm.data_ptr(), NG, HID, EPS, s) + fvk.quantize_fp8_static(self.nrm.data_ptr(), self.a8h.data_ptr(), self._site_ptr(0, 0), NG * HID, s) + else: + fvk.rms_norm(self.Hb.data_ptr(), self.Wn[(0, "in_ln")].data_ptr(), self.nrm.data_ptr(), NG, HID, EPS, s) + slim_done = False + for li in range(NL): + qkv_scale = self._site_ptr(li, 0) if fused else None + if not fused: + self._quant( + self.nrm, self.a8h, + self.qkv_fused or self._layer_needs_fp8(li, ("q", "k", "v")), li, 0) + if self.qkv_fused: + wide = (H + 2 * KVH) * D + self._proj(self.nrm, self.a8h, li, "qkv", self.QKVt, wide, qkv_scale) + fvk.cosmos3_edge_qk_norm_rope_strided_bf16( + self.QKVt.data_ptr(), self.QKVt.data_ptr() + 2 * H * D, + self.Wn[(li, "q_norm")].data_ptr(), self.Wn[(li, "k_norm")].data_ptr(), + self.cos_f.data_ptr(), self.sin_f.data_ptr(), + self.Qb.data_ptr(), self.Kj[li, 0, NU:NJ].data_ptr(), + NG, H, KVH, wide, wide, EPS, s) + self.Vj[li, 0, NU:NJ].view(NG, KVH * D).copy_(self.QKVt[:, (H + KVH) * D :]) + else: + self._proj(self.nrm, self.a8h, li, "q", self.Qt, H * D, qkv_scale) + self._proj(self.nrm, self.a8h, li, "k", self.Kt, KVH * D, qkv_scale) + self._proj(self.nrm, self.a8h, li, "v", self.Vj[li, 0, NU:NJ].view(NG, KVH * D), KVH * D, qkv_scale) + fvk.cosmos3_edge_qk_norm_rope_bf16( + self.Qt.data_ptr(), self.Kt.data_ptr(), + self.Wn[(li, "q_norm")].data_ptr(), self.Wn[(li, "k_norm")].data_ptr(), + self.cos_f.data_ptr(), self.sin_f.data_ptr(), + self.Qb.data_ptr(), self.Kj[li, 0, NU:NJ].data_ptr(), + NG, H, KVH, D, D, EPS, s) + if fused and self.slim_last and li == NL - 1: + # Last layer: only the NA action rows are consumed by the head. + # K/V (computed above) still cover all rows; Q/attention/o/FFN + # shrink to M=NA. Site scales calibrated at full M are upper + # bounds for the row subset. + fvk.cosmos3_edge_gather_rows_bf16( + self.Qb.data_ptr(), self.Qs.data_ptr(), self.action_indexes.data_ptr(), NA, HID, s) + self._fa4_fwd( + self.Qs.view(1, NA, H, D), self.Kj[li], self.Vj[li], + softmax_scale=D ** -0.5, causal=False, pack_gqa=True, + out=self.attn_s, lse=self.lse_s) + attn_s2d = self.attn_s.view(NA, HID) + fvk.quantize_fp8_static( + attn_s2d.data_ptr(), self.a8h.data_ptr(), self._site_ptr(li, 1), NA * HID, s) + self._proj(attn_s2d, self.a8h, li, "o", self.ob, HID, self._site_ptr(li, 1)) + fvk.cosmos3_edge_gather_rows_bf16( + self.Hb.data_ptr(), self.action_hidden.data_ptr(), self.action_indexes.data_ptr(), NA, HID, s) + if self.ffn_fp4: + self.fp4.cosmos3_edge_res_rms_fp4_sfa_bf16( + self.action_hidden.data_ptr(), self.ob.data_ptr(), + self.Wn[(li, "post_ln")].data_ptr(), + self.a4h.data_ptr(), self.sfa_h.data_ptr(), NA, HID, EPS, s) + self._fp4_gemm(self.a4h, self.sfa_h, li, "up", self.up16, NA, FF, HID, s) + self.fp4.cosmos3_edge_relu2_fp4_sfa_fp16( + self.up16.data_ptr(), self.a4f.data_ptr(), self.sfa_f.data_ptr(), NA, FF, s) + self._fp4_gemm(self.a4f, self.sfa_f, li, "down", self.dn16, NA, HID, FF, s) + self.dn[:NA].copy_(self.dn16[:NA]) + else: + fvk.residual_add_rms_norm_fp8( + self.action_hidden.data_ptr(), self.ob.data_ptr(), + self.Wn[(li, "post_ln")].data_ptr(), + self.a8h.data_ptr(), NA, HID, EPS, self._site_ptr(li, 2), s) + self._proj(self.nrm2[:NA], self.a8h, li, "up", self.up, FF, self._site_ptr(li, 2)) + fvk.cosmos3_edge_relu2_to_fp8_static_bf16( + self.up.data_ptr(), self.a8f.data_ptr(), self._site_ptr(li, 3), NA * FF, s) + self._proj(self.up[:NA], self.a8f, li, "down", self.dn, HID, self._site_ptr(li, 3)) + fvk.residual_add(self.action_hidden.data_ptr(), self.dn.data_ptr(), NA * HID, s) + slim_done = True + break + self._fa4_fwd( + self.Qb.view(1, NG, H, D), self.Kj[li], self.Vj[li], + softmax_scale=D ** -0.5, causal=False, pack_gqa=True, + out=self.attn, lse=self.lse) + attn2d = self.attn.view(NG, HID) + if fused: + fvk.quantize_fp8_static(attn2d.data_ptr(), self.a8h.data_ptr(), self._site_ptr(li, 1), NG * HID, s) + self._proj(attn2d, self.a8h, li, "o", self.ob, HID, self._site_ptr(li, 1)) + if self.ffn_fp4: + self._ffn_fp4_block(li, self.ob, s) + else: + fvk.residual_add_rms_norm_fp8( + self.Hb.data_ptr(), self.ob.data_ptr(), self.Wn[(li, "post_ln")].data_ptr(), + self.a8h.data_ptr(), NG, HID, EPS, self._site_ptr(li, 2), s) + self._proj(self.nrm2, self.a8h, li, "up", self.up, FF, self._site_ptr(li, 2)) + fvk.cosmos3_edge_relu2_to_fp8_static_bf16( + self.up.data_ptr(), self.a8f.data_ptr(), self._site_ptr(li, 3), NG * FF, s) + self._proj(self.up, self.a8f, li, "down", self.dn, HID, self._site_ptr(li, 3)) + if li + 1 < NL: + fvk.residual_add_rms_norm_fp8( + self.Hb.data_ptr(), self.dn.data_ptr(), self.Wn[(li + 1, "in_ln")].data_ptr(), + self.a8h.data_ptr(), NG, HID, EPS, self._site_ptr(li + 1, 0), s) + else: + fvk.residual_add(self.Hb.data_ptr(), self.dn.data_ptr(), NG * HID, s) + else: + self._quant(attn2d, self.a8h, self._layer_needs_fp8(li, ("o",)), li, 1) + self._proj(attn2d, self.a8h, li, "o", self.ob, HID) + if self.ffn_fp4: + self._ffn_fp4_block(li, self.ob, s) + else: + fvk.residual_add_rms_norm( + self.Hb.data_ptr(), self.ob.data_ptr(), self.Wn[(li, "post_ln")].data_ptr(), + self.nrm2.data_ptr(), NG, HID, EPS, s) + self._quant(self.nrm2, self.a8h, self._layer_needs_fp8(li, ("up",)), li, 2) + self._proj(self.nrm2, self.a8h, li, "up", self.up, FF) + fvk.relu2_inplace_bf16(self.up.data_ptr(), NG * FF, s) + self._quant(self.up, self.a8f, self._layer_needs_fp8(li, ("down",)), li, 3) + self._proj(self.up, self.a8f, li, "down", self.dn, HID) + if li + 1 < NL: + fvk.residual_add_rms_norm( + self.Hb.data_ptr(), self.dn.data_ptr(), self.Wn[(li + 1, "in_ln")].data_ptr(), + self.nrm.data_ptr(), NG, HID, EPS, s) + else: + fvk.residual_add(self.Hb.data_ptr(), self.dn.data_ptr(), NG * HID, s) + + if not slim_done: + fvk.cosmos3_edge_gather_rows_bf16( + self.Hb.data_ptr(), self.action_hidden.data_ptr(), self.action_indexes.data_ptr(), NA, HID, s) + fvk.rms_norm( + self.action_hidden.data_ptr(), self.norm_g.data_ptr(), self.action_norm.data_ptr(), NA, HID, EPS, s) + self.gemm.bf16_nn( + self.action_norm.data_ptr(), self.action_out_w.data_ptr(), self.action_out.data_ptr(), NA, AD, HID, s) + fvk.cosmos3_edge_add_bias_zero_action_tail_bf16( + self.action_out.data_ptr(), self.action_out_bias.data_ptr(), NA, AD, self.raw_action_dim, s) + fvk.cosmos3_edge_fill_flat_velocity_bf16( + self.action_out.data_ptr(), self.velocity.data_ptr(), self.velocity.numel(), NA * AD, s) + return self.velocity + + # ---- whole-denoise loop: forward + native UniPC per step ---- + def run_loop(self) -> torch.Tensor: + # UniPC state buffers are allocated once; their contents are fully + # rewritten in step order on every run (step 0 reads no history). + if self.unipc.prev_m1 is None: + self.unipc.reset(self.latent) + for step in range(self.num_steps): + velocity = self.forward_step(step, self.latent) + self.unipc.step(self.latent, velocity, step) + return self.latent + + def calibrate(self, noise_flat_f32: torch.Tensor, *, margin: float = 1.25) -> None: + """One dynamic-quant denoise pass to record per-site scale ceilings. + + The recorded ceilings (times ``margin``) become the static scales for + the fused quant chain; FP8 E4M3 saturates gracefully on the rare + excursion past the ceiling. + """ + if self.quant != "fp8" or self.bf16_projs: + return + self.calibrated = False + self.latent.copy_(noise_flat_f32.to(device=DEV, dtype=torch.float32)) + self.run_loop() + torch.cuda.synchronize() + self.site_scale.mul_(margin) + self.calibrated = True + + def capture(self, warmup_noise: torch.Tensor | None = None) -> None: + if warmup_noise is not None: + self.latent.copy_(warmup_noise.to(device=DEV, dtype=torch.float32)) + st = torch.cuda.Stream() + st.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(st): + for _ in range(2): + self.run_loop() + torch.cuda.current_stream().wait_stream(st) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.run_loop() + + def denoise(self, noise_flat_f32: torch.Tensor) -> torch.Tensor: + """Full 30-step denoise from flat noise; returns the final flat latent.""" + self.latent.copy_(noise_flat_f32.to(device=DEV, dtype=torch.float32)) + if self.graph is not None: + self.graph.replay() + else: + self.run_loop() + return self.latent diff --git a/flash_rt/models/cosmos3_edge/static_engine.py b/flash_rt/models/cosmos3_edge/static_engine.py new file mode 100644 index 00000000..62944f81 --- /dev/null +++ b/flash_rt/models/cosmos3_edge/static_engine.py @@ -0,0 +1,346 @@ +"""Static-shape Cosmos3-Edge denoise engine scaffold. + +This module is the bridge between the math reference and the native Thor +kernel engine. It owns the fixed AV inverse-dynamics geometry, static vision +conditioning, action token slots, and cached und/text K/V. The individual math +calls still delegate to the Torch reference; later P1 steps replace those calls +with fvk/FA2/FP4 kernels behind this stable boundary. +""" + +from __future__ import annotations + +import torch + +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import SPEC +from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump +from flash_rt.models.cosmos3_edge.dump_replay import EDGE_ACTION_MODEL_SHAPE, EDGE_FLAT_DIM +from flash_rt.models.cosmos3_edge.layer_ref import ( + EdgeTransformerFvkLinearReference, + EdgeTransformerTorchReference, + EdgeUndKVCache, +) +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights + + +class EdgeStaticBufferEngine: + """Static-buffer, static-und-cache Edge AV denoise engine. + + The public methods expose the contract needed by the future optimized + backend: feed a flat denoise latent plus timestep, receive a flat velocity. + The vision velocity segment is zero for the current AV inverse-dynamics + workload; only the action tail is computed. + """ + + def __init__( + self, + weights: EdgeTransformerWeights, + boundary_dump: EdgeBoundaryDump, + *, + device: str | torch.device = "cuda", + dtype: torch.dtype = torch.bfloat16, + ): + self.weights = weights + self.boundary_dump = boundary_dump + self.device = torch.device(device) + self.dtype = dtype + self.reference = EdgeTransformerFvkLinearReference(weights, device=self.device, dtype=self.dtype) + if self.reference.gemm is None: + self.reference = EdgeTransformerTorchReference(weights, device=self.device, dtype=self.dtype) + self.gemm = None + self.fill_flat_velocity = None + self.add_bias_zero_action_tail = None + self.scatter_rows = None + self.gather_rows = None + self.copy_action_tail = None + self.add_action_bias_timestep = None + if self.device.type == "cuda": + try: + import flash_rt.flash_rt_kernels as fvk + + self.gemm = fvk.GemmRunner() + self.fill_flat_velocity = getattr(fvk, "cosmos3_edge_fill_flat_velocity_bf16", None) + self.add_bias_zero_action_tail = getattr( + fvk, + "cosmos3_edge_add_bias_zero_action_tail_bf16", + None, + ) + self.scatter_rows = getattr(fvk, "cosmos3_edge_scatter_rows_bf16", None) + self.gather_rows = getattr(fvk, "cosmos3_edge_gather_rows_bf16", None) + self.copy_action_tail = getattr(fvk, "cosmos3_edge_copy_action_tail_f32_to_bf16", None) + self.add_action_bias_timestep = getattr( + fvk, + "cosmos3_edge_add_action_bias_timestep_bf16", + None, + ) + except Exception: + self.gemm = None + self.fill_flat_velocity = None + self.add_bias_zero_action_tail = None + self.scatter_rows = None + self.gather_rows = None + self.copy_action_tail = None + self.add_action_bias_timestep = None + + self.boundary_dump.validate_geometry() + self.base = int(self.boundary_dump.layer0_input_causal.shape[0]) + self.vision_indexes = ( + self.boundary_dump.tensors["s00/vfm_in/vision/sequence_indexes"].to(device=self.device) - self.base + ).contiguous() + self.action_indexes = ( + self.boundary_dump.tensors["s00/vfm_in/action/sequence_indexes"].to(device=self.device) - self.base + ).contiguous() + self.domain_id = int(self.boundary_dump.tensors["s00/vfm_in/action/domain_id/0"].item()) + self.raw_action_dim = int(self.boundary_dump.tensors["s00/vfm_in/action/raw_action_dim/0"].item()) + + self.full = torch.zeros_like(self.boundary_dump.layer0_input_full, device=self.device, dtype=self.dtype) + self.velocity = torch.zeros(EDGE_FLAT_DIM, device=self.device, dtype=self.dtype) + self.graph_flat_noise = torch.empty(EDGE_FLAT_DIM, device=self.device, dtype=torch.float32) + self.graph_timestep = torch.empty(1, 1, device=self.device, dtype=torch.int64) + self.graph = None + self.timestep_embed_cache: torch.Tensor | None = None + self.action_input = torch.empty(EDGE_ACTION_MODEL_SHAPE, device=self.device, dtype=self.dtype) + self.action_encoded = torch.empty(EDGE_ACTION_MODEL_SHAPE[0], SPEC.hidden_size, device=self.device, dtype=self.dtype) + self.action_hidden = torch.empty_like(self.action_encoded) + self.action_output = torch.empty(EDGE_ACTION_MODEL_SHAPE, device=self.device, dtype=self.dtype) + self.action_in_w = weights.load_tensor( + "action_proj_in.fc.weight", + device=self.device, + dtype=self.dtype, + )[self.domain_id].view(SPEC.action_dim, SPEC.hidden_size).contiguous() + self.action_in_bias = weights.load_tensor( + "action_proj_in.bias.weight", + device=self.device, + dtype=self.dtype, + )[self.domain_id].contiguous() + self.action_modality = weights.load_tensor( + "action_modality_embed", + device=self.device, + dtype=self.dtype, + ).contiguous() + self.action_static_bias = (self.action_in_bias + self.action_modality).contiguous() + self.action_out_w = weights.load_tensor( + "action_proj_out.fc.weight", + device=self.device, + dtype=self.dtype, + )[self.domain_id].view(SPEC.hidden_size, SPEC.action_dim).contiguous() + self.action_out_bias = weights.load_tensor( + "action_proj_out.bias.weight", + device=self.device, + dtype=self.dtype, + )[self.domain_id].contiguous() + self._install_static_vision_tokens() + self.und_cache: EdgeUndKVCache = self.reference.precompute_und_kv_cache(self.boundary_dump) + + def precompute_timestep_embeds(self, timesteps: tuple[int, ...] | list[int]) -> None: + if not timesteps: + self.timestep_embed_cache = None + return + t = torch.tensor(tuple(int(v) for v in timesteps), device=self.device, dtype=torch.int64) + self.timestep_embed_cache = self.reference.timestep_embed(t).contiguous() + + def _install_static_vision_tokens(self) -> None: + vision_noisy = self.boundary_dump.tensors["s00/vfm_in/vision/noisy_frame_indexes/0"].to(device=self.device) + vision_hidden = self.reference.encode_vision_tokens( + self.boundary_dump.vision_tokens, + torch.zeros((), device=self.device), + vision_noisy, + ) + self.full[self.vision_indexes] = vision_hidden + + def _stream(self) -> int: + return torch.cuda.current_stream().cuda_stream + + def _encode_action_tokens( + self, + action: torch.Tensor, + timestep: torch.Tensor, + *, + action_loaded: bool = False, + timestep_embed: torch.Tensor | None = None, + ) -> torch.Tensor: + if not action_loaded: + self.action_input.copy_(action.to(device=self.device, dtype=self.dtype)) + if self.gemm is None: + self.action_encoded.copy_(self.action_input @ self.action_in_w) + else: + self.gemm.bf16_nn( + self.action_input.data_ptr(), + self.action_in_w.data_ptr(), + self.action_encoded.data_ptr(), + EDGE_ACTION_MODEL_SHAPE[0], + SPEC.hidden_size, + SPEC.action_dim, + self._stream(), + ) + if self.add_action_bias_timestep is not None: + if timestep_embed is None: + timestep_embed = self.reference.timestep_embed(timestep.reshape(-1)[:1]).contiguous() + self.add_action_bias_timestep( + self.action_encoded.data_ptr(), + self.action_static_bias.data_ptr(), + timestep_embed.data_ptr(), + EDGE_ACTION_MODEL_SHAPE[0], + SPEC.hidden_size, + self._stream(), + ) + else: + self.action_encoded.add_(self.action_static_bias) + if timestep_embed is None: + timestep_embed = self.reference.timestep_embed( + timestep.reshape(-1)[:1].expand(EDGE_ACTION_MODEL_SHAPE[0]) + ) + elif timestep_embed.shape[0] == 1: + timestep_embed = timestep_embed.expand(EDGE_ACTION_MODEL_SHAPE[0], SPEC.hidden_size) + self.action_encoded.add_(timestep_embed) + return self.action_encoded + + def _decode_action_velocity(self, full_out: torch.Tensor) -> torch.Tensor: + if self.gather_rows is not None: + self.gather_rows( + full_out.data_ptr(), + self.action_hidden.data_ptr(), + self.action_indexes.data_ptr(), + EDGE_ACTION_MODEL_SHAPE[0], + SPEC.hidden_size, + self._stream(), + ) + else: + self.action_hidden.copy_(full_out[self.action_indexes]) + if self.gemm is None: + self.action_output.copy_(self.action_hidden @ self.action_out_w) + else: + self.gemm.bf16_nn( + self.action_hidden.data_ptr(), + self.action_out_w.data_ptr(), + self.action_output.data_ptr(), + EDGE_ACTION_MODEL_SHAPE[0], + SPEC.action_dim, + SPEC.hidden_size, + self._stream(), + ) + if self.add_bias_zero_action_tail is not None: + self.add_bias_zero_action_tail( + self.action_output.data_ptr(), + self.action_out_bias.data_ptr(), + EDGE_ACTION_MODEL_SHAPE[0], + EDGE_ACTION_MODEL_SHAPE[1], + self.raw_action_dim, + self._stream(), + ) + else: + self.action_output.add_(self.action_out_bias) + self.action_output[:, self.raw_action_dim :] = 0 + return self.action_output + + def full_sequence_for_step( + self, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + *, + timestep_index: int | None = None, + ) -> torch.Tensor: + action_loaded = False + if ( + self.copy_action_tail is not None + and flat_noise.device == self.device + and flat_noise.dtype == torch.float32 + ): + self.copy_action_tail( + flat_noise.data_ptr(), + self.action_input.data_ptr(), + flat_noise.numel(), + EDGE_ACTION_MODEL_SHAPE[0] * EDGE_ACTION_MODEL_SHAPE[1], + self._stream(), + ) + action_loaded = True + action = self.action_input + else: + action = flat_noise.to(device=self.device, dtype=self.dtype)[-60 * SPEC.action_dim :].reshape( + EDGE_ACTION_MODEL_SHAPE + ) + timestep_embed = None + if self.timestep_embed_cache is not None and timestep_index is not None: + timestep_embed = self.timestep_embed_cache[timestep_index : timestep_index + 1] + action_encoded = self._encode_action_tokens( + action, + timestep, + action_loaded=action_loaded, + timestep_embed=timestep_embed, + ) + if self.scatter_rows is not None: + self.scatter_rows( + action_encoded.data_ptr(), + self.full.data_ptr(), + self.action_indexes.data_ptr(), + EDGE_ACTION_MODEL_SHAPE[0], + SPEC.hidden_size, + self._stream(), + ) + else: + self.full[self.action_indexes] = action_encoded + return self.full + + def action_velocity_for_step( + self, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + *, + timestep_index: int | None = None, + ) -> torch.Tensor: + full = self.full_sequence_for_step(flat_noise, timestep, timestep_index=timestep_index) + full_out = self.reference.forward_gen_with_und_cache(self.boundary_dump, full, self.und_cache) + return self._decode_action_velocity(full_out) + + def flat_velocity_for_step( + self, + flat_noise: torch.Tensor, + timestep: torch.Tensor, + *, + timestep_index: int | None = None, + ) -> torch.Tensor: + action_velocity = self.action_velocity_for_step(flat_noise, timestep, timestep_index=timestep_index) + action_flat = action_velocity.reshape(EDGE_ACTION_MODEL_SHAPE[0] * EDGE_ACTION_MODEL_SHAPE[1]) + if self.fill_flat_velocity is not None: + self.fill_flat_velocity( + action_flat.data_ptr(), + self.velocity.data_ptr(), + self.velocity.numel(), + action_flat.numel(), + self._stream(), + ) + return self.velocity + self.velocity.zero_() + self.velocity[-action_flat.numel() :] = action_flat.to(dtype=self.velocity.dtype) + return self.velocity + + def capture_velocity_graph(self, flat_noise: torch.Tensor, timestep: torch.Tensor) -> None: + if self.device.type != "cuda": + raise RuntimeError("Cosmos3-Edge velocity graph capture requires CUDA") + if not getattr(self.reference, "graph_attention_available", False): + raise RuntimeError( + "Cosmos3-Edge velocity graph capture requires graph-safe native " + "attention. On Thor, set FLASHRT_COSMOS3_EDGE_FA4_FWD=1 to use " + "the lower-level FA4 forward entry that supports the current " + "opt-in graph replay path." + ) + self.graph_flat_noise.copy_(flat_noise.to(device=self.device, dtype=torch.float32)) + self.graph_timestep.copy_(timestep.to(device=self.device, dtype=torch.int64).reshape(1, 1)) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + self.flat_velocity_for_step(self.graph_flat_noise, self.graph_timestep) + torch.cuda.current_stream().wait_stream(stream) + + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.flat_velocity_for_step(self.graph_flat_noise, self.graph_timestep) + + def replay_velocity_graph(self, flat_noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + if self.graph is None: + self.capture_velocity_graph(flat_noise, timestep) + self.graph_flat_noise.copy_(flat_noise.to(device=self.device, dtype=torch.float32)) + self.graph_timestep.copy_(timestep.to(device=self.device, dtype=torch.int64).reshape(1, 1)) + self.graph.replay() + return self.velocity diff --git a/flash_rt/models/cosmos3_edge/static_unipc.py b/flash_rt/models/cosmos3_edge/static_unipc.py new file mode 100644 index 00000000..5dc20c07 --- /dev/null +++ b/flash_rt/models/cosmos3_edge/static_unipc.py @@ -0,0 +1,226 @@ +"""Static UniPC scheduler step for the Cosmos3-Edge 30-step AV path.""" + +from __future__ import annotations + +from dataclasses import dataclass +import os + +import torch + +from flash_rt.models.cosmos3_edge.dump_replay import EDGE_NUM_TRAIN_TIMESTEPS, EDGE_SHIFT +from flash_rt.models.cosmos3_video.fm_solvers_unipc import FlowUniPCMultistepScheduler + + +@dataclass(frozen=True) +class EdgeUniPCCoefficients: + sigma: float + corrector_order: int + predictor_order: int + c_sample: float + c_last: float + c_prev_m1: float + c_prev_m2: float + c_curr_m: float + p_sample: float + p_curr_m: float + p_prev_m1: float + + +def _lambda(sigma: torch.Tensor) -> torch.Tensor: + if float(sigma.item()) == 0.0: + return torch.tensor(float("inf"), dtype=torch.float32) + alpha = torch.tensor(1.0, dtype=torch.float32) - sigma.to(dtype=torch.float32) + return torch.log(alpha) - torch.log(sigma.to(dtype=torch.float32)) + + +def _predictor_coefficients(sigmas: torch.Tensor, step: int, order: int) -> tuple[float, float, float]: + sigma_t = sigmas[step + 1].to(dtype=torch.float32) + sigma_s0 = sigmas[step].to(dtype=torch.float32) + alpha_t = torch.tensor(1.0, dtype=torch.float32) - sigma_t + h = _lambda(sigma_t) - _lambda(sigma_s0) + hh = -h + h_phi_1 = torch.expm1(hh) + b_h = torch.expm1(hh) + sample_coeff = float((sigma_t / sigma_s0).item()) + curr_m_coeff = float((-alpha_t * h_phi_1).item()) + prev_m1_coeff = 0.0 + + if order == 2: + rk = (_lambda(sigmas[step - 1]) - _lambda(sigma_s0)) / h + coeff = -alpha_t * b_h * torch.tensor(0.5, dtype=torch.float32) / rk + prev_m1_coeff += float(coeff.item()) + curr_m_coeff += float((-coeff).item()) + elif order != 1: + raise ValueError(f"unsupported UniPC predictor order: {order}") + + return sample_coeff, curr_m_coeff, prev_m1_coeff + + +def _corrector_coefficients(sigmas: torch.Tensor, step: int, order: int) -> tuple[float, float, float, float]: + sigma_t = sigmas[step].to(dtype=torch.float32) + sigma_s0 = sigmas[step - 1].to(dtype=torch.float32) + alpha_t = torch.tensor(1.0, dtype=torch.float32) - sigma_t + h = _lambda(sigma_t) - _lambda(sigma_s0) + hh = -h + h_phi_1 = torch.expm1(hh) + b_h = torch.expm1(hh) + + last_coeff = float((sigma_t / sigma_s0).item()) + prev_m1_coeff = float((-alpha_t * h_phi_1).item()) + prev_m2_coeff = 0.0 + curr_m_coeff = 0.0 + + if order == 1: + coeff = -alpha_t * b_h * torch.tensor(0.5, dtype=torch.float32) + curr_m_coeff += float(coeff.item()) + prev_m1_coeff += float((-coeff).item()) + elif order == 2: + rk = (_lambda(sigmas[step - 2]) - _lambda(sigma_s0)) / h + rks = torch.tensor([rk.item(), 1.0], dtype=torch.float32) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + matrix_rows = [] + rhs = [] + for power in range(1, order + 1): + matrix_rows.append(torch.pow(rks, power - 1)) + rhs.append(h_phi_k * factorial_i / b_h) + factorial_i *= power + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + rhos = torch.linalg.solve(torch.stack(matrix_rows), torch.tensor(rhs, dtype=torch.float32)) + coeff = -alpha_t * b_h + prev_m2_coeff += float((coeff * rhos[0] / rk).item()) + prev_m1_coeff += float((coeff * rhos[0] * (-1.0 / rk)).item()) + curr_m_coeff += float((coeff * rhos[1]).item()) + prev_m1_coeff += float((-coeff * rhos[1]).item()) + else: + raise ValueError(f"unsupported UniPC corrector order: {order}") + + return last_coeff, prev_m1_coeff, prev_m2_coeff, curr_m_coeff + + +def _precompute_coefficients(num_steps: int, *, shift: float) -> tuple[torch.Tensor, tuple[EdgeUniPCCoefficients, ...]]: + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=EDGE_NUM_TRAIN_TIMESTEPS, + shift=1.0, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(num_steps, device="cpu", shift=shift) + sigmas = scheduler.sigmas.to(dtype=torch.float32) + coeffs: list[EdgeUniPCCoefficients] = [] + for step in range(num_steps): + if step == 0: + corrector_order = 0 + c_sample, c_last, c_prev_m1, c_prev_m2, c_curr_m = 1.0, 0.0, 0.0, 0.0, 0.0 + else: + corrector_order = 1 if step == 1 else 2 + c_last, c_prev_m1, c_prev_m2, c_curr_m = _corrector_coefficients( + sigmas, + step, + corrector_order, + ) + c_sample = 0.0 + + predictor_order = min(2, num_steps - step) + predictor_order = min(predictor_order, step + 1) + p_sample, p_curr_m, p_prev_m1 = _predictor_coefficients(sigmas, step, predictor_order) + coeffs.append( + EdgeUniPCCoefficients( + sigma=float(sigmas[step].item()), + corrector_order=corrector_order, + predictor_order=predictor_order, + c_sample=c_sample, + c_last=c_last, + c_prev_m1=c_prev_m1, + c_prev_m2=c_prev_m2, + c_curr_m=c_curr_m, + p_sample=p_sample, + p_curr_m=p_curr_m, + p_prev_m1=p_prev_m1, + ) + ) + return scheduler.timesteps, tuple(coeffs) + + +class EdgeStaticUniPCScheduler: + """One-launch native UniPC update for fixed Cosmos3-Edge dump geometry.""" + + def __init__(self, num_steps: int, *, device: torch.device, shift: float = EDGE_SHIFT): + self.device = torch.device(device) + self.timesteps, self.coefficients = _precompute_coefficients(num_steps, shift=shift) + self.timesteps = self.timesteps.to(device=self.device, dtype=torch.int64) + self.native_step = None + if self.device.type == "cuda" and os.environ.get("FLASHRT_COSMOS3_EDGE_NATIVE_UNIPC", "1") != "0": + try: + import flash_rt.flash_rt_kernels as fvk + + self.native_step = getattr(fvk, "cosmos3_edge_unipc_step_f32_bf16", None) + except Exception: + self.native_step = None + self.prev_m1: torch.Tensor | None = None + self.prev_m2: torch.Tensor | None = None + self.last_sample: torch.Tensor | None = None + self.current_m: torch.Tensor | None = None + self.current_last_sample: torch.Tensor | None = None + + @property + def native_available(self) -> bool: + return self.native_step is not None + + def reset(self, sample: torch.Tensor) -> None: + self.prev_m1 = torch.empty_like(sample) + self.prev_m2 = torch.empty_like(sample) + self.last_sample = torch.empty_like(sample) + self.current_m = torch.empty_like(sample) + self.current_last_sample = torch.empty_like(sample) + + def step(self, sample: torch.Tensor, velocity: torch.Tensor, step_index: int) -> torch.Tensor: + if self.native_step is None: + raise RuntimeError("native UniPC step binding is not available") + if sample.dtype != torch.float32 or velocity.dtype != torch.bfloat16: + raise TypeError(f"expected sample=float32 and velocity=bf16, got {sample.dtype} and {velocity.dtype}") + if sample.device.type != "cuda" or velocity.device.type != "cuda": + raise TypeError("native UniPC step requires CUDA tensors") + if self.prev_m1 is None or self.prev_m2 is None or self.last_sample is None: + self.reset(sample) + assert self.prev_m1 is not None + assert self.prev_m2 is not None + assert self.last_sample is not None + assert self.current_m is not None + assert self.current_last_sample is not None + + coeff = self.coefficients[step_index] + prev_m1_ptr = self.prev_m1.data_ptr() if coeff.corrector_order >= 1 or coeff.predictor_order >= 2 else 0 + prev_m2_ptr = self.prev_m2.data_ptr() if coeff.corrector_order >= 2 else 0 + last_ptr = self.last_sample.data_ptr() if coeff.corrector_order >= 1 else 0 + self.native_step( + sample.data_ptr(), + velocity.contiguous().data_ptr(), + prev_m1_ptr, + prev_m2_ptr, + last_ptr, + sample.data_ptr(), + self.current_m.data_ptr(), + self.current_last_sample.data_ptr(), + sample.numel(), + coeff.sigma, + coeff.corrector_order, + coeff.predictor_order, + coeff.c_sample, + coeff.c_last, + coeff.c_prev_m1, + coeff.c_prev_m2, + coeff.c_curr_m, + coeff.p_sample, + coeff.p_curr_m, + coeff.p_prev_m1, + torch.cuda.current_stream().cuda_stream, + ) + + old_prev_m2 = self.prev_m2 + old_last = self.last_sample + self.prev_m2 = self.prev_m1 + self.prev_m1 = self.current_m + self.current_m = old_prev_m2 + self.last_sample = self.current_last_sample + self.current_last_sample = old_last + return sample diff --git a/flash_rt/models/cosmos3_edge/vae_native.py b/flash_rt/models/cosmos3_edge/vae_native.py new file mode 100644 index 00000000..92981eee --- /dev/null +++ b/flash_rt/models/cosmos3_edge/vae_native.py @@ -0,0 +1,304 @@ +"""Native Wan2.2 VAE encode helpers for Cosmos3-Edge Thor. + +This module is intentionally narrow: it only fuses the ResidualBlock +``RMS_norm -> SiLU`` pair into FlashRT's existing BF16 NCDHW kernel and +leaves CausalConv3d on the official cuDNN BF16 path. That gives a low-risk +first native VAE subgraph while preserving the upstream cache contract. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import flash_rt.flash_rt_kernels as fvk + + +def _native_rms_silu_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + if x.dtype != torch.bfloat16 or x.device.type != "cuda": + raise TypeError("native Wan VAE RMS+SiLU expects a CUDA BF16 tensor") + if x.dim() != 5: + raise ValueError(f"native Wan VAE RMS+SiLU expects NCDHW, got {tuple(x.shape)}") + b, c, t, h, w = (int(dim) for dim in x.shape) + x_c = x if x.is_contiguous() else x.contiguous() + gamma_c = gamma.reshape(-1).contiguous() + out = torch.empty_like(x_c) + rc = fvk.bf16_rms_silu_ncdhw( + int(x_c.data_ptr()), + int(gamma_c.data_ptr()), + int(out.data_ptr()), + 0, + 0, + b, + c, + t, + h, + w, + 1e-12, + torch.cuda.current_stream().cuda_stream, + ) + if rc != 0: + raise RuntimeError(f"bf16_rms_silu_ncdhw failed rc={rc}") + return out + + +def _native_avgdown3d_bf16(module: Any, x: torch.Tensor) -> torch.Tensor: + if x.dtype != torch.bfloat16 or x.device.type != "cuda": + raise TypeError("native Wan VAE AvgDown3D expects a CUDA BF16 tensor") + if x.dim() != 5: + raise ValueError(f"native Wan VAE AvgDown3D expects NCDHW, got {tuple(x.shape)}") + b, c, t, h, w = (int(dim) for dim in x.shape) + factor_t = int(module.factor_t) + factor_s = int(module.factor_s) + out_c = int(module.out_channels) + group_size = int(module.group_size) + pad_t = (factor_t - (t % factor_t)) % factor_t + if h % factor_s != 0 or w % factor_s != 0: + raise ValueError("native Wan VAE AvgDown3D requires divisible spatial dimensions") + x_c = x if x.is_contiguous() else x.contiguous() + out = torch.empty( + (b, out_c, (t + pad_t) // factor_t, h // factor_s, w // factor_s), + device=x.device, + dtype=x.dtype, + ) + rc = fvk.cosmos3_edge_avgdown3d_bf16( + int(x_c.data_ptr()), + int(out.data_ptr()), + b, + c, + t, + h, + w, + out_c, + factor_t, + factor_s, + group_size, + torch.cuda.current_stream().cuda_stream, + ) + if rc is not None and rc != 0: + raise RuntimeError(f"cosmos3_edge_avgdown3d_bf16 failed rc={rc}") + return out + + +def _find_triplet( + layers: list[nn.Module], + start: int, + *, + rms_cls: type[Any], + causal_conv_cls: type[Any], +) -> tuple[int, int] | None: + if start >= len(layers) or not isinstance(layers[start], rms_cls): + return None + silu_idx = start + 1 + if silu_idx >= len(layers) or not isinstance(layers[silu_idx], nn.SiLU): + return None + conv_idx = silu_idx + 1 + while conv_idx < len(layers) and isinstance(layers[conv_idx], nn.Dropout): + conv_idx += 1 + if conv_idx >= len(layers) or not isinstance(layers[conv_idx], causal_conv_cls): + return None + return silu_idx, conv_idx + + +def install_wan_vae_encode_native_rms_silu() -> dict[str, Any]: + """Patch Wan2.2 encoder ResidualBlock.forward with native RMS+SiLU. + + Returns a small stats dict for traceability. The patch is process-local and + idempotent for the imported upstream classes. + """ + from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import ( + CausalConv3d, + RMS_norm, + ResidualBlock, + _update_cache_and_apply, + ) + + if getattr(ResidualBlock, "_flashrt_cosmos3_edge_native_rms_silu", False): + return {"patched": False, "reason": "already_patched"} + + original_forward = ResidualBlock.forward + + def patched_forward(self: Any, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]): + h = self.shortcut(x) + layers = list(self.residual) + i = 0 + while i < len(layers): + triplet = _find_triplet(layers, i, rms_cls=RMS_norm, causal_conv_cls=CausalConv3d) + if triplet is None: + layer = layers[i] + if isinstance(layer, CausalConv3d) and feat_cache is not None: + x = _update_cache_and_apply(x, layer, feat_cache, feat_idx) + else: + x = layer(x) + i += 1 + continue + + _silu_idx, conv_idx = triplet + rms = layers[i] + conv = layers[conv_idx] + try: + s = _native_rms_silu_ncdhw(x, rms.gamma) + except Exception: + s = layers[i + 1](rms(x)) + for dropout_idx in range(i + 2, conv_idx): + s = layers[dropout_idx](s) + if feat_cache is not None: + x = _update_cache_and_apply(s, conv, feat_cache, feat_idx) + else: + x = conv(s) + i = conv_idx + 1 + return x + h + + ResidualBlock._flashrt_cosmos3_edge_original_forward = original_forward + ResidualBlock.forward = patched_forward + ResidualBlock._flashrt_cosmos3_edge_native_rms_silu = True + return {"patched": True, "target": "Wan2.2 ResidualBlock RMS_norm+SiLU"} + + +def install_wan_vae_encode_native_avgdown3d() -> dict[str, Any]: + """Patch Wan2.2 encoder AvgDown3D shortcut pooling to FlashRT CUDA.""" + from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import AvgDown3D + + if getattr(AvgDown3D, "_flashrt_cosmos3_edge_native_avgdown3d", False): + return {"patched": False, "reason": "already_patched"} + + original_forward = AvgDown3D.forward + + def patched_forward(self: Any, x: torch.Tensor) -> torch.Tensor: + try: + return _native_avgdown3d_bf16(self, x) + except Exception: + return original_forward(self, x) + + AvgDown3D._flashrt_cosmos3_edge_original_forward = original_forward + AvgDown3D.forward = patched_forward + AvgDown3D._flashrt_cosmos3_edge_native_avgdown3d = True + return {"patched": True, "target": "Wan2.2 AvgDown3D BF16"} + + +def _causal_conv3d_t1_as_conv2d(conv: Any, x: torch.Tensor) -> torch.Tensor | None: + if x.dim() != 5 or int(x.shape[2]) != 1: + return None + if tuple(int(item) for item in conv.stride) != (1, 1, 1): + return None + if tuple(int(item) for item in conv.dilation) != (1, 1, 1): + return None + if int(conv.groups) != 1: + return None + weight = conv.weight + if weight.dim() != 5: + return None + + k_t = int(weight.shape[2]) + temporal_index = int(conv._padding[4]) + if temporal_index < 0 or temporal_index >= k_t: + return None + + pad_w_l, pad_w_r, pad_h_t, pad_h_b = (int(item) for item in conv._padding[:4]) + x2 = x[:, :, 0, :, :] + if pad_w_l == pad_w_r and pad_h_t == pad_h_b: + y2 = F.conv2d( + x2, + weight[:, :, temporal_index, :, :], + conv.bias, + stride=(int(conv.stride[1]), int(conv.stride[2])), + padding=(pad_h_t, pad_w_l), + dilation=(int(conv.dilation[1]), int(conv.dilation[2])), + groups=int(conv.groups), + ) + else: + y2 = F.conv2d( + F.pad(x2, (pad_w_l, pad_w_r, pad_h_t, pad_h_b)), + weight[:, :, temporal_index, :, :], + conv.bias, + stride=(int(conv.stride[1]), int(conv.stride[2])), + padding=0, + dilation=(int(conv.dilation[1]), int(conv.dilation[2])), + groups=int(conv.groups), + ) + return y2.unsqueeze(2).contiguous() + + +def install_wan_vae_encode_t1_conv2d() -> dict[str, Any]: + """Patch Wan2.2 CausalConv3d prime chunks to equivalent Conv2d. + + For a no-cache, single-frame chunk, the causal temporal pad means a + kernel-3 CausalConv3d only consumes the last temporal filter slice. This + opt-in keeps steady cached chunks on the official Conv3d path and only + rewrites the cheap but frequent prime/single-frame sites for A/B. + """ + from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import CausalConv3d + + if getattr(CausalConv3d, "_flashrt_cosmos3_edge_t1_conv2d", False): + return {"patched": False, "reason": "already_patched"} + + original_forward = CausalConv3d.forward + + def patched_forward(self: Any, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + if cache_x is None: + y = _causal_conv3d_t1_as_conv2d(self, x) + if y is not None: + return y + return original_forward(self, x, cache_x) + + CausalConv3d._flashrt_cosmos3_edge_original_forward = original_forward + CausalConv3d.forward = patched_forward + CausalConv3d._flashrt_cosmos3_edge_t1_conv2d = True + return {"patched": True, "target": "Wan2.2 CausalConv3d T=1 no-cache Conv2d"} + + +def _causal_conv3d_channels_last3d_320(conv: Any, x: torch.Tensor, cache_x: torch.Tensor | None) -> torch.Tensor | None: + if cache_x is None or x.dim() != 5 or x.dtype != torch.bfloat16 or x.device.type != "cuda": + return None + b, c, t, h, w = (int(dim) for dim in x.shape) + if (c, h, w) != (320, 120, 208) or t <= 1: + return None + weight = conv.weight + if weight.dim() != 5 or tuple(int(dim) for dim in weight.shape) != (320, 320, 3, 3, 3): + return None + if tuple(int(item) for item in conv.stride) != (1, 1, 1): + return None + if tuple(int(item) for item in conv.dilation) != (1, 1, 1): + return None + if int(conv.groups) != 1: + return None + + padding = list(int(item) for item in conv._padding) + if padding != [1, 1, 1, 1, 2, 0]: + return None + x_eff = x + if padding[4] > 0: + x_eff = torch.cat([cache_x.to(device=x.device, dtype=x.dtype), x_eff], dim=2) + padding[4] -= int(cache_x.shape[2]) + x_eff = F.pad(x_eff, padding).contiguous(memory_format=torch.channels_last_3d) + weight_cl = weight.contiguous(memory_format=torch.channels_last_3d) + return F.conv3d(x_eff, weight_cl, conv.bias).contiguous() + + +def install_wan_vae_encode_channels_last3d_conv320() -> dict[str, Any]: + """Patch steady 320-channel Wan2.2 CausalConv3d sites to channels_last_3d. + + The formal probe shows per-conv conversion is only positive on the 320x + steady cached sites. This opt-in patch leaves 160/640-channel sites on the + official path and keeps outputs contiguous for the upstream cache contract. + """ + from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import CausalConv3d + + if getattr(CausalConv3d, "_flashrt_cosmos3_edge_channels_last3d_conv320", False): + return {"patched": False, "reason": "already_patched"} + + original_forward = CausalConv3d.forward + + def patched_forward(self: Any, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + y = _causal_conv3d_channels_last3d_320(self, x, cache_x) + if y is not None: + return y + return original_forward(self, x, cache_x) + + CausalConv3d._flashrt_cosmos3_edge_channels_last3d_original_forward = original_forward + CausalConv3d.forward = patched_forward + CausalConv3d._flashrt_cosmos3_edge_channels_last3d_conv320 = True + return {"patched": True, "target": "Wan2.2 CausalConv3d 320ch channels_last_3d"} diff --git a/flash_rt/models/cosmos3_edge/weights.py b/flash_rt/models/cosmos3_edge/weights.py new file mode 100644 index 00000000..8bdda359 --- /dev/null +++ b/flash_rt/models/cosmos3_edge/weights.py @@ -0,0 +1,91 @@ +"""Lazy Cosmos3-Edge transformer weight access.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import torch +from safetensors import safe_open + +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import ( + SPEC, + iter_expected_shapes, + layer_key, + load_transformer_weight_map, +) + + +@dataclass(frozen=True) +class WeightRef: + key: str + shard: Path + shape: tuple[int, ...] + + +class EdgeTransformerWeights: + """Index-backed loader for the Edge diffusion transformer shards. + + The P1/P2 engine needs explicit, stable tensor names without loading the + entire checkpoint at construction time. This class validates the index once + and opens only the shard needed for each requested tensor. + """ + + def __init__(self, checkpoint: str | Path): + self.checkpoint = Path(checkpoint) + self.transformer_dir = self.checkpoint / "transformer" + self.weight_map = load_transformer_weight_map(self.checkpoint) + self.expected_shapes = dict(iter_expected_shapes()) + missing = sorted(set(self.expected_shapes).difference(self.weight_map)) + if missing: + raise ValueError(f"missing Edge transformer weights: {missing[:8]}") + + def ref(self, key: str) -> WeightRef: + if key not in self.expected_shapes: + raise KeyError(f"unknown Cosmos3-Edge transformer weight: {key}") + shard = self.transformer_dir / self.weight_map[key] + return WeightRef(key=key, shard=shard, shape=self.expected_shapes[key]) + + def tensor_shape(self, key: str) -> tuple[int, ...]: + ref = self.ref(key) + with safe_open(str(ref.shard), framework="pt", device="cpu") as f: + shape = tuple(f.get_slice(key).get_shape()) + if shape != ref.shape: + raise ValueError(f"{key} shape mismatch: expected {ref.shape}, got {shape}") + return shape + + def load_tensor( + self, + key: str, + *, + device: str | torch.device = "cpu", + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + ref = self.ref(key) + with safe_open(str(ref.shard), framework="pt", device="cpu") as f: + tensor = f.get_tensor(key) + if tuple(tensor.shape) != ref.shape: + raise ValueError(f"{key} shape mismatch: expected {ref.shape}, got {tuple(tensor.shape)}") + if dtype is not None: + tensor = tensor.to(dtype=dtype) + if torch.device(device).type != "cpu": + tensor = tensor.to(device=device, non_blocking=True) + return tensor.contiguous() + + def layer_tensor( + self, + layer: int, + suffix: str, + *, + device: str | torch.device = "cpu", + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + return self.load_tensor(layer_key(layer, suffix), device=device, dtype=dtype) + + def validate_indexed_shapes(self) -> None: + for key in self.expected_shapes: + self.tensor_shape(key) + + @property + def num_layers(self) -> int: + return SPEC.num_layers diff --git a/tests/test_cosmos3_edge_boundary_dump.py b/tests/test_cosmos3_edge_boundary_dump.py new file mode 100644 index 00000000..2e55b868 --- /dev/null +++ b/tests/test_cosmos3_edge_boundary_dump.py @@ -0,0 +1,220 @@ +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") + +from flash_rt.models.cosmos3_edge.boundary_dump import ( # noqa: E402 + EDGE_BOUNDARY_ACTION_TOKENS, + EDGE_BOUNDARY_GEN_TOKENS, + EDGE_BOUNDARY_UND_TOKENS, + EDGE_BOUNDARY_VISION_TOKENS, + EDGE_HEAD_DIM, + EDGE_HIDDEN_SIZE, + EdgeBoundaryDump, +) + + +def _local_boundary_dump() -> Path | None: + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + ] + return next((path for path in candidates if path.exists()), None) + + +def _local_prelayer_boundary_dump() -> Path | None: + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors" + ), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors" + ), + ] + return next((path for path in candidates if path.exists()), None) + + +def _local_slim_prepare_dump() -> Path | None: + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0/" + "official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt" + ), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt" + ), + ] + return next((path for path in candidates if path.exists()), None) + + +def _local_text_embedding_shard() -> Path | None: + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge/" + "transformer/diffusion_pytorch_model-00001-of-00002.safetensors" + ), + Path("/work/models/Cosmos3-Edge/transformer/diffusion_pytorch_model-00001-of-00002.safetensors"), + ] + return next((path for path in candidates if path.exists()), None) + + +def _local_cosmos3_edge_model_root() -> Path | None: + candidates = [ + Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), + Path("/work/models/Cosmos3-Edge"), + ] + return next((path for path in candidates if path.exists()), None) + + +def test_cosmos3_edge_boundary_dump_geometry_when_present(): + dump_path = _local_boundary_dump() + if dump_path is None: + pytest.skip("local Cosmos3-Edge boundary dump is not available") + + dump = EdgeBoundaryDump(dump_path) + dump.validate_geometry() + shapes = dump.shapes + + assert shapes.und_tokens == EDGE_BOUNDARY_UND_TOKENS + assert shapes.gen_tokens == EDGE_BOUNDARY_GEN_TOKENS + assert shapes.vision_tokens == EDGE_BOUNDARY_VISION_TOKENS + assert shapes.action_tokens == EDGE_BOUNDARY_ACTION_TOKENS + assert shapes.hidden_size == EDGE_HIDDEN_SIZE + assert shapes.head_dim == EDGE_HEAD_DIM + assert dump.rope_cos_full.shape == (EDGE_BOUNDARY_GEN_TOKENS, EDGE_HEAD_DIM) + + +def test_cosmos3_edge_boundary_dump_from_tensors_when_present(): + dump_path = _local_boundary_dump() + if dump_path is None: + pytest.skip("local Cosmos3-Edge boundary dump is not available") + + file_dump = EdgeBoundaryDump(dump_path) + memory_dump = EdgeBoundaryDump.from_tensors(dict(file_dump.tensors), path="") + memory_dump.validate_geometry() + + assert memory_dump.path == Path("") + assert memory_dump.shapes == file_dump.shapes + assert torch.equal(memory_dump.rope_cos_full, file_dump.rope_cos_full) + + +def test_cosmos3_edge_prepare_payload_derives_step0_vfm_boundary_when_present(): + boundary_path = _local_prelayer_boundary_dump() + prepare_path = _local_slim_prepare_dump() + if boundary_path is None or prepare_path is None: + pytest.skip("local Cosmos3-Edge pre-layer boundary/prepare dumps are not available") + + from safetensors.torch import load_file + + from flash_rt.models.cosmos3_edge.action_only_official import ( # noqa: E402 + _derive_step0_vfm_boundary_from_prepare_payload, + ) + + prepare = torch.load(prepare_path, map_location="cpu", weights_only=False) + derived = _derive_step0_vfm_boundary_from_prepare_payload(prepare["payload"], seed=0) + boundary = load_file(str(boundary_path), device="cpu") + + assert len(derived) == 27 + for key, tensor in derived.items(): + assert key in boundary + expected = boundary[key] + assert tuple(tensor.shape) == tuple(expected.shape), key + assert torch.equal(tensor.cpu(), expected.cpu()), key + + +def test_cosmos3_edge_prepare_payload_derives_step0_causal_seq_when_weights_present(): + boundary_path = _local_prelayer_boundary_dump() + prepare_path = _local_slim_prepare_dump() + embedding_path = _local_text_embedding_shard() + if boundary_path is None or prepare_path is None or embedding_path is None: + pytest.skip("local Cosmos3-Edge pre-layer boundary/prepare/model dumps are not available") + + from safetensors import safe_open + from safetensors.torch import load_file + + from flash_rt.models.cosmos3_edge.action_only_official import ( # noqa: E402 + _derive_step0_vfm_boundary_from_prepare_payload, + ) + + prepare = torch.load(prepare_path, map_location="cpu", weights_only=False) + with safe_open(str(embedding_path), framework="pt", device="cpu") as handle: + embedding = handle.get_tensor("embed_tokens.weight") + derived = _derive_step0_vfm_boundary_from_prepare_payload( + prepare["payload"], + seed=0, + text_embedding_weight=embedding, + ) + boundary = load_file(str(boundary_path), device="cpu") + + assert len(derived) == 28 + assert torch.equal(derived["s00/lm_in/causal_seq"].cpu(), boundary["s00/lm_in/causal_seq"].cpu()) + + +def test_cosmos3_edge_full_only_seq_derivation_probe_when_weights_present(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the full_only_seq derivation probe") + boundary_path = _local_prelayer_boundary_dump() + model_root = _local_cosmos3_edge_model_root() + if boundary_path is None or model_root is None: + pytest.skip("local Cosmos3-Edge pre-layer boundary/model root is not available") + + from safetensors.torch import load_file + + from flash_rt.models.cosmos3_edge import EdgeTransformerWeights + from flash_rt.models.cosmos3_edge.layer_ref import EdgeTransformerTorchReference + + dump = EdgeBoundaryDump(boundary_path) + boundary = load_file(str(boundary_path), device="cpu") + weights = EdgeTransformerWeights(model_root) + ref = EdgeTransformerTorchReference(weights, device="cuda", dtype=torch.bfloat16) + full_probe = ref.full_sequence_for_step( + dump, + dump.tensors["steps/00/noise_x"].to("cuda"), + dump.tensors["steps/00/timestep"].to("cuda"), + ).cpu() + + expected = boundary["s00/lm_in/full_only_seq"] + assert torch.equal(full_probe, expected) + + +def test_cosmos3_edge_slim_prepare_derives_executable_boundary_when_present(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the executable boundary derivation probe") + boundary_path = _local_prelayer_boundary_dump() + prepare_path = _local_slim_prepare_dump() + model_root = _local_cosmos3_edge_model_root() + if boundary_path is None or prepare_path is None or model_root is None: + pytest.skip("local Cosmos3-Edge pre-layer boundary/prepare/model root is not available") + + from safetensors.torch import load_file + + from flash_rt.models.cosmos3_edge.action_only_official import ( # noqa: E402 + _derive_step0_executable_boundary_from_prepare_artifact, + ) + + derived = _derive_step0_executable_boundary_from_prepare_artifact( + prepare_path, + model_root, + seed=0, + device="cuda", + ) + boundary = load_file(str(boundary_path), device="cpu") + dump = EdgeBoundaryDump.from_tensors(derived, path="") + dump.validate_geometry() + + assert len(derived) == 31 + assert torch.equal(derived["s00/layers/00/input/causal_seq"], derived["s00/lm_in/causal_seq"]) + assert torch.equal(derived["s00/layers/00/input/full_only_seq"], derived["s00/lm_in/full_only_seq"]) + assert torch.equal(derived["s00/lm_in/full_only_seq"], boundary["s00/lm_in/full_only_seq"]) diff --git a/tests/test_cosmos3_edge_dump_replay.py b/tests/test_cosmos3_edge_dump_replay.py new file mode 100644 index 00000000..9d0bc327 --- /dev/null +++ b/tests/test_cosmos3_edge_dump_replay.py @@ -0,0 +1,172 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") + +from flash_rt.models.cosmos3_edge.dump_replay import ( # noqa: E402 + EDGE_ACTION_MODEL_SHAPE, + EDGE_FLAT_DIM, + EDGE_VISION_SHAPE, + EdgeDenoiseDump, + EdgeDenoiseReplay, +) +from flash_rt.models.cosmos3_edge.static_unipc import EdgeStaticUniPCScheduler # noqa: E402 + + +def test_cosmos3_edge_split_flat_geometry(): + class _SyntheticDump: + final_vision = torch.empty(EDGE_VISION_SHAPE) + vision_dim = final_vision.numel() + + def split_flat(self, flat): + flat = flat.reshape(-1) + return ( + flat[: self.vision_dim].reshape(self.final_vision.shape), + flat[self.vision_dim :].reshape(EDGE_ACTION_MODEL_SHAPE), + ) + + dump = _SyntheticDump() + flat = torch.arange(EDGE_FLAT_DIM, dtype=torch.float32) + vision, action = dump.split_flat(flat) + + assert vision.shape == EDGE_VISION_SHAPE + assert action.shape == EDGE_ACTION_MODEL_SHAPE + assert action[0, 0].item() == dump.vision_dim + + +def test_cosmos3_edge_local_dump_replays_scheduler_when_present(): + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors" + ), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/tensors.safetensors" + ), + ] + dump_path = next((path for path in candidates if path.exists()), None) + if dump_path is None: + pytest.skip("local Cosmos3-Edge P0 dump is not available") + + dump = EdgeDenoiseDump(dump_path) + dump.validate_geometry() + + device = "cuda" if torch.cuda.is_available() else "cpu" + result = EdgeDenoiseReplay(dump, device=device).replay() + final_vision_diff = (result.final_parts.vision.float() - dump.final_vision.float()).abs().max().item() + + assert result.timesteps[0] == 999 + assert result.timesteps[-1] == 256 + assert result.final_parts.vision.shape == EDGE_VISION_SHAPE + assert result.final_parts.action_model.shape == EDGE_ACTION_MODEL_SHAPE + assert final_vision_diff < 1e-6 + assert result.max_input_abs_diff < (1e-6 if device == "cuda" else 6e-3) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="native static UniPC replay requires CUDA") +def test_cosmos3_edge_static_unipc_replays_scheduler_when_present(): + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors" + ), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/tensors.safetensors" + ), + ] + dump_path = next((path for path in candidates if path.exists()), None) + if dump_path is None: + pytest.skip("local Cosmos3-Edge P0 dump is not available") + + dump = EdgeDenoiseDump(dump_path) + reference = EdgeDenoiseReplay(dump, device="cuda").replay(check_inputs=False).final_flat.cuda() + static_scheduler = EdgeStaticUniPCScheduler(dump.num_steps, device=torch.device("cuda")) + if not static_scheduler.native_available: + pytest.skip("native UniPC step binding is not built") + + latent = dump.step_noise(0).to(device="cuda") + static_scheduler.reset(latent) + for step in range(dump.num_steps): + latent = static_scheduler.step(latent, dump.step_velocity(step).to(device="cuda"), step) + torch.cuda.synchronize() + + assert torch.allclose(latent, reference, rtol=0, atol=5e-5) + + +def test_cosmos3_edge_flashrt_engine_is_public(): + from flash_rt.models.cosmos3_edge import EdgeDenoiseFlashRT + + assert EdgeDenoiseFlashRT.__name__ == "EdgeDenoiseFlashRT" + + +def test_cosmos3_edge_prepare_slim_can_derive_condition_reference(): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _derive_condition_reference_from_prepare_payload, + _derive_initial_noise_from_prepare_payload, + _prepare_payload_with_derived_initial_noise, + _prepare_payload_with_derived_condition_reference, + _slim_prepare_payload, + ) + + import numpy as np + + sequence_plans = [SimpleNamespace(has_vision=True, has_action=True, has_sound=False)] + x0_vision = torch.arange(2 * 3, dtype=torch.float32).reshape(1, 2, 1, 1, 3) + x0_action = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.bfloat16) + gen_data_clean = SimpleNamespace( + raw_state_vision=[torch.empty(1, 3, 2, 4, 4)], + x0_tokens_vision=[x0_vision], + x0_tokens_action=[x0_action], + x0_tokens_sound=None, + raw_action_dim=[torch.tensor(2)], + num_vision_items_per_sample=None, + ) + seed = [7] + mask_vision = torch.tensor([1, 0, 1, 0, 1, 0], dtype=torch.float32).reshape(x0_vision.shape) + mask_action = torch.zeros_like(x0_action, dtype=torch.float32) + condition_mask = [torch.cat([mask_vision.reshape(-1), mask_action.reshape(-1)])] + pure_vision = torch.from_numpy( + np.random.RandomState(seed[0]).standard_normal(tuple(x0_vision.shape)).astype("float32") + ) + pure_action = torch.from_numpy( + np.random.RandomState(seed[0]).standard_normal(tuple(x0_action.shape)).astype("float32") + ).to(torch.bfloat16) + expected_noise_action = pure_action.clone() + expected_noise_action[:, 2:] = 0 + initial_noise = [ + torch.cat( + [ + (mask_vision * x0_vision + (1.0 - mask_vision) * pure_vision).reshape(-1), + expected_noise_action.reshape(-1).to(torch.float32), + ] + ) + ] + expected_action = x0_action.to(torch.float32).clone() + expected_action[:, 2:] = 0 + condition_reference = [torch.cat([x0_vision.reshape(-1), expected_action.reshape(-1)])] + payload = (sequence_plans, gen_data_clean, [], [], initial_noise, condition_reference, condition_mask) + + derived = _derive_condition_reference_from_prepare_payload(payload) + assert torch.equal(derived[0], condition_reference[0]) + derived_noise = _derive_initial_noise_from_prepare_payload(payload, seed) + assert torch.equal(derived_noise[0], initial_noise[0]) + + slim = _slim_prepare_payload( + payload, + no_raw_state_vision=True, + derive_condition_reference=True, + derive_initial_noise=True, + ) + assert slim[1].raw_state_vision is None + assert slim[4] is None + assert slim[5] is None + + restored = _prepare_payload_with_derived_initial_noise(slim, seed) + restored = _prepare_payload_with_derived_condition_reference(restored) + assert torch.equal(restored[4][0], initial_noise[0]) + assert torch.equal(restored[5][0], condition_reference[0]) diff --git a/tests/test_cosmos3_edge_kernel_bindings.py b/tests/test_cosmos3_edge_kernel_bindings.py new file mode 100644 index 00000000..ad9b1840 --- /dev/null +++ b/tests/test_cosmos3_edge_kernel_bindings.py @@ -0,0 +1,596 @@ +"""Cosmos3-Edge native kernel canaries.""" + +from __future__ import annotations + +import importlib + +import pytest +import torch +import torch.nn.functional as F + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="relu2 kernel canary requires CUDA") +@pytest.mark.parametrize("numel", [10000, 10001]) +def test_cosmos3_edge_relu2_inplace_bf16_matches_torch(numel: int): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "relu2_inplace_bf16"): + pytest.skip("relu2_inplace_bf16 binding is not built") + + torch.manual_seed(numel) + x = torch.randn(numel, device="cuda", dtype=torch.bfloat16) * 3 + expected = torch.relu(x).square() + actual = x.clone() + fvk.relu2_inplace_bf16( + actual.data_ptr(), + actual.numel(), + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="qk norm+RoPE kernel canary requires CUDA") +def test_cosmos3_edge_qk_norm_rope_bf16_matches_two_step_kernels(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_qk_norm_rope_bf16"): + pytest.skip("cosmos3_edge_qk_norm_rope_bf16 binding not built") + + rows, q_heads, k_heads, head_dim = 17, 16, 8, 128 + torch.manual_seed(1234) + q = torch.randn(rows, q_heads, head_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn(rows, k_heads, head_dim, device="cuda", dtype=torch.bfloat16) + q_w = torch.randn(head_dim, device="cuda", dtype=torch.bfloat16) + k_w = torch.randn(head_dim, device="cuda", dtype=torch.bfloat16) + cos = torch.randn(rows, head_dim, device="cuda", dtype=torch.bfloat16) + sin = torch.randn(rows, head_dim, device="cuda", dtype=torch.bfloat16) + + # Reference is the fp64 math (per-head RMSNorm -> rotate-half RoPE). The + # fused kernel keeps fp32 through the whole chain and rounds to bf16 once, + # so it must sit within bf16-output rounding of the exact math (it is + # closer to it than the two-step norm->bf16->rope kernel chain). + def _ref64(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + x64 = x.double() + rms = torch.rsqrt(x64.pow(2).mean(-1, keepdim=True) + 1e-5) + n = x64 * rms * w.double() + half = head_dim // 2 + rot = torch.cat([-n[..., half:], n[..., :half]], dim=-1) + return n * cos.double().unsqueeze(1) + rot * sin.double().unsqueeze(1) + + q_expected = _ref64(q, q_w) + k_expected = _ref64(k, k_w) + stream = torch.cuda.current_stream().cuda_stream + + q_actual = torch.empty_like(q) + k_actual = torch.empty_like(k) + fvk.cosmos3_edge_qk_norm_rope_bf16( + q.data_ptr(), + k.data_ptr(), + q_w.data_ptr(), + k_w.data_ptr(), + cos.data_ptr(), + sin.data_ptr(), + q_actual.data_ptr(), + k_actual.data_ptr(), + rows, + q_heads, + k_heads, + head_dim, + head_dim, + 1e-5, + stream, + ) + torch.cuda.synchronize() + + q_max = (q_actual.double() - q_expected).abs().max().item() + k_max = (k_actual.double() - k_expected).abs().max().item() + q_cos = F.cosine_similarity(q_actual.double().flatten(), q_expected.flatten(), dim=0).item() + k_cos = F.cosine_similarity(k_actual.double().flatten(), k_expected.flatten(), dim=0).item() + assert q_max <= 0.09 # bf16 output ulp at these randn magnitudes + assert k_max <= 0.09 + assert q_cos > 0.99999 + assert k_cos > 0.99999 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="FP8 conv3d canary requires CUDA") +def test_cosmos3_edge_fp8_conv3d_v17_binding_matches_torch_small_shape(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + missing = [ + name + for name in ( + "fp8_conv3d_v17_ndhwc_bf16out", + "bf16_ndhwc_to_ncdhw_transpose", + ) + if not hasattr(fvk, name) + ] + if missing: + pytest.skip(f"required FP8 conv3d binding(s) not built: {missing}") + + b, t_cache, t_new, h, w, ci, co = 1, 2, 3, 4, 5, 32, 8 + torch.manual_seed(1357) + cache = (torch.randn(b, t_cache, h, w, ci, device="cuda") * 0.25).to(torch.float8_e4m3fn) + new = (torch.randn(b, t_new, h, w, ci, device="cuda") * 0.25).to(torch.float8_e4m3fn) + weight = (torch.randn(co, 3, 3, 3, ci, device="cuda") * 0.08).to(torch.float8_e4m3fn) + bias = torch.randn(co, device="cuda", dtype=torch.bfloat16) * 0.05 + alpha = 0.75 + + y_ndhwc = torch.empty(b, t_new, h, w, co, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.current_stream().cuda_stream + rc = fvk.fp8_conv3d_v17_ndhwc_bf16out( + cache.data_ptr(), + new.data_ptr(), + weight.data_ptr(), + y_ndhwc.data_ptr(), + bias.data_ptr(), + b, + t_cache, + t_new, + h, + w, + ci, + co, + alpha, + stream, + ) + assert rc == 0 + actual = torch.empty(b, co, t_new, h, w, device="cuda", dtype=torch.bfloat16) + rc = fvk.bf16_ndhwc_to_ncdhw_transpose( + y_ndhwc.data_ptr(), + actual.data_ptr(), + b, + co, + t_new, + h, + w, + stream, + ) + assert rc == 0 + torch.cuda.synchronize() + + x_ref = torch.cat( + [ + cache.float().permute(0, 4, 1, 2, 3), + new.float().permute(0, 4, 1, 2, 3), + ], + dim=2, + ) + w_ref = weight.float().permute(0, 4, 1, 2, 3).contiguous() + expected = F.conv3d(x_ref, w_ref, bias=None, padding=(0, 1, 1)) + expected = (expected * alpha + bias.float().view(1, co, 1, 1, 1)).to(torch.bfloat16) + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="BF16 conv3d canary requires CUDA") +def test_cosmos3_edge_bf16_conv3d_v0_binding_matches_torch_small_shape(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + missing = [ + name + for name in ( + "bf16_conv3d_v0_ndhwc_bf16out", + "bf16_ndhwc_to_ncdhw_transpose", + ) + if not hasattr(fvk, name) + ] + if missing: + pytest.skip(f"required BF16 conv3d binding(s) not built: {missing}") + + b, t_cache, t_new, h, w, ci, co = 1, 2, 3, 4, 5, 32, 16 + torch.manual_seed(2468) + cache = (torch.randn(b, t_cache, h, w, ci, device="cuda") * 0.25).to(torch.bfloat16) + new = (torch.randn(b, t_new, h, w, ci, device="cuda") * 0.25).to(torch.bfloat16) + weight = (torch.randn(co, 3, 3, 3, ci, device="cuda") * 0.08).to(torch.bfloat16) + bias = torch.randn(co, device="cuda", dtype=torch.bfloat16) * 0.05 + alpha = 0.75 + + y_ndhwc = torch.empty(b, t_new, h, w, co, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.current_stream().cuda_stream + rc = fvk.bf16_conv3d_v0_ndhwc_bf16out( + cache.data_ptr(), + new.data_ptr(), + weight.data_ptr(), + y_ndhwc.data_ptr(), + bias.data_ptr(), + b, + t_cache, + t_new, + h, + w, + ci, + co, + alpha, + stream, + ) + assert rc == 0 + actual = torch.empty(b, co, t_new, h, w, device="cuda", dtype=torch.bfloat16) + rc = fvk.bf16_ndhwc_to_ncdhw_transpose( + y_ndhwc.data_ptr(), + actual.data_ptr(), + b, + co, + t_new, + h, + w, + stream, + ) + assert rc == 0 + torch.cuda.synchronize() + + x_ref = torch.cat( + [ + cache.permute(0, 4, 1, 2, 3), + new.permute(0, 4, 1, 2, 3), + ], + dim=2, + ) + w_ref = weight.permute(0, 4, 1, 2, 3).contiguous() + expected = F.conv3d(x_ref, w_ref, bias=None, padding=(0, 1, 1)) + expected = (expected.float() * alpha + bias.float().view(1, co, 1, 1, 1)).to(torch.bfloat16) + + diff = (actual.float() - expected.float()).abs() + assert diff.max().item() <= 0.03125 + assert F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0).item() > 0.9999 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="flat velocity fill kernel canary requires CUDA") +def test_cosmos3_edge_fill_flat_velocity_bf16_matches_torch(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_fill_flat_velocity_bf16"): + pytest.skip("cosmos3_edge_fill_flat_velocity_bf16 binding is not built") + + flat_dim = 1201920 + action_numel = 60 * 64 + torch.manual_seed(4321) + action = torch.randn(action_numel, device="cuda", dtype=torch.bfloat16) + actual = torch.empty(flat_dim, device="cuda", dtype=torch.bfloat16) + expected = torch.zeros_like(actual) + expected[-action_numel:] = action + + fvk.cosmos3_edge_fill_flat_velocity_bf16( + action.data_ptr(), + actual.data_ptr(), + actual.numel(), + action.numel(), + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="action bias tail kernel canary requires CUDA") +def test_cosmos3_edge_add_bias_zero_action_tail_bf16_matches_torch(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_add_bias_zero_action_tail_bf16"): + pytest.skip("cosmos3_edge_add_bias_zero_action_tail_bf16 binding is not built") + + rows, cols, valid_cols = 60, 64, 9 + torch.manual_seed(9876) + x = torch.randn(rows, cols, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(cols, device="cuda", dtype=torch.bfloat16) + expected = x.clone() + expected.add_(bias) + expected[:, valid_cols:] = 0 + actual = x.clone() + + fvk.cosmos3_edge_add_bias_zero_action_tail_bf16( + actual.data_ptr(), + bias.data_ptr(), + rows, + cols, + valid_cols, + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="row scatter/gather kernel canary requires CUDA") +def test_cosmos3_edge_row_scatter_gather_bf16_matches_torch(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + missing = [ + name + for name in ( + "cosmos3_edge_scatter_rows_bf16", + "cosmos3_edge_gather_rows_bf16", + ) + if not hasattr(fvk, name) + ] + if missing: + pytest.skip(f"required kernel binding(s) not built: {missing}") + + rows, hidden, full_rows = 60, 2048, 128 + torch.manual_seed(2468) + row_indices = torch.randperm(full_rows, device="cuda", dtype=torch.int64)[:rows].contiguous() + src = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16) + dst = torch.randn(full_rows, hidden, device="cuda", dtype=torch.bfloat16) + expected_dst = dst.clone() + expected_dst[row_indices] = src + actual_dst = dst.clone() + stream = torch.cuda.current_stream().cuda_stream + + fvk.cosmos3_edge_scatter_rows_bf16( + src.data_ptr(), + actual_dst.data_ptr(), + row_indices.data_ptr(), + rows, + hidden, + stream, + ) + gathered = torch.empty_like(src) + fvk.cosmos3_edge_gather_rows_bf16( + actual_dst.data_ptr(), + gathered.data_ptr(), + row_indices.data_ptr(), + rows, + hidden, + stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual_dst, expected_dst) + assert torch.equal(gathered, src) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="action tail copy kernel canary requires CUDA") +def test_cosmos3_edge_copy_action_tail_f32_to_bf16_matches_torch(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_copy_action_tail_f32_to_bf16"): + pytest.skip("cosmos3_edge_copy_action_tail_f32_to_bf16 binding is not built") + + flat_dim = 1201920 + action_numel = 60 * 64 + torch.manual_seed(1357) + flat_noise = torch.randn(flat_dim, device="cuda", dtype=torch.float32) + expected = flat_noise[-action_numel:].to(dtype=torch.bfloat16) + actual = torch.empty(action_numel, device="cuda", dtype=torch.bfloat16) + + fvk.cosmos3_edge_copy_action_tail_f32_to_bf16( + flat_noise.data_ptr(), + actual.data_ptr(), + flat_noise.numel(), + action_numel, + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="action bias+timestep kernel canary requires CUDA") +def test_cosmos3_edge_add_action_bias_timestep_bf16_matches_torch(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_add_action_bias_timestep_bf16"): + pytest.skip("cosmos3_edge_add_action_bias_timestep_bf16 binding is not built") + + rows, hidden = 60, 2048 + torch.manual_seed(8642) + x = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16) + static_bias = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + timestep = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + expected = x.clone() + expected.add_(static_bias) + expected.add_(timestep.expand(rows, hidden)) + actual = x.clone() + + fvk.cosmos3_edge_add_action_bias_timestep_bf16( + actual.data_ptr(), + static_bias.data_ptr(), + timestep.data_ptr(), + rows, + hidden, + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="BF16 add kernel canary requires CUDA") +@pytest.mark.parametrize("shape", [(60, 2048), (257,)]) +def test_cosmos3_edge_add_bf16_matches_torch(shape: tuple[int, ...]): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_add_bf16"): + pytest.skip("cosmos3_edge_add_bf16 binding is not built") + + torch.manual_seed(sum(shape)) + a = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + b = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + expected = a + b + actual = torch.empty_like(expected) + + fvk.cosmos3_edge_add_bf16( + a.data_ptr(), + b.data_ptr(), + actual.data_ptr(), + actual.numel(), + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +def _avgdown3d_reference( + x: torch.Tensor, + *, + out_channels: int, + factor_t: int, + factor_s: int, + group_size: int, +) -> torch.Tensor: + pad_t = (factor_t - x.shape[2] % factor_t) % factor_t + x_pad = torch.nn.functional.pad(x, (0, 0, 0, 0, pad_t, 0)) + b, c, t, h, w = x_pad.shape + factor = factor_t * factor_s * factor_s + return ( + x_pad.view( + b, + c, + t // factor_t, + factor_t, + h // factor_s, + factor_s, + w // factor_s, + factor_s, + ) + .permute(0, 1, 3, 5, 7, 2, 4, 6) + .contiguous() + .view(b, c * factor, t // factor_t, h // factor_s, w // factor_s) + .view(b, out_channels, group_size, t // factor_t, h // factor_s, w // factor_s) + .mean(dim=2) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="AvgDown3D kernel canary requires CUDA") +@pytest.mark.parametrize( + ("shape", "out_channels", "factor_t", "factor_s", "group_size"), + [ + ((1, 160, 5, 8, 12), 160, 1, 2, 4), + ((1, 160, 5, 8, 12), 320, 2, 2, 4), + ((1, 320, 6, 6, 8), 640, 2, 2, 4), + ((1, 640, 3, 4, 6), 640, 1, 1, 1), + ], +) +def test_cosmos3_edge_avgdown3d_bf16_matches_torch( + shape: tuple[int, ...], + out_channels: int, + factor_t: int, + factor_s: int, + group_size: int, +): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_avgdown3d_bf16"): + pytest.skip("cosmos3_edge_avgdown3d_bf16 binding is not built") + + torch.manual_seed(sum(shape) + out_channels + factor_t * 17 + factor_s) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + expected = _avgdown3d_reference( + x, + out_channels=out_channels, + factor_t=factor_t, + factor_s=factor_s, + group_size=group_size, + ) + actual = torch.empty_like(expected) + + fvk.cosmos3_edge_avgdown3d_bf16( + x.data_ptr(), + actual.data_ptr(), + shape[0], + shape[1], + shape[2], + shape[3], + shape[4], + out_channels, + factor_t, + factor_s, + group_size, + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="UniPC step kernel canary requires CUDA") +def test_cosmos3_edge_unipc_step_f32_bf16_matches_torch_formula(): + try: + fvk = importlib.import_module("flash_rt.flash_rt_kernels") + except Exception as exc: # pragma: no cover - build/env dependent + pytest.skip(f"flash_rt_kernels not importable: {exc}") + if not hasattr(fvk, "cosmos3_edge_unipc_step_f32_bf16"): + pytest.skip("cosmos3_edge_unipc_step_f32_bf16 binding is not built") + + torch.manual_seed(9753) + numel = 4099 + sample = torch.randn(numel, device="cuda", dtype=torch.float32) + velocity = torch.randn(numel, device="cuda", dtype=torch.bfloat16) + prev_m1 = torch.randn(numel, device="cuda", dtype=torch.float32) + prev_m2 = torch.randn(numel, device="cuda", dtype=torch.float32) + prev_last = torch.randn(numel, device="cuda", dtype=torch.float32) + sigma = 0.713 + coeffs = { + "c_sample": 0.0, + "c_last": 0.41, + "c_prev_m1": -0.37, + "c_prev_m2": 0.08, + "c_curr_m": 0.88, + "p_sample": 0.73, + "p_curr_m": -0.22, + "p_prev_m1": 0.19, + } + current_m = sample - sigma * velocity + corrected = ( + coeffs["c_last"] * prev_last + + coeffs["c_prev_m1"] * prev_m1 + + coeffs["c_prev_m2"] * prev_m2 + + coeffs["c_curr_m"] * current_m + ) + expected_next = coeffs["p_sample"] * corrected + coeffs["p_curr_m"] * current_m + coeffs["p_prev_m1"] * prev_m1 + + actual_next = sample.clone() + actual_m = torch.empty_like(sample) + actual_last = torch.empty_like(sample) + fvk.cosmos3_edge_unipc_step_f32_bf16( + actual_next.data_ptr(), + velocity.data_ptr(), + prev_m1.data_ptr(), + prev_m2.data_ptr(), + prev_last.data_ptr(), + actual_next.data_ptr(), + actual_m.data_ptr(), + actual_last.data_ptr(), + numel, + sigma, + 2, + 2, + coeffs["c_sample"], + coeffs["c_last"], + coeffs["c_prev_m1"], + coeffs["c_prev_m2"], + coeffs["c_curr_m"], + coeffs["p_sample"], + coeffs["p_curr_m"], + coeffs["p_prev_m1"], + torch.cuda.current_stream().cuda_stream, + ) + torch.cuda.synchronize() + + assert torch.allclose(actual_m, current_m, rtol=0, atol=1e-6) + assert torch.allclose(actual_last, corrected, rtol=0, atol=1e-6) + assert torch.allclose(actual_next, expected_next, rtol=0, atol=1e-6) diff --git a/tests/test_cosmos3_edge_layer_ref.py b/tests/test_cosmos3_edge_layer_ref.py new file mode 100644 index 00000000..7d363069 --- /dev/null +++ b/tests/test_cosmos3_edge_layer_ref.py @@ -0,0 +1,122 @@ +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") + +from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump # noqa: E402 +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import SPEC # noqa: E402 +from flash_rt.models.cosmos3_edge.layer_ref import ( # noqa: E402 + EdgeLayer0TorchReference, + EdgeTransformerFvkLinearReference, + _attention, +) +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights # noqa: E402 + + +def _local_paths() -> tuple[Path, Path] | None: + candidates = [ + ( + Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + ), + ( + Path("/work/models/Cosmos3-Edge"), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + ), + ] + return next((paths for paths in candidates if paths[0].exists() and paths[1].exists()), None) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="layer0 reference requires CUDA for this shape") +def test_cosmos3_edge_layer0_torch_reference_matches_boundary_dump(): + paths = _local_paths() + if paths is None: + pytest.skip("local Cosmos3-Edge checkpoint/boundary dump is not available") + checkpoint, dump_path = paths + + dump = EdgeBoundaryDump(dump_path) + dump.validate_geometry() + ref = EdgeLayer0TorchReference(EdgeTransformerWeights(checkpoint), device="cuda") + + with torch.no_grad(): + causal, full = ref.forward(dump) + + target_causal = dump.layer0_output_causal.to(device="cuda", dtype=torch.bfloat16) + target_full = dump.layer0_output_full.to(device="cuda", dtype=torch.bfloat16) + + causal_cos = torch.nn.functional.cosine_similarity(causal.float().flatten(), target_causal.float().flatten(), dim=0) + full_cos = torch.nn.functional.cosine_similarity(full.float().flatten(), target_full.float().flatten(), dim=0) + + assert causal_cos.item() > 0.9999 + assert full_cos.item() > 0.9999 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Thor native attention reference check requires CUDA") +def test_cosmos3_edge_native_attention_matches_reference_at_supported_shape(): + ref = EdgeTransformerFvkLinearReference(object(), device="cuda") + if ref.ctx is None: + pytest.skip("Thor attention_mha_bf16 is not available") + + generator = torch.Generator(device="cuda").manual_seed(1234) + q = torch.randn(128, SPEC.num_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + k = torch.randn(256, SPEC.num_kv_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + v = torch.randn(256, SPEC.num_kv_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + + with torch.no_grad(): + out = ref.attention(q, k, v, is_causal=False) + target = _attention(q, k, v, is_causal=False) + cos = torch.nn.functional.cosine_similarity(out.float().flatten(), target.float().flatten(), dim=0) + rel_l2 = torch.linalg.vector_norm((out - target).float()) / torch.linalg.vector_norm(target.float()) + + assert cos.item() > 0.9999 + assert rel_l2.item() < 0.01 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Thor native attention fallback check requires CUDA") +def test_cosmos3_edge_native_attention_matches_reference_at_small_shape(): + ref = EdgeTransformerFvkLinearReference(object(), device="cuda") + if ref.ctx is None: + pytest.skip("Thor attention_mha_bf16 is not available") + generator = torch.Generator(device="cuda").manual_seed(5678) + q = torch.randn(32, SPEC.num_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + k = torch.randn(48, SPEC.num_kv_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + v = torch.randn(48, SPEC.num_kv_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + + with torch.no_grad(): + out = ref.attention(q, k, v, is_causal=False) + target = _attention(q, k, v, is_causal=False) + + cos = torch.nn.functional.cosine_similarity(out.float().flatten(), target.float().flatten(), dim=0) + rel_l2 = torch.linalg.vector_norm((out - target).float()) / torch.linalg.vector_norm(target.float()) + assert cos.item() > 0.9999 + assert rel_l2.item() < 0.01 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Thor FA4 fwd reference check requires CUDA") +def test_cosmos3_edge_fa4_fwd_attention_matches_reference(monkeypatch): + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_FA4_FWD", "1") + ref = EdgeTransformerFvkLinearReference(object(), device="cuda") + if ref._get_fa4_fwd() is None: + pytest.skip("Thor FA4 fwd path is not available") + + generator = torch.Generator(device="cuda").manual_seed(9012) + q = torch.randn(64, SPEC.num_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + k = torch.randn(128, SPEC.num_kv_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + v = torch.randn(128, SPEC.num_kv_heads, SPEC.head_dim, device="cuda", dtype=torch.bfloat16, generator=generator) + + with torch.no_grad(): + out = ref.attention(q, k, v, is_causal=False) + target = _attention(q, k, v, is_causal=False) + + cos = torch.nn.functional.cosine_similarity(out.float().flatten(), target.float().flatten(), dim=0) + rel_l2 = torch.linalg.vector_norm((out - target).float()) / torch.linalg.vector_norm(target.float()) + assert cos.item() > 0.9999 + assert rel_l2.item() < 0.01 diff --git a/tests/test_cosmos3_edge_step0_reference.py b/tests/test_cosmos3_edge_step0_reference.py new file mode 100644 index 00000000..3111be75 --- /dev/null +++ b/tests/test_cosmos3_edge_step0_reference.py @@ -0,0 +1,175 @@ +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") + +from flash_rt.models.cosmos3_edge.boundary_dump import EdgeBoundaryDump # noqa: E402 +from flash_rt.models.cosmos3_edge.denoise_ref import EdgeDenoiseTorchReference # noqa: E402 +from flash_rt.models.cosmos3_edge.dump_replay import EdgeDenoiseDump # noqa: E402 +from flash_rt.models.cosmos3_edge.layer_ref import EdgeTransformerTorchReference # noqa: E402 +from flash_rt.models.cosmos3_edge.static_engine import EdgeStaticBufferEngine # noqa: E402 +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights # noqa: E402 + + +def _local_paths() -> tuple[Path, Path] | None: + candidates = [ + ( + Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + ), + ( + Path("/work/models/Cosmos3-Edge"), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0_boundary_step0/tensors.safetensors" + ), + ), + ] + return next((paths for paths in candidates if paths[0].exists() and paths[1].exists()), None) + + +def _local_denoise_dump_path() -> Path | None: + candidates = [ + Path( + "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" + "dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors" + ), + Path( + "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" + "edge_av_inverse_0/tensors.safetensors" + ), + ] + return next((path for path in candidates if path.exists()), None) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="full step0 reference requires CUDA") +def test_cosmos3_edge_step0_full_sequence_reconstruction_matches_boundary(): + paths = _local_paths() + if paths is None: + pytest.skip("local Cosmos3-Edge checkpoint/boundary dump is not available") + checkpoint, dump_path = paths + + dump = EdgeBoundaryDump(dump_path) + ref = EdgeTransformerTorchReference(EdgeTransformerWeights(checkpoint), device="cuda") + + with torch.no_grad(): + full = ref.full_sequence_for_step(dump, dump.tensors["steps/00/noise_x"], dump.tensors["steps/00/timestep"]) + + target = dump.layer0_input_full.to(device="cuda", dtype=torch.bfloat16) + cos = torch.nn.functional.cosine_similarity(full.float().flatten(), target.float().flatten(), dim=0) + + assert cos.item() > 0.9999 + assert (full.float() - target.float()).abs().max().item() < 1e-2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="full step0 reference requires CUDA") +def test_cosmos3_edge_step0_torch_reference_matches_action_velocity(): + paths = _local_paths() + if paths is None: + pytest.skip("local Cosmos3-Edge checkpoint/boundary dump is not available") + checkpoint, dump_path = paths + + dump = EdgeBoundaryDump(dump_path) + ref = EdgeTransformerTorchReference(EdgeTransformerWeights(checkpoint), device="cuda") + + with torch.no_grad(): + action_velocity = ref.action_velocity_for_step( + dump, + dump.tensors["steps/00/noise_x"], + dump.tensors["steps/00/timestep"], + ) + und_cache = ref.precompute_und_kv_cache(dump) + cached_action_velocity = ref.action_velocity_for_step_with_und_cache( + dump, + dump.tensors["steps/00/noise_x"], + dump.tensors["steps/00/timestep"], + und_cache, + ) + + official = dump.tensors["steps/00/velocity"] + official_action = official[-60 * 64 :].reshape(60, 64).to(device="cuda", dtype=torch.bfloat16) + cos = torch.nn.functional.cosine_similarity(action_velocity.float().flatten(), official_action.float().flatten(), dim=0) + + assert cos.item() > 0.9999 + assert (action_velocity.float() - official_action.float()).abs().max().item() < 2e-2 + assert (cached_action_velocity.float() - action_velocity.float()).abs().max().item() == 0.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="reference denoise requires CUDA") +def test_cosmos3_edge_reference_denoise_one_step_matches_dump(): + paths = _local_paths() + denoise_path = _local_denoise_dump_path() + if paths is None or denoise_path is None: + pytest.skip("local Cosmos3-Edge checkpoint/dumps are not available") + checkpoint, boundary_path = paths + + ref = EdgeDenoiseTorchReference( + EdgeDenoiseDump(denoise_path), + EdgeBoundaryDump(boundary_path), + EdgeTransformerWeights(checkpoint), + device="cuda", + ) + + with torch.no_grad(): + result = ref.run(max_steps=1) + + assert result.steps_run == 1 + assert result.timesteps == (999,) + assert result.max_input_abs_diff == 0.0 + assert result.max_velocity_abs_diff < 2e-2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="static engine requires CUDA") +def test_cosmos3_edge_static_engine_step0_contract(): + paths = _local_paths() + if paths is None: + pytest.skip("local Cosmos3-Edge checkpoint/boundary dump is not available") + checkpoint, boundary_path = paths + + dump = EdgeBoundaryDump(boundary_path) + engine = EdgeStaticBufferEngine(EdgeTransformerWeights(checkpoint), dump, device="cuda") + if not getattr(engine.reference, "native_attention_available", False): + pytest.skip("native attention extension is not available") + + with torch.no_grad(): + full = engine.full_sequence_for_step(dump.tensors["steps/00/noise_x"], dump.tensors["steps/00/timestep"]) + velocity = engine.flat_velocity_for_step(dump.tensors["steps/00/noise_x"], dump.tensors["steps/00/timestep"]) + + target_full = dump.layer0_input_full.to(device="cuda", dtype=torch.bfloat16) + official = dump.tensors["steps/00/velocity"].to(device="cuda", dtype=torch.bfloat16) + vision_dim = official.numel() - 60 * 64 + + assert torch.nn.functional.cosine_similarity(full.float().flatten(), target_full.float().flatten(), dim=0) > 0.9999 + assert (full.float() - target_full.float()).abs().max().item() < 1e-2 + assert velocity[:vision_dim].float().abs().max().item() == 0.0 + assert (velocity[vision_dim:].float() - official[vision_dim:].float()).abs().max().item() < 2e-2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="static engine CUDA graph requires CUDA") +def test_cosmos3_edge_static_engine_velocity_graph_matches_eager(): + paths = _local_paths() + if paths is None: + pytest.skip("local Cosmos3-Edge checkpoint/boundary dump is not available") + checkpoint, boundary_path = paths + + dump = EdgeBoundaryDump(boundary_path) + engine = EdgeStaticBufferEngine(EdgeTransformerWeights(checkpoint), dump, device="cuda") + if not getattr(engine.reference, "graph_attention_available", False): + pytest.skip("graph-safe native attention extension is not available") + + with torch.no_grad(): + eager = engine.flat_velocity_for_step( + dump.tensors["steps/00/noise_x"], + dump.tensors["steps/00/timestep"], + ).clone() + graph = engine.replay_velocity_graph( + dump.tensors["steps/00/noise_x"], + dump.tensors["steps/00/timestep"], + ).clone() + + assert (graph.float() - eager.float()).abs().max().item() == 0.0 diff --git a/tests/test_cosmos3_edge_thor_spec.py b/tests/test_cosmos3_edge_thor_spec.py new file mode 100644 index 00000000..ff973c9d --- /dev/null +++ b/tests/test_cosmos3_edge_thor_spec.py @@ -0,0 +1,32 @@ +from pathlib import Path + +import pytest + +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import ( + GEN_ATTENTION_KEYS, + SPEC, + UND_ATTENTION_KEYS, + iter_expected_shapes, + layer_key, + validate_transformer_index, +) + + +def test_cosmos3_edge_spec_counts_and_tower_names(): + shapes = dict(iter_expected_shapes()) + + assert SPEC.num_layers == 28 + assert len(shapes) == 549 + assert shapes[layer_key(0, "self_attn.to_q.weight")] == (2048, 2048) + assert shapes[layer_key(0, "self_attn.add_k_proj.weight")] == (1024, 2048) + assert shapes[layer_key(27, "self_attn.k_norm_und_for_gen.weight")] == (128,) + + assert "self_attn.add_q_proj.weight" in GEN_ATTENTION_KEYS + assert "self_attn.to_q.weight" in UND_ATTENTION_KEYS + + +def test_cosmos3_edge_local_checkpoint_index_matches_when_present(): + checkpoint = "/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge" + if not Path(checkpoint).exists(): + pytest.skip("local Cosmos3-Edge checkpoint is not available") + validate_transformer_index(checkpoint) diff --git a/tests/test_cosmos3_edge_weights.py b/tests/test_cosmos3_edge_weights.py new file mode 100644 index 00000000..09238f14 --- /dev/null +++ b/tests/test_cosmos3_edge_weights.py @@ -0,0 +1,49 @@ +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") + +from flash_rt.frontends.torch._cosmos3_edge_thor_spec import layer_key # noqa: E402 +from flash_rt.models.cosmos3_edge.weights import EdgeTransformerWeights # noqa: E402 + + +def _local_checkpoint() -> Path | None: + candidates = [ + Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), + Path("/work/models/Cosmos3-Edge"), + ] + return next((path for path in candidates if path.exists()), None) + + +def test_cosmos3_edge_weight_loader_refs_expected_shard(): + checkpoint = _local_checkpoint() + if checkpoint is None: + pytest.skip("local Cosmos3-Edge checkpoint is not available") + + weights = EdgeTransformerWeights(checkpoint) + ref = weights.ref(layer_key(0, "self_attn.to_q.weight")) + + assert ref.shape == (2048, 2048) + assert ref.shard.name in { + "diffusion_pytorch_model-00001-of-00002.safetensors", + "diffusion_pytorch_model-00002-of-00002.safetensors", + } + + +def test_cosmos3_edge_weight_loader_reads_selected_tensors(): + checkpoint = _local_checkpoint() + if checkpoint is None: + pytest.skip("local Cosmos3-Edge checkpoint is not available") + + weights = EdgeTransformerWeights(checkpoint) + + norm = weights.load_tensor("norm.weight") + gen_q = weights.layer_tensor(0, "self_attn.to_q.weight", dtype=torch.bfloat16) + und_k = weights.layer_tensor(0, "self_attn.add_k_proj.weight", dtype=torch.bfloat16) + + assert norm.shape == (2048,) + assert gen_q.shape == (2048, 2048) + assert und_k.shape == (1024, 2048) + assert gen_q.dtype == torch.bfloat16 diff --git a/tests/test_load_model_use_fp8_kwarg.py b/tests/test_load_model_use_fp8_kwarg.py index f371ab9a..0fb167f2 100644 --- a/tests/test_load_model_use_fp8_kwarg.py +++ b/tests/test_load_model_use_fp8_kwarg.py @@ -1,4 +1,5 @@ import ast +import json from pathlib import Path import sys import types @@ -470,6 +471,1141 @@ def infer(self, *args, **kwargs): } +def test_load_model_accepts_cosmos3_edge_thor_config(): + from flash_rt.api import load_model + + class Cosmos3EdgeFrontend: + seen = None + + def __init__(self, checkpoint, num_views=1, hardware=None): + type(self).seen = { + "checkpoint": checkpoint, + "num_views": num_views, + "hardware": hardware, + } + + def set_prompt(self, *args, **kwargs): + return None + + def infer(self, *args, **kwargs): + return None + + with patch("flash_rt.hardware.resolve_pipeline_class", + return_value=Cosmos3EdgeFrontend): + model = load_model( + "unused-checkpoint", + config="cosmos3_edge", + framework="torch", + hardware="thor", + num_views=1, + ) + + assert isinstance(model._pipe, Cosmos3EdgeFrontend) + assert Cosmos3EdgeFrontend.seen == { + "checkpoint": "unused-checkpoint", + "num_views": 1, + "hardware": "thor", + } + + +def test_cosmos3_edge_infer_exposes_torch_reference_backend_args(): + import inspect + + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + sig = inspect.signature(Cosmos3EdgeTorchFrontendThor.infer) + assert "reference_dump" in sig.parameters + assert "boundary_dump" in sig.parameters + assert "live_dump_out" in sig.parameters + assert "live_flashrt_handoff" in sig.parameters + assert "live_boundary_out" in sig.parameters + assert "live_boundary_in" in sig.parameters + assert "live_boundary_prepare_in" in sig.parameters + assert "live_boundary_prepare_live" in sig.parameters + assert "live_handoff_trace_out" in sig.parameters + assert "upstream_trace_out" in sig.parameters + assert "vae_encode_dump_out" in sig.parameters + assert "vae_latent_in" in sig.parameters + assert "vae_encode_dump_input" in sig.parameters + assert "vae_encode_profile_out" in sig.parameters + assert "vae_native_rms_silu" in sig.parameters + assert "vae_t1_conv2d" in sig.parameters + assert "vae_native_avgdown3d" in sig.parameters + assert "vae_channels_last3d_conv320" in sig.parameters + assert "vae_compile_encode" in sig.parameters + assert "vae_compile_trace_out" in sig.parameters + assert "prepare_dump_out" in sig.parameters + assert "prepare_replay_in" in sig.parameters + assert "prepare_inventory_out" in sig.parameters + assert "prepare_slim_no_raw_state_vision" in sig.parameters + assert "prepare_slim_derive_condition_reference" in sig.parameters + assert "prepare_slim_derive_initial_noise" in sig.parameters + assert "live_prelayer_bootstrap" in sig.parameters + assert "live_warm_request" in sig.parameters + assert "cache_warmup_vae" in sig.parameters + assert "cache_warmup_prepare" in sig.parameters + assert "warmup" in sig.parameters + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + with pytest.raises(ValueError, match="official.*official_action_only.*replay.*torch_ref.*flashrt"): + pipe.infer(output_dir="/tmp/unused", backend="bad") + with pytest.raises(ValueError, match="warmup must be non-negative"): + pipe.infer(output_dir="/tmp/unused", backend="official_action_only", warmup=-1) + with pytest.raises(ValueError, match="backend='flashrt' requires reference_dump"): + pipe.infer(output_dir="/tmp/unused", backend="flashrt") + with pytest.raises(ValueError, match="live_dump_out.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", live_dump_out="/tmp/live.safetensors") + with pytest.raises(ValueError, match="upstream_trace_out.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", upstream_trace_out="/tmp/upstream.json") + with pytest.raises(ValueError, match="VAE/prepare boundary.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_encode_dump_out="/tmp/vae.safetensors") + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_encode_profile_out="/tmp/vae_profile.json") + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_native_rms_silu=True) + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_t1_conv2d=True) + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_native_avgdown3d=True) + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_channels_last3d_conv320=True) + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", vae_compile_encode=True) + with pytest.raises(ValueError, match="VAE/prepare boundary/profile.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", prepare_dump_out="/tmp/prepare.pt") + with pytest.raises(ValueError, match="VAE/prepare boundary/profile/inventory.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", prepare_inventory_out="/tmp/prepare.json") + with pytest.raises(ValueError, match="VAE/prepare boundary/profile/inventory.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", prepare_slim_no_raw_state_vision=True) + with pytest.raises(ValueError, match="VAE/prepare boundary/profile/inventory.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", prepare_slim_derive_initial_noise=True) + with pytest.raises(ValueError, match="vae_compile_encode cannot be combined"): + pipe.infer( + output_dir="/tmp/unused", + backend="official_action_only", + vae_compile_encode=True, + vae_native_rms_silu=True, + ) + with pytest.raises(ValueError, match="vae_compile_encode cannot be combined"): + pipe.infer( + output_dir="/tmp/unused", + backend="official_action_only", + vae_compile_encode=True, + vae_native_avgdown3d=True, + ) + with pytest.raises(ValueError, match="cache_warmup_vae cannot be combined"): + pipe.infer( + output_dir="/tmp/unused", + backend="official_action_only", + cache_warmup_vae=True, + vae_latent_in="/tmp/vae.safetensors", + ) + with pytest.raises(ValueError, match="live FlashRT handoff.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", live_flashrt_handoff=True) + with pytest.raises(ValueError, match="live FlashRT handoff.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", live_boundary_in="/tmp/boundary.safetensors") + with pytest.raises(ValueError, match="live FlashRT handoff.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", live_boundary_prepare_live=True) + with pytest.raises(ValueError, match="live_warm_request.*official_action_only"): + pipe.infer(output_dir="/tmp/unused", backend="official", live_warm_request=True) + with pytest.raises(ValueError, match="live_dump_out cannot be combined"): + pipe.infer( + output_dir="/tmp/unused", + backend="official_action_only", + live_dump_out="/tmp/live.safetensors", + live_flashrt_handoff=True, + ) + with pytest.raises(ValueError, match="live_boundary_prepare_live cannot be combined"): + pipe.infer( + output_dir="/tmp/unused", + backend="official_action_only", + live_boundary_in="/tmp/boundary.safetensors", + live_boundary_prepare_live=True, + ) + + +def test_cosmos3_edge_live_warm_request_builds_official_warmup_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_warm_request=True, + live_handoff_trace_out="trace.json", + upstream_trace_out="upstream.json", + cache_warmup_prepare=True, + prepare_slim_no_raw_state_vision=True, + prepare_slim_derive_initial_noise=True, + ) + + cmd = seen["cmd"] + assert cmd[2] == "flash_rt.models.cosmos3_edge.action_only_official" + assert "--flashrt-live-prelayer-bootstrap" in cmd + assert cmd[cmd.index("--warmup") + 1] == "1" + assert cmd[cmd.index("--flashrt-live-handoff-trace-out") + 1] == str(tmp_path / "trace.json") + assert cmd[cmd.index("--flashrt-upstream-trace-out") + 1] == str(tmp_path / "upstream.json") + assert "--flashrt-cache-warmup-vae" in cmd + assert "--flashrt-cache-warmup-prepare" in cmd + assert "--flashrt-prepare-slim-no-raw-state-vision" in cmd + assert "--flashrt-prepare-slim-derive-initial-noise" in cmd + assert result["live_warm_request"] is True + assert result["live_prelayer_bootstrap"] is True + assert result["cache_warmup_vae"] is True + assert result["cache_warmup_prepare"] is True + assert result["prepare_slim_no_raw_state_vision"] is True + assert result["prepare_slim_derive_initial_noise"] is True + assert result["warmup"] == 1 + assert result["live_handoff_trace_out"] == str(tmp_path / "trace.json") + assert result["upstream_trace_out"] == str(tmp_path / "upstream.json") + + +def test_cosmos3_edge_live_boundary_in_builds_official_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_boundary_in="prelayer_boundary.safetensors", + live_handoff_trace_out="boundary_in_trace.json", + ) + + cmd = seen["cmd"] + assert cmd[2] == "flash_rt.models.cosmos3_edge.action_only_official" + assert cmd[cmd.index("--flashrt-live-boundary-in") + 1] == str(tmp_path / "prelayer_boundary.safetensors") + assert cmd[cmd.index("--flashrt-live-handoff-trace-out") + 1] == str(tmp_path / "boundary_in_trace.json") + assert "--flashrt-live-prelayer-bootstrap" not in cmd + assert result["live_boundary_in"] == str(tmp_path / "prelayer_boundary.safetensors") + assert result["live_handoff_trace_out"] == str(tmp_path / "boundary_in_trace.json") + + +def test_cosmos3_edge_live_boundary_prepare_in_builds_official_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_boundary_prepare_in="slim_prepare.pt", + live_boundary_out="derived_boundary.safetensors", + live_handoff_trace_out="prepare_boundary_trace.json", + ) + + cmd = seen["cmd"] + assert cmd[2] == "flash_rt.models.cosmos3_edge.action_only_official" + assert cmd[cmd.index("--flashrt-live-boundary-prepare-in") + 1] == str(tmp_path / "slim_prepare.pt") + assert cmd[cmd.index("--flashrt-live-boundary-out") + 1] == str(tmp_path / "derived_boundary.safetensors") + assert cmd[cmd.index("--flashrt-live-handoff-trace-out") + 1] == str(tmp_path / "prepare_boundary_trace.json") + assert "--flashrt-live-boundary-in" not in cmd + assert result["live_boundary_prepare_in"] == str(tmp_path / "slim_prepare.pt") + assert result["live_boundary_out"] == str(tmp_path / "derived_boundary.safetensors") + assert result["live_handoff_trace_out"] == str(tmp_path / "prepare_boundary_trace.json") + + +def test_cosmos3_edge_live_boundary_prepare_live_builds_official_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_boundary_prepare_live=True, + live_boundary_out="derived_boundary.safetensors", + live_handoff_trace_out="prepare_live_trace.json", + upstream_trace_out="prepare_live_upstream.json", + prepare_slim_no_raw_state_vision=True, + prepare_slim_derive_condition_reference=True, + prepare_slim_derive_initial_noise=True, + ) + + cmd = seen["cmd"] + assert cmd[2] == "flash_rt.models.cosmos3_edge.action_only_official" + assert "--flashrt-live-boundary-prepare-live" in cmd + assert cmd[cmd.index("--flashrt-live-boundary-out") + 1] == str(tmp_path / "derived_boundary.safetensors") + assert cmd[cmd.index("--flashrt-live-handoff-trace-out") + 1] == str(tmp_path / "prepare_live_trace.json") + assert cmd[cmd.index("--flashrt-upstream-trace-out") + 1] == str(tmp_path / "prepare_live_upstream.json") + assert "--flashrt-live-boundary-in" not in cmd + assert "--flashrt-live-boundary-prepare-in" not in cmd + assert "--flashrt-prepare-slim-no-raw-state-vision" in cmd + assert "--flashrt-prepare-slim-derive-condition-reference" in cmd + assert "--flashrt-prepare-slim-derive-initial-noise" in cmd + assert result["live_boundary_prepare_live"] is True + assert result["live_boundary_out"] == str(tmp_path / "derived_boundary.safetensors") + assert result["live_handoff_trace_out"] == str(tmp_path / "prepare_live_trace.json") + assert result["upstream_trace_out"] == str(tmp_path / "prepare_live_upstream.json") + assert result["prepare_slim_no_raw_state_vision"] is True + assert result["prepare_slim_derive_condition_reference"] is True + assert result["prepare_slim_derive_initial_noise"] is True + + +def test_cosmos3_edge_live_boundary_prepare_live_allows_warm_prepare_cache_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_boundary_prepare_live=True, + cache_warmup_vae=True, + cache_warmup_prepare=True, + warmup=1, + prepare_slim_no_raw_state_vision=True, + prepare_slim_derive_condition_reference=True, + prepare_slim_derive_initial_noise=True, + ) + + cmd = seen["cmd"] + assert "--flashrt-live-boundary-prepare-live" in cmd + assert "--flashrt-cache-warmup-vae" in cmd + assert "--flashrt-cache-warmup-prepare" in cmd + assert cmd[cmd.index("--warmup") + 1] == "1" + assert "--flashrt-live-boundary-in" not in cmd + assert "--flashrt-live-boundary-prepare-in" not in cmd + assert result["live_boundary_prepare_live"] is True + assert result["cache_warmup_vae"] is True + assert result["cache_warmup_prepare"] is True + assert result["warmup"] == 1 + + +def test_cosmos3_edge_vae_boundary_builds_official_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_prelayer_bootstrap=True, + vae_encode_dump_out="vae.safetensors", + vae_encode_dump_input=True, + vae_encode_profile_out="vae_profile.json", + vae_native_rms_silu=True, + vae_t1_conv2d=True, + vae_native_avgdown3d=True, + vae_channels_last3d_conv320=True, + ) + + cmd = seen["cmd"] + assert cmd[cmd.index("--flashrt-vae-encode-dump-out") + 1] == str(tmp_path / "vae.safetensors") + assert "--flashrt-vae-encode-dump-input" in cmd + assert cmd[cmd.index("--flashrt-vae-encode-profile-out") + 1] == str(tmp_path / "vae_profile.json") + assert "--flashrt-vae-native-rms-silu" in cmd + assert "--flashrt-vae-t1-conv2d" in cmd + assert "--flashrt-vae-native-avgdown3d" in cmd + assert "--flashrt-vae-channels-last3d-conv320" in cmd + assert "--flashrt-cache-warmup-vae" not in cmd + assert result["vae_encode_dump_out"] == str(tmp_path / "vae.safetensors") + assert result["vae_encode_dump_input"] is True + assert result["vae_encode_profile_out"] == str(tmp_path / "vae_profile.json") + assert result["vae_native_rms_silu"] is True + assert result["vae_t1_conv2d"] is True + assert result["vae_native_avgdown3d"] is True + assert result["vae_channels_last3d_conv320"] is True + + +def test_cosmos3_edge_vae_compile_builds_official_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_prelayer_bootstrap=True, + vae_compile_encode=True, + vae_compile_trace_out="vae_compile.json", + ) + + cmd = seen["cmd"] + assert "--flashrt-vae-compile-encode" in cmd + assert cmd[cmd.index("--flashrt-vae-compile-trace-out") + 1] == str(tmp_path / "vae_compile.json") + assert result["vae_compile_encode"] is True + assert result["vae_compile_trace_out"] == str(tmp_path / "vae_compile.json") + + +def test_cosmos3_edge_prepare_boundary_builds_official_cmd(tmp_path, monkeypatch): + from flash_rt.frontends.torch import cosmos3_edge_thor + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + input_json = tmp_path / "input.json" + input_json.write_text('{"samples":[]}\n', encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + seen = {} + + def fake_run(cmd, **kwargs): + seen["cmd"] = cmd + seen["cwd"] = kwargs["cwd"] + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(cosmos3_edge_thor.subprocess, "run", fake_run) + monkeypatch.chdir(tmp_path) + + pipe = Cosmos3EdgeTorchFrontendThor("unused-checkpoint", hardware="thor") + pipe.set_prompt(input_json=str(input_json)) + result = pipe.infer( + output_dir=str(tmp_path / "out"), + backend="official_action_only", + cosmos_root=str(tmp_path), + config_file=str(config_file), + live_prelayer_bootstrap=True, + prepare_dump_out="prepare.pt", + prepare_replay_in="prepare_in.pt", + prepare_inventory_out="prepare_inventory.json", + prepare_slim_no_raw_state_vision=True, + prepare_slim_derive_condition_reference=True, + prepare_slim_derive_initial_noise=True, + ) + + cmd = seen["cmd"] + assert cmd[cmd.index("--flashrt-prepare-dump-out") + 1] == str(tmp_path / "prepare.pt") + assert cmd[cmd.index("--flashrt-prepare-replay-in") + 1] == str(tmp_path / "prepare_in.pt") + assert cmd[cmd.index("--flashrt-prepare-inventory-out") + 1] == str(tmp_path / "prepare_inventory.json") + assert "--flashrt-prepare-slim-no-raw-state-vision" in cmd + assert "--flashrt-prepare-slim-derive-condition-reference" in cmd + assert "--flashrt-prepare-slim-derive-initial-noise" in cmd + assert result["prepare_dump_out"] == str(tmp_path / "prepare.pt") + assert result["prepare_replay_in"] == str(tmp_path / "prepare_in.pt") + assert result["prepare_inventory_out"] == str(tmp_path / "prepare_inventory.json") + assert result["prepare_slim_no_raw_state_vision"] is True + assert result["prepare_slim_derive_condition_reference"] is True + assert result["prepare_slim_derive_initial_noise"] is True + + +def test_cosmos3_edge_action_only_output_writer(tmp_path): + from flash_rt.frontends.torch.cosmos3_edge_thor import ( + Cosmos3EdgeTorchFrontendThor, + ) + + output_path = Cosmos3EdgeTorchFrontendThor._write_action_only_output( + [[1.0, 2.0, 3.0]], + tmp_path, + ) + payload = json.loads(output_path.read_text(encoding="utf-8")) + + assert payload["status"] == "success" + assert payload["outputs"][0]["content"]["action"] == [[1.0, 2.0, 3.0]] + assert payload["outputs"][0]["files"] == [] + + +def test_cosmos3_edge_official_action_only_rewriter(tmp_path): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _rewrite_sample_outputs_action_only, + ) + + sample_dir = tmp_path / "sample" + sample_dir.mkdir() + (sample_dir / "vision.mp4").write_bytes(b"fake") + sample_outputs = sample_dir / "sample_outputs.json" + sample_outputs.write_text( + json.dumps( + { + "status": "success", + "outputs": [ + { + "content": {"action": [[1.0, 2.0]]}, + "files": ["vision.mp4"], + } + ], + } + ), + encoding="utf-8", + ) + + assert _rewrite_sample_outputs_action_only(tmp_path) == 1 + payload = json.loads(sample_outputs.read_text(encoding="utf-8")) + assert payload["outputs"][0]["files"] == [] + assert payload["outputs"][0]["content"]["action"] == [[1.0, 2.0]] + assert not (sample_dir / "vision.mp4").exists() + + +def test_cosmos3_edge_official_action_only_live_dump_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_live_dump_out, + ) + + argv, path = _extract_live_dump_out( + ["-i", "sample.json", "--flashrt-live-dump-out", "/tmp/live.safetensors", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert path == Path("/tmp/live.safetensors") + + argv, path = _extract_live_dump_out( + ["-i", "sample.json", "--flashrt-live-dump-out=/tmp/live2.safetensors"] + ) + assert argv == ["-i", "sample.json"] + assert path == Path("/tmp/live2.safetensors") + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_DUMP_OUT", "/tmp/live3.safetensors") + argv, path = _extract_live_dump_out(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert path == Path("/tmp/live3.safetensors") + + +def test_cosmos3_edge_official_action_only_live_handoff_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_live_handoff_args, + ) + + argv, enabled, boundary, boundary_in, boundary_prepare_in, prepare_live, trace, prelayer = _extract_live_handoff_args( + [ + "-i", + "sample.json", + "--flashrt-live-flashrt-handoff", + "--flashrt-live-prelayer-bootstrap", + "--flashrt-live-boundary-out", + "/tmp/boundary.safetensors", + "--flashrt-live-boundary-in", + "/tmp/prelayer_boundary.safetensors", + "--flashrt-live-handoff-trace-out", + "/tmp/trace.json", + "--seed", + "0", + ] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + assert prelayer is True + assert boundary == Path("/tmp/boundary.safetensors") + assert boundary_in == Path("/tmp/prelayer_boundary.safetensors") + assert boundary_prepare_in is None + assert prepare_live is False + assert trace == Path("/tmp/trace.json") + + argv, enabled, boundary, boundary_in, boundary_prepare_in, prepare_live, trace, prelayer = _extract_live_handoff_args( + [ + "-i", + "sample.json", + "--flashrt-live-boundary-out=/tmp/boundary2.safetensors", + "--flashrt-live-boundary-in=/tmp/prelayer_boundary2.safetensors", + "--flashrt-live-boundary-prepare-in=/tmp/slim_prepare.pt", + "--flashrt-live-boundary-prepare-live", + "--flashrt-live-handoff-trace-out=/tmp/trace2.json", + ] + ) + assert argv == ["-i", "sample.json"] + assert enabled is True + assert prelayer is False + assert boundary == Path("/tmp/boundary2.safetensors") + assert boundary_in == Path("/tmp/prelayer_boundary2.safetensors") + assert boundary_prepare_in == Path("/tmp/slim_prepare.pt") + assert prepare_live is True + assert trace == Path("/tmp/trace2.json") + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_FLASHRT_HANDOFF", "1") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_PRELAYER_BOOTSTRAP", "1") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_OUT", "/tmp/boundary3.safetensors") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_IN", "/tmp/prelayer_boundary3.safetensors") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_PREPARE_IN", "/tmp/slim_prepare3.pt") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_BOUNDARY_PREPARE_LIVE", "1") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_LIVE_HANDOFF_TRACE_OUT", "/tmp/trace3.json") + argv, enabled, boundary, boundary_in, boundary_prepare_in, prepare_live, trace, prelayer = _extract_live_handoff_args(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + assert prelayer is True + assert boundary == Path("/tmp/boundary3.safetensors") + assert boundary_in == Path("/tmp/prelayer_boundary3.safetensors") + assert boundary_prepare_in == Path("/tmp/slim_prepare3.pt") + assert prepare_live is True + assert trace == Path("/tmp/trace3.json") + + +def test_cosmos3_edge_live_handoff_trace_native_scheduler_summary(tmp_path): + from flash_rt.models.cosmos3_edge.action_only_official import _LiveFlashRTHandoff + + trace = tmp_path / "trace.json" + handoff = _LiveFlashRTHandoff(Path("/tmp/checkpoint"), None, None, None, None, trace, prelayer_bootstrap=True) + handoff.trace["flashrt_velocity_calls"] = [{"step": 0, "s": 0.25}, {"step": 1, "s": 0.35}] + handoff.trace["native_scheduler"]["runs"].append( + { + "num_steps": 30, + "shift": 10.0, + "total_s": 0.42, + "scheduler_step_total_s": 0.03, + } + ) + + handoff.save_trace() + + payload = json.loads(trace.read_text(encoding="utf-8")) + assert payload["summary"]["native_scheduler_enabled"] is True + assert payload["summary"]["native_scheduler_run_count"] == 1 + assert payload["summary"]["native_scheduler_step_count"] == 30 + assert payload["summary"]["native_scheduler_step_total_s"] == pytest.approx(0.03) + + +def test_cosmos3_edge_official_action_only_upstream_trace_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_upstream_trace_out, + ) + + argv, path = _extract_upstream_trace_out( + ["-i", "sample.json", "--flashrt-upstream-trace-out", "/tmp/upstream.json", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert path == Path("/tmp/upstream.json") + + argv, path = _extract_upstream_trace_out( + ["-i", "sample.json", "--flashrt-upstream-trace-out=/tmp/upstream2.json"] + ) + assert argv == ["-i", "sample.json"] + assert path == Path("/tmp/upstream2.json") + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_UPSTREAM_TRACE_OUT", "/tmp/upstream3.json") + argv, path = _extract_upstream_trace_out(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert path == Path("/tmp/upstream3.json") + + +def test_cosmos3_edge_official_action_only_vae_boundary_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_encode_boundary_args, + ) + + argv, dump_out, latent_in, dump_input = _extract_vae_encode_boundary_args( + [ + "-i", + "sample.json", + "--flashrt-vae-encode-dump-out", + "/tmp/vae.safetensors", + "--flashrt-vae-latent-in=/tmp/latent.safetensors", + "--flashrt-vae-encode-dump-input", + "--seed", + "0", + ] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert dump_out == Path("/tmp/vae.safetensors") + assert latent_in == Path("/tmp/latent.safetensors") + assert dump_input is True + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_ENCODE_DUMP_OUT", "/tmp/env_vae.safetensors") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_LATENT_IN", "/tmp/env_latent.safetensors") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_ENCODE_DUMP_INPUT", "1") + argv, dump_out, latent_in, dump_input = _extract_vae_encode_boundary_args(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert dump_out == Path("/tmp/env_vae.safetensors") + assert latent_in == Path("/tmp/env_latent.safetensors") + assert dump_input is True + + +def test_cosmos3_edge_official_action_only_vae_profile_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_encode_profile_out, + ) + + argv, path = _extract_vae_encode_profile_out( + ["-i", "sample.json", "--flashrt-vae-encode-profile-out", "/tmp/profile.json", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert path == Path("/tmp/profile.json") + + argv, path = _extract_vae_encode_profile_out( + ["-i", "sample.json", "--flashrt-vae-encode-profile-out=/tmp/profile2.json"] + ) + assert argv == ["-i", "sample.json"] + assert path == Path("/tmp/profile2.json") + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_ENCODE_PROFILE_OUT", "/tmp/profile3.json") + argv, path = _extract_vae_encode_profile_out(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert path == Path("/tmp/profile3.json") + + +def test_cosmos3_edge_official_action_only_vae_profile_native_candidate_summary(tmp_path): + from flash_rt.models.cosmos3_edge.action_only_official import _VAEEncodeProfiler + + out = tmp_path / "vae_profile.json" + profiler = _VAEEncodeProfiler(out) + profiler.encode_calls.append({"encode_call": 0, "input": None, "output": None, "s": 1.0}) + profiler.modules.extend( + [ + {"name": "encoder.block.conv", "class": "CausalConv3d"}, + {"name": "encoder.conv1", "class": "CausalConv3d"}, + ] + ) + profiler.events.extend( + [ + { + "encode_call": 0, + "name": "encoder.conv1", + "class": "CausalConv3d", + "input": [ + {"shape": [1, 12, 1, 240, 416]}, + None, + ], + "output": {"shape": [1, 160, 1, 240, 416]}, + "s": 0.02, + "parameter_numel": 52000, + }, + { + "encode_call": 0, + "name": "encoder.block.conv", + "class": "CausalConv3d", + "input": [ + {"shape": [1, 160, 24, 240, 416]}, + {"shape": [1, 160, 2, 240, 416]}, + ], + "output": {"shape": [1, 160, 24, 240, 416]}, + "s": 0.11, + "parameter_numel": 691360, + }, + ] + ) + + profiler.save() + + payload = json.loads(out.read_text(encoding="utf-8")) + candidates = payload["native_candidate_summary"] + assert candidates[0]["candidate_type"] == "steady_cached_causal_conv3d" + assert candidates[0]["name"] == "encoder.block.conv" + assert candidates[0]["input_shape"] == [1, 160, 24, 240, 416] + assert candidates[0]["cache_shape"] == [1, 160, 2, 240, 416] + assert candidates[1]["candidate_type"] == "prime_t1_no_cache" + + shape_summary = payload["shape_summary"] + steady_shape = next(item for item in shape_summary if item["name"] == "encoder.block.conv") + assert steady_shape["count"] == 1 + assert steady_shape["total_s"] == pytest.approx(0.11) + + +def test_cosmos3_edge_official_action_only_vae_native_rms_silu_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_native_rms_silu, + ) + + argv, enabled = _extract_vae_native_rms_silu( + ["-i", "sample.json", "--flashrt-vae-native-rms-silu", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + + argv, enabled = _extract_vae_native_rms_silu(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is False + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_NATIVE_RMS_SILU", "1") + argv, enabled = _extract_vae_native_rms_silu(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + + +def test_cosmos3_edge_official_action_only_vae_t1_conv2d_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_t1_conv2d, + ) + + argv, enabled = _extract_vae_t1_conv2d(["-i", "sample.json", "--flashrt-vae-t1-conv2d", "--seed", "0"]) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + + argv, enabled = _extract_vae_t1_conv2d(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is False + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_T1_CONV2D", "1") + argv, enabled = _extract_vae_t1_conv2d(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + + +def test_cosmos3_edge_official_action_only_vae_native_avgdown3d_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_native_avgdown3d, + ) + + argv, enabled = _extract_vae_native_avgdown3d( + ["-i", "sample.json", "--flashrt-vae-native-avgdown3d", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + + argv, enabled = _extract_vae_native_avgdown3d(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is False + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_NATIVE_AVGDOWN3D", "1") + argv, enabled = _extract_vae_native_avgdown3d(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + + +def test_cosmos3_edge_official_action_only_vae_channels_last3d_conv320_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_channels_last3d_conv320, + ) + + argv, enabled = _extract_vae_channels_last3d_conv320( + ["-i", "sample.json", "--flashrt-vae-channels-last3d-conv320", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + + argv, enabled = _extract_vae_channels_last3d_conv320(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is False + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_CHANNELS_LAST3D_CONV320", "1") + argv, enabled = _extract_vae_channels_last3d_conv320(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + + +def test_cosmos3_edge_official_action_only_vae_compile_encode_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_vae_compile_encode, + ) + + argv, enabled, trace = _extract_vae_compile_encode( + ["-i", "sample.json", "--flashrt-vae-compile-encode", "--flashrt-vae-compile-trace-out", "/tmp/aot.json"] + ) + assert argv == ["-i", "sample.json"] + assert enabled is True + assert trace == Path("/tmp/aot.json") + + argv, enabled, trace = _extract_vae_compile_encode( + ["-i", "sample.json", "--flashrt-vae-compile-trace-out=/tmp/aot2.json"] + ) + assert argv == ["-i", "sample.json"] + assert enabled is False + assert trace == Path("/tmp/aot2.json") + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_COMPILE_ENCODE", "1") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_VAE_COMPILE_TRACE_OUT", "/tmp/aot3.json") + argv, enabled, trace = _extract_vae_compile_encode(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + assert trace == Path("/tmp/aot3.json") + + +def test_cosmos3_edge_official_action_only_prepare_boundary_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_prepare_boundary_args, + ) + + argv, dump, replay, inventory, slim, derive_ref, derive_noise = _extract_prepare_boundary_args( + [ + "-i", + "sample.json", + "--flashrt-prepare-dump-out", + "/tmp/prepare.pt", + "--flashrt-prepare-replay-in=/tmp/prepare_in.pt", + "--flashrt-prepare-inventory-out", + "/tmp/prepare.json", + "--flashrt-prepare-slim-no-raw-state-vision", + "--flashrt-prepare-slim-derive-condition-reference", + "--flashrt-prepare-slim-derive-initial-noise", + ] + ) + assert argv == ["-i", "sample.json"] + assert dump == Path("/tmp/prepare.pt") + assert replay == Path("/tmp/prepare_in.pt") + assert inventory == Path("/tmp/prepare.json") + assert slim is True + assert derive_ref is True + assert derive_noise is True + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_PREPARE_DUMP_OUT", "/tmp/env_prepare.pt") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_PREPARE_REPLAY_IN", "/tmp/env_prepare_in.pt") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_PREPARE_INVENTORY_OUT", "/tmp/env_prepare.json") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_PREPARE_SLIM_NO_RAW_STATE_VISION", "1") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_PREPARE_SLIM_DERIVE_CONDITION_REFERENCE", "1") + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_PREPARE_SLIM_DERIVE_INITIAL_NOISE", "1") + argv, dump, replay, inventory, slim, derive_ref, derive_noise = _extract_prepare_boundary_args(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert dump == Path("/tmp/env_prepare.pt") + assert replay == Path("/tmp/env_prepare_in.pt") + assert inventory == Path("/tmp/env_prepare.json") + assert slim is True + assert derive_ref is True + assert derive_noise is True + + +def test_cosmos3_edge_prepare_payload_inventory_tensors(): + from dataclasses import dataclass + + import torch + + from flash_rt.models.cosmos3_edge.action_only_official import ( + _prepare_payload_inventory, + _slim_prepare_payload_no_raw_state_vision, + ) + + @dataclass + class ToyGenData: + raw_state_vision: list + x0_tokens_vision: list + + gen_data = ToyGenData( + raw_state_vision=[torch.zeros((1, 3, 2), dtype=torch.float32)], + x0_tokens_vision=[torch.zeros((1, 1, 2), dtype=torch.float32)], + ) + payload = ( + [], + gen_data, + [], + [], + [torch.zeros((8,), dtype=torch.float32)], + [torch.zeros((8,), dtype=torch.float32)], + [torch.zeros((8,), dtype=torch.bool)], + ) + inventory = _prepare_payload_inventory(payload, signature=("sig",), source="unit") + + assert inventory["schema"] == "flashrt_cosmos3_edge_prepare_inventory_v1" + assert inventory["tensor_count"] == 5 + assert inventory["tensor_bytes"] == 1 * 3 * 2 * 4 + 1 * 1 * 2 * 4 + 8 * 4 + 8 * 4 + 8 + assert inventory["field_names"] == [ + "sequence_plans", + "gen_data_clean", + "cond_text_tokens", + "uncond_text_tokens", + "initial_noise", + "condition_reference", + "condition_mask", + ] + by_top = {item["field"]: item for item in inventory["tensor_bytes_by_top_level"]} + assert by_top["gen_data_clean"]["tensor_count"] == 2 + assert by_top["condition_mask"]["bytes"] == 8 + + slim = _slim_prepare_payload_no_raw_state_vision(payload) + assert slim[1].raw_state_vision is None + assert payload[1].raw_state_vision is not None + slim_inventory = _prepare_payload_inventory(slim, signature=("sig",), source="unit") + assert slim_inventory["tensor_count"] == 4 + paths = {item["path"] for item in slim_inventory["largest_tensors"]} + assert "gen_data_clean.raw_state_vision[0]" not in paths + + +def test_cosmos3_edge_official_action_only_warmup_vae_cache_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_warmup_vae_cache, + ) + + argv, enabled = _extract_warmup_vae_cache( + ["-i", "sample.json", "--flashrt-cache-warmup-vae", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + + argv, enabled = _extract_warmup_vae_cache(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is False + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_CACHE_WARMUP_VAE", "1") + argv, enabled = _extract_warmup_vae_cache(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + + +def test_cosmos3_edge_official_action_only_warmup_prepare_cache_arg_parsing(monkeypatch): + from flash_rt.models.cosmos3_edge.action_only_official import ( + _extract_warmup_prepare_cache, + ) + + argv, enabled = _extract_warmup_prepare_cache( + ["-i", "sample.json", "--flashrt-cache-warmup-prepare", "--seed", "0"] + ) + assert argv == ["-i", "sample.json", "--seed", "0"] + assert enabled is True + + argv, enabled = _extract_warmup_prepare_cache(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is False + + monkeypatch.setenv("FLASHRT_COSMOS3_EDGE_CACHE_WARMUP_PREPARE", "1") + argv, enabled = _extract_warmup_prepare_cache(["-i", "sample.json"]) + assert argv == ["-i", "sample.json"] + assert enabled is True + + +def test_cosmos3_edge_upstream_trace_marks_expected_handoff_exception(tmp_path): + from flash_rt.models.cosmos3_edge.action_only_official import _UpstreamTrace + + class _FlashRTVelocityReady(RuntimeError): + pass + + trace = _UpstreamTrace(tmp_path / "trace.json") + with pytest.raises(_FlashRTVelocityReady): + with trace.record("Cosmos3VFMNetwork.forward", expected_exceptions=("_FlashRTVelocityReady",)): + raise _FlashRTVelocityReady("handoff") + + trace.save() + payload = json.loads((tmp_path / "trace.json").read_text(encoding="utf-8")) + assert payload["events"][0]["ok"] is True + assert payload["events"][0]["metadata"]["expected_exception"] == "_FlashRTVelocityReady" + assert payload["summary"][0]["ok_count"] == 1 + + +def test_cosmos3_edge_warmup_vae_cache_signature_includes_content(): + torch = pytest.importorskip("torch") + from flash_rt.models.cosmos3_edge.action_only_official import _WarmupVAECache + + a = torch.zeros(2, 3, dtype=torch.float32) + b = a.clone() + c = a.clone() + c[-1, -1] = 1.0 + + assert _WarmupVAECache._signature(a) == _WarmupVAECache._signature(b) + assert _WarmupVAECache._signature(a) != _WarmupVAECache._signature(c) + + +def test_cosmos3_edge_warmup_prepare_cache_signature_includes_content(): + torch = pytest.importorskip("torch") + from flash_rt.models.cosmos3_edge.action_only_official import _WarmupPrepareCache + + batch_a = {"video": [torch.zeros(2, 3, dtype=torch.float32)], "caption": ["pick"]} + batch_b = {"video": [batch_a["video"][0].clone()], "caption": ["pick"]} + batch_c = {"video": [batch_a["video"][0].clone()], "caption": ["pick"]} + batch_c["video"][0][-1, -1] = 1.0 + + assert _WarmupPrepareCache._signature(batch_a, [0], False) == _WarmupPrepareCache._signature(batch_b, [0], False) + assert _WarmupPrepareCache._signature(batch_a, [0], False) != _WarmupPrepareCache._signature(batch_c, [0], False) + assert _WarmupPrepareCache._signature(batch_a, [0], False) != _WarmupPrepareCache._signature(batch_a, [1], False) + + def test_wan22_infer_exposes_teacache_parameters(): import inspect from flash_rt.frontends.torch.wan22_rtx import Wan22TorchFrontendRtx From fd7f777002eff9bb5690442859c992644c20e8da Mon Sep 17 00:00:00 2001 From: LiangSu8899 Date: Tue, 21 Jul 2026 08:01:06 -0400 Subject: [PATCH 2/3] feat(cosmos3-edge): TeaCache step caching and the v1 Thor performance table Fixed compute-step schedules baked into the whole-denoise CUDA graph: skipped steps reuse the last computed velocity while the native UniPC still advances all 30 steps. --teacache-computes K picks K evenly spaced computed steps (always step 0 and the final step); --teacache-steps takes an explicit list. For AV inverse dynamics the action velocity field is dominated by the static clean-video conditioning, so every swept schedule (15/10/6/4/3/2 computes) stays within quantization-level error of the official action; the 3- and 2-compute tiers also pass on a second-seed dump. Thor, 15 back-to-back iterations, official eager denoise 33.14s, gates cos>=0.999 / rel_l2<3%: fp8, no cache 5.792s (5.72x) rel_l2 0.59% fp8 + TeaCache 15 2.877s (11.52x) rel_l2 0.56% fp8 + TeaCache 6 1.153s (28.74x) rel_l2 0.57% fp8 + TeaCache 3 0.576s (57.56x) rel_l2 0.54% fp8 + TeaCache 2 0.384s (86.36x) rel_l2 0.59% fp8 + fp4-ffn + TC2 0.362s (91.57x) rel_l2 2.34% The 2-compute tiers are an aggressive benchmark operating point, not yet a dataset-level task-success claim. Full table in docs/cosmos3_edge_thor.md. --- benchmarks/cosmos3_edge_thor_denoise.py | 25 ++++++++++ docs/cosmos3_edge_thor.md | 46 ++++++++++++++----- flash_rt/models/cosmos3_edge/denoise_ref.py | 3 ++ flash_rt/models/cosmos3_edge/pipeline_thor.py | 20 +++++++- 4 files changed, 81 insertions(+), 13 deletions(-) diff --git a/benchmarks/cosmos3_edge_thor_denoise.py b/benchmarks/cosmos3_edge_thor_denoise.py index cc46b05c..6364e3bf 100644 --- a/benchmarks/cosmos3_edge_thor_denoise.py +++ b/benchmarks/cosmos3_edge_thor_denoise.py @@ -152,6 +152,18 @@ def main() -> None: ap.add_argument("--bf16-projs", default="", help="comma-separated projections kept in bf16 (fp8 engine)") ap.add_argument("--ffn-fp4", action="store_true", help="NVFP4 W4A4 FFN (up/down) via flash_rt_fp4") ap.add_argument("--no-slim-last", action="store_true", help="disable the M=60 slim last layer") + ap.add_argument( + "--teacache-computes", + type=int, + default=0, + help="TeaCache: number of computed denoise steps, evenly spaced over the schedule " + "(always includes step 0 and the final step); 0 disables", + ) + ap.add_argument( + "--teacache-steps", + default="", + help="TeaCache: explicit comma-separated computed step indices (overrides --teacache-computes)", + ) ap.add_argument("--json-out", default=None) ap.add_argument( "--profile-loop-breakdown", @@ -206,6 +218,16 @@ def main() -> None: if args.device == "cuda": torch.cuda.reset_peak_memory_stats() + teacache_steps: tuple[int, ...] | None = None + if args.teacache_steps: + teacache_steps = tuple(int(v) for v in args.teacache_steps.split(",") if v != "") + elif args.teacache_computes > 1: + n_steps = 30 + k = args.teacache_computes + teacache_steps = tuple(sorted({round(i * (n_steps - 1) / (k - 1)) for i in range(k)})) + if teacache_steps is not None and args.engine == "bf16-eager": + raise SystemExit("TeaCache requires the quantized static engine (--engine fp8/quant-bf16)") + t0 = time.perf_counter() if args.engine == "bf16-eager": runner = EdgeDenoiseFlashRT( @@ -225,6 +247,7 @@ def main() -> None: bf16_projs=tuple(p for p in args.bf16_projs.split(",") if p), ffn_fp4=args.ffn_fp4, slim_last=not args.no_slim_last, + teacache_steps=teacache_steps, use_cuda_graphs=not args.no_quant_graph, ) if args.device == "cuda": @@ -292,6 +315,8 @@ def main() -> None: "engine": args.engine, "bf16_projs": args.bf16_projs, "quant_graph": bool(args.engine != "bf16-eager" and not args.no_quant_graph), + "teacache_steps": list(teacache_steps) if teacache_steps is not None else None, + "teacache_computes": len(teacache_steps) if teacache_steps is not None else 30, "checkpoint": str(checkpoint), "reference_dump": str(reference_dump), "boundary_dump": str(boundary_dump), diff --git a/docs/cosmos3_edge_thor.md b/docs/cosmos3_edge_thor.md index 9856b2c5..0f61c997 100644 --- a/docs/cosmos3_edge_thor.md +++ b/docs/cosmos3_edge_thor.md @@ -62,22 +62,44 @@ selected by the caller. - The last layer runs slim at M=60: only the action rows feed the head, so final-layer Q/attention/o-proj/FFN shrink to the 60 action tokens while K/V still cover the full sequence (math-identical; `--no-slim-last` to disable). - -Measured on Thor (15 back-to-back iterations, hot regime, `av_inverse_0`, -official eager denoise baseline 33.14s; gate cos >= 0.999, rel_l2 < 3%, -both engines re-validated on the second-seed dump): - -| engine | denoise P50 | speedup | final action cos | rel_l2 | -|---|---|---|---|---| -| `--engine bf16-eager` (per-op path) | 11.998s | 2.76x | 0.9999959 | 0.29% | -| `--engine fp8` | 5.792s | 5.72x | 0.9999834 | 0.59% | -| `--engine fp8 --ffn-fp4` | 5.456s | 6.07x | 0.9997067 | 2.42% | +- TeaCache step caching: a fixed compute-step schedule baked into the denoise + graph. Skipped steps reuse the last computed velocity while the native UniPC + still advances all 30 steps. `--teacache-computes K` computes K evenly spaced + steps (always step 0 and step 29); `--teacache-steps` takes an explicit list. + For this AV inverse-dynamics workload the action velocity field is dominated + by the static clean-video conditioning, so even the 2-compute schedule stays + within quantization-level error of the official action. + +## Cosmos3-Edge AV inverse dynamics on Jetson AGX Thor — performance table (v1) + +All latency numbers are direct measurements of the same 30-step denoise +boundary (`av_inverse_0`, 15 back-to-back iterations, hot regime). They do not +include tokenization, VAE preprocessing, or response wrapping. Accuracy gates: +final-action cosine >= 0.999 and relative L2 < 3% against the official output; +the TeaCache 3/2-compute rows and both engines were re-validated on a +second-seed dump. + +| Configuration | Denoise P50 | Speedup vs Official | Cosine | Rel-L2 | Peak mem (GiB) | +|---|---|---|---|---|---| +| Official Cosmos3-Edge (eager) | 33.14 s | 1.00x | reference | reference | — | +| FlashRT FP8, no step cache | 5.792 s | **5.72x** | 0.999983 | 0.59% | 7.85 | +| FlashRT FP8 + NVFP4 FFN, no step cache | 5.456 s | **6.07x** | 0.999707 | 2.42% | 7.54 | +| FlashRT FP8 + TeaCache, 15 compute steps | 2.877 s | **11.52x** | 0.999985 | 0.56% | 7.85 | +| FlashRT FP8 + TeaCache, 6 compute steps | 1.153 s | **28.74x** | 0.999984 | 0.57% | 7.85 | +| FlashRT FP8 + TeaCache, 3 compute steps | 0.576 s | **57.56x** | 0.999986 | 0.54% | 7.85 | +| FlashRT FP8 + TeaCache, 2 compute steps | 0.384 s | **86.36x** | 0.999983 | 0.59% | 7.85 | +| FlashRT FP8 + NVFP4 FFN + TeaCache 2 | 0.362 s | **91.57x** | 0.999727 | 2.34% | 7.54 | + +The 2-compute rows are an aggressive benchmark operating point, not yet a +dataset-level task-success claim. ```bash python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 --iters 15 \ - --warmup-steps 30 --enforce-gates # FP8 default + --warmup-steps 30 --enforce-gates # FP8, no cache python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 --ffn-fp4 \ - --iters 15 --warmup-steps 30 --enforce-gates # max-quantization variant + --iters 15 --warmup-steps 30 --enforce-gates # max quantization +python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 \ + --teacache-computes 3 --iters 15 --warmup-steps 30 --enforce-gates ``` FP8 is the recommended default (large precision margin); `--ffn-fp4` trades diff --git a/flash_rt/models/cosmos3_edge/denoise_ref.py b/flash_rt/models/cosmos3_edge/denoise_ref.py index 4f8b1825..270fc99f 100644 --- a/flash_rt/models/cosmos3_edge/denoise_ref.py +++ b/flash_rt/models/cosmos3_edge/denoise_ref.py @@ -174,6 +174,7 @@ def __init__( bf16_projs: tuple[str, ...] = (), ffn_fp4: bool = False, slim_last: bool = True, + teacache_steps: tuple[int, ...] | None = None, use_cuda_graphs: bool = True, ): from flash_rt.models.cosmos3_edge.pipeline_thor import CosmosEdgeThor @@ -199,6 +200,8 @@ def __init__( self.static_engine = None self.static_scheduler = self.engine.unipc self.engine.calibrate(self.denoise_dump.step_noise(0)) + if teacache_steps is not None: + self.engine.set_teacache(teacache_steps) if self.use_cuda_graphs: self.engine.capture(warmup_noise=self.denoise_dump.step_noise(0)) diff --git a/flash_rt/models/cosmos3_edge/pipeline_thor.py b/flash_rt/models/cosmos3_edge/pipeline_thor.py index ab156a5a..9593f77b 100644 --- a/flash_rt/models/cosmos3_edge/pipeline_thor.py +++ b/flash_rt/models/cosmos3_edge/pipeline_thor.py @@ -235,6 +235,7 @@ def qf8(w_nk: torch.Tensor): # per-tensor FP8 E4M3; GEMM wants B as [K,N] self.cos_f = t["s00/layers/00/rope/cos/full_only_seq"].to(device=DEV, dtype=BF).contiguous() self.sin_f = t["s00/layers/00/rope/sin/full_only_seq"].to(device=DEV, dtype=BF).contiguous() + self.compute_steps: frozenset[int] | None = None self.unipc = EdgeStaticUniPCScheduler(self.num_steps, device=torch.device(DEV), shift=shift) if not self.unipc.native_available: raise RuntimeError("native UniPC step binding is required") @@ -455,6 +456,20 @@ def forward_step(self, step: int, latent: torch.Tensor) -> torch.Tensor: self.action_out.data_ptr(), self.velocity.data_ptr(), self.velocity.numel(), NA * AD, s) return self.velocity + def set_teacache(self, compute_steps: tuple[int, ...] | list[int] | None) -> None: + """Fixed compute-step schedule: skipped steps reuse the last computed + velocity (the static velocity buffer) while UniPC still advances every + step. Step 0 must always compute. Call before ``capture``.""" + if self.graph is not None: + raise RuntimeError("set_teacache must be called before graph capture") + if compute_steps is None: + self.compute_steps = None + return + steps = sorted(set(int(v) for v in compute_steps)) + if not steps or steps[0] != 0 or steps[-1] >= self.num_steps: + raise ValueError(f"invalid TeaCache compute schedule: {steps}") + self.compute_steps = frozenset(steps) + # ---- whole-denoise loop: forward + native UniPC per step ---- def run_loop(self) -> torch.Tensor: # UniPC state buffers are allocated once; their contents are fully @@ -462,7 +477,10 @@ def run_loop(self) -> torch.Tensor: if self.unipc.prev_m1 is None: self.unipc.reset(self.latent) for step in range(self.num_steps): - velocity = self.forward_step(step, self.latent) + if self.compute_steps is None or step in self.compute_steps: + velocity = self.forward_step(step, self.latent) + else: + velocity = self.velocity self.unipc.step(self.latent, velocity, step) return self.latent From e110b448d464c61a99ec9fbc67fa87e4d718dc9a Mon Sep 17 00:00:00 2001 From: LiangSu8899 Date: Tue, 21 Jul 2026 08:31:36 -0400 Subject: [PATCH 3/3] docs(cosmos3-edge): publish Thor usage and benchmark guidance --- README.md | 18 +- USAGE.md | 33 +- benchmarks/cosmos3_edge_thor_denoise.py | 20 +- docs/cosmos3_edge_thor.md | 1505 +++----------------- docs/stable_api.md | 11 +- tests/test_cosmos3_edge_boundary_dump.py | 65 +- tests/test_cosmos3_edge_dump_replay.py | 33 +- tests/test_cosmos3_edge_layer_ref.py | 24 +- tests/test_cosmos3_edge_step0_reference.py | 40 +- tests/test_cosmos3_edge_thor_spec.py | 6 +- tests/test_cosmos3_edge_weights.py | 11 +- 11 files changed, 289 insertions(+), 1477 deletions(-) diff --git a/README.md b/README.md index 6f84eaa6..bac951b2 100644 --- a/README.md +++ b/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). | diff --git a/USAGE.md b/USAGE.md index dbc5324a..4bd0e28d 100644 --- a/USAGE.md +++ b/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 diff --git a/benchmarks/cosmos3_edge_thor_denoise.py b/benchmarks/cosmos3_edge_thor_denoise.py index 6364e3bf..b913eece 100644 --- a/benchmarks/cosmos3_edge_thor_denoise.py +++ b/benchmarks/cosmos3_edge_thor_denoise.py @@ -6,15 +6,13 @@ official ``benchmark.json`` only to print the recorded eager baseline next to the FlashRT numbers. -Run inside the Thor container and isolated venv: +Run inside a Thor environment with the FlashRT extensions available: - cd /work/official/flashrt-public - source /work/.venv_cosmos_thor/bin/activate python benchmarks/cosmos3_edge_thor_denoise.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --reference-dump dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors \ - --boundary-dump dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors \ - --official-benchmark dev_scratch_cosmos3_thor/edge_av_inverse_0/official_outputs/benchmark.json + --checkpoint /path/to/Cosmos3-Edge \ + --reference-dump /path/to/denoise/tensors.safetensors \ + --boundary-dump /path/to/step0_boundary/tensors.safetensors \ + --official-benchmark /path/to/official/benchmark.json """ from __future__ import annotations @@ -121,18 +119,18 @@ def _profile_loop_breakdown(runner: Any, *, device: str) -> dict[str, Any]: def main() -> None: ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--checkpoint", default="/work/models/Cosmos3-Edge") + ap.add_argument("--checkpoint", required=True) ap.add_argument( "--reference-dump", - default="dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors", + required=True, ) ap.add_argument( "--boundary-dump", - default="dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors", + required=True, ) ap.add_argument( "--official-benchmark", - default="dev_scratch_cosmos3_thor/edge_av_inverse_0/official_outputs/benchmark.json", + default=None, ) ap.add_argument("--device", default="cuda") ap.add_argument("--warmup-steps", type=int, default=1) diff --git a/docs/cosmos3_edge_thor.md b/docs/cosmos3_edge_thor.md index 0f61c997..587888da 100644 --- a/docs/cosmos3_edge_thor.md +++ b/docs/cosmos3_edge_thor.md @@ -1,1392 +1,221 @@ -# Cosmos3-Edge Thor Baseline +# Cosmos3-Edge AV on Jetson AGX Thor -`config="cosmos3_edge"` is the Thor baseline runner for NVIDIA's official -Cosmos Framework inference path. It is intentionally upstream-first: run this -baseline, capture latency/outputs, then port the FlashRT optimized denoise/action -path behind the same config. +FlashRT provides an optimized 30-step denoise engine for the Cosmos3-Edge AV +inverse-dynamics policy on Jetson AGX Thor (`sm_110`). The integration keeps +the official Cosmos preprocessing and action contract while replacing the +denoise boundary with a static, quantized FlashRT graph. -Current status: +This page covers the supported path, measured performance, usage, and +correctness requirements. The latency table measures denoise only. It does not +include tokenization, video preprocessing, VAE encode, or response handling. -- Official Cosmos Framework baseline is runnable from FlashRT. -- P0 denoise dump is captured for `av_inverse_0`. -- P1 scheduler/latent replay scaffold is wired as `backend="replay"`. -- P1 Torch reference validates the native transformer math path through all 28 - layers for step 0 action velocity, and through the full 30-step UniPC denoise - loop as `backend="torch_ref"`. -- `backend="flashrt"` is wired to the current optimized eager denoise engine: - static AV buffers, static und K/V cache, BF16 FlashRT GEMMs/RMSNorm/relu2, - fused qk-norm+RoPE, cached timestep embeddings, native action input tail copy, - native flat velocity fill, native action bias+timestep add, native action - bias+tail clear, native action row scatter/gather, and FA4 native gen - attention. -- `backend="official_action_only"` runs the official no-dump Cosmos Framework - path, but skips inverse-dynamics vision decode/save and rewrites - `sample_outputs.json` to action-only. -- `official_action_only` also has an opt-in `live_dump_out` capture mode that - records a denoise safetensors dump from the live official run without editing - the external Cosmos Framework checkout. -- `official_action_only` has an opt-in live FlashRT handoff prototype: step 0 - runs official `_get_velocity` to capture the live boundary, then later - denoise steps short-circuit `_get_velocity` through FlashRT velocity. -- `--live-prelayer-bootstrap` is a stronger handoff mode: it captures the live - boundary before official decoder layers, aborts the official transformer, and - returns FlashRT velocity for all 30 denoise steps. -- `--live-warm-request` is the measured request path for the live handoff: it - enables pre-layer bootstrap and one official warmup request so cold FA4/engine - setup is paid before the timed batch. -- In live handoff mode, the Cosmos UniPC sampler is patched in-process to use - FlashRT's native `cosmos3_edge_unipc_step_f32_bf16` path for the fixed - single-sample Edge AV schedule when available. Set - `FLASHRT_COSMOS3_EDGE_LIVE_NATIVE_UNIPC=0` to fall back to the official - Python scheduler. +## Performance -## Quantized whole-graph denoise engine (`models/cosmos3_edge/pipeline_thor.py`) +The following results use the same `av_inverse_0` 30-step denoise boundary and +15 back-to-back hot iterations. Accuracy is measured against the official +final action with the gates `cosine >= 0.999` and `relative L2 < 3%`. Both +non-cached engines and the 3- and 2-compute TeaCache schedules were also +checked with a second seed. -`CosmosEdgeThor` is the optimized denoise engine: pre-quantized gen-tower -weights, static buffers, per-layer joint und+gen K/V with the static und prefix -pre-installed, FA4 attention, native UniPC, and the entire 30-step denoise -captured in **one CUDA graph**. It reads no environment variables; precision is -selected by the caller. - -- `quant="fp8"` (default): per-tensor FP8 E4M3 weights for all six gen-tower - projections, calibrated static activation scales (one dynamic-scale denoise - pass records per-site amax ceilings via `fp8_accumulate_scale_max`, margin - 1.25), fused `residual_add_rms_norm_fp8` / `quantize_fp8_static` / - `cosmos3_edge_relu2_to_fp8_static_bf16` quant chain. -- `ffn_fp4=True` (opt-in): NVFP4 W4A4 up/down FFN GEMMs via `flash_rt_fp4` - (`cutlass_fp4_sq_fp16` + fused `cosmos3_edge_res_rms_fp4_sfa_bf16` and - `cosmos3_edge_relu2_fp4_sfa_fp16` quantizers, dynamic per-16-block scales). -- The gen-stream fused qk-norm+RoPE kernel is the warp-per-head register - variant (one warp per `(row, head)` vector, `__shfl_xor` rope-partner - exchange, no shared memory or barriers). -- The last layer runs slim at M=60: only the action rows feed the head, so - final-layer Q/attention/o-proj/FFN shrink to the 60 action tokens while K/V - still cover the full sequence (math-identical; `--no-slim-last` to disable). -- TeaCache step caching: a fixed compute-step schedule baked into the denoise - graph. Skipped steps reuse the last computed velocity while the native UniPC - still advances all 30 steps. `--teacache-computes K` computes K evenly spaced - steps (always step 0 and step 29); `--teacache-steps` takes an explicit list. - For this AV inverse-dynamics workload the action velocity field is dominated - by the static clean-video conditioning, so even the 2-compute schedule stays - within quantization-level error of the official action. - -## Cosmos3-Edge AV inverse dynamics on Jetson AGX Thor — performance table (v1) - -All latency numbers are direct measurements of the same 30-step denoise -boundary (`av_inverse_0`, 15 back-to-back iterations, hot regime). They do not -include tokenization, VAE preprocessing, or response wrapping. Accuracy gates: -final-action cosine >= 0.999 and relative L2 < 3% against the official output; -the TeaCache 3/2-compute rows and both engines were re-validated on a -second-seed dump. - -| Configuration | Denoise P50 | Speedup vs Official | Cosine | Rel-L2 | Peak mem (GiB) | -|---|---|---|---|---|---| -| Official Cosmos3-Edge (eager) | 33.14 s | 1.00x | reference | reference | — | +| Configuration | Denoise P50 | Speedup vs Official | Cosine | Rel-L2 | Peak memory (GiB) | +|---|---:|---:|---:|---:|---:| +| Official Cosmos3-Edge (eager) | 33.14 s | 1.00x | reference | reference | not recorded | | FlashRT FP8, no step cache | 5.792 s | **5.72x** | 0.999983 | 0.59% | 7.85 | | FlashRT FP8 + NVFP4 FFN, no step cache | 5.456 s | **6.07x** | 0.999707 | 2.42% | 7.54 | -| FlashRT FP8 + TeaCache, 15 compute steps | 2.877 s | **11.52x** | 0.999985 | 0.56% | 7.85 | -| FlashRT FP8 + TeaCache, 6 compute steps | 1.153 s | **28.74x** | 0.999984 | 0.57% | 7.85 | -| FlashRT FP8 + TeaCache, 3 compute steps | 0.576 s | **57.56x** | 0.999986 | 0.54% | 7.85 | -| FlashRT FP8 + TeaCache, 2 compute steps | 0.384 s | **86.36x** | 0.999983 | 0.59% | 7.85 | -| FlashRT FP8 + NVFP4 FFN + TeaCache 2 | 0.362 s | **91.57x** | 0.999727 | 2.34% | 7.54 | - -The 2-compute rows are an aggressive benchmark operating point, not yet a -dataset-level task-success claim. - -```bash -python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 --iters 15 \ - --warmup-steps 30 --enforce-gates # FP8, no cache -python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 --ffn-fp4 \ - --iters 15 --warmup-steps 30 --enforce-gates # max quantization -python benchmarks/cosmos3_edge_thor_denoise.py --engine fp8 \ - --teacache-computes 3 --iters 15 --warmup-steps 30 --enforce-gates -``` - -FP8 is the recommended default (large precision margin); `--ffn-fp4` trades -~6% latency for a still-passing but thinner gate margin. Measured negatives, -kept default-off or not landed: a fused QKV wide GEMM (cuBLASLt N=4096 tactic -regression, behind `qkv_fused=`) and SageAttention2 int8 gen attention (the -SM80-era int8 mma core is both numerically wrong and slower than FA4 on -SM110). The largest remaining denoise levers are a relu2 epilogue fused into -the up GEMM and a Thor-native int8/fp8 attention kernel. - -## Quickstart - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_flashrt \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework -``` - -The example must run inside the same container/venv used for the official -baseline, with `cosmos-framework` importable. `--vae-path` is optional only when -the container can download `Wan2.2_VAE.pth` itself. - -To validate the no-dump action-only official path without decoded vision output: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_action_only_official \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only -``` - -To capture a live official denoise dump while keeping the action-only output: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_dump \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-dump-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_denoise.safetensors -``` +| FlashRT FP8 + TeaCache, 15 computes | 2.877 s | **11.52x** | 0.999985 | 0.56% | 7.85 | +| FlashRT FP8 + TeaCache, 6 computes | 1.153 s | **28.74x** | 0.999984 | 0.57% | 7.85 | +| FlashRT FP8 + TeaCache, 3 computes | 0.576 s | **57.56x** | 0.999986 | 0.54% | 7.85 | +| FlashRT FP8 + TeaCache, 2 computes | 0.384 s | **86.36x** | 0.999983 | 0.59% | 7.85 | +| FlashRT FP8 + NVFP4 FFN + TeaCache, 2 computes | **0.362 s** | **91.57x** | 0.999727 | 2.34% | 7.54 | -This capture mode is for boundary validation, not baseline timing: enabling the -official velocity postprocess hook makes Cosmos run an extra uncond velocity -branch even when guidance is 1.0. +The production optimization claim is the no-cache result: **5.72x with FP8** +or **6.07x with FP8 plus NVFP4 FFN**. Those rows execute all 30 velocity +evaluations and isolate the benefit of the FlashRT engine from step caching. -To exercise the live FlashRT handoff prototype: +TeaCache is an optional approximation and its usable schedule is workload +dependent. The values above are benchmark observations for two validated +samples, not a dataset-level task-success claim. Camera placement, robot +dynamics, task distribution, conditioning strength, and user fine-tuning can +all change how quickly the velocity field varies between steps. Qualify each +TeaCache schedule against the target checkpoint and deployment data before +using it. The 2-compute rows are aggressive benchmark operating points. -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_handoff \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-flashrt-handoff \ - --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_handoff_boundary.safetensors -``` +For the measured workload, **FP8 + TeaCache with 3 computes** is the suggested +starting point: 0.576 s, 57.56x over official eager denoise, and 0.54% relative +L2. Start with more computes for a new task and reduce the count only after +task-level validation. -This is a correctness/architecture prototype. It still pays official step 0, -FlashRT engine initialization, and cold FA4 setup inside -`generate_samples_from_batch`; the current archived run reports 26.10s -`generate_batch`. `--live-boundary-out` is only for debug artifact archival; the -runtime can construct the boundary in memory, which measured 26.20s and the same -action metrics. +## Optimization structure -To skip the official step-0 transformer as well: +The no-cache 6x-class path combines several changes rather than relying on one +kernel: -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_trace.json -``` - -The archived pre-layer run proves the official decoder was aborted -(`prelayer_aborted=true`) and returns action within the gate, but one-shot -latency is still dominated by FA4 first-use compilation: first FlashRT velocity -8.50s, steady FlashRT velocity 0.398s/step. - -To archive the smaller native-denoise/VLM input contract at the pre-layer -boundary: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_boundary \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary_upstream_trace.json -``` +- Static AV buffers and a pre-installed und-conditioned K/V prefix remove + repeated allocation and invariant work. +- FP8 E4M3 weights and activations cover the six generated-stream projections. + Activation scales are calibrated once with an amax margin. +- The optional NVFP4 path replaces the FFN up/down projections with W4A4 + CUTLASS kernels. +- Fused residual, RMSNorm, quantization, ReLU2, Q/K normalization, and RoPE + kernels reduce launches and intermediate traffic. +- FA4 handles generated-stream attention and a native UniPC kernel advances + the scheduler. +- The final layer computes only the 60 action rows while retaining full K/V + context. +- All 30 denoise steps are captured in one CUDA Graph. -Archived `official_action_only_live_prelayer_boundary.safetensors` metrics: +TeaCache adds a fixed compute schedule to that graph. A skipped denoise step +still runs native UniPC but reuses the most recent velocity buffer. The schedule +contains no runtime branch and requires no additional kernel. -- 60.36 MiB file, 31 tensors / 60.36 MiB tensor footprint -- contains `lm_in/full_only_seq`, layer-0 input, RoPE, VFM tokens, action - tokens, timestep, and `noise_x` -- does not store layer-0 output or step-0 velocity; the older step-0 handoff - boundary was 87.77 MiB / 36 tensors because it archived those outputs -- run action gate: cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 -- `generate_batch`: 26.70s; native UniPC scheduler: 30 steps +## Requirements -This is the input contract for a future native VLM/prefill-to-denoise engine. -It still relies on official upstream code to synthesize `lm_in/full_only_seq`, -RoPE, VFM tokens, and action tokens from the raw request. +- Jetson AGX Thor with CUDA compute capability 11.0. +- A FlashRT build containing the Thor kernels and FA4 backend. +- The public Cosmos3-Edge checkpoint. +- An importable NVIDIA Cosmos Framework checkout for the official baseline and + live preprocessing path. +- The Wan2.2 VAE checkpoint used by Cosmos3-Edge. -To replay that pre-layer boundary as the live denoise engine input: +Use arguments or environment-specific configuration to point at these assets; +FlashRT does not assume a host directory layout. ```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_boundary_in \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-boundary-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_upstream_trace.json +export COSMOS_EDGE_CHECKPOINT=/path/to/Cosmos3-Edge +export COSMOS_EDGE_INPUT=/path/to/av_inverse_input.json +export COSMOS_FRAMEWORK_ROOT=/path/to/cosmos-framework +export WAN_VAE_CHECKPOINT=/path/to/Wan2.2_VAE.pth +export COSMOS_EDGE_OUTPUT=/path/to/output ``` -Archived `official_action_only_live_boundary_in_outputs` metrics: - -- `OmniInference.generate_batch`: 26.03s -- official velocity calls: 0; FlashRT velocity calls: 30 -- boundary load/validate: 0.8ms; engine construct: 1.29s -- native UniPC scheduler: 1 run / 30 steps, scheduler step total 0.0071s -- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 - -This is still artifact-backed. Its purpose is to make the native prefill target -executable: once native VAE/SigLIP/VLM code can synthesize the same boundary -tensors live, the denoise engine no longer needs official transformer capture. - -The boundary-in path can also be combined with the current warm repeated-request -prepare cache: +Build the extensions for Thor using the normal repository build flow: ```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_boundary_in_warm_prepare_cache \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-boundary-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors \ - --live-warm-request \ - --cache-warmup-prepare \ - --prepare-slim-no-raw-state-vision \ - --prepare-slim-derive-condition-reference \ - --prepare-slim-derive-initial-noise \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_warm_request_prepare_cache_slim_derive_noise_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_boundary_in_warm_request_prepare_cache_slim_derive_noise_upstream_trace.json +cmake -B build -S . -DGPU_ARCH=110 +cmake --build build -j"$(nproc)" --target flash_rt_kernels flash_rt_fp4 ``` -Archived metrics for this service-style path: - -- measured `OmniInference.generate_batch`: 12.08s; denoise: 12.07s -- warmup `OmniInference.generate_batch`: 26.08s -- official velocity calls: 0; FlashRT velocity calls: 60 -- measured `_prepare_inference_data`: 0.042s; prepare hit/store: 1/1 -- measured VAE encode calls: 0 -- native UniPC scheduler: 2 runs / 60 steps, scheduler step total 0.0134s -- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 - -The helper `_derive_step0_vfm_boundary_from_prepare_payload` formalizes the -part of the pre-layer boundary that is already implied by the slim prepare -contract. From `official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt` -and seed 0 plus `embed_tokens.weight` it derives 28 step-0 -VFM/noise/pack/position/RoPE/text-causal tensors exactly: - -- `steps/00/noise_x` -- VFM vision tokens, condition mask, sequence indexes, and empty no-loss/noisy - metadata -- VFM action noisy tokens, condition mask, sequence/timestep/loss/noisy indexes, - domain id, and raw action dim -- fixed packed metadata: causal/full indices, offsets, sample offsets, and - `position_ids` -- layer-0 RoPE cos/sin tensors for causal and full-only sequences -- causal LM hidden state from `cond_text_tokens + eos + <|vision_start|>` and - `embed_tokens.weight` - -Verifier metrics: 28 tensors / 10.65 MiB derived exactly. The layer-0 input -tensors are checked as aliases of `lm_in` (2 tensors / 25.10 MiB). The remaining -1 tensor / 24.61 MiB is the full-only LM packed hidden state. The verifier now -probes that tensor too and requires all 6300 rows to be bit-exact; the action -rows require the reference action projection/timestep path to run with CUDA -matmul TF32 disabled. This narrows the remaining native VLM/prefill task to -synthesizing the same prepare/prefill state from live raw request data rather -than resolving an unresolved packed-hidden-state formula gap. - -To run the live denoise engine from the 9.18 MiB slim prepare artifact instead -of the 60.36 MiB pre-layer boundary artifact: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_prepare_boundary_in \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --prepare-replay-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ - --live-boundary-prepare-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ - --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_boundary_in_derived_boundary.safetensors \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_boundary_in_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_boundary_in_upstream_trace.json -``` - -Archived `official_action_only_live_prepare_boundary_in_outputs` metrics: - -- input artifact: 9.18 MiB slim-derived-noise prepare payload -- derived debug boundary: 62.65 MiB / 31 tensors; larger than the 60.36 MiB - pre-layer artifact only because the debug save clones shared tensor views -- boundary derive: 0.258s; engine construct: 1.074s -- `OmniInference.generate_batch`: 22.24s; denoise: 22.22s; decode: 0.0084s -- official velocity calls: 0; FlashRT velocity calls: 30 -- measured `_prepare_inference_data`: 0.101s; measured VAE encode calls: 0 -- native UniPC scheduler: 1 run / 30 steps -- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 - -This removes the 60.36 MiB executable boundary artifact from the runtime input -contract for replay-style testing. It is still artifact-backed prepare replay: -raw-video no-artifact live inference still needs native VAE/SigLIP/VLM prefill. - -To run the same executable-boundary synthesis from this request's in-memory -prepared state, without reading a prepare artifact: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-boundary-prepare-live \ - --live-boundary-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_derived_boundary.safetensors \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_upstream_trace.json \ - --prepare-slim-no-raw-state-vision \ - --prepare-slim-derive-condition-reference \ - --prepare-slim-derive-initial-noise -``` - -Archived `official_action_only_live_prepare_live_outputs` metrics: - -- runtime source: in-memory `_prepare_inference_data` payload, not a prepare or - pre-layer artifact -- derived debug boundary: 62.65 MiB / 31 tensors; debug save clones shared views -- boundary derive: 0.143s; engine construct: 1.267s -- `OmniInference.generate_batch`: 25.93s; denoise: 25.91s; decode: 0.0088s -- official velocity calls: 0; FlashRT velocity calls: 30 -- measured `_prepare_inference_data`: 3.88s; measured VAE encode: 3.66s; text tokenization: 30.8ms; pack: 2.65ms -- native UniPC scheduler: 1 run / 30 steps -- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 - -This is the first no-prepare-artifact live denoise handoff. It still uses the -official raw-video prepare path, so strict native E2E still requires native -VAE/SigLIP/VLM prefill to produce the same prepared state. - -The same bridge can be combined with upstream Wan2.2 AOT encode using -`--vae-compile-encode`. This is not a FlashRT native VAE implementation, but it -is the current no-prepare-artifact service-warmup target line for native VAE -work: - -- `official_action_only_live_prepare_live_vae_compile_encode_outputs` -- AOT setup compile/load: 149.39s; loaded AOT functions: 5 -- measured `OmniInference.generate_batch`: 24.33s; denoise: 24.32s; decode: - 0.0085s -- measured `_prepare_inference_data`: 2.65s; measured VAE encode: 2.43s -- boundary derive: 0.146s; engine construct: 1.714s -- official velocity calls: 0; FlashRT velocity calls: 30; native UniPC 30 steps -- final action cos 0.9999948 / rel_l2 0.322% / max_abs 0.00802 - -For the current no-prepare/pre-layer-artifact service floor, combine the live -prepare bridge with one warmup request plus the VAE and prepare caches, and keep -`--live-boundary-out` unset: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_warm_cache_nosave_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --warmup 1 \ - --cache-warmup-vae \ - --cache-warmup-prepare \ - --live-boundary-prepare-live \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_warm_cache_nosave_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prepare_live_warm_cache_nosave_upstream_trace.json \ - --prepare-slim-no-raw-state-vision \ - --prepare-slim-derive-condition-reference \ - --prepare-slim-derive-initial-noise -``` - -Archived `official_action_only_live_prepare_live_warm_cache_nosave_outputs` -metrics: - -- measured `OmniInference.generate_batch`: 13.33s; denoise: 13.32s; decode: - 0.00020s -- warmup `OmniInference.generate_batch`: 25.77s -- measured `_prepare_inference_data`: 0.075s; measured VAE encode calls: 0 -- official velocity calls: 0; FlashRT velocity calls: 60; native UniPC 60 steps -- measured native UniPC run: 13.23s -- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 - -A debug-save variant with `--live-boundary-out` measured 22.04s because the -archival path perturbs the service run, so production/service floor measurements -should leave that flag off. - -The `--live-warm-request` path moves cold FA4/engine setup and repeat VAE encode -out of the measured batch for the same sample by enabling pre-layer bootstrap, -one official warmup request, and a FlashRT warmup VAE latent cache. FlashRT -patches warmup to clone `data_batch` first, because upstream mutates video -tensors in-place during preprocessing. - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_warm_request_vae_cache \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-warm-request \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_vae_cache_handoff_trace.json -``` - -Archived `official_action_only_live_warm_request_vae_cache_outputs` metrics: - -- `OmniInference.generate_batch`: 12.12s -- `OmniMoTModel.generate_samples_from_batch`: 12.11s -- `OmniMoTModel.decode`: 0.00012s -- live native UniPC scheduler: 2 runs / 60 steps, native scheduler step total - 0.0128s -- final action cos 0.9999958 / rel_l2 0.291% / max_abs 0.00687 - -The warmup batch still pays cold cost: `[warmup] OmniInference.generate_batch` -25.95s. The VAE cache key includes shape, dtype, device, element count, and a -small sampled content fingerprint so same-shape different-video tensors do not -reuse the warmup latent silently. - -For P4 upstream breakdown, add `--upstream-trace-out`: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_warm_request_vae_cache \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-warm-request \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_vae_cache_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_vae_cache_upstream_trace.json -``` - -Archived uncached upstream trace metrics: - -- measured `OmniMoTModel.encode`: 3.41s -- measured `OmniMoTModel.get_data_and_condition`: 3.49s -- measured `OmniMoTModel._prepare_inference_data`: 3.52s -- input batch preparation: 0.57s total -- warmup official `Cosmos3VFMNetwork.forward`: 9.79s - -With warmup VAE cache enabled by `--live-warm-request`, the measured upstream -trace changes to: - -- measured `OmniMoTModel.encode`: 0.00051s -- measured `OmniMoTModel._prepare_inference_data`: 0.1068s -- measured cache hits: 1, warmup cache stores: 1, measured cache misses: 0 -- warmup `OmniMoTModel.encode`: 3.65s -- warmup prelayer handoff trace events are marked as expected/ok - -For the no-warmup VAE replacement boundary, dump the measured VAE encode latent -and replay it: - -```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_vae_boundary \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --vae-encode-dump-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_encode.safetensors \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_boundary_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_boundary_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_vae_latent_replay \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --vae-latent-in dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_encode.safetensors \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_latent_replay_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_latent_replay_upstream_trace.json -``` - -Archived VAE boundary/replay metrics: - -- boundary dump: `official_action_only_live_prelayer_vae_encode.safetensors`, - 4.57 MiB, tensor `vae_encode/output` shape `[1,48,16,30,52]`, plus input - signature metadata -- boundary run measured `OmniMoTModel.encode`: 3.72s; - `_prepare_inference_data`: 3.86s; `generate_batch`: 26.22s -- latent replay measured `OmniMoTModel.encode`: 0.0108s; - `vae_latent_replay_hit`: 0.0050s; `_prepare_inference_data`: 0.141s; - `generate_batch`: 22.55s -- latent replay final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 +## Official baseline -This is not the final no-dump VAE path. It is the correctness and latency -boundary a native Wan2.2 VAE encoder must match to remove the remaining -no-warmup official encode. - -For a wider no-warmup prepare/prefill replacement boundary, dump and replay the -full `_prepare_inference_data` return tuple: +Run the official action-only path first. This establishes the output and +end-to-end baseline without decoding or saving generated video: ```bash python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_boundary_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare.pt \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_boundary_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_boundary_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_replay_v2_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare.pt \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_replay_v2_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_replay_v2_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory_replay_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory.json \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_inventory_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_inventory.json \ - --prepare-slim-no-raw-state-vision \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_inventory.json \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_dump_replay_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_inventory.json \ - --prepare-slim-no-raw-state-vision \ - --prepare-slim-derive-condition-reference \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_inventory.json \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_dump_replay_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-dump-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_inventory.json \ - --prepare-slim-no-raw-state-vision \ - --prepare-slim-derive-condition-reference \ - --prepare-slim-derive-initial-noise \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_upstream_trace.json - -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_outputs \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --prepare-replay-in /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt \ - --prepare-inventory-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_inventory.json \ - --live-handoff-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_handoff_trace.json \ - --upstream-trace-out /work/official/flashrt-public/dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump_replay_upstream_trace.json + --checkpoint "$COSMOS_EDGE_CHECKPOINT" \ + --input-json "$COSMOS_EDGE_INPUT" \ + --output-dir "$COSMOS_EDGE_OUTPUT/official" \ + --vae-path "$WAN_VAE_CHECKPOINT" \ + --cosmos-root "$COSMOS_FRAMEWORK_ROOT" \ + --backend official_action_only ``` -Archived prepare boundary/replay metrics: - -- prepare artifact: `official_action_only_live_prelayer_prepare.pt`, 297 MiB, - with input signature and the 7-tuple consumed by `generate_samples_from_batch` -- dump run measured `_prepare_inference_data`: 4.40s; `prepare_boundary_dump` - event: 3.87s -- replay measured `_prepare_inference_data`: 0.232s; `prepare_replay_hit`: - 0.182s; `generate_batch`: 22.66s -- replay final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 -- inventory replay writes - `official_action_only_live_prelayer_prepare_inventory.json` using schema - `flashrt_cosmos3_edge_prepare_inventory_v1`: 11 tensors, 297.13 MiB total. - The largest tensor is `gen_data_clean.raw_state_vision[0]` - `[1,3,61,480,832]` at 278.79 MiB; the denoise seed/reference/mask vectors - are `initial_noise[0]`, `condition_reference[0]`, and `condition_mask[0]`, - each `[1201920]` float32 at 4.58 MiB. This is the current machine-readable - native prefill contract. -- inventory replay measured `_prepare_inference_data`: 0.207s; - `prepare_replay_hit`: 0.159s; `generate_batch`: 22.75s; final action - cos 0.9999958 / rel_l2 0.291% / max_abs 0.00687. -- `--prepare-slim-no-raw-state-vision` proves the post-prepare denoise path - does not need the 278.79 MiB RGB `raw_state_vision` tensor once - `temporal_positions_vision` and `x0_tokens_vision` are present. The archived - slim dump is 18.35 MiB, 10 tensors / 18.34 MiB total, with `x0_tokens_vision`, - action state, and the three 4.58 MiB noise/reference/mask vectors intact. - The slim dump run keeps action cos 0.9999958 / rel_l2 0.291%; replaying the - slim `.pt` gives `prepare_replay_hit` 0.026s, `_prepare_inference_data` - 0.075s, `generate_batch` 22.62s, and unchanged action metrics. -- `--prepare-slim-derive-condition-reference` proves `condition_reference` is - exactly derivable from `gen_data_clean.x0_tokens_vision/action` and - `raw_action_dim` after raw video has been dropped. The derived slim artifact - stores `condition_reference=None`, is 13.76 MiB, and has 9 tensors / - 13.75 MiB. Replay restores `condition_reference` before entering the official - sampler, so the runtime payload inventory returns to 10 tensors / 18.34 MiB. - The archived replay has `prepare_replay_hit` 0.026s, - `_prepare_inference_data` 0.074s, `generate_batch` 22.56s, final action cos - 0.9999959 / rel_l2 0.291% / max_abs 0.00687, and native UniPC 1 run / - 30 steps. -- `--prepare-slim-derive-initial-noise` further proves `initial_noise` is - bit-exactly derivable from the request seed, `x0_tokens_vision/action`, - `condition_mask`, and `raw_action_dim`. The slim-derived-noise artifact stores - `raw_state_vision=None`, `condition_reference=None`, and `initial_noise=None`; - it is 9.18 MiB with 8 tensors / 9.17 MiB, consisting of the prepared latent - state plus `condition_mask`. Replay restores both `initial_noise` and - `condition_reference`, so runtime inventory returns to 10 tensors / - 18.34 MiB. The archived replay has `prepare_replay_hit` 0.051s, - `_prepare_inference_data` 0.100s, `generate_batch` 22.59s, final action cos - 0.9999959 / rel_l2 0.291% / max_abs 0.00687, and native UniPC 1 run / - 30 steps. +The full official backend remains available with `--backend official` when +decoded vision output is needed. -This is a full prepare/prefill correctness boundary for future native -VAE/SigLIP/VLM-prefill replacement. It is not a production path: the artifact is -input-specific and still loaded via `torch.save`/`torch.load`. The derived slim -result shows the native denoise contract can be reduced to latent/action state, -temporal positions, condition mask, request seed/signature metadata, and enough -metadata to reconstruct condition references and initial noise. +## FlashRT denoise benchmark -For the first native Wan2.2 VAE encode subgraph, enable native BF16 -`RMS_norm -> SiLU` inside encode `ResidualBlock`s: +The optimized benchmark consumes the captured denoise and step-0 boundary +fixtures. All asset paths are explicit so a run cannot silently use a stale +local artifact. ```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_prelayer_vae_native_rms_silu \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-prelayer-bootstrap \ - --vae-native-rms-silu \ - --vae-encode-profile-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_native_rms_silu_profile.json \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_native_rms_silu_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_vae_native_rms_silu_upstream_trace.json +export COSMOS_EDGE_REFERENCE=/path/to/denoise/tensors.safetensors +export COSMOS_EDGE_BOUNDARY=/path/to/step0_boundary/tensors.safetensors +export COSMOS_EDGE_OFFICIAL_BENCHMARK=/path/to/official/benchmark.json ``` -Archived native RMS+SiLU metrics: - -- profile inventory artifact: - `official_action_only_live_prelayer_vae_encode_profile.json` shows one - `WanVAE_.encode` call from BF16 `[1,3,61,480,832]` to BF16 - `[1,48,16,30,52]`, 27 `CausalConv3d`, 22 `RMS_norm`, 21 `SiLU`, and - top time in the early high-resolution 160-channel residual convs. -- baseline profile-only `WanVAE_.encode`: 3.778s; native RMS+SiLU profile - `WanVAE_.encode`: 3.252s. The refreshed profile artifacts now include - `shape_summary` and `native_candidate_summary`; the top native candidate is - `encoder.downsamples.0.downsamples.0.residual.2` as - `steady_cached_causal_conv3d` on `[1,160,24,240,416]`. -- upstream measured `OmniMoTModel.encode`: 4.197s -> 3.295s; measured - `_prepare_inference_data`: 4.348s -> 3.436s. -- native RMS+SiLU run `generate_batch`: 25.63s; decode 0.0088s; live native - UniPC scheduler 1 run / 30 steps; final action cos 0.9999955 / - rel_l2 0.299% / max_abs 0.00929. - -The remaining VAE bottleneck is still the early CausalConv3d chain: -`encoder.downsamples.0` and `encoder.downsamples.1` dominate the profile. The -next native VAE step should fuse/quantize the `RMS_norm -> SiLU -> CausalConv3d` -triplet for those fixed shapes, reusing the Motus VAE FP8/FP4 conv swap -structure but recalibrated against the Cosmos3-Edge action gate. - -An additional no-go probe rewrites no-cache single-frame Wan2.2 -`CausalConv3d` calls to mathematically equivalent `Conv2d` via -`--vae-t1-conv2d`. This is correct but slower on Thor when combined with native -RMS+SiLU: - -- `WanVAE_.encode`: 3.291s -> 3.507s -- measured `OmniMoTModel.encode`: 3.295s -> 3.511s -- `OmniInference.generate_batch`: 25.63s -> 25.92s -- final action cos 0.9999960 / rel_l2 0.282% / max_abs 0.00799 - -The T=1 Conv2d rewrite is therefore kept default-off as an opt-in probe. The -steady `T=24` cached chunks are still the important target. - -Another opt-in native VAE probe, `--vae-native-avgdown3d`, replaces Wan2.2 -`AvgDown3D` shortcut pooling with `cosmos3_edge_avgdown3d_bf16`. The kernel -canary is bit-exact against the upstream reshape/permute/group-mean semantics, -including the temporal front-pad case. In the real no-warmup pre-layer run, -local `AvgDown3D` profile time moves from about 77ms to 67ms, but the complete -path is not better when combined with native RMS+SiLU: - -- `WanVAE_.encode`: 3.252s -> 3.504s -- measured `OmniMoTModel.encode`: 3.295s -> 3.508s -- `OmniInference.generate_batch`: 25.63s -> 26.03s -- final action cos 0.9999955 / rel_l2 0.299% / max_abs 0.00929 - -This is kept as an opt-in correctness scaffold and default-off/no-go. The -worthwhile native VAE target remains the high-resolution CausalConv3d chain. - -The steady cached `T=24` Conv3d-to-Conv2d-slices decomposition has also been -probed via `dev_log_cosmos3_thor/probe_vae_conv2d_decompose.py`. It is local -math-equivalent enough for the action gate, but not a useful default path: - -- stage0 160-channel `[1,160,24,240,416]`: cos 0.999995 / - rel_l2 0.303%, but 0.185s vs BF16 cuDNN Conv3d 0.112s. -- stage1 320-channel `[1,320,24,120,208]`: cos 0.999995 / - rel_l2 0.311%, but 0.099s vs BF16 cuDNN Conv3d 0.086s. -- stage2 640-channel `[1,640,12,60,104]`: cos 0.999996 / - rel_l2 0.293%, 0.0294s vs 0.0296s, only ~1.01x and not the dominant site. - -This keeps the decomposition default-off too. The remaining viable path is a -single Thor-specific CausalConv3d kernel for the early steady chunks, not a -multi-call Torch/cuDNN decomposition. - -A first Thor-specific BF16 MMA CausalConv3d bring-up kernel is also wired as -`bf16_conv3d_v0_ndhwc_bf16out` for `sm_110a`. It mirrors the Motus v17 virtual -cache concat/direct causal-output structure, but uses BF16 tensor cores and -stays probe-only. The small-shape canary passes, and the real VAE shape probe -is accurate enough for the local gate, but still slower than cuDNN: - -- stage0 160-channel `[1,160,24,240,416]`: cos 0.999996 / - rel_l2 0.282%, kernel-only 0.275s vs BF16 cuDNN 0.112s. -- stage1 320-channel `[1,320,24,120,208]`: cos 0.999996 / - rel_l2 0.284%, kernel-only 0.214s vs BF16 cuDNN 0.087s. -- stage2 640-channel `[1,640,12,60,104]`: cos 0.999996 / - rel_l2 0.301%, kernel-only 0.089s vs BF16 cuDNN 0.029s. - -The v0 kernel is therefore default-off/no-go. It is useful as a native -correctness scaffold, but the next attempt needs better tiling/persistence or -triplet fusion, not this one-kernel v0 swap. - -The first steady cached chunk also has a cache-shape corner case: upstream -passes a one-frame cache into `CausalConv3d.forward`, which then prepends one -temporal zero via `F.pad`. A cache2 normalization probe confirms that replacing -that with an explicit `[zero, cache]` two-frame cache is bit-exact, but not a -useful cuDNN optimization by itself: - -- stage0 160-channel: prebuilt cache2 Conv3d 0.11231s vs 0.11247s reference; - building cache2 in Python plus Conv3d is 0.11308s. -- stage1 320-channel: prebuilt cache2 Conv3d 0.08665s vs 0.08669s reference; - building cache2 in Python plus Conv3d is 0.08697s. -- stage2 640-channel: prebuilt cache2 Conv3d 0.02971s vs 0.02960s reference; - building cache2 in Python plus Conv3d is 0.02976s. -- native `update_cache2_ncdhw_bf16` matches Python cache2 exactly, but still - adds 0.30-1.19ms per call at these shapes. - -This remains a native-kernel interface fact, not a standalone default path. -Cache2 is only worth standardizing if a future fused CausalConv3d/triplet -kernel consumes it directly and wins the full-site A/B. - -The `vae_channels_last3d_probe.json` artifact tests cuDNN BF16 Conv3d with -`channels_last_3d` tensors at the same three steady sites. It is bit-exact and -shows a useful preconverted-layout signal, but not a default-ready per-conv -rewrite: - -- stage0 160-channel: total conversion path 0.119s vs 0.112s reference - (0.95x), preconverted-layout conv 0.073s (1.55x). -- stage1 320-channel: total conversion path 0.076s vs 0.086s reference - (1.13x), preconverted-layout conv 0.053s (1.62x). -- stage2 640-channel: total conversion path 0.031s vs 0.030s reference - (0.95x), preconverted-layout conv 0.026s (1.16x). - -Conclusion: `channels_last_3d` is worth pursuing only if the VAE encode segment -can keep activations and weights in that layout across adjacent residual blocks -or if a fused native triplet consumes the layout directly. Per-conv conversion -stays default-off. - -An opt-in full-run patch, `--vae-channels-last3d-conv320`, applies this only to -steady cached 320-channel CausalConv3d sites and returns contiguous tensors to -preserve the upstream cache contract. Combined with `--vae-native-rms-silu`, it -is correct but slower in the real no-warmup pre-layer run: measured -`OmniMoTModel.encode` 4.04s vs 3.29s for native RMS+SiLU alone, -`_prepare_inference_data` 4.18s, `generate_batch` 26.41s, final action cos -0.9999955 / rel_l2 0.299% / max_abs 0.00929. The profile shows the 320-channel -cached sites become the top regression at 0.13-0.24s each, so the patch remains -default-off/no-go. - -The follow-up `vae_resblock_channels_last3d_probe.json` artifact tests a -block-local layout propagation variant across three real steady -`ResidualBlock`s. It keeps the second half of each block in `channels_last_3d` -instead of returning to contiguous after every conv. The probe remains correct, -but all three blocks are slower than the official block: - -- stage0 160-channel block: 0.441s vs 0.323s reference (0.73x). -- stage1 320-channel block: 0.259s vs 0.222s reference (0.85x). -- stage2 640-channel block: 0.077s vs 0.072s reference (0.93x). - -This closes the small-step `channels_last_3d` path: block-local propagation is -not enough. Future VAE work should use a fused RMS/SiLU/CausalConv3d triplet or -a stronger Thor-specific steady CausalConv3d kernel that consumes a stable -layout directly. - -For a non-native service-warmup comparison, `--vae-compile-encode` enables the -upstream Wan2.2 chunk-level AOT `compile_encode` path for 480p 16:9. The -archived run compiled 5 loadable AOT functions in 133.7s during setup, then -reduced measured no-warmup pre-layer VAE encode: - -- measured `OmniMoTModel.encode`: 4.197s eager baseline -> 2.428s AOT -- measured `_prepare_inference_data`: 2.554s -- `OmniInference.generate_batch`: 24.43s, decode 0.0085s -- final action cos 0.9999958 / rel_l2 0.293% / max_abs 0.00553 - -This is useful as an opt-in long-lived-service baseline and a latency target for -native VAE work. It is not the final FlashRT native implementation because it -pays large startup compile cost and still relies on Torch Inductor/AOT. - -The direct Motus FP8 conv3d reuse path has been probed on Thor: +FP8, all 30 computes: ```bash -python dev_log_cosmos3_thor/probe_vae_triplet_fp8.py \ - --json-out dev_scratch_cosmos3_thor/edge_av_inverse_0/vae_triplet_fp8_probe_site0.json +python benchmarks/cosmos3_edge_thor_denoise.py \ + --checkpoint "$COSMOS_EDGE_CHECKPOINT" \ + --reference-dump "$COSMOS_EDGE_REFERENCE" \ + --boundary-dump "$COSMOS_EDGE_BOUNDARY" \ + --official-benchmark "$COSMOS_EDGE_OFFICIAL_BENCHMARK" \ + --engine fp8 \ + --iters 15 \ + --warmup-steps 30 \ + --enforce-gates ``` -The SM120-origin v17/v18 FP8 conv3d kernels compile for `sm_110a`; CMake now -exposes `fp8_conv3d_v17_ndhwc_bf16out`, -`fp8_conv3d_v17_anyco_ndhwc_bf16out`, and -`fp8_conv3d_v18_ncdhw_res_bf16out` on Thor. The small-shape kernel canary -matches torch exactly, but the real Cosmos VAE triplet A/B is negative: - -- stage0 160-channel `[1,160,24,240,416]`: FP8 v18 cos 0.999868 / - rel_l2 1.69%, but 0.245s vs BF16 cuDNN 0.113s. -- stage1 320-channel `[1,320,24,120,208]`: FP8 v18 cos 0.999816 / - rel_l2 1.92%, but 0.182s vs BF16 cuDNN 0.086s. -- stage2 640-channel `[1,640,12,60,104]`: FP8 v18 cos 0.999931 / - rel_l2 1.38%, but 0.082s vs BF16 cuDNN 0.030s. - -Therefore the Motus v17/v18 FP8 conv3d path is kept default-off for -Cosmos3-Edge. The first Thor-specific BF16 v0 kernel above is also slower, so -the next VAE conv step needs a better Thor-specific BF16/FP8 CausalConv3d -tiling/fusion strategy, not a direct Motus conv swap or the v0 BF16 probe. - -To also reuse the whole prepared inference state from the warmup request: +FP8 plus NVFP4 FFN, all 30 computes: ```bash -python examples/cosmos3_edge_thor_baseline.py \ - --checkpoint /work/models/Cosmos3-Edge \ - --input-json /work/.tmp_cosmos_edge_inputs/av_inverse_0.json \ - --output-dir /work/.tmp_cosmos_edge_outputs_live_warm_request_prepare_cache \ - --vae-path /work/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth \ - --cosmos-root /work/external/cosmos-framework \ - --backend official_action_only \ - --live-warm-request \ - --cache-warmup-prepare \ - --live-handoff-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_prepare_cache_handoff_trace.json \ - --upstream-trace-out dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_warm_request_prepare_cache_upstream_trace.json +python benchmarks/cosmos3_edge_thor_denoise.py \ + --checkpoint "$COSMOS_EDGE_CHECKPOINT" \ + --reference-dump "$COSMOS_EDGE_REFERENCE" \ + --boundary-dump "$COSMOS_EDGE_BOUNDARY" \ + --official-benchmark "$COSMOS_EDGE_OFFICIAL_BENCHMARK" \ + --engine fp8 \ + --ffn-fp4 \ + --iters 15 \ + --warmup-steps 30 \ + --enforce-gates ``` -Archived `official_action_only_live_warm_request_prepare_cache_outputs` metrics: - -- `OmniInference.generate_batch`: 11.98s -- `OmniMoTModel.generate_samples_from_batch`: 11.97s -- `OmniMoTModel.decode`: 0.00015s -- measured `OmniMoTModel._prepare_inference_data`: 0.0185s -- measured prepare cache hits: 1, stores: 1, misses: 0 -- measured VAE encode calls: 0 -- live native UniPC scheduler: 2 runs / 60 steps, native scheduler step total - 0.0133s -- final action cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687 -- Adding `--prepare-slim-no-raw-state-vision` slims the cached prepare payload - by dropping raw RGB state after warmup prepare. Archived slim cache metrics: - `generate_batch` 12.08s, denoise 12.07s, measured prepare 0.0146s, prepare - hit/store 1/1 with `slim_no_raw_state_vision=true`, measured VAE encode calls - 0, native UniPC 2 runs / 60 steps, final action cos 0.9999958 / rel_l2 0.291%. -- Adding `--prepare-slim-derive-condition-reference` on top stores the warmup - prepare payload without raw RGB and without `condition_reference`; the - measured cache hit derives `condition_reference` before the official sampler. - Archived derived slim cache metrics: `generate_batch` 12.00s, denoise 11.99s, - measured prepare 0.0149s, prepare hit/store 1/1 with both slim flags recorded, - measured VAE encode calls 0, native UniPC 2 runs / 60 steps, final action cos - 0.9999959 / rel_l2 0.291%. -- Adding `--prepare-slim-derive-initial-noise` on top stores the repeated-request - prepare payload without raw RGB, stored `condition_reference`, or stored - `initial_noise`; the measured cache hit derives noise from seed/x0/mask and - then derives the reference. Archived slim-derived-noise cache metrics: - `generate_batch` 12.05s, denoise 12.04s, measured prepare 0.0382s, prepare - hit/store 1/1 with all three slim flags recorded, measured VAE encode calls 0, - native UniPC 2 runs / 60 steps, final action cos 0.9999959 / rel_l2 0.291%. - -## P1 Replay Scaffold - -The replay scaffold validates the denoise boundary without rerunning the -transformer. It reads the P0 dump, replays the recorded velocity tensors through -FlashRT's local UniPC scheduler, and splits the final flat latent into: +## TeaCache -- vision latent: `[1,48,16,30,52]` -- action model latent: `[60,64]` - -Run it inside the container/venv: +`--teacache-computes K` selects `K` evenly spaced velocity evaluations and +always includes steps 0 and 29. `--teacache-steps` accepts an explicit +comma-separated compute set and takes precedence. ```bash -python - <<'PY' -from flash_rt.frontends.torch.cosmos3_edge_thor import Cosmos3EdgeTorchFrontendThor - -pipe = Cosmos3EdgeTorchFrontendThor("/work/models/Cosmos3-Edge", hardware="thor") -result = pipe.infer( - backend="replay", - output_dir="/work/.tmp_cosmos_edge_replay", - reference_dump=( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/tensors.safetensors" - ), - replay_device="cuda", +python benchmarks/cosmos3_edge_thor_denoise.py \ + --checkpoint "$COSMOS_EDGE_CHECKPOINT" \ + --reference-dump "$COSMOS_EDGE_REFERENCE" \ + --boundary-dump "$COSMOS_EDGE_BOUNDARY" \ + --official-benchmark "$COSMOS_EDGE_OFFICIAL_BENCHMARK" \ + --engine fp8 \ + --teacache-computes 3 \ + --iters 15 \ + --warmup-steps 30 \ + --enforce-gates +``` + +For qualification, retain a no-cache output for every evaluated seed and +report both action-space metrics and downstream task success. Passing the +cosine and relative-L2 gates is necessary for this benchmark but does not by +itself establish policy quality. + +## Python API + +The stable model registration exposes the official pipeline through +`flash_rt.load_model`: + +```python +import flash_rt + +model = flash_rt.load_model( + checkpoint="/path/to/Cosmos3-Edge", + config="cosmos3_edge", + framework="torch", + hardware="thor", ) -print(result) -PY -``` - -Expected validation facts for the current dump: - -- `num_steps`: 30 -- `flat_dim`: 1201920 -- `timesteps`: `999 ... 256` -- CUDA replay `max_input_abs_diff`: 0.0 - -## P1 Torch Reference - -The native bring-up path now has two local correctness gates: - -- `EdgeLayer0TorchReference`: layer 0 causal/full outputs match the official - boundary dump exactly after correcting the public checkpoint tower mapping. -- `EdgeTransformerTorchReference`: all 28 layers plus the action head match the - step-0 official action velocity. -- `EdgeDenoiseTorchReference`: the full 30-step UniPC loop computes action - velocity live from the Torch reference instead of replaying recorded velocity. -- `EdgeDenoiseFlashRT`: the public optimized eager engine used by - `backend="flashrt"`. The velocity path is the static FlashRT engine rather - than the two-tower reference replay path, and the fixed UniPC latent update - uses the native scheduler step when available. The frontend writes an - action-only `sample_outputs.json` for this backend and intentionally does not - emit decoded vision files. -- `EdgeStaticBufferEngine`: owns the static AV geometry, static vision tokens, - static und K/V cache, action slots, and the flat velocity contract. Its action - in/out projections use `GemmRunner.bf16_nn` when `flash_rt_kernels` is - available. The action input bias and modality embedding are precombined once - during static engine initialization. -- `EdgeTransformerFvkLinearReference`: routes cached-path BF16 linear projections - through `GemmRunner.bf16_nn` and cached-path RMSNorm through `fvk.rms_norm`; - cached-path main q/k RoPE through `qwen36_partial_rope_qk_bf16`, gen q/k - norm+RoPE through `cosmos3_edge_qk_norm_rope_bf16`, and FFN activation through - `relu2_inplace_bf16`. On Thor, non-causal gen attention uses the vendored FA4 - CuTe path when the `thor-fa4` extra is installed. `flash_rt_fa2` is not - expected on Thor. -- Hot-path weights/transposes are resident in the static engine cache. The - current checkpoint keeps 532 bf16 tensors, 4 fp32 timestep tensors, 4 boundary - tensors, and 336 transposed GEMM weights resident. -- `EdgeStaticBufferEngine.flat_velocity_for_step` uses - `cosmos3_edge_fill_flat_velocity_bf16` when the binding is available. This - replaces the Torch `zero_()` plus action-tail assignment for the flat velocity - output buffer. -- `EdgeStaticBufferEngine._decode_action_velocity` uses - `cosmos3_edge_add_bias_zero_action_tail_bf16` when the binding is available. - This fuses action output bias add and invalid action-dimension clearing. -- `EdgeStaticBufferEngine` uses `cosmos3_edge_scatter_rows_bf16` and - `cosmos3_edge_gather_rows_bf16` when available for action token row movement - between the compact action buffer and full sequence buffer. -- `EdgeStaticBufferEngine.full_sequence_for_step` uses - `cosmos3_edge_copy_action_tail_f32_to_bf16` when available for CUDA float32 - latents, avoiding the Torch action-tail slice and BF16 cast. -- `EdgeStaticBufferEngine._encode_action_tokens` uses - `cosmos3_edge_add_action_bias_timestep_bf16` when available. The timestep - embedding is computed once for the scalar timestep and broadcast natively - across the 60 action rows. -- `EdgeDenoiseFlashRT` precomputes all dump timesteps into - `EdgeStaticBufferEngine.timestep_embed_cache`, so measured denoise steps reuse - resident BF16 timestep embeddings instead of recomputing the Torch - cos/sin/MLP path. - -Install the FA4 dependency set in the isolated container venv before measuring -the optimized eager path: - -```bash -cd /work/official/flashrt-public -source /work/.venv_cosmos_thor/bin/activate -python -m pip install -e ".[thor-fa4]" -``` - -Run the full reference denoise check inside the container/venv: - -```bash -python - <<'PY' -from flash_rt.frontends.torch.cosmos3_edge_thor import Cosmos3EdgeTorchFrontendThor - -pipe = Cosmos3EdgeTorchFrontendThor("/work/models/Cosmos3-Edge", hardware="thor") -result = pipe.infer( - backend="flashrt", - output_dir="/work/.tmp_cosmos_edge_flashrt", - reference_dump=( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/tensors.safetensors" - ), - boundary_dump=( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - replay_device="cuda", +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, ) -print(result) -PY ``` -Current 30-step optimized eager result with FA4 attention: - -- und-cache precompute/init: 1.44s -- 30-step run wall time after one warm step: P50 11.998s over 10 iters -- speedup vs official `OmniMoTModel.generate_samples_from_batch`: 2.762x -- `max_input_abs_diff`: 0.00711 -- `max_velocity_abs_diff`: 0.046875 -- final action vs official: cos 0.9999959, rel_l2 0.286%, max_abs 0.00734 -- peak allocated VRAM during warm run: 11.02 GiB -- estimated action-only E2E with official upstream, skipped vision decode, and - FlashRT denoise: 14.52s, or 3.26x vs official 47.32s E2E -- hard gate run: `flashrt_benchmark_gate_10iter.json` with `gates_passed=true` -- native flat velocity fill canary: - `flashrt_benchmark_fill_velocity_native_3iter.json`, P50 11.97s, speedup - 2.77x, with the same action accuracy gate. -- native action bias+tail-clear canary: - `flashrt_benchmark_action_tail_native_3iter.json`, P50 12.02s, speedup - 2.76x, with the same action accuracy gate. -- native action row scatter/gather canary: - `flashrt_benchmark_action_rows_native_3iter.json`, P50 12.03s, speedup - 2.75x, with the same action accuracy gate. -- native action input tail-copy canary: - `flashrt_benchmark_action_input_native_3iter.json`, P50 12.02s, speedup - 2.76x, with the same action accuracy gate. -- native action bias+timestep canary: - `flashrt_benchmark_action_bias_timestep_native_3iter.json`, P50 12.03s, - speedup 2.75x, with the same action accuracy gate. -- timestep cache canary: - `flashrt_benchmark_timestep_cache_3iter.json`, P50 12.04s, speedup 2.75x, - with the same action accuracy gate. -- native UniPC scheduler canary: - `flashrt_benchmark_native_unipc_3iter.json`, P50 12.00s, speedup 2.76x, - final action cos 0.9999959, rel_l2 0.286%. This is enabled by default when - the binding is available. -- opt-in native residual add canary: - `flashrt_benchmark_residual_add_native_3iter.json`, P50 12.24s, speedup - 2.71x, with the same action accuracy gate. This kernel is correct but is not - default because the extra launches are slower than Torch add in the current - eager stack. -- opt-in CUDA graph run: - `flashrt_benchmark_graph_10iter.json` with `use_cuda_graphs=true`, - `graph_attention=true`, `native_scheduler=true`, P50 12.01s, speedup 2.76x, - final action cos 0.9999962, rel_l2 0.279%. Graph replay is verified, but - eager remains the default because it is marginally faster on this Thor run. -- graph loop breakdown: - `flashrt_benchmark_graph_loop_profile_3iter.json` adds CUDA event timing for - one profiled 30-step graph run. Velocity graph replay accounts for 11.966s of - the 11.970s step total; native UniPC scheduler accounts for only 0.0041s - total, with p50 0.141ms/step. Capturing the Python 30-step loop as one larger - graph is therefore not the next useful latency target; the remaining time is - inside per-step velocity compute. -- additional seed check: seed=1 dump - `dev_scratch_cosmos3_thor/edge_av_inverse_1/tensors.safetensors`, final - action cos 0.9999956 / rel_l2 0.300% / max_abs 0.00923. This run skips the - speedup gate because the official seed=1 denoise baseline was a single faster - run at 29.16s; it is used as the anti-overfit accuracy check. - -Reproduce the number: - -```bash -python benchmarks/cosmos3_edge_thor_denoise.py \ - --iters 10 \ - --warmup-steps 1 \ - --enforce-gates \ - --json-out dev_scratch_cosmos3_thor/edge_av_inverse_0/flashrt_benchmark_gate_10iter.json -``` - -Reproduce the opt-in CUDA graph number: - -```bash -FLASHRT_COSMOS3_EDGE_FA4_FWD=1 \ -python benchmarks/cosmos3_edge_thor_denoise.py \ - --iters 10 \ - --warmup-steps 1 \ - --use-cuda-graphs \ - --enforce-gates \ - --json-out dev_scratch_cosmos3_thor/edge_av_inverse_0/flashrt_benchmark_graph_10iter.json -``` - -Verify the existing delivery artifacts without rerunning GPU benchmarks: - -```bash -python dev_log_cosmos3_thor/verify_delivery.py -``` - -Attention notes: - -- FA4 (`flash_rt.hardware.thor.fa4_backend`) is the default native eager - attention path for gen/full GQA. It is accurate at the Cosmos3 shape and avoids - the long-context failure seen in the cuBLAS MHA fallback. FA4 accepts BF16 - Q/K/V directly, so the hot path avoids per-layer BF16→FP16 casts. -- `fmha_strided_full` remains behind `FLASHRT_COSMOS3_EDGE_FMHA=1`; it is - accurate on some large random shapes but increased full step-0 velocity error - in the Cosmos3 stack, so it is not the default. -- `attention_mha_bf16` remains behind `FLASHRT_COSMOS3_EDGE_MHA=1`; its - `softmax_bf16` path is not reliable for Cosmos3's long `S_kv=6425`. -- Velocity-only CUDA graph capture is available behind - `FLASHRT_COSMOS3_EDGE_FA4_FWD=1` plus `use_cuda_graphs=True`. The failure was - narrowed to the final gen norm bypassing the weight cache; after routing - `norm_moe_gen.weight` through the resident cache, the 28-layer - `flat_velocity` probe and the 10-iter denoise graph benchmark both pass. Graph - replay is opt-in rather than default because the measured P50 is 12.01s vs - 12.00s for eager in the current artifact set. -- The graph loop profile shows native scheduler and Python-loop glue are no - longer material at the 12s denoise scale: scheduler event time is 4.1ms across - all 30 steps. This closes the immediate "capture the outer loop" question as - a no-go for latency; further denoise work must reduce the per-layer velocity - kernels themselves. -- A forced graph experiment with SDPA attention also exited with code 139, so - there is no SDPA graph fallback enabled in this Python eager stack. -- `FLASHRT_COSMOS3_EDGE_FA4_FWD=1` switches eager attention to the lower-level - FA4 `_flash_attn_fwd` entry with preallocated output/LSE buffers. It is also - required for the opt-in graph path. As an eager-only experiment, seed0 3-iter - P50 was 12.008s with the same action accuracy, which does not beat the - default gate result. -- A/B attempts to replace per-layer `torch.cat` with manual static K/V copies - and precompute action timestep embeddings did not improve P50 latency on Thor, - so those overrides are not in the default path. The earlier two-op - `relu_inplace_bf16 + square_` experiment was also not useful, but the fused - `relu2_inplace_bf16` kernel is now enabled by default. Seed0 3-iter A/B: - native relu2 P50 12.274s vs Torch relu2 P50 12.926s, with identical final - action metrics. -- `cosmos3_edge_qk_norm_rope_bf16` fuses gen q/k RMSNorm and RoPE and is enabled - by default. Seed0 3-iter A/B: fused P50 11.938s vs two-step native path P50 - 12.255s, with identical final action metrics. -- `cosmos3_edge_fill_flat_velocity_bf16` is enabled by default when - `flash_rt_kernels` exposes the binding. It fills the `[1201920]` flat velocity - buffer in one native launch, zeroing the vision segment and copying the - `[60,64]` action tail. The CUDA canary in - `tests/test_cosmos3_edge_kernel_bindings.py` checks exact equality with the - Torch construction. -- `cosmos3_edge_add_bias_zero_action_tail_bf16` is enabled by default when - available. It adds the action output bias for the valid action dimensions and - clears columns `raw_action_dim:64` in one native launch. The CUDA canary checks - exact equality with the prior Torch `add_` plus slice-clear construction. -- `cosmos3_edge_scatter_rows_bf16` and `cosmos3_edge_gather_rows_bf16` are - enabled by default when available. They replace the action-row advanced - indexing operations used to install encoded action tokens into the full - sequence and gather decoded action hidden states back out. The CUDA canary - checks exact equality with PyTorch indexed row copy. -- `cosmos3_edge_copy_action_tail_f32_to_bf16` is enabled by default when - available and the latent is CUDA float32. It copies the action tail from the - flat denoise latent into the BF16 action input buffer in one native launch. - The CUDA canary checks exact equality with PyTorch tail slicing plus BF16 - conversion. -- `cosmos3_edge_add_action_bias_timestep_bf16` is enabled by default when - available. It fuses the action static bias add and timestep embedding add, - preserving the prior two-step BF16 rounding order. The timestep embedding is - computed as one row instead of 60 duplicate rows; this changes final action - metrics only in the last few decimals and remains well inside the reference - gate. -- `EdgeDenoiseFlashRT` now precomputes the 30 fixed timestep embeddings from - the dump schedule. The measured denoise loop indexes this cache by step, so - the per-step action encode path no longer calls the Torch timestep - `cos/sin/linear/silu/linear` sequence. -- `cosmos3_edge_unipc_step_f32_bf16` is enabled by default when available. It - folds the fixed 30-step UniPC x0 conversion, corrector, and predictor update - into one native launch per denoise step. The kernel preserves the prior - BF16 `sigma * velocity` rounding order before updating the float32 latent. -- `cosmos3_edge_add_bf16` is available as an opt-in native residual add - experiment with `FLASHRT_COSMOS3_EDGE_NATIVE_ADD=1`. The CUDA canary checks - exact equality with Torch BF16 add. It remains disabled by default because the - measured 3-iter canary regressed P50 to 12.24s versus the 12.00s default. - -P2 NVFP4 canary: - -- SM110 exposes the NVFP4-named bindings in `flash_rt_kernels`, and minimal - `quantize_bf16_to_nvfp4_swizzled` plus `fp4_w4a16_gemm_sm120_bf16out` - canaries execute. -- A real Cosmos3 layer-0 gen `add_q_proj` canary with correct `[N,K]` weight - layout and `bf16_weight_to_nvfp4_swizzled` alpha produced cos 0.9974, - rel_l2 10.3%, and was slower than BF16 GEMM when activation quantization was - included. This does not pass the P2 accuracy/perf gate, so NVFP4 is not wired - into the default backend. - -P4 E2E note: - -- Official benchmark for `av_inverse_0` reports `OmniMoTModel.decode` at - 11.65s. Since inverse dynamics only emits action in the final JSON, this is a - strong candidate for an E2E skip in a future no-dump path. -- The external Cosmos Framework path confirms the skip opportunity: - `inference.py` unconditionally does `outputs.pop("vision") -> decode_vision` - and saves `vision.mp4` before writing `content["action"]`; the action payload - itself is taken from `outputs["action"]` and does not depend on decoded vision. - The external checkout remains unmodified. -- FlashRT-owned `backend="official_action_only"` validates that skip on the - real no-dump official path: `OmniMoTModel.decode` is 0.00025s, generated - action is bit-exact to the official baseline JSON, `files=[]`, and no - `vision*` artifact remains. This still uses the official eager denoise path; - the remaining P4 gap is wiring the FlashRT denoise engine to live upstream - conditioning instead of captured dumps. -- The opt-in live capture path produced - `official_action_only_live_denoise.safetensors` from the real official run: - 30 steps, valid `EdgeDenoiseDump` geometry, action-only JSON with `files=[]`, - decode 0.00048s, and selected tensors match the original reference dump with - max_abs 0.0. This proves a FlashRT-owned live capture boundary, but not yet - live FlashRT denoise handoff. -- The live FlashRT handoff prototype produced - `official_action_only_live_handoff_boundary.safetensors` and an action-only - run with `OmniInference.generate_batch` 26.10s, decode 0.0083s, files `[]`, - final action cos 0.9999959 / rel_l2 0.292% / max_abs 0.00729 vs official - baseline. This removes captured denoise dumps from the online path after - step 0; remaining work is removing the official step-0 bootstrap and reducing - cold handoff overhead. -- The pre-layer bootstrap variant removes the official step-0 decoder forward: - `prelayer_aborted=true`, 30 FlashRT velocity calls, `generate_batch` 26.24s, - decode 0.0090s, final action cos 0.9999955 / rel_l2 0.308% / max_abs 0.00780. - It confirms the remaining one-shot latency is cold FA4/engine setup rather - than official decoder compute. -- The archived pre-layer boundary contract - `official_action_only_live_prelayer_boundary.safetensors` is the smaller - native-denoise/VLM input boundary: 60.36 MiB, 31 tensors, no layer-0 output, - and no step-0 velocity. The older step-0 handoff boundary was 87.77 MiB / - 36 tensors because it also stored layer-0 output and velocity. The run keeps - the same action gate (cos 0.9999959 / rel_l2 0.291% / max_abs 0.00687) and - native UniPC covers all 30 steps. This boundary defines what native prefill - must synthesize live from raw input before the denoise engine can be fully - detached from official upstream code. -- `--live-boundary-in` can now use that 60.36 MiB boundary artifact to - initialize the live FlashRT denoise engine directly. The archived replay run - has `official_velocity_call_count=0`, 30 FlashRT velocities, native UniPC 30 - steps, `generate_batch` 26.03s, and the same action gate. It is a runnable - native prefill target contract, not the final no-artifact live path. -- Combined with `--live-warm-request --cache-warmup-prepare` and all slim derive - flags, boundary-in gives the current artifact/cache service path: - measured `generate_batch` 12.08s, official velocity calls 0, FlashRT velocity - calls 60, measured VAE encode calls 0, measured prepare 0.042s, and the same - action gate. -- The prepare-to-boundary derive helper proves 28 VFM/noise/pack/position/RoPE/text - tensors / 10.65 MiB are already implied by the 9.18 MiB slim prepare contract - plus `embed_tokens.weight`, - and verifies 2 layer-0 input tensors / 25.10 MiB as `lm_in` aliases. The - remaining pre-layer boundary gap is 1 tensor / 24.61 MiB: - `lm_in/full_only_seq`; the verifier now requires the derived full 6300-row - tensor to be bit-exact, including the 60 action rows. -- `--live-boundary-prepare-in` now derives the executable 31-tensor boundary - in-process from the same 9.18 MiB slim prepare artifact plus checkpoint - weights, then drives the live FlashRT denoise engine with 0 official velocity - calls. The archived run measured `generate_batch` 22.24s, denoise 22.22s, - boundary derive 0.258s, engine construct 1.074s, native UniPC 30 steps, and - the same action gate. The optional `--live-boundary-out` debug artifact is - 62.65 MiB because shared views are cloned for safetensors archival; it is not - the runtime input. -- With `--live-warm-request`, the real measured pre-layer batch drops to - `generate_batch` 12.12s / denoise 12.11s while preserving the action gate; - the trace shows 60 FlashRT velocity calls and native UniPC scheduler coverage - for both warmup and measured sampler runs. Only the first warmup velocity call - pays cold compile, and a measured VAE cache hit removes the previous 3.41s - encode from the real request. The warmup VAE cache key includes a sampled - content fingerprint in addition to tensor metadata. -- The opt-in `--cache-warmup-prepare` path reuses the full prepared inference - state from warmup for a repeated request: `generate_batch` 11.98s / denoise - 11.97s, measured `_prepare_inference_data` 0.0185s, prepare hit/store 1/1, - measured misses 0, and measured VAE encode calls 0. The live sampler now uses - the native FlashRT UniPC step for this fixed Edge schedule: 2 native scheduler - runs / 60 steps, 0.0133s total scheduler step time. With - `--prepare-slim-no-raw-state-vision`, the same cache stores a slim post-prepare - payload without raw RGB frames: measured `generate_batch` 12.08s, prepare - 0.0146s, and trace hit/store events both mark - `slim_no_raw_state_vision=true`. With - `--prepare-slim-derive-condition-reference`, the cached payload also omits - `condition_reference` and restores it from `x0_tokens_vision/action` on hit: - measured `generate_batch` 12.00s, prepare 0.0149s, and both cache events mark - `slim_derive_condition_reference=true`. With - `--prepare-slim-derive-initial-noise`, the cache also omits `initial_noise` - and restores it from seed/x0/mask metadata on hit: measured `generate_batch` - 12.05s, prepare 0.0382s, and cache events mark all three slim flags. This is - a service-style repeated-request prefill cache, not the final no-warmup live - VAE/SigLIP/VLM replacement. -- The opt-in upstream trace without VAE cache identified measured VAE - encode/get-data as the next P4 target; the current warmup cache moves that - repeated work out of the measured request for repeated-sample/service-style - warm paths. Native live VAE/SigLIP/VLM prefill remains the broader no-warmup - P4 target. - -Important mapping correction from the layer-0 gate: - -- `to_q/to_k/to_v/to_out` = und/causal tower -- `add_q_proj/add_k_proj/add_v_proj/to_add_out` = gen/full tower -- `norm_added_q/norm_added_k` belong to gen/full -- `k_norm_und_for_gen` is required for gen attending to und keys - -## Optimization Port - -The Cosmos3-Nano FlashRT structure is reusable, but not drop-in compatible: +`backend="flashrt"` is the fixture-backed correctness and denoise benchmark +route. Supply `reference_dump` and `boundary_dump` to that backend. The live +official-to-FlashRT handoff remains an integration path for development; it is +not the source of the denoise-only numbers above. -- Reuse: hardware dispatch, frontend/pipeline split, UniPC loop shape, static - conditioning cache, CUDA graph capture, TeaCache step reuse, and model-local - kernels. -- Change: Edge uses hidden size 2048, 28 layers, 16/8 GQA heads, `relu2`, and - Nemotron Dense VL reasoner/processor instead of Nano's 4096/36-layer Qwen path. -- Thor work: add a model-local `pipeline_thor.py`, SM110 joint-attention/mrope - kernels, per-layer BF16 reference taps, then FP8/NVFP4 calibration. +## Scope and limitations -Keep this official runner as the accuracy and benchmark gate while the optimized -Thor pipeline is filled in. +- The performance table is denoise P50 in a hot process, not cold-start or raw + request end-to-end latency. +- The no-cache engine is the portable optimization result. TeaCache schedules + require case-by-case validation and can differ substantially after + fine-tuning. +- Official preprocessing, including the Wan2.2 VAE encode, is outside the + optimized denoise boundary. On the measured sample VAE encode was about + 3.7 s, so it is the next major end-to-end optimization target. +- The current route is registered only for `framework="torch"` and + `hardware="thor"`; unsupported hardware fails during model resolution. diff --git a/docs/stable_api.md b/docs/stable_api.md index ca6d67fa..d6b9a528 100644 --- a/docs/stable_api.md +++ b/docs/stable_api.md @@ -113,11 +113,12 @@ Returns a `VLAModel` wrapping the appropriate frontend for the detected compare_ref=..., return_metadata=...)`, returning the denoised vision latent; `predict()` is not part of this API. Precision is selected with `load_model(..., use_fp8=True|False)`. See `docs/cosmos3_video_usage.md`. -- `config="cosmos3_edge"` is a Thor official-baseline runner for NVIDIA - Cosmos Framework. It exposes `set_prompt(input_json=...)` or - `set_prompt(sample=...)`, then `infer(output_dir=..., vae_path=..., - cosmos_root=...)`. It is the accuracy/latency gate before the optimized Thor - Edge pipeline replaces the reference execution. See `docs/cosmos3_edge_thor.md`. +- `config="cosmos3_edge"` is the Thor integration for NVIDIA Cosmos Framework. + It exposes `set_prompt(input_json=...)` or `set_prompt(sample=...)`, then + `infer(output_dir=..., vae_path=..., cosmos_root=...)`. The official backend + establishes the end-to-end reference; the fixture-backed FlashRT backend + validates and benchmarks the optimized denoise boundary. See + `docs/cosmos3_edge_thor.md`. - `config="groot_n17"` is registered for `framework="torch"` on `hardware in {"thor", "rtx_sm120", "rtx_sm89"}`. On RTX, `rtx_sm120` resolves through the historical shared RTX registration and diff --git a/tests/test_cosmos3_edge_boundary_dump.py b/tests/test_cosmos3_edge_boundary_dump.py index 2e55b868..1a0a0534 100644 --- a/tests/test_cosmos3_edge_boundary_dump.py +++ b/tests/test_cosmos3_edge_boundary_dump.py @@ -1,3 +1,4 @@ +import os from pathlib import Path import pytest @@ -16,66 +17,36 @@ ) +def _path_from_env(name: str) -> Path | None: + raw = os.environ.get(name) + if not raw: + return None + path = Path(raw).expanduser() + return path if path.exists() else None + + def _local_boundary_dump() -> Path | None: - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - ] - return next((path for path in candidates if path.exists()), None) + return _path_from_env("COSMOS3_EDGE_BOUNDARY_DUMP") def _local_prelayer_boundary_dump() -> Path | None: - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors" - ), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/official_action_only_live_prelayer_boundary.safetensors" - ), - ] - return next((path for path in candidates if path.exists()), None) + return _path_from_env("COSMOS3_EDGE_PRELAYER_BOUNDARY_DUMP") def _local_slim_prepare_dump() -> Path | None: - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0/" - "official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt" - ), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/official_action_only_live_prelayer_prepare_slim_derive_noise_dump.pt" - ), - ] - return next((path for path in candidates if path.exists()), None) + return _path_from_env("COSMOS3_EDGE_SLIM_PREPARE_DUMP") def _local_text_embedding_shard() -> Path | None: - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge/" - "transformer/diffusion_pytorch_model-00001-of-00002.safetensors" - ), - Path("/work/models/Cosmos3-Edge/transformer/diffusion_pytorch_model-00001-of-00002.safetensors"), - ] - return next((path for path in candidates if path.exists()), None) + root = _local_cosmos3_edge_model_root() + if root is None: + return None + path = root / "transformer" / "diffusion_pytorch_model-00001-of-00002.safetensors" + return path if path.exists() else None def _local_cosmos3_edge_model_root() -> Path | None: - candidates = [ - Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), - Path("/work/models/Cosmos3-Edge"), - ] - return next((path for path in candidates if path.exists()), None) + return _path_from_env("COSMOS3_EDGE_CHECKPOINT") def test_cosmos3_edge_boundary_dump_geometry_when_present(): diff --git a/tests/test_cosmos3_edge_dump_replay.py b/tests/test_cosmos3_edge_dump_replay.py index 9d0bc327..0338b195 100644 --- a/tests/test_cosmos3_edge_dump_replay.py +++ b/tests/test_cosmos3_edge_dump_replay.py @@ -1,3 +1,4 @@ +import os from pathlib import Path from types import SimpleNamespace @@ -16,6 +17,14 @@ from flash_rt.models.cosmos3_edge.static_unipc import EdgeStaticUniPCScheduler # noqa: E402 +def _local_denoise_dump() -> Path | None: + raw = os.environ.get("COSMOS3_EDGE_REFERENCE_DUMP") + if not raw: + return None + path = Path(raw).expanduser() + return path if path.exists() else None + + def test_cosmos3_edge_split_flat_geometry(): class _SyntheticDump: final_vision = torch.empty(EDGE_VISION_SHAPE) @@ -38,17 +47,7 @@ def split_flat(self, flat): def test_cosmos3_edge_local_dump_replays_scheduler_when_present(): - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors" - ), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/tensors.safetensors" - ), - ] - dump_path = next((path for path in candidates if path.exists()), None) + dump_path = _local_denoise_dump() if dump_path is None: pytest.skip("local Cosmos3-Edge P0 dump is not available") @@ -69,17 +68,7 @@ def test_cosmos3_edge_local_dump_replays_scheduler_when_present(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="native static UniPC replay requires CUDA") def test_cosmos3_edge_static_unipc_replays_scheduler_when_present(): - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors" - ), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/tensors.safetensors" - ), - ] - dump_path = next((path for path in candidates if path.exists()), None) + dump_path = _local_denoise_dump() if dump_path is None: pytest.skip("local Cosmos3-Edge P0 dump is not available") diff --git a/tests/test_cosmos3_edge_layer_ref.py b/tests/test_cosmos3_edge_layer_ref.py index 7d363069..571de334 100644 --- a/tests/test_cosmos3_edge_layer_ref.py +++ b/tests/test_cosmos3_edge_layer_ref.py @@ -1,3 +1,4 @@ +import os from pathlib import Path import pytest @@ -16,23 +17,12 @@ def _local_paths() -> tuple[Path, Path] | None: - candidates = [ - ( - Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - ), - ( - Path("/work/models/Cosmos3-Edge"), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - ), - ] - return next((paths for paths in candidates if paths[0].exists() and paths[1].exists()), None) + checkpoint_raw = os.environ.get("COSMOS3_EDGE_CHECKPOINT") + boundary_raw = os.environ.get("COSMOS3_EDGE_BOUNDARY_DUMP") + if not checkpoint_raw or not boundary_raw: + return None + paths = (Path(checkpoint_raw).expanduser(), Path(boundary_raw).expanduser()) + return paths if paths[0].exists() and paths[1].exists() else None @pytest.mark.skipif(not torch.cuda.is_available(), reason="layer0 reference requires CUDA for this shape") diff --git a/tests/test_cosmos3_edge_step0_reference.py b/tests/test_cosmos3_edge_step0_reference.py index 3111be75..2cd09898 100644 --- a/tests/test_cosmos3_edge_step0_reference.py +++ b/tests/test_cosmos3_edge_step0_reference.py @@ -1,3 +1,4 @@ +import os from pathlib import Path import pytest @@ -14,37 +15,20 @@ def _local_paths() -> tuple[Path, Path] | None: - candidates = [ - ( - Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - ), - ( - Path("/work/models/Cosmos3-Edge"), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0_boundary_step0/tensors.safetensors" - ), - ), - ] - return next((paths for paths in candidates if paths[0].exists() and paths[1].exists()), None) + checkpoint_raw = os.environ.get("COSMOS3_EDGE_CHECKPOINT") + boundary_raw = os.environ.get("COSMOS3_EDGE_BOUNDARY_DUMP") + if not checkpoint_raw or not boundary_raw: + return None + paths = (Path(checkpoint_raw).expanduser(), Path(boundary_raw).expanduser()) + return paths if paths[0].exists() and paths[1].exists() else None def _local_denoise_dump_path() -> Path | None: - candidates = [ - Path( - "/home/heima-thor/suliang/nvidia_fp8_80ms/official/flashrt-public/" - "dev_scratch_cosmos3_thor/edge_av_inverse_0/tensors.safetensors" - ), - Path( - "/work/official/flashrt-public/dev_scratch_cosmos3_thor/" - "edge_av_inverse_0/tensors.safetensors" - ), - ] - return next((path for path in candidates if path.exists()), None) + raw = os.environ.get("COSMOS3_EDGE_REFERENCE_DUMP") + if not raw: + return None + path = Path(raw).expanduser() + return path if path.exists() else None @pytest.mark.skipif(not torch.cuda.is_available(), reason="full step0 reference requires CUDA") diff --git a/tests/test_cosmos3_edge_thor_spec.py b/tests/test_cosmos3_edge_thor_spec.py index ff973c9d..5ced1d86 100644 --- a/tests/test_cosmos3_edge_thor_spec.py +++ b/tests/test_cosmos3_edge_thor_spec.py @@ -1,3 +1,4 @@ +import os from pathlib import Path import pytest @@ -26,7 +27,8 @@ def test_cosmos3_edge_spec_counts_and_tower_names(): def test_cosmos3_edge_local_checkpoint_index_matches_when_present(): - checkpoint = "/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge" - if not Path(checkpoint).exists(): + raw = os.environ.get("COSMOS3_EDGE_CHECKPOINT") + checkpoint = Path(raw).expanduser() if raw else None + if checkpoint is None or not checkpoint.exists(): pytest.skip("local Cosmos3-Edge checkpoint is not available") validate_transformer_index(checkpoint) diff --git a/tests/test_cosmos3_edge_weights.py b/tests/test_cosmos3_edge_weights.py index 09238f14..1e505ca2 100644 --- a/tests/test_cosmos3_edge_weights.py +++ b/tests/test_cosmos3_edge_weights.py @@ -1,3 +1,4 @@ +import os from pathlib import Path import pytest @@ -10,11 +11,11 @@ def _local_checkpoint() -> Path | None: - candidates = [ - Path("/home/heima-thor/suliang/nvidia_fp8_80ms/models/Cosmos3-Edge"), - Path("/work/models/Cosmos3-Edge"), - ] - return next((path for path in candidates if path.exists()), None) + raw = os.environ.get("COSMOS3_EDGE_CHECKPOINT") + if not raw: + return None + path = Path(raw).expanduser() + return path if path.exists() else None def test_cosmos3_edge_weight_loader_refs_expected_shard():