feat: cap shivini device-memory pool via env var and --max-device-allocation flag

[error2215]

What ❔

Adds an opt-in cap on shivini's GPU device-memory pool, exposed two ways:

• A new env var ZKOS_WRAPPER_MAX_DEVICE_ALLOCATION (Kubernetes-style sizes: 32Gi, 48GiB, 64G, or raw bytes).
• A new global CLI flag --max-device-allocation on the wrapper binary, which seeds the same value before any GPU code runs.

When a cap is provided, both ProverContext::create_with_config sites — wrapper/src/wrapper/gpu.rs::ensure_stark_context (the cached stark_context shared across phases 1+2) and the inline path in wrapper/src/lib.rs (prove_risc_wrapper_with_snark) — thread it through ProverContextConfig::with_maximum_device_allocation. When neither is set, behavior is identical to today (shivini grabs all free device memory).

New modules:

• wrapper/src/gpu_config.rs — env-var constant, parse_byte_size (K8s-style suffix parser), set_max_device_allocation / max_device_allocation (cached via OnceLock), unit tests. Lives outside the gpu feature gate so non-GPU builds still parse the CLI value and so library callers can seed the cap without depending on shivini.
• wrapper/src/gpu/context.rs — single helper apply_env_overrides(base), gated behind gpu.

The CLI setter takes precedence over the env var when both are provided; the resolved value is cached in a OnceLock so the env var is parsed at most once.

Why ❔

The default shivini pool grabs ~47 GB up front for the SNARK precomputation, and the subsequent get_light_setup / gpu_setup_and_vk_from_base_setup_vk_params_and_hints needs another large device buffer for the LDE table and Merkle caps — total ~60–70 GB of free VRAM on a single contiguous device.

This is fine on H200 (141 GB HBM3e), but breaks on:

• NVIDIA RTX PRO 6000 Blackwell (96 GB) under MIG, where the largest commonly-used slice is 2g.48gb — the pool itself barely fits and the next allocation OOMs in shivini with ErrorMemoryAllocation.
• H100 80 GB, which is borderline and OOMs in many shared-GPU configurations.

The failure surfaces as a panic out of shivini at the SNARK setup stage:

allocated 1048576 bytes on device
allocated 50398756864 bytes on device      (~46.9 GB, succeeds)
allocated 65536 bytes on host
allocated 2147483648 bytes on host         (2 GB pinned)
thread 'main' panicked ... ErrorMemoryAllocation

GPU	VRAM	Result
H200	141 GB	works
H100 80 GB	80 GB	OOM (borderline)
RTX PRO 6000 Blackwell, full card	96 GB	works
RTX PRO 6000 Blackwell, MIG 2g.48gb	48 GB	OOM
RTX PRO 6000 Blackwell, MIG 1g.24gb	24 GB	OOM

FRI proving fits comfortably on the smaller slices; only the SNARK pipeline trips this.

Capping the pool (e.g. --max-device-allocation 21Gi or ZKOS_WRAPPER_MAX_DEVICE_ALLOCATION=21Gi) leaves headroom for the setup buffers, and SNARK proving completes on smaller GPUs / MIG slices. The flag is opt-in; operators running on a full H200 do not need to change anything.

The two-channel design (env var + CLI) is deliberate: downstream services that embed zkos-wrapper as a library (running outside the wrapper binary, e.g. in Kubernetes-managed snark-prover pods) need the env var; operators driving the binary directly want a flag. Both feed the same OnceLock.

Is this a breaking change?

• Yes
• No — purely additive. New public surface: one constant, two free functions, one new module (gpu_config), one new submodule (gpu::context), one new global CLI flag, one new env var. Default behavior when both are absent is byte-identical to before.

Checklist

• PR title corresponds to the body of PR (we generate changelog entries from PRs).
• Tests for the changes have been added / updated. (parse_byte_size unit tests in wrapper/src/gpu_config.rs — 4 tests covering raw bytes, decimal suffixes, binary suffixes, and rejected garbage.)
• Documentation comments have been added / updated. (Module-level //! docs on both new modules; /// help string on the CLI flag explaining unit formats and precedence.)
• Code has been formatted. (cargo fmt --check clean on the project-pinned nightly-2026-02-10 toolchain.)