Gemma 4 E4B (NVIDIA Nemotron 3.5 Bolt E4B): per-layer inputs and shared KV, Intel Arc first

Checkpoint

nvidia/NVIDIA-Nemotron-3.5-Bolt-E4B-BF16, revision 49843eed8db1f1c616c47af75cfd36e07aec0667. This is NVIDIA's post-trained Gemma 4 E4B (Gemma4ForConditionalGeneration, text config gemma4_text), released on 2026-10-23 for robots and tool use.

What the native Gemma 4 family refuses today

zig/src/families/gemma4/config.zig rejects this checkpoint at three gates:

Line	Gate	E4B value
186	`GemmaPerLayerInputsNotServed`	`hidden_size_per_layer_input: 256`, `vocab_size_per_layer_input: 262144`
187	`GemmaSharedKvNotServed`	`num_kv_shared_layers: 18`
188	`GemmaKeysNotValues`	`attention_k_eq_v: false`

The rest of the geometry also differs from the 12B and the 31B:

42 layers, hidden size 2,560, intermediate size 10,240;
8 query heads and 2 KV heads;
256-wide heads with a 512-token sliding window, and 512-wide global heads;
final_logit_softcapping: 30, vocabulary of 262,144, tied embeddings.

Proposed work

E-series support in the family. Config, loader and the two missing mechanisms:
- per-layer inputs: embedding table, projection and gate;
- shared KV: layers 24-41 read the cache of layers 22 and 23.
These parts are backend-independent, so the CUDA (Proposal: Gemma 4 on CUDA, built on the shared CUDA operators (shared extensions first) #1048) and Metal paths can reuse them.
A reproducible converter from BF16 to MLX affine 4-bit, group 64. A local prototype matches mx.quantize (mlx 0.32.4) byte for byte on all 345 quantized tensors.
The Intel Arc (XPU, Level Zero) backend first: text only, one stream, then the server with tools, prefix cache and lanes. Each step comes with the backend's usual receipts: forced-token comparison against a BF16 transformers reference, FP64 projection checks, and window/seam bit-identity tests.

Out of scope for now:

audio and vision;
the DFlash drafter: its GGUF metadata and its model card name different licenses, so I'm keeping it local until that's clarified;
CUDA kernels, left to whoever owns the Gemma CUDA work.

Question

Is anyone already working on the E-series (E2B/E4B) parts, per-layer inputs or shared KV, on any backend? If so, I'm happy to pair or split along the lines above.

/atc claim

Update: I'd also like to cover CUDA for the E-series (DGX Spark), on top of the Gemma 4 CUDA path from #1048, unless someone already owns it. Happy to wait for a go or to pair.

@Mirrdhyn

Thanks @Mirrdhyn, and sorry for the wait: yes, go ahead, CUDA included.

One request on order. First the family support the E-series needs, shared by every backend: config and loader, per-layer inputs (the embedding table, projection and gate), shared KV for layers 24 to 41, and separate K and V, with host tests against a float64 reference. Then CUDA on top of the Gemma 4 CUDA series that just landed (#1174 to #1186); please coordinate with @jschmied, who wrote it. Intel Arc can follow. Exactness is the gate as always: drafted, resumed and concurrent replies equal plain ones.

- Synthema, Ash's AI Assistant.

@Mirrdhyn

@Mirrdhyn glad to pair on the CUDA side. A few pointers into the Gemma 4 CUDA family:

Attention. E4B's heads are already covered by window_attention: 256-wide with a 512-token window (cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities #1123) and 512-wide global over a linear cache (cuda window_attention: head dim 512 over a linear cache, for global attention #1124). Its 8/2 GQA is within what those kernels take.
Where the three refusals land on CUDA.
- Per-layer inputs: the embedding, projection and gate need weights in cuda_weights.zig and a step in cuda_forward.zig.
- Shared KV: layers 24-41 reading 22/23 means cuda_forward.zig's per-layer caches index a source layer instead of owning one, and the cache bytes (Model.openBytes in cuda_forward.zig, which cuda_native.zig's Plan sums for admission) shrink by 18 layers.
- Separate K and V: global layers take V from k_proj's output today (cuda_weights.zig:34 and :216). E4B needs their v_proj loaded and used, as the sliding layers already do.
Snapshots. cuda_snap.zig copies every layer's cache today, so with shared KV it should copy only the owning layers, or resume could double count.

If it helps, split it this way: you do the family parts shared by every backend (config, loader, the f64 reference), and I take the CUDA pieces on top. Or the other way round. Your call.

@jschmied

Thanks both. Agreed on the order: family parts first, CUDA on the Gemma 4 CUDA series, Intel Arc after.

@jschmied, I'll take your first split: I do the family parts shared by every backend, you take the CUDA pieces on top. In families/gemma4:

the three gates in config.zig (186-188) open for the E-series: per-layer inputs, shared KV, separate K and V;
the loader for the per-layer embedding table, projection and gate, and for the global layers' v_proj; like transformers, it skips the shared layers' unused k_proj, v_proj and k_norm;
the shared-KV map: which layer owns the cache each layer reads, so a backend sizes and snapshots only the owning layers (your cuda_snap.zig point);
host tests against a float64 reference;
the BF16 to MLX affine 4/64 converter, which matches mx.quantize on all 345 quantized tensors.

Some context on where this comes from: I already have local E4B prototypes on Intel Arc and CUDA (GB10), with Metal in progress, each a standalone path written before the family could take the E-series. They won't become PRs as they are; what they bring here is the mechanisms already checked end to end, oracles, and receipts to cross-check the family parts and your CUDA pieces.

Three points to settle before I write it, since every backend inherits the answer:

Per-module precision. On E4B, 4/64 everywhere is not the only recipe worth serving: with the per-layer gate and projections kept in BF16, moving the attention projections to 8 bits lowers the KL against BF16 further in my forced-token runs; I'll bring the numbers for both recipes with the PR. I'd have the loader read mlx-lm's per-module overrides in quantization (a module set to false stays BF16, a module with its own bits), so mlx-lm and the converter write the same checkpoints. Fine with you, or should the family stay 4/64 only for now?
Rounding steps on the per-layer path. The family follows MLX's bf16 steps (forward.zig:188). For the per-layer inputs (projection, scale, sum with the table row, gate) I'd do the same, so the Metal path can be checked against mlx-lm, and write the float64 reference to those steps. My CUDA prototype keeps the projection in fp32 there; I'd drop that for the family's rule unless you prefer otherwise.
Backends without the E-series yet. Once the config admits E4B, a backend that doesn't compute per-layer inputs or shared KV must still refuse it rather than load it. I'd move the three refusals from the config to each backend's admission, lifted backend by backend as each one lands with its receipts. Metal included: I can bring it in the same series or right after, as you prefer.

I'll post here when the family PR is up.

Sounds good: you take the family parts, I take the CUDA pieces on top. My view on the three points:

Per-module precision: yes, read mlx-lm's per-module overrides in quantization (false stays BF16, a module with its own bits), so mlx-lm and the converter write the same checkpoints. It changes which weights are stored at which width, not the activation formats, so I don't see a conflict with the family's rules. The KL numbers for both recipes in the PR would be good to have.
Rounding steps: follow the family's MLX bf16 steps on the per-layer path too, and write the float64 reference to those steps. One order for every backend is easier to check. On CUDA I'll follow the family's steps rather than keeping the projection in fp32.
Refusals: agreed. Each backend's admission refuses the E-series until it computes per-layer inputs and shared KV, and the refusal is lifted backend by backend with its receipts. For Metal, right after the family PR keeps the first PR smaller, but either is fine with me.

I'll start the CUDA pieces (per-layer inputs, shared-KV sizing and snapshots of the owning layers only) once your family PR is up.

For the ATC board, if you like: post atc propose 1222 pair and I'll accept it, so the radar shows the CUDA pieces as yours too. It changes nothing to what we agreed above.

/atc propose pair

/atc accept 6411720610

@jschmied

The family parts are up: #1400 (config, per-backend refusals), #1401 (tensor plan, per-module overrides), #1402 (per-layer inputs on the host at MLX's bf16 steps, float64 reference), and #1403 (converter, with the KL of both recipes). @jschmied, #1402 also adds cuda_layer_input.zig, so a CUDA test can import the host steps. Metal comes next.

gemma4: admit the E4B geometry in the family config, each backend ref…

gemma4: plan the E4B's per-layer input, separate global value and sha…

gemma4: compute the E-series per-layer inputs on the host at MLX's bf…

/atc renew

gemma4: compute the E-series per-layer inputs on the host at MLX's bf…

@Mirrdhyn

Status: #1400 and #1401 shipped in 1.1.7 and #1402 in 1.1.8. We didn't take #1403, as it would make lossy conversions of our own. Next come Metal from @Mirrdhyn and the CUDA pieces from @jschmied, each lifting its own backend's refusal with its receipts.

/atc renew

@Mirrdhyn

Status: #1552 has landed. The XPU backend now serves Gemma 4 E4B on the Intel Arc Pro B70 from mlx_lm.convert 4-bit
checkpoints, with the per-layer input linears kept in BF16 and, as an option, 8-bit attention projections. It keeps
prompt states, runs up to 8 streams with shared forwards, and answers /v1/decisions. On the B70 it reads 2,048-token
prompts at 6,772 tok/s and decodes one stream at 135.6 tok/s. 4 streams reach 426 tok/s together, each reply equal to
the stream's reply alone.

Left: Metal from @Mirrdhyn and the CUDA pieces from @jschmied, each lifting its own backend's refusal with its
receipts. The audio tower and the drafter are not in yet. Thanks @Mirrdhyn.

Navigation Menu

Repository navigation

Gemma 4 E4B (NVIDIA Nemotron 3.5 Bolt E4B): per-layer inputs and shared KV, Intel Arc first #1222

Checkpoint

What the native Gemma 4 family refuses today

Proposed work

Question

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Gemma 4 E4B (NVIDIA Nemotron 3.5 Bolt E4B): per-layer inputs and shared KV, Intel Arc first #1222

Description

Checkpoint

What the native Gemma 4 family refuses today

Proposed work

Question

Activity

Mirrdhyn commented on Oct 23, 2026

Mirrdhyn commented on Oct 23, 2026

ashhart commented on Oct 24, 2026

jschmied commented on Oct 24, 2026

Mirrdhyn commented on Oct 26, 2026

jschmied commented on Oct 26, 2026

Mirrdhyn commented on Oct 26, 2026

jschmied commented on Oct 26, 2026

Mirrdhyn commented on Oct 26, 2026

Mirrdhyn commented on Oct 26, 2026

Mirrdhyn commented on Oct 27, 2026

ashhart commented on Oct 28, 2026

jschmied commented on Oct 30, 2026

ashhart commented on Oct 31, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions