[Main] fix(absorbed-mla): handle delayed weight gradients

[Wohox]

Summary

Port of 99008e2 (landed on dev via #6020) to main, reduced to the part that actually takes effect on main, plus the corresponding unit test.

AbsorbedMLASelfAttention consumes linear_kv_up_proj.weight directly (_get_kv_up_weights) instead of running a GEMM through the module, so that module never executes a forward that can enqueue Transformer Engine's delayed-wgrad closure. Under --delay-wgrad-compute, backward_dw() then popped a closure that was never queued.

Changes

• Build linear_kv_up_proj with a shallow config copy that sets delay_wgrad_compute = False, leaving its gradient on the plain autograd path.
• Drop linear_kv_up_proj.backward_dw() from _backward_kv_proj() accordingly.
• Add tests/unit_tests/transformer/experimental_attention_variant/test_attention_delay_wgrad.py: under delay_wgrad_compute=True the forward-executing MLA linears (linear_q_down_proj, linear_q_up_proj, linear_kv_down_proj, linear_proj) stay deferred until backward_dw(), while linear_kv_up_proj.weight.grad is already present after loss.backward() and is left bit-identical by the flush. An eager (delay_wgrad_compute=False) reference run pins that deferring changes only when the wgrad GEMM runs, never the gradient values.

Every projection that executes a GEMM continues to defer its weight gradient.

Deltas versus the dev commit

• The dev commit also forwards the flush to core_attention from AbsorbedMLASelfAttention.backward_dw(). That is omitted here: DSAttention does not define backward_dw() on main, so the call would be a guarded no-op.
• The dev hunk also covers the split K/V up-projection variant (linear_k_up_proj / linear_v_up_proj), which does not exist on main. Only the combined path is carried over.
• The test is ported from the dev file of the same name, restricted to the absorbed-MLA coverage that applies on main (the CSA/dsv4_hybrid class and the DSA indexer-flush assertions are dev-only).

Known gap, not addressed here

On main, DSAIndexer's linear_wq_b / linear_wk / linear_weights_proj are built from the shared config, so they defer their wgrads under delay_wgrad_compute — and nothing flushes them, since DSAttention has no backward_dw() and MLASelfAttention.backward_dw() does not traverse core_attention. That is pre-existing on main and orthogonal to this fix; porting DSAttention.backward_dw() from dev is a separate change.

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