markbaseengine/TEXT_NAN_FIX_PLAN.md

# TEXT NaN修复方案（基于Audio经验）

## 问题分析

### Buffer重用冲突链
**TEXT Layer forward流程**:
```
1. Attention阶段（使用temps.h多次）:
   - Line 84-87: input → temps.h (RMSNorm #1)
   - Line 89-91: temps.h → temps.q
   - Line 105-106: temps.h → temps.k
   - Line 108-109: temps.h → temps.v
   - Line 171-172: temps.attn → temps.h (覆盖 #1)
   - Line 180-182: input → temps.h (RMSNorm #2，覆盖 #1)
   - Line 185-187: temps.h → temps.ns

2. FFN阶段（使用temps.h）:
   - Line 53-54: temps.gate → temps.h (down proj)

3. PostFFN阶段（使用temps.h多次）:
   - Line 207-209: input → temps.h (RMSNorm #3，覆盖 #2)
   - Line 225-227: temps.gating → temps.h (覆盖 #3)
   - Line 235-238: temps.h → input (最终输出)
```

**关键问题**:
- `temps.h`被多次重用：5次写入（attention 3次，FFN 1次，postFFN 2次）
- `input`被多次覆盖：residual add + 最终输出
- 类似Audio的多轮操作竞争buffer

### Audio修复对比
**Audio问题**:
```
多轮操作竞争tempBuffer:
- applySubsampleConv → tempBuffer
- applyInputProjection → subsampleBuf (修复)
- applyDepthwiseConv1D → tempBuffer冲突
- applySiLU → tempBuffer冲突
- applyResidualAdd → tempBuffer冲突
```

**Audio修复方案**:
```swift
// 创建layerBuffer（67MB）给audio layers专用
let layerBuffer = engine.device.makeBuffer(length: 67 * 1024 * 1024, options: .storageModeShared)!

// 所有audio layer操作使用layerBuffer，避免竞争tempBuffer
applyRMSNorm(..., output: layerBuffer)
applyDepthwiseConv1D(..., output: layerBuffer)
applySiLU(..., output: layerBuffer)
applyResidualAdd(..., output: layerBuffer)
```

## TEXT修复方案

### 方案1: Buffer隔离（推荐）
**创建attention专用buffer**:
```swift
// 在ForwardTemps中添加attentionBuffer
public struct ForwardTemps {
    public let q: MTLBuffer
    public let k: MTLBuffer
    public let v: MTLBuffer
    public let h: MTLBuffer      // FFN专用
    public let attnH: MTLBuffer  // NEW: Attention专用
    public let gate: MTLBuffer
    public let up: MTLBuffer
    public let attn: MTLBuffer
    public let gating: MTLBuffer
    public let ns: MTLBuffer
    public let io: MTLBuffer
    
    public init(...) throws {
        q = try buf(nHeads * maxHeadDim)
        k = try buf(nKvHeads * maxHeadDim)
        v = try buf(nKvHeads * maxHeadDim)
        h = try buf(hiddenSize)            // FFN专用
        attnH = try buf(hiddenSize)        // NEW: Attention专用
        gate = try buf(maxIntermediate)
        up = try buf(maxIntermediate)
        attn = try buf(nHeads * maxHeadDim)
        gating = try buf(256)
        ns = try buf(max(hiddenSize, nHeads * maxHeadDim))
        io = try buf(hiddenSize)
    }
}
```

**修改LayerOptimized.swift**:
```swift
// Attention阶段使用attnH（不覆盖h）
try rmsNorm(..., output: temps.attnH)      // Line 87: attnH #1
try quantizedMatmul(..., input: temps.attnH, output: temps.q)  // Line 91
try quantizedMatmul(..., input: temps.attnH, output: temps.k)  // Line 106
try quantizedMatmul(..., input: temps.attnH, output: temps.v)  // Line 109
try quantizedMatmul(..., output: temps.attnH)  // Line 172: attnH #2

// Residual add使用h（暂存）
try eltwiseAdd(..., b: temps.attnH, output: temps.h)  // Line 177: h暂存

// Post-attention norm使用attnH
try rmsNorm(..., input: temps.h, output: temps.attnH)  // Line 182: attnH #3

// FFN使用h（独立）
try quantizedMatmul(..., output: temps.h)  // Line 54: FFN专用h
try eltwiseAdd(..., b: temps.h, output: input)  // Line 57: residual

// PostFFN使用h
try rmsNorm(..., output: temps.h)  // Line 209: h专用
try eltwiseAddScaled(..., output: input)  // Line 238: 最终输出
```

**预期效果**:
- 避免attention和FFN竞争h buffer
- 类似Audio修复，减少NaN风险
- 内存增加：2560 Floats = 10KB（微不足道）

### 方案2: Input保护（简化）
**创建inputCopy buffer**:
```swift
// 在Layer forward开始时保存input副本
let inputCopy = temps.io  // 使用现有io buffer
let blit = cmdBuf.makeBlitCommandEncoder()!
blit.copy(from: input, sourceOffset: 0, to: inputCopy, destinationOffset: 0, size: hiddenSize * 4)
blit.endEncoding()

// 后续操作使用inputCopy（保护原始input）
try eltwiseAdd(..., a: inputCopy, b: temps.h, output: input)  // Line 177
try eltwiseAdd(..., a: inputCopy, b: temps.h, output: input)  // Line 57
```

**预期效果**:
- 保护原始input不被过早覆盖
- 使用现有io buffer（无额外内存）
- 简化修改

### 方案3: 顺序修复（保守）
**检查每个kernel的NaN输出**:
```swift
// 在attention后检查h
let cmdBufDebug = engine.commandQueue.makeCommandBuffer()!
try rmsNorm(..., output: temps.h, cmdBuf: cmdBufDebug)
cmdBufDebug.commit()
cmdBufDebug.waitUntilCompleted()
checkNaN(temps.h, "After RMSNorm")

// 在quantizedMatmul后检查
// ...逐个检查定位NaN源头
```

**预期效果**:
- 精确定位哪个kernel产生NaN
- 保守修复，最小修改
- 调试时间长

## 推荐修复路径

### 优先级排序
1. **方案1**（Buffer隔离）- 推荐，彻底修复
2. **方案2**（Input保护）- 简化，快速修复
3. **方案3**（顺序检查）- 调试，精确定位

### 实施步骤（方案1）
1. 修改ForwardTemps.swift添加attnH
2. 修改LayerOptimized.swift使用attnH
3. 测试验证NaN消除
4. 性能测试确认无影响

### 时间预估
- 方案1: ~30分钟（修改+测试）
- 方案2: ~15分钟（快速修复）
- 方案3: ~60分钟（深度调试）

## 验证方法

### 测试代码
```swift
// 在Layer forward后检查
let layerOutputPtr = input.contents().assumingMemoryBound(to: Float.self)
let layerNaNCount = Array(UnsafeBufferPointer(start: layerOutputPtr, count: hiddenSize)).filter { $0.isNaN }.count
print("Layer output NaN: \(layerNaNCount)/\(hiddenSize)")
assert(layerNaNCount == 0, "Layer produced NaN!")
```

### 测试模型
- E2B: 已加载成功，可用于测试
- E4B: Layer 34缺失，跳过
- 12B: Layer 1缺失，跳过

## 总结

**TEXT NaN问题性质**:
- Buffer重用冲突（类似Audio）
- Temps.h多次覆盖
- Input多次residual

**修复方向**:
- Buffer隔离（Audio经验）
- 减少重用
- 专用buffer

**预期效果**:
- NaN消除（类似Audio修复）
- 性能无损
- 内存微增（10KB）

**下一步**:
1. 选择修复方案（推荐方案1）
2. 实施修改
3. 测试验证
4. 完成TEXT NaN修复

---

**创建时间**: Day 3 Session（~4小时）
**基于**: Audio NaN修复经验（67%就绪）
**目标**: TEXT 95%就绪（零NaN）