File size: 6,254 Bytes

ef6eb55

# PAX-Coder Demo

Verified GPU kernel generation demonstration with realistic Lean 4 proofs, PTX kernels, and Futhark specs.

## Quick Start

```bash

# Mock mode (no model download required)

python3 demo.py



# Non-interactive mode (useful for CI/scripting)

python3 demo.py --no-pause



# Use actual Ollama model (if running locally)

python3 demo.py --live



# Target H100 (Hopper, sm_90) instead of RTX 3080 (Ampere, sm_86)

python3 demo.py --arch sm_90

```

## What It Shows

The demo demonstrates PAX-Coder end-to-end with 5 kernel categories:

| Category | What's Proven | Example Output |
|----------|---------------|-----------------|
| **FP16** | IEEE-754 binary16 rounding error bound | Lean 4 proof + PTX `cvt.rn.f16.f32` |
| **GEMM** | 128×128 matrix multiply correctness | `mma.sync` kernel + index partition proof |
| **Pipeline** | 3-stage async GEMM throughput bound | `cp.async` + happens-before proof |
| **Epilogue** | Bias+GeLU fusion numerical stability | In-register computation proof |
| **Warp** | Tree-reduction warp shuffle correctness | `shfl.sync.xor` + divergence-free guarantee |

For each category, you see:
- **Prompt**: What was asked
- **Lean 4 Proof**: Machine-checked correctness (zero sorry)
- **PTX Kernel**: Hand-rolled `mma.sync`, `ldmatrix`, `cp.async` code
- **Futhark Spec**: Functional reference implementation
- **PAX Certificate**: Which proof obligations (PO1–PO8) are satisfied

## Command-Line Options

```

--live              Use actual Ollama instance (localhost:11434)

--arch {sm_86,sm_90}

                    Target GPU (default: sm_86 / RTX 3080)

--no-rich           Disable colored terminal output (plain text)

--speed SPEED       Streaming animation speed multiplier (default: 1.0)

--no-pause          Skip pauses between demos (for CI/automation)

--help              Show this help message

```

## Requirements

### Minimal (Mock Mode)
- Python 3.10+
- Standard library only

### Optional (Live Mode + Rich Output)
- Ollama running at `localhost:11434` with `Snapkitty/pax-coder-7b` model
- `rich` library: `pip install rich`

```bash

# Install rich for prettier output

pip install rich



# Run Ollama locally for --live mode

ollama run Snapkitty/pax-coder-7b

```

## Example Output

Running `python3 demo.py --no-pause` will generate ~680 lines showing:

1. **PAX-Coder banner** with legal information
2. **5 kernel demos** (fp16, gemm, pipeline, epilogue, warp)
3. **VRAM usage breakdown** (~8.1 GB on RTX 3080)
4. **Call-to-action** for Sovereign Node Key
5. **Links** to GitHub, HuggingFace, documentation

### Sample Output Structure

```

================================================================================

DEMO 1/5: FP16

================================================================================



📋 PROMPT:

Prove that IEEE-754 binary16 rounding error is bounded by 0.5 ulp...



🔍 LEAN 4 PROOF:

theorem fp16_rounding_bound (x : Float)...

  nlinarith [ulp_nonneg (roundToFP16 x), ...]



⚙️  PTX KERNEL:

// IEEE-754 binary16 RNE conversion

.target sm_86

cvt.rn.f16.f32 h_out, f_in;



🌐 FUTHARK SPEC:

def round_fp16 (x : f32) : f16 = f16.from_f32 x



✓ PAX CERTIFICATE: [PO4 | PO5 | PO7]

```

## Proof Obligations (PO1–PO8)

| PO | Invariant | Example |
|----|-----------|---------|
| PO1 | Index space partition | Coverage + disjointness proven |
| PO2 | Address space separation | `shared ∩ global = ∅` |
| PO3 | SIMT reconvergence | Before every barrier |
| PO4 | Happens-before SPO | Strict partial order proven |
| PO5 | Permission sum ≤ 1 | Fractional permissions at every address |
| PO6 | Barrier permission conservation | Preserved across `__syncthreads` |
| PO7 | Data-race freedom | No concurrent writes to same address |
| PO8 | Termination + correctness | Kernel always terminates correctly |

Each PAX-Coder output lists which POs are satisfied by that kernel.

## Architecture Targets

### Ampere (sm_86) — RTX 3080 — Default

- `mma.sync.aligned.m16n8k8.f32` (FP32 accumulate)

- `mma.sync.aligned.m16n8k16.f32` (FP16 input)

- `ldmatrix.sync.aligned.m8n8.x4.b16`

- `cp.async.ca.shared.global` + `cp.async.wait_group`
- Shared memory: 48 KB (or 100 KB dynamic)

### Hopper (sm_90) — H100 — `--arch sm_90`
- TMA (Tensor Memory Accelerator) multicast
- `cp.async.bulk` (pipelined async copy)
- Cluster sync primitives
- Thread blocks per cluster

## Running on Different GPUs

```bash

# Default: RTX 3080 Ampere (sm_86)

python3 demo.py



# H100 Hopper (sm_90)

python3 demo.py --arch sm_90



# With actual model (requires Ollama)

ollama run Snapkitty/pax-coder-7b "Write verified GEMM for sm_90"

python3 demo.py --live --arch sm_90

```

## For CI/Automation

```bash

# Non-interactive, plain text, full output to file

python3 demo.py --no-pause --no-rich > pax_demo.log 2>&1



# Check all POs are satisfied

python3 demo.py --no-pause 2>&1 | grep "PAX CERTIFICATE"

```

## Troubleshooting

### `ModuleNotFoundError: No module named 'rich'`
Rich is optional. Run `pip install rich` or use `--no-rich` for plain output.

### `ModuleNotFoundError: No module named 'requests'`
Only needed for `--live` mode. Install with `pip install requests`.

### Ollama connection refused
Ensure Ollama is running: `ollama serve`
Then in another terminal: `ollama run Snapkitty/pax-coder-7b`

### Unicode/Encoding errors on Windows
The script handles UTF-8 automatically. If issues persist, try `--no-rich`.

---

## Learn More

- **GitHub**: https://github.com/SNAPKITTYWEST/pax-coder
- **HuggingFace**: https://huggingface.co/Snapkitty/pax-coder-7b
- **Architecture Doc**: https://github.com/SNAPKITTYWEST/pax-coder/tree/main/PAX
- **User Guide**: https://github.com/SNAPKITTYWEST/pax-coder/tree/main/docs

## License

PAX-Coder is tri-licensed:
- **BSL-1.1** (Business Source License 1.1) — restricts until 2028-08-08
- **AGPL-3.0** (GNU Affero General Public License 3.0) — starting 2028-08-08
- **MPL-2.0** (Mozilla Public License 2.0) — alternative terms

Copyright: Ahmad Ali Parr, Bel Esprit D'Accord Irrevocable Trust

---

*Evidence or Silence — 2026*