MarkBase Admin
ac75faa0cc
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
131 lines
3.3 KiB
Markdown
131 lines
3.3 KiB
Markdown
# 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（最高優先）
|
||
```swift
|
||
// 優化後：
|
||
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同步（次要）
|
||
```swift
|
||
// 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重構
|
||
```swift
|
||
// 新結構：
|
||
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（進階）
|
||
```swift
|
||
// 合併多個操作成單個 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 |