Skip to content

Commit d56423d

Browse files
wangzhaodewangzhaode
committed
[CPU:Perf] Optimize SpacemiT IME2 dynamic prefill
Discussed-in: Merge-Request 29377173 , URL: https://code.alibaba-inc.com/AliNN/AliNNPrivate/codereview/29377173 GitOrigin-RevId: bd9b59eee7a85e7130ded4795023e3ee8c6abaa5
1 parent e8a3e6d commit d56423d

4 files changed

Lines changed: 202 additions & 1 deletion

File tree

‎skills/riscv-cpu-optimize/SKILL.md‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@ description: MNN RISC-V CPU 算子与 LLM 性能优化工作流，覆盖标准 R
6868
| 现象 | 优先检查 |
6969
|---|---|
7070
| prefill 慢 | pack/动态量化遍数、M tile、权重复用、矩阵单元利用率 |
71+
| 首次 prefill 慢，后续换 shape 又慢 | Geometry 是否重建子 Op/Execution、不可变权重是否重复 reorder |
72+
| 短 prompt 汇总速度异常低 | 是否混入首次 prefill 冷启动；逐条记录 cold/hot 时间和 token 数 |
7173
| decode 慢 | packed weight 字节数、持续带宽、dispatch/barrier、epilogue |
7274
| kernel 快但模型不快 | 调用次数、Attention/KV、layout conversion、线程池 |
7375
| 增加线程反而慢 | 共享矩阵单元、内存带宽、核拓扑、同步成本 |
@@ -82,6 +84,30 @@ decode tokens/s 上限 ≈ 持续有效带宽 / 每 token 必读权重与元数
8284

8385
不要把接口峰值带宽、稀疏 TOPS 或单条指令峰值当成模型可达到的吞吐。
8486

87+
#### 动态 shape 与首次预热
88+
89+
动态输入变慢时，先区分毫秒级 activation arena 调整与秒级不可变权重准备。若 fused op 在 Geometry
90+
层分解，shape 变化可能重新生成子 `Op`；当 Execution cache 以 `Op*` 为键时，这会导致缓存失效、
91+
重建 Execution，并再次 reorder 全部权重。用构造次数、子 Op 地址和 weight-reorder 计时闭合证据链，
92+
不要先用大块预留内存或固定 padding 掩盖问题。
93+
94+
当算子拓扑、权重和参数不变，只有 Tensor shape/binding 变化时，优先实现
95+
`GeometryComputer::onRecompute`：
96+
97+
1. 校验子命令、临时 Tensor 和输入输出数量仍匹配；不匹配就返回 `false` 走完整重建。
98+
2. 更新临时 Tensor 的 shape、类型、layout，以及 Command 的输入输出绑定。
99+
3. 保留原 Command、`BufferStorage`、子 `Op` 指针和 Execution，使已重排权重继续复用。
100+
4. 分别验证首次 cold prefill、至少两个从未出现过的新 shape，以及主 tile 两侧的 tail。
101+
102+
不要把 decode tuning 当成 prefill warmup。当前 `Llm::tuning(OP_ENCODER_NUMBER)` 使用 M=1 decode
103+
Module；prefill 可能使用另一 Module/Pipeline 和独立 Execution cache。若产品允许启动预热，可在接收
104+
首个用户请求前执行一次 prefill-only forward 并 reset KV；长期方案是让不同 Module 共享不可变的
105+
packed-weight Resource。
106+
107+
不要用 `M % tile == 0` 门禁整个 vendor fast path。让主 kernel 处理完整 tile，tail kernel 使用相同的
108+
packed-A/B ABI 处理剩余行。Chunk prefill 只作为长输入的独立实验；固定到 1024/2048 等长度会增加
109+
短输入计算量，不能替代 Execution 复用。
110+
85111
### 2. 建立正确性阶梯
86112

87113
按以下层级逐步对拍：
@@ -147,6 +173,7 @@ padding 或二进制偶然布局维持成绩。
147173
- vendor 路径短生成；
148174
- 标准 RVV 短生成；
149175
- 跨 prefill 分支阈值的长 prompt；
176+
- 首次 cold prefill 与连续多个不同 shape 的逐请求计时；
150177
- 目标模型 prefill 与 decode；
151178
- 不满足门禁时的 fallback；
152179
- `git diff --check` 和厂商目录外的污染检查。

‎source/backend/cpu/riscv/rvv/spacemit_ime2/MNNSpacemitIme2ConvInt8Executor.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ bool SpacemitIme2ConvInt8Executor::onSetupLinearFastPath(const std::vector<Tenso
141141

142142
if (mResourceInt8->mWeightBits != 4 || !mResourceInt8->mDynamicQuant || !mUseBatchQuan || !mIm2ColBasedInt8 ||
143143
mMixedKernel || mOnlineReorderWeightSme || m4BitPtq || !mLinear1x1 || dynamicQuantOption == 2 ||
144-
inputBlockNum != 1 || realCount <= static_cast<size_t>(dstXUnit) || realCount % 4 != 0 ||
144+
inputBlockNum != 1 || realCount <= static_cast<size_t>(dstXUnit) ||
145145
realCount != static_cast<size_t>(inputPlane) || dstBytes != 4 || gcore->bytes != 4 || gcore->pack != 4 ||
146146
srcDepthQuad == 0 || srcDepthQuad % 2 != 0 || dstDepthQuad == 0 || dstDepthQuad % 8 != 0 || threadCount <= 0 ||
147147
mBlockNum <= 0 || srcUnit != 16 ||

‎source/backend/cpu/riscv/rvv/spacemit_ime2/MNNSpacemitIme2GemmI8I4Local.cpp‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1534,6 +1534,110 @@ static void MNNSpacemitIme2GemmI8I4HpM1(size_t blk_len, const uint8_t* quant_a_p
15341534
}
15351535
}
15361536

1537+
static void MNNSpacemitIme2GemmI8I4HpM1Residual(size_t blk_len, const uint8_t* quant_a_ptr, const uint8_t* quant_b_data,
1538+
const uint8_t* quant_b_zp, float* c_ptr, size_t count_m, size_t count_n,
1539+
size_t k_blks, size_t ldc) {
1540+
// Full prefill tiles use the M4 kernel below. The packer stores the final 1-3 rows in the same single-row HP
1541+
// layout used here, so the scheduler can consume them one at a time without repacking or duplicating packed B.
1542+
(void)count_m;
1543+
(void)ldc;
1544+
if (blk_len != 258 || quant_a_ptr == nullptr || quant_b_data == nullptr || quant_b_zp != nullptr ||
1545+
c_ptr == nullptr || count_n == 0 || count_n % 32 != 0 || k_blks == 0) {
1546+
return;
1547+
}
1548+
constexpr size_t NB_COLS = 32;
1549+
constexpr size_t B_SUB_STRIDE = sizeof(_Float16) * NB_COLS + 16 * NB_COLS;
1550+
constexpr size_t B_RESIDUAL_SUB_STRIDE = B_SUB_STRIDE + sizeof(_Float16) * NB_COLS;
1551+
constexpr size_t B_SUPER_STRIDE = 8 * B_RESIDUAL_SUB_STRIDE;
1552+
const size_t b_tile_stride = k_blks * B_SUPER_STRIDE;
1553+
1554+
for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
1555+
uint8_t* b_data = (uint8_t*)quant_b_data + (ni / NB_COLS) * b_tile_stride;
1556+
int8_t* a_data = (int8_t*)quant_a_ptr;
1557+
float* dst_c = c_ptr + ni;
1558+
1559+
asm volatile(
1560+
"vsetvli t0, x0, e16, m1 \n\t"
1561+
"vxor.vv v31, v31, v31 \n\t"
1562+
"mv t4, %[BK] \n\t"
1563+
"li t0, 0x4c00 \n\t"
1564+
"fmv.h.x fa0, t0 \n\t"
1565+
1566+
".align 4 \n\t"
1567+
"_M1R_BLK_LOOP%=: \n\t"
1568+
"li t5, 8 \n\t"
1569+
"addi t6, %[A], 288 \n\t"
1570+
"flh ft1, (t6) \n\t"
1571+
"addi t6, %[A], 272 \n\t"
1572+
1573+
"vsetvli t0, x0, e16, m1 \n\t"
1574+
"vxor.vv v16, v18, v18 \n\t"
1575+
"vxor.vv v17, v18, v18 \n\t"
1576+
"vxor.vv v18, v18, v18 \n\t"
1577+
"vxor.vv v19, v18, v18 \n\t"
1578+
1579+
"_M1R_INNER_BLK_LOOP%=: \n\t"
1580+
"flh fa1, (t6) \n\t"
1581+
"addi t6, t6, 2 \n\t"
1582+
"flh ft0, (%[A]) \n\t"
1583+
"addi %[A], %[A], 2 \n\t"
1584+
1585+
"vsetvli t0, x0, e8, mf4 \n\t"
1586+
"vle8.v v3, (%[A]) \n\t"
1587+
"addi %[A], %[A], 32 \n\t"
1588+
1589+
"vsetvli t0, x0, e16, mf2 \n\t"
1590+
"vle16.v v8, (%[B]) \n\t"
1591+
"addi %[B], %[B], 64 \n\t"
1592+
"vl4r.v v4, (%[B]) \n\t"
1593+
"addi %[B], %[B], 512 \n\t"
1594+
"vle16.v v12, (%[B]) \n\t"
1595+
"addi %[B], %[B], 64 \n\t"
1596+
"vfmul.vf v8, v8, ft0 \n\t"
1597+
"vfmul.vf v9, v8, fa0 \n\t"
1598+
"vfmul.vf v10, v8, fa1 \n\t"
1599+
"vfwmacc.vf v31, ft1, v10 \n\t"
1600+
"vfmul.vf v10, v12, fa1 \n\t"
1601+
"vfwmacc.vf v31, ft1, v10 \n\t"
1602+
1603+
"vsetvli t0, x0, e8, m1 \n\t"
1604+
"vpack.vv v0, v8, v9, 3 \n\t"
1605+
"vsrl.vi v28, v3, 4 \n\t"
1606+
1607+
"vsetvli t0, x0, e16, m1 \n\t"
1608+
"vnpack4.vv v2, v3, v3, 3 \n\t"
1609+
"vnpack4.vv v3, v28, v28, 3 \n\t"
1610+
"vmadotsu.hp v16, v3, v4, v0, 4, i4 \n\t"
1611+
"vmadotsu.hp v17, v3, v5, v0, 5, i4 \n\t"
1612+
"vmadotsu.hp v18, v3, v6, v0, 6, i4 \n\t"
1613+
"vmadotsu.hp v19, v3, v7, v0, 7, i4 \n\t"
1614+
"vmadotu.hp v16, v2, v4, v0, 0, i4 \n\t"
1615+
"vmadotu.hp v17, v2, v5, v0, 1, i4 \n\t"
1616+
"vmadotu.hp v18, v2, v6, v0, 2, i4 \n\t"
1617+
"vmadotu.hp v19, v2, v7, v0, 3, i4 \n\t"
1618+
1619+
"addi t5, t5, -1 \n\t"
1620+
"bgtz t5, _M1R_INNER_BLK_LOOP%= \n\t"
1621+
"vpack.vv v8, v16, v17, 1 \n\t"
1622+
"vpack.vv v12, v18, v19, 1 \n\t"
1623+
"vpack.vv v20, v8, v12, 2 \n\t"
1624+
1625+
"vsetvli t0, x0, e16, mf2 \n\t"
1626+
"addi t4, t4, -1 \n\t"
1627+
"vfwmacc.vf v31, ft1, v20 \n\t"
1628+
"addi %[A], t6, 2 \n\t"
1629+
"bgtz t4, _M1R_BLK_LOOP%= \n\t"
1630+
1631+
"vsetvli t0, x0, e32, m1 \n\t"
1632+
"vse32.v v31, (%[DST]) \n\t"
1633+
: [A] "+r"(a_data), [B] "+r"(b_data)
1634+
: [DST] "r"(dst_c), [BK] "r"(k_blks)
1635+
: "cc", "memory", "t0", "t1", "t2", "t3", "t4", "t5", "t6", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
1636+
"v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22",
1637+
"v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "ft0", "ft1");
1638+
}
1639+
}
1640+
15371641
static void MNNSpacemitIme2GemmI8I4HpM4(size_t blk_len, const uint8_t* quant_a_ptr, const uint8_t* quant_b_data,
15381642
const uint8_t* quant_b_zp, float* c_ptr, size_t count_m, size_t count_n,
15391643
size_t k_blks, size_t ldc) {
@@ -2299,6 +2403,11 @@ extern "C" __attribute__((aligned(64))) size_t MNNSpacemitIme2GemmI8I4Local(size
22992403
MNNSpacemitIme2GemmI8I4HpM4(blkLen, quantAPtr, quantBData, quantBZp, cPtr, countM, countN, kBlocks, ldc);
23002404
return 4;
23012405
}
2406+
if (countM > 0 && quantBZp == nullptr) {
2407+
MNNSpacemitIme2GemmI8I4HpM1Residual(blkLen, quantAPtr, quantBData, quantBZp, cPtr, countM, countN, kBlocks,
2408+
ldc);
2409+
return 1;
2410+
}
23022411
return 0;
23032412
}
23042413
if (blkLen == 260) {

‎source/geometry/GeometryFusedProj.cpp‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,15 @@
1919
namespace MNN {
2020
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
2121
class GeometryFusedProj : public GeometryComputer {
22+
static void _updateTensorShape(Tensor* tensor, const Tensor* source, int channel = -1) {
23+
TensorUtils::copyShape(source, tensor, true);
24+
tensor->buffer().type = source->getType();
25+
if (channel >= 0 && tensor->dimensions() >= 2) {
26+
tensor->setLength(1, channel);
27+
}
28+
TensorUtils::setLinearLayout(tensor);
29+
}
30+
2231
static std::shared_ptr<Command> _makeCmd(std::shared_ptr<BufferStorage> storage,
2332
const std::vector<Tensor*>& inputs,
2433
const std::vector<Tensor*>& outputs) {
@@ -74,6 +83,62 @@ class GeometryFusedProj : public GeometryComputer {
7483
}
7584

7685
public:
86+
// Keep the generated commands and their executions across dynamic shapes. Only tensor bindings and temporary
87+
// shapes change; recreating the commands would reorder every immutable projection weight again.
88+
virtual bool onRecompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
89+
Context& context, CommandBuffer& res) const override {
90+
if (_keepWhole(context, op, inputs.size(), outputs.size())) {
91+
return false;
92+
}
93+
auto param = op->main_as_FusedLinearParam();
94+
if (param == nullptr || param->convs() == nullptr || inputs.empty()) {
95+
return false;
96+
}
97+
const int numConvs = (int)param->convs()->size();
98+
const bool isGateUp = param->act_silu_mul();
99+
const bool hasLn = param->has_ln();
100+
const int numProjOut = isGateUp ? 1 : numConvs;
101+
const int expectedCommands = (hasLn ? 1 : 0) + (isGateUp ? 3 : numConvs);
102+
const int expectedExtras = (hasLn ? 1 : 0) + (isGateUp ? 2 : 0);
103+
if ((isGateUp && numConvs != 2) || (!isGateUp && (numConvs < 3 || numConvs > 4)) ||
104+
inputs.size() < (hasLn ? 2 : 1) || outputs.size() < numProjOut + (hasLn ? 1 : 0) ||
105+
res.command.size() != expectedCommands || res.extras.size() != expectedExtras) {
106+
return false;
107+
}
108+
109+
int commandIndex = 0;
110+
int extraIndex = 0;
111+
Tensor* projInput = hasLn ? inputs[1] : inputs[0];
112+
if (hasLn) {
113+
auto normalized = res.extras[extraIndex++].get();
114+
_updateTensorShape(normalized, projInput);
115+
auto& command = res.command[commandIndex++];
116+
command->inputs = {inputs[0], inputs[1]};
117+
command->outputs = {outputs[numProjOut], normalized};
118+
projInput = normalized;
119+
}
120+
if (!isGateUp) {
121+
for (int i = 0; i < numConvs; ++i) {
122+
auto& command = res.command[commandIndex++];
123+
command->inputs = {projInput};
124+
command->outputs = {outputs[i]};
125+
}
126+
return true;
127+
}
128+
129+
auto gate = res.extras[extraIndex++].get();
130+
auto up = res.extras[extraIndex].get();
131+
_updateTensorShape(gate, projInput, param->convs()->GetAs<Convolution2D>(0)->common()->outputCount());
132+
_updateTensorShape(up, projInput, param->convs()->GetAs<Convolution2D>(1)->common()->outputCount());
133+
res.command[commandIndex]->inputs = {projInput};
134+
res.command[commandIndex++]->outputs = {gate};
135+
res.command[commandIndex]->inputs = {projInput};
136+
res.command[commandIndex++]->outputs = {up};
137+
res.command[commandIndex]->inputs = {up, gate};
138+
res.command[commandIndex]->outputs = {outputs[0]};
139+
return true;
140+
}
141+
77142
virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
78143
Context& context, CommandBuffer& res) const override {
79144
if (_keepWhole(context, op, inputs.size(), outputs.size())) {

0 commit comments

Comments
 (0)