Expose speculative draft proposals through runtime APIs

[Tianlei Wu (tianleiwu)]

Summary

Expose speculative draft proposals through the runtime APIs and verify them within the dynamic batching engine. The verify path preserves request transaction boundaries and uses one batched device argmax for contiguous logits rows, with the existing host path retained as a compatibility fallback.

Stack

Depends on #2491.

This runtime change does not depend on an ONNX Runtime version bump.

Changes

• Add C, C++, and Python APIs for submitting draft token proposals and querying the per-step draft limit.
• Plan paged block-table capacity for target tokens plus the proposed draft block.
• Preserve the pre-sampling sequence boundary when committing accepted draft tokens, including drafted stop tokens.
• Select verify rows with a batched device argmax when their logits are contiguous.
• Fall back to the existing host-copy argmax path when the device or layout cannot use the batched path.
• Add engine coverage for speculative request behavior and drafted stop-token completion.

Testing

cmake --build build/Linux/Debug --parallel $(nproc)
./build/Linux/Debug/engine_unit_tests

286 tests passed.

The built Python module also exposes Request.set_draft_tokens and Engine.max_draft_tokens_per_step.

[Copilot]

Pull request overview

This PR exposes speculative draft token proposals through the public runtime APIs (C/C++/Python) and integrates draft verification into the dynamic batching engine, including a fast-path that verifies contiguous draft rows via a single device argmax and a stop-token safeguard to preserve turn-completion semantics.

Changes:

• Add runtime APIs to set per-step speculative draft proposals and query engine draft support/limits.
• Update dynamic scheduling/resource planning to account for draft-driven step length and cache needs.
• Optimize draft verification with a contiguous-row device argmax path, retain host fallback, and add engine test coverage for drafted stop-token behavior.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
test/cpp/engine/engine_step_tests.cpp	Adds a regression test ensuring drafted stop tokens do not bypass normal end-of-turn handling.
src/python/python.cpp	Exposes Request.set_draft_tokens and Engine.max_draft_tokens_per_step in Python bindings.
src/ort_genai.h	Adds C++ RAII wrapper methods for setting draft tokens and querying max draft tokens per step.
src/ort_genai_c.h	Adds C API declarations + documentation for draft-token proposals and engine draft-limit query.
src/ort_genai_c.cpp	Implements the new C API entry points for draft proposals and max-draft query.
src/engine/scheduler.cpp	Plans cache capacity for base token + drafts before final token budgeting is applied.
src/engine/scheduled_requests.cpp	Adds a device argmax fast-path for draft verification and stop-token rejection in draft prefix acceptance.
src/engine/request.h	Adds stop-token query and refactors staging boundary tracking for transaction sampling.
src/engine/request.cpp	Implements stop-token detection and refactors transaction sampling boundary handling to match draft semantics.
src/engine/decoders/varlen_decoder_io.cpp	Improves the error message when draft verification is attempted with non-per-token logits.

Suppressed comments (1)

src/ort_genai_c.cpp:1470

• The OGA_TRY macro is placed on the same line as the function signature here, unlike the rest of this file (where OGA_TRY starts on its own line). This looks like an accidental formatting regression and makes the C API wrapper harder to read.

OgaResult* OgaRequestHasUnseenTokens(const OgaRequest* request, bool* out) {  OGA_TRY
  *out = request->HasUnseenTokens();
  return nullptr;
  OGA_CATCH
}

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.