diff --git a/CMakeLists.txt b/CMakeLists.txt index 32db2ee5..12736537 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1692,6 +1692,32 @@ target_include_directories(flash_rt_kernels PRIVATE ${CUTLASS_DIR}/tools/util/include ) +# ── 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 + 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() + +# 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(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) + 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 54d3c25a..93dfe435 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,137 @@ 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) { + 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), + 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) { + 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), + 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) { + 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), + 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) { + 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, + 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) { + 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), + 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 00000000..5596d5d9 --- /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 00000000..e75e2dad --- /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 00000000..b6161fdb --- /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 00000000..7898d347 --- /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 00000000..4f2a1c6a --- /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 00000000..72dd3836 --- /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 00000000..3a2752bc --- /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 00000000..497b230a --- /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_orin_sm87.md b/docs/hyvla05_orin_sm87.md new file mode 100644 index 00000000..1986ecfb --- /dev/null +++ b/docs/hyvla05_orin_sm87.md @@ -0,0 +1,285 @@ +# Hy-Embodied-0.5-VLA on Jetson Orin SM87 + +> Orin SM87 adaptation of the existing HyVLA Thor path. SM87 has no native +> FP8/FP4 tensor cores, so Thor FP8/NVFP4 kernels are not compiled or used. +> The Orin frontend keeps the validated HyVLA IO/prefix/graph path and maps the +> lower-precision GEMM slots to SM80-family INT8 W8A8 rowwise CUTLASS kernels. + +## Platform + +| Field | Value | +|---|---| +| Device | Jetson AGX Orin / SM87 | +| GPU family | Ampere, SM87 | +| Native FP8 / FP4 | No | +| Build target | `-DGPU_ARCH=87` | +| Attention | PyTorch SDPA efficient backend baseline | +| Low-bit GEMM | SM80 INT8 rowwise W8A8 (`ENABLE_SM80_INT8_CUTLASS`) | + +## Dispatch + +`flash_rt/hardware/__init__.py` registers: + +```python +("hyvla", "torch", "rtx_sm87") -> ( + "flash_rt.frontends.torch.hyvla_orin", + "HyVLATorchFrontendOrin", +) +``` + +The key is also included in `_SM87_ALLOWED`, so +`flash_rt.load_model(ckpt, config="hyvla", framework="torch")` resolves on +Orin. + +## Files + +| File | Purpose | +|---|---| +| `flash_rt/frontends/torch/hyvla_orin.py` | Orin frontend; inherits Thor tokenizer/preprocess/prefix/graph orchestration; disables Thor-only FP8/FP4 fused options; materializes INT8 per-row weights. | +| `flash_rt/models/hyvla/pipeline_orin.py` | Orin pipeline subclass; inherits BF16 math and overrides the lower-precision GEMM slot with INT8 W8A8 rowwise CUTLASS; fused ViT forward (pending-add+LN, efficient SDPA). | +| `csrc/kernels/hyvla_vit_fuse.cu` | Fused ViT residual-add + LayerNorm kernel (SM87 + SM110 builds). | +| `tests/test_orin_hyvla05_e2e_check.py` | BF16 baseline vs default-INT8 fixed-noise action cosine gate (>= 0.999) plus state/noise/prompt input boundaries. | +| `tests/test_orin_hyvla05_graphsafe.py` | Graph-vs-eager, replay-stability, and fused-vs-unfused (attention-prep, ViT add+LN) gates. | +| `tests/test_orin_hyvla05_arch_gate.py` | SM87 fail-fast hardware gate and FP4 rejection (mocked CUDA, no device needed). | + +## Precision policy + +| Component | Orin v2 precision (current default) | +|---|---| +| Embeddings / token assembly | BF16 | +| HYViT2 + merger | BF16 graph (INT8 opt-in, fails gate — see dead ends) | +| MoT VLM prefill QKV/O | BF16 (outputs feed the KV cache read by 10 denoise steps) | +| MoT VLM prefill FFN (gate/up/down) | **INT8 W8A8 by default** (`use_int8_vlm_ffn`) | +| Expert denoise QKV/O/FFN GEMMs | **INT8 W8A8 by default** (`use_int8_exp`) | +| RMSNorm / RoPE / QK-Norm | BF16 | +| Attention | BF16/SDPA | +| State/time/action head | BF16/FP32 update as in Thor path | +| FP4 / Thor FP8 fused kernels | Unsupported on SM87 | + +`HyVLATorchFrontendOrin(..., use_fp8=True)` maps to the precision-safe Orin +INT8 tier: expert denoise INT8 **plus prefill FFN INT8** (gate/up/down; QKV/O +stay BF16 to protect the KV cache). Pass `use_fp8=False, use_int8=False` for +the BF16 baseline. Diagnostic flags: `use_int8_vlm=True` (all-tower prefill +INT8) and `use_int8_vit=True` (ViT INT8) both fail the 0.999 gate — see dead +ends. + +## Measured gates + +Measured on the downloaded checkpoint at +`/path/to/checkpoint/Hy-Embodied-0.5-VLA-RoboTwin`, fixed +noise, three random 6-frame cameras, prompt `pick up the bottle`. + +| Config | Reference action cosine | MAE | Graph safety | +|---|---:|---:|---| +| BF16 baseline | 0.999911 | 2.45e-3 | PASS (`graph==eager`, replay max diff 0) | +| **BF16 + ViT fusion (Stage 4)** | **0.999931** | 1.75e-3 | PASS | +| INT8 default + Stage 3 fused kernels | 0.999531 | 6.46e-3 | PASS (`graph==eager`, replay max diff 0) | +| **INT8 + ViT fusion (Stage 4, current default)** | **0.999442** | 6.69e-3 | PASS (`graph==eager`, replay max diff 0) | +| INT8 default (expert + prefill FFN) | 0.999580 | 6.12e-3 | PASS | +| INT8 expert-only (Stage 1) | 0.999681 | 5.56e-3 | PASS | +| all-tower INT8 (`use_int8_vlm=True`) | 0.998843 | 8.77e-3 | graph-safe but below gate | +| ViT INT8 (`use_int8_vit=True`) | 0.997459 | 1.38e-2 | FAIL — dead end | + +Stage latency, sequential median of 20 iterations: + +| Config | E2E `predict_actions` | ViT+merger graph | Prefix assembly* | Prefill+denoise graph | +|---|---:|---:|---:|---:| +| BF16 baseline | 403.6 ms | 170.7 ms | 7.6 ms | 205.2 ms | +| INT8 Stage 1 (expert-only) | 384.4 ms | 173.9 ms | 7.2 ms | 187.5 ms | +| INT8 Stage 2 (expert + prefill FFN + prefix cache) | 353.6 ms | 171.4 ms | — | 167.3 ms | +| INT8 Stage 3 (Stage 2 + fused norm/rope kernels) | 293.2 ms | 170.9 ms | 7.1 ms | 107.0 ms | +| **INT8 Stage 4 (Stage 3 + ViT fusion, default)** | **277.5 ms** | 156.9 ms | 7.3 ms | 106.9 ms | +| BF16 Stage 4 (ViT fusion) | 315.2 ms | 156.5 ms | 7.8 ms | 143.8 ms | +| BF16 Stage 2 (prefix cache only) | 391.9 ms | 172.6 ms | — | 204.9 ms | + +\* stageprof's assembly line re-runs the uncached reproduction for +measurement; inside `predict_actions` the Orin frontend caches the +prompt/camera-static artifacts (segment mask, permutation, bf16-rounded RoPE +tables, suffix mask) per `(prompt, num_cam)`, which removed ~12 ms/call from +E2E in both BF16 and INT8 tiers (403.6→391.9 and 365.7→353.6). + +Stage 2 levers: prefill FFN INT8 took the main graph 187.5→167.3 ms; the +static-prefix cache removed the per-call mask/RoPE recomputation. Both +preserve the 0.999 action-cosine gate. + +Stage 3 levers (main graph 167.3→107.0 ms, E2E 353.6→293.2 ms): + +1. **Kernel-backed RMSNorm shim** — torch 2.3 has no `F.rms_norm`; the eager + Python fallback was a hot elementwise chain. Dispatching the existing + `fvk.rms_norm` (bit-equal math: fp32 sum-of-squares, single bf16 round) + removed ~38 ms from the main graph. +2. **Fused residual-add + RMSNorm** (`fvk.residual_add_rms_norm`) in the + expert denoise loop: one kernel replaces add + norm, and the fp32 + accumulation before rounding keeps drift at ≤1 ULP vs torch. +3. **Fused RoPE + QK-Norm + KV-write megakernel** — Thor's + `hyvla_rope_qknorm_kvwrite_bf16` is plain CUDA and was compiled for SM87 + (CMake gate `FLASHRT_HAVE_HYVLA_ORIN`; `csrc/bindings.cpp` needed the + HyVLA def block moved out of the `ENABLE_NVFP4` guard, which silently + excluded it on SM87). Collapses ~11 launches per attention block × 352 + blocks/step into one kernel. Auto-enabled in the Orin frontend when the + symbol is present (`use_fused=True`). + +All three preserve graph safety (max diff 0) and the 0.999 gate (cos drops +0.999580→0.999531 from the fused RoPE rounding order, still comfortably +above gate). + +Stage 4 levers (ViT 170.9→156.9 ms, E2E 293.2→277.5 ms): + +1. **Memory-efficient SDPA for ViT spatial attention** — the q/k/v slices + from the packed QKV GEMM are strided; the flash backend force-copies all + three (plus the output) contiguous, ~25 ms of `copy_` per ViT pass. The + memory-efficient backend reads the strides natively (2.89→0.62 ms per + attention segment in µbench). Same finding as Thor's efficient-SDPA lever: + accuracy *improves* vs flash (BF16 E2E cos 0.999911→0.999931). +2. **Fused residual-add + LayerNorm** — new kernel + `hyvla_vit_add_layer_norm_bf16` (`csrc/kernels/hyvla_vit_fuse.cu`): + in-place bf16-rounded residual add (bit-equal to torch add) + LayerNorm + matching this repo's `layer_norm_kernel`. The Orin `vit_forward` override + carries each block's MLP output as `pending` and fuses it into the next + block's entry LN; spacetime blocks keep the torch path (their entry LN + also adds the time positional embedding). LN deviates from torch's Welford + reduction by ≤1 bf16 ULP on ~2e-5 of elements — unbiased rounding noise, + no gate impact. + +Note: the "drop history frames after the final spacetime block" lever from +Thor is already inherited via the base pipeline (blocks 24-26 run on 3 +frames instead of 18) — the Stage 3 ViT estimate of 171 ms already included +it. + +## Dead ends (Stage 2, measured) + +| Scheme | Result | Mechanism | +|---|---|---| +| **ViT INT8** (`use_int8_vit`) | FAIL: E2E cos 0.997459; only −7 ms (166.7 vs 173.9 ms) | ViT merged-output cos only 0.999006 vs BF16 (MAE 3.4e-2); INT8 GEMM saves ~52 ms but per-GEMM activation quant (+14 ms) and extra elementwise (+28 ms) eat most of it; ViT errors propagate through 32-layer prefill + 10 denoise steps. MLP-only variant worse (merged cos 0.998638) | +| **All-tower prefill INT8** (`use_int8_vlm`) | cos 0.998843 < 0.999 | QKV INT8 error lands in the KV cache and compounds over 10 denoise steps; FFN-only variant avoids the KV path and passes | +| ViT INT8 t64x128 tile | slower than 128×128 at ViT shapes | µbench: (588,3456,1152) 0.556 vs 0.436 ms | +| **ViT elementwise fusion via existing kernels** | rejected pre-integration | `bias_gelu_bf16_strict` uses tanh-approx gelu (the reference uses exact erf; bit-mismatch max 0.0156); `bias_residual_layer_norm_bf16` LN drifts 1-2 ULP vs `F.layer_norm` on ViT shapes — both would erode the 0.999 gate | + +**Reference eager anchor (measured on this Orin, same fixed inputs, warmup 3 + +median of 5): 3137.3 ms.** FlashRT speedups vs the reference eager path: + +| Config | E2E | Speedup vs reference eager | +|---|---:|---:| +| **INT8 + ViT fusion (Stage 4, default)** | **277.5 ms** | **11.3×** | +| INT8 Stage 3 (fused norm/rope kernels) | 293.2 ms | 10.7× | +| INT8 Stage 2 (expert + prefill FFN + prefix cache) | 353.6 ms | 8.87× | +| INT8 Stage 2 before prefix cache | 365.7 ms | 8.58× | +| INT8 Stage 1 (expert-only) | 384.4 ms | 8.16× | +| BF16 Stage 2 | 391.9 ms | 8.00× | +| BF16 baseline | 403.6 ms | 7.77× | + +Cross-platform comparison (Thor numbers from `hyvla05_thor_sm110.md`): + +| Platform | Reference eager | Native BF16 (graph) | Production | Speedup | +|---|---:|---:|---:|---:| +| Thor SM110 | ~930 ms | 248.6 ms | 158.3 ms (FP8+fused+autotune) | 5.9× | +| Orin SM87 | 3137 ms | 315.2 ms | 277.5 ms (INT8 + fused norm/rope/ViT kernels) | 11.3× | + +The reference eager path is much slower on Orin than Thor (3.1 s vs 0.93 s) because the +eager PyTorch path is launch/dispatch-bound on 16 SMs and falls back to the +slow SDPA math backend; the CUDA-Graph capture collapses most of that. The +Orin absolute latency (~278 ms) is ~1.8× Thor's because SM87 lacks FP8 and +has lower bandwidth (~204 vs 243 GB/s) and compute; the remaining gap is +dominated by the ViT stage (~157 ms, 57% of E2E), whose ~100 ms of BF16 +GEMMs run near the cuBLAS compute ceiling on both platforms. + +## Roofline (Stage 4, measured anchors) + +Hardware anchors measured on this unit (warm, unlocked clocks): + +| Anchor | Value | +|---|---:| +| >>L2 read bandwidth (1 GB buffer) | **97 GB/s** | +| bf16 cuBLAS peak (4096³, warm) | 29.9 TFLOPS | +| bf16 at ViT qkv shape (3528,1152,3456) | 9.5 TFLOPS | +| bf16 at FFN shapes (3528,·,4304/1152) | 21–24 TFLOPS | +| INT8 CUTLASS peak (4096³) | 38.7 TOPS | +| INT8 at prefill FFN (240,2048,12288) | 18.9 TOPS | +| INT8 at denoise M=41 shapes | 6–8 TOPS, weight-streaming at ~71–95 GB/s | + +Stage floors vs measured (median of 20): + +| Stage | Dominant regime | Floor | Measured | Headroom | +|---|---|---:|---:|---:| +| ViT+merger (156.9 ms) | GEMM compute: ~106 ms cuBLAS bf16 (profiler) | ~150 ms | 156.9 ms | ~5% | +| Prefill graph | FFN INT8 ≈47 ms (18.9 TOPS) + QKV/O bf16 + attn + quant | ~65 ms | — | — | +| Denoise graph (×10 steps) | INT8 weight streaming: 369 MB/step ≈ 97 GB/s → ~38 ms floor + quant/attn | ~70 ms | prefill+denoise combined: 106.9 ms | — | +| Main graph total | | ~135 ms | 106.9 ms graph replay (faster than eager-shape µbench sum) | ~0–20% | +| E2E | | ~230 ms | 277.5 ms | ~17% | + +Reading: the ViT stage sits within ~5% of its measured cuBLAS+bandwidth +floor — its qkv GEMM (9.5 TFLOPS, N=3456 K=1152) is the single slowest +shape and alone costs ~71 ms of the 106 ms GEMM budget. The main graph's +denoise portion is weight-streaming-bound at M=41 (INT8 weights read at +71–95 GB/s, i.e. already at the measured memory ceiling); the prefill +portion is compute-bound in the FFN INT8 GEMMs. The ~17% E2E headroom +splits between prefix assembly (7.3 ms eager), graph-replay gaps, and +elementwise tails (GELU ~9 ms, spacetime mixes ~14 ms) that are already +near-bandwidth. Conclusion: no >1 ms framework-side lever remains; the +quantifiable floor for this model shape on SM87 is ~230–250 ms. + +## Assets + +Official model repository: see the Hy-Embodied-0.5-VLA project page on +Hugging Face. + +Checkpoint target: + +```bash +/path/to/checkpoint/Hy-Embodied-0.5-VLA-RoboTwin +``` + +Direct Hugging Face access may be reset; use `HF_ENDPOINT=https://hf-mirror.com` +when needed. + +## Build + +```bash +cmake -B build_orin_sm87 -S . \ + -DGPU_ARCH=87 \ + -DFLASHRT_ENABLE_HYVLA=ON \ + -DFA2_ARCH_NATIVE_ONLY=ON \ + -DFA2_HDIMS='128;256' \ + -DFA2_DTYPES='bf16' +cmake --build build_orin_sm87 -j4 +``` + +For SM87, CMake enables `ENABLE_SM80_INT8_CUTLASS` and builds the existing +rowwise INT8 CUTLASS kernels into `flash_rt_kernels`. + +## Verification + +Checkpoint-gated precision and graph-safety gates (skipped automatically when +`FLASHRT_HYVLA_CHECKPOINT` is unset): + +```bash +FLASHRT_HYVLA_CHECKPOINT=/path/to/Hy-Embodied-0.5-VLA-RoboTwin \ + PYTHONPATH=. python -m pytest \ + tests/test_orin_hyvla05_e2e_check.py \ + tests/test_orin_hyvla05_graphsafe.py -v +``` + +Hardware fail-fast and dispatch gates run without a device or checkpoint: + +```bash +PYTHONPATH=. python -m pytest \ + tests/test_orin_hyvla05_arch_gate.py \ + tests/test_orin_hyvla05_dispatch.py -v +``` + +## Current caveats + +- The fused RoPE/QKNorm/KV-write megakernel (plain CUDA) is compiled for SM87 + and enabled by default (`use_fused=True`); Thor-only fused FP8 quant, FFN + megakernels, and NVFP4 remain deliberately disabled. +- INT8 uses dynamic per-row activation scales and BF16 GEMM outputs; static + calibration and FP16-output epilogues are possible next steps if profiling + shows the quant kernels on the critical path. +- `predict_actions` caches prompt/camera-static prefix artifacts; the first + call per prompt pays the full mask/RoPE build (~7 ms), later calls reuse it. +- Remaining ceiling: ViT (~157 ms, 57% of E2E) is the largest stage — ~100 ms + of BF16 GEMMs near the cuBLAS compute ceiling plus near-bandwidth-bound + elementwise (GELU, spacetime ops). INT8 there fails the precision gate, so + further ViT cuts need model-side changes (fewer history frames, smaller + ViT) rather than framework-side work. +- The reference E2E oracle requires the official HyVLA repo and checkpoint assets. diff --git a/docs/hyvla05_thor_sm110.md b/docs/hyvla05_thor_sm110.md new file mode 100644 index 00000000..bcebb79b --- /dev/null +++ b/docs/hyvla05_thor_sm110.md @@ -0,0 +1,173 @@ +# 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 -B build -S . -DGPU_ARCH=110 -DFLASHRT_ENABLE_HYVLA=ON +> cmake --build build --target flash_rt_kernels # fused kernels (first build) +> 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 + +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_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 (recorded values, all vs. HF/transformers eager, same fixed noise) + +| 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 (`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) + +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) + +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) + +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 5d533a2e..bdbdf627 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` @@ -255,6 +259,8 @@ based on `use_fp8` / `use_fp16`; `rtx_sm89` resolves directly to the dedicated SM89 frontend class. Wan2.2 TI2V-5B is registered for `(config="wan22_ti2v_5b", framework="torch", arch="rtx_sm120")`. +Hy-Embodied-0.5-VLA is registered for `(config="hyvla", +framework="torch", arch in {"thor", "rtx_sm87"})`. ### `_PIPELINE_MAP` diff --git a/flash_rt/api.py b/flash_rt/api.py index e7d61380..cb09a7bd 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,26 @@ 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) ── + _hyvla_fp4 = False 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 +734,15 @@ 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 + _hyvla_fp4 = True + 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, ) @@ -766,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/configs/hyvla.yaml b/flash_rt/configs/hyvla.yaml new file mode 100644 index 00000000..1577fdf6 --- /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 c15eca03..4c0604ed 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 00000000..5228c2a0 --- /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_orin.py b/flash_rt/frontends/torch/hyvla_orin.py new file mode 100644 index 00000000..919eb0bf --- /dev/null +++ b/flash_rt/frontends/torch/hyvla_orin.py @@ -0,0 +1,293 @@ +"""HyVLA frontend for Jetson Orin SM87. + +SM87 has no native FP8/FP4 tensor cores, so the Thor FP8/NVFP4 paths are +mapped to the existing SM80-family INT8 W8A8 rowwise kernels. The public IO, +preprocessing, prefix assembly, graph cache, and weight spec are inherited from +``HyVLATorchFrontendThor``. +""" + +from __future__ import annotations + +import logging + +import numpy as np +import torch +import torch.nn.functional as F + +from flash_rt.frontends.torch.hyvla_thor import HyVLATorchFrontendThor, _BF16 +from flash_rt.models.hyvla.pipeline_orin import HyVLAOrinBF16Pipeline + +logger = logging.getLogger(__name__) + +INT8_QUANT_MAX = 127.0 +INT8_QUANT_EPS = 1e-12 + + +def _quantize_per_row_int8(w_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + w_f32 = w_bf16.float().contiguous() + scale = torch.clamp( + w_f32.abs().amax(dim=1) / INT8_QUANT_MAX, + min=INT8_QUANT_EPS, + ).to(device=w_bf16.device, dtype=torch.float32).contiguous() + q = torch.clamp( + torch.round(w_f32 / scale[:, None]), + -127, 127, + ).to(torch.int8).contiguous() + return q, scale + + +class HyVLATorchFrontendOrin(HyVLATorchFrontendThor): + _REQUIRED_CAPABILITY = (8, 7) + _ARCH_NAME = "Jetson Orin SM87" + + def __init__(self, checkpoint_dir: str, *, hardware: str = "rtx_sm87", + use_fp8: bool = True, use_fp8_vit: bool = False, + use_fused: bool = True, use_fp4: bool = False, + use_fused_quant: bool = False, use_autotune: bool = False, + use_ffn_mega: bool = False, use_int8: bool | None = None, + use_int8_vlm: bool = False, # opt-in: fails the 0.999 E2E cosine gate + use_int8_vlm_ffn: bool | None = None, # validated default tier (cosine >= 0.999) + use_int8_exp: bool = True, # validated default tier (cosine >= 0.999) + use_int8_vit: bool | None = None, # opt-in: fails the 0.999 E2E cosine gate + vit_int8_parts: tuple = ("qkv", "proj", "fc1", "fc2"), # only used by opt-in use_int8_vit + **kwargs): + if use_fp4: + raise RuntimeError( + "HyVLATorchFrontendOrin does not support FP4: SM87 has no " + "native FP4 tensor cores. Use the default INT8 W8A8 path or " + "pass use_fp8=False, use_int8=False for BF16.") + if use_fp8_vit: + logger.warning( + "HyVLA Orin ignores Thor's use_fp8_vit; ViT quantization on " + "SM87 uses the INT8 path (use_int8_vit).") + if use_fused_quant: + logger.warning("HyVLA Orin disables Thor-only fused FP8 quantization.") + if use_ffn_mega: + logger.warning("HyVLA Orin disables Thor-only FP8 FFN megakernels.") + if use_autotune: + logger.warning("HyVLA Orin INT8 path does not use Thor FP8 autotune.") + + self.use_int8 = bool(use_fp8) if use_int8 is None else bool(use_int8) + self.use_int8_vlm = bool(use_int8_vlm) + self.use_int8_exp = bool(use_int8_exp) + self.use_int8_vit = False if use_int8_vit is None else bool(use_int8_vit) + self.vit_int8_parts = tuple(vit_int8_parts) + vlm_ffn = self.use_int8 if use_int8_vlm_ffn is None else bool(use_int8_vlm_ffn) + if use_fp8 and self.use_int8: + logger.info( + "HyVLA Orin maps use_fp8=True to SM87 INT8 W8A8 rowwise GEMMs " + "(expert tower + prefill FFN by default; use_int8_vlm=True " + "quantizes prefill QKV/O too, use_int8_vit=True is opt-in and " + "fails the 0.999 E2E gate).") + + super().__init__( + checkpoint_dir, + hardware=hardware, + 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, + ) + if self.use_int8_vit: + self._quantize_vit_int8() + self.pipe = HyVLAOrinBF16Pipeline(self) + import flash_rt.flash_rt_kernels as fvk + # ViT levers: memory-efficient SDPA reads the strided QKV slices + # directly (the flash backend force-copies them), and the fused + # residual-add+LayerNorm kernel collapses the ViT elementwise pairs. + self.pipe._vit_eff_sdpa = True + self.pipe._vit_fuse_ln = hasattr(fvk, "hyvla_vit_add_layer_norm_bf16") + if use_fused: + if hasattr(fvk, "hyvla_rope_qknorm_kvwrite_bf16"): + self.pipe._fused_attn = True + else: + logger.warning( + "HyVLA Orin: fused RoPE/QKNorm/KV-write kernel not in " + "this build; falling back to the torch attention-prep path.") + if self.use_int8: + self._quantize_int8() + self._vlm_fp8_ready = self.use_int8_vlm + self._exp_fp8_ready = self.use_int8_exp + self.pipe.enable_int8() + if vlm_ffn and not self.use_int8_vlm: + self._quantize_vlm_ffn_int8() + self.pipe._vlm_ffn_int8 = True + + def _quantize_int8(self): + def q_list(src): + qs, ss = [], [] + for w in src: + q, s = _quantize_per_row_int8(w) + qs.append(q) + ss.append(s) + return qs, ss + + if self.use_int8_exp: + 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 + else: + self._exp_fp8_ready = False + + if self.use_int8_vlm: + 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 + else: + self._vlm_fp8_ready = False + torch.cuda.synchronize() + + def _quantize_vlm_ffn_int8(self): + """INT8-quantize only the VLM prefill FFN weights (gate/up + down, + vision and text branches). QKV/O stay BF16 so the KV cache written by + prefill — and read by all 10 denoise steps — keeps full precision.""" + def q_list(src): + qs, ss = [], [] + for w in src: + q, s = _quantize_per_row_int8(w) + qs.append(q) + ss.append(s) + return qs, ss + + 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_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) + torch.cuda.synchronize() + logger.info( + "HyVLA Orin: VLM prefill FFN (gate/up/down) quantized to INT8; " + "QKV/O remain BF16.") + + def _quantize_vit_int8(self): + """Replace the 27-block ViT GEMM weights (qkv/proj/fc1/fc2) in-place + with per-row INT8 + FP32 scale wrappers consumed by + ``HyVLAOrinBF16Pipeline._vit_*``. LayerNorm weights/biases and the + patch-embed conv stay BF16.""" + from flash_rt.models.hyvla.pipeline_orin import _W8Int8 + + part_attrs = { + "qkv": "_vit_qkv_w", + "proj": "_vit_proj_w", + "fc1": "_vit_fc1_w", + "fc2": "_vit_fc2_w", + } + for part in self.vit_int8_parts: + if part not in part_attrs: + raise ValueError( + f"unknown vit_int8_parts entry {part!r}; " + f"valid: {sorted(part_attrs)}") + attr = part_attrs[part] + src = getattr(self, attr) + wrapped = [] + for w in src: + q, s = _quantize_per_row_int8(w) + wrapped.append(_W8Int8(q, s)) + setattr(self, attr, wrapped) + del src + torch.cuda.synchronize() + logger.info( + "HyVLA Orin: ViT INT8 W8A8 enabled for parts %s.", + list(self.vit_int8_parts)) + + + @torch.no_grad() + def predict_actions(self, images, prompt=None, state=None, noise=None, + use_graph=True): + """Orin variant of the Thor ``predict_actions`` with a static-prefix + cache: everything that depends only on (prompt, num_cam) — the segment + mask, the [vision|text] permutation, the bf16-rounded RoPE tables + (fp64 math), and the suffix mask — is computed once per prompt and + reused across frames. Only the image-dependent ``prefix_embs`` are + rebuilt each call.""" + if prompt is not None and prompt != self._prompt: + self.set_prompt(prompt) + if self._lang_tokens is None: + raise RuntimeError("call set_prompt() before predict_actions()") + 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) + merged = self._vit_merge(imgs5, use_graph=use_graph) + + 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: + 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 + + key = (self._prompt, merged.shape[0]) + stat = getattr(self, "_static_prefix_cache", None) + if stat is None or stat["key"] != key: + (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()) + att2d = att2d[:, perm][:, :, perm] + prefix_pos = prefix_pos[:, perm] + pcos, psin = self._rope_cos_sin(prefix_pos) + 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) + stat = {"key": key, "perm": perm, "n_vis": n_vis, "S_p": S_p, + "pmask": att2d[:, None], "pcos": pcos, "psin": psin, + "smask": smask, "scos": scos, "ssin": ssin} + self._static_prefix_cache = stat + + (prefix_embs, _pad_masks, _att_masks, _mm_prefix, + _idx_ranges, _full_ranges) = self._assemble_prefix(merged) + prefix_embs = prefix_embs[:, stat["perm"]] + + 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) + 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(stat["S_p"], stat["n_vis"], prefix_embs, + stat["pmask"], stat["pcos"], stat["psin"], + stat["smask"], stat["scos"], stat["ssin"], + state_t, noise_t, use_graph=use_graph) + return x_t.float().cpu().numpy() + + +__all__ = ["HyVLATorchFrontendOrin"] diff --git a/flash_rt/frontends/torch/hyvla_thor.py b/flash_rt/frontends/torch/hyvla_thor.py new file mode 100644 index 00000000..d27a6f9e --- /dev/null +++ b/flash_rt/frontends/torch/hyvla_thor.py @@ -0,0 +1,670 @@ +"""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: + #: 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, + 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._require_arch() + 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 + 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) + 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 + # Standard tokenizer path only — never execute checkpoint-provided + # Python code (no trust_remote_code). + self._tokenizer = AutoTokenizer.from_pretrained( + self.checkpoint_dir) + 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) + if self._lang_tokens is None: + raise RuntimeError("call set_prompt() before predict_actions()") + 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: + 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 + + (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) + 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, + 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 5d231ead..13e81b60 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -102,6 +102,12 @@ 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"), + ("hyvla", "torch", "rtx_sm87"): + ("flash_rt.frontends.torch.hyvla_orin", "HyVLATorchFrontendOrin"), + # ── GROOT N1.6 ── ("groot", "torch", "thor"): ("flash_rt.frontends.torch.groot_thor", "GrootTorchFrontendThor"), @@ -196,6 +202,7 @@ def detect_arch() -> str: _SM87_ALLOWED = { ("pi05", "torch", "rtx_sm87"), ("qwen3_vl", "torch", "rtx_sm87"), + ("hyvla", "torch", "rtx_sm87"), } diff --git a/flash_rt/models/hyvla/__init__.py b/flash_rt/models/hyvla/__init__.py new file mode 100644 index 00000000..8aa927ab --- /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_orin.py b/flash_rt/models/hyvla/pipeline_orin.py new file mode 100644 index 00000000..13a6002f --- /dev/null +++ b/flash_rt/models/hyvla/pipeline_orin.py @@ -0,0 +1,434 @@ +"""HyVLA forward path for Jetson Orin SM87. + +The BF16 math is inherited from the Thor correctness path. When enabled, the +GEMM slots that Thor names ``fp8`` are backed by SM87 INT8 W8A8 rowwise kernels. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +try: + from torch.nn.attention import sdpa_kernel, SDPBackend +except ImportError: # pragma: no cover - torch < 2.2 + sdpa_kernel = None + SDPBackend = None + +import flash_rt.flash_rt_kernels as fvk +import flash_rt.models.hyvla.pipeline_thor as _thor_pipeline +from flash_rt.models.hyvla.pipeline_thor import HyVLAThorBF16Pipeline, _rot_half + + +if not hasattr(_thor_pipeline.F, "rms_norm"): + def _rms_norm_torch(x, weight, eps): + y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + if weight is not None: + y = y * weight.float() + return y.to(x.dtype) + + def _rms_norm(x, normalized_shape, weight=None, eps=1e-5): + # torch 2.3 lacks F.rms_norm. The single-launch fvk.rms_norm kernel + # implements the exact reference math (fp32 sum-of-squares -> rsqrt -> + # weight multiply -> bf16 rounding) and is bit-equal to the torch + # expansion on contiguous rows; fall back for anything exotic. + if (x.is_cuda and x.dtype == torch.bfloat16 and weight is not None + and weight.dtype == torch.bfloat16 + and len(normalized_shape) == 1 + and normalized_shape[0] == x.shape[-1]): + xc = x if x.is_contiguous() else x.contiguous() + wc = weight if weight.is_contiguous() else weight.contiguous() + rows = xc.numel() // xc.shape[-1] + out = torch.empty_like(xc) + fvk.rms_norm(xc.data_ptr(), wc.data_ptr(), out.data_ptr(), + rows, xc.shape[-1], eps, + torch.cuda.current_stream().cuda_stream) + return out.reshape(x.shape) + return _rms_norm_torch(x, weight, eps) + + _thor_pipeline.F.rms_norm = _rms_norm + + +class _W8Int8: + """Per-output-row INT8 weight + FP32 scale pair for the ViT INT8 path.""" + + __slots__ = ("w", "s") + + def __init__(self, w, s): + self.w = w + self.s = s + + +def _vit_int8_linear(x, w8, bias): + """x (..., K) bf16 @ W8_int8.T -> (..., N) bf16 + bias via SM87 rowwise INT8.""" + wq, ws = w8.w, w8.s + N, K = wq.shape + orig = x.shape + xc = x.reshape(-1, K).contiguous() + M = xc.shape[0] + st = torch.cuda.current_stream().cuda_stream + a8 = torch.empty(M, K, dtype=torch.int8, device=x.device) + act_scale = torch.empty(M, dtype=torch.float32, device=x.device) + fvk.quantize_int8_rowwise( + xc.data_ptr(), a8.data_ptr(), act_scale.data_ptr(), M, K, st) + out = torch.empty(M, N, dtype=torch.bfloat16, device=x.device) + status = fvk.cutlass_int8_rowwise_bf16out( + a8.data_ptr(), wq.data_ptr(), act_scale.data_ptr(), ws.data_ptr(), + out.data_ptr(), M, N, K, st) + if status != 0: + raise RuntimeError( + f"cutlass_int8_rowwise_bf16out failed: status={status} " + f"shape=({M},{N},{K})") + out = out + bias + return out.reshape(*orig[:-1], N) + + +class HyVLAOrinBF16Pipeline(HyVLAThorBF16Pipeline): + def enable_fp8(self): + raise RuntimeError( + "HyVLA Orin does not support Thor FP8 GEMMs; use enable_int8().") + + def enable_fp4(self): + raise RuntimeError( + "HyVLA Orin does not support FP4 because SM87 has no native FP4 tensor cores.") + + def enable_int8(self): + self._fp8 = True + self._orin_int8 = True + self.gemm = None + + def autotune_gemms(self, shapes, num_algos=16): + return 0 + + def _attn(self, q, k, v, mask): + out = super()._attn(q, k, v, mask) + return torch.nan_to_num(out) + + # ------------------------------------------------------------------ + # ViT INT8 GEMM sites (SM87). Weight lists are replaced in-place by + # the frontend with _W8Int8 wrappers when use_int8_vit is enabled; + # otherwise the parent BF16 F.linear path is kept. + # ------------------------------------------------------------------ + def _vit_qkv(self, h): + if isinstance(self._vit_qkv_w_cur, _W8Int8): + qkv = _vit_int8_linear(h, self._vit_qkv_w_cur, self._vit_qkv_b_cur) + bk, N, _ = h.shape + 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] + return super()._vit_qkv(h) + + def _vit_spatial_attn(self, q, k, v): + bk, _, N, _ = q.shape + if getattr(self, "_vit_eff_sdpa", False): + # The q/k/v slices out of the packed QKV GEMM are strided; the + # flash backend force-copies them contiguous (4 copy_ per block, + # ~25% of ViT time). The memory-efficient backend accepts the + # strides directly and reads the strided layout natively. + if sdpa_kernel is not None: + ctx = sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION) + else: + ctx = torch.backends.cuda.sdp_kernel( + enable_flash=False, enable_math=False, enable_mem_efficient=True) + with ctx: + out = F.scaled_dot_product_attention(q, k, v, scale=self.vit_scale) + else: + out = F.scaled_dot_product_attention(q, k, v, scale=self.vit_scale) + out = out.transpose(1, 2).reshape(bk, N, -1) + if isinstance(self._vit_proj_w_cur, _W8Int8): + return _vit_int8_linear(out, self._vit_proj_w_cur, self._vit_proj_b_cur) + return F.linear(out, self._vit_proj_w_cur, self._vit_proj_b_cur) + + def _vit_mlp(self, x): + if isinstance(self._vit_fc1_w_cur, _W8Int8): + x = _vit_int8_linear(x, self._vit_fc1_w_cur, self._vit_fc1_b_cur) + x = F.gelu(x) + return _vit_int8_linear(x, self._vit_fc2_w_cur, self._vit_fc2_b_cur) + return super()._vit_mlp(x) + + # ------------------------------------------------------------------ + # Fused ViT forward: residual-add + LayerNorm pairs collapse into one + # hyvla_vit_add_layer_norm_bf16 launch. The previous block's MLP + # output is carried as ``pending`` and fused into the next block's + # entry LN; the post-attention add is fused into the pre-MLP LN. + # Spacetime blocks keep the torch path (their entry LN also adds the + # time positional embedding). + # ------------------------------------------------------------------ + def _vit_add_ln(self, x, add, lnw, lnb): + bk, n, d = x.shape + out = torch.empty_like(x) + fvk.hyvla_vit_add_layer_norm_bf16( + x.data_ptr(), add.data_ptr(), lnw.data_ptr(), lnb.data_ptr(), + out.data_ptr(), bk * n, d, self.vit_eps, + torch.cuda.current_stream().cuda_stream) + return out + + def _vit_ln(self, x, lnw, lnb): + bk, n, d = x.shape + out = torch.empty_like(x) + fvk.layer_norm(x.data_ptr(), lnw.data_ptr(), lnb.data_ptr(), + out.data_ptr(), bk * n, d, self.vit_eps, + torch.cuda.current_stream().cuda_stream) + return out + + @torch.no_grad() + def vit_forward(self, imgs): + if not getattr(self, "_vit_fuse_ln", False): + return super().vit_forward(imgs) + 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) + hh = ww = x.shape[-1] + n = hh * ww + d = x.shape[1] + x = x.flatten(2).transpose(1, 2) + x = (x + self._vit_pos_embed_rescale(hh, ww, x.dtype)).contiguous() + last_st = max(self.vit_spacetime_ids) if K > 1 else -1 + sliced = False + pending = None + for li in range(27): + Wcur = self.W + self._vit_qkv_w_cur = Wcur._vit_qkv_w[li]; self._vit_qkv_b_cur = Wcur._vit_qkv_b[li] + self._vit_proj_w_cur = Wcur._vit_proj_w[li]; self._vit_proj_b_cur = Wcur._vit_proj_b[li] + self._vit_fc1_w_cur = Wcur._vit_fc1_w[li]; self._vit_fc1_b_cur = Wcur._vit_fc1_b[li] + self._vit_fc2_w_cur = Wcur._vit_fc2_w[li]; self._vit_fc2_b_cur = Wcur._vit_fc2_b[li] + self._vit_f8 = self._fp8 and getattr(Wcur, "_vit_fp8_ready", False) + if self._vit_f8: + self._vit_qkv_w8c = Wcur._vit_qkv_w8[li]; self._vit_qkv_wsc = Wcur._vit_qkv_ws[li] + self._vit_proj_w8c = Wcur._vit_proj_w8[li]; self._vit_proj_wsc = Wcur._vit_proj_ws[li] + self._vit_fc1_w8c = Wcur._vit_fc1_w8[li]; self._vit_fc1_wsc = Wcur._vit_fc1_ws[li] + self._vit_fc2_w8c = Wcur._vit_fc2_w8[li]; self._vit_fc2_wsc = Wcur._vit_fc2_ws[li] + ln1w, ln1b = Wcur._vit_ln1_w[li], Wcur._vit_ln1_b[li] + ln2w, ln2b = Wcur._vit_ln2_w[li], Wcur._vit_ln2_b[li] + + if li in self.vit_spacetime_ids and K > 1: + if pending is not None: + x = x + pending + pending = None + b, kf = bk // K, K + 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: + if pending is not None: + h = self._vit_add_ln(x, pending, ln1w, ln1b) + else: + h = self._vit_ln(x, ln1w, ln1b) + q, k, v = self._vit_qkv(h) + attn_out = self._vit_spatial_attn(q, k, v) + + h2 = self._vit_add_ln(x, attn_out, ln2w, ln2b) + pending = self._vit_mlp(h2) + if li == last_st and li < 26: + x = (x + pending).contiguous() + pending = None + x = x.view(num_cam, K, n, d)[:, -1].contiguous() + sliced = True + if pending is not None: + x = x + pending + if not sliced: + x = x.view(num_cam, K, n, d)[:, -1] + return x + + def _int8_rowwise_gemm(self, x, wq, ws): + """x (..., K) bf16 -> (M, N) bf16 via dynamic per-row INT8 W8A8.""" + N, K = wq.shape + xc = x.reshape(-1, K).contiguous() + M = xc.shape[0] + st = torch.cuda.current_stream().cuda_stream + a8 = torch.empty(M, K, dtype=torch.int8, device=x.device) + act_scale = torch.empty(M, dtype=torch.float32, device=x.device) + fvk.quantize_int8_rowwise( + xc.data_ptr(), a8.data_ptr(), act_scale.data_ptr(), M, K, st) + out = torch.empty(M, N, dtype=torch.bfloat16, device=x.device) + status = fvk.cutlass_int8_rowwise_bf16out( + a8.data_ptr(), wq.data_ptr(), act_scale.data_ptr(), ws.data_ptr(), + out.data_ptr(), M, N, K, st) + if status != 0: + raise RuntimeError( + f"cutlass_int8_rowwise_bf16out failed: status={status} " + f"shape=({M},{N},{K})") + return out + + def _fp8_gemm(self, x, w8, ws): + if not getattr(self, "_orin_int8", False): + raise RuntimeError("HyVLA Orin lower-precision GEMM requires enable_int8().") + return self._int8_rowwise_gemm(x, w8, ws) + + # ------------------------------------------------------------------ + # Prefill FFN-only INT8: QKV / O stay BF16 (their outputs feed the KV + # cache read by 10 denoise steps), gate/up/down run INT8 W8A8. + # ------------------------------------------------------------------ + def _block_ffn8(self, hs, n_vis, w_text, w_vis, qk_w, mask, cos, sin, + kbuf, vbuf, off, ffnv, ffnt): + S = hs.shape[1] + D = hs.shape[2] + hd, nh, nkv = self.head_dim, self.n_heads, self.n_kv + + 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) + qkv = torch.cat([hs_v @ w_vis[0].t(), hs_t @ w_text[0].t()], 0) + + if getattr(self, "_fused_attn", False): + q = torch.empty(1, nh, S, hd, dtype=torch.bfloat16, device=hs.device) + self._rope_qknorm_kvwrite(qkv, cos, sin, qk_w, q, kbuf, vbuf, S, off) + 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] + 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: + 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) + att = self._attn(q, kbuf[:, :, : off + S], vbuf[:, :, : off + S], mask) + att = att.transpose(1, 2).reshape(1, S, self.q_dim) + + 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) + gu = torch.cat([self._int8_rowwise_gemm(hs_v, ffnv[0], ffnv[1]), + self._int8_rowwise_gemm(hs_t, ffnt[0], ffnt[1])], 0) + g, u = gu.chunk(2, -1) + act = F.silu(g) * u + dn = torch.cat([self._int8_rowwise_gemm(act[:n_vis], ffnv[2], ffnv[3]), + self._int8_rowwise_gemm(act[n_vis:], ffnt[2], ffnt[3])], 0) + return hs + dn[None] + + @torch.no_grad() + def prefill(self, prefix_embs, n_vis, pmask, pcos, psin, kbuf, vbuf): + if getattr(self, "_vlm_ffn_int8", False): + W = self.W + hs = prefix_embs + for li in range(32): + text, vis = self._vlm_w(li) + qk = (W._qk_norm_q[li], W._qk_norm_k[li]) + ffnv = (W._vlm_gu_v8[li], W._vlm_gu_v_ws[li], + W._vlm_d_v8[li], W._vlm_d_v_ws[li]) + ffnt = (W._vlm_gu_t8[li], W._vlm_gu_t_ws[li], + W._vlm_d_t8[li], W._vlm_d_t_ws[li]) + hs = self._block_ffn8(hs, n_vis, text, vis, qk, pmask, + pcos, psin, kbuf[li], vbuf[li], 0, + ffnv, ffnt) + return hs + return super().prefill(prefix_embs, n_vis, pmask, pcos, psin, kbuf, vbuf) + + # ------------------------------------------------------------------ + # Expert denoise with fused residual-add + RMSNorm. The down-proj + # output of layer N is carried as ``pending`` and fused into layer + # N+1's input norm (and into the final norm after layer 31), turning + # (add + norm) pairs into single launches across all 32x10 blocks. + # ------------------------------------------------------------------ + def _res_add_rms_norm(self, residual, x, weight): + rows = residual.shape[-2] + dim = residual.shape[-1] + out = torch.empty_like(residual) + fvk.residual_add_rms_norm( + residual.data_ptr(), x.data_ptr(), weight.data_ptr(), + out.data_ptr(), rows, dim, self.rms_eps, + torch.cuda.current_stream().cuda_stream) + return out + + def _rope_qknorm_kvwrite(self, qkv, cos, sin, qk_w, q, kbuf, vbuf, S, off): + """Fused RoPE(q,k)+QK-Norm+KV-write (1 launch). ``qkv`` must be a + contiguous (S, (nq+2*nkv)*hd) bf16 tensor; writes ``q`` and the + GQA-pre-expanded KV cache rows at ``off``.""" + hd = self.head_dim + S_tot = kbuf.shape[2] + kv_rep = kbuf.shape[1] // self.n_kv + fvk.hyvla_rope_qknorm_kvwrite_bf16( + qkv.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, self.n_heads, self.n_kv, hd, S_tot, off, self.rms_eps, + kv_rep, torch.cuda.current_stream().cuda_stream) + + def _exp_block_r(self, hs, w, qk_w, mask, cos, sin, kbuf, vbuf, off, + fp8w, pending): + S = hs.shape[1] + D = hs.shape[2] + hd, nh, nkv = self.head_dim, self.n_heads, self.n_kv + + if pending is None: + hs_n = F.rms_norm(hs, (D,), w[4], self.rms_eps) + else: + hs_n = self._res_add_rms_norm(hs, pending, w[4]) + + qkv = self._int8_rowwise_gemm(hs_n[0], fp8w[0], fp8w[1]) + if getattr(self, "_fused_attn", False): + q = torch.empty(1, nh, S, hd, dtype=torch.bfloat16, device=hs.device) + self._rope_qknorm_kvwrite(qkv, cos, sin, qk_w, q, kbuf, vbuf, S, off) + 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] + 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: + 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) + att = self._attn(q, kbuf[:, :, : off + S], vbuf[:, :, : off + S], mask) + att = att.transpose(1, 2).reshape(1, S, self.q_dim) + + o = self._int8_rowwise_gemm(att[0], fp8w[2], fp8w[3]) + hs_n2 = self._res_add_rms_norm(hs, o[None], w[5]) + gu = self._int8_rowwise_gemm(hs_n2[0], fp8w[4], fp8w[5]) + g, u = gu.chunk(2, -1) + act = F.silu(g) * u + dn = self._int8_rowwise_gemm(act, fp8w[6], fp8w[7]) + return hs, dn[None] + + @torch.no_grad() + def denoise(self, state, x_t, time_embs, smask, scos, ssin, + kbuf, vbuf, S_p, num_steps=10): + W = self.W + if not (getattr(self, "_orin_int8", False) + and getattr(W, "_exp_fp8_ready", False)): + return super().denoise(state, x_t, time_embs, smask, scos, ssin, + kbuf, vbuf, S_p, num_steps=num_steps) + 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) + pending = None + for li in range(32): + exp = self._exp_w(li) + qk = (W._qk_norm_q[li], W._qk_norm_k[li]) + hs, pending = self._exp_block_r( + hs, exp, qk, smask, scos, ssin, + kbuf[li], vbuf[li], S_p, self._exp_w_fp8(li), pending) + hs_n = self._res_add_rms_norm(hs, pending, W._exp_final_norm_w) + v_t = F.linear(hs_n[:, -x_t.shape[1]:], W._aout_w, W._aout_b) + x_t.add_(dt * v_t.to(x_t.dtype)) + return x_t + + +__all__ = ["HyVLAOrinBF16Pipeline"] diff --git a/flash_rt/models/hyvla/pipeline_thor.py b/flash_rt/models/hyvla/pipeline_thor.py new file mode 100644 index 00000000..b3b47e2a --- /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"] diff --git a/tests/test_hyvla_arch_gate.py b/tests/test_hyvla_arch_gate.py new file mode 100644 index 00000000..12647842 --- /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 00000000..ec268dde --- /dev/null +++ b/tests/test_hyvla_fp4_routing.py @@ -0,0 +1,102 @@ +"""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: + """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_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 +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, monkeypatch, caplog): + # Orin has no FP4 tensor cores: the route must degrade to the INT8 path + # 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) + 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) diff --git a/tests/test_hyvla_kernel_contracts.py b/tests/test_hyvla_kernel_contracts.py new file mode 100644 index 00000000..5d1a66f2 --- /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 00000000..5ba5b5de --- /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 00000000..9cd163ad --- /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" diff --git a/tests/test_orin_hyvla05_arch_gate.py b/tests/test_orin_hyvla05_arch_gate.py new file mode 100644 index 00000000..405a8f9f --- /dev/null +++ b/tests/test_orin_hyvla05_arch_gate.py @@ -0,0 +1,64 @@ +"""HyVLA Orin hardware-gate (fail-fast) tests — torch.cuda is mocked, no GPU needed.""" + +import pytest + +torch = pytest.importorskip("torch") + +try: + import flash_rt.frontends.torch.hyvla_orin as orin_mod +except ImportError as exc: # pragma: no cover + pytest.skip(f"hyvla_orin frontend not importable: {exc}", allow_module_level=True) + + +class _Probe: + """Run _require_arch against mocked CUDA state.""" + + _cls = orin_mod.HyVLATorchFrontendOrin + + 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().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: (11, 0)) + with pytest.raises(RuntimeError, match="requires Jetson Orin SM87"): + _Probe().run() + + +def test_accepts_sm87(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, 7)) + _Probe().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().run() + + +def test_fp4_rejected_before_any_cuda_work(): + # SM87 has no FP4 tensor cores; the constructor must fail fast, + # before checkpoint loading or CUDA allocation. + with pytest.raises(RuntimeError, match="does not support FP4"): + orin_mod.HyVLATorchFrontendOrin("/nonexistent/fake-ckpt", use_fp4=True) + + +def test_missing_prompt_raises_runtime_error(): + fe = orin_mod.HyVLATorchFrontendOrin.__new__(orin_mod.HyVLATorchFrontendOrin) + # Minimal state to reach the prompt contract check only. + fe._prompt = None + fe._lang_tokens = None + with pytest.raises(RuntimeError, match="set_prompt"): + orin_mod.HyVLATorchFrontendOrin.predict_actions(fe, images=None) diff --git a/tests/test_orin_hyvla05_dispatch.py b/tests/test_orin_hyvla05_dispatch.py new file mode 100644 index 00000000..ac46e131 --- /dev/null +++ b/tests/test_orin_hyvla05_dispatch.py @@ -0,0 +1,14 @@ +"""Dispatch smoke for HyVLA on Jetson Orin SM87.""" + +import pytest + +pytest.importorskip("numpy") +pytest.importorskip("torch") + + +def test_hyvla_orin_dispatch_resolves(): + from flash_rt.hardware import resolve_pipeline_class + + cls = resolve_pipeline_class("hyvla", "torch", "rtx_sm87") + assert cls.__module__ == "flash_rt.frontends.torch.hyvla_orin" + assert cls.__name__ == "HyVLATorchFrontendOrin" diff --git a/tests/test_orin_hyvla05_e2e_check.py b/tests/test_orin_hyvla05_e2e_check.py new file mode 100644 index 00000000..e4af111e --- /dev/null +++ b/tests/test_orin_hyvla05_e2e_check.py @@ -0,0 +1,111 @@ +"""Fixed-noise precision gate for HyVLATorchFrontendOrin (requires SM87 + checkpoint). + +Loads the BF16 baseline first, captures the fixed-noise reference action, +frees it, then loads the default INT8 W8A8 tier and asserts action cosine +>= 0.999 against the reference — the documented Orin precision gate. Also +covers the public input boundaries inherited by the Orin override. + +Set FLASHRT_HYVLA_CHECKPOINT to the Hy-Embodied-0.5-VLA directory to run. +""" + +import os + +import pytest + +torch = pytest.importorskip("torch") +np = pytest.importorskip("numpy") + +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_orin import HyVLATorchFrontendOrin +except ImportError as exc: # pragma: no cover + pytest.skip(f"hyvla_orin frontend not importable: {exc}", allow_module_level=True) + +PROMPT = "pick up the bottle" + + +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 + + +def _cos(a, b): + a = a.ravel().astype(np.float64) + b = b.ravel().astype(np.float64) + return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)) + + +@pytest.fixture(scope="module") +def ref_and_int8(): + # Sequential loads: Orin unified memory should not hold two 9 GB weight + # copies at once. Capture the BF16 reference, free it, then load INT8. + fe_bf16 = HyVLATorchFrontendOrin(CKPT, use_fp8=False, use_int8=False) + fe_bf16.set_prompt(PROMPT) + images, state, noise = _inputs(fe_bf16) + a_bf16 = fe_bf16.predict_actions(images, state=state, noise=noise, + use_graph=False) + del fe_bf16 + torch.cuda.empty_cache() + + fe_int8 = HyVLATorchFrontendOrin(CKPT) # default INT8 W8A8 tier + fe_int8.set_prompt(PROMPT) + return a_bf16, fe_int8 + + +def test_default_int8_tier_enables_validated_fusion(ref_and_int8): + _, fe_int8 = ref_and_int8 + pipe = fe_int8.pipe + assert pipe._fused_attn, \ + "default tier must use the fused RoPE/QKNorm/KV-write kernel" + assert pipe._vit_fuse_ln or pipe._vit_eff_sdpa, \ + "default tier must enable at least one ViT fusion lever" + assert pipe._vlm_ffn_int8, "default tier must INT8-quantize the prefill FFN" + + +def test_int8_vs_bf16_fixed_noise_cosine(ref_and_int8): + a_bf16, fe_int8 = ref_and_int8 + images, state, noise = _inputs(fe_int8) + a_int8 = fe_int8.predict_actions(images, state=state, noise=noise, + use_graph=False) + assert np.isfinite(a_int8).all() + assert a_int8.shape == (1, fe_int8.chunk, fe_int8.max_action_dim) + cos = _cos(a_int8, a_bf16) + assert cos >= 0.999, f"INT8 vs BF16 cosine {cos:.6f} < 0.999 gate" + + +def test_eager_is_deterministic_with_fixed_noise(ref_and_int8): + _, fe_int8 = ref_and_int8 + images, state, noise = _inputs(fe_int8) + a1 = fe_int8.predict_actions(images, state=state, noise=noise, use_graph=False) + a2 = fe_int8.predict_actions(images, state=state, noise=noise, use_graph=False) + assert np.array_equal(a1, a2), "eager path must be deterministic for fixed noise" + + +def test_oversized_state_rejected(ref_and_int8): + _, fe_int8 = ref_and_int8 + images, _, noise = _inputs(fe_int8) + big_state = np.zeros((1, fe_int8.max_state_dim + 1), dtype=np.float32) + with pytest.raises(ValueError, match="max_state_dim"): + fe_int8.predict_actions(images, state=big_state, noise=noise, use_graph=False) + + +def test_wrong_noise_size_rejected(ref_and_int8): + _, fe_int8 = ref_and_int8 + images, state, _ = _inputs(fe_int8) + bad_noise = np.zeros((1, 7), dtype=np.float32) + with pytest.raises(ValueError, match="noise must have"): + fe_int8.predict_actions(images, state=state, noise=bad_noise, use_graph=False) diff --git a/tests/test_orin_hyvla05_graphsafe.py b/tests/test_orin_hyvla05_graphsafe.py new file mode 100644 index 00000000..b526b82a --- /dev/null +++ b/tests/test_orin_hyvla05_graphsafe.py @@ -0,0 +1,111 @@ +"""Graph-safety gate for HyVLATorchFrontendOrin (requires SM87 + checkpoint). + +Verifies the 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). + 3. fused == unfused — the fused RoPE/QKNorm/KV-write attention-prep and + the ViT fused add+LayerNorm produce the same actions (within FP noise) + as the torch fallback paths. + +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 pytest + +torch = pytest.importorskip("torch") +np = pytest.importorskip("numpy") + +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_orin import HyVLATorchFrontendOrin +except ImportError as exc: # pragma: no cover + pytest.skip(f"hyvla_orin 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 + + +def _cos(a, b): + a = a.ravel().astype(np.float64) + b = b.ravel().astype(np.float64) + return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)) + + +@pytest.fixture(scope="module") +def frontend(): + fe = HyVLATorchFrontendOrin(CKPT) # default INT8 W8A8 tier + 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 = _cos(a_eager, a_graph) + 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" + + +def test_fused_matches_unfused_attention_prep(frontend): + images, state, noise = _inputs(frontend) + if not frontend.pipe._fused_attn: + pytest.skip("fused RoPE/QKNorm/KV-write kernel not in this build") + a_fused = frontend.predict_actions(images, state=state, noise=noise, + use_graph=False) + frontend.pipe._fused_attn = False + try: + a_unfused = frontend.predict_actions(images, state=state, noise=noise, + use_graph=False) + finally: + frontend.pipe._fused_attn = True + cos = _cos(a_fused, a_unfused) + assert cos >= 0.999, f"fused vs unfused attention-prep cosine {cos} < 0.999" + + +def test_fused_matches_unfused_vit_layer_norm(frontend): + images, state, noise = _inputs(frontend) + if not frontend.pipe._vit_fuse_ln: + pytest.skip("hyvla_vit_add_layer_norm_bf16 not in this build") + a_fused = frontend.predict_actions(images, state=state, noise=noise, + use_graph=False) + frontend.pipe._vit_fuse_ln = False + try: + a_unfused = frontend.predict_actions(images, state=state, noise=noise, + use_graph=False) + finally: + frontend.pipe._vit_fuse_ln = True + cos = _cos(a_fused, a_unfused) + assert cos >= 0.999, f"fused vs unfused ViT add+LN cosine {cos} < 0.999"