Motivation

This PR adds a feature to dump both auxiliary and last hidden states after a request is finished.
It is mainly intended for SpecForge offline training and this feature can also be useful for collecting data from real user requests for training purposes.

The existing --enable-return-hidden-states option only returned the last hidden state and caused slowdowns due to frequent device-to-host copies.

Modifications

• hidden states (aux and last) are concat and copied to host at only once after the request is done.
• dump is performed in a round robin way across TP ranks.
• added a dump stream for concat and DtoH comm of hidden states which can overlap with the main inference job
• added a dump process pool that actually saves data to disk

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.