怎样在安卓端利用 Vulkan 编写 4-bit 专用算子：直接在 GPU 上进行解量化与计算融合

前言

在安卓端侧部署大语言模型（LLM）时，内存带宽往往是最大的瓶颈。4-bit 量化（INT4）可以将权重体积减少 75%，但如果我们在推理前先用 CPU 或 GPU 算子将其还原为 FP16，会产生额外的显存读写开销。

本文将教你如何编写一个高效的 Vulkan Compute Shader，直接在计算循环中完成 INT4 解量化与 GEMV（矩阵向量乘法）的融合。

1. 数据打包策略

在 Vulkan 中，我们通常使用 uint32_t 类型的 Buffer 来存储打包后的权重。一个 32 位整数可以存放 8 个 4-bit 权重。

布局示例：
– 输入权重：w0, w1, w2, w3, w4, w5, w6, w7
– 存储方式：packed_val = (w7 << 28) | … | (w1 << 4) | w0

2. 核心 GLSL 算子实现

这里的关键是在循环内部使用位运算（Bitwise Operators）实时提取权重，并应用量化参数（Scale 和 ZeroPoint）。

#version 450
layout(local_size_x = 64) in;

// 绑定缓冲区
layout(set = 0, binding = 0) buffer InputBuffer { float vec_in[]; };
layout(set = 0, binding = 1) buffer WeightBuffer { uint packed_weights[]; };
layout(set = 0, binding = 2) buffer QuantParam { float scales[]; float zeros[]; };
layout(set = 0, binding = 3) buffer OutputBuffer { float vec_out[]; };

layout(push_constant) uniform Params {
    uint K; // 维度
    uint N; // 输出通道
} params;

void main() {
    uint row = gl_GlobalInvocationID.x;
    if (row >= params.N) return;

    float sum = 0.0;
    uint k_blocks = params.K / 8; // 一个uint包含8个权重

    float scale = scales[row];
    float zero = zeros[row];

    for (uint i = 0; i < k_blocks; ++i) {
        // 1. 读取一次 uint32 包含 8 个权重
        uint packed_val = packed_weights[row * k_blocks + i];

        // 2. 循环展开提取 4-bit 并计算
        for (uint j = 0; j < 8; ++j) {
            // 提取低 4 位
            uint val = (packed_val >> (j * 4)) & 0x0Fu;

            // 解量化: (q - zero) * scale
            float w = (float(val) - zero) * scale;

            // 累加计算
            sum += vec_in[i * 8 + j] * w;
        }
    }

    vec_out[row] = sum;
}

3. 操作要点与优化

1. 位运算技巧：使用 (packed >> shift) & 0x0F 是处理 INT4 的标准姿势。注意在 Vulkan 中必须使用 uint 类型以避免符号位干扰。
2. 计算融合：千万不要先写一个单独的解压算子存入临时 Buffer。融合算子的意义在于：原本需要读取 4 个字节才能获得 1 个 FP32 权重，现在读取 4 个字节就能获得 8 个权重，带宽利用率理论提升 8 倍。
3. 对齐要求：确保你的 K 维度是 8 的倍数。如果不满足，需要在数据预处理阶段进行 Padding。

4. 在安卓端调用

在 NCNN 或 MNN 等国产框架中，你可以通过自定义算子（Custom Op）接口，将上述 GLSL 编译为 Spir-V 字节码，并配置对应的 VkDescriptorSet。

实操建议：
– 使用 shared memory 缓存 vec_in 以减少重复读取。
– 配合 subgroup 指令（如果硬件支持）进行跨线程求和，进一步压榨 Adreno 或 Mali GPU 的性能。

总结

通过在 Vulkan GLSL 中直接嵌入解量化逻辑，我们能够以极小的计算代价换取巨大的带宽节省。这不仅是让 7B 参数模型在手机上跑起来的关键，更是提升生成速度（Tokens/s）的核心手段。