[Distributed] Separate prefill and decode

[kwen2501]

Prefill phase and decoding phase are now done separately.
Decoding phase will leverage KV cache and speed improves.

Requires a PyTorch side fix to work. Will change the pin version.

Requires pytorch PR pytorch/pytorch#136243 to land.
Otherwise, you may hit issue described in: pytorch/pytorch#136225

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 1 Unrelated Failure

As of commit 3512101 with merge base e27e162 ():

NEW FAILURE - The following job has failed:

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

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-mps / macos-job (gh) (trunk failure)
  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.

[lessw2020]

Removing this limitation is probably our most important next step.

[lessw2020]

Great work, looks great. Left comment but removing the all prompt same len is likely most important next step. Anyway this pr is a big step forward, nice job!