From d48d6f62ee7b089b27745a64f8e3a19a64c77c62 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Thu, 9 Jul 2026 08:44:39 -0700 Subject: [PATCH 1/8] feat(distributed): add block-diagonal varlen context parallelism for packed sequences Add a self-contained CP implementation for packed (multi-document) sequences where masking must stay block-causal per document, which the load-balanced DTensor context_parallel path cannot express: - batch: contiguous sequence sharding + per-step train context, pluggable via the model-owned _cp_make_batch_fn hook in cp_utils.make_cp_batch_and_ctx - runtime: drop-in SDPA that exchanges K/V across the CP group and runs per-document causal attention on local queries, with collective-safe path selection and fallback consensus across ranks - kernels: FlashAttention/TransformerEngine varlen kernels driven by per-step precomputed cu_seqlens (zero per-layer host syncs), a dense masked-SDPA fallback, host-side validation of all kernel indices, and a long-left-prefix guard for boundary documents - exchange: differentiable K/V collectives -- fused [K;V] all-gather with reduce-scatter backward (default), plus opt-in needed-only delivery via left-neighbor halo p2p or all-to-all-v for documents spanning >2 ranks - packed: the cp_size==1 degenerate path that routes stock SDPA calls of a packed sequence through the same varlen kernel - state: knob normalization and activation-checkpointing-safe step state (readable from the autograd recompute worker thread) Co-Authored-By: Claude Fable 5 Signed-off-by: Yuhe Zhang --- .../distributed/blockdiag_cp/__init__.py | 69 ++ .../distributed/blockdiag_cp/batch.py | 281 ++++++++ .../distributed/blockdiag_cp/exchange.py | 374 ++++++++++ .../distributed/blockdiag_cp/kernels.py | 657 ++++++++++++++++++ .../distributed/blockdiag_cp/packed.py | 197 ++++++ .../distributed/blockdiag_cp/runtime.py | 362 ++++++++++ .../distributed/blockdiag_cp/state.py | 206 ++++++ 7 files changed, 2146 insertions(+) create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/__init__.py create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/batch.py create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/exchange.py create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/kernels.py create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/packed.py create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/runtime.py create mode 100644 nemo_automodel/components/distributed/blockdiag_cp/state.py diff --git a/nemo_automodel/components/distributed/blockdiag_cp/__init__.py b/nemo_automodel/components/distributed/blockdiag_cp/__init__.py new file mode 100644 index 0000000000..a103623aa7 --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/__init__.py @@ -0,0 +1,69 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Block-diagonal (per-document) varlen context parallelism for packed sequences. + +Packed-sequence training concatenates many documents into one long sequence; a +correct attention mask is block-causal per document. The stock DTensor +``context_parallel`` path assumes a single causal document, so packed VLM/LLM +batches need a CP implementation that reshards by contiguous chunks and rebuilds +per-document masking on every rank. + +This package provides that implementation, split by responsibility: + +- :mod:`.state` -- runtime knob normalization + activation-checkpoint-safe step state. +- :mod:`.kernels` -- dense-mask and varlen (FlashAttention / TransformerEngine) kernels. +- :mod:`.exchange` -- differentiable K/V collectives (all-gather, left-halo, all-to-all-v). +- :mod:`.runtime` -- the SDPA entry point and collective-safe path selection. +- :mod:`.batch` -- batch padding/sharding and the per-step train context. +- :mod:`.packed` -- the cp_size==1 packed-sequence varlen SDPA hook. + +Integration follows the model-owned CP convention of +:func:`nemo_automodel.components.distributed.cp_utils.make_cp_batch_and_ctx`: a model +attaches :func:`make_cp_blockdiag_batch_and_ctx` to the batch as ``_cp_make_batch_fn`` +and routes its softmax attention through :func:`cp_blockdiag_sdpa` while the returned +context is active. +""" + +from nemo_automodel.components.distributed.blockdiag_cp.batch import make_cp_blockdiag_batch_and_ctx +from nemo_automodel.components.distributed.blockdiag_cp.kernels import precompute_blockdiag_varlen_meta +from nemo_automodel.components.distributed.blockdiag_cp.packed import ( + cp1_packed_varlen_backend, + disable_cp1_packed_varlen, + enable_cp1_packed_varlen, +) +from nemo_automodel.components.distributed.blockdiag_cp.runtime import cp_blockdiag_sdpa +from nemo_automodel.components.distributed.blockdiag_cp.state import ( + configure_cp_varlen, + cp_attn_fire_count, + cp_varlen_runtime_config, + normalize_attn_backend, + normalize_kv_exchange, + reset_cp_attn_fire_count, +) + +__all__ = [ + "configure_cp_varlen", + "cp_attn_fire_count", + "cp_varlen_runtime_config", + "cp_blockdiag_sdpa", + "cp1_packed_varlen_backend", + "disable_cp1_packed_varlen", + "enable_cp1_packed_varlen", + "make_cp_blockdiag_batch_and_ctx", + "normalize_attn_backend", + "normalize_kv_exchange", + "precompute_blockdiag_varlen_meta", + "reset_cp_attn_fire_count", +] diff --git a/nemo_automodel/components/distributed/blockdiag_cp/batch.py b/nemo_automodel/components/distributed/blockdiag_cp/batch.py new file mode 100644 index 0000000000..a7791515fb --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/batch.py @@ -0,0 +1,281 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Batch padding, sequential sharding, and block-diagonal CP context setup.""" + +from __future__ import annotations + +import contextlib +from typing import Any, Callable, ContextManager + +import torch +from torch.distributed.device_mesh import DeviceMesh + +from nemo_automodel.components.distributed.blockdiag_cp import kernels +from nemo_automodel.components.distributed.blockdiag_cp import state as state_module + + +def _cp_blockdiag_doc_ids(batch: dict, seq_len: int, device, batch_size: int) -> torch.Tensor: + """Resolve per-position document ids ``[B, S]`` (0 == padding) for the mask. + + Prefers the collator's ``_packed_seq_ids`` (1-based document index per token, + present when a pack holds >1 document). Otherwise falls back to the 4-D + block-causal ``attention_mask`` diagonal (valid positions) or, lacking both, + treats the whole sequence as a single document. + + Args: + batch: The training batch; may contain ``_packed_seq_ids`` ``[B, S]`` + (int document index per token) or ``attention_mask`` (``[B, 1, S, S]`` + block-causal bool, or ``[B, S]`` validity/indexed mask). + seq_len: ``S``, the (unpadded) sequence length. + device: Device for the returned tensor. + batch_size: ``B``, used for the all-ones fallback. + + Returns: + Per-position document ids ``[B, S]`` (int64, 0 == padding). + """ + seq_ids = batch.get("_packed_seq_ids", None) + if seq_ids is not None: + return seq_ids.to(device=device, dtype=torch.long) + attn = batch.get("attention_mask", None) + if attn is not None and attn.dim() == 4: + # [B, 1, S, S] block-causal bool -> diagonal gives per-position validity. + valid = attn[:, 0].diagonal(dim1=-2, dim2=-1) # [B, S] + return valid.to(device=device, dtype=torch.long) + if attn is not None and attn.dim() == 2: + # [B, S] standard validity mask (0/1 or bool) or indexed packing ids. + return attn.to(device=device, dtype=torch.long) + return torch.ones(batch_size, seq_len, device=device, dtype=torch.long) + + +def make_cp_blockdiag_batch_and_ctx( + cp_mesh: DeviceMesh, + tp_mesh: DeviceMesh | None, + batch: dict[str, Any], + *, + loss_mask: torch.Tensor | None = None, + padding_token_id: int = 0, +) -> tuple[Callable[[], ContextManager], dict[str, Any]]: + """Sequentially shard a pre-embedded batch for block-diagonal CP. + + Pads the sequence to a multiple of the CP world size, slices each + sequence-aligned tensor to this rank's contiguous chunk (a differentiable + slice for ``inputs_embeds``, so gradients flow back to the trainable vision + tower / embedding table), and returns ``(train_ctx, batch)`` where entering + ``train_ctx`` activates the per-document CP SDPA state for the step. + + Softmax attention must route through + :func:`~nemo_automodel.components.distributed.blockdiag_cp.runtime.cp_blockdiag_sdpa` + while this context is active. A model opts in by attaching this callable to + the batch as ``_cp_make_batch_fn`` (the model-owned CP hook honored by + :func:`nemo_automodel.components.distributed.cp_utils.make_cp_batch_and_ctx`) + and rebinding its attention's SDPA call for the step. + + Args: + cp_mesh: The context-parallel device (sub)mesh. + tp_mesh: Accepted for ``_cp_make_batch_fn`` signature compatibility; + unused (block-diagonal CP shards only the sequence dimension). + batch: The training batch. Must contain pre-embedded ``inputs_embeds`` + ``[B, S, H]`` (multimodal token replacement happens pre-shard) and is + mutated in place: ``attention_mask`` is dropped, ``padding_mask`` + ``[B, S]`` (bool, True == pad) is added, and every sequence-aligned + tensor is padded then sliced to this rank's ``[row_offset, + row_offset + S_full/cp)`` chunk. + loss_mask: Optional per-token loss mask ``[B, S]``; padded with 0 and + sharded like the other sequence-aligned tensors (stored back into + ``batch["loss_mask"]``). + padding_token_id: Accepted for signature compatibility; unused (the + batch is pre-embedded, so there are no token ids to pad). + + Returns: + ``(train_ctx, batch)``: a zero-arg callable returning the per-step + context manager, and the sharded batch. + """ + from contextlib import nullcontext + + from torch.nn.attention import SDPBackend, sdpa_kernel + + world = cp_mesh.size() + if world <= 1: + return nullcontext, batch + + rank = cp_mesh.get_local_rank() + group = cp_mesh.get_group() + + assert "inputs_embeds" in batch, "block-diagonal CP requires pre-embedded 'inputs_embeds' in the batch" + + ie = batch["inputs_embeds"] + B, S = ie.shape[0], ie.shape[1] + device = ie.device + + # Resolve document ids BEFORE dropping attention_mask: for single-document / + # text-only packs ``_packed_seq_ids`` is absent and the per-position validity + # (real vs padding) comes from the 4D block-causal ``attention_mask`` diagonal. + doc_ids = _cp_blockdiag_doc_ids(batch, S, device, B) # [B, S] on full (unpadded) seq + + # The per-rank block-diagonal mask is rebuilt inside the CP SDPA from + # ``doc_ids``; drop the (now stale, full-length) 4D mask so the model's own + # mask machinery does not fire on the sharded local sequence. + batch.pop("attention_mask", None) + + # Preserve the padding signal for MoE routing / load-balance statistics after + # the full attention_mask is dropped (0 == padding in doc_ids). Built at the + # unpadded length S so it shards through ``_shard`` like every other + # sequence-aligned tensor -- the CP-pad tail is filled with + # ``PAD_FILL["padding_mask"]=True`` there, matching ``doc_ids.eq(0)`` on the pad. + batch["padding_mask"] = doc_ids.eq(0) + + if loss_mask is not None: + batch["loss_mask"] = loss_mask + + pad_len = (-S) % world + if pad_len: + doc_ids = torch.cat( + [doc_ids, torch.zeros(B, pad_len, device=device, dtype=doc_ids.dtype)], + dim=1, + ) + S_full = S + pad_len + local_len = S_full // world + row_offset = rank * local_len + + # Defense-in-depth (do not remove): the per-document SDPA all-gathers K/V over + # ``group`` (``exchange._AllGatherSeqDiff``) to rebuild the full sequence, while + # the sequence is sharded into ``world == cp_mesh.size()`` chunks above. The + # all-gather group MUST therefore have exactly ``world`` ranks. On torch 2.8, + # ``device_mesh["cp"].get_group()`` was mis-resolved on a mesh that carries a + # flattened ``dp_shard_cp`` dim: with dp>1 it returned a dp*cp-sized group, so + # the all-gather over-gathered -> ``key_full`` was dp* too long -> a cryptic + # shape crash deep in the backward AC-recompute. Later torch versions resolve + # it correctly. Fail loud and early here instead of crashing in backward. + _group_world = torch.distributed.get_world_size(group) + assert _group_world == world, ( + f"block-diagonal CP: K/V all-gather group world ({_group_world}) != " + f"cp_mesh.size() ({world}). The cp sub-group is mis-resolved -- a DeviceMesh " + f"flatten/slice bug seen on torch 2.8 when dp>1; upgrade torch." + ) + + # Block-diagonal CP shards ONE packed sequence per rank and assumes local batch + # B==1 (packing collapses many samples into a single sequence). B>1 is + # unsupported: the deepstack visual-embed sharding below indexes batch row 0 + # (``vpm[0, ...]``), so B>1 would die with a cryptic size mismatch deep in the + # model. Fail loud here. + _cp_bsz = batch["inputs_embeds"].shape[0] + if _cp_bsz != 1: + raise ValueError( + f"block-diagonal context parallelism requires local_batch_size=1 (one packed " + f"sequence per rank), got batch size {_cp_bsz}. Enable packing and set " + f"the local batch size to 1." + ) + + # Per-tensor padding fill: each tensor's "ignore" value is semantic, not + # dtype-derived (mirrors make_cp_batch_and_ctx's PAD_FILL). + PAD_FILL = { + "labels": -100, + "_packed_seq_ids": 0, + "loss_mask": 0, + "padding_mask": True, + "visual_pos_masks": False, # deepstack VLMs: pad positions are non-visual + } + + def _shard(key: str, tensor: torch.Tensor) -> torch.Tensor: + """Pad ``tensor`` on its sequence dim then slice this rank's contiguous chunk. + + Args: + key: Batch key, selects the pad fill sentinel and the sequence dim + (``position_ids`` may be ``[3, B, S]`` -> seq dim 2; else dim 1). + tensor: A sequence-aligned batch tensor (``[B, S, ...]`` or + ``[3, B, S]`` for mRoPE position ids). + + Returns: + The local shard ``[..., local_len, ...]`` along the sequence dim. + """ + seq_dim = 2 if (key == "position_ids" and tensor.dim() == 3) else 1 + if pad_len: + pad_shape = list(tensor.shape) + pad_shape[seq_dim] = pad_len + if tensor.dtype.is_floating_point: + fill = torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device) + else: + fill = torch.full( + pad_shape, + PAD_FILL.get(key, 0), + dtype=tensor.dtype, + device=tensor.device, + ) + tensor = torch.cat([tensor, fill], dim=seq_dim) + sl = [slice(None)] * tensor.dim() + sl[seq_dim] = slice(row_offset, row_offset + local_len) + return tensor[tuple(sl)] + + # Deepstack VLMs (e.g. Qwen3-VL): ``_deepstack_visual_embeds`` is a LIST of + # [n_visual, H] tensors indexed by visual token (sequence order), NOT + # sequence-aligned. Slice each to this rank's contiguous visual-token range so + # the per-shard deepstack merge (hidden_states[visual_pos_masks] += embeds) + # sees matching counts. Computed from the FULL (padded) visual_pos_masks + # before it is sharded. + if "_deepstack_visual_embeds" in batch and "visual_pos_masks" in batch: + vpm = batch["visual_pos_masks"] # [B, S] bool (packing -> B == 1) + if pad_len: + vpm = torch.cat( + [vpm, torch.zeros(vpm.shape[0], pad_len, dtype=vpm.dtype, device=vpm.device)], + dim=1, + ) + v_start = int(vpm[0, :row_offset].sum()) + n_local = int(vpm[0, row_offset : row_offset + local_len].sum()) + batch["_deepstack_visual_embeds"] = [d[v_start : v_start + n_local] for d in batch["_deepstack_visual_embeds"]] + + seq_aligned = ( + "inputs_embeds", + "labels", + "position_ids", + "_packed_seq_ids", + "loss_mask", + "padding_mask", + "visual_pos_masks", + ) + for key in seq_aligned: + if key in batch and isinstance(batch[key], torch.Tensor): + batch[key] = _shard(key, batch[key]) + + runtime_config = state_module.cp_varlen_runtime_config() + step_state = { + "group": group, + "doc_ids": doc_ids, + "row_offset": row_offset, + "seq_dim": 2, + # Per-step snapshot of the runtime config, so cp_blockdiag_sdpa selects its + # path from STATE (a function of (qkv, state)) rather than reading module + # globals on the hot path. configure_cp_varlen() remains the config holder; + # this captures it once per step (see _resolve_cp_varlen_config). + "attn_backend": runtime_config["attn_backend"], + "kv_exchange": runtime_config["kv_exchange"], + # Precompute the varlen cu_seqlens once per step (step-constant in + # doc_ids/row_offset/local_len) so the per-layer softmax SDPA does zero + # GPU->CPU .item() syncs. The flash/te varlen path reads this; the dense + # fallback ignores it. + "varlen_meta": kernels.precompute_blockdiag_varlen_meta(doc_ids, row_offset, local_len, device), + } + + @contextlib.contextmanager + def _ctx(): + token = state_module._CP_BLOCKDIAG_STATE.set(step_state) + # EFFICIENT/MATH support arbitrary masks (flash does not); both are correct + # on the local (non-DTensor) tensors produced by the K/V exchange. + with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]): + try: + yield + finally: + state_module._CP_BLOCKDIAG_STATE.reset(token) + + return _ctx, batch diff --git a/nemo_automodel/components/distributed/blockdiag_cp/exchange.py b/nemo_automodel/components/distributed/blockdiag_cp/exchange.py new file mode 100644 index 0000000000..b41692e0c7 --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/exchange.py @@ -0,0 +1,374 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable K/V collectives and needed-only exchange plans.""" + +from __future__ import annotations + +import torch + +from nemo_automodel.components.distributed.blockdiag_cp.kernels import _cp_blockdiag_varlen, _varlen_seg_for_rank + + +class _AllGatherSeqDiff(torch.autograd.Function): + """All-gather a per-rank sequence shard along ``seq_dim`` over ``group``. + + Forward concatenates every rank's shard ``[B, H, L, D]`` in rank order + (sequential sharding), producing the full sequence ``[B, H, L*world, D]``. + Backward reduce-scatters the incoming gradient: each K/V position is read by + every rank's queries, so its gradient is the sum over ranks of the per-rank + grad slice; reduce-scatter(SUM) hands each rank the summed gradient for the + shard it owns. All shards have equal length ``L`` (the sequence is padded to + a multiple of the CP world size before sharding). + """ + + @staticmethod + def forward(ctx, x: torch.Tensor, group, seq_dim: int) -> torch.Tensor: + """All-gather ``x`` (this rank's shard) into the full sequence along ``seq_dim``.""" + ctx.group = group + ctx.seq_dim = seq_dim + world = torch.distributed.get_world_size(group) + ctx.world = world + x = x.contiguous() + gathered = [torch.empty_like(x) for _ in range(world)] + torch.distributed.all_gather(gathered, x, group=group) + return torch.cat(gathered, dim=seq_dim) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + """Reduce-scatter(SUM) the full-sequence gradient back to this rank's shard.""" + chunks = [c.contiguous() for c in grad_out.chunk(ctx.world, dim=ctx.seq_dim)] + local = torch.empty_like(chunks[0]) + torch.distributed.reduce_scatter(local, chunks, op=torch.distributed.ReduceOp.SUM, group=ctx.group) + return local, None, None + + +def _compute_blockdiag_kv_plan(doc_ids_full: torch.Tensor, world: int, local_len: int, dev) -> dict: + """Global left-halo exchange plan, computed from replicated ``doc_ids``. + + For each rank r the block-diagonal path needs the contiguous K/V range + ``[s_first_r, real_end_r)`` = its boundary document's straddle (owned by ranks < r) + plus its own local chunk. ``back_r = r*local_len - s_first_r`` is the straddle + length. If ``back_r <= local_len`` for ALL ranks, every straddle fits entirely in + the LEFT neighbor (rank r-1), so a single neighbor p2p suffices (``use_halo``). If + any document spans >2 ranks (``back_r > local_len`` -- e.g. one doc across the whole + sequence) the caller falls back to the general all-to-all-v. Every rank computes + this from the same replicated ``doc_ids`` and therefore agrees on ``use_halo`` and + on all send/recv counts without communication. + + Args: + doc_ids_full: Replicated per-position document ids ``[B, S_full]`` (row 0 + used) or ``[S_full]`` (0 == padding) on the full padded sequence. + world: CP world size. + local_len: Per-rank local sequence length ``L`` (``S_full == world * L``). + dev: Device used for the intermediate segmentation tensors. + + Returns: + A plan dict with per-rank lists ``recv`` (tokens received from r-1 == + back_r), ``send`` (tokens sent to r+1 == back_{r+1}; ``send[-1]=0``, + ``recv[0]=0``), ``nreal``, ``s_first``, ``real_end``, and the boolean + ``use_halo``. + """ + dids = doc_ids_full if doc_ids_full.dim() == 1 else doc_ids_full[0] + backs, nreals, s_first, real_end = [], [], [], [] + for rr in range(world): + seg = _varlen_seg_for_rank(dids, rr * local_len, local_len, dev) + if seg is None: + backs.append(0) + nreals.append(0) + s_first.append(rr * local_len) # empty needed range + real_end.append(rr * local_len) + else: + backs.append(rr * local_len - seg["s_first"]) + nreals.append(seg["n_real"]) + s_first.append(seg["s_first"]) + real_end.append(seg["real_end"]) + use_halo = all(b <= local_len for b in backs) + send = backs[1:] + [0] # rank rr sends back_{rr+1} to rr+1 + return { + "use_halo": use_halo, + "recv": backs, + "send": send, + "nreal": nreals, + "s_first": s_first, + "real_end": real_end, + } + + +class _LeftHaloExchange(torch.autograd.Function): + """Uniform-size single-step neighbor exchange for node-local halo mode. + + Every rank sends the last ``halo_size`` tokens of its chunk ``[B, C, L, D]`` to + ``next_peer`` = (rank+1)%world and receives ``halo_size`` tokens from + ``prev_peer`` = (rank-1)%world, send-before-recv, in one ``batch_isend_irecv``. + The caller slices the last ``recv_count`` of the received block as its real + straddle (rank 0's wraparound block and any over-send are ignored). + Differentiable: backward is the reverse exchange (send grad to prev, recv from + next). + """ + + @staticmethod + def forward(ctx, x_local, group, halo_size, prev_peer, next_peer): + """Exchange the chunk suffix; ``x_local`` is ``[B, C, L, D]``, returns ``[B, C, halo_size, D]``.""" + ctx.group = group + ctx.halo_size = halo_size + ctx.prev_peer = prev_peer + ctx.next_peer = next_peer + ctx.shape = tuple(x_local.shape) + B, C, L, D = x_local.shape + send_buf = x_local[:, :, L - halo_size :, :].contiguous() # suffix sent to next + recv_buf = torch.empty(B, C, halo_size, D, dtype=x_local.dtype, device=x_local.device) + ops = [ + torch.distributed.P2POp(torch.distributed.isend, send_buf, next_peer, group=group), + torch.distributed.P2POp(torch.distributed.irecv, recv_buf, prev_peer, group=group), + ] + for req in torch.distributed.batch_isend_irecv(ops): + req.wait() + return recv_buf # [B, C, halo_size, D] from prev rank's chunk suffix + + @staticmethod + def backward(ctx, grad_recv): + """Reverse exchange: route ``grad_recv`` ``[B, C, halo_size, D]`` back to the sending suffix.""" + B, C, L, D = ctx.shape + hs = ctx.halo_size + grad_send = torch.empty(B, C, hs, D, dtype=grad_recv.dtype, device=grad_recv.device) + ops = [ + torch.distributed.P2POp(torch.distributed.isend, grad_recv.contiguous(), ctx.prev_peer, group=ctx.group), + torch.distributed.P2POp(torch.distributed.irecv, grad_send, ctx.next_peer, group=ctx.group), + ] + for req in torch.distributed.batch_isend_irecv(ops): + req.wait() + grad_x = grad_recv.new_zeros(B, C, L, D) + grad_x[:, :, L - hs :, :] = grad_send # grad for the suffix we sent to next + return grad_x, None, None, None, None + + +def _blockdiag_halo_attention(query, key, value, doc_ids, group, plan, gm, row_offset, scale, backend): + """Needed-only block-diagonal CP attention via the left-halo exchange. + + Each rank attends its local queries against ``[boundary-doc halo from rank-1] + + [local chunk]`` instead of the full all-gathered sequence. The caller must fail + fast when this returns ``None``: the halo collective has already executed, so a + rank-local all-gather fallback would fork the CP group's collective order. + + Args: + query: Local query shard ``[B, Hq, L, D]``. + key: Local key shard ``[B, Hkv, L, D]``. + value: Local value shard ``[B, Hkv, L, D]``. + doc_ids: Replicated per-position document ids ``[B, S_full]`` (0 == pad). + group: The CP process group. + plan: Replicated halo plan from :func:`_compute_blockdiag_kv_plan` + (augmented with ``rank``/``world``). + gm: This rank's per-step varlen metadata (global ``s_first``/``real_end``). + row_offset: Global position of this rank's first local query row. + scale: Softmax scale (``None`` -> kernel default). + backend: Varlen kernel backend, ``"flash"`` or ``"te"``. + + Returns: + Attention output ``[B, Hq, L, D]``, or ``None`` when the kernel is + unavailable. + """ + rank = plan["rank"] + world = plan["world"] + recv_count = plan["recv"][rank] + # Uniform halo size across the whole ring (max straddle) so every rank posts + # the same tensor shape. + max_halo = max(plan["recv"]) if plan["recv"] else 0 + + Hkv = key.shape[1] + kv_local = torch.cat((key, value), dim=1) # [B, 2*Hkv, L, D] + if max_halo > 0: + prev_peer = torch.distributed.get_global_rank(group, (rank - 1) % world) + next_peer = torch.distributed.get_global_rank(group, (rank + 1) % world) + recv_block = _LeftHaloExchange.apply(kv_local, group, max_halo, prev_peer, next_peer) # [B,2Hkv,max_halo,D] + halo = recv_block[:, :, max_halo - recv_count :, :] # last recv_count == this rank's real straddle + else: + halo = kv_local[:, :, :0, :] # no straddle anywhere -> empty + kv_needed = torch.cat((halo, kv_local), dim=2) # [B,2Hkv, recv_count+L, D] == global [s_first, (r+1)*L) + key_needed = kv_needed[:, :Hkv] + value_needed = kv_needed[:, Hkv:] + + n_real = gm["n_real"] + if n_real == 0: + # All-pad rank: no real queries. Keep the halo subgraph live (0-weighted touch) + # so its backward fires symmetrically with the neighbor (counts are 0 here, so + # both sides post nothing -- no hang). Output stays zeros. + out = torch.zeros_like(query) + return out + 0.0 * (key_needed.sum() + value_needed.sum()).to(out.dtype) + + # cu_seqlens are relative to s_first already; only the SLICE base changes: key_needed + # IS the [s_first, real_end) slice, so index it at [0, recv+n_real). + local_meta = dict(gm) + local_meta["s_first"] = 0 + local_meta["real_end"] = recv_count + n_real + return _cp_blockdiag_varlen( + query, + key_needed, + value_needed, + doc_ids, + row_offset, + scale, + backend=backend, + meta=local_meta, + ) + + +def _needed_kv_a2a_plan(plan: dict, rank: int, world: int, local_len: int, dev): + """Per-(src,dst) split sizes + local gather index for the general needed-only + all-to-all-v exchange (the >2-rank / single-doc case the halo can't cover). + + Each rank r needs the contiguous global range ``[s_first_r, real_end_r)``. As a + SENDER, this rank s sends to each dst r the intersection of its owned chunk + ``[s*L,(s+1)*L)`` with r's needed range (a token may go to several dsts, so the + send buffer duplicates it once per destination). As a receiver, the pieces arrive + in source order and concatenate into the contiguous needed range. All counts derive + from the replicated plan, so every rank agrees. + + Args: + plan: Replicated plan from :func:`_compute_blockdiag_kv_plan`. + rank: This rank's index within the CP group. + world: CP world size. + local_len: Per-rank local sequence length ``L``. + dev: Device for the produced ``send_index`` tensor. + + Returns: + ``(in_splits, out_splits, send_index)``: per-destination send counts, + per-source receive counts, and the 1D int64 local token indices (with + duplication) to gather into the send buffer. + """ + s_first = plan["s_first"] + real_end = plan["real_end"] + base = rank * local_len + in_splits, out_splits = [], [] + send_idx = [] + for r in range(world): + lo = max(base, s_first[r]) + hi = min(base + local_len, real_end[r]) + n = max(0, hi - lo) + in_splits.append(n) + if n > 0: + send_idx.append(torch.arange(lo - base, hi - base, device=dev, dtype=torch.long)) + for s in range(world): + lo = max(s * local_len, s_first[rank]) + hi = min((s + 1) * local_len, real_end[rank]) + out_splits.append(max(0, hi - lo)) + send_index = torch.cat(send_idx) if send_idx else torch.empty(0, dtype=torch.long, device=dev) + return in_splits, out_splits, send_index + + +class _NeededKVExchange(torch.autograd.Function): + """Differentiable multi-cast all-to-all-v K/V delivery. + + Delivers each rank exactly the K/V range it attends to, zero-redundantly (a + source token is sent once per rank that needs it). Forward maps the local + chunk ``[B, C, L, D]`` to the received needed range ``[B, C, R_recv, D]``. + Backward scatter-ADDs the per-destination grads back to the source token + (index_add), matching what the all-gather's reduce_scatter would sum. + """ + + @staticmethod + def forward(ctx, x_local, group, in_splits, out_splits, send_index): + """All-to-all-v ``x_local`` ``[B, C, L, D]`` into the needed range ``[B, C, R_recv, D]``.""" + ctx.group = group + ctx.in_splits = in_splits + ctx.out_splits = out_splits + ctx.save_for_backward(send_index) + ctx.local_len = x_local.shape[2] + B, C, L, D = x_local.shape + # gather (with duplication) the tokens to send, tokens-first for all_to_all_single + xp = x_local.permute(2, 0, 1, 3).contiguous() # [L, B, C, D] + send_buf = xp.index_select(0, send_index).contiguous() # [Tsend, B, C, D] + recv_buf = send_buf.new_empty((sum(out_splits), B, C, D)) + torch.distributed.all_to_all_single( + recv_buf, + send_buf, + output_split_sizes=out_splits, + input_split_sizes=in_splits, + group=group, + ) + return recv_buf.permute(1, 2, 0, 3).contiguous() # [B, C, Rrecv, D] + + @staticmethod + def backward(ctx, grad_out): + """Reverse all-to-all-v of ``grad_out`` ``[B, C, R_recv, D]``; multi-cast grads accumulate.""" + (send_index,) = ctx.saved_tensors + B, C, R, D = grad_out.shape + gp = grad_out.permute(2, 0, 1, 3).contiguous() # [Rrecv, B, C, D] + grad_send = gp.new_empty((sum(ctx.in_splits), B, C, D)) + # reverse direction: swap in/out splits + torch.distributed.all_to_all_single( + grad_send, + gp, + output_split_sizes=ctx.in_splits, + input_split_sizes=ctx.out_splits, + group=ctx.group, + ) + grad_x = gp.new_zeros((ctx.local_len, B, C, D)) + grad_x.index_add_(0, send_index, grad_send) # multi-cast tokens accumulate + return grad_x.permute(1, 2, 0, 3).contiguous(), None, None, None, None + + +def _blockdiag_a2a_attention(query, key, value, doc_ids, group, plan, gm, row_offset, scale, backend): + """Needed-only block-diagonal CP attention via the general all-to-all-v exchange. + + Handles documents spanning >2 ranks (single long doc / unpacked long context) + that the left-halo can't, still zero-redundantly (vs the full all-gather + fallback). Same argument contract as :func:`_blockdiag_halo_attention`. + + Args: + query: Local query shard ``[B, Hq, L, D]``. + key: Local key shard ``[B, Hkv, L, D]``. + value: Local value shard ``[B, Hkv, L, D]``. + doc_ids: Replicated per-position document ids ``[B, S_full]`` (0 == pad). + group: The CP process group. + plan: Replicated plan from :func:`_compute_blockdiag_kv_plan` + (augmented with ``rank``/``world``). + gm: This rank's per-step varlen metadata (global ``s_first``/``real_end``). + row_offset: Global position of this rank's first local query row. + scale: Softmax scale (``None`` -> kernel default). + backend: Varlen kernel backend, ``"flash"`` or ``"te"``. + + Returns: + Attention output ``[B, Hq, L, D]``, or ``None`` when the kernel is + unavailable. + """ + rank = plan["rank"] + world = plan["world"] + local_len = query.shape[2] + in_splits, out_splits, send_index = _needed_kv_a2a_plan(plan, rank, world, local_len, query.device) + + Hkv = key.shape[1] + kv_local = torch.cat((key, value), dim=1) # [B, 2*Hkv, L, D] + kv_needed = _NeededKVExchange.apply(kv_local, group, in_splits, out_splits, send_index) # [B,2Hkv,need,D] + key_needed = kv_needed[:, :Hkv] + value_needed = kv_needed[:, Hkv:] + + n_real = gm["n_real"] + if n_real == 0: + out = torch.zeros_like(query) + return out + 0.0 * (key_needed.sum() + value_needed.sum()).to(out.dtype) + needed_len = plan["real_end"][rank] - plan["s_first"][rank] + local_meta = dict(gm) + local_meta["s_first"] = 0 + local_meta["real_end"] = needed_len + return _cp_blockdiag_varlen( + query, + key_needed, + value_needed, + doc_ids, + row_offset, + scale, + backend=backend, + meta=local_meta, + ) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py new file mode 100644 index 0000000000..77a9734540 --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py @@ -0,0 +1,657 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dense and varlen block-diagonal attention kernels.""" + +from __future__ import annotations + +import logging + +import torch + +logger = logging.getLogger(__name__) + +_CP_FLASH_DETERMINISTIC = False +_CP_FLASH_WARNED = False +_CP_FLASH_ENGAGED = False +_CP_VARLEN_SHAPE_LOGGED = False +_CP_FLASH_LONG_SEGMENT_WARNED = False +_TE_DPA_CACHE = {} + + +def _varlen_backend_unavailable_reason( + backend: str, + dtype: torch.dtype, + device: torch.device, +) -> str | None: + """Return a cheap, non-kernel-launching reason a varlen kernel cannot run. + + Needed-only KV exchange cannot safely discover these failures after a + rank-local kernel call: an all-padding rank skips the call while a real-token + rank may fail and otherwise diverge in collective order. The runtime calls + this on every CP rank and reaches a small consensus before exchanging K/V. + Shape/data-dependent kernel failures are still handled by a post-call + consensus in :mod:`nemo_automodel.components.distributed.blockdiag_cp.runtime`. + + Args: + backend: Varlen backend name (``"flash"`` or ``"te"``). + dtype: Query/key/value dtype the kernel would run with. + device: Device the kernel would run on. + + Returns: + A human-readable reason string, or ``None`` when the kernel can run. + """ + if dtype not in (torch.float16, torch.bfloat16): + return f"varlen CP attention needs fp16/bf16, got {dtype}" + if device.type != "cuda": + return f"varlen CP attention requires CUDA, got device={device}" + if backend == "flash": + try: + from flash_attn import flash_attn_varlen_func # noqa: F401 + except Exception as exc: + return f"flash_attn varlen kernel is unavailable ({type(exc).__name__})" + return None + if backend == "te": + try: + from transformer_engine.pytorch import DotProductAttention # noqa: F401 + except Exception as exc: + return f"TransformerEngine varlen kernel is unavailable ({type(exc).__name__})" + return None + return f"unsupported varlen backend={backend}" + + +def _varlen_metadata_unavailable_reason( + meta: dict | None, + *, + query_len: int, + key_len: int, + device: torch.device, +) -> str | None: + """Validate every index consumed by a varlen CUDA kernel. + + FlashAttention trusts ``cu_seqlens`` and the caller-provided maxima. A bad + terminal offset therefore does not reliably raise a Python exception; it can + become an asynchronous illegal memory access and poison the whole CUDA + context. Validate the compact, per-step metadata on the host before the + first attention layer launches a kernel. The result is cached in ``meta``; + the cache object is intentionally shared by the shallow metadata copies used + by halo/A2A, so this costs one small device-to-host copy per step and shape, + not once per layer. + + Args: + meta: Per-step varlen metadata as produced by + :func:`precompute_blockdiag_varlen_meta` (``cu_q``/``cu_k`` are 1D + int32 device tensors of per-document cumulative offsets). + query_len: Local query length ``L`` (rows this rank attends with). + key_len: Key length ``S`` visible to this rank (full sequence for + all-gather, needed range for halo/A2A). + device: Device the ``cu_q``/``cu_k`` tensors must live on. + + Returns: + A human-readable reason string when the metadata is unsafe, else ``None``. + """ + if not isinstance(meta, dict): + return "missing varlen metadata" + + try: + n_real = int(meta["n_real"]) + except (KeyError, TypeError, ValueError) as exc: + return f"invalid n_real ({type(exc).__name__})" + + query_len = int(query_len) + key_len = int(key_len) + if not 0 <= n_real <= query_len: + return f"n_real={n_real} is outside query length {query_len}" + if n_real == 0: + return None + + required = ("s_first", "real_end", "cu_q", "cu_k", "max_q", "max_k") + missing = [name for name in required if name not in meta] + if missing: + return f"missing varlen metadata fields: {missing}" + + try: + s_first = int(meta["s_first"]) + real_end = int(meta["real_end"]) + max_q = int(meta["max_q"]) + max_k = int(meta["max_k"]) + except (TypeError, ValueError) as exc: + return f"non-integral varlen scalar metadata ({type(exc).__name__})" + + signature = ( + n_real, + query_len, + key_len, + s_first, + real_end, + max_q, + max_k, + str(device), + ) + cache = meta.setdefault("_validation_cache", {}) + if signature in cache: + return cache[signature] + + reason = None + if not 0 <= s_first <= real_end <= key_len: + reason = f"K/V slice [{s_first}, {real_end}) is outside key length {key_len}" + else: + cu_q = meta["cu_q"] + cu_k = meta["cu_k"] + for name, cu in (("cu_q", cu_q), ("cu_k", cu_k)): + if not isinstance(cu, torch.Tensor): + reason = f"{name} is not a tensor" + break + if cu.dtype != torch.int32: + reason = f"{name} must be int32, got {cu.dtype}" + break + if cu.device != device: + reason = f"{name} is on {cu.device}, expected {device}" + break + if cu.ndim != 1 or cu.numel() < 2: + reason = f"{name} must be a 1D tensor with at least two entries" + break + if not cu.is_contiguous(): + reason = f"{name} must be contiguous" + break + if reason is None and cu_q.numel() != cu_k.numel(): + reason = f"cu_q/cu_k segment counts differ: {cu_q.numel() - 1} vs {cu_k.numel() - 1}" + + if reason is None: + # One synchronized copy validates all offsets. Never let an unchecked + # CUDA index tensor reach FlashAttention merely to save this per-step + # microsecond-scale guard. + q_offsets = meta["cu_q"].detach().cpu().tolist() + k_offsets = meta["cu_k"].detach().cpu().tolist() + q_lens = [b - a for a, b in zip(q_offsets, q_offsets[1:])] + k_lens = [b - a for a, b in zip(k_offsets, k_offsets[1:])] + if q_offsets[0] != 0 or k_offsets[0] != 0: + reason = "cu_q and cu_k must start at zero" + elif any(length <= 0 for length in q_lens + k_lens): + reason = "cu_q/cu_k must be strictly increasing" + elif q_offsets[-1] != n_real: + reason = f"cu_q[-1]={q_offsets[-1]} does not equal n_real={n_real}" + elif k_offsets[-1] != real_end - s_first: + reason = f"cu_k[-1]={k_offsets[-1]} does not equal K/V slice length {real_end - s_first}" + elif k_lens[0] < q_lens[0] or k_lens[1:] != q_lens[1:]: + reason = "only the first local document may have more K/V than Q tokens" + elif max_q != max(q_lens) or max_k != max(k_lens): + reason = f"max_seqlen mismatch: supplied q/k={max_q}/{max_k}, actual={max(q_lens)}/{max(k_lens)}" + + cache[signature] = reason + return reason + + +def _te_varlen_dpa(num_q_heads, num_kv_heads, head_dim, scale, device, dtype): + """Cached TE DotProductAttention module for thd/varlen block-diagonal CP attention.""" + key = (num_q_heads, num_kv_heads, head_dim, round(float(scale), 10), str(device), str(dtype)) + m = _TE_DPA_CACHE.get(key) + if m is None: + from transformer_engine.pytorch import DotProductAttention + + # bottom_right: the local query rows are a SUFFIX of their document's keys + # (the left-straddling doc has sk > sq), so causal must align to the + # bottom-right corner -- matches flash's bottom-right and the dense + # global-position causal. padding_causal (top-left) is WRONG here. + m = DotProductAttention( + num_attention_heads=num_q_heads, + kv_channels=head_dim, + num_gqa_groups=num_kv_heads, + attn_mask_type="padding_causal_bottom_right", + qkv_format="thd", + softmax_scale=scale, + ).to(device) + _TE_DPA_CACHE[key] = m + return m + + +def _flash_varlen_with_long_prefix_guard( + q_packed: torch.Tensor, + k_packed: torch.Tensor, + v_packed: torch.Tensor, + *, + cu_q: torch.Tensor, + cu_k: torch.Tensor, + max_q: int, + max_k: int, + local_query_len: int, + scale, + meta: dict, +): + """Run FlashAttention without an asymmetric-prefix varlen launch. + + In CP, only the first local document can have ``K > Q``: it may include a + left halo from the preceding rank. Some FlashAttention builds have produced + an asynchronous illegal access for long packs with this layout. + Fixed-sequence Flash accepts the same bottom-right causal Q/K layout without + consulting ``cu_seqlens``. Peel off just that one asymmetric segment, then + process all remaining (ordinary ``K == Q``) documents in one varlen call. + + Normal packs stay on the original single-call path. A guarded pack adds at + most one kernel launch, independent of its number of packed documents. + + Args: + q_packed: Packed local queries ``[n_real, Hq, D]`` (``n_real`` = real + local query tokens, ``Hq`` = query heads, ``D`` = head dim). + k_packed: Packed keys ``[T_k, Hkv, D]`` covering the needed K/V range. + v_packed: Packed values ``[T_k, Hkv, D]``; same layout as ``k_packed``. + cu_q: Per-document cumulative query offsets, 1D int32 ``[n_docs + 1]``. + cu_k: Per-document cumulative key offsets, 1D int32 ``[n_docs + 1]``. + max_q: Maximum per-document query segment length. + max_k: Maximum per-document key segment length. + local_query_len: This rank's local sequence length (for diagnostics). + scale: Softmax scale (``None`` -> kernel default ``D**-0.5``). + meta: Per-step metadata carrying ``first_q``/``first_k``/``max_tail``. + + Returns: + Packed attention output ``[n_real, Hq, D]``. + """ + from flash_attn import flash_attn_func, flash_attn_varlen_func + + first_q = int(meta.get("first_q", 0)) + first_k = int(meta.get("first_k", 0)) + if first_q <= 0: + first_q = int(cu_q[1].item()) + if first_k <= 0: + first_k = int(cu_k[1].item()) + + # Keep every symmetric segment on the original one-call fast path. Peel + # off any left-straddling prefix, not just a >local-shard instance: the + # extra launch is bounded to one per layer/rank, and this avoids relying + # on an unknown size threshold inside the external FlashAttention build. + if first_k <= first_q: + return flash_attn_varlen_func( + q_packed, + k_packed, + v_packed, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max_q, + max_seqlen_k=max_k, + softmax_scale=scale, + causal=True, + deterministic=_CP_FLASH_DETERMINISTIC, + ) + + global _CP_FLASH_LONG_SEGMENT_WARNED + if not _CP_FLASH_LONG_SEGMENT_WARNED: + logger.warning( + "Flash varlen long-prefix guard engaged: first_q=%d first_k=%d local_len=%d; " + "using fixed Flash for the boundary document and varlen Flash for the remaining documents", + first_q, + first_k, + local_query_len, + ) + _CP_FLASH_LONG_SEGMENT_WARNED = True + + first_out = flash_attn_func( + q_packed[:first_q].unsqueeze(0), + k_packed[:first_k].unsqueeze(0), + v_packed[:first_k].unsqueeze(0), + softmax_scale=scale, + causal=True, + deterministic=_CP_FLASH_DETERMINISTIC, + ).squeeze(0) + + if cu_q.numel() == 2: + return first_out + + tail_cu_q = (cu_q[1:] - first_q).contiguous() + tail_cu_k = (cu_k[1:] - first_k).contiguous() + tail_max = int(meta.get("max_tail", 0)) + if tail_max <= 0: + tail_max = int((tail_cu_q[1:] - tail_cu_q[:-1]).max().item()) + tail_out = flash_attn_varlen_func( + q_packed[first_q:], + k_packed[first_k:], + v_packed[first_k:], + cu_seqlens_q=tail_cu_q, + cu_seqlens_k=tail_cu_k, + max_seqlen_q=tail_max, + max_seqlen_k=tail_max, + softmax_scale=scale, + causal=True, + deterministic=_CP_FLASH_DETERMINISTIC, + ) + return torch.cat((first_out, tail_out), dim=0) + + +def _cp_blockdiag_mask( + doc_ids: torch.Tensor, + row_offset: int, + local_len: int, + full_len: int, + batch_size: int, +) -> torch.Tensor: + """Per-document causal attention mask for block-diagonal CP, shape ``[B, 1, L, S]``. + + ``doc_ids`` is the full (all-rank, padded) per-position document index ``[B, S]`` + (0 == padding). Query rows are this rank's local positions + ``[row_offset, row_offset+local_len)``; key columns span the full sequence. A + query attends to a key iff they share a document, neither is padding, and the + key is causally visible (global query position >= key position) -- identical to + the block-causal mask a non-CP packed run would build for real tokens. The + diagonal is always allowed so a query row is never fully masked (which the + math/efficient SDPA backend turns into NaN); for padding rows this is a + harmless self-edge whose output is dropped by the -100 labels. + + Args: + doc_ids: Per-position document ids ``[B, S]`` or ``[S]`` (0 == padding), + where ``B`` = batch and ``S`` = full padded sequence length. + row_offset: Global position of this rank's first local query row. + local_len: ``L``, the number of local query rows. + full_len: ``S``, the number of key columns (full padded sequence). + batch_size: ``B``, used to expand a 1D ``doc_ids``. + + Returns: + Boolean allow-mask ``[B, 1, L, S]`` (True == may attend). + """ + if doc_ids.dim() == 1: + doc_ids = doc_ids.unsqueeze(0).expand(batch_size, -1) + device = doc_ids.device + L, S = local_len, full_len + row_doc = doc_ids[:, row_offset : row_offset + L] # [B, L] + col_doc = doc_ids # [B, S] + same_doc = row_doc.unsqueeze(2) == col_doc.unsqueeze(1) # [B, L, S] + not_pad = (row_doc.unsqueeze(2) > 0) & (col_doc.unsqueeze(1) > 0) # [B, L, S] + row_pos = torch.arange(row_offset, row_offset + L, device=device).view(1, L, 1) + col_pos = torch.arange(S, device=device).view(1, 1, S) + causal = row_pos >= col_pos # [1, L, S] + # Always allow the diagonal (q_pos == k_pos) so every query attends to >=1 key even + # in all-pad/empty rows -- prevents NaN/hang. + self_diag = row_pos == col_pos # [1, L, S] + allow = (same_doc & not_pad & causal) | self_diag + return allow.unsqueeze(1) # [B, 1, L, S] + + +def _varlen_seg_for_rank(dids: torch.Tensor, row_offset: int, local_len: int, dev) -> dict | None: + """Per-rank block-diagonal varlen segmentation for ONE packed sequence. + + Factoring the segmentation out of the per-layer hot path lets + :func:`precompute_blockdiag_varlen_meta` run it once per step so the + attention layers do zero GPU->CPU host syncs. + + Args: + dids: Full (padded) per-position document id vector ``[S]`` (0 == pad). + row_offset: Global position of this rank's first local query row. + local_len: ``L``, the local query chunk length. + dev: Device for the produced ``cu_q``/``cu_k`` tensors. + + Returns: + The cu_seqlens / slice metadata dict for this rank's local query chunk + ``[row_offset, row_offset + local_len)`` (keys ``n_real``, ``s_first``, + ``real_end``, ``cu_q``, ``cu_k``, ``max_q``, ``max_k``, ``first_q``, + ``first_k``, ``max_tail``), or ``None`` if the chunk is entirely padding. + """ + local = dids[row_offset : row_offset + local_len] # [L] + n_real = int((local > 0).sum().item()) + if n_real == 0: + return None + real_dids = local[:n_real] # padding is a tail -> real rows are the prefix + + # per-document local query segment lengths (run-length encode) + if n_real == 1: + seg_q = torch.ones(1, dtype=torch.long, device=dev) + else: + bnd = torch.nonzero(real_dids[1:] != real_dids[:-1], as_tuple=False).flatten() + 1 + edges = torch.cat( + [ + torch.zeros(1, dtype=torch.long, device=dev), + bnd.to(torch.long), + torch.tensor([n_real], dtype=torch.long, device=dev), + ] + ) + seg_q = edges[1:] - edges[:-1] + + # left-straddle: how far the first local doc extends before row_offset + back = 0 + if row_offset > 0: + d0 = local[0] + prefix = dids[:row_offset] + ne = torch.nonzero(prefix != d0, as_tuple=False) + back = row_offset if ne.numel() == 0 else int(row_offset - 1 - ne.max().item()) + + seg_k = seg_q.clone() + seg_k[0] = seg_k[0] + back + s_first = row_offset - back + real_end = row_offset + n_real + + cu_q = torch.zeros(seg_q.numel() + 1, dtype=torch.int32, device=dev) + cu_q[1:] = torch.cumsum(seg_q, 0).to(torch.int32) + cu_k = torch.zeros(seg_k.numel() + 1, dtype=torch.int32, device=dev) + cu_k[1:] = torch.cumsum(seg_k, 0).to(torch.int32) + max_q = int(seg_q.max().item()) + max_k = int(seg_k.max().item()) + max_tail = int(seg_q[1:].max().item()) if seg_q.numel() > 1 else 0 + return { + "n_real": n_real, + "s_first": s_first, + "real_end": real_end, + "cu_q": cu_q, + "cu_k": cu_k, + "max_q": max_q, + "max_k": max_k, + "first_q": int(seg_q[0].item()), + "first_k": int(seg_k[0].item()), + "max_tail": max_tail, + } + + +def precompute_blockdiag_varlen_meta(doc_ids: torch.Tensor, row_offset: int, local_len: int, device) -> dict: + """Precompute this rank's varlen segmentation once per step. + + The block-diagonal varlen cu_seqlens depend only on ``(doc_ids, row_offset, + local_len)`` -- all step-constant -- yet rebuilding them inline on every + attention layer (forward + AC recompute) costs 3-4 host syncs per rebuild. + :func:`~nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx` + calls this once and stashes the result in the CP state; ``_cp_blockdiag_varlen(meta=...)`` + then runs with zero ``.item()`` syncs. + + Assumes one packed sequence per rank (B == 1, the CP contract). + + Args: + doc_ids: Per-position document ids ``[B, S]`` (row 0 used) or ``[S]`` + (0 == padding), covering the full padded sequence. + row_offset: Global position of this rank's first local query row. + local_len: ``L``, the local query chunk length. + device: Device for the produced ``cu_q``/``cu_k`` tensors. + + Returns: + A dict consumed directly by ``_cp_blockdiag_varlen``; ``{"n_real": 0}`` + for an all-padding chunk (the varlen path then takes its grad-preserving + 0-weighted touch). + """ + dids = (doc_ids if doc_ids.dim() == 1 else doc_ids[0]).to(device=device, dtype=torch.long) + seg = _varlen_seg_for_rank(dids, row_offset, local_len, device) + return seg if seg is not None else {"n_real": 0} + + +def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale=None, backend="flash", meta=None): + """Block-diagonal CP attention via varlen (flash or TE) -- no dense ``[B,1,L,S]`` mask. + + Equivalent to ``_cp_blockdiag_mask`` + SDPA for real (non-padding) query rows. + The sequence is sharded contiguously and packing puts padding (doc id 0) only + as a contiguous tail, so the local query's real rows are a prefix ``[0, n_real)``. + For each document the local rows touch, we emit a varlen segment: q segment = + the local rows in that doc, k segment = ``[doc_start, last_local_q+1)`` (the only + doc with ``sk > sq`` straddles the left boundary). Bottom-right causal alignment + reproduces same-doc + global-causal exactly. Padding query rows are returned as + zeros (their loss is masked by -100 labels). + + Args: + query: Local query shard ``[B, Hq, L, D]`` (``B`` = batch, ``Hq`` = query + heads, ``L`` = local sequence length, ``D`` = head dim). + key_full: Keys ``[B, Hkv, S, D]`` covering the K/V range indexed by + ``meta``'s ``[s_first, real_end)`` slice (full sequence for the + all-gather path, needed range for halo/A2A). + value_full: Values ``[B, Hkv, S, D]``; same layout as ``key_full``. + doc_ids: Per-position document ids ``[B, S_full]`` or ``[S_full]`` + (0 == padding) on the full padded sequence. + row_offset: Global position of this rank's first local query row. + scale: Softmax scale (``None`` -> kernel default ``D**-0.5``). + backend: ``"flash"`` (flash_attn_varlen_func) or ``"te"`` + (TransformerEngine DotProductAttention thd). + meta: Optional per-step segmentation from + :func:`precompute_blockdiag_varlen_meta`; rebuilt inline when absent. + + Returns: + Attention output ``[B, Hq, L, D]``, or ``None`` to signal "fall back to + the dense path" (kernel import failed, or a non-half dtype). + """ + global _CP_FLASH_WARNED + if query.dtype not in (torch.float16, torch.bfloat16): + if not _CP_FLASH_WARNED: + logger.warning( + "varlen CP attention needs fp16/bf16, got %s; reporting the unavailable varlen path to the caller", + query.dtype, + ) + _CP_FLASH_WARNED = True + return None + try: + if backend == "flash": + from flash_attn import flash_attn_varlen_func # noqa: F401 + except Exception: + if not _CP_FLASH_WARNED: + logger.warning("flash_attn import failed; reporting the unavailable varlen path to the caller") + _CP_FLASH_WARNED = True + return None + + B, Hq, L, D = query.shape + Hkv = key_full.shape[1] + if doc_ids.dim() == 1: + doc_ids = doc_ids.unsqueeze(0).expand(B, -1) + dev = query.device + out = torch.zeros_like(query) + + try: + for b in range(B): + # Use the per-step precomputed meta when available (B==1, the CP + # contract) so the hot path does zero .item() host syncs; otherwise + # rebuild inline for an unsupported multi-row/direct caller or when + # metadata was not armed -- identical math. + if meta is not None and B == 1: + seg = meta + else: + seg = _varlen_seg_for_rank(doc_ids[b], row_offset, L, dev) + if seg is None: + seg = {"n_real": 0} + n_real = seg["n_real"] + if n_real == 0: + # All-padding local chunk: no kernel call. But the output MUST stay + # attached to key_full/value_full, otherwise this CP rank skips the + # all-gather's backward reduce_scatter while real-token ranks fire it + # -> collective desync -> NCCL hang in backward (observed on cross-node + # cp>1 where a whole rank-chunk lands in the pad tail; the dense path + # never hangs because its self_diag diagonal always routes through + # value_full). A 0-weighted touch keeps autograd symmetric. + out[b] = out[b] + 0.0 * (key_full[b].sum() + value_full[b].sum()).to(out.dtype) + continue # output stays zeros numerically; grad path preserved + + s_first = seg["s_first"] + real_end = seg["real_end"] + cu_q = seg["cu_q"] + cu_k = seg["cu_k"] + max_q = seg["max_q"] + max_k = seg["max_k"] + + metadata_reason = _varlen_metadata_unavailable_reason( + seg, + query_len=L, + key_len=key_full.shape[2], + device=dev, + ) + if metadata_reason is not None: + raise ValueError(f"unsafe varlen metadata: {metadata_reason}") + + if key_full.shape != value_full.shape: + raise ValueError(f"K/V shapes differ: {tuple(key_full.shape)} vs {tuple(value_full.shape)}") + if key_full.device != dev or value_full.device != dev: + raise ValueError("Q/K/V must be on the same device") + if key_full.dtype != query.dtype or value_full.dtype != query.dtype: + raise ValueError("Q/K/V must have the same dtype") + if Hq % Hkv != 0: + raise ValueError(f"query heads {Hq} are not divisible by KV heads {Hkv}") + if D != key_full.shape[-1] or D != value_full.shape[-1]: + raise ValueError("Q/K/V head dimensions differ") + + global _CP_VARLEN_SHAPE_LOGGED + if not _CP_VARLEN_SHAPE_LOGGED: + logger.info( + "first varlen shape: backend=%s q_tokens=%d kv_tokens=%d segments=%d max_q=%d max_k=%d " + "heads=%d/%d head_dim=%d slice=[%d,%d)", + backend, + n_real, + real_end - s_first, + cu_q.numel() - 1, + max_q, + max_k, + Hq, + Hkv, + D, + s_first, + real_end, + ) + _CP_VARLEN_SHAPE_LOGGED = True + + q_packed = query[b, :, :n_real, :].transpose(0, 1).contiguous() # [n_real, Hq, D] + k_packed = key_full[b, :, s_first:real_end, :].transpose(0, 1).contiguous() # [n_real+back, Hkv, D] + v_packed = value_full[b, :, s_first:real_end, :].transpose(0, 1).contiguous() + + if backend == "te": + dpa = _te_varlen_dpa(Hq, Hkv, D, scale if scale is not None else D**-0.5, dev, query.dtype) + o = dpa( + q_packed, + k_packed, + v_packed, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_k, + max_seqlen_q=max_q, + max_seqlen_kv=max_k, + ) # [n_real, Hq*D] or [n_real, Hq, D] + o = o.view(n_real, Hq, D) + else: + o = _flash_varlen_with_long_prefix_guard( + q_packed, + k_packed, + v_packed, + cu_q=cu_q, + cu_k=cu_k, + max_q=max_q, + max_k=max_k, + local_query_len=L, + scale=scale, + meta=seg, + ) # [n_real, Hq, D] + out[b, :, :n_real, :] = o.transpose(0, 1) + except torch.cuda.OutOfMemoryError: + # Never signal fallback on OOM: the dense-mask fallback path allocates + # strictly more than the varlen kernel that just failed, so it is + # doomed and misattributes the failure site. Surface the real OOM site. + raise + except Exception as e: + if not _CP_FLASH_WARNED: + logger.warning( + "CP varlen backend '%s' failed (%s: %s); reporting failure to the caller", + backend, + type(e).__name__, + str(e)[:100], + ) + _CP_FLASH_WARNED = True + return None + + global _CP_FLASH_ENGAGED + if not _CP_FLASH_ENGAGED: + logger.info("block-diagonal CP attention using VARLEN backend=%s (no dense [B,1,L,S] mask)", backend) + _CP_FLASH_ENGAGED = True + return out diff --git a/nemo_automodel/components/distributed/blockdiag_cp/packed.py b/nemo_automodel/components/distributed/blockdiag_cp/packed.py new file mode 100644 index 0000000000..7ce8e1665c --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/packed.py @@ -0,0 +1,197 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""cp_size==1 packed-sequence varlen SDPA integration. + +Same varlen kernel as the CP block-diagonal path, but for ``cp_size == 1`` (no +sequence sharding). The whole packed sequence lives on one rank, so it +degenerates to ``_cp_blockdiag_varlen`` with ``row_offset=0`` (q == k == full +sequence; square; ``cu_q == cu_k``). This gives packed-sequence block-diagonal +attention to models whose softmax attention only dispatches to sdpa: the model +calls :func:`enable_cp1_packed_varlen` at the start of its forward, which (1) +stashes ``doc_ids`` and (2) routes ``F.scaled_dot_product_attention`` to the +flash/te varlen kernel -- like the CP path for cp>1. + +IMPORTANT -- state must survive activation-checkpointing RECOMPUTE in backward: +the patch is installed once (process-global) and the state is set per-forward and +NOT reset, so the AC worker thread re-running a layer's forward during backward +reads the SAME doc_ids and reproduces the exact varlen output. (A context manager +that reset on exit caused the forward to use varlen but the recompute to fall back +to dense -> AC "saved vs recomputed metadata" shape mismatch.) ``_ThreadSharedVar`` +makes the state visible in the autograd worker thread. The outer model clears +stale state before each new forward (:func:`disable_cp1_packed_varlen`), then the +text backend re-arms it for the current packed batch, so vision/unpacked +attention cannot inherit it. +""" + +from __future__ import annotations + +import torch + +from nemo_automodel.components.distributed.blockdiag_cp import kernels, runtime, state + +_PACKED_STATE = state._ThreadSharedVar() +_PACKED_SDPA_INSTALLED = False + + +def _packed_varlen_sdpa( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + enable_gqa=False, + **kwargs, +): + """SDPA replacement for cp1 packed runs: block-diagonal varlen from doc_ids. + + Only intervenes for a full packed-sequence call whose batch and sequence + dims match the armed ``doc_ids``; any other SDPA call (or state unset) + passes straight through to stock SDPA. + + Args: + query: Queries ``[B, Hq, S, D]`` over the full (unsharded) packed + sequence (``B`` = batch, ``Hq`` = query heads, ``S`` = sequence + length, ``D`` = head dim). + key: Keys ``[B, Hkv, S, D]``. + value: Values ``[B, Hkv, S, D]``. + attn_mask: Forwarded to stock SDPA on pass-through; ignored on the + varlen path (masking is rebuilt from ``doc_ids``). + dropout_p: Dropout probability (pass-through / dense fallback only). + is_causal: Forwarded on pass-through; the varlen path is always + per-document causal. + scale: Softmax scale (``None`` -> ``D**-0.5``). + enable_gqa: Grouped-query attention flag as passed by HF's sdpa path. + **kwargs: Ignored; accepted for SDPA signature compatibility. + + Returns: + Attention output ``[B, Hq, S, D]``. + """ + packed_state = _PACKED_STATE.get() + # Only intervene for the full packed sequence whose length matches doc_ids; + # any other SDPA call (or state unset) passes straight through. + doc_ids = None if packed_state is None else packed_state.get("doc_ids") + expected_shape = tuple(doc_ids.shape) if isinstance(doc_ids, torch.Tensor) and doc_ids.dim() == 2 else None + if ( + packed_state is None + or expected_shape is None + or query.dim() != 4 + or key.dim() != 4 + or value.dim() != 4 + or query.shape[0] != expected_shape[0] + or key.shape[0] != expected_shape[0] + or value.shape[0] != expected_shape[0] + or query.shape[2] != expected_shape[1] + or key.shape[2] != expected_shape[1] + or value.shape[2] != expected_shape[1] + ): + return runtime._ORIGINAL_SDPA( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) + # query/key/value: [B, Hq, S, D] / [B, Hkv, S, D]; full (unsharded) sequence. + out = kernels._cp_blockdiag_varlen( + query, + key, + value, + packed_state["doc_ids"], + 0, + scale, + packed_state["backend"], + meta=packed_state.get("varlen_meta"), + ) + if out is None: # backend unavailable / unsupported -> safe fallback + B = query.shape[0] + L = query.shape[2] + S = key.shape[2] + n_q_heads = query.shape[1] + n_kv_heads = key.shape[1] + if enable_gqa and n_kv_heads != n_q_heads: + n_rep = n_q_heads // n_kv_heads + key = key.repeat_interleave(n_rep, dim=1) + value = value.repeat_interleave(n_rep, dim=1) + enable_gqa = False + allow = kernels._cp_blockdiag_mask(packed_state["doc_ids"], 0, L, S, B) + return runtime._ORIGINAL_SDPA( + query, + key, + value, + attn_mask=allow, + dropout_p=dropout_p, + is_causal=False, + scale=scale, + enable_gqa=enable_gqa, + ) + return out + + +def enable_cp1_packed_varlen(doc_ids: torch.Tensor, backend: str) -> None: + """Arm cp1 packed block-diagonal varlen for the rest of this step. + + Idempotently installs the global SDPA patch and sets the per-forward + doc_ids/backend state. It remains armed through backward's + activation-checkpoint recomputation; the next outer model forward clears it + via :func:`disable_cp1_packed_varlen`. + + Args: + doc_ids: Per-position document ids ``[1, S]`` or ``[S]`` (0 == padding) + over the full packed sequence. + backend: Varlen kernel backend, ``"flash"`` or ``"te"``. + """ + global _PACKED_SDPA_INSTALLED + if not _PACKED_SDPA_INSTALLED: + import torch.nn.functional as F_module + + F_module.scaled_dot_product_attention = _packed_varlen_sdpa + _PACKED_SDPA_INSTALLED = True + # Segmentation depends only on the packed document ids, not on the layer's + # Q/K/V tensors. Compute it once per outer forward so every attention + # layer (and activation-checkpoint recompute) reuses the same CUDA + # cu_seqlens without repeating several GPU->CPU ``.item()`` synchronizations. + # Packing currently requires one packed row per local batch. Preserve the + # old inline fallback for an unexpected multi-row direct caller rather than + # silently applying row 0's metadata to every row. + varlen_meta = None + if isinstance(doc_ids, torch.Tensor) and (doc_ids.dim() == 1 or (doc_ids.dim() == 2 and doc_ids.shape[0] == 1)): + varlen_meta = kernels.precompute_blockdiag_varlen_meta( + doc_ids, + row_offset=0, + local_len=int(doc_ids.shape[-1]), + device=doc_ids.device, + ) + _PACKED_STATE.set({"doc_ids": doc_ids, "backend": backend, "varlen_meta": varlen_meta}) + + +def disable_cp1_packed_varlen() -> None: + """Disarm stale cp1 state before starting a new outer model forward. + + Decoder-layer activation-checkpoint recomputation happens before the next + outer forward, so the preceding step's state remains available for backward + and is cleared before vision or any other attention in the next batch runs. + """ + _PACKED_STATE.set(None) + + +def cp1_packed_varlen_backend() -> str | None: + """The configured cp1 packed varlen backend ('te'/'flash'), or None if disabled (dense).""" + backend = state._CP_ATTN_BACKEND + return backend if backend in ("te", "flash") else None diff --git a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py new file mode 100644 index 0000000000..3983910d57 --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py @@ -0,0 +1,362 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Block-diagonal SDPA routing and collective-safe fallback policy.""" + +from __future__ import annotations + +import logging +import os + +import torch + +from nemo_automodel.components.distributed.blockdiag_cp import exchange, kernels +from nemo_automodel.components.distributed.blockdiag_cp import state as state_module + +logger = logging.getLogger(__name__) +_KV_EXCHANGE_PATH_LOGGED = False +_KV_XNODE_LOGGED = False + + +def _cp_group_spans_nodes(group) -> bool: + """True if the CP group spans more than one node. + + The needed-only exchanges (halo neighbor-p2p, a2a all-to-all-v) are verified when + the CP group is node-local (cp_size <= GPUs/node, the common mesh layout where cp + is the fast dim). Cross-node needed-only exchange is still fabric-sensitive, so + the production default falls back to NCCL all-gather when the CP group spans nodes. + Heuristic: cp_size > the local node's world size (LOCAL_WORLD_SIZE / + NPROC_PER_NODE / visible GPUs). + + Set ``NEMO_CP_ALLOW_XNODE=1`` to force the needed-only path cross-node for targeted + validation on a specific cluster. + """ + if os.environ.get("NEMO_CP_ALLOW_XNODE") == "1": + return False + world = torch.distributed.get_world_size(group) + local_ws = int(os.environ.get("LOCAL_WORLD_SIZE") or os.environ.get("NPROC_PER_NODE") or 0) + if local_ws <= 0: + try: + local_ws = torch.cuda.device_count() + except Exception: + local_ws = world + spans = local_ws > 0 and world > local_ws + global _KV_XNODE_LOGGED + if spans and not _KV_XNODE_LOGGED: + _KV_XNODE_LOGGED = True + logger.warning( + "needed-only KV exchange disabled: CP group (cp_size=%d) spans nodes " + "(local_world_size=%d) -> using all-gather; node-local CP is unaffected.", + world, + local_ws, + ) + return spans + + +def _resolve_cp_varlen_config(state: dict) -> tuple[str, str]: + """Return ``(attn_backend, kv_exchange)`` for this step. + + Prefers the per-step snapshot threaded into ``state`` by the batch/ctx builder (so + :func:`cp_blockdiag_sdpa` is a function of ``(qkv, state)`` on the production path), + and falls back to the runtime configuration owned by + :mod:`~nemo_automodel.components.distributed.blockdiag_cp.state` for callers that + set the step state manually (e.g. parity tests). + """ + return ( + state.get("attn_backend", state_module._CP_ATTN_BACKEND), + state.get("kv_exchange", state_module._CP_KV_EXCHANGE), + ) + + +def _needed_only_preflight( + state: dict, + group, + backend: str, + query_dtype: torch.dtype, + device: torch.device, + *, + query_len: int, + key_len: int, +) -> tuple[bool, str | None]: + """Collectively decide whether every CP rank can run the varlen kernel.""" + cache_key = (backend, str(query_dtype), str(device), int(query_len), int(key_len)) + cached = state.get("_needed_only_preflight") + if cached is not None and cached[0] == cache_key: + return cached[1], cached[2] + + local_reason = kernels._varlen_backend_unavailable_reason(backend, query_dtype, device) + if local_reason is None: + local_reason = kernels._varlen_metadata_unavailable_reason( + state.get("varlen_meta"), + query_len=query_len, + key_len=key_len, + device=device, + ) + all_available = local_reason is None + if torch.distributed.is_initialized() and torch.distributed.get_world_size(group) > 1: + available = torch.tensor(int(all_available), dtype=torch.int32, device=device) + torch.distributed.all_reduce(available, op=torch.distributed.ReduceOp.MIN, group=group) + all_available = bool(available.item()) + + reason = None + if not all_available: + reason = local_reason or "varlen kernel is unavailable on at least one CP rank" + state["_needed_only_preflight"] = (cache_key, all_available, reason) + return all_available, reason + + +def _needed_only_kernel_succeeded_on_all_ranks(out, group, device) -> bool: + """Make a rank-local varlen result safe to act on collectively. + + In particular, an all-padding rank returns a zero tensor without invoking + FlashAttention/TE. It must nevertheless learn when a peer's real-token + kernel returned ``None`` before either rank advances to another collective. + """ + succeeded = out is not None + if torch.distributed.is_initialized() and torch.distributed.get_world_size(group) > 1: + status = torch.tensor(int(succeeded), dtype=torch.int32, device=device) + torch.distributed.all_reduce(status, op=torch.distributed.ReduceOp.MIN, group=group) + succeeded = bool(status.item()) + return succeeded + + +def _select_kv_exchange_path( + state, + group, + doc_ids, + local_len, + device, + offset, + *, + query_dtype: torch.dtype | None = None, +): + """Decide this step's KV-exchange path and WHY. + + Every downgrade to all-gather names its cause (mode, kernel, missing varlen meta, + or cross-node topology) so a silent fall-through can't hide a misconfiguration. + + Args: + state: The per-step CP state dict (memoizes the plan and preflight). + group: The CP process group. + doc_ids: Replicated per-position document ids ``[B, S_full]`` (0 == pad). + local_len: Per-rank local sequence length ``L``. + device: Device for plan tensors and the preflight all-reduce. + offset: Global position of this rank's first local query row. + query_dtype: Query dtype for the kernel preflight (skip when ``None``). + + Returns: + ``(path, plan, reason)`` where ``path`` is ``"halo"``/``"a2a"``/``"allgather"``, + ``plan`` is the cached block-diagonal KV plan for halo/a2a (else ``None``), + and ``reason`` names why the path was chosen. + """ + attn_backend, kv_exchange = _resolve_cp_varlen_config(state) + if kv_exchange not in ("halo", "a2a"): + return "allgather", None, f"mode={kv_exchange}" + if attn_backend not in ("flash", "te"): + return "allgather", None, f"needed-only requires flash/te kernel, got {attn_backend}" + if state.get("varlen_meta") is None: + return "allgather", None, "no varlen_meta (dense / non-varlen step)" + if _cp_group_spans_nodes(group): + return "allgather", None, "CP group spans nodes (no GPU<->NIC P2P)" + if query_dtype is not None: + available, unavailable_reason = _needed_only_preflight( + state, + group, + attn_backend, + query_dtype, + device, + query_len=local_len, + key_len=doc_ids.shape[-1], + ) + if not available: + return "allgather", None, f"needed-only preflight failed: {unavailable_reason}" + plan = state.get("kv_plan") + if plan is None: + world = torch.distributed.get_world_size(group) + plan = exchange._compute_blockdiag_kv_plan(doc_ids, world, local_len, device) + plan["rank"] = offset // local_len if local_len else 0 + plan["world"] = world + state["kv_plan"] = plan # reused across layers + AC recompute this step + if kv_exchange == "halo" and plan["use_halo"]: + return "halo", plan, "neighbor p2p (doc<=2 ranks)" + return "a2a", plan, "all-to-all-v (doc spans >2 ranks)" + + +def cp_blockdiag_sdpa( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + **kwargs, +) -> torch.Tensor: + """Block-diagonal context-parallel SDPA. + + Drop-in replacement for ``torch.nn.functional.scaled_dot_product_attention`` + while the context returned by + :func:`~nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx` + is active; a plain pass-through to stock SDPA otherwise. K/V are exchanged + across the CP group (all-gather, or a needed-only halo/a2a exchange) and one + local attention runs the local queries against the delivered keys with a + per-document causal mask. The passed ``attn_mask`` / ``is_causal`` are ignored + on the CP path -- masking is rebuilt from the document ids so packed sequences + never attend across document boundaries. + + Args: + query: This rank's LOCAL query shard ``[B, Hq, L, D]`` (``B`` = batch, + ``Hq`` = query heads, ``L`` = local sequence length, ``D`` = head dim). + key: This rank's LOCAL key shard ``[B, Hkv, L, D]``. + value: This rank's LOCAL value shard ``[B, Hkv, L, D]``. + attn_mask: Ignored on the CP path (forwarded to stock SDPA otherwise). + dropout_p: Dropout probability (dense fallback path only). + is_causal: Ignored on the CP path (forwarded to stock SDPA otherwise). + scale: Softmax scale (``None`` -> ``D**-0.5``). + enable_gqa: Grouped-query attention flag as passed by HF's sdpa path. + **kwargs: Ignored; accepted for SDPA signature compatibility. + + Returns: + Attention output ``[B, Hq, L, D]`` for this rank's local rows. + """ + step_state = state_module._CP_BLOCKDIAG_STATE.get() + if step_state is None: + # No CP active (e.g. cp_size==1 fall-through) -- behave like stock SDPA. + return _ORIGINAL_SDPA( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) + + state_module._CP_ATTN_FIRE_COUNT[0] += 1 + group = step_state["group"] + doc_ids = step_state["doc_ids"] # [B, S] full (padded) document ids; 0 == padding + offset = step_state["row_offset"] # global position of this rank's first local row + seq_dim = 2 + attn_backend, _ = _resolve_cp_varlen_config(step_state) + + # Needed-only KV exchange (opt-in via kv_exchange='halo' or 'a2a') lets each rank + # attend only local + required boundary K/V instead of the full all-gathered sequence. + # The plan is computed once per step from replicated doc_ids and memoized in state. + gm = step_state.get("varlen_meta") + # Decide the KV-exchange path (needed-only halo/a2a vs full all-gather) explicitly, + # logging why, including downgrades to all-gather for mode, kernel, missing meta, or + # a cross-node CP group where needed-only exchange is disabled by default. + path, plan, reason = _select_kv_exchange_path( + step_state, + group, + doc_ids, + query.shape[seq_dim], + query.device, + offset, + query_dtype=query.dtype, + ) + global _KV_EXCHANGE_PATH_LOGGED + if not _KV_EXCHANGE_PATH_LOGGED: + _KV_EXCHANGE_PATH_LOGGED = True + # cp_size is informational; the all-gather path may run without a live process + # group (e.g. the in-process parity tests pass group=None), so look it up safely. + cp_sz = torch.distributed.get_world_size(group) if torch.distributed.is_initialized() else -1 + logger.info( + "KV exchange path=%s (%s; cp_size=%d, local_len=%d)", + path, + reason, + cp_sz, + query.shape[seq_dim], + ) + if path in ("halo", "a2a"): + if path == "halo": + out = exchange._blockdiag_halo_attention( + query, key, value, doc_ids, group, plan, gm, offset, scale, attn_backend + ) + else: + out = exchange._blockdiag_a2a_attention( + query, key, value, doc_ids, group, plan, gm, offset, scale, attn_backend + ) + if _needed_only_kernel_succeeded_on_all_ranks(out, group, query.device): + assert out is not None + return out + # The needed-only KV collective has already executed. An all-padding + # rank may have produced a local zero while a peer's real-token kernel + # failed, so the consensus above intentionally makes every rank raise. + raise RuntimeError( + f"needed-only CP {path} exchange completed, but the {attn_backend} " + "varlen kernel failed on at least one CP rank; all ranks are " + "refusing an unsafe post-exchange full-allgather fallback" + ) + + # Use one all-gather of stacked [K;V] instead of two separate collectives + # (halves collective launch/latency; the differentiable cat + the single + # reduce_scatter on backward split cleanly back into the K and V grads). + n_kv_heads_local = key.shape[1] + kv_full = exchange._AllGatherSeqDiff.apply(torch.cat((key, value), dim=1), group, seq_dim) + key_full = kv_full[:, :n_kv_heads_local] + value_full = kv_full[:, n_kv_heads_local:] + + if attn_backend in ("flash", "te"): + out = kernels._cp_blockdiag_varlen( + query, + key_full, + value_full, + doc_ids, + offset, + scale, + backend=attn_backend, + meta=step_state.get("varlen_meta"), + ) + if out is not None: + return out + # out is None -> varlen backend unavailable / unsupported: fall through to dense. + + # GQA: HF's sdpa path passes enable_gqa=True with UN-repeated K/V when the + # attention_mask is None (which it is here -- we rebuild masking ourselves). + # But the memory-efficient SDPA backend does not support enable_gqa, so with + # our custom 4D mask (flash excluded) the dispatcher would silently fall back + # to the MATH kernel and materialize the full [B, H_q, L, S] fp32 score matrix + # (OOM at long context). Expand K/V to the query head count here and disable + # enable_gqa so the efficient kernel handles the masked attention. + n_q_heads = query.shape[1] + n_kv_heads = key_full.shape[1] + if enable_gqa and n_kv_heads != n_q_heads: + n_rep = n_q_heads // n_kv_heads + key_full = key_full.repeat_interleave(n_rep, dim=1) + value_full = value_full.repeat_interleave(n_rep, dim=1) + enable_gqa = False + + B = query.shape[0] + L = query.shape[seq_dim] + S = key_full.shape[seq_dim] + + allow = kernels._cp_blockdiag_mask(doc_ids, offset, L, S, B) # [B, 1, L, S] + + return _ORIGINAL_SDPA( + query, + key_full, + value_full, + attn_mask=allow, + dropout_p=dropout_p, + is_causal=False, + scale=scale, + enable_gqa=enable_gqa, + ) + + +# Captured once at import so the routed function can delegate when CP is inactive. +_ORIGINAL_SDPA = torch.nn.functional.scaled_dot_product_attention diff --git a/nemo_automodel/components/distributed/blockdiag_cp/state.py b/nemo_automodel/components/distributed/blockdiag_cp/state.py new file mode 100644 index 0000000000..de40363b53 --- /dev/null +++ b/nemo_automodel/components/distributed/blockdiag_cp/state.py @@ -0,0 +1,206 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime knobs and activation-checkpoint-safe state for block-diagonal CP. + +The knob normalization (synonym maps and defaults) lives here, in a +dependency-free leaf, so every consumer (the kernel-side +:func:`configure_cp_varlen` entry point and any policy-side config parser) +derives the accepted values from the same table and can never disagree on +synonyms or defaults. +""" + +from __future__ import annotations + +from typing import Any + +# Canonical value -> accepted synonyms (each canonical is also a synonym of itself). +ATTN_BACKEND_SYNONYMS: dict[str, tuple[str, ...]] = { + "flash": ("flash", "flash_attn", "flash_attention", "flash_attention_2", "fa2"), + "te": ("te", "transformer_engine", "transformerengine"), + "dense": ("dense", "sdpa", "torch_sdpa"), +} +KV_EXCHANGE_SYNONYMS: dict[str, tuple[str, ...]] = { + "allgather": ("allgather", "all_gather"), + "halo": ("halo", "neighbor", "needed", "needed_only"), + "a2a": ("a2a", "a2av", "all_to_all", "alltoall", "needed_a2a"), +} + +ATTN_BACKEND_DEFAULT = "flash" +KV_EXCHANGE_DEFAULT = "allgather" + +# Canonical user-facing values (single source for config schemas + docs). +ATTN_BACKEND_VALUES: tuple[str, ...] = tuple(ATTN_BACKEND_SYNONYMS) +KV_EXCHANGE_VALUES: tuple[str, ...] = tuple(KV_EXCHANGE_SYNONYMS) + +_ATTN_BACKEND_REVERSE = {syn: canon for canon, syns in ATTN_BACKEND_SYNONYMS.items() for syn in syns} +_KV_EXCHANGE_REVERSE = {syn: canon for canon, syns in KV_EXCHANGE_SYNONYMS.items() for syn in syns} + + +def _clean(value: Any) -> str: + return str(value).strip().lower().replace("-", "_") + + +def normalize_attn_backend(value: Any) -> str: + """Canonicalize an attention-backend knob to one of ``ATTN_BACKEND_VALUES``. + + ``None``/``""``/``auto``/``default`` map to the default (``flash``). + + Args: + value: The raw user-facing knob value (string-like or ``None``). + + Returns: + The canonical backend name (``"flash"``, ``"te"``, or ``"dense"``). + + Raises: + ValueError: If the value is not a recognized backend or synonym. + """ + if value is None: + return ATTN_BACKEND_DEFAULT + cleaned = _clean(value) + if cleaned in ("", "auto", "default"): + return ATTN_BACKEND_DEFAULT + canon = _ATTN_BACKEND_REVERSE.get(cleaned) + if canon is not None: + return canon + raise ValueError( + f"Unsupported block-diagonal CP attention backend '{value}'. Use one of: {', '.join(ATTN_BACKEND_VALUES)}." + ) + + +def normalize_kv_exchange(value: Any) -> str: + """Canonicalize a KV-exchange knob to one of ``KV_EXCHANGE_VALUES``. + + ``None``/``""``/``auto``/``default`` map to the default (``allgather``). + + Args: + value: The raw user-facing knob value (string-like or ``None``). + + Returns: + The canonical exchange name (``"allgather"``, ``"halo"``, or ``"a2a"``). + + Raises: + ValueError: If the value is not a recognized exchange mode or synonym. + """ + if value is None: + return KV_EXCHANGE_DEFAULT + cleaned = _clean(value) + if cleaned in ("", "auto", "default"): + return KV_EXCHANGE_DEFAULT + canon = _KV_EXCHANGE_REVERSE.get(cleaned) + if canon is not None: + return canon + raise ValueError( + f"Unsupported block-diagonal CP kv_exchange '{value}'. Use one of: {', '.join(KV_EXCHANGE_VALUES)}." + ) + + +_CP_ATTN_BACKEND = ATTN_BACKEND_DEFAULT +# KV delivery for the block-diagonal path: +# "allgather" - full O(S) K/V all-gather to every rank (default; always correct) +# "halo" - needed-only left-neighbor halo exchange: each rank fetches only its +# boundary document's straddle from rank r-1 and attends local+halo. +# O(S/cp) per-rank KV + ~doc-sized comm instead of O(S). Auto-falls back +# to all-to-all-v when a document spans >2 ranks; the decision is computed +# from the replicated doc_ids so all ranks agree without communication. +# "a2a" - needed-only all-to-all-v (general case, handles docs spanning >2 ranks). +_CP_KV_EXCHANGE = KV_EXCHANGE_DEFAULT + + +def configure_cp_varlen(*, attn_backend: str = "flash", kv_exchange: str = "allgather") -> None: + """Configure the block-diagonal CP attention path from parsed runtime config. + + Args: + attn_backend: Varlen kernel selection; any synonym accepted by + :func:`normalize_attn_backend` (default ``"flash"``). + kv_exchange: K/V delivery mode; any synonym accepted by + :func:`normalize_kv_exchange` (default ``"allgather"``). + """ + backend = normalize_attn_backend(attn_backend) + exchange = normalize_kv_exchange(kv_exchange) + + global _CP_ATTN_BACKEND, _CP_KV_EXCHANGE + changed = (_CP_ATTN_BACKEND, _CP_KV_EXCHANGE) != (backend, exchange) + _CP_ATTN_BACKEND = backend + _CP_KV_EXCHANGE = exchange + if changed: + # Keep the one-shot kernel marker in its owning module. Import lazily to + # avoid a state <-> kernels import cycle during module initialization. + from nemo_automodel.components.distributed.blockdiag_cp import kernels + + kernels._CP_FLASH_ENGAGED = False + + +def cp_varlen_runtime_config() -> dict[str, str]: + """Return the currently configured block-diagonal CP runtime settings. + + Returns: + A dict with keys ``"attn_backend"`` and ``"kv_exchange"`` holding the + canonical configured values. + """ + return { + "attn_backend": _CP_ATTN_BACKEND, + "kv_exchange": _CP_KV_EXCHANGE, + } + + +class _ThreadSharedVar: + """A ``contextvars.ContextVar``-style get/set/reset holder visible ACROSS threads. + + The block-diagonal CP state must be readable inside the autograd worker thread + that runs activation-checkpointing recompute during backward. A real + ``ContextVar`` is per-thread and reads its default (None) there, which would + silently drop the CP state mid-recompute -- the softmax SDPA would fall back to + local-only attention. Training steps are sequential, so a single shared slot + with token-based restore (supporting nesting) is safe; backward only ever READS + the slot. + """ + + def __init__(self): + self._value = None + + def get(self): + """Return the current value (``None`` when unset).""" + return self._value + + def set(self, value): + """Set the value; returns a token (the previous value) for :meth:`reset`.""" + token = self._value + self._value = value + return token + + def reset(self, token): + """Restore the value captured by a previous :meth:`set`.""" + self._value = token + + +# Per-step state read by the block-diagonal SDPA (set by make_cp_blockdiag_batch_and_ctx). +# Holds: {"group", "doc_ids" [B, S_full] int, "row_offset" int, "seq_dim": 2, ...}. +_CP_BLOCKDIAG_STATE = _ThreadSharedVar() + +# Sentinel: counts block-diagonal CP attention invocations (the CP path actually +# running) so a recipe can fail loud if cp>1 but the hook never fired. A plain +# list cell suffices: the forward runs single-threaded; AC-recompute increments in +# backward are post-check and only matter as ">0". +_CP_ATTN_FIRE_COUNT: list[int] = [0] + + +def reset_cp_attn_fire_count() -> None: + """Zero the CP-attention fire counter (call before each forward).""" + _CP_ATTN_FIRE_COUNT[0] = 0 + + +def cp_attn_fire_count() -> int: + """Number of block-diagonal CP attention calls since the last reset.""" + return _CP_ATTN_FIRE_COUNT[0] From 67c556f38d2c773e01ed8c7dc12d378a62387422 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Thu, 9 Jul 2026 08:44:40 -0700 Subject: [PATCH 2/8] test(distributed): parity tests for block-diagonal varlen CP - CPU unit tests: block-causal mask construction, CP-vs-full attention parity (simulated ranks, incl. GQA), all-gather forward/backward, all-padding-rank gradient attachment, varlen metadata validation, halo/a2a exchange-plan geometry, KV path-selection downgrades, the flash long-prefix guard, and the cp1 packed hook contract - 1-GPU parity tests (skipped without CUDA): flash and TE varlen outputs vs the dense-mask SDPA reference across multi-doc packs with padding tails, GQA layouts, rank-straddling documents, single documents spanning all ranks, and head_dim 256; bit-identity of the per-step precomputed metadata against inline segmentation; the cp_size==1 packed-sequence SDPA hook Co-Authored-By: Claude Fable 5 Signed-off-by: Yuhe Zhang --- .../distributed/test_blockdiag_cp.py | 599 ++++++++++++++++++ .../test_blockdiag_cp_varlen_gpu.py | 198 ++++++ 2 files changed, 797 insertions(+) create mode 100644 tests/unit_tests/distributed/test_blockdiag_cp.py create mode 100644 tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py diff --git a/tests/unit_tests/distributed/test_blockdiag_cp.py b/tests/unit_tests/distributed/test_blockdiag_cp.py new file mode 100644 index 0000000000..54d7154303 --- /dev/null +++ b/tests/unit_tests/distributed/test_blockdiag_cp.py @@ -0,0 +1,599 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for the block-diagonal context-parallelism core. + +These are the regression guards for CP *numerical correctness*: the headline test +proves that the block-diagonal CP attention (per-rank local queries against +all-gathered K/V with a per-document causal mask) is equivalent to plain full +attention with a block-causal mask -- i.e. cp_size>1 does not change the math vs +cp_size=1. No GPU / no distributed init required: the all-gather is mocked to +identity (the test passes the full K/V) and the CP ranks are simulated +in-process. +""" + +import sys +import types + +import pytest +import torch + +from nemo_automodel.components.distributed.blockdiag_cp import batch as bd_batch +from nemo_automodel.components.distributed.blockdiag_cp import exchange as bd_exchange +from nemo_automodel.components.distributed.blockdiag_cp import kernels as bd_kernels +from nemo_automodel.components.distributed.blockdiag_cp import packed as bd_packed +from nemo_automodel.components.distributed.blockdiag_cp import runtime as bd_runtime +from nemo_automodel.components.distributed.blockdiag_cp import state as bd_state + + +def _doc_ids(): + # 2 documents (positions 0-2 and 3-5) + 2 padding positions (6-7); length 8 so it + # shards evenly for world in {2, 4}. + return torch.tensor([[1, 1, 1, 2, 2, 2, 0, 0]], dtype=torch.long) + + +class _IdentityGather: + """Stand-in for _AllGatherSeqDiff whose apply() returns its input unchanged, so each + simulated CP rank sees the FULL (already-gathered) K/V we pass in.""" + + @staticmethod + def apply(x, group, seq_dim): + return x + + +def _run_blockdiag_cp(Q, K, V, doc_ids, world, enable_gqa=False): + """Simulate ``world`` CP ranks in-process; returns the concatenated local outputs. + + Args: + Q: Full queries ``[B, Hq, S, D]`` (sliced per simulated rank). + K: Full keys ``[B, Hkv, S, D]`` (the identity all-gather passes them through). + V: Full values ``[B, Hkv, S, D]``. + doc_ids: Per-position document ids ``[B, S]`` (0 == padding). + world: Number of simulated CP ranks. + enable_gqa: Forwarded to the SDPA under test. + + Returns: + Concatenated per-rank outputs ``[B, Hq, S, D]``. + """ + B, Hq, S, D = Q.shape + assert S % world == 0 + L = S // world + orig = bd_exchange._AllGatherSeqDiff + bd_exchange._AllGatherSeqDiff = _IdentityGather # identity all-gather: we pass full K/V below + outs = [] + try: + for r in range(world): + state = { + "group": None, + "doc_ids": doc_ids, + "row_offset": r * L, + "seq_dim": 2, + } + token = bd_state._CP_BLOCKDIAG_STATE.set(state) + try: + q_local = Q[:, :, r * L : (r + 1) * L, :] + outs.append(bd_runtime.cp_blockdiag_sdpa(q_local, K, V, enable_gqa=enable_gqa)) + finally: + bd_state._CP_BLOCKDIAG_STATE.reset(token) + finally: + bd_exchange._AllGatherSeqDiff = orig + return torch.cat(outs, dim=2) + + +def _full_attention(Q, K, V, doc_ids, enable_gqa=False): + """Reference: plain full attention ``[B, Hq, S, D]`` with the block-causal (per-document) mask. + + Args: + Q: Full queries ``[B, Hq, S, D]``. + K: Full keys ``[B, Hkv, S, D]`` (repeat-interleaved when ``enable_gqa``). + V: Full values ``[B, Hkv, S, D]``. + doc_ids: Per-position document ids ``[B, S]`` (0 == padding). + enable_gqa: Expand K/V heads to match Q before the reference SDPA. + """ + B, Hq, S, D = Q.shape + mask = bd_kernels._cp_blockdiag_mask(doc_ids, 0, S, S, B) # [B, 1, S, S] + K2, V2 = K, V + if enable_gqa and K.shape[1] != Hq: + n = Hq // K.shape[1] + K2 = K.repeat_interleave(n, dim=1) + V2 = V.repeat_interleave(n, dim=1) + return bd_runtime._ORIGINAL_SDPA(Q, K2, V2, attn_mask=mask, is_causal=False) + + +def test_knob_normalization_synonyms_and_defaults(): + assert bd_state.normalize_attn_backend(None) == "flash" + assert bd_state.normalize_attn_backend("AUTO") == "flash" + assert bd_state.normalize_attn_backend("flash_attention_2") == "flash" + assert bd_state.normalize_attn_backend("transformer_engine") == "te" + assert bd_state.normalize_attn_backend("sdpa") == "dense" + assert bd_state.normalize_kv_exchange(None) == "allgather" + assert bd_state.normalize_kv_exchange("all-gather") == "allgather" + assert bd_state.normalize_kv_exchange("needed_only") == "halo" + assert bd_state.normalize_kv_exchange("all_to_all") == "a2a" + with pytest.raises(ValueError, match="attention backend"): + bd_state.normalize_attn_backend("bogus") + with pytest.raises(ValueError, match="kv_exchange"): + bd_state.normalize_kv_exchange("bogus") + + +def test_configure_cp_varlen_roundtrip(): + prev = bd_state.cp_varlen_runtime_config() + try: + bd_state.configure_cp_varlen(attn_backend="te", kv_exchange="halo") + assert bd_state.cp_varlen_runtime_config() == {"attn_backend": "te", "kv_exchange": "halo"} + finally: + bd_state.configure_cp_varlen(**prev) + + +def test_blockdiag_mask_shards_concat_to_full(): + """Per-rank masks (query rows sliced) concatenated == the full block-causal mask.""" + doc_ids = _doc_ids() + S = doc_ids.shape[1] + full = bd_kernels._cp_blockdiag_mask(doc_ids, 0, S, S, 1) # [1, 1, S, S] + for world in (2, 4): + L = S // world + shards = [bd_kernels._cp_blockdiag_mask(doc_ids, r * L, L, S, 1) for r in range(world)] + recombined = torch.cat(shards, dim=2) + assert recombined.shape == full.shape + assert torch.equal(recombined, full), f"mask mismatch at world={world}" + + +def test_blockdiag_mask_expected_matrix(): + doc_ids = _doc_ids() + expected = torch.tensor( + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 0, 0, 0, 0], + [0, 0, 0, 1, 1, 0, 0, 0], + [0, 0, 0, 1, 1, 1, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 1], + ], + dtype=torch.bool, + ).view(1, 1, 8, 8) + + got = bd_kernels._cp_blockdiag_mask(doc_ids, 0, 8, 8, 1) + assert torch.equal(got, expected) + + +@pytest.mark.parametrize("world", [2, 4]) +def test_blockdiag_sdpa_parity_vs_full_attention(world): + """cp=world block-diagonal attention == cp=1 full attention on identical inputs.""" + torch.manual_seed(0) + B, H, S, D = 1, 4, 8, 16 + Q = torch.randn(B, H, S, D, dtype=torch.float32) + K = torch.randn(B, H, S, D, dtype=torch.float32) + V = torch.randn(B, H, S, D, dtype=torch.float32) + doc_ids = _doc_ids() + + cp_out = _run_blockdiag_cp(Q, K, V, doc_ids, world=world, enable_gqa=False) + full_out = _full_attention(Q, K, V, doc_ids, enable_gqa=False) + + assert cp_out.shape == full_out.shape == (B, H, S, D) + max_diff = (cp_out - full_out).abs().max().item() + assert max_diff < 1e-5, f"world={world} CP-vs-full SDPA max_diff={max_diff}" + + +@pytest.mark.parametrize("world", [2, 4]) +def test_blockdiag_sdpa_parity_gqa(world): + """GQA path (n_kv_heads < n_q_heads): CP repeat-interleave matches full attention.""" + torch.manual_seed(1) + B, Hq, Hkv, S, D = 1, 4, 2, 8, 16 + Q = torch.randn(B, Hq, S, D, dtype=torch.float32) + K = torch.randn(B, Hkv, S, D, dtype=torch.float32) + V = torch.randn(B, Hkv, S, D, dtype=torch.float32) + doc_ids = _doc_ids() + + cp_out = _run_blockdiag_cp(Q, K, V, doc_ids, world=world, enable_gqa=True) + full_out = _full_attention(Q, K, V, doc_ids, enable_gqa=True) + + assert cp_out.shape == full_out.shape == (B, Hq, S, D) + max_diff = (cp_out - full_out).abs().max().item() + assert max_diff < 1e-5, f"GQA world={world} CP-vs-full SDPA max_diff={max_diff}" + + +def test_blockdiag_sdpa_noop_without_state(): + """With no CP state set, cp_blockdiag_sdpa is a plain pass-through to stock SDPA.""" + torch.manual_seed(2) + Q = torch.randn(1, 2, 4, 8) + out = bd_runtime.cp_blockdiag_sdpa(Q, Q, Q, is_causal=True) + ref = bd_runtime._ORIGINAL_SDPA(Q, Q, Q, is_causal=True) + assert torch.allclose(out, ref, atol=1e-6) + + +def test_blockdiag_batch_synthesizes_and_shards_padding_mask(monkeypatch): + class _Mesh: + def size(self): + return 2 + + def get_local_rank(self): + return 1 + + def get_group(self): + return object() + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) + + batch = { + "inputs_embeds": torch.randn(1, 5, 3), + "labels": torch.arange(5).view(1, 5), + "_packed_seq_ids": torch.tensor([[1, 1, 2, 0, 0]], dtype=torch.long), + } + + ctx, sharded = bd_batch.make_cp_blockdiag_batch_and_ctx(_Mesh(), None, batch) + + # rank 1 of 2 on a padded length-6 sequence -> local rows [3, 6): pad-only tail. + assert torch.equal(sharded["padding_mask"], torch.tensor([[True, True, True]])) + assert torch.equal(sharded["labels"], torch.tensor([[3, 4, -100]])) + with ctx(): + state = bd_state._CP_BLOCKDIAG_STATE.get() + assert state is not None + assert state["row_offset"] == 3 + assert torch.equal(state["doc_ids"], torch.tensor([[1, 1, 2, 0, 0, 0]])) + assert state["varlen_meta"]["n_real"] == 0 # all-padding local chunk + assert bd_state._CP_BLOCKDIAG_STATE.get() is None # ctx exit restores the slot + + +def test_blockdiag_batch_shards_loss_mask(monkeypatch): + class _Mesh: + def size(self): + return 2 + + def get_local_rank(self): + return 0 + + def get_group(self): + return object() + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) + + batch = { + "inputs_embeds": torch.randn(1, 6, 3), + "labels": torch.arange(6).view(1, 6), + "_packed_seq_ids": torch.tensor([[1, 1, 1, 2, 2, 0]], dtype=torch.long), + } + loss_mask = torch.tensor([[1, 1, 1, 1, 1, 0]], dtype=torch.long) + + _, sharded = bd_batch.make_cp_blockdiag_batch_and_ctx(_Mesh(), None, batch, loss_mask=loss_mask) + + assert torch.equal(sharded["loss_mask"], torch.tensor([[1, 1, 1]])) + assert sharded["inputs_embeds"].shape == (1, 3, 3) + + +def test_blockdiag_batch_rejects_batch_size_gt_one(monkeypatch): + class _Mesh: + def size(self): + return 2 + + def get_local_rank(self): + return 0 + + def get_group(self): + return object() + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) + + batch = { + "inputs_embeds": torch.randn(2, 4, 3), + "labels": torch.zeros(2, 4, dtype=torch.long), + } + with pytest.raises(ValueError, match="local_batch_size=1"): + bd_batch.make_cp_blockdiag_batch_and_ctx(_Mesh(), None, batch) + + +def test_blockdiag_doc_ids_from_2d_attention_mask(): + batch = {"attention_mask": torch.tensor([[1, 1, 1, 0, 0]], dtype=torch.long)} + got = bd_batch._cp_blockdiag_doc_ids(batch, seq_len=5, device=torch.device("cpu"), batch_size=1) + assert torch.equal(got, torch.tensor([[1, 1, 1, 0, 0]], dtype=torch.long)) + + +@pytest.mark.parametrize("world", [2, 3, 4]) +def test_allgather_seqdiff_forward_backward(monkeypatch, world): + """_AllGatherSeqDiff: forward concatenates shards over the group; backward + reduce-scatters the summed gradient. Validates world>2 (the all-gather is generic).""" + + def fake_all_gather(out_list, x, group=None): + for o in out_list: + o.copy_(x) # simulate every rank holding the same shard + + def fake_reduce_scatter(local, chunks, op=None, group=None): + local.copy_(sum(chunks)) + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world) + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + monkeypatch.setattr(torch.distributed, "reduce_scatter", fake_reduce_scatter) + + x = torch.randn(1, 2, 4, 8, requires_grad=True) + out = bd_exchange._AllGatherSeqDiff.apply(x, None, 2) + assert out.shape == (1, 2, 4 * world, 8) + assert torch.allclose(out[:, :, :4], x.detach()) # first shard == this rank's x + + out.sum().backward() + assert x.grad is not None and x.grad.shape == x.shape + assert torch.equal(x.grad, torch.full_like(x, world)) + + +def test_varlen_all_padding_rank_keeps_grad_attachment(): + """An all-padding CP rank must keep its output attached to K/V. + + Otherwise that rank skips the all-gather's backward reduce_scatter while + real-token ranks fire it -> collective desync -> NCCL hang in backward. + """ + doc_ids = torch.tensor([[1, 1, 0, 0]], dtype=torch.long) + q = torch.randn(1, 2, 2, 8, dtype=torch.bfloat16, requires_grad=True) + k = torch.randn(1, 2, 4, 8, dtype=torch.bfloat16, requires_grad=True) + v = torch.randn(1, 2, 4, 8, dtype=torch.bfloat16, requires_grad=True) + + out = bd_kernels._cp_blockdiag_varlen(q, k, v, doc_ids, row_offset=2, backend="te") + + assert out is not None + assert out.requires_grad + assert torch.equal(out, torch.zeros_like(out)) + out.float().sum().backward() + assert k.grad is not None and v.grad is not None + assert torch.equal(k.grad, torch.zeros_like(k.grad)) + assert torch.equal(v.grad, torch.zeros_like(v.grad)) + + +def test_varlen_metadata_guard_accepts_valid_left_straddle(): + meta = { + "n_real": 5, + "s_first": 0, + "real_end": 8, + "cu_q": torch.tensor([0, 2, 5], dtype=torch.int32), + "cu_k": torch.tensor([0, 5, 8], dtype=torch.int32), + "max_q": 3, + "max_k": 5, + } + + reason = bd_kernels._varlen_metadata_unavailable_reason( + meta, + query_len=8, + key_len=8, + device=torch.device("cpu"), + ) + + assert reason is None + assert meta["_validation_cache"] + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("cu_q", torch.tensor([0, 2, 6], dtype=torch.int32), "n_real"), + ("cu_k", torch.tensor([0, 5, 9], dtype=torch.int32), "slice length"), + ("max_k", 4, "max_seqlen mismatch"), + ], +) +def test_varlen_metadata_guard_rejects_kernel_oob_inputs(field, value, message): + meta = { + "n_real": 5, + "s_first": 0, + "real_end": 8, + "cu_q": torch.tensor([0, 2, 5], dtype=torch.int32), + "cu_k": torch.tensor([0, 5, 8], dtype=torch.int32), + "max_q": 3, + "max_k": 5, + } + meta[field] = value + + reason = bd_kernels._varlen_metadata_unavailable_reason( + meta, + query_len=8, + key_len=8, + device=torch.device("cpu"), + ) + + assert reason is not None + assert message in reason + + +def test_precompute_meta_matches_mask_segmentation(): + """The precomputed cu_seqlens reproduce the per-rank straddle geometry.""" + # doc 1 spans [0, 6) so rank 1 (rows [4, 8)) has a 4-token left straddle. + doc_ids = torch.tensor([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 0], dtype=torch.long) + meta = bd_kernels.precompute_blockdiag_varlen_meta(doc_ids, row_offset=4, local_len=4, device=torch.device("cpu")) + assert meta["n_real"] == 4 + assert meta["s_first"] == 0 # doc 1 starts at position 0 + assert meta["real_end"] == 8 + assert torch.equal(meta["cu_q"], torch.tensor([0, 2, 4], dtype=torch.int32)) + assert torch.equal(meta["cu_k"], torch.tensor([0, 6, 8], dtype=torch.int32)) + assert meta["first_q"] == 2 and meta["first_k"] == 6 + + # all-padding chunk -> sentinel meta + pad_meta = bd_kernels.precompute_blockdiag_varlen_meta( + doc_ids, row_offset=11, local_len=1, device=torch.device("cpu") + ) + assert pad_meta == {"n_real": 0} + + +def test_kv_plan_halo_vs_a2a_decision(): + """Docs fitting in <=2 ranks -> halo; a doc spanning >2 ranks -> a2a fallback.""" + dev = torch.device("cpu") + # world=4, local_len=4: doc boundaries at 6 and 12 -> every straddle <= 4. + doc_ids = torch.tensor([1] * 6 + [2] * 6 + [3] * 4, dtype=torch.long) + plan = bd_exchange._compute_blockdiag_kv_plan(doc_ids, world=4, local_len=4, dev=dev) + assert plan["use_halo"] + # rank 1 straddles doc 1 back to position 0 (back=4); rank 2 straddles doc 2 by 2. + assert plan["recv"] == [0, 4, 2, 0] + assert plan["send"] == [4, 2, 0, 0] + assert plan["nreal"] == [4, 4, 4, 4] + + # one doc across the whole sequence: rank 3 straddles 12 > local_len -> no halo. + single = torch.ones(16, dtype=torch.long) + plan_single = bd_exchange._compute_blockdiag_kv_plan(single, world=4, local_len=4, dev=dev) + assert not plan_single["use_halo"] + assert plan_single["recv"] == [0, 4, 8, 12] + assert plan_single["s_first"] == [0, 0, 0, 0] + + +def test_needed_kv_a2a_plan_split_symmetry(): + """Every rank's a2a send splits mirror the receivers' out splits and cover the needed ranges.""" + dev = torch.device("cpu") + world, local_len = 4, 4 + doc_ids = torch.ones(world * local_len, dtype=torch.long) # single 16-token doc + plan = bd_exchange._compute_blockdiag_kv_plan(doc_ids, world=world, local_len=local_len, dev=dev) + + in_splits_all, out_splits_all = [], [] + for r in range(world): + in_splits, out_splits, send_index = bd_exchange._needed_kv_a2a_plan(plan, r, world, local_len, dev) + in_splits_all.append(in_splits) + out_splits_all.append(out_splits) + assert send_index.numel() == sum(in_splits) + # each rank receives exactly its needed range [s_first, real_end) + assert sum(out_splits) == plan["real_end"][r] - plan["s_first"][r] + for src in range(world): + for dst in range(world): + assert in_splits_all[src][dst] == out_splits_all[dst][src] + + +def test_select_kv_exchange_path_downgrades_name_reasons(): + """The path selector names every downgrade to all-gather.""" + state = {"attn_backend": "flash", "kv_exchange": "allgather", "varlen_meta": {"n_real": 1}} + path, plan, reason = bd_runtime._select_kv_exchange_path( + state, None, torch.ones(1, 8, dtype=torch.long), 4, torch.device("cpu"), 0 + ) + assert path == "allgather" and plan is None and "mode=allgather" in reason + + state = {"attn_backend": "dense", "kv_exchange": "halo", "varlen_meta": {"n_real": 1}} + path, _, reason = bd_runtime._select_kv_exchange_path( + state, None, torch.ones(1, 8, dtype=torch.long), 4, torch.device("cpu"), 0 + ) + assert path == "allgather" and "flash/te" in reason + + state = {"attn_backend": "flash", "kv_exchange": "halo", "varlen_meta": None} + path, _, reason = bd_runtime._select_kv_exchange_path( + state, None, torch.ones(1, 8, dtype=torch.long), 4, torch.device("cpu"), 0 + ) + assert path == "allgather" and "varlen_meta" in reason + + +def test_flash_long_prefix_guard_peels_only_boundary_segment(monkeypatch): + calls = [] + + def fake_fixed(q, k, v, **kwargs): + calls.append(("fixed", q.shape[1], k.shape[1], kwargs["causal"])) + return q + 1 + + def fake_varlen(q, k, v, **kwargs): + calls.append(("varlen", q.shape[0], k.shape[0], kwargs["max_seqlen_q"], kwargs["max_seqlen_k"])) + return q + 2 + + flash_attn = types.ModuleType("flash_attn") + flash_attn.flash_attn_func = fake_fixed + flash_attn.flash_attn_varlen_func = fake_varlen + monkeypatch.setitem(sys.modules, "flash_attn", flash_attn) + + q = torch.zeros(6, 4, 8, dtype=torch.bfloat16) + k = torch.zeros(11, 2, 8, dtype=torch.bfloat16) + v = torch.zeros_like(k) + cu_q = torch.tensor([0, 2, 6], dtype=torch.int32) + cu_k = torch.tensor([0, 7, 11], dtype=torch.int32) + out = bd_kernels._flash_varlen_with_long_prefix_guard( + q, + k, + v, + cu_q=cu_q, + cu_k=cu_k, + max_q=4, + max_k=7, + local_query_len=4, + scale=0.5, + meta={"first_q": 2, "first_k": 7, "max_tail": 4}, + ) + + assert calls == [("fixed", 2, 7, True), ("varlen", 4, 4, 4, 4)] + assert torch.equal(out[:2], torch.ones_like(out[:2])) + assert torch.equal(out[2:], torch.full_like(out[2:], 2)) + + +def test_flash_long_prefix_guard_keeps_normal_single_varlen_call(monkeypatch): + calls = [] + + def fake_varlen(q, k, v, **kwargs): + calls.append((q.shape[0], k.shape[0])) + return q + + flash_attn = types.ModuleType("flash_attn") + flash_attn.flash_attn_func = lambda *args, **kwargs: pytest.fail("fixed Flash must not run for a normal segment") + flash_attn.flash_attn_varlen_func = fake_varlen + monkeypatch.setitem(sys.modules, "flash_attn", flash_attn) + + q = torch.zeros(6, 4, 8, dtype=torch.bfloat16) + k = torch.zeros(6, 2, 8, dtype=torch.bfloat16) + out = bd_kernels._flash_varlen_with_long_prefix_guard( + q, + k, + k, + cu_q=torch.tensor([0, 2, 6], dtype=torch.int32), + cu_k=torch.tensor([0, 2, 6], dtype=torch.int32), + max_q=4, + max_k=4, + local_query_len=4, + scale=None, + meta={}, + ) + + assert calls == [(6, 6)] + assert out is q + + +def test_cp1_packed_precomputes_varlen_metadata_once_per_forward(monkeypatch): + """Every CP1 attention layer reuses the metadata armed by the outer forward.""" + doc_ids = torch.tensor([[1, 1, 2, 2, 0]], dtype=torch.long) + meta = {"n_real": 4, "sentinel": object()} + precompute_calls = [] + forwarded_meta = [] + + def fake_precompute(got_doc_ids, row_offset, local_len, device): + precompute_calls.append((got_doc_ids, row_offset, local_len, device)) + return meta + + def fake_varlen(query, key, value, doc_ids_arg, row_offset, scale, backend, *, meta=None): + forwarded_meta.append(meta) + return query + + monkeypatch.setattr(bd_packed.kernels, "precompute_blockdiag_varlen_meta", fake_precompute) + monkeypatch.setattr(bd_packed.kernels, "_cp_blockdiag_varlen", fake_varlen) + + bd_packed.enable_cp1_packed_varlen(doc_ids, "flash") + try: + qkv = torch.randn(1, 2, 5, 4) + assert bd_packed._packed_varlen_sdpa(qkv, qkv, qkv) is qkv + assert bd_packed._packed_varlen_sdpa(qkv, qkv, qkv) is qkv + finally: + bd_packed.disable_cp1_packed_varlen() + + assert len(precompute_calls) == 1 + got_doc_ids, row_offset, local_len, device = precompute_calls[0] + assert got_doc_ids is doc_ids + assert row_offset == 0 + assert local_len == doc_ids.shape[-1] + assert device == doc_ids.device + assert forwarded_meta == [meta, meta] + + +def test_cp1_packed_passes_through_non_matching_shapes(): + """SDPA calls that do not match the armed doc_ids (e.g. vision attention) pass through.""" + doc_ids = torch.tensor([[1, 1, 2, 2, 0]], dtype=torch.long) + bd_packed.enable_cp1_packed_varlen(doc_ids, "flash") + try: + q = torch.randn(1, 2, 7, 4) # seq len 7 != 5 -> pass-through + out = bd_packed._packed_varlen_sdpa(q, q, q, is_causal=True) + ref = bd_runtime._ORIGINAL_SDPA(q, q, q, is_causal=True) + assert torch.allclose(out, ref, atol=1e-6) + finally: + bd_packed.disable_cp1_packed_varlen() diff --git a/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py b/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py new file mode 100644 index 0000000000..3a006b6a4e --- /dev/null +++ b/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py @@ -0,0 +1,198 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""1-GPU parity: varlen block-diagonal CP attention (flash AND te) == dense-mask SDPA. + +The CP ranks are simulated in-process (each rank's local query chunk attends the +full K/V, exactly what the all-gather delivers), so this needs a single GPU. It +sweeps multi-document packs with padding tails, GQA head layouts, left-straddling +documents (the ``sk > sq`` boundary case), single documents spanning every rank, +and head_dim 256, for every simulated rank offset. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from nemo_automodel.components.distributed.blockdiag_cp.kernels import ( + _cp_blockdiag_mask, + _cp_blockdiag_varlen, + precompute_blockdiag_varlen_meta, +) +from nemo_automodel.components.distributed.blockdiag_cp.packed import ( + _PACKED_STATE, + enable_cp1_packed_varlen, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="varlen parity requires CUDA") + + +def _dense_ref(query, key_full, value_full, doc_ids, row_offset): + """The exact ops the dense CP path runs: GQA-repeat K/V + [B,1,L,S] mask + SDPA. + + Args: + query: Local queries ``[B, Hq, L, D]``. + key_full: Full-sequence keys ``[B, Hkv, S, D]``. + value_full: Full-sequence values ``[B, Hkv, S, D]``. + doc_ids: Per-position document ids ``[B, S]`` (0 == padding). + row_offset: Global position of the first local query row. + + Returns: + Reference attention output ``[B, Hq, L, D]``. + """ + B, Hq, L, D = query.shape + Hkv, S = key_full.shape[1], key_full.shape[2] + k = key_full + v = value_full + if Hkv != Hq: + n_rep = Hq // Hkv + k = key_full.repeat_interleave(n_rep, dim=1) + v = value_full.repeat_interleave(n_rep, dim=1) + allow = _cp_blockdiag_mask(doc_ids, row_offset, L, S, B) # [B,1,L,S] bool + return F.scaled_dot_product_attention(query, k, v, attn_mask=allow) + + +def _make_doc_ids(seg_lens, pad, S, device): + """Build ``[1, S]`` document ids: 1-based per-document runs + a zero pad tail.""" + ids = [] + for d, n in enumerate(seg_lens, start=1): + ids += [d] * n + ids += [0] * pad + assert len(ids) == S, (len(ids), S) + return torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0) # [1,S] + + +# (name, seg_lens, pad, S, cp, Hq, Hkv, D) +_CASES = [ + ("multidoc+pad cp2", [40, 30, 25], 33, 128, 2, 8, 8, 128), + ("multidoc+pad cp4", [40, 30, 25], 33, 128, 4, 8, 8, 128), + ("multidoc+pad cp8", [40, 30, 25], 33, 128, 8, 8, 8, 128), + ("gqa cp4", [50, 40, 38], 0, 128, 4, 8, 2, 128), + ("gqa cp8 +pad", [50, 40, 20], 18, 128, 8, 8, 2, 128), + ("single-doc cp4", [128], 0, 128, 4, 8, 8, 128), + ("straddle cp8", [70, 20, 20], 18, 128, 8, 8, 8, 128), + ("hd256 gqa straddle cp8", [70, 20, 20], 18, 128, 8, 16, 2, 256), + ("heavy-pad cp4", [20, 10], 34, 64, 4, 8, 8, 128), +] + + +def _backend_or_skip(backend: str) -> None: + if backend == "flash": + pytest.importorskip("flash_attn") + else: + pytest.importorskip("transformer_engine.pytorch") + + +@requires_cuda +@pytest.mark.parametrize("backend", ["flash", "te"]) +@pytest.mark.parametrize(("name", "seg_lens", "pad", "S", "cp", "Hq", "Hkv", "D"), _CASES) +def test_varlen_blockdiag_parity(backend, name, seg_lens, pad, S, cp, Hq, Hkv, D): + """Varlen output matches the dense-mask reference on every simulated CP rank.""" + _backend_or_skip(backend) + device = "cuda" + dtype = torch.bfloat16 + atol = 2e-2 + torch.manual_seed(0) + B = 1 + doc_ids = _make_doc_ids(seg_lens, pad, S, device) + key_full = torch.randn(B, Hkv, S, D, device=device, dtype=dtype) + value_full = torch.randn(B, Hkv, S, D, device=device, dtype=dtype) + L = S // cp + for r in range(cp): + off = r * L + q = torch.randn(B, Hq, L, D, device=device, dtype=dtype) + ref = _dense_ref(q, key_full, value_full, doc_ids, off) # [B,Hq,L,D] + got = _cp_blockdiag_varlen(q, key_full, value_full, doc_ids, off, backend=backend) + assert got is not None, f"{backend} path returned None (import/dtype fallback)" + # compare only REAL local rows (doc id > 0); padding rows are zeros by design + local = doc_ids[0, off : off + L] + real = local > 0 + if real.sum() == 0: + # all-padding shard: output must be all zeros + md = got.abs().max().item() + assert md == 0.0, f"{name} r{r}: all-pad shard not zero ({md})" + continue + rr = real.nonzero().flatten() + d = (ref[0, :, rr, :].float() - got[0, :, rr, :].float()).abs().max().item() + assert d < atol, f"[{backend}] {name} r{r} off{off}: max|diff|={d:.4f} >= {atol}" + + +@requires_cuda +@pytest.mark.parametrize("backend", ["flash"]) +@pytest.mark.parametrize(("name", "seg_lens", "pad", "S", "cp", "Hq", "Hkv", "D"), _CASES) +def test_precomputed_meta_bit_identical(backend, name, seg_lens, pad, S, cp, Hq, Hkv, D): + """Feeding a precomputed meta yields BIT-IDENTICAL output to the inline path. + + The segmentation depends only on (doc_ids, row_offset, local_len) -- all + step-constant -- so the per-step precompute must reproduce the exact same + kernel launch (same cu_seqlens) as the inline per-call segmentation. + """ + _backend_or_skip(backend) + device = "cuda" + torch.manual_seed(0) + doc_ids = _make_doc_ids(seg_lens, pad, S, device) + key_full = torch.randn(1, Hkv, S, D, device=device, dtype=torch.bfloat16) + value_full = torch.randn(1, Hkv, S, D, device=device, dtype=torch.bfloat16) + L = S // cp + for r in range(cp): + off = r * L + q = torch.randn(1, Hq, L, D, device=device, dtype=torch.bfloat16) + base = _cp_blockdiag_varlen(q, key_full, value_full, doc_ids, off, backend=backend) + meta = precompute_blockdiag_varlen_meta(doc_ids, off, L, device) + fast = _cp_blockdiag_varlen(q, key_full, value_full, doc_ids, off, backend=backend, meta=meta) + assert base is not None and fast is not None, f"{name} r{r}: None output" + assert torch.equal(base, fast), ( + f"[{backend}] {name} r{r} off{off}: precomputed-meta output differs " + f"max|diff|={(base.float() - fast.float()).abs().max().item()}" + ) + + +@requires_cuda +@pytest.mark.parametrize("backend", ["flash", "te"]) +@pytest.mark.parametrize( + ("seg_lens", "pad", "S", "Hq", "Hkv", "D"), + [ + ([40, 30, 25], 33, 128, 8, 8, 128), + ([50, 40, 20], 18, 128, 16, 2, 256), # GQA 8:1, head_dim 256 + ], +) +def test_cp1_packed_hook_parity(backend, seg_lens, pad, S, Hq, Hkv, D): + """cp1 packed path: F.sdpa under the armed varlen hook == dense block-diagonal SDPA. + + Exercises the full hook chain (enable_cp1_packed_varlen global patch -> + _packed_varlen_sdpa -> _cp_blockdiag_varlen at row_offset=0), as used by + packed cp_size==1 runs. The whole packed sequence is on one rank + (q == k == full). The hook persists (no ctx); state is reset after for + test isolation. + """ + _backend_or_skip(backend) + device = "cuda" + dtype = torch.bfloat16 + atol = 2e-2 + torch.manual_seed(0) + doc_ids = _make_doc_ids(seg_lens, pad, S, device) # [1, S] + q = torch.randn(1, Hq, S, D, device=device, dtype=dtype) + k = torch.randn(1, Hkv, S, D, device=device, dtype=dtype) + v = torch.randn(1, Hkv, S, D, device=device, dtype=dtype) + ref = _dense_ref(q, k, v, doc_ids, 0) # [1, Hq, S, D] + enable_cp1_packed_varlen(doc_ids, backend) + try: + got = F.scaled_dot_product_attention(q, k, v, scale=D**-0.5, enable_gqa=(Hkv != Hq)) + finally: + _PACKED_STATE.set(None) # disarm (patch stays installed but passes through) + real = (doc_ids[0] > 0).nonzero().flatten() + d = (ref[0, :, real, :].float() - got[0, :, real, :].float()).abs().max().item() + assert d < atol, f"[cp1-hook {backend}] max|diff|={d:.4f} >= {atol}" From 5c53d089df13fe149e74000e99629cad6cc92f69 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Thu, 9 Jul 2026 10:43:47 -0700 Subject: [PATCH 3/8] fix(distributed): replace blockdiag CP batch asserts with real exceptions The pre-embedded-input contract check and the defense-in-depth CP-group world-size guard vanish under python -O as asserts. Raise ValueError / RuntimeError instead, matching the adjacent local_batch_size guard. Co-Authored-By: Claude Fable 5 Signed-off-by: Yuhe Zhang --- .../components/distributed/blockdiag_cp/batch.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/batch.py b/nemo_automodel/components/distributed/blockdiag_cp/batch.py index a7791515fb..0415619406 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/batch.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/batch.py @@ -113,7 +113,8 @@ def make_cp_blockdiag_batch_and_ctx( rank = cp_mesh.get_local_rank() group = cp_mesh.get_group() - assert "inputs_embeds" in batch, "block-diagonal CP requires pre-embedded 'inputs_embeds' in the batch" + if "inputs_embeds" not in batch: + raise ValueError("block-diagonal CP requires pre-embedded 'inputs_embeds' in the batch") ie = batch["inputs_embeds"] B, S = ie.shape[0], ie.shape[1] @@ -159,11 +160,12 @@ def make_cp_blockdiag_batch_and_ctx( # shape crash deep in the backward AC-recompute. Later torch versions resolve # it correctly. Fail loud and early here instead of crashing in backward. _group_world = torch.distributed.get_world_size(group) - assert _group_world == world, ( - f"block-diagonal CP: K/V all-gather group world ({_group_world}) != " - f"cp_mesh.size() ({world}). The cp sub-group is mis-resolved -- a DeviceMesh " - f"flatten/slice bug seen on torch 2.8 when dp>1; upgrade torch." - ) + if _group_world != world: + raise RuntimeError( + f"block-diagonal CP: K/V all-gather group world ({_group_world}) != " + f"cp_mesh.size() ({world}). The cp sub-group is mis-resolved -- a DeviceMesh " + f"flatten/slice bug seen on torch 2.8 when dp>1; upgrade torch." + ) # Block-diagonal CP shards ONE packed sequence per rank and assumes local batch # B==1 (packing collapses many samples into a single sequence). B>1 is From bf6b00689b734ddcf99f66ba8dfd4c203ead9d69 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Thu, 9 Jul 2026 10:44:09 -0700 Subject: [PATCH 4/8] refactor(distributed): scope cp1 packed varlen SDPA patch via attention hooks enable_cp1_packed_varlen permanently monkeypatched torch.nn.functional.scaled_dot_product_attention process-wide, leaving a shape-heuristic dispatch latched for any coincidentally-shaped SDPA call. Rescope it with the repo's bounded pattern (cp_utils.attach_cp_sdpa_hooks): attach_cp1_packed_varlen_hooks installs/restores the patch via forward pre/post hooks (always_call=True) on the checkpoint-wrapped inner self_attn module, which also covers activation-checkpointing recompute in backward. enable/disable now only arm the per-forward doc_ids/backend state; the latching _PACKED_SDPA_INSTALLED global is gone. Also shrink the package __all__ to the genuine integration entry points (make_cp_blockdiag_batch_and_ctx, cp_blockdiag_sdpa, configure_cp_varlen, and the cp1 packed hook API); knob normalization, the varlen metadata precompute, and the fire counters stay module-internal (tests import them directly). Model wiring is a follow-up PR. Tests: new CPU unit test pins the scoping contract (process-wide SDPA untouched, patch live inside hooked forwards and during AC recompute, restored after); the GPU cp1 parity test now runs the full hook chain through a hooked module instead of the global patch. Co-Authored-By: Claude Fable 5 Signed-off-by: Yuhe Zhang --- .../distributed/blockdiag_cp/__init__.py | 27 +++--- .../distributed/blockdiag_cp/packed.py | 85 +++++++++++++------ .../distributed/test_blockdiag_cp.py | 70 +++++++++++++++ .../test_blockdiag_cp_varlen_gpu.py | 44 +++++++--- 4 files changed, 174 insertions(+), 52 deletions(-) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/__init__.py b/nemo_automodel/components/distributed/blockdiag_cp/__init__.py index a103623aa7..b535f63cd8 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/__init__.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/__init__.py @@ -33,37 +33,32 @@ :func:`nemo_automodel.components.distributed.cp_utils.make_cp_batch_and_ctx`: a model attaches :func:`make_cp_blockdiag_batch_and_ctx` to the batch as ``_cp_make_batch_fn`` and routes its softmax attention through :func:`cp_blockdiag_sdpa` while the returned -context is active. +context is active. For ``cp_size == 1`` packed runs, the model scopes the varlen SDPA +patch to its attention forwards with :func:`attach_cp1_packed_varlen_hooks` and arms +the per-forward state with :func:`enable_cp1_packed_varlen` / +:func:`disable_cp1_packed_varlen`. + +Only these integration entry points are exported; everything else (knob +normalization, varlen metadata precompute, fire counters, kernels) is an internal +detail of the package's modules. Model wiring lands in follow-up PRs. """ from nemo_automodel.components.distributed.blockdiag_cp.batch import make_cp_blockdiag_batch_and_ctx -from nemo_automodel.components.distributed.blockdiag_cp.kernels import precompute_blockdiag_varlen_meta from nemo_automodel.components.distributed.blockdiag_cp.packed import ( + attach_cp1_packed_varlen_hooks, cp1_packed_varlen_backend, disable_cp1_packed_varlen, enable_cp1_packed_varlen, ) from nemo_automodel.components.distributed.blockdiag_cp.runtime import cp_blockdiag_sdpa -from nemo_automodel.components.distributed.blockdiag_cp.state import ( - configure_cp_varlen, - cp_attn_fire_count, - cp_varlen_runtime_config, - normalize_attn_backend, - normalize_kv_exchange, - reset_cp_attn_fire_count, -) +from nemo_automodel.components.distributed.blockdiag_cp.state import configure_cp_varlen __all__ = [ + "attach_cp1_packed_varlen_hooks", "configure_cp_varlen", - "cp_attn_fire_count", - "cp_varlen_runtime_config", "cp_blockdiag_sdpa", "cp1_packed_varlen_backend", "disable_cp1_packed_varlen", "enable_cp1_packed_varlen", "make_cp_blockdiag_batch_and_ctx", - "normalize_attn_backend", - "normalize_kv_exchange", - "precompute_blockdiag_varlen_meta", - "reset_cp_attn_fire_count", ] diff --git a/nemo_automodel/components/distributed/blockdiag_cp/packed.py b/nemo_automodel/components/distributed/blockdiag_cp/packed.py index 7ce8e1665c..71b9de56ca 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/packed.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/packed.py @@ -18,21 +18,29 @@ sequence sharding). The whole packed sequence lives on one rank, so it degenerates to ``_cp_blockdiag_varlen`` with ``row_offset=0`` (q == k == full sequence; square; ``cu_q == cu_k``). This gives packed-sequence block-diagonal -attention to models whose softmax attention only dispatches to sdpa: the model -calls :func:`enable_cp1_packed_varlen` at the start of its forward, which (1) -stashes ``doc_ids`` and (2) routes ``F.scaled_dot_product_attention`` to the -flash/te varlen kernel -- like the CP path for cp>1. +attention to models whose softmax attention only dispatches to sdpa. + +Integration mirrors :func:`nemo_automodel.components.distributed.cp_utils.attach_cp_sdpa_hooks`: + +1. :func:`attach_cp1_packed_varlen_hooks` registers forward pre/post hooks on every + ``self_attn`` module (the checkpoint-wrapped INNER module, so the hooks also fire + during activation-checkpointing recompute). The pre-hook swaps + ``F.scaled_dot_product_attention`` for :func:`_packed_varlen_sdpa`; the post-hook + (``always_call=True``) restores stock SDPA. The patch is therefore live only + while a hooked attention forward is running -- never process-wide. +2. The model arms the per-forward state with :func:`enable_cp1_packed_varlen` + (doc_ids + backend) at the start of its forward and clears stale state before + the NEXT forward with :func:`disable_cp1_packed_varlen`. IMPORTANT -- state must survive activation-checkpointing RECOMPUTE in backward: -the patch is installed once (process-global) and the state is set per-forward and -NOT reset, so the AC worker thread re-running a layer's forward during backward -reads the SAME doc_ids and reproduces the exact varlen output. (A context manager -that reset on exit caused the forward to use varlen but the recompute to fall back -to dense -> AC "saved vs recomputed metadata" shape mismatch.) ``_ThreadSharedVar`` -makes the state visible in the autograd worker thread. The outer model clears -stale state before each new forward (:func:`disable_cp1_packed_varlen`), then the -text backend re-arms it for the current packed batch, so vision/unpacked -attention cannot inherit it. +the state is set per-forward and NOT reset at the end of the step, so the AC +worker thread re-running a layer's forward during backward reads the SAME doc_ids +and reproduces the exact varlen output. (A per-step reset caused the forward to +use varlen but the recompute to fall back to dense -> AC "saved vs recomputed +metadata" shape mismatch.) ``_ThreadSharedVar`` makes the state visible in the +autograd worker thread. Because the outer model clears stale state before each +new forward, vision/unpacked attention cannot inherit it -- and the shape check +in :func:`_packed_varlen_sdpa` passes any non-matching call straight through. """ from __future__ import annotations @@ -42,7 +50,6 @@ from nemo_automodel.components.distributed.blockdiag_cp import kernels, runtime, state _PACKED_STATE = state._ThreadSharedVar() -_PACKED_SDPA_INSTALLED = False def _packed_varlen_sdpa( @@ -144,25 +151,55 @@ def _packed_varlen_sdpa( return out +def attach_cp1_packed_varlen_hooks(model: torch.nn.Module) -> None: + """Scope the cp1 packed varlen SDPA patch to the model's attention forwards. + + Registers a forward pre-hook / post-hook pair on every ``self_attn`` module + that installs :func:`_packed_varlen_sdpa` as ``F.scaled_dot_product_attention`` + for the duration of that module's forward and restores stock SDPA afterwards + (``always_call=True``, so a raising forward cannot leak the patch). Hooks are + attached to the checkpoint-wrapped INNER module because CheckpointWrapper's + recompute bypasses ``__call__`` on the wrapper -- this is what keeps the + varlen path active during activation-checkpointing recompute in backward. + + Same bounded-patch pattern as + :func:`nemo_automodel.components.distributed.cp_utils.attach_cp_sdpa_hooks`. + Outside these hooks, ``F.scaled_dot_product_attention`` is untouched. + + Args: + model: The model whose ``self_attn`` submodules route softmax attention + through ``F.scaled_dot_product_attention``. + """ + import torch.nn.functional as F_module + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper + + def _pre_hook(module, args, kwargs): + F_module.scaled_dot_product_attention = _packed_varlen_sdpa + return args, kwargs + + def _post_hook(module, inputs, output): + F_module.scaled_dot_product_attention = runtime._ORIGINAL_SDPA + + for name, module in model.named_modules(): + if name.endswith("self_attn"): + target = module._checkpoint_wrapped_module if isinstance(module, CheckpointWrapper) else module + target.register_forward_pre_hook(_pre_hook, with_kwargs=True) + target.register_forward_hook(_post_hook, always_call=True) + + def enable_cp1_packed_varlen(doc_ids: torch.Tensor, backend: str) -> None: """Arm cp1 packed block-diagonal varlen for the rest of this step. - Idempotently installs the global SDPA patch and sets the per-forward - doc_ids/backend state. It remains armed through backward's - activation-checkpoint recomputation; the next outer model forward clears it - via :func:`disable_cp1_packed_varlen`. + Sets the per-forward doc_ids/backend state read by the SDPA patch that + :func:`attach_cp1_packed_varlen_hooks` scopes to the attention forwards. The + state remains armed through backward's activation-checkpoint recomputation; + the next outer model forward clears it via :func:`disable_cp1_packed_varlen`. Args: doc_ids: Per-position document ids ``[1, S]`` or ``[S]`` (0 == padding) over the full packed sequence. backend: Varlen kernel backend, ``"flash"`` or ``"te"``. """ - global _PACKED_SDPA_INSTALLED - if not _PACKED_SDPA_INSTALLED: - import torch.nn.functional as F_module - - F_module.scaled_dot_product_attention = _packed_varlen_sdpa - _PACKED_SDPA_INSTALLED = True # Segmentation depends only on the packed document ids, not on the layer's # Q/K/V tensors. Compute it once per outer forward so every attention # layer (and activation-checkpoint recompute) reuses the same CUDA diff --git a/tests/unit_tests/distributed/test_blockdiag_cp.py b/tests/unit_tests/distributed/test_blockdiag_cp.py index 54d7154303..a91374465c 100644 --- a/tests/unit_tests/distributed/test_blockdiag_cp.py +++ b/tests/unit_tests/distributed/test_blockdiag_cp.py @@ -597,3 +597,73 @@ def test_cp1_packed_passes_through_non_matching_shapes(): assert torch.allclose(out, ref, atol=1e-6) finally: bd_packed.disable_cp1_packed_varlen() + + +def test_cp1_packed_hooks_scope_patch_and_cover_ac_recompute(monkeypatch): + """The SDPA patch is live only inside hooked attention forwards (incl. AC recompute). + + ``attach_cp1_packed_varlen_hooks`` must (1) leave the process-wide + ``F.scaled_dot_product_attention`` untouched outside a hooked forward, (2) route + matching SDPA calls inside the attention forward through the varlen path, and + (3) fire again during activation-checkpointing recompute in backward (hooks sit + on the checkpoint-wrapped inner module). + """ + import torch.nn.functional as F + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper + + doc_ids = torch.tensor([[1, 1, 2, 2, 0]], dtype=torch.long) + varlen_calls = [] + patched_at_forward_entry = [] + + def fake_varlen(query, key, value, doc_ids_arg, row_offset, scale, backend, *, meta=None): + varlen_calls.append(backend) + return query * 2.0 + + monkeypatch.setattr(bd_packed.kernels, "_cp_blockdiag_varlen", fake_varlen) + + class _Attn(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(4, 4, bias=False) + + def forward(self, x): + """Project ``x`` ``[B, S, E=4]`` to ``[B, H=2, S, D=2]``, run SDPA, return ``[B, S, E]``.""" + # Record the patch state at forward entry: AC recompute may early-stop the + # body once every saved tensor is rebuilt, but it always enters here. + patched_at_forward_entry.append(F.scaled_dot_product_attention is bd_packed._packed_varlen_sdpa) + B, S, _ = x.shape + q = self.proj(x).view(B, S, 2, 2).transpose(1, 2) # [B, H, S, D] + o = F.scaled_dot_product_attention(q, q, q) + return o.transpose(1, 2).reshape(B, S, 4) + + class _Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = checkpoint_wrapper(_Attn()) + + def forward(self, x): + """Run the checkpoint-wrapped attention on ``x`` ``[B, S, E]``.""" + return self.self_attn(x) + + model = _Block() + bd_packed.attach_cp1_packed_varlen_hooks(model) + bd_packed.enable_cp1_packed_varlen(doc_ids, "flash") + try: + # Arming state does NOT patch process-wide SDPA. + assert F.scaled_dot_product_attention is bd_runtime._ORIGINAL_SDPA + x = torch.randn(1, 5, 4, requires_grad=True) + out = model(x) + assert patched_at_forward_entry == [True] # patch was live inside the attention forward + assert varlen_calls == ["flash"] # ...and routed the matching SDPA call to varlen + assert F.scaled_dot_product_attention is bd_runtime._ORIGINAL_SDPA # restored after forward + out.sum().backward() # AC recompute re-enters the inner forward -> hooks fire again + assert patched_at_forward_entry == [True, True] + assert x.grad is not None + assert F.scaled_dot_product_attention is bd_runtime._ORIGINAL_SDPA + # A doc_ids-shape-matching SDPA call OUTSIDE any hooked forward is untouched. + q = torch.randn(1, 2, 5, 2) + n_calls = len(varlen_calls) + assert torch.allclose(F.scaled_dot_product_attention(q, q, q), bd_runtime._ORIGINAL_SDPA(q, q, q)) + assert len(varlen_calls) == n_calls + finally: + bd_packed.disable_cp1_packed_varlen() diff --git a/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py b/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py index 3a006b6a4e..026d0c264d 100644 --- a/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py +++ b/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py @@ -27,15 +27,17 @@ import torch import torch.nn.functional as F +from nemo_automodel.components.distributed.blockdiag_cp import ( + attach_cp1_packed_varlen_hooks, + disable_cp1_packed_varlen, + enable_cp1_packed_varlen, +) from nemo_automodel.components.distributed.blockdiag_cp.kernels import ( _cp_blockdiag_mask, _cp_blockdiag_varlen, precompute_blockdiag_varlen_meta, ) -from nemo_automodel.components.distributed.blockdiag_cp.packed import ( - _PACKED_STATE, - enable_cp1_packed_varlen, -) +from nemo_automodel.components.distributed.blockdiag_cp.runtime import _ORIGINAL_SDPA requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="varlen parity requires CUDA") @@ -170,13 +172,13 @@ def test_precomputed_meta_bit_identical(backend, name, seg_lens, pad, S, cp, Hq, ], ) def test_cp1_packed_hook_parity(backend, seg_lens, pad, S, Hq, Hkv, D): - """cp1 packed path: F.sdpa under the armed varlen hook == dense block-diagonal SDPA. + """cp1 packed path: F.sdpa under the scoped varlen hooks == dense block-diagonal SDPA. - Exercises the full hook chain (enable_cp1_packed_varlen global patch -> - _packed_varlen_sdpa -> _cp_blockdiag_varlen at row_offset=0), as used by - packed cp_size==1 runs. The whole packed sequence is on one rank - (q == k == full). The hook persists (no ctx); state is reset after for - test isolation. + Exercises the full hook chain (attach_cp1_packed_varlen_hooks pre/post hooks -> + _packed_varlen_sdpa -> _cp_blockdiag_varlen at row_offset=0), as used by packed + cp_size==1 runs. The whole packed sequence is on one rank (q == k == full). The + SDPA patch must be live only during the hooked attention forward and restored + afterwards. """ _backend_or_skip(backend) device = "cuda" @@ -188,11 +190,29 @@ def test_cp1_packed_hook_parity(backend, seg_lens, pad, S, Hq, Hkv, D): k = torch.randn(1, Hkv, S, D, device=device, dtype=dtype) v = torch.randn(1, Hkv, S, D, device=device, dtype=dtype) ref = _dense_ref(q, k, v, doc_ids, 0) # [1, Hq, S, D] + + class _SdpaSelfAttn(torch.nn.Module): + def forward(self, q, k, v, scale, enable_gqa): + """Stock-SDPA attention over ``q`` ``[B, Hq, S, D]`` / ``k``/``v`` ``[B, Hkv, S, D]``.""" + return F.scaled_dot_product_attention(q, k, v, scale=scale, enable_gqa=enable_gqa) + + class _Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _SdpaSelfAttn() + + def forward(self, q, k, v, scale, enable_gqa): + """Route ``q``/``k``/``v`` ``[B, H, S, D]`` through the hooked attention module.""" + return self.self_attn(q, k, v, scale, enable_gqa) + + model = _Block() + attach_cp1_packed_varlen_hooks(model) enable_cp1_packed_varlen(doc_ids, backend) try: - got = F.scaled_dot_product_attention(q, k, v, scale=D**-0.5, enable_gqa=(Hkv != Hq)) + got = model(q, k, v, D**-0.5, Hkv != Hq) finally: - _PACKED_STATE.set(None) # disarm (patch stays installed but passes through) + disable_cp1_packed_varlen() + assert F.scaled_dot_product_attention is _ORIGINAL_SDPA # patch scoped to the hooked forward real = (doc_ids[0] > 0).nonzero().flatten() d = (ref[0, :, real, :].float() - got[0, :, real, :].float()).abs().max().item() assert d < atol, f"[cp1-hook {backend}] max|diff|={d:.4f} >= {atol}" From f35613f9e29802d15354b8bd238cbc29b4291df4 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Thu, 9 Jul 2026 10:44:28 -0700 Subject: [PATCH 5/8] refactor(distributed): route varlen kernel probes through safe_import helpers Replace the bare try/except import probes for flash-attn and TransformerEngine with module-level safe_import_from / safe_import_te from nemo_automodel.shared.import_utils, per the repo's optional-dependency convention. The kernel symbols themselves are still imported at the call sites so per-call stubs (tests) and lazy TE extension loading keep working. Co-Authored-By: Claude Fable 5 Signed-off-by: Yuhe Zhang --- .../distributed/blockdiag_cp/kernels.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py index 77a9734540..b5bf6fa87a 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py @@ -20,8 +20,15 @@ import torch +from nemo_automodel.shared.import_utils import safe_import_from, safe_import_te + logger = logging.getLogger(__name__) +# Availability probes only -- the kernel symbols themselves are imported at the +# call sites so a per-call stub (tests) or lazy TE extension load keeps working. +HAS_FLASH_VARLEN, _ = safe_import_from("flash_attn", "flash_attn_varlen_func") +HAS_TE, _ = safe_import_te() + _CP_FLASH_DETERMINISTIC = False _CP_FLASH_WARNED = False _CP_FLASH_ENGAGED = False @@ -57,16 +64,12 @@ def _varlen_backend_unavailable_reason( if device.type != "cuda": return f"varlen CP attention requires CUDA, got device={device}" if backend == "flash": - try: - from flash_attn import flash_attn_varlen_func # noqa: F401 - except Exception as exc: - return f"flash_attn varlen kernel is unavailable ({type(exc).__name__})" + if not HAS_FLASH_VARLEN: + return "flash_attn varlen kernel is unavailable" return None if backend == "te": - try: - from transformer_engine.pytorch import DotProductAttention # noqa: F401 - except Exception as exc: - return f"TransformerEngine varlen kernel is unavailable ({type(exc).__name__})" + if not HAS_TE: + return "TransformerEngine varlen kernel is unavailable" return None return f"unsupported varlen backend={backend}" @@ -518,12 +521,9 @@ def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale ) _CP_FLASH_WARNED = True return None - try: - if backend == "flash": - from flash_attn import flash_attn_varlen_func # noqa: F401 - except Exception: + if backend == "flash" and not HAS_FLASH_VARLEN: if not _CP_FLASH_WARNED: - logger.warning("flash_attn import failed; reporting the unavailable varlen path to the caller") + logger.warning("flash_attn is unavailable; reporting the unavailable varlen path to the caller") _CP_FLASH_WARNED = True return None From abb90e29e362ebfdc8ec74dd708148df0a4649ea Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Thu, 9 Jul 2026 10:44:28 -0700 Subject: [PATCH 6/8] test(distributed): 2-rank torchrun parity for blockdiag CP KV exchanges Add the first real-process-group coverage for the differentiable K/V collectives (_AllGatherSeqDiff, _LeftHaloExchange, _NeededKVExchange) and the halo/a2a attention paths: a 2-rank NCCL torchrun runner (run_blockdiag_cp_2rank.py, launched from pytest via L2_CP_BlockDiag_Varlen_Test.sh) drives the full production path (configure_cp_varlen -> make_cp_blockdiag_batch_and_ctx on a real DeviceMesh -> cp_blockdiag_sdpa) forward + backward and compares local outputs, the input-embedding grad slice, and cross-rank-summed per-parameter grads of a tiny q/k/v/o attention module against a single-process dense block-causal reference. Matrix: kv_exchange allgather / halo / a2a (both needed-only paths selectable at cp=2 via the knob; the selector's decision is pinned so a silent downgrade to all-gather cannot fake coverage) x backend dense (fp32), flash, TE (bf16), over three doc layouts: a document straddling the rank boundary, an entire rank chunk of padding (collective no-hang), and a single document spanning both ranks. Co-Authored-By: Claude Fable 5 Signed-off-by: Yuhe Zhang --- .../L2_CP_BlockDiag_Varlen_Test.sh | 23 ++ .../run_blockdiag_cp_2rank.py | 262 ++++++++++++++++++ .../test_blockdiag_cp_parity.py | 34 +++ 3 files changed, 319 insertions(+) create mode 100755 tests/functional_tests/context_parallel/L2_CP_BlockDiag_Varlen_Test.sh create mode 100644 tests/functional_tests/context_parallel/run_blockdiag_cp_2rank.py create mode 100644 tests/functional_tests/context_parallel/test_blockdiag_cp_parity.py diff --git a/tests/functional_tests/context_parallel/L2_CP_BlockDiag_Varlen_Test.sh b/tests/functional_tests/context_parallel/L2_CP_BlockDiag_Varlen_Test.sh new file mode 100755 index 0000000000..8149e41b61 --- /dev/null +++ b/tests/functional_tests/context_parallel/L2_CP_BlockDiag_Varlen_Test.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail # Exit immediately if a command exits with a non-zero status + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0,1" + +# 2-rank block-diagonal varlen CP parity (allgather / halo / a2a KV exchange) +python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run \ + tests/functional_tests/context_parallel/run_blockdiag_cp_2rank.py diff --git a/tests/functional_tests/context_parallel/run_blockdiag_cp_2rank.py b/tests/functional_tests/context_parallel/run_blockdiag_cp_2rank.py new file mode 100644 index 0000000000..a6b117f6db --- /dev/null +++ b/tests/functional_tests/context_parallel/run_blockdiag_cp_2rank.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""2-rank (real NCCL) forward+backward parity for block-diagonal varlen CP. + +Runs the full production path -- ``configure_cp_varlen`` -> +``make_cp_blockdiag_batch_and_ctx`` (real DeviceMesh / process group) -> +``cp_blockdiag_sdpa`` -- for every KV-exchange collective (``allgather`` full +K/V all-gather with reduce-scatter backward, ``halo`` left-neighbor p2p, and +``a2a`` needed-only all-to-all-v, both selectable at cp=2 via the +``kv_exchange`` knob) and both varlen kernels (flash / TE) plus the dense +fallback. Each case is checked against a single-process dense block-causal +reference computed from identical full tensors on every rank: local outputs, +the input-embedding gradient slice, and the cross-rank-summed per-parameter +gradients of a tiny q/k/v/o attention module. Scenarios cover a document +straddling the rank boundary, an entire rank chunk of padding (the collective +no-hang case), and a single document spanning both ranks. + +Run:: + + torchrun --standalone --nproc-per-node=2 \ + tests/functional_tests/context_parallel/run_blockdiag_cp_2rank.py +""" + +from __future__ import annotations + +import sys +from typing import Callable + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.device_mesh import init_device_mesh + +from nemo_automodel.components.distributed.blockdiag_cp import ( + configure_cp_varlen, + cp_blockdiag_sdpa, + make_cp_blockdiag_batch_and_ctx, +) +from nemo_automodel.components.distributed.blockdiag_cp import kernels as bd_kernels +from nemo_automodel.components.distributed.blockdiag_cp import runtime as bd_runtime +from nemo_automodel.components.distributed.blockdiag_cp import state as bd_state + +B, S, E, HQ, HKV, D = 1, 256, 128, 8, 4, 64 + +# (name, per-document token counts, zero-id pad tail); lengths sum to S. +SCENARIOS = [ + ("straddle+pad", [100, 80, 44], 32), # doc 2 straddles the rank boundary at 128 + ("rank1-all-pad", [120], 136), # rank 1's whole chunk is padding + ("single-doc", [256], 0), # one document spans both ranks +] + +# (attn_backend, kv_exchange, dtype); dense exercises the masked-SDPA fallback in +# fp32 for a tight tolerance, flash/TE are the bf16 varlen kernels. +CASES = [("dense", "allgather", torch.float32)] + [ + (backend, exch, torch.bfloat16) for backend in ("flash", "te") for exch in ("allgather", "halo", "a2a") +] + + +class _TinyAttn(nn.Module): + """Minimal q/k/v/o-projection attention block for CP-vs-dense grad parity.""" + + def __init__(self, dtype: torch.dtype, device: torch.device): + super().__init__() + self.q_proj = nn.Linear(E, HQ * D, bias=False, dtype=dtype, device=device) + self.k_proj = nn.Linear(E, HKV * D, bias=False, dtype=dtype, device=device) + self.v_proj = nn.Linear(E, HKV * D, bias=False, dtype=dtype, device=device) + self.o_proj = nn.Linear(HQ * D, E, bias=False, dtype=dtype, device=device) + + def forward(self, x: torch.Tensor, sdpa_fn: Callable) -> torch.Tensor: + """Project, attend (GQA), and re-project one sequence chunk. + + Args: + x: Hidden states ``[B, L, E]`` (``B`` = batch, ``L`` = this chunk's + sequence length -- full ``S`` for the reference, local shard for + CP -- ``E`` = embedding dim). + sdpa_fn: SDPA-compatible callable applied to ``q`` ``[B, HQ, L, D]`` + and ``k``/``v`` ``[B, HKV, L, D]``. + + Returns: + Output hidden states ``[B, L, E]``. + """ + b, s, _ = x.shape + q = self.q_proj(x).view(b, s, HQ, D).transpose(1, 2) # [B, HQ, L, D] + k = self.k_proj(x).view(b, s, HKV, D).transpose(1, 2) # [B, HKV, L, D] + v = self.v_proj(x).view(b, s, HKV, D).transpose(1, 2) + o = sdpa_fn(q, k, v, enable_gqa=True) # [B, HQ, L, D] + return self.o_proj(o.transpose(1, 2).reshape(b, s, HQ * D)) + + +def _dense_blockdiag_sdpa(doc_ids: torch.Tensor) -> Callable: + """Single-process dense block-causal SDPA reference over the full sequence. + + Args: + doc_ids: Per-position document ids ``[1, S]`` (0 == padding) on the full + packed sequence. + + Returns: + An SDPA-compatible callable: ``q`` ``[B, Hq, S, D]``, ``k``/``v`` + ``[B, Hkv, S, D]`` -> ``[B, Hq, S, D]`` under the per-document causal mask. + """ + + def fn(q, k, v, enable_gqa=False): + if enable_gqa and k.shape[1] != q.shape[1]: + n_rep = q.shape[1] // k.shape[1] + k = k.repeat_interleave(n_rep, dim=1) + v = v.repeat_interleave(n_rep, dim=1) + allow = bd_kernels._cp_blockdiag_mask(doc_ids, 0, q.shape[2], k.shape[2], q.shape[0]) + return F.scaled_dot_product_attention(q, k, v, attn_mask=allow) + + return fn + + +def _make_doc_ids(seg_lens: list[int], pad: int, device) -> torch.Tensor: + """Build ``[1, S]`` document ids: 1-based per-document runs + a zero pad tail.""" + ids = [] + for d, n in enumerate(seg_lens, start=1): + ids += [d] * n + ids += [0] * pad + assert len(ids) == S, (len(ids), S) + return torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0) + + +def _rel_diff(got: torch.Tensor, ref: torch.Tensor) -> float: + """Max abs difference between same-shape tensors, scaled by max |ref| (>= 1e-3).""" + d = (got.float() - ref.float()).abs().max().item() + return d / max(ref.float().abs().max().item(), 1e-3) + + +def _run_case(cp_mesh, device, backend, exch, dtype, seg_lens, pad): + """Run one (backend, kv_exchange, scenario) parity case on this rank. + + Returns: + ``(ok, detail)``: whether every check passed locally, and a per-metric + relative-diff summary string. + """ + rank = dist.get_rank() + world = cp_mesh.size() + local_len = S // world + off = rank * local_len + tol_out = 1e-3 if dtype is torch.float32 else 2e-2 + tol_grad = 2e-3 if dtype is torch.float32 else 4e-2 + + doc_ids = _make_doc_ids(seg_lens, pad, device) # [1, S] + torch.manual_seed(1234) # identical params on every rank and for the reference + model = _TinyAttn(dtype, device) + torch.manual_seed(7) + x_data = torch.randn(B, S, E, device=device, dtype=dtype) + torch.manual_seed(11) + g = torch.randn(B, S, E, device=device, dtype=dtype) # upstream grad for out [B, S, E] + g[:, doc_ids[0] == 0, :] = 0 # padding rows carry no upstream gradient + + # ---- single-process dense reference (full sequence, same weights) ---- + x_ref = x_data.clone().requires_grad_(True) + out_ref = model(x_ref, _dense_blockdiag_sdpa(doc_ids)) + (out_ref * g).sum().backward() + ref_param_grads = {n: p.grad.detach().float().clone() for n, p in model.named_parameters()} + ref_x_grad = x_ref.grad.detach().float().clone() + model.zero_grad(set_to_none=True) + + # ---- CP run through the production batch/ctx + SDPA path ---- + configure_cp_varlen(attn_backend=backend, kv_exchange=exch) + x_cp = x_data.clone().requires_grad_(True) + batch = {"inputs_embeds": x_cp, "_packed_seq_ids": doc_ids.clone()} + ctx, sharded = make_cp_blockdiag_batch_and_ctx(cp_mesh, None, batch) + x_loc = sharded["inputs_embeds"] # [B, local_len, E], differentiable slice of x_cp + with ctx(): + step_state = bd_state._CP_BLOCKDIAG_STATE.get() + # Pin the KV-exchange path actually taken (memoized; reused by the sdpa call). + path, _, why = bd_runtime._select_kv_exchange_path( + step_state, + step_state["group"], + step_state["doc_ids"], + local_len, + device, + step_state["row_offset"], + query_dtype=dtype, + ) + expected = exch if backend != "dense" else "allgather" + if path != expected: + return False, f"kv path {path} != expected {expected} ({why})" + out_loc = model(x_loc, cp_blockdiag_sdpa) # [B, local_len, E] + (out_loc * g[:, off : off + local_len]).sum().backward() + + ok = True + msgs = [f"path={path}"] + real = doc_ids[0, off : off + local_len] > 0 + if real.any(): + rel = _rel_diff(out_loc.detach()[0][real], out_ref.detach()[0, off : off + local_len][real]) + ok &= rel < tol_out + msgs.append(f"out={rel:.1e}") + rel = _rel_diff(x_cp.grad[:, off : off + local_len], ref_x_grad[:, off : off + local_len]) + ok &= rel < tol_grad + msgs.append(f"dx={rel:.1e}") + for n, p in model.named_parameters(): + # An all-padding rank's queries are disconnected from its zero output (only + # K/V keep the 0-weighted grad attachment), so q_proj may have no grad there. + g32 = torch.zeros_like(p, dtype=torch.float32) if p.grad is None else p.grad.detach().float().clone() + dist.all_reduce(g32, op=dist.ReduceOp.SUM) # rank contributions sum to the full-batch grad + rel = _rel_diff(g32, ref_param_grads[n]) + ok &= rel < tol_grad + msgs.append(f"d{n.split('.')[0]}={rel:.1e}") + model.zero_grad(set_to_none=True) + return ok, " ".join(msgs) + + +def main() -> int: + """Run the case matrix; return 0 iff every case passed on every rank.""" + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + print("[skip] block-diagonal CP 2-rank parity requires 2 CUDA devices", file=sys.stderr) + return 0 + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + if world != 2: + raise SystemExit(f"expected torchrun --nproc-per-node=2, got world_size={world}") + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("cp",)) + cp_mesh = mesh["cp"] + + failures = [] + for backend, exch, dtype in CASES: + reason = None + if backend != "dense": + # Environment-derived, so every rank skips (or runs) the case together. + reason = bd_kernels._varlen_backend_unavailable_reason(backend, dtype, device) + if reason is not None: + if rank == 0: + print(f"[skip] {backend}/{exch}: {reason}") + continue + for name, seg_lens, pad in SCENARIOS: + ok, detail = _run_case(cp_mesh, device, backend, exch, dtype, seg_lens, pad) + flag = torch.tensor([1 if ok else 0], device=device) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + passed = bool(flag.item()) + if rank == 0: + print(f"[{'ok' if passed else 'FAIL'}] {backend:<5} {exch:<9} {name:<14} {detail}") + if not passed: + failures.append((backend, exch, name)) + dist.barrier() + if rank == 0: + print("[ALL CASES PASSED]" if not failures else f"[{len(failures)} CASES FAILED] {failures}") + dist.destroy_process_group() + return 0 if not failures else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/functional_tests/context_parallel/test_blockdiag_cp_parity.py b/tests/functional_tests/context_parallel/test_blockdiag_cp_parity.py new file mode 100644 index 0000000000..b8c10012de --- /dev/null +++ b/tests/functional_tests/context_parallel/test_blockdiag_cp_parity.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Functional test for block-diagonal varlen context parallelism. + +Launches a real 2-rank torchrun job (run_blockdiag_cp_2rank.py) that checks +forward outputs, input gradients, and per-parameter gradients for every +KV-exchange collective (allgather / halo / a2a) and attention backend +(dense / flash / TE) against a single-process dense block-causal reference. +""" + +from tests.utils.test_utils import run_test_script + +TEST_FOLDER = "context_parallel" +CP_BLOCKDIAG_VARLEN_TEST_FILENAME = "L2_CP_BlockDiag_Varlen_Test.sh" + + +class TestBlockDiagContextParallel: + """Test suite for block-diagonal (packed-sequence) context parallelism.""" + + def test_cp_blockdiag_varlen_2rank_parity(self): + """2-rank fwd+bwd parity for allgather/halo/a2a KV exchange vs a dense reference.""" + run_test_script(TEST_FOLDER, CP_BLOCKDIAG_VARLEN_TEST_FILENAME) From dff34c258e4ca75e324fbb6849dae6645fd9fac6 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Mon, 13 Jul 2026 10:11:29 -0700 Subject: [PATCH 7/8] test(distributed): stress FlashAttention boundary guard and carry CP diagnostics Add a single-GPU shape-stress sweep for the block-diagonal CP FlashAttention boundary guard. The sweep drives production-sized query/key segment lengths that bracket FlashAttention tile boundaries plus maximally packed layouts with non-tile-aligned tails, each through a non-reentrant activation checkpoint that replays the attention forward during backward, and asserts finite forward and gradient outputs. Explicit synchronizations pin any delayed illegal-address report to the exact case that launched the offending kernel. Carry compact host-only forensic context so an asynchronous kernel fault named at the post-exchange consensus can report the exact shape without adding a hot-path synchronization: the varlen validator retains a validation snapshot of its shape metadata, the needed-only halo/all-to-all exchanges record a diagnostic describing the chosen path, and the collective success check prints that diagnostic to stderr before re-raising an async CUDA error. Signed-off-by: Yuhe Zhang --- .../distributed/blockdiag_cp/exchange.py | 27 +++ .../distributed/blockdiag_cp/kernels.py | 19 ++ .../distributed/blockdiag_cp/runtime.py | 30 ++- .../test_blockdiag_cp_flash_boundary_gpu.py | 201 ++++++++++++++++++ 4 files changed, 274 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py diff --git a/nemo_automodel/components/distributed/blockdiag_cp/exchange.py b/nemo_automodel/components/distributed/blockdiag_cp/exchange.py index b41692e0c7..a51af66b70 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/exchange.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/exchange.py @@ -212,6 +212,19 @@ def _blockdiag_halo_attention(query, key, value, doc_ids, group, plan, gm, row_o local_meta = dict(gm) local_meta["s_first"] = 0 local_meta["real_end"] = recv_count + n_real + gm.setdefault( + "_needed_only_diagnostic", + { + "path": "halo", + "rank": rank, + "world": world, + "row_offset": row_offset, + "recv_count": recv_count, + "max_halo": max_halo, + "key_needed_len": int(key_needed.shape[2]), + "kernel_meta": gm.get("_validation_snapshot"), + }, + ) return _cp_blockdiag_varlen( query, key_needed, @@ -362,6 +375,20 @@ def _blockdiag_a2a_attention(query, key, value, doc_ids, group, plan, gm, row_of local_meta = dict(gm) local_meta["s_first"] = 0 local_meta["real_end"] = needed_len + gm.setdefault( + "_needed_only_diagnostic", + { + "path": "a2a", + "rank": rank, + "world": world, + "row_offset": row_offset, + "needed_len": needed_len, + "in_splits": list(in_splits), + "out_splits": list(out_splits), + "key_needed_len": int(key_needed.shape[2]), + "kernel_meta": gm.get("_validation_snapshot"), + }, + ) return _cp_blockdiag_varlen( query, key_needed, diff --git a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py index b5bf6fa87a..927d7d18eb 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py @@ -179,6 +179,25 @@ def _varlen_metadata_unavailable_reason( k_offsets = meta["cu_k"].detach().cpu().tolist() q_lens = [b - a for a, b in zip(q_offsets, q_offsets[1:])] k_lens = [b - a for a, b in zip(k_offsets, k_offsets[1:])] + # Retain a host-only forensic snapshot. If a CUDA kernel reports an + # asynchronous illegal address at the following collective consensus, + # the CUDA context is already poisoned and copying cu_seqlens there is + # no longer possible. This compact snapshot lets that error name the + # exact shape without adding another synchronization on the hot path. + meta["_validation_snapshot"] = { + "n_real": n_real, + "query_len": query_len, + "key_len": key_len, + "s_first": s_first, + "real_end": real_end, + "max_q": max_q, + "max_k": max_k, + "first_q": int(meta.get("first_q", 0)), + "first_k": int(meta.get("first_k", 0)), + "max_tail": int(meta.get("max_tail", 0)), + "cu_q": q_offsets, + "cu_k": k_offsets, + } if q_offsets[0] != 0 or k_offsets[0] != 0: reason = "cu_q and cu_k must start at zero" elif any(length <= 0 for length in q_lens + k_lens): diff --git a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py index 3983910d57..05c0563e8a 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py @@ -116,7 +116,7 @@ def _needed_only_preflight( return all_available, reason -def _needed_only_kernel_succeeded_on_all_ranks(out, group, device) -> bool: +def _needed_only_kernel_succeeded_on_all_ranks(out, group, device, *, diagnostic=None) -> bool: """Make a rank-local varlen result safe to act on collectively. In particular, an all-padding rank returns a zero tensor without invoking @@ -125,7 +125,26 @@ def _needed_only_kernel_succeeded_on_all_ranks(out, group, device) -> bool: """ succeeded = out is not None if torch.distributed.is_initialized() and torch.distributed.get_world_size(group) > 1: - status = torch.tensor(int(succeeded), dtype=torch.int32, device=device) + try: + status = torch.tensor(int(succeeded), dtype=torch.int32, device=device) + except Exception as exc: + # CUDA launches are asynchronous: a bad attention kernel commonly + # surfaces here rather than at its Python call. Use a plain stderr + # print because rank-filtered logging may suppress nonzero ranks. + import sys + + try: + global_rank = torch.distributed.get_rank() + except Exception: + global_rank = -1 + print( + "CP_NEEDED_ONLY_ASYNC_ERROR " + f"global_rank={global_rank} device={device} " + f"error={exc!r} diagnostic={diagnostic!r}", + file=sys.stderr, + flush=True, + ) + raise torch.distributed.all_reduce(status, op=torch.distributed.ReduceOp.MIN, group=group) succeeded = bool(status.item()) return succeeded @@ -290,7 +309,12 @@ def cp_blockdiag_sdpa( out = exchange._blockdiag_a2a_attention( query, key, value, doc_ids, group, plan, gm, offset, scale, attn_backend ) - if _needed_only_kernel_succeeded_on_all_ranks(out, group, query.device): + if _needed_only_kernel_succeeded_on_all_ranks( + out, + group, + query.device, + diagnostic=gm.get("_needed_only_diagnostic") if gm is not None else None, + ): assert out is not None return out # The needed-only KV collective has already executed. An all-padding diff --git a/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py b/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py new file mode 100644 index 0000000000..e64afb68ea --- /dev/null +++ b/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py @@ -0,0 +1,201 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""1-GPU stress test for the FlashAttention boundary-segment guard. + +The left-straddling boundary document (``sk > sq``) is the layout that provoked +asynchronous illegal-address reports in some FlashAttention builds, which +:func:`_flash_varlen_with_long_prefix_guard` peels onto a fixed-shape kernel. This +sweep exercises a large set of production-sized shapes whose query/key segment +lengths bracket FlashAttention tile boundaries, plus maximally packed layouts with +many non-tile-aligned tails. Every case runs the guard through a non-reentrant +activation checkpoint (the wrapper that first surfaced the failure, replaying the +attention forward during backward) and asserts finite forward and gradient +outputs. Explicit synchronizations pin any delayed illegal-address report to the +exact case that launched the offending kernel. +""" + +from __future__ import annotations + +import gc + +import pytest +import torch +from torch.utils.checkpoint import checkpoint + +from nemo_automodel.components.distributed.blockdiag_cp.kernels import ( + _flash_varlen_with_long_prefix_guard, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="flash boundary stress requires CUDA") + +H_Q = 32 +H_KV = 8 +HEAD_DIM = 128 +LOCAL_LEN = 8192 + + +def _boundary_cases() -> list[tuple[int, int]]: + # (local query tokens in the boundary document, tokens inherited from the + # left CP rank). Values bracket FlashAttention tile boundaries and include + # the very asymmetric layouts omitted by the original three regressions. + queries = [ + 1, + 15, + 16, + 17, + 31, + 32, + 33, + 63, + 64, + 65, + 95, + 127, + 128, + 129, + 255, + 256, + 257, + 511, + 512, + 513, + 1023, + 1024, + 1025, + 2047, + 2048, + 2049, + 4095, + 4096, + 4097, + 8191, + 8192, + ] + backs = [1, 63, 64, 127, 128, 255, 256, 511, 512, 1023, 2048, 4095, 4096, 8191, 8192] + cases = [(q, backs[index % len(backs)]) for index, q in enumerate(queries)] + cases.extend((q, back) for q in (1, 17, 65, 129, 513, 1025, 2049, 4097, 8191) for back in (8191, 8192)) + return cases + + +def _tail_lengths(total: int, segments: int) -> list[int]: + assert total >= segments > 0 + base, remainder = divmod(total, segments) + return [base + (index < remainder) for index in range(segments)] + + +def _all_cases() -> list[tuple[str, int, int, list[int]]]: + cases = [(f"boundary_q{q}_back{back}", q, q + back, []) for q, back in _boundary_cases()] + # Production permits at most 512 packed documents. Exercise the maximum + # segment count as well as mixed non-tile-aligned tails after an asymmetric + # first segment. + for first_q, back, tail_segments in ( + (1, 8192, 511), + (17, 8191, 511), + (129, 4097, 127), + (2049, 8192, 31), + (4097, 4095, 7), + ): + tails = _tail_lengths(LOCAL_LEN - first_q, tail_segments) + cases.append( + ( + f"tail_q{first_q}_back{back}_segments{tail_segments + 1}", + first_q, + first_q + back, + tails, + ) + ) + return cases + + +def _run_case(name: str, first_q: int, first_k: int, tails: list[int], device: torch.device) -> None: + q_lengths = [first_q, *tails] + k_lengths = [first_k, *tails] + q_total = sum(q_lengths) + k_total = sum(k_lengths) + assert 0 < q_total <= LOCAL_LEN + assert first_k >= first_q + + generator = torch.Generator(device=device).manual_seed(1729 + first_q * 17 + first_k * 31 + len(tails)) + query = torch.randn( + q_total, + H_Q, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + generator=generator, + requires_grad=True, + ) + key = torch.randn( + k_total, + H_KV, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + generator=generator, + requires_grad=True, + ) + value = torch.randn_like(key, requires_grad=True) + cu_q = torch.tensor([0, *torch.tensor(q_lengths).cumsum(0).tolist()], dtype=torch.int32, device=device) + cu_k = torch.tensor([0, *torch.tensor(k_lengths).cumsum(0).tolist()], dtype=torch.int32, device=device) + + def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return _flash_varlen_with_long_prefix_guard( + q, + k, + v, + cu_q=cu_q, + cu_k=cu_k, + max_q=max(q_lengths), + max_k=max(k_lengths), + local_query_len=LOCAL_LEN, + scale=HEAD_DIM**-0.5, + meta={ + "first_q": first_q, + "first_k": first_k, + "max_tail": max(tails) if tails else 0, + }, + ) + + # Match the production text-attention wrapper: the failure that motivated + # this stress test was observed while a non-reentrant checkpoint was + # replaying an attention forward during backward. + output = checkpoint( + attention, + query, + key, + value, + use_reentrant=False, + preserve_rng_state=False, + determinism_check="none", + ) + torch.cuda.synchronize(device) + assert output.shape == query.shape, name + assert torch.isfinite(output).all(), name + output.float().square().mean().backward() + torch.cuda.synchronize(device) + for tensor in (query, key, value): + assert tensor.grad is not None, name + assert torch.isfinite(tensor.grad).all(), name + + +@requires_cuda +def test_flash_boundary_shape_stress(): + """Sweep the boundary-guard shapes on a single GPU with checkpointed replay.""" + pytest.importorskip("flash_attn") + device = torch.device("cuda", torch.cuda.current_device()) + for case in _all_cases(): + _run_case(*case, device) + gc.collect() + torch.cuda.empty_cache() From 1e3b7c9f5aa584b7cbfbe95ad5ff5c63306317d4 Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Mon, 20 Jul 2026 08:15:05 -0700 Subject: [PATCH 8/8] fix(distributed): harden packed block-diagonal CP Signed-off-by: Yuhe Zhang --- .../distributed/blockdiag_cp/__init__.py | 2 +- .../distributed/blockdiag_cp/batch.py | 2 +- .../distributed/blockdiag_cp/exchange.py | 10 +- .../distributed/blockdiag_cp/kernels.py | 46 ++++++- .../distributed/blockdiag_cp/packed.py | 8 +- .../distributed/blockdiag_cp/runtime.py | 52 ++++++-- .../distributed/blockdiag_cp/state.py | 2 +- .../distributed/test_blockdiag_cp.py | 121 ++++++++++++++++-- .../test_blockdiag_cp_flash_boundary_gpu.py | 3 +- .../test_blockdiag_cp_varlen_gpu.py | 56 +++++++- 10 files changed, 260 insertions(+), 42 deletions(-) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/__init__.py b/nemo_automodel/components/distributed/blockdiag_cp/__init__.py index b535f63cd8..dc92166098 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/__init__.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/nemo_automodel/components/distributed/blockdiag_cp/batch.py b/nemo_automodel/components/distributed/blockdiag_cp/batch.py index 0415619406..5747501a05 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/batch.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/batch.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/nemo_automodel/components/distributed/blockdiag_cp/exchange.py b/nemo_automodel/components/distributed/blockdiag_cp/exchange.py index a51af66b70..1f6008401e 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/exchange.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/exchange.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -154,7 +154,7 @@ def backward(ctx, grad_recv): return grad_x, None, None, None, None -def _blockdiag_halo_attention(query, key, value, doc_ids, group, plan, gm, row_offset, scale, backend): +def _blockdiag_halo_attention(query, key, value, doc_ids, group, plan, gm, row_offset, scale, backend, dropout_p): """Needed-only block-diagonal CP attention via the left-halo exchange. Each rank attends its local queries against ``[boundary-doc halo from rank-1] + @@ -174,6 +174,7 @@ def _blockdiag_halo_attention(query, key, value, doc_ids, group, plan, gm, row_o row_offset: Global position of this rank's first local query row. scale: Softmax scale (``None`` -> kernel default). backend: Varlen kernel backend, ``"flash"`` or ``"te"``. + dropout_p: Attention dropout probability. Returns: Attention output ``[B, Hq, L, D]``, or ``None`` when the kernel is @@ -233,6 +234,7 @@ def _blockdiag_halo_attention(query, key, value, doc_ids, group, plan, gm, row_o row_offset, scale, backend=backend, + dropout_p=dropout_p, meta=local_meta, ) @@ -332,7 +334,7 @@ def backward(ctx, grad_out): return grad_x.permute(1, 2, 0, 3).contiguous(), None, None, None, None -def _blockdiag_a2a_attention(query, key, value, doc_ids, group, plan, gm, row_offset, scale, backend): +def _blockdiag_a2a_attention(query, key, value, doc_ids, group, plan, gm, row_offset, scale, backend, dropout_p): """Needed-only block-diagonal CP attention via the general all-to-all-v exchange. Handles documents spanning >2 ranks (single long doc / unpacked long context) @@ -351,6 +353,7 @@ def _blockdiag_a2a_attention(query, key, value, doc_ids, group, plan, gm, row_of row_offset: Global position of this rank's first local query row. scale: Softmax scale (``None`` -> kernel default). backend: Varlen kernel backend, ``"flash"`` or ``"te"``. + dropout_p: Attention dropout probability. Returns: Attention output ``[B, Hq, L, D]``, or ``None`` when the kernel is @@ -397,5 +400,6 @@ def _blockdiag_a2a_attention(query, key, value, doc_ids, group, plan, gm, row_of row_offset, scale, backend=backend, + dropout_p=dropout_p, meta=local_meta, ) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py index 927d7d18eb..b95901ddd3 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,6 +34,7 @@ _CP_FLASH_ENGAGED = False _CP_VARLEN_SHAPE_LOGGED = False _CP_FLASH_LONG_SEGMENT_WARNED = False +_CP_TE_DROPOUT_WARNED = False _TE_DPA_CACHE = {} @@ -249,6 +250,7 @@ def _flash_varlen_with_long_prefix_guard( max_k: int, local_query_len: int, scale, + dropout_p: float, meta: dict, ): """Run FlashAttention without an asymmetric-prefix varlen launch. @@ -274,6 +276,7 @@ def _flash_varlen_with_long_prefix_guard( max_k: Maximum per-document key segment length. local_query_len: This rank's local sequence length (for diagnostics). scale: Softmax scale (``None`` -> kernel default ``D**-0.5``). + dropout_p: Attention dropout probability. meta: Per-step metadata carrying ``first_q``/``first_k``/``max_tail``. Returns: @@ -301,6 +304,7 @@ def _flash_varlen_with_long_prefix_guard( cu_seqlens_k=cu_k, max_seqlen_q=max_q, max_seqlen_k=max_k, + dropout_p=dropout_p, softmax_scale=scale, causal=True, deterministic=_CP_FLASH_DETERMINISTIC, @@ -321,6 +325,7 @@ def _flash_varlen_with_long_prefix_guard( q_packed[:first_q].unsqueeze(0), k_packed[:first_k].unsqueeze(0), v_packed[:first_k].unsqueeze(0), + dropout_p=dropout_p, softmax_scale=scale, causal=True, deterministic=_CP_FLASH_DETERMINISTIC, @@ -342,6 +347,7 @@ def _flash_varlen_with_long_prefix_guard( cu_seqlens_k=tail_cu_k, max_seqlen_q=tail_max, max_seqlen_k=tail_max, + dropout_p=dropout_p, softmax_scale=scale, causal=True, deterministic=_CP_FLASH_DETERMINISTIC, @@ -499,7 +505,18 @@ def precompute_blockdiag_varlen_meta(doc_ids: torch.Tensor, row_offset: int, loc return seg if seg is not None else {"n_real": 0} -def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale=None, backend="flash", meta=None): +def _cp_blockdiag_varlen( + query, + key_full, + value_full, + doc_ids, + row_offset, + scale=None, + backend="flash", + *, + dropout_p=0.0, + meta=None, +): """Block-diagonal CP attention via varlen (flash or TE) -- no dense ``[B,1,L,S]`` mask. Equivalent to ``_cp_blockdiag_mask`` + SDPA for real (non-padding) query rows. @@ -524,14 +541,16 @@ def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale scale: Softmax scale (``None`` -> kernel default ``D**-0.5``). backend: ``"flash"`` (flash_attn_varlen_func) or ``"te"`` (TransformerEngine DotProductAttention thd). + dropout_p: Attention dropout probability. meta: Optional per-step segmentation from :func:`precompute_blockdiag_varlen_meta`; rebuilt inline when absent. Returns: Attention output ``[B, Hq, L, D]``, or ``None`` to signal "fall back to - the dense path" (kernel import failed, or a non-half dtype). + the dense path" when the requested varlen kernel path is unavailable or + unsupported for the supplied inputs. """ - global _CP_FLASH_WARNED + global _CP_FLASH_WARNED, _CP_TE_DROPOUT_WARNED if query.dtype not in (torch.float16, torch.bfloat16): if not _CP_FLASH_WARNED: logger.warning( @@ -540,6 +559,15 @@ def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale ) _CP_FLASH_WARNED = True return None + if backend == "te" and dropout_p > 0.0: + if not _CP_TE_DROPOUT_WARNED: + logger.warning( + "TransformerEngine THD varlen attention does not support dropout_p=%s; " + "reporting the unavailable varlen path so the caller preserves dropout via dense SDPA", + dropout_p, + ) + _CP_TE_DROPOUT_WARNED = True + return None if backend == "flash" and not HAS_FLASH_VARLEN: if not _CP_FLASH_WARNED: logger.warning("flash_attn is unavailable; reporting the unavailable varlen path to the caller") @@ -628,7 +656,14 @@ def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale v_packed = value_full[b, :, s_first:real_end, :].transpose(0, 1).contiguous() if backend == "te": - dpa = _te_varlen_dpa(Hq, Hkv, D, scale if scale is not None else D**-0.5, dev, query.dtype) + dpa = _te_varlen_dpa( + Hq, + Hkv, + D, + scale if scale is not None else D**-0.5, + dev, + query.dtype, + ) o = dpa( q_packed, k_packed, @@ -650,6 +685,7 @@ def _cp_blockdiag_varlen(query, key_full, value_full, doc_ids, row_offset, scale max_k=max_k, local_query_len=L, scale=scale, + dropout_p=dropout_p, meta=seg, ) # [n_real, Hq, D] out[b, :, :n_real, :] = o.transpose(0, 1) diff --git a/nemo_automodel/components/distributed/blockdiag_cp/packed.py b/nemo_automodel/components/distributed/blockdiag_cp/packed.py index 71b9de56ca..1b8561ff0b 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/packed.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/packed.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -77,7 +77,7 @@ def _packed_varlen_sdpa( value: Values ``[B, Hkv, S, D]``. attn_mask: Forwarded to stock SDPA on pass-through; ignored on the varlen path (masking is rebuilt from ``doc_ids``). - dropout_p: Dropout probability (pass-through / dense fallback only). + dropout_p: Dropout probability. is_causal: Forwarded on pass-through; the varlen path is always per-document causal. scale: Softmax scale (``None`` -> ``D**-0.5``). @@ -124,6 +124,7 @@ def _packed_varlen_sdpa( 0, scale, packed_state["backend"], + dropout_p=dropout_p, meta=packed_state.get("varlen_meta"), ) if out is None: # backend unavailable / unsupported -> safe fallback @@ -200,6 +201,9 @@ def enable_cp1_packed_varlen(doc_ids: torch.Tensor, backend: str) -> None: over the full packed sequence. backend: Varlen kernel backend, ``"flash"`` or ``"te"``. """ + if doc_ids.dim() == 1: + doc_ids = doc_ids.unsqueeze(0) + # Segmentation depends only on the packed document ids, not on the layer's # Q/K/V tensors. Compute it once per outer forward so every attention # layer (and activation-checkpoint recompute) reuses the same CUDA diff --git a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py index 05c0563e8a..3fac33fa4f 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -129,20 +129,18 @@ def _needed_only_kernel_succeeded_on_all_ranks(out, group, device, *, diagnostic status = torch.tensor(int(succeeded), dtype=torch.int32, device=device) except Exception as exc: # CUDA launches are asynchronous: a bad attention kernel commonly - # surfaces here rather than at its Python call. Use a plain stderr - # print because rank-filtered logging may suppress nonzero ranks. - import sys - + # surfaces here rather than at its Python call. Include the global + # rank and kernel metadata so the originating rank stays visible. try: global_rank = torch.distributed.get_rank() except Exception: global_rank = -1 - print( - "CP_NEEDED_ONLY_ASYNC_ERROR " - f"global_rank={global_rank} device={device} " - f"error={exc!r} diagnostic={diagnostic!r}", - file=sys.stderr, - flush=True, + logger.exception( + "CP_NEEDED_ONLY_ASYNC_ERROR global_rank=%d device=%s error=%r diagnostic=%r", + global_rank, + device, + exc, + diagnostic, ) raise torch.distributed.all_reduce(status, op=torch.distributed.ReduceOp.MIN, group=group) @@ -159,6 +157,7 @@ def _select_kv_exchange_path( offset, *, query_dtype: torch.dtype | None = None, + dropout_p: float = 0.0, ): """Decide this step's KV-exchange path and WHY. @@ -173,6 +172,7 @@ def _select_kv_exchange_path( device: Device for plan tensors and the preflight all-reduce. offset: Global position of this rank's first local query row. query_dtype: Query dtype for the kernel preflight (skip when ``None``). + dropout_p: Attention dropout probability. Returns: ``(path, plan, reason)`` where ``path`` is ``"halo"``/``"a2a"``/``"allgather"``, @@ -184,6 +184,8 @@ def _select_kv_exchange_path( return "allgather", None, f"mode={kv_exchange}" if attn_backend not in ("flash", "te"): return "allgather", None, f"needed-only requires flash/te kernel, got {attn_backend}" + if attn_backend == "te" and dropout_p > 0.0: + return "allgather", None, "TE THD varlen dropout requires the dense fallback" if state.get("varlen_meta") is None: return "allgather", None, "no varlen_meta (dense / non-varlen step)" if _cp_group_spans_nodes(group): @@ -241,7 +243,7 @@ def cp_blockdiag_sdpa( key: This rank's LOCAL key shard ``[B, Hkv, L, D]``. value: This rank's LOCAL value shard ``[B, Hkv, L, D]``. attn_mask: Ignored on the CP path (forwarded to stock SDPA otherwise). - dropout_p: Dropout probability (dense fallback path only). + dropout_p: Dropout probability. is_causal: Ignored on the CP path (forwarded to stock SDPA otherwise). scale: Softmax scale (``None`` -> ``D**-0.5``). enable_gqa: Grouped-query attention flag as passed by HF's sdpa path. @@ -286,6 +288,7 @@ def cp_blockdiag_sdpa( query.device, offset, query_dtype=query.dtype, + dropout_p=dropout_p, ) global _KV_EXCHANGE_PATH_LOGGED if not _KV_EXCHANGE_PATH_LOGGED: @@ -303,11 +306,31 @@ def cp_blockdiag_sdpa( if path in ("halo", "a2a"): if path == "halo": out = exchange._blockdiag_halo_attention( - query, key, value, doc_ids, group, plan, gm, offset, scale, attn_backend + query, + key, + value, + doc_ids, + group, + plan, + gm, + offset, + scale, + attn_backend, + dropout_p, ) else: out = exchange._blockdiag_a2a_attention( - query, key, value, doc_ids, group, plan, gm, offset, scale, attn_backend + query, + key, + value, + doc_ids, + group, + plan, + gm, + offset, + scale, + attn_backend, + dropout_p, ) if _needed_only_kernel_succeeded_on_all_ranks( out, @@ -343,6 +366,7 @@ def cp_blockdiag_sdpa( offset, scale, backend=attn_backend, + dropout_p=dropout_p, meta=step_state.get("varlen_meta"), ) if out is not None: diff --git a/nemo_automodel/components/distributed/blockdiag_cp/state.py b/nemo_automodel/components/distributed/blockdiag_cp/state.py index de40363b53..d3dba14853 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/state.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/state.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/unit_tests/distributed/test_blockdiag_cp.py b/tests/unit_tests/distributed/test_blockdiag_cp.py index a91374465c..2ee079a205 100644 --- a/tests/unit_tests/distributed/test_blockdiag_cp.py +++ b/tests/unit_tests/distributed/test_blockdiag_cp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -214,6 +214,35 @@ def test_blockdiag_sdpa_noop_without_state(): assert torch.allclose(out, ref, atol=1e-6) +def test_blockdiag_sdpa_forwards_dropout_to_varlen(monkeypatch): + """The CP wrapper preserves the caller's dropout probability on the varlen path.""" + forwarded_dropout = [] + + def fake_varlen(query, key, value, doc_ids, row_offset, scale, backend, *, dropout_p=0.0, meta=None): + """Return ``query`` ``[B, Hq, L, D]`` while recording the routed dropout probability.""" + forwarded_dropout.append(dropout_p) + return query + + monkeypatch.setattr(bd_exchange, "_AllGatherSeqDiff", _IdentityGather) + monkeypatch.setattr(bd_kernels, "_cp_blockdiag_varlen", fake_varlen) + state = { + "group": None, + "doc_ids": torch.tensor([[1, 1, 2, 2]], dtype=torch.long), + "row_offset": 0, + "attn_backend": "flash", + "kv_exchange": "allgather", + "varlen_meta": {"n_real": 4}, + } + token = bd_state._CP_BLOCKDIAG_STATE.set(state) + try: + qkv = torch.randn(1, 2, 4, 8, dtype=torch.bfloat16) + assert bd_runtime.cp_blockdiag_sdpa(qkv, qkv, qkv, dropout_p=0.25) is qkv + finally: + bd_state._CP_BLOCKDIAG_STATE.reset(token) + + assert forwarded_dropout == [0.25] + + def test_blockdiag_batch_synthesizes_and_shards_padding_mask(monkeypatch): class _Mesh: def size(self): @@ -480,16 +509,37 @@ def test_select_kv_exchange_path_downgrades_name_reasons(): ) assert path == "allgather" and "varlen_meta" in reason + state = {"attn_backend": "te", "kv_exchange": "halo", "varlen_meta": {"n_real": 1}} + path, _, reason = bd_runtime._select_kv_exchange_path( + state, + None, + torch.ones(1, 8, dtype=torch.long), + 4, + torch.device("cpu"), + 0, + dropout_p=0.1, + ) + assert path == "allgather" and "dropout" in reason + def test_flash_long_prefix_guard_peels_only_boundary_segment(monkeypatch): calls = [] def fake_fixed(q, k, v, **kwargs): - calls.append(("fixed", q.shape[1], k.shape[1], kwargs["causal"])) + calls.append(("fixed", q.shape[1], k.shape[1], kwargs["causal"], kwargs["dropout_p"])) return q + 1 def fake_varlen(q, k, v, **kwargs): - calls.append(("varlen", q.shape[0], k.shape[0], kwargs["max_seqlen_q"], kwargs["max_seqlen_k"])) + calls.append( + ( + "varlen", + q.shape[0], + k.shape[0], + kwargs["max_seqlen_q"], + kwargs["max_seqlen_k"], + kwargs["dropout_p"], + ) + ) return q + 2 flash_attn = types.ModuleType("flash_attn") @@ -512,10 +562,11 @@ def fake_varlen(q, k, v, **kwargs): max_k=7, local_query_len=4, scale=0.5, + dropout_p=0.25, meta={"first_q": 2, "first_k": 7, "max_tail": 4}, ) - assert calls == [("fixed", 2, 7, True), ("varlen", 4, 4, 4, 4)] + assert calls == [("fixed", 2, 7, True, 0.25), ("varlen", 4, 4, 4, 4, 0.25)] assert torch.equal(out[:2], torch.ones_like(out[:2])) assert torch.equal(out[2:], torch.full_like(out[2:], 2)) @@ -524,7 +575,7 @@ def test_flash_long_prefix_guard_keeps_normal_single_varlen_call(monkeypatch): calls = [] def fake_varlen(q, k, v, **kwargs): - calls.append((q.shape[0], k.shape[0])) + calls.append((q.shape[0], k.shape[0], kwargs["dropout_p"])) return q flash_attn = types.ModuleType("flash_attn") @@ -544,10 +595,11 @@ def fake_varlen(q, k, v, **kwargs): max_k=4, local_query_len=4, scale=None, + dropout_p=0.125, meta={}, ) - assert calls == [(6, 6)] + assert calls == [(6, 6, 0.125)] assert out is q @@ -556,14 +608,26 @@ def test_cp1_packed_precomputes_varlen_metadata_once_per_forward(monkeypatch): doc_ids = torch.tensor([[1, 1, 2, 2, 0]], dtype=torch.long) meta = {"n_real": 4, "sentinel": object()} precompute_calls = [] - forwarded_meta = [] + forwarded = [] def fake_precompute(got_doc_ids, row_offset, local_len, device): precompute_calls.append((got_doc_ids, row_offset, local_len, device)) return meta - def fake_varlen(query, key, value, doc_ids_arg, row_offset, scale, backend, *, meta=None): - forwarded_meta.append(meta) + def fake_varlen( + query, + key, + value, + doc_ids_arg, + row_offset, + scale, + backend, + *, + dropout_p=0.0, + meta=None, + ): + """Return ``query`` ``[B, Hq, S, D]`` while recording per-forward kernel options.""" + forwarded.append((meta, dropout_p)) return query monkeypatch.setattr(bd_packed.kernels, "precompute_blockdiag_varlen_meta", fake_precompute) @@ -572,8 +636,8 @@ def fake_varlen(query, key, value, doc_ids_arg, row_offset, scale, backend, *, m bd_packed.enable_cp1_packed_varlen(doc_ids, "flash") try: qkv = torch.randn(1, 2, 5, 4) - assert bd_packed._packed_varlen_sdpa(qkv, qkv, qkv) is qkv - assert bd_packed._packed_varlen_sdpa(qkv, qkv, qkv) is qkv + assert bd_packed._packed_varlen_sdpa(qkv, qkv, qkv, dropout_p=0.2) is qkv + assert bd_packed._packed_varlen_sdpa(qkv, qkv, qkv, dropout_p=0.2) is qkv finally: bd_packed.disable_cp1_packed_varlen() @@ -583,7 +647,26 @@ def fake_varlen(query, key, value, doc_ids_arg, row_offset, scale, backend, *, m assert row_offset == 0 assert local_len == doc_ids.shape[-1] assert device == doc_ids.device - assert forwarded_meta == [meta, meta] + assert forwarded == [(meta, 0.2), (meta, 0.2)] + + +def test_cp1_packed_1d_doc_ids_use_block_diagonal_attention(): + """A documented ``[S]`` doc-id vector must not fall through to ordinary causal SDPA.""" + torch.manual_seed(7) + doc_ids = torch.tensor([1, 1, 2, 2], dtype=torch.long) + qkv = torch.randn(1, 2, 4, 8) + allow = bd_kernels._cp_blockdiag_mask(doc_ids, 0, 4, 4, 1) + ref = bd_runtime._ORIGINAL_SDPA(qkv, qkv, qkv, attn_mask=allow) + ordinary = bd_runtime._ORIGINAL_SDPA(qkv, qkv, qkv, is_causal=True) + + bd_packed.enable_cp1_packed_varlen(doc_ids, "flash") + try: + got = bd_packed._packed_varlen_sdpa(qkv, qkv, qkv, is_causal=True) + finally: + bd_packed.disable_cp1_packed_varlen() + + torch.testing.assert_close(got, ref) + assert not torch.allclose(got, ordinary) def test_cp1_packed_passes_through_non_matching_shapes(): @@ -615,7 +698,19 @@ def test_cp1_packed_hooks_scope_patch_and_cover_ac_recompute(monkeypatch): varlen_calls = [] patched_at_forward_entry = [] - def fake_varlen(query, key, value, doc_ids_arg, row_offset, scale, backend, *, meta=None): + def fake_varlen( + query, + key, + value, + doc_ids_arg, + row_offset, + scale, + backend, + *, + dropout_p=0.0, + meta=None, + ): + """Return doubled ``query`` ``[B, Hq, S, D]`` for hook/recompute observability.""" varlen_calls.append(backend) return query * 2.0 diff --git a/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py b/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py index e64afb68ea..ac20e2302f 100644 --- a/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py +++ b/tests/unit_tests/distributed/test_blockdiag_cp_flash_boundary_gpu.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -161,6 +161,7 @@ def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor max_k=max(k_lengths), local_query_len=LOCAL_LEN, scale=HEAD_DIM**-0.5, + dropout_p=0.0, meta={ "first_q": first_q, "first_k": first_k, diff --git a/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py b/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py index 026d0c264d..417f21f6c0 100644 --- a/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py +++ b/tests/unit_tests/distributed/test_blockdiag_cp_varlen_gpu.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,7 @@ disable_cp1_packed_varlen, enable_cp1_packed_varlen, ) +from nemo_automodel.components.distributed.blockdiag_cp import packed as bd_packed from nemo_automodel.components.distributed.blockdiag_cp.kernels import ( _cp_blockdiag_mask, _cp_blockdiag_varlen, @@ -132,6 +133,59 @@ def test_varlen_blockdiag_parity(backend, name, seg_lens, pad, S, cp, Hq, Hkv, D assert d < atol, f"[{backend}] {name} r{r} off{off}: max|diff|={d:.4f} >= {atol}" +@requires_cuda +def test_flash_varlen_dropout_is_applied_and_backward_is_finite(): + """Nonzero Flash dropout changes varlen attention and preserves finite Q/K/V gradients.""" + _backend_or_skip("flash") + device = "cuda" + dtype = torch.bfloat16 + torch.manual_seed(23) + doc_ids = torch.ones(1, 64, dtype=torch.long, device=device) + query = torch.randn(1, 4, 64, 64, device=device, dtype=dtype, requires_grad=True) + key = torch.randn(1, 4, 64, 64, device=device, dtype=dtype, requires_grad=True) + value = torch.randn(1, 4, 64, 64, device=device, dtype=dtype, requires_grad=True) + + without_dropout = _cp_blockdiag_varlen(query, key, value, doc_ids, 0, backend="flash") + torch.manual_seed(29) + with_dropout = _cp_blockdiag_varlen(query, key, value, doc_ids, 0, backend="flash", dropout_p=0.25) + + assert without_dropout is not None and with_dropout is not None + assert not torch.equal(without_dropout, with_dropout) + with_dropout.float().square().mean().backward() + for tensor in (query, key, value): + assert tensor.grad is not None + assert torch.isfinite(tensor.grad).all() + + +@requires_cuda +def test_te_varlen_dropout_uses_dense_block_diagonal_fallback(): + """TE THD dropout falls back before launch and exactly matches masked dense SDPA.""" + _backend_or_skip("te") + device = "cuda" + dtype = torch.bfloat16 + torch.manual_seed(41) + doc_ids = _make_doc_ids([20, 12], 0, 32, device) + query = torch.randn(1, 4, 32, 64, device=device, dtype=dtype, requires_grad=True) + key = torch.randn(1, 4, 32, 64, device=device, dtype=dtype, requires_grad=True) + value = torch.randn(1, 4, 32, 64, device=device, dtype=dtype, requires_grad=True) + allow = _cp_blockdiag_mask(doc_ids, 0, 32, 32, 1) + + enable_cp1_packed_varlen(doc_ids, "te") + try: + torch.manual_seed(43) + got = bd_packed._packed_varlen_sdpa(query, key, value, dropout_p=0.25) + finally: + disable_cp1_packed_varlen() + torch.manual_seed(43) + ref = _ORIGINAL_SDPA(query, key, value, attn_mask=allow, dropout_p=0.25) + + torch.testing.assert_close(got, ref) + got.float().square().mean().backward() + for tensor in (query, key, value): + assert tensor.grad is not None + assert torch.isfinite(tensor.grad).all() + + @requires_cuda @pytest.mark.parametrize("backend", ["flash"]) @pytest.mark.parametrize(("name", "seg_lens", "pad", "S", "cp", "Hq", "Hkv", "D"), _CASES)