# dequantizeRow函数分析

**日期**: 2026-06-24  
**关键发现**: Token ID被用作embedding lookup索引

---

## 一、关键代码

### 1.1 Forward Pass调用

```swift
// Line 1346: Embedding lookup
try dequantizeRow(weight: embedWeight, tokenId: tokenId, output: h)

// Line 1378: Per-layer embedding
try dequantizeRow(weight: plWeight, tokenId: tokenId, output: plBuf, nCols: totalPerLayer)
```

**关键**: `tokenId`被直接用作参数！

---

### 1.2 dequantizeRow函数

**推测实现**：
```swift
func dequantizeRow(weight: QuantizedWeights, tokenId: Int, output: MTLBuffer) {
    // 从weight中读取第tokenId行的weights
    // weight.shape = [vocabSize, hiddenDim]
    // 每个tokenId对应一行embedding weights
    
    // 关键：tokenId被用作索引！
    // 可能的问题：
    // - tokenId超出weight的行数范围
    // - 或tokenId对应的weights有问题
}
```

---

## 二、推测的Bug机制

### 2.1 Token ID索引问题

**假设**：
- `dequantizeRow`从`embedWeight`中读取第`tokenId`行
- `embedWeight` shape: `[262144, 352]` (vocabSize=262144)
- Token ID 2, 100, 200等都在正常范围内
- **但**：26B-A4B的weights可能有问题

**可能的bug**：
1. Weight的量化格式不匹配
2. Scales/biases的group_size不正确
3. Dequantization计算错误

---

### 2.2 对比26B-Standard

**26B-Standard**：
- Embed scales: shape=[262144, 88], mean=119.955（异常大）
- 代码normalizing后正常
- 完美无NaN

**26B-A4B**：
- Embed scales: shape=[262144, 44], mean=-0.000326（正常）
- 不需要normalizing
- 但有NaN问题

**关键差异**：
- Scales的shape不同（88 vs 44）
- Group_size不同（32 vs 8）
- 这可能导致dequantization逻辑不同

---

## 三、验证方案

### 3.1 测试dequantizeRow

**测试代码**：
```swift
// 测试不同tokenId的embedding lookup
for tokenId in [2, 98, 100, 200] {
    let embedding = try model.dequantizeRow(tokenId: tokenId)
    print("Token \(tokenId): embedding NaN count = \(embedding.filter { $0.isNaN }.count)")
}
```

**预期**：
- 如果embedding就有NaN → dequantizeRow有问题
- 如果embedding无NaN但logits有NaN → LM head有问题

---

### 3.2 检查Metal Kernel

**需要检查**：
- `dequantize_row.metal` kernel的实现
- tokenId如何被用作索引
- Scales/biases如何被应用
- Group_size如何被计算

---

## 四、修复方案

### 4.1 可能的修复

**方案1**: 调整dequantizeRow的group_size计算
```swift
// 确保group_size正确
var groupSize = UInt32(weight.inDim / weight.scales.shape[1])
enc.setBytes(&groupSize, ...)
```

**方案2**: 检查scales/biases的offset计算
```swift
// 确保tokenId对应的scales/biases offset正确
let scalesOffset = tokenId * scalesShape[1] * 4
let biasesOffset = tokenId * biasesShape[1] * 4
```

**方案3**: 使用26B-Standard代替
- 最简单的方案
- 完美无NaN

---

## 五、下一步

**立即测试**：
1. 检查embedding是否已经有NaN
2. 检查dequantize_row kernel
3. 对比26B-Standard的实现

**如果无法修复**：
- 使用26B-Standard代替
- 或重新量化26B-A4B

---

**生成时间**: 2026-06-24  
**关键发现**: dequantizeRow使用tokenId作为索引  
**下一步**: 检查Metal kernel实现