Repository navigation
[CUDA] Harden INT3 QMoE validation and correctness regressions - #33578
David Fan (jiafatom) wants to merge 10 commits into
Conversation
Signed-off-by: David Fan <jiafa@microsoft.com>
|
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. |
There was a problem hiding this comment.
🟡 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
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
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.
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
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
RunOptionson every pluginOnRunEnd, and don't return beforeOnRunEndruns; - make the failed-capture cleanup best-effort (non-throwing) so it cannot mask the original run error;
- minor:
sync_streamon 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.
|
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. |
|
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. |
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
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()); |
There was a problem hiding this comment.
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 | |||
There was a problem hiding this comment.
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.

Latest Review Hardening (c8dc3c5)
Verification for c8dc3c5
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.
Previous Review Follow-up (d8ec089)
Validation Executed for This Commit
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.
Historical Executed Validation
/tmp/int3-validation-hardening/Release; Python import identity was explicitly checked against this new build. No old comparison ORT binaries or objects were reused.--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.Local Environment Notes
ee3ccbd2b2344299d3a4506c2954a47b2181a485rather than the repository pin. No dependency pin is changed in this PR. Local default-pin qualification is not claimed.ALLOW_RELEASED_ONNX_OPSET_ONLY=0produced 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.