【Bug已解決】[CUDA] Qwen3.6-35B-A3B Throughput Optimization 解決方案一、現(xiàn)象長什么樣把 Qwen3.6-35B-A3B一個(gè) 35B 參數(shù)、每次激活約 3B 的 MoE 模型導(dǎo)出成 ONNX在 ONNX Runtime 的 CUDA EP 上做推理吞吐明顯低于預(yù)期比如同樣一張 H100對(duì)比廠商優(yōu)化后的參考實(shí)現(xiàn)只有 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]) # 默認(rèn)配置下MoE 專家計(jì)算與 GQA 注意力沒充分融合吞吐偏低最小信號(hào)吞吐遠(yuǎn)低于參考實(shí)現(xiàn) GPU 利用率波動(dòng)大kernel 啟動(dòng)密集 專家路由 各 expert 計(jì)算沒融成高效 kernel注意結(jié)果正確只是慢。這是針對(duì) Qwen3.6-35B-A3B 這個(gè)具體 MoE 結(jié)構(gòu)的吞吐優(yōu)化問題。二、背景Qwen3.35B-A3B 是 MoE 結(jié)構(gòu)總參數(shù) 35B但每層只有約 3B 參數(shù)被激活8 個(gè) expert 里選 2~3 個(gè)。它還有幾個(gè)對(duì) CUDA 吞吐極關(guān)鍵的特征MoE 路由 分組專家計(jì)算token 按 router 分?jǐn)?shù)被分發(fā)到不同 expert每個(gè) expert 是一次大矩陣乘。若沒融合路由、分組、各 expert 計(jì)算被拆成幾十上百個(gè)小 kernellaunch 開銷爆炸。GQA分組查詢注意力KV 頭遠(yuǎn)少于 Q 頭。需要GroupQueryAttention融合 kernel且正確傳num_kv_heads。可能的長上下文 / 高并發(fā)batch 大時(shí)kernel 是否能針對(duì)常見形狀特化CUDA graph很關(guān)鍵。ONNX Runtime 的 CUDA EP 要讓這個(gè)模型跑滿吞吐需要(a) MoE 融合把路由專家計(jì)算合成一個(gè)高效 kernel 或緊密調(diào)度的序列(b) GQA 融合(c) CUDA graph 捕獲固定結(jié)構(gòu)(d) 動(dòng)態(tài)維度覆蓋讓 kernel 特化。默認(rèn)配置下這些沒全開于是吞吐掉一截。三、根因根因是CUDA EP 的關(guān)鍵優(yōu)化對(duì) Qwen3.6-35B-A3B 沒全激活且模型導(dǎo)出形態(tài)不利于融合MoE 未充分融合導(dǎo)出時(shí) router 各 expert 是標(biāo)準(zhǔn)MatMul/Gather/ConcatORT 的 MoE 融合 pass 沒匹配上大量小 kernel 串行。GQA 融合屬性/結(jié)構(gòu)不匹配GroupQueryAttention融合要求num_kv_heads/head_size齊全且導(dǎo)出保留Attention節(jié)點(diǎn)若拆成裸MatMul則無法融合。CUDA graph 沒開默認(rèn)enable_cuda_graphfalse每輪推理重新錄制命令launch 開銷大。動(dòng)態(tài)維度未覆蓋沒給batch/seq設(shè)自由維度邊界kernel 無法特化。不是結(jié)果錯(cuò)融合與圖捕獲沒開導(dǎo)致 kernel 碎片化、GPU 利用率低、吞吐低。所以這不是數(shù)值錯(cuò)而是融合與圖捕獲未激活MoE 大模型吞吐被 launch 開銷拖垮。四、最小可運(yùn)行復(fù)現(xiàn)下面用 Python 模擬“MoE 融合與否對(duì) kernel 啟動(dòng)次數(shù)的影響”Qwen3 風(fēng)格每層 8 expert 選 2import numpy as np def run_unfused(num_tokens, num_experts8, top_k2): 未融合每個(gè) 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時(shí)未融合 ~16000 次 launch、融合 1 次上萬倍 launch 差距——這正是 MoE 不融合時(shí)吞吐崩塌的簡化模型實(shí)際加速?zèng)]這么夸張但量級(jí)說明問題。五、解決方案第一層最小直接修復(fù)最小修復(fù)打開 CUDA 優(yōu)化開關(guān)并讓導(dǎo)出形態(tài)可被融合識(shí)別針對(duì) 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])導(dǎo)出時(shí)保留 MoE 子圖與Attention/GroupQueryAttention節(jié)點(diǎn)用支持 MoE 導(dǎo)出的工具讓 ORT 的 MoE 融合 pass 能匹配。這一層立刻把吞吐抬上去。六、解決方案第二層結(jié)構(gòu)性改進(jìn)把“Qwen3.6-35B-A3B 在 CUDA 上的優(yōu)化配置”收口成唯一的配置對(duì)象OrtQwenThroughputPolicy部署讀它from dataclasses import dataclass, field from typing import Dict, Tuple dataclass(frozenTrue) class OrtQwenThroughputPolicy: Qwen3.6-35B-A3B 在 CUDA EP 上的吞吐優(yōu)化單一事實(shí)來源。 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 # 模型結(jié)構(gòu)提示用于診斷 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、覆蓋動(dòng)態(tài)維度以特化 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融合與圖捕獲配置固化避免“忘了開某個(gè)開關(guān)又變慢”。七、解決方案第三層斷言 / CI 守護(hù)把“Qwen3 MoE 吞吐優(yōu)化開關(guān)生效”做成斷言。下面用 pytest 風(fēng)格守護(hù)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) 動(dòng)態(tài)維度覆蓋(3) MoE/GQA 子圖保留(4) 優(yōu)化等級(jí)為 ALL。CI 跑通即代表吞吐優(yōu)化路徑激活。八、排查清單遇到 Qwen3 MoE 在 CUDA 上吞吐低看融合是否生效session 里有沒有FusedMatMul/GroupQueryAttention/MoE節(jié)點(diǎn)。開優(yōu)化等級(jí)ORT_ENABLE_ALL別留ORT_DISABLE_ALL。開 CUDA graphenable_cuda_graphtrue。覆蓋動(dòng)態(tài)維度給batch/seq設(shè)邊界幫助 kernel 特化。檢查導(dǎo)出形態(tài)MoE 子圖、GQA 節(jié)點(diǎn)有沒有被展開成裸算子。統(tǒng)一策略對(duì)象用OrtQwenThroughputPolicy固化。CI 守護(hù)斷言關(guān)鍵開關(guān)開啟、融合子圖保留。九、小結(jié)[CUDA] Qwen3.6-35B-A3B Throughput Optimization的根因是CUDA EP 的關(guān)鍵優(yōu)化MoE 融合、GroupQueryAttention融合、CUDA graph、動(dòng)態(tài)維度特化默認(rèn)沒激活且模型導(dǎo)出形態(tài)可能把 MoE/GQA 拆成無法被融合識(shí)別的裸算子導(dǎo)致 kernel 碎片化、launch 密集、GPU 利用率低、吞吐只有預(yù)期的零頭。最小修復(fù)是打開ORT_ENABLE_ALL、啟用 CUDA graph 與 TF32、覆蓋動(dòng)態(tài)維度并保證導(dǎo)出時(shí)保留 MoE/GQA 子圖結(jié)構(gòu)性改進(jìn)是用唯一的OrtQwenThroughputPolicy固化配置CI 用四組斷言守護(hù)“融合開關(guān)生效、維度覆蓋、優(yōu)化等級(jí)為 ALL”。記住ORT 跑 MoE 大模型融合和圖捕獲要顯式打開否則就是一堆小 kernel 在空轉(zhuǎn)——這和 GPT-OSS 這類 MoE 的優(yōu)化思路一致但 Qwen3 的 A3B 激活比更稀疏融合收益更大。