[distributed] add batch decoding

[lessw2020]

This PR:

1. enables batch decoding such that multiple prompts can be run and processed at the same time.
2. moves the dst/src calcs for pp comms out of the decoding loop
3. updates final display to show each respective prompt and result.
4. Tested with both PP=1 and PP=2

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 5f71825 with merge base 8b45633 ():

NEW FAILURE - The following job has failed:

• pull / test-mps / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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

[kwen2501]

If pp_degree = 1, the code as is would cause prompt_length to be incremented twice per iteration.

That's why in PR #1149 I consolidated the increments in one place.

Would be nice to test with --pp=1 flag.

[kwen2501]

LGTM!