[wave] NSA: end-to-end training integration with autograd

[harsh-nod]

Parent

Part of #1243 — DeepSeek NSA kernels for MI350

Description

Wire up all NSA forward and backward kernels into a PyTorch-compatible autograd Function that enables end-to-end training.

Deliverables

1. NSAFunction(torch.autograd.Function) with:

   • forward(): runs the full NSA inference pipeline ([wave] NSA: end-to-end inference pipeline #1251), saves tensors for backward
   • backward(): orchestrates all backward kernels in correct order

2. Backward orchestration order:

   dO → gating backward → dO_cmp, dO_slc, dO_swa, dg_*  (#1254)
        ↓                    ↓                    ↓
   compressed attn bwd   selection attn bwd    sliding window bwd
   → dQ_cmp, dK_cmp,    → dQ_slc, dK_slc,   → dQ_swa, dK_swa,
     dV_cmp (#1253)        dV_slc (#1252)       dV_swa (FA bwd)
        ↓
   mean-pool bwd
   → dK_pool, dV_pool
   
   Aggregate: dQ = dQ_cmp + dQ_slc + dQ_swa
              dK = dK_pool + dK_slc + dK_swa
              dV = dV_pool + dV_slc + dV_swa
              dg_cmp, dg_slc, dg_swa (from gating bwd)
   

3. Memory optimization:

   • Activation checkpointing: recompute compressed KV and block_indices in backward instead of saving
   • Gradient accumulation support for pipeline/data parallelism

4. Convenience wrapper: nsa_attention(q, k, v, g_cmp, g_slc, g_swa, config) -> o that calls NSAFunction.apply

Testing

• torch.autograd.gradcheck with FP64 reference
• Compare training loss curves: NSA vs dense attention on a small model (should match within tolerance)

Depends on

• All backward kernel tickets: [wave] NSA: selection attention backward kernel (dQ, dK, dV) #1252, [wave] NSA: compressed attention backward kernel #1253, [wave] NSA: gating parameter gradients & top-k STE backward #1254
• Inference pipeline: [wave] NSA: end-to-end inference pipeline #1251