perf: batchable scan-based Gemma inference engine (restrict-compliant)

[koenvanderveen]

What

Adds gemma_inference_restrict_scan.py: a variant of the restrict-compliant Gemma-3 inference engine that runs the layer stack through flax.linen.scan (one compiled layer body) instead of a Python for-loop that jax.jit unrolled into num_layers copies.

Why

The unrolled graph kept every layer's weights + fp32 matmul working-copies live at once (~24 bytes/param, ~12× the bf16 weights). On 27B this peaked at ~738 GB — it OOM'd on the largest RAM GCP box and could not be batched at all.

Scan keeps only one layer's working set live → 27B peak ~230 GB regardless of batch size, so 27B both fits and batches on a CPU machine.

Measured on 27B (n2d-highmem-96, CPU)

batch	decode tok/s	s/example (100-tok gen)	peak RAM
1 (unrolled, old)	0.06	1633 s	738 GB
128 (scan)	3.75	27 s	227 GB
256 (scan)	7.15	14.5 s	228 GB
512 (scan)	12.76	8.3 s	228 GB
1024 (scan)	17.75	6.3 s	235 GB

Decode is memory-bandwidth bound; throughput scales ~linearly with batch until the knee at ~512 (~195× faster per example than unbatched). Beyond 512 it turns compute-bound with a very expensive prefill compile.

syft-restrict

• The scan primitive (flax.linen.scan) and weight restack live in public wrappers (build_scanned_blocks, stack_layer_params); the private Block is passed by name. Same allow_functions → policy_id unchanged (d84d1e21530fa500, verified).
• Greedy output byte-identical to the opt engine on 270m and 1b.

Notes / follow-ups

• x = jnp.float32(x) is required because nn.scan's carry needs a fixed dtype (the unrolled loop silently ran layer 0 in bf16, the rest fp32). jnp.float32 is already allow-listed → no policy change.
• Before production: regenerate the certificate/obfuscated file for owner re-approval (source hash changed, policy_id did not), and spot-check byte-identity directly on 27B.
• The engine is standalone; to drive it via gemma_bench.py (which lives on rasswanth/optimize-inference), place it alongside the bench in notebooks/enclave/gemma/.