[distributed] prefill (single and multi-prompt) and single prompt generation and decoding

[lessw2020]

This PR:
1 - enables single or batch prompt prefill and decoding.
2 - adds three supporting functions for this (encode string, create padded prompt, batch decode next tokens).
Single:

Batch - prompts:

and batch decode:

3 - enables single prompt decoding up to max_tokens with formatted output:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1133

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9242dfb with merge base 7a4f0d1 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[kwen2501]

Thank for the big effort. LGTM. Would appreciate it if the code can be cleaned up more (please see comments in line)

[kwen2501]

Minor comment:
For PR description, would be easier to preserve info in git history if using text than images.

[Jack-Khuu]

Rebasing to pick up CI fix