[CuTe, Bwd] Return zero gradients for empty Q/K workloads

[guoriyue]

Summary

• extend the reviewed empty-workload backward shortcut from learnable-sink attention to standard Q/K/V attention;
• return mathematically correct zero dQ/dK/dV before scratch allocation, preprocessing, JIT selection, or attention-kernel launch;
• zero and reuse caller-provided gradient buffers;
• preserve the existing learnable-sink dSink behavior.

Problem

CuTe forward already defines empty workloads: empty Q returns an empty output, while empty K returns zero output with -inf LSE. Standard backward did not share that contract.

On RTX 5090 / SM120 with BF16, unmodified main gives:

• dense Sq=0: backward preprocessing rejects an inferred empty-tensor stride;
• dense Sk=0: the main backward path reaches a zero-work CUDA launch and returns runtime error 9;
• batch=0: the same dense preprocessing/launch failures;
• focused baseline result: 8 failed, 6 passed.

The six passing baseline cases are the varlen combinations: their current autograd route already happens to return successfully. They remain regression coverage for the shared mathematical contract; this PR does not claim those six were broken.

Why the early return is correct

When there are no query rows, no loss term contributes to K or V, so dK=dV=0 and dQ is empty. When there are no K/V rows, the output is independent of Q, so dQ=0 and dK/dV are empty.

The existing zero-gradient logic was incorrectly nested under learnable_sink. Hoisting it to the common Q/K/V dispatch boundary applies the same rule to ordinary attention before any backward implementation detail sees an empty tensor.

This is useful for ragged or distributed training where one microbatch, sequence, or rank can legitimately receive zero tokens.

Validation

RTX 5090 / SM120:

• unmodified main: 8 failed, 6 passed;
• patch: all 14 dense, varlen, causal/non-causal, batch-empty, and preallocated-buffer cases pass;
• all 14 pass under CuTe FakeTensor mode;
• preallocated dQ/dK/dV buffers are returned by identity and zeroed;
• dedicated cache assertions prove preprocess, main backward, and postprocess compile caches remain empty;
• 22 existing empty-Q/empty-K forward regressions pass;
• git diff --check passes.

Scope

zeros_like or in-place zeroing may still perform the device work required to materialize returned gradients; the claim is specifically that no attention scratch space is allocated and no backward attention/preprocess/postprocess kernel is compiled or launched. Sparse MLA/QV backward uses a separate implementation and remains out of scope.