问题定义
输入两个 Tensor:\(input[Batch, Ci, Ih, Iw]\) (经典 NCHW 布局)、\(kernel[Co, Ci, Kh, Kw]\)(为适应 Coalesced Access 可重排为 \([Co, Kh, Kw, Ci]\)),其中前者为 FP32/FP16 类型,后者为 INT4/INT8 类型。输出卷积结果。
可以发现,对于 \(Iw = Ih = Ow = Oh = Kw = Kh = Dw = Dh = Sw = Sh = 1\)(输入、输出、卷积核尺寸、扩张率 Dilute、步长 Stride 均为 1x1),\(Pw = Ph = 0\)(不填充)的情况,卷积退化为标准卷积乘法 \([Batch, Ci] \times [Ci, Co]\),以下简记为 \(A[N, K] \times B[K, M]\)。(当然,即使不是如此,也可以通过 Im2Col 或 Implicit Conv 等方式将其转换成可直接使用加速库的矩阵乘法)。
对于 \(B\)(在 LLM 推理中通常为量化后权重矩阵),采用 Per-block 量化,即在 \(K\) 维度上,每 \(block\_size\) 个元素共享一套量化参数(\(scale\) 和 \(zero\_point\)),\(M\) 维度不共享。因此可以理解为,将 \(B\) 的每列按 \(block\_size\) 划分为多个竖条,总块数为 \(K / block\_size \times M\)。