Add ZAYA1-base contrib model (MoE with CCA attention)

[jimburtoft]

Summary

• Adds Zyphra's ZAYA1-base (8.84B total / 800M active, MoE with CCA attention) to contrib
• Full NxDI implementation with custom CCA attention, non-linear MLP router, and MoD (Mixture of Depths)
• Includes vLLM-neuron serving support with near-linear throughput scaling

Model Details

• Source: Zyphra/ZAYA1-base (Apache 2.0)
• Architecture: 80 layers alternating attention (CCA) and MoE, 16 real + 1 skip expert, partial RoPE
• Parameters: 8.84B total / 800M active per token
• Instance: trn2.3xlarge (TP=2)
• Prerequisite: Zyphra's custom transformers fork (pip install transformers @ git+https://github.com/Zyphra/transformers.git@zaya)

Performance

Metric	Neuron trn2 (TP=2)	GPU A10G	Advantage
TKG (batch=1)	22.3 tok/s	6.2 tok/s	3.6x
TKG (batch=4)	86.3 tok/s	23.9 tok/s	3.6x
TTFT	75.6 ms	184.7 ms	2.4x

vLLM serving: 72.3 tok/s at concurrency=4 with near-linear scaling.

Key Technical Contributions

• ManualConv1d: Replaces nn.Conv1d to avoid NCC_ITEN404 crash after all-gather
• SPMDRank: Per-rank tensor extraction for SPMD tracing compatibility
• CCA state caching: Conv states and previous hidden states persisted via input_output_aliases
• Static expert dispatch: Mask-based dispatch replacing unsupported torch.bincount
• SDK 2.29 validated: Tested PASS with 21.1 tok/s

Testing

• 7 files: README.md, src/{init,modeling_zaya}.py, test/{init,integration/{init,test_model},unit/init}.py
• Integration tests: smoke test, prefill accuracy (" Paris"), batch independence (4 prompts), token matching
• Validated on trn2.3xlarge with SDK 2.27, 2.28, and 2.29