Preserve input validation under Python optimization

flash_mla/flash_mla_interface.py

@@ -5,6 +5,13 @@
 
 import flash_mla.cuda as flash_mla_cuda
 
+
+def _check_argument(condition: bool, message: str, error_type=ValueError) -> None:
+    """Validate public API inputs even when Python assertions are disabled."""
+    if not condition:
+        raise error_type(message)
+
+
 @dataclasses.dataclass
 class FlashMLASchedMeta:
     """
@@ -103,20 +110,33 @@ def flash_mla_with_kvcache(
     """
     sched_meta = tile_scheduler_metadata
     indices_in_kvcache = indices
-    assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta"
-    assert num_splits is None, "num_splits must be None"
+    _check_argument(
+        isinstance(sched_meta, FlashMLASchedMeta),
+        "tile_scheduler_metadata must be of type FlashMLASchedMeta",
+        TypeError,
+    )
+    _check_argument(num_splits is None, "num_splits must be None")
 
     topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None
     extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None
     extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
 
+    if topk is not None:
+        _check_argument(not causal, "causal must be False when sparse attention is enabled")
+        _check_argument(is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled")
+    else:
+        _check_argument(
+            indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None,
+            "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used.",
+        )
+        _check_argument(
+            block_table is not None and cache_seqlens is not None,
+            "block_table and cache_seqlens must be provided when dense attention is used.",
+        )
+
     if not sched_meta.have_initialized:
-        # Sanity check. We only perform sanity check during the first invocation to save CPU time.
-        if indices_in_kvcache is not None:
-            assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)"
-
         # Initialize the tile scheduler metadata during the first invocation.
         sched_meta.have_initialized = True
         sched_meta.config = FlashMLASchedMeta.Config(
@@ -136,22 +156,20 @@ def flash_mla_with_kvcache(
     else:
         # Check whether the input arguments are consistent with sched_meta
         helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta."
-        assert sched_meta.config is not None
-        assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg
-        assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg
-        assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg
-        assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg
-        assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg
-        assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg
-        assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg
-        assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg
-        assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg
-        assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg
+        _check_argument(sched_meta.config is not None, "initialized sched_meta has no config", RuntimeError)
+        _check_argument(sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg)
+        _check_argument(sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg)
+        _check_argument(sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg)
+        _check_argument(sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg)
+        _check_argument(sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg)
+        _check_argument(sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg)
+        _check_argument(sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg)
+        _check_argument(sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg)
+        _check_argument(sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg)
+        _check_argument(sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg)
 
     if topk is not None:
         # Sparse attention
-        assert not causal, "causal must be False when sparse attention is enabled"
-        assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled"
         out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd(
             q, k_cache, indices_in_kvcache, topk_length, attn_sink,
             sched_meta.tile_scheduler_metadata, sched_meta.num_splits,
@@ -160,8 +178,6 @@ def flash_mla_with_kvcache(
         )
     else:
         # Dense attention
-        assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used."
-        assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used."
         out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd(
             q, k_cache, head_dim_v,
             cache_seqlens, block_table,
@@ -383,8 +399,8 @@ def flash_attn_varlen_func(
     deterministic: bool = False,
     is_varlen: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    assert dropout_p == 0.0
-    assert not deterministic
+    _check_argument(dropout_p == 0.0, "dropout_p must be 0.0")
+    _check_argument(not deterministic, "deterministic must be False")
     return FlashAttnVarlenFunc.apply(
         q, k, v,
         cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv,
@@ -403,8 +419,8 @@ def flash_attn_varlen_qkvpacked_func(
     deterministic: bool = False,
     is_varlen: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    assert dropout_p == 0.0
-    assert not deterministic
+    _check_argument(dropout_p == 0.0, "dropout_p must be 0.0")
+    _check_argument(not deterministic, "deterministic must be False")
     return FlashAttnVarlenFunc.apply(
         qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:],
         cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
@@ -426,8 +442,8 @@ def flash_attn_varlen_kvpacked_func(
     deterministic: bool = False,
     is_varlen: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    assert dropout_p == 0.0
-    assert not deterministic
+    _check_argument(dropout_p == 0.0, "dropout_p must be 0.0")
+    _check_argument(not deterministic, "deterministic must be False")
     return FlashAttnVarlenFunc.apply(
         q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:],
         cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv,

tests/test_flash_mla_input_validation.py

@@ -0,0 +1,101 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def load_interface(cuda_backend):
+    """Load the Python wrapper with an isolated fake CUDA extension."""
+    package = types.ModuleType("flash_mla")
+    package.__path__ = []
+    package.cuda = cuda_backend
+
+    saved_package = sys.modules.get("flash_mla")
+    saved_cuda = sys.modules.get("flash_mla.cuda")
+    sys.modules["flash_mla"] = package
+    sys.modules["flash_mla.cuda"] = cuda_backend
+    try:
+        path = Path(__file__).parents[1] / "flash_mla" / "flash_mla_interface.py"
+        spec = importlib.util.spec_from_file_location("flash_mla_interface_under_test", path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+    finally:
+        if saved_package is None:
+            sys.modules.pop("flash_mla", None)
+        else:
+            sys.modules["flash_mla"] = saved_package
+        if saved_cuda is None:
+            sys.modules.pop("flash_mla.cuda", None)
+        else:
+            sys.modules["flash_mla.cuda"] = saved_cuda
+
+
+class InputValidationTests(unittest.TestCase):
+    def setUp(self):
+        self.backend = types.ModuleType("flash_mla.cuda")
+        self.backend.calls = 0
+
+        def sparse_decode_fwd(q, *args):
+            self.backend.calls += 1
+            return q, q, None, None
+
+        self.backend.sparse_decode_fwd = sparse_decode_fwd
+        self.interface = load_interface(self.backend)
+
+    def test_invalid_sparse_options_do_not_reach_backend(self):
+        sched_meta, _ = self.interface.get_mla_metadata()
+        q = torch.empty(1, 1, 64, 576)
+        k_cache = torch.empty(1, 64, 1, 656)
+        indices = torch.empty(1, 1, 64, dtype=torch.int32)
+
+        with self.assertRaisesRegex(ValueError, "causal must be False"):
+            self.interface.flash_mla_with_kvcache(
+                q,
+                k_cache,
+                None,
+                None,
+                512,
+                sched_meta,
+                causal=True,
+                is_fp8_kvcache=False,
+                indices=indices,
+            )
+
+        self.assertFalse(sched_meta.have_initialized)
+        with self.assertRaisesRegex(ValueError, "is_fp8_kvcache must be True"):
+            self.interface.flash_mla_with_kvcache(
+                q,
+                k_cache,
+                None,
+                None,
+                512,
+                sched_meta,
+                is_fp8_kvcache=False,
+                indices=indices,
+            )
+
+        self.assertFalse(sched_meta.have_initialized)
+        self.assertEqual(self.backend.calls, 0)
+
+    def test_legacy_num_splits_placeholder_is_rejected(self):
+        sched_meta, _ = self.interface.get_mla_metadata()
+        with self.assertRaisesRegex(ValueError, "num_splits must be None"):
+            self.interface.flash_mla_with_kvcache(None, None, None, None, 512, sched_meta, num_splits=1)
+
+    def test_unsupported_prefill_options_are_rejected(self):
+        calls = (
+            (self.interface.flash_attn_varlen_func, (None, None, None, None, None, None, None)),
+            (self.interface.flash_attn_varlen_qkvpacked_func, (None, None, None, None)),
+            (self.interface.flash_attn_varlen_kvpacked_func, (None, None, None, None, None, None, None)),
+        )
+        for function, args in calls:
+            with self.subTest(function=function.__name__), self.assertRaisesRegex(ValueError, "dropout_p must be 0.0"):
+                function(*args, dropout_p=0.1)
+
+
+if __name__ == "__main__":
+    unittest.main()