人工智能大模型模型优化模型量化模型压缩【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址https://gitcode.com/GitHub_Trending/te/Model-Optimizer点击查看免费下载本文档深入解析 Model-Optimizer 仓库中的 Conv3D Implicit GEMM CUDA 内核conv/README.md 的配套实现它以隐式 GEMM 形式将三维卷积映射到 BF16 WMMA Tensor Core 上执行并在 GEMM 内部融合 FP4E2M1伪量化用于加速视频生成模型如 Wan2.2VAE 中的 Conv3D 层。读完本文你将掌握该内核的调用方式、精度格式、源码级实现原理、与 ModelOpt NVFP4 PTQ 流水线的自动集成机制以及如何运行配套测试与基准验证其正确性和性能。1. 内核定位为什么视频生成模型的 Conv3D 需要专用内核视频生成模型如 Wan2.2的 VAE 解码/编码路径中包含大量 3D 卷积层Conv3D。与 2D 卷积相比Conv3D 的计算密度更高且激活张量形状为[N, Cin, D, H, W]空间维度帧、高、宽显著放大参与卷积的乘加次数。Model-Optimizer 中的modelopt.torch.kernels.quantization.conv模块为此类场景提供了一个专用内核Conv3D Implicit GEMM——使用 BF16 WMMAm16n16k16Tensor Core 执行矩阵乘并可选地在 GEMM 内部融合 FP4E2M1激活伪量化。从仓库结构看该模块位于 modelopt/torch/kernels/quantization/conv/与同目录下的attention/、gemm/、common/FP8、NVFP4 量化工具等内核模块并列。需要特别说明的是该内核不是独立使用的黑盒它已被集成进modelopt.torch.quantization的_QuantConv3d当用户通过 ModelOpt PTQ 对nn.Conv3d层应用 NVFP4 量化时推理路径会自动切换到该内核。官方文档中明确说明目前只在视频生成模型如 Wan2.2的 VAE Conv3D 层上做过测试。2. 运行环境要求根据 conv/README.md使用该内核需要满足以下条件要求说明GPU 架构SM80Ampere 或更新BF16 WMMA Tensor Core 自 Ampere 起可用PyTorch支持 JIT C 扩展编译torch.utils.cpp_extension分组卷积不支持groups必须为 1运行时还有双重保护详见第 4 节架构守卫CUDA 内核源码用#if __CUDA_ARCH__ 800保证在旧架构上编译为空桩Python 侧的_get_cuda_module()则通过torch.cuda.get_device_capability()在运行时拒绝 SM 小于 8 的设备。3. 数据类型与精度设计该内核在不同阶段使用不同的精度其设计意图是输入输出精度透明、内部计算用 Tensor Core、累加保持高精度阶段精度输入 / 输出张量FP32、FP16 或 BF16dtype 被原样保留内部计算通过 WMMA m16n16k16 Tensor Core 以 BF16 执行累加FP32 累加器FP4 激活量化E2M1 值FP8 E4M3 scale从 implicit_gemm_cuda.py 的 Python 包装器可以看到输入x、权重w、bias 与act_amax在进入 CUDA 内核前会被统一转为 FP32 contiguous 张量内核内部计算完成后输出再通过y.to(x.dtype)恢复原始 dtype。这意味着该内核对 FP32/FP16/BF16 输入都是透明可用的。4. 源码结构三个文件的分工该内核由三个文件组成conv/README.md 的 Files 一节文件职责implicit_gemm_cuda.pyPython API 与 JIT 编译入口implicit_gemm_kernel.cuCUDA 内核BF16 WMMA FP4 量化implicit_gemm_binding.cppPyTorch C 扩展绑定Python 包装器在 implicit_gemm_cuda.py 中通过torch.utils.cpp_extension.load()完成 JIT 编译编译名称为conv3d_implicit_gemm_cuda_v20_wmma并携带以下编译选项-O3 --use_fast_math -lineinfo --ptxas-options-v -stdc17C 绑定层implicit_gemm_binding.cpp通过PYBIND11_MODULE暴露两个函数conv3d_implicit_gemm_cuda与fp4_fake_quant_cuda。5. 核心实现Implicit GEMM 的代码级剖析5.1 从 Conv3D 到 GEMM 的映射传统卷积可直接调用 cuDNN 的 im2col GEMM 流程但显式 im2col 会引入巨大的全局内存开销。该内核采用implicit GEMM方式不物理展开 im2col 矩阵而是在内核加载 A-tile激活 tile时现场计算每个元素的空间坐标并直接从x_pad中取值。从 implicit_gemm_kernel.cu 可以看到M 维索引m_idx被解构为(n_batch, od_val, oh_val, ow_val)利用ODHW、OHW、OW预计算常量做整数除法K 维索引k_idx被解构为(c, kd_v, kh_v, kw_v)最终元素地址通过od_val * sd kd_v * dd等公式计算得出。这一过程完全没有查找表文档将其列为关键优化点之一On-the-fly spatial index computation (no global memory lookup tables)。5.2 Tile 配置、Warp 划分与共享内存布局内核的模板参数与 tile 配置在 implicit_gemm_kernel.cu 与 C 分发逻辑implicit_gemm_binding.cpp中定义参数值BLOCK_M / BLOCK_N / BLOCK_K64 / 64 / 256WARPS_M / WARPS_N2 / 4共 8 warp 256 线程每 warp 计算输出 tile32×16WARP_TILES_M2WARP_TILES_N1 个 16×16 WMMA 片段WMMA 指令wmma::mma_syncm16n16k16BF16 输入、FP32 累加共享内存As[BLOCK_M][BK_STRIDE]M-majorBs[BLOCK_K][BN_STRIDE]K-major全部 BF16共享内存总量64*(2568)*2 256*(648)*2 70,656字节约 69 KB两个重要细节Bank 冲突消解BK_STRIDE BLOCK_K 8、BN_STRIDE BLOCK_N 8即每行额外填充 8 个 BF16 元素16 字节既保持 16 字节对齐又打破 bank 冲突implicit_gemm_kernel.cu。共享内存复用输出阶段将同一块动态共享内存 reinterpret 为 FP32 staging buffer各 warp 先把 WMMA 累加片段store_matrix_sync到共享内存再由全部线程协作拷贝到全局内存并在拷贝时叠加 bias 与边界检查implicit_gemm_kernel.cu。5.3 L2 友好的 block 调度内核采用 swizzled grid 调度L2_SWIZZLE_GROUP 8通过将 M 维 block 按组重排使相邻 block 访问的激活/权重数据在 L2 缓存中具有更好的空间局部性implicit_gemm_kernel.cu。C 侧在计算 grid 大小时也按 8 的组粒度做了 paddingimplicit_gemm_binding.cpp。5.4 架构守卫编译期与运行期的双重保护编译期整个内核体被#if defined(__CUDA_ARCH__) __CUDA_ARCH__ 800包裹implicit_gemm_kernel.cu。因为 PyTorch 默认的-gencode列表可能包含 sm_75而 sm_75 没有 BF16 WMMA 片段所以旧架构上会编译为空桩。运行期_get_cuda_module()通过torch.cuda.get_device_capability()读取 GPU 算力并用_MIN_SM_MAJOR 8在分发前拦截不支持的设备implicit_gemm_cuda.py。6. FP4 量化路径E2M1 值与 FP8 E4M3 scale6.1 E2M1 量化表FP4E2M11 符号位 2 指数位 1 尾数位只有 8 个可表示值{0, 0.5, 1, 1.5, 2, 3, 4, 6}。内核中的量化函数fp4_quantize_value使用无分支的谓词选择链实现implicit_gemm_kernel.cuq (scaled 5.0f) ? 4.0f : 6.0f; q (scaled 3.5f) ? 3.0f : q; q (scaled 2.5f) ? 2.0f : q; q (scaled 1.75f) ? 1.5f : q; q (scaled 1.25f) ? 1.0f : q; q (scaled 0.75f) ? 0.5f : q; q (scaled 0.25f) ? 0.0f : q;即边界点为 0.25 / 0.75 / 1.25 / 1.75 / 2.5 / 3.5 / 5.0。测试用例 test_implicit_gemm.py 中有对应的 Python 参考实现_py_fp4_fake_quant_ref并在TestFP4FakeQuantValues中对边界上下取值做了逐值验证。6.2 scale 计算与 FP8 E4M3 round-trip每个 FP4 block 的 scale 计算分三步implicit_gemm_kernel.cufloat scaled block_max / (6.0f * global_scale); // global_scale act_amax / (6.0f * 448.0f) scaled fminf(scaled, 448.0f); float quantized fp8_e4m3_round_trip(scaled); // scale 量化为 FP8 E4M3 return quantized * global_scale;fp8_e4m3_round_trip直接操作 IEEE 754 位模式implicit_gemm_kernel.cu将 FP32 的尾数截断到 3 bit、处理进位与指数越界max 448、min normal exponent -6实现 E4M3 的完整 round-trip。若计算出的 scale 过小 1e-5则回退为 1.0以避免除零implicit_gemm_kernel.cu。6.3 寄存器融合加载即量化这是该内核的核心优化之一FP4 激活量化在 A-tile 加载时于寄存器内完成Register-fused FP4 quantization (quantize during A-tile load, eliminates sync)。其流程是implicit_gemm_kernel.cuPass 1每个 lane 从全局内存加载ELEMS_PER_LANE个值到寄存器数组Pass 2对每个 FP4 子块先经warp_reduce_max归约得到block_max再计算 scale / inv_scale此时不需要额外的__syncthreads因为归约在 warp 内完成Pass 3对每个值计算符号、量化到 E2M1 最近层级、再乘回 scale直接以 BF16 写入共享内存。由于量化在 A-tile 加载过程中就地完成激活数据不需要额外的全局内存往返no extra global memory pass。static_assert(BLOCK_K % FP4_BLOCK_SIZE 0)与FP4_BLOCK_SIZE 16保证了子块划分与 K-tile 对齐implicit_gemm_kernel.cu。注意一个细节该路径只量化激活权重由调用方预先量化好传入测试注释中也明确写道 the fused kernel does NOT quantize weights。6.4 独立的fp4_fake_quant内核除了融合路径模块还提供了独立的fp4_fake_quant函数使用与 GEMM 内 A-tile 加载完全相同的 CUDA device 函数fp4_quantize_value、warp_reduce_max、quantize_scale_fp8、fp8_e4m3_round_trip。其内核fp4_fake_quant_kernel将输入视为[num_blocks, block_size]每个 warp 处理一行一个 FP4 block同样两遍完成implicit_gemm_kernel.cu。测试中它被用作融合路径的非融合参考与跨实现一致性验证对象。7. Python API 与实战用法7.1conv3d_implicit_gemm_cuda导入路径from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda参数签名implicit_gemm_cuda.py参数说明x输入张量[N, Cin, D, H, W]w权重张量[Cout, Cin, kD, kH, kW]bias可选 bias[Cout]stride卷积步长(D, H, W)也接受单个 intpadding卷积填充(D, H, W)也接受单个 int 或 6 元组dilation卷积膨胀(D, H, W)也接受单个 intact_amax激活绝对值最大值标量张量quant_actTrue时必填quant_act是否对激活启用 FP4 伪量化fp4_block_sizeFP4 量化 block 大小仅允许16、32、64、128、256需要留意几个 Python 侧行为非法fp4_block_size会抛出ValueErrorimplicit_gemm_cuda.py输入必须是 5D 张量且w的Cin必须等于x的Cin否则报分组卷积不支持的错误implicit_gemm_cuda.pypadding 由F.pad预先在主机侧完成支持不对称 padding权重会被 reshape 为[K, Cout]的[Cout, K]转置布局以匹配 B-tile 的 K-major 读取输出先以[M, Cout]平坦布局返回再view回[N, Cout, OD, OH, OW]并转回输入 dtype因此输出形状与torch.nn.functional.conv3d完全一致。7.2 官方 README 的完整示例以下示例原样继承自 conv/README.mdimport torch from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op x torch.randn(1, 128, 21, 60, 106, devicecuda) w torch.randn(512, 128, 3, 3, 3, devicecuda) block_size 128 # Without FP4 activation quantization (drop-in-style Conv3D call) out conv3d_implicit_gemm_cuda(x, w, stride(1, 1, 1), padding(1, 1, 1)) # Optional FP4 block quantization of weights along the GEMM K dimension. # The kernels A-tile (activations) is quantized along K Cin*kD*kH*kW, # so weights must be flattened to [Cout, K] before quantizing to match. Cout, Cin w.shape[:2] K Cin * w.shape[2] * w.shape[3] * w.shape[4] w_flat w.reshape(Cout, K) w_q_flat dynamic_block_quantize_op( w_flat, block_size, w_flat.abs().max().unsqueeze(0), 4, # num_bits 2, # exponent_bits 8, # scale_num_bits 4, # scale_exponent_bits ) w_q w_q_flat.reshape_as(w) # With FP4 activation fake quantization out_q conv3d_implicit_gemm_cuda( x, w_q, stride(1, 1, 1), padding(1, 1, 1), act_amaxx.abs().max().unsqueeze(0), quant_actTrue, fp4_block_sizeblock_size, # 16, 32, 64, 128, or 256 )其中dynamic_block_quantize_op是 ModelOpt 高层动态块量化 API在 tensor_quant.py 中被导出并注册为torch.library算子tensorrt::dynamic_block_quantize_op参数含义为num_bits4E2M1 共 4 bit、exponent_bits2、scale_num_bits8scale 用 FP8、scale_exponent_bits4E4M3。这里的关键点正如代码注释所强调内核的 A-tile 沿 GEMM K 维K Cin*kD*kH*kW做 block 量化因此权重也必须先展平成[Cout, K]再量化才能与激活的量化维度对齐。7.3fp4_fake_quant导入路径from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import fp4_fake_quant参数说明x输入张量任意形状numel必须能被block_size整除global_amax标量张量——计算 scale 所用的全局绝对值最大值block_size每个 FP4 量化块包含的元素数默认16该函数内部将输入转为 FP32断言numel % block_size 0调用fp4_fake_quant_cuda后恢复原形状与原 dtypeimplicit_gemm_cuda.py。8. 与 ModelOpt 量化体系的自动集成8.1_QuantConv3d与 QuantModuleRegistry_QuantConv3d通过装饰器注册进QuantModuleRegistryQuantModuleRegistry.register({nn.Conv3d: nn.Conv3d})因此当nn.Conv3d被 ModelOpt PTQ 替换为量化模块后NVFP4 配置会自动走这条新路径。完整实现位于 quant_conv.py。8.2 自动分发的判定条件_should_use_implicit_gemm()要求同时满足存在input_quantizer与weight_quantizer两者都是 NVFP4 动态块量化器——_is_nvfp4_quantizer的判定为num_bits (2, 1)、block_sizes[scale_bits] (4, 3)、block_sizes[type] dynamicquant_conv.pygroups 1分组卷积自动回退 cuDNN。8.3 校准期 / 推理期 / 训练期的三态行为_QuantConv3d.forward的分发逻辑quant_conv.py校准期input_quantizer._if_calib and not _if_quant只收集 amax走更快的 cuDNN 路径训练模式self.training打印 warning 并回退 cuDNN 量化路径该内核只支持推理不支持 QAT 反向传播量化推理期调用_implicit_gemm_forward()即act_amax input_quantizer._get_amax(input)后进入融合内核输出还会经过output_quantizer。权重在_nvfp4_quantize_weight_along_k中被展平为[Cout, -1]沿 K 维量化再还原形状quant_conv.py与 README 手动示例中的做法一致。fp4_block_size取自input_quantizer.block_sizes.get(-1, 16)quant_conv.py。8.4 与 PTQ 配置文件的衔接仓库中的 NVFP4 配方示例如 diffusers/nvfp4_fp8_conv.yaml展示了在真实 PTQ 流程中如何配置nn.Linear/nn.Conv2d的量化格式对nn.Conv3d应用 NVFP4 配置后_QuantConv3d即会按上述逻辑自动接管推理路径。更完整的模型级 NVFP4 配置示例可参考 modelopt_recipes/README.md。9. 正确性测试与验证方法9.1 运行测试按 conv/README.mdGPU 环境下执行python -m pytest tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py -v测试文件 test_implicit_gemm.py 约 1500 行覆盖范围非常完整可作为内核行为的权威参考基础正确性对照torch.nn.functional.conv3d最小形状、单/多通道、bias、batch1stride / padding / dilation均匀与非对称配置如stride(1,2,2)、padding(0,1,2)、dilation(1,2,2)核尺寸1×1×1、1×3×3、5×5×5真实形状Wan2.2 形状x[1,128,21,60,106]、w[512,128,3,3,3]、大 Cout/Cin边界情况M / Cout / K 不整除 tile如K189不整除 256、输出 1×1×1确定性多次调用输出逐位相同。9.2 精度容差误差随 √K 增长由于输入在进入 WMMA 前被舍入到 BF16且累加顺序与 cuDNN 的 FP32 路径不同测试采用了随 K 维缩放的容差ATOL 1.0scaled_atol ATOL * (K / 1000.0) ** 0.5并要求平均误差小于scaled_atol * 0.1test_implicit_gemm.py。这与 README 中BF16 舍入误差随 K 维累积最大绝对误差约按 √K 增长的说明相互印证。9.3 FP4 路径的多层交叉验证FP4 融合路径与独立fp4_fake_quant的测试设计尤其严密融合 vs 非融合参考fp4(im2col) fp4(weight)BF16 矩阵乘对比融合内核输出block 大小单调性error(16) error(64) error(256)带 1.2× 松弛验证更细粒度 block 降低量化误差独立fp4_fake_quant对拍精确 E2M1 值、负数、边界上下取整、scale 缩放、零 block、多 block与仓库内其他 FP4 实现一致性Tritonfp4_fake_quant_block、cuda_ext_mx.fused_amax_convert、高层dynamic_block_quantize_op三者逐一allclose比对TestFP4FakeQuantVsModelopt并覆盖 FP32/FP16/BF16 三种输入 dtype。10. 性能基准工具模块内附带延迟基准脚本 bench_implicit_gemm.py对比 cuDNNF.conv3d、非量化 implicit GEMM、FP4 量化 implicit GEMM 三者的中位延迟默认 warmup 20 次、迭代 100 次。其命令行参数参数默认说明--shapesall形状集合small、medium、wan22、stride、all--warmup20预热迭代次数--iters100计时迭代次数--fp4-block-size128FP4 block 大小仅128或256预置的wan22形状组模拟了 Wan2.2 VAE 的真实 Conv3D 配置如wan22_128x512N1, Cin128, D21, H60, W106, Cout512, kernel3×3×3。输出中GEMM/cuDNN与FP4/cuDNN两列比值 1.0 表示慢于 cuDNN 1.0 表示更快。需要注意的是该基准只报告相对比值仓库中并未固化任何胜出结论具体性能取决于 GPU 型号与形状。11. 已知限制与注意事项根据 conv/README.md当前状态为Integrated已注册进QuantModuleRegistryNVFP4 Conv3D 自动分发但存在以下限制首次调用有 JIT 编译延迟约数秒内核通过torch.utils.cpp_extension.load()按需编译verboseTrue会打印编译过程不支持分组卷积groups 1ModelOpt 端到端流程中_QuantConv3d会自动回退 cuDNN 路径数值精度BF16 舍入误差随 K 维累积最大绝对误差约按√K增长相对 cuDNN FP32 路径更小的fp4_block_size能降低量化误差仅推理整个调用链用torch.no_grad()包裹不适用于 QAT 反向传播_QuantConv3d在训练模式下也会回退FP4 只量化激活权重需调用方预先量化README 示例或_QuantConv3d内部都会沿 K 维完成。12. 文件索引路径作用modelopt/torch/kernels/quantization/conv/README.md本内核的官方文档本文的主体依据modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.pyPython API、JIT 编译、形状/参数校验modelopt/torch/kernels/quantization/conv/implicit_gemm_kernel.cuCUDA 内核BF16 WMMA FP4 量化 L2 swizzlemodelopt/torch/kernels/quantization/conv/implicit_gemm_binding.cppPyTorch C 扩展绑定与内核分发modelopt/torch/kernels/quantization/conv/bench_implicit_gemm.pycuDNN vs GEMM vs GEMMFP4 延迟基准modelopt/torch/quantization/nn/modules/quant_conv.py_QuantConv3d集成NVFP4 自动分发与三态行为modelopt/torch/quantization/tensor_quant.pydynamic_block_quantize_op高层量化算子tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py完整测试正确性、容差、FP4 交叉验证、确定性modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yamlNVFP4/FP8 PTQ 配方示例赞分享人工智能大模型模型优化模型量化模型压缩【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址https://gitcode.com/GitHub_Trending/te/Model-Optimizer点击查看免费下载相关推荐Model-Optimizer 模型内核微基准指南在 meta tensor 上规划并运行 BF16/FP8/NVFP4 GEMM 与 fused-MoE 基准Model Optimizer 模型内核微基准指南在 meta tensor 上规划并运行 BF16/FP8/NVFP4 GEMM 与 fused MoE 基人工智能大模型模型优化模型量化模型压缩CUDA Samples 之 cudaTensorCoreGemm 深度解析基于 WMMA API 与 Tensor Core 的 GEMM 实现CUDA Samples 之 cudaTensorCoreGemm 深度解析基于 WMMA API 与 Tensor Core 的 GEMM 实现 本篇文章以示例工程cuda-samples 之 tf32TensorCoreGemm基于 WMMA 与异步拷贝的 TF32 Tensor Core GEMM 深度解析cuda samples 之 tf32TensorCoreGemm基于 WMMA 与异步拷贝的 TF32 Tensor Core GEMM 深度解析 导读 本示例工程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
