visible-cx/Gemma-4-31B-CoreAI

Gemma 4 31B — Core AI (.aimodel)

google/gemma-4-31B-it-qat-q4_0-unquantized converted to Core AI .aimodel bundles for Apple silicon by visible-cx. These are derivative artifacts: Google's QAT-trained weights rounded onto the int4 grid they were trained for and re-expressed as a Core AI graph with a custom Metal flash-decode attention kernel. They load through Core AI on macOS and are not usable by PyTorch, GGUF or MLX.

⚠️ UNQUALIFIED. Neither bundle has ever been executed, on any machine. They converted cleanly and their producer fingerprints pass; nothing beyond that is established. No token has been generated. Read Status before assuming anything about behaviour.

⚠️ Expected known issue — long unguided generations (unverified on this model). On the sibling E2B/E4B exports the Core AI runtime retains one compiled specialization per sequence-length signature, costing roughly 81 MB of GPU allocations per generated token until process exit, and the decode graphs here share the length-dependent structure that causes it. Guided decoding is not a defence — shape reuse is. Until a shape-stable re-export exists for this model, treat these bundles as suitable for bounded, repeating shapes and short generations only. Measured details: the gemma-4-E2B-CoreAI card.

Contents

Path	Bytes	Files	Context	Functions	Status
gpu-pipelined/gemma4_31b_qat_decode_int4linsym_msdpa_g8	20,122,108,736	8	4096	main (decode only)	UNQUALIFIED
hybrid-pf64/gemma4_31b_qat_decode_int4linsym_msdpa_g8_pf64	20,123,256,366	8	4096	main + prefill	EXPERIMENTAL

gpu-pipelined/, file by file:

File	Bytes
…aimodel/main.mlirb	20,089,915,701
…aimodel/main.hash	32
…aimodel/metadata.json	105
metadata.json (bundle manifest)	656
tokenizer/tokenizer.json	32,169,626
tokenizer/tokenizer_config.json	3,730
tokenizer/chat_template.jinja	18,683
tokenizer/generation_config.json	203

hybrid-pf64/ is the same layout with main.mlirb at 20,091,063,302 B — 1,147,601 B more than the decode-only build, which is the whole cost of a second entrypoint over deduplicated weights.

There is no ctx8192/ctx16384 folder. --max-ctx sets language.max_context_length in the manifest and nothing else, so a wider window is a manifest edit rather than a conversion. Whether it is affordable is a separate question; see the envelope below.

Stop token: eos_token = "<turn|>" (id 106), the turn terminator Gemma 4 emits. A host that stops on the raw upstream <eos> instead will overrun every reply.

Provenance

Base checkpoint	google/gemma-4-31B-it-qat-q4_0-unquantized — 31.3B bf16 parameters, two shards, 62,578,507,874 B, ungated
Zoo recipe	gemma-4-31b, status = "verified" — export_gemma4_12b_decode_pipelined.py int4lin --lin-sym --metal-sdpa
Toolchain base	apple/coreai-models @ b1cb71b8522d99408059fa0b98b8742171bcb0b8 + the coreai-model-zoo python overlay
Toolchain	coreai-torch 0.4.1, coreai-core 1.0.0b2, coreai-opt 0.2.1, torch 2.9.0
Producer fingerprint	coreai-core 1.0.0b2 — verified on the inner asset metadata.json
Asset creation date	20260817T112551Z (gpu-pipelined/), 20260818T132515Z (hybrid-pf64/)
Weight format	int4, per-block-32, symmetric absmax (--lin-sym) — the ggml q4_0 grid the QAT checkpoint was trained on
Vocab	262,144

"QAT-unquantized" means QAT-trained, stored at full width. The int4 rounding happens at export, onto the grid the training already targeted — this is not post-training quantization of a full-precision model.

To reproduce this conversion you need a ≥192 GB host. The export peaks at 148.9 GB RSS during the in-RAM int4 quantization pass (841 s wall, Linux x86_64), falling back to 60–90 GB for the export itself. A 128 GB box under a 75% kill rule aborts at 96 GB. The hybrid-pf64/ build peaked at 147.3 GB, so the second trace costs no extra peak.

--metal-sdpa is mandatory here, not an optimisation

_msdpa_g8 in the bundle name is the custom Metal flash-decode SDPA kernel at its default split factor 8. The 31B's full-attention layers have a 32 heads × 512 fp16 query tensor that overflows MPSGraph's GPU decode scratch heap; with the stock SDPA the model crashes at the first token.

That kernel is structurally q=1 — it opens with b, h, s, hd = query.shape # s == 1, reshapes the query to [h, hd], and computes no causal mask because a decode step's single query attends the whole grown cache. That is why gpu-pipelined/ is decode-only and prompts are prefilled one token at a time, and it is the constraint hybrid-pf64/ was built to work around.

Requirements

• Apple silicon Mac, Core AI runtime. iPhone and iPad are unsupported for this model.

• Engine contract: 2 inputs (input_ids, position_ids) → logits, with one growing KV pair. No static inputs and no per-step mask input — verified by reading the graph bytecode: the declared symbols are input_ids, position_ids, keyCache, valueCache, logits, prefill and the kernel. There is no mask input.

• States: keyCache / valueCache Float16, 60 × 1 × 16 × ? × 512. Dynamic sequence dim → GrowingKVCache (initial 256, doubling), not a static allocation at the manifest maximum.

• KV cost: 1,966,080 bytes per token of context (fp16) — 8.05 GB at 4096, 16.1 GB at 8192, 32.2 GB at 16384. The steepest KV curve of any bundle in this org.

• Memory envelope:

  Weights resident (main.mlirb)	20.09 GB
  + KV at 4096	~28.1 GB
  + KV at 8192	~36.2 GB

  On a 16 GB Mac, Metal's recommendedMaxWorkingSetSize is ≈ 10.7 GB, so the weights alone exceed that machine's entire working set by roughly 2× before a single token of KV. Minimum practical machine memory: 64 GB. A 32 GB machine is not a supported target for this bundle at any context.

• The bundle manifest declares runtime_env COREAI_CHUNK_THRESHOLD=1.

Measurements

None. No measurement of any kind has been made on Apple silicon, and no token has been generated from either bundle. Every figure in this card is a byte count, an export-host reading, a box-side torch comparison, or arithmetic from state shapes and weight formats.

The only throughput figures that exist are the model zoo's, on hardware not used here:

Source	Machine	Decode	Prefill
Core AI model zoo (upstream, same recipe)	M4 Max	17.2 tok/s	22.1 tok/s

Read published Gemma 4 Q4 figures of 40–50 tok/s with care: those use Gemma 4's MTP speculative drafters, which Core AI does not implement. Compare plain-decode to plain-decode.

The memory envelope above is a bundle-bytes proxy, and this project has one measurement saying that class of proxy is unsafe. Across six bundles compiled zero-prompt under a watchdog, the compiled graph ran 0.999× to 1.35× the bundle on disk, with no ordering of bundle bytes predicting the ordering of blobs — and on the one large bundle where the proxy was checked it was 1.35 GiB optimistic. This model has never been on a machine that could compile it, so its 21.84 GiB forecast requirement is a number of exactly that kind. Treat it as a lower bound with unknown error, not as a budget.

What has been gated, box-side

The hybrid-pf64/ bundle's kernel-vs-composite numerics, at this model's exact full-layer decode shapes (32 query heads, 4 global KV heads replicated across 16 cache slots, head_dim 512):

	vs the MPSGraph composite	vs an fp32 ground truth
max abs, S=137	1.98e-2	9.7e-4 — PSNR 75.4 dB
max abs, S=1024	1.17e-2	75.1 dB class
the composite itself, vs fp32	—	53–60 dB

The divergence from the composite is the composite's own fp16 error: against fp32 the Metal kernel is the more accurate of the two. The 12B shape (1 global KV head, 8 slots) is the control and gives the same figures. This gates the block-GQA mapping and the scale=1.0 convention, which the 31B exercises for the first time — the 12B's single global head makes that mapping degenerate. It does not gate the Metal source (no GPU on the export box) and does not replace a Mac oracle.

Trace accounting on hybrid-pf64/: 10 full layers took the metal branch at S=1 and the composite branch at S=64, 10/10 both ways — a bundle where either count is 0 fails the export.

hybrid-pf64/ — decode kernel and prefill composite in one bundle

The first Gemma 4 31B bundle with a prefill entrypoint, from one copy of the weights.

main (S=1) runs the full-attention layers on the custom flash-decode Metal kernel gemma4_dense_full_sdpa_occ8 — the same kernel the zoo's published _msdpa_g8 decode bundles carry, unmodified. prefill (S=64) runs those same layers on the MPSGraph SDPA composite, which has a query axis and applies the intra-chunk causal mask. Sliding layers use the composite in both entrypoints, as they always have.

Why this is possible at all. Two catalog facts said the 31B could have decode or prefill but not both: the Metal kernel is valid only at s == 1, and without the kernel the full layers overflow the decode scratch heap. Both are true, and the dichotomy still does not follow — the SDPA implementation is a property of the entrypoint, not of the bundle. main and prefill are two separate torch.export traces over one set of weights, each at a static query width, so a wrapper that dispatches on query.shape[2] is resolved at trace time and each entrypoint gets the implementation that is correct for it.

The three things a Mac must answer, in order:

1. Does main load and decode? This is the bundle's whole reason to exist: the full layers no longer touch MPSGraph SDPA, so the scratch-heap crash should not occur. If it still crashes, the crash was never about the SDPA lowering and the catalog account of it is wrong.
2. Does prefill survive 31B dims? Unknown and unknowable off-device. The documented crash is on the decode path; an S=64 chunk allocates a different, larger shape through the same MPSGraph lowering. Predicted outcome: it works (the 12B dense2in-pf64 rides the same composite at S=64) — but if it crashes, the honest fix is to drop the prefill function, not the kernel. Decode is the leg that must survive.
3. Chunk parity. A prompt walked in 64-token chunks must give the same logits as the same prompt walked token by token. The two entrypoints use different attention implementations, so this is a real question here in a way it is not for a single-kernel bundle. The box-side gate says the two agree to fp16 noise at S=1; nothing has checked S=64.

Usage

Swift Package Manager, via CoreAIKit — a community package, not affiliated with Apple, requiring macOS 27 beta:

.package(url: "https://github.com/john-rocky/coreai-kit", branch: "main")
// target dependency: .product(name: "CoreAIKit", package: "coreai-kit")

import CoreAIKit

let model = ModelID(
    "visible-cx/Gemma-4-31B-CoreAI",
    path: "gpu-pipelined/gemma4_31b_qat_decode_int4linsym_msdpa_g8")

let chat = try await ChatSession(model: model)   // ~20 GB download on first use

Unlike the E2B/E4B bundles in this org these are plain 2-input graphs — no PLE table sidecar, no static input buffers — so either engine will load them, and the sequential engine is available for grammar-constrained decoding. Set configuration.engineVariant = .sequential for that. Expect a very large first-run download and a long cold compile; budget disk for the compiled graph on top of the 20 GB bundle.

Integrity

Core AI .aimodel bundles are not byte-reproducible: the exporter is not deterministic even against itself. Verify by digesting the exact published bytes rather than by rebuilding. Every bundle carries main.hash, the raw 32 bytes of sha256(main.mlirb); on the Hub the same value is recoverable from the LFS oid without fetching the 20 GB file.

Status

Artifact	Status
gpu-pipelined/gemma4_31b_qat_decode_int4linsym_msdpa_g8	UNQUALIFIED — converted and fingerprinted, never executed. Established: the checkpoint is the QAT source, the recipe is the zoo's verified one, the export completed, and the producer fingerprint passes. Not established: any oracle, any parity gate, any throughput, any residency measurement, any guided-decoding check.
hybrid-pf64/…_msdpa_g8_pf64	EXPERIMENTAL — not qualified. Adds an S=64 prefill entrypoint over the same weights, with a box-side numerics gate on the kernel-vs-composite mapping (above). No Mac-side oracle, no device benchmark. Do not route production traffic to it.

A qualifying run would need, in order: a decode oracle against an fp32 reference (the gate that catches int4 or kernel-level error, which no fingerprint can see); a residency measurement covering both dirty footprint and wired growth; and grammar-constrained decoding against a real schema, which has never been exercised on any msdpa bundle.

License

Google publishes the upstream QAT checkpoint under Apache-2.0 with a license_link to the Gemma 4 license, and this repo mirrors that declaration. Use is governed by those terms and by the Gemma Prohibited Use Policy; the obligations travel with any redistribution of these bundles. The contribution here is the conversion, not the weights.