Sync with Microsoft ONNX Runtime - 22092026 - #1305
Merged
Merged
Conversation
…soft#32619) ### Description - **Backends:** Remove the FlashAttention-only restriction for `is_causal=0`; support non-causal MEA, paged decode, XQA speculative decoding, and latent attention. - **Masking:** Allow each query to attend through the full live KV sequence while keeping left windows anchored to its own position. - **Coverage:** Add regression cases for head size 128 with 16/32/64-token pages, ragged batches, and backend-specific masking; update CUDA documentation. ### Motivation and Context DFlash2’s non-causal drafter uses head size 128, for which FlashAttention requires pages of at least 128 tokens. A valid 64-token page therefore fails the drafter’s PagedAttention node and forces target-only decoding for the rest of the run. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>
### Summary Strengthens attribute/input validation in three contrib-op shape-inference functions in bert_defs.cc so malformed or malicious models are rejected at `Graph::Resolve` instead of triggering out-of-bounds reads or signed-integer overflow during shape inference. ### Changes `CausalConvWithState`: validate the ndim attribute is in [1, 3] and cross-check tensor ranks against it (weight == ndim+2, channels-first input == ndim+2, channels-last input >= 3) before the spatial-dim loop indexes input.dim(2+i). Previously only a rank >= 2 guard existed, so ndim=2/3 with a low-rank input read past the shape's dimensions. `GatedDeltaNet`: require the head counts/sizes to be positive and add step-by-step overflow guards before computing the state_update capsule width, preventing signed int64 overflow (UB) in state_update_capacity * (num_heads_v + num_heads_k*head_size_qk + num_heads_v*head_size_v). `GroupQueryAttention` / `SparseAttention`: check the parsed `total_sequence_length` initializer is non-empty before indexing data[0].
Enumerate every device-supported configuration and let each operation apply its own performance-preference order. This prevents the global supported-config table order from acting as performance policy and allows future vendor- and workload-specific selection. Centralize required subgroup-size validation in the selector. Request an explicit subgroup size only when SubgroupSizeControl is available, preserving support for adapters with a fixed subgroup size. Add dedicated WebGPU math tests for subgroup-size compatibility and operation preference ordering.
### Description
Add the schema and validation contract needed for mixed-width QMoE
expert weights.
- Add optional `fc1_expert_weight_bits`, `fc2_expert_weight_bits`, and
`fc3_expert_weight_bits` attributes, inheriting from
`expert_weight_bits` when omitted.
- Pass effective bit widths through shared validation and compute packed
byte dimensions as `logical_elements * bits / 8`.
- Require byte-aligned weight rows while allowing zero-point rows to pad
to a whole byte.
- Preserve legacy single-`pack_size` helper overloads and existing 2/4/8
validation diagnostics.
- Wire effective widths through CPU, CUDA, and WebGPU QMoE providers.
- Require matching FC1/FC3 widths for fused SwiGLU.
- Bypass legacy prepacking and return an explicit `NOT_IMPLEMENTED`
status for mixed-width execution. Execution support will follow
separately.
- Preserve the existing uniform-width behavior and schema version.
### Attribute semantics
`expert_weight_bits` remains the legacy global value and the fallback
for each optional projection override:
```text
effective_fc1_bits = fc1_expert_weight_bits if present, otherwise expert_weight_bits
effective_fc2_bits = fc2_expert_weight_bits if present, otherwise expert_weight_bits
effective_fc3_bits = fc3_expert_weight_bits if present, otherwise expert_weight_bits
```
If an override is omitted, it is never left uninitialized: it inherits
`expert_weight_bits`, whose schema default is 4. If all three overrides
are present, each override takes precedence for its corresponding
projection. All effective values are currently restricted to 2, 4, or 8
bits. Fused SwiGLU additionally requires effective FC1 and FC3 widths to
match.
Examples:
- `expert=4`, no overrides -> effective widths `(4, 4, 4)`, existing
uniform 4-bit path.
- `expert=4`, `fc1=2`, `fc2=4`, `fc3=2` -> effective widths `(2, 4, 2)`,
mixed-width contract path.
- `expert=4`, `fc1=2`, `fc2=2`, `fc3=2` -> effective widths `(2, 2, 2)`,
but currently still treated as an override/mixed configuration because
the effective widths differ from the legacy global value.
- For a uniform 2-bit model, use `expert_weight_bits=2` and omit the
overrides.
### Current execution boundary
For this contract-only change, `is_mixed_width` means that at least one
effective projection width differs from `expert_weight_bits`:
```text
is_mixed_width =
effective_fc1_bits != expert_weight_bits ||
effective_fc2_bits != expert_weight_bits ||
effective_fc3_bits != expert_weight_bits
```
This intentionally rejects even `expert=4, fc1=2, fc2=2, fc3=2`. Shape
and zero-point validation use the effective 2-bit widths, while the
existing execution and dispatch paths still use the legacy global width.
Allowing execution in that state could interpret 2-bit data as 4-bit
data. The explicit rejection prevents silent incorrect results.
Follow-up execution support must make prepacking, scale/zero-point
interpretation, packed layouts, workspace sizing, and kernel dispatch
projection-aware. Once dispatch consumes effective widths, uniformly
overridden configurations can be normalized to one effective width,
while genuinely mixed configurations can dispatch FC1/FC2/FC3
independently.
### Testing
- Built `onnxruntime_provider_test` with CUDA enabled.
- Added CPU contract and invalid-shape tests.
- Added row-wise and block-wise mixed-width zero-point validation,
including invalid FC2 shapes.
- Added availability-gated CUDA and WebGPU provider contract tests.
- Preserved the legacy INT2 invalid-hidden-size diagnostic.
- Covered arbitrary 3/5/6-bit packed-byte calculations in the shared
helper.
- Repository lintrunner passes for the changed files.
Related to microsoft#32657.
ai-fw-intg
requested review from
Jaswanth51,
ankitm3k,
jatinwadhwa921 and
vthaniel
September 21, 2026 20:38
hdharpure9922
self-requested a review
September 22, 2026 05:51
hdharpure9922
approved these changes
Sep 22, 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
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.