Mark auxiliary LSE output non-differentiable

[Mr-Neutr0n]

Problem

The SM100 prefill autograd function returns both attention output and an auxiliary log-sum-exp tensor. PyTorch currently marks both outputs as differentiable when any input requires gradients, but backward explicitly discards dlse. This advertises an LSE gradient path that the implementation does not support.

Fix

Mark the returned LSE tensor non-differentiable with ctx.mark_non_differentiable. LSE remains saved for computing the supported attention-output gradient.

Regression coverage

The new CPU-only test loads the wrapper with an isolated fake CUDA backend and verifies that:

• the attention output remains differentiable
• LSE has neither requires_grad nor a gradient function

Validation

• python -m unittest tests.test_flash_mla_autograd_unit
• ruff check tests/test_flash_mla_autograd_unit.py
• Python compilation and git diff --check

This is independent of #201 and #202; it changes only the SM100 prefill autograd contract.

Prepared with OpenAI Codex assistance; I reproduced the unsupported gradient path and reviewed the change.