feat(laguna): support packed THD context parallelism

[akoumpa]

What does this PR do ?

Add native packed THD plus context-parallel training support for Laguna, including full and sliding-window attention semantics.

This is stacked on the Laguna XS 2.1 recipe PR so this diff contains only the THD/CP implementation and validation recipe.

Changelog

• Add Laguna-native THD forward handling and CP capability declarations.
• Route packed CP through block-diagonal SDPA while preserving document and sliding-window boundaries.
• Preserve the framework native TE THD input sharder for CP1 baselines.
• Extend block-diagonal CP metadata and masks for THD lengths and local attention windows.
• Add CPU float32 logits/gradient parity coverage and CP mask/runtime tests.
• Add an EP8 + CP2 + 1K THD Laguna XS validation recipe and update documentation.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Validation

• CPU float32 packed/per-document logits and all-parameter gradient parity passed for full and sliding-window attention.
• 48 Laguna/block-diagonal CP tests passed.
• 104 capability and mesh-validation tests passed.
• 27 docs/recipe coverage checks passed; 1 skipped.
• Ruff and import-linter passed.
• cw-dfw Slurm job 16473638: 100-step, eight-GPU EP8 + CP2 run completed.
• W&B CP2 run: https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/905zourb
• Train loss: 2.9746 at step 0 to 2.1269 at step 99.
• Validation loss: 2.8546 at step 49 to 2.6684 at step 99.
• Peak memory: 42.90 GiB.

Controlled CP1/CP2 loss parity

• cw-dfw Slurm job 16481738: both matched 100-step runs completed on the same eight-GPU allocation.
• Both runs use the same frozen 3,392-pack THD dataset, global batch size 8, seed, checkpoint, optimizer, and schedule. Dataset SHA-256: 6bf1bf3f0a969ad6a7120e980d0a11c452be89814396c25386962fca6d973e56.
• All 100 per-step label-token counts match exactly.
• CP1 W&B (native TE THD): https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/0y2cigaj
• CP2 W&B (block-diagonal SDPA THD): https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/5rgo3ay8
• Mean absolute loss difference: 0.00946 (0.44%); maximum: 0.02258 (0.92%, step 0); RMSE: 0.01039; Pearson r: 0.999795.
• Step 0: CP1 2.95205, CP2 2.97463; step 99: CP1 2.10886, CP2 2.12573.
• Peak memory: CP1 39.02 GiB; CP2 42.90 GiB.
• The attention backend is the necessary controlled difference: TE handles packed THD at CP1, while Laguna uses block-diagonal SDPA for packed sliding-window CP2 because TE does not support that combination.

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

[akoumpa]

/ok to test ea7baf2

[github-actions]

🌿 Preview your docs: https://nvidia-preview-preview-0ee3a52b69d9.docs.buildwithfern.com/nemo/automodel

[akoumpa]

cw-dfw matched loss parity completed on Slurm job 16481738. Across 100 identical global batches, mean absolute loss difference is 0.00946 (0.44%), max is 0.02258 (0.92%), Pearson r is 0.999795, and all label-token counts match.

• CP1: https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/0y2cigaj
• CP2: https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/5rgo3ay8

[akoumpa]

/ok to test a0c66d6

[akoumpa]

Conflict repair + validation update

Rebased the two THD+CP commits onto current main after #3639 merged. The PR is now conflict-free at head 10e70902c1995fdcfcc9634137bdf7e59b11357d.

The matched cw-dfw validation remains:

• Slurm job: 16481738
• Frozen 3,392-pack dataset SHA-256: 6bf1bf3f0a969ad6a7120e980d0a11c452be89814396c25386962fca6d973e56
• CP=1 W&B: https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/0y2cigaj
• CP=2 W&B: https://wandb.ai/Nemo-automodel/laguna-xs-2p1-sft/runs/5rgo3ay8
• 100 identical global batches; all label-token counts match
• Mean absolute loss delta: 0.00946 (0.44%)
• Max absolute loss delta: 0.02258 (0.92%)
• RMSE: 0.01039; Pearson r: 0.999795
• Step 0: CP=1 2.95205, CP=2 2.97463
• Step 99: CP=1 2.10886, CP=2 2.12573

Post-rebase local checks: 48 Laguna/block-diagonal CP tests passed; Ruff passed; recipe/model-coverage tests passed (5 passed, 1 skipped).

[akoumpa]

/ok to test 10e7090

[jgerh]

Completed tech pubs review of docs/model-coverage/llm/poolside/laguna.mdx and provided a few copyedits

[akoumpa]

/ok to test 0ee3a52

[akoumpa]

/ok to test 0ee3a52