[Akshay Sonawane (apsonawane)]

Summary

Adds the com.microsoft::FoldMid2 and com.microsoft::FoldTail (version 1) CUDA contrib operators for speculative decoding and $S \in [2, 4]$ DFlash2 verification (squashing #33550 and removing unused MtpMoe):

• com.microsoft::FoldMid2 (onnxruntime/contrib_ops/cuda/llm/speculative_fold.{h,cc,cu}):
  • Commits accepted speculative prefix tokens (head_rows 2 through 4, head_capacity = 262144), prepares head token IDs / cache slots / rotary slices, and updates control/context state (MayInplace or in-place 7-output DFlash2 mode) and the mapped host result ring.
• com.microsoft::FoldTail (onnxruntime/contrib_ops/cuda/llm/speculative_fold.{h,cc,cu}):
  • Aggregates downstream validation status bits, optionally updates the configured draft-control slot (ctl_slot in {-1, 2, 8, 9}), and publishes the result-ring sequence word.

[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.