[HaochenYuan]

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

What does this PR do?

Run a synthetic local forward/backward pass on every pipeline stage before the first scheduled training iteration. This initializes lazy kernels concurrently across PP stages without pipeline P2P dependencies.This prewarm does not capture or replay CUDA Graphs. Normal graph capture still happens later in the training schedule.

Performance

On a full-scale DSV4-Flash run using 64 GB200 GPUs (TP=1, PP=4, VPP=4, CP=16, EP=16, 64K sequence length), startup time through the completion of the first training iteration decreased from 866.0s to 370.8s, a 2.34x speedup.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[hxbai]

/claude strict-review

[buptzyb]

Prewarm THD metadata should alias Q/K the way real batches do (interaction with dev)

With for_pipeline_prewarm=True, the four assignments right after this block (L1278-1281, just outside the diff, so no suggestion block) give cu_seqlens_kv, cu_seqlens_q_padded and cu_seqlens_kv_padded each their own clone(). Real self-attention batches share one tensor between Q and KV, and one between the two *_padded fields:

• megatron/core/datasets/data_schedule.py:840-845 at dev, and likewise get_thd_batch_on_this_cp_rank (megatron/core/utils.py:2881-2886 at dev).
• packed_seq_params.py:695-701 at dev keeps that alias on purpose when it transforms the metadata.

Why it matters. Since #6206, the cuDNN DSA indexer uses its general packed-THD scorer only when it can prove, without a device sync, that the Q and K boundaries match. For CUDA tensors that means the same tensor or the same view (_packed_thd_kernel_applicable, experimental_attention_variant/dsa_cudnn_kernels.py:849-867 at dev; _get_packed_thd_metadata, :189-212, supplies cu_seqlens_q_padded and cu_seqlens_kv_padded to it). With the clones, the multi-entry synthetic pack ([0, max_T, max_T, ...]) resolves to UNFUSED_BOUNDS ("packed THD without compatible cu_seqlens metadata", dsa_scoring_plan.py:155-163 at dev) instead of PACKED_THD_GENERAL:

• Prewarm warms a different kernel. It runs the unfused per-head fallback scorer, while real batches, which alias, use the cuDNN general packed-THD indexer. That kernel's first call, and any one-time setup it does, moves into the first training step. Avoiding that is the point of prewarm. We haven't measured how much it costs.
• A misleading warning. report_unfused_scoring_once (dsa_scoring_plan.py:184-206 at dev) logs "DSA indexer scoring falls back to the unfused per-head loop ..." on every rank by design, but no training step takes that path. It also uses up the once-per-reason slot, so a later real fallback with the same reason would not be reported.

In our runs with this PR, every rank logs that warning exactly once, during prewarm, and the training steps take the fused path. I also checked the gate on CPU with this PR merged onto dev: the current prewarm inputs resolve to UNFUSED_BOUNDS, and the inputs from the change below resolve to PACKED_THD_GENERAL. A single-sequence synthetic pack (thd_max_packed_sequences unset or 1) usually takes the single-sequence segment scorer, which does not use this metadata.

Numerics: no change. Prewarm throws its results away: grads are zeroed, and buffers, FP8 state, RNG, the DSA tracker and the loss scales are restored (parallel_prewarm.py:257 and :287-307). The change only affects which tensors the synthetic sample shares; CUDA Graph capture inputs stay as they are.

Suggestion: replace L1278-1281 so that only prewarm aliases, and capture keeps four independent buffers:

                static_inputs["cu_seqlens_q"] = cu_seqlens
                if for_pipeline_prewarm:
                    # Match real self-attention batches, which share one tensor between Q and KV
                    # (and one padded tensor between both *_padded). Fused DSA scoring relies on
                    # that aliasing to recognize self-attention without a device sync.
                    padded_cu_seqlens = cu_seqlens.clone()
                    static_inputs["cu_seqlens_kv"] = cu_seqlens
                    static_inputs["cu_seqlens_q_padded"] = padded_cu_seqlens
                    static_inputs["cu_seqlens_kv_padded"] = padded_cu_seqlens
                else:
                    # CUDA Graph capture keeps independent static buffers.
                    static_inputs["cu_seqlens_kv"] = cu_seqlens.clone()
                    static_inputs["cu_seqlens_q_padded"] = cu_seqlens.clone()
                    static_inputs["cu_seqlens_kv_padded"] = cu_seqlens.clone()

After the rebase, dev's _uses_graph_dynamic_dsa_route() block stays right after this.

A CPU unit test for the change: it fails on the current head and passes with the change, both on the head and merged onto dev. It is much narrower than the state-restore test @yaox12 asked for, but it needs no GPU or process group.

tests/unit_tests/pipeline_parallel/test_parallel_prewarm.py

# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

from types import SimpleNamespace

import torch

from megatron.core.transformer.enums import AttnMaskType
from megatron.core.transformer.transformer_layer import TransformerLayer


def _make_thd_transformer_layer(*, cuda_graph_impl):
    layer = object.__new__(TransformerLayer)
    torch.nn.Module.__init__(layer)
    layer.config = SimpleNamespace(
        bf16=True,
        context_parallel_size=1,
        create_attention_mask_in_dataloader=False,
        cuda_graph_impl=cuda_graph_impl,
        cuda_graph_modules=[],
        fp16=False,
        hidden_size=4,
        max_seqlen_per_dp_cp_rank=4,
        sequence_packing_scheduler=object(),
        sequence_parallel=False,
        tensor_model_parallel_size=1,
        thd_max_packed_sequences=2,
    )
    layer.self_attention = SimpleNamespace(attn_mask_type=AttnMaskType.causal)
    layer.is_moe_layer = False
    return layer


def test_pipeline_prewarm_aliases_thd_qk_metadata_only_for_prewarm(monkeypatch):
    monkeypatch.setattr(torch.cuda, 'current_device', lambda: torch.device('cpu'))

    prewarm_inputs = _make_thd_transformer_layer(cuda_graph_impl='none').get_layer_static_inputs(
        4, 1, for_pipeline_prewarm=True
    )
    assert prewarm_inputs['cu_seqlens_q'] is prewarm_inputs['cu_seqlens_kv']
    assert prewarm_inputs['cu_seqlens_q_padded'] is prewarm_inputs['cu_seqlens_kv_padded']
    assert prewarm_inputs['cu_seqlens_q'] is not prewarm_inputs['cu_seqlens_q_padded']

    graph_inputs = _make_thd_transformer_layer(cuda_graph_impl='local').get_layer_static_inputs(
        4, 1, for_pipeline_prewarm=False
    )
    qk_metadata = [
        graph_inputs['cu_seqlens_q'],
        graph_inputs['cu_seqlens_kv'],
        graph_inputs['cu_seqlens_q_padded'],
        graph_inputs['cu_seqlens_kv_padded'],
    ]
    assert len({tensor.data_ptr() for tensor in qk_metadata}) == 4

[HaochenYuan]

Applied the suggested aliasing for prewarm: Q and KV share one tensor, and their padded metadata shares a second tensor. CUDA Graph capture retains four independently writable buffers.

Added CPU regression tests covering prewarm with none, local, and transformer_engine graph configurations. They also verify that the aliases survive PackedSeqParams construction and that graph-capture buffers remain independent.