fix(gemma4): drive the TP ranks concurrently so long prompts do not stall

CLAUDE.md

@@ -13,7 +13,7 @@ Every model line is behind a cargo feature; only `qwen3` is a default feature, s
 | Qwen3-4B / 8B | `pegainfer-qwen3` | `qwen3` (default) | Full attention, TP support |
 | Qwen3.5-4B / 9B / 27B · Qwen3.8-27B | `pegainfer-qwen35` | `--features qwen35` (needs build-time Python + Triton) | Hybrid Gated DeltaNet + full attention. Qwen3.8 shares the line: same `model_type`, same text geometry — see `docs/models/qwen35/support-qwen38.md` |
 | DeepSeek-V2-Lite | `pegainfer-deepseek-v2-lite` | `--features deepseek-v2-lite` | MoE + EP, 2-GPU |
-| Gemma 4 | `pegainfer-gemma4` | `--features gemma4` | Sliding-window + global full attention, single GPU, batched decode, opt-in chunked prefill |
+| Gemma 4 | `pegainfer-gemma4` | `--features gemma4` | Sliding-window + global full attention, dense sizes at TP2 (one rank per device), batched decode, opt-in chunked prefill (single rank only) |
 | Kimi-K2 | `pegainfer-kimi-k2` | `--features kimi-k2` | MLA + MoE + Marlin INT4, 8-GPU EP |
 | GLM5.2 | `pegainfer-glm52` | `--features glm52` | MLA + MoE + FP8, 8-GPU EP (bring-up) |
 | Kimi-K3 | `pegainfer-k3` | `--features k3` | Hybrid KDA + MLA, latent MoE + MXFP4, EP (bring-up — single-rank decode wired) |

README.md

@@ -147,7 +147,7 @@ Only **`qwen3`** is enabled by default, including in the prebuilt binary. Build
 | --- | --- | --- | --- |
 | **Qwen3 · dense 0.6B to 32B** | Full attention, GQA | `qwen3` · default | Greedy + sampling, tensor parallel, prefix cache, KV offload; DFlash / DSpark on 4B. [Model page](https://pegainfer.org/models/qwen3-4b/) |
 | **Qwen3.5 · dense 0.8B to 27B** | Gated DeltaNet + full attention | `qwen35` | Text-only BF16; build-time Triton. Qwen3.8-27B rides the same feature. [Model page](https://pegainfer.org/models/qwen35/) · [Qwen3.8 support record](docs/models/qwen35/support-qwen38.md) |
-| **Gemma 4 · 12B and 26B-A4B** | Sliding-window + global attention; NVFP4 routed experts on 26B | `gemma4` | Text-only, single GPU, batched decode and optional chunked prefill. [Model page](https://pegainfer.org/models/gemma4/) |
+| **Gemma 4 · 12B and 26B-A4B** | Sliding-window + global attention; NVFP4 routed experts on 26B | `gemma4` | Text-only, batched decode and optional chunked prefill; the dense sizes also run at tensor parallel 2, one rank per device. [Model page](https://pegainfer.org/models/gemma4/) |
 | **DeepSeek-V2-Lite** | MLA + MoE | `deepseek-v2-lite` | 2-GPU EP2 correctness path. [Status and limits](docs/models/deepseek-v2-lite/status.md) |
 | **Kimi-K2 / K2.5** | MLA + MoE, Marlin INT4 | `kimi-k2` | 8-GPU expert parallelism. [Model roadmap](docs/models/kimi-k2/roadmap.md) |
 | **GLM-5.2** | Sparse MLA + MoE, FP8 | `glm52` | Blackwell; EP decode, TP4 prefill, native MTP speculative decoding, P/D disaggregation. Bring-up. [Model page](https://pegainfer.org/models/glm52/) |

docs/benchmarks/gemma4-31b-tp2-l20.md

@@ -0,0 +1,69 @@
+# Gemma 4 31B bf16 tensor-parallel-2 on two L20s
+
+TL;DR: the stock 31B bf16 checkpoint serves on two 48 GiB L20s (sm_89, PCIe) as a two-rank eager engine — 29.91 GiB of weights per rank, a `PEGAINFER_MAX_CONTEXT=8192` envelope at 8 decode slots, cold load 30 s, first token 57 ms, 49.4 ms per decode step, and 143 tok/s aggregate at 8 concurrent. The checkpoint does not fit one card (57 GiB), so there is no in-box single-rank baseline for it; numerical parity rests on the 31B-geometry synthetic gate (`models/gemma4/tp.md`).
+
+## Rig
+
+| | |
+| --- | --- |
+| checkpoint | `gemma-4-31b-it`, bf16, text tower only (356 vision/audio tensors skipped) |
+| hardware | two NVIDIA L20 (sm_89, 44.99 GiB usable each), PCIe, one p8s node |
+| build | `--features gemma4`, `PEGAINFER_CUDA_SM=89`, CUDA 12.9 |
+| runtime | `--tp-size=2 --cuda-graph=false`, `NCCL_PROTO=LL128`, pod NCCL 2.18.3 |
+| envelope | `PEGAINFER_MAX_CONTEXT=8192`, `PEGAINFER_DECODE_SLOTS=8` |
+| checkpoint source | node-local disk (a PVC read is ~3.5 min per rank) |
+
+## Envelope
+
+| | per rank |
+| --- | --- |
+| weights (device) | 29.91 GiB |
+| global KV pool | 2.50 GiB (1025 pages x 2.50 MiB) |
+| resident after weights | 14.18 GiB free |
+| load wall (cold) | 30 s (31.4 s in a second cold run) |
+
+The 8192 x 16-slot default point does not quite fit: 29.91 GiB of weights plus 14.38 GiB of pools (9.38 local + 5.00 global, all binary GiB) is 44.29 GiB, and the ~0.90 GiB implied by the "14.18 GiB free" row above — CUDA context, step arena, cuBLAS workspace, NCCL buffers — puts it at ≈45.2 GiB against the 44.99 GiB usable. It misses by ~0.2 GiB, not by a wide margin. 8 slots is the point that fits. `PEGAINFER_MAX_CONTEXT`/`PEGAINFER_DECODE_SLOTS` trade context against concurrency — see `models/gemma4/tp.md` for the other measured envelopes and the page arithmetic.
+
+## Method
+
+Raw `/v1/completions`, greedy, one connection per request. Time to first token is measured directly as the wall of a `max_tokens=1` request; per-step time is `(t(65) - t(1)) / 64`, which cancels the first-token cost. One warmup request is issued first to absorb the one-time buffer/pool warm; concurrency is a thread pool of `N` over `N` requests of 64 tokens each, wall-timed as a batch.
+
+## Results
+
+Prompt `The capital of France is`, 5 TTFT samples and 3 TPOT samples:
+
+| metric | value |
+| --- | --- |
+| first token (TTFT) | 57 ms (58, 57, 57, 57, 57) |
+| decode step (TPOT) | 49.4 ms (49.4, 49.41, 49.40) |
+
+Throughput, 64 tokens per request:
+
+| concurrency | wall | aggregate | per request |
+| ---: | ---: | ---: | ---: |
+| 1 | 3.17 s | 20.2 tok/s | 3.17 s |
+| 4 | 3.47 s | 73.8 tok/s | 0.87 s |
+| 8 | 3.57 s | 143.5 tok/s | 0.45 s |
+
+A decode step costs ~49 ms for one row and ~56 ms for eight (3.57 s / 64 steps), so batch width is nearly free until the collective and the wide-batch GEMMs bite.
+
+## Eager vs captured decode
+
+The default is captured (`--cuda-graph=true`, as at TP1); both modes measured back to back on the pair:
+
+| mode | TTFT | step | c8 aggregate |
+| --- | ---: | ---: | ---: |
+| eager (`--cuda-graph=false`) | 56.5 ms | 49.4 ms | 143.7 tok/s |
+| captured (default) | 57.0 ms | 51.5 ms | 146.0 tok/s |
+
+A wash: the collective, not launch count, dominates the step, so capture buys little and costs ~4% per step at c1. Both start, serve and shut down cleanly; the numbers in the tables above are the eager run.
+
+## Correctness
+
+The distribution gate at 31B geometry measured exact (`models/gemma4/tp.md`): the synthetic checkpoint carrying this checkpoint's shapes came out bit-identical between one rank and two (48/48 picks, worst picked-token gap `0.0000`) — a measurement, since the gate asserts against `LOGBROB_LINE` rather than against exactness. On the real checkpoint the sane answers come out through the chat template — `What is the capital of France?` → `Paris`; `Name three primary colors.` → `The three primary colors are red, yellow, and blue.` (the sampled defaults; greedy repeats, which is the model's published non-default behavior, not a TP artifact).
+
+## Notes
+
+- Raw completions without `<bos>` degenerate, as the tokenizer contract states; the chat endpoint (which carries the template's BOS) is the coherent path.
+- The first GEMM of a cold two-rank process would hang on this environment's NCCL before the module-loading warm was added — see "Cold start" in `models/gemma4/tp.md`. The load figure above is with that warm in place.
+- Every rank holds the whole 262144-entry vocabulary projection (~2.8 GiB at 31B), and the extra ranks compute a full batch of logits per decode step that is then discarded. The step numbers above include that waste; see "Known bounds" in `models/gemma4/tp.md`.