[contrib] Add Gemma 2 9B NeuronX port

[dhwanw]

Description

Adds NeuronX port of Gemma 2 9B (model_type=gemma2) to the contrib models collection.

Model Information

Field	Value
Model	google/gemma-2-9b
Architecture	Gemma2ForCausalLM (decoder-only, GQA, 4 RMSNorm layers per block)
Parameters	9B
TP Degree	8
Precision	BF16

Checklist

• Model compiles successfully on Neuron
• Token matching validated (84.69% greedy, 99.38% teacher-forced)
• Performance profiled (47.1 tok/s)
• README with architecture details, usage, validation results
• Integration tests included

Folder Structure

contrib/models/gemma-2-9b/
├── README.md
├── src/
│   ├── __init__.py
│   └── modeling_gemma2.py
└── test/
    └── integration/
        └── test_model.py

Testing

• Token Match (greedy): 84.69% (10 prompts, 32 tokens each)
• Token Match (teacher-forced): 99.38%
• Throughput: 47.1 tok/s (TP=8, BS=1, seq_len=128)
• 8/10 prompts at 100% greedy match

Note: Uses custom Gemma2NeuronConfig with attn_cls=NeuronGemma2Attention. Attention logit softcapping disabled (NKI kernel limitation). Sliding window disabled (head_dim=256 exceeds NKI limit of 128). Final logit softcapping (30.0) omitted from forward pass since tanh is monotonic and doesn't affect greedy argmax.

Compatibility

• Neuron SDK: 2.22+
• Instance: trn1.32xlarge

🤖 Generated with Claude Code