[codex] Preserve real and padded THD sequence lengths

[xiaoyao0115]

What does this PR do?

This PR fixes THD tail-padding metadata in two cases:

1. Dummy sequence padding: preserves existing gaps between valid and padded sequence boundaries when appending the dummy sequence.
2. Non-dummy padding: keeps cu_seqlens unchanged and extends cu_seqlens_padded[-1] to the actual padded tensor size.

For non-dummy padding, tensors are padded before context-parallel slicing to ensure the correct token distribution.

The default dummy-sequence behavior remains unchanged.

Why?

Using physical padded boundaries as valid-token boundaries makes padding appear to be real data and prevents Transformer Engine from inferring gaps between sequences correctly.

Impact

THD attention still receives physical storage boundaries, while routing and loss paths retain the real token boundaries. The default dummy-sequence behavior remains unchanged. The extend_last policy is opt-in.

The sequence-packing test now validates deterministic metadata locally on every rank and checks padding presence across the CP group. This avoids incorrectly comparing independent DP replicas or requiring every CP slice to contain padding.

Validation

• BASE_REF=dev CHECK_ONLY=true SKIP_DOCS=true bash tools/autoformat.sh
• tests/unit_tests/test_sequence_packing.py: 29 passed
• tests/unit_tests/transformer/test_thd_cuda_graph.py: 28 passed, 2 skipped pending the required TE change
• tests/unit_tests/transformer/moe/test_token_dispatcher.py: 19 passed, 92 skipped by configuration

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HaochenYuan]

Suggestion: replace pad_packed_seq_by_appending_dummy_seq: bool with an explicit tail-policy enum:

--thd-tail-padding-policy {
    extend_last,
    zero_valid_dummy,
    local_hidden_only
}
# - extend_last (default): extend only cu_seqlens_padded[-1]; keep cu_seqlens unchanged.  
# - zero_valid_dummy: append the original real endpoint to cu_seqlens and the target endpoint to cu_seqlens_padded. This preserves the original sequence partition but consumes one sequence slot.  
# - local_hidden_only (experimental): pad local hidden after CP slicing without changing metadata. Restrict this to P2P and tail-safe implementations such as the DSV4 custom RoPE; reject A2A and generic TE fused THD RoPE.

This avoids silently changing the meaning of the existing boolean option and makes the trade-offs explicit

[xiaoyao0115]

/ok to test d314db2

[xiaoyao0115]

/ok to test 101e3d4

[xiaoyao0115]

/ok to test 13ecbe2

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Enterprise

Run ID: 52284b83-3b82-4454-994f-985481be4722

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[xiaoyao0115]

/ok to test 05e6f6d

[xiaoyao0115]

/ok to test 93daff4

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30519040434

[xiaoyao0115]

/ok to test b4b68ba

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30526345696