# Router Scale Normalization Fix Applied

## Fix Date
2026-06-20 22:16

## Problem
**routerScale = 31.25** (raw value, too large)
- Causes softmax overflow in MoE router computation
- Similar to 26B-Standard scales issue

## Solution
**Normalize routerScale by hiddenSize**
```swift
// Model.swift:516-519 (modified)
let rawRouterScale = rsFloats.first ?? 1.0
routerScale = rawRouterScale / Float(hiddenSize)
```

**Effect**:
- Before: routerScale = 31.25
- After: routerScale = 31.25 / 2816 = 0.01105
- Result: Stable softmax, no overflow

## Why This Works

**Similar to 26B-Standard fix**:
```
26B-Standard scales:
  - Raw: ~120
  - Fix: Divide by hiddenSize (120/2816 = 0.0426)
  - Result: Fixed NaN, model works

26B-A4B routerScale:
  - Raw: 31.25
  - Fix: Divide by hiddenSize (31.25/2816 = 0.01105)
  - Expected: Fix generation hanging
```

**Router computation flow**:
1. Router logits (raw): [numExperts]
2. Scale logits: logits * routerScale
3. Softmax: exp(scaled_logits) / sum

**If routerScale too large**:
- scaled_logits = logits * 31.25
- exp(scaled_logits) can overflow
- NaN in softmax
- Generation hangs

**If routerScale normalized**:
- scaled_logits = logits * 0.01105
- exp(scaled_logits) stable
- Softmax works
- Generation succeeds

## Code Changes

**File**: `/Users/accusys/MarkBase12B/Sources/G12B/Model.swift`

**Lines**: 516-519

**Before**:
```swift
let rsData = try rsReader.read(tensor: rsDesc)
let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
routerScale = rsFloats.first ?? 1.0  // Raw value
```

**After**:
```swift
let rsData = try rsReader.read(tensor: rsDesc)
let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
let rawRouterScale = rsFloats.first ?? 1.0
// Normalize router scale by hidden_size (similar to scales normalization for 26B-Standard)
// This prevents softmax overflow in MoE router computation
routerScale = rawRouterScale / Float(hiddenSize)
```

## Testing

**Next step**: Test generation with normalized routerScale

**Expected**:
- ✅ Generation works (no hang)
- ✅ No NaN in router computation
- ✅ Stable softmax
- ✅ Valid token generation

**Test command**:
```bash
swift test --filter MoEDebugTests/test26BA4BSimpleGenerationDebug
```

## Pattern Recognition

**Normalization pattern discovered**:
1. **26B-Standard scales**: Divide by hiddenSize
2. **26B-A4B routerScale**: Divide by hiddenSize
3. **Pattern**: Raw scale values need normalization by hiddenSize

**General rule**:
```
If scale value is loaded from tensor and seems large (>10):
  → Normalize by dividing by hiddenSize
  → Prevents numerical overflow
  → Matches model training normalization
```

## Confidence

**Confidence level**: ⭐⭐⭐⭐⭐ (Very High)

**Reasons**:
- ✅ Same pattern as 26B-Standard fix (proven to work)
- ✅ Router scale purpose is to scale logits before softmax
- ✅ Large scale values cause overflow
- ✅ Normalization prevents overflow

## If Fix Works

**Implications**:
- ✅ 26B-A4B MoE model will work
- ✅ First MoE model successfully running
- ✅ MoE implementation validated
- ✅ Pattern for fixing MoE numerical issues

**Comparison**:
```
26B-Standard: 40 tok/s (Dense, already works)
26B-A4B MoE: Expected 20-30 tok/s (MoE, should work after fix)
31B-IT: 11.7 tok/s (Dense, already works)
```

## If Fix Doesn't Work

**Next debugging steps**:
1. Check expert scales normalization
2. Add NaN checks in router computation
3. Test router forward pass separately
4. Check Metal kernels

**But**: High confidence this fix will work

## Summary

**✅ Fix applied**: Router scale normalization

**📊 Expected result**: Generation works (no hang)

**🔧 Pattern**: Normalize scale values by hiddenSize

**⏱️ Next**: Test generation with fix

---

**Status**: Fix applied, ready for testing