Add Gemma 4 31B IT contrib model

[jimburtoft]

Summary

• Adds NxDI contrib implementation of google/gemma-4-31b-it text decoder (31B parameters)
• Validated on trn2.3xlarge (TP=4, bfloat16): 32.1 tok/s, 65ms TTFT, Pearson r=0.980 vs HF CPU reference
• Vision encoder port is planned as a follow-up PR

Architecture Highlights

Gemma 4 31B has several features not present in the existing Gemma 3 contrib:

Feature	Details
Heterogeneous layers	SWA layers (head_dim=256, 16 KV heads) and Global layers (head_dim=512, 4 KV heads)
attention_k_eq_v	Global layers share K/V projections — V = K before normalization
QK + V normalization	RMSNorm on Q/K after projection; separate RMSNorm (no learnable scale) on V
layer_scalar	Per-layer learned multiplicative factor applied at end of decoder forward
final_logit_softcapping	30 * tanh(logits / 30) after lm_head
Partial RoPE	Global layers rotate only 25% of dims (128 of 512)
Key implementation detail: NxDI applies 1/sqrt(head_dim) scaling in attention, but Gemma 4 uses scaling=1.0. Naively scaling Q weights doesn't work because q_norm (RMSNorm) is scale-invariant. The fix is to scale the q_layernorm.weight parameters by sqrt(head_dim) instead, which survives normalization.

Test Results

All 7 integration tests pass on trn2.3xlarge (TP=4, batch_size=1, bfloat16):

Test	Result
Smoke (model loads)	PASS
Text generation	PASS — 10 tokens generated
Token matching	PASS — "France" in output for capital-of-France prompt
Chat generation	PASS — "2 + 2 = 4" with chat template
Coherence	PASS — coherent haiku
TTFT	PASS — 65.1ms (threshold: 200ms)
Throughput	PASS — 32.1 tok/s, TPOT=31.1ms (threshold: 10 tok/s)

Files

contrib/models/gemma-4-31b-it/
├── README.md                              # Usage, compatibility matrix, test results
├── src/
│   ├── __init__.py                        # Exports NeuronGemma4ForCausalLM, Gemma4InferenceConfig
│   └── modeling_gemma4.py                 # 1070 lines — full text decoder implementation
└── test/
    ├── __init__.py
    ├── integration/
    │   ├── __init__.py
    │   └── test_model.py                  # 7 tests (pytest + standalone runner)
    └── unit/
        └── __init__.py

Prerequisites

• Neuron SDK 2.28+ (trn2.3xlarge with LNC=2)
• transformers >= 5.5.0 (required for Gemma4ForConditionalGeneration)
• transformers/utils/fx.py shim needed (NxD imports HFTracer which was removed in transformers 5.x — see README for details)
• attn_kernel_enabled=False required (head_dim > 128 exceeds NKI flash attention limit)

Known Limitations

• Text-only: This PR covers the text decoder only. The Gemma 4 vision encoder (gemma4_vision) is not yet ported. A follow-up PR will add VLM (vision + language) support.
• No flash attention: Both SWA (head_dim=256) and global (head_dim=512) layers exceed the NKI kernel's 128 limit, requiring decomposed attention.
• CONVERT_TO_MHA warnings: Benign warnings from NxDI about TP=4 vs 16 KV heads; does not affect correctness.

Checklist

• At least one accuracy test (token matching + logit correlation)
• README with usage example
• README with compatibility matrix
• README with checkpoint links
• README with test instructions