ZAYA1 Model Architecture

• ZAYA1 model architecture is a mixture-of-experts transformer that integrates MI300X-aware tuning, custom convolutional attention, and expert routing to optimize large-scale training.
• The design incorporates per-layer residual scaling, rotary embeddings, and specialized AMD-specific kernels to maximize throughput and minimize latency.
• It achieves a competitive balance between dense and MoE components, yielding strong evaluation results across tasks like reasoning, mathematics, and coding.

The ZAYA1 model architecture is a mixture-of-experts (MoE) transformer designed for large-scale training on AMD MI300X GPUs with Pollara interconnect. ZAYA1-base incorporates a suite of systems and modeling innovations tailored to the AMD hardware stack, including MI300X-aware dimensioning, custom convolutional attention mechanisms, per-layer residual scaling, and expert routing. The architecture achieves a competitive balance of training throughput and inference latency with strong evaluation results across tasks, establishing the maturity of AMD’s distributed compute environment for state-of-the-art pretraining (Anthony et al., 21 Nov 2025).

1. Overall Model Structure

ZAYA1-base is built with $L=40$ transformer layers and an embedding dimension $h=2048$. The vocabulary size is $v=262{,}272$, chosen to be divisible by 64 for optimized device throughput. Each transformer layer contains an MoE block comprising $E=16$ experts, with a top-$k=1$ expert selected per token at each routing step. This yields $8.3$ billion total parameters (considering all experts) but an “active” parameter count of $760$ million (the dense backbone plus one expert per token path).

The forward path through each transformer layer $\ell$ follows this sequence:

1. Residual-scaled RMSNorm $\rightarrow$ Compressed Convolutional Attention (CCA) $\rightarrow$ residual add
2. Residual-scaled RMSNorm $h=2048$0 ZAYA1 Router gating $h=2048$1 expert MLP (MoE) $h=2048$2 residual add
3. Final RMSNorm

Residual scaling is implemented on every residual path via per-channel learnable gates.

2. Transformer Layer Components

Attention and Token Path

CCA attention receives input $h=2048$3 and projects it to queries, keys, and values with the following details:

• $h=2048$4 total attention heads, each with head dimension $h=2048$5
• Query heads: $h=2048$6 ($h=2048$7)
• Key/value heads: $h=2048$8 ($h=2048$9)

Projections:

• $v=262{,}272$0
• $v=262{,}272$1 with analogous dimensions (with $v=262{,}272$2 and $v=262{,}272$3 handling half of key/value each)

CCA then applies a convolutional stage:

• Depthwise conv1d ($v=262{,}272$4) plus grouped conv1d (groups $v=262{,}272$5, $v=262{,}272$6) along the sequence
• FlashAttention operates in a compressed latent space of size $v=262{,}272$7
• Rotary position embeddings (RoPE) are applied to half the channels of each head, supporting 4k–1M context extension

Outputs are projected back via $v=262{,}272$8, followed by RMSNorm (with $v=262{,}272$9) and per-head key temperature.

MoE and Routing

MoE routing in ZAYA1 involves the following operations for each token:

• Down-projection: $E=16$0, where $E=16$1
• Exponential Depth Averaging (EDA): $E=16$2 (with learned scalar $E=16$3)
• Outputs go to a 3-layer MLP (GeLU activations), yielding logits $E=16$4
• Post-softmax, each token's expert is selected as $E=16$5, with bias vector $E=16$6

The chosen expert’s MLP has weights:

• First FC: $E=16$7, with $E=16$8 (hidden expansion factor $E=16$9)
• Activation: SwiGLU across pre-activation width $k=1$0
• Second FC: $k=1$1 where $k=1$2
• Followed by residual addition and RMSNorm

3. MI300X-Aware Sizing Principles

The architecture’s sizing rules and GEMM shapes are directly informed by MI300X hardware characteristics:

• All core dimensions ($k=1$3) are set as multiples of 64, maximizing rocBLAS/hipBLASLt performance
• Microbatch product $k=1$4 is divisible by 64, and $k=1$5 is integer to avoid padding overhead
• MLP expansion factor is fixed ($k=1$6)
• MoE per-layer parameter count: $k=1$7
• Convolutional and attention kernel sizes, e.g., $k=1$8, are chosen based on MI300X TFLOPs heatmaps to maximize utilization

These practices are derived from explicit MI300X benchmarking, targeting “hot” performance regions for compute and memory transfers.

4. AMD-Specific Kernels and Communication

The model stack incorporates several AMD-specific optimizations:

Component	Optimization/Detail
CCA conv kernels	Tuned for MI300X HBM2 bandwidth and warp size
Custom HIP kernels	Multi-tensor Muon optimizer kernels; fused residual-add + RMSNorm kernels (two-stage)
Communication	Gradient-fusion buffer sizes saturate Pollara 400 Gbps at break-even; ZeRO-1/context-parallel worlds aligned to xGMI hardware node boundaries

The optimization of collective communication primitives (all-reduce, reduce-scatter, all-gather, broadcast) as well as kernel fusion is critical for training throughput on MI300X + Pollara platforms.

5. Parameter and Compute Profile

Per-layer parameter and FLOPs breakdown, with $k=1$9 (no tensor/data parallelism) (Anthony et al., 21 Nov 2025):

Component	Parameter Count (per layer)	FLOPs per token (approx.)
Attention Q,K,V,O	$8.3$0 M	$8.3$1 kM
CCA convs + RoPE	—	$8.3$2 GFLOPs
Router down-proj	$8.3$3 M	$8.3$4 kM
Router MLP (2)	$8.3$5 M	$8.3$6 kM
Router logits	$8.3$7 M	$8.3$8k
Expert FC1	$8.3$9 M	$760$0 kM
Expert FC2	$760$1 M	$760$2 kM
Residual scaling	negligible ($760$30.004 M)	$760$40.1k

Total per-layer parameters: $760$5 M Total per-layer FLOPs per token: $760$6 k A forward pass over $760$7 tokens, $760$8, totals $760$9 M FLOPs per layer; all 40 layers give $\ell$0 G FLOPs per sample. Inference latency is dominated by expert MLPs (60%), attention kernels (30%), and routing/norms (10%).

6. Special Architectural Components

• Embeddings: Token embeddings $\ell$1, tied with the LM head.
• Normalization: All RMSNorm (no learnable bias); router MLP uses standard LayerNorm before GeLU.
• Activation Functions: GeLU in router blocks; SwiGLU within expert MLPs.
• Rotary Embeddings: RoPE is applied to half of each head’s channels only, supporting long-context extrapolation.
• Residual Scaling: Per-layer parameterized by $\ell$2 and bias $\ell$3:

$\ell$4

• CCA Compression: Query compression $\ell$5, key/value compression $\ell$6, denoted "CCGQA" in model documentation.

7. Comparative Performance and Context

ZAYA1-base achieves performance at or above leading models of similar and larger active scale (Qwen3-4B, Gemma3-12B) and outperforms Llama-3-8B and OLMoE on benchmarks targeting reasoning, mathematics, and coding. The empirical findings suggest that the combination of tailored architecture and hardware-aware engineering enables the AMD stack to match or exceed the competitiveness of established foundation model pretraining environments (Anthony et al., 21 Nov 2025).