Gemma 4 Optimization with Agentic Kernel Engineering

Sep 29, 2026 · 6 min read

OpenRelay is a hardware-agnostic inference platform. We run models on NVIDIA and AMD alike, and we are a little obsessive about one thing: getting every TFLOP and every gigabyte of HBM a chip has into the workloads our customers run, at the latency they hold us to.

A spec sheet tells you what a GPU could do. What it actually does depends on a pile of kernels, lookup tables and flags, most of them tuned for some other model. Closing that gap is daily work for us, and these days agents do most of it. Here is one evening of that work: Gemma 4 31B on AMD MI355X.

A new chip, the same model

When we moved Gemma 4 31B from NVIDIA B200 to AMD MI355X, getting it running took a day. Blackwell's NVFP4 weights don't run on MI355X, so we switched to AMD's Quark MXFP4 checkpoint and kept Google's speculative-decoding drafter. Memory was the easy part: weights and drafter come to about 20 GB, which leaves most of the chip's 288 GB of HBM for KV cache, and long prompts eat KV cache.

Compute was another story. The first working config did 27.8K prefill tokens per second per GPU, which sounds fine until you do the arithmetic: the matrix math was running at under a fifth of the chip's MXFP4 peak. Held to our latency target, p95 time to first token under four seconds, one GPU topped out at 2.0 requests per second.

Agentic kernel engineering

This is the kind of gap we hand to agents. The loop they run is deliberately boring.

A research agent reads the kernel libraries' source and open pull requests and ranks what is likely to matter. A tuning agent works on real GPUs, one change at a time, and nothing it measures counts unless the model still scores at least 96.5% on GSM8K. It writes everything down, mistakes included. It caught one this time: a round of benchmarks that had quietly measured the wrong engine. It threw those numbers out and fixed its launcher so that can't happen again.

For Gemma 4 we gave it four MI355X GPUs and an evening.

What it found

Seven minutes in, the first profile came back with attention at 43% of prefill time. AMD's kernel library, AITER, picks attention tile sizes from a lookup table per GPU, and the MI355X table had no prefill entry for Gemma 4's large attention heads. So prefill ran on a tile sized for decoding, which re-read the same keys and values over and over.

Two new entries in that table made attention 2.6x faster and prefill a third faster, with the model's answers unchanged.

The second find was a GPU waiting on its CPU. With speculative decoding, each request needs three tokens per decode step, and the engine had only captured CUDA graphs up to 64 tokens. Any batch over 21 requests launched its kernels one at a time from the CPU. One flag fixed it.

The rest were smaller: tuned matrix-multiply configs for Gemma's exact shapes, a fused activation-and-quantize kernel, native RMSNorm so the compiler can fuse around it, and a kernel that was recompiling itself for every new batch size in the middle of serving, 190 milliseconds at a time. By the end of the evening, one MI355X served 1.5x the traffic under the same latency target, tokens streamed at 22 ms instead of 40 to 43, and GSM8K read 98.8% before and after.

How much of the chip we use now

Matrix math is now 61% of prefill time, which is where you want a GPU spending it, and those kernels run at 55 to 61% of the MXFP4 peak at the clocks the chip holds under load. Better kernels might buy another 1.2 to 1.3x. Past that, physics wins.

We are open-sourcing all of it

Everything the agent changed is in OpenRelayInc/inference-recipes: the Dockerfile, the tuned tables, the patches, and the benchmark harness behind every number in this post. The fixes that belong upstream are pull requests to AMD's AITER and to vLLM, each re-measured on their current main branch.

One of our fixes didn't make the cut. AITER's main branch has a new MI355X attention kernel that already runs Gemma 4's sliding-window layers faster than our tuned tile did (0.55 ms against our 0.61), so we left that one out and cheered instead.

For the kernel engineers

How we measured

One MI355X, the same harness before and after. The baseline is our first working MI355X config on stock vllm/vllm-openai-rocm:v0.30.0, not B200. We never ran this harness on B200, so nothing here compares vendors. Load tests are open-loop, 300 seconds per rate, two seeds, with requests averaging 9.5K tokens in and 380 out. Prefill is measured with 16 concurrent 10K-token prompts.

The tile

Gemma 4 has 50 sliding-window layers with 256-dim heads and 10 global layers with 512-dim heads and 8 query heads per KV head. On the global layers a 16-row tile is 2 tokens per program. The entry, keyed on long queries so decode keeps its tile:

"D_GEQ_512.Q_GEQ_256.DT_fp8_fp8": {"BLOCK_M": 128, "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64,
                                   "num_warps": 4, "num_stages": 1, "waves_per_eu": 1}

Graphs and the compile storm

Two draft tokens make each decode request 3 tokens, so capture sizes that stop at 64 leave batches over 21 requests ungraphed. Capturing to 256 costs 7 seconds of startup and 3.1 GiB. Separately, AITER declared a padded row count (scaleM_pad) as a Triton constexpr, so every new batch size compiled a new kernel mid-serving. As a runtime argument, 14 batch sizes need 4 compiles instead of 12, with bit-identical output.

What didn't help: Triton and hipBLASLt MXFP4 GEMMs (2.5x slower than AITER's assembly kernels), vLLM's QK-norm and RoPE fusion passes (they don't match Gemma 4's layout), CUDA graphs up to 16K tokens (no gain for 14 GiB), 128 concurrent sequences (twice the time per token for 1 to 3% more throughput), and a third draft token (same capacity).

Try it

Three of the changes are flags on the stock image. On vLLM after v0.30.0, also set VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1: without it, MXFP4 GEMMs fall back to Triton and Gemma 4 prefill drops from 27.7K to 18.9K tokens per second on a recent nightly.

--max-num-batched-tokens 8192
--compilation-config '{"cudagraph_capture_sizes":[1,2,4,8,16,24,32,48,64,80,96,112,128,144,160,176,192,224,256]}'
--kernel-config '{"ir_op_priority":{"rms_norm":["native"]}}'

If you would like to build the agents that do this every day, we are hiring.