markbaseengine/26B_A4B_ROUTER_SCALE_ANALYSIS.md

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:

// 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:

// 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):

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):

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