[dev] [DeepSeek-v4] Context Parallel support

docker/Dockerfile.ci.dev

@@ -2,7 +2,7 @@
 # syntax=docker/dockerfile:1.3-labs
 
 ARG FROM_IMAGE_NAME
 FROM ${FROM_IMAGE_NAME} as main
 ENV PIP_CONSTRAINT=""
 ENV DEBIAN_FRONTEND=noninteractive
 ARG UV_VERSION=0.7.2
@@ -31,7 +31,7 @@
     curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
 EOF
 
-RUN ln -sf /usr/local/cuda/targets/x86_64-linux/include/cuda \
+RUN ln -sf /usr/local/cuda/include/cccl/cuda \
     /usr/local/include/cuda
 RUN find /usr/local/cuda -name "utility" 2>/dev/null | head -5 && \
     ls /usr/local/cuda/targets/x86_64-linux/include/ | head -20
@@ -40,6 +40,8 @@
 COPY megatron/core/__init__.py /workspace/megatron/core/
 COPY megatron/core/package_info.py /workspace/megatron/core/
 ARG IMAGE_TYPE=dev
+ARG CUDNN_FRONTEND_COMMIT=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d
+ARG FLASH_MLA_COMMIT=b7643bd54521f563b839b98289b5cd048c062ba2
 RUN --mount=type=cache,target=/root/.cache/uv \
     bash -ex <<"EOF"
     export NVTE_CUDA_ARCHS="80;90;100"
@@ -61,7 +63,12 @@
         --no-install-package nvidia-cusolver-cu12 \
         --no-install-package nvidia-cusparse-cu12 \
         --no-install-package nvidia-cusparselt-cu12 \
-        --no-install-package nvidia-nccl-cu12
+        --no-install-package nvidia-nccl-cu12 \
+        --no-install-package nvidia-cudnn-frontend
+    uv pip install --no-build-isolation \
+        "flash-mla @ git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_COMMIT}" \
+        "nvidia-cudnn-frontend[cutedsl] @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_COMMIT}" \
+        "torch-c-dlpack-ext==0.1.5"
 EOF
 
 # Install DeepEP
@@ -83,10 +90,15 @@
 EOF
 
 COPY assets/ /opt/data/
+# Reinstall CuTe DSL cleanly after upgrading the version provided by the base image.
+RUN uv pip install --reinstall --no-deps \
+    nvidia-cutlass-dsl==4.5.0 \
+    nvidia-cutlass-dsl-libs-base==4.5.0 \
+    nvidia-cutlass-dsl-libs-cu13==4.5.0
 ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python
 
 ##### For NVIDIANS only #####
 FROM main as jet
 ARG JET_API_VERSION
 ENV PATH="$PATH:/opt/jet/bin"
 RUN --mount=type=secret,id=JET_INDEX_URLS bash -ex <<"EOF"

megatron/core/datasets/data_schedule.py

@@ -573,8 +573,8 @@ def get_batch_on_this_rank_for_sequence_packing(
         data_iterator (Iterator): The data iterator to get the batch from.
         mtp_on_this_rank (bool): Whether to use multi-token prediction.
         vp_stage (Optional[int]): The stage of the pipeline.
-        config: Model parallel config used for optional THD packed-sequence padding.
-            When None or config.pad_packed_seq_alignment is None, no padding is applied.
+        config: Model config used for CP partitioning and optional THD packed-sequence padding.
+            When None, CP partitioning defaults to zigzag and no padding is applied.
     Returns:
         tuple of (tokens, labels, loss_mask, attention_mask, position_ids,
         packed_seq_params, padding_mask)
@@ -630,6 +630,34 @@ def get_batch_on_this_rank_for_sequence_packing(
             group_size=local_cp_size_val
         )
 
+    cp_partition_mode = config.cp_partition_mode if config is not None else "zigzag"
+    contiguous_cp_local_target_len = None
+    pad_alignment = (
+        getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None
+    )
+    alignment = target_len = max_num_seqs = None
+    if pad_alignment is not None:
+        alignment, target_len, max_num_seqs = get_thd_padding_kwargs(
+            pad_alignment,
+            getattr(config, 'max_seqlen_per_dp_cp_rank', None),
+            getattr(config, 'thd_max_packed_sequences', None),
+            getattr(config, 'cuda_graph_impl', 'none') != 'none',
+        )
+    if (
+        is_tp_rank_0
+        and cp_group.size() > 1
+        and cp_partition_mode == "contiguous"
+        and pad_alignment is not None
+    ):
+        if target_len is not None:
+            contiguous_cp_local_target_len = target_len
+        else:
+            # Fix the local width before slicing so later padding cannot shift rank origins.
+            assert alignment is not None
+            total_rows = int(batch['cu_seqlens_padded'][-1].item())
+            local_rows = (total_rows + cp_group.size() - 1) // cp_group.size()
+            contiguous_cp_local_target_len = ((local_rows + alignment - 1) // alignment) * alignment
+
     # Build padding_mask before CP slicing while tensors still have the full
     # packed length represented by cu_seqlens_padded[-1].
     if is_tp_rank_0:
@@ -644,7 +672,18 @@ def get_batch_on_this_rank_for_sequence_packing(
         cp_slice_keys = ['padding_mask']
         if is_first_or_last_stage or mtp_on_this_rank:
             cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask'])
-        get_cp_slice_for_thd(batch, cp_group, keys=cp_slice_keys)
+        partition_total_tokens = (
+            contiguous_cp_local_target_len * cp_group.size()
+            if contiguous_cp_local_target_len is not None
+            else None
+        )
+        get_cp_slice_for_thd(
+            batch,
+            cp_group,
+            keys=cp_slice_keys,
+            cp_partition_mode=cp_partition_mode,
+            partition_total_tokens=partition_total_tokens,
+        )
 
     # Broadcast cu_seqlens_size because we need it to create placeholder for cu_seqlens and
     # cu_seqlens_padded for non TP 0 ranks.
@@ -658,12 +697,17 @@ def get_batch_on_this_rank_for_sequence_packing(
     # Broadcast total_tokens because padding_mask is prepared on every PP stage.
     # Tokens/labels/loss_mask/position_ids use the same length on stages that own them.
     if is_tp_rank_0:
-        # Under VPP, the last PP stage has labels but no tokens, so derive
-        # total_tokens from cu_seqlens_padded, which is present on every
-        # stage. cu_seqlens_padded keeps the pre-CP packed length; divide
-        # by cp_size to match the already CP-sliced sequence tensors.
-        cp_world = cp_group.size()
-        total_tokens = (batch['cu_seqlens_padded'][-1].to(torch.int32) // cp_world).reshape(1)
+        if contiguous_cp_local_target_len is not None:
+            total_tokens = torch.tensor(
+                [contiguous_cp_local_target_len], dtype=torch.int32, device=dev
+            )
+        else:
+            # Under VPP, the last PP stage has labels but no tokens, so derive
+            # total_tokens from cu_seqlens_padded, which is present on every
+            # stage. cu_seqlens_padded keeps the pre-CP packed length; divide
+            # by cp_size to match the already CP-sliced sequence tensors.
+            cp_world = cp_group.size()
+            total_tokens = (batch['cu_seqlens_padded'][-1].to(torch.int32) // cp_world).reshape(1)
     else:
         total_tokens = torch.empty(1, dtype=torch.int32, device=dev)
     broadcast_tensor(total_tokens, tp_src_rank, tp_group)
@@ -778,21 +822,13 @@ def get_batch_on_this_rank_for_sequence_packing(
         max_seqlen_kv=max_seqlen,
         local_cp_size=local_cp_size,
         cp_group=cp_group,
+        cp_partition_mode=cp_partition_mode,
         pad_between_seqs=False,
     )
 
     # Pad the already-packed THD tensors at the end when requested. CUDA Graph
     # additionally pads cu_seqlens tensors to thd_max_packed_sequences + 1 entries.
-    pad_alignment = (
-        getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None
-    )
     if pad_alignment is not None and packed_seq_params is not None:
-        alignment, target_len, max_num_seqs = get_thd_padding_kwargs(
-            pad_alignment,
-            getattr(config, 'max_seqlen_per_dp_cp_rank', None),
-            getattr(config, 'thd_max_packed_sequences', None),
-            getattr(config, 'cuda_graph_impl', 'none') != 'none',
-        )
         tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = (
             pad_sequence_for_thd(
                 tokens,

megatron/core/datasets/data_schedule_utils.py

@@ -2,7 +2,7 @@
 
 from functools import lru_cache
 from math import ceil, log2
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Dict, List, Literal, Optional, Sequence, Tuple
 
 import torch
 
@@ -12,16 +12,26 @@
 _DYNAMIC_CP_WORKLOAD_CAP_DELTA = 0.05
 
 
-def get_cp_slice_for_thd(batch, cp_group, keys: Optional[Sequence[str]] = None):
+def get_cp_slice_for_thd(
+    batch,
+    cp_group,
+    keys: Optional[Sequence[str]] = None,
+    cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag",
+    partition_total_tokens: Optional[int] = None,
+):
     """Partition sequence data for context parallelism in THD format.
 
-    Uses TE's THD partitioned indices to split the packed sequence across CP ranks.
+    ``zigzag`` uses TE's THD partitioned indices. ``contiguous`` splits the flattened rows into
+    ``cp_size`` equal slices and assigns one slice to each rank.
     Only keys present in the batch are sliced.
 
     Args:
         batch: Dict with packed sequence data.
         cp_group: Context parallel process group.
         keys: Sequence data keys to slice. Defaults to the original THD data tensors.
+        cp_partition_mode: How to assign packed rows to CP ranks.
+        partition_total_tokens: Optional total used to tail-pad tensors selected by ``keys``
+            before slicing. Existing cu_seqlens metadata is left unchanged.
     """
     cp_size = cp_group.size()
     if cp_size <= 1:
@@ -33,10 +43,44 @@ def get_cp_slice_for_thd(batch, cp_group, keys: Optional[Sequence[str]] = None):
     # Use cu_seqlens_padded[-1] for total_tokens instead of batch['tokens'].size(0):
     # under VPP, the last PP stage has labels/loss_mask but no tokens, so
     # batch['tokens'] is None on that stage. cu_seqlens_padded is always populated.
-    total_tokens = int(cu_seqlens[-1].item())
-    index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank)
+    total_tokens = (
+        int(cu_seqlens[-1].item())
+        if partition_total_tokens is None
+        else int(partition_total_tokens)
+    )
     if keys is None:
         keys = ('tokens', 'position_ids', 'labels', 'loss_mask')
+    if partition_total_tokens is not None:
+        for key in keys:
+            if key not in batch or batch[key] is None:
+                continue
+            pad_len = total_tokens - batch[key].numel()
+            if pad_len < 0:
+                raise RuntimeError(
+                    f"partition_total_tokens={total_tokens} is smaller than {key} length "
+                    f"{batch[key].numel()}."
+                )
+            if pad_len > 0:
+                pad_value = True if key == 'padding_mask' else 0
+                batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)])
+    if cp_partition_mode == "contiguous":
+        if total_tokens % cp_size != 0:
+            raise RuntimeError(
+                f"Contiguous CP slicing requires total_tokens={total_tokens} to be divisible by "
+                f"cp_size={cp_size}."
+            )
+        local_rows = total_tokens // cp_size
+        # Rank r takes [r * local_rows, (r + 1) * local_rows).
+        row_slice = slice(cp_rank * local_rows, (cp_rank + 1) * local_rows)
+        for key in keys:
+            if key in batch and batch[key] is not None:
+                batch[key] = batch[key][row_slice]
+        return
+
+    if cp_partition_mode != "zigzag":
+        raise ValueError(f"Unsupported CP partition mode: {cp_partition_mode}")
+
+    index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank)
     for key in keys:
         if key in batch and batch[key] is not None:
             batch[key] = batch[key].index_select(0, index)

megatron/core/extensions/transformer_engine.py

@@ -1748,9 +1748,10 @@ def __init__(
             self.kept_packed_seq_params.discard("cu_seqlens_q_padded")
             self.kept_packed_seq_params.discard("cu_seqlens_kv_padded")
 
-        # total_tokens and seq_idx are only for Mamba and should not be forwarded to TE attention.
+        # These fields are MCore-only and should not be forwarded to TE attention.
         self.kept_packed_seq_params.discard("total_tokens")
         self.kept_packed_seq_params.discard("seq_idx")
+        self.kept_packed_seq_params.discard("cp_partition_mode")
 
         if config.qk_clip or config.log_max_attention_logit:
             # qk-clip is only supported in TE 2.9.0 and later

megatron/core/fusions/fused_mla_yarn_rope_apply.py

@@ -86,6 +86,7 @@ def _mla_rope_fwd_inplace_kernel(
     batch_size,
     seq_num,
     cu_seqlens_q,
+    position_ids,
     stride_x_seq,
     stride_x_nheads,
     stride_cos_seq,
@@ -112,7 +113,9 @@ def _mla_rope_fwd_inplace_kernel(
     pid_m = tl.program_id(axis=0)
     pid_head = tl.program_id(axis=1)
 
-    if cu_seqlens_q is None:
+    if position_ids is not None:
+        token_idx = tl.load(position_ids + pid_m)
+    elif cu_seqlens_q is None:
         token_idx = pid_m // batch_size
     else:
         token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size)
@@ -181,6 +184,7 @@ def _mla_rope_bwd_inplace_kernel(
     batch_size,
     seq_num,
     cu_seqlens_q,
+    position_ids,
     stride_x_seq,
     stride_x_nheads,
     stride_cos_seq,
@@ -205,7 +209,9 @@ def _mla_rope_bwd_inplace_kernel(
     pid_m = tl.program_id(axis=0)
     pid_head = tl.program_id(axis=1)
 
-    if cu_seqlens_q is None:
+    if position_ids is not None:
+        token_idx = tl.load(position_ids + pid_m)
+    elif cu_seqlens_q is None:
         token_idx = pid_m // batch_size
     else:
         token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size)
@@ -270,6 +276,7 @@ def forward(
         rotary_interleaved=False,
         inverse=False,
         remove_interleaving=False,
+        position_ids=None,
     ):
         """
         Forward function for _FusedMLARoPEInplace.
@@ -288,13 +295,16 @@ def forward(
         seq_num = None
         if cu_seqlens_q is None:
             # sbhd
+            assert position_ids is None
             max_seqlen, batch_size, nheads, headdim = q.shape
             q = q.view(-1, nheads, headdim)
             total_seqlen = q.shape[0]
         else:
             # thd
             total_seqlen, nheads, headdim = q.shape
             seq_num = len(cu_seqlens_q) - 1
+            if position_ids is not None:
+                assert position_ids.shape == (total_seqlen,)
         assert q.stride(-1) == 1
         assert cos.stride(-1) == 1
         assert sin.stride(-1) == 1
@@ -312,6 +322,7 @@ def forward(
             batch_size,
             seq_num,
             cu_seqlens_q,
+            position_ids,
             q.stride(0),
             q.stride(1),
             cos.stride(0),
@@ -321,7 +332,8 @@ def forward(
             INVERSE=inverse,
             REMOVE_INTERLEAVING=remove_interleaving,
         )
-        ctx.save_for_backward(cos, sin)
+        ctx.save_for_backward(cos, sin, *(() if position_ids is None else (position_ids,)))
+        ctx.has_position_ids = position_ids is not None
         ctx.nope_dim = nope_dim
         ctx.emb_dim = emb_dim
         ctx.cu_seqlens_q = cu_seqlens_q
@@ -343,7 +355,11 @@ def backward(ctx, grad):
             grad: [seq_len, batch_size, head_num, nope_dim + emb_dim]
                 or [total_seq_len, head_num, nope_dim + emb_dim]
         """
-        cos, sin = ctx.saved_tensors
+        if ctx.has_position_ids:
+            cos, sin, position_ids = ctx.saved_tensors
+        else:
+            cos, sin = ctx.saved_tensors
+            position_ids = None
         max_seqlen = None
         batch_size = None
         seq_num = None
@@ -353,6 +369,8 @@ def backward(ctx, grad):
             total_seqlen = grad.shape[0]
         else:
             seq_num = len(ctx.cu_seqlens_q) - 1
+            if ctx.has_position_ids:
+                grad = grad.contiguous()
             total_seqlen, nheads, headdim = grad.shape
         assert grad.stride(-1) == 1
 
@@ -367,6 +385,7 @@ def backward(ctx, grad):
             batch_size,
             seq_num,
             ctx.cu_seqlens_q,
+            position_ids,
             grad.stride(0),
             grad.stride(1),
             cos.stride(0),
@@ -378,7 +397,7 @@ def backward(ctx, grad):
         )
         if ctx.cu_seqlens_q is None:
             grad = grad.view(max_seqlen, batch_size, nheads, headdim)
-        return grad, None, None, None, None, None, None, None, None, None, None
+        return grad, None, None, None, None, None, None, None, None, None, None, None
 
 
 def fused_mla_rope_inplace(
@@ -393,6 +412,7 @@ def fused_mla_rope_inplace(
     rotary_interleaved: bool = False,
     inverse: bool = False,
     remove_interleaving: bool = False,
+    position_ids: Optional[torch.Tensor] = None,
 ):
     """
     Fused RoPE applied inplace to the trailing emb_dim elements of a tensor,
@@ -414,6 +434,8 @@ def fused_mla_rope_inplace(
         rotary_interleaved: whether to apply RoPE interleaved, only supports False for now
         inverse: if True, apply the inverse rotation
         remove_interleaving: if True, output RoPE dims in non-interleaved layout
+        position_ids: optional THD row positions. When supplied, these positions
+            replace the built-in CP row-to-position mapping.
 
     Returns:
         t: inplace modified input tensor
@@ -430,6 +452,7 @@ def fused_mla_rope_inplace(
         rotary_interleaved,
         inverse,
         remove_interleaving,
+        position_ids,
     )