# Mixtral MoE Block

> Mixtral decoder block — GQA + Sparse MoE (8 experts, top-2) + RMSNorm + RoPE

Pick when you have MoE training infra and want best quality per active param. Routing instability and memory cost (still scales with total params) are the trade-off.

- Category: NLP/LLM
- Layers: 9
- Parameters: 1.45B
- Input shape (batchless): 1 × 4096 × 4096
- Output shape: 1 × 4096 × 4096
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/mixtral-block/model.json
- Open on the canvas: https://neurarch.com/?template=mixtral-block

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | hidden_states | Input | shape=[1, 4096, 4096] | 1 × 4096 × 4096 |
| 2 | input_norm | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 3 | self_attn | Grouped Query Attn | embedDim=4096, numHeads=32, numKVHeads=8 | 1 × 4096 × 4096 |
| 4 | rotary_emb | RoPE |  |  |
| 5 | attn_residual | Add |  | 1 × 4096 × 4096 |
| 6 | post_attn_norm | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 7 | block_sparse_moe | MoE Layer | embedDim=4096, numExperts=8, topK=2 | 1 × 4096 × 4096 |
| 8 | moe_residual | Add |  | 1 × 4096 × 4096 |
| 9 | hidden_out | Output |  | 1 × 4096 × 4096 |

## Verifier findings

- **info** `moe-no-aux-loss` at `block_sparse_moe`: MoE layers require an auxiliary router z-loss + load-balance loss during training to prevent expert collapse. This is not visible in the architecture diagram but must be in the training loop. Fix: Add a note on this layer. Typical aux_loss coefficient: 1e-2 (Mixtral/Switch Transformer).

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class MixtralMoEBlock(nn.Module):
    def __init__(self):
        super().__init__()

        self.rmsNorm_1 = nn.RMSNorm(4096)
        self.groupedQueryAttention_1 = nn.ModuleDict({
            'q_proj': nn.Linear(4096, 4096,        bias=False),   # 32 heads × 128
            'k_proj': nn.Linear(4096, 1024, bias=False),   # 8 KV heads × 128
            'v_proj': nn.Linear(4096, 1024, bias=False),
            'o_proj': nn.Linear(4096, 4096,        bias=False),
        })  # GQA: 32Q / 8KV heads (requires F.scaled_dot_product_attention)
        self.rmsNorm_2 = nn.RMSNorm(4096)
        self.moeLayer_1 = nn.ModuleDict({
            'router': nn.Linear(4096, 8, bias=False),
            'experts': nn.ModuleList([
                nn.Sequential(
                    nn.Linear(4096, 14336, bias=False), nn.SiLU(),
                    nn.Linear(14336, 4096, bias=False),
                ) for _ in range(8)
            ]),
        })  # MoE top-2

    @staticmethod
    def _moe_forward(moe, x, top_k=2):
        """Top-k routed MoE forward over a {'router', 'experts'} ModuleDict."""
        scores = moe['router'](x).softmax(dim=-1)
        top_w, top_i = scores.topk(top_k, dim=-1)
        top_w = top_w / top_w.sum(dim=-1, keepdim=True)
        flat_x = x.reshape(-1, x.size(-1))
        flat_i = top_i.reshape(-1, top_k)
        flat_w = top_w.reshape(-1, top_k)
        out = torch.zeros_like(flat_x)
        for e_idx in flat_i.unique():
            hit = flat_i == e_idx
            rows = hit.any(dim=-1)
            w = (flat_w * hit).sum(dim=-1)[rows].unsqueeze(-1)
            out[rows] += w * moe['experts'][int(e_idx)](flat_x[rows])
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md