【Bug已解决】[CUDA] GPT-OSS-20B Throughput Optimization 解决方案一、现象长什么样把开源 MoE 模型GPT-OSS-20B导出成 ONNX在 ONNX Runtime 的 CUDA EP 上做批推理吞吐量远低于预期比如同样一张 A100对比直接用 PyTorch FlashAttention 只有 30%~50% 的吞吐实测吞吐~180 tokens/s/batch预期 ~500 tokens/s/batch GPU 利用率忽高忽低kernel 启动次数极多ncu 看 launch 密集最小触发import onnxruntime as ort so ort.SessionOptions() so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess ort.InferenceSession(gpt-oss-20b.onnx, so, providers[CUDAExecutionProvider]) # 默认 CUDA 选项下MoE 专家计算与 GQA 注意力没有被充分融合 for batch in dataloader: sess.run(None, batch)注意模型能跑、结果正确只是太慢。这是性能优化类问题根因在 CUDA EP 对该 MoE GQA 结构的融合/调度没激活最优路径。二、背景GPT-OSS-20B 是 ~20B 参数的 MoE 模型结构上有两类对吞吐极其关键的算子MoE混合专家每层有若干 expert FFN 一个 routertoken 按 router 分数被分发到不同 expert。若没有 MoE 融合 kernelORT 会把“路由 分组 各 expert 计算 合并”拆成一长串小 kernellaunch 开销爆炸。GQA分组查询注意力query 头多、KV 头少。ORT 有专门的GroupQueryAttention融合算子但需要正确开启并匹配属性尤其attention_bias、num_kv_heads。ONNX Runtime 的 CUDA EP 要通过若干开关把这些结构“识别并融合”graph_optimization_level ORT_ENABLE_ALL否则融合 pass 不跑CUDA provider 选项里的enable_cuda_graph、use_tf32、cuda_graph_enable_partial对 MoEORT 的QMoE/FusedMatMul以及contrib里的 MoE 融合节点需要模型导出时保留可被识别的 MoE 子图模式动态维度覆盖给batch_size/seq_len设自由维度上下界让 CUDA graph / kernel 特化。默认配置下融合没完全生效于是吞吐掉一半以上。三、根因根因是CUDA EP 的几个关键优化开关没被激活且模型导出形态不利于融合MoE 未融合模型导出时把 router expert 拆成了标准MatMul/Gather/ConcatORT 的 MoE 融合识别TopK 按专家分组的MatMul链没匹配上于是一堆小 kernel 串行launch 瓶颈明显。GQA 融合属性不匹配GroupQueryAttention融合要求num_kv_heads/num_heads/head_size等属性齐全导出时若把 QKV 拆成三个独立MatMul而没有Attention/GroupQueryAttention节点ORT 没法融合成单一 fused kernel。CUDA graph 没开默认enable_cuda_graphfalse每次推理都重新录制并提交命令launch 开销大而 GPT-OSS-20B 这种固定结构非常适合 CUDA graph 整图重放。动态维度未覆盖没给batch/seq设自由维度边界kernel 无法针对常见形状特化反复走通用路径。所以不是算错而是融合与图捕获没打开导致 kernel 碎片化、launch 密集、GPU 利用率低。四、最小可运行复现下面用 Python 模拟“融合与否对 kernel 启动次数的影响”不依赖真实 GPU但精准复现吞吐模型import time def run_unfused(num_tokens: int, num_experts: int 8): 未融合每个 token × 每个专家一次 MatMul极多 launch。 launches 0 for _ in range(num_tokens): for _ in range(num_experts): launches 1 # 一次 kernel launch return launches def run_fused(num_tokens: int, num_experts: int 8): 融合整批一次 fused MoE kernellaunch 数恒定。 return 1 # 一个融合 kernel 处理全部 if __name__ __main__: for n in (128, 512, 2048): unfused run_unfused(n) fused run_fused(n) # 用 launch 数近似“开销”launch 越多越慢 print(ftokens{n}: 未融合 launch{unfused}, 融合 launch{fused}, f加速比≈{unfused/fused:.0f}x)跑出来tokens2048时未融合 16384 次 launch、融合 1 次约 16000x 的 launch 差距——这正是 MoE 不融合时吞吐崩塌的简化模型。实际加速没这么夸张融合 kernel 内部也有成本但差距量级说明问题。五、解决方案第一层最小直接修复最小修复打开所有该开的 CUDA 优化开关并让模型导出形态可被融合识别。import onnxruntime as ort so ort.SessionOptions() so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL # 给动态维度设边界帮助 kernel 特化与 CUDA graph so.add_free_dimension_override_by_name(batch_size, 1, 64) so.add_free_dimension_override_by_name(seq_len, 1, 8192) # CUDA provider 选项 cuda_opts { device_id: 0, enable_cuda_graph: True, use_tf32: True, cuda_graph_enable_partial: True, max_batch_size: 64, } so.append_session_config_entry(gpu_graph_capture, 1) provider (CUDAExecutionProvider, cuda_opts) sess ort.InferenceSession(gpt-oss-20b.onnx, so, providers[provider])同时导出时保留 MoE 结构用optimum的 MoE 导出或确保TopK分组MatMul模式不被提前展开让 ORT 的 MoE 融合 pass 能匹配。这一层立刻把吞吐抬上去。六、解决方案第二层结构性改进把“GPT-OSS 类 MoE 模型在 CUDA 上的优化配置”收口成唯一的配置对象OrtGptOssThroughputPolicy所有部署读它from dataclasses import dataclass, field from typing import Dict, Tuple dataclass(frozenTrue) class OrtGptOssThroughputPolicy: GPT-OSS-20B 在 CUDA EP 上的吞吐优化单一事实来源。 optimization_level: str ORT_ENABLE_ALL enable_cuda_graph: bool True use_tf32: bool True cuda_graph_partial: bool True # 动态维度边界让 kernel 针对常见形状特化 free_dim_overrides: Tuple[Tuple[str, int, int], ...] ( (batch_size, 1, 64), (seq_len, 1, 8192), ) # MoE 融合要求保留的子图模式导出时不能展开 keep_moe_subgraph: bool True # GQA 融合要求导出时保留 Attention/GroupQueryAttention 节点 keep_gqa_node: bool True # 是否启用 QMoE 量化专家进一步提速 qmoe_block_quant: bool False 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 融合 MoE GQA、开 CUDA graph、覆盖动态维度以特化 kernel POLICY OrtGptOssThroughputPolicy() def build_session_options(policy: OrtGptOssThroughputPolicy POLICY) - dict: return { opt_level: policy.optimization_level, cuda_opts: policy.cuda_provider_options(), free_dims: policy.free_dim_overrides, }所有部署读同一份POLICY融合与图捕获配置不再散落各处避免“忘了开某个开关又变慢”。七、解决方案第三层断言 / CI 守护把“优化开关生效、融合被识别”做成断言。下面用 pytest 风格守护import pytest def test_cuda_graph_enabled(policy): opts policy.cuda_provider_options() assert opts[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_and_gqa_kept(policy): assert policy.keep_moe_subgraph is True assert policy.keep_gqa_node is True def test_optimization_is_all(policy): assert policy.optimization_level ORT_ENABLE_ALL这四组断言锁住(1) CUDA graph 开(2) 动态维度覆盖(3) MoE/GQA 子图保留可被融合(4) 优化等级为 ALL。CI 跑通即代表吞吐优化路径已激活。八、排查清单遇到 ORT CUDA 上 MoE 大模型吞吐低先确认融合是否生效用session.get_providers()和 ORT 的图优化日志看有没有FusedMatMul/GroupQueryAttention/MoE节点。打开优化等级graph_optimization_levelORT_ENABLE_ALL别留ORT_DISABLE_ALL。开 CUDA graphenable_cuda_graphtrue适合固定结构。覆盖动态维度给batch/seq设自由维度上下界帮助 kernel 特化。检查导出形态MoE 子图、GQA 节点有没有被提前展开成裸MatMul导致无法融合。统一策略对象用OrtGptOssThroughputPolicy固化。CI 守护断言关键开关开启、融合子图保留。九、小结[CUDA] GPT-OSS-20B Throughput Optimization的根因是 CUDA EP 的关键优化开关MoE 融合、GroupQueryAttention融合、CUDA graph、动态维度特化默认没激活且模型导出形态可能把 MoE/GQA 拆成无法被融合识别的裸算子导致 kernel 碎片化、launch 密集、GPU 利用率低、吞吐只有预期的零头。最小修复是打开ORT_ENABLE_ALL、启用 CUDA graph 与 TF32、覆盖动态维度并保证导出时保留 MoE/GQA 子图结构性改进是用唯一的OrtGptOssThroughputPolicy固化配置CI 用四组断言守护“融合开关生效、维度覆盖、优化等级为 ALL”。记住ORT 跑大模型融合和图捕获要显式打开否则就是一堆小 kernel 在空转。