Heterogeneous TP on gemma-4: running prefill on half the GPUs

In this post, Ariff Nazhan, our Senior AI Infrastructure Engineer, walks through how the team halved the GPUs used to process AI prompts, freeing four GPUs for speech and retrieval models on the same hardware. Reusing that capacity lets us run more of the services behind our customer-facing AI agents without adding GPUs.

We serve gemma-4-31B for our text and voice agents with NVIDIA Dynamo and vLLM, using disaggregated serving. Prefill engines process the prompt, decode engines generate the tokens, and the KV cache, the model's working memory for a request, moves from one to the other over RDMA. From 12 October our two B300 nodes ran one prefill engine and one decode engine, each at TP8 on a whole node. TP, tensor parallelism, is the number of GPUs one engine is split across.

The prefill node was mostly idle. We wanted to run prefill on four GPUs and give the other four to other models, but vLLM would not hand gemma-4's KV cache from a TP4 prefill to a TP8 decode. This post covers why, the patch we wrote for it, and what the prefill GPUs look like now.

Eight prefill GPUs, mostly idle

DCGM reports how busy each GPU is. Over the week the TP8 prefill engine ran, its eight GPUs averaged 7.7%. In the hours with real traffic, at least 5,000 prompt tokens per second, they averaged 22.5%, and the busiest hour reached 44%.

Prefill GPU utilisation, TP8 on eight GPUsHourly mean of the eight prefill GPUs, and the prompt load they served. Only hours with traffic; nights and quiet hours are cut out. Times in MYT.

Week average

7.7%7.7%

128 hours, nights included

Busy hours

22.5%22.5%

38 hours with at least 5,000 prompt tokens/s

Busiest hour

44%44%

at 22,400 prompt tokens/s

Show the data

Prefill GPU utilisation, hour by hour. Over the week the eight TP8 prefill GPUs averaged 7.7%. In 38 busy hours, at least 5,000 prompt tokens per second, they averaged 22.5%, and the busiest hour reached 44.0% at 22,391 prompt tokens per second.

Hour ending (MYT)	Utilisation	Prompt tokens/s
12 Oct, 13:00	0.8%	340
12 Oct, 14:00	0.8%	494
12 Oct, 15:00	0.4%	943
12 Oct, 16:00	0.8%	638
12 Oct, 17:00	0.5%	763
12 Oct, 18:00	0.6%	366
12 Oct, 19:00	0.4%	168
12 Oct, 20:00	4.5%	3,005
12 Oct, 21:00	19.9%	17,754
12 Oct, 22:00	30.9%	32,020
12 Oct, 23:00	0.6%	1,396
13 Oct, 00:00	19.0%	19,437
13 Oct, 01:00	14.6%	17,600
13 Oct, 02:00	10.8%	9,512
13 Oct, 03:00	0.0%	6
13 Oct, 04:00	0.0%	0
13 Oct, 05:00	0.0%	0
13 Oct, 06:00	0.0%	0
13 Oct, 07:00	0.0%	30
13 Oct, 08:00	40.8%	53,466
13 Oct, 09:00	0.0%	2
13 Oct, 10:00	0.1%	6
13 Oct, 11:00	0.0%	50
13 Oct, 12:00	0.4%	131
13 Oct, 13:00	0.0%	54
13 Oct, 14:00	0.2%	62
13 Oct, 15:00	0.2%	80
13 Oct, 16:00	1.0%	370
13 Oct, 17:00	0.1%	127
13 Oct, 18:00	0.2%	72
13 Oct, 19:00	0.0%	5
13 Oct, 20:00	0.0%	0
13 Oct, 21:00	0.0%	0
13 Oct, 22:00	0.0%	0
13 Oct, 23:00	0.0%	0
14 Oct, 00:00	0.0%	0
14 Oct, 01:00	0.0%	0
14 Oct, 02:00	0.0%	2
14 Oct, 03:00	0.0%	0
14 Oct, 04:00	0.0%	0
14 Oct, 05:00	0.0%	0
14 Oct, 06:00	0.0%	0
14 Oct, 07:00	0.0%	0
14 Oct, 08:00	0.6%	0
14 Oct, 09:00	0.0%	0
14 Oct, 10:00	0.0%	26
14 Oct, 11:00	0.0%	80
14 Oct, 12:00	0.3%	72
14 Oct, 13:00	0.0%	0
14 Oct, 14:00	0.0%	1
14 Oct, 15:00	0.5%	0
14 Oct, 16:00	0.5%	139
14 Oct, 17:00	0.1%	50
14 Oct, 18:00	0.0%	34
14 Oct, 19:00	0.0%	3
14 Oct, 20:00	20.5%	23,614
14 Oct, 21:00	19.7%	20,866
14 Oct, 22:00	39.8%	47,157
14 Oct, 23:00	0.3%	135
15 Oct, 00:00	20.1%	21,927
15 Oct, 01:00	30.4%	37,186
15 Oct, 02:00	5.2%	3,671
15 Oct, 03:00	0.0%	0
15 Oct, 04:00	0.0%	0
15 Oct, 05:00	0.0%	0
15 Oct, 06:00	2.5%	3,660
15 Oct, 07:00	35.5%	41,971
15 Oct, 08:00	28.0%	31,354
15 Oct, 09:00	43.5%	50,670
15 Oct, 10:00	7.8%	8,726
15 Oct, 11:00	19.6%	21,910
15 Oct, 12:00	20.3%	24,701
15 Oct, 13:00	0.1%	66
15 Oct, 14:00	30.5%	34,424
15 Oct, 15:00	36.7%	38,099
15 Oct, 16:00	19.5%	18,261
15 Oct, 17:00	0.6%	487
15 Oct, 18:00	0.0%	12
15 Oct, 19:00	0.1%	108
15 Oct, 20:00	0.0%	0
15 Oct, 21:00	0.0%	62
15 Oct, 22:00	0.1%	26
15 Oct, 23:00	0.0%	23
16 Oct, 00:00	0.1%	41
16 Oct, 01:00	0.0%	0
16 Oct, 02:00	0.3%	144
16 Oct, 03:00	0.5%	464
16 Oct, 04:00	0.1%	416
16 Oct, 05:00	0.1%	51
16 Oct, 06:00	0.3%	299
16 Oct, 07:00	0.1%	159
16 Oct, 08:00	0.8%	696
16 Oct, 09:00	4.7%	5,563
16 Oct, 10:00	15.5%	16,643
16 Oct, 11:00	20.2%	18,336
16 Oct, 12:00	21.2%	18,815
16 Oct, 13:00	14.8%	19,527
16 Oct, 14:00	14.9%	12,551
16 Oct, 15:00	13.4%	11,884
16 Oct, 16:00	15.3%	16,532
16 Oct, 17:00	12.9%	11,755
16 Oct, 18:00	7.6%	8,910
16 Oct, 19:00	2.5%	2,492
16 Oct, 20:00	7.1%	2,475
16 Oct, 21:00	8.1%	2,310
16 Oct, 22:00	5.4%	1,768
16 Oct, 23:00	2.8%	769
17 Oct, 00:00	1.5%	450
17 Oct, 01:00	1.2%	230
17 Oct, 02:00	0.7%	276
17 Oct, 03:00	0.8%	158
17 Oct, 04:00	2.5%	512
17 Oct, 05:00	6.5%	1,017
17 Oct, 06:00	16.1%	2,018
17 Oct, 07:00	2.8%	822
17 Oct, 08:00	3.1%	906
17 Oct, 09:00	11.3%	2,890
17 Oct, 10:00	24.0%	6,785
17 Oct, 11:00	19.6%	6,074
17 Oct, 12:00	23.7%	8,800
17 Oct, 13:00	32.6%	16,276
17 Oct, 14:00	21.3%	8,459
17 Oct, 15:00	44.0%	22,391
17 Oct, 16:00	16.6%	6,046
17 Oct, 17:00	24.1%	6,626
17 Oct, 18:00	16.5%	4,349
17 Oct, 19:00	7.7%	2,663
17 Oct, 20:00	3.8%	1,477

The decode node was the busy one. In the same busy hours its eight GPUs averaged 66%. Decode keeps a request on the GPUs for every output token, while prefill handles each prompt in one pass. So we wanted to halve the prefill engine to TP4 and keep decode at TP8, which keeps each request's token rate and the decode engine's KV cache where they were.

The layout we wantedOne request through disaggregated serving

1. Clientprompt
3. Frontendroutes the request
5. Prefill engineTP4 · 4 GPUs
6. KV cacheNIXL over RDMA
7. Decode engineTP8 · 8 GPUs

tokens stream back to the client

Mismatched TP answered 200 with the wrong text

vLLM's KV connector supports different TP sizes on the two sides in general, so we tried it. On gemma-4 it did not work. On vLLM 0.28 every pod came up Ready and every request returned HTTP 200. Our recall check caught it: nine prompts that each hide a code in a long numbered list and ask for it back.

Prefill TP8 feeding decode TP4, stock vLLM 0.28What a load test would have seen, and what the recall check saw

• HTTP status200on every request
• Recall check0 of 9codes returned correctly
• KV transfers failed40 of 40rejected at the handshake

A load test that only counts status codes would have scored these layouts as healthy. Every attempt so far, on both of our fleets:

vLLM 0.31 fails loudly where 0.28 failed silently. That is better, but the layout still does not run.

Why gemma-4 breaks the KV handoff

gemma-4-31B mixes two kinds of attention layer. 50 are sliding-window layers with 16 KV heads of size 256. 10 are global layers with 4 KV heads of size 512. Tensor parallelism shards KV heads across GPUs, so each GPU holds a slice of every layer's cache.

vLLM's NIXL connector, which moves the cache from prefill to decode, sized every layer's slice from one model-wide head count of 16. For the sliding layers that is right. For the global layers it breaks at TP8.

KV heads per GPUOne sliding-window layer and one global layer

Sliding-window layer16 KV heads of size 256 · 50 such layers

TP4 · 4 GPUs

0123456789101112131415

TP8 · 8 GPUs

0123456789101112131415

TP4 → TP8: 4 heads per GPU become 2. The slice halves.

Global layer4 KV heads of size 512 · 10 such layers

TP4 · 4 GPUs

g0g1g2g3

TP8 · 8 GPUs

g0g0g1g1g2g2g3g3

same head on both GPUs

TP4 → TP8: 1 head per GPU stays 1. The slice does not halve, but the connector, assuming 16 heads, expected it to.

A sliding-window layer has 16 KV heads. At TP4 each GPU holds 4, and at TP8 each holds 2. The slice halves, which is what the connector expects.

A global layer has only 4 KV heads. At TP4 that is one whole head per GPU.

Four heads cannot be split eight ways, so at TP8 each head is copied onto two GPUs and every GPU keeps one whole head.

The slice per GPU is the same size at TP4 and at TP8. The connector, assuming 16 heads everywhere, expected it to halve.

Before reading, the decode engine checks that each remote slice is the size it expects. With 16 heads assumed everywhere, it expects the prefill's TP4 slice to be twice its own. For the first global layer, layer 5, the two are equal, and the handshake fails:

AssertionError: SPLIT region 5: remote P KV block_len 131072 must equal
local 131072 * remote_heads 4 // local_heads 2 // block_size_ratio 1

There is a second problem underneath. vLLM's hybrid KV cache manager saves memory on models like gemma-4 by letting a sliding layer and a global layer share cache pages. A shared page cannot be re-sharded layer by layer, so any fix also needs the hybrid manager turned off.

What the patch does

Each side now advertises how many KV heads per GPU every region of its cache holds, and the handshake check and the read plan use that per-region count. A sliding region with 4 heads per prefill GPU and 2 per decode GPU is read in halves, as before. A global region with one head on both sides is read whole, by both decode GPUs that share the head.

How two decode GPUs read one prefill GPUPrefill TP4 into decode TP8, with the patch

Two limits apply. Decode TP must be at least prefill TP when layers have different head counts; the reverse raises an error instead of producing wrong output. Both engines also need --disable-hybrid-kv-cache-manager; with the hybrid manager on, the patch stops with a clear error.

Testing it away from production

We tested on our H20 fleet before touching the B300 cluster. The test pair ran on two nodes with its own etcd, NATS and Dynamo frontend on separate ports, so the production frontend and its workers never saw it. Each run used the same flags with stock and with patched vLLM 0.31.0.

A TP2 prefill feeding a TP4 decode passed too, on stock vLLM as well. At TP4 the four global heads still split evenly, one per GPU, so the old arithmetic holds. The problem only appears once a layer has fewer heads than the decode engine has GPUs.

Shipping it through our Dynamo fork

We carry the fix in our Dynamo fork as a patch file against vLLM 0.31.0, together with its tests, and build the serving image from that line. The change is in PR #1 of our Dynamo fork.

Running on B300

Since 17 October the cluster runs one TP4 prefill engine on four GPUs of the first node and one TP8 decode engine on all eight GPUs of the second, with the KV cache crossing between the nodes over RoCE. The four GPUs the prefill engine gave back now run our speech and retrieval models.

Where the engines runTwo B300 nodes, eight GPUs each

Node 1prefill

01234567

KV cache over RoCE: prefill GPU n is read by decode GPUs 2n and 2n + 1

Node 2decode

01234567

• Prefill engine, TP4
• Decode engine, TP8
• Given back: speech on 4 and 5, embedding and reranking on 6 and 7

What we checked after the rollout:

• Each of the eight decode GPUs logs TransferTopology(tp_ratio=2, local_tp=8, remote_tp=4).
• Neither engine logs a SPLIT region or handshake error.
• The recall check passes 9 of 9 on prompts of about 12,500 tokens.
• Our gemma-4 checks pass 4 of 4: tool-call turns with thinking on and off, JSON output, and the build header.

Prefill utilisation after the switch

Four GPUs now do the prefill work that eight did. In the busy hours since the switch, again at least 5,000 prompt tokens per second, the four prefill GPUs averaged 30.4%, and the busiest hour reached 49%.

Utilisation against loadEach dot is one hour. Dashed lines are the fitted rate for each layout.

• TP8 prefill, 128 hours
• TP4 prefill, 25 hours

Show the data

Prefill GPU utilisation, hour by hour. At the same prompt rate a TP4 prefill GPU is about 2.9 times as busy as a TP8 one: 2.75 against 0.95 utilisation points per 1,000 prompt tokens per second.

Hour ending (MYT)	Prefill GPUs	Utilisation	Prompt tokens/s
12 Oct, 13:00	8	0.8%	340
12 Oct, 14:00	8	0.8%	494
12 Oct, 15:00	8	0.4%	943
12 Oct, 16:00	8	0.8%	638
12 Oct, 17:00	8	0.5%	763
12 Oct, 18:00	8	0.6%	366
12 Oct, 19:00	8	0.4%	168
12 Oct, 20:00	8	4.5%	3,005
12 Oct, 21:00	8	19.9%	17,754
12 Oct, 22:00	8	30.9%	32,020
12 Oct, 23:00	8	0.6%	1,396
13 Oct, 00:00	8	19.0%	19,437
13 Oct, 01:00	8	14.6%	17,600
13 Oct, 02:00	8	10.8%	9,512
13 Oct, 03:00	8	0.0%	6
13 Oct, 04:00	8	0.0%	0
13 Oct, 05:00	8	0.0%	0
13 Oct, 06:00	8	0.0%	0
13 Oct, 07:00	8	0.0%	30
13 Oct, 08:00	8	40.8%	53,466
13 Oct, 09:00	8	0.0%	2
13 Oct, 10:00	8	0.1%	6
13 Oct, 11:00	8	0.0%	50
13 Oct, 12:00	8	0.4%	131
13 Oct, 13:00	8	0.0%	54
13 Oct, 14:00	8	0.2%	62
13 Oct, 15:00	8	0.2%	80
13 Oct, 16:00	8	1.0%	370
13 Oct, 17:00	8	0.1%	127
13 Oct, 18:00	8	0.2%	72
13 Oct, 19:00	8	0.0%	5
13 Oct, 20:00	8	0.0%	0
13 Oct, 21:00	8	0.0%	0
13 Oct, 22:00	8	0.0%	0
13 Oct, 23:00	8	0.0%	0
14 Oct, 00:00	8	0.0%	0
14 Oct, 01:00	8	0.0%	0
14 Oct, 02:00	8	0.0%	2
14 Oct, 03:00	8	0.0%	0
14 Oct, 04:00	8	0.0%	0
14 Oct, 05:00	8	0.0%	0
14 Oct, 06:00	8	0.0%	0
14 Oct, 07:00	8	0.0%	0
14 Oct, 08:00	8	0.6%	0
14 Oct, 09:00	8	0.0%	0
14 Oct, 10:00	8	0.0%	26
14 Oct, 11:00	8	0.0%	80
14 Oct, 12:00	8	0.3%	72
14 Oct, 13:00	8	0.0%	0
14 Oct, 14:00	8	0.0%	1
14 Oct, 15:00	8	0.5%	0
14 Oct, 16:00	8	0.5%	139
14 Oct, 17:00	8	0.1%	50
14 Oct, 18:00	8	0.0%	34
14 Oct, 19:00	8	0.0%	3
14 Oct, 20:00	8	20.5%	23,614
14 Oct, 21:00	8	19.7%	20,866
14 Oct, 22:00	8	39.8%	47,157
14 Oct, 23:00	8	0.3%	135
15 Oct, 00:00	8	20.1%	21,927
15 Oct, 01:00	8	30.4%	37,186
15 Oct, 02:00	8	5.2%	3,671
15 Oct, 03:00	8	0.0%	0
15 Oct, 04:00	8	0.0%	0
15 Oct, 05:00	8	0.0%	0
15 Oct, 06:00	8	2.5%	3,660
15 Oct, 07:00	8	35.5%	41,971
15 Oct, 08:00	8	28.0%	31,354
15 Oct, 09:00	8	43.5%	50,670
15 Oct, 10:00	8	7.8%	8,726
15 Oct, 11:00	8	19.6%	21,910
15 Oct, 12:00	8	20.3%	24,701
15 Oct, 13:00	8	0.1%	66
15 Oct, 14:00	8	30.5%	34,424
15 Oct, 15:00	8	36.7%	38,099
15 Oct, 16:00	8	19.5%	18,261
15 Oct, 17:00	8	0.6%	487
15 Oct, 18:00	8	0.0%	12
15 Oct, 19:00	8	0.1%	108
15 Oct, 20:00	8	0.0%	0
15 Oct, 21:00	8	0.0%	62
15 Oct, 22:00	8	0.1%	26
15 Oct, 23:00	8	0.0%	23
16 Oct, 00:00	8	0.1%	41
16 Oct, 01:00	8	0.0%	0
16 Oct, 02:00	8	0.3%	144
16 Oct, 03:00	8	0.5%	464
16 Oct, 04:00	8	0.1%	416
16 Oct, 05:00	8	0.1%	51
16 Oct, 06:00	8	0.3%	299
16 Oct, 07:00	8	0.1%	159
16 Oct, 08:00	8	0.8%	696
16 Oct, 09:00	8	4.7%	5,563
16 Oct, 10:00	8	15.5%	16,643
16 Oct, 11:00	8	20.2%	18,336
16 Oct, 12:00	8	21.2%	18,815
16 Oct, 13:00	8	14.8%	19,527
16 Oct, 14:00	8	14.9%	12,551
16 Oct, 15:00	8	13.4%	11,884
16 Oct, 16:00	8	15.3%	16,532
16 Oct, 17:00	8	12.9%	11,755
16 Oct, 18:00	8	7.6%	8,910
16 Oct, 19:00	8	2.5%	2,492
16 Oct, 20:00	8	7.1%	2,475
16 Oct, 21:00	8	8.1%	2,310
16 Oct, 22:00	8	5.4%	1,768
16 Oct, 23:00	8	2.8%	769
17 Oct, 00:00	8	1.5%	450
17 Oct, 01:00	8	1.2%	230
17 Oct, 02:00	8	0.7%	276
17 Oct, 03:00	8	0.8%	158
17 Oct, 04:00	8	2.5%	512
17 Oct, 05:00	8	6.5%	1,017
17 Oct, 06:00	8	16.1%	2,018
17 Oct, 07:00	8	2.8%	822
17 Oct, 08:00	8	3.1%	906
17 Oct, 09:00	8	11.3%	2,890
17 Oct, 10:00	8	24.0%	6,785
17 Oct, 11:00	8	19.6%	6,074
17 Oct, 12:00	8	23.7%	8,800
17 Oct, 13:00	8	32.6%	16,276
17 Oct, 14:00	8	21.3%	8,459
17 Oct, 15:00	8	44.0%	22,391
17 Oct, 16:00	8	16.6%	6,046
17 Oct, 17:00	8	24.1%	6,626
17 Oct, 18:00	8	16.5%	4,349
17 Oct, 19:00	8	7.7%	2,663
17 Oct, 20:00	8	3.8%	1,477
18 Oct, 15:00	4	11.6%	6,412
18 Oct, 16:00	4	15.3%	5,676
18 Oct, 17:00	4	16.0%	5,254
18 Oct, 18:00	4	13.9%	4,362
18 Oct, 19:00	4	6.7%	2,342
18 Oct, 20:00	4	7.1%	2,437
18 Oct, 21:00	4	11.4%	3,636
18 Oct, 22:00	4	9.3%	2,538
18 Oct, 23:00	4	16.9%	1,809
19 Oct, 00:00	4	10.6%	1,292
19 Oct, 01:00	4	2.6%	788
19 Oct, 02:00	4	2.9%	1,128
19 Oct, 03:00	4	3.8%	1,174
19 Oct, 04:00	4	3.4%	1,477
19 Oct, 05:00	4	3.7%	1,287
19 Oct, 06:00	4	14.5%	2,223
19 Oct, 07:00	4	3.7%	1,240
19 Oct, 08:00	4	3.8%	1,436
19 Oct, 09:00	4	15.3%	4,390
19 Oct, 10:00	4	36.1%	13,943
19 Oct, 11:00	4	44.0%	17,264
19 Oct, 12:00	4	48.9%	17,096
19 Oct, 13:00	4	33.0%	11,996
19 Oct, 14:00	4	36.9%	12,084
19 Oct, 15:00	4	31.7%	9,530

Halving the GPUs accounts for a factor of 2 of that 2.9. We have not explained the rest yet. The two windows carry different traffic, the after window is only a day long, and the hybrid KV manager is now off. Even at its busiest hour the TP4 engine stays under half of its GPUs' time, so four GPUs are enough for the traffic we see today.

Summary

Our TP8 prefill engine used eight GPUs at 22.5% in busy hours. We wanted it on four, but gemma-4's two kinds of attention layer meant stock vLLM could not hand its KV cache from a TP4 prefill to a TP8 decode, and on the version we started with it failed silently. We ported a fix for the KV connector, tested it on a separate stack, and shipped it in our Dynamo fork. Prefill now runs on four GPUs, and the other four run our speech and retrieval models.

If you are doing the same thing, here is what we'd recommend:

• Look at GPU utilisation per role. Our prefill and decode nodes were sized the same and used very differently.
• Check the answer, not the status code. A recall check in front of every test caught a layout that returned 200 with noise.
• Read the per-layer attention shapes of your model before relying on a disaggregation feature. Mismatched TP is supported for models whose layers all have the same number of KV heads, and gemma-4's do not.
• Carry fixes as patch files with their tests, pinned to the version they were ported to, so they are easy to drop later.

Sources

• Release v1.6.0-nightly-20261016-gemma4-hetero-tp.1
• PR #1: the 1.6.0 line with the vLLM patch
• The patch: vllm-0.31.0-nixl-per-region-kv-heads.patch
• SCICOM_CHANGELOGS.md: every Scicom change to Dynamo