[Distributed] Add lanes to KV cache

[kwen2501]

KV cache is extended to have multiple lanes, each letting a separate batch pass through, achieving pipeline parallelism.

# The number of cache lanes is the same as the maximum number of
# micro-batches that can be "in flight" in parallel -- imagine each
# micro-batch takes 1 "pipeline lane," they need distinct KV cache spaces.
# When decoding is done for certain micro-batches, we can reuse the KV cache
# lanes.

Major changes

1. setup_caches will take one kwarg cache_lanes (default to 1).

def setup_caches(self, max_batch_size, max_seq_length, cache_lanes: int = 1)

2. attention.kv_cache is now a nn.ModuleList, containing multiple KVCache's, each corresponding to a lane.

3. We now pass kwargs = {"input_pos": input_pos, "cache_lane": lane} to the step() function. Removing the temporary helper function model.setup_input_pos.

Requires pytorch/pytorch#136416 to support pass-in of kwargs.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 9514b54 with merge base 8d01d9b ():
💚 Looks good so far! There are no failures yet. 💚

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

[lessw2020]

syntax error - this should be 'decoder' and not 'decorder'.

[lessw2020]

same, syntax error - this should be 'decoder' and not 'decorder'.

[lessw2020]

same, syntax error - this should be 'decoder' and not 'decorder'.

[lessw2020]

last one, syntax error - this should be 'decoder' and not 'decorder'.

[lessw2020]

nice addition!
minor note that 'decorder' should be 'decoder' in the code for ease of understanding/syntax.