Skip to content

[CUDA] Harden INT3 QMoE validation and correctness regressions - #33578

Open
David Fan (jiafatom) wants to merge 10 commits into
mainfrom
jiafa/int3-qmoe-validation-hardening
Open

David Fan (jiafatom) wants to merge 10 commits into
mainfrom
jiafa/int3-qmoe-validation-hardening

Conversation

@jiafatom

@jiafatom David Fan (jiafatom) commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Latest Review Hardening (c8dc3c5)

  • Avoid RunOptions copies for successful runs and EPs without OnRunEnd. Failure-only internal configuration uses a shared key, avoids overwrite warnings, and falls back to invoking teardown if option preparation throws.
  • Failed-capture cleanup in both graph managers is best-effort: warn on unexpected EndCapture/destroy errors, clear CUDA last error, and do not replace the original execution failure with a cleanup exception. The plugin failure path honors sync_stream.
  • Added repository mock-plugin regression tests for success/failure forwarding, caller-options isolation, missing callbacks, and teardown status propagation; added an InferenceSession execution-failure/recovery regression for the 3-argument OnRunEnd contract.

Verification for c8dc3c5

  • Current-source A100 plugin graph-manager harness passed unexpected EndCapture warning/no-throw, sticky-error clearing, repeated failed same-ID capture discard, invalidated-capture recovery, and subsequent successful replay. Temporary harness only; diagnostics use stubs, CUDA calls are real.
  • Full host bridge, legacy graph manager, plugin graph manager, CudaEp, and mock-plugin test translation units passed -Werror syntax compilation. Formatting and whitespace checks passed.
  • New unit tests were not executed locally. Compiling the complete InferenceSession test translation unit is blocked by an existing mismatch between current tensor type metadata and cached ONNX dependency headers (tensorprotoutils.h array initializer); no product/dependency workaround was applied.
  • Linux/Windows CUDA and CUDA Plugin EP build/test jobs passed on the previous head d8ec089. Those are historical results, not qualification of c8dc3c5. Current-head hosted CI and full-runtime/plugin qualification remain pending.

Summary

Follow-up to merged #33569, tracked by #32657. Reuse the merged kernels and canonical packed model format; do not reintroduce superseded #33555 private prepack or experimental benchmark artifacts.

  • Validate finite/nonnegative scales by default, including models with omitted zero points; remove the ORT_VALIDATE_INT3_CONTENT gate.
  • Cache successful validation only when all supplied scale/zero-point metadata is immutable. Revalidate runtime-fed or overridable metadata on each non-capturing invocation.
  • Reject dynamic metadata during CUDA graph capture. Require previously validated immutable metadata for capture, with non-capturing warmup guidance.
  • Add formal Python CUDA regressions and document validation, capture, and scratch support boundaries. Existing INT2/INT4-only kernels and INT3 compute kernels are unchanged.

Previous Review Follow-up (d8ec089)

  • Clear the expected sticky CUDA error after an invalidated failed capture in both graph managers.
  • Propagate execution success through PluginExecutionProvider to CUDA plugin OnRunEnd using a callback-local RunOptions copy and an internal run config entry; no OrtEp C ABI change or caller RunOptions mutation. Discard failed captures without instantiation/cache/replay and do not count failed warmup runs.
  • Zero the bound output before each positive capture/replay check, and verify non-capturing recovery after three rejected same-ID capture attempts.
  • Clarify automatic ORT-managed warmup versus explicit external-capture warmup.

Validation Executed for This Commit

  • A100 GPU harness compiled the current plugin graph manager source and passed: three failed same-ID captures are discarded, invalidated capture clears cudaGetLastError(), normal stream work recovers, and a subsequent same-ID successful capture/replay writes the expected output. The temporary harness is not part of this PR; it links the installed ORT library for C API access and stubs only stack-trace diagnostics. This is manager-level GPU validation, not a rebuilt full-runtime/plugin integration test.
  • Complete current host-bridge and CUDA plugin callback translation units passed -Werror syntax compilation using existing build dependency headers and the plugin-specific CMake macros. Plugin graph manager syntax compilation also passed.
  • Ruff lint/format, touched C++ clang-format checks, and git diff --check passed.
  • Full QMoE runtime/plugin integration and sanitizer suites were NOT rerun for this commit: the previous isolated build/venv are no longer available. Hosted CI and full plugin integration qualification remain pending. The results below are historical and must not be interpreted as current-head qualification.

Acceptance Checklist

The checklist below records prior implementation coverage. Runtime code has subsequently changed in the legacy and plugin failed-capture paths; current-head verification is listed separately above.

  • Default validation: omitted zero points reject negative/NaN/Inf scales in either projection; zero scales remain valid; immutable initializers are covered.
  • Dynamic content: same-Session valid -> invalid -> valid mutations for scales and zero points; overridable initializers do not inherit immutable validation caching.
  • Width/dtype/block coverage: FP16/BF16, FC1/FC2 pairs (3,3), (3,4), (4,3), (2,3), (3,2), block32/64/128 where applicable, independent packed-code/dequantization reference with dtype-aware tolerances.
  • Explicit zero points: reject non-4 INT3 codes and nonzero padding, preserving valid asymmetric INT2/INT4 zero-point behavior in mixed projections.
  • ORT-managed capture: warmed-up immutable metadata supports capture/replay; dynamic scales/zero points reject capture; dynamic packed weights with immutable metadata remain supported.
  • Scratch: one-row budget minus one byte rejects; exact boundary succeeds; multirow limits and partial tiles match the reference; zero-token inputs validate content without token scratch.
  • Transitions: same-Session decode -> prefill -> empty -> partial tile -> decode remains correct under a bounded two-row tile.
  • Newly built CUDA provider focused tests, existing compatibility regressions, and targeted Compute Sanitizer checks pass, subject to the environment notes below.
  • Direct external capture: an unvalidated immutable initializer is rejected on a user-provided CUDA stream, followed by successful non-capturing recovery, warmed external capture, and replay with three input values. Both implicit and explicit zero points are covered.
  • Repeated calls: all three baseline pairs (3,3), (3,4), (4,3), FP16/BF16, constant and dynamic scales; 12 calls per mode with alternating inputs/scales and independent-reference comparisons.
  • Concurrent host Run calls with CUDA Graph disabled: three workers, six calls per worker, shared and independent Sessions, constant and dynamic scale/zero-point metadata, all three baseline pairs and FP16/BF16. Each worker owns its binding and input/output buffers. This checks correctness under overlapping host calls, not guaranteed simultaneous GPU kernel execution.

Historical Executed Validation

  • Independent CUDA13/SM80 build completed in /tmp/int3-validation-hardening/Release; Python import identity was explicitly checked against this new build. No old comparison ORT binaries or objects were reused.
  • Full formal QMoE CUDA test file after the follow-up: 195 passed, 13 skipped, 289 subtests passed. This consists of 77 INT3 validation tests with 285 subtests and 118 existing compatibility cases with 4 subtests. Ten skips are unsupported CUDA row-wise asymmetric configurations; three are opt-in benchmarks, not claimed coverage.
  • New follow-up slice: 32 passed, 12 subtests passed (2 external-capture, 6 repeated-call and 24 concurrency cases).
  • Complete INT3 test-class Compute Sanitizer memcheck: 77 passed, 285 subtests passed; ERROR SUMMARY: 0 errors, with --error-exitcode 99, in 1939.78 seconds. Includes external capture, repeated calls, concurrency, mutation, scratch and transitions. This supersedes the earlier 7-test targeted memcheck; other non-INT3 workloads are not claimed as sanitizer-covered.
  • Focused sequential racecheck (FP16, pair (3,3), repeated constant/dynamic scales): 1 passed, 2 subtests passed; 0 hazards, 0 errors, 0 warnings.
  • Focused shared-Session dynamic-metadata concurrency racecheck: the test passed, but Compute Sanitizer emitted Internal Sanitizer Error: failure to track a kernel launch; future records might be incomplete, followed by a summary of 1 error. This run is inconclusive/tool-limited, not a racecheck pass and does not establish a code-level race. Device racecheck also does not establish host-thread race freedom.
  • Existing packed-format/standard-operator reference suite: 223 passed.
  • New fixture self-check: 30 bit-pair/block/dtype configurations, pack/unpack roundtrips, finite reference outputs, and transition shapes.
  • Ruff lint/format and git whitespace checks passed; source worktree remains clean.
  • Hosted checks previously reported at the earlier head: 98 SUCCESS, 1 NEUTRAL (CodeQL), 0 failures. New commits require their own hosted checks; those historical results are not represented as current-head CI qualification.

Local Environment Notes

  • Tests ran on an A100 (SM80), CUDA13, Python3.12, Torch2.12.1+cu130. The formal test file and its helper were copied unchanged to an isolated temporary directory to avoid importing the unbuilt source Python package.
  • Dependency source caches were reused, including ONNX source ee3ccbd2b2344299d3a4506c2954a47b2181a485 rather than the repository pin. No dependency pin is changed in this PR. Local default-pin qualification is not claimed.
  • Installed Python ONNX1.23.2 stamps default opset28 in existing fixtures, while this ORT build accepts released opsets through27 by default. The initial compatibility run failed at model loading (104 failed, 14 passed, 13 skipped), before the failing models could execute QMoE. Rerunning with the process-local setting ALLOW_RELEASED_ONNX_OPSET_ONLY=0 produced the 118-pass compatibility result above. No product default, test source, schema, or dependency pin was changed to work around this environment mismatch.

Scope And Remaining Limits

Scope includes host metadata validation, execution-result forwarding, legacy and plugin CUDA failed-capture cleanup, the existing formal CUDA test file, and operator documentation. No new kernel, private prepack, benchmark harness, raw experiment data, CPU INT3 support, model export, quality evaluation, performance claim, or broad concurrency/capture qualification is included. The empty-CUDA-graph warning in the negative external-capture test is expected: metadata is rejected before any kernel is launched.

The previously disclosed direct unvalidated-initializer external-capture gap is now covered rather than waived. Remaining deployment qualification includes long-duration stress, concurrent same-Session CUDA Graph usage (not tested or claimed here), plugin runtime qualification, and non-A100 architectures such as SM86/SM89/Hopper. A100 evidence must not be generalized to those platforms. These broader deployment scopes are deferred, not marked as passed; this PR is ready for review within the stated A100 test scope.

@jiafatom

Copy link
Copy Markdown
Contributor Author

Documentation follow-up in d375a9b: replaced the bit-width-specific "INT3 CUDA Support and Validation" section with topic-based sections for integer execution support, runtime content validation, CUDA graph capture, and token scratch configuration. Supported width combinations and bit-specific zero-point rules are tables; current validation/capture/scratch contracts remain explicitly scoped to the raw packed path containing INT3, not generalized to other integer paths. No implementation or test changes. Section-order/uniqueness and key-contract assertions plus git diff --check passed. Previously reported GPU results apply to implementation commit 88b27d2; GPU tests were not rerun for this documentation-only commit. Fetched origin/main before submission; no additional main commits needed integration.

@jiafatom
David Fan (jiafatom) marked this pull request as ready for review October 25, 2026 14:20
Copilot AI balanced review requested due to automatic review settings October 25, 2026 14:20

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Rejected captures can still be cached and replayed on later runs, bypassing QMoE validation and returning stale output.

1 open issue
What changed in this PR

Hardens CUDA INT3 QMoE metadata validation on top of the merged packed kernels, without changing compute kernels or the canonical model format.

Changes:

  • Makes scale validation mandatory and restricts validation caching to immutable metadata.
  • Rejects dynamic or unvalidated metadata during CUDA graph capture.
  • Adds CUDA regressions and documents validation, capture, and scratch-memory boundaries.
File Description
onnxruntime/​test/​python/​transformers/​test_qmoe_cuda.py Adds validation, reference-correctness, capture, scratch-limit, and transition regressions.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc Removes opt-in validation and adds capture-time metadata guards.
docs/​contrib_ops/​cuda/​moe_qmoe.md Documents supported INT3 combinations and runtime constraints.

Review effort: Balanced

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

@tianleiwu Tianlei Wu (tianleiwu) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall the mandatory INT3 QMoE metadata validation, initializer-only validation caching, and failed-capture teardown in CUDAExecutionProvider::OnRunEnd / CUDAGraphManager::CaptureEnd look solid.

Left a few inline follow-up suggestions regarding sticky CUDA error cleanup on invalidated captures, CUDA plugin EP (CudaEp / CudaGraphManager) parity for failed captures, and resetting the output buffer in the positive test_cuda_graph_content_policy cases.

Comment thread docs/contrib_ops/cuda/moe_qmoe.md Resolved

@tianleiwu Tianlei Wu (tianleiwu) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the follow-up. I re-checked my earlier threads against d8ec089: the sticky-error clearing in both graph managers, run-success propagation to the plugin EP (no OrtEp ABI change), output.zero_() before each capture/replay comparison, and the warmup wording in the docs are all addressed, and those threads are resolved.

The remaining items are non-blocking hardening suggestions on the new failed-capture path (inline):

  • avoid copying RunOptions on every plugin OnRunEnd, and don't return before OnRunEnd runs;
  • make the failed-capture cleanup best-effort (non-throwing) so it cannot mask the original run error;
  • minor: sync_stream on the plugin failure path, and a shared constant for the internal config key.

Test coverage note: the Python tests exercise the in-tree CUDA EP, but nothing in the repo covers the plugin OnRunEnd translation (ep.internal.run_succeeded) or the new 3-argument OnRunEnd contract in InferenceSession; the PR text says this was only checked with a local harness. The hosted CUDA / CUDA Plugin EP jobs are still pending on this head, so please confirm they pass before merge.

@jiafatom

Copy link
Copy Markdown
Contributor Author

Addressed the four latest hardening comments in c8dc3c5 and resolved their threads. Added repository coverage for plugin OnRunEnd forwarding and the InferenceSession success/failure contract. Prior-head Linux/Windows CUDA and CUDA Plugin EP build/test checks have now passed. Current-head verification and the local ONNX-cache limitation are recorded separately at the top of the PR description; new-head CI must pass before merge.

@jiafatom

Copy link
Copy Markdown
Contributor Author

Fixed the Exceptions Disabled MinSizeRel build failure in the latest commit: the new plugin OnRunEnd fallback used native try/catch, which fails under -fno-exceptions. Replaced it with ORT_TRY / ORT_CATCH / ORT_HANDLE_EXCEPTION and retained teardown/status handling. The complete current host-bridge translation unit passes -Werror syntax compilation both with exceptions enabled and with -fno-exceptions plus ORT_NO_EXCEPTIONS, ONNX_NO_EXCEPTIONS, JSON_NOEXCEPTION, and MLAS_NO_EXCEPTION. clang-format and git diff --check also pass. This is focused local compilation, not a full MinSizeRel rebuild; current-head CI still needs to complete.

@tianleiwu Tianlei Wu (tianleiwu) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the follow-up. I read c8dc3c568d and 1e75de94b3 closely and found no correctness problem. The failed-capture cleanup is warn-only and non-throwing in both graph managers (it ends the capture, clears the non-sticky error, and caches/replays nothing). The plugin OnRunEnd resets the thread-local capture flag, honors sync_stream, and never lets a cleanup error replace the run's own error. The run-success forwarding adds no public C API surface (the OrtEp struct and its OnRunEnd callback signature are untouched, the flag travels as an internal run-config entry, the caller's RunOptions is copied only on failure, no locks or shared state), and other EPs keep working through the default 3-arg virtual. All earlier concerns are addressed; one small inline suggestion remains.

Validation: hosted CI is green at head (102 checks succeeded, CodeQL neutral; Linux/Windows CUDA, CUDA Plugin EP with and without cuDNN, TensorRT, MinSizeRel without exceptions). Locally on A100 (sm_80, CUDA 13.0, Linux; C++ gtests not run) test_qmoe_cuda.py passes on the in-tree CUDA EP (195 passed, 13 pre-existing skips) and on the CUDA Plugin EP (194 passed; TestQMoEInt3Validation is 77/77 on both; the one plugin failure is test_int4_fp16_qmoe_cuda_graph, which is not touched by this PR and fails in OrtValue.ortvalue_from_numpy(..., 'cuda') because the plugin-only build has no in-tree CUDA provider). These Python tests are not run by any CI job.

Recommend approving after verification on the impacted devices: sm_90 (H100) and Blackwell (sm_100/sm_120), running the INT3 QMoE and CUDA-graph tests with both CUDA EPs, plus Windows with a GPU. Verified: hosted CI and local A100/Linux. Not verified: sm_86/89/90/100/120, Windows locally.

for (auto* xp : exec_providers_to_stop) {
bool synchronize_execution_providers = run_options.config_options.GetConfigOrDefault(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "0") == "0";
auto status = xp->OnRunEnd(synchronize_execution_providers, run_options);
auto status = xp->OnRunEnd(synchronize_execution_providers, run_options, retval.IsOK());

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

retval.IsOK() is re-evaluated for every EP in this loop, and ORT_CHECK_AND_SET_RETVAL(status) (if (retval.IsOK()) retval = status;) flips retval as soon as one EP's OnRunEnd fails. Every EP after that one is then told run_succeeded == false although graph execution succeeded, so a CUDA EP would discard a capture it completed correctly (the Run fails anyway and the next Run simply re-captures, so this is only a wasted capture). Capturing the flag once before the loop, e.g. const bool run_succeeded = retval.IsOK();, keeps the meaning 'graph execution succeeded' for all EPs. PartialRun (L3646-3649, training only) has the same pattern.

@@ -0,0 +1,5 @@
#pragma once

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: the other 21 headers in core/session/plugin_ep/ start with the standard // Copyright (c) Microsoft Corporation. All rights reserved. / // Licensed under the MIT License. header, and this new file also has no trailing newline. lintrunner passes at this head, so this is only for consistency.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants