feat(distributed): block-diagonal varlen CP for packed sequences

[yuhezhang-ai]

Summary

Adds nemo_automodel/components/distributed/blockdiag_cp/, a self-contained context-parallel
implementation for packed (multi-document) sequences, where correct masking is
block-causal per document. Softmax attention runs as a FlashAttention/TransformerEngine
varlen kernel driven by per-step precomputed cu_seqlens (zero per-layer host syncs),
with a dense masked-SDPA fallback that is numerically identical. K/V delivery across the CP
group is a fused differentiable [K;V] all-gather by default, with opt-in needed-only
exchanges (left-neighbor halo p2p, or all-to-all-v when a document spans more than two
ranks) that cut per-rank K/V from O(S) to O(S/cp) + straddle.

Part of #2985.

Scope and follow-up

This PR intentionally lands the standalone block-diagonal CP primitive and its correctness
tests. It does not modify any model, recipe, or example config. A follow-up Amber upstream
PR will wire it into the Qwen3.5 VLM path using a real configuration with context
parallelism and packed sequences, building on this PR and #2990. That integration will
attach make_cp_blockdiag_batch_and_ctx as the model's _cp_make_batch_fn and route its
attention through cp_blockdiag_sdpa.

Motivation

Sequence packing is the standard way to train VLMs/LLMs on mixed-length documents without
padding waste, and long-context packed training needs CP. The stock DTensor
context_parallel path (ring attention over load-balanced shards) assumes one causal
document per sequence; a packed sequence attended that way leaks attention across document
boundaries. TE's THD path covers input_ids-based LLM batches but not pre-embedded VLM
batches (inputs_embeds, deepstack visual embeds, mRoPE position ids).

This PR provides CP that:

• keeps every document's attention exactly block-causal (bit-parity with the non-CP mask),
• shards pre-embedded VLM batches (differentiable inputs_embeds slice, so gradients flow
  back to the trainable vision tower),
• stays collective-safe under activation-checkpoint recompute and degenerate layouts (a
  rank whose whole chunk is padding must still participate in every collective — see the
  gradient-attachment tests),
• and reuses the model-owned CP hook introduced with the existing CP-for-VLM
  infrastructure (_cp_make_batch_fn in cp_utils.make_cp_batch_and_ctx), so no changes
  to cp_utils are needed.

Design overview

blockdiag_cp/
  state.py     knob normalization (attn_backend: flash|te|dense, kv_exchange:
               allgather|halo|a2a) + AC-recompute-safe per-step state
  batch.py     make_cp_blockdiag_batch_and_ctx: pad to cp multiple, contiguous
               per-rank sequence slice of every seq-aligned tensor, per-step
               varlen-metadata precompute, train context
  runtime.py   cp_blockdiag_sdpa: drop-in SDPA; selects the KV path per step
               (with named downgrade reasons), runs varlen or dense attention;
               collective preflight + post-kernel consensus so no rank can
               diverge in collective order
  kernels.py   varlen segmentation + host-side validation of every index a
               varlen CUDA kernel consumes; flash long-left-prefix guard;
               cached TE DotProductAttention; dense [B,1,L,S] mask builder
  exchange.py  differentiable collectives: fused [K;V] all-gather (backward =
               reduce-scatter SUM), left-halo neighbor p2p, needed-only
               all-to-all-v with index_add backward
  packed.py    cp_size==1 degenerate path: same varlen kernel for packed
               sequences on one rank (row_offset=0), armed per forward and
               stable across AC recompute

Key correctness properties, each pinned by a test:

1. Parity: per-rank block-diagonal attention concatenated over ranks equals full
   attention under the block-causal mask (fp32 CPU, and bf16 flash/TE vs dense on GPU).
2. Boundary documents: a document straddling a rank boundary yields an asymmetric
   (K > Q) first varlen segment with bottom-right causal alignment; a guard peels it
   into a fixed-shape Flash call to avoid a known varlen failure mode on long prefixes.
3. Collective safety: all-padding ranks keep their output attached to the gathered
   K/V (0-weighted touch) so backward reduce-scatter fires symmetrically; needed-only
   paths reach an all-reduce consensus before and after the kernel so every rank raises
   together instead of forking the collective order.
4. Kernel-input hygiene: every cu_seqlens/offset consumed by a varlen kernel is
   validated on the host once per step (FlashAttention trusts these; a bad terminal
   offset is an async illegal memory access, not a Python exception).

Proof

• tests/unit_tests/distributed/test_blockdiag_cp.py (CPU): mask construction,
  CP-vs-full parity for world 2/4 incl. GQA, all-gather fwd/bwd, all-pad-rank gradient
  attachment, metadata validation accept/reject, halo/a2a plan geometry, KV path
  downgrade reasons, long-prefix guard, dropout forwarding, 1-D document ids, and the
  cp1 packed hook contract.
• tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py (1 GPU, skipped
  without CUDA): flash and TE varlen vs dense-mask SDPA across multi-doc packs with
  padding tails, GQA (incl. 8:1, head_dim 256), rank-straddling docs, single doc
  spanning all ranks, heavy padding; bit-identity of precomputed metadata vs inline
  segmentation; the cp_size==1 packed hook; native Flash dropout and TE's correct
  dense block-diagonal dropout fallback.
• tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py (1 GPU,
  skipped without CUDA): production-sized FlashAttention boundary shapes under
  checkpointed forward replay and backward.
• tests/functional_tests/context_parallel/test_blockdiag_cp_parity.py (2 GPUs,
  scheduled L2): real NCCL forward/backward parity for dense, FlashAttention, and TE
  across all-gather, halo, and all-to-all-v exchange. It checks local outputs, input
  gradients, and every Q/K/V/output-projection parameter gradient, including a document
  straddling the rank boundary and an all-padding rank.

The implementation has been validated in internal large-scale runs (up to 128 GPUs),
including long-context (32k+) packed VLM training with FSDP2 and multi-node CP, where
loss curves match the cp_size=1 baseline.

Notes

• The needed-only (halo/a2a) exchanges default off (kv_exchange="allgather") and
  auto-downgrade with a logged reason (non-varlen backend, missing metadata, or a CP
  group spanning nodes; NEMO_CP_ALLOW_XNODE=1 overrides for validation).

🤖 Generated with Claude Code

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

[yuhezhang-ai]

/ok to test abb90e2

[yuhezhang-ai]

/ok to test dff34c2

[yuhezhang-ai]

/ok to test 1e3b7c9