markbaseengine/OPTIMIZATION_PLAN.md

TEXT Generation優化計畫

問題分析

• 單個 forward pass 有 11 個 waitUntilCompleted() 呼叫
• 每個呼叫阻塞 CPU 等待 GPU 完成
• E4B: 11.3秒/token, 12B: 5.8秒/token（應該 <1秒）

根本原因

目前流程：
Embedding → wait()
Scale → wait()
PerLayer → wait()
Layer 0 → wait()
Layer 1 → wait()
...
Layer 42 → wait()
LM Head → wait()
Readback → wait()

總共: 11+ 次同步等待

優化策略

1. Batch Commands（最高優先）

// 優化後：
let cmdBuf = engine.commandQueue.makeCommandBuffer()!

// 所有操作加入同一個 command buffer
try dequantizeRow(...)      // 不等待
try scaleBuffer(...)         // 不等待
for layer in layers {        // 不等待
    try layer.forward(...)
}
try lmHead.forward(...)      // 不等待

// 最後才等待
cmdBuf.commit()
cmdBuf.waitUntilCompleted()  // 只等待一次！

預期改善：

• 從 11次等待 → 1次等待
• 減少 GPU-CPU同步開銷
• 預估速度提升 10倍以上

2. 移除 Per-Layer同步（次要）

// Line 1120-1135: Per-layer norm loop有同步等待
for layerIdx in 0..<numHiddenLayers {
    try rmsNorm(...)  // ← 應該批次處理
    cmdBufNorm.commit()
    cmdBufNorm.waitUntilCompleted()  // ← 移除這個
}

優化：

• 用單個 kernel批次處理所有 layers的 norm
• 或用 separate command buffer不等待

3. Forward Pass重構

// 新結構：
public func forwardOptimized(tokenId: Int, position: Int) throws -> [Float] {
    let cmdBuf = engine.commandQueue.makeCommandBuffer()!
    
    // 所有 GPU操作批次執行
    try embeddingPhase(cmdBuf)      // Embedding + Scale + PerLayer
    try layersPhase(cmdBuf)          // All layers in one batch
    try lmHeadPhase(cmdBuf)          // LM Head + Final Norm
    
    cmdBuf.commit()
    cmdBuf.waitUntilCompleted()      // 只等待一次
    
    return readbackLogits()
}

4. Kernel Fusion（進階）

// 合併多個操作成單個 kernel
kernel void embedding_scale_norm(
    device float* embedding,
    device float* scale,
    device float* norm_weight,
    device float* output,
    uint id [[thread_position_in_grid]]
) {
    // 一次執行：dequantize + scale + norm
    float val = dequantize(embedding[id]);
    val *= scale[0];
    val = rms_norm(val, norm_weight[id]);
    output[id] = val;
}

實作步驟

Step 1: 修改 Model.swift forward function

1. 移除所有中間的 waitUntilCompleted()
2. 只在最後保留一個 waitUntilCompleted()
3. 所有操作加入同一個 command buffer

Step 2: 測試驗證

1. 確保數值正確（無NaN）
2. 測量時間改善

Step 3: Kernel Fusion（optional）

1. 建立組合 kernel
2. 進一步減少 kernel launch overhead

預期成果

Metric	目前	優化後（預估）	改善倍數
E4B token生成	11.3秒	~1秒	10倍
12B token生成	5.8秒	~0.5秒	10倍
waitUntilCompleted呼叫	11次	1次	11倍
GPU-CPU同步	高	低	顯著

限制

• 需要修改大量 forward pass邏輯
• 需要確保數值穩定性
• 可能影響 debugging能力

替代方案（如果 batch困難）

1. 使用 asynchronous completion handlers
2. Pipeline多個 forward passes
3. Pre-compute KV cache for common tokens