Preserve input validation under Python optimization

[Mr-Neutr0n]

Problem

The public Python wrappers use assert for required input validation. Python removes those statements under python -O, so invalid inputs can reach the CUDA extension.

A reproduced example is sparse decode with causal=True and is_fp8_kvcache=False: current main calls the backend in optimized mode instead of rejecting the unsupported combination. Reuse-consistency checks, the legacy num_splits guard, and unsupported prefill options are removed in the same way.

Invalid first calls can also set have_initialized before the later mode check fails, leaving scheduler state behind for a request that never ran.

Fix

• replace public assert statements with explicit argument exceptions that survive optimization
• use TypeError for the scheduler object type, ValueError for unsupported arguments and reuse mismatches, and RuntimeError for an impossible internal scheduler state
• validate dense/sparse mode requirements before committing first-call scheduler state
• preserve the accepted inputs and backend call signatures

Regression coverage

The new CPU-only tests isolate the wrapper from the CUDA extension and verify under both regular and optimized Python that:

• invalid sparse options never reach the backend and do not initialize scheduler state
• a non-None legacy num_splits placeholder is rejected
• nonzero dropout is rejected by all three prefill entry points

Validation

• uv run --with torch --with numpy python -m unittest tests/test_flash_mla_input_validation.py
• uv run --with torch --with numpy python -O -m unittest tests/test_flash_mla_input_validation.py
• uvx ruff check tests/test_flash_mla_input_validation.py
• python3 -m compileall -q flash_mla tests/test_flash_mla_input_validation.py
• verified no public assert statements remain in flash_mla_interface.py
• git diff --check

This is independent of #201: that PR makes scheduler initialization transactional when the backend itself raises; this PR ensures invalid arguments cannot bypass validation or commit state before a backend call.

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

[Mr-Neutr0n]

Rebased onto current main (post Ascend/V4.1 sparse-only decode). Conflict resolution kept the original goal: public input checks now raise exceptions that survive python -O instead of assert, adapted to today’s sparse-required API/messages (no revived dense path).

Verified mergeable vs main. CPU-only regression tests pass under both normal and -O:

• python -m unittest tests.test_flash_mla_input_validation
• python -O -m unittest tests.test_flash_mla_input_validation

Independent of #201 (scheduler rollback on backend failure).