[Akshay Sonawane (apsonawane)]

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

• Tests added/updated
• No breaking changes (or documented in description)
• Documentation updated (if applicable)

[Copilot]

🟡 Changes recommended

Unresolved memory-safety, state-propagation, uninitialized-output, and documentation-validation issues block approval.

6 open issues

• Regenerate operator documentation from the updated build · New
• Initialize the complete public scratch output buffer · New
• Handle non-finite router logits without generating invalid expert indices · New
• Return updated fold state through explicit outputs · New
• Bound rotary-cache accesses to the actual cache extent · New
• Validate ring metadata before state commit or record access · New

What changed in this PR

Adds CUDA contrib operators for multi-token prediction and speculative decoding directly on main, replacing the earlier stacked PR.

Changes:

• Implements NVFP4 MtpMoe projections and FoldMid2/FoldTail state updates.
• Adds schemas, CUDA registrations, and operator documentation.
• Adds internal kernel tests; fresh GPU execution remains outstanding.

File	Description
onnxruntime/​test/​contrib_ops/​cuda_kernels/​speculative_fold_test.cc	Tests acceptance, state updates, and muted calls.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​mtp_moe_test.cc	Tests routing and reference projection arithmetic.
onnxruntime/​core/​graph/​contrib_ops/​ms_opset.h	Registers the three schemas.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Defines operator contracts and types.
onnxruntime/​contrib_ops/​cuda/​llm/​speculative_fold.h	Declares fold launchers.
onnxruntime/​contrib_ops/​cuda/​llm/​speculative_fold.cu	Implements speculative state kernels.
onnxruntime/​contrib_ops/​cuda/​llm/​speculative_fold.cc	Adds fold validation and CUDA wrappers.
onnxruntime/​contrib_ops/​cuda/​llm/​mtp_moe.h	Defines dimensions and launcher interfaces.
onnxruntime/​contrib_ops/​cuda/​llm/​mtp_moe.cu	Implements routing and expert projections.
onnxruntime/​contrib_ops/​cuda/​llm/​mtp_moe.cc	Adds validation and output allocation.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers CUDA implementations.
docs/​ContribOperators.md	Documents the new operators.
cmake/​onnxruntime_unittests.cmake	Includes internal tests and kernel sources.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Accepted row counts exceed CUDA launch limits, and documentation ordering conflicts with generator validation.

2 open issues

• Preserve support-level grouping in generated table of contents · New
• Validate or tile row dispatch beyond CUDA grid limits · New

6 resolved since last review

• Validate ring metadata before state commit or record access
• Bound rotary-cache accesses to the actual cache extent
• Return updated fold state through explicit outputs
• Handle non-finite router logits without generating invalid expert indices
• Initialize the complete public scratch output buffer
• Regenerate operator documentation from the updated build

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Pushed 24625bf to address both outstanding review threads and the common root cause of all six failed CUDA/TensorRT/plugin pipelines.

• Fixed speculative-fold session test compilation: explicitly include the status assertion header and reuse typed output names for the gsl::span-based Run overload in all three calls.
• Restore experimental operators to the correct support-level group in the generated documentation TOC.
• Validate the documented MtpMoe 1–65535 row range before output allocation and before any low-level CUDA access; add boundary regressions.

Validation: full-PR lintrunner apply/check-only passes, git diff --check passes, affected schema/kernel documentation matches the official Python renderer with source-derived metadata, entire TOC support-level/name ordering passes, and static validation placement checks pass.

No local builds or compilation were run. The new native tests were not executed locally. Fresh CI is triggered by this push and is not claimed green.

[Copilot]

🟡 Changes recommended

State propagation and unknown-batch shape inference have correctness issues, and generated documentation does not match the schema.

3 open issues

• Require both updated-state outputs and preserve committed state · New
• Preserve unknown batch dimension when next_tokens lacks shape · New
• Regenerate documentation to match the current schema · New

2 resolved since last review

• Validate or tile row dispatch beyond CUDA grid limits
• Preserve support-level grouping in generated table of contents

_{Review effort: Balanced}

[kunal-vaishnavi]

From Copilot:

Two-pass source review of commit 22d6117. This PR merged after the review was completed, so I am posting this finding here for a follow-up fix.

[P1] Prevent a reused ring slot from looking committed while its payload is overwritten

Location: onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu, line 98.

With the schema-valid ptrs[1] = 1, the next commit overwrites words 1-7 while word 0 still contains the previous sequence number. A concurrent host poller can therefore read sequence 1 and payload from commit 2; reading word 0 again does not help if the writer has not published sequence 2 yet. The system fence before the final sequence store orders initial publication, but does not protect an already-published slot from reuse. An H200 probe using this store/fence order observed sequence=1, token=2, length=2, sequence_again=1.

Please add a slot-ownership/consumer-acknowledgment protocol, or a writing-generation protocol with the corresponding acquire/revalidation rules for the host reader, before overwriting a published record. If preventing reuse until consumption is a caller obligation, make that an explicit required contract instead of promising that reading word 0 first always yields a complete record. Add a mapped-host-memory regression that polls concurrently across ring wraparound; the current tests inspect records after synchronization.

Validation: the standalone CUDA probe reproduced the publication mechanism on H200; it was not a build or session-level run of the unmodified ORT operator. No full PR-specific ORT test suites were run. This finding is separate from the already-addressed ring-metadata and returned-state comments.