DEV Community: xbill

Teaching Kiro to Paint: A Stateful Image-Editing Skill Built on Gemini's Interactions API and MCP

xbill — Thu, 23 Jul 2026 00:31:36 +0000

TL;DR: nb2lite-skill-kiro wraps Google's gemini-3.1-flash-lite-image model (NB2Lite) in a tiny FastMCP server and packages it as a Kiro skill. You type "generate an image of a cyberpunk kitchen" into Kiro, and it just... does it. Then you say "add a neon RAMEN sign" and it edits the same image without re-prompting the whole scene. Oh, and the cover image of this article? Generated by the thing the article is about — dogfooding all the way down. More on that at the end.

Background: why another image tool?

Most image-generation workflows are stateless. You send a prompt, you get pixels back, and the model immediately forgets everything. Want to tweak the result? You re-describe the entire scene and pray the character, lighting, and composition survive the round trip. (Narrator: they don't.)

Google's NB2Lite — the friendly nickname for gemini-3.1-flash-lite-image — takes a different approach. It's a high-efficiency image model with sub-2-second generations, solid text rendering in 25+ languages, and — the headline feature — support for the stateful Interactions API, which lets you iterate on an image across multiple turns while the model keeps the visual context server-side.

This repo glues that capability into Kiro, so your coding agent can generate and iteratively refine images as a natural part of a session. It ships as two things in one repo:

A Model Context Protocol (MCP) server (nb2lite-agent, a single-file FastMCP app in server.py) exposing exactly four tools.
A Kiro skill (nb2lite-image) that teaches Kiro when and how to use those tools well.

The Interactions API: images with a memory

The Interactions API is Gemini's stateful endpoint. The core loop looks like this:

You call client.interactions.create(...) with a prompt and store=True.
The response includes an interaction_id — a handle to the turn's visual context, persisted on Google's servers.
On the next call, you pass previous_interaction_id, and the model edits the existing canvas — preserving character, style, lighting, and pixel continuity.

So instead of this (stateless suffering):

"A watercolor fox in a forest at dawn, mist, soft light, wearing a red scarf, three birch trees on the left, and now also holding a lantern"

...you write this:

"Add a lantern in its paw."

That's it. The stored context holds the rest.

A few practical details the server handles for you:

Every turn returns a new interaction ID. Chain the latest one; editing from a stale ID silently forks your session from an older state (a subtle and very annoying bug if you roll this by hand).
Aspect ratio is chosen at generation time (1:1, 16:9, 9:16, 4:3, 3:4) and inherited on stateful edits — changing it mid-session degrades pixel continuity, so the edit tool deliberately doesn't accept one.
Thinking levels: low (default, fast drafts) or high (complex rendering, accurate text layout, character composition). The generic API spec also lists minimal and medium, but the live API rejects them for this model with an HTTP 400 — the server saves you from discovering that the hard way.

What is MCP, in one minute

The Model Context Protocol is an open standard for connecting AI assistants to tools and data. Before it, giving a model access to some service meant writing a bespoke integration for each assistant — N assistants × M services, everyone reinventing the same plumbing. MCP collapses that: a tool author writes one MCP server that exposes typed tools, and any MCP-capable client (Kiro, Claude Code, Claude Desktop, and a growing list of others) can discover and call them with no per-client glue code.

An MCP server is usually a small local process that speaks JSON-RPC over stdio. The client launches it, asks "what tools do you have?", and from then on the model can call them like functions.

The nb2lite-agent server exposes exactly four:

Tool	What it does
`generate_image`	Text → 1k image. Saves locally, returns the path + an interaction ID.
`edit_image`	Stateful edit: takes the previous interaction ID + a description of only the change.
`edit_local_image`	Uploads any local image file inline (base64) and applies an edit — your entry point for existing files.
`get_help`	Reports live config: API key status, active model, output directory, full tool reference.

Images land on disk as gen_<timestamp>_<uuid8>.jpg (or edit_/edit_local_ prefixed) — the UUID suffix keeps concurrent generations from clobbering each other. Errors come back as 🔴 ... text strings rather than protocol errors, so the agent can read and react to them.

And what's a Kiro skill?

If MCP is the hands (the tools Kiro can physically call), a skill is the muscle memory — a markdown file (SKILL.md) plus bundled resources that load into Kiro's context and teach it the workflow: which tool to reach for, in what order, with which constraints.

Kiro skills live in .kiro/skills/<skill-name>/ inside a project. When Kiro detects a trigger phrase that matches the skill's description, it activates the skill automatically and starts using the guidance encoded there.

For nb2lite-image, the skill encodes things like:

Call get_help first when diagnosing setup issues — if the API key is missing, nothing else will work.
Keep edit prompts incremental: describe the change, not the scene.
Always chain the latest interaction ID.
Generations are billable — batch related edits and prefer thinking_level: low for drafts.

The skill also bundles the MCP server itself (mcp/server.py), its requirements, an installer script, and a vendored copy of the Interactions API developer guide — so it's self-contained: install the skill, and you have everything needed to also stand up the server.

Installing it: the "I just want it to work" edition

You need three things: Python 3.10+, Kiro, and a Gemini API key (free from Google AI Studio). Pick one of the paths below.

Path A: Clone and bootstrap (this repo)

# 1. Get the code
git clone https://github.com/xbill9/nb2lite-skill-kiro.git
cd nb2lite-skill-kiro

# 2. One-command setup: installs deps, registers the MCP server
#    in .mcp.json, and prompts for your API key (stored in ~/gemini.key)
./init.sh

# 3. Restart Kiro in this directory and approve the server when prompted.

That's genuinely it. init.sh is safe to rerun if anything looks off.

Path B: Install into your project

From a clone of the repo:

make init TARGET=/path/to/your/project ARGS='--output-dir ./images'

This copies the skill into <project>/.kiro/skills/nb2lite-image/ and writes the nb2lite-agent entry into that project's .mcp.json. It reuses ~/gemini.key if you've set one up. Restart Kiro in the target project, approve the server, done.

Path C: Manual registration

If you'd rather wire it up yourself:

# Install deps
pip install -r requirements.txt

# Register the MCP server in your project's .mcp.json
# (add this entry under "mcpServers")

{
  "mcpServers": {
    "nb2lite-agent": {
      "command": "python",
      "args": [".kiro/skills/nb2lite-image/mcp/server.py"],
      "env": {
        "GEMINI_API_KEY": "your-key-here",
        "IMAGE_OUTPUT_DIR": "./images"
      }
    }
  }
}

Then copy the skill files into .kiro/skills/nb2lite-image/ and restart Kiro.

Path D: Docker (nothing on the host but Docker)

The server is published as xbill9/nb2lite-agent. Add this to your project's .mcp.json:

{
  "mcpServers": {
    "nb2lite-agent": {
      "command": "docker",
      "args": [
        "run", "--rm", "-i",
        "-e", "GEMINI_API_KEY",
        "-v", "/abs/path/to/project:/abs/path/to/project",
        "-w", "/abs/path/to/project",
        "xbill9/nb2lite-agent"
      ],
      "env": { "GEMINI_API_KEY": "your-key-here" }
    }
  }
}

The -v "$PWD:$PWD" -w "$PWD" mount matters: the server saves images to disk and reads local files for edit_local_image, so the container must see your project at the same absolute path as the host.

Troubleshooting, the whole guide

MCP server doesn't appear → restart Kiro in the project directory.
Tools return 🔴 GEMINI_API_KEY is not set → run source set_env.sh (or export the key) and restart.
Anything else → ask Kiro to call get_help; it reports the live config.

Examples: a session in practice

Once installed, you talk to it in plain English. A real flow looks like:

You: "Generate a cozy cabin in a snowy forest at dusk, 16:9."

Kiro calls:

generate_image(
    prompt="A cozy log cabin in a snowy forest at dusk, warm light in the windows",
    aspect_ratio="16:9",
    thinking_level="low",
)
# 🟢 Saved to: ./gen_1784759001_a1b2c3d4.jpg
# Interaction ID: v1_ChdpRU5...

You: "Nice. Add smoke curling from the chimney."

edit_image(
    previous_interaction_id="v1_ChdpRU5...",
    edit_prompt="add gentle smoke curling from the chimney",
)
# 🟢 Saved to: ./edit_1784759050_e5f6a7b8.jpg
# Interaction ID: v1_Xk9mPq2...   ← a NEW id; the next edit chains this one

You: "Now make it night, with aurora in the sky."

Same tool, newest ID, and the cabin, trees, and chimney smoke all stay put — only the sky changes. No re-prompting, no continuity roulette.

And for images that didn't come from the model at all:

You: "Take ./whiteboard-sketch.png and render it as a clean 3D product mockup."

edit_local_image(
    image_path="./whiteboard-sketch.png",
    edit_prompt="render this hand-drawn sketch as a high-fidelity 3D product mockup",
    aspect_ratio="4:3",
)

It returns an interaction ID too — so follow-up refinements switch to edit_image and go stateful from there.

Dogfooding: about that cover image 🐕🍖

If the term is new to you: "eating your own dog food" means using your own product for real work, not just demoing it. It's the difference between "this should work" and "I ship with this every day." If a tool is good enough for your users, it should be good enough for you — and if it isn't, you'll be the first to feel the pain and fix it.

This repo dogfoods itself at every layer:

The skill is active inside its own repository — open Kiro in a clone and the nb2lite-image skill and nb2lite-agent server are already wired up, so every development session doubles as an integration test.
The integration tests (make test) drive the same four MCP tools an end user would, against the live API.
And now, the cover image of this article was generated by the exact skill the article describes, from inside a Kiro session in this repo. One tool call, first attempt, no retouching:

generate_image(
    prompt="A sleek dark-themed tech cover image for a developer tool called "
           "'nb2lite-image'. The scene shows a futuristic AI workspace: a glowing "
           "terminal/IDE interface with code on the left side, and on the right, a "
           "vibrant image generation panel showing colorful AI-generated artwork "
           "materializing from pixels. In the center foreground, the Kiro logo — a "
           "stylized Kiro ghost with a neon blue glow — acts as the bridge between "
           "code and image, with light rays connecting it to both sides. Bold text "
           "reads 'nb2lite-image' at the top and 'Powered by Gemini • Built for Kiro' "
           "at the bottom. Color palette: deep navy background, electric blue and "
           "purple accents, white text, with sparkling particle effects. High-tech, "
           "modern, developer-aesthetic, cinematic lighting.",
    aspect_ratio="16:9",
    thinking_level="high",
)
# 🟢 Image successfully saved!
# • Saved to: cover-image.jpg
# • Interaction ID: v1_ChdtbDloYW9iV0JaZlZqTWNQcHJDTXVBOBIXbWw5aGFvYldCWmZWak1jUHByQ011QTg

(That exact output is committed to the repo as cover-image.jpg, receipts and all.)

Worth noticing:

The text rendered correctly. "nb2lite-image" at the top and the full tagline at the bottom came out crisp and typo-free — that's what thinking_level: "high" buys you on text-heavy layouts.
Kiro is literally in the picture. The glowing "K" logo at center stage isn't decoration — it represents the actual agent orchestrating the generation. The skill generated an image that explains its own architecture.
The dual-monitor composition tells the whole story. Code editor on the left, AI-generated artwork on the right, Kiro bridging the two — that's the exact workflow the article describes, visualized in one frame.
If I wanted the accent color changed, I wouldn't regenerate — I'd edit_image with that interaction ID and say "shift the blue accents to purple." That's the whole point.

Dogfooding is the cheapest credibility there is: no cherry-picked gallery, no "results may vary" fine print — the tool's real output is literally the first thing you saw when you opened this article. If the skill had flubbed the lettering or mangled the layout, you'd be looking at the evidence right now. Instead, the article ships with its own proof baked into the header.

Links

Repo: github.com/xbill9/nb2lite-skill-kiro (Apache-2.0)
Docker image: hub.docker.com/r/xbill9/nb2lite-agent
Interactions API reference: ai.google.dev/api/interactions-api
Model Context Protocol: modelcontextprotocol.io

This is a third-party community project, not affiliated with or endorsed by Anthropic or Google. Bring your own Gemini API key — and remember generations are billable, so draft on low and save high for the money shot.

Teaching Claude Code to Paint: A Stateful Image-Editing Skill Built on Gemini's Interactions API and MCP

xbill — Wed, 22 Jul 2026 22:44:23 +0000

TL;DR: nb2lite-skill-claude wraps Google's gemini-3.1-flash-lite-image model in a tiny FastMCP server and packages it as a Claude Code skill. You type "generate an image of a cyberpunk kitchen" into Claude Code, and it just... does it. Then you say "add a neon RAMEN sign" and it edits the same image without re-prompting the whole scene. Oh, and the cover image of this article? Generated by the thing the article is about — dogfooding all the way down. More on that at the end.

Background: why another image tool?

Google's Nano Banana 2 Lite — the friendly nickname for gemini-3.1-flash-lite-image — takes a different approach. It's a high-efficiency image model with sub-2-second generations, solid text rendering in 25+ languages, and — the headline feature — support for the stateful Interactions API, which lets you iterate on an image across multiple turns while the model keeps the visual context server-side.

This repo glues that capability into Claude Code, so your coding agent can generate and iteratively refine images as a natural part of a session. It ships as two things in one repo:

A Model Context Protocol (MCP) server (nb2lite-agent, a single-file FastMCP app in server.py) exposing exactly four tools.
A Claude Code skill (nb2lite-image) that teaches Claude when and how to use those tools well.

The Interactions API: images with a memory

The Interactions API is Gemini's stateful endpoint. The core loop looks like this:

You call client.interactions.create(...) with a prompt and store=True.
The response includes an interaction_id — a handle to the turn's visual context, persisted on Google's servers.
On the next call, you pass previous_interaction_id, and the model edits the existing canvas — preserving character, style, lighting, and pixel continuity.

So instead of this (stateless suffering):

"A watercolor fox in a forest at dawn, mist, soft light, wearing a red scarf, three birch trees on the left, and now also holding a lantern"

...you write this:

"Add a lantern in its paw."

That's it. The stored context holds the rest.

A few practical details the server handles for you:

Every turn returns a new interaction ID. Chain the latest one; editing from a stale ID silently forks your session from an older state (a subtle and very annoying bug if you roll this by hand).
Aspect ratio is chosen at generation time (1:1, 16:9, 9:16, 4:3, 3:4) and inherited on stateful edits — changing it mid-session degrades pixel continuity, so the edit tool deliberately doesn't accept one.
Thinking levels: low (default, fast drafts) or high (complex rendering, accurate text layout, character composition). The generic API spec also lists minimal and medium, but the live API rejects them for this model with an HTTP 400 — the server saves you from discovering that the hard way.

What is MCP, in one minute

The Model Context Protocol is an open standard for connecting AI assistants to tools and data. Before it, giving a model access to some service meant writing a bespoke integration for each assistant — N assistants × M services, everyone reinventing the same plumbing. MCP collapses that: a tool author writes one MCP server that exposes typed tools, and any MCP-capable client (Claude Code, Claude Desktop, and a growing list of others) can discover and call them with no per-client glue code.

An MCP server is usually a small local process that speaks JSON-RPC over stdio. The client launches it, asks "what tools do you have?", and from then on the model can call them like functions.

The nb2lite-agent server exposes exactly four:

Tool	What it does
`generate_image`	Text → 1k image. Saves locally, returns the path + an interaction ID.
`edit_image`	Stateful edit: takes the previous interaction ID + a description of only the change.
`edit_local_image`	Uploads any local image file inline (base64) and applies an edit — your entry point for existing files.
`get_help`	Reports live config: API key status, active model, output directory, full tool reference.

And what's a Claude Code skill?

If MCP is the hands (the tools Claude can physically call), a skill is the muscle memory — a markdown file (SKILL.md) plus bundled resources that load into Claude's context and teach it the workflow: which tool to reach for, in what order, with which constraints.

For nb2lite-image, the skill encodes things like:

Call get_help first when diagnosing setup issues — if the API key is missing, nothing else will work.
Keep edit prompts incremental: describe the change, not the scene.
Always chain the latest interaction ID.
Generations are billable — batch related edits and prefer thinking_level: low for drafts.

Installing it: the "I just want it to work" edition

You need three things: Python 3.10+, Claude Code, and a Gemini API key (free from Google AI Studio). Pick one of the paths below.

Path A: The plugin marketplace (fewest keystrokes)

Inside Claude Code, type:

/plugin marketplace add xbill9/nb2lite-skill-claude
/plugin install nb2lite-image@nb2lite-skill-claude

This installs the skill and auto-registers the MCP server. The plugin manifest carries no API key (as it should!) — the server reads GEMINI_API_KEY from your environment, so make sure it's exported before launching Claude Code.

Path B: Clone and bootstrap (this repo)

# 1. Get the code
git clone https://github.com/xbill9/nb2lite-skill-claude.git
cd nb2lite-skill-claude

# 2. One-command setup: installs deps, registers the MCP server
#    in .mcp.json, and prompts for your API key (stored in ~/gemini.key)
./init.sh

# 3. Restart Claude Code in this directory and approve the server
#    when prompted. Verify with:
/mcp        # should list nb2lite-agent

That's genuinely it. init.sh is safe to rerun if anything looks off.

Path C: Install into your project

From a clone of the repo:

make init TARGET=/path/to/your/project ARGS='--output-dir ./images'

This copies the skill into <project>/.claude/skills/nb2lite-image/ and writes the nb2lite-agent entry into that project's .mcp.json. It reuses ~/gemini.key if you've set one up. Restart Claude Code in the target project, approve the server, done.

Path D: Docker (nothing on the host but Docker)

The server is published as xbill9/nb2lite-agent:

claude mcp add nb2lite-agent --env GEMINI_API_KEY="$(cat ~/gemini.key)" -- \
  docker run --rm -i -e GEMINI_API_KEY -v "$PWD:$PWD" -w "$PWD" xbill9/nb2lite-agent

Troubleshooting, the whole guide

/mcp doesn't list the server → restart Claude Code in the project directory.
Tools return 🔴 GEMINI_API_KEY is not set → run source set_env.sh (or export the key) and restart.
Anything else → ask Claude to call get_help; it reports the live config.

Examples: a session in practice

Once installed, you talk to it in plain English. A real flow looks like:

You: "Generate a cozy cabin in a snowy forest at dusk, 16:9."

Claude calls:

generate_image(
    prompt="A cozy log cabin in a snowy forest at dusk, warm light in the windows",
    aspect_ratio="16:9",
    thinking_level="low",
)
# 🟢 Saved to: ./gen_1784759001_a1b2c3d4.jpg
# Interaction ID: v1_ChdpRU5...

You: "Nice. Add smoke curling from the chimney."

edit_image(
    previous_interaction_id="v1_ChdpRU5...",
    edit_prompt="add gentle smoke curling from the chimney",
)
# 🟢 Saved to: ./edit_1784759050_e5f6a7b8.jpg
# Interaction ID: v1_Xk9mPq2...   ← a NEW id; the next edit chains this one

You: "Now make it night, with aurora in the sky."

Same tool, newest ID, and the cabin, trees, and chimney smoke all stay put — only the sky changes. No re-prompting, no continuity roulette.

And for images that didn't come from the model at all:

You: "Take ./whiteboard-sketch.png and render it as a clean 3D product mockup."

edit_local_image(
    image_path="./whiteboard-sketch.png",
    edit_prompt="render this hand-drawn sketch as a high-fidelity 3D product mockup",
    aspect_ratio="4:3",
)

It returns an interaction ID too — so follow-up refinements switch to edit_image and go stateful from there.

Dogfooding: about that cover image 🐕🍖

This repo dogfoods itself at every layer:

The skill is active inside its own repository — open Claude Code in a clone and the nb2lite-image skill and nb2lite-agent server are already wired up, so every development session doubles as an integration test.
The integration tests (make test) drive the same four MCP tools an end user would, against the live API.
And now, the cover image of this article was generated by the exact skill the article describes, from inside a Claude Code session in this repo. One tool call, first attempt, no retouching:

generate_image(
    prompt="A wide tech blog cover illustration: a friendly robot artist "
           "painting a glowing galaxy on an easel, while a chain of connected "
           "frames behind it shows the same picture evolving step by step "
           "(day sky, then sunset, then storm with lightning). Flat vector "
           "style, deep indigo background, neon cyan and orange accents. "
           "Title text 'NB2Lite + MCP', subtitle 'Stateful image editing "
           "as a Claude Code skill'. Crisp, accurate lettering.",
    aspect_ratio="16:9",
    thinking_level="high",
)
# 🟢 Image successfully saved!
# • Saved to: gen_1784759177_cbab8b65.jpg
# • Interaction ID: v1_ChdpRU5hb2o3SWMzV2pNY1AtUFgy...

(That exact output is committed to the repo as devto-cover.jpg, receipts and all.)

Worth noticing:

The text rendered correctly. "NB2Lite + MCP" and the full subtitle came out crisp and typo-free — that's what thinking_level: "high" buys you on text-heavy layouts.
The model illustrated its own pitch. The chain of frames (day → sunset → storm → galaxy) is the stateful edit loop — the image explains the Interactions API better than a diagram I'd have drawn by hand.
If I wanted the accent color changed, I wouldn't regenerate — I'd edit_image with that interaction ID and say "make the orange accents magenta." That's the whole point.

Gemma 4 E2B on a Single TPU v6e Chip: A Serving Deep Dive

xbill — Tue, 21 Jul 2026 03:48:40 +0000

Measured 2026-07-21 on vllm/vllm-tpu:nightly (vLLM 0.23.1rc1.dev1076), a GCE flex-start ct6e-standard-1t (one TPU v6e chip, 32 GB HBM) in europe-west4-a.

TL;DR

The plain google/gemma-4-E2B-it serves well on one v6e chip; none of its QAT siblings load at all. The 2-billion-parameter "efficient" Gemma 4 sustains 213 tok/s for a single user with a 16 ms first token, scales to ~2,200 output tok/s across concurrent streams, handles OpenAI-style function calling — including parallel calls and refusal to hallucinate calls — without a miss, and answers simple vision questions accurately in ~200 ms.

The QAT variants are a different story: the int4 compressed-tensors export hits an unimplemented quantization path, and the bf16 QAT export trips a loader bug — the Gemma 4 implementation demands per-layer norms that E2B's KV-sharing architecture legitimately doesn't have. Filed upstream as tpu-inference #3225.

One capability coupling to know about: with a reasoning parser configured, schema enforcement only engages when thinking is enabled — thinking-off requests sail through unconstrained with a 200 status. Config interaction, not TPU limitation; details below.

1. Getting to a serving endpoint

The host is a GCE flex-start VM — capacity granted on request, billed until deleted, hard-stopped at a 4-hour max run, $1.35/chip-hour. A startup script installs Docker, pulls vllm/vllm-tpu:nightly, fetches the Hugging Face token from Secret Manager via the metadata server, and launches vLLM.

Boot timeline: VM RUNNING at t+0 (200 GB boot disk — the 10 GB default cannot hold the vLLM image) → Docker installed ~t+1:00 → image pulled ~t+6:00 → weights downloaded, XLA compiled, health green ~t+8:30.

Two environment quirks worth knowing:

Direct SSH silently times out on some networks even when the VPC allows tcp:22 — the block is upstream of the VPC. IAP tunneling (gcloud compute ssh --tunnel-through-iap) rides over HTTPS and works; so does tunneling the API port with gcloud compute start-iap-tunnel <vm> 8000.
vLLM auto-selects an fp8_e5m2 KV cache on v6e — the largest memory consumer is 8-bit before any weight quantization enters the picture.

Serving flags: --max-model-len 65536 --gpu-memory-utilization 0.9 --max_num_batched_tokens 4096 --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4, bf16 weights, tensor-parallel 1.

2. Three ways the QAT checkpoints fail to load

Checkpoint / path	Failure	Verdict
`-qat-w4a16-ct` · JAX	int4 compressed-tensors scheme unimplemented for E2B's `per_layer_model_projection`	✕ no load
`-qat-q4_0-unquantized` · JAX	`k_norm.weight` "missing" for layers 15–34	✕ no load
`-qat-q4_0-unquantized` · torchax	identical missing-weights error via `MODEL_IMPL_TYPE=vllm`	✕ no load
`gemma-4-E2B-it` (plain) · JAX	loads and serves	✓ serves

The forensics point at the loader, not the checkpoint. Reading the safetensors headers of both repos: the plain export ships self_attn.k_norm for all 35 layers; the QAT export ships it only for the 15 non-KV-shared layers. Both configs are identical — including num_kv_shared_layers: 20. Layers 15–34 reuse K/V from lower layers and have no k-norm of their own, so the QAT export is the architecturally honest one; the plain checkpoint only loads because it carries those (unused) tensors anyway. Proposed fix in #3225: skip instantiating K/V-side parameters for KV-shared layers.

Until it lands: serve the plain checkpoint. At 2B parameters (~5 GB bf16 against 32 GB HBM), 4-bit weights buy little here anyway — memory pressure lives in the KV cache, which is already fp8.

3. What one chip is worth: the concurrency sweep

Same workload at every level — 1,024-token prompts, 128-token completions, vllm bench serve, random dataset.

Concurrency	Req/s	Output tok/s	Total tok/s	TTFT med	TTFT p99	TPOT med	Per-stream tok/s
1	1.64	209	1,884	16 ms	17 ms	4.7 ms	213
8	9.44	1,209	10,878	27 ms	99 ms	6.2 ms	161
32	12.78	1,636	14,721	155 ms	189 ms	17.5 ms	57
64	16.72	2,140	19,262	122 ms	349 ms	25.3 ms	39
100 (burst)	17.31	2,215	19,938	833 ms	1,573 ms	36.8 ms	27

Reading the curve:

Prefill is effectively free at low load. A 1,024-token prompt reaches first token in 16 ms — roughly 64K prefill tok/s for a single stream.
c=8 is nearly free concurrency: six times the throughput of a single stream for +1.5 ms/token; each of 8 users still sees ~160 tok/s.
The knee is between 32 and 64. Going 32→64 buys +31% throughput for +45% per-token latency; 64→burst buys +3.5% for another +45%.
Capacity-planning number: run at ≤64 concurrent streams for smooth per-user experience; the ceiling is ~17 req/s at this workload shape.

(Single run per configuration — more meaningful here than usual: a kernel study on this same stack, cited in §7, measured run-to-run cv ≤ 0.3% under greedy decoding with static shapes.)

4. Function calling at 2B scale

Served with --tool-call-parser gemma4 --enable-auto-tool-choice, probed with two OpenAI-style tools at temperature 0. Five scenarios, five clean results:

Probe	Behavior	Latency
Simple call	correct tool, inferred the optional `unit` arg from phrasing	166 ms
Result synthesis	fed the tool result back → clean natural-language answer	140 ms
No-tool restraint	answered directly, no spurious call	97 ms
Parallel calls	both tool calls emitted in one turn, correct args each	150 ms
Underspecified	asked "What city are you interested in?" instead of hallucinating a call	44 ms

A 2B model producing well-formed tool_calls JSON, choosing correctly between calling and answering, batching parallel calls, and asking for missing arguments — at double-digit-millisecond latency. For high-volume, low-complexity agent steps, the quality floor is higher than the parameter count suggests.

5. Structured output works — but only with thinking on

Probe	Observed behavior	Verdict
`json_schema`, thinking off (default)	free prose with a 200 status; `strict: true`, `guided_json`, `structured_outputs` spellings all equally unenforced	✕ silently skipped
`json_object`, thinking off	fenced code block, array where an object was asked, invented enum value	± prompt-level
`json_schema` + `enable_thinking`	exact schema conformance — bare JSON object, typed integer, "ASAP" correctly mapped into the `high` enum	✓ enforced
Reasoning, default	no reasoning traces on any prompt	off by default
Reasoning, `enable_thinking`	parser cleanly splits thinking trace from a terse answer; ~2.4× completion tokens	✓ works

The mechanism: with --reasoning-parser gemma4 configured, vLLM defers grammar enforcement until the reasoning section ends. Thinking off → no reasoning terminator → the grammar never engages, and unconstrained prose ships with a 200 status. Enable thinking ("chat_template_kwargs": {"enable_thinking": true}) and the same request is enforced exactly.

Operational guidance: pair structured output with enable_thinking: true under a reasoning parser — or drop --reasoning-parser from servers that don't need it. And keep client-side validation regardless: the silent-skip failure mode means a trusting client can't tell an enforced response from a lucky one.

6. Vision at 2B: accurate and nearly free

One restart with --limit-mm-per-prompt '{"image":4,"audio":1}' makes it a vision server. COCO validation images, base64 data URIs, temperature 0:

Probe	Answer (abridged)	Latency
Describe (two cats)	"Two tabby cats… on a bright pink surface… a remote control visible"	197 ms
Count + attributes	2 animals, both cats, remote + blanket identified	421 ms
Scene (bear)	"A bear lying down in a grassy outdoor environment"	329 ms
Room inventory	wall-mounted TV, shelving, furniture correctly enumerated	874 ms

An image costs ~280 prompt tokens and adds almost nothing over a text request. Two notes: server-side fetching of external image URLs proved flaky (intermittent 422s) — base64 data URIs are the reliable path — and the first multimodal request after boot can 422 while the processor warms; retry once.

7. fp8 KV cache, HBM anatomy, and a related kernel result

fp8 vs bf16 KV: six greedy prompts (explanation, code, listing, translation, arithmetic, summarization — 889 completion tokens) run under the default fp8_e5m2 cache, then re-run after a restart with --kv-cache-dtype bfloat16. Result: 6 of 6 outputs byte-identical. On this (small, greedy) sample the compression is genuinely free — take the fp8 cache.

Where the 32 GB goes (bf16-KV boot, 65,536 max context): usable HBM reports as 31.24 GiB; at 0.9 utilization vLLM works within 28.12 GiB — roughly 5.5–6 GiB weights, 16.3 GiB KV cache (8,713 blocks × 128 tokens × 15 layers × 128 KiB), ~6 GiB workspace.

The interesting physics: E2B's KV sharing is directly visible in the allocator — only 15 of 35 layers hold KV tensors, each with a single 256-dim KV head, so a token costs ~15 KiB of KV in bf16 (~7.5 KiB under fp8). That yields ~1.1 million tokens of resident KV — seventeen full 65K-context conversations on-chip — which is why the sweep saturated on compute, never memory. Also: of the 404-second engine init, 329 seconds is XLA compilation — the dominant term in the ~10-minute cold start.

Related work: a kernel-substitution study by Zimbres (DOI 10.5281/zenodo.21404069) shows the RPA v3 kernel's decode block-size heuristic costs 27.7–68.7% of large-batch throughput on 27B/31B models on v6e. E2B sits at the low-exposure end of that effect: at our c=64 operating point, attention is ~9% of memory traffic (vs ~41% in their regime) precisely because of the KV-sharing design above. Testing the override on E2B — smallest model, one chip, KV sharing — is a natural follow-up.

8. Cost breakdown

Rate verified against Google's published Dynamic Workload Scheduler pricing: $1.35 per chip-hour for v6e flex-start (europe-west4, us-east1, us-east5, asia-northeast1).

Operating point	Output tok/s	$ / M output tokens
Saturation (burst)	2,215	$0.17
Sweet spot (c=64)	2,140	$0.18
Interactive (c=8)	1,209	$0.31
Single stream	209	$1.79

Breakeven vs. the API: against Gemini 2.5 Flash-Lite at $0.40/M output tokens, self-hosting wins once you sustain ~940 output tok/s — roughly the c=8 operating point held continuously.
Cold start (~8.5–10.5 min, mostly XLA compile) costs $0.19–0.24 per provisioning.
A full 4-hour session costs $5.40 and, at saturation, delivers ~30M output tokens — about $12 worth at Flash-Lite prices. Flex-start fits batch bursts, not idle always-on endpoints.

9. Reproduction

# Serve (on the TPU VM)
docker run --name vllm-gemma4 --privileged --net=host -d \
  -v /dev/shm:/dev/shm --shm-size 10gb -e HF_HOME=/dev/shm \
  -e HF_TOKEN=$(gcloud secrets versions access latest --secret=hf-token) \
  vllm/vllm-tpu:nightly vllm serve google/gemma-4-E2B-it \
  --tensor-parallel-size 1 --max-model-len 65536 \
  --gpu-memory-utilization 0.9 --max_num_batched_tokens 4096 \
  --disable_chunked_mm_input --enable-auto-tool-choice \
  --tool-call-parser gemma4 --reasoning-parser gemma4 \
  --limit-mm-per-prompt '{"image":4,"audio":1}'   # {"image":0,"audio":0} for text-only
# KV comparison: add --kv-cache-dtype bfloat16 (default is fp8_e5m2 on v6e)

# Benchmark (per concurrency level C)
vllm bench serve --backend vllm --model google/gemma-4-E2B-it \
  --dataset-name random --num-prompts 100 \
  --random-input-len 1024 --random-output-len 128 --max-concurrency C

# Reach the endpoint from a network that blocks direct traffic
gcloud compute start-iap-tunnel vllm-gemma4-e2b 8000 \
  --local-host-port=localhost:8000 --zone=europe-west4-a

Environment: vLLM 0.23.1rc1.dev1076+g5c342876a (vllm-tpu:nightly, tpu-inference backend) · TPU v6e-1 (ct6e-standard-1t, 32 GB HBM, GCE flex-start) · bf16 weights, fp8_e5m2 KV cache. Single benchmark run per configuration; treat deltas under ~10% as noise.

tpu-management: a Claude Code skill for running Gemma 4 on Cloud TPUs

xbill — Tue, 21 Jul 2026 00:10:45 +0000

Getting a Gemma 4 model serving on a Cloud TPU involves a surprising amount of ceremony: finding a zone that actually has v6e capacity, requesting flex-start VMs, sizing boot disks, wiring Hugging Face tokens through Secret Manager, picking the right vLLM flags for TPU, and remembering to tear everything down before the billing meter embarrasses you.

I packaged all of that into tpu-management — a Claude Code skill plus an MCP server (tpu-devops) — so you can just ask Claude to do it.

What it does

The skill teaches Claude the full TPU serving lifecycle, and the MCP server gives it ~40 tools to execute it:

Find capacity: sweep every zone with quota for an available TPU (find_tpu), check dimensioned quotas (get_zones_with_available_quota), estimate cost before creating anything.
Provision: create flex-start TPU VMs (v6e/v5p) or legacy queued resources (v5e), with a startup script that installs Docker, pulls vllm/vllm-tpu:nightly, fetches the HF token from Secret Manager, and serves the model automatically.
Serve & verify: manage the vLLM container, watch boot via the serial console, verify model health, and query the deployed Gemma 4 endpoint directly.
Diagnose: pull vLLM/docker/system/Cloud Logging logs — including Gemma-4-powered log triage (analyze_cloud_logging), where the self-hosted model analyzes its own infrastructure's logs.
Benchmark: run vllm bench serve against the deployment and get TTFT/throughput/P95 numbers back.
Tear down: destroy resources safely, with guardrails (the skill insists on confirming teardown, and warns that flex-start bills until deletion).

It also encodes the hard-won details you'd otherwise learn at 2 a.m.: the default 10 GB boot disk that the vLLM TPU image overflows, the TPUS_PER_TPU_FAMILY quota that's invisible to gcloud regions describe, why you watch the serial console instead of SSH, and the exact vLLM flags Gemma 4 needs on TPU (--tool-call-parser gemma4, --disable_chunked_mm_input, and friends).

Install

The fastest path is the Claude Code plugin marketplace — one command gets you the skill and the MCP server:

/plugin marketplace add xbill9/tpu-skill-claude
/plugin install tpu-management@tpu-skill-claude

Prefer the classic routes? Clone and install:

git clone https://github.com/xbill9/tpu-skill-claude
cd tpu-skill-claude
make skill-install          # skill only, all projects
./project-setup.sh --global # skill + MCP server registration

Or skip the clone entirely and unzip the packaged skill:

curl -L -o /tmp/tpu-management-skill.zip \
  https://github.com/xbill9/tpu-skill-claude/raw/main/dist/tpu-management-skill.zip
mkdir -p ~/.claude/skills && unzip -o /tmp/tpu-management-skill.zip -d ~/.claude/skills/

You'll need an authenticated gcloud CLI, the TPU API enabled, and a Hugging Face token (the agent's save_hf_token tool stores it in Secret Manager for you). Configuration is all environment variables: GOOGLE_CLOUD_PROJECT, MODEL_NAME, ACCELERATOR_TYPE, TENSOR_PARALLEL_SIZE.

What it looks like in practice

Once installed, you talk to Claude Code like this:

"Find me a v6e-1 with available quota and serve gemma-4-12B on it"

Claude checks existing resources first (the skill's rule: never create before checking what exists), sweeps zones for capacity, creates the VM with the proven flags, follows the serial console until vLLM logs "Application startup complete," and hands you a working OpenAI-compatible endpoint. Then:

"Benchmark it and tear it down when done"

The skill ships with a full reference guide (references/tpu-guide.md) : flex-start zones per TPU family, quota metrics and how to request increases, and troubleshooting — so Claude can answer capacity and billing questions even without the MCP server connected.

Under the hood

The repo doubles as both a Claude Code plugin marketplace and a plain skill repo. A Makefile keeps the three distribution formats (project skill, plugin layout, zip) in sync from one set of sources, and the MCP server is a single-file FastMCP app — easy to read, easy to fork for your own model/accelerator combos.

Issues and PRs welcome: https://github.com/xbill9/tpu-skill-claude

Gemma4 DevOps In Action

xbill — Mon, 20 Jul 2026 14:20:31 +0000

In the last entry I got Gemma-4's 128-expert MoE running on an inf2.24xlarge and signed off with a
cliffhanger: fitting it on a 2-core box "needs fp4 — a separate expedition." This is that expedition.
It ended nowhere near where I thought it would: not with fp4, and not on the 8xlarge I was
aiming for, but on the smallest, cheapest Inferentia2 instance AWS sells — a single inf2.xlarge
with 16 GB of host RAM — running a 26B-parameter model. Here's the refinement trail, dead ends
included.

The finished port (see the previous article) served the 26B-A4B on a 24xlarge: 12 NeuronCores, 192 GB
HBM, ~$6.49/hr on-demand. It worked, it was correct, and it was overkill for one user asking one
question at a time. The whole point of a "4B-active" MoE is that it's cheap to run; it shouldn't need
a datacenter-class box to host. The goal: get it onto the bottom-tier inf2.xlarge — 2 cores, 32 GB
HBM, and only 16 GB of host RAM — at ~$0.76/hr. That's 8.6× cheaper.

	24xlarge (shipped)	inf2.xlarge (goal)
NeuronCores	12 (used 8)	2
HBM	192 GB	32 GB
Host RAM	768 GB	16 GB
On-demand	~$6.49/hr	~$0.76/hr

Two walls stood in the way: the model doesn't fit 32 GB of HBM at bf16, and it doesn't fit 16 GB of host
RAM the normal way. I had to break both.

Wall 1: the memory math, and why "A4B" doesn't save you

The experts are ~93% of the weight. At bf16 the 128 experts are ~45.6 GB — replicate that across TP=2
and you're at ~22.8 GB per core against a 16 GB budget. Top-8 routing reduces compute, not
footprint: all 128 experts must be resident. So the experts alone blow the box.

My first move was the obvious one — int8 the experts. NxD has QuantizedColumnParallel /
QuantizedRowParallel, and (from the prior article) int8-per-channel on this model is numerically
perfect: token-for-token identical to fp32. Experts drop to ~11.4 GB/rank. Should fit.

It didn't. Could not load the model status=4 Allocation Failure — ~3–4 GB over the 16 GB core. And
inf2 has no 4-core SKU to split the difference. My conclusion at the time, which I even wrote down: to
close that last 3–4 GB you need fp4 experts.

The fp4 rabbit hole (two of them, both dead ends)

I chased fp4 twice.

First, QuantizedColumnParallel advertises F4E2M1FN_X4. But its state_dict weight is a packed
uint16 microscaling format (32 fp4 → 8 uint16) produced by from_float/preshard_hook — nothing
like int8's five-line quantizer. Then I read the AWS docs plainly: NxD Inference quantization supports
INT8 and FP8 only. FP8 is 8-bit — same size as int8, so it doesn't help. QuantizedDtype.F4E2M1FN_X4
exists as a constant but isn't a wired-up inference path.

Second, on the newest Neuron 2.31 stack (neuronx-cc 2.26, NxD 0.19) fp4 looked less absent —
there's BLOCKWISE_SYMMETRIC, block_size, from_float. So I tried again, and hit the floor twice:

The clean QuantizedColumnParallel fp4 forward calls blockwise_scale_dequantize, which is a CPU-only torch reference (assert device == cpu) — it can't trace to a NeuronCore.
The real device path wants the blockwise_mm NKI kernel, which NxD imports from neuronxcc.nki._private.blockwise_mm — but 2.26 ships it under _pre_prod_kernels. Mispathed, pre-prod, import fails.
Swapping my DenseExperts for NxD's ExpertMLPs (which has fp4 packing) didn't help either: its blockwise NKI path kernel_asserts, and its use_torch_block_wise=True path crashes torch-xla's shape inference on a data-dependent index op. MoE routing that won't trace — the exact problem the dense trick solved for compute, now back for quantization.

fp4 on inf2 is a wall in this SDK. I'd spent real time being sure it was the only door. It wasn't
even a door.

The breakthrough: I was quantizing the wrong thing

Here's the reframe that unlocked it. I'd been staring at the experts because they're huge. But the
experts were already int8 when I OOM'd. The 3–4 GB I needed wasn't in the experts — it was in the
weights I'd left in bf16. Chiefly one: the tied lm_head, replicated at 1.48 GB per rank.

The fix wasn't smaller experts. It was an all-int8 squeeze of everything still fat:

int8 and shard the lm_head. QuantizedColumnParallel(gather_output=True) slices the 262k-row head across ranks and quantizes it: 1.48 GB → ~0.37 GB/rank. This was the single biggest lever, and it had nothing to do with the experts.
int8 the shared dense MLP (the other half of the dual-path FFN): another ~0.27 GB/rank.
Keep the int8 experts.

Per-rank resident dropped to ~12–13 GB — comfortably under 16 GB. And the output was still right:

Compiler status PASS · MB_TRACED
DEVICE_PARIS True
DEV GEN: 'The capital of France is Paris.'
prefill 232 ms · neff 26.7 GB

The 26B MoE ran on a 2-core inf2.8xlarge. The "not achievable without fp4" conclusion was right
that fp4 was out — and wrong about everything else.

Wall 2: the 8xlarge and the xlarge have the same HBM — but not the same host RAM

Here's the subtlety that trips everyone. inf2.8xlarge and inf2.xlarge have the same 2 NeuronCores
and 32 GB HBM. The difference is host RAM: 128 GB vs 16 GB. My squeeze fit the device. Would
it fit the host on the cheapest box?

Compiling won't: the ModelBuilder trace holds an fp32 discovery model + a structure model + an fp32
checkpoint dict — ~180 GB peak. That OOMs even the 8xlarge's 128 GB (I added a 100 GB swapfile to get
the compile through, ~40 min). But compiling and deploying are different jobs. Deploying a saved
neff doesn't need the model in host RAM at all — the transformer weights live on the device.

So the xlarge deploy is deliberately slim (deploy_sqz.py / optb_server_sqz.py):

Structure (which layers are KV-shared, head dims, sliding window) comes from a meta-device model instantiation — accelerate.init_empty_weights(), ~0 host RAM.
Word embeddings come from a standalone 1.48 GB embed_tokens.pt table, loaded once on the host.
The neff carries every transformer weight on the device (torch.jit.load + one initialize_with_saved_weights).
A 40 GB swapfile absorbs the one-time neff-load peak.

Result on a stock 16 GB inf2.xlarge:

NEFF_LOADED in 112s
DEV GEN: 'The capital of France is Paris.'
PREFILL 250 ms

Host RSS peaked ~11 GB; swap was barely touched. Compile needs 180 GB; deploy needs a laptop's worth.
A 26B model, on a box you'd hesitate to run a 7B on.

The bug that turned Paris into a wall of spaces

Except the first slim run didn't say Paris. It said:

DEV GEN: '                                                              '
DEVICE_PARIS False

The neff was proven (it made Paris at compile). So the deploy was feeding it something wrong. Repeated
identical tokens is the fingerprint of degenerate, near-uniform logits — the classic symptom of
embeddings at the wrong magnitude.

Gemma-4 doesn't use a plain embedding. It uses Gemma4TextScaledWordEmbedding, whose forward is:

return super().forward(input_ids) * self.embed_scale   # embed_scale = hidden_size ** 0.5

The ×√hidden_size (= √2816 ≈ 53) happens inside the embedding, and the model forward does not
re-normalize — it just does hidden_states = inputs_embeds. The neff was traced on scaled embeds. My
slim host path used a plain torch.nn.Embedding and fed unscaled ones — 53× too small. Everything
downstream washed out to noise.

One line:

ie = emb(ids) * (hidden_size ** 0.5)   # match Gemma4TextScaledWordEmbedding

DEV GEN: 'The capital of France is Paris.'. This is the same trap that bit the E2B slim server; when you
move embedding lookup to the host, you inherit whatever the model's embedding class was doing.

Shipping it: 512/128, published, and a clean-pull proof

I recompiled a production 512/128 bucket build (512 total / 128 prompt tokens) — same recipe, ~40 min
on an 8xlarge with swap, neff 26.8 GB, re-validated on the xlarge (loads in 112 s, ~10 GB host peak,
Paris). Then I wrapped it in a slim OpenAI-compatible server and published:

Docker Hub xbill9/gemma4-optb-26b:xlarge — a thin image; the entrypoint pulls the ~28 GB of artifacts from the HF repo on first start.
HF xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (public) — neff + embedding table + config + server + Dockerfile.

The final proof was the one that matters to anyone but me: a cold, clean pull. Fresh spot
inf2.xlarge, add swap, docker run, walk away.

READY in 490.9s — slim int8-squeeze, ModelBuilder TP=2, MAX=512 BUCKET=128
{"role":"assistant","content":"The capital of France is Paris."}
{"response":"AWS Inferentia is a purpose-built machine learning accelerator
             designed to provide high-performance, low-cost inference..."}

(No HF_TOKEN needed — the repo is public. The 490 s is one-time cold-EBS neff load, then it serves.)
Terminated the box; net ongoing spend, zero.

The honest caveat

Decode is ~6 tok/s. That's the price of the DenseExperts trick from the last article: computing all
128 experts every token instead of the sparse top-8 — ~16× the necessary expert FLOPs, on 2 cores. It's
correct and it's cheap to host; it is not fast. Making sparse routing actually trace under
ModelBuilder is the real next expedition, and (see the fp4 detour) "the vendor primitive exists" is not
the same as "it traces." For single-user, cost-sensitive, latency-tolerant serving of a 26B on a $0.76/hr
box, the tradeoff is the right one. For a chatbot under load, it isn't yet.

Takeaways

When something's 3 GB over, don't assume it's the obvious 3 GB. I burned days on fp4 experts; the fix was quantizing the replicated lm_head I'd never looked at. Profile the residency, don't eyeball the architecture.
"Not supported" beats "not documented well." fp4 has constants and partial code paths in NxD 0.19 that make it look one commit away. It isn't — the kernels are CPU-only refs or pre-prod. Read the feature guide's supported list and believe it.
Compile-fit and deploy-fit are different budgets. A model that needs 180 GB of host RAM to trace can deploy on 16 GB, because the weights live on the device. Slim host loading + swap is what turns a 24xlarge model into an xlarge product.
Moving embedding lookup to the host means inheriting the embedding class. Gemma's ×√H scale is inside Gemma4TextScaledWordEmbedding; miss it and you get a confident wall of whitespace.
Ship the clean-pull proof. "Works on my validated box" isn't a deliverable; "docker run on a fresh spot instance says Paris" is.

Artifacts

Docker Hub: xbill9/gemma4-optb-26b:xlarge
HF: xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (public)
Recipe: tp_mb_moe_sqz.py (all-int8 squeeze: int8 experts + sharded int8 lm_head + int8 shared MLP) · deploy_sqz.py / optb_server_sqz.py (slim host-embedding deploy)

A 26B mixture-of-experts, on the cheapest accelerator instance AWS rents, for the price of a large coffee
per day. Written with AI assistance in a Claude Code session; every log line quoted is from a real run on
real hardware.

Inferntia2 DevOps in Action

xbill — Mon, 20 Jul 2026 14:17:17 +0000

In the last entry I got Gemma-4's 128-expert MoE running on an inf2.24xlarge and signed off with a
cliffhanger: fitting it on a 2-core box "needs fp4 — a separate expedition." This is that expedition.
It ended nowhere near where I thought it would: not with fp4, and not on the 8xlarge I was
aiming for, but on the smallest, cheapest Inferentia2 instance AWS sells — a single inf2.xlarge
with 16 GB of host RAM — running a 26B-parameter model. Here's the refinement trail, dead ends
included.

	24xlarge (shipped)	inf2.xlarge (goal)
NeuronCores	12 (used 8)	2
HBM	192 GB	32 GB
Host RAM	768 GB	16 GB
On-demand	~$6.49/hr	~$0.76/hr

Two walls stood in the way: the model doesn't fit 32 GB of HBM at bf16, and it doesn't fit 16 GB of host
RAM the normal way. I had to break both.

Wall 1: the memory math, and why "A4B" doesn't save you

The fp4 rabbit hole (two of them, both dead ends)

I chased fp4 twice.

The clean QuantizedColumnParallel fp4 forward calls blockwise_scale_dequantize, which is a CPU-only torch reference (assert device == cpu) — it can't trace to a NeuronCore.
The real device path wants the blockwise_mm NKI kernel, which NxD imports from neuronxcc.nki._private.blockwise_mm — but 2.26 ships it under _pre_prod_kernels. Mispathed, pre-prod, import fails.
Swapping my DenseExperts for NxD's ExpertMLPs (which has fp4 packing) didn't help either: its blockwise NKI path kernel_asserts, and its use_torch_block_wise=True path crashes torch-xla's shape inference on a data-dependent index op. MoE routing that won't trace — the exact problem the dense trick solved for compute, now back for quantization.

fp4 on inf2 is a wall in this SDK. I'd spent real time being sure it was the only door. It wasn't
even a door.

The breakthrough: I was quantizing the wrong thing

The fix wasn't smaller experts. It was an all-int8 squeeze of everything still fat:

int8 and shard the lm_head. QuantizedColumnParallel(gather_output=True) slices the 262k-row head across ranks and quantizes it: 1.48 GB → ~0.37 GB/rank. This was the single biggest lever, and it had nothing to do with the experts.
int8 the shared dense MLP (the other half of the dual-path FFN): another ~0.27 GB/rank.
Keep the int8 experts.

Per-rank resident dropped to ~12–13 GB — comfortably under 16 GB. And the output was still right:

Compiler status PASS · MB_TRACED
DEVICE_PARIS True
DEV GEN: 'The capital of France is Paris.'
prefill 232 ms · neff 26.7 GB

The 26B MoE ran on a 2-core inf2.8xlarge. The "not achievable without fp4" conclusion was right
that fp4 was out — and wrong about everything else.

Wall 2: the 8xlarge and the xlarge have the same HBM — but not the same host RAM

So the xlarge deploy is deliberately slim (deploy_sqz.py / optb_server_sqz.py):

Structure (which layers are KV-shared, head dims, sliding window) comes from a meta-device model instantiation — accelerate.init_empty_weights(), ~0 host RAM.
Word embeddings come from a standalone 1.48 GB embed_tokens.pt table, loaded once on the host.
The neff carries every transformer weight on the device (torch.jit.load + one initialize_with_saved_weights).
A 40 GB swapfile absorbs the one-time neff-load peak.

Result on a stock 16 GB inf2.xlarge:

NEFF_LOADED in 112s
DEV GEN: 'The capital of France is Paris.'
PREFILL 250 ms

Host RSS peaked ~11 GB; swap was barely touched. Compile needs 180 GB; deploy needs a laptop's worth.
A 26B model, on a box you'd hesitate to run a 7B on.

The bug that turned Paris into a wall of spaces

Except the first slim run didn't say Paris. It said:

DEV GEN: '                                                              '
DEVICE_PARIS False

Gemma-4 doesn't use a plain embedding. It uses Gemma4TextScaledWordEmbedding, whose forward is:

return super().forward(input_ids) * self.embed_scale   # embed_scale = hidden_size ** 0.5

One line:

ie = emb(ids) * (hidden_size ** 0.5)   # match Gemma4TextScaledWordEmbedding

Shipping it: 512/128, published, and a clean-pull proof

Docker Hub xbill9/gemma4-optb-26b:xlarge — a thin image; the entrypoint pulls the ~28 GB of artifacts from the HF repo on first start.
HF xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (public) — neff + embedding table + config + server + Dockerfile.

The final proof was the one that matters to anyone but me: a cold, clean pull. Fresh spot
inf2.xlarge, add swap, docker run, walk away.

READY in 490.9s — slim int8-squeeze, ModelBuilder TP=2, MAX=512 BUCKET=128
{"role":"assistant","content":"The capital of France is Paris."}
{"response":"AWS Inferentia is a purpose-built machine learning accelerator
             designed to provide high-performance, low-cost inference..."}

(No HF_TOKEN needed — the repo is public. The 490 s is one-time cold-EBS neff load, then it serves.)
Terminated the box; net ongoing spend, zero.

The honest caveat

Takeaways

When something's 3 GB over, don't assume it's the obvious 3 GB. I burned days on fp4 experts; the fix was quantizing the replicated lm_head I'd never looked at. Profile the residency, don't eyeball the architecture.
"Not supported" beats "not documented well." fp4 has constants and partial code paths in NxD 0.19 that make it look one commit away. It isn't — the kernels are CPU-only refs or pre-prod. Read the feature guide's supported list and believe it.
Compile-fit and deploy-fit are different budgets. A model that needs 180 GB of host RAM to trace can deploy on 16 GB, because the weights live on the device. Slim host loading + swap is what turns a 24xlarge model into an xlarge product.
Moving embedding lookup to the host means inheriting the embedding class. Gemma's ×√H scale is inside Gemma4TextScaledWordEmbedding; miss it and you get a confident wall of whitespace.
Ship the clean-pull proof. "Works on my validated box" isn't a deliverable; "docker run on a fresh spot instance says Paris" is.

Artifacts

Docker Hub: xbill9/gemma4-optb-26b:xlarge
HF: xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (public)
Recipe: tp_mb_moe_sqz.py (all-int8 squeeze: int8 experts + sharded int8 lm_head + int8 shared MLP) · deploy_sqz.py / optb_server_sqz.py (slim host-embedding deploy)

Five Gemma-4 models, one accelerator: what porting E2B 31B to AWS Inferentia2 taught me

xbill — Fri, 17 Jul 2026 13:39:29 +0000

I ported the whole Gemma-4 family — E2B, E4B, 12B, 31B, and the 26B-A4B MoE — to run on AWS
Inferentia2. Each has its own write-up in this series; this is the map. What's shared, what's different,
how the recipe evolved from "trace it and pray" to a single-rank-compile pipeline that carries a
30-billion-parameter dense model and a 128-expert MoE — and the one bug that appears, in some costume, in
**every* port.*

The whole family at a glance

	E2B	E4B	12B	31B	26B-A4B
Params (total / active)	~5B / 2B	~8B / 4B	12B	31B	26B / ~4B
Dense or sparse	dense	dense	dense	dense	MoE (128 exp, top-8)
Per-Layer Embeddings	✅	✅	—	—	—
Cross-layer KV-sharing	✅	✅	—	—	—
MatFormer ("effective" size)	✅	✅	—	—	—
Model class	`Gemma4ForConditionalGeneration`	same	`…Unified` (encoder-free)	`Gemma4ForConditionalGeneration`	same
Hardware	1 core (`inf2.xlarge`)	`inf2.8xlarge`	`inf2.8xlarge`	`inf2.24xlarge`	`inf2.24xlarge`
Tensor parallelism	TP=1	TP=2	TP=2	TP=8	TP=8
Recipe	`torch_neuronx.trace`	`tp_alias`	ModelBuilder	ModelBuilder	ModelBuilder + MoE
Correctness	SEQ_MATCH ✅	SEQ_MATCH ✅	SEQ_MATCH ✅	SEQ_MATCH ✅	SEQ_MATCH ✅
Speed	~44 tok/s	~34 tok/s	prefill ~101 ms	prefill ~115 ms	prefill ~77 ms

All five decode token-for-token identically to their CPU fp32 reference, and all five answer "The
capital of France is Paris." That equality bar — not "looks fluent" — is the thread that ties the
series together, for a reason I'll come back to.

What every port shares

1. The vendor stack can't do any of them. optimum-neuron and the Neuron vLLM backend have no
Gemma-4 model class at all, and their graph builders can't express the architecture (KV-sharing on the
E-family, MoE on the 26B). The public endpoint I started from loaded something and served
fluent-looking gibberish. Every one of these is a hand-rolled port of transformers-5.13's own forward.

2. The same skeleton underneath. Softcap 30, a 262,144-token vocabulary, tied embeddings, GQA,
and interleaved sliding + global attention show up across the family. So does the serving shape:
host-side word embeddings, a device prefill + decode with an on-device KV cache, greedy
decode compared against a host fp32 reference.

3. The same environment quirks. The aws_neuronx_venv_pytorch_2_8_nxd_inference venv +
pip install transformers==5.13.0, a transformers.utils.fx shim so neuronx_distributed imports under
the new transformers, and export PATH=$VENV/bin:/opt/aws/neuron/bin (that PATH bit fixes a
libneuronpjrt-path import error that will waste your afternoon).

4. Softcap is a host-only concern. 30·tanh(x/30) is monotonic, so it never changes the greedy
argmax. Every port drops it from the device graph (no tanh over 262K logits) and re-applies it
host-side only when sampling needs real probabilities.

What makes each one different

E2B / E4B — the "effective" models. These are the hard-architecture ports. Per-Layer
Embeddings (a separate embedding contribution per layer), cross-layer KV-sharing (many layers don't
compute their own K/V — they borrow a neighbor's), and MatFormer nesting (the "2B/4B" is a
sub-network of a bigger model). The whole reason this project started: KV-sharing is exactly what the
vendor graph builder cannot express — but it traces fine as live graph dependencies when you just
torch_neuronx.trace the real forward. E2B fits one 16 GB core; E4B is ~2× and needs TP=2, which forces
the GQA question below.

12B — the encoder-free multimodal one. Ships as Gemma4UnifiedForConditionalGeneration (multimodal
class) with no audio/vision encoder loaded — you take model.language_model and move on. Drops PLE
and KV-sharing, but its sliding_window=1024 attention overflows Neuron's fused-attention SBUF, so
you force the eager attention implementation to make it compile.

31B — the dense giant. No PLE, no KV-sharing — architecturally simple. The hard part is scale:
it doesn't fit fewer than 8 cores, and the hand-driven tensor-parallel tracer that served E4B OOMs or
deadlocks at 31B. This is where NxD ModelBuilder becomes mandatory (compile one rank, shard
weights per-rank). A ~39-minute compile, 182 GB host peak, and a mixed-attention layout where the 10
global layers (nkv=4) get replicated instead of sharded.

26B-A4B — the MoE. The only sparse one, and the weirdest: a shared dense MLP running in parallel
with a 128-expert MoE, four feed-forward layernorms, a router with its own scale + per-expert scale.
The expert forward is a data-dependent gather/scatter that won't trace — so you compute all 128
experts, mask to the top-8 (identical math, static shape). "A4B" means ~4B active params but all 128
experts (~49 GB) are resident — it saves compute, not memory, and needs the same 24xlarge as the
2×-larger 31B.

The recipe evolved — and that's the real story

You can read the family as one recipe maturing under pressure:

torch_neuronx.trace (E2B). Trace the real HF forward on one core. Proves KV-sharing/PLE compile at all. Doesn't scale past one core.
tp_alias (E4B). Hand-driven tensor parallelism: on-device aliased KV buffers + a CPU-seed prefill, sharding GQA down to 1 KV head/rank. Serves two cores. Doesn't survive contact with 31B (the all-ranks-in-process tracer OOMs; serializing it deadlocks).
NxD ModelBuilder (12B, 31B, 26B). Compile one rank, load weights per-rank from a checkpoint_loader. This is the recipe that finally carries everything — the dense giant and the MoE — because the per-rank compile is what makes 8-way TP tractable.

Each step was forced by the next model refusing to run on the previous step's recipe. The MoE then
extends ModelBuilder one more notch (all-experts-dense experts as two auto-shardable parallel linears +
an SPMD-rank scatter).

The bug that's in every port

Here's the thread from the top. Every one of these ports, at some point, produced output that was
wrong while the CPU reference and my unit tests were perfect — and every time, the cause was the same
shape: the device faithfully computed a broken input, and only a coherence check caught it.

31B: a manually-built chat prompt tokenized <start_of_turn> into literal characters (Gemma-4's turn tokens are actually <|turn>/<turn|>). Device output matched the CPU reference token for token — because both faithfully processed the same garbage. SEQ_MATCH True on nonsense.
26B MoE: the device output was empty while CPU was perfect. A standalone tensor-parallel scatter baked rank 0's slice into the single-rank ModelBuilder trace, so every rank weighted the wrong experts. My math and sharding logic unit-tested identical to HF — the bug lived entirely in "how the primitive behaves under the trace."
E-family: the original "gibberish" endpoint — a real model loaded and producing confident nonsense because the harness around it (PLE path, embedding scale, KV wiring) was wrong, not the accelerator.

The lesson that survives all five: device == reference proves the accelerator is faithful. It does
not prove the reference is correct. A shared upstream bug — chat template, PLE, a runtime-rank scatter
— passes an equality check in dead silence. So you assert the output is sensible ("Paris"), not only
that device equals host. It's why every result in this series is quoted as an actual sentence, not just a
green SEQ_MATCH.

If you only take four things

Trace the model's real forward. The features that break vendor graph builders (KV-sharing, MoE routing) often trace fine as ordinary graph dependencies. Don't reimplement what you can trace.
Let the recipe escalate with the model. trace → hand-rolled TP → ModelBuilder. Reach for the single-rank-compile path before the manual one OOMs, not after.
"Active params" and "effective size" are compute/marketing numbers, not memory numbers. MatFormer and MoE don't shrink the HBM you must reserve — budget for the whole thing.
Equality with a reference is necessary, not sufficient. Always also check the output is coherent. The most expensive bugs in this whole family were the ones that passed SEQ_MATCH.

Artifacts

Docker Hub: xbill9/gemma4-optb{,-e4b,-31b,-26b,-26b-int8}
Hugging Face: xbill9/gemma-4-{E2B,E4B,31B,26B-A4B}-it-inferentia2 (+ -int8)
Per-model deep-dives: the other posts in this series.

Written with AI assistance (the ports, the debugging, and this write-up were done across Claude Code
sessions); every number quoted is from a real run on real hardware.

Porting Gemma-4 12B (the encoder-free multimodal one) to AWS Inferentia2

xbill — Fri, 17 Jul 2026 13:30:09 +0000

The middle child of the series. The 12B is the first Gemma-4 that's neither a MatFormer "effective"
model (E2B/E4B) nor a giant (31B/26B) — but it has its own two surprises: it's packaged as a
**multimodal* model with an encoder that isn't there, and its attention overflows a Neuron hardware
buffer in a way the smaller models never did. This is the short, clean port that only needed three
fixes on top of the E4B recipe — once I found them.*


Model	`google/gemma-4-12B-it` — dense, `model_type: gemma4_unified`
Hardware	AWS Inferentia2 — one `inf2.8xlarge` (1 chip / 2 cores / 32 GB HBM), TP=2
Result	Greedy decode token-for-token identical to the CPU fp32 reference (SEQ_MATCH True); prefill ~101 ms; `"…Paris."`
Footprint	~12 GB / rank, bf16, KV 256/64

Surprise #1: a multimodal class with no encoder

google/gemma-4-12B-it doesn't load as Gemma4ForConditionalGeneration like the E-family — it's
Gemma4UnifiedForConditionalGeneration, model_type: gemma4_unified. "Unified" means it's the
multimodal architecture… except this checkpoint ships no audio or vision encoder. It's an
encoder-free multimodal model: all the multimodal plumbing in the class, none of the encoders.

For an Inferentia text port that's actually good news — you take model.language_model and ignore
the rest, exactly as with the other sizes. But you have to know to reach past the unified wrapper, and
the config's vision/audio fields are present-but-inert, which makes "is this multimodal or not?" a
five-minute detour before you can even instantiate the text path.

The base recipe (from E4B → 12B)

The 12B reuses the E4B ModelBuilder device-prefill recipe almost wholesale: single-rank compile +
per-rank weight loading, on-device KV cache, host-side word embeddings, a bucketed device prefill, and
greedy decode. What it drops and what it changes:

No Per-Layer Embeddings. hidden_size_per_layer_input = 0 — feed plain inputs_embeds, delete the PLE path (same simplification the 31B enjoys).
No cross-layer KV-sharing. Every layer owns its KV.
Mixed sliding + global attention, sliding_window = 1024, GQA (16 q / 8 kv), softcap 30.

Then three fixes, each of which cost a compile cycle to find.

Fix #1: the global layers shard at nkv=1

The global (full_attention) layers have 1 KV head (num_key_value_heads-per-global = 1). Under
TP=2 you can shard them — but you shard the query heads and keep the single KV head grouped
correctly, groups = nq // TP. Get the grouping wrong and the global layers silently attend to the
wrong keys. (This is the same class of "how do the global layers shard" problem the 31B has, but the
31B's globals have nkv=4 and must be replicated — 12B's nkv=1 can shard. Same question, opposite
answer, per model.)

Fix #2: drop the softcap on-device

Gemma caps its final logits with 30 · tanh(logits / 30). On the host reference you apply it. On the
device you don't — softcap is monotonic, so it doesn't change the argmax, and computing a tanh
over a 262,144-token vocabulary on-device is a lot of transcendental for a no-op. Drop it from the
graph; the server re-applies it host-side only if it needs real probabilities for sampling.

Fix #3: force eager attention — sliding_window=1024 overflows fused-attn SBUF

This is the one that actually blocks you. Neuron's compiler wants to fuse attention (flash-style) for
speed. At sliding_window = 1024, that fused kernel's working set overflows SBUF (the on-chip
scratch buffer) and the compile fails. The fix is blunt: force the eager attention implementation
(attn_implementation="eager") so attention lowers as plain matmuls that fit. You give up the fused
kernel's speed, but it compiles and runs correctly, which beats a fast graph that doesn't exist. (The
smaller-window E-family models don't hit this — it's specifically the 1024-window layers that blow the
buffer.)

Result

With those three, the 12B compiles and its device greedy decode is token-for-token identical to the
CPU fp32 reference — SEQ_MATCH True, "…Paris.", ~101 ms first-token prefill, ~12 GB/rank in bf16
on a single inf2.8xlarge at TP=2. A clean port, hiding behind two non-obvious gotchas.

Takeaways

"Unified"/multimodal classes may be encoder-free — grab model.language_model and move on. Don't let the multimodal wrapper convince you it's a bigger problem than it is.
Softcap is a host-only concern for greedy decode. Monotonic ⇒ argmax-invariant ⇒ don't pay for it on-device over a 262K vocab.
SBUF overflow ⇒ force eager attention. When a large sliding_window fails to compile the fused kernel, the eager path is the escape hatch: slower, but it exists.
The same architectural question gets a different answer per model. "How do the global attention layers shard?" is nkv=1 → shard query heads here and nkv=4 → replicate on the 31B. Re-derive it; don't copy it.

Artifacts

Recipe/subdir: nxd-gemma-12b-inf2-work/optb/
Compiled model: s3://xbill-gemma4-patches-2b/optb-12b/mb_12b_256.pt

Written with AI assistance (the port, the debugging, and this write-up were done in a Claude Code
session); every result quoted is from a real run on real hardware.

My Inferentia port matched its reference token-for-token — and still output garbage

xbill — Fri, 17 Jul 2026 13:29:49 +0000

I ported Google's biggest open Gemma-4 — the 30-billion-parameter dense model — to AWS Inferentia2.
The compiled model passed my strictest check: its output was token-for-token identical to the
CPU reference. It was also complete gibberish. Both facts were true at the same time, and the reason
is a lesson worth more than the port.

This is the sequel to porting the smaller Gemma-4 models (E2B/E4B/12B). Those were hard because of
architecture — Per-Layer Embeddings, cross-layer KV-sharing, MatFormer nesting. The 31B throws all
of that away and hands you a different kind of hard: scale. And a trap.


Model	`google/gemma-4-31B-it` — dense, 60 layers, `model_type: gemma4`
Hardware	`inf2.24xlarge` — 6 chips / 12 cores / 384 GB host, TP=8
Software	`torch-neuronx` 2.8.0 · `neuronx-distributed` · `transformers` 5.13.0
Result	Device greedy decode == CPU fp32 reference (SEQ_MATCH True); prefill ~115 ms; `"The capital of France is Paris."`
Compile	Single-rank NxD ModelBuilder, ~39 min, host peak 182 GB, neffs 108 GB

The setup: the same model, eight times too big

31B is architecturally boring, and that's the point. enable_moe_block: False — dense.
hidden_size_per_layer_input: 0 — no PLE. num_kv_shared_layers: 0 — every layer owns its KV. It's
a plain decoder. It just happens to be a plain decoder with 30 billion parameters, which is about
60 GB in bf16, which does not fit on one 16 GB NeuronCore. The moment you reach for tensor
parallelism to spread it across 8 cores, everything about how you build the graph changes — and the
recipe that served the 4B and 12B models walks straight off a cliff.

The itch: the recipe that scaled to 12B dies at 31B

For E4B and 12B I hand-drove neuronx_distributed's parallel_model_trace — I owned the ranks, the
KV cache, the whole loop. It's clean at 4–12B. At 31B it failed two different ways, and both were
instructive:

All 8 ranks trace in one process → OOM. The model is ~15 GB/rank in bf16. Eight simultaneous fp32 compile graphs blew past 300 GB and killed the 384 GB box. That's not the model being big. That's the tracer being 8× parallel.
Serialize the ranks (max_parallel_compilations=1) → deadlock. Memory drops to 83 GB, a single rank reaches Compiler status PASS, and then the multi-rank weight-collection rendezvous hangs — worker blocked on pipe_read, CPU idle, 25 minutes, nothing.

The N=22 of this story. You can route around it — but the thing you route around is the thing that
owns you. The honest read: at 30B, hand-driving the tracer is the bug.

The pivot: stop hand-driving, let ModelBuilder do it

NxD's ModelBuilder is the designed large-model path. It compiles one rank, then loads
weights per-rank from a checkpoint_loader — no 8× compile, no queue deadlock. The 12B port had
already used it; I should have reached for it before the manual recipe fell over, not after.

The whole port is three callables and a wrap module:

mb = ModelBuilder(router=None, tp_degree=8,
                  checkpoint_loader=checkpoint_loader, compiler_workdir="/data/mb_wd")
mb.add("prefill", inst, [prefill_example], compiler_args=CARGS)
mb.add("decode",  inst, [decode_example],  compiler_args=CARGS)
model = mb.trace(initialize_model_weights=True)   # compile ONE rank, shard weights per rank
torch.jit.save(model, "/data/mb_31b_256.pt")      # 108 GB: graph + all 8 ranks' weights

The wrinkle: one attention layout is a lie

The one piece of architecture that does survive is nasty. 31B interleaves two attention layouts
with different head dimensions and KV-head counts:

layer type	count	q heads	kv heads	head_dim	v_proj
`sliding_attention`	50	32	16	256	present
`full_attention` (global)	10	32	4	512	None (`attention_k_eq_v` — V reuses K)

The global layers have 4 KV heads, and 4 < TP=8. You cannot shard 4 heads across 8 ranks. The
rule that works — shard the 50 sliding layers, replicate the 10 global layers:

nkv = a.k_proj.out_features // a.head_dim
if nkv % TP == 0:                       # sliding: 16 kv -> 2/rank, shard everything
    a.q_proj = col(a.q_proj); a.k_proj = col(a.k_proj)
    a.v_proj = col(a.v_proj); a.o_proj = row(a.o_proj)
else:                                    # global: 4 kv < TP=8 -> leave q/k/v/o full on every rank
    pass

Because the sharded layers' o_proj is row-parallel (it all-reduces back to a full hidden state), and
the replicated layers take a full hidden state in and out, the two layouts compose — the residual
stream is full-width at every boundary. nkv < TP is your signal: replicate, don't shard.

The trap under the trap: a buffer that isn't a parameter

This one cost an afternoon and produced zero error message — just a model whose cosine similarity to
the reference was ~0. Gemma-4 scales every layer's output by a learned layer_scalar. It's registered
as a buffer, not a parameter — and ModelBuilder's sharded checkpoint loader moves parameters
only. So all 60 layer_scalars silently defaulted to 1.0, and 60 layers of wrong-by-a-constant
compounded into noise.

# read the buffers straight from the safetensors and copy them in — the loader won't
for i, lyr in enumerate(layers):
    if hasattr(lyr, "layer_scalar") and i in lsv:
        lyr.layer_scalar.copy_(lsv[i])

Moral #1: buffers aren't parameters. Anything you register_buffer — per-layer scalars, some
norms — needs an explicit copy, or your model silently uses config defaults.

It compiles

build_module: 50 sharded-attn layers, 10 replicated-attn (global) layers, TP=8
  loaded 60 layer_scalar buffers
Sharding weights for ranks: 0...7  ->  Done Sharding weights in 172.6s
Finished building model in 2365.3 seconds (~39 min)
MB_TRACED -> MB_SAVED -> RUN_EXIT 0

Single-rank compile, all 8 ranks' weights sharded, host peak 182 GB on the 384 GB box — no OOM.
The manual path's 300 GB was never the model; it was the tracer. 108 GB of neffs, banked to S3.

The reveal: a perfect match to a broken answer

Reload the neffs, run the prompt, compare to the CPU fp32 reference. Here's the part that stops you
cold:

CPU GEN: '<start_of_turn>model\n<start_of_turn>model\n<start_of_turn>model...'
DEV GEN: '<start_of_turn>model\n<start_of_turn>model\n<start_of_turn>model...'
SEQ_MATCH True

SEQ_MATCH True. The device reproduced the CPU reference token for token. My compile was
numerically flawless. And the output was garbage — the model just echoing turn markers forever.

The tell is that both were wrong, identically. When your device output matches a broken reference
exactly, the bug isn't in the accelerator. It's upstream of both, in something they share. Here:
the prompt.

The hunt: the tokenizer was lying too

Two compounding facts about the Gemma-4 snapshot:

The chat template ships as a separate chat_template.jinja (the "Google Gemma 4 Canonical Chat Template", 18.7 KB, with a thinking-token scaffold) — not embedded in tokenizer_config.json. So apply_chat_template() raises "no chat template set."
Gemma-4's turn markers are <|turn> (105) and <turn|> (106) — not <start_of_turn>.

So the "obvious" fallback — hand-write "<start_of_turn>user\n..." and tokenize — does this:

"<start_of_turn>"  ->  ['<', 'start', '_', 'of', '_', 'turn', '>']     # 7 literal chars
convert_tokens_to_ids("<start_of_turn>")  ->  3   # unknown token

The model gets a malformed prompt, predicts < as the likeliest next token, and loops. The CPU
reference, fed the same malformed prompt, does the exact same thing. SEQ_MATCH True the whole time,
on nonsense.

The fix is to fetch and apply the real template:

if not getattr(tok, "chat_template", None):
    tok.chat_template = open(hf_hub_download(REPO, "chat_template.jinja", token=hf_tok)).read()
d = tok.apply_chat_template([{"role": "user", "content": PROMPT}], add_generation_prompt=True)
prompt = d if isinstance(d, list) else d["input_ids"]   # returns a BatchEncoding, not a plain dict

The payoff

Correct prompt, same neffs, same reload:

DEV GEN: 'The capital of France is Paris.'
DEVICE_PARIS True
DEVICE PREFILL first-token: 115 ms

Moral #2 — the one this whole post is for: SEQ_MATCH == True proves the accelerator faithfully
reproduces your reference. It does not prove your reference is correct. A shared upstream bug
passes an equality check in dead silence. Validate that the output is sensible, not only that device
equals host.

There's also a small, load-time gotcha worth its own line: torch.jit.load restores the graph and the
weights, but the weights aren't on the cores until you initialize them — and the call wants a
start-rank tensor:

model = torch.jit.load("/data/mb_31b_256.pt")
model.nxd_model.initialize_with_saved_weights(torch.tensor([0], dtype=torch.int32))
# without this: RuntimeError "This model is not initialized ..."

The side quest: racing a spot market that had nothing

inf2.24xlarge is 96 vCPUs — the entire default spot quota. One box, no headroom. During this work,
spot capacity for it was zero across us-west-2, us-east-1, and us-east-2 for 60+ minutes at a
stretch, and the boxes that did appear were often reclaimed in ~10–15 minutes. A 39-minute compile
does not fit in a 15-minute window. What made it possible:

A multi-region capacity poller — an instant EC2 Fleet firing across every AZ in all three regions on a loop, grabbing the first inf2.24xlarge that blinked into existence.
Durable state in same-region S3. Weights downloaded from HF once, mirrored to S3, never fetched again. Neffs pushed to S3 the instant the trace succeeds — so a reclaim after compile costs nothing.
Compile-only / device-only modes to bank the 39-minute compile immediately and skip the 121 GB fp32 CPU-reference generation when I only needed to eyeball the device output.

Moral #3: design the run for reclaims. When one instance is your whole quota and capacity is
intermittent, the gap between "impossible" and "done" is where your durable state lives.

The findings, in one table

#	The thing	Where it bites	How I knew	The fix
1	Hand-driven `parallel_model_trace` OOMs (300 GB, 8× in-process)	trace	box killed; MPC=1 → deadlock instead	Switch to `ModelBuilder` (compile 1 rank, shard per-rank)
2	Global layers have `nkv=4 < TP=8`	attention sharding	`4 % 8 != 0`	Replicate the 10 global layers; shard the 50 sliding
3	`layer_scalar` is a buffer, loader skips it	silent	cosine ≈ 0, no error	Copy the buffers from safetensors by hand
4	`torch.jit.load` leaves weights off the cores	reload	`RuntimeError: not initialized`	`initialize_with_saved_weights(tensor([0]))`
5	Chat template ships separately; turn tokens are `<\	turn>`/`<turn\	>`	prompt
—	`libneuronpjrt-path` "file not found" on `import neuronx_distributed`	env	—	It's just PATH: `export PATH=$VENV/bin:/opt/aws/neuron/bin:$PATH`

What I'm not claiming

Not a throughput result. I have first-token prefill latency (~115 ms) and a correctness proof. Decode tok/s, batching, and the production 512/128 buckets are open work.
Not a server. The validated model still needs MB_LOAD + initialize_with_saved_weights + the chat-template path wired into a persistent endpoint.
The compile is single-rank-fast but wall-clock-slow (~39 min). That's a real operational cost, not something I've optimized away.

Artifacts

Compiled TP=8 model + weights: s3://xbill-gemma4-31b-usw2/ (neffs/mb_31b_256.pt, weights/).
Recipe: tp_mb.py — one file: ModelBuilder trace, MB_LOAD reload, SKIP_VALIDATE/DEVICE_ONLY modes, and the chat-template prompt path.

Environment, for the reproducers: the aws_neuronx_venv_pytorch_2_8_nxd_inference venv +
pip install transformers==5.13.0, a transformers.utils.fx shim for the tfm-5.13 import surface, and
the PATH export above.

Written with AI assistance (the debugging, the porting, and this write-up were done in a Claude Code
session); every log line quoted here is from a real run on real hardware.

Porting a 128-expert MoE (Gemma-4 26B-A4B) to AWS Inferentia2 — where every rank weighted the wrong experts

xbill — Fri, 17 Jul 2026 13:29:28 +0000

Of the five Gemma-4 models I ported to AWS Inferentia2, the 26B-A4B was the one that was supposed to
be impossible: a Mixture of Experts. It compiled on the first try — then produced nothing. The
device output was empty. The CPU reference was perfect. My unit tests all passed. The bug was in a
place I'd have sworn couldn't have one.

This is the MoE entry in the series. The dense models (E2B/E4B/12B/31B) are hard because of
architecture; the 26B is hard because it's sparse — 128 experts, top-8 routing — and none of the
AWS vendor stack has ever traced an MoE for Gemma-4.


Model	`google/gemma-4-26B-A4B-it` — MoE, ~4B active / 26B total, `model_type: gemma4`
Hardware	`inf2.24xlarge` — 12 NeuronCores / 192 GB HBM, TP=8
Result	Device greedy decode == CPU fp32 reference (SEQ_MATCH True); prefill 77 ms; `"The capital of France is Paris."`
Artifacts	Docker Hub `xbill9/gemma4-optb-26b` · HF `xbill9/gemma-4-26B-A4B-it-inferentia2`

"A4B" is a lie your memory budget will not forgive

The name says 26B-A4B: ~4 billion active parameters per token. That sounds like a 4B model. It
is not. All 128 experts (~49 GB) must be resident in HBM — top-8 routing only reduces compute, not
footprint. So you budget for a 26B, not a 4B, and it needs the 24xlarge's 192 GB, same class as the
2×-larger dense 31B. (This bites again later when we try to shrink it onto a smaller box.)

The architecture: a dual-path FFN nobody warned me about

I inspected the transformers-5.13 module on a meta device before writing a line of port code, and it's
weirder than "MLP replaced by MoE." Every one of the 30 layers runs a shared dense MLP in parallel
with the 128-expert MoE, combined and passed through four feed-forward layernorms:

residual (post-attention)
├─ dense: pre_feedforward_layernorm  → mlp (2112)                → post_feedforward_layernorm_1  ─┐
└─ moe:   pre_feedforward_layernorm_2 → router → 128 experts (704) → post_feedforward_layernorm_2 ─┤
                                            hidden = dense + moe ────────────────────────────────────┘
   → post_feedforward_layernorm → + residual → × layer_scalar

Router (Gemma4TextRouter): its own RMSNorm + a scale + a per_expert_scale, then softmax(128, fp32) → top-8 → renormalize → × per_expert_scale.
Experts (Gemma4TextExperts): a fused gate_up_proj [128,1408,2816] + down_proj [128,2816,704], and its forward is a sparse gather/scatter loop (torch.where, index_add_).
Attention is the same mixed sliding/global layout as the 31B (which I'd already solved).

That sparse expert loop is the crux: it's data-dependent and will not trace to a static Neuron
graph.

The traceable trick: compute all experts, mask to top-8

You can't put a per-token, per-expert gather/scatter in a static HLO graph. So don't. Compute all 128
experts on every token, weight each by its router weight — which is zero for the non-top-8 — and
sum. Because a non-selected expert contributes 0 × expert(x) = 0, this is mathematically identical
to HF's sparse top-8, but it's a fixed-shape sequence of matmuls the compiler loves. Wasteful in FLOPs
(you compute 128 experts to use 8), correct in every bit, and traceable.

And the beautiful part: the whole thing collapses onto two standard parallel linears that
ModelBuilder already knows how to shard — no custom 3D-parameter sharding:

gate_up_proj [128,1408,2816] → [180224, 2816] as one ColumnParallelLinear (rank r gets experts 16*r*…16*r*+15)
down_proj [128,2816,704] → [2816, 90112] as one RowParallelLinear (input-sharded → all-reduce)

I keep HF's dual-path layer forward and its router unchanged (router replicated; its top_k_weights
already carry the renorm + per-expert scale) and swap only self.experts → my DenseExperts. Before
spending a cent on hardware, I unit-tested the math and the TP=8 sharding against HF on CPU:
MAXDIFF ≈ 2e-6, cosine 1.0. The math was right.

The bug: SEQ_MATCH-perfect math, empty device output

First device run: it compiled (the first-of-its-kind MoE trace on Neuron — 30 MoE layers,
MB_TRACED). And the output was ''. Empty. First token = end-of-turn, immediately.

CPU reference: "The capital of France is Paris." ✅
My all-experts-dense math: unit-tested identical to HF ✅
My TP=8 sharding logic: unit-tested identical to HF ✅
Device: nothing ❌

When your math is proven, your sharding logic is proven, and the device is still wrong, the bug is in
the gap between "how I think NxD's primitives behave" and "how they actually behave under the trace."

Here it was. To weight rank r's experts by the router, I scattered the dense weight matrix per-rank
with scatter_to_tensor_model_parallel_region. That function picks its slice using
get_tensor_model_parallel_rank() — a Python int. ModelBuilder compiles one rank and replicates
the graph across all 8 at runtime. So the trace baked rank 0's slice (Wd[:, 0:16]) into the graph,
and at runtime every rank weighted experts 0–15 while its gate_up computed its own experts
(16–31, 32–47, …). Total misalignment → garbage → the model confidently ends the turn.

The auto-sharded ColumnParallel/RowParallel don't have this problem because their weights are
pre-sliced at load time — the forward graph is rank-agnostic. A standalone collective on a runtime
activation is not.

The fix is the exact mechanism NxD's own MoE module uses (enable_spmd_rank): register an SPMDRank
module (a [1] int32 parameter checkpoint-loaded from arange(TP), so each rank gets its own number)
and scatter with the runtime-rank-aware variant:

Wl = scatter_to_process_group_spmd(Wd, 1, self.spmd_rank.get_rank())  # each rank gets ITS experts

Recompile. CPU GEN == DEV GEN == "The capital of France is Paris.", SEQ_MATCH True, 77 ms
prefill. The last and hardest of the five Gemma-4 variants was on Inferentia.

The general lesson, now taped to my monitor: any standalone tensor-parallel collective on a
runtime activation needs the SPMD-rank variant under ModelBuilder — auto-sharded parallel linears are
fine because their weights are pre-sliced; runtime tensors are not.

A packaging trap worth 20 wasted minutes

Publishing, the server wouldn't load the model: PytorchStreamReader ... failed finding central directory. My background "mirror to S3" loop had uploaded the 64.6 GB neff while torch.jit.save was
still writing it — a truncated 21.9 GB partial — and I'd terminated the box before the final clean
upload finished. Fix: never mirror an artifact mid-write; assert zipfile.is_zipfile(...) then do one
clean upload. (The "real" neff being 3× the size of the corrupt partial was the tell.)

Can it run on a smaller box? Yes, no, and "deeper than it looks"

The experts are ~93% of the weight, so quantizing them is the obvious lever. I built an int8-expert
variant (QuantizedColumnParallel/QuantizedRowParallel, per-channel symmetric) — two NxD autograd
quirks to fix on the way (.clone() the down output; monkeypatch scale_dequantize to be out-of-place,
it does x *= scale in place on an async-comm view). The result was genuinely good:

int8 is numerically PERFECT — SEQ_MATCH True, token-for-token identical to fp32.
int8 is genuinely smaller — neff 41.8 GB vs 64.6 GB bf16, exactly the 22.8 GB expert saving (the compiler keeps it int8, doesn't upcast).

But it fits only the 24xlarge. On a 2-core box (inf2.8xlarge/inf2.xlarge, 32 GB HBM, TP=2) it
Allocation Failures — experts drop to 11.4 GB/rank, but lm_head + the compiled graph + Neuron's
runtime reserve push it ~3–4 GB over the 16 GB core. And inf2 has no 4-core SKU to split the
difference. Fitting a 2-core box needs fp4 experts — which, unlike int8's five-line quantizer, is a
deep NxD microscaling integration (packed uint16, from_float hooks) with real accuracy risk. That's
a separate expedition.

So the shipped set is bf16 (24xlarge) and a validated int8 (24xlarge, smaller/faster, bf16-exact
quality) — both published — with fp4-for-2-cores documented as the next frontier.

Takeaways

A data-dependent op won't trace? Make it fixed-shape. All-experts-dense trades FLOPs for a static graph and exact correctness. Optimize routing later.
Unit-test math AND sharding logic on CPU before you pay for a NeuronCore. It's what let me say with certainty "the bug is in the primitives, not my code."
Runtime-rank is the MoE-on-ModelBuilder gotcha. Weights get pre-sliced; runtime activations need the SPMD-rank scatter. If the device is empty while CPU + unit tests are perfect, look here.
"Active params" is a compute number, not a memory number. MoE saves FLOPs, not HBM.
int8 per-channel was free accuracy-wise — a smaller, faster serving artifact at zero quality cost. The sub-8-bit step is where the real work (and risk) begins.

Artifacts

Docker Hub: xbill9/gemma4-optb-26b (bf16) · xbill9/gemma4-optb-26b-int8
HF: xbill9/gemma-4-26B-A4B-it-inferentia2 (bf16) · -int8
Recipe: tp_mb_moe.py (the DenseExperts + SPMDRank scatter) · tp_mb_moe_int8.py

Written with AI assistance (the port, the debugging, and this write-up were done in a Claude Code
session); every log line quoted is from a real run on real hardware.

Five Gemma-4 models, one accelerator: what porting E2B 31B to AWS Inferentia2 taught me

xbill — Fri, 17 Jul 2026 13:29:07 +0000

The whole family at a glance

	E2B	E4B	12B	31B	26B-A4B
Params (total / active)	~5B / 2B	~8B / 4B	12B	31B	26B / ~4B
Dense or sparse	dense	dense	dense	dense	MoE (128 exp, top-8)
Per-Layer Embeddings	✅	✅	—	—	—
Cross-layer KV-sharing	✅	✅	—	—	—
MatFormer ("effective" size)	✅	✅	—	—	—
Model class	`Gemma4ForConditionalGeneration`	same	`…Unified` (encoder-free)	`Gemma4ForConditionalGeneration`	same
Hardware	1 core (`inf2.xlarge`)	`inf2.8xlarge`	`inf2.8xlarge`	`inf2.24xlarge`	`inf2.24xlarge`
Tensor parallelism	TP=1	TP=2	TP=2	TP=8	TP=8
Recipe	`torch_neuronx.trace`	`tp_alias`	ModelBuilder	ModelBuilder	ModelBuilder + MoE
Correctness	SEQ_MATCH ✅	SEQ_MATCH ✅	SEQ_MATCH ✅	SEQ_MATCH ✅	SEQ_MATCH ✅
Speed	~44 tok/s	~34 tok/s	prefill ~101 ms	prefill ~115 ms	prefill ~77 ms

What every port shares

What makes each one different

The recipe evolved — and that's the real story

You can read the family as one recipe maturing under pressure:

torch_neuronx.trace (E2B). Trace the real HF forward on one core. Proves KV-sharing/PLE compile at all. Doesn't scale past one core.
tp_alias (E4B). Hand-driven tensor parallelism: on-device aliased KV buffers + a CPU-seed prefill, sharding GQA down to 1 KV head/rank. Serves two cores. Doesn't survive contact with 31B (the all-ranks-in-process tracer OOMs; serializing it deadlocks).
NxD ModelBuilder (12B, 31B, 26B). Compile one rank, load weights per-rank from a checkpoint_loader. This is the recipe that finally carries everything — the dense giant and the MoE — because the per-rank compile is what makes 8-way TP tractable.

The bug that's in every port

31B: a manually-built chat prompt tokenized <start_of_turn> into literal characters (Gemma-4's turn tokens are actually <|turn>/<turn|>). Device output matched the CPU reference token for token — because both faithfully processed the same garbage. SEQ_MATCH True on nonsense.
26B MoE: the device output was empty while CPU was perfect. A standalone tensor-parallel scatter baked rank 0's slice into the single-rank ModelBuilder trace, so every rank weighted the wrong experts. My math and sharding logic unit-tested identical to HF — the bug lived entirely in "how the primitive behaves under the trace."
E-family: the original "gibberish" endpoint — a real model loaded and producing confident nonsense because the harness around it (PLE path, embedding scale, KV wiring) was wrong, not the accelerator.

If you only take four things

Trace the model's real forward. The features that break vendor graph builders (KV-sharing, MoE routing) often trace fine as ordinary graph dependencies. Don't reimplement what you can trace.
Let the recipe escalate with the model. trace → hand-rolled TP → ModelBuilder. Reach for the single-rank-compile path before the manual one OOMs, not after.
"Active params" and "effective size" are compute/marketing numbers, not memory numbers. MatFormer and MoE don't shrink the HBM you must reserve — budget for the whole thing.
Equality with a reference is necessary, not sufficient. Always also check the output is coherent. The most expensive bugs in this whole family were the ones that passed SEQ_MATCH.

Artifacts

Docker Hub: xbill9/gemma4-optb{,-e4b,-31b,-26b,-26b-int8}
Hugging Face: xbill9/gemma-4-{E2B,E4B,31B,26B-A4B}-it-inferentia2 (+ -int8)
Per-model deep-dives: the other posts in this series.

Written with AI assistance (the ports, the debugging, and this write-up were done across Claude Code
sessions); every number quoted is from a real run on real hardware.

Smash Story: The Demo Script That Out-Debugged My Test Suite

xbill — Wed, 15 Jul 2026 23:58:42 +0000

This is a Smash Stories submission for the DEV Summer Bug Smash: a debugging story about the gap between "all tests pass" and "it actually works" — and the unlikely hero that closed it.

The setup

The project is a small MCP (Model Context Protocol) server that wraps Google's gemini-3.1-flash-lite-image model. It exposes image generation and stateful image editing as four tools that any MCP-speaking agent can call — Claude Code, a Google ADK agent, and a Rust CLI all consume the same ~300-line Python server. (Full architecture write-up here.)

By every signal a developer normally trusts, it was healthy:

✅ 10/10 unit tests passing
✅ ruff + mypy clean
✅ Used successfully through AI agents for days
✅ Published as a Docker image

Then I wrote a demo script. It found a production bug in under a minute of runtime.

The smash

demo.sh walks the stack live: discover the tools, generate an image, then do a stateful edit. To keep the demo cheap, step 2 requested the lowest quality tier the server documents:

cargo run --quiet -- generate "a tiny robot chef cooking ramen" 16:9 minimal

First run:

🔴 Image generation failed: Error code: 400 - {'error': {'message':
"'minimal' is not a supported thinking level for this model.
Allowed values are: low, high.", 'code': 'invalid_request'}}

Wait. The server's own validation had approved minimal before sending it. Here's that validation:

# server.py — as shipped
SUPPORTED_THINKING_LEVELS = {"minimal", "low", "medium", "high"}

@mcp.tool()
def generate_image(
    prompt: str, aspect_ratio: str = "1:1", thinking_level: str = "medium"
) -> str:
    ...

Four allowed values. The live API accepts two: low and high. And look at the default — medium. That's the real smash-worthy find:

Every live call that didn't explicitly override thinking_level was a guaranteed HTTP 400. The validation layer wasn't validating the API's contract — it was validating a stale memory of it.

Why ten green tests never noticed

The suite mocks the Gemini client, as unit tests should:

@patch("server._get_client")
def test_generate_image_success(self, mock_get_client):
    mock_client.interactions.create.return_value = mock_interaction
    result = generate_image(prompt="test", thinking_level="medium")
    self.assertIn("🟢 Image successfully saved!", result)

The mock returns success for any input — including inputs the real API rejects. The tests correctly proved "the server forwards medium faithfully." Faithfully forwarding an invalid value is still a bug; it's just invisible from inside the mock boundary.

Two conditions had to align for this to ship:

A local allowlist duplicated a remote-owned contract. SUPPORTED_THINKING_LEVELS was a cached copy of a fact only the API owns. Cached copies drift.
Every previous live caller happened to override the default. Agents kept requesting high for quality — so the broken default and the two phantom values were never exercised. f(x) being called a hundred times tells you nothing about f().

The fix

Two lines of production code, plus the part that actually takes discipline — locking the discovery in so it can't regress:

-SUPPORTED_THINKING_LEVELS = {"minimal", "low", "medium", "high"}
+SUPPORTED_THINKING_LEVELS = {"low", "high"}

-    prompt: str, aspect_ratio: str = "1:1", thinking_level: str = "medium"
+    prompt: str, aspect_ratio: str = "1:1", thinking_level: str = "low"

# New regression test: the live API only accepts low/high for this
# model; medium must now be rejected locally with a readable error.
result = generate_image(prompt="test", thinking_level="medium")
self.assertIn("Unsupported thinking level 'medium'", result)

Then the sweep (three tool signatures, docstrings, the server's self-describing get_help, every doc that repeated the wrong values) and a rebuild + push of the published Docker image, which had been shipping the bug to anyone who pulled it.

Before / after

	Before	After
Live call with default params	HTTP 400, every time	🟢 image saved
`thinking_level="minimal"` / `"medium"`	Approved locally, rejected remotely	Rejected locally with the allowed values named
Test suite	10/10 green (bug invisible)	11 assertions incl. contract regression test
Published image `xbill9/nb2lite-mcp`	Shipped the broken default	Rebuilt, pushed, verified live

Elapsed time from first failure to fixed-image-on-Docker-Hub: about ten minutes — because the failing tool call came back as readable text (Allowed values are: low, high) instead of a stack trace. Error messages that name the fix are half the debugging.

What I learned

Mocked tests verify your code; they cannot verify the contract. Keep one cheap live smoke call in the loop — mine now lives in the demo script itself (DEMO_FAST=1 ./demo.sh).
Test your defaults specifically. Defaults are the values nobody passes explicitly, so nobody exercises them — they're where contract drift hides longest.
A local allowlist of remote-owned values is a drift time bomb. If you pre-validate for better agent-facing errors (worth it!), pin a regression test to the values you've observed the API reject.
A demo script is the cheapest end-to-end test you'll ever write. Real credentials, real API, the real happy path — exactly the layer mocks can't reach. Mine paid for itself on its very first execution, before any audience saw it.

Best Use of Google AI

The whole project is built on Google AI, end to end:

The server wraps gemini-3.1-flash-lite-image through the Interactions API — the stateful sessions (store=True + previous_interaction_id) are what make multi-turn image editing work: the demo's edit step adds a neon RAMEN sign to the exact image generated moments earlier, with pixel-level continuity.
One of the three consumers is a Google ADK agent (LlmAgent on gemini-2.5-flash) that imports the server's tools over MCP via MCPToolset — Gemini calling Gemini, with the bug fix sitting in between.
The bug itself was a contract mismatch against the live Gemini API, and the fix was verified against it — the 400 error's precise, actionable message (Allowed values are: low, high) is what made this a ten-minute smash.