diff --git a/astro.config.mjs b/astro.config.mjs index 6e2b25c..edb5ed1 100644 --- a/astro.config.mjs +++ b/astro.config.mjs @@ -121,7 +121,7 @@ export default defineConfig({ { label: 'Qwen3 Dense', link: '/models/qwen3-4b/' }, { label: 'Qwen3.5 Dense', link: '/models/qwen35/' }, { label: 'GLM-5.2', link: '/models/glm52/' }, - { label: 'Gemma 4 12B', link: '/models/gemma4/' }, + { label: 'Gemma 4', link: '/models/gemma4/' }, ], }, ], diff --git a/public/models/gemma4/perf-26b.svg b/public/models/gemma4/perf-26b.svg new file mode 100644 index 0000000..c889ff9 --- /dev/null +++ b/public/models/gemma4/perf-26b.svg @@ -0,0 +1,2 @@ + +Time to first token (s)020406080 s1.075.9118.6069.570.99-7.6%4.87-17.5%13.62-26.8%40.75-41.4%PegaInfervLLMTime per output token (ms)0481216 ms8.739.9311.1313.418.67-0.7%9.42-5.1%10.42-6.4%12.39-7.6%10.6K40K81.7K163Kprompt tokens diff --git a/src/content/docs/index.mdx b/src/content/docs/index.mdx index 389a410..cd673f7 100644 --- a/src/content/docs/index.mdx +++ b/src/content/docs/index.mdx @@ -88,4 +88,4 @@ Decode is captured as a CUDA graph with pre-allocated buffers, keeping per-token | DeepSeek-V4 | MoE + compressor + indexer, 8-GPU | | DeepSeek-V2-Lite | MoE + expert parallelism, 2-GPU | | Kimi-K2 | MLA + MoE + Marlin INT4, 8-GPU expert parallelism | -| [Gemma 4 12B](/models/gemma4/) | Sliding-window local + global attention, dual paged KV | +| [Gemma 4](/models/gemma4/) | Sliding-window local + global attention, dual paged KV; 12B dense and 26B-A4B routed experts in NVFP4 | diff --git a/src/content/docs/models/gemma4.md b/src/content/docs/models/gemma4.md index dd9adcd..77610bc 100644 --- a/src/content/docs/models/gemma4.md +++ b/src/content/docs/models/gemma4.md @@ -1,12 +1,14 @@ --- -title: Gemma 4 12B -description: "Running Gemma 4 12B on pegainfer: build with the gemma4 feature, launch, the memory envelope a slot count buys, long-context profiles, and heterogeneous attention notes." +title: Gemma 4 +description: "Running Gemma 4 on pegainfer: the 12B dense line and the 26B-A4B routed line, build with the gemma4 feature, launch, the memory envelope a slot count buys, long-context profiles, and heterogeneous attention notes." tableOfContents: minHeadingLevel: 2 maxHeadingLevel: 3 --- -Gemma 4 uses two kinds of attention in one model. Five of every six layers use sliding-window local attention at head dim 256. Every sixth layer, and the last one, uses global attention at head dim 512 with a single KV head. pegainfer serves the 12B text stack behind the `gemma4` cargo feature. Each kind gets its own paged KV pool, and decode runs as bucketed CUDA graphs with per-iteration scheduling. +Gemma 4 uses two kinds of attention in one model. Five of every six layers use sliding-window local attention at head dim 256. Every sixth layer, and the last one, uses global attention at head dim 512 with a single KV head. pegainfer serves the Gemma 4 text stack behind the `gemma4` cargo feature. Each kind gets its own paged KV pool, and decode runs as bucketed CUDA graphs with per-iteration scheduling. + +Two checkpoints are served today. 12B is the dense line. 26B-A4B is the routed one: 128 routed experts with top-8 routing beside a dense branch on every MoE layer, served from the published NVFP4 checkpoint with the expert weights kept packed on the device. Everything below applies to both lines unless a section says otherwise. ## Build @@ -37,6 +39,16 @@ curl -s localhost:8000/v1/completions -H 'Content-Type: application/json' \ Gemma 4's published generation defaults are sampled, not greedy (`temperature 1.0`, `top_k 64`, `top_p 0.95`), so a greedy comparison has to pass `"temperature": 0` explicitly. +The routed line launches the same way, pointed at the NVFP4 checkpoint: + +```bash +target/release/pegainfer \ + --model-path models/Gemma-4-26B-A4B-NVFP4 \ + --port 8000 +``` + +The loader validates the routed layer's manifest before the first upload and fails closed on a missing, extra, mis-shaped or mis-typed tensor. The experts are repacked once at load for the Marlin NVFP4 GEMM and read in their stored format from then on; the dense projections stay BF16. + ## Memory footprint The KV pools are allocated for every decode slot at startup, so the footprint is fixed before the first request arrives. At the defaults (16 slots, an 8192-token ceiling) 12B BF16 sits at 32.3 GiB while idle. That is 22.18 GiB of weights, 9.27 GiB of pools, and the rest CUDA context, RoPE tables, step buffers and the captured decode graphs. A 32 GiB card cannot start that configuration. @@ -53,6 +65,8 @@ These figures were measured on a 48 GiB card, which is the hardware used through | 128K ceiling x 8 slots | 43.5 GiB | | 262K ceiling x 4 slots | 42.6 GiB | +26B-A4B on the long-context profile below held 29,917 MiB at its peak through the measurement further down, weights and pools included. That is a peak under load rather than the idle resident figure this table reports. + ## Scheduling One engine loop admits whatever the pools can hold, up to the slot count. Then every active request advances one token in a single batched decode step, sharing one pass over the weights. @@ -104,25 +118,31 @@ At the full ceiling a 49K-token prompt returns its first token in about 13 s, an Smaller segments are nearly free here. With two ~49K prompts arriving together, a live stream's p99 gap falls from 2613–2673 ms at `N=8192` to 381–403 ms at `N=1024`, and admission latency does not move. At `N=512` the per-step cost shows up: the stall drops to 214–224 ms, but admission and wall time cost about 13% more. `N=1024` is the recommended long-context profile; `512` is the tail-protective option. +The long-context profile both measurements below use is: + +```bash +PEGAINFER_MAX_CONTEXT=262144 PEGAINFER_MIX_CHUNK_TOKENS=2496 PEGAINFER_DECODE_SLOTS=2 +``` + ## Performance -Measured against vLLM 0.23.1rc1 on the same card. vLLM's own `vllm bench serve` client drove both engines, so the metric definitions, the percentile maths and the request pacing are vLLM's code, used the same way on both sides. +Both lines were measured against vLLM 0.23.1rc1 on the same card. vLLM's own `vllm bench serve` client drove both engines, so the metric definitions, the percentile maths and the request pacing are vLLM's code, used the same way on both sides. -Single GPU (sm_89, x86_64), 49,140 MiB, driver 570.211.01, CUDA 12.9. Checkpoint `google/gemma-4-12B-it` at revision `707f0a3b8a3c`. BF16 weights and KV on both engines. pegainfer `e38dfdd8`. +Single GPU (sm_89, x86_64), 49,140 MiB, driver 570.211.01, CUDA 12.9. -![Time to first token and time per output token versus prompt length, PegaInfer against vLLM, at 10.6K / 40K / 81.7K / 163K prompt tokens](/models/gemma4/perf.svg) +The protocol is the same for both lines. One request in flight, four prompt lengths, four rounds interleaved so each engine leads twice, and the prompt-length order rotating each round so no length always runs in the same thermal position. Each engine started fresh per phase, and startup time enters no metric. The seed is the round number, so both engines within a round see identical prompts and different rounds use different ones. One discarded warm-up per length per boot, then three kept requests — 12 per engine per length, 96 in all. A cell counted only if every generation was exactly 256 tokens, and all 32 passed on both lines. Prefix caching is off on both sides, and random prompts share no prefix by construction. -*One request in flight, four prompt lengths. Medians of four interleaved rounds, 12 kept requests per point.* +vLLM ran at its defaults with `--max-model-len 262144` and `--gpu-memory-utilization 0.92`, which restate what it would have chosen anyway, plus `--no-enable-prefix-caching` and `--language-model-only` for symmetry. No `VLLM_*` environment variable was set. It picked a 2496-token chunked prefill width for itself, and ran with `torch.compile`, CUDA graphs and FlashInfer sampling. pegainfer ran the long-context profile above, with the prefill width set to vLLM's choice. -### Method +### 12B against vLLM -vLLM ran at its defaults with four flags: `--max-model-len 262144` and `--gpu-memory-utilization 0.92`, both of which restate what it would have chosen anyway, plus `--no-enable-prefix-caching` and `--language-model-only` for symmetry. No `VLLM_*` environment variable was set. It then picked a 2496-token chunked prefill width for itself, and ran with `torch.compile`, CUDA graphs and FlashInfer sampling. pegainfer ran the documented long-context profile: `PEGAINFER_MAX_CONTEXT=262144`, `PEGAINFER_MIX_CHUNK_TOKENS=2496`, `PEGAINFER_DECODE_SLOTS=2`. +Checkpoint `google/gemma-4-12B-it` at revision `707f0a3b8a3c`. BF16 weights and KV on both engines. pegainfer `e38dfdd8`. -The two widths are close but not equal. pegainfer was given vLLM's 2496 rather than a width of its own choosing, but its steps round down to whole 128-row tiles, so the effective width was 2432. vLLM will not go below 2496 for this model, so the two cannot be set to the same number. The smaller width means more steps on the pegainfer side. +![Time to first token and time per output token versus prompt length, PegaInfer against vLLM, at 10.6K / 40K / 81.7K / 163K prompt tokens](/models/gemma4/perf.svg) -Four rounds, boot order alternating so each engine leads twice, prompt-length order rotating each round so no length always runs in the same thermal position. Each engine started fresh per phase; startup time enters no metric. The seed is the round number, so both engines within a round see identical prompts. One discarded warm-up per length per boot, then three kept requests — 12 per engine per length, 96 in all. A cell counted only if every generation was exactly 256 tokens; all 32 passed. Prefix caching off on both sides, and random prompts share no prefix by construction. +*One request in flight, four prompt lengths. Medians of four interleaved rounds, 12 kept requests per point.* -### Results +The two prefill widths are close but not equal. pegainfer was given vLLM's 2496 rather than a width of its own choosing, but its steps round down to whole 128-row tiles, so the effective width was 2432. vLLM will not go below 2496 for this model, so the two cannot be set to the same number. The smaller width means more steps on the pegainfer side. The figure plots medians. The table below adds end-to-end latency, and the paired deltas computed inside each round, where both engines saw the same prompts. Negative means pegainfer is faster. @@ -139,9 +159,41 @@ The shortest prompt is the one to be careful with. At 10,602 tokens the two engi Peak GPU memory, sampled every 100 ms, was 34,354 MiB against vLLM's 47,308 MiB. vLLM's number is what its 0.92 utilization setting reserves up front, not what it actually used. +### 26B-A4B against vLLM + +The ModelOpt NVFP4 export of `google/gemma-4-26B-A4B-it`. pegainfer `e7a41975`. + +This run adds one flag to the four above: `--kv-cache-dtype bfloat16`. The checkpoint's quantization config declares an FP8 KV cache, and vLLM honours it by storing the cache as fp8_e4m3 by default. pegainfer stores BF16, so the like-for-like comparison sets vLLM to BF16 too, and the second table below reports vLLM at its own default, because that default costs it real time at depth. + +![Time to first token and time per output token versus prompt length, pegainfer against vLLM at like-for-like KV precision, at 10.6K / 40K / 81.7K / 163K prompt tokens](/models/gemma4/perf-26b.svg) + +*One request in flight, four prompt lengths, the KV cache at BF16 on both engines. Medians of four interleaved rounds, 12 kept requests per point.* + +| prompt tokens | E2EL pegainfer | E2EL vLLM | E2EL Δ | paired Δ TTFT | paired Δ TPOT | paired Δ E2EL | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 10,602 | 3.20 s | 3.30 s | −2.9% | −4.7 … −8.8% | −0.0 … −0.8% | −1.5 … −3.4% | +| 40,002 | 7.27 s | 8.44 s | −13.8% | −16.4 … −18.4% | −4.7 … −5.8% | −12.8 … −14.5% | +| 81,653 | 16.28 s | 21.45 s | −24.1% | −26.2 … −27.5% | −6.1 … −7.0% | −23.6 … −24.8% | +| 163,336 | 43.91 s | 72.98 s | −39.8% | −41.3 … −41.4% | −7.0 … −7.7% | −39.7 … −39.8% | + +Round-to-round spread is at most 3.4%, on vLLM at the shortest prompt; pegainfer stays within 1.9% everywhere and 0.7% at the longest prompt. + +**vLLM at its own default.** The same protocol with vLLM storing its KV cache as fp8_e4m3, the four shared flags only, against pegainfer at BF16 KV, measured the same morning on the commit one frontend change earlier (`ea02a9f7`; the change since removes a per-request cost of about 40 ms on pegainfer's side, visible only at the shortest prompt): + +| prompt tokens | E2EL pegainfer | E2EL vLLM (fp8 KV) | E2EL Δ | paired Δ TTFT | paired Δ TPOT | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 10,602 | 3.23 s | 3.12 s | +3.4% | +5.6 … +12.7% | +2.3 … +2.7% | +| 40,002 | 7.31 s | 7.28 s | +0.4% | −0.8 … +2.7% | +1.1 … +1.5% | +| 81,653 | 16.24 s | 16.51 s | −1.6% | −3.2 … −2.0% | +3.5 … +4.3% | +| 163,336 | 43.93 s | 45.91 s | −4.3% | −5.7 … −4.9% | +8.0 … +10.3% | + +pegainfer's numbers are the same in both tables; the difference between them is what the FP8 cache buys vLLM's Triton attention at depth, where its 163K time to first token drops from 69.6 s to 43.0 s. At equal precision the gap grows with prompt length. At vLLM's default the two engines are close, with vLLM ahead at 10.6K and pegainfer ahead from 81.7K on. + +Peak GPU memory, sampled every 100 ms, was 29,917 MiB against vLLM's 45,641 MiB. vLLM's number is what its 0.92 utilization setting reserves up front, not what it actually used. + ### Caveats -Part of the gap at long prompts comes from the kernels, not from the engines as a whole. vLLM serves this checkpoint through its Triton attention backend, which its configuration layer selects because FlashAttention does not cover Gemma 4's two head dimensions in this version. pegainfer has kernels written for those two dimensions. So the deeper points compare what each engine ships for this model today, not two implementations of the same kernel, and we did not measure how much of the gap that accounts for. +Part of the gap at long prompts comes from the kernels, not from the engines as a whole. vLLM serves both checkpoints through its Triton attention backend, which its configuration layer selects because FlashAttention does not cover Gemma 4's two head dimensions in this version. pegainfer has kernels written for those two dimensions. On the routed line vLLM also decompresses the experts through its Marlin NvFp4 path, because the card has no native FP4; pegainfer's expert GEMM is a Marlin path too. So the deeper points compare what each engine ships for these models today, not two implementations of the same kernel, and we did not measure how much of the gap that accounts for. The two engines also get here differently. vLLM covers 262,144 positions out of the box, while pegainfer refuses a prompt this long until someone sets the profile above. That extra setup is on pegainfer's side. Within that profile, the one setting we could have picked in our own favour is the prefill chunk width, and we used vLLM's. @@ -184,6 +236,7 @@ Decode steps run at power-of-two batch sizes and replay as CUDA graphs captured - **Global layers have no `v_proj` tensor.** `attention_k_eq_v` is built into the checkpoint rather than decided at runtime: V is taken from the K projection's output. The two split straight after that. K goes through `k_norm` and then RoPE, V through a scale-free norm and no RoPE. Both are still written out, so the saving on global KV is in how the cache is stored, not a pointer alias. - **The proportional RoPE is not ordinary partial RoPE.** Only 128 of the global head's 512 dimensions rotate, but the frequency denominator stays the full head dimension. - **The text graph also needs** scaled token embeddings, RMSNorm before and after both attention and the feed-forward block, a scale-free V norm, tied embeddings, GELU-tanh, a residual layer scalar, and final logit softcapping. +- **The routed line adds a parallel expert branch**, not a shared-expert variant: every MoE layer computes its dense branch and its 128 routed experts at top-8, then combines the two. The attention and KV stack is the one above. - 12B ships one `model.safetensors` with no index file. Its vision and audio tensors sit under `model.vision_embedder.*`, `model.embed_vision.*` and `model.embed_audio.*`. Serving is text-only: those tokens are rejected before embedding and suppressed before sampling. - `eos_token_id` appears three times in the checkpoint's config files, with three different values. `generation_config.json` is the one the engine follows. @@ -193,4 +246,6 @@ Decode steps run at power-of-two batch sizes and replay as CUDA graphs captured - **No prefix sharing between requests.** Two live requests with the same prefix each pay for it. The conversation cache above works across turns of one conversation, not across concurrent requests. - **Prompts prefill whole by default**, unless the chunked walk limits them. - **KV capacity is not reported to the frontend**, so its capacity metrics stay empty for this model. -- 26B-A4B and 31B are on the roadmap, not served today. +- **26B-A4B serves the NVFP4 checkpoint only.** BF16 experts are not served. +- **Text only.** Multimodal inputs are not supported, and the checkpoints' vision tower is skipped at load. +- 31B is on the roadmap, not served today.