onnxruntime

1b6e8e06 - [CUDA] Add SpeculativePrefixCommit contrib operator (#33474)

[CUDA] Add SpeculativePrefixCommit contrib operator (#33474) ## Description Adds the `com.microsoft::SpeculativePrefixCommit` (opset 1) CUDA contrib operator for single-launch device-side speculative prefix verification, control/context state commit, and result-ring publication across $3 \le S \le 8$ verification rows ($B \ge 1$). The operator runs entirely on the device with zero host synchronization or dynamic allocation so it can be captured and replayed inside a CUDA Graph. ## Summary of Changes ### `com.microsoft::SpeculativePrefixCommit` Schema & CUDA Kernel | File | Change | |------|--------| | `onnxruntime/core/graph/contrib_ops/bert_defs.cc`, `ms_opset.h` | Registers `com.microsoft::SpeculativePrefixCommit(1)` with attribute `head_rows` (`3..8`, default `8`), 7 inputs (`next_tokens`, `input_ids`, `status`, `ctl`, `ctx_commit`, `prefixes`, `ptrs`), and 3 required outputs (`head_c`, `updated_ctl`, `updated_ctx_commit`), enforcing rank 1 (`[rows]`) or rank 2 (`[batch, rows]`) on `next_tokens` and propagating `head_c` shape `[batch]` only when `next_tokens` has a known shape | | `onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc`, `onnxruntime/contrib_ops/cuda/llm/speculative_fold.{h,cc,cu}` | Implements `SpeculativePrefixCommit` (`MayInplace(3, 1)` and `MayInplace(4, 2)`), launching one thread per batch item (`PrefixCommitKernel`) to match verifier predictions against draft tokens (`input_ids[b, 1:rows]`), write the updated 16-word control state (`updated_ctl`) and 2-word context pair (`updated_ctx_commit`) in-place or to disjoint output buffers, and publish an 8-word record to the per-request result ring after `__threadfence_system()` | ### Tests and Documentation - `onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc`, `cmake/onnxruntime_unittests.cmake`: Kernel-level tests (`FoldContribTest.*`) covering all accepted prefix lengths for `head_rows` in `{3, 4, 8}`, in-place and disjoint state output buffers, multi-batch per-item state and result-ring progression, muted prefill (`ctl[3] != 0`), invalid metadata / ring validation flags (`ctl[5]`), verifier status aggregation, and CUDA Graph capture/replay. - `onnxruntime/test/contrib_ops/speculative_fold_op_test.cc`: Session-level tests (`SpeculativeFoldTest.*`) covering batched $S=8$ commits through `InferenceSession`, multi-run state progression carrying `updated_ctl` / `updated_ctx_commit` across `Run` calls, muted prefill with variable row counts, shape inference with known and unknown `next_tokens` shapes, and rejection of invalid `next_tokens` ranks, missing required state outputs, or partially overlapping buffers. - `docs/ContribOperators.md`, `docs/OperatorKernels.md`, `docs/contrib_ops/cuda/speculative_fold.md`: Regenerates schema/kernel tables and documents the operator interface, state aliasing contract, result-ring layout, error flags, and graph wiring. ## Testing - Run provider unit tests: `onnxruntime_provider_test --gtest_filter="FoldContribTest.*:SpeculativeFoldTest.*"` - End-to-end verification on `qwen_flash_next_int4_int8_rc3_contrib_dflash2` ($S=8$ DFlash2 speculative decoding). ## Checklist - [x] Tests added/updated - [x] No breaking changes (or documented in description) - [x] Documentation updated (if applicable) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>