MNN Vulkan 新特性集成实战：Cooperative Matrix / Subgroup 扩展从示例代码到落地

MNN Vulkan 新特性集成实战Cooperative Matrix / Subgroup 扩展从示例代码到落地【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN本篇技术指南聚焦 MNN 中「方向 C集成 Vulkan 新特性」这一独立优化轨道完整讲解如何把用户提供的 Vulkan 新特性示例代码如 VK_KHR_cooperative_matrix、GL_KHR_shader_subgroup_arithmetic 等扩展适配进 MNN Vulkan 后端覆盖输入收集、示例理解、兼容性评估、shader 适配、特性检测、fallback 设计与多设备验证全流程。读完你将掌握一条可复用的新特性集成流水线并能基于VulkanDevice的设备能力查询接口coopMat / subgroup在 MNN 中安全地启用 GPU 新特性而不错伤不支持的设备。一、新特性集成轨道定位与前置条件在 MNN 的 Vulkan 优化技能体系中见 skills/vulkan-optimize/SKILL.md共有三个并行优化方向方向 A「模型级优化」、方向 B「指定算子优化」以及本指南所属的方向 C「新特性集成」。三者是并列的独立轨道方向 C 的触发条件是用户提供了 Vulkan 新特性示例代码或参考实现并希望将其适配集成进 MNN例如用 cooperative matrix 加速 GEMM、用 subgroup shuffle 免 barrier 等。方向 C 的参考文档即本指南skills/vulkan-optimize/new-feature.md整体复杂度评估为中-高其核心约束有两个必须有示例代码。集成新特性的前提是用户提供.compshader host 代码或一个完整可运行的 demo。没有示例时应当主动要求补上——不同 GPU vendor 对同一特性的行为可能完全不同没有参考实现就无从对齐语义。必须设计 fallback。新特性尤其 coop matrix目前主要只在 Adreno 上可用集成时必须保证不支持的设备Mali、Apple、旧驱动回退到原路径且功能正确。开始前必须确认以下输入清单缺任何一项都需主动要求补齐□ 示例代码.comp host或完整可运行 demo □ 特性说明解决什么问题如 coop matrix 加速 GEMM、subgroup shuffle 免 barrier □ 目标算子用在 MNN 哪个算子 □ 目标平台哪些 GPU 支持Adreno / Mali / Applecoop 目前主要 Adreno □ 预期收益性能 / 精度 / 功能二、第一步理解示例代码拿到示例后需要同时从 Host 端与 Device 端两个层面拆解才能判断它能否映射到 MNN 的既有抽象上。2.1 Host 端要回答的问题用了哪些 Vulkan API / 设备能力查询pipeline / descriptor / push constant 是怎么构建的local_size/ dispatch grid 如何配置spec constant 传了什么值2.2 Device 端.comp要回答的问题用了哪些新内置coopMatMulAdd/subgroupAdd等、新 layoutconstant_id、扩展 require数据类型是FLOAT宏还是裸类型内存作用域 / 同步memory_scope如何处理2.3 关键逻辑核心计算流程是什么新特性在哪个环节起作用示例自身是否带有 fallback重要提示coop / subgroup 相关特性要求SPIR-V ≥ 1.3因此编译示例时必须带--target-env vulkan1.1。如示例可独立运行应先在目标设备上验证示例本身正确编译 spirv-val校验再谈集成——示例不成立后面的一切都没有意义。三、第二步评估兼容性——一张表对齐 MNN 与示例集成前必须逐项核对示例做法与 MNN 做法的差异这是决定改动量的关键。下表是 MNN Vulkan 后端的标准适配维度适配维度MNN 做法示例做法改动数据排布NC4HW4??数据类型FLOAT/FLOAT4fp16/fp32 宏??Buffer/Image编译期MNN_VULKAN_IMAGE??内存管理vkBn-onAcquireBuffer/getMemoryPool??pipeline 构建vkBn-getPipeline(name, types, localSize, spec)??参数传递descriptor set writeBuffer push constant / uniform??local_sizegetSubgroupSize()动态??3.1 Buffer / Image 是编译期二选一MNN 的 Vulkan 后端分两棵完全独立的代码树由编译期宏MNN_VULKAN_IMAGE决定MNN_VULKAN_IMAGEOFF→source/backend/vulkan/buffer/*LLM 全链路走 bufferMNN_VULKAN_IMAGEON→source/backend/vulkan/image/*。改 shader 前必须先grep MNN_VULKAN_IMAGE project/android/build_64/CMakeCache.txt确认当前构建走的是哪个后端并保证改的路径与构建选的后端一致否则改动不会生效。3.2 设备特性检测VulkanDevice是唯一入口MNN 将设备能力查询集中在 VulkanDevice.hpp 中。与方向 C 直接相关的查询接口如下auto coop vkBn-getDevice().getCoopMatInfo(); // supportCoopMat, selectedFP16CoopMatShape const auto sg vkBn-getDevice().getSubgroupInfo(); // size, stages, ops uint32_t sgSize vkBn-getDevice().getSubgroupSize();从 VulkanDevice.hpp 的源码结构可以看到三组关键数据结构SubgroupInfosizesubgroup 大小、stages支持的 shader stage、opsVkSubgroupFeatureFlags如 arithmetic / ballot / shuffle 能力、quadAllStages。getSubgroupSize()即返回mSubgroupInfo.size。CoopMatInfosupportCoopMat是否支持 cooperative matrix、enabledCoopMatFeatures启用的 feature 位、fp32CoopMatShape/fp16CoopMatShape/s8CoopMatShape各类型支持的全部 {M, N, K} 形状、以及selectedFP32CoopMatShape/selectedFP16CoopMatShape/selectedS8CoopMatShape选中的最优形状{M, N, K}三元组。其中supportS8S8S32专用于 W8A8 conv1x1 prefill 路径的 S8S8→S32 cooperative matrix。这些数据来自 VulkanDevice.cpp 的checkCoopMat运行时探测流程通过vkGetInstanceProcAddr获取vkGetPhysicalDeviceFeatures2带 KHR 兜底检查VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME扩展是否存在不存在直接返回通过VkPhysicalDeviceCooperativeMatrixFeaturesKHR查询cooperativeMatrixfeature 是否为VK_TRUE调用vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR枚举全部VkCooperativeMatrixPropertiesKHR过滤scope VK_SCOPE_SUBGROUP_KHR且非 saturating accumulation 的条目按 A/B/C/Result 类型分别归类为 FP16FLOAT16×FLOAT16→FLOAT16、FP32FLOAT32×FLOAT32→FLOAT32、S8S8S32SINT8×SINT8→SINT32三类形状对每一类按M*N*K总面积取最大者作为selected*形状最后置supportCoopMat truesupportS8S8S32依据 s8 形状列表是否为空。Fallback 必须设计支持特性 → 走新路径不支持 → 走原路径保证功能正确。以 coop matrix 为例MNN 的通行 gate 条件是gpuType()ADRENO supportCoopMat见下文 §5 的 conv1x1 选路源码。四、第三步实施集成——五步递进每步编译验证集成过程严格按序逐步进行每一步都要编译验证避免问题堆积到最后无法定位。第 1 步特性检测复用VulkanDevice已有的查询能力必要时新增查询在 host 代码中写if (supported) 建新 pipeline的分支逻辑。特性检测的正确性直接影响后续一切选路务必先单测验证检测结果与真机实际能力一致。第 2 步shader 适配为 MNN.comp把示例 shader 改写为 MNN 的.comp规范用FLOAT/FLOAT4宏fp16/fp32 自动切换而不是裸类型适配 NC4HW4 数据排布用 spec constantlayout(constant_idN)传 COOP_M/N/K、activation 等编译期可调参数在头部写清扩展 require 与特性注释coop / subgroup / memory_scope 必须进--target-env vulkan1.1——这是最容易踩的陷阱见 §7 常见问题表第一行缺了它 glslangValidator 或 spirv-opt 会直接拒绝编译。第 3 步外科式重生成 shader 数组MNN 的 GLSL 源不会被构建系统直接编译运行时读的是AllShader.h/cpp中由makeshader.py.comp→glslangValidator -V→spirv-opt -O→xxd嵌入生成的 SPIR-V 字节数组脚本位于 source/backend/vulkan/buffer/compiler/makeshader.py 等三处。⚠️严禁直接跑全量 makeshader本机 glslang/spirv-opt 与仓库版本常不同跑全量会把所有 shader 数组重编/重排污染AllShader.cpp出现几万行无关 diff。正确做法是只重生成改动的那几个数组对每个改动的 shaderfp32 fp16 变体用与 makeshader 相同的管线单独编译得到const unsigned char glsl_name_comp[]{...}; unsigned int glsl_name_comp_lenN;再用脚本精确替换AllShader.cpp中对应段落。新增 shader 还需要三处注册AllShader.h追加 extern 声明AllShader.cpp替换/追加字节数组VulkanShaderMap.cpp把 name → 数组 的映射注册进 map。以仓库中已有的 coop shader 为例VulkanShaderMap.cpp 的注册形式为mMaps.insert(std::make_pair(glsl_attention_prefill_coop_qkv_comp, std::make_pair(glsl_attention_prefill_coop_qkv_comp,glsl_attention_prefill_coop_qkv_comp_len))); mMaps.insert(std::make_pair(glsl_attention_prefill_coop_qkv_FP16_comp, std::make_pair(glsl_attention_prefill_coop_qkv_FP16_comp,glsl_attention_prefill_coop_qkv_FP16_comp_len)));可以看到仓库里已经沉淀了attention_prefill_coop_qkv、attention_prefill_coop_qk、matmul_coop、matmul_coop_rm、dynamic_w8a8_coop_gemm、int4_weight_to_coop、int8_weight_to_coop、attention_prefill_coop_scale_oacc等 coop 系列 shader说明这条集成流水线在仓库中已有大量实际落地样本可参照。此外新.comp文件必须在对应的macro.json中登记useFP16决定是否生成_FP16变体可选地声明macro列表按宏名批量生成变体。格式参见 buffer/execution/glsl/macro.json{ blit.comp: { useFP16: true, macro: [ C4 ] }, reduce_int.comp: { useFP16: false, macro: [ VMAX, VMIN, MEAN, PROD, SUM ] } }改完确认无污染grep -c glsl_name_comp_len AllShader.cpp并用git diff --stat确认 diff 中只有目标数组无关_len行不应出现。第 4 步host 选路在onCreate/onEncode中写if (feature supported) 用新 pipeline else 原路径两条路径都走 createSet 绑定用独立 pipeline/分支隔离不要修改原 shader防止 fallback 被破坏。仓库中 conv1x1 的选路是一个典型样本。在 VulkanConvolution.cpp 的onCreate中auto coopMatInfo extra-getDevice().getCoopMatInfo(); const char* disableCoopMatEnv getenv(MNN_VULKAN_DISABLE_COOPMAT); const bool disableCoopMat disableCoopMatEnv ! nullptr std::string(disableCoopMatEnv) 1; const bool useCoopMat !disableCoopMat coopMatInfo.supportCoopMat supportSubgroupArithmetic extra-gpuType() VulkanRuntime::ADRENO; // CoopMat path only supports int4/int8 weight. For 2/3-bit, go to ...随后在useInt8Conv is1x1分支内perChannelAsym S8S8S32→VulkanConv1x1CoopA8W8A8 prefill 路径否则 →VulkanConv1x1CoopCoopMat 只支持 int4/int8再 fallback 到VulkanConv1x1Generalnative int8/int4/int2/int3。同时仓库还提供了MNN_VULKAN_DISABLE_COOPMAT1环境变量作为运行时开关便于不重编二进制就临时关闭 coop 路径做 A/B——这是新特性集成时可以复用的工程实践。attention prefill 的选路则走 VulkanAttention.cpprearrange_q → init_state → (per k-block: qk → softmax → qkv) → finalizecoop QKV / subgroup softmax 由设备能力选路每个onEncode决定本次 dispatch 哪些 shader。集成新特性时务必把目标 shape 代入选路逻辑确认真正落到新路径避免改了没被调度。第 5 步编译验证cd project/android/build_64 cmake .. -DMNN_VULKANON -DMNN_VULKAN_IMAGEOFF -DMNN_BUILD_LLMON \ -DMNN_SUPPORT_TRANSFORMER_FUSEON -DMNN_LOW_MEMORYON -DMNN_ARM82ON make llm_demo -j8 # SEP_BUILDOFF → 会重编 libMNN.so注意make MNN不会连带重编 Vulkan 静态库必须用make llm_demo改 shader 后AllShader.cpp变了也会触发重编。推送到真机并运行换 shader 后必须先清 pipeline cache否则 stale cache 直接 segfaultadb push libMNN.so llm_demo /data/local/tmp/MNN/ adb -s serial shell cd /data/local/tmp/MNN rm -f tmp/mnn_cachefile.bin \ LD_LIBRARY_PATH. ./llm_demo model/config_vk.json prompt.txt ndecode五、第四步验证——正确性、性能、兼容性三管齐下正确性三层 oracle数值层dump tensor 对比op 层MNNV2Basic.out单层对比端到端跑完整模型关 sampler 随机性temperature:0.0, greedyCPU/Vulkan 同 prompt 前 N token 应一致fp16 误差内。数学等价的改动应与 baseline逐 token 一致。支持和不支持特性的设备都要测——fallback 路径同样必须验证正确。若 CPU oracle 不可用低 bit 路径 CPU 本身就乱用冻结 baseline 二进制做 GPU-baseline vs GPU-opt 的逐 token 贪心对比。注意 Mac MoltenVK 行为不代表 Android最终验收必须 Android 真机。性能交替 A/B 而非顺序对比手机 GPU 有两大测量陷阱冷启动首跑不算数首跑含 auto-tune pipeline 编译清 cache 后第一次仍是冷的先 warm 再测跨时段热漂移约 ±8–10%先测 base 再测 opt 不可比必须交替 A/B——base 和 opt 两份二进制或同一二进制不同 env背靠背配对base→opt→base→opt看每轮配对里 opt 是否稳定胜出。换 shader 做 A/B 时每轮跑前rm tmp/mnn_cachefile.bintok/s 不含 load清 cache 只影响 load 时间A/B 仍公平。带MNN_GPU_TIME_PROFILEON的 build 只用于看相对占比会放大绝对耗时并改变调度最终收益以不带 profile 的干净 build 端到端 tok/s 为准。兼容性全量 op 回归确保无回归不支持特性的设备Mali / 旧驱动走 fallback 无性能退步。六、第五步文档与提交在 shader 头部注释中写明特性 / 适用设备 / 原理 / fallback。提交信息遵循仓库约定格式[Vulkan:Feature] Add 特性 for 算子 - Runtime feature detection via VulkanDevice - New shader using 特性, --target-env vulkan1.1 - Fallback path for unsupported devices (Mali/old driver) - Tested on 设备, 加速比七、通过标准与常见问题排查通过标准示例充分理解能解释原理/核心逻辑特性检测正确runtime 能判设备是否支持shader 适配 MNNFLOAT 宏 / NC4HW4 / spec constant / target-envfallback 存在且正确不支持设备走原路径shader 外科式重生成 三处注册AllShader.cpp 无污染支持 不支持设备都验证通过有交替 A/B 性能数据常见问题速查表问题原因修复shader 编译失败 / spirv-opt 拒绝缺--target-env vulkan1.1coop/subgroup/memory_scope 必带运行 segfaultpipeline cache stale / descriptor layout mismatchrm mnn_cachefile.bin核对 binding/push constant size host↔shader支持设备上性能反降特性 overhead 收益如 coop K 维太小限制特定 shape 才走新路径如 coop-QK headDim128 属负收益fallback 被破坏改动影响原路径用独立 pipeline/分支隔离别改原 shader数值乱用错 layoutcoop 只 int4/8dispatcher 显式 gate见 conv1x1 选路中isLowBit23判断八、更多参考Vulkan 优化 Skill 总览方向 C 与方向 A/B 的关系、编译与真机运行、shader 修改流程防污染、正确性三层 oracle、交替 A/B 方法论。优化手册coop matrix、subgroup 归约、epilogue 融合等技巧与陷阱的唯一权威来源§1 性能分析方法论、§2 优化技巧、§3 常见陷阱、§4 Packed weight 设计、§5 速查表、§6 候选方向。基准建立、kernel 优化、集成回归方向 A/B 的配套流程。源码参考VulkanDevice.hpp特性查询接口、VulkanDevice.cppcoopMat 运行时探测、VulkanConvolution.cppconv1x1 coop 选路样本、VulkanShaderMap.cppshader 注册样本、macro.json变体登记格式。【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考