[dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention

[hxbai]

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

What does this PR do ?

Add Packed Sequence (THD) support for DSv4 Hybrid Attention.

⚠️ 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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hxbai]

/claude strict-review

[claude]

[IMPORTANT Correctness] Dense-loss backward for THD is missing max_seqlen_q and max_seqlen_k kwargs.

The forward path at lines 1185-1192 builds dense_attn_kwargs with all four cu_seqlens + max_seqlen fields and passes them to both _compute_dense_indexer_score and _compute_dense_attn_score. However, the backward's dense_bwd_kwargs here only passes cu_seqlens_q and cu_seqlens_k, omitting max_seqlen_q and max_seqlen_k.

cuDNN THD-mode wrappers typically require max_seqlen for internal tile sizing. If dense_indexer_backward_wrapper follows the same convention as the forward wrappers, this will either silently produce wrong gradients or raise at runtime.

Suggested fix:

if is_thd:
    dense_bwd_kwargs = dict(
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_compressed_idx,
        max_seqlen_q=int(max_seqlen_q),
        max_seqlen_k=int(max_seqlen_compressed_idx),
    )

[hxbai]

fixed

[claude]

[IMPORTANT Correctness] Sparse-loss backward: THD indexer gradient reshape relies on .view() to collapse the fake-BSHD B=1 dim returned by the cuDNN wrapper.

The cuDNN indexer_backward_wrapper was called with fake-BSHD inputs shaped (1, total_q, ...) (line 1238-1243), so ig["d_index_q"] is (1, total_q, idx_nh, idx_hd). The .view(total_q, idx_nh, idx_hd) here works because 1 * total_q == total_q, but it's semantically fragile — if the wrapper ever pads or returns a different leading dimension, .view() will silently reinterpret memory.

Consider using .squeeze(0) to explicitly remove the batch dimension, consistent with the pattern used elsewhere in this file (e.g., q_thd = q_indexer.squeeze(1) in csa.py):

precomputed_grad_q_indexer = ig["d_index_q"].squeeze(0)
precomputed_grad_k_indexer = ig["d_index_k"].squeeze(0)
precomputed_grad_weights = ig["d_weights"].squeeze(0)

[hxbai]

modified

[claude]

[IMPORTANT Correctness] When total_comp_capacity > cu_seqlens_compressed[-1] (capacity-padded for CUDA graph capture), invalid compressed rows are routed to total_kv + src_comp (line 326). However, out is sized total_kv_full = total_kv + compressed_kv_thd.shape[0], so the maximum valid destination index is total_kv_full - 1.

For the last invalid compressed row, dst_comp = total_kv + (total_comp_capacity - 1), which equals total_kv_full - 1 — that's exactly in-bounds. But this only works because out is sized using compressed_kv_thd.shape[0] (capacity) rather than cu_seqlens_compressed[-1] (true count). This tight coupling between the out allocation at line 288 and this fallback logic is fragile.

Consider adding a bounds assertion:

assert (total_kv + src_comp.max().item()) < total_kv_full, "padding dest OOB"

Or document the invariant that out is always capacity-sized (not true-count-sized) explicitly.

[hxbai]

Added a comment to clarify

[ko3n1g]

/ok to test e979a98

[hxbai]

/ok to test 7f3c633

[hxbai]

/ok to test 215de5e

[hxbai]

/ok to test 7e38374

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28212100264

[hxbai]

/ok to test e1bdccc

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28225017892