Fix scheduler state after backend failure

[Mr-Neutr0n]

Problem

flash_mla_with_kvcache currently sets have_initialized and stores the first-call configuration before the CUDA backend has generated scheduling metadata.

If that backend call raises, the FlashMLASchedMeta object is left in a partially initialized state: it has a committed configuration but no metadata. A retry with corrected input shapes is then rejected by the stale configuration checks before it can reach the backend.

Fix

Build the prospective first-call configuration locally, invoke the backend, and commit the configuration and initialized flag only after the backend returns successfully. Existing metadata reuse and consistency checks are unchanged after a successful call.

Regression coverage

The new CPU-only wrapper test fault-injects the CUDA extension:

1. the first scheduling call raises;
2. the scheduler remains completely uninitialized;
3. a retry with a different, corrected batch shape reaches the backend and succeeds;
4. only that successful call commits the scheduler configuration.

Before this fix, step 2 fails and the retry is rejected without a second backend call.

Validation

• uv run --with torch --with numpy python -m unittest tests/test_flash_mla_interface_unit.py
• uvx ruff check tests/test_flash_mla_interface_unit.py
• python3 -m compileall -q flash_mla tests/test_flash_mla_interface_unit.py
• git diff --check

Prepared with OpenAI Codex assistance; I reproduced the failure and reviewed the change.