[Models] Add MomentumDeltaNet implementation from arXiv:2605.05838 (ICML2026)

[bhaochen]

Summary

Add MomentumDeltaNet (MDN) implementation as proposed in MDN: Parallelizing Stepwise Momentum for Delta Linear Attention (https://arxiv.org/abs/2605.05838, ICML2026).

Stage 2 (this update): Replace degenerate mu->0 delta-rule reuse with complete stepwise momentum formulation (log_alpha/log_mu/beta/eta/p, dual state [S,M], M_t=mu*M+eta*k@w, S_t=alpha*S-beta*M) via Triton kernels (utils, wy_fast, chunk_delta_h/o + solve_tril) using offset+arange (no tl.make_block_ptr/tl.advance) and fused recurrent Triton path. Add naive.py PyTorch reference for testing. Port full parameter set in MomentumDeltaNet (a/m/e/b_proj, A/Mu/log_factor/D, dt/mu bias, tau, min_log_mu).

Fix all P1/P2 feedback from second review 5087504235:

• Remove dead chunk_indices chain (hardcoded None in 7 sites) in chunk.py:181
• Keep chunk_momentum_delta_rule/fused_recurrent_momentum_delta_rule as primary (degenerate chunk_mode_rule kept as compat alias)
• Remove P0 #n numbering from comments/tests, keep self-contained descriptions
• Remove dead scale parametrize in tests/ops/test_momentum_delta.py:16
• Remove dead w,A computation in fused_recurrent.py:53 (torch.empty_like(v) instead of prepare_wy_repr_fwd)
• Add assert mode in ['chunk','fused_recurrent'] in momentum_deltanet.py:88
• Inline use_qk_l2norm at call sites
• Compress header blocks to ≤3 lines
• Single quotes in __all__

Test plan

• grep -r make_block_ptr fla/ops/momentum_delta_rule 0
• pytest tests/ops/test_momentum_delta.py 14 passed (degenerate chunk vs fused + varlen + grad count + l2 + layer l2)
• Full momentum Triton vs naive bf16 max 0.00099, chunk vs fused 0.0014, bwd finite (normalized q/k/p, l2 + p*alpha)
• MomentumDeltaNet cuda bf16 T=64/128/256 eval + train bwd finite, short_conv/varlen fallback to degenerate works
• Full chunk varlen currently NotImplemented for Triton path (degenerate covers varlen), fused varlen supported for inference

Benchmark

New kernels reuse delta_rule patterns but add momentum terms; no make_block_ptr so Triton main compatible. Performance expected similar to chunk_delta_rule (same grid, extra log/exp in cumsum/pkt). Detailed H100 before/after to be added in follow-up.

Breaking changes

None for degenerate path. Full momentum adds new params (A_log, Mu_log, etc.) — old chunk_mode_rule checkpoints load with strict=False or via degenerate alias.

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
• Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
• This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.

If you ticked the "minor" box above

N/A

[zhiyuan1i]

Three P0 issues that break training on this PR:

1. Backward gradient count mismatch: chunk_momentum_delta_rule and fused_recurrent_momentum_delta_rule return 11 gradients (copied from chunk_delta_rule's template which has scale), but their forwards have 10 and 8 inputs respectively. First backward call crashes with "incorrect number of gradients". dh0 also lands in the output_final_state slot, so initial_state gets None.

2. fused_recurrent bwd passes raw v instead of u: the kernel expects u = v - k·h (what the fwd saves), but receives the original v. All of dk/db/dq are wrong. The layer routes q_len≤64 through this path.

3. Default qk_norm='l2' never takes effect: the op call doesn't pass use_qk_l2norm_in_kernel, so q/k are unnormalized under default config — silently diverging from DeltaNet family math.

Also: no tests at all. Any naive fwd/bwd parity check would catch all three. And there's no momentum coefficient anywhere — if the paper's momentum cancels algebraically, that needs a comment; otherwise the implementation doesn't match the paper.

Please fix the gradient returns, pass u to bwd, wire up l2norm, and add fwd/bwd + varlen numerical tests.

[zhiyuan1i]

Style issues (in addition to the P0 correctness issues from the previous review):

1. P1: chunk.py:186,198 — chunk_indices is computed and save_for_backward'd but never used: fwd/bwd hardcode chunk_indices=None everywhere (7 sites). Dead code chain, and each sub-kernel recomputes indices.

2. P1: Naming chunk_mode_rule/fused_recurrent_mode_rule is uninformative and misleading. Should be chunk_momentum_delta_rule/fused_recurrent_momentum_delta_rule per repo convention.

3. P1: Comments like "(fixes P0 #3)" and "Regression for P0 #1/#2/#3" reference review round numbers that are meaningless after merge. Remove the numbers, keep self-contained descriptions.

4. P1: test_momentum_delta.py:17,39 — scale parametrize is never used in the test body (op hardcodes 1/sqrt(D)), dead parameter.

5. P1: fused_recurrent.py:55 — w, u, A = prepare_wy_repr_fwd(...) computes w and A then immediately discards them. Dead computation.

6. P2: momentum_deltanet.py __init__ missing assert mode in ['chunk', 'fused_recurrent'] (delta_net.py:126 has it).

7. P2: use_qk_l2norm temp variable used only twice — inline at call site like delta_net.py:256.

8. P2: Header comment blocks (13 lines in chunk.py, 9 lines in momentum_deltanet.py) too long — compress to 3 lines max.

9. P2: __init__.py uses double quotes, repo convention is single quotes.

Also the biggest concern: this is named MomentumDeltaNet but the implementation is standard delta rule (mu→0, alpha=1 degenerate). If the paper's momentum cancels algebraically, that needs a prominent comment; otherwise the naming doesn't match the semantics.