microsoft / microsoft/onnxruntime
[Feature Request] Jagged batches support (NJT) for Transformer / SDPA inference
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the feature request
PyTorch / HF (previously branded as BetterTransformer) now have some support for NJT representation:
- https://github.com/onnx/onnx/issues/6525
This allows to have efficient inference in context of continuous batching.
Does ORT have such NJT-enabled SDPA kernels built-in? (i.e. FlashAttention kernels efficiently supporting block-diagonal masks)
Thanks!
### Describe scenario use case
N/A
Contributor guide
Research direction
No files or tests are named. Start by locating ONNX Runtime's existing SDPA and FlashAttention inference kernels, then compare their mask handling with PyTorch/Hugging Face NJT and block-diagonal-mask support. Done would require a clear determination of built-in support, or a scoped implementation and validation plan for continuous batching.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, pytorch
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100