[QUESTION] In-depth explanation and verification of roll_tensor logic for Context Parallelism in MTP

I'm currently studying the Multi-Token Prediction (MTP) implementation, specifically the `roll_tensor` function located in `megatron/core/transformer/multi_token_prediction.py`. I have a question regarding its logic when operating with Context Parallelism (`CP > 1`).

To understand its behavior, I manually traced the function with a concrete example, and the result was counter-intuitive. I would be very grateful if you could verify my trace and clarify the underlying logic.

**My Trace of `roll_tensor(shifts=-1)` with CP=2:**

Let's assume a global sequence `[1, 2, 3, 4, 5, 6, 7, 8]` is split across two CP ranks as follows:
* **CP Rank 0 holds:** `[1, 2, 3, 4]`
* **CP Rank 1 holds:** `[5, 6, 7, 8]`

My goal is to understand the output on each rank after the `roll_tensor` call.

**On CP Rank 0 (local_rank=0):**
1. **Initial Tensor:** `[1, 2, 3, 4]`
2. **Chunking:** `[1, 2]` and `[3, 4]`
3. **Local Roll:** `[2, 1]` and `[4, 3]`
4. **Send/Recv Prep:** `tensor_send_list` is `[1, 3]`. Rank 0 will send `3` (from the 2nd chunk) to Rank 1 and expects to receive a value from Rank 1 to fill the end of its 1st chunk.
5. **Communication:** Let's assume Rank 0 receives `5` from Rank 1.
6. **Splicing:**
* The 1st chunk `[2, 1]` becomes `[2, 5]`.
* The 2nd chunk `[4, 3]` becomes `[4, 0]` (since for `rank=0`, `tensor_recv_list[1]` is hardcoded to 0).
7. **Final Concatenation on Rank 0:** `torch.cat([2, 5], [4, 0])` results in **`[2, 5, 4, 0]`**.

**On CP Rank 1 (local_rank=1, the last rank):**
1. **Initial Tensor:** `[5, 6, 7, 8]`
2. **Chunking:** `[5, 6]` and `[7, 8]`
3. **Local Roll:** `[6, 5]` and `[8, 7]`
4. **Send/Recv Prep:** `tensor_send_list` is `[5, 7]`. Rank 1 will send `5` (from the 1st chunk) to Rank 0 and expects to receive a value from Rank 0 to fill the end of its 2nd chunk.
5. **Communication:** Let's assume Rank 1 receives `3` from Rank 0. The special rule for the last rank also applies: `tensor_recv_list[0] = tensor_send_list[1]`, so `tensor_recv_list[0]` becomes `7`.
6. **Splicing:**
* The 1st chunk `[6, 5]` becomes `[6, 7]`.
* The 2nd chunk `[8, 7]` becomes `[8, 3]`.
7. **Final Concatenation on Rank 1:** `torch.cat([6, 7], [8, 3])` results in **`[6, 7, 8, 3]`**.

**My Question:**

The traced outputs (`[2, 5, 4, 0]` on Rank 0 and `[6, 7, 8, 3]` on Rank 1) are not what I would intuitively expect from a simple distributed roll. My intuitive expectation for the final state would be `[2, 3, 4, 5]` on Rank 0 and `[6, 7, 8, 0]` on Rank 1.

This discrepancy strongly suggests that the data layout from `get_batch_on_this_cp_rank()` is more complex than a simple sequential split.

Could you please confirm if my trace is correct? If so, could you explain the data layout that `roll_tensor` is designed to work with, which leads to this specific behavior?

Finally, what is the purpose of splitting the tensor into two chunks (`tensor.chunk(2, dim=dims)`) at the beginning of the function? Is it related to handling micro-batches?

Thank you for your time and for helping clarify this complex but important detail!