From 9fc4ca99aa158d8ec7144e92582f6e16157856b8 Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:37:09 +0800 Subject: [PATCH 1/3] feat(hyvla): native FlashRT Thor SM110 port of Hy-Embodied-0.5-VLA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hy-Embodied-0.5-VLA (HunYuan MoT dual-tower + flow matching) on Jetson Thor SM110 with runtime dynamic per-tensor FP8, fused megakernels, and optional NVFP4 FFN. New kernels (csrc/kernels/, SM110-gated): - hyvla_fused_thor: RoPE+QK-Norm+KV-write megakernel (also SM87-portable) - hyvla_vit_fuse: ViT residual-add + LayerNorm (also SM87-portable) - hyvla_quant_fp8_thor: single-CTA dynamic FP8 quantize (graph-safe) - hyvla_ffn_fp8_thor: FFN gate/up+SiLU and down+residual fused GEMM Architecture (per docs/adding_new_model.md): - models/hyvla/pipeline_thor.py — Thor compute path - frontends/torch/hyvla_thor.py — HyVLATorchFrontendThor - frontends/torch/_hyvla_thor_spec.py — declarative weight spec - _PIPELINE_MAP: ("hyvla","torch","thor") one-to-one - load_model(config="hyvla") dispatches through the standard VLA path Performance (real image, Thor SM110): - E2E predict: 159.6 ms (vs reference eager ~930 ms, ~5.8x) - Action cosine vs reference: 0.999706 Docs: docs/hyvla05_thor_sm110.md (English), docs/stable_api.md updated. --- CMakeLists.txt | 23 + csrc/bindings.cpp | 89 +++ csrc/kernels/hyvla_ffn_fp8_thor.cu | 224 +++++++ csrc/kernels/hyvla_ffn_fp8_thor.cuh | 14 + csrc/kernels/hyvla_fused_thor.cu | 95 +++ csrc/kernels/hyvla_fused_thor.cuh | 10 + csrc/kernels/hyvla_quant_fp8_thor.cu | 51 ++ csrc/kernels/hyvla_quant_fp8_thor.cuh | 11 + csrc/kernels/hyvla_vit_fuse.cu | 78 +++ csrc/kernels/hyvla_vit_fuse.cuh | 8 + docs/hyvla05_thor_sm110.md | 169 +++++ docs/stable_api.md | 6 +- flash_rt/api.py | 38 +- flash_rt/configs/hyvla.yaml | 65 ++ flash_rt/executors/torch_weights.py | 7 + flash_rt/frontends/torch/_hyvla_thor_spec.py | 196 ++++++ flash_rt/frontends/torch/hyvla_thor.py | 628 +++++++++++++++++++ flash_rt/hardware/__init__.py | 4 + flash_rt/models/hyvla/__init__.py | 1 + flash_rt/models/hyvla/pipeline_thor.py | 572 +++++++++++++++++ 20 files changed, 2278 insertions(+), 11 deletions(-) create mode 100644 csrc/kernels/hyvla_ffn_fp8_thor.cu create mode 100644 csrc/kernels/hyvla_ffn_fp8_thor.cuh create mode 100644 csrc/kernels/hyvla_fused_thor.cu create mode 100644 csrc/kernels/hyvla_fused_thor.cuh create mode 100644 csrc/kernels/hyvla_quant_fp8_thor.cu create mode 100644 csrc/kernels/hyvla_quant_fp8_thor.cuh create mode 100644 csrc/kernels/hyvla_vit_fuse.cu create mode 100644 csrc/kernels/hyvla_vit_fuse.cuh create mode 100644 docs/hyvla05_thor_sm110.md create mode 100644 flash_rt/configs/hyvla.yaml create mode 100644 flash_rt/frontends/torch/_hyvla_thor_spec.py create mode 100644 flash_rt/frontends/torch/hyvla_thor.py create mode 100644 flash_rt/models/hyvla/__init__.py create mode 100644 flash_rt/models/hyvla/pipeline_thor.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 32db2ee53..e042cc095 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1692,6 +1692,29 @@ target_include_directories(flash_rt_kernels PRIVATE ${CUTLASS_DIR}/tools/util/include ) +# ── Hy-Embodied-0.5-VLA Thor SM110 kernels ── +if(GPU_ARCH STREQUAL "110") + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/hyvla_fused_thor.cu + csrc/kernels/hyvla_vit_fuse.cu + csrc/kernels/hyvla_quant_fp8_thor.cu + csrc/kernels/hyvla_ffn_fp8_thor.cu) + target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_HYVLA_THOR=1) + message(STATUS "Hy-VLA Thor SM110 kernels: ENABLED") +endif() + +# ── Hy-Embodied-0.5-VLA Orin SM87 kernels ── +# The fused RoPE+QKNorm+KV-write megakernel is plain CUDA (no SM110-only +# instructions) and removes ~11 launches per attention block; the FP8 +# Thor kernels stay SM110-only. +if(GPU_ARCH STREQUAL "87") + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/hyvla_fused_thor.cu + csrc/kernels/hyvla_vit_fuse.cu) + target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_HYVLA_ORIN=1) + message(STATUS "Hy-VLA Orin SM87 fused attention-prep kernel: ENABLED") +endif() + target_link_libraries(flash_rt_kernels PRIVATE CUDA::cublas CUDA::cublasLt diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 54d3c25ab..c15483614 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -156,6 +156,14 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/gated_deltanet_qwen36.cuh" #endif #include "kernels/qwen3_qkv_post_proc.cuh" +#if defined(FLASHRT_HAVE_HYVLA_THOR) || defined(FLASHRT_HAVE_HYVLA_ORIN) +#include "kernels/hyvla_fused_thor.cuh" +#include "kernels/hyvla_vit_fuse.cuh" +#ifdef FLASHRT_HAVE_HYVLA_THOR +#include "kernels/hyvla_quant_fp8_thor.cuh" +#include "kernels/hyvla_ffn_fp8_thor.cuh" +#endif +#endif #ifdef FLASHRT_HAVE_NVFP4_SWIZZLE #include "kernels/silu_mul_to_nvfp4_swizzled.cuh" #include "kernels/fp4_swiglu_compact_sm120.cuh" @@ -8112,6 +8120,87 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; py::arg("k_stride_head"), py::arg("stream") = 0); #endif +#if defined(FLASHRT_HAVE_HYVLA_THOR) || defined(FLASHRT_HAVE_HYVLA_ORIN) + m.def("hyvla_rope_qknorm_kvwrite_bf16", + [](uintptr_t qkv, uintptr_t cos, uintptr_t sin, uintptr_t qn_w, + uintptr_t kn_w, uintptr_t q_out, uintptr_t kbuf, uintptr_t vbuf, + int S, int nq, int nkv, int hd, int S_tot, int off, float eps, + int kv_rep, uintptr_t stream) { + hyvla_rope_qknorm_kvwrite_bf16( + reinterpret_cast(qkv), + reinterpret_cast(cos), + reinterpret_cast(sin), + reinterpret_cast(qn_w), + reinterpret_cast(kn_w), + reinterpret_cast(q_out), + reinterpret_cast(kbuf), + reinterpret_cast(vbuf), + S, nq, nkv, hd, S_tot, off, eps, kv_rep, to_stream(stream)); + }, + py::arg("qkv"), py::arg("cos"), py::arg("sin"), py::arg("qn_w"), + py::arg("kn_w"), py::arg("q_out"), py::arg("kbuf"), py::arg("vbuf"), + py::arg("S"), py::arg("nq"), py::arg("nkv"), py::arg("hd"), + py::arg("S_tot"), py::arg("off"), py::arg("eps") = 1e-5f, + py::arg("kv_rep") = 1, py::arg("stream") = 0, + "Hy-VLA fused RoPE(q,k)+QK-Norm(q,k)+KV-write megakernel (bf16). " + "kv_rep>1 stores the KV cache pre-expanded for GQA."); + + m.def("hyvla_vit_add_layer_norm_bf16", + [](uintptr_t residual, uintptr_t x_add, uintptr_t ln_weight, + uintptr_t ln_bias, uintptr_t out, int rows, int dim, float eps, + uintptr_t stream) { + hyvla_vit_add_layer_norm_bf16( + reinterpret_cast(residual), + reinterpret_cast(x_add), + reinterpret_cast(ln_weight), + reinterpret_cast(ln_bias), + reinterpret_cast(out), rows, dim, eps, + to_stream(stream)); + }, + py::arg("residual"), py::arg("x_add"), py::arg("ln_weight"), + py::arg("ln_bias"), py::arg("out"), py::arg("rows"), py::arg("dim"), + py::arg("eps") = 1e-6f, py::arg("stream") = 0, + "Hy-VLA ViT fused residual-add (bf16 round, in-place) + LayerNorm."); + +#ifdef FLASHRT_HAVE_HYVLA_THOR + m.def("hyvla_quant_fp8_dyn_bf16", + [](uintptr_t x, uintptr_t out, uintptr_t scale, int n, uintptr_t stream) { + hyvla_quant_fp8_dyn_bf16( + reinterpret_cast(x), + reinterpret_cast(out), + reinterpret_cast(scale), n, to_stream(stream)); + }, + py::arg("x"), py::arg("out"), py::arg("scale"), py::arg("n"), + py::arg("stream") = 0, + "Hy-VLA single-CTA dynamic per-tensor FP8 quant (small M, graph-safe)."); + + m.def("hyvla_ffn_gu_silu_bf16", + [](uintptr_t x, uintptr_t gu, uintptr_t act, int M, int K, int Nout, + uintptr_t sx, float sgu, uintptr_t stream) { + hyvla_ffn_gu_silu_bf16( + reinterpret_cast(x), reinterpret_cast(gu), + reinterpret_cast(act), M, K, Nout, + reinterpret_cast(sx), sgu, to_stream(stream)); + }, + py::arg("x"), py::arg("gu"), py::arg("act"), py::arg("M"), py::arg("K"), + py::arg("Nout"), py::arg("sx"), py::arg("sgu"), py::arg("stream") = 0, + "Hy-VLA FFN kernel A: gu-GEMM gate/up + silu_mul -> bf16 act (Thor)."); + + m.def("hyvla_ffn_dn_res_bf16", + [](uintptr_t a, uintptr_t dn, uintptr_t res, uintptr_t y, int M, int K, int N, + uintptr_t sa, float sdn, uintptr_t stream) { + hyvla_ffn_dn_res_bf16( + reinterpret_cast(a), reinterpret_cast(dn), + reinterpret_cast(res), reinterpret_cast(y), + M, K, N, reinterpret_cast(sa), sdn, to_stream(stream)); + }, + py::arg("a"), py::arg("dn"), py::arg("res"), py::arg("y"), py::arg("M"), + py::arg("K"), py::arg("N"), py::arg("sa"), py::arg("sdn"), + py::arg("stream") = 0, + "Hy-VLA FFN kernel B: dn-GEMM + residual -> bf16 (Thor)."); +#endif +#endif + #ifdef ENABLE_LINGBOT #include "kernels/lingbot_bindings.inc" // m.def("lingbot_...", &lingbot_...) for the LingBot-VLA model #endif diff --git a/csrc/kernels/hyvla_ffn_fp8_thor.cu b/csrc/kernels/hyvla_ffn_fp8_thor.cu new file mode 100644 index 000000000..5596d5d92 --- /dev/null +++ b/csrc/kernels/hyvla_ffn_fp8_thor.cu @@ -0,0 +1,224 @@ +// FlashRT — Hy-VLA denoise FFN megakernel (Thor SM110, plain FP8 MMA). +// +// The occupancy-preserving whole-FFN fusion (kept all 20 SMs busy) that the +// single-CTA quant PoC lacked. Two persistent tiled FP8 GEMMs, m16n8k32 e4m3 +// (the plain mma.sync — Thor rejects sm_120's .kind::f8f6f4), cp.async 2-stage. +// Adapted from the sm_120 action_ffn_v6t template: +// A: gu-GEMM (K=1024 -> 2*Nout=4096) with gate/up dual-accumulator, then +// silu(gate)*up -> bf16 act (M, Nout=2048). +// B: dn-GEMM (K=2048 -> N=1024) + residual -> bf16 y. +// Dynamic FP8: activation scale is a device pointer read at launch (graph-safe, +// like fp8_nn_dev); weight scale is a host constant (fixed at quantize time). +// The act between A and B is dynamically requantized by the caller +// (quantize_fp8_device) so no grid barrier / static act scale is needed. +// Verified vs a float FFN reference: A cos 0.999999, B cos 0.9995. +// +// Static shared memory (<48KB) so NO cudaFuncSetAttribute is needed — the +// launch has no host attribute call, keeping the captured path clean. +#include +#include +#include + +namespace { + +constexpr int NUM_WARPS = 4; +constexpr int THREADS = NUM_WARPS * 32; +constexpr int M_ROWS = 16; +constexpr int BLOCK_N = 32; + +__device__ __forceinline__ void mma_e4m3(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.m16n8k32.row.col.f32.e4m3.e4m3.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__ void cp16(uint32_t s, const uint8_t* g) { + int b = (g == nullptr) ? 0 : 16; + asm volatile("cp.async.ca.shared.global [%0],[%1],16,%2;\n" :: "r"(s), "l"(g), "r"(b)); +} +__device__ __forceinline__ uint32_t sma(const void* p) { + return (uint32_t)__cvta_generic_to_shared(p); +} +__device__ __forceinline__ float siluf(float x) { return x / (1.0f + __expf(-x)); } + +// ── Kernel A: gu-GEMM gate/up + silu_mul -> bf16 act ── +template +__global__ void __launch_bounds__(THREADS, 8) ffn_A( + const __nv_fp8_e4m3* __restrict__ x, const __nv_fp8_e4m3* __restrict__ gu, + __nv_bfloat16* __restrict__ act, int M, int K, int Nout, + const float* __restrict__ sx_ptr, float sgu) { + constexpr int PAD = BK + 16; + __shared__ __align__(16) uint8_t As[2 * M_ROWS * PAD]; + __shared__ __align__(16) uint8_t Bg[2 * BLOCK_N * PAD]; + __shared__ __align__(16) uint8_t Bu[2 * BLOCK_N * PAD]; + const int cta = blockIdx.x, m_base = blockIdx.y * M_ROWS, t = threadIdx.x; + if (cta >= Nout / BLOCK_N) return; + const int nb = cta * BLOCK_N, warp = t / 32, lane = t % 32, l = lane % 4, h = lane / 4; + constexpr int NA = (BLOCK_N / 8 + NUM_WARPS - 1) / NUM_WARPS; + constexpr int KA = BK / 32; + auto issue = [&](int st, int kb) { + constexpr int AR = THREADS / (BK / 16), AI = (M_ROWS + AR - 1) / AR; + #pragma unroll + for (int it = 0; it < AI; ++it) { + int idx = it * THREADS + t, ra = idx / (BK / 16), ko = (idx & (BK / 16 - 1)) * 16; + if (ra < M_ROWS) { + const uint8_t* s = nullptr; int rg = m_base + ra; + if (rg < M && kb + ko < K) s = (const uint8_t*)&x[rg * K + kb + ko]; + cp16(sma(&As[st * M_ROWS * PAD + ra * PAD + ko]), s); + } + } + constexpr int BT = BLOCK_N * BK / 16, BI = (BT + THREADS - 1) / THREADS; + #pragma unroll + for (int it = 0; it < BI; ++it) { + int idx = it * THREADS + t, rb = idx / (BK / 16), ko = (idx & (BK / 16 - 1)) * 16; + if (rb < BLOCK_N) { + int ng = nb + rb; const uint8_t* sg = nullptr; const uint8_t* su = nullptr; + if (ng < Nout && kb + ko < K) { + sg = (const uint8_t*)&gu[ng * K + kb + ko]; + su = (const uint8_t*)&gu[(Nout + ng) * K + kb + ko]; + } + cp16(sma(&Bg[st * BLOCK_N * PAD + rb * PAD + ko]), sg); + cp16(sma(&Bu[st * BLOCK_N * PAD + rb * PAD + ko]), su); + } + } + }; + float ag[NA][4] = {0}, au[NA][4] = {0}; + int is = 0, ki = 0; issue(0, 0); asm volatile("cp.async.commit_group;\n" ::); is = 1; ki = BK; + int stg = 0; + for (int kb = 0; kb < K; kb += BK) { + if (ki < K) issue(is, ki); asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group 1;\n" ::); __syncthreads(); + #pragma unroll + for (int kk = 0; kk < KA; ++kk) { + int k0 = kk * 32 + 4 * l, k2 = k0 + 16, r0 = h, r1 = h + 8; + uint32_t A0 = *(uint32_t*)&As[stg * M_ROWS * PAD + r0 * PAD + k0]; + uint32_t A1 = *(uint32_t*)&As[stg * M_ROWS * PAD + r1 * PAD + k0]; + uint32_t A2 = *(uint32_t*)&As[stg * M_ROWS * PAD + r0 * PAD + k2]; + uint32_t A3 = *(uint32_t*)&As[stg * M_ROWS * PAD + r1 * PAD + k2]; + #pragma unroll + for (int na = 0; na < NA; ++na) { + int cn = warp * NA * 8 + na * 8 + h; + uint32_t G0 = *(uint32_t*)&Bg[stg * BLOCK_N * PAD + cn * PAD + k0]; + uint32_t G1 = *(uint32_t*)&Bg[stg * BLOCK_N * PAD + cn * PAD + k2]; + uint32_t U0 = *(uint32_t*)&Bu[stg * BLOCK_N * PAD + cn * PAD + k0]; + uint32_t U1 = *(uint32_t*)&Bu[stg * BLOCK_N * PAD + cn * PAD + k2]; + mma_e4m3(ag[na][0], ag[na][1], ag[na][2], ag[na][3], A0, A1, A2, A3, G0, G1); + mma_e4m3(au[na][0], au[na][1], au[na][2], au[na][3], A0, A1, A2, A3, U0, U1); + } + } + __syncthreads(); // WAR: finish reading stage `stg` before it is reissued + stg ^= 1; is ^= 1; ki += BK; + } + asm volatile("cp.async.wait_all;\n" ::); + float ds = (*sx_ptr) * sgu; + #pragma unroll + for (int na = 0; na < NA; ++na) { + int cb = nb + warp * NA * 8 + na * 8 + 2 * l; + #pragma unroll + for (int j = 0; j < 4; ++j) { + int row = (j < 2) ? h : (h + 8), col = cb + (j & 1), rg = m_base + row; + if (rg < M && col < Nout) { + float g = ag[na][j] * ds, u = au[na][j] * ds; + act[rg * Nout + col] = __float2bfloat16(siluf(g) * u); + } + } + } +} + +// ── Kernel B: dn-GEMM + residual -> bf16 ── +template +__global__ void __launch_bounds__(THREADS, 8) ffn_B( + const __nv_fp8_e4m3* __restrict__ a, const __nv_fp8_e4m3* __restrict__ dn, + const __nv_bfloat16* __restrict__ res, __nv_bfloat16* __restrict__ y, + int M, int K, int N, const float* __restrict__ sa_ptr, float sdn) { + constexpr int PAD = BK + 16; + __shared__ __align__(16) uint8_t As[2 * M_ROWS * PAD]; + __shared__ __align__(16) uint8_t Bs[2 * BLOCK_N * PAD]; + const int cta = blockIdx.x, m_base = blockIdx.y * M_ROWS, t = threadIdx.x; + if (cta >= N / BLOCK_N) return; + const int nb = cta * BLOCK_N, warp = t / 32, lane = t % 32, l = lane % 4, h = lane / 4; + constexpr int NA = (BLOCK_N / 8 + NUM_WARPS - 1) / NUM_WARPS; + constexpr int KA = BK / 32; + auto issue = [&](int st, int kb) { + constexpr int AR = THREADS / (BK / 16), AI = (M_ROWS + AR - 1) / AR; + #pragma unroll + for (int it = 0; it < AI; ++it) { + int idx = it * THREADS + t, ra = idx / (BK / 16), ko = (idx & (BK / 16 - 1)) * 16; + if (ra < M_ROWS) { + const uint8_t* s = nullptr; int rg = m_base + ra; + if (rg < M && kb + ko < K) s = (const uint8_t*)&a[rg * K + kb + ko]; + cp16(sma(&As[st * M_ROWS * PAD + ra * PAD + ko]), s); + } + } + constexpr int BT = BLOCK_N * BK / 16, BI = (BT + THREADS - 1) / THREADS; + #pragma unroll + for (int it = 0; it < BI; ++it) { + int idx = it * THREADS + t, rb = idx / (BK / 16), ko = (idx & (BK / 16 - 1)) * 16; + if (rb < BLOCK_N) { + const uint8_t* s = nullptr; int ng = nb + rb; + if (ng < N && kb + ko < K) s = (const uint8_t*)&dn[ng * K + kb + ko]; + cp16(sma(&Bs[st * BLOCK_N * PAD + rb * PAD + ko]), s); + } + } + }; + float acc[NA][4] = {0}; + int is = 0, ki = 0; issue(0, 0); asm volatile("cp.async.commit_group;\n" ::); is = 1; ki = BK; + int stg = 0; + for (int kb = 0; kb < K; kb += BK) { + if (ki < K) issue(is, ki); asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.wait_group 1;\n" ::); __syncthreads(); + #pragma unroll + for (int kk = 0; kk < KA; ++kk) { + int k0 = kk * 32 + 4 * l, k2 = k0 + 16, r0 = h, r1 = h + 8; + uint32_t A0 = *(uint32_t*)&As[stg * M_ROWS * PAD + r0 * PAD + k0]; + uint32_t A1 = *(uint32_t*)&As[stg * M_ROWS * PAD + r1 * PAD + k0]; + uint32_t A2 = *(uint32_t*)&As[stg * M_ROWS * PAD + r0 * PAD + k2]; + uint32_t A3 = *(uint32_t*)&As[stg * M_ROWS * PAD + r1 * PAD + k2]; + #pragma unroll + for (int na = 0; na < NA; ++na) { + int cn = warp * NA * 8 + na * 8 + h; + uint32_t B0 = *(uint32_t*)&Bs[stg * BLOCK_N * PAD + cn * PAD + k0]; + uint32_t B1 = *(uint32_t*)&Bs[stg * BLOCK_N * PAD + cn * PAD + k2]; + mma_e4m3(acc[na][0], acc[na][1], acc[na][2], acc[na][3], A0, A1, A2, A3, B0, B1); + } + } + __syncthreads(); // WAR: finish reading stage `stg` before it is reissued + stg ^= 1; is ^= 1; ki += BK; + } + asm volatile("cp.async.wait_all;\n" ::); + float ds = (*sa_ptr) * sdn; + #pragma unroll + for (int na = 0; na < NA; ++na) { + int cb = nb + warp * NA * 8 + na * 8 + 2 * l; + #pragma unroll + for (int j = 0; j < 4; ++j) { + int row = (j < 2) ? h : (h + 8), col = cb + (j & 1), rg = m_base + row; + if (rg < M && col < N) + y[rg * N + col] = __float2bfloat16(acc[na][j] * ds + __bfloat162float(res[rg * N + col])); + } + } +} + +} // namespace + +extern "C" void hyvla_ffn_gu_silu_bf16( + const void* x_fp8, const void* gu_w_fp8, void* act_bf16, + int M, int K, int Nout, const void* sx_ptr, float sgu, cudaStream_t stream) { + int mt = (M + M_ROWS - 1) / M_ROWS; + dim3 grid(Nout / BLOCK_N, mt); + ffn_A<128><<>>( + (const __nv_fp8_e4m3*)x_fp8, (const __nv_fp8_e4m3*)gu_w_fp8, + (__nv_bfloat16*)act_bf16, M, K, Nout, (const float*)sx_ptr, sgu); +} + +extern "C" void hyvla_ffn_dn_res_bf16( + const void* act_fp8, const void* dn_w_fp8, const void* residual, void* y_bf16, + int M, int K, int N, const void* sa_ptr, float sdn, cudaStream_t stream) { + int mt = (M + M_ROWS - 1) / M_ROWS; + dim3 grid(N / BLOCK_N, mt); + ffn_B<256><<>>( + (const __nv_fp8_e4m3*)act_fp8, (const __nv_fp8_e4m3*)dn_w_fp8, + (const __nv_bfloat16*)residual, (__nv_bfloat16*)y_bf16, + M, K, N, (const float*)sa_ptr, sdn); +} diff --git a/csrc/kernels/hyvla_ffn_fp8_thor.cuh b/csrc/kernels/hyvla_ffn_fp8_thor.cuh new file mode 100644 index 000000000..e75e2dadc --- /dev/null +++ b/csrc/kernels/hyvla_ffn_fp8_thor.cuh @@ -0,0 +1,14 @@ +// FlashRT — Hy-VLA denoise FFN megakernel declaration (Thor SM110, plain FP8 MMA). +#pragma once +#include + +// A: gu-GEMM (K->2*Nout) with gate/up dual-accumulator + silu_mul -> bf16 act (M,Nout). +// descale = (*sx) * sgu ; gu_w is fp8 (2*Nout, K): rows[0:Nout)=gate, [Nout:2Nout)=up. +extern "C" void hyvla_ffn_gu_silu_bf16( + const void* x_fp8, const void* gu_w_fp8, void* act_bf16, + int M, int K, int Nout, const void* sx_ptr, float sgu, cudaStream_t stream); + +// B: dn-GEMM (K->N) + residual -> bf16 y (M,N). descale = (*sa) * sdn ; dn_w fp8 (N,K). +extern "C" void hyvla_ffn_dn_res_bf16( + const void* act_fp8, const void* dn_w_fp8, const void* residual, void* y_bf16, + int M, int K, int N, const void* sa_ptr, float sdn, cudaStream_t stream); diff --git a/csrc/kernels/hyvla_fused_thor.cu b/csrc/kernels/hyvla_fused_thor.cu new file mode 100644 index 000000000..b6161fdb5 --- /dev/null +++ b/csrc/kernels/hyvla_fused_thor.cu @@ -0,0 +1,95 @@ +// ================================================================ +// FlashRT — Hy-VLA fused attention-prep megakernel (Thor SM110) +// +// One launch replaces ~11 tiny torch ops per attention block: +// split(qkv) → RoPE(q) → RoPE(k) → QK-Norm(q) → QK-Norm(k) +// → write K,V into the layer KV-cache at row `off`. +// +// Hy-VLA order is RoPE-FIRST then QK-Norm (RMSNorm over head_dim), +// the reverse of the existing qwen3 fused kernel — hence a bespoke +// kernel. rotate_half (NeoX) convention; cos/sin are per-position +// (shared across heads). GQA: nq query heads, nkv KV heads. +// +// Layouts (all bf16, contiguous): +// qkv : (S, (nq+2*nkv)*hd) +// cos/sin: (S, hd) qn_w/kn_w: (hd) +// q_out : (nq, S, hd) == (1,nq,S,hd) for SDPA +// kbuf/vbuf: (nkv, S_tot, hd) (one layer slice); write row off+s +// ================================================================ + +#include "common.cuh" +#include + +__global__ void hyvla_rope_qknorm_kvwrite_bf16_kernel( + const __nv_bfloat16* __restrict__ qkv, + const __nv_bfloat16* __restrict__ cos, + const __nv_bfloat16* __restrict__ sin, + const __nv_bfloat16* __restrict__ qn_w, + const __nv_bfloat16* __restrict__ kn_w, + __nv_bfloat16* __restrict__ q_out, + __nv_bfloat16* __restrict__ kbuf, + __nv_bfloat16* __restrict__ vbuf, + int S, int nq, int nkv, int hd, int S_tot, int off, float eps, int kv_rep) +{ + const int s = blockIdx.x; + const int tid = threadIdx.x; // 0..hd-1 + if (tid >= hd) return; + const int Dqkv = (nq + 2 * nkv) * hd; + const int half = hd >> 1; + const int partner = (tid < half) ? (tid + half) : (tid - half); + const float sgn = (tid < half) ? -1.0f : 1.0f; // rotate_half sign + __shared__ float shared[64]; + + const float cs = to_f32<__nv_bfloat16>(cos[s * hd + tid]); + const float sn = to_f32<__nv_bfloat16>(sin[s * hd + tid]); + const float qnw = to_f32<__nv_bfloat16>(qn_w[tid]); + const float knw = to_f32<__nv_bfloat16>(kn_w[tid]); + + // Q heads: RoPE then QK-Norm -> q_out + for (int h = 0; h < nq; ++h) { + __syncthreads(); + const __nv_bfloat16* base = qkv + s * Dqkv + h * hd; + float x = to_f32<__nv_bfloat16>(base[tid]); + float xp = to_f32<__nv_bfloat16>(base[partner]); + float roped = x * cs + sgn * xp * sn; + float ss = block_reduce_sum(roped * roped, shared); + float rms = rsqrtf(ss / hd + eps); + q_out[h * S * hd + s * hd + tid] = from_f32<__nv_bfloat16>(roped * rms * qnw); + } + // K heads: RoPE then QK-Norm -> kbuf[off+s] + for (int kh = 0; kh < nkv; ++kh) { + __syncthreads(); + const __nv_bfloat16* base = qkv + s * Dqkv + nq * hd + kh * hd; + float x = to_f32<__nv_bfloat16>(base[tid]); + float xp = to_f32<__nv_bfloat16>(base[partner]); + float roped = x * cs + sgn * xp * sn; + float ss = block_reduce_sum(roped * roped, shared); + float rms = rsqrtf(ss / hd + eps); + __nv_bfloat16 kval = from_f32<__nv_bfloat16>(roped * rms * knw); + for (int rr = 0; rr < kv_rep; ++rr) + kbuf[(kh * kv_rep + rr) * S_tot * hd + (off + s) * hd + tid] = kval; + } + // V heads: raw copy -> vbuf[off+s] (replicated kv_rep times for GQA) + for (int kh = 0; kh < nkv; ++kh) { + const __nv_bfloat16* base = qkv + s * Dqkv + (nq + nkv) * hd + kh * hd; + __nv_bfloat16 vval = base[tid]; + for (int rr = 0; rr < kv_rep; ++rr) + vbuf[(kh * kv_rep + rr) * S_tot * hd + (off + s) * hd + tid] = vval; + } +} + +extern "C" void hyvla_rope_qknorm_kvwrite_bf16( + const void* qkv, const void* cos, const void* sin, + const void* qn_w, const void* kn_w, + void* q_out, void* kbuf, void* vbuf, + int S, int nq, int nkv, int hd, int S_tot, int off, float eps, + int kv_rep, cudaStream_t stream) +{ + dim3 grid(S); + dim3 block(hd); + hyvla_rope_qknorm_kvwrite_bf16_kernel<<>>( + (const __nv_bfloat16*)qkv, (const __nv_bfloat16*)cos, (const __nv_bfloat16*)sin, + (const __nv_bfloat16*)qn_w, (const __nv_bfloat16*)kn_w, + (__nv_bfloat16*)q_out, (__nv_bfloat16*)kbuf, (__nv_bfloat16*)vbuf, + S, nq, nkv, hd, S_tot, off, eps, kv_rep < 1 ? 1 : kv_rep); +} diff --git a/csrc/kernels/hyvla_fused_thor.cuh b/csrc/kernels/hyvla_fused_thor.cuh new file mode 100644 index 000000000..7898d3473 --- /dev/null +++ b/csrc/kernels/hyvla_fused_thor.cuh @@ -0,0 +1,10 @@ +// FlashRT — Hy-VLA fused attention-prep megakernel declaration. +#pragma once +#include + +extern "C" void hyvla_rope_qknorm_kvwrite_bf16( + const void* qkv, const void* cos, const void* sin, + const void* qn_w, const void* kn_w, + void* q_out, void* kbuf, void* vbuf, + int S, int nq, int nkv, int hd, int S_tot, int off, float eps, + int kv_rep, cudaStream_t stream); diff --git a/csrc/kernels/hyvla_quant_fp8_thor.cu b/csrc/kernels/hyvla_quant_fp8_thor.cu new file mode 100644 index 000000000..4f2a1c6a4 --- /dev/null +++ b/csrc/kernels/hyvla_quant_fp8_thor.cu @@ -0,0 +1,51 @@ +// FlashRT — Hy-VLA single-CTA dynamic per-tensor FP8 quant (Thor SM110). +// +// Collapses quantize_fp8_device's 4 graph nodes (memset + absmax + +// compute_scale + quantize) into ONE single-block kernel. For the expert +// denoise tower (M=41) every activation tensor is <=84K elements, so a single +// CTA holds it: pass 1 reduces the per-tensor amax, pass 2 casts to e4m3. No +// cross-block atomics -> fully deterministic -> trivially CUDA-graph-safe. +// +// Numerics are bit-identical to quantize_fp8_device: +// amax = max|x| over n>>1 bf16 pairs (odd tail dropped, as in reference) +// scale = max(amax/448, 1e-12) +// out[i]= e4m3( clamp(x[i] * (1/scale), +-448) ) +#include "common.cuh" +#include + +__global__ void hyvla_quant_fp8_dyn_bf16_kernel( + const __nv_bfloat16* __restrict__ x, __nv_fp8_e4m3* __restrict__ out, + float* __restrict__ scale, int n) +{ + const __nv_bfloat162* x2 = reinterpret_cast(x); + const int n2 = n >> 1; + __shared__ float red[32]; + + float local_max = 0.0f; + for (int i = threadIdx.x; i < n2; i += blockDim.x) { + __nv_bfloat162 v = x2[i]; + local_max = fmaxf(local_max, + fmaxf(fabsf(to_f32<__nv_bfloat16>(v.x)), + fabsf(to_f32<__nv_bfloat16>(v.y)))); + } + float amax = block_reduce_max(local_max, red); // broadcast to all lanes + float sc = fmaxf(amax / 448.0f, 1e-12f); + if (threadIdx.x == 0) *scale = sc; + const float inv_s = 1.0f / sc; + + for (int i = threadIdx.x; i < n2; i += blockDim.x) { + __nv_bfloat162 v = x2[i]; + float v0 = to_f32<__nv_bfloat16>(v.x) * inv_s; + float v1 = to_f32<__nv_bfloat16>(v.y) * inv_s; + out[2 * i] = __nv_fp8_e4m3(fminf(fmaxf(v0, -448.0f), 448.0f)); + out[2 * i + 1] = __nv_fp8_e4m3(fminf(fmaxf(v1, -448.0f), 448.0f)); + } +} + +extern "C" void hyvla_quant_fp8_dyn_bf16( + const void* x, void* out, float* scale, int n, cudaStream_t stream) +{ + hyvla_quant_fp8_dyn_bf16_kernel<<<1, 512, 0, stream>>>( + reinterpret_cast(x), + reinterpret_cast<__nv_fp8_e4m3*>(out), scale, n); +} diff --git a/csrc/kernels/hyvla_quant_fp8_thor.cuh b/csrc/kernels/hyvla_quant_fp8_thor.cuh new file mode 100644 index 000000000..72dd3836e --- /dev/null +++ b/csrc/kernels/hyvla_quant_fp8_thor.cuh @@ -0,0 +1,11 @@ +// FlashRT — Hy-VLA single-CTA dynamic per-tensor FP8 quant declaration. +#pragma once +#include + +// Dynamic per-tensor FP8 (e4m3) quant of a small bf16 tensor in ONE launch. +// Matches quantize_fp8_device numerics exactly but collapses its +// memset+absmax+compute_scale+quantize (4 nodes) into a single deterministic +// single-block kernel — for the denoise expert tower where M<=64 keeps the +// whole tensor comfortably inside one CTA. Graph-safe (no atomics, 1 block). +extern "C" void hyvla_quant_fp8_dyn_bf16( + const void* x, void* out, float* scale, int n, cudaStream_t stream); diff --git a/csrc/kernels/hyvla_vit_fuse.cu b/csrc/kernels/hyvla_vit_fuse.cu new file mode 100644 index 000000000..3a2752bc2 --- /dev/null +++ b/csrc/kernels/hyvla_vit_fuse.cu @@ -0,0 +1,78 @@ +// ================================================================ +// FlashRT — Hy-VLA Orin ViT fusion kernels +// +// hyvla_vit_add_layer_norm_bf16: +// residual += x_add (bf16 round, in-place — matches torch add) +// out = LayerNorm(residual) +// Fuses the ViT post-attention residual add with the following LayerNorm +// (and, across blocks, the previous block's MLP residual with the entry +// LayerNorm), removing one full read+write pass per site. +// +// Precision contract: the add rounds to bf16 exactly like torch's +// elementwise add; the LayerNorm is bit-identical to this repo's +// layer_norm_kernel (fp32 two-pass mean/var, rsqrtf, single bf16 round). +// ================================================================ + +#include "hyvla_vit_fuse.cuh" +#include "common.cuh" + +__global__ void hyvla_vit_add_layer_norm_bf16_kernel( + __nv_bfloat16* __restrict__ residual, + const __nv_bfloat16* __restrict__ x_add, + const __nv_bfloat16* __restrict__ ln_weight, + const __nv_bfloat16* __restrict__ ln_bias, + __nv_bfloat16* __restrict__ out, + int dim, float eps) { + extern __shared__ float partial[]; + + int row = blockIdx.x; + using T2 = __nv_bfloat162; + T2* res2 = reinterpret_cast(residual + (size_t)row * dim); + const T2* add2 = reinterpret_cast(x_add + (size_t)row * dim); + const T2* w2 = reinterpret_cast(ln_weight); + const T2* b2 = reinterpret_cast(ln_bias); + T2* out2 = reinterpret_cast(out + (size_t)row * dim); + int dim2 = dim >> 1; + + // Pass 1: residual = bf16(residual + x_add) to global, sum for mean. + // Re-reading residual from global in passes 2/3 keeps this bit-equal + // to running torch add then layer_norm_kernel sequentially. + float local_sum = 0.0f; + for (int i = threadIdx.x; i < dim2; i += blockDim.x) { + T2 rv = res2[i], av = add2[i]; + __nv_bfloat16 r0 = from_f32<__nv_bfloat16>(to_f32(rv.x) + to_f32(av.x)); + __nv_bfloat16 r1 = from_f32<__nv_bfloat16>(to_f32(rv.y) + to_f32(av.y)); + res2[i] = make_packed2<__nv_bfloat16>(r0, r1); + local_sum += to_f32(r0) + to_f32(r1); + } + float mean = block_reduce_sum(local_sum, partial) / dim; + + float local_var = 0.0f; + for (int i = threadIdx.x; i < dim2; i += blockDim.x) { + T2 val = res2[i]; + float d0 = to_f32(val.x) - mean, d1 = to_f32(val.y) - mean; + local_var += d0 * d0 + d1 * d1; + } + float inv_std = rsqrtf(block_reduce_sum(local_var, partial) / dim + eps); + + for (int i = threadIdx.x; i < dim2; i += blockDim.x) { + T2 val = res2[i], wv = w2[i], bv = b2[i]; + float n0 = (to_f32(val.x) - mean) * inv_std * to_f32(wv.x) + to_f32(bv.x); + float n1 = (to_f32(val.y) - mean) * inv_std * to_f32(wv.y) + to_f32(bv.y); + out2[i] = make_packed2<__nv_bfloat16>( + from_f32<__nv_bfloat16>(n0), from_f32<__nv_bfloat16>(n1)); + } +} + +extern "C" void hyvla_vit_add_layer_norm_bf16( + void* residual, const void* x_add, + const void* ln_weight, const void* ln_bias, + void* out, int rows, int dim, float eps, cudaStream_t stream) { + int smem = 256 * sizeof(float); + hyvla_vit_add_layer_norm_bf16_kernel<<>>( + reinterpret_cast<__nv_bfloat16*>(residual), + reinterpret_cast(x_add), + reinterpret_cast(ln_weight), + reinterpret_cast(ln_bias), + reinterpret_cast<__nv_bfloat16*>(out), dim, eps); +} diff --git a/csrc/kernels/hyvla_vit_fuse.cuh b/csrc/kernels/hyvla_vit_fuse.cuh new file mode 100644 index 000000000..497b230a9 --- /dev/null +++ b/csrc/kernels/hyvla_vit_fuse.cuh @@ -0,0 +1,8 @@ +// FlashRT — Hy-VLA Orin ViT fusion kernel declarations. +#pragma once +#include + +extern "C" void hyvla_vit_add_layer_norm_bf16( + void* residual, const void* x_add, + const void* ln_weight, const void* ln_bias, + void* out, int rows, int dim, float eps, cudaStream_t stream); diff --git a/docs/hyvla05_thor_sm110.md b/docs/hyvla05_thor_sm110.md new file mode 100644 index 000000000..f2f92040e --- /dev/null +++ b/docs/hyvla05_thor_sm110.md @@ -0,0 +1,169 @@ +# Hy-Embodied-0.5-VLA — Thor SM110 (Authoritative Document) + +> **Production configuration: `flash_rt.load_model(ckpt, config="hyvla", framework="torch")` → +> `HyVLATorchFrontendThor(ckpt, use_fp8=True, use_fused=True[, use_autotune=True])`** +> **E2E 159.6 ms (`+use_autotune` 158.3 ms; HF/transformers eager ~930 ms, 5.8–5.9x), +> action cosine 0.999706 vs. HF/transformers eager (same fixed noise), CUDA-graph bitwise reproducible.** +> +> Framework-native: the runtime **does not import any upstream training code** and **does not use torch.compile/Inductor**. +> Composition: CUDA Graph (prefill + 10 denoise steps in a single graph) + dynamic per-tensor FP8 (calibration-free, graph-safe) +> + fused megakernel (rope+qknorm+kvwrite, KV pre-expanded for GQA) + memory-efficient SDPA +> + ViT trailing-stage history-frame drop + per-shape FP8 GEMM autotune. +> +> ```bash +> cmake --build build --target flash_rt_kernels # fused kernels (first build) +> PYTHONPATH=. python3 tests/test_thor_hyvla05_e2e_check.py --fp8 --fused # precision gate cos >= 0.999 +> PYTHONPATH=. python3 tests/test_thor_hyvla05_graphsafe.py --fp8 --fused # graph-safety bitwise gate +> PYTHONPATH=. python3 tests/test_thor_hyvla05_stageprof.py --fp8 --fused # per-stage latency +> ``` + +## Key Takeaway + +Under single-request latency, all three major stages (ViT / prefill / denoise) on Thor (20 SMs, 243 GB/s) +are **memory/latency-bound at batch=1, achieving only ~35% of the ideal roofline — this is already the +practical floor for a per-operator hand-written path**. Further gains come not from low-bit quantization +or single-block fusion (both measured ineffective/harmful), but from **whole-layer/whole-model fusion** +(an early Inductor prototype `--v4` in this repo reached 98.8 ms via automatic fusion — the automated +version of this approach — but suffered 30–60 s recompilation per prompt, cache bloat degradation, and +silent dynamo fallback; it is not framework-grade. The production equivalent = hand-written persistent +megakernels or TRT/MLIR-TRT), **reducing flow steps (distillation)**, or **async chunking to hide +latency** — all of which require model or engineering-form changes. +Industry anchor: NVIDIA achieves 44 ms / 23 Hz on Pi0.5 with hand-written kernels + MLIR-TRT +(with a lighter visual workload than this model). + +## Components + +| File | Role | +|---|---| +| `flash_rt/frontends/torch/hyvla_thor.py` | `HyVLATorchFrontendThor`: tokenization, image preprocessing, prefix assembly, segmented prefix mask + prefix-LM suffix mask, bf16-round RoPE table, time embedding, CUDA Graph cache (keyed by `(S_p,n_vis)`), FP8/FP4/FFN-mega weight quantization, `set_prompt`/`infer`/`predict_actions`. Flags: `use_fp8/use_fused/use_autotune` (production) + `use_fp8_vit/use_fp4/use_fused_quant/use_ffn_mega` (opt-in diagnostics, off by default) | +| `flash_rt/models/hyvla/pipeline_thor.py` | `HyVLAThorBF16Pipeline`: `vit_forward` (27 layers including 6 spacetime layers), `merger_forward`, `prefill` (32-layer MoT), `denoise` (32-layer expert, 10-step Euler), `_fp8_gemm` (+per-shape autotune), `_ffn_mega_bf16`, efficient-SDPA `_attn` | +| `flash_rt/frontends/torch/_hyvla_thor_spec.py` | Declarative `ModelWeightSpec` (ViT/merger/VLM dual-branch/expert/action head/tied lm_head) | +| `csrc/kernels/hyvla_fused_thor.cu`(+`.cuh`) | Production fused kernel `hyvla_rope_qknorm_kvwrite_bf16` (split+rope+qknorm+kv-write = 1 launch, `kv_rep` pre-expands GQA) | +| `csrc/kernels/hyvla_quant_fp8_thor.cu` | Single-CTA dynamic FP8 quantization (`use_fused_quant` diagnostic; measured net loss, see dead ends) | +| `csrc/kernels/hyvla_ffn_fp8_thor.cu` | FFN megakernel (gu+silu_mul / dn+residual, `use_ffn_mega` diagnostic; measured neutral, see dead ends); reusable Thor plain-FP8-MMA GEMM reference | +| `flash_rt/{executors/torch_weights.py,hardware/__init__.py,api.py,configs/hyvla.yaml}` | `ToBf16` transform / `_PIPELINE_MAP` registration / config allowlist / metadata | +| `tests/test_thor_hyvla05_{baseline,tower_check,vit_check,e2e_check}.py` | Precision oracles + segmented/full-chain precision gates (`--fp8 --fused --fp4`) | +| `tests/test_thor_hyvla05_{bench,stageprof,graphsafe}.py` | Latency benchmark / per-stage profiling / graph-safety bitwise gate | +| `tests/test_thor_hyvla05_{gemm_ceiling,vit_prof,fp4_sf_check}.py` | GEMM ceiling / ViT internal profiling / NVFP4 SF validation | + +## Precision Gates (all vs. HF/transformers eager, same fixed noise; `e2e_check`/`tower_check`/`vit_check`) + +| Checkpoint | cosine | +|---|---| +| ViT + merger (147 visual tokens) | 0.999844 | +| Dual-tower prefill + denoise + action head (fed the original model's prefix) | 0.999978 | +| Full-chain native BF16 (load_model path) | 0.999910 | +| **Full-chain production (fp8 + fused + efficient-SDPA + autotune)** | **0.999706** | + +Graph-safety gate (`graphsafe.py`): `use_graph=True vs False` bitwise identical (`max|delta|=0`) + stable replay, +verified for both the ViT graph and the main graph. + +## Key Mechanisms / Correctness Pitfalls (Highest Reuse Value) + +1. **Architecture = Pi0 action head + SigLIP-so400m-isomorphic ViT + HunYuan MoT dual tower**. ViT 1152h/hd72/patch16, learned pos_embed (128x128 bilinear rescale) + 6 spacetime causal temporal attention layers (block {3,7,11,15,19,23}, 6 history frames). VLM 2048h/6144i/32L/GQA 16Q-4KV/hd128; expert 1024h/2048i/32L. +2. **RMSNorm is pure-weight (no `1+w`), fp32 upcast, eps=1e-5** (ViT LayerNorm has bias, eps=1e-6). Using `1+w` incorrectly → cos drops to ~0.5. +3. **QK-Norm is applied after RoPE** (RMSNorm over hd=128), opposite to off-the-shelf norm-then-rope kernels → requires a custom fused kernel. +4. **RoPE `inv_freq` must round-trip through bf16**: the original model stores `rotary_emb.inv_freq` in bf16 via `.to(bf16)`, and rounding accumulates with position; using full-precision inv_freq → action cos only reaches ~0.96. This was the most time-consuming correctness issue to locate. +5. **MoT `_v` static routing**: prefix is reordered by modality into contiguous `[vision|text]` slices → two pointer-offset GEMMs per operator, no gather; the expert tower routes entirely through `_v`, reusing the VLM's QK-norm weights. +6. **Determinism / reference pitfall (biggest trap)**: monkeypatching `sample_noise` to inject fixed noise **did not take effect** (run-to-run max|delta|=2.6e-2 was the signal) — the reference became random noise → E2E falsely measured 0.958. Switching to **explicitly passing noise** to `sample_actions` immediately yielded 0.9999. **Always verify run-to-run delta is approximately 0 before trusting fixed-noise results.** + +## Performance Optimization Ladder (all measured, same prompt/image, warmup + median) + +| Stage | E2E | Notes | +|---|---|---| +| HF/transformers eager (anchor) | ~930 ms | | +| Native BF16 eager | 303 ms | 3.07x, no upstream overhead | +| + CUDA Graph (prefill+denoise in one graph) | 248.6 ms | 3.74x, graph==eager bitwise | +| + Dynamic per-tensor FP8 (denoise, in-graph) | 225.2 ms | 4.13x, cos 0.999833 (FP8 only pays off in-graph) | +| + FP8 prefill (MoT dual-branch) | 207.0 ms | 4.49x | +| + Fused megakernel (rope+qknorm+kvwrite) | 195.5 ms | 4.76x, ~11 torch ops → 1 launch | +| + efficient-SDPA (GQA expansion + memory-efficient backend) | 172.9 ms | 5.38x, attention 41 → ~12 ms, **cos actually improved to 0.999706** (bool-mask + enable_gqa forces the slow math backend) | +| + ViT trailing-stage history-frame drop | 166.1 ms | 5.60x, per-frame independence after last spacetime block → 18 → 3 frames, mathematically equivalent | +| + KV pre-expansion for GQA (megakernel `kv_rep`) | 159.6 ms | 5.83x, zero repeat_interleave in attention | +| **+ Per-shape FP8 GEMM autotune** | **158.3 ms** | **5.9x**, per-shape `autotune_fp8_nn_dev` before capture (following the motus pattern), bitwise identical | + +## Roofline / Hardware Utilization (all measured anchors) + +Thor SM110: 20 SMs; HBM **243 GB/s**; bf16 GEMM achievable **~110 TFLOPS** (cuBLAS large square matrices; +raw-mma silicon >=378 → cuBLAS bf16 only utilizes ~30% of silicon); fp8 `fp8_nn_dev` **~270 TFLOPS**. + +| Component | GEMM FLOP | Bound by | Ideal ceiling | Measured | Utilization | +|---|---|---|---|---|---| +| ViT+merger (bf16, 3 cameras x 6 frames) | 2718 G | compute | 24.7 ms | 71.5 | **35%** | +| prefill (fp8, M=240, x1) | 756 G | memory/weight | 12.7 ms | 33 | **38%** | +| denoise (fp8, M=41, x10 steps) | 333 G | memory/weight+KV | 18.2 ms | 52 | **35%** | + +All three stages reach only ~35% of the ideal ceiling, but **all sit at the practical floor for a small +GPU + small batch** (corroborated by *Demystifying VLA Inference*, arXiv 2602.18397: at batch=1 the action +head is strictly memory-bound, and the backbone becomes memory-bound at the low-bandwidth edge; high-end +GPUs achieve 73–82% of theoretical peak in measured kernels). The unified root cause: **small M cannot +saturate** — a single M=41 fp8 GEMM only achieves 34% of BW (41 rows too thin, poor L2 reuse, per-SM +ramp/tail dominates). This explains all dead ends below. + +## Dead Ends (all measured; check before writing code) + +| Approach | Result | Mechanism | +|---|---|---| +| **FP8-ViT** | Net loss (ViT graph 71.5 → 85.4 ms) | ViT utilization only 35%, not GEMM-compute-bound; FP8 only adds quantization kernel overhead | +| **Single-CTA fused quantization** (4 kernels → 1, `use_fused_quant`) | +2.6 ms regression | Single CTA occupies 1 SM; original 4 kernels spread blocks across all SMs, in-graph launch cost is ~0; fewer nodes gives no benefit, lower occupancy is harmful | +| **FFN megakernel** (`use_ffn_mega`, gu+silu / dn+res two kernels) | E2E neutral (+0.5 ms) | M=41 FFN is latency-bound, cuBLASLt is already good enough; fused silu/mul/residual is negligible | +| **FFN single-kernel fused** (grid-barrier, see two-case comparison) | 14% slower than split | Barrier serializes globally + persistent grid reduces occupancy, exceeding saved launch/HBM round-trips | +| **Elementwise fusion** (residual+norm, silu_mul) | ~0 (reverted) | At M=41, ~42K elements are too cheap, already amortized by graph | +| **FP4 (W4A16 on Thor)** | denoise 1.14x (=FP8), prefill large GEMM 2.67% but only 21% of total | Thor has no native FP4 MMA, see below | +| **grid-barrier 192 CTA launch** | Deadlock | Exceeds co-residency capacity (20 SMs ~160) → late-arriving CTAs never launch. Must use persistent grid launched at capacity | +| bool-mask + `enable_gqa` SDPA | ~41 ms (slow math fallback) | SM110 silently falls back to math backend; fixed with GQA-expansion + memory-efficient | + +**CUDA Graph safety rules (learned from bugs):** `_fp8_gemm` transients must use `torch.empty` each time +(entering the graph's private pool) — shared scratch aliases the two outputs of MoT `torch.cat`; during +capture all allocations come from the private pool and are overwritten between replays; pointer kernels +must pass `current_stream().cuda_stream` (passing 0 = default stream = not captured); FP8 weights require +(K,N) layout, `cutlass_fp4` SF requires swizzled 128x4 (wrong layout produces no error but cos collapses). + +## Why FP4 Does Not Work (confirmed via DGX Spark SM121 reports + local ptxas verification) + +- **Thor has no native FP4 tensor core MMA**: native NVFP4 requires `tcgen05.mma` + TMEM, **only available on datacenter SM100 (B200)**. + Measured: `mma.sync.kind::f8f6f4` is rejected by ptxas on sm_110; standard `mma.sync.m16n8k32 e4m3` (FP8) works. + → On edge Blackwell (Thor/Spark), FP4 can only be **W4A16**: decompress back to bf16/fp8 before computing — **zero compute speedup, only saves weight bandwidth**. +- Bandwidth savings only materialize when **weight-bandwidth-bound and saturating BW**: denoise M=41 is latency-bound (34% BW) → FP4=FP8=1.14x. + On Spark, NVFP4 measured 65 tok/s is actually **slower** than FP8 at 91 (decompression + small smem overhead). +- The only meaningful case = prefill large GEMM (M=240, gu 2.67x), but prefill is only 21% of E2E and ViT (bf16) cannot benefit. + **Therefore FP4 provides approximately no overall benefit for this model on Thor, consistent with industry edge-Blackwell findings.** Code is correct, gated behind `use_fp4` for archival. + +## Two Megakernel Implementation Comparison + +Isolated denoise FFN @ M=41 latency benchmark, three approaches: + +| Implementation | Latency | Correctness | +|---|---|---| +| per-op (cuBLASLt `fp8_nn_dev` autotuned + torch silu_mul + quant) | **61.9 us** | baseline | +| Case A (2-kernel split, boundary quantization) | **62.5 us (1.01x, tied)** | cos 0.99987 vs per-op | +| Case B (single-kernel fused + grid-barrier) | **69.8 us (1.13x, slower)** | cos 0.9994 | + +**Two independent implementations reach the same conclusion: at M=41, fusion cannot improve utilization and is even harmful.** True megakernel gains come from whole-model fusion +(Mirage/MPK, Hazy "no-bubbles" — eliminating launches + bubbles + HBM round-trips across multiple layers) or raising M (batching), not from single-block fusion. +The building blocks are in place and have high reuse value: plain-FP8-MMA layout verified on Thor at cos 0.999998, grid-barrier graph is capturable +(state pre-zeroed + self-resetting), 2-stage cp.async loop tail requires `__syncthreads` (otherwise WAR race). + +## Paths Toward <100 ms (all require architectural changes, not incremental) + +1. **Reduce flow steps** (10 → 2-4, Consistency Policy / distillation) — denoise 52 ms is the only stage that can be truly cut; requires retraining. +2. **Whole-model megakernel** (Mirage-style) or **TRT/MLIR-TRT + hand-written kernels + Q/DQ removal** (NVIDIA's actual stack for Pi0.5 → 44 ms) — the Inductor prototype `--v4` in this repo already demonstrated **98.8 ms** is achievable via automatic fusion (ViT 51.5 / denoise 29), but it recompiles per prompt, suffers 15% cache-bloat degradation, silent dynamo fallback, and is not registered in `_PIPELINE_MAP` — not framework-grade. +3. **Async action chunking** (Pi real-time chunking) — overlap inference with execution, hide latency at the system level; 158 ms does not block control. +4. Reduce history frames (ViT's largest cost) / smaller backbone / layer skipping. + +## Early Inductor Prototype (`--v4`, Historical Reference, Not Production) + +`tests/test_thor_hyvla05_fast.py --v4 --compile-vit`: `torch.compile` whole-block fusion + static FP8 + +grouped-bmm/FA2 + full-graph capture, E2E **98.8 ms (8.6–9.1x)**, cos 0.9985–0.9998. It proved the benefit +of the "whole-layer fusion" approach (ViT 51.5, denoise 29 are both below native), but due to **30–60 s +recompilation per prompt, 15% Inductor cache-bloat degradation, silent dynamo fallback, and no integration +with `load_model`**, it is not used in production. Its six measurement methodology lessons (L2 residency +illusion / no-fusion baseline illusion / nsys sum != wall clock / synchronization floor / Inductor cache +bloat / dynamo silent downgrade) and the FP4-in-Inductor verdict have been distilled into internal documentation. + +## Hardware and Model Profile (Measured) + +NVIDIA Thor cc11.0 (SM110), 20 SMs, 125.7 GB unified memory, L2 32 MB, CUDA 13.0 / driver 580, +L4T R38.2, torch 2.9.0a0, transformers 5.10.2, GPU full clock 1575 MHz. +Model: VLM tower 2048h/32L/GQA 16Q-4KV/hd128/Dff6144 (dual text+vision projections per layer); expert tower +1024h/inter2048/32L; ViT HYViT2-400M 27L/1152h/patch16/6-frame spacetime (stride 4); chunk 40; 10-step Euler. diff --git a/docs/stable_api.md b/docs/stable_api.md index 5d533a2ee..f9dacd896 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" | "cosmos3_edge" + config: str = "pi05", # "pi05" | "pi0" | "groot" | "groot_n17" | "pi0fast" | "motus" | "wan22_ti2v_5b" | "cosmos3_video" | "cosmos3_edge" | "hyvla" device=None, # reserved # Pi0-FAST-specific: decode_cuda_graph: bool = False, @@ -130,6 +130,10 @@ Returns a `VLAModel` wrapping the appropriate frontend for the detected frontend; `rtx_sm89` resolves directly to its dedicated SM89 frontend. `use_fp16=True, use_fp8=False` requests the explicit RTX reference frontend for the selected hardware. +- `config="hyvla"` (Hy-Embodied-0.5-VLA) is registered for + `framework="torch"` on `hardware="thor"`. Thor uses runtime dynamic + FP8 with fused megakernels and optional NVFP4 FFN (`use_fp4=True`). + See `docs/hyvla05_thor_sm110.md`. ### `flash_rt.VLAModel` diff --git a/flash_rt/api.py b/flash_rt/api.py index e7d613809..436176d82 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -461,12 +461,12 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, "Supported: pi0, pi05, llm, mllm") elif config not in ("pi05", "groot", "groot_n17", "pi0", "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video", - "cosmos3_edge", "nexn2", "qwen36_moe"): + "cosmos3_edge", "nexn2", "qwen36_moe", "hyvla"): raise ValueError( f"Unknown config: {config}. " f"Supported: pi05, groot, groot_n17, pi0, pi0fast, motus, " f"wan22_ti2v_5b, cosmos3_video, cosmos3_edge, nexn2, " - f"qwen36_moe") + f"qwen36_moe, hyvla") if framework not in ("torch", "jax", "jetson_pi"): raise ValueError( f"Unknown framework: {framework}. Supported: torch, jax, jetson_pi") @@ -699,14 +699,25 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, logger.info("GROOT N1.7 Thor NVFP4 tier enabled") use_fp4 = False # do not fall through to the Pi0.5 FP4 routing - # ── FP4 routing (Pi0.5 torch + Pi0.5 JAX on Thor) ── + # ── FP4 routing (Pi0.5 torch + Pi0.5 JAX on Thor, HyVLA torch on Thor) ── if use_fp4: - if config != "pi05" or framework not in ("torch", "jax") or arch != "thor": - logger.warning( - "use_fp4=True is only supported for config='pi05' with " - "framework in ('torch', 'jax') on Thor; got config='%s' " - "framework='%s' hardware='%s'. Falling back to FP8.", - config, framework, arch) + _fp4_ok = ( + (config == "pi05" and framework in ("torch", "jax") and arch == "thor") + or (config == "hyvla" and framework == "torch" and arch == "thor") + ) + if not _fp4_ok: + if config == "hyvla" and arch == "rtx_sm87": + logger.warning( + "use_fp4=True is not supported for config='hyvla' on " + "Jetson Orin SM87 (no native FP4 tensor cores). " + "Falling back to the HyVLA INT8/BF16 path.") + else: + logger.warning( + "use_fp4=True is only supported for config='pi05' with " + "framework in ('torch', 'jax') on Thor, or config='hyvla' " + "with framework='torch' on Thor; got config='%s' " + "framework='%s' hardware='%s'. Falling back to FP8.", + config, framework, arch) use_fp4 = False else: try: @@ -722,7 +733,14 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, use_fp4 = False if use_fp4: - if framework == "torch": + if config == "hyvla": + from flash_rt.frontends.torch.hyvla_thor import ( + HyVLATorchFrontendThor, + ) + pipe_cls = HyVLATorchFrontendThor + logger.info("HyVLA Thor FP4 tier enabled") + use_fp4 = False # routed; skip Pi0.5 path below + elif framework == "torch": from flash_rt.frontends.torch.pi05_thor_fp4 import ( Pi05TorchFrontendThorFP4, ) diff --git a/flash_rt/configs/hyvla.yaml b/flash_rt/configs/hyvla.yaml new file mode 100644 index 000000000..1577fdf63 --- /dev/null +++ b/flash_rt/configs/hyvla.yaml @@ -0,0 +1,65 @@ +# FlashRT — Hy-Embodied-0.5-VLA configuration (metadata; runtime dims are code constants) +name: hyvla +arch: hy +vision_type: hyvit2_siglip # SigLIP-so400m isomorphic + 6-frame spacetime +backbone: hunyuan_vl_mot # Mixture-of-Transformers dual-tower +vlm_hidden: 2048 +expert_hidden: 1024 +vlm_layers: 32 +expert_layers: 32 +decoder_steps: 10 # flow-matching Euler steps +action_dim: 32 +state_dim: 32 +chunk_size: 40 + +# Attention (GQA) +num_heads: 16 +num_kv_heads: 4 +head_dim: 128 +rope_type: ntk_alpha # base = 10000 * 1000^(128/126) = 1.1158847e7 +qk_norm: rms_after_rope # RMSNorm over head_dim, applied AFTER RoPE +rms_norm_eps: 1.0e-5 + +# MoT routing: every proj + norm has a text and a vision (_v) copy, routed per-token. +# Expert tower runs 100% through the _v branch; QK-norm is a single shared copy. +mot_routing: static_vision_text # sorted [vision|text] contiguous slices + +# FFN +activation: silu # SwiGLU +intermediate_vlm: 6144 +intermediate_expert: 2048 + +# Vision encoder (HYViT2-400M) +vision: + patch_size: 16 + image_size: 224 + num_patches: 196 # 14x14, merged 2x2 -> 49 tokens/camera + merged_tokens_per_cam: 49 + hidden_dim: 1152 + num_layers: 27 + num_heads: 16 + head_dim: 72 + layernorm_eps: 1.0e-6 + spacetime_layer_stride: 4 # blocks {3,7,11,15,19,23} attend over 6 history frames + history_frames: 6 + num_cameras: 3 + +# Production config (2026-08-04): use_fp8=True, use_fused=True +# -> E2E 172.9 ms (the reference eager path ~930 ms, 5.38x), action cos 0.999706. +# One CUDA graph over prefill + 10-step denoise; dynamic per-tensor FP8 +# (calibration-free, graph-safe) on prefill+denoise GEMMs; bf16 ViT + +# action head; custom fused megakernel (rope+qknorm+kvwrite); and the +# memory-efficient SDPA backend. +# Gated experimental (default off, both measured): +# use_fp8_vit — FP8 ViT is a NET LOSS (+14 ms: quant/bias passes outweigh +# the large-M GEMM saving). +# use_fp4 — NVFP4 W4A4 prefill FFN is correct (cos 0.993) but only +# ~neutral in E2E at this model's M sizes. +quant: + enabled: true + dtype: fp8_dynamic_per_tensor + prefill: fp8 + denoise: fp8 + vision: bf16 + action_head: bf16 + calibration: none # dynamic amax computed on GPU each forward diff --git a/flash_rt/executors/torch_weights.py b/flash_rt/executors/torch_weights.py index c15eca032..4c0604edc 100644 --- a/flash_rt/executors/torch_weights.py +++ b/flash_rt/executors/torch_weights.py @@ -43,6 +43,7 @@ _FP16 = torch.float16 _FP32 = torch.float32 +_BF16 = torch.bfloat16 # ════════════════════════════════════════════════════════════════════ @@ -337,6 +338,11 @@ def apply(self, x, ctx): return x.to(_FP32) +class ToBf16: + def apply(self, x, ctx): + return x.to(_BF16) + + class T: """``.T.contiguous()`` — CUTLASS FP8 col-major path (encoder GEMMs).""" @@ -481,6 +487,7 @@ def finalize(self) -> None: "FusedGateUp", "ToFp16", "ToFp32", + "ToBf16", "T", "tT", "InterleaveQK", diff --git a/flash_rt/frontends/torch/_hyvla_thor_spec.py b/flash_rt/frontends/torch/_hyvla_thor_spec.py new file mode 100644 index 000000000..5228c2a05 --- /dev/null +++ b/flash_rt/frontends/torch/_hyvla_thor_spec.py @@ -0,0 +1,196 @@ +"""Declarative weight spec for HyVLATorchFrontendThor (BF16 baseline). + +Hy-Embodied-0.5-VLA is a Mixture-of-Transformers dual-tower VLA: + + * ViT ``dual_tower.vlm.model.visual.vision_tower`` — 27 SigLIP-so400m + blocks (fused ``attn.qkv`` + bias, LayerNorm **with bias**, patch-16), + 6 of which (indices 3,7,11,15,19,23) run extra 6-frame spacetime + attention at runtime (adopt-by-reference — no extra params). + * merger ``dual_tower.vlm.model.visual.merger`` — proj1 → 2x2 + NormalizedDwPooler → GELU → proj2 (196→49 tokens/cam @ 2048). + * VLM tower ``dual_tower.vlm.model.language_model.model`` — 32 layers, + hidden 2048, GQA 16Q/4KV hd128, SwiGLU inter 6144, RMSNorm eps 1e-5, + **QK-Norm (RMSNorm over head_dim) applied AFTER RoPE**. Every proj + + input/post_attention_layernorm has a text and a vision (``_v``) copy + (MoT routing); QK-norm is a single shared copy. + * expert tower ``dual_tower.expert.model`` — 32 layers, hidden 1024, + same head geometry, SwiGLU inter 2048. At inference runs 100% through + the ``_v`` branch and reuses the VLM tower's QK-norm weights. + * action head (structurally = Pi0): action_in_proj / action_out_proj / + action_time_mlp_in/out / state_proj. + * tied embedding: no ``embed_tokens`` in the checkpoint — the input + embedding table is ``lm_head.weight`` (``tie_word_embeddings``). + +The BF16 baseline keeps the whole model in **BF16** (no FP8 / no ``Quant``). +Weights are stored in the exact layout the reference implementation +consumes: fused QKV / Gate-Up as ``[N, D]`` (out, in) row-major so the +forward computes ``x @ w.t()``; single projections (o_proj, down_proj) +as ``[N, D]``. + +All checkpoint keys are uniformly ``model.``-prefixed → strip it once at +the source. Spec-side keys therefore start ``dual_tower.*`` / ``action_*`` +/ ``state_proj.*``. +""" + +from __future__ import annotations + +import torch + +from flash_rt.executors.weight_loader import Item, LayerBlock, ModelWeightSpec +from flash_rt.executors.torch_weights import ( + Attr, + Cat, + TensorList, + ToBf16, +) + +_BF16 = torch.bfloat16 + + +# ════════════════════════════════════════════════════════════════════ +# ViT block (27 layers) — SigLIP-so400m isomorphic, LayerNorm WITH bias +# ════════════════════════════════════════════════════════════════════ + +def _vit_block() -> LayerBlock: + bp = "dual_tower.vlm.model.visual.vision_tower.blocks.{i}" + items = [ + Item("vit_ln1_w", f"{bp}.norm1.weight", [ToBf16()], TensorList("_vit_ln1_w")), + Item("vit_ln1_b", f"{bp}.norm1.bias", [ToBf16()], TensorList("_vit_ln1_b")), + Item("vit_ln2_w", f"{bp}.norm2.weight", [ToBf16()], TensorList("_vit_ln2_w")), + Item("vit_ln2_b", f"{bp}.norm2.bias", [ToBf16()], TensorList("_vit_ln2_b")), + # attn.qkv is already fused in the checkpoint: [3456, 1152] (+ bias). + Item("vit_qkv_w", f"{bp}.attn.qkv.weight", [ToBf16()], TensorList("_vit_qkv_w")), + Item("vit_qkv_b", f"{bp}.attn.qkv.bias", [ToBf16()], TensorList("_vit_qkv_b")), + Item("vit_proj_w", f"{bp}.attn.proj.weight", [ToBf16()], TensorList("_vit_proj_w")), + Item("vit_proj_b", f"{bp}.attn.proj.bias", [ToBf16()], TensorList("_vit_proj_b")), + Item("vit_fc1_w", f"{bp}.mlp.fc1.weight", [ToBf16()], TensorList("_vit_fc1_w")), + Item("vit_fc1_b", f"{bp}.mlp.fc1.bias", [ToBf16()], TensorList("_vit_fc1_b")), + Item("vit_fc2_w", f"{bp}.mlp.fc2.weight", [ToBf16()], TensorList("_vit_fc2_w")), + Item("vit_fc2_b", f"{bp}.mlp.fc2.bias", [ToBf16()], TensorList("_vit_fc2_b")), + ] + return LayerBlock(prefix_fmt="", num_layers=27, items=items, name="vit") + + +# ════════════════════════════════════════════════════════════════════ +# VLM language tower (32 layers) — text + vision (_v) branches +# ════════════════════════════════════════════════════════════════════ + +def _vlm_block() -> LayerBlock: + dp = "dual_tower.vlm.model.language_model.model.layers.{i}" + sa = f"{dp}.self_attn" + + def branch(suffix: str, tag: str) -> list[Item]: + # suffix "" = text branch, "_v" = vision branch. + mlp = f"{dp}.mlp{'_v' if suffix else ''}" + return [ + Item(f"vlm_qkv{tag}", + Cat([f"{sa}.q_proj{suffix}.weight", + f"{sa}.k_proj{suffix}.weight", + f"{sa}.v_proj{suffix}.weight"], dim=0, dtype=_BF16), + [], TensorList(f"_vlm_qkv{tag}")), + Item(f"vlm_o{tag}", f"{sa}.o_proj{suffix}.weight", + [ToBf16()], TensorList(f"_vlm_o{tag}")), + Item(f"vlm_gu{tag}", + Cat([f"{mlp}.gate_proj.weight", f"{mlp}.up_proj.weight"], + dim=0, dtype=_BF16), + [], TensorList(f"_vlm_gu{tag}")), + Item(f"vlm_d{tag}", f"{mlp}.down_proj.weight", + [ToBf16()], TensorList(f"_vlm_d{tag}")), + Item(f"vlm_ln_in{tag}", f"{dp}.input_layernorm{suffix}.weight", + [ToBf16()], TensorList(f"_vlm_ln_in{tag}")), + Item(f"vlm_ln_post{tag}", f"{dp}.post_attention_layernorm{suffix}.weight", + [ToBf16()], TensorList(f"_vlm_ln_post{tag}")), + ] + + items = branch("", "_t") + branch("_v", "_v") + [ + # Shared QK-norm (single copy, no _v twin). + Item("qk_norm_q", f"{sa}.query_layernorm.weight", + [ToBf16()], TensorList("_qk_norm_q")), + Item("qk_norm_k", f"{sa}.key_layernorm.weight", + [ToBf16()], TensorList("_qk_norm_k")), + ] + return LayerBlock(prefix_fmt="", num_layers=32, items=items, name="vlm") + + +# ════════════════════════════════════════════════════════════════════ +# Expert tower (32 layers) — _v branch only +# ════════════════════════════════════════════════════════════════════ + +def _expert_block() -> LayerBlock: + dp = "dual_tower.expert.model.layers.{i}" + sa = f"{dp}.self_attn" + mlp = f"{dp}.mlp_v" + items = [ + Item("exp_qkv", + Cat([f"{sa}.q_proj_v.weight", + f"{sa}.k_proj_v.weight", + f"{sa}.v_proj_v.weight"], dim=0, dtype=_BF16), + [], TensorList("_exp_qkv_v")), + Item("exp_o", f"{sa}.o_proj_v.weight", [ToBf16()], TensorList("_exp_o_v")), + Item("exp_gu", + Cat([f"{mlp}.gate_proj.weight", f"{mlp}.up_proj.weight"], + dim=0, dtype=_BF16), + [], TensorList("_exp_gu_v")), + Item("exp_d", f"{mlp}.down_proj.weight", [ToBf16()], TensorList("_exp_d_v")), + Item("exp_ln_in", f"{dp}.input_layernorm_v.weight", + [ToBf16()], TensorList("_exp_ln_in_v")), + Item("exp_ln_post", f"{dp}.post_attention_layernorm_v.weight", + [ToBf16()], TensorList("_exp_ln_post_v")), + ] + return LayerBlock(prefix_fmt="", num_layers=32, items=items, name="expert") + + +# ════════════════════════════════════════════════════════════════════ +# Singletons — merger, action head, final norms, tied embed, ViT patch +# ════════════════════════════════════════════════════════════════════ + +def _singletons() -> list[Item]: + mg = "dual_tower.vlm.model.visual.merger" + vt = "dual_tower.vlm.model.visual.vision_tower" + return [ + # --- ViT patch embed + learned pos embed (rescaled at runtime) --- + Item("vit_patch_w", f"{vt}.patch_embed.proj.weight", [ToBf16()], Attr("_vit_patch_w")), + Item("vit_patch_b", f"{vt}.patch_embed.proj.bias", [ToBf16()], Attr("_vit_patch_b")), + Item("vit_pos", f"{vt}.pos_embed", [ToBf16()], Attr("_vit_pos_embed")), + # --- merger --- + Item("mg_proj1_w", f"{mg}.proj1.weight", [ToBf16()], Attr("_mg_proj1_w")), + Item("mg_proj1_b", f"{mg}.proj1.bias", [ToBf16()], Attr("_mg_proj1_b")), + Item("mg_proj2_w", f"{mg}.proj2.weight", [ToBf16()], Attr("_mg_proj2_w")), + Item("mg_proj2_b", f"{mg}.proj2.bias", [ToBf16()], Attr("_mg_proj2_b")), + Item("mg_pred0_w", f"{mg}.pooler.predictor.0.weight", [ToBf16()], Attr("_mg_pred0_w")), + Item("mg_pred0_b", f"{mg}.pooler.predictor.0.bias", [ToBf16()], Attr("_mg_pred0_b")), + Item("mg_pred2_w", f"{mg}.pooler.predictor.2.weight", [ToBf16()], Attr("_mg_pred2_w")), + Item("mg_pred2_b", f"{mg}.pooler.predictor.2.bias", [ToBf16()], Attr("_mg_pred2_b")), + # --- action head --- + Item("ain_w", "action_in_proj.weight", [ToBf16()], Attr("_ain_w")), + Item("ain_b", "action_in_proj.bias", [ToBf16()], Attr("_ain_b")), + Item("aout_w", "action_out_proj.weight", [ToBf16()], Attr("_aout_w")), + Item("aout_b", "action_out_proj.bias", [ToBf16()], Attr("_aout_b")), + Item("atmlp_in_w", "action_time_mlp_in.weight", [ToBf16()], Attr("_atmlp_in_w")), + Item("atmlp_in_b", "action_time_mlp_in.bias", [ToBf16()], Attr("_atmlp_in_b")), + Item("atmlp_out_w", "action_time_mlp_out.weight", [ToBf16()], Attr("_atmlp_out_w")), + Item("atmlp_out_b", "action_time_mlp_out.bias", [ToBf16()], Attr("_atmlp_out_b")), + Item("state_w", "state_proj.weight", [ToBf16()], Attr("_state_w")), + Item("state_b", "state_proj.bias", [ToBf16()], Attr("_state_b")), + # --- expert final RMSNorm --- + Item("exp_final_norm", "dual_tower.expert.model.norm.weight", + [ToBf16()], Attr("_exp_final_norm_w")), + # --- tied input embedding = lm_head.weight --- + Item("embed", "dual_tower.vlm.model.language_model.lm_head.weight", + [ToBf16()], Attr("_embed_weight")), + ] + + +def build_spec() -> ModelWeightSpec: + return ModelWeightSpec( + framework="torch", + singletons=_singletons(), + blocks=[ + _vit_block(), + _vlm_block(), + _expert_block(), + ], + ) + + +__all__ = ["build_spec"] diff --git a/flash_rt/frontends/torch/hyvla_thor.py b/flash_rt/frontends/torch/hyvla_thor.py new file mode 100644 index 000000000..a3f4a8fde --- /dev/null +++ b/flash_rt/frontends/torch/hyvla_thor.py @@ -0,0 +1,628 @@ +"""HyVLATorchFrontendThor — native FlashRT frontend for Hy-Embodied-0.5-VLA +on Thor SM110 (BF16 baseline). + +Reachable via ``flash_rt.load_model(ckpt, config="hyvla", framework="torch")``. +Owns the full inference IO path with **no reference training-code import at +runtime**: tokenizer (checkpoint ``AutoTokenizer`` + hard-coded hy special +token ids), image preprocessing (resize-with-pad + ``*2-1``), prefix assembly +(BOS / hy_User / per-camera vision blocks / language / hy_Assistant), the +segmented prefix mask + prefix-LM suffix mask, the NTK-alpha RoPE tables, and +the flow-matching time embeddings. The heavy math lives in +``flash_rt.models.hyvla.pipeline_thor`` (ViT+merger, MoT prefill, expert +denoise), validated at cosine ≥ 0.9998 vs the HF reference eager path. + +The BF16 baseline uses plain torch matmuls and SDPA attention; the +optimized path swaps GEMMs/norms for ``fvk`` pointer kernels + a single +CUDA graph + FP8. +""" + +from __future__ import annotations + +import json +import math +import pathlib +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F + +from flash_rt.executors.torch_weights import SafetensorsSource, WeightLoader +from flash_rt.frontends.torch._hyvla_thor_spec import build_spec +from flash_rt.models.hyvla.pipeline_thor import HyVLAThorBF16Pipeline + +_BF16 = torch.bfloat16 + +# hy special token ids (verified from the checkpoint tokenizer_config). +_TOK_BOS = 120000 +_TOK_HY_USER = 120006 +_TOK_VISION_START = 120684 # <｜hy_place▁holder▁no▁666｜> +_TOK_VISION_END = 120685 # <｜hy_place▁holder▁no▁667｜> +_TOK_VISION_SPLIT = 120689 # <｜hy_place▁holder▁no▁671｜> +_TOK_HY_ASSISTANT = "<｜hy_Assistant｜>" + + +def _resize_with_pad(img, height=224, width=224, pad_value=-1.0, mode="bilinear"): + """Pi0-style resize with aspect-preserving center pad. (B,C,H,W).""" + ch, cw = img.shape[2:] + if (ch, cw) == (height, width): + return img + ratio = max(cw / width, ch / height) + rh, rw = int(ch / ratio), int(cw / ratio) + resized = F.interpolate(img, size=(rh, rw), mode=mode, align_corners=False) + ph, pw = max(0, height - rh), max(0, width - rw) + t, l = ph // 2, pw // 2 + return F.pad(resized, (l, pw - l, t, ph - t), value=pad_value) + + +def _camera_tensor(image): + t = torch.as_tensor(np.asarray(image)) + scale_uint8 = t.dtype == torch.uint8 + if t.ndim == 3: + if t.shape[-1] == 3: + t = t.permute(2, 0, 1) + elif t.shape[0] != 3: + raise ValueError(f"camera image must have 3 channels, got shape {tuple(t.shape)}") + t = t.unsqueeze(0) + elif t.ndim == 4: + if t.shape[-1] == 3: + t = t.permute(0, 3, 1, 2) + elif t.shape[1] != 3: + raise ValueError(f"camera frames must have 3 channels, got shape {tuple(t.shape)}") + else: + raise ValueError(f"camera image must be rank 3 or 4, got shape {tuple(t.shape)}") + t = t.contiguous().float() + if scale_uint8: + t = t / 255.0 + return t + + +class HyVLATorchFrontendThor: + def __init__(self, checkpoint_dir: str, *, hardware: str = "thor", + use_fp8: bool = False, use_fp8_vit: bool = False, + use_fused: bool = False, use_fp4: bool = False, + use_fused_quant: bool = False, use_autotune: bool = False, + use_ffn_mega: bool = False, **kwargs): + self.checkpoint_dir = str(checkpoint_dir) + self.device = "cuda" + self.use_fp8 = bool(use_fp8) + self.use_fp8_vit = bool(use_fp8_vit) + self.use_fused = bool(use_fused) + self.use_fp4 = bool(use_fp4) + self.use_fused_quant = bool(use_fused_quant) + self.use_autotune = bool(use_autotune) + self.use_ffn_mega = bool(use_ffn_mega) + cfg_path = pathlib.Path(self.checkpoint_dir) / "config.json" + with open(cfg_path) as f: + cfg = json.load(f) + self.cfg = cfg + self.num_steps = int(cfg.get("num_steps", 10)) + self.chunk = int(cfg.get("n_action_steps", cfg.get("chunk_size", 40))) + self.max_action_dim = int(cfg.get("max_action_dim", 32)) + self.max_state_dim = int(cfg.get("max_state_dim", 32)) + self.tok_max_len = int(cfg.get("tokenizer_max_length", 64)) + self.proj_width = int(cfg.get("proj_width", 1024)) + # Camera key order (must match the prefix assembly order). + self.image_keys = list(cfg.get("image_features", {}).keys()) + + txt = cfg.get("vlm_config_dict", {}).get("text_config", {}) + self.rope_theta = float(txt.get("rope_theta", 10000.0)) + self.head_dim = int(txt.get("head_dim", 128)) + alpha = float(txt.get("rope_scaling", {}).get("alpha", 1000.0)) + self._rope_base = self.rope_theta * alpha ** (self.head_dim / (self.head_dim - 2)) + self.n_kv = int(txt.get("num_key_value_heads", 4)) + self.n_heads = int(txt.get("num_attention_heads", 16)) + self.d_vlm = int(txt.get("hidden_size", 2048)) + + self._load_weights() + self.pipe = HyVLAThorBF16Pipeline(self) + self._graph_cache = {} # (S_p, n_vis) -> {"graph", "buf"} + self._vit_graph_cache = {} # (num_cam, K) -> {"graph", "img", "out"} + self._tokenizer = None + self._prompt = None + self._lang_tokens = None + self._lang_masks = None + self._precompute_time_embs() + if self.use_fp8: + self._quantize_fp8() + self.pipe.enable_fp8() + if self.use_fused_quant: + # denoise expert tower runs at M=41; collapse the 4-node + # quantize_fp8_device into the single-CTA fused quant there. + self.pipe._small_quant_m = 64 + if self.use_ffn_mega: + self._quantize_ffn_mega() + self.pipe._ffn_mega = True + if self.use_fused: + self.pipe._fused_attn = True + if self.use_fp4: + self._quantize_fp4() + self.pipe.enable_fp4() + + # ------------------------------------------------------------------ + def _load_weights(self): + sf = pathlib.Path(self.checkpoint_dir) / "model.safetensors" + src = SafetensorsSource(str(sf), device=self.device, strip_prefix="model.") + WeightLoader(source=src, target=self, spec=build_spec()).run() + + def _quantize_fp8(self): + """Quantize expert-tower AND VLM-tower (text+vision) GEMM weights to + graph-safe FP8: fp8 (K,N) tensor + precomputed device fp32 per-tensor + scale, consumed by pipeline._fp8_gemm (dynamic-activation FP8).""" + import flash_rt.flash_rt_kernels as fvk + st = torch.cuda.current_stream().cuda_stream + + def q(w_bf16): + wkn = w_bf16.t().contiguous() # (N,K) -> (K,N) + K, N = wkn.shape + w8 = torch.empty(K, N, dtype=torch.uint8, device=self.device) + ws = torch.empty(1, dtype=torch.float32, device=self.device) + fvk.quantize_fp8_device(wkn.data_ptr(), w8.data_ptr(), ws.data_ptr(), K * N, st) + return w8, ws + + def q_list(src): + w8s, wss = [], [] + for w in src: + w8, ws = q(w) + w8s.append(w8); wss.append(ws) + return w8s, wss + + # Expert tower (uniform _v) + self._exp_qkv8, self._exp_qkv_ws = q_list(self._exp_qkv_v) + self._exp_o8, self._exp_o_ws = q_list(self._exp_o_v) + self._exp_gu8, self._exp_gu_ws = q_list(self._exp_gu_v) + self._exp_d8, self._exp_d_ws = q_list(self._exp_d_v) + self._exp_fp8_ready = True + + # VLM tower (vision + text branches) + self._vlm_qkv_v8, self._vlm_qkv_v_ws = q_list(self._vlm_qkv_v) + self._vlm_o_v8, self._vlm_o_v_ws = q_list(self._vlm_o_v) + self._vlm_gu_v8, self._vlm_gu_v_ws = q_list(self._vlm_gu_v) + self._vlm_d_v8, self._vlm_d_v_ws = q_list(self._vlm_d_v) + self._vlm_qkv_t8, self._vlm_qkv_t_ws = q_list(self._vlm_qkv_t) + self._vlm_o_t8, self._vlm_o_t_ws = q_list(self._vlm_o_t) + self._vlm_gu_t8, self._vlm_gu_t_ws = q_list(self._vlm_gu_t) + self._vlm_d_t8, self._vlm_d_t_ws = q_list(self._vlm_d_t) + self._vlm_fp8_ready = True + + # ViT tower (27 blocks) — measured NET LOSS on Thor (quant+bias passes + # outweigh large-M GEMM savings; confirms the FP8-ViT dead-end). + # Opt-in only. + if self.use_fp8_vit: + self._vit_qkv_w8, self._vit_qkv_ws = q_list(self._vit_qkv_w) + self._vit_proj_w8, self._vit_proj_ws = q_list(self._vit_proj_w) + self._vit_fc1_w8, self._vit_fc1_ws = q_list(self._vit_fc1_w) + self._vit_fc2_w8, self._vit_fc2_ws = q_list(self._vit_fc2_w) + self._vit_fp8_ready = True + torch.cuda.synchronize() + + def _quantize_ffn_mega(self): + """Quantize the expert-tower FFN weights (gu, dn) to FP8 in (N,K) layout + for the denoise FFN megakernel (hyvla_ffn_gu_silu_bf16 / _dn_res_bf16). + + The megakernel reads weight rows K-contiguous (N,K) — the ORIGINAL + orientation — unlike _fp8_gemm's (K,N). Weight scale is read to a host + float (constant across forwards); the activation scale stays dynamic + (device pointer). Done once at load (the .item() sync is fine here).""" + import flash_rt.flash_rt_kernels as fvk + st = torch.cuda.current_stream().cuda_stream + + def q_nk(w_bf16): # (N,K) bf16 -> (N,K) fp8 uint8 + host float scale + N, K = w_bf16.shape + wc = w_bf16.contiguous() + w8 = torch.empty(N, K, dtype=torch.uint8, device=self.device) + ws = torch.empty(1, dtype=torch.float32, device=self.device) + fvk.quantize_fp8_device(wc.data_ptr(), w8.data_ptr(), ws.data_ptr(), N * K, st) + return w8, float(ws.item()) + + gu8, gus, dn8, dns = [], [], [], [] + for w in self._exp_gu_v: + a, b = q_nk(w); gu8.append(a); gus.append(b) + for w in self._exp_d_v: + a, b = q_nk(w); dn8.append(a); dns.append(b) + self._exp_gu_mk, self._exp_gu_mk_s = gu8, gus + self._exp_d_mk, self._exp_d_mk_s = dn8, dns + self._exp_inter = self._exp_gu_v[0].shape[0] // 2 # 2*inter -> inter + self._exp_ffn_mega_ready = True + torch.cuda.synchronize() + + def _quantize_fp4(self): + """Quantize VLM prefill FFN weights (gu, down; text + vision branches) + to NVFP4 via the flash_rt_fp4 family (packed 4-bit + swizzled UE4M3 SF), + consumed by cutlass_fp4_sq_fp16. Prefill runs at M=240 where FP4 wins + (gu 2.67x) and once (no Euler compounding), so it is the right FP4 target. + Activation quant uses the SAME F4 family (see pipeline._fp4_gemm_f4) — + mixing F4 weights with fvk fused-quant SF is the swizzle-mismatch trap.""" + import flash_rt.flash_rt_fp4 as F4 + st = torch.cuda.current_stream().cuda_stream + + def q_nvfp4(w_bf16): + N, K = w_bf16.shape + assert K % 64 == 0, f"cutlass_fp4_sq_fp16 needs K%64==0, got K={K}" + w16 = w_bf16.to(torch.float16).contiguous() + packed = torch.empty(N, K // 2, dtype=torch.uint8, device=self.device) + sf = torch.empty(F4.sfa_size_bytes(N, K, True), dtype=torch.uint8, device=self.device) + F4.quantize_fp4_dynamic_sfa_fp16(w16.data_ptr(), packed.data_ptr(), + sf.data_ptr(), N, K, True, st) + return packed, sf + + def q_list(src): + ps, ss = [], [] + for w in src: + p, s = q_nvfp4(w); ps.append(p); ss.append(s) + return ps, ss + + self._vlm_gu_v4, self._vlm_gu_v4sf = q_list(self._vlm_gu_v) + self._vlm_d_v4, self._vlm_d_v4sf = q_list(self._vlm_d_v) + self._vlm_gu_t4, self._vlm_gu_t4sf = q_list(self._vlm_gu_t) + self._vlm_d_t4, self._vlm_d_t4sf = q_list(self._vlm_d_t) + self._vlm_gu_N = self._vlm_gu_v[0].shape[0] # 2*inter (gate+up merged) + self._vlm_D = self._vlm_gu_v[0].shape[1] # VLM hidden (gu K) + self._vlm_inter = self._vlm_d_v[0].shape[1] # inter (down K) + self._vlm_fp4_ready = True + torch.cuda.synchronize() + + def _precompute_time_embs(self): + embs = [] + for s in range(self.num_steps): + t = 1.0 + s * (-1.0 / self.num_steps) + embs.append(self._sinusoidal_time(t)) + self._time_embs = torch.stack(embs).to(self.device, _BF16) # (steps,1,proj_width) + + def _sinusoidal_time(self, t, min_period=4e-3, max_period=4.0): + dim = self.proj_width + frac = torch.linspace(0.0, 1.0, dim // 2, dtype=torch.float64) + period = min_period * (max_period / min_period) ** frac + scaling = 1.0 / period * 2 * math.pi + tt = torch.tensor([t], dtype=torch.float64) + sin_in = scaling[None, :] * tt[:, None] + return torch.cat([torch.sin(sin_in), torch.cos(sin_in)], dim=1) # (1,dim) + + # ------------------------------------------------------------------ + def _rope_cos_sin(self, positions): + """positions (1,S) long -> cos,sin (1,1,S,head_dim) bf16 (rotate_half). + + The original model is cast to bf16, so its ``rotary_emb.inv_freq`` + buffer is bf16-quantized; that rounding compounds over positions, so + we must round our inv_freq through bf16 to match the reference tables + exactly. + """ + hd = self.head_dim + inv_freq = 1.0 / (self._rope_base ** ( + torch.arange(0, hd, 2, dtype=torch.float64, device=self.device) / hd)) + inv_freq = inv_freq.to(_BF16).to(torch.float64) # match reference bf16 storage + freqs = positions.to(torch.float64)[..., None] * inv_freq[None, None, :] + emb = torch.cat([freqs, freqs], dim=-1) # (1,S,hd) + return emb.cos()[:, None].to(_BF16), emb.sin()[:, None].to(_BF16) + + @property + def tokenizer(self): + if self._tokenizer is None: + from transformers import AutoTokenizer + self._tokenizer = AutoTokenizer.from_pretrained( + self.checkpoint_dir, trust_remote_code=True) + return self._tokenizer + + def _tokenize(self, prompt): + task = prompt.strip().replace("_", " ").replace("\n", " ") + if not task.endswith(_TOK_HY_ASSISTANT): + task = f"{task}{_TOK_HY_ASSISTANT}" + out = self.tokenizer([task], padding="max_length", padding_side="right", + truncation=True, max_length=self.tok_max_len, + return_tensors="pt", add_special_tokens=False) + return (out["input_ids"].to(self.device), + out["attention_mask"].to(self.device).bool()) + + def set_prompt(self, prompt, state=None): + self._prompt = prompt + self._lang_tokens, self._lang_masks = self._tokenize(prompt) + + # ------------------------------------------------------------------ + def _embed_ids(self, ids): + return F.embedding(ids, self._embed_weight) + + def _preprocess_images(self, images): + """Normalize public inputs to (num_cam,K,3,H,W) in [0,1].""" + if isinstance(images, dict): + keys = [k for k in self.image_keys if k in images] + if not keys: + keys = [k for k in ("image", "wrist_image", "wrist_image_right") if k in images] + if not keys: + raise ValueError("images dict does not contain configured camera keys") + images = [images[k] for k in keys] + + if isinstance(images, (list, tuple)): + if not images: + raise ValueError("images list must have at least one camera") + images = torch.stack([_camera_tensor(im) for im in images], 0) + else: + images = torch.as_tensor(np.asarray(images)) + scale_uint8 = images.dtype == torch.uint8 + if images.ndim == 5: + if images.shape[-1] == 3: + images = images.permute(0, 1, 4, 2, 3) + elif images.shape[2] != 3: + raise ValueError( + f"images must have channel dimension of size 3, got {tuple(images.shape)}") + elif images.ndim == 4: + images = torch.stack([_camera_tensor(im) for im in images], 0) + else: + raise ValueError( + f"images must be list/dict or rank 4/5 tensor, got {tuple(images.shape)}") + images = images.contiguous().float() + if scale_uint8: + images = images / 255.0 + + out = [] + for cam in range(images.shape[0]): + im = images[cam].to(self.device, _BF16) + im = _resize_with_pad(im, 224, 224, pad_value=0.0) + im = im * 2.0 - 1.0 + out.append(im[None]) + return out + + @torch.no_grad() + def _assemble_prefix(self, merged): + """merged: (num_cam, 49, 2048) merged vision tokens. Returns prefix tensors.""" + dev = self.device + + embs = [self._embed_ids(torch.tensor([[_TOK_BOS]], device=dev)), + self._embed_ids(torch.tensor([[_TOK_HY_USER]], device=dev))] + att = [1, 1] + mm = [False, False] + pad = [torch.ones((1, 2), dtype=torch.bool, device=dev)] + idx_ranges, full_ranges = [], [] + + vstart = self._embed_ids(torch.tensor([[_TOK_VISION_START]], device=dev)) + vend = self._embed_ids(torch.tensor([[_TOK_VISION_END]], device=dev)) + vsplit = self._embed_ids(torch.tensor([[_TOK_VISION_SPLIT]], device=dev)) + + for ci in range(merged.shape[0]): + img_emb = merged[ci][None] # (1,49,2048) + g = int(img_emb.shape[1] ** 0.5) # 7 + embs.append(vstart); att.append(1); mm.append(False) + pad.append(torch.ones((1, 1), dtype=torch.bool, device=dev)) + grid = img_emb.view(1, g, g, -1) + split_exp = vsplit.unsqueeze(1).expand(1, g, 1, grid.shape[-1]) + with_split = torch.cat([grid, split_exp], dim=2).reshape(1, -1, grid.shape[-1]) + embs.append(with_split) + row_len = g + 1 + total = g * row_len + start = len(att) + idx_ranges.extend([(start + r * row_len, start + r * row_len + g) for r in range(g)]) + full_ranges.append((start, start + total)) + att.extend([1] * total) + mm.extend(([True] * g + [False]) * g) + pad.append(torch.ones((1, total), dtype=torch.bool, device=dev)) + embs.append(vend); att.append(1); mm.append(False) + pad.append(torch.ones((1, 1), dtype=torch.bool, device=dev)) + + lang_emb = self._embed_ids(self._lang_tokens) # (1,64,2048) + embs.append(lang_emb) + pad.append(self._lang_masks) + n_lang = lang_emb.shape[1] + att.extend([1] * n_lang) + mm.extend([False] * n_lang) + + prefix_embs = torch.cat(embs, dim=1) + pad_masks = torch.cat(pad, dim=1).bool() + att_masks = torch.tensor(att, dtype=torch.bool, device=dev)[None] + mm_prefix = torch.tensor(mm, dtype=torch.bool, device=dev)[None] + return prefix_embs, pad_masks, att_masks, mm_prefix, idx_ranges, full_ranges + + @staticmethod + def _make_att_2d(pad_masks, att_masks): + cumsum = torch.cumsum(att_masks.long(), dim=1) + att2d = cumsum[:, None, :] <= cumsum[:, :, None] + pad2d = pad_masks[:, None, :] * pad_masks[:, :, None] + return att2d & pad2d + + def _apply_segment_mask(self, att2d, idx_ranges, full_ranges): + dev = att2d.device + all_idx = [i for (s, e) in idx_ranges for i in range(s, e)] + if all_idx: + idx = torch.tensor(all_idx, device=dev) + att2d[:, idx[:, None], idx[None, :]] = False + for fs, fe in full_ranges: + img_idx = [i for (s, e) in idx_ranges if s >= fs and e <= fe for i in range(s, e)] + if img_idx: + idx = torch.tensor(img_idx, device=dev) + att2d[:, idx[:, None], idx[None, :]] = True + return att2d + + # ------------------------------------------------------------------ + # ViT + merger CUDA-graph (BF16, unchanged precision). Keyed on + # (num_cam, K); removes per-block launch overhead of the 27-block ViT. + # ------------------------------------------------------------------ + def _vit_merge(self, imgs5, use_graph=True): + """imgs5 (num_cam,K,3,224,224) bf16 [-1,1] -> merged (num_cam,49,2048).""" + if not use_graph: + return self.pipe.merger_forward(self.pipe.vit_forward(imgs5)) + key = tuple(imgs5.shape[:2]) + g = self._vit_graph_cache.get(key) + if g is None: + g = self._build_vit_graph(imgs5.shape) + self._vit_graph_cache[key] = g + g["img"].copy_(imgs5) + g["graph"].replay() + return g["out"].clone() + + def _build_vit_graph(self, shape): + dev = self.device + img = torch.zeros(shape, dtype=_BF16, device=dev) + out = self.pipe.merger_forward(self.pipe.vit_forward(img)).clone() # shape probe + + def body(): + out.copy_(self.pipe.merger_forward(self.pipe.vit_forward(img))) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + body() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + body() + return {"graph": graph, "img": img, "out": out} + + # ------------------------------------------------------------------ + # CUDA-graph capture of prefill + 10-step denoise (BF16, unchanged + # precision). Keyed on (S_p, n_vis); ViT + prefix assembly stay eager + # and feed the static input buffers before each replay. + # ------------------------------------------------------------------ + def _captured_body(self, b, n_vis, S_p): + self.pipe.prefill(b["pe"], n_vis, b["pmask"], b["pcos"], b["psin"], + b["kbuf"], b["vbuf"]) + self.pipe.denoise(b["state"], b["x"], self._time_embs, + b["smask"], b["scos"], b["ssin"], + b["kbuf"], b["vbuf"], S_p, num_steps=self.num_steps) + + def _build_graph(self, key): + S_p, n_vis = key + dev = self.device + L, nkv, hd = 32, self.n_kv, self.head_dim + # KV cache is stored PRE-EXPANDED to all query heads: the megakernel + # replicates each KV head kv_rep times, so attention reads it directly + # and skips the per-call repeat_interleave over the whole cache + # (measured ~36us/call x 640 calls on the 281-row cache). + n_kvc = self.n_heads + D, S_s = self.d_vlm, 1 + self.chunk + z = lambda *s, dt=_BF16: torch.zeros(*s, dtype=dt, device=dev) + b = { + "pe": z(1, S_p, D), + "pmask": z(1, 1, S_p, S_p, dt=torch.bool), + "pcos": z(1, 1, S_p, hd), "psin": z(1, 1, S_p, hd), + "smask": z(1, 1, S_s, S_p + S_s, dt=torch.bool), + "scos": z(1, 1, S_s, hd), "ssin": z(1, 1, S_s, hd), + "state": z(1, self.max_state_dim), + "x": z(1, self.chunk, self.max_action_dim, dt=torch.float32), + "kbuf": z(L, 1, n_kvc, S_p + S_s, hd), + "vbuf": z(L, 1, n_kvc, S_p + S_s, hd), + } + # Per-shape FP8 GEMM autotune BEFORE capture (graph-safe: mutates only + # the GemmRunner algo cache). A dry eager body records the exact (M,N,K) + # set the captured path will hit, then we tune each on self.pipe.gemm. + if self.use_autotune and self.use_fp8: + self.pipe._gemm_shapes = set() + self._captured_body(b, n_vis, S_p) + shapes = self.pipe._gemm_shapes + self.pipe._gemm_shapes = None + self.pipe.autotune_gemms(shapes) + torch.cuda.synchronize() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + self._captured_body(b, n_vis, S_p) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + self._captured_body(b, n_vis, S_p) + return {"graph": graph, "buf": b} + + def _graph_forward(self, S_p, n_vis, prefix_embs, pmask, pcos, psin, + smask, scos, ssin, state_t, noise_t, use_graph=True): + dev = self.device + if not use_graph: + L, nkv, hd = 32, self.n_kv, self.head_dim + kbuf = torch.zeros(L, 1, self.n_heads, S_p + 1 + self.chunk, hd, dtype=_BF16, device=dev) + vbuf = torch.zeros_like(kbuf) + x = noise_t.clone().float() + self.pipe.prefill(prefix_embs.clone(), n_vis, pmask, pcos, psin, kbuf, vbuf) + return self.pipe.denoise(state_t, x, self._time_embs, smask, scos, ssin, + kbuf, vbuf, S_p, num_steps=self.num_steps) + g = self._graph_cache.get((S_p, n_vis)) + if g is None: + g = self._build_graph((S_p, n_vis)) + self._graph_cache[(S_p, n_vis)] = g + b = g["buf"] + b["pe"].copy_(prefix_embs); b["pmask"].copy_(pmask) + b["pcos"].copy_(pcos); b["psin"].copy_(psin) + b["smask"].copy_(smask); b["scos"].copy_(scos); b["ssin"].copy_(ssin) + b["state"].copy_(state_t); b["x"].copy_(noise_t) + g["graph"].replay() + return b["x"].clone() + + # ------------------------------------------------------------------ + @torch.no_grad() + def predict_actions(self, images, prompt=None, state=None, noise=None, + use_graph=True): + """Full native forward. images: (num_cam,K,3,H,W) [0,1] array/tensor. + Returns raw action chunk (1, chunk, max_action_dim) as numpy fp32.""" + if prompt is not None and prompt != self._prompt: + self.set_prompt(prompt) + assert self._lang_tokens is not None, "call set_prompt() first" + dev = self.device + + if not torch.is_tensor(images): + images = torch.as_tensor(np.asarray(images)) + cam_imgs = self._preprocess_images(images) + imgs5 = torch.cat(cam_imgs, 0) # (num_cam,K,3,224,224) + merged = self._vit_merge(imgs5, use_graph=use_graph) + + # state -> (1, max_state_dim) + if state is None: + state_t = torch.zeros(1, self.max_state_dim, device=dev, dtype=_BF16) + else: + st = torch.as_tensor(np.asarray(state), device=dev, dtype=_BF16).reshape(1, -1) + if st.shape[1] < self.max_state_dim: + st = F.pad(st, (0, self.max_state_dim - st.shape[1])) + state_t = st + + (prefix_embs, pad_masks, att_masks, mm_prefix, + idx_ranges, full_ranges) = self._assemble_prefix(merged) + att2d = self._make_att_2d(pad_masks, att_masks) + att2d = self._apply_segment_mask(att2d, idx_ranges, full_ranges) + prefix_pos = torch.cumsum(pad_masks.long(), dim=1) - 1 + + mm = mm_prefix[0] + perm = torch.cat([torch.nonzero(mm).squeeze(-1), torch.nonzero(~mm).squeeze(-1)]) + n_vis = int(mm.sum().item()) + prefix_embs = prefix_embs[:, perm] + att2d = att2d[:, perm][:, :, perm] + prefix_pos = prefix_pos[:, perm] + + pcos, psin = self._rope_cos_sin(prefix_pos) + pmask = att2d[:, None] + + S_p = pad_masks.shape[1] + S_s = 1 + self.chunk + suffix_pad = torch.ones(1, S_s, dtype=torch.bool, device=dev) + suffix_att = torch.tensor([1, 1] + [0] * (self.chunk - 1), + dtype=torch.bool, device=dev)[None] + suffix_att2d = self._make_att_2d(suffix_pad, suffix_att) + prefix_pad_2d = pad_masks[:, perm][:, None, :].expand(1, S_s, S_p) + smask = torch.cat([prefix_pad_2d, suffix_att2d], 2)[:, None] + suffix_pos = pad_masks.long().sum(-1)[:, None] + torch.cumsum(suffix_pad.long(), 1) - 1 + scos, ssin = self._rope_cos_sin(suffix_pos) + + if noise is None: + noise_t = torch.randn(1, self.chunk, self.max_action_dim, + dtype=torch.float32, device=dev) + else: + noise_t = torch.as_tensor(np.asarray(noise), device=dev, dtype=torch.float32) + + x_t = self._graph_forward(S_p, n_vis, prefix_embs, pmask, pcos, psin, + smask, scos, ssin, state_t, noise_t, + use_graph=use_graph) + return x_t.float().cpu().numpy() + + # ------------------------------------------------------------------ + @torch.no_grad() + def infer(self, obs): + images = obs.get("images") + if images is None: + images = {k: obs[k] for k in self.image_keys if k in obs} + if not images: + images = [obs[k] for k in ("image", "wrist_image", "wrist_image_right") if k in obs] + state = obs.get("state") + noise = obs.get("noise") + actions = self.predict_actions(images, prompt=obs.get("prompt"), + state=state, noise=noise) + return {"actions": actions[0]} + + +__all__ = ["HyVLATorchFrontendThor"] diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 5d231ead6..0006aec7b 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -102,6 +102,10 @@ def detect_arch() -> str: ("pi0", "jax", "rtx_sm89"): ("flash_rt.frontends.jax.pi0_rtx", "Pi0JaxFrontendRtx"), + # ── Hy-Embodied-0.5-VLA (HunYuan MoT dual-tower + flow matching) ── + ("hyvla", "torch", "thor"): + ("flash_rt.frontends.torch.hyvla_thor", "HyVLATorchFrontendThor"), + # ── GROOT N1.6 ── ("groot", "torch", "thor"): ("flash_rt.frontends.torch.groot_thor", "GrootTorchFrontendThor"), diff --git a/flash_rt/models/hyvla/__init__.py b/flash_rt/models/hyvla/__init__.py new file mode 100644 index 000000000..8aa927ab0 --- /dev/null +++ b/flash_rt/models/hyvla/__init__.py @@ -0,0 +1 @@ +"""Hy-Embodied-0.5-VLA — Thor SM110 and Orin compute paths.""" diff --git a/flash_rt/models/hyvla/pipeline_thor.py b/flash_rt/models/hyvla/pipeline_thor.py new file mode 100644 index 000000000..b3b47e2a0 --- /dev/null +++ b/flash_rt/models/hyvla/pipeline_thor.py @@ -0,0 +1,572 @@ +"""Hy-Embodied-0.5-VLA forward path on Thor SM110 (BF16 baseline). + +This module owns the model-specific forwards (repo contract §0 rule 1/3 — +they may NOT live in ``hardware/thor/shared_primitives.py``): + + * ``vit_forward`` — 27-block HYViT2 incl. the 6 spacetime blocks, + per camera, over 6 history frames. + * ``merger_forward`` — proj1 → 2x2 NormalizedDwPooler → GELU → proj2. + * ``prefill_forward`` — 32-layer MoT VLM tower over the sorted + ``[vision|text]`` prefix; fills the per-layer + KV cache the denoise loop reads. + * ``denoise_forward`` — 32-layer expert tower (``_v`` only) + action + head, 10 flow-matching Euler steps. + +The BF16 baseline is **correctness-first**: heavy GEMMs are plain +``torch`` bf16 matmuls and attention is ``F.scaled_dot_product_attention`` +(the numerically-exact reference path, GQA 16/4 + materialized mask). +The math mirrors the validated reference implementation +(action-chunk cosine ≥ 0.999 vs the HF reference eager path). The +optimized path swaps the GEMMs/norms for ``fvk`` pointer kernels, the +two masked-attention bodies for fused kernels, and captures one CUDA +graph — the forward signatures and the ``[vision|text]`` static-routing +layout are chosen so that swap is local (no structural change). No +external training code is imported. + +Key model constants (verified from the checkpoint + reference source): + D_vlm=2048, D_exp=1024, n_heads=16, n_kv=4, head_dim=128, + q_dim=2048, kv_dim=512, inter_vlm=6144, inter_exp=2048, + rms_eps=1e-5 (all towers + QK-norm + final norm), + chunk=40, num_steps=10, dt=-1/steps, t: 1.0→0.1. + QK-Norm (RMSNorm over head_dim) is applied AFTER RoPE (rotate_half). +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +import flash_rt.flash_rt_kernels as fvk + + +def _rot_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +class HyVLAThorBF16Pipeline: + """BF16 forward for the Hy-VLA dual-tower + action head. + + Weights are read from ``W`` (the frontend), which carries the + declarative-spec attributes (``_vlm_qkv_t`` lists, ``_ain_w`` etc.). + Buffers (KV cache, x_t, time embeddings) are pre-allocated by the + frontend and passed in, so the forward performs no Python-level + allocation — the graph-capture safety precondition. + """ + + def __init__(self, W): + self.W = W + self.n_heads = 16 + self.n_kv = 4 + self.head_dim = 128 + self.q_dim = self.n_heads * self.head_dim # 2048 + self.kv_dim = self.n_kv * self.head_dim # 512 + self.rms_eps = 1e-5 + self._fp8 = False + self.gemm = None + self._fused_attn = False + self._fp4 = False + self._F4 = None + # When M <= this, use the single-CTA fused dynamic FP8 quant + # (hyvla_quant_fp8_dyn_bf16, 1 launch) instead of quantize_fp8_device + # (4 nodes). 0 disables. Set by the frontend for the denoise tower. + self._small_quant_m = 0 + # When a set(), _fp8_gemm records the (M,N,K) it sees so the frontend + # can autotune the cuBLASLt FP8 algo per shape before graph capture. + self._gemm_shapes = None + # Fuse the expert denoise FFN (gu+silu_mul, dn+residual) into two + # occupancy-preserving persistent megakernels (hyvla_ffn_*). + self._ffn_mega = False + # ViT (HYViT2-400M) + self.vit_heads = 16 + self.vit_hd = 72 + self.vit_scale = self.vit_hd ** -0.5 + self.vit_eps = 1e-6 + self.vit_time_base = 100.0 + self.vit_spacetime_ids = set(range(3, 27, 4)) # {3,7,11,15,19,23} + + # ══════════════════════════════════════════════════════════════════ + # FP8 (dynamic per-tensor, graph-safe) — expert denoise GEMMs + # ══════════════════════════════════════════════════════════════════ + def enable_fp8(self): + from flash_rt.core.context import FvkContext + self.gemm = FvkContext().gemm + self._fp8 = True + + def enable_fp4(self): + import flash_rt.flash_rt_fp4 as _F4 + self._F4 = _F4 + self._fp4 = True + + def autotune_gemms(self, shapes, num_algos=16): + """Per-shape cuBLASLt FP8 algo autotune (motus pattern). Runs BEFORE + graph capture on ``self.gemm`` — the SAME GemmRunner _fp8_gemm calls, so + the tuned algo is cached (keyed on (M,N,K)) and every captured + ``fp8_nn_dev`` for that shape picks it up automatically. Dummy buffers: + only the algo is timed; real scale pointers are set per call.""" + if self.gemm is None or not hasattr(self.gemm, "autotune_fp8_nn_dev"): + return 0 + dev = "cuda" + n = 0 + for (M, N, K) in sorted(shapes): + A = torch.empty(M, K, dtype=torch.uint8, device=dev) + B = torch.empty(K, N, dtype=torch.uint8, device=dev) + D = torch.empty(M, N, dtype=torch.bfloat16, device=dev) + sa = torch.ones(1, dtype=torch.float32, device=dev) + sb = torch.ones(1, dtype=torch.float32, device=dev) + self.gemm.autotune_fp8_nn_dev(A.data_ptr(), B.data_ptr(), D.data_ptr(), + M, N, K, sa.data_ptr(), sb.data_ptr(), num_algos) + n += 1 + torch.cuda.synchronize() + return n + + def _fp4_gemm_f4(self, x, w4, wsf, N, K): + """F4-family W4A4 GEMM (Thor). x (M,K) bf16 -> (M,N) bf16. + + Activation is dynamically quantized to NVFP4 with the SAME swizzled SF + family as the weight (both from flash_rt_fp4), then cutlass_fp4_sq_fp16 + (A@Bᵀ, weight stored (N,K/2)+SF). Verified cos 0.990/GEMM. Graph-safe.""" + F4 = self._F4 + st = torch.cuda.current_stream().cuda_stream + xc = x.reshape(-1, K).contiguous().to(torch.float16) + M = xc.shape[0] + xp = torch.empty(M, K // 2, dtype=torch.uint8, device=x.device) + xsf = torch.empty(F4.sfa_size_bytes(M, K, False), dtype=torch.uint8, device=x.device) + F4.quantize_fp4_dynamic_sfa_fp16(xc.data_ptr(), xp.data_ptr(), xsf.data_ptr(), M, K, False, st) + out = torch.empty(M, N, dtype=torch.float16, device=x.device) + F4.cutlass_fp4_sq_fp16(xp.data_ptr(), xsf.data_ptr(), w4.data_ptr(), wsf.data_ptr(), + out.data_ptr(), M, N, K, 1.0, 0.0, st) + return out.to(torch.bfloat16) + + def _fp8_gemm(self, x, w8, ws): + """x (M,K) bf16 -> (M,N) bf16 via dynamic-scale FP8 (graph-safe). + + w8 is the fp8 weight stored (K,N); activation amax is computed on-GPU + each call (device scale). cuBLASLt device-scale FP8 GEMM is + CUDA-graph-capturable on Thor (verified).""" + K, N = w8.shape + xc = x.reshape(-1, K).contiguous() + M = xc.shape[0] + if self._gemm_shapes is not None: + self._gemm_shapes.add((M, N, K)) + st = torch.cuda.current_stream().cuda_stream + a8 = torch.empty(M, K, dtype=torch.uint8, device=x.device) + dsa = torch.empty(1, dtype=torch.float32, device=x.device) + if 0 < self._small_quant_m and M <= self._small_quant_m: + fvk.hyvla_quant_fp8_dyn_bf16(xc.data_ptr(), a8.data_ptr(), + dsa.data_ptr(), M * K, st) + else: + fvk.quantize_fp8_device(xc.data_ptr(), a8.data_ptr(), dsa.data_ptr(), M * K, st) + out = torch.empty(M, N, dtype=torch.bfloat16, device=x.device) + self.gemm.fp8_nn_dev(a8.data_ptr(), w8.data_ptr(), out.data_ptr(), + M, N, K, dsa.data_ptr(), ws.data_ptr(), st) + return out + + def _fp8_gemm_bias(self, x, w8, ws, bias): + """FP8 GEMM (…,K)->(…,N) + bias, for the biased ViT projections.""" + orig = x.shape + out = self._fp8_gemm(x.reshape(-1, orig[-1]), w8, ws) # (M,N) bf16 + out = out + bias + return out.reshape(*orig[:-1], out.shape[-1]) + + def _ffn_mega_bf16(self, hs_post, D, norm_w, mk): + """Expert denoise FFN via the two persistent megakernels. + + hs_post (1,S_s,D) bf16 is the post-attention residual stream (input to + the FFN AND the residual). mk = (gu8, sgu, dn8, sdn). Returns the new + hidden (1,S_s,D). Dynamic FP8: activation amax on-GPU each call + (graph-safe); weight scale is the baked host float. Scratch tensors go + to the graph private pool (per-call torch.empty, alias-safe).""" + gu8, sgu, dn8, sdn = mk + dev = hs_post.device + st = torch.cuda.current_stream().cuda_stream + S_s = hs_post.shape[1] + Nout = gu8.shape[0] // 2 # inter (gate+up merged -> inter) + INTER = dn8.shape[1] # dn K + hs_n = F.rms_norm(hs_post, (D,), norm_w, self.rms_eps)[0].contiguous() + x8 = torch.empty(S_s, D, dtype=torch.uint8, device=dev) + sx = torch.empty(1, dtype=torch.float32, device=dev) + fvk.quantize_fp8_device(hs_n.data_ptr(), x8.data_ptr(), sx.data_ptr(), S_s * D, st) + act = torch.empty(S_s, Nout, dtype=torch.bfloat16, device=dev) + fvk.hyvla_ffn_gu_silu_bf16(x8.data_ptr(), gu8.data_ptr(), act.data_ptr(), + S_s, D, Nout, sx.data_ptr(), sgu, st) + a8 = torch.empty(S_s, INTER, dtype=torch.uint8, device=dev) + sa = torch.empty(1, dtype=torch.float32, device=dev) + fvk.quantize_fp8_device(act.data_ptr(), a8.data_ptr(), sa.data_ptr(), S_s * INTER, st) + y = torch.empty(S_s, D, dtype=torch.bfloat16, device=dev) + fvk.hyvla_ffn_dn_res_bf16(a8.data_ptr(), dn8.data_ptr(), hs_post[0].contiguous().data_ptr(), + y.data_ptr(), S_s, INTER, D, sa.data_ptr(), sdn, st) + return y[None] + + # ══════════════════════════════════════════════════════════════════ + # ViT (vision tower) + merger — BF16 + # ══════════════════════════════════════════════════════════════════ + def _vit_qkv(self, h): + """h (bk, N, 1152) -> q,k,v each (bk, heads, N, 72).""" + bk, N, _ = h.shape + if getattr(self, "_vit_f8", False): + qkv = self._fp8_gemm_bias(h, self._vit_qkv_w8c, self._vit_qkv_wsc, self._vit_qkv_b_cur) + else: + qkv = F.linear(h, self._vit_qkv_w_cur, self._vit_qkv_b_cur) + qkv = qkv.reshape(bk, N, 3, self.vit_heads, self.vit_hd).permute(2, 0, 3, 1, 4) + return qkv[0], qkv[1], qkv[2] + + def _vit_spatial_attn(self, q, k, v): + """(bk, heads, N, 72) full non-causal attention -> proj. (bk, N, 1152).""" + bk, _, N, _ = q.shape + out = F.scaled_dot_product_attention(q, k, v, scale=self.vit_scale) + out = out.transpose(1, 2).reshape(bk, N, -1) + if getattr(self, "_vit_f8", False): + return self._fp8_gemm_bias(out, self._vit_proj_w8c, self._vit_proj_wsc, self._vit_proj_b_cur) + return F.linear(out, self._vit_proj_w_cur, self._vit_proj_b_cur) + + def _vit_time_pe(self, kf, device, dtype): + """Fixed sinusoidal e(t), base 100, e(0)=0. (kf, 1152).""" + dim = self.vit_heads * self.vit_hd + t = torch.arange(kf, dtype=torch.float32, device=device).unsqueeze(1) + inv_freq = torch.exp(torch.arange(0, dim, 2, dtype=torch.float32, device=device) + * (-torch.log(torch.tensor(self.vit_time_base)) / dim)) + pe = torch.empty(kf, dim, dtype=torch.float32, device=device) + pe[:, 0::2] = torch.sin(t * inv_freq) + pe[:, 1::2] = torch.cos(t * inv_freq) - 1.0 + return pe.to(dtype) + + def _vit_time_mix(self, q, k, v, b, kf): + """Causal-in-time softmax over K frames folded onto V. (bk,H,N,d).""" + bk, heads, n, d = v.shape + rs = lambda t: t.view(b, kf, heads, n, d).permute(0, 3, 2, 1, 4).reshape(b * n, heads, kf, d) + q_t, k_t, v_t = rs(q), rs(k), rs(v) + scores = (q_t @ k_t.transpose(-2, -1)) * self.vit_scale + mask = torch.triu(torch.ones(kf, kf, device=scores.device, dtype=torch.bool), 1) + scores = scores.masked_fill(mask, float("-inf")) + vm = scores.softmax(dim=-1).to(v_t.dtype) @ v_t + return vm.view(b, n, heads, kf, d).permute(0, 3, 2, 1, 4).reshape(bk, heads, n, d) + + def _vit_mlp(self, x): + if getattr(self, "_vit_f8", False): + x = self._fp8_gemm_bias(x, self._vit_fc1_w8c, self._vit_fc1_wsc, self._vit_fc1_b_cur) + x = F.gelu(x) + return self._fp8_gemm_bias(x, self._vit_fc2_w8c, self._vit_fc2_wsc, self._vit_fc2_b_cur) + x = F.linear(x, self._vit_fc1_w_cur, self._vit_fc1_b_cur) + x = F.gelu(x) + return F.linear(x, self._vit_fc2_w_cur, self._vit_fc2_b_cur) + + def _vit_block(self, x, li, num_frames): + """One ViT block; spacetime when li in vit_spacetime_ids.""" + W = self.W + self._vit_qkv_w_cur = W._vit_qkv_w[li]; self._vit_qkv_b_cur = W._vit_qkv_b[li] + self._vit_proj_w_cur = W._vit_proj_w[li]; self._vit_proj_b_cur = W._vit_proj_b[li] + self._vit_fc1_w_cur = W._vit_fc1_w[li]; self._vit_fc1_b_cur = W._vit_fc1_b[li] + self._vit_fc2_w_cur = W._vit_fc2_w[li]; self._vit_fc2_b_cur = W._vit_fc2_b[li] + self._vit_f8 = self._fp8 and getattr(W, "_vit_fp8_ready", False) + if self._vit_f8: + self._vit_qkv_w8c = W._vit_qkv_w8[li]; self._vit_qkv_wsc = W._vit_qkv_ws[li] + self._vit_proj_w8c = W._vit_proj_w8[li]; self._vit_proj_wsc = W._vit_proj_ws[li] + self._vit_fc1_w8c = W._vit_fc1_w8[li]; self._vit_fc1_wsc = W._vit_fc1_ws[li] + self._vit_fc2_w8c = W._vit_fc2_w8[li]; self._vit_fc2_wsc = W._vit_fc2_ws[li] + ln1w, ln1b = W._vit_ln1_w[li], W._vit_ln1_b[li] + ln2w, ln2b = W._vit_ln2_w[li], W._vit_ln2_b[li] + bk, n, d = x.shape + + if li in self.vit_spacetime_ids and num_frames > 1: + b, kf = bk // num_frames, num_frames + pe = self._vit_time_pe(kf, x.device, x.dtype) + h = F.layer_norm(x.view(b, kf, n, d) + pe.view(1, kf, 1, d), + (d,), ln1w, ln1b, self.vit_eps).view(bk, n, d) + q, k, v = self._vit_qkv(h) + v = self._vit_time_mix(q, k, v, b, kf) + attn_out = self._vit_spatial_attn(q, k, v) + else: + h = F.layer_norm(x, (d,), ln1w, ln1b, self.vit_eps) + q, k, v = self._vit_qkv(h) + attn_out = self._vit_spatial_attn(q, k, v) + + x = x + attn_out + x = x + self._vit_mlp(F.layer_norm(x, (d,), ln2w, ln2b, self.vit_eps)) + return x + + def _vit_pos_embed_rescale(self, h, w, dtype): + """Bilinear-rescale learned pos_embed (128x128) to (h,w). (1, h*w, 1152).""" + pos = self.W._vit_pos_embed # (1, 16384, 1152) + g = int(pos.shape[1] ** 0.5) # 128 + if (h, w) == (g, g): + return pos + pe2d = pos[0].T.contiguous().view(1, -1, g, g).float() + pe2d = F.interpolate(pe2d, (h, w), mode="bilinear", align_corners=False) + return pe2d.view(-1, h * w).T.contiguous()[None].to(dtype) + + @torch.no_grad() + def vit_forward(self, imgs): + """imgs (num_cam, K, 3, 224, 224) bf16 in [-1,1] -> (num_cam, 196, 1152).""" + W = self.W + num_cam, K = imgs.shape[0], imgs.shape[1] + bk = num_cam * K + x = imgs.reshape(bk, 3, 224, 224) + x = F.conv2d(x, W._vit_patch_w, W._vit_patch_b, stride=16) # (bk,1152,14,14) + hh = ww = x.shape[-1] + x = x.flatten(2).transpose(1, 2) # (bk,196,1152) + x = x + self._vit_pos_embed_rescale(hh, ww, x.dtype) + # History frames only feed the spacetime time-mix (causal over K, with + # the current frame last). After the FINAL spacetime block the remaining + # blocks are per-frame independent and only the current frame reaches the + # merger — so drop the history there. Numerically identical, and those + # blocks then do 1/K of the work. + last_st = max(self.vit_spacetime_ids) if K > 1 else -1 + sliced = False + for li in range(27): + x = self._vit_block(x, li, K) + if li == last_st and li < 26: + x = x.view(num_cam, K, hh * ww, -1)[:, -1] # (num_cam,N,D) + sliced = True + if not sliced: + x = x.view(num_cam, K, hh * ww, -1)[:, -1] # current frame + return x + + @torch.no_grad() + def merger_forward(self, x, grid=14): + """x (num_cam, 196, 1152) -> (num_cam, 49, 2048). NormalizedDwPooler 2x2.""" + W = self.W + B = x.shape[0] + h = w = grid + x = x.reshape(B, h, w, -1) + x = F.linear(x, W._mg_proj1_w, W._mg_proj1_b) # (B,14,14,2048) + C = x.shape[-1] + new_x = (x.reshape(B, h // 2, 2, w // 2, 2, C) + .permute(0, 1, 3, 2, 4, 5).reshape(B, h // 2, w // 2, 4, C)) + pooled = new_x.mean(-2, keepdim=True).expand(-1, -1, -1, 4, -1) + fused = torch.cat([new_x, pooled], dim=-1) # (B,7,7,4,4096) + score = F.linear(fused, W._mg_pred0_w, W._mg_pred0_b) + score = F.gelu(score) + score = F.linear(score, W._mg_pred2_w, W._mg_pred2_b) # (B,7,7,4,2048) + x = (new_x * score.softmax(dim=-2)).sum(dim=-2) # (B,7,7,2048) + x = F.gelu(x) + x = F.linear(x, W._mg_proj2_w, W._mg_proj2_b) + return x.reshape(B, -1, C) + + # ------------------------------------------------------------------ + def _attn(self, q, k, v, mask): + """GQA attention with a materialized bool mask. + + q (1, n_heads, S, hd) ; k/v (1, n_kv, S_kv, hd) ; mask (1,1,S,S_kv). + Expand KV heads to n_heads and drop enable_gqa so SDPA can pick the + memory-efficient backend (the bool-mask + enable_gqa combo forces the + slow math backend — measured ~41ms of E2E).""" + if k.shape[1] != q.shape[1]: + r = q.shape[1] // k.shape[1] + k = k.repeat_interleave(r, dim=1) + v = v.repeat_interleave(r, dim=1) + with torch.nn.attention.sdpa_kernel( + [torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, + torch.nn.attention.SDPBackend.MATH]): + return F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + + # ------------------------------------------------------------------ + def _block(self, hs, n_vis, w_text, w_vis, qk_w, mask, cos, sin, + kbuf, vbuf, off, fp8w=None, fp8v=None, fp8t=None, + fp4v=None, fp4t=None, ffn_mk=None): + """One MoT transformer block over sorted ``[vision|text]`` tokens. + + ``w_text``/``w_vis`` are per-branch weight tuples + ``(qkv, o, gu, d, ln_in, ln_post)``. When ``w_text is None`` every + token uses ``w_vis`` (the all-vision expert suffix). ``fp8w`` (expert) + or ``fp8v``/``fp8t`` (prefill vision/text branches) = + ``(qkv8,qkv_ws,o8,o_ws,gu8,gu_ws,d8,d_ws)`` enable graph-safe dynamic + FP8 for the GEMMs. Writes rope+norm'd K/V into ``kbuf``/``vbuf`` at row + ``off`` and attends over ``[:off+S]``. + """ + B, S, D = hs.shape + hd, nh, nkv = self.head_dim, self.n_heads, self.n_kv + _fp8 = self._fp8 and fp8w is not None and w_text is None + _fp8p = self._fp8 and fp8v is not None and w_text is not None + _fp4p = self._fp4 and fp4v is not None and w_text is not None + + if w_text is None: + hs_n = F.rms_norm(hs, (D,), w_vis[4], self.rms_eps) + qkv = self._fp8_gemm(hs_n[0], fp8w[0], fp8w[1]) if _fp8 else hs_n[0] @ w_vis[0].t() + else: + hs_v = F.rms_norm(hs[0, :n_vis], (D,), w_vis[4], self.rms_eps) + hs_t = F.rms_norm(hs[0, n_vis:], (D,), w_text[4], self.rms_eps) + if _fp8p: + qkv = torch.cat([self._fp8_gemm(hs_v, fp8v[0], fp8v[1]), + self._fp8_gemm(hs_t, fp8t[0], fp8t[1])], 0) + else: + qkv = torch.cat([hs_v @ w_vis[0].t(), hs_t @ w_text[0].t()], 0) + + if getattr(self, "_fused_attn", False): + st = torch.cuda.current_stream().cuda_stream + S_tot = kbuf.shape[2] + kv_rep = kbuf.shape[1] // nkv # 4 when the cache is pre-expanded + q = torch.empty(1, nh, S, hd, dtype=torch.bfloat16, device=hs.device) + fvk.hyvla_rope_qknorm_kvwrite_bf16( + qkv.contiguous().data_ptr(), + cos.reshape(S, hd).contiguous().data_ptr(), + sin.reshape(S, hd).contiguous().data_ptr(), + qk_w[0].data_ptr(), qk_w[1].data_ptr(), + q.data_ptr(), kbuf.data_ptr(), vbuf.data_ptr(), + S, nh, nkv, hd, S_tot, off, self.rms_eps, kv_rep, st) + else: + q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], -1) + q = q.view(S, nh, hd).transpose(0, 1)[None] + k = k.view(S, nkv, hd).transpose(0, 1)[None] + v = v.view(S, nkv, hd).transpose(0, 1)[None] + + # RoPE (rotate_half) THEN QK-Norm (RMSNorm over head_dim). + q = q * cos + _rot_half(q) * sin + k = k * cos + _rot_half(k) * sin + q = F.rms_norm(q, (hd,), qk_w[0], self.rms_eps) + k = F.rms_norm(k, (hd,), qk_w[1], self.rms_eps) + + if kbuf.shape[1] != nkv: # pre-expanded cache + r = kbuf.shape[1] // nkv + k = k.repeat_interleave(r, dim=1) + v = v.repeat_interleave(r, dim=1) + kbuf[:, :, off:off + S].copy_(k) + vbuf[:, :, off:off + S].copy_(v) + k_use = kbuf[:, :, : off + S] + v_use = vbuf[:, :, : off + S] + + att = self._attn(q, k_use, v_use, mask) + att = att.transpose(1, 2).reshape(1, S, self.q_dim) + + if w_text is None: + o = self._fp8_gemm(att[0], fp8w[2], fp8w[3]) if _fp8 else att[0] @ w_vis[1].t() + hs = hs + o[None] + if self._ffn_mega and ffn_mk is not None and _fp8: + hs = self._ffn_mega_bf16(hs, D, w_vis[5], ffn_mk) + else: + hs_n = F.rms_norm(hs, (D,), w_vis[5], self.rms_eps) + gu = self._fp8_gemm(hs_n[0], fp8w[4], fp8w[5]) if _fp8 else hs_n[0] @ w_vis[2].t() + g, u = gu.chunk(2, -1) + act = F.silu(g) * u + dn = self._fp8_gemm(act, fp8w[6], fp8w[7]) if _fp8 else act @ w_vis[3].t() + hs = hs + dn[None] + else: + if _fp8p: + o_v = self._fp8_gemm(att[0, :n_vis], fp8v[2], fp8v[3]) + o_t = self._fp8_gemm(att[0, n_vis:], fp8t[2], fp8t[3]) + else: + o_v = att[0, :n_vis] @ w_vis[1].t() + o_t = att[0, n_vis:] @ w_text[1].t() + hs = hs + torch.cat([o_v, o_t], 0)[None] + hs_v = F.rms_norm(hs[0, :n_vis], (D,), w_vis[5], self.rms_eps) + hs_t = F.rms_norm(hs[0, n_vis:], (D,), w_text[5], self.rms_eps) + if _fp4p: + N_gu = fp4v[4] + gu = torch.cat([self._fp4_gemm_f4(hs_v, fp4v[0], fp4v[1], N_gu, D), + self._fp4_gemm_f4(hs_t, fp4t[0], fp4t[1], N_gu, D)], 0) + elif _fp8p: + gu = torch.cat([self._fp8_gemm(hs_v, fp8v[4], fp8v[5]), + self._fp8_gemm(hs_t, fp8t[4], fp8t[5])], 0) + else: + gu = torch.cat([hs_v @ w_vis[2].t(), hs_t @ w_text[2].t()], 0) + g, u = gu.chunk(2, -1) + act = F.silu(g) * u + if _fp4p: + inter, Dh = fp4v[5], fp4v[6] + dn = torch.cat([self._fp4_gemm_f4(act[:n_vis], fp4v[2], fp4v[3], Dh, inter), + self._fp4_gemm_f4(act[n_vis:], fp4t[2], fp4t[3], Dh, inter)], 0) + elif _fp8p: + dn = torch.cat([self._fp8_gemm(act[:n_vis], fp8v[6], fp8v[7]), + self._fp8_gemm(act[n_vis:], fp8t[6], fp8t[7])], 0) + else: + dn = torch.cat([act[:n_vis] @ w_vis[3].t(), + act[n_vis:] @ w_text[3].t()], 0) + hs = hs + dn[None] + return hs + + # ------------------------------------------------------------------ + def _vlm_w(self, li): + W = self.W + text = (W._vlm_qkv_t[li], W._vlm_o_t[li], W._vlm_gu_t[li], + W._vlm_d_t[li], W._vlm_ln_in_t[li], W._vlm_ln_post_t[li]) + vis = (W._vlm_qkv_v[li], W._vlm_o_v[li], W._vlm_gu_v[li], + W._vlm_d_v[li], W._vlm_ln_in_v[li], W._vlm_ln_post_v[li]) + return text, vis + + def _vlm_w_fp8(self, li): + W = self.W + if not getattr(W, "_vlm_fp8_ready", False): + return None, None + vis = (W._vlm_qkv_v8[li], W._vlm_qkv_v_ws[li], W._vlm_o_v8[li], W._vlm_o_v_ws[li], + W._vlm_gu_v8[li], W._vlm_gu_v_ws[li], W._vlm_d_v8[li], W._vlm_d_v_ws[li]) + text = (W._vlm_qkv_t8[li], W._vlm_qkv_t_ws[li], W._vlm_o_t8[li], W._vlm_o_t_ws[li], + W._vlm_gu_t8[li], W._vlm_gu_t_ws[li], W._vlm_d_t8[li], W._vlm_d_t_ws[li]) + return vis, text + + def _vlm_w_fp4(self, li): + W = self.W + if not getattr(W, "_vlm_fp4_ready", False): + return None, None + vis = (W._vlm_gu_v4[li], W._vlm_gu_v4sf[li], W._vlm_d_v4[li], W._vlm_d_v4sf[li], + W._vlm_gu_N, W._vlm_inter, W._vlm_D) + text = (W._vlm_gu_t4[li], W._vlm_gu_t4sf[li], W._vlm_d_t4[li], W._vlm_d_t4sf[li], + W._vlm_gu_N, W._vlm_inter, W._vlm_D) + return vis, text + + def _exp_w(self, li): + W = self.W + return (W._exp_qkv_v[li], W._exp_o_v[li], W._exp_gu_v[li], + W._exp_d_v[li], W._exp_ln_in_v[li], W._exp_ln_post_v[li]) + + def _exp_w_fp8(self, li): + W = self.W + if not getattr(W, "_exp_fp8_ready", False): + return None + return (W._exp_qkv8[li], W._exp_qkv_ws[li], W._exp_o8[li], W._exp_o_ws[li], + W._exp_gu8[li], W._exp_gu_ws[li], W._exp_d8[li], W._exp_d_ws[li]) + + def _exp_ffn_mk(self, li): + W = self.W + if not getattr(W, "_exp_ffn_mega_ready", False): + return None + return (W._exp_gu_mk[li], W._exp_gu_mk_s[li], W._exp_d_mk[li], W._exp_d_mk_s[li]) + + # ------------------------------------------------------------------ + @torch.no_grad() + def prefill(self, prefix_embs, n_vis, pmask, pcos, psin, kbuf, vbuf): + """Run the 32-layer MoT VLM tower; fills kbuf/vbuf rows [0:S_p].""" + hs = prefix_embs + for li in range(32): + text, vis = self._vlm_w(li) + qk = (self.W._qk_norm_q[li], self.W._qk_norm_k[li]) + fp8v, fp8t = self._vlm_w_fp8(li) + fp4v, fp4t = self._vlm_w_fp4(li) + hs = self._block(hs, n_vis, text, vis, qk, pmask, pcos, psin, + kbuf[li], vbuf[li], 0, fp8v=fp8v, fp8t=fp8t, + fp4v=fp4v, fp4t=fp4t) + return hs + + # ------------------------------------------------------------------ + @torch.no_grad() + def denoise(self, state, x_t, time_embs, smask, scos, ssin, + kbuf, vbuf, S_p, num_steps=10): + """32-layer expert tower + action head, ``num_steps`` Euler steps. + + ``x_t`` (1, chunk, 32) fp32 is updated in place and returned. + ``time_embs`` (num_steps, 1, D_exp) bf16 precomputed by frontend. + """ + W = self.W + S_s = 1 + x_t.shape[1] + dt = -1.0 / num_steps + state_emb = (F.linear(state.to(torch.bfloat16), W._state_w, W._state_b))[:, None] + for s in range(num_steps): + action_emb = F.linear(x_t.to(torch.bfloat16), W._ain_w, W._ain_b) + t_emb = time_embs[s].expand_as(action_emb) + ate = torch.cat([action_emb, t_emb], 2) + ate = F.linear(ate, W._atmlp_in_w, W._atmlp_in_b) + ate = F.silu(ate) + ate = F.linear(ate, W._atmlp_out_w, W._atmlp_out_b) + hs = torch.cat([state_emb, ate], 1) + for li in range(32): + exp = self._exp_w(li) + qk = (W._qk_norm_q[li], W._qk_norm_k[li]) + hs = self._block(hs, S_s, None, exp, qk, smask, scos, ssin, + kbuf[li], vbuf[li], S_p, fp8w=self._exp_w_fp8(li), + ffn_mk=self._exp_ffn_mk(li)) + hs = F.rms_norm(hs, (hs.shape[-1],), W._exp_final_norm_w, self.rms_eps) + v_t = F.linear(hs[:, -x_t.shape[1]:], W._aout_w, W._aout_b) + x_t.add_(dt * v_t.to(x_t.dtype)) # in-place: static buffer for CUDA-graph replay + return x_t + + +__all__ = ["HyVLAThorBF16Pipeline"] From 06e5526d82c671b76a8e67626b453484827405c5 Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:06:25 +0800 Subject: [PATCH 2/3] fix(hyvla): model build option, FP4 routing fix, safe contracts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback: Build isolation: - New FLASHRT_ENABLE_HYVLA CMake option (OFF by default). Both arch blocks (SM110 full set, SM87 portable shared kernels) now require the option together with the matching GPU_ARCH; sources, compile definitions, and pybind symbols gate together. HyVLA OFF leaves the default SM87/SM110 build and symbol surface unchanged. FP4 routing correctness: - load_model(config="hyvla", use_fp4=True) previously dropped use_fp4 before constructing the frontend. A dedicated _hyvla_fp4 flag now carries the route and the hyvla kwarg branch passes use_fp4=True explicitly; the validated fp8 production tier selects use_fused=True and use_autotune when the frontend accepts them. Hardware fail-fast and safe loading: - HyVLATorchFrontendThor validates torch.cuda capability (11,0) before loading weights or allocating CUDA state; a documented FLASHRT_HYVLA_FORCE_ARCH env override skips the probe for development only. The check is inheritable by the Orin subclass. - Tokenizer loading no longer uses trust_remote_code — checkpoint- provided Python code is never executed during model loading. Public input contracts (deterministic under python -O): - predict_actions: missing-prompt assert replaced with RuntimeError; state wider than max_state_dim and wrong-element-count noise raise ValueError before any graph lookup or copy. - All five fused-kernel pybind APIs validate their launch contract (positive dims, hd==128, off+S<=S_tot, eps>0, kv_rep>=1, K%16 and N/Nout%32 alignment) and raise ValueError instead of launching invalid grids. Reproducible in-tree validation: - tests/test_hyvla_thor_dispatch.py — registration one-to-one gate - tests/test_hyvla_arch_gate.py — mocked-capability fail-fast tests - tests/test_hyvla_fp4_routing.py — proves use_fp4/use_fused reach the frontend (stubbed construction, mocked flash_rt_fp4) - tests/test_hyvla_kernel_contracts.py — negative-dimension contract tests for all five kernels - tests/test_hyvla_thor_graphsafe.py — graph==eager + replay stability gate with deterministic seed-0 inputs (checkpoint via env var) - docs/hyvla05_thor_sm110.md updated to reference only committed tests and document the build option. --- CMakeLists.txt | 11 ++-- csrc/bindings.cpp | 50 +++++++++++++++++ docs/hyvla05_thor_sm110.md | 36 ++++++------ flash_rt/api.py | 13 +++++ flash_rt/frontends/torch/hyvla_thor.py | 48 +++++++++++++++- tests/test_hyvla_arch_gate.py | 53 ++++++++++++++++++ tests/test_hyvla_fp4_routing.py | 76 +++++++++++++++++++++++++ tests/test_hyvla_kernel_contracts.py | 77 ++++++++++++++++++++++++++ tests/test_hyvla_thor_dispatch.py | 27 +++++++++ tests/test_hyvla_thor_graphsafe.py | 71 ++++++++++++++++++++++++ 10 files changed, 439 insertions(+), 23 deletions(-) create mode 100644 tests/test_hyvla_arch_gate.py create mode 100644 tests/test_hyvla_fp4_routing.py create mode 100644 tests/test_hyvla_kernel_contracts.py create mode 100644 tests/test_hyvla_thor_dispatch.py create mode 100644 tests/test_hyvla_thor_graphsafe.py diff --git a/CMakeLists.txt b/CMakeLists.txt index e042cc095..127365371 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1692,8 +1692,12 @@ target_include_directories(flash_rt_kernels PRIVATE ${CUTLASS_DIR}/tools/util/include ) -# ── Hy-Embodied-0.5-VLA Thor SM110 kernels ── -if(GPU_ARCH STREQUAL "110") +# ── Hy-Embodied-0.5-VLA model kernels (opt-in) ── +# HyVLA sources compile only when FLASHRT_ENABLE_HYVLA=ON together with the +# matching architecture gate. HyVLA OFF leaves the default SM87/SM110 build +# and symbol surface unchanged. +option(FLASHRT_ENABLE_HYVLA "Build HyVLA model kernels" OFF) +if(FLASHRT_ENABLE_HYVLA AND GPU_ARCH STREQUAL "110") target_sources(flash_rt_kernels PRIVATE csrc/kernels/hyvla_fused_thor.cu csrc/kernels/hyvla_vit_fuse.cu @@ -1703,11 +1707,10 @@ if(GPU_ARCH STREQUAL "110") message(STATUS "Hy-VLA Thor SM110 kernels: ENABLED") endif() -# ── Hy-Embodied-0.5-VLA Orin SM87 kernels ── # The fused RoPE+QKNorm+KV-write megakernel is plain CUDA (no SM110-only # instructions) and removes ~11 launches per attention block; the FP8 # Thor kernels stay SM110-only. -if(GPU_ARCH STREQUAL "87") +if(FLASHRT_ENABLE_HYVLA AND GPU_ARCH STREQUAL "87") target_sources(flash_rt_kernels PRIVATE csrc/kernels/hyvla_fused_thor.cu csrc/kernels/hyvla_vit_fuse.cu) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index c15483614..93dfe4351 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -8126,6 +8126,25 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; uintptr_t kn_w, uintptr_t q_out, uintptr_t kbuf, uintptr_t vbuf, int S, int nq, int nkv, int hd, int S_tot, int off, float eps, int kv_rep, uintptr_t stream) { + if (S <= 0 || nq <= 0 || nkv <= 0 || S_tot <= 0) + throw py::value_error( + "hyvla_rope_qknorm_kvwrite_bf16 requires S>0, nq>0, " + "nkv>0, S_tot>0"); + if (hd != 128) + throw py::value_error( + "hyvla_rope_qknorm_kvwrite_bf16 supports hd==128 only, " + "got " + std::to_string(hd)); + if (off < 0 || off + S > S_tot) + throw py::value_error( + "hyvla_rope_qknorm_kvwrite_bf16: invalid offset window " + "off=" + std::to_string(off) + " S=" + std::to_string(S) + + " S_tot=" + std::to_string(S_tot)); + if (!(eps > 0.f)) + throw py::value_error( + "hyvla_rope_qknorm_kvwrite_bf16 requires eps>0"); + if (kv_rep < 1) + throw py::value_error( + "hyvla_rope_qknorm_kvwrite_bf16 requires kv_rep>=1"); hyvla_rope_qknorm_kvwrite_bf16( reinterpret_cast(qkv), reinterpret_cast(cos), @@ -8149,6 +8168,13 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; [](uintptr_t residual, uintptr_t x_add, uintptr_t ln_weight, uintptr_t ln_bias, uintptr_t out, int rows, int dim, float eps, uintptr_t stream) { + if (rows <= 0 || dim <= 0 || (dim & 1) != 0) + throw py::value_error( + "hyvla_vit_add_layer_norm_bf16 requires rows>0 and a " + "positive even dim"); + if (!(eps > 0.f)) + throw py::value_error( + "hyvla_vit_add_layer_norm_bf16 requires eps>0"); hyvla_vit_add_layer_norm_bf16( reinterpret_cast(residual), reinterpret_cast(x_add), @@ -8165,6 +8191,8 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; #ifdef FLASHRT_HAVE_HYVLA_THOR m.def("hyvla_quant_fp8_dyn_bf16", [](uintptr_t x, uintptr_t out, uintptr_t scale, int n, uintptr_t stream) { + if (n <= 0) + throw py::value_error("hyvla_quant_fp8_dyn_bf16 requires n>0"); hyvla_quant_fp8_dyn_bf16( reinterpret_cast(x), reinterpret_cast(out), @@ -8177,6 +8205,17 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; m.def("hyvla_ffn_gu_silu_bf16", [](uintptr_t x, uintptr_t gu, uintptr_t act, int M, int K, int Nout, uintptr_t sx, float sgu, uintptr_t stream) { + if (M <= 0 || K <= 0 || Nout <= 0) + throw py::value_error( + "hyvla_ffn_gu_silu_bf16 requires M>0, K>0, Nout>0"); + if (K % 16 != 0) + throw py::value_error( + "hyvla_ffn_gu_silu_bf16 requires K%16==0, got K=" + + std::to_string(K)); + if (Nout % 32 != 0) + throw py::value_error( + "hyvla_ffn_gu_silu_bf16 requires Nout%32==0, got Nout=" + + std::to_string(Nout)); hyvla_ffn_gu_silu_bf16( reinterpret_cast(x), reinterpret_cast(gu), reinterpret_cast(act), M, K, Nout, @@ -8189,6 +8228,17 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; m.def("hyvla_ffn_dn_res_bf16", [](uintptr_t a, uintptr_t dn, uintptr_t res, uintptr_t y, int M, int K, int N, uintptr_t sa, float sdn, uintptr_t stream) { + if (M <= 0 || K <= 0 || N <= 0) + throw py::value_error( + "hyvla_ffn_dn_res_bf16 requires M>0, K>0, N>0"); + if (K % 16 != 0) + throw py::value_error( + "hyvla_ffn_dn_res_bf16 requires K%16==0, got K=" + + std::to_string(K)); + if (N % 32 != 0) + throw py::value_error( + "hyvla_ffn_dn_res_bf16 requires N%32==0, got N=" + + std::to_string(N)); hyvla_ffn_dn_res_bf16( reinterpret_cast(a), reinterpret_cast(dn), reinterpret_cast(res), reinterpret_cast(y), diff --git a/docs/hyvla05_thor_sm110.md b/docs/hyvla05_thor_sm110.md index f2f92040e..bcebb79b3 100644 --- a/docs/hyvla05_thor_sm110.md +++ b/docs/hyvla05_thor_sm110.md @@ -11,10 +11,12 @@ > + ViT trailing-stage history-frame drop + per-shape FP8 GEMM autotune. > > ```bash +> cmake -B build -S . -DGPU_ARCH=110 -DFLASHRT_ENABLE_HYVLA=ON > cmake --build build --target flash_rt_kernels # fused kernels (first build) -> PYTHONPATH=. python3 tests/test_thor_hyvla05_e2e_check.py --fp8 --fused # precision gate cos >= 0.999 -> PYTHONPATH=. python3 tests/test_thor_hyvla05_graphsafe.py --fp8 --fused # graph-safety bitwise gate -> PYTHONPATH=. python3 tests/test_thor_hyvla05_stageprof.py --fp8 --fused # per-stage latency +> python -m pytest tests/test_hyvla_thor_dispatch.py tests/test_hyvla_arch_gate.py \ +> tests/test_hyvla_fp4_routing.py tests/test_hyvla_kernel_contracts.py -q +> FLASHRT_HYVLA_CHECKPOINT=/path/to/Hy-Embodied-0.5-VLA \ +> python -m pytest tests/test_hyvla_thor_graphsafe.py -q # graph==eager + replay stability > ``` ## Key Takeaway @@ -42,11 +44,11 @@ Industry anchor: NVIDIA achieves 44 ms / 23 Hz on Pi0.5 with hand-written kernel | `csrc/kernels/hyvla_quant_fp8_thor.cu` | Single-CTA dynamic FP8 quantization (`use_fused_quant` diagnostic; measured net loss, see dead ends) | | `csrc/kernels/hyvla_ffn_fp8_thor.cu` | FFN megakernel (gu+silu_mul / dn+residual, `use_ffn_mega` diagnostic; measured neutral, see dead ends); reusable Thor plain-FP8-MMA GEMM reference | | `flash_rt/{executors/torch_weights.py,hardware/__init__.py,api.py,configs/hyvla.yaml}` | `ToBf16` transform / `_PIPELINE_MAP` registration / config allowlist / metadata | -| `tests/test_thor_hyvla05_{baseline,tower_check,vit_check,e2e_check}.py` | Precision oracles + segmented/full-chain precision gates (`--fp8 --fused --fp4`) | -| `tests/test_thor_hyvla05_{bench,stageprof,graphsafe}.py` | Latency benchmark / per-stage profiling / graph-safety bitwise gate | -| `tests/test_thor_hyvla05_{gemm_ceiling,vit_prof,fp4_sf_check}.py` | GEMM ceiling / ViT internal profiling / NVFP4 SF validation | +| `tests/test_hyvla_thor_dispatch.py`, `test_hyvla_arch_gate.py`, `test_hyvla_fp4_routing.py` | Registration, SM110 fail-fast, and FP4/fused routing gates (no GPU required) | +| `tests/test_hyvla_kernel_contracts.py` | Host-side contract validation of the fused-kernel pybind APIs | +| `tests/test_hyvla_thor_graphsafe.py` | graph==eager equivalence + replay stability gate (needs Thor + checkpoint) | -## Precision Gates (all vs. HF/transformers eager, same fixed noise; `e2e_check`/`tower_check`/`vit_check`) +## Precision Gates (recorded values, all vs. HF/transformers eager, same fixed noise) | Checkpoint | cosine | |---|---| @@ -55,8 +57,8 @@ Industry anchor: NVIDIA achieves 44 ms / 23 Hz on Pi0.5 with hand-written kernel | Full-chain native BF16 (load_model path) | 0.999910 | | **Full-chain production (fp8 + fused + efficient-SDPA + autotune)** | **0.999706** | -Graph-safety gate (`graphsafe.py`): `use_graph=True vs False` bitwise identical (`max|delta|=0`) + stable replay, -verified for both the ViT graph and the main graph. +Graph-safety gate (`tests/test_hyvla_thor_graphsafe.py`): `use_graph=True vs False` +equivalence + stable replay, covering the ViT graph and the main graph. ## Key Mechanisms / Correctness Pitfalls (Highest Reuse Value) @@ -153,13 +155,15 @@ The building blocks are in place and have high reuse value: plain-FP8-MMA layout ## Early Inductor Prototype (`--v4`, Historical Reference, Not Production) -`tests/test_thor_hyvla05_fast.py --v4 --compile-vit`: `torch.compile` whole-block fusion + static FP8 + -grouped-bmm/FA2 + full-graph capture, E2E **98.8 ms (8.6–9.1x)**, cos 0.9985–0.9998. It proved the benefit -of the "whole-layer fusion" approach (ViT 51.5, denoise 29 are both below native), but due to **30–60 s -recompilation per prompt, 15% Inductor cache-bloat degradation, silent dynamo fallback, and no integration -with `load_model`**, it is not used in production. Its six measurement methodology lessons (L2 residency -illusion / no-fusion baseline illusion / nsys sum != wall clock / synchronization floor / Inductor cache -bloat / dynamo silent downgrade) and the FP4-in-Inductor verdict have been distilled into internal documentation. +An early `torch.compile` whole-block fusion + static FP8 + grouped-bmm/FA2 + +full-graph capture prototype reached E2E **98.8 ms (8.6–9.1x)**, cos 0.9985–0.9998. +It proved the benefit of the "whole-layer fusion" approach (ViT 51.5, denoise 29 +are both below native), but due to **30–60 s recompilation per prompt, 15% +Inductor cache-bloat degradation, silent dynamo fallback, and no integration +with `load_model`**, it is not used in production. The six measurement +methodology lessons it exposed (L2 residency illusion / no-fusion baseline +illusion / nsys sum != wall clock / synchronization floor / Inductor cache +bloat / dynamo silent downgrade) informed the production path design. ## Hardware and Model Profile (Measured) diff --git a/flash_rt/api.py b/flash_rt/api.py index 436176d82..cb09a7bd4 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -700,6 +700,7 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, use_fp4 = False # do not fall through to the Pi0.5 FP4 routing # ── FP4 routing (Pi0.5 torch + Pi0.5 JAX on Thor, HyVLA torch on Thor) ── + _hyvla_fp4 = False if use_fp4: _fp4_ok = ( (config == "pi05" and framework in ("torch", "jax") and arch == "thor") @@ -738,6 +739,7 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, HyVLATorchFrontendThor, ) pipe_cls = HyVLATorchFrontendThor + _hyvla_fp4 = True logger.info("HyVLA Thor FP4 tier enabled") use_fp4 = False # routed; skip Pi0.5 path below elif framework == "torch": @@ -784,6 +786,17 @@ 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 == "hyvla": + # The routed FP4 tier must reach the frontend explicitly; the + # generic kwarg set never forwards use_fp4. + if _hyvla_fp4 and "use_fp4" in sig.parameters: + kwargs["use_fp4"] = True + # FP8 tier = the validated production config (fp8 + fused + # megakernels); select it explicitly when the frontend accepts it. + if use_fp8 and "use_fused" in sig.parameters: + kwargs["use_fused"] = True + if "use_autotune" in sig.parameters: + kwargs["use_autotune"] = bool(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. diff --git a/flash_rt/frontends/torch/hyvla_thor.py b/flash_rt/frontends/torch/hyvla_thor.py index a3f4a8fde..d27a6f9ec 100644 --- a/flash_rt/frontends/torch/hyvla_thor.py +++ b/flash_rt/frontends/torch/hyvla_thor.py @@ -78,6 +78,31 @@ def _camera_tensor(image): class HyVLATorchFrontendThor: + #: Development override for the hardware gate (e.g. running the weight + #: loader on a non-Thor box). Any non-empty value skips the capability + #: probe; kernels still require the real hardware at runtime. Not a + #: supported production path. + _FORCE_ARCH_ENV = "FLASHRT_HYVLA_FORCE_ARCH" + _REQUIRED_CAPABILITY = (11, 0) + _ARCH_NAME = "Jetson Thor SM110" + + def _require_arch(self): + import os + + if os.environ.get(self._FORCE_ARCH_ENV, ""): + return # explicit documented dev override: skip the probe + if not torch.cuda.is_available(): + raise RuntimeError( + f"HyVLA frontend requires a CUDA device ({self._ARCH_NAME}); " + "CUDA is not available.") + cap = torch.cuda.get_device_capability() + if cap != self._REQUIRED_CAPABILITY: + raise RuntimeError( + f"HyVLA frontend requires {self._ARCH_NAME} (capability " + f"{self._REQUIRED_CAPABILITY}), found capability {cap}. Set " + f"{self._FORCE_ARCH_ENV}=1 to bypass this check for " + "development only.") + def __init__(self, checkpoint_dir: str, *, hardware: str = "thor", use_fp8: bool = False, use_fp8_vit: bool = False, use_fused: bool = False, use_fp4: bool = False, @@ -85,6 +110,7 @@ def __init__(self, checkpoint_dir: str, *, hardware: str = "thor", use_ffn_mega: bool = False, **kwargs): self.checkpoint_dir = str(checkpoint_dir) self.device = "cuda" + self._require_arch() self.use_fp8 = bool(use_fp8) self.use_fp8_vit = bool(use_fp8_vit) self.use_fused = bool(use_fused) @@ -238,7 +264,9 @@ def _quantize_fp4(self): def q_nvfp4(w_bf16): N, K = w_bf16.shape - assert K % 64 == 0, f"cutlass_fp4_sq_fp16 needs K%64==0, got K={K}" + if K % 64 != 0: + raise ValueError( + f"cutlass_fp4_sq_fp16 needs K%64==0, got K={K}") w16 = w_bf16.to(torch.float16).contiguous() packed = torch.empty(N, K // 2, dtype=torch.uint8, device=self.device) sf = torch.empty(F4.sfa_size_bytes(N, K, True), dtype=torch.uint8, device=self.device) @@ -299,8 +327,10 @@ def _rope_cos_sin(self, positions): def tokenizer(self): if self._tokenizer is None: from transformers import AutoTokenizer + # Standard tokenizer path only — never execute checkpoint-provided + # Python code (no trust_remote_code). self._tokenizer = AutoTokenizer.from_pretrained( - self.checkpoint_dir, trust_remote_code=True) + self.checkpoint_dir) return self._tokenizer def _tokenize(self, prompt): @@ -554,7 +584,8 @@ def predict_actions(self, images, prompt=None, state=None, noise=None, Returns raw action chunk (1, chunk, max_action_dim) as numpy fp32.""" if prompt is not None and prompt != self._prompt: self.set_prompt(prompt) - assert self._lang_tokens is not None, "call set_prompt() first" + if self._lang_tokens is None: + raise RuntimeError("call set_prompt() before predict_actions()") dev = self.device if not torch.is_tensor(images): @@ -568,6 +599,10 @@ def predict_actions(self, images, prompt=None, state=None, noise=None, state_t = torch.zeros(1, self.max_state_dim, device=dev, dtype=_BF16) else: st = torch.as_tensor(np.asarray(state), device=dev, dtype=_BF16).reshape(1, -1) + if st.shape[1] > self.max_state_dim: + raise ValueError( + f"state has {st.shape[1]} dims, max_state_dim is " + f"{self.max_state_dim}") if st.shape[1] < self.max_state_dim: st = F.pad(st, (0, self.max_state_dim - st.shape[1])) state_t = st @@ -604,6 +639,13 @@ def predict_actions(self, images, prompt=None, state=None, noise=None, dtype=torch.float32, device=dev) else: noise_t = torch.as_tensor(np.asarray(noise), device=dev, dtype=torch.float32) + want = self.chunk * self.max_action_dim + if noise_t.numel() != want: + raise ValueError( + f"noise must have {want} elements " + f"(chunk={self.chunk} x max_action_dim={self.max_action_dim}), " + f"got {noise_t.numel()}") + noise_t = noise_t.reshape(1, self.chunk, self.max_action_dim) x_t = self._graph_forward(S_p, n_vis, prefix_embs, pmask, pcos, psin, smask, scos, ssin, state_t, noise_t, diff --git a/tests/test_hyvla_arch_gate.py b/tests/test_hyvla_arch_gate.py new file mode 100644 index 000000000..126478421 --- /dev/null +++ b/tests/test_hyvla_arch_gate.py @@ -0,0 +1,53 @@ +"""HyVLA hardware-gate (fail-fast) tests — torch.cuda is mocked, no GPU needed.""" + +import os + +import pytest + +torch = pytest.importorskip("torch") + +try: + import flash_rt.frontends.torch.hyvla_thor as hy_mod +except ImportError as exc: # pragma: no cover + pytest.skip(f"hyvla_thor frontend not importable: {exc}", allow_module_level=True) + + +class _Probe: + """Run _require_arch against mocked CUDA state.""" + + def __init__(self, available, capability): + self._cls = hy_mod.HyVLATorchFrontendThor + self._available = available + self._capability = capability + + def run(self): + obj = object.__new__(self._cls) + return self._cls._require_arch(obj) + + +def test_rejects_when_cuda_unavailable(monkeypatch): + monkeypatch.delenv("FLASHRT_HYVLA_FORCE_ARCH", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="CUDA is not available"): + _Probe(False, None).run() + + +def test_rejects_wrong_capability(monkeypatch): + monkeypatch.delenv("FLASHRT_HYVLA_FORCE_ARCH", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 9)) + with pytest.raises(RuntimeError, match="requires Jetson Thor SM110"): + _Probe(True, (8, 9)).run() + + +def test_accepts_sm110(monkeypatch): + monkeypatch.delenv("FLASHRT_HYVLA_FORCE_ARCH", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (11, 0)) + _Probe(True, (11, 0)).run() # must not raise + + +def test_documented_env_override_skips_probe(monkeypatch): + monkeypatch.setenv("FLASHRT_HYVLA_FORCE_ARCH", "1") + # No CUDA mocking: the override must return before touching torch.cuda. + _Probe(False, None).run() diff --git a/tests/test_hyvla_fp4_routing.py b/tests/test_hyvla_fp4_routing.py new file mode 100644 index 000000000..0fd2e293b --- /dev/null +++ b/tests/test_hyvla_fp4_routing.py @@ -0,0 +1,76 @@ +"""load_model routing tests for HyVLA: the FP4 tier must reach the frontend.""" + +import sys +import types + +import pytest + +torch = pytest.importorskip("torch") + +try: + import flash_rt.frontends.torch.hyvla_thor as hy_mod +except ImportError as exc: # pragma: no cover + pytest.skip(f"hyvla_thor frontend not importable: {exc}", allow_module_level=True) + +import flash_rt # noqa: E402 + + +class _RecordingFrontend: + last_kwargs = None + + def __init__(self, checkpoint, **kwargs): + _RecordingFrontend.last_kwargs = dict(kwargs) + self.checkpoint = checkpoint + + +@pytest.fixture +def stubbed(monkeypatch): + """Stub the frontend class and the flash_rt_fp4 extension.""" + monkeypatch.setattr(hy_mod, "HyVLATorchFrontendThor", _RecordingFrontend) + fake_fp4 = types.ModuleType("flash_rt.flash_rt_fp4") + fake_fp4.has_nvfp4 = lambda: True + monkeypatch.setitem(sys.modules, "flash_rt.flash_rt_fp4", fake_fp4) + _RecordingFrontend.last_kwargs = None + yield _RecordingFrontend + + +def test_use_fp4_reaches_the_frontend(stubbed): + flash_rt.load_model("/nonexistent/fake-ckpt", config="hyvla", + framework="torch", hardware="thor", use_fp4=True) + kw = stubbed.last_kwargs + assert kw is not None, "frontend was never constructed" + assert kw.get("use_fp4") is True, \ + f"use_fp4=True did not reach the frontend; kwargs={kw}" + + +def test_fp8_tier_selects_fused_production_config(stubbed): + flash_rt.load_model("/nonexistent/fake-ckpt", config="hyvla", + framework="torch", hardware="thor", + use_fp8=True, use_fp4=False) + kw = stubbed.last_kwargs + assert kw is not None + assert kw.get("use_fp8") is True + assert kw.get("use_fused") is True, \ + "the validated fp8 production tier must enable the fused megakernels" + + +def test_default_route_does_not_enable_fp4(stubbed): + flash_rt.load_model("/nonexistent/fake-ckpt", config="hyvla", + framework="torch", hardware="thor") + kw = stubbed.last_kwargs + assert kw is not None + assert kw.get("use_fp4") in (None, False) + + +def test_hyvla_orin_fp4_falls_back_with_warning(stubbed, caplog): + # Orin has no FP4 tensor cores: the route must degrade to the INT8 path + # instead of silently claiming FP4. + try: + flash_rt.load_model("/nonexistent/fake-ckpt", config="hyvla", + framework="torch", hardware="rtx_sm87", + use_fp4=True) + except Exception: + pytest.skip("Orin frontend not importable in this environment") + msgs = [r.message for r in caplog.records] + assert any("SM87" in m and "FP4" in m.upper() for m in msgs), \ + f"expected an SM87 FP4 fallback warning, got {msgs}" diff --git a/tests/test_hyvla_kernel_contracts.py b/tests/test_hyvla_kernel_contracts.py new file mode 100644 index 000000000..5d1a66f26 --- /dev/null +++ b/tests/test_hyvla_kernel_contracts.py @@ -0,0 +1,77 @@ +"""Host-side contract validation of the HyVLA fused-kernel pybind APIs. + +Validation throws before any CUDA work, so these tests pass dummy pointers +and assert the Python-visible ValueError. Requires the built module (CUDA +host toolchain), but not a running device for the negative cases. +""" + +import pytest + +pytest.importorskip("torch") + +try: + from flash_rt import flash_rt_kernels as fvk +except ImportError as exc: # pragma: no cover + pytest.skip(f"flash_rt_kernels is not built: {exc}", allow_module_level=True) + +if not hasattr(fvk, "hyvla_rope_qknorm_kvwrite_bf16"): + pytest.skip("HyVLA kernels require FLASHRT_ENABLE_HYVLA", allow_module_level=True) + + +def test_rope_qknorm_rejects_wrong_head_dim(): + with pytest.raises(ValueError, match="hd==128"): + fvk.hyvla_rope_qknorm_kvwrite_bf16( + 0, 0, 0, 0, 0, 0, 0, 0, S=8, nq=4, nkv=1, hd=64, S_tot=8, off=0) + + +def test_rope_qknorm_rejects_bad_offset_window(): + with pytest.raises(ValueError, match="invalid offset"): + fvk.hyvla_rope_qknorm_kvwrite_bf16( + 0, 0, 0, 0, 0, 0, 0, 0, S=8, nq=4, nkv=1, hd=128, S_tot=8, off=4) + + +def test_rope_qknorm_rejects_nonpositive_shapes(): + with pytest.raises(ValueError): + fvk.hyvla_rope_qknorm_kvwrite_bf16( + 0, 0, 0, 0, 0, 0, 0, 0, S=0, nq=4, nkv=1, hd=128, S_tot=8, off=0) + with pytest.raises(ValueError, match="eps"): + fvk.hyvla_rope_qknorm_kvwrite_bf16( + 0, 0, 0, 0, 0, 0, 0, 0, S=8, nq=4, nkv=1, hd=128, S_tot=8, off=0, + eps=0.0) + with pytest.raises(ValueError, match="kv_rep"): + fvk.hyvla_rope_qknorm_kvwrite_bf16( + 0, 0, 0, 0, 0, 0, 0, 0, S=8, nq=4, nkv=1, hd=128, S_tot=8, off=0, + kv_rep=0) + + +def test_vit_add_layer_norm_rejects_odd_dim(): + with pytest.raises(ValueError, match="even"): + fvk.hyvla_vit_add_layer_norm_bf16(0, 0, 0, 0, 0, rows=4, dim=127) + + +def test_vit_add_layer_norm_rejects_nonpositive_rows(): + with pytest.raises(ValueError): + fvk.hyvla_vit_add_layer_norm_bf16(0, 0, 0, 0, 0, rows=0, dim=128) + + +@pytest.mark.skipif(not hasattr(fvk, "hyvla_quant_fp8_dyn_bf16"), + reason="Thor-only kernel") +def test_quant_fp8_dyn_rejects_nonpositive_n(): + with pytest.raises(ValueError, match="n>0"): + fvk.hyvla_quant_fp8_dyn_bf16(0, 0, 0, 0) + + +@pytest.mark.skipif(not hasattr(fvk, "hyvla_ffn_gu_silu_bf16"), + reason="Thor-only kernel") +def test_ffn_gu_silu_rejects_misaligned_shapes(): + with pytest.raises(ValueError, match="Nout%32"): + fvk.hyvla_ffn_gu_silu_bf16(0, 0, 0, M=1, K=1024, Nout=1000, sx=0, sgu=1.0) + with pytest.raises(ValueError, match="K%16"): + fvk.hyvla_ffn_gu_silu_bf16(0, 0, 0, M=1, K=1000, Nout=1024, sx=0, sgu=1.0) + + +@pytest.mark.skipif(not hasattr(fvk, "hyvla_ffn_dn_res_bf16"), + reason="Thor-only kernel") +def test_ffn_dn_res_rejects_misaligned_shapes(): + with pytest.raises(ValueError, match="N%32"): + fvk.hyvla_ffn_dn_res_bf16(0, 0, 0, 0, M=1, K=1024, N=1000, sa=0, sdn=1.0) diff --git a/tests/test_hyvla_thor_dispatch.py b/tests/test_hyvla_thor_dispatch.py new file mode 100644 index 000000000..5ba5b5ded --- /dev/null +++ b/tests/test_hyvla_thor_dispatch.py @@ -0,0 +1,27 @@ +"""Dispatch and registration smoke tests for HyVLA (no GPU required).""" + + +def test_hyvla_thor_dispatch_resolves(): + from flash_rt.hardware import resolve_pipeline_class + + cls = resolve_pipeline_class("hyvla", "torch", "thor") + assert cls.__name__ == "HyVLATorchFrontendThor" + assert cls.__module__ == "flash_rt.frontends.torch.hyvla_thor" + + +def test_hyvla_pipeline_map_is_one_to_one(): + from flash_rt.hardware import _PIPELINE_MAP + + entries = {k: v for k, v in _PIPELINE_MAP.items() if k[0] == "hyvla"} + assert ("hyvla", "torch", "thor") in entries + classes = [v[1] for v in entries.values()] + assert len(classes) == len(set(classes)), "multiple tuples share a class" + + +def test_hyvla_is_a_supported_load_model_config(): + import inspect + + import flash_rt + + doc = inspect.getsource(flash_rt.load_model) + assert '"hyvla"' in doc diff --git a/tests/test_hyvla_thor_graphsafe.py b/tests/test_hyvla_thor_graphsafe.py new file mode 100644 index 000000000..9cd163ad2 --- /dev/null +++ b/tests/test_hyvla_thor_graphsafe.py @@ -0,0 +1,71 @@ +"""Graph-safety gate for HyVLATorchFrontendThor (requires Thor SM110 + checkpoint). + +Verifies the two invariants the CUDA-graph capture relies on: + + 1. graph == eager — replaying the captured graph produces the same action + chunk as the un-captured eager path for identical inputs. + 2. replay-stable — replaying the graph twice on the same static inputs + yields identical output (no transient-buffer aliasing). + +Inputs are fully synthetic and deterministic (seed-0 torch.rand images/state, +RandomState(0) noise), so the gate is reproducible with only the checkpoint. +Set FLASHRT_HYVLA_CHECKPOINT to the Hy-Embodied-0.5-VLA directory to run. +""" + +import os + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +CKPT = os.environ.get("FLASHRT_HYVLA_CHECKPOINT", "") +if not CKPT or not os.path.isdir(CKPT): + pytest.skip( + "set FLASHRT_HYVLA_CHECKPOINT to the Hy-Embodied-0.5-VLA checkpoint " + "directory to run this gate", allow_module_level=True) + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + +try: + from flash_rt.frontends.torch.hyvla_thor import HyVLATorchFrontendThor +except ImportError as exc: # pragma: no cover + pytest.skip(f"hyvla_thor frontend not importable: {exc}", allow_module_level=True) + + +def _inputs(fe): + """Deterministic synthetic inputs (identical across runs and machines).""" + g = torch.Generator(device="cpu").manual_seed(0) + img = torch.rand(1, 6, 3, 224, 224, generator=g) + state = torch.rand(1, fe.max_state_dim, generator=g) * 0.1 + images = torch.stack([img[0], img[0].clone(), img[0].clone()], 0) + noise = np.random.RandomState(0).randn( + 1, fe.chunk, fe.max_action_dim).astype(np.float32) + return images.numpy(), state.numpy(), noise + + +@pytest.fixture(scope="module") +def frontend(): + fe = HyVLATorchFrontendThor(CKPT, use_fp8=True, use_fused=True) + fe.set_prompt("pick up the bottle") + return fe + + +def test_graph_matches_eager(frontend): + images, state, noise = _inputs(frontend) + a_eager = frontend.predict_actions(images, state=state, noise=noise, + use_graph=False) + a_graph = frontend.predict_actions(images, state=state, noise=noise, + use_graph=True) + cos = float(np.dot(a_eager.ravel(), a_graph.ravel()) / + (np.linalg.norm(a_eager) * np.linalg.norm(a_graph) + 1e-12)) + assert cos >= 0.9999, f"graph vs eager cosine {cos} < 0.9999" + + +def test_replay_is_stable(frontend): + images, state, noise = _inputs(frontend) + a1 = frontend.predict_actions(images, state=state, noise=noise, use_graph=True) + a2 = frontend.predict_actions(images, state=state, noise=noise, use_graph=True) + assert np.array_equal(a1, a2), \ + "replayed graph output differs run-to-run on identical static inputs" From a70aa84caf28ebab4dfd69611b63196fff2b7780 Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:23:19 +0800 Subject: [PATCH 3/3] test(hyvla): fix FP4 routing stub signature from Orin production run Production validation on Orin SM87 found the routing tests failed with kwargs={'num_views': 2}: load_model feature-detects accepted kwargs via inspect.signature(pipe_cls), and the recording stub declared only **kwargs, so use_fp4/use_fp8/use_fused were never forwarded. The real routing path is unaffected (HyVLATorchFrontendThor declares all named parameters). - Stub now mirrors the real HyVLATorchFrontendThor constructor signature exactly. - Orin FP4 fallback test stubs HyVLATorchFrontendOrin as well so it validates the warning and the dropped use_fp4 without touching a checkpoint (skips only where the Orin frontend is not importable). --- tests/test_hyvla_fp4_routing.py | 40 +++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/tests/test_hyvla_fp4_routing.py b/tests/test_hyvla_fp4_routing.py index 0fd2e293b..ec268dde1 100644 --- a/tests/test_hyvla_fp4_routing.py +++ b/tests/test_hyvla_fp4_routing.py @@ -16,11 +16,27 @@ class _RecordingFrontend: + """Stub mirroring HyVLATorchFrontendThor's constructor signature. + + load_model feature-detects accepted kwargs via inspect.signature(pipe_cls), + so the named tier parameters must be declared exactly like the real + frontend or load_model will not forward them. + """ last_kwargs = None - def __init__(self, checkpoint, **kwargs): - _RecordingFrontend.last_kwargs = dict(kwargs) - self.checkpoint = checkpoint + def __init__(self, checkpoint_dir, *, hardware="thor", + use_fp8=False, use_fp8_vit=False, + use_fused=False, use_fp4=False, + use_fused_quant=False, use_autotune=False, + use_ffn_mega=False, **kwargs): + _RecordingFrontend.last_kwargs = { + "hardware": hardware, "use_fp8": use_fp8, + "use_fp8_vit": use_fp8_vit, "use_fused": use_fused, + "use_fp4": use_fp4, "use_fused_quant": use_fused_quant, + "use_autotune": use_autotune, "use_ffn_mega": use_ffn_mega, + **kwargs, + } + self.checkpoint_dir = checkpoint_dir @pytest.fixture @@ -62,15 +78,25 @@ def test_default_route_does_not_enable_fp4(stubbed): assert kw.get("use_fp4") in (None, False) -def test_hyvla_orin_fp4_falls_back_with_warning(stubbed, caplog): +def test_hyvla_orin_fp4_falls_back_with_warning(stubbed, monkeypatch, caplog): # Orin has no FP4 tensor cores: the route must degrade to the INT8 path - # instead of silently claiming FP4. + # instead of silently claiming FP4. Stub the Orin frontend as well so the + # test never touches a real checkpoint. try: + import flash_rt.frontends.torch.hyvla_orin as hy_orin + except ImportError: + pytest.skip("hyvla_orin frontend not importable in this environment") + monkeypatch.setattr(hy_orin, "HyVLATorchFrontendOrin", _RecordingFrontend) + _RecordingFrontend.last_kwargs = None + + import logging + with caplog.at_level(logging.WARNING): flash_rt.load_model("/nonexistent/fake-ckpt", config="hyvla", framework="torch", hardware="rtx_sm87", use_fp4=True) - except Exception: - pytest.skip("Orin frontend not importable in this environment") msgs = [r.message for r in caplog.records] assert any("SM87" in m and "FP4" in m.upper() for m in msgs), \ f"expected an SM87 FP4 fallback warning, got {msgs}" + # The constructed frontend must NOT have received use_fp4=True. + kw = _RecordingFrontend.last_kwargs + assert kw is None or kw.get("use_fp4") in (None, False)