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