---
title: ZAYA1 Model Architecture
url: https://www.emergentmind.com/topics/zaya1-model-architecture
type: topic
---

# ZAYA1 Model Architecture

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 [2511.17127].

## 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 $\rightarrow$ ZAYA1 Router gating $\rightarrow$ expert MLP (MoE) $\rightarrow$ 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 $x_\ell \in \mathbb{R}^{B \times S \times h}$ and projects it to queries, keys, and values with the following details:
- $a=16$ total attention heads, each with head dimension $d_h = h / a = 128$
- Query heads: $a_q = 8$ ($c_q = 1/2$)
- Key/value heads: $g = 2$ ($c_{kv} = 1/8$)

Projections:
- $W_Q \in \mathbb{R}^{h \times (a_q \cdot d_h)}$
- $W_K, W_{V1}, W_{V2}$ with analogous dimensions (with $V1$ and $V2$ handling half of key/value each)

CCA then applies a convolutional stage:
- Depthwise conv1d ($k_0=2$) plus grouped conv1d (groups $a_q + g$, $k_1=2$) along the sequence
- FlashAttention operates in a compressed latent space of size $(a_q \cdot d_h)$
- Rotary position embeddings (RoPE) are applied to half the channels of each head, supporting 4k–1M context extension

Outputs are projected back via $W_O \in \mathbb{R}^{(a_q \cdot d_h) \times h}$, followed by RMSNorm (with $\epsilon = 10^{-5}$) and per-head key temperature.

### MoE and Routing

MoE routing in ZAYA1 involves the following operations for each token:
- Down-projection: $W_{\text{down}} \in \mathbb{R}^{h \times D}$, where $D=256$
- Exponential Depth Averaging (EDA): $r_\ell = W_{\text{down}} x_\ell + \gamma \cdot r_{\ell-1}$ (with learned scalar $\gamma$)
- Outputs go to a 3-layer MLP (GeLU activations), yielding logits $W_{\text{gate},\ell} \in \mathbb{R}^{D \times E}$
- Post-softmax, each token's expert is selected as $e_{\text{idx}} = \arg\max_j(s_\ell + b_\ell)$, with bias vector $b_\ell \in \mathbb{R}^E$

The chosen expert’s MLP has weights:
- First FC: $W_{e1} \in \mathbb{R}^{h \times f}$, with $f = 4096$ (hidden expansion factor $\alpha=2$)
- Activation: SwiGLU across pre-activation width $f$
- Second FC: $W_{e2} \in \mathbb{R}^{f_o \times h}$ where $f_o = f / 2 = 2048$
- 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 ($h, d_h, D, f, f_o$) are set as multiples of 64, maximizing rocBLAS/hipBLASLt performance
- Microbatch product $b \cdot s \cdot h$ is divisible by 64, and $(b \cdot a) / t$ is integer to avoid padding overhead
- MLP expansion factor is fixed ($f = 2h, f_o = h$)
- MoE per-layer parameter count: $h \cdot f + f_o \cdot h$
- Convolutional and attention kernel sizes, e.g., $2048 \times 1024$, 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 $t=1$ (no tensor/data parallelism) [2511.17127]:

| Component           | Parameter Count (per layer)     | FLOPs per token (approx.)      |
|---------------------|-------------------------------|-------------------------------|
| Attention Q,K,V,O   | $\approx 5.2$ M               | $\approx 9.5$ kM              |
| CCA convs + RoPE    | —                             | $\approx 0.02$ GFLOPs         |
| Router down-proj    | $\approx 0.52$ M              | $\approx 1.05$ kM             |
| Router MLP (2)      | $\approx 0.13$ M              | $\approx 0.13$ kM             |
| Router logits       | $0.004$ M                     | $4.1$k                        |
| Expert FC1          | $\approx 8.39$ M              | $\approx 16.8$ kM             |
| Expert FC2          | $\approx 4.19$ M              | $\approx 8.4$ kM              |
| Residual scaling    | negligible ($\sim$0.004 M)     | $\sim$0.1k                    |

Total per-layer parameters: $\approx 18.4$ M  
Total per-layer FLOPs per token: $\approx 36$ k  
A forward pass over $S=1024$ tokens, $b=1$, totals $\approx 37$ M FLOPs per layer; all 40 layers give $\approx 1.5$ 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 $E_{\text{tok}} \in \mathbb{R}^{v \times h}$, 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 $\alpha, \beta \in \mathbb{R}^h$ and bias $b_\ell$:
  $$
  \text{ResScale}_\alpha(x) = \alpha \odot x + b_\ell
  $$
- **CCA Compression**: Query compression $2\times$, key/value compression $8\times$, 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 [2511.17127].

Source: https://www.emergentmind.com/topics/zaya1-model-architecture