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修复方案:

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

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

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

// 在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输出:

// 在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分钟（深度调试）

验证方法

测试代码

// 在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）