【Bug已解决】[CUDA] Qwen3.6-35B-A3B Throughput Optimization 解决方案

【Bug已解决】[CUDA] Qwen3.6-35B-A3B Throughput Optimization 解决方案一、现象长什么样把 Qwen3.6-35B-A3B一个 35B 参数、每次激活约 3B 的 MoE 模型导出成 ONNX在 ONNX Runtime 的 CUDA EP 上做推理吞吐明显低于预期比如同样一张 H100对比厂商优化后的参考实现只有 40%~60% 的吞吐import onnxruntime as ort so ort.SessionOptions() so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess ort.InferenceSession(qwen3.6-35b-a3b.onnx, so, providers[CUDAExecutionProvider]) # 默认配置下MoE 专家计算与 GQA 注意力没充分融合吞吐偏低最小信号吞吐远低于参考实现 GPU 利用率波动大kernel 启动密集 专家路由 各 expert 计算没融成高效 kernel注意结果正确只是慢。这是针对 Qwen3.6-35B-A3B 这个具体 MoE 结构的吞吐优化问题。二、背景Qwen3.35B-A3B 是 MoE 结构总参数 35B但每层只有约 3B 参数被激活8 个 expert 里选 2~3 个。它还有几个对 CUDA 吞吐极关键的特征MoE 路由 分组专家计算token 按 router 分数被分发到不同 expert每个 expert 是一次大矩阵乘。若没融合路由、分组、各 expert 计算被拆成几十上百个小 kernellaunch 开销爆炸。GQA分组查询注意力KV 头远少于 Q 头。需要GroupQueryAttention融合 kernel且正确传num_kv_heads。可能的长上下文 / 高并发batch 大时kernel 是否能针对常见形状特化CUDA graph很关键。ONNX Runtime 的 CUDA EP 要让这个模型跑满吞吐需要(a) MoE 融合把路由专家计算合成一个高效 kernel 或紧密调度的序列(b) GQA 融合(c) CUDA graph 捕获固定结构(d) 动态维度覆盖让 kernel 特化。默认配置下这些没全开于是吞吐掉一截。三、根因根因是CUDA EP 的关键优化对 Qwen3.6-35B-A3B 没全激活且模型导出形态不利于融合MoE 未充分融合导出时 router 各 expert 是标准MatMul/Gather/ConcatORT 的 MoE 融合 pass 没匹配上大量小 kernel 串行。GQA 融合属性/结构不匹配GroupQueryAttention融合要求num_kv_heads/head_size齐全且导出保留Attention节点若拆成裸MatMul则无法融合。CUDA graph 没开默认enable_cuda_graphfalse每轮推理重新录制命令launch 开销大。动态维度未覆盖没给batch/seq设自由维度边界kernel 无法特化。不是结果错融合与图捕获没开导致 kernel 碎片化、GPU 利用率低、吞吐低。所以这不是数值错而是融合与图捕获未激活MoE 大模型吞吐被 launch 开销拖垮。四、最小可运行复现下面用 Python 模拟“MoE 融合与否对 kernel 启动次数的影响”Qwen3 风格每层 8 expert 选 2import numpy as np def run_unfused(num_tokens, num_experts8, top_k2): 未融合每个 token 被选中的 expert 各一次 MatMul极多 launch。 launches 0 for _ in range(num_tokens): for _ in range(top_k): launches 1 return launches def run_fused(num_tokens, num_experts8, top_k2): 融合整批一次 fused MoE kernel。 return 1 if __name__ __main__: for n in (512, 2048, 8192): unfused run_unfused(n) fused run_fused(n) print(ftokens{n}: 未融合 launch{unfused}, 融合 launch{fused}, f差距≈{unfused/fused:.0f}x)跑出来tokens8192时未融合 ~16000 次 launch、融合 1 次上万倍 launch 差距——这正是 MoE 不融合时吞吐崩塌的简化模型实际加速没这么夸张但量级说明问题。五、解决方案第一层最小直接修复最小修复打开 CUDA 优化开关并让导出形态可被融合识别针对 Qwen3.6-35B-A3Bimport onnxruntime as ort so ort.SessionOptions() so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL so.add_free_dimension_override_by_name(batch_size, 1, 128) so.add_free_dimension_override_by_name(seq_len, 1, 32768) cuda_opts { device_id: 0, enable_cuda_graph: True, use_tf32: True, cuda_graph_enable_partial: True, max_batch_size: 128, } provider (CUDAExecutionProvider, cuda_opts) sess ort.InferenceSession(qwen3.6-35b-a3b.onnx, so, providers[provider])导出时保留 MoE 子图与Attention/GroupQueryAttention节点用支持 MoE 导出的工具让 ORT 的 MoE 融合 pass 能匹配。这一层立刻把吞吐抬上去。六、解决方案第二层结构性改进把“Qwen3.6-35B-A3B 在 CUDA 上的优化配置”收口成唯一的配置对象OrtQwenThroughputPolicy部署读它from dataclasses import dataclass, field from typing import Dict, Tuple dataclass(frozenTrue) class OrtQwenThroughputPolicy: Qwen3.6-35B-A3B 在 CUDA EP 上的吞吐优化单一事实来源。 optimization_level: str ORT_ENABLE_ALL enable_cuda_graph: bool True use_tf32: bool True cuda_graph_partial: bool True free_dim_overrides: Tuple[Tuple[str, int, int], ...] ( (batch_size, 1, 128), (seq_len, 1, 32768), ) # MoE 融合要求保留的子图 keep_moe_subgraph: bool True keep_gqa_node: bool True # 模型结构提示用于诊断 model_kind: str qwen3_moe_35b_a3b def cuda_provider_options(self) - Dict: return { enable_cuda_graph: self.enable_cuda_graph, use_tf32: self.use_tf32, cuda_graph_enable_partial: self.cuda_graph_partial, } def describe(self) - str: return 融合 MoEGQA、开 CUDA graph、覆盖动态维度以特化 Qwen3 MoE kernel POLICY OrtQwenThroughputPolicy() def build_options(policy: OrtQwenThroughputPolicy POLICY) - dict: return { opt: policy.optimization_level, cuda: policy.cuda_provider_options(), free: policy.free_dim_overrides, }所有部署读同一份POLICY融合与图捕获配置固化避免“忘了开某个开关又变慢”。七、解决方案第三层断言 / CI 守护把“Qwen3 MoE 吞吐优化开关生效”做成断言。下面用 pytest 风格守护import pytest def test_cuda_graph_on(policy): assert policy.cuda_provider_options()[enable_cuda_graph] is True def test_free_dims_covered(policy): names [d[0] for d in policy.free_dim_overrides] assert batch_size in names and seq_len in names def test_moe_gqa_kept(policy): assert policy.keep_moe_subgraph is True assert policy.keep_gqa_node is True def test_opt_level_all(policy): assert policy.optimization_level ORT_ENABLE_ALL这四组断言锁住(1) CUDA graph 开(2) 动态维度覆盖(3) MoE/GQA 子图保留(4) 优化等级为 ALL。CI 跑通即代表吞吐优化路径激活。八、排查清单遇到 Qwen3 MoE 在 CUDA 上吞吐低看融合是否生效session 里有没有FusedMatMul/GroupQueryAttention/MoE节点。开优化等级ORT_ENABLE_ALL别留ORT_DISABLE_ALL。开 CUDA graphenable_cuda_graphtrue。覆盖动态维度给batch/seq设边界帮助 kernel 特化。检查导出形态MoE 子图、GQA 节点有没有被展开成裸算子。统一策略对象用OrtQwenThroughputPolicy固化。CI 守护断言关键开关开启、融合子图保留。九、小结[CUDA] Qwen3.6-35B-A3B Throughput Optimization的根因是CUDA EP 的关键优化MoE 融合、GroupQueryAttention融合、CUDA graph、动态维度特化默认没激活且模型导出形态可能把 MoE/GQA 拆成无法被融合识别的裸算子导致 kernel 碎片化、launch 密集、GPU 利用率低、吞吐只有预期的零头。最小修复是打开ORT_ENABLE_ALL、启用 CUDA graph 与 TF32、覆盖动态维度并保证导出时保留 MoE/GQA 子图结构性改进是用唯一的OrtQwenThroughputPolicy固化配置CI 用四组断言守护“融合开关生效、维度覆盖、优化等级为 ALL”。记住ORT 跑 MoE 大模型融合和图捕获要显式打开否则就是一堆小 kernel 在空转——这和 GPT-OSS 这类 MoE 的优化思路一致但 Qwen3 的 A3B 激活比更稀疏融合收益更大。