NVIDIA / NVIDIA/Megatron-LM

Feature Request: Dynamic Inference Context Support for T5

Open
#3,016 2 comments 0 reactions 1 assignee Claimed by @janEbert View on GitHub
community-request enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.