Sync with Microsoft ONNX Runtime - 10062026 - #1128
Merged
Merged
Conversation
) ### Description: This PR adds shape validation in `GemmNodeGroupSelector::Check` to prevent incorrect DQ-Gemm-Q to QGemm fusion when DequantizeLinear nodes have incompatible scale/zero_point shapes. ### Problem The QGemm op requires a_zero_point and b_zero_point as mandatory inputs, and only supports scalar or 1D scale/zero_point. Previously, the selector did not validate these constraints, which could lead to runtime crashes when incorrect fusion occurs. ### Changes - Added `IsScalarOr1DWithSizeOneOrN()` helper to validate that a NodeArg's shape is scalar or 1D with size 1 or N. Rejects unknown shapes to avoid incorrect fusion. - Added checks in `GemmNodeGroupSelector::Check` to maintain consistency with QGemm checks. - Added test `Gemm_NoQGemmFusionWithHighRankScaleZp` to verify that DQ nodes with rank-2 scale/zero_point are not fused into QGemm.
…dels (microsoft#28484) ### Description Adds a per-head Q/K RMS normalization prologue to the WebGPU `GroupQueryAttention` contrib op so that Qwen3-style models can fold the standalone `SimplifiedLayerNormalization` dispatches on Q and K into the attention kernel. The MS-domain `GroupQueryAttention` schema gains two optional inputs and one attribute: * `q_norm_weight` (input 14, optional, shape `(head_size,)`) * `k_norm_weight` (input 15, optional, shape `(head_size,)`) * `qk_norm_epsilon` (FLOAT attribute, default `1e-6`) A new Level 2 optimizer pass `GroupQueryAttentionPreNormFusion` rewrites the `Reshape -> SimplifiedLayerNormalization -> Reshape -> GroupQueryAttention` pattern produced by the Qwen3 ONNX export into a single `GroupQueryAttention` node that carries the norm weights directly. The pass is scoped to the WebGPU EP only. ### WebGPU runtime paths | sequence_length | Behavior | |---|---| | 1 (decode) | RMSNorm is folded into `FusedQKRotaryEmbedding`. | | > 1 (prefill) | Falls back to two standalone `SimplifiedLayerNorm` dispatches into scratch `qNorm`/`kNorm` tensors, then runs the unfused `FusedQKRotaryEmbedding`. Matches the pre-fusion graph timing exactly so prefill cannot regress. | ### Gating The optimizer only rewrites nodes that the WebGPU kernel can handle: * `do_rotary == 1` (the fused decode path runs inside the rotary kernel; the prefill fallback also runs only on the `do_rotary=1` branch) * Non-packed QKV (K at slot 1, V at slot 2 wired) * Already-fused nodes (q/k_norm_weight present) are skipped * The pattern must use a single norm weight per side with matching epsilon and `axis=-1`, and the surrounding reshapes must collapse to `(... head_size)` / expand back to `(... hidden_size)`. ### Defense-in-depth on other EPs The CPU, CUDA, and JSEP `GroupQueryAttention` kernels reject hand-authored models that wire inputs 14/15 with EP-specific error messages, since none of those runtimes implement the prologue: * CPU: `GroupQueryAttention (CPU): q_norm_weight / k_norm_weight inputs are not supported...` * CUDA: `GroupQueryAttention (CUDA): q_norm_weight / k_norm_weight inputs are not supported...` * JSEP: `GroupQueryAttention (JSEP): q_norm_weight / k_norm_weight inputs are not supported...` ### Motivation and Context Perf improvement on **D3D12 backend** for **Qwen3-1.7B** model on a **Windows** machine with **RTX 5060Ti** card <img width="426" height="205" alt="image" src="https://github.com/user-attachments/assets/83e3a0f1-4c2f-4c0f-845d-d5342091d004" />
…rosoft#28849) ### Description Upgrades the `cudnn_frontend` dependency from **1.12.0 → 1.24.0** and wires the updated cuDNN SDPA (scaled dot-product attention) kernels into the CUDA `MultiHeadAttention` and `GroupQueryAttention` operators. On SM≥90 (Hopper/Blackwell), cuDNN SDPA is auto-preferred for FP16/BF16 ahead of Flash Attention / cutlass FMHA, which significantly improves GQA prefill throughput. ### Key Changes | Area | Change | |---|---| | Dependency | `cmake/deps.txt`: `cudnn_frontend` 1.12.0 → 1.24.0. | | Build | `cmake/external/cudnn_frontend.cmake`: mark cudnn_frontend headers as `SYSTEM` includes so v1.24's unused static helper does not trip `-Werror=unused-function`. | | SDPA wrapper | `cudnn_fmha/cudnn_flash_attention.cc`: migrate to the v1.24 API — `set_generate_stats(false)` (replaces deprecated `set_is_inference`), diagonal-band causal masking (`set_diagonal_alignment` + `set_diagonal_band_right_bound`/`set_diagonal_band_left_bound`), and synthesize the missing `seq_len_q`/`seq_len_kv` side that v1.24 now requires when a padding mask is used. | | MHA | `multihead_attention.{cc,h}`: enable cuDNN SDPA for FP16 **and** BF16; compute cuDNN eligibility before Flash and prefer it on SM≥90 unless the user pinned a kernel. | | GQA | `group_query_attention.{cc,h}`, `group_query_attention_impl.cu`, `attention_data.h`: add a cuDNN SDPA path (non-quantized FP16/BF16, no softcap/smooth-softmax/head-sink/local-window, BNSH KV cache), dispatched after XQA and before Flash/MEA/unfused. | | Kernel selection | `attention_kernel_options.{cc,h}`: track explicit `sdpa_kernel` selection and honor an explicit `ORT_ENABLE_CUDNN_FLASH_ATTENTION=0` so it disables the SM≥90 auto path. | ### Kernel Priority - **SM≥90, FP16/BF16:** cuDNN SDPA is auto-preferred unless the user explicitly selects a kernel via the `sdpa_kernel` provider option or sets `ORT_ENABLE_CUDNN_FLASH_ATTENTION=0`. - **GQA decode:** XQA remains highest priority where eligible; cuDNN SDPA outranks Flash/MEA/unfused for the remaining eligible cases. - `ORT_ENABLE_CUDNN_FLASH_ATTENTION=0` disables cuDNN entirely (including the auto path); `=1` force-enables it; the `sdpa_kernel` provider option overrides env vars. ### Benchmark Results Measured with `onnxruntime/test/python/transformers/benchmark_gqa.py` on **NVIDIA H200 (SM 9.0)**, CUDA 13.0 / cuDNN 9.19, `Llama3-8B`-shaped GQA (b1, 32 query heads, 8 KV heads, head size 128, FP16). - **Baseline** = Flash Attention (`ORT_ENABLE_CUDNN_FLASH_ATTENTION=0`) - **This PR** = cuDNN SDPA (default on SM≥90) The prefill (prompt) phase shows the largest gains for the dense variants: **ORT-GQA-Dense — prompt latency (ms, lower is better)** | seq_len | Baseline (Flash) | This PR (cuDNN) | Speedup | |---:|---:|---:|---:| | 64 | 0.082 | 0.054 | 1.52× | | 128 | 0.186 | 0.056 | 3.33× | | 256 | 0.184 | 0.065 | 2.83× | | 512 | 0.247 | 0.074 | 3.32× | | 1024 | 0.295 | 0.122 | 2.42× | | 2048 | 0.681 | 0.250 | 2.72× | | 4096 | 1.294 | 0.699 | 1.85× | | 8192 | 3.864 | 1.256 | 3.08× | **ORT-GQA-Dense-PackedQKV — prompt latency (ms, lower is better)** | seq_len | Baseline (Flash) | This PR (cuDNN) | Speedup | |---:|---:|---:|---:| | 64 | 0.197 | 0.053 | 3.69× | | 128 | 0.160 | 0.055 | 2.93× | | 256 | 0.213 | 0.060 | 3.53× | | 512 | 0.226 | 0.073 | 3.09× | | 1024 | 0.333 | 0.291 | 1.15× | | 2048 | 0.595 | 0.252 | 2.36× | | 4096 | 1.312 | 0.697 | 1.88× | | 8192 | 5.014 | 1.259 | 3.98× | Prefill is ~**1.5×–4×** faster across sequence lengths for both dense variants. Decode (token) latency is unchanged within run-to-run noise. ### Testing - `onnxruntime_provider_test --gtest_filter='GroupQueryAttentionTest.*:MultiHeadAttentionTest.*'` — GQA 44/44, MHA 18/18 pass on H200. - Broader attention regression (`AttentionTest.*:PackedMultiHeadAttentionTest.*:DecoderMaskedMultiHeadAttentionTest.*`) — 143/143 pass. - Verified the new path on cuDNN 9.8 and 9.19 (decode `s_q==1` causal-mask edge case handled for cuDNN ≤ 9.9). - Verified `ORT_ENABLE_CUDNN_FLASH_ATTENTION=0` disables the SM≥90 auto path (0 cuDNN selections) while the default run selects cuDNN. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
…nels (follow-up microsoft#28261) (microsoft#28308) ## Description Follow-up to microsoft#28261 (RVV CPU EP). This PR adds four RVV MLAS kernels on top of the existing `if(HAS_RISCV64_RVV)` block in `cmake/onnxruntime_mlas.cmake`: 1. **INT8 GEMM** (`riscv64/qgemm_kernel_rvv.cpp`) — `vwmulu.vv` / `vwaddu.wv` widening; wired through `MLAS_PLATFORM` 4-signedness dispatch (LARCH64 idiom). 2. **M=1 SGEMM routing** — extends the ARM64/WASM `MlasGemvFloatKernel` `#elif` to RISCV64; kernel comes from existing `SgemvKernelScalar.cpp` (rv-gcc autovec). 3. **Activation kernels** (`riscv64/activation_kernel_rvv.cpp`) — RVV `Erf`, `Tanh`, `Logistic`, `ComputeExpF32`, `Silu`, `GeluErf`. 4. **INT8 GEMV M=1 fast path** — `MlasGemmQuantTryGemvKernel<...RVV>` specialization (mirrors AVX2 `qgemm_kernel_avx2.cpp:131`); U8×S8 only. Rebased onto current `main`. Later commits fold in review feedback: per-signedness dispatch structs, the per-matrix `ZeroPointB` accumulator, NaN round-trip in the activation kernels, a full-range two-step `exp`, and infinity handling in SiLU/GELU. The correctness fixes are verified on SpacemiT K3 (riscv64, RVV) with `onnxruntime_mlas_test`. ## Performance K3 (SpacemiT X100, VLEN=256, 4 threads, p50 ms over 30 reps, real BAAI/bge-* ONNX inputs). Baseline = `microsoft/onnxruntime` post-microsoft#28261 (`62f742f1aa`). Cross-built with `riscv64-linux-gnu-g++` 15.2 + `-march=rv64gcv -mabi=lp64d`. ### FP32 transformer encoders | Model | Upstream | This PR | Speedup | | ------------------------------ | --------: | --------: | :-----: | | BAAI/bge-small-zh-v1.5 | 66.3 ms | 63.8 ms | 1.04× | | BAAI/bge-base-zh-v1.5 | 404.4 ms | 393.5 ms | 1.03× | | BAAI/bge-reranker-base | 403.7 ms | 391.7 ms | 1.03× | ### INT8 quantized | Model | Upstream | This PR | Speedup | | ------------------------------ | --------: | --------: | :-----: | | BAAI/bge-small-zh-v1.5 INT8 | 301.3 ms | 131.3 ms | 2.29× | | BAAI/bge-base-zh-v1.5 INT8 | 1958.8 ms | 669.2 ms | **2.93×** | | BAAI/bge-reranker-base INT8 | 1956.8 ms | 668.6 ms | **2.93×** | ### INT8 GEMV M=1 (kernel-level micro-bench, 1 thread) | Shape | scalar | autovec | this PR | vs autovec | | ----------- | -------: | -------: | -------: | :--------: | | K=N=384 | 1.45 GOPS | 2.34 GOPS | 16.68 GOPS | **7.13×** | | K=N=768 | 1.46 GOPS | 2.34 GOPS | 6.17 GOPS | 2.64× | | K=N=4096 | 1.46 GOPS | 2.33 GOPS | 1.92 GOPS | 0.82× (memory-bound) | The GEMV path triggers only when `RangeCountM == 1` and zero-points are zero (`qgemm.h:331` gate); BERT seq=128 encoders do not exercise it, so its contribution is not visible in the e2e tables above. --------- Signed-off-by: qiurui144 <happyqiurui@163.com>
…osoft#28721) WebGPU Pad kernel int64 / int32 truncation can read to oob read
GatherBlockQuantized shader could read oob for ill formed indicies. The correct solution would be to throw and exception but the indicies input sits on gpu and we don't want to take the performance hit by copying it to cpu. Instead, we only check inside the shader to not read oob.
…h MatMulNBits) (microsoft#28749) ## Summary This PR lets the CUDA `com.microsoft::QMoE` operator prepack **raw** int4/int8 expert weights into the CUTLASS `fpA_intB` layout **inside ORT's `PrePack()` hook**, instead of requiring callers to run the layout transform offline via `pack_weights_for_cuda_mixed_gemm`. This makes integer QMoE symmetric with `MatMulNBits::PrePack_B`, and lets exporters ship the schema-conformant `[E, N, K/pack]` quantized weights produced by `quantize_matmul_{4,8}bits` directly, with no offline pre-pack step. The behaviour is **opt-in and backward compatible**: a new `weights_prepacked` attribute is a tri-state and defaults to `-1` (auto), which the CUDA EP treats as "already CUTLASS-prepacked" (today's behaviour). `1` forces the prepacked interpretation explicitly, and only `weights_prepacked=0` triggers the new in-`PrePack` layout transform. ## What changed - **New `weights_prepacked` attribute** on the QMoE schema (tri-state, default `-1`/auto). `-1` lets each execution provider pick its own backward-compatible default; the CUDA EP treats auto as prepacked. `1` = the int4/int8 `fc1`/`fc2` initializers are already in the CUTLASS `fpA_intB` layout (today's behaviour). `0` = the initializers are raw `[E, N, K/pack]` tensors and the kernel runs the layout transform itself in `PrePack()`. - **`PrePackIntExpertWeights`** — loops over the `E` experts and applies the per-expert transpose + CUTLASS `fpA_intB` row-permutation / column-interleave / bias / pair-interleave transform on the GPU, mirroring `pack_weights_for_cuda_mixed_gemm`. Architecture-aware packing per `docs/contrib_ops/cuda/moe_qmoe.md` §7 (SM90 is its own layout group; all other supported arches share the SM80 layout; SM75+ required). - **`PrePack()` dispatch** for the int weight slots (2 and 5) when `quant_type == "int"` and `weights_prepacked == 0`. The source initializers are released after their shapes are cached (`fc1/fc2_weights_shape_`), so peak weight memory stays ~1×. - **`ComputeInternal`** prefers the prepacked GPU buffers when the PrePack hook populated them (gated on `int_weights_consumed_by_prepack`), and otherwise falls through to the raw initializer pointers (e.g. for sessions that set `session.disable_prepacking`). ## Schema note This **does add a schema attribute** (`weights_prepacked`) to QMoE. It is backward compatible because the default (`-1`/auto) is interpreted by the CUDA EP as prepacked, preserving the existing offline-prepacked behaviour, but it is a schema surface-area change. ## Diff scope | File | Change | |---|---| | `onnxruntime/core/graph/contrib_ops/contrib_defs.cc` | New `weights_prepacked` schema attribute + docs | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.h` | New private method + prepack buffer / cached-shape members | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` | `PrePackIntExpertWeights` + PrePack dispatch + ComputeInternal hookup | | `onnxruntime/test/python/transformers/test_qmoe_cuda.py` | CUDA smoke test for the raw-weight `weights_prepacked=0` path | FP4 / FP8 / WFP4AFP8 paths are untouched, and there is no behaviour change for callers that pre-prepacked their weights. ## Testing - `onnxruntime_providers_cuda` builds and links cleanly (nvcc 13.2 / sm_90). - `TestQMoEIntPrePackSmoke` (`test_qmoe_cuda.py`) builds a QMoE graph with raw int4 weights and `weights_prepacked=0`, runs it through the CUDA kernel, and asserts the output is finite with a plausible magnitude. Verified on H200 (SM90); node placement confirmed on `CUDAExecutionProvider` via profiling. - Existing int4 QMoE parity tests (`phi3` / `swiglu`, fp16) pass on CUDA — no regression in the default `weights_prepacked=-1` (auto, prepacked) path. > Note: this is a smoke test, not a numerical parity check. The existing offline > pre-pack harness hardcodes `force_arch=80` and produces incorrect output on > SM≥90, so a bit-parity comparison against it is intentionally omitted until > that harness honours the runtime SM. --------- Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
ai-fw-intg
requested review from
Jaswanth51,
ankitm3k,
jatinwadhwa921 and
vthaniel
June 9, 2026 20:33
ankitm3k
approved these changes
Jun 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Automated daily backmerge from ORT main to ovep-develop. No conflicts detected. Do NOT squash or rebase - use merge commit only.