| # PAX Architecture | |
| **Verified GPU Computing via Lean 4 + PTX + Futhark** | |
| Ahmad Ali Parr · 2026 | |
| --- | |
| ## Overview | |
| PAX (Proof-Carrying Architecture for eXecution) is a sovereign GPU computing framework that | |
| generates formally verified CUDA kernels for NVIDIA Ampere (sm_86) and Hopper (sm_90). | |
| Every kernel PAX produces ships with: | |
| 1. **Lean 4 theorems** — machine-checked correctness proofs (zero sorry) | |
| 2. **PTX implementation** — hand-rolled mma.sync / cp.async code | |
| 3. **Futhark functional spec** — compiler-verifiable reference | |
| 4. **WORM audit receipt** — Blake3+Ed25519 sealed output | |
| --- | |
| ## Axioms | |
| ### Axiom 1 — Index Space Primacy | |
| Every thread accesses exactly one element of a formally defined, non-overlapping index space. | |
| The partition must be proven: coverage (every element assigned) and disjointness (no element shared). | |
| ### Axiom 2 — Permission Necessity | |
| Every memory access requires a fractional permission. Sum of permissions at any address ≤ 1. | |
| Reads require shared permission; writes require exclusive permission. | |
| ### Axiom 3 — Synchronization as State Transition | |
| Every barrier (`__syncthreads`, `cp.async.wait_group`) is a state transition in the | |
| happens-before partial order. No memory access is valid without a prior HB edge. | |
| ### Axiom 4 — Warp Distinctness | |
| Each warp executes SIMT without divergence on the critical mma.sync path. | |
| Divergence is permitted only on boundary checks (row/col bounds). | |
| ### Axiom 5 — Verification Non-Negotiability | |
| No kernel ships without a machine-checked proof of its critical path. | |
| sorries in proof files = blocked deployment. | |
| --- | |
| ## Proof Obligations (PO1–PO8) | |
| | PO | Name | Axiom | Lean 4 Theorem | | |
| |----|------|-------|----------------| | |
| | PO1 | Index space partition | 1 | `partition_coverage`, `partition_disjoint` | | |
| | PO2 | Address space separation | 2 | `shared_global_disjoint` | | |
| | PO3 | SIMT reconvergence | 4 | `warp_reconverges_before_barrier` | | |
| | PO4 | Happens-before SPO | 3 | `hb_strict_partial_order` | | |
| | PO5 | Permission sum ≤ 1 | 2 | `permission_sum_bound` | | |
| | PO6 | Barrier permission conservation | 3 | `barrier_conserves_permissions` | | |
| | PO7 | Data-race freedom | 2,3 | `no_data_race` | | |
| | PO8 | Termination + correctness | 5 | `kernel_terminates`, `kernel_correct` | | |
| --- | |
| ## HyperKitty Constraint DAG | |
| ```xml | |
| <DAG> | |
| <Node id="🧠Input"/> | |
| <Node id="📚Memory"/> | |
| <Node id="🔍Retrieval"/> | |
| <Node id="⚙Transform"/> | |
| <Node id="⚖Constraint"/> | |
| <Node id="🔐Proof"/> | |
| <Node id="🌐Output"/> | |
| <Edge from="🧠Input" to="📚Memory"/> | |
| <Edge from="📚Memory" to="🔍Retrieval"/> | |
| <Edge from="🔍Retrieval" to="⚙Transform"/> | |
| <Edge from="⚙Transform" to="⚖Constraint"/> | |
| <Edge from="⚖Constraint" to="🔐Proof"/> | |
| <Edge from="🔐Proof" to="🌐Output"/> | |
| </DAG> | |
| ``` | |
| Formalized in `PAX/ConstraintDAG.lean` as a verified Lean 4 inductive type. | |
| Proven acyclic, single-source (Input), single-sink (Output). | |
| --- | |
| ## Kernel Categories | |
| ### FP16 Rounding (fp16) | |
| IEEE-754 binary16 round-to-nearest-even. Proven: `|round(x) - x| ≤ 0.5 ulp`. | |
| Matches hardware `__float2half_rn` and PTX `cvt.rn.f16.f32`. | |
| ### GEMM (gemm) | |
| `mma.sync.aligned.m16n8k8` FP16→FP32. Tile: 128×128 work-group, 32×64 warp, 16×8 MMA. | |
| Proven: `wmma_gemm = gemm_spec` for all FP16 inputs in normal range. | |
| ### Pipeline (pipeline) | |
| 3-stage `cp.async` double buffer. Proven: achieved throughput ≥ (1-1/3) × min(compute_bw, memory_bw). | |
| HB edges: `HB(copy[s], compute[s])` and `HB(compute[s], copy[s+1])`. | |
| ### Epilogue (epilogue) | |
| `Bias + GeLU` and `Residual + GeLU` in-register fusion. | |
| Proven: `|GeLU_approx(x) - GeLU_exact(x)| ≤ 0.001` for `x ∈ [-8, 8]`. | |
| ### Warp (warp) | |
| `shfl.sync.xor.b32` butterfly reduction. Proven correct for dot product and softmax max. | |
| --- | |
| ## Hardware Target | |
| | Property | Value | | |
| |----------|-------| | |
| | GPU | NVIDIA RTX 3080 | | |
| | Architecture | Ampere sm_86 | | |
| | VRAM | 10 GB GDDR6X | | |
| | Tensor Cores | 3rd gen (m16n8k8 FP16→FP32) | | |
| | Async Copy | `cp.async.ca.shared.global` | | |
| | Max Shared Mem | 48 KB/block (or 100 KB with dynamic) | | |
| Secondary target: H100 sm_90 (TMA cluster multicast, `cp.async.bulk`). | |
| --- | |
| ## Build | |
| ```bash | |
| # Lean 4 proofs | |
| cd PAX && lake build | |
| # PTX kernels | |
| nvcc -arch=sm_86 -ptx src/rtx_gemm_ptx.cu -o build/pax_gemm.ptx | |
| nvcc -arch=sm_86 src/rtx_gemm_ptx.cu -o build/pax_gemm.so --shared | |
| # Futhark spec | |
| futhark cuda src/pax_kernel.fut -o build/pax_kernel | |
| # Fine-tune PAX-Coder | |
| python3 export_training_data.py | |
| python3 finetune_pax_coder.py | |
| ``` | |