CUDA: bit-exact fused kernels for BF16 and FoldQuant, 23-40% faster on Jetson AGX Orin by quangnd-vr · Pull Request #48 · VinRobotics/vla.cpp

quangnd-vr · 2026-10-22T20:25:13Z

What

A faster CUDA path that returns the same bits as the plain ggml graph, for every arch except ACT. It runs at the default precision (F32 activations, the checkpoint's weights, no flash attention) and needs no flag.

Fused nodes (src/layers/fused_ops.h, src/cuda/vla_cuda_fused.cu): GGML_OP_CUSTOM nodes that replace chains of ggml ops with one kernel each. They cover LayerNorm / RMSNorm with affine or modulation, bias + activation, bias + residual (optionally gated), the Q/K/V split with RoPE into the attention layout, and the SiLU / GELU gates. They round where ggml rounds (__fmul_rn / __fadd_rn) and keep ggml-cuda's reduction trees.
gemm_exact: a BF16-weight GEMM runs on the fastest cuBLASLt algorithm, or gemm_skinny kernel, that reproduces ggml's own cublasGemmEx call bit for bit. The choice is checked per shape at first use, on two random operand sets, outside graph capture. It is only used where ggml itself calls cuBLAS (more than 16 rows).
attn_exact: ggml's attention chain (two cuBLAS GEMMs around soft_max_f32), with the softmax and permute fused. Where an FFMA or TF32 tensor-core kernel reproduces that chain for a shape, that kernel is used instead.
BitVLA: the ternary GEMM runs on mma.sync with fused epilogues, and the row kernels are rewritten. int8 x ternary products are summed in int32, so the result is exact in any order.
FoldQuant checkpoints (π0.5, GR00T N1.5 / N1.6 / N1.7): the same fused path runs around the integer sites. The act node that quantizes a site's input also runs the op before it when it is the only reader: the GELU / SiLU gate (glu_act) and the LayerNorm / adaRMS (norm_act). The prologue's CTA-per-row kernel covers rows up to K=16384, and the integer GEMM takes wider N tiles for prefill sites; int32 sums are exact in any order.

VLA_FUSED_OPS=0 runs the old graph. The other switches, the coverage per arch and how the bits are kept are documented in docs/backend/cuda-exact.md.

The branch is rebased on 5499b77 (#50, #52). EncStack keeps the SYCL cfg.fused path and the reduced-precision path ahead of this one. The pi0.5 expert takes this path only without a language-bucket mask, a KV slot or flash attention.

Nothing in ggml or llama.cpp is touched.

Why

Jetson AGX Orin, vla-bench with the benchmark protocol (3 processes × (3 warmup + 20 reps), best mean), 1012c05 vs this branch, both built the same way:

Model	`1012c05` ms	this branch ms	change
BitVLA	134.2	80.0	-40%
VLA-JEPA	97.7	61.4	-37%
VLA-Adapter	126.6	79.7	-37%
Evo-1	434.3	299.0	-31%
GR00T N1.7	133.8	94.0	-30%
TurboVLA	41.0	29.9	-27%
GR00T N1.5	131.2	98.0	-25%
SmolVLA	216.3	165.2	-24%
π0.5	353.2	267.1	-24%
π0	349.6	269.4	-23%
GR00T N1.6	133.1	102.1	-23%
OpenVLA-OFT	385.0	297.7	-23%
Octo-Small	21.2-24.9	18.7-21.6	host-bound, within noise

The calibrated FoldQuant LIBERO checkpoints (vrfai/*-w4a4, *-w4a4-res8), same protocol with jetson_clocks on:

Model	Config	`1012c05` ms	this branch ms	change
GR00T N1.7	w4a4 / w4a4-res8	84.3 / 85.3	55.5 / 56.4	-34%
GR00T N1.5	w4a4 / w4a4-res8	96.9 / 98.7	67.4 / 69.2	-30%
GR00T N1.6	w4a4 / w4a4-res8	106.7 / 108.1	71.9 / 72.8	-33%
π0.5	w4a4 / w4a4-res8	216.0 / 234.7	142.9 / 152.0	-34% / -35%

Device memory grows by 0 to 207 MiB (cuBLASLt workspace, per-shape attention and staging buffers).

Because the actions do not change, success rates measured at the defaults still hold. At the defaults, most archs are now faster than with the flags that change numerics. The full report is in docs/benchmark/jetson-agx-orin.md.

Verified

Builds clean under -Wall -Wextra (first-party code)
ctest passes (15/15, including the new fused_cuda_ops, which compares every fused node with the ggml chain it replaces, bytewise; for glu_act and norm_act, against the ggml ops followed by fq_act)
CPU-only build with ASan + UBSan and -Werror (the CI job): builds, ctest passes
Numeric output unchanged: vla_predict_check is identical between 1012c05 and this branch for all 13 archs in 20 view configurations (1, 2 and 3 views). On GR00T N1.7 it also stays identical with each switch turned off. Beyond that, replays of synthetic requests and of requests recorded from LIBERO-Object runs (GR00T N1.5 / N1.6 / N1.7, SmolVLA, Evo-1, BitVLA, VLA-Adapter, OpenVLA-OFT, TurboVLA) return bit-identical actions.
FoldQuant: the action md5 of the 8 calibrated LIBERO files and of a π0.5 SO101 model at W8A8, W4A4 and W4A4 with INT8 o/down is the same on 1012c05 and this branch, at 1 and 2 views.

Archs and backends tested: all 13 archs with checkpoints (ACT has none public) on CUDA, Jetson AGX Orin 64GB, JetPack 6.2, CUDA 12.6, sm_87. Other backends are untouched; the new nodes are emitted only on CUDA.

GGML_OP_CUSTOM nodes (layers/fused_ops.h, cuda/vla_cuda_fused.cu) that replace chains of ggml ops with one kernel each and give the same bits: LayerNorm / RMSNorm with affine or modulation, bias + GELU / GELU (erf), bias + residual with or without a gate, SiLU / GELU gates, Q/K/V split with RoPE into the attention layout, patchify. gemm_exact (mm_exact): a BF16-weight GEMM on the fastest cuBLASLt algorithm or gemm_skinny kernel that reproduces ggml's cublasGemmEx call bit for bit, checked per shape at first use on two random operand sets; otherwise that call itself. Only where ggml uses cuBLAS (> 16 rows). attn_exact: attention as ggml's chain (two cuBLAS GEMMs, soft_max_f32, permute) with the softmax emulated in one kernel, or a fused FFMA or TF32 kernel where it reproduces that chain (checked per shape). Nothing uses these yet except mm_act once a model turns the nodes on. tests/test_fused_cuda_ops compares every node with the ggml chain it replaces, bitwise. Verified: test_fused_cuda_ops on Jetson AGX Orin (CUDA 12.6, sm_87).

…odes EncStack (SigLIP and GR00T's VL self-attention blocks), the Qwen3 LM, the Qwen3-VL tower and the DiT head take the fused nodes when the model turns them on. EncStack keeps ggml's graph for several views: ggml multiplies 3D activations with one batched cuBLAS call, which one GEMM over all views does not reproduce. GR00T N1.7 also runs its vision tower inside the main graph (all views, patchify on the device), so the tower outputs never leave the device. VLA_GR00T_FUSED_VIT=0 restores the separate vision graph. rope_2d takes a view dimension; with one view it emits the same nodes. Verified on Jetson AGX Orin: bit-identical actions for every model on synthetic requests, and on requests recorded from LIBERO-Object runs (GR00T N1.5, N1.6 with 2 views, N1.7 with 1 and 2 views).

The Gemma prefix and action expert run RMSNorm (pi0.5: adaRMS) into the BF16 GEMM operand, bias-free NEOX RoPE + head layout with the K/V cache, exact attention with MQA and the suffix mask, GEGLU into the BF16 operand, and pi0.5's gated residuals in one kernel. In F32 the SigLIP layers are EncStack's block op for op, so the fused blocks stand in. Verified on Jetson AGX Orin: bit-identical actions on synthetic requests.

SmolVLA's SigLIP layers are EncStack's block op for op and run on the fused blocks (1024-token attention on the TF32 kernel). Evo-1: InternViT (fused QKV split, bias + layer scale + residual, bias + GELU erf), the Qwen2 LM (GQA, padding mask, SiLU gate) and the DiT. VLA_FUSED_OPS=vit, lm or dit limits Evo-1's fused paths to one part. Verified on Jetson AGX Orin: bit-identical actions on synthetic requests and on requests recorded from LIBERO-Object runs (Evo-1 with 3 views).

DINOv2 + SigLIP towers and projector (OpenVLA-OFT, VLA-Adapter); the Llama-2 and Qwen2.5 LMs (RMSNorm into the BF16 operand, bias + NEOX RoPE, exact attention, VLA-Adapter with its causal mask, SiLU gate). BF16 weights only. TurboVLA: DINOv3 layers (bias + half-split RoPE + head layout in one pass; the rotation takes sin itself, uploaded next to the signed table), exact attention over the views, and fused LayerNorm and residual epilogues in the fusion layers, text enhancer and ACT decoder. Octo: the block transformer (with its additive mask) and the diffusion head's LayerNorm and residuals. Verified on Jetson AGX Orin: bit-identical actions on synthetic requests, and on requests recorded from LIBERO-Object runs for OpenVLA-OFT, VLA-Adapter and TurboVLA. Octo's recorded requests do not replay deterministically on the original either, so only its synthetic set is checked.

The ternary GEMM runs on mma.sync.m16n8k32 (int8 x the 2-bit codes as uint8, -2*rowsum(A) applied once per row), with a cp.async pipeline and the weights repacked into the fragment layout at load. Products are int8 x {-1, 0, 1} summed in int32, exact in any order; the epilogue (scales, bias, gelu_tanh, residual, head layout, repeat_kv) keeps the original kernels' bf16 roundings. q/k/v run in one launch. Softmax, the norms with the act_quant that follows them, and RoPE with repeat_kv are rewritten with the original kernels' reduction order and FMA contractions. Image embeddings stay on the device for the LM input. VLA_BITVLA_MMA=0 runs the original kernels; VLA_BITVLA_DEVICE_EMBEDS=0 copies the embeddings through the host. Verified on Jetson AGX Orin: bit-identical actions on synthetic requests and on requests recorded from LIBERO-Object runs; per-layer LM dumps identical to the original kernels.

docs/backend/cuda-exact.md describes the fused nodes, gemm_exact and attn_exact, how they keep ggml's output, what each arch runs fused, the switches and how to verify. ARCHITECTURE, USAGE and the README link it. The AGX Orin report is re-run at e6653e8 with its protocol and compared with 1012c05 built the same way. getrusage does not see the CUDA buffers on this board, so the memory column is now the peak cudaMalloc + cuMemCreate total, counted with an LD_PRELOAD shim. Verified: vla_predict_check gives identical actions at 1012c05 and e6653e8 for every arch, 20 view configurations; ctest passes.

…sed ops VLA_GR00T_FUSED_VIT defaulted on even with VLA_FUSED_OPS=0 or off CUDA. The batched plain ggml graph is not bit-identical to the per-view one, so with two or more views the fallback differed from main. It now defaults to the vision fused ops being active. Checked on Jetson AGX Orin: vla_predict_check matches 1012c05 for 1, 2 and 3 views at the defaults, VLA_FUSED_OPS=0/vit/dit and VLA_GR00T_FUSED_VIT=0; ctest 15/15; latency unchanged.

…on the GPU pi0.5 writes the state into the prompt, so most calls bring a new prompt length and every GEMM over it is chosen again. That took about 2.3 s per new length on Orin. - A cuBLASLt pick for the same N, K and views within 64 rows is tried first and kept untimed when it is bit-exact on both operand sets. VLA_GEMM_REUSE=0 turns this off. - The random test operands come from a jump-ahead LCG kernel. The values match the host loop byte for byte, so the selection itself is unchanged. New-length call on pi0.5 SO101 bf16: 2596 ms -> 555 ms p50. Actions stay bit-identical to 1012c05 (80/80 over a prompt-length sweep), verify_all passes, steady-state latencies unchanged.

…UDA path A Gemma layer whose sites are all FoldQuant now takes gemma_attn_fused and gemma_mlp_fused: q/k/v come from the FoldQuant GEMMs, the attention core, RoPE, GEGLU, adaRMS and the gated residuals use the exact fused nodes, and the o / down sites keep their act nodes and residual epilogues. Actions unchanged on the pi0.5 SO101 w4a4 / w4a4od8 / w8a8 files. Orin, 2 views 224, 48 tokens: w4a4 214.9 -> 177.3 ms, w4a4od8 235.0 -> 197.8 ms, w8a8 265.6 -> 228.0 ms.

…=16384 Rows past 8192 (the pi0.5 prefix down input) used the warp-per-row kernel. They now take act_row_kernel with 16 warps a row; the values are formed in the same order, so the blobs are unchanged (test_foldquant_cuda_op passes). M=528, K=16384 alone on Orin: A8 660 -> 319 us, A4 504 -> 308 us.

gelu_mul_act (layers/fused_ops.h, built by fq_act_geglu in layers/fq_linear.h) computes gelu(g) * u as k_gelu_mul does and quantizes the product as the site's act node would, so the [T, ffn] product is never written out and read back. The act chain is spelled with IEEE ops without flush-to-zero, since vla_cuda_fused.cu builds with -use_fast_math while the FoldQuant kernels build with -fmad=false. Used by the pi0.5 Gemma MLPs; VLA_FQ_NO_ACT_FUSE=1 keeps the two nodes. test_fused_cuda_ops compares the blob bytes against geglu + fq_act for A8 and A4, rotation blocks 1/32/64 and ascale before or after the butterfly. pi0.5 SO101 actions unchanged; Orin, 2 views 224, 48 tokens: w4a4 177.7 -> 164.4 ms, w8a8 226.3 -> 212.4 ms.

Two mma tilings join the candidates: 64x128x128 over 8 warps (variant 17) and 128x128x128 with 64x64 warp tiles, 2 stages (variant 18). Prefill sites (M > 64) with N a multiple of 128 take 17, and W8 sites past 512 rows take 18; the rest keep 128x64x128. The accumulation is integer, so the outputs are unchanged (test_foldquant_cuda_op passes). Measured in the models on Orin (prefill GEMM time): pi0.5 2 views W8 73 -> 54 ms, W4 37 -> 32 ms; pi0.5 1 view W8 67 -> 61 ms, W4 39 -> 36 ms; GR00T N1.7 W8A8 / W4A4 at 1-2 views unchanged or up to 1 ms faster.

…pe on FQ Qwen3 layers glu_act generalizes gelu_mul_act to gelu(g) * u, silu(g) * u and gelu(g), each quantized as the act node of the site it feeds. Qwen3 FQ FFNs use silu_mul, the DiT FQ FFN uses gelu; the DiT plain LayerNorm and the Qwen3 q/k norm + IMROPE use the existing fused nodes on FQ layers. GR00T N1.7 FQ (W8A8, W4A4) actions are bit-identical to upstream; 2-view latency drops 95.8 -> 91.4 ms (W8) and 79.9 -> 76.2 ms (W4).

The cuBLAS row threshold only matters for float layers, where ggml picks a different BF16 GEMM below 17 rows. FoldQuant layers run the in-tree INT GEMMs, so the pi0.5 LIBERO expert (10 rows) now takes the fused path. Bit-identical: pi0.5 LIBERO w4a4 152.8 -> 144.9 ms on AGX Orin.

…xpert adaRMS norm_act runs the norm feeding a FoldQuant site inside its act node, with ggml's 1024-thread reduction order, so the normed rows are never written out. The N1.7 cross-attention q site writes its output in head layout. Both are bit-identical (test_fused_cuda_ops, verify_all, action md5 of every calibrated and SO101 checkpoint). AGX Orin, w4a4: N1.7 58.0 -> 55.6 ms, N1.6 73.1 -> 71.9, N1.5 68.3 -> 67.5, pi0.5 145.0 -> 142.7.

vla_cuda_fused.cu compiles with -use_fast_math; device-linked together with the other vla_cuda_ops sources, CUDA 12.9's nvlink rejects the mix ("Cicc option values for '-ftz' do not match"), as it did for bitvla. Same compile flags, so the same code.

hungho77

Thanks for this, the speedups look great. I went through the diff against 5499b77 and have a few concerns, mostly around the process-global caches in vla_cuda_fused.cu. These are from reading the code, I haven't run them, so please correct me if I'm missing something.

Suggested priority:

Before merge: the per-shape plan/scratch caches growing without bound (memory and step stalls on Orin), and the BitVLA repack set returning wrong actions after a reload.
Should fix: caches not keyed by device/stream, the row-less bf16_operand(W) calls in the DiT/Qwen3 paths, the process-wide fused_ops_active() flag, mm_exact flattening batched inputs, and attn_plan under stream capture.
Can be follow-ups: k_out length in gemma_attn_fused, the VLA_FUSED_OPS parsing, and the duplicated quantize chain in k_glu_act.

Details are inline.

- GEMM/attention scratch (path B, BF16 staging, cuBLASLt workspace) is one buffer per (device, stream) and slot instead of one per shape, found with the stream's handles through a thread-local lookup; plans are per device. pi0.5 with a prompt length that changes every call: +447 MiB -> +29 MiB over 21 lengths. - Attention tries a nearby shape's path A/C choice, checked bit-exact against path B, without timing, as GEMMs do; the bit-exact checks compare on the device. New-length request 445 -> 357 ms (repeat 298). - Without a plan under graph capture, attn_exact runs path B. - BitVLA: a freed model's weights leave the repacked set, so a reload in the same process repacks again (it returned wrong actions before). - Fused nodes only for weights in a CUDA buffer (fused_ops_for); the FoldQuant glu_act in the DiT and Qwen3 follows cfg.fused_ops. - bf16_operand with the operand rows in the DiT, Qwen3 and vision towers. - mm_exact leaves batched operands to ggml_mul_mat. - gemma_attn_fused: k_out holds the new tokens only when a cache is given. - Evo-1: VLA_FUSED_OPS parts matched exactly. Actions unchanged: bit-identical to 1012c05 on all archs (1/2/3 views), verify_all, test_fused_cuda_ops (36 cases), FoldQuant md5. Fixed-shape latency unchanged (A/B against the previous head).

quangnd-vr · 2026-10-27T15:58:02Z

Thanks for the careful read. I went through each point with a reproduction or a
measurement on the Orin; details inline. Fixed: the per-shape caches (memory and
new-shape stalls), the BitVLA repack set, device/stream keying, the capture
fallback, row-less bf16_operand, the process-wide fused flag, batched operands
in mm_exact, k_out with a cache, and the VLA_FUSED_OPS parsing in Evo-1.
k_glu_act keeps its own copy, with numbers in its thread. Pushed as ac6f53f. All checks were rerun:
actions are still bit-identical to 1012c05, also with the branch merged
into the current main (06ffae9).

hungho77 self-requested a review October 26, 2026 07:19

quangnd-vr added 18 commits October 26, 2026 15:40

Document the FoldQuant fused path and its Jetson AGX Orin latency

cf5cb26

CUDA fused ops: drop unused float2 reduction helpers

eff98a7

quangnd-vr force-pushed the perf/cuda-exact-kernels branch from 0616aa0 to eff98a7 Compare October 26, 2026 09:30

quangnd-vr changed the title ~~CUDA: bit-exact fused kernels, 23-40% faster on Jetson AGX Orin~~ CUDA: bit-exact fused kernels for BF16 and FoldQuant, 23-40% faster on Jetson AGX Orin Oct 26, 2026

hungho77 reviewed Oct 27, 2026

View reviewed changes

hungho77 merged commit f464653 into VinRobotics:main Oct 27, 2026
7 checks passed

Navigation Menu

Repository navigation

CUDA: bit-exact fused kernels for BF16 and FoldQuant, 23-40% faster on Jetson AGX Orin - #48

CUDA: bit-exact fused kernels for BF16 and FoldQuant, 23-40% faster on Jetson AGX Orin#48
hungho77 merged 20 commits into
VinRobotics:mainfrom
quangnd-vr:perf/cuda-exact-kernels

quangnd-vr commented Oct 22, 2026 •

edited

Loading

Uh oh!

hungho77 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

quangnd-vr commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

quangnd-vr commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why

Verified

Uh oh!

hungho77 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

quangnd-vr commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

quangnd-vr commented Oct 22, 2026 •

edited

Loading