Feature Request: Dynamic Inference Context Support for T5
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
The current inference implementation for T5 models only supports static batching via `StaticInferenceEngine`. It would be valuable to extend [`DynamicInferenceContext`](https://github.com/NVIDIA/Megatron-LM/blob/5247a1f4600c258c8f5d40c3bd934f210dbf3698/megatron/core/inference/contexts/dynamic_context.py#L193) to support encoder-decoder architectures like T5 similar to how it currently works for GPT inference via `DynamicInferenceEngine`.
## Current Behavior
The T5 inference wrapper does not support KV caching or dynamic batching
## Desired Behavior
Enable KV caching (encoder, decoder self-attention, and cross-attention) with in-flight dynamic batching for T5.
## Related Code References
- Dynamic context implementation: [`megatron/core/inference/contexts/dynamic_context.py#L193`](https://github.com/NVIDIA/Megatron-LM/blob/5247a1f4600c258c8f5d40c3bd934f210dbf3698/megatron/core/inference/contexts/dynamic_context.py#L193)
- Dynamic inference engine:
[`megatron/core/inference/engines/dynamic_engine.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/inference/engines/dynamic_engine.py)
Contributor guide
Assessment
This issue has not been assessed yet.