【Bug已解决】[Feature Request][CUDA] QMoE: support block-wise quantized expert weights (block_size ＞ 0) …

【Bug已解决】[Feature Request][CUDA] QMoE: support block-wise quantized expert weights (block_size 0) 解决方案一、现象长什么样用 ONNX Runtime 的 QMoE量化 MoE在 CUDA 上跑专家权重被block-wise 量化block_size 0比如每 128 个权重一组一个 scale的模型。要么加载失败要么跑出错误结果——当前 CUDA QMoE 只支持per-tensor / per-channel量化不支持 block-wiseimport onnxruntime as ort # 专家权重是 block-wise 量化block_size128每 128 个元素一个 scale sess ort.InferenceSession(qmoe_blockwise.onnx, providers[CUDAExecutionProvider]) # 报错或不支持CUDA QMoE 未实现 block_size 0 的反量化路径最小信号block_size 0per-channel/per-tensor- CUDA QMoE 正常 block_size 0block-wise - CUDA QMoE 不支持/结果错注意这是功能缺口feature request不是崩溃。CPU 端可能已支持 block-wise但 CUDA 端没实现对应反量化内核。二、背景MoE 专家权重量化时scale 的粒度决定了精度与速度的平衡per-tensor整个权重一个 scale最快但精度最低。per-channelper-output-row每个输出通道一个 scale精度较好。block-wiseblock_size N如 128每 N 个连续权重共享一个 scale。N 越小精度越高每组动态范围更一致N128 是常见甜点。block-wise 在 GPTQ/AWQ 等权重量化里非常普遍。block-wise 的反量化要求内核按块取 scale权重被切成长度为block_size的块每块有一个 scale和可能的 zero_point反量化时W_deq[i] W_q[i] * scale[block_index(i)] zp[block_index(i)]。CUDA QMoE 当前的反量化内核只实现了“全局/按行 scale”的索引方式scale 形状是[num_experts, inter]之类按输出通道取没有实现“按块偏移取 scale”——即scale的形状变成了[num_experts, inter * hidden / block_size]内核需要用i / block_size作为 scale 索引而旧内核用的是i / hidden_size行索引。于是 block-wise 模型要么不被识别、要么 scale 取错 - 结果错。三、根因根因是CUDA QMoE 反量化内核只支持按行per-channel取 scale没有实现按块block-wise,block_size0取 scale 的索引逻辑scale 索引方式不对per-channel 时 scale 索引 rowblock-wise 时 scale 索引 i / block_size。CUDA 内核写死了按row取遇到 block-wise 的[num_experts, inter*hidden/block_size]形状 scale 时索引越界或错位反量化用错 scale。block_size 参数未透传内核没接收/使用block_size属性默认按 0per-channel处理。只影响 block-wiseper-tensor/per-channel 路径照常工作所以问题只在block_size0时暴露。CPU 端若已实现 block-wise则更凸显 CUDA 端缺口。所以这不是数值错在支持的粒度下而是CUDA QMoE 缺 block-wise 反量化的实现导致该粒度模型不可用/出错。四、最小可运行复现下面用 NumPy 模拟“block-wise 反量化scale 索引按块而非按行”的差异import numpy as np def dequant_per_channel(wq, scale): 旧 CUDA 内核按行输出通道取 scale。 # scale 形状 [inter,] 每输出行一个 return (wq.astype(np.float32)) * scale.reshape(-1, 1) def dequant_blockwise(wq, scale, block_size128): 正确 block-wise按块取 scale。 out np.zeros_like(wq, dtypenp.float32) for i in range(0, wq.shape[1], block_size): blk slice(i, i block_size) s scale[:, i // block_size].reshape(-1, 1) out[:, blk] wq[:, blk].astype(np.float32) * s return out if __name__ __main__: inter, hidden 4, 256 block_size 128 wq np.random.randint(-8, 8, size(inter, hidden), dtypenp.int8) # block-wise scale 形状[inter, hidden//block_size] scale np.random.rand(inter, hidden // block_size).astype(np.float32) 0.5 correct dequant_blockwise(wq, scale, block_size) # 旧内核如果用 per-channel scale 形状会直接报错形状不匹配这里演示索引不同 print(block-wise 反量化输出形状:, correct.shape) # 验证第 0 块用 scale[:,0]第 1 块用 scale[:,1] assert np.allclose(correct[:, :block_size], wq[:, :block_size].astype(np.float32) * scale[:, 0].reshape(-1, 1))跑出来block-wise 反量化按i // block_size取 scale每块用各自 scale旧内核按行取会形状不匹配或取错。这复现了“block-wise 需要按块索引 scale”的机制。五、解决方案第一层最小直接修复最小修复给 CUDA QMoE 反量化内核加上 block-wise 路径按i/block_size取 scale透传block_size参数。对使用者临时规避是把模型重导出为 per-channel/per-tensor 量化牺牲一点 block-wise 精度换 CUDA 支持# 导出时把 QMoE 专家权重的量化粒度从 block_size128 改成 per-channel # 用量化工具把 block-wise scale 上采样/合并成 per-channel scale # 或在推理时强制 CPU EP 跑若 CPU 已支持 block-wise import onnxruntime as ort sess ort.InferenceSession(qmoe_blockwise.onnx, providers[CPUExecutionProvider])对 ORT 仓库侧修复是改 CUDA QMoE 内核dequant时若block_size 0用element_idx / block_size作为 scale 索引并正确加载[num_experts, inter, hidden/block_size]形状的 scale 张量。这一层立刻让 block-wise 模型在 CUDA 上可用且正确。六、解决方案第二层结构性改进把“QMoE 支持的量化粒度”收口成唯一的配置对象OrtQmoeBlockwiseCudaPolicy加载与内核选择读它from dataclasses import dataclass, field from typing import Tuple, Literal dataclass(frozenTrue) class OrtQmoeBlockwiseCudaPolicy: CUDA QMoE 量化粒度支持的单一事实来源。 # 已支持的量化粒度 supported_granularities: Tuple[str, ...] (per_tensor, per_channel, block_wise) # block-wise 的默认块大小 block_size: int 128 # 各粒度对应的 scale 索引方式 scale_index_mode: Tuple[str, ...] (scalar, row, block) # 受影响 EP ep: str CUDAExecutionProvider def scale_index(self, granularity: str) - str: m {per_tensor: scalar, per_channel: row, block_wise: block} return m[granularity] def describe(self) - str: return CUDA QMoE 支持 per_tensor/per_channel/block_wise按块索引 scale POLICY OrtQmoeBlockwiseCudaPolicy() def plan_qmoe(granularity: str, policy: OrtQmoeBlockwiseCudaPolicy POLICY) - str: assert granularity in policy.supported_granularities return policy.scale_index(granularity)所有 QMoE 加载与内核选择读同一份POLICYblock-wise 成为一等公民新增粒度必须实现对应 scale 索引。七、解决方案第三层断言 / CI 守护把“block-wise 在 CUDA 上正确反量化”做成断言。下面用 pytest 风格守护复用第四节逻辑import numpy as np def test_blockwise_supported(policy): assert block_wise in policy.supported_granularities def test_block_scale_index(policy): assert policy.scale_index(block_wise) block def test_dequant_blockwise_correct(): inter, hidden, bs 4, 256, 128 wq np.random.randint(-8, 8, size(inter, hidden), dtypenp.int8) scale np.random.rand(inter, hidden // bs).astype(np.float32) 0.5 out dequant_blockwise(wq, scale, bs) assert np.allclose(out[:, :bs], wq[:, :bs].astype(np.float32) * scale[:, 0].reshape(-1, 1)) def test_block_size_positive(policy): assert policy.block_size 0这四组断言锁住(1) block-wise 已支持(2) scale 按块索引(3) block-wise 反量化数值正确(4) block_size 为正。CI 跑通即代表 CUDA QMoE block-wise 被正确支持。八、排查清单遇到 CUDA QMoE block-wise 模型不支持/结果错确认量化粒度block_size 0即 block-wiseper-channel 正常 - 锁定 block 路径缺失。看 scale 索引CUDA 内核是不是按行取 scale没按i/block_size取。查 block_size 是否透传内核有没有接收block_size参数。临时规避重导出为 per-channel或暂用 CPU EP若已支持 block-wise。根本修复CUDA 内核加 block-wise 路径按块索引 scale。统一策略对象用OrtQmoeBlockwiseCudaPolicy固化粒度支持。CI 守护断言 block-wise 支持、scale 按块索引、数值正确。九、小结[Feature Request][CUDA] QMoE: support block-wise quantized expert weights (block_size 0)的根因是CUDA QMoE 反量化内核只实现了 per-tensor / per-channel按行取 scale的路径没有实现 block-wiseblock_size0按i/block_size取 scale的索引逻辑block_size参数也未透传导致 block-wise 量化如 GPTQ/AWQ 常见的块量化的模型在 CUDA 上不可用或结果错。最小修复是给 CUDA QMoE 内核加 block-wise 反量化路径按块索引 scale、透传block_size临时规避是重导出为 per-channel 或暂用 CPU EP结构性改进是用唯一的OrtQmoeBlockwiseCudaPolicy固化粒度支持CI 用四组断言守护“block-wise 支持、scale 按块索引、数值正确、block_size 为正”。记住block-wise 量化的精髓是按块取 scaleCUDA 内核必须实现i/block_size索引否则精度与可用性都丢。