feat(deepseek_v4): faithful THD context-parallel path for DS4 CSA (bounded activation memory)

[ISEEKYAN]

What

Implements the THD-packed context-parallel path for DeepSeek-V4 Compressed Sparse Attention (CSA) in Megatron-Lite, as a faithful port of Megatron Core's CompressedSparseAttention._forward_thd_cp (NVIDIA/dev, NVIDIA#5087). This is the bounded-memory CP path for DS4 CSA — an alternative to the streaming-softmax approach in #6, built by importing Core's DSv4 THD-CP helpers zero-copy rather than re-deriving the sparse math.

Design

• mcore = NVIDIA/dev (container-provided). This PR does not vendor or modify megatron/core; the lite CSA imports the DSv4 THD-CP helpers directly: csa_cp_utils (prepare_cp_compressor_input, exchange_cp_boundary_hidden, compute_cp_indexer_topk, _thd_cp_position_ids), csa_cp_layout_kernels.build_attention_indices, and the THD sparse-attn kernels (dsa_sparse_attn(is_thd=True), FusedIndexerSparseAttnFromTopkFunc, unfused_compressed_sparse_attn).
• THD is the primary CP route. protocol.py no longer pops packed_seq_params and model.py no longer discards the THD signal — cu_seqlens is threaded end-to-end (DeepseekV4Model → DeepseekV4Layer → DeepseekV4CSAAttention → CompressedSparseAttention).
• The BSHD dense-softmax fallback is deleted (it was the O(S²/cp) activation-memory source); the only BSHD routes are the existing CP=1 fused sparse kernels.
• Kernel / indexer-loss knobs are implementation config (constructor args, mirroring GLM-5's DSA), not part of the HF DeepseekV4Config; fused DSA is the production default, unfused is a debug fallback.

GPU verification (cw H100/SM90, full evidence)

Validated with a standalone torchrun CP1-vs-CP4 parity harness on the lite CompressedSparseAttention (real DS4 geometry head_dim=512; variable-length multi-sample THD packs with sequence boundaries straddling CP shard cuts and non-ratio-aligned lengths):

• Forward parity: CP4 == CP1 to the fused kernel's own run-to-run non-determinism floor (~1e-3); bit-identical on the unfused/debug path.
• Backward parity: CP4 == CP1 within bf16 reduction-order rounding (input-grad ~1 bf16 ULP; fp32 spot-check confirms pure FP noise, not an algorithmic asymmetry); loss identical; no densify OOM at 16384 on the fused path.
• Memory (the point): CP=4 bounds per-rank activation memory — CP1/CP4 ratio 2.0× @4k → 3.4× @32k (approaching 4× as sequence length grows); the fused path is ~3× leaner in absolute terms than the unfused densifying path (CP1@32768: 12.1 GiB vs 35.9 GiB).
• q_padding_mask / indexer-loss path: runs green at CP1 and CP4.

Environment note

The fused DSA CP indexer requires cudnn-frontend ≥ 1.27 (cudnn.DSA.indexer_forward_wrapper must accept q_causal_offsets); 1.25.0 blocks the fused CP path. FlashMLA sparse-prefill requires the real DS4 head geometry (head_dim=512 / d_v=512).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request implements the DeepSeek-V4 Compressed Sparse Attention (CSA) THD context-parallel (CP) path in Megatron-Lite, enabling sequence-parallel training with packed sequences and removing the old BSHD dense fallback. It also adds corresponding unit tests to verify the layout, differentiability, and CP=1 correctness. The review feedback highlights several robustness improvements to prevent potential AttributeErrors, specifically by guarding against empty compressed_group_ids and using getattr to safely access attributes on packed_seq_params.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

To prevent potential AttributeError when compressed_group_ids is empty (e.g., on short sequences or specific CP configurations where some ranks have 0 compressed blocks), we should guard the indexer logic by checking if compressed_group_ids.shape[0] > 0.

Suggested change

	if indexer is not None:
	if indexer is not None and compressed_group_ids.shape[0] > 0:

[gemini-code-assist]

Similarly, if compressed_group_ids is empty, self.compressor._forward_thd will return None, None, causing an AttributeError when calling .squeeze(1). Guard this block to ensure it only runs when there are compressed groups.

Suggested change

	compressed_kv_local, _ = self.compressor._forward_thd(
	hidden_compact,
	cu_seqlens,
	max_seqlen_q=max_seqlen_q,
	compressed_group_ids=compressed_group_ids,
	)
	compressed_kv_rank_major = gather_from_sequence_parallel_region(
	compressed_kv_local.squeeze(1), group=cp_group
	)
	if compressed_group_ids.shape[0] > 0:
	compressed_kv_local, _ = self.compressor._forward_thd(
	hidden_compact,
	cu_seqlens,
	max_seqlen_q=max_seqlen_q,
	compressed_group_ids=compressed_group_ids,
	)
	compressed_kv_rank_major = gather_from_sequence_parallel_region(
	compressed_kv_local.squeeze(1), group=cp_group
	)

[gemini-code-assist]

Using getattr to access attributes on packed_seq_params is safer and more robust than direct attribute access, especially when integrating with external frameworks or custom runtimes where packed_seq_params might be a different class or dictionary.

    def _thd_cu_seqlens(self, packed_seq_params: Any) -> torch.Tensor:
        """Padded cu_seqlens (falling back to unpadded) for the THD packed layout."""
        padded = getattr(packed_seq_params, "cu_seqlens_q_padded", None)
        return padded if padded is not None else getattr(packed_seq_params, "cu_seqlens_q", None)

[gemini-code-assist]

Use getattr to safely retrieve max_seqlen_q from packed_seq_params to avoid potential AttributeError.

Suggested change

	max_seqlen_q = int(packed_seq_params.max_seqlen_q)
	max_seqlen_q = int(getattr(packed_seq_params, "max_seqlen_q", 0))

[gemini-code-assist]

Safely retrieve unpadded and padded cu_seqlens using getattr to prevent potential AttributeError if packed_seq_params is missing these attributes.

Suggested change

	cu_seqlens_q_unpadded = None
	if (
	packed_seq_params.cu_seqlens_q is not None
	and packed_seq_params.cu_seqlens_q_padded is not None
	and packed_seq_params.cu_seqlens_q.data_ptr()
	!= packed_seq_params.cu_seqlens_q_padded.data_ptr()
	):
	cu_seqlens_q_unpadded = packed_seq_params.cu_seqlens_q
	cu_seqlens_q = getattr(packed_seq_params, "cu_seqlens_q", None)
	cu_seqlens_q_padded = getattr(packed_seq_params, "cu_seqlens_q_padded", None)
	cu_seqlens_q_unpadded = None
	if (
	cu_seqlens_q is not None
	and cu_seqlens_q_padded is not None
	and cu_seqlens_q.data_ptr() != cu_seqlens_q_padded.data_ptr()
	):
	cu_seqlens_q_unpadded = cu_seqlens_q