File size: 4,660 Bytes
ef6eb55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | # 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
```
|