ONNX Runtime CUDA MatMulNBits 算子深度解析权重量化矩阵乘的分派链路、GEMV 内核与 fpA_intB 路径【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime导读本文基于 ONNX Runtime 官方算子文档 docs/contrib_ops/cuda/matmul_nbits.md系统剖析com.microsoft::MatMulNBits在 CUDA Execution Provider 中的完整实现从算子 Schema 与权重内存布局到fpA_intB CUTLASS → 融合 GEMV 快速路径 → 反量化 cuBLAS GEMM 回退的分层分派链再到 MoE Router 特化内核、bias 处理与全部可调环境变量。读完本文你将能够理解 weight-only 量化线性层含 MoE 路由与 LM head在 GPU 上如何被高效执行掌握权重预打包prepacking两种布局的选择依据并能据此调优推理路径与排查性能瓶颈。MatMulNBits 计算Y A · dequant(B)ᵀ ( bias)其中B是N × K的 4/8 bit 权重量化矩阵采用 block-wise分组scale 与可选的 zero point。它是 weight-only 量化线性层的基础算子也是 MoE Router、LM Head 等大模型组件的性能关键。1. 算子 Schema属性与输入1.1 属性表Attribute含义K输入特征维度A的列数即逻辑B的列数。N输出特征维度逻辑B的行数。bits量化位宽4或8。block_size沿K方向的量化组大小16 / 32 / 64 / 128每组一个 scale以及可选的 zero point。accuracy_level内部处理A时的最低精度级别默认0表示未设置。weight_prepackedCUDA fpA_intB 权重布局选择器详见 §2.10默认表示B为标准 MatMulNBits 布局、运行时可预打包1表示B已按 CUDA SM80 fpA_intB 布局预打包2表示B已按 CUDA SM90HopperfpA_intB 布局预打包由原生 SM90 内核消费要求 SM90 设备且block_size ∈ {64, 128}。原生 SM90 内核在 Windows/MSVC 构建中不编译CUDA 13 host stub 的 over-aligned TMA 参数触发 MSVCC2719详见 moe_qmoe.md §14.1此类构建下默认的0/1布局在 Hopper 上运行 SM80 兼容内核。1.2 输入表InputIndexNotesA0ActivationsFP16 / BF16 / FP32形状[M, K]。B1Packed 4/8-bit weights。scales2每组一个 scale元素类型与A相同。zero_points3可选。打包整数对称量化默认或与A相同类型。g_idx/reorder_idx4可选的 group/reorder indexact-order 量化。bias5可选的[N]bias加到输出上。M是展平后的token 数M 1即 decode / GEMV 场景正是快速内核的目标。从源码看构造时这些属性在 matmul_nbits.h 中通过info.GetAttrint64_t(K, K_)等调用强制读取并校验block_size_ 0has_zero_points_/has_g_idx_/has_bias_则根据输入数量与InputDefs()是否存在来判断同时记录zero_points是否与 scale 同类型is_zero_points_scale_same_type_。算子注册matmul_nbits.cc覆盖三个 T1 变体float、MLFloat16、BFloat16且 CUDA EP 额外注册了 FP32 变体但 FP32 永远不会被判定为 fpA_intB 兼容。2. 权重格式对于 4-bitB以[N, ceil(K / block_size), block_size / 2]字节存储每个 expert/输出行n是连续的K/2字节每个字节打包两个 4-bit 权重其前在概念上对应scales张量中的K / block_size个 scale。量化默认是列向分组column_wise_quant_blk_ true行向布局把K的块交错在N之间无法沿N切分这会禁用分块回退路径。对称 4-bit 权重存储值0..15反量化为(q − 8) · scale当没有zero_points输入时快速内核硬编码 zero point 为8。这一默认值同样体现在 fpA_intB 预打包逻辑中matmul_nbits.cc 定义了kDefaultZeroPoint4Bit 8.0f与kDefaultZeroPoint8Bit 128.0f。2.1 CUDA fpA_intB 预打包布局weight_prepackedweight_prepacked1表示输入B的张量形状与字节数与标准 MatMulNBits 的B完全相同但字节已被重排为 CUDA fpA_intB SM80 weight-only 布局。在 ORT 预打包阶段CUDA EP 会把这些字节直接交给 fpA_intB 内核而不做额外的 GPU 拷贝当B已驻留设备例如已固定到 GPU 的 constant initializer时整个预打包步骤被跳过PrePack_B中is_packed false直接返回当B驻留主机时照常传输到设备但不再执行运行时权重转置 / mixed-GEMM 预处理步骤。这一点在 matmul_nbits.cc 中有明确实现预打包权重会先做字节数校验若 tensor 已在 GPU 上则跳过拷贝否则仅做一次cudaMemcpyAsync即完成不经过unpack_uint4_transposed_to_int8_direct_cuda/transpose_uint8_matrix_and_convert_to_int8preprocess_weights_for_mixed_gemm_cuda的运行时重排流水线。通过 Python 暴露的离线 CUDA 打包器可产生该布局from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant prepacked_flat _cuda_quant.pack_weights_for_cuda_mixed_gemm( q_weight.reshape(N, -1), N, K, bits, 80 ) prepacked_b np.asarray(prepacked_flat, dtypenp.int8).view(np.uint8).reshape(q_weight.shape)最后一个参数是目标打包架构80→ SM80 布局由 SM80 CUTLASS 内核消费包括在新 GPU 上通过兼容路径消费节点上设置weight_prepacked190→ 原生 SM90Hopper布局节点上设置weight_prepacked2。weight_prepacked2选择原生 SM90Hopper TMA/WGMMAmixed-GEMM 内核及其 Hopper 权重布局。它要求计算能力 9.0 的设备且block_size ∈ {64, 128}SM90 内核要求group_size是 64 元素 Hopper K tile 的倍数因此block_size32走非 Hopper 内核。在 SM90 设备上运行时预打包weight_prepacked0与 SM80 预打包weight_prepacked1的权重继续路由到 SM80 CUTLASS 内核/布局。值得注意的实现细节原生 SM90 内核是否被编译由COMPILE_HOPPER_TMA_GEMMS宏决定Windows/MSVC 构建未定义而ValidateSm90PrepackedWeightSupport(sm, block_size)被刻意提取为纯 host 函数见 matmul_nbits_sm90_validation.h可在无 Hopper GPU 的机器上用合成(sm, block_size)做单元测试。IsNativeSm90FpAIntBGemmCompiled()则以非 inline 单一定义放在matmul_nbits.cc中避免各翻译单元因宏不一致产生 ODR 冲突。3. 分派链Dispatch ChainMatMulNBitsT::ComputeInternal总是先尝试最廉价且适用的路径失败后逐级回退到更通用的实现从 matmul_nbits.cc 的ComputeInternal源码可以精确对应这张图BF16 门槛BFloat16 在计算能力 8.0 的设备上直接返回NOT_IMPLEMENTEDfpA_intB 路径仅在USE_FPA_INTB_GEMM构建下如果构造期判定has_fpA_intB_gemm_先经 profiler 挑选最优 tactic见 §6命中后直接返回融合快速路径门槛仅当reorder_idx nullptr且zero_points若有为打包整数非T类型时才尝试TryMatMulNBitsbias 兜底重试若快速路径因 bias 失败则以bias nullptr重试成功后用独立的LaunchMatMulNBitsBiasAdd内核补加 bias反量化 GEMM都不满足时如M 1、存在reorder_idx、typed zero points、形状不支持执行反量化后接 cuBLAS GEMM最后同样补加 bias。快速融合路径的尝试条件可以总结为一句没有reorder_idx且zero_points若有是打包整数。这与文档和源码中的(reorder_idx_data nullptr) (!zero_points || !zero_points-IsDataTypeT())判定完全一致。4. 快速路径 —— 融合 GEMVTryMatMulNBitsmatmul_nbits.cuh按位宽分派bits 8→TryMatMul8Bits不支持 bias若设置 bias 直接返回falsebits 4→TryMatMul4Bits。TryMatMul4Bitsmatmul_4bits.cu首先应用所有融合内核共用的守卫条件n % kColsPerThreadBlock (8) 0 and k % 8 0 and m 16即融合路径只处理单个 token 或小批量。守卫通过后按M1内核§4.1、小 M 批量内核§4.2、Router 特化§4.3依次选择。4.1 通用 4-bit M1 GEMV 内核MatMulFloat4BitsKernelM1T, block_size, has_zero_point实现在 matmul_4bits_m1_impl.cuh由 dtype 专属的翻译单元实例化matmul_4bits_half.cu/matmul_4bits_bfloat16.cu/matmul_4bits_float.cuLaunchgrid(ceil(N / 8), 1)block(warpSize, 8)—— 每个输出列一个 warp每 block 8 个 warp即kColsPerThreadBlock 8。Scales / zero points每个线程块仅一次载入共享内存若总需求超过每块共享内存上限launcher 返回false。内层循环每个 lane 每次迭代消费 8 个打包 int4 权重随后做 warp 归约三级 unroll 16/4/1 加 remainder最后 warp shuffle 归约。支持的block_size16 / 32 / 64 / 128。Bias不支持 —— 当bias ! nullptr时内核向TryMatMul4Bits返回falsebias 的处理见 §7。4.2 小 M 批量 GEMV2 ≤ M ≤ 16当2 M 16时TryMatMul4Bits先尝试寄存器分块的 half/BF16 批量内核TryMatMulBatched4Bits源自matmul_4bits_batched.cuh再尝试共享内存版MatMulFloatInt4KernelSmallM。这些路径将每个打包权重字只反量化一次并在多个激活行之间复用。源码注释揭示了设计动机matmul_4bits.cu单行 M1 内核每个输出行启动一个 block各行会独立地重复读取并反量化全部B权重流量与反量化工作量随M线性放大而对2 M capkSmallMMax 16每个反量化权重可对寄存器中CtaM行激活做累加把权重流量削减到ceil(M/CtaM)分之一 —— 这与 TensorRT-LLMweightOnlyBatchedGemv/ AWQ / llama.cpp MMVQ 在小批量下的设计一致。half/bf16 走寄存器分块内核CtaM x CtaNtile2 宽累加器float 没有张量核 GEMM 回退因此同一范围内使用共享内存小 M 内核。4.3 Router GEMV 特化MatMulFloatInt4RouterKernelT, BlockSize是面向 MoE-Router GEMVoutput(1, N) A(1, K) · dequant(B(N, K)) bias(N)的特化内核由IsSupportedRouterGemvShape选中条件为zero_points nullptr对称量化、M 1block_size ∈ {32, 64}且K % block_size 0(N, K)精确匹配 gated router 形状ModelNexpertsKhidden sizegpt-oss-20b322880设计要点均可在 matmul_4bits.cu 源码中核实每个 expert 列一个 warp每 block 8 个 warpkColsPerThreadBlock 8grid 为(N / 8, 1)。N通过 grid 在运行时传入、K作为运行时 kernel 参数传入因此每个(T, BlockSize)只需一次实例化即可服务所有 router 形状只有BlockSize是模板参数它决定 scale 步长。不使用共享内存scale 直接从全局内存读取这些极小尺寸下天然 L2 驻留省去了通用内核的 staging__syncthreads。Bias 融合lane 0 在 store 前加上bias[n_id]。组大小 32 vs 64两者都支持因为kPerIter 256可被两者整除保证每迭代的 scale 步长精确。相对 3264 把 scale 加载次数减半按模型精度/延迟权衡选择量化粒度即可。每个 expert 一个 scale 的 per-row / per-column 量化不是受支持的 MatMulNBits 布局被刻意排除。与通用内核保持相同的不变式三级 unroll16/4/1 remainder、warp-shuffle 归约。若要新增另一个 router只需在IsSupportedRouterGemvShape中加入它的(N, K)只要N % 8 0且K % block_size 0无需改动内核。该特化可通过环境变量ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION关闭见 §8。5. 回退路径 —— 反量化 GEMM当 fpA_intB 路径与融合 GEMV 都不适用时例如M 1、存在reorder_idx、typed zero points 或不支持的形状算子将B反量化到 scratch buffer再执行稠密 cuBLAS GEMMDequantizeBlockwise4b/DequantizeBlockwise8b或列向专用 helper把打包权重展开为T写入N × K_padded的 scratch buffer其中K_padded ceil(K / block_size) · block_size随后单个 cuBLASGEMMtransb true产出Y。分块变体Chunked。当N很大时物化完整的N × K_padded反量化矩阵可能主导显存。实现在以下条件同时满足时沿N把反量化GEMM 切成多块列向量化column-wise且无 reorder_idx且force_chunked_被设置或scratch 256 MB且N 2 × chunk_target_rows。chunk_target_rows默认为 32768可配置见 §8。源码注释给出的设计理由matmul_nbits.cc很清晰256 MB 是 scratch 预算超过后完整N*K_paddedbuffer 会在 8–12 GB 消费级 GPU 上主导显存并带来 OOM 风险N 2*chunk_target_rows默认 65536保证至少两块摊薄每块 cuBLAS 调用开销而 32768 行/块足以打满 SM同时保持 scratch ≲128 MB。分块循环中每块用cublasGemmHelper计算C[:, n_start:n_end] A[M,K] B_chunk[chunk_n, K]^T输出按完整Ldc N步长写入matmul_nbits.cc。6. fpA_intB_gemm 路径CUTLASS weight-onlyonnxruntime_USE_FPA_INTB_GEMM是主构建开关CUDA 构建默认ON。其默认内核集合刻意保持精简排除原生 Hopper 内核覆盖 SM80 布局的 RC 模型契约FP16 activations INT4 或 INT8 weightsscale-only 量化block_size32无 zero-point、bias、g_idxN % (bits8 ? 32 : 64) 0、K % block_size 0、sm_ 75未打包权重或weight_prepacked1SM80 布局。设置onnxruntime_USE_FPA_INTB_GEMM_FULLON可构建遗留的完整内核矩阵。完整模式额外支持 BF16、block_size 64 与 128、zero-point、bias以及原生 SM90 布局weight_prepacked2原生 SM90 内核只支持block_size ∈ {64, 128}见 §2.1。CheckFpAIntBEligibilitymatmul_nbits.cc是 fpA_intB / CUTLASS weight-only-GEMM 资格判定的唯一权威来源构造器与 Level-1 workspace 估算共用二者永不冲突。其核心逻辑compact 构建仅接受 FP16 输入input0_elem_type FLOAT16完整构建接受 FP16/BF16FP32 变体永远不兼容预打包权重强制开启 fpA_intB 路径prepacked || (fpa_intb_option ! 0)compact 模式base_okblock_size 32 (nbits 4 || nbits 8) !has_zero_points !has_g_idx !has_bias weight_prepacked ! 2 N % (nbits8 ? 32 : 64) 0 K % block_size 0 device_sm 75完整模式放宽为block_size ∈ {32,64,128}、允许 zero-point/bias但依旧排除g_idxweight_prepacked2额外要求device_sm 90 block_size ∈ {64, 128}。当通过ORT_FPA_INTB_GEMM启用后符合条件的 MatMulNBits 节点使用源自 TensorRT-LLM 的 CUTLASS weight-only 内核。weight 与 scale 输入完整模式下还有 zero-point必须是 ORT 可预打包的 constant initializer。运行时由 profiler 挑选最优 tactic小M可能使用专用 CUDA GEMV 内核bestTactic-enableCudaKernel否则用 CUTLASS grouped GEMM。compact fpA_intB GEMV 不支持 zero-point 或 bias。此路径激活时优先于 §3 中除它之外的所有路径。weight_prepacked0时CUDA EP 在 ORT 预打包阶段把标准 MatMulNBits 权重 initializer 预处理为 fpA_intB 布局weight_prepacked1时 initializer 视为已预处理仅做字节数检查后直接拷贝对应 matmul_nbits.cc 的PrePack_B实现。预打包权重是严格约束的若 ORT 未以onnxruntime_USE_FPA_INTB_GEMMON构建任何非零weight_prepacked值都会在 kernel 构造时抛出异常任何非零weight_prepacked都强制启用 fpA_intB 路径因此启用开关session configep.cuda.fpa_intb_gemm或环境变量ORT_FPA_INTB_GEMM对预打包权重被忽略—— 布局选择在导出时已固定运行时无法关闭非零weight_prepacked要求输入A为 FP16compact 默认构建或 FP16/BF16完整构建因为只有 CUDA fpA_intB 路径消费该布局weight_prepacked必须与所选内核期望的布局一致1是 SM80 布局2是原生 SM90Hopper布局2额外要求计算能力 9.0 设备与block_size ∈ {64, 128}否则被拒绝。另一个易踩的坑源码注释特别强调matmul_nbits.ccGEMV 内核必须用打包时所针对的架构FpAIntBPackingSmForKernel()而非裸设备 SM 启动因为 GEMV 的 interleave 布局与架构相关 —— 架构在 [90,100) 使用ColumnMajorInterleavedForHopper而 SM80 打包使用ColumnMajorInterleaved。若在 SM90 上误传设备 SM会按 Hopper interleave 读取 SM80 打包权重产生错误结果。FpAIntBPackingSmForKernel()复用EffectiveFpAIntBWorkspaceSm()matmul_nbits.h仅当device_sm 90 weight_prepacked 2时为 90否则为 80保证 workspace 估算与运行时布局选择永远一致。CUDA graph 捕获下也有专门处理捕获期间禁止惰性 profilingprofiling 会启动内核、记录/同步事件、分配/释放 scratchComputeInternal通过isCapturing(stream)检测后改查已 profiling 的 bucket捕获前的 warmup 会填充这些 bucket若找不到可用 tactic 会明确报错提示捕获前先做一次 warmup 推理。7. Bias 处理只有 Router 特化§4.3在 GEMV 内部融合 bias。对其他所有快速路径形状TryMatMul4Bits/TryMatMul8Bits在存在 bias 时返回false。随后ComputeInternal按两步处理对应 matmul_nbits.cc以bias nullptr重试TryMatMulNBits成功后用独立的MatMulNBitsBiasAdd内核LaunchMatMulNBitsBiasAdd补加 bias累加在 float 中以保证 half/bfloat16 精度若快速路径仍不适用落入反量化GEMM 回退§5其本身忽略 bias随后用同一个 bias-add 内核补加。bias-add 内核本身是一个 grid-stride 循环matmul_4bits.cuoutput[idx] output[idx] bias[col]其中col idx % n并以 float 中间量累加后转回Tgrid 上限被钳制在 CUDAgridDim.x上限2³¹−1以内超出部分由 grid-stride 循环覆盖。8. 环境变量VariableType / defaultEffectORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATIONbool,0禁用 Router GEMV 特化§4.3相关形状回退到通用 GEMV / 反量化路径。适合 A/B 基准对比。ORT_FPA_INTB_GEMMint/string,0启用 CUTLASS weight-only 路径§6。0或off禁用其他值启用。ORT_MATMULNBITS_FORCE_CHUNKEDint,0无视尺寸启发式强制使用分块反量化GEMM 回退§5。ORT_MATMULNBITS_CHUNK_SIZEint64,32768分块回退中每块的目标行数。值 1时重置为默认值。环境变量通过 ORT 跨平台的ParseEnvironmentVariableWithDefaulthelper 读取Windows 上安全而非std::getenv。此外matmul_nbits.h 中还定义了一组session config键它们优先于同名环境变量config wins并且同时被内置 CUDA EP 与 CUDA plugin EP 读取ep.cuda.fpa_intb_gemm↔ORT_FPA_INTB_GEMM0/off关闭其余开启解析逻辑见ParseFpAIntBEnabled不区分大小写且空串/0/off视为关闭ep.cuda.fpa_intb_profile_m↔ORT_FPA_INTB_PROFILE_M初始 profile M buckets空则用 profiler 默认 bucket 集。还有两个调试/诊断开关ORT_FPA_INTB_DEBUG1会在运行时打印所选 tactic、kernel 路径GEMV CUDA 内核 vs CUTLASS GEMM、权重格式runtime-prepacked-SM80 / offline-prepacked-SM80 / offline-prepacked-SM90 / raw、设备与打包 SM用于追踪 SM90 布局选择问题matmul_nbits.ccORT_LLM_VERBOSE 1时还会输出每个(m, n, k, group_size)的最佳 tactic。9. 测试CUDA EP 内部测试通过CUDA_EP_Unittest运行见 cuda_provider_test.cc。从onnxruntime_provider_test运行./onnxruntime_provider_test --gtest_filterCUDA_EP_Unittest.*该 wrapper 执行内部 CUDA-UT 共享库覆盖 fpA_intB_gemm_kernel_test.cc 下的 fpA_intB / MatMulNBits groupwise GEMM 测试以及 matmul_nbits_sm90_validation_test.cc 下的 SM90 校验测试后者利用ValidateSm90PrepackedWeightSupport是纯 host 函数的特点可在无 Hopper GPU 的机器上以合成(sm, block_size)验证拒绝逻辑。Python 算子测试onnxruntime/test/python/transformers含 QMoE / GEMV profiling 辅助脚本如profile_qmoe_gemv.sh。CUDA 预打包权重一致性测试test_op_matmulnbits_prepacked_cuda.py。这些测试用onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm(..., 80)生成weight_prepacked1的 initializer并与运行时 fpA_intB 预打包在 int4/int8、GEMV/GEMM 形状的M下做输出对比。不支持的预打包配置的构造失败测试位于 matmul_4bits_test.cc。GEMV profiling 的基线与方法论记录在 qmoe_gemv_experiments.md。对比 Router 特化与通用路径对同一模型分别以ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION1与默认值运行即可。修改内核后的重建提醒修改任何.cu内核后需重建 CUDA providerninja onnxruntime_providers_cuda并重跑相关测试注意仓库构建说明中 nvcc 增量构建的注意事项。总结MatMulNBits 在 CUDA EP 上的实现是一套典型的分层分派 专用内核 兜底回退设计构造期依据属性与设备确定 fpA_intB 资格含预打包布局的严格校验运行期按fpA_intB GEMV/GEMM → 融合 4/8-bit GEMVM1 / 小 M 批量 / Router 特化→ 反量化 cuBLAS GEMM全 N 或分块的优先级逐级尝试。理解这一链路后你可以通过weight_prepacked提前固定布局、通过ORT_FPA_INTB_GEMM/ep.cuda.fpa_intb_gemm控制 CUTLASS 路径、通过ORT_MATMULNBITS_*系列环境变量调节分块回退行为并用ORT_FPA_INTB_DEBUG与ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION做可量化的 A/B 排查 —— 这些能力共同支撑起大模型 weight-only 量化推理在 CUDA 上的高效落地。【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考