Repository navigation
Conversation
`--features transposed-matmul` (`build(..., transposed_matmul=True)`) rewrites each decoder `MatMul(x, Transpose(weight))` into `com.microsoft::FusedMatMul(x, weight, transB=1)`. The projections -- the packed Q/K/V, `o_proj`, the MLP and the LM head -- then keep the checkpoint's `[out, in]` weight instead of the `[in, out]` one FoldTransposedInitializerPass would store, and the transpose happens inside the GEMM. At batch 1 that is the faster layout for cuBLAS. Unlike `Gemm`, `FusedMatMul` takes the rank-3 activations a decoder layer feeds. The rewrite is a graph pass rather than a different `Linear`, because the GQA and PackQKV rewrites match the plain `MatMul`; it runs after fusion and lowering, so they still fire, and before weights are loaded. A `Transpose` is only taken when it reads a weight: an initializer, or the axis-0 `Concat` of initializers that PackQKV emits, which FoldConcatInitializersPass still folds once weights are loaded. Transposed activations, shared transposes and other permutations are left alone. `EpCapabilities.supports_transposed_matmul` gates it and only CUDA sets it; elsewhere the request warns and is ignored, as `fp8-kv-cache` does. Only decoder components are rewritten. Tested: the pass on single and packed weights against numpy, and a tiny fp16 Llama built for CUDA keeps GQA and packed Q/K/V, has no plain `MatMul` or weight `Transpose` left after loading, and matches the `MatMul` export on ORT CPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Please let us know the actual use case. Thanks. We don't want to blindly optimize the graph just for measured ORT perf. |
|
Thanks for the comment. I’m running Qwen3-Omni, which mobius doesn’t support yet, at batch 1 on a single GPU. The tuning that took is too large to review as one change, so I’m splitting it into separate PRs; Qwen3-Omni itself will follow as its own PR. Profiling against vLLM on an RTX PRO 6000 showed the projections’ weight layout was the main gap. This change alone took about 2 ms off a decode step and brought the export level with vLLM’s int8 path on the same GPU (5.70 vs 5.71 ms per step). |
|
Thanks for the explanation! |
|
cc @gramalingam @titaiwangms another case where we may want to retain the old layout and not fold transpose. |
Summary
--features transposed-matmul(build(..., transposed_matmul=True)) rewrites each decoderMatMul(x, Transpose(weight))intocom.microsoft::FusedMatMul(x, weight, transB=1)o_proj, the MLP and the LM head) then keep the checkpoint's[out, in]weight instead of the[in, out]oneFoldTransposedInitializerPasswould store, and the transpose happens inside the GEMM; at batch 1 cuBLAS is faster on that layoutGemm,FusedMatMultakes the rank-3 activations a decoder layer feedsEpCapabilities.supports_transposed_matmul); other EPs warn and ignore the request, asfp8-kv-cachedoes. Only decoder components are rewrittenNothing changes for exports that do not pass
transposed-matmul.How the rewrite is placed
It is a graph pass rather than a different
Linear, because the GQA and PackQKV rewrites match the plainMatMul. It runs after fusion and lowering, so they still fire, and before weights are loaded. ATransposeis only taken when it reads a weight: an initializer, or the axis-0Concatof initializers that PackQKV emits, whichFoldConcatInitializersPassstill folds once weights are loaded. Transposed activations, transposes with another consumer, transposes that are graph outputs and other permutations are left alone, and the rewritten node keeps the output name (sologitsstayslogits).Performance
Batch 1, fp16,
--features static-cachewith and withouttransposed-matmul, ORT 1.29 (CUDA EP), no other process on the GPU. Decode is one step under CUDA graph replay with an 8192-slot cache, the median of five blocks of 200 replays; prefill is one forward without a graph, the median of 20 runs. Each figure is the mean of rounds 2-3 of three rounds that alternate between the two exports.ms, RTX PRO 6000 Blackwell Max-Q (300 W):
static-cache+ transposed-matmulms, A100 80GB PCIe:
static-cache+ transposed-matmultransposed-matmulexport is the slower one in every round; at 2048 on Blackwell the rounds disagree on the direction, so that figure is mostly drift.It is opt-in because it only helps the batch-1 decode step and costs prefill a little.
Validation
MatMuland, after weights are loaded, noConcatorTransposeleft, and matches theMatMulexport on ORT CPU--features transposed-matmulthrough tobuildNot yet measured: batch sizes above 1, models other than the two above.
🤖 Generated with Claude Code