From 508b53daa000a5a824d292a32febf01ba25d5e6f Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n Date: Sat, 1 Aug 2026 00:40:43 +0530 Subject: [PATCH] Preserve input validation under Python optimization --- flash_mla/flash_mla_interface.py | 70 ++++++++++------ tests/test_flash_mla_input_validation.py | 101 +++++++++++++++++++++++ 2 files changed, 144 insertions(+), 27 deletions(-) create mode 100644 tests/test_flash_mla_input_validation.py diff --git a/flash_mla/flash_mla_interface.py b/flash_mla/flash_mla_interface.py index a3740b0f..1e5ede55 100644 --- a/flash_mla/flash_mla_interface.py +++ b/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,8 +110,12 @@ 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 @@ -112,11 +123,20 @@ def flash_mla_with_kvcache( 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, diff --git a/tests/test_flash_mla_input_validation.py b/tests/test_flash_mla_input_validation.py new file mode 100644 index 00000000..16639053 --- /dev/null +++ b/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()