# 26B-A4B Router Scale Analysis - Potential Issue Found

## Discovery Date
2026-06-20 22:13

## ✅ Router Structure Test: PASSED

### Router Components Verified
```
Layer 0 Router:
  ✓ routerProj: present (8-bit, inDim=2816, outDim=128)
  ✓ routerScale: 31.25 ⚠️ POTENTIAL ISSUE
  ✓ perExpertScale: present [128 values]
  ✓ topK: 8

Expert Components:
  ✓ expertGate: present (128 experts, 704 output, 2816 input, 4-bit)
  ✓ expertUp: present (same structure)
  ✓ expertDown: present (same structure)
```

### ⚠️ Key Finding: routerScale = 31.25

**Potential Issue**: Router scale value is 31.25, which might need normalization

**Comparison with 26B-Standard**:
```
26B-Standard scales issue:
  - Original: scales ~120
  - Problem: Too large, caused numerical issues
  - Fix: Normalize by hidden_size (120/2816 = 0.0426)
  - Result: Fixed NaN issues

26B-A4B router scale:
  - Current: routerScale = 31.25
  - Question: Is this already normalized? Or needs normalization?
  - Potential fix: Divide by hidden_size? (31.25/2816 = 0.011)
```

### Router Scale Purpose

In MoE models, router scale is used to scale router logits before softmax:
```swift
// Layer.swift:837 (moeForward)
var scaled = routerData.map { $0 * routerScale }
```

**Effect**:
- If routerScale is too large → softmax overflow
- If routerScale is too small → softmax underflow
- Both cause numerical instability or NaN

### Analysis

**Router computation flow**:
1. Router proj: input [hidden_size] → output [num_experts]
2. Raw logits: ~some range
3. Scale logits: logits * routerScale
4. Softmax: exp(scaled_logits) / sum

**If routerScale=31.25 is too large**:
- scaled_logits could overflow exp() function
- NaN in softmax computation
- Generation hangs or crashes

### Hypothesis

**routerScale might need normalization**:
```swift
// Possible fix in Model.swift
let routerScale = rsFloats.first ?? 1.0
let normalizedRouterScale = routerScale / Float(hiddenSize)

// Use normalizedRouterScale in Layer
```

**Or**: routerScale is already correct and issue is elsewhere

### Testing Required

1. **Check router computation values**:
   - What are raw router logits?
   - What are scaled logits?
   - Do they overflow?

2. **Try normalization**:
   - Divide routerScale by hidden_size
   - Test if generation works

3. **Check softmax implementation**:
   - Is it handling overflow correctly?
   - Are there NaN checks?

### Related Code

**Router scale loading** (Model.swift:508-519):
```swift
if let rsDesc = allTensors.first(where: { $0.name == "\(prefix).router.scale" }) {
    let rsData = try rsReader.read(tensor: rsDesc)
    let rsFloats = SafeTensorsReader.bf16ToFloat32(rsData)
    routerScale = rsFloats.first ?? 1.0  // Gets first value
}
```

**Router scale usage** (Layer.swift:837):
```swift
var scaled = routerData.map { $0 * routerScale }
```

### Comparison with Other Models

| Model | MoE | routerScale | Notes |
|-------|-----|-------------|-------|
| 26B-Standard | No | N/A | Uses scales normalization (120/2816) |
| 31B-IT | No | N/A | Dense, no router |
| **26B-A4B** | Yes | **31.25** | Needs investigation |

### Next Steps

**Immediate**:
1. ✅ Run generation test (currently in progress)
2. If hangs → try router scale normalization
3. Test with routerScale / hiddenSize

**If normalization fixes**:
- Add normalization to Model.swift
- Similar to scales normalization fix
- Document in validation report

**If normalization doesn't fix**:
- Check other potential issues
- Expert selection logic
- Metal kernels
- Forward pass sequence

### Files

**Test code**:
- `/Users/accusys/MarkBase12B/Tests/G12BTests/MoEDebugTests.swift`

**Test output**:
- `/Users/accusys/MarkBase12B/MOE_ROUTER_STRUCTURE_TEST.log`

**Model**:
- `/Users/accusys/MarkBase12B/models/gemma-4-26b-a4b-it-4bit/`

**Router scale tensor**:
- `language_model.model.layers.0.router.scale`
- Shape: [2816] bf16
- Value: 31.25 (first element)

---

## Summary

**✅ Router structure is correct and complete**

**⚠️ Potential issue**: routerScale=31.25 might need normalization

**🔧 Possible fix**: Divide by hiddenSize (31.25/2816 = 0.011)

**📊 Test result**: Router structure test passed, generation test in progress