# DeepSeek-V4 Attention Training Systems

Status: first upstream-backed systems pass; inspected on 2026-07-19

This chapter studies the training implementation around DeepSeek-V4's CSA,
HCA, and window-only attention layers. It focuses on the system contracts that
are easy to miss when reading only the model equations:

- dense-to-sparse phase changes;
- auxiliary indexer loss and gradient ownership;
- packed variable-length sequences;
- Context Parallelism (CP), especially compression across rank boundaries;
- fused forward/backward kernels and their runtime glue;
- mixed precision, recomputation, checkpointing, and current limitations.

The short answer to the first implementation question is:

| Upstream Megatron `dev` path | Packed / varlen | CP | CP + packed | Qualification |
| --- | --- | --- | --- | --- |
| Legacy V3.2 `DSAttention` | **No end-to-end support** | no dedicated end-to-end path found | no | its public top-level indexer path still rejects `packed_seq_params`; THD helpers in `dsa.py` are also reused by V4 |
| V4 CSA, ratio 4 | yes | only through THD | yes | merged; contiguous CP partition required |
| V4 HCA, ratio 128 | yes | only through THD | yes | merged; no lightning indexer, but compressor and dense compressed-history path are covered |
| V4 window-only, ratio 0 | yes | only through THD | yes | merged; covers local-window boundary exchange without a compressor |

Thus “DSA kernels support THD metadata” and “the legacy DSA model path supports
packed training” are not equivalent statements. The former is true in current
upstream; the latter is not.

The upstream snapshot used for code claims is Megatron-LM `dev` commit
[`ceaa6c5d7ef`](https://github.com/NVIDIA/Megatron-LM/commit/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d),
dated 2026-07-18. The primary status source is the
[DeepSeek-V4 training tracker](https://github.com/NVIDIA/Megatron-LM/issues/4468).
“Merged” below means merged into `dev`, not necessarily released on `main` or
validated at the full published 1M-token scale.

## 1. Three Different Meanings Of “Supported”

For this feature, support should be recorded at three levels:

1. **Operator/API support:** tensors, metadata, forward, and backward paths
   exist for the requested layout.
2. **Composition support:** the operator works with packed data, CP, CUDA
   Graphs, mixed precision, MTP, optimizer state, and the actual heterogeneous
   layer schedule.
3. **Recipe validation:** a complete training recipe reaches the target
   context length with acceptable throughput, memory, convergence, and resume
   behavior.

As of the snapshot:

- V4 CSA/HCA functional support, THD packed support, fusion kernels, and the
  dedicated THD-CP path are merged into `dev`.
- The tracker still marks **long-context training** as “In Progress” and asks
  for validation of the 64K-to-1M curriculum.
- The tracker marks the dedicated FP8 indexer work as “In Progress” and the
  V4-specific FP4 QAT recipe as only partially supported.

It is therefore fair to say that the main implementation skeleton exists, but
too strong to say that every production-scale combination is mature.

## 2. The Training Lifecycle Changes The Operator Graph

The [initial CSA/HCA PR](https://github.com/NVIDIA/Megatron-LM/pull/4458)
exposes three configurations. They correspond to materially different work,
not just different loss coefficients.

| Phase | Indexer built? | Main CSA attention | Indexer teacher | Dominant long-context cost |
| --- | --- | --- | --- | --- |
| dense mode | no | all completed ratio-4 compressed entries + local window | none | approximately $O(L^2/4)$ main attention |
| sparse warm-up | yes | top-k compressed entries + local window | dense distribution over all visible compressed entries | sparse main path plus approximately $O(L^2/4)$ teacher-score work |
| sparse mode | yes | top-k compressed entries + local window | selected-set teacher | approximately $O(Lk)$ main and auxiliary work, plus index search |

Here $L$ is the original-token sequence length and $k$ is the number of
selected compressed entries. HCA has no indexer: it always reads every
completed ratio-128 compressed entry plus the local window.

The important correction is that “dense warm-up” can refer to two different
things:

- `--csa-dense-mode` removes the CSA indexer and sends every completed
  compressed entry into the main attention candidate list.
- omitting `--dsa-indexer-use-sparse-loss` while the indexer is enabled keeps
  the **main output sparse**, but computes a dense teacher distribution for the
  indexer loss.

The V4 report establishes the high-level dense-to-sparse curriculum. The exact
flags and loss implementation here are Megatron's public realization of that
curriculum; they are not proof that DeepSeek's internal trainer uses identical
code.

## 3. Gradient Ownership Around Discrete Top-k

The top-k indices are discrete, so the language-model loss does not train the
selection decision through a differentiable top-k. Instead, the public V4 path
retains the DSA-style auxiliary objective.

For query $t$, let $a_{t,j}$ be the main-attention teacher logit for compressed
candidate $j$, and let $i_{t,j}$ be the lightning-indexer logit. Abstractly,

$$
\mathcal L_I
=\sum_t D_{KL}
\left(
  \operatorname{Softmax}(a_{t,\mathcal C_t})
  \;\|\;
  \operatorname{Softmax}(i_{t,\mathcal C_t})
\right),
$$

where $\mathcal C_t$ is either the full visible compressed history in the
dense-loss phase or the selected set in sparse-loss mode.

Current Megatron explicitly detaches:

- the hidden-state and low-rank query inputs entering the indexer;
- the main query and compressed KV used to construct the teacher.

The resulting ownership is:

```text
language-model loss
  -> main Q/KV projections
  -> main CSA/HCA compressor
  -> sparse/dense compressed attention
  -> inverse RoPE and output projections

indexer KL loss
  -> indexer Q projection
  -> indexer head-weight projection
  -> separate indexer compressor
  -X-> backbone hidden state through the detached indexer input
  -X-> main attention through the detached teacher
```

The indexer compressor is a separate learned module; it is not the main CSA
compressor with a narrower view. This matters for parameter groups,
checkpoint conversion, optimizer state, DDP readiness, and precision policy.

The auxiliary scalar is attached to the main output through an autograd helper
so its backward runs during the normal model backward. A subtle distributed
contract follows: even when the configured loss coefficient is zero, indexer
parameters may still need zero gradients to be marked ready for overlapped
DDP. The later
[#5809 fix](https://github.com/NVIDIA/Megatron-LM/pull/5809) exists specifically
for this V4-CP case.

## 4. What “Varlen” Means In Megatron

The V4 implementation uses Megatron's THD packed layout:

```text
tokens from sample A | tokens from sample B | ... | capacity padding
<-------------------- total-token axis ---------------------------->

cu_seqlens_q / cu_seqlens_kv
  delimit samples and prevent cross-sample attention or compression
```

This supports different sample lengths within one packed microbatch. It also
supports sample lengths not divisible by the compression ratio:

$$
n_b^{cmp}=\left\lfloor\frac{L_b}{r}\right\rfloor.
$$

Only complete groups are compressed. The final $L_b \bmod r$ tokens do not
create a padded compressed entry; they remain reachable through the local
window. This matches autoregressive inference, where an incomplete compression
buffer is also not yet committed.

THD support is therefore more than flattening a tensor. Every one of these
operations must reset at a packed-sample boundary:

- ratio-4 overlapping compression;
- ratio-128 non-overlapping compression;
- RoPE positions for main and indexer operands;
- causal visibility of completed compressed groups;
- local-window indices;
- indexer top-k and its dense/sparse KL target;
- selected-index lowering into the concatenated physical KV buffer.

The merged
[#5011 PR](https://github.com/NVIDIA/Megatron-LM/pull/5011) adds this V4 THD
path. Current unit tests cover end-to-end forward and backward for ratios 0,
4, and 128, multi-segment ragged packs, non-divisible tails, fused/unfused
RoPE, and SBHD-versus-single-segment-THD parity.

### 4.1 Static capacity versus dynamic metadata

CUDA Graph capture cannot perform a device-to-host synchronization just to
discover a new packed length. The current code therefore requires
`max_seqlen_q`/`max_seqlen_kv`, allocates fixed-capacity buffers, and lets
device-resident `cu_seqlens` describe which rows are real on a replay.

This is a useful distinction:

- different segment boundaries inside an already captured capacity are
  supported and tested;
- arbitrary capacity or shape changes still require a compatible graph bucket
  or a new capture.

## 5. Why V4 CP Is A Dedicated Algorithm

Standard dense-attention CP can ring-exchange K/V chunks while each rank owns a
slice of Q. CSA/HCA add two complications:

1. a compressed entry can depend on tokens immediately to the left of a CP
   boundary;
2. a local query can select a compressed entry owned by any CP rank.

The current V4 path uses **contiguous** CP token partitions and performs the
following steps.

```text
rank-local hidden/Q/KV rows
  |
  +-- P2P receive left boundary hidden/KV rows
  |     supports local window and compression groups crossing the boundary
  |
  +-- compact the complete compressed groups owned by this rank
  |     a group is owned by the rank containing its final token
  |
  +-- all-gather compressed indexer K across CP ranks       [CSA only]
  +-- all-gather compressed main KV across CP ranks         [CSA/HCA]
  |
  +-- local Q scores the full compressed index memory       [CSA only]
  +-- lower logical top-k ids to physical rank-major rows
  |
  +-- attend to:
        left boundary raw KV + rank-local raw KV + gathered compressed KV
```

This is not a fallback that all-gathers the full hidden sequence. The raw
activation-sized work remains local apart from a fixed left-boundary exchange;
only compressed global state is replicated.

### 5.1 Boundary exchange

Each rank exchanges a fixed left boundary of

$$
D_{boundary}=\max(D_{window},D_{compress}).
$$

In the current implementation, $D_{window}=128$,
$D_{compress}=8$ for overlapping ratio-4 CSA, and $D_{compress}=128$ for HCA.
Backward sends the corresponding boundary gradients to the owning rank.

The implementation currently requires the number of local token rows to be at
least this boundary width. That becomes a practical lower bound on per-rank
sequence allocation and can limit extreme CP sizes on short examples.

### 5.2 Compressed ownership and fixed capacity

A compression group is assigned to the CP rank containing its last token.
That rank receives enough left context to compute the group exactly, packs its
owned groups into a fixed-capacity local buffer, and participates in an
all-gather.

The gathered buffer is rank-major for efficient fixed-size communication, but
packed sequences are logically sequence-major. A device-side map translates
between these layouts. Padding slots exist for capacity and alignment, but are
excluded from sample visibility and loss normalization.

### 5.3 Local Q, global compressed K

CSA index search keeps Q local and supplies the indexer with the full gathered
compressed-K memory. Per-query causal offsets restore the query's original
position inside its packed sequence. This offset support is part of the
required cuDNN Frontend wrapper contract for fused CP.

HCA skips the indexer and inserts all visible gathered ratio-128 entries into
the attention candidate list. Window-only layers skip both compressor and
indexer but still require correct raw boundary rows.

### 5.4 What CP saves and what it replicates

Ignoring constants, each rank holds:

$$
O(L/P)\quad\text{local Q/raw activations}
$$

plus

$$
O(L/r)\quad\text{gathered compressed global K/KV},
$$

where $P$ is CP size and $r$ is 4 or 128.

Consequences:

- CP reduces the dominant token-wise activation footprint and work.
- It does **not** shard the final compressed memory seen by every rank.
- CSA's ratio-4 gathered main and indexer memories can become a material
  replicated footprint at very long context.
- HCA's ratio-128 all-gather is much smaller and more regular.
- Communication grows with compressed state rather than full hidden state,
  which is the main opportunity created by the architecture.

This is the most important system-level caveat behind the phrase “CP support.”

### 5.5 TODO: avoid full ratio-4 replication at 1M context

The current all-gather is a reasonable first implementation because its shape
and communication schedule are static, but ratio-4 compression is not small at
1M context. With 512 BF16 main-KV dimensions and a 128-dimensional BF16
indexer key, one CSA layer has approximately

$$
\frac{10^6}{4}\times(512+128)\times2\ \text{bytes}
\approx305\ \text{MiB}
$$

of gathered compressed state on **every** CP rank. This total does not shrink
with CP size. Twenty-one CSA layers would represent about 6.3 GiB of such
payload if all layer states were simultaneously resident; actual peak depends
on pipeline placement, saved tensors, offload, and recomputation.

Track the following possible improvements rather than treating full
replication as the final design:

1. **Shard compressed main KV after distributed top-k.** Each rank scores its
   local indexer-K shard, the ranks merge local top-k candidates, and selected
   main KV is fetched from its owner.
2. **Replicate only indexer K.** Keep the narrower global search memory local
   everywhere, but use owner-aware all-to-all for selected 512-dimensional
   main KV.
3. **Hierarchical replication.** All-gather within an NVLink domain, shard
   across nodes, and perform a second-level top-k merge or selected-KV fetch.
4. **Query-block routing reuse or a selected-KV software cache.** Amortize
   remote fetches when nearby queries select overlapping compressed entries,
   while treating any cross-query index sharing that changes model semantics
   as an architecture experiment rather than a transparent runtime change.
5. **Communication-aware activation policy.** Recompute or offload gathered
   buffers only when the saved memory exceeds the cost of replaying collectives.

None is automatically better. A naive on-demand scheme can send
$O((L/P)k)$ KV references or payloads and be much worse than one regular
$O(L/r)$ all-gather. Evaluation must measure:

- unique selected entries per query block, not only per-query top-k;
- score/id merge bytes and selected-KV request/response bytes;
- forward and backward owner routing;
- load imbalance from hot compressed entries;
- CUDA Graph compatibility and fixed-capacity padding;
- overlap with indexer compute and sparse attention;
- the CP size and context length where sharding first beats replication.

The key research question is whether the top-k selection has enough reuse and
locality to turn irregular selected-KV exchange into less traffic than the
current predictable all-gather.

## 6. Exact Current CP/Varlen Contract

The dedicated V4 forward path enforces:

- `qkv_format == "thd"` whenever CP size is greater than one;
- `cp_partition_mode == "contiguous"`;
- a sequence-packing scheduler for static or dynamic CP;
- self-attention only in the THD-CP core path;
- dummy packed batch dimension equal to one;
- TP size equal to one for the entire V4 hybrid attention implementation.

The CP tests use CP sizes 2 and 4 and compare local outputs, hidden gradients,
and CP-reduced parameter gradients against a full-sequence CP=1 reference for:

- ratio 0 window-only layers;
- ratio 4 CSA with sparse and dense indexer-loss targets;
- ratio 128 HCA;
- fused and unfused RoPE;
- fused and unfused DSA paths where available;
- training and ratio-4 evaluation;
- CUDA Graph forward/backward replay;
- changed padded segment boundaries at replay time;
- peak allocated memory and graph-time scaling guards.

This is unusually useful unit coverage. It still does not replace an end-to-end
64K, 128K, or 1M training run with the full MoE/MTP/optimizer stack. The
official tracker preserves that distinction by leaving long-context training
“In Progress.”

## 7. Why Legacy DSA Is Different

Current upstream contains THD-aware helpers in `dsa.py` and `dsa_kernels.py`:

- packed top-k index globalization;
- THD forward/backward indexer-loss helpers;
- per-segment cumulative-length handling;
- causal offsets used by the V4 CP indexer.

However, the legacy `DSAttention` model path is still SBHD-oriented:

- `DSAIndexer.forward_with_scores` explicitly asserts that
  `packed_seq_params is None`;
- its main fallback attention constructs dense SBHD masks and batched matrix
  multiplications;
- there is no V3.2-specific CP boundary exchange, global selected-KV ownership,
  or CP parity test analogous to the V4 suite.

So the appropriate status is:

> Low-level DSA machinery has gained THD capabilities as part of the V4 work,
> but upstream does not currently expose legacy V3.2 DSA as an end-to-end
> packed/CP training path.

This also explains why searching only for `PackedSeqParams` in `dsa.py` gives
the wrong answer.

## 8. Kernel Layer Versus Runtime Layer

The implementation arrived in layers:

| Layer | Main upstream item | Responsibility |
| --- | --- | --- |
| functional model | [#4458](https://github.com/NVIDIA/Megatron-LM/pull/4458) | CSA/HCA modules, layer ratios, dense/warm-up/sparse modes, PyTorch fallback |
| fused operators | [#4894](https://github.com/NVIDIA/Megatron-LM/pull/4894) | integrate cuDNN Frontend 1.24.0 and FlashMLA-backed V4 sparse kernels |
| packed runtime | [#5011](https://github.com/NVIDIA/Megatron-LM/pull/5011) | THD metadata propagation, per-segment compression/indexing/attention |
| CP runtime | [#5087](https://github.com/NVIDIA/Megatron-LM/pull/5087) | contiguous partition, boundary exchange, compact layout, compressed all-gather, CP backward |
| hardening | tracker bug-fix list | dense loss/RoPE, precision, inverse-RoPE aliasing, MTP+CP, zero-loss DDP grads, mixed optimizer states |

The current fused stack itself is split:

```text
FlashMLA sparse forward
  -> selected MQA attention forward

cuDNN Frontend DSA namespace
  -> indexer scoring and radix top-k
  -> sparse/dense indexer-score recomputation
  -> sparse-attention backward
  -> indexer backward

Megatron runtime glue
  -> packed/CP layouts and logical-to-physical indices
  -> boundary P2P and compressed all-gathers
  -> loss attachment/reduction and parameter-gradient contracts
  -> heterogeneous layer selection and CUDA Graph capacity
```

This division explains why “the kernel is available” is insufficient. A
correct sparse-attention kernel does not know where a ratio-4 group crossing a
CP boundary lives, which packed sample owns an index, or how an auxiliary loss
must be reduced.

### 8.1 Hardware and dependency gates

In the inspected config, `apply_dsa_kernel_fusion` requires:

- SM90 or later;
- the `flash_mla` sparse-forward package from the NVIDIA development branch;
- a compatible cuDNN Frontend DSA package;
- causal-offset-capable cuDNN wrappers for ratio-4 CSA with CP.

The top-k list is padded to a kernel alignment of 128 on SM90 and 64 on SM100
for the current DSA dimensions. Padding is a runtime/kernel tax even when the
logical selected width is not aligned.

On SM90, fused ratio-4 CSA with a **dense** indexer loss is rejected because
that cuDNN Frontend dense-DSA path is considered unreliable. The supported
choices are sparse indexer loss or the unfused fallback. This is a concrete
example of support varying by phase and GPU generation.

### 8.2 Training kernel map

The fused training implementation is not “ordinary projections plus one
FlashAttention kernel.” FlashMLA owns the selected-attention forward, but the
complete forward/backward graph crosses several operator families:

| Stage | Current upstream solution | Why it is not just glue |
| --- | --- | --- |
| Q/KV and indexer projections | Megatron/Transformer Engine GEMMs plus the model's RMSNorm/RoPE path | multiple main/indexer branches, protected BF16 submodules under FP8 training, and per-layer positional modes |
| C4/C128 compression | V4 compressor modules plus THD/CP layout helpers | learned-score reduction must reset at sample boundaries and obtain left-boundary rows under CP |
| indexer score + top-k | cuDNN Frontend `DSA` score kernels plus TRT-LLM radix top-k | causal offsets, ragged lengths, fixed top-k capacity, alignment padding, and local-to-global id lowering are part of correctness |
| selected attention forward | FlashMLA `flash_mla_sparse_fwd` | consumes a flat global index list, returns output and LSE, and uses architecture-specific top-k alignment |
| indexer teacher/loss | cuDNN DSA sparse or dense score helpers plus FP32 loss math | teacher operands are detached; dense-loss warm-up and sparse-loss mode have different asymptotic work |
| main attention backward | cuDNN DSA `sparse_attention_backward_wrapper` | consumes saved output/LSE and produces Q, KV, and attention-sink gradients without materializing a dense attention matrix |
| indexer backward | cuDNN DSA sparse/dense indexer backward wrappers | discrete top-k has no LM-loss gradient, so this is a separate auxiliary-gradient path |
| inverse RoPE and output projection | V4 hybrid-attention module plus ordinary/fused projection operators | inverse rotation must preserve the unrotated output saved for sparse backward |
| distributed/layout runtime | Megatron THD and CP code | boundary P2P, compressed all-gather, rank-major packing, global-id translation, CUDA Graph capacity, and DDP readiness live outside FlashMLA |

The concrete fused training path is unusually asymmetric:

```text
forward
  indexer score + radix top-k
    -> combine compressed ids and window ids
    -> FlashMLA selected-attention forward
    -> recompute sparse/dense teacher scores for the indexer loss
    -> eagerly run the indexer-loss backward kernel

ordinary autograd backward
  saved output + LSE + dO
    -> cuDNN sparse-attention backward
    -> return the already-computed indexer gradients with the final loss scale
```

Eagerly computing the indexer backward during forward is possible because its
auxiliary loss has detached backbone and teacher inputs and a known unit loss
gradient. The wrapper defers the final loss-scale multiplication until normal
backward. This avoids retaining a second large autograd subgraph, but means
that a forward-only kernel trace for training already contains work named
“backward.” The inspected implementation is in
[`dsa_kernels.py`](https://github.com/NVIDIA/Megatron-LM/blob/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py).

### 8.3 What remains outside the fused attention core

The public fusion should not be read as one persistent kernel covering the
entire V4 attention layer. In particular:

- compressor projection/reduction, Q/KV up-projections, inverse RoPE, and
  grouped output projection remain distinct operator boundaries;
- THD segment maps, window/compressed-index concatenation, top-k padding, and
  CP logical-to-physical translation are runtime operations around the core;
- CP collectives are scheduled by Megatron rather than hidden inside
  FlashMLA or cuDNN DSA;
- dense curriculum phases intentionally use different score paths and cannot
  be represented by the sparse forward alone;
- backward has different dataflow and ownership from forward, so an inference
  decode kernel cannot simply be reused as the whole training operator.

This leaves meaningful kernel/runtime research even after sparse forward is
highly optimized: compressor fusion, temporary-layout elimination, score/top-k
fusion, recomputation policy, communication overlap, and a better division of
work between FlashMLA, cuDNN DSA, and Megatron.

## 9. Precision Is Per Component, Not One Attention Dtype

The current upstream policy is heterogeneous:

| Component | Current public training policy |
| --- | --- |
| main/indexer activation path | normally BF16, with selected GEMMs eligible for the enclosing FP8 recipe |
| CSA indexer Q projection | eligible for FP8 model initialization/execution |
| indexer head-weight projection | kept out of FP8; BF16 |
| main and indexer compressor GEMMs | kept out of FP8; BF16 |
| compressor positional parameter | FP32 parameter |
| attention sink | FP32 parameter |
| softmax, fallback scores, and KL reductions | FP32 math in the inspected reference/fallback code |
| published FP4 indexer/QAT path | not yet a completed Megatron V4 recipe |

The distinction between “an FP8 indexer GEMM can occur under an FP8 recipe”
and “the optimized FP8 indexer kernel project is complete” matters. The
official tracker still marks the latter in progress.

Keeping compressor and indexer subparts at different dtypes creates optimizer
state groups that cannot be assumed homogeneous. The tracker links
[#5308](https://github.com/NVIDIA/Megatron-LM/pull/5308) for preserving
compressor/indexer precision and
[#5835](https://github.com/NVIDIA/Megatron-LM/pull/5835) for distributed
optimizer save/load with mixed-dtype parameter groups.

## 10. Activation Memory And Recompute

Sparse execution removes the need to materialize the full main attention
matrix, but training still needs or recomputes several nonstandard objects:

- main and indexer compressed outputs;
- top-k indices and physical-row maps;
- FlashMLA log-sum-exp state;
- boundary hidden/KV rows;
- gathered compressed K/KV;
- indexer prediction and teacher scores for the auxiliary loss;
- inverse-RoPE input/output state;
- compressor softmax and normalization intermediates.

The fused backward deliberately recomputes indexer and attention scores through
cuDNN wrappers instead of retaining every dense score tensor. That is
especially important in the dense-loss warm-up phase, whose logical teacher is
quadratic over original versus compressed sequence positions.

Current constraints include:

- generic `checkpoint_core_attention` is rejected by V4 hybrid attention;
- QKV-linear offload is rejected;
- selective `mla_up_proj` recomputation is implemented and has a dynamic-CP
  parity test;
- core-attention activation offload has separate integration, but should not
  be confused with arbitrary checkpointing of the whole CP communication path.

An outer recompute boundary that encloses boundary exchange or compressed
all-gathers would replay communication during backward. Any future recompute
policy should therefore be evaluated in end-to-end step time, not only saved
bytes.

The inverse-RoPE operation also has an aliasing contract: fused sparse backward
retains the raw attention output, so applying inverse rotation in-place can
corrupt saved state. The
[#5526 fix](https://github.com/NVIDIA/Megatron-LM/pull/5526) changed this to
preserve private output storage.

## 11. Parallelism Consequences

### 11.1 Tensor Parallelism

The current V4 hybrid attention path hard-asserts TP=1. This is one of the
largest remaining integration constraints: attention cannot presently use TP
to shard its query heads, shared KV path, indexer, or grouped output
projection. Large training runs must obtain attention-side sequence scaling
primarily from CP while MoE layers may use their own EP/TP folding.

### 11.2 Context Parallelism

CP is architecture-specific here. Its useful scaling comes from local Q/raw
activations and sparse core work, while compressed global buffers are
replicated after all-gather. Performance should be reported with both:

$$
\text{tokens/s/GPU}
\quad\text{and}\quad
\text{compressed all-gather bytes per layer}.
$$

Otherwise a run can look memory-scalable while becoming communication-bound.

### 11.3 Pipeline and virtual pipeline parallelism

CSA, HCA, and window-only layers have different parameter, activation, and
communication costs. Equal layer counts do not imply balanced PP stages. A
useful planner must preserve the global layer pattern and balance estimated
cost, especially because CSA layers add an indexer and HCA layers have a
growing dense compressed candidate list.

The implementation enables heterogeneous distributed checkpoints and carries
per-layer compression ratios into ordinary and MTP layers. Resume validation
must check the layer schedule, indexer existence, precision-specific optimizer
state, and global layer numbering rather than only tensor shapes.

### 11.4 MTP interaction

MTP duplicates attention-layer contracts in predictor layers and interacts
with packed sequence shifting and CP partitioning. The tracker records a later
fix for MTP with contiguous CP and another for the layer/indexer accounting.
This is a runtime-composition issue, not a sparse-kernel issue.

## 12. The Bug History Is A Design Checklist

The tracker links fixes for:

- dense indexer loss and per-layer RoPE type;
- compressor/indexer high precision under FP8 training;
- fused inverse-RoPE output aliasing;
- MTP with contiguous CP partitions;
- zero-loss indexer gradients under CP and overlapped DDP;
- mixed-dtype distributed optimizer save/load.

These are not incidental cleanup items. Together they identify the main review
invariants for future implementations:

1. phase flags must change the intended main and auxiliary graphs;
2. RoPE mode is per layer and per compressed coordinate system;
3. packed sample boundaries must reset both compressors;
4. selected logical ids must map to the correct CP-owned physical rows;
5. saved sparse-attention output must survive inverse RoPE;
6. every conditional parameter must participate in the distributed-gradient
   readiness protocol;
7. optimizer/checkpoint code must preserve mixed-dtype groups;
8. MTP token shifting must use the same contiguous CP convention.

## 13. TODO: GPU Profiling And Roofline Validation

No GPU profiling was performed for this pass. The tables below are a future
measurement plan, not measured performance claims. Until H100/B200 access is
available, the established result is the upstream operator/runtime map above.

The upstream code is mature enough for targeted profiling. A useful first
matrix is:

| Axis | Suggested values |
| --- | --- |
| GPU | H100/SM90, B200/SM100 |
| layer | window-only, CSA, HCA |
| phase | dense mode, dense-loss warm-up, sparse-loss training |
| layout | SBHD fixed-length, THD ragged, THD ragged + CP |
| sequence | 8K, 32K, 64K, 128K; extend only after stability |
| CP | 1, 2, 4, 8 where local rows remain valid |
| kernel | fused, unfused reference |
| precision | BF16 baseline, FP8 recipe with protected submodules |

Record at least:

- end-to-end tokens/s/GPU and step-time breakdown;
- peak allocated and reserved memory;
- boundary P2P and compressed all-gather bytes/time;
- compressor, indexer/top-k, sparse forward, score recompute, and backward time;
- top-k padding waste on SM90 versus SM100;
- dense-loss warm-up overhead relative to sparse loss;
- activation recompute/offload savings and communication replay;
- numerical parity, indexer loss, top-k overlap, and gradient norms;
- checkpoint/resume parity across a phase transition.

The most interesting roofline question is whether ratio-4 CSA CP moves from
compute-bound sparse attention to communication- or indexer-bound execution as
$P$ grows, because its globally gathered compressed state shrinks with $r$ but
not with $P$.

## 14. Current Conclusions

### Established in public upstream

- Megatron `dev` has functional and fused V4 CSA/HCA training paths.
- THD packed/varlen support is merged for ratios 0, 4, and 128.
- THD + contiguous CP support is merged and has forward/backward parity tests
  for CP=2 and CP=4 across the three layer types.
- CP exchanges only a fixed raw left boundary, then all-gathers compressed
  indexer K and/or main KV.
- The public CSA implementation detaches indexer and teacher inputs and trains
  the indexer with a DSA-style auxiliary loss.
- TP is currently fixed at one, generic core-attention checkpointing is not
  supported, and the fused path has architecture/dependency gates.
- Legacy V3.2 `DSAttention` is not wired as an end-to-end THD/CP model path,
  despite sharing increasingly THD-aware low-level helpers.

### Still open or requiring measurement

- 64K-to-1M full-recipe correctness, convergence, and throughput;
- scaling of ratio-4 compressed all-gathers beyond CP=4;
- the 1M-context replication TODO in Section 5.5, including whether distributed
  top-k plus owner-routed selected KV can beat the regular all-gather;
- dynamic-CP scheduling quality for the V4 layer mixture;
- end-to-end PP/VPP/EP/CP/MTP balance;
- completed FP8 indexer kernels and V4 FP4-QAT recipe;
- production determinism and top-k tie behavior at large scale;
- whether a sharded or hierarchical compressed-memory exchange beats the
  current all-gather at very long context;
- whether TP>1 can be added without losing the shared-KV/index contract or
  creating excessive collective overhead.

## 15. Source And Code Map

- [DeepSeek-V4 report](https://arxiv.org/html/2606.19348v1): architecture and
  published training curriculum.
- [Megatron V4 tracker](https://github.com/NVIDIA/Megatron-LM/issues/4468):
  current merged/in-progress status and hardening PRs.
- [Megatron MoE roadmap](https://github.com/NVIDIA/Megatron-LM/issues/4815):
  broader `dev` roadmap; V4 remains labelled ongoing there.
- [CSA/HCA functional PR #4458](https://github.com/NVIDIA/Megatron-LM/pull/4458).
- [Fusion-kernel PR #4894](https://github.com/NVIDIA/Megatron-LM/pull/4894).
- [THD packed PR #5011](https://github.com/NVIDIA/Megatron-LM/pull/5011).
- [V4 CP PR #5087](https://github.com/NVIDIA/Megatron-LM/pull/5087).
- [`csa.py` at the inspected revision](https://github.com/NVIDIA/Megatron-LM/blob/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d/megatron/core/transformer/experimental_attention_variant/csa.py): compressor, indexer, THD, and CP attention paths.
- [`csa_cp_utils.py` at the inspected revision](https://github.com/NVIDIA/Megatron-LM/blob/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py): boundary exchange, compact ownership, and local-Q/global-K index layout.
- [`dsa_kernels.py` at the inspected revision](https://github.com/NVIDIA/Megatron-LM/blob/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py): FlashMLA/cuDNN wrapper split and fused backward.
- [`deepseek_v4_hybrid_attention.py` at the inspected revision](https://github.com/NVIDIA/Megatron-LM/blob/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py): runtime validation, CP boundary plumbing, inverse RoPE, and output projection.
- [V4 THD-CP tests at the inspected revision](https://github.com/NVIDIA/Megatron-LM/blob/ceaa6c5d7ef81e8408d2ae04ad63ec3c23905c4d/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py): CP parity, graph, memory, and timing guards.