[xuantengh]

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

What does this PR do?

• Bumps Megatron dependency https://github.com/fla-org/flash-linear-attention to the release 0.5.1
• Introduce GDN2 chunk kernel from FLA, and integrate GDN2 as a subclass of the _GDNBase

Convergence: see the wandb run.

Fixes #4955

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[xuantengh]

/claude review

[claude]

Light review — one inline nit and two items to double-check.

1. uv.lock not updated. pyproject.toml bumps flash-linear-attention from ~=0.4.0 to flash-linear-attention[cuda]>=0.5.1, but uv.lock is not part of this PR. Per the contribution workflow, uv.lock must be regenerated (uv lock inside the container) when pyproject.toml dependencies change, otherwise uv sync --locked fails in CI. Please confirm the lock file is committed (it may just be missing from the diff view).

2. chunk_gdn2 not reset to None on import failure — see inline comment on the except ImportError block.

Otherwise the change looks clean and well-tested: test_module_construction, test_gpu_forward_backward, the parametrized _prepare_input_for_gated_delta_rule coverage, and the dedicated test_parallel_gated_delta_net2_correctness give good coverage of the new GDN2 path across TP/SP/CP.

[xuantengh]

/claude review

[claude]

test_jit_compiled_helpers is broken — mismatch with the actual return contract of _prepare_input_for_gated_delta_rule.

The refactored method returns a single dict that includes the output gate under the key "gate", and forward() relies on exactly that (gate = kernel_inputs.pop("gate") at gated_delta_net.py:519). But the test unpacks the result as a 2-tuple:

# tests/unit_tests/ssm/test_gated_delta_net.py:352
kernel_inputs, gate_out = gdn._prepare_input_for_gated_delta_rule(
    qkv, gate, gate_feats, A_log_mock, dt_bias_mock, batch, seq_len
)

Unpacking a dict into two names raises ValueError: too many values to unpack, so this test errors out. Even after fixing the unpack, assert set(kernel_inputs.keys()) == expected_keys (line 356) would still fail because expected_keys omits "gate", which the real dict contains.

Suggested reconciliation to match the implementation (which forward() depends on):

kernel_inputs = gdn._prepare_input_for_gated_delta_rule(
    qkv, gate, gate_feats, A_log_mock, dt_bias_mock, batch, seq_len
)
gate_out = kernel_inputs["gate"]

assert set(kernel_inputs.keys()) == expected_keys | {"gate"}

Relatedly, the docstring is inaccurate: the signature is -> dict[str, torch.Tensor] and it returns one dict, but the Returns: block still describes tuple[dict[str, Tensor], Tensor] with the gate as a separate tuple element. Please update it to describe the single dict (with gate as a key).

[xuantengh]

/claude review

[xuantengh]

/ok to test b22c5cf

[claude]

The matching except ImportError block (lines 51-56) sets causal_conv1d, l2norm, and chunk_gated_delta_rule to None, but does not set chunk_gdn2 = None. If FLA is installed but too old to have fla.ops.gdn2.chunk (or FLA is absent entirely), this import fails and chunk_gdn2 is never bound. The GDN2 guard in __init__ (assert chunk_gdn2 is not None) then raises NameError: name 'chunk_gdn2' is not defined instead of the intended friendly AssertionError message about requiring flash-linear-attention >= 0.5.1.

Please add chunk_gdn2 = None to the except ImportError block:

except ImportError:
    causal_conv1d = None
    l2norm = None
    chunk_gated_delta_rule = None
    chunk_gdn2 = None

    HAVE_FLA = False

[claude]

Light review — the changes add a GDN2 variant to the gated delta net layer, with good test coverage (parametrized construction, forward/backward, correctness, and jit-helper tests for both variants).

One correctness issue flagged inline: the new chunk_gdn2 import isn't defaulted to None in the except ImportError block, so the GDN2 availability guard would raise NameError instead of the intended AssertionError when FLA is missing or too old.

Everything else looks consistent — config plumbing, FLOPs accounting, and sharded_state_dict all handle the new variant symmetrically.

[santhnm2]

Approved for transformer changes

[santhnm2]

I think it's ok to keep because the GPT experimental attention path is still active even if no new features are being added; we can remove if and when the GPT path is deleted entirely

[xuantengh]

/ok to test 63aebd9

[xuantengh]

/ok to test a5051b2

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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