Extends flex fast-inference to `unsloth/gemma-4-26B-A4B-it` (30 layers,
128 experts top-k 8, H=2816, ~3.8B active of 25.2B). Mirrors the
FlexGptOssInference / FlexMoEInference template with Gemma 4-specific
wiring:
- Dual dense MLP + MoE per decoder layer (Gemma4TextMLP alongside
Gemma4TextExperts; outputs summed before the residual add, then
multiplied by the per-layer `layer_scalar` buffer).
- Per-layer sliding-window dispatch (25 sliding @ 1024 tokens, 5 full
attention) via twin BlockMask built once per generate() entry.
- Two-tier RoPE: sliding layers use rope_theta=10K with full head_dim
rotation; full-attn layers use rope_type=proportional with theta=1M
and partial_rotary_factor=0.25 (the inv_freq's zero-padded tail makes
the generic rotate_half a no-op on the unrotated dims, so a single
rotary helper covers both).
- Per-head Q/K/V RMSNorm applied before RoPE / KV write.
- attention_k_eq_v=True on full-attn layers (v_proj is None): value is
the raw k_proj output, followed only by v_norm (with_scale=False).
- Rebind Gemma4TextExperts.forward to forward_native_grouped_mm so
decode uses the grouped_mm backend (the slow Python loop in the stock
forward is neither fast nor CUDA-graph-capturable). The routing
weights already include per_expert_scale via Gemma4TextRouter.forward,
so no extra folding is needed.
- CUDA graph capture and UNSLOTH_FLEX_COMPILE_WALKER=1 inherit from the
MoE template (single bucket ladder, single pool across buckets).
Arch detection:
- `_detect_arch` distinguishes dense vs MoE Gemma 4 via
`text_config.num_experts > 1` (same class name covers both variants).
- `bind_peft_model` extends the MoE no-deepcopy shortcut to gemma4_moe.
Validation (B200, bf16, 3 chat prompts, 64 tokens):
- Dense 31B sanity check: coherent completion via existing
FlexGemma4Inference.
- Flex (cudagraph) vs HF naive: 24/24, 2/2, 4/4 tokens bitwise match.
- Merge parity: 7/7 cases bitwise (rank 16 / 64, 1 + 2 adapters, bf16
+ fp32, E=128, 2I=1408, H=2816 / H, I=704).
- Throughput bs=8/16/32/48: 510 / 1030 / 1564 / 916 tok/s vs HF naive
134 tok/s at bs=8 (3.8x-11.7x).
- GRPO smoke: DAPO-Math-17k, seed 3407, max_steps=5 — stable (see
follow-up comment on PR).
Out of scope: E2B/E4B KV-shared + per-layer-input variants (guarded
with NotImplementedError); bnb-4bit stacked experts (no such class
ships for Gemma 4 today).
749fb75bd0