[aroshanghias-nvd]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Complete the packed-MTP boundary correction from #6912 for CP>1. Padded physical boundaries locate each rank's chunks, but logical sequence lengths determine whether a shifted destination has a real next token. The CP path currently keeps padding-origin values and default all-ones mask entries alive past the logical end.

For seven logical tokens in eight physical slots, one left shift of an all-ones mask currently produces [1,1,1,1,1,1,1,0]; the valid result is [1,1,1,1,1,1,0,0]. Nonzero padding similarly leaks into the last logical target. This adds 13 lines to zero invalid destinations in each rank's two physical chunks, preserving the existing physical-indexing and CP1 corrections.

Related: #6912. This is a follow-up for logical-end validity, not a duplicate physical-boundary fix.

Validation

• Added a regression covering CP1/2/4, padded/unpadded multiple sequences, tiny logical lengths, two tensor rows, float/bool masks, integer IDs, zero/nonzero padding, six successive shifts, sum enabled/disabled, and input/metadata immutability. Assertion outcomes are coordinated across ranks so a failing base cannot strand peers in communication.
• Native CUDA/NCCL execution of the full imported module on four GB200 GPUs with Torch 2.11.0+cu130 and pytest 8.4.2: unchanged base 12 expected failures / 42 passes; exact patch 54 passes, on every rank.
• Broader MTP selection: 94 passes (including those 54), 3 failures, 4 skips. The three embedding-scatter failures reproduce on the unchanged base: Torch LayerNorm rejects sequence parallelism during construction, before rolling. The four skipped HSM cases require eight ranks.
• Additional native process_mtp_loss checks passed all four configurations: default masks, derived labels, five MTP depths, both calculate_per_token_loss settings, and TP2/CP2 or TP1/CP4, including CP groups [0,2] and [1,3]. Labels and auxiliary gradients matched an independent logical-sequence reference. These used synthetic output/loss callbacks, not a full-model training run.
• Black 26.3.0, AST scope checks, and git diff --check pass.

Native validation used base f6c33bde4cdc10805ee7157db746d47b758f5f02. The PR was then based on main 4fe0daffe4fc45efc231efc7b5bfc3018d81d516; both changed files were identical between those bases, and their final bytes remain identical to the natively tested files. The full newer tree has not been rerun. The test harness bypassed only the unrelated autouse dataset-download fixture.

Patched module SHA256: a0749e517ae82c443271b329783fcdd4c0fac15df7a29235ea4fc750d4ab115c.
Test file SHA256: 5b43d53e0ee47ccb1f28d1313a98eef9e094962b38870c87b35fe402a119a1e1.

The additional per-sequence allocation/masking cost has not been benchmarked. This PR makes no training-quality or TMPE-remediation claim.

Review readiness

The current head is ee7b0088816963f2d712420413b1598cc318b892. A delegated code review covered every added line and the surrounding roll/caller contracts; no further source correction was identified. The exact source hashes, raw per-rank native XML, and caller receipts were rechecked. Targeted Black 26.3.0, AST parsing, git diff --check, and an independent integer coordinate audit passed again.

Repository CI has not run on this head: only DCO/DCO gate and reviewer-notification checks are successful. The runner-vetting bot requests additional validation. GitHub reports the SSH signature as unverified (unknown_key), and no upstream pull-request/7278 mirror was present at the readiness check. These observations do not establish the sole cause of missing CI. A trusted vetter can approve the exact head with /ok to test ee7b0088816963f2d712420413b1598cc318b892. Marking the PR ready begins review/CI routing; it does not certify passing CI or merge readiness.

Pre-checks

• Added relevant unit tests.
• Added full-model functional tests (not included; native caller checks described above).
• Ran the full repository autoformatter (targeted Black and diff checks passed).

[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.

[aroshanghias-nvd]

@Phlip79 @Mellonta This is ready for technical review at ee7b0088816963f2d712420413b1598cc318b892. Every added line and the surrounding packed-roll/caller contracts have been reviewed; no further source changes were identified. Both changed files still match the exact natively tested bytes: 54/54 targeted cases passed on every rank, versus 12 failures on the unchanged base. Caller checks passed TP2/CP2 and TP1/CP4 with five MTP depths and both per-token loss settings. The existing three LayerNorm failures reproduce on the base; four larger-world tests were skipped.

Repository CI has not run on this head, and copy-pr-bot requests runner vetting. Could a trusted vetter approve the exact head with /ok to test ee7b0088816963f2d712420413b1598cc318b892? GitHub reports the existing SSH signature as unverified (unknown_key); the source commit does not need a rewrite to request explicit-SHA vetting. There is no full-model validation or performance benchmark, and this PR makes no training-quality or TMPE-remediation claim.

[Mellonta]

@aroshanghias-nvd btw, where did you see this issue? is there a concrete use case?