Skip to content

Split prefill/decode within single mxr - #5128

Open
turneram wants to merge 9 commits into
developfrom
merge-prefill-decode
Open

Split prefill/decode within single mxr#5128
turneram wants to merge 9 commits into
developfrom
merge-prefill-decode

Conversation

@turneram

Copy link
Copy Markdown
Contributor

Motivation

Add split prefill/decode support for dynamic-sequence ONNX models. This enables efficient static specialization of decode (sequence_length=1) and prefill (sequence_length=max) workloads while preserving a single runtime-facing program.

Technical Details

  • Added ONNX parser options for split prefill/decode and symbolic shapes.
  • Generates static decode and prefill submodules behind a runtime select_module dispatcher.
  • Exposed the options through the C++, Python file/buffer, and driver interfaces.
  • Preserved symbolic dimensions through optimization and serialization.
  • Improved evaluation and output allocation for fused multi-output tuple parameters, including loop and scan outputs.
  • Added reference, GPU, serialization, GQA, multi-I/O, driver, Python, and unsupported-length validation.
  • Added developer documentation covering architecture, usage, and validation.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

  • Added: New functionality.
  • Changed: Changes to existing functionality.
  • Removed: Functionality or support that has been removed. (Compared to a previous release)
  • Optimized: Component performance that has been optimized or improved.
  • Resolved Issues: Known issues from a previous version that have been resolved.
  • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

@turneram
turneram requested review from a team and causten as code owners August 11, 2026 17:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant