[Model] Gemma 4 M2 26B-A4B

[FeathBow]

Description

Parent: #758. Crate: pegainfer-gemma4 · feature: gemma4.

M2 adds the first Gemma MoE to the green 12B text stack: exact router normalization and scaling, fp32 softmax with top-8 renormalization and per-expert scale, the parallel dense branch beside the routed expert path on every MoE layer, device-resident routing and grouping with skewed and empty experts handled, and a graph-safe fixed-capacity scratch.

M0 pointed the 26B quantization question at a W8A16 or int8-per-channel candidate. What this line serves is the published NVFP4 checkpoint through a Marlin path: expert weights stay packed on the device and the routed GEMM consumes NVFP4 directly instead of widening to bf16. BF16 experts remain the parent's stated target and are not served by this work, so the parent's target table needs NVFP4 recorded as the first validated 26B path.

Proposed Plan

• Load the checkpoint: the routed layer's manifest entries, the config it is keyed on, and NVFP4 tensor decode onto the device, with validation that fails closed on missing, extra, mis-shaped and mis-typed tensors before the first upload.
• Compute the experts in their stored format: the Marlin NVFP4 GEMM, the block alignment pass and the weight repack, gated against a reference dequantization rather than against final logits.
• Assemble the routed block — router, dense branch, routed branch, combine — and gate it with a brick oracle that compares router pre-softmax logits, dense output, routed output and the combined result separately, which is what M0 asked for and what a final-logits comparison cannot do.
• Wire the block into serving: decode graph capture for routed steps and the overlap lane, gated on HTTP liveness, cancellation and pressure recovery.
• Reduce the step cost with measured changes, each carrying its own A/B on one machine class with the full flag set, and each stating whether it is bit-invariant and showing it: the MoE kernels themselves, the layer and norm fusions, the per-step launch and upload overhead, an opt-in FP8 KV pool for the sliding family, and an opt-in admission door that lands with the measurement explaining why it is not the default.
• Land the changes this work needs in code that predates this line: moving the Marlin launch machinery somewhere both consumers can reach, giving the paged-KV launch bodies one home so the hd512 and prep twins stop being hand-copied, retiring the probe forward so the numeric gates drive the serving path, collapsing the duplicated first-token, retirement and admission contracts, and routing the hand-derived page offsets through the checked helper. Each is behavior-preserving and says so.
• Draw the shared-versus-local MoE boundary in writing, as the parent requires, and extract only what a second consumer already exercises. On current evidence that is the Marlin launch machinery, the alignment pass, the repack and the top-k row reduction; router normalization and scaling, the fp32 softmax and top-8 renormalization, per-expert scales, the NVFP4 weight layout and the dense-branch combine stay model-local.
• Run and retain the same-host comparison against a pinned vLLM version, record peak observed device memory for the served envelope, and publish the recipe and hardware-matrix entry.

Acceptance Criteria

• The router and expert path compared component-wise against a reference implementation, not attributed from final logits alone.
• Full-model logits against the frozen HF fixtures for the single-token, short and window-edge cases, and greedy token-exact continuation against the reference.
• Ragged batch replay and reorder identity, HTTP liveness, cancellation and pressure recovery all green on the served checkpoint.
• Peak observed device memory recorded for the served envelope, alongside the retained benchmark snapshot and its thresholds, with both the closed and open loop present.
• The shared-versus-local MoE boundary written down, and the parent's target table corrected to record NVFP4 as the first validated 26B path.
• A recipe and hardware-matrix entry for the line, with unsupported options failing explicitly.