[jingqiny-99]

Summary

This PR makes DSA indexer scoring correct for general packed THD inputs and removes repeated layout construction from the packed context-parallel hot path.

The central semantic correction is that packed sequence cardinality is independent of context parallelism: CP=1 can contain multiple sequences, and CP>1 can contain either one or many. Scorer selection now follows the actual THD metadata and kernel preconditions instead of using CP size as a proxy for sequence count.

The final branch preserves existing configuration and backend-hook compatibility. It does not claim full CUDA Graph support for packed DSA+CP. Such configurations remain constructible and receive an experimental-support warning; capture/replay support is follow-up work.

What changed

1. Plan-driven general THD scoring

• Add a total scoring-plan resolver with explicit outcomes for plain causal scoring, packed single-segment scoring, general packed THD scoring, unfused bounds, and decline.
• Run multi-sequence CP=1 scoring as one cuDNN THD call using the original packed Q/K/W tensors and cu-seqlens.
• Run CP>1 general packed scoring from the same THD metadata, using per-sequence zigzag query segments and gathered-key reorder metadata.
• Keep the single-sequence specialization only when its structural preconditions hold.
• Pass explicit key-position state through the production dispatcher so unsupported layouts reach the plan's DECLINE outcome before a cuDNN launch.
• Emit one diagnostic when a valid layout intentionally uses the existing unfused bounds path.

This removes both invalid implications: CP=1 no longer means “one sequence”, and CP>1 no longer means “multiple sequences”.

2. Host-built packed-CP layout metadata

• Record compact host Q/K cu-seqlens during the existing THD route-prebuild stage.
• Build all-rank query positions and gathered-key reorder tables from those host integers.
• Construct the reorder as a direct inverse permutation instead of a full argsort.
• Copy each finished table to the target device with a direct tensor transfer.
• Memoize layouts on the per-microbatch PackedSeqParams, so later DSA layers reuse the first build.
• Use the computed Q/K output-cover flags in both positions_and_reorder memo sites, so equal geometry has an equal cache key.

Independent CPU goldens cover CP1/2/4, both/all ranks, multiple sequences, zero-length entries, all-empty packs, padding, distinct Q/K boundaries, and every Q/K cover-flag combination. A production-forward regression test also enters both memo sites and proves that the second site reuses the first result.

3. Compatibility and invariant hardening

• Keep the existing backend-hook kwargs use_local_indexer_varlen and single_packed_thd_sequence on their CP-only contract.
• Offer CP-independent layout facts as separate optional kwargs.
• Move callable-signature filtering into megatron.core.utils.filter_kwargs_for_callable; exact signatures, **kwargs, positional-only parameters, opaque callables, and the identity-safe signature cache have direct tests.
• Preserve the pre-existing opaque full-hook varlen_is_plain_causal contract.
• Preserve valid self-attention Q/K metadata aliases through THD padding and inherited metadata.
• Reject inconsistent scorer devices, dtypes, views, TP slices, and gradient arity before producing incorrect results.
• Pass configured cp_comm_type into attention layers for Dynamic CP even when configured CP is 1; the effective microbatch CP group may still be larger than 1.
• Keep that effective Dynamic CP group inside selective core-attention checkpoint replay, then restore the caller's group after both the original call and recomputation.
• Restore the non-optional starts/ends autograd contract after the caller's concrete-bounds guard.

4. CUDA Graph scope

Packed DSA+CP graph configurations are no longer rejected by an empirical hardware/version matrix. The final check is structural and warning-only:

• static packed DSA warns when configured CP > 1 and CUDA graphs are enabled;
• Dynamic CP packed DSA warns whenever CUDA graphs are enabled, including configured CP=1, because the effective per-microbatch group can be larger;
• the warning is independent of scorer backend and capture module scope;
• eager execution (cuda_graph_impl="none") is unchanged.

The evidence is diagnostic, not an allow/deny table. On the validation GB200 image, cuDNN static CP2 eager execution passes, while local attention capture reaches graph warmup and fails when the cloned packed metadata is empty. Backend none was previously measured only for static CP2/local/attention and also failed during graph warmup. TileLang was unavailable in that image and remains unmeasured. No unmeasured combination is presented as known-good, and none is hard-rejected by this PR.

Full capture/replay support for changing packed layouts is intentionally deferred.

Validation

All GPU work below ran on Lyris GB200 in the standard gb200 partition (no backfill), using dsv4-gb200-torch2604-nvrx060.

• Final consolidated targeted suite: 164 passed — config compatibility/warnings, scoring plans, hook ABI, generic callable filtering, independent host-layout goldens and cache lifetime, cross-site memo reuse, Dynamic-CP communication and selective-recompute group lifetime, explicit-position dispatch, plus real cuDNN CP2 parity on both zigzag ranks and three seeds against an independent PyTorch forward/backward reference.
• Profiler-free model probes:
  • static CP2 eager: pass across one- and multi-sequence boundary changes;
  • Dynamic CP configured CP2 / effective CP1 eager: pass;
  • Dynamic CP configured CP1 / effective CP2 eager: pass after the communication-propagation fix;
  • static CP2 local-attention graph: configuration constructs and warns, eager preflight runs, then capture fails at graph warmup on empty cloned metadata (known deferred boundary).
• Selected H2D A/B: direct transfer reduced helper latency by 3.0%–17.9% across the measured CP2/4/8/16 layouts versus pin-then-transfer, with bit-identical outputs. This is a helper-level measurement, not an end-to-end iteration percentage.

Only the conclusion-relevant H2D nsys report was retained; correctness probes were run without profiling.

Commit map

1. Define explicit indexer scoring plans.
2. Preserve optional backend-hook compatibility.
3. Preserve THD Q/K metadata aliases.
4. Prebuild packed CP layout metadata.
5. Support general packed THD indexer scoring.
6. Add the initial packed DSA+CP graph check.
7. Preserve inherited THD metadata aliases.
8. Harden packed cuDNN scorer invariants.
9. Refine graph and hook compatibility from review feedback.
10. Final review follow-up: remove empirical hard rejection, fix memo/H2D/dispatcher contracts, generalize the hook helper, cover Dynamic CP configured-CP1 and selective recomputation, and expand independent/GB200 tests.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jingqiny-99]

/ok to test bd85b78

[jingqiny-99]

/ok to test 4888213

[jingqiny-99]

/ok to test 9fd365a

[jingqiny-99]

/ok to test 7986148

[buptzyb]

Should we merge this helper with maybe_move_tensor_to_cpu? https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/moe/moe_utils.py#L1238

I think maybe_move_tensor_to_cpu is in a bad place. Moving it here would be better.

BTW looks like maybe_move_tensor_to_cpu is using non_blocking always. We may need a new argument synchronization: bool = True, which determines whether we add a stream sync after the non-blocking copy.

Ignore me if you think this is a bad idea.

[jingqiny-99]

I think you are right, the maybe_move_tensor_to_cpu is in a strange place, I think it's good to move it here, and merge them in the future, because in my opinion, it may become a large helper (move tensor to arbitrary device and maybe do some synchronization if need), and it isn't transparent to its callers.

[jingqiny-99]

/ok to test 88a285f

[jingqiny-99]

/ok to test cff3e90

[buptzyb]

Do we need values.pin_memory().to(device) to avoid pageable H2D?

[jingqiny-99]

I think it is right, but in my test case, removing pin_memory will save copy time. So I will add an argument here to decide whether run values.pin_memory() here, and its default value will be False.

[buptzyb]

Is it because your values are already on pin? Non-pin h2d is dangerous since pagable copy is blocking and non-graphable IIRC.

[jingqiny-99]

soga..Thanks for teaching me, and sorry for my limited knowledge. I will pin the memory by default.

[buptzyb]

Happen to read an article about a hang issue caused by pagable memcpy: https://mp.weixin.qq.com/s/uOJALDKb588mFynaeJQJoA

[jingqiny-99]

/ok to test 9659ab5

[jingqiny-99]

/ok to test b7cd2d5

[buptzyb]

I have to admit I cannot understand every single detail of this PR. But the code organization and functionality generally LGTM. One concern is the main branch is also making their refactoring to the dsa code (like #7211), so the main-dev divergency would be larger and larger. Beg @hxbai for a final review, thanks!