feat(structures): ggml host adapter — llama.cpp distribution for pi0.5 on Thor (SM110) - #185
Open
LiangSu8899 wants to merge 17 commits into
Open
feat(structures): ggml host adapter — llama.cpp distribution for pi0.5 on Thor (SM110)#185LiangSu8899 wants to merge 17 commits into
LiangSu8899 wants to merge 17 commits into
Conversation
…oundaries Up fuses bias + tanh-GELU with FP4 packed + SFD output; Down fuses bias + fp32 residual add with fp32 output (AlignmentCD=4), plus a beta=0 wrapper for plain bias GEMM. The CUTLASS workspace is cached per shape so steady-state calls are graph-capture safe. New symbols only; the existing fp16-boundary siglip_ffn kernels are untouched.
Relocates the ggml-cuda FlashRT layer here as the third host adapter: pure-CUDA kernels (repack, fused norm/modulation/rope, activation quantize, NVFP4 GEMM wrapper) behind a C header, plus the ggml-facing dispatch half (subgraph window matchers and caches) compiled inside the host's build tree. Fused-epilogue GEMMs are consumed from csrc/gemm/fp4.
…ml pi0.5 host Adds the jetson_pi_edge_pi05 pipeline binding (vla_tick_pipeline, complete hot-path coverage, 11 segments mapping the host's fused windows onto catalog structures) and a qualification runner: manifest validation and structure-version pins offline, plus an opt-in on-device gate that holds the steady-state action chunk to a stored golden (exact match; the adapter is bitwise deterministic after warmup). Verified all gates green on the live host and red under a deliberate pin bump and a kernel-set change.
…late kernel
{RMS_NORM, MUL(w), ADD(mul, norm)} == rms_norm(x)*(1+w) runs as one kernel
via ada_rms_mod with scale = w and shift = 0 (cached zero vector). ggml's
own fused rms_norm cannot express this form because the add operand is the
norm output itself. Opt-out via GGML_FLASHRT_NO_RMS_GEMMA.
Fuses the prefix layers' q/k/v projections, RoPE, scale, KV padding to the FA KQ stride, and the permuted f16 materializations into the fused QKV GEMM plus qkv_post. qkv_post_full additionally re-emits the rope'd K and plain V as f32 rows so graph-tail persistent-KV stores keep reading their tensors; pad rows are zeroed to preserve the graph's PAD semantics. Bitwise-identical action output; opt-out via GGML_FLASHRT_NO_QKV_PREFILL.
…16 K/V directly gemm_bias_f16out converts once from the fp32 accumulator in the epilogue, bit-equal to an fp32-out GEMM followed by an fp16 cast. The ggml adapter's vision QKV window uses it for the K/V legs when their cast tensors are handed in, eliminating the two full-tensor cast copies per layer.
FLASH_ATTN_EXT with q_tokens <= 16 over a single padded f16 KV head runs as a batched QK^T GEMM (GQA heads share the stride-0 K operand), a masked fp32-accurate softmax, and a batched PV GEMM whose fp32 result lands directly in the node's [hd, n_head, n_tok] layout via a strided C — the same decomposition the torch pipeline uses at these shapes, where it beats the stream-k flash kernel plus its fixup pass. Opt-out via GGML_FLASHRT_NO_DEC_ATTN.
The rms_gemma and decomposed decode attention windows change the action output at the ulp-amplification level (both judged against the FP16 reference on gripper-active real frames); the golden follows the verified current state. All gates green.
The GQA query heads share the one KV head, so the per-head batched QK^T and PV GEMMs collapse into single GEMMs over all n_head*n_tok query rows. Ordering the rows t-major makes the PV output land contiguously in the flash-attention node's [hd, n_head, n_tok] destination, removing the strided-C store penalty. Numerically identical to the batched form.
…tention qkv_post optionally stores the rope'd+scaled Q rows as f16 alongside the f32 tensor; its [M, Nq] layout is exactly the t-major gather order the decomposed decode attention consumes. A one-shot grow-only slot hands the buffer to the following attention window, which then skips its gather kernel (one launch less per layer). The f16 values round from the same fp32 registers the gather would read, so the results are bit-identical.
The host now offloads its action input projection, which moves the final actions at the usual chaotic-amplification scale; distances to the f16 reference were re-judged before this refresh.
The prefill graph ends with one tiny f32->f16 row copy per layer and KV tensor into the persistent encoder-KV buffers. Batch a run of them into a single launch; every copy is still performed with identical rounding, so results are bit-exact.
Export the vendored FA4 SM100-compatible forward (CuTe-DSL) ahead of time for sm_110a at the padded head_dim-80 shape and run the vision FLASH_ATTN_EXT through it. The padded Q/K/V/dst buffers all share one dense (B, S, H, D) linear layout, so the f32 boundaries reduce to dense elementwise converts and the f16 K/V pass straight through. A small shim file provides the _cuda* runtime aliases the AOT object expects, keeping the build free of any CuTe-DSL runtime dependency; the module loads once outside CUDA graph capture and the wrapper is capture-safe. Includes the regeneration script and provenance notes. Refreshes the e2e action golden: the vision attention numerics move, and the distance to the f16 reference improves.
… padding The vision FA node is followed by a head-de-pad view and a CONT whose strided copy costs more than the attention itself. Absorb the pair: the FA4 f16 output converts directly into the CONT's packed f32 destination, skipping the padded f32 store and the strided copy. The same elements go through the same f16->f32 conversion, so results are bit-identical.
Second AOT export of the vendored FA4 forward at the prefill shape (head_dim 256, GQA with one KV head, full attention). The prefill is a prefix-LM: its mask is row-uniform pad-only and the real KV length equals the query count, so passing that length as the KV dynamic shape reproduces the mask exactly and the padded tail is never read. The f32 graph boundaries again reduce to dense elementwise converts. Refreshes the e2e action golden. On the real-frame judge the overall distance to the f16 reference moves -0.0085 (within the swing this metric shows for any attention change; the gripper dimension improves +0.024); GGML_FLASHRT_NO_PREFILL_FA4=1 falls back per site.
Documents building a host against the adapter, the runtime switch matrix, the three validation layers (operator tests, qualification gates, benchmark/parity methodology), the layer architecture with its capture-safety and fusion-window invariants, and the AOT FA4 module mechanics. Adds measured Thor numbers to the overview.
The ggml adapter's build resolves CUTLASS from third_party/cutlass by default; pin it (NVIDIA/cutlass v4.4.2) as a submodule so a recursive clone is self-contained instead of relying on a local checkout.
Member
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the ggml host adapter — the third host of the
flash_rt/structurescatalog after the vllm/sglang runtime adapters — bringing the FlashRT
NVFP4 kernels and fused structures to the llama.cpp family as a
build-time source integration: the host's CMake compiles the adapter
inside its own tree, only C symbols cross the boundary, and the host
ships a single self-contained binary plus GGUF weights. No Python, no
torch, no FlashRT runtime on the device.
Structures stay single-source: the NVFP4 tcgen05 GEMMs and fused
epilogues are consumed from
csrc/gemm/fp4/, and FlashAttention-4 isconsumed as ahead-of-time-compiled modules exported from the vendored
CuTe-DSL forward (
fa4_aot/, regeneration script included) — the hostbuild needs no CuTe-DSL toolchain. The
fr_*.cufiles are translationonly: ggml wire-format repack, activation quantize, fused
norm/RoPE/modulation producers, and the
ggml_tensorwindow predicates.Also in this branch:
bindings/jetson_pi_edge_pi05.yaml(complete-hot-pathbinding for the reference host) and
qualification/(manifest + structureversion pins + exact e2e action golden), plus adapter documentation
(
README.md,USAGE.md,TESTING.md,DEVELOPMENT.md) and CUTLASSv4.4.2 tracked as a submodule so a recursive clone is self-contained.
Measured performance (Jetson AGX Thor, pi0.5, 2 camera views, P50 warm)
llama_encode+llama_decodeFor reference, the FlashRT torch frontend runs the same checkpoint at
36.4 ms end-to-end on the same device.
Testing
The reference host integration (fuse-hook call sites, CMake wiring, pi0
host fixes) lives on the Jetson-PI-Edge side; a ready-to-build branch
with this adapter mounted as a submodule:
git clone --recursive -b feat/flashrt-thor-kernels https://github.com/LiangSu8899/Jetson-PI-Edge cd Jetson-PI-Edge cmake -B build -DGGML_CUDA=ON -DGGML_CUDA_FLASHRT=ON -DCMAKE_BUILD_TYPE=Release cmake --build build --target llama-server -jThen quantize the published BF16 GGUFs with
llama-quantize ... NVFP4and run
llama-serverexactly as before — the host's HTTP protocol andclients are unchanged. Validation layers (operator tests, the
qualification gates, the benchmark/parity methodology) are documented in
flash_rt/structures/adapters/ggml/TESTING.md; every fused window has aGGML_FLASHRT_NO_*runtime kill switch for window-by-window bisection.The gap to the torch frontend (~6 ms) is accounted for and the
optimization line continues on this branch; updates land here and the
host branch follows via submodule bumps.