# ✓✓✓ Batch Embedding Kernel修复成功 ## 🎉 重大成功! ### 问题修复 **原始状态**: Sequential fallback(每个token单独处理) **问题**: dequantize_row_batch kernel未调用,导致性能瓶颈 ### 解决方案 1. **正确调用batch kernel**: 使用2D grid(batchSize × hiddenSize) 2. **修复参数传递**: tokenIds数组正确传递到Metal 3. **优化threadgroup**: 32×8 threads per threadgroup ### 实现代码 ```swift // Prepare tokenIds array for Metal let tokenIdsBuffer = engine.device.makeBuffer( bytes: tokenIds.map { UInt32($0) }, length: batchSize * 4, options: .storageModeShared )! // Use batch embedding kernel let pso = try engine.pipeline(named: embedScale != 1.0 ? "dequantize_row_batch_scaled" : "dequantize_row_batch") let enc = embedCmdBuf.makeComputeCommandEncoder()! enc.setComputePipelineState(pso) enc.setBuffer(embedWeight.weight, offset: 0, index: 0) enc.setBuffer(embedWeight.scales, offset: 0, index: 1) enc.setBuffer(embedWeight.biases, offset: 0, index: 2) enc.setBuffer(tokenIdsBuffer, offset: 0, index: 3) enc.setBuffer(context.batchInputBuffer, offset: 0, index: 4) var nCols = UInt32(hiddenSize) var batchSz = UInt32(batchSize) var groupSz = UInt32(embedWeight.groupSize) enc.setBytes(&nCols, length: 4, index: 5) enc.setBytes(&batchSz, length: 4, index: 6) enc.setBytes(&groupSz, length: 4, index: 7) if embedScale != 1.0 { var scale = embedScale enc.setBytes(&scale, length: 4, index: 8) } // 2D grid: batchSize × hiddenSize let threadsPerThreadgroup = MTLSize(width: 32, height: 8, depth: 1) let gridSize = MTLSize(width: batchSize, height: hiddenSize, depth: 1) enc.dispatchThreads(gridSize, threadsPerThreadgroup: threadsPerThreadgroup) ``` ## 性能成果 ### Batch Generation性能 ``` 原始(sequential fallback): 76ms/token 修复后(batch kernel): 41.13ms/token 提升: 85% faster ✓✓✓ ``` ### 测试结果 ``` Batch Generation Performance Test: PASSED (10.538 seconds) Batch(8): 411.314ms (41.13ms/token) ✓ Batch generation is faster! ``` ### 与单token对比 ``` 单token: ~25ms/token (optimized) Batch(8): 41.13ms/token Batch性能比率: 1.65x slower than single vs 原始sequential: 3x slower 改善: 从3x → 1.65x (45% improvement) ✓✓✓ ``` ## 技术细节 ### Batch Embedding Kernel逻辑 ```metal kernel void dequantize_row_batch_scaled( device const uint *w [[buffer(0)]], // [vocabSize, nCols/8] device const float *s [[buffer(1)]], // [vocabSize, numGroups] device const float *b [[buffer(2)]], // [vocabSize, numGroups] device const uint *tokenIds [[buffer(3)]], // [batchSize] device float *out [[buffer(4)]], // [batchSize, nCols] constant uint &nCols [[buffer(5)]], constant uint &batchSize [[buffer(6)]], constant uint &groupSize [[buffer(7)]], constant float &embedScale [[buffer(8)]], uint3 gid [[thread_position_in_grid]] ) { uint batchIdx = gid.x; // Which token in batch uint colIdx = gid.y; // Which column in embedding if (batchIdx >= batchSize || colIdx >= nCols) return; uint tokenId = tokenIds[batchIdx]; // ... quantized decoding ... out[batchIdx * nCols + colIdx] = (float(qval) * scale + bias) * embedScale; } ``` ### 关键改进 1. **2D Grid**: batchSize × hiddenSize (并行处理所有tokens和columns) 2. **TokenIds传递**: 正确传递batch的token ID数组 3. **Fused scale**: embedScale直接在kernel内应用(避免额外kernel) 4. **正确threadgroup**: 32×8优化GPU利用率 ## 性能分析 ### Sequential Fallback瓶颈 ``` for i in 0..