【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 模型，结构上有两类对吞吐极其关键的算子：

1. MoE（混合专家）：每层有若干 expert FFN + 一个 router，token 按 router 分数被分发到不同 expert。若没有 MoE 融合 kernel，ORT 会把“路由 + 分组 + 各 expert 计算 + 合并”拆成一长串小 kernel，launch 开销爆炸。
2. 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；
• 对 MoE：ORT 的QMoE/FusedMatMul以及contrib里的 MoE 融合节点（需要模型导出时保留可被识别的 MoE 子图模式）；
• 动态维度覆盖：给batch_size/seq_len设自由维度上下界，让 CUDA graph / kernel 特化。

默认配置下，融合没完全生效，于是吞吐掉一半以上。

三、根因

根因是CUDA EP 的几个关键优化开关没被激活，且模型导出形态不利于融合：

1. MoE 未融合：模型导出时把 router + expert 拆成了标准MatMul/Gather/Concat，ORT 的 MoE 融合（识别TopK+ 按专家分组的MatMul链）没匹配上，于是一堆小 kernel 串行，launch 瓶颈明显。
2. GQA 融合属性不匹配：GroupQueryAttention融合要求num_kv_heads/num_heads/head_size等属性齐全；导出时若把 QKV 拆成三个独立MatMul而没有Attention/GroupQueryAttention节点，ORT 没法融合成单一 fused kernel。
3. CUDA graph 没开：默认enable_cuda_graph=false，每次推理都重新录制并提交命令，launch 开销大；而 GPT-OSS-20B 这种固定结构非常适合 CUDA graph 整图重放。
4. 动态维度未覆盖：没给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 kernel（launch 数恒定）。""" return 1 # 一个融合 kernel 处理全部 if __name__ == "__main__": for n in (128, 512, 2048): unfused = run_unfused(n) fused = run_fused(n) # 用 launch 数近似“开销”：launch 越多越慢 print(f"tokens={n}: 未融合 launch={unfused}, 融合 launch={fused}, " f"加速比≈{unfused/fused:.0f}x")

跑出来：tokens=2048时未融合 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(frozen=True) 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 大模型吞吐低：

1. 先确认融合是否生效：用session.get_providers()和 ORT 的图优化日志看有没有FusedMatMul/GroupQueryAttention/MoE节点。
2. 打开优化等级：graph_optimization_level=ORT_ENABLE_ALL，别留ORT_DISABLE_ALL。
3. 开 CUDA graph：enable_cuda_graph=true，适合固定结构。
4. 覆盖动态维度：给batch/seq设自由维度上下界，帮助 kernel 特化。
5. 检查导出形态：MoE 子图、GQA 节点有没有被提前展开成裸MatMul，导致无法融合。
6. 统一策略对象：用OrtGptOssThroughputPolicy固化。
7. 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 在空转。