[Aravind-11]

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

What does this PR do?

Adds fine-grained selective recompute and activation offloading to Gated DeltaNet (gdn) and Gated DeltaNet 2 (gdn2), matching what GatedDeltaProduct already supports with gdp_in_proj / gdp_qkv. Today the GDN family can only recompute its output norm (gdn_norm_out).

Roadmap: #5676, "GDN memory optimization" (fine-grained activation offloading and selective recompute for in-projection, conv1d, and gated delta rule logic).

New options

Option	Kind	What it frees
gdn_in_proj	--recompute-modules	Input projection output and its CP-to-HP all-to-all. Only hidden_states is kept.
gdn_qkv	--recompute-modules	Causal conv output and the prepared q/k/v.
gdn_qkv	--offload-modules	Moves the causal-conv input to the host.

Rejected with a clear error:

• Any of these on a non-GDN model.
• gdn_qkv (recompute or offload) with gdn_pre_gated_delta_rule_fusion, which fuses the conv and q/k/v prep with the gates.
• gdn_qkv offload together with gdn_in_proj recompute. The same rule exists for gdp_qkv / gdp_in_proj.

Implementation

• Shared base: the causal conv, previously duplicated in gdn.py and gdn2.py, moves into _GDNBase._causal_conv_qkv.
• Prep split: _prepare_input_for_gated_delta_rule is split into _prepare_qkv_for_gated_delta_rule and _prepare_gates_for_gated_delta_rule. The original method keeps its signature and behaviour, and pre_gated_delta_rule is unchanged for callers.
• Recompute units: _project_input wraps in_proj + all-to-all and _conv_and_prepare_qkv_maybe_recompute wraps conv + q/k/v prep, each in a CheckpointWithoutOutput (and the gdn_qkv offload group). This follows GatedDeltaProductMixer._gdp_chunk_forward.
• Fused-gate lifetime: fused pre-GDR returns the gate as a strided view of the input projection and does not materialize a contiguous copy. Discarding the projection before the norm reads that view corrupts the forward. The unfused path copies the gate with .contiguous(), which is why a CPU replay of the unfused module still matched. gdn_in_proj with fusion stays allowed; gdn_qkv with fusion stays rejected.
• Hook placement: gdn_qkv recompute is hooked on the raw gated-delta-rule output, and the gdn_qkv offload group is committed on that same tensor. gdn_in_proj is discarded on the norm output, after the gated norm has read its gate. forward_post_core_attn also registers that same restore on the output projection, before the gdn_norm_out hook and without a second discard. Hooks on one tensor run FIFO, and hooks on the output projection run before hooks on the norm output, so the projection (and a fused gate view into it) is restored before either consumer reads it. The new discards stay inside forward_pre_attn_and_core_attn, so these options keep two-stage attention working. gdn_norm_out still does not, because its checkpoint is discarded on the output projection.

Testing

Earlier CPU checks:

• Exact replay: loss, input grad and every parameter grad are bit-identical to no-recompute. Covers gdn/gdn2 × fp32/bf16 × batch 1/2 × {gdn_in_proj, gdn_qkv, both, both + gdn_norm_out}: 32/32 cases. Each requested checkpoint released its outputs in forward. Runs the deterministic torch-native kernels, with CUDA RNG / TE hooks stubbed.
• Hook order is load-bearing: registering the q/k/v hook before the in_proj hook fails immediately on reading the freed projection.
• No regression without the options: output and all grads are bit-identical to main for gdn/gdn2 × fp32/bf16.
• Activations saved for backward: bf16, hidden 2048, 16 key / 32 value heads of dim 128, seq 1024. These numbers use the torch-native delta-rule kernel, which saves many fp32 intermediates. With FLA kernels the baseline is smaller, so the relative savings should be larger.

	gdn	gdn2
baseline	558.9 MiB	743.1 MiB
gdn_in_proj	−72.4 MiB	−120.0 MiB
gdn_qkv	−32.1 MiB	−32.1 MiB
gdn_in_proj + gdn_qkv + gdn_norm_out	−112.6 MiB (−20.1%)	−160.2 MiB (−21.6%)

Unit tests in tests/unit_tests/ssm/test_gated_delta_net_recompute.py:

• GPU recompute exactness: the existing exact-replay test is parametrized over the new modules (GPU + FLA, TP/SP/CP grid) and asserts that every checkpoint discards its outputs once.
• Strided fused gate: test_in_proj_recompute_with_strided_gate_view builds the fused pre-GDR gate layout (a strided view of qkvzba, shared storage with the real gate) and checks gdn_in_proj alone and together with gdn_norm_out. GDN2 skips it. The real fused kernel is covered by test_in_proj_recompute_with_real_pre_gated_delta_rule_fusion, which skips unless causal_conv1d backward is installed.
• gdn_qkv offload: warmup offload bytes are non-zero, and a second forward/backward after PipelineOffloadManager.reset() stays bit-identical. Combinations are offload alone, with gdn_qkv, with gdn_norm_out, and with both. gdn_in_proj stays rejected next to offload.
• Config validation: CPU tests for the accepted and rejected combinations.

GPU run of that file, commit 732159d2154c242ee812d00213bd764013c25acc, on one NVIDIA A100-SXM4-40GB (torch 2.11.0+cu130, CUDA 13.0, Transformer Engine 2.20.0, flash-linear-attention 0.5.1, world size 1):

pytest -k "1-False-1 or test_gdn_" tests/unit_tests/ssm/test_gated_delta_net_recompute.py
30 passed, 4 skipped, 120 deselected in 396.18s

The 30 passes are the tp=1 / sequence-parallel off / cp=1 slice: 10 selective-recompute cases, 2 strided-gate cases, 8 offload cases, and the 10 CPU config tests. Outputs and grads matched with torch.equal. The 4 skips are the two GDN2 strided-gate cases, the GDN2 real-fusion case, and the GDN1 real-fusion case (causal_conv1d backward is not installed). The 120 deselected cases are the same matrix at tp=2 or cp=2.

The Colab process disabled torch.compile and kept a single Triton autotune config so the first-shape compile fit in the assignment. The commit leaves both enabled. The comparison is eager against eager on the same kernels.

Still pending:

• tp=2, sequence parallel, and cp=2
• the real fused pre-GDR kernel (needs causal_conv1d backward)
• a training run with gdn_qkv offload
• a peak-memory measurement with FLA kernels

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel …, I have added or updated its bit-exact determinism test (no kernel changes)
• I have added proper typing to my code
• I have added relevant documentation (transformer_config.py docstrings, fine-grained offloading guide)
• I have run the autoformatter.sh on my PR (BASE_REF=main CHECK_ONLY=true; black, isort, pylint, and ruff clean on the changed Python)

[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.