docs/ARCHITECTURE.md

# Kiến trúc Nexus Coder / Nexus Coder Architecture

## Tổng quan / Overview

Nexus Coder v0.1 sử dụng kiến trúc **Mixture of Experts (MoE) Transformer** tương tự Mixtral 8x7B và DeepSeek-V3.

Nexus Coder v0.1 uses a **Mixture of Experts (MoE) Transformer** architecture similar to Mixtral 8x7B and DeepSeek-V3.

## Các thành phần / Components

### 1. Token Embedding
- Vocab size: 32,000
- Hidden size: 2,048
- Tokens được nhúng thành vector 2048 chiều

### 2. Grouped Query Attention (GQA)
- 16 query heads
- 4 KV heads (ratio 4:1)
- Head dimension: 128
- Giảm 4x memory cho KV cache so với MHA truyền thống

### 3. Rotary Position Embedding (RoPE)
- Base: 10,000
- Hỗ trợ tối đa 50,000 positions
- Cho phép model hiểu vị trí tương đối giữa các tokens

### 4. RMSNorm
- Thay thế LayerNorm truyền thống
- Không có bias, không trừ mean
- Nhanh hơn ~10-20%

### 5. SwiGLU Activation
- `SiLU(gate(x)) * up(x)`
- Hiệu quả hơn ReLU/GELU
- Có 3 ma trận: gate, up, down (3 * hidden * intermediate params)

### 6. Mixture of Experts (MoE) - Cốt lõi
- **24 experts** tổng cộng (mỗi expert là một SwiGLU FFN)
- **3 active experts** mỗi token (top-3 routing)
- Router: linear layer (hidden_size → num_experts)
- Load balancing loss: auxiliary loss để tránh expert collapse

#### Routing Algorithm
```
1. Router tính gate_logits = W_router @ x
2. routing_weights = softmax(gate_logits)
3. top_k_weights, top_k_indices = topk(routing_weights, k=3)
4. Normalize top_k_weights
5. Mỗi token đi qua 3 expert được chọn
6. Output = sum(weight_i * expert_i(x))
```

## Tính toán tham số / Parameter Math

```
Embedding:       vocab_size × hidden = 32000 × 2048 = 65.5M
Per layer attn:  2048² + 2×(2048×512) + 2048² = 10.5M (Q, K, V, O with GQA)
Per expert:      3 × 2048 × 5632 = 34.6M (gate + up + down)
Per layer MoE:   24 × 34.6M = 830M (total)
                 3 × 34.6M = 104M (active)
Per layer total: 10.5M + 830M = 840.5M
12 layers:     10,086M
LM head:        65.5M
────────────────────────────────────
TOTAL:          10,223M ≈ 10.22B ✓
ACTIVE:         65.5 + 12×(10.5 + 104) + 65.5 = 1,503M ≈ 1.50B ✓
```

## Workflow

### Training Workflow
1. Tokenize input text → token IDs
2. Embed tokens → hidden states [B, L, H]
3. For each layer:
   - Pre-norm → Attention → residual
   - Pre-norm → MoE (router + experts) → residual
4. Final norm → LM head → logits
5. Compute cross-entropy loss + aux loss
6. Backpropagation

### Inference Workflow
1. Tokenize prompt
2. Forward pass through all layers
3. Get logits for last position
4. Apply temperature, top-k, top-p
5. Sample next token
6. Append to sequence, repeat

## Tối ưu / Optimizations

- **KV Cache**: Cache K, V từ các token trước để tăng tốc generation
- **GQA**: Giảm memory và computation cho attention
- **Pre-norm**: Ổn định hơn post-norm trong training
- **Mixed Precision**: Hỗ trợ fp16/bf16 để tiết kiệm memory
- **Gradient Checkpointing**: Đánh đổi compute lấy memory (chưa implement trong v0.1)