diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index b0088c4798e..4bb2359bdb9 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/extension/llm/custom_ops/custom_ops.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -16,22 +17,42 @@ import torch +from executorch.extension.llm.custom_ops.custom_sdpa_schema import ( + custom_sdpa_fake, + CustomSDPANativeOpUnavailable, +) from torch._inductor.lowering import lowerings as L, register_lowering from torch.library import impl aten = torch.ops.aten -try: - op = torch.ops.llama.sdpa_with_kv_cache.default - assert op is not None - op2 = torch.ops.llama.fast_hadamard_transform.default - assert op2 is not None - op3 = torch.ops.llama.quantized_moe_ffn.default - assert op3 is not None -except: - # This is needed to ensure that custom ops are registered - from executorch.extension.pybindings import portable_lib # noqa # usort: skip + +def _native_custom_ops_are_registered() -> bool: + try: + return all( + op is not None + for op in ( + torch.ops.llama.sdpa_with_kv_cache.default, + torch.ops.llama.fast_hadamard_transform.default, + torch.ops.llama.quantized_moe_ffn.default, + torch.ops.llama.custom_sdpa.default, + ) + ) + except (AttributeError, RuntimeError): + return False + + +def _ensure_native_custom_ops_loaded() -> None: + if _native_custom_ops_are_registered(): + return + + try: + import executorch.extension.pybindings.portable_lib # noqa: F401 # usort: skip + except ImportError as exc: + raise CustomSDPANativeOpUnavailable( + "portable_lib is unavailable, so native LLM custom ops cannot be loaded" + ) from exc # Ideally package is installed in only one location but usage of # PYATHONPATH can result in multiple locations. @@ -42,18 +63,45 @@ logging.info(f"Looking for libcustom_ops_aot_lib.so in {package_path}") libs = list(package_path.glob("**/*custom_ops_aot_lib.*")) + if not libs: + raise CustomSDPANativeOpUnavailable( + f"custom_ops_aot_lib was not found under {package_path}" + ) + if len(libs) != 1: + raise RuntimeError(f"Expected 1 library but got {len(libs)}") - assert len(libs) == 1, f"Expected 1 library but got {len(libs)}" logging.info(f"Loading custom ops library: {libs[0]}") torch.ops.load_library(libs[0]) - op = torch.ops.llama.sdpa_with_kv_cache.default - assert op is not None - op2 = torch.ops.llama.fast_hadamard_transform.default - assert op2 is not None - op3 = torch.ops.llama.quantized_moe_ffn.default - assert op3 is not None + if not _native_custom_ops_are_registered(): + raise RuntimeError("Native LLM custom op registration is incomplete") + + +_ensure_native_custom_ops_loaded() custom_ops_lib = torch.library.Library("llama", "IMPL") +impl(custom_ops_lib, "custom_sdpa", "Meta")(custom_sdpa_fake) + + +def custom_sdpa( + query, + key, + value, + start_pos, + attn_mask=None, + drpout_p=0.0, + is_causal=False, + scale=None, +): + return torch.ops.llama.custom_sdpa( + query, + key, + value, + start_pos, + attn_mask, + drpout_p, + is_causal, + scale, + ) def _validate_params( @@ -107,6 +155,19 @@ def _validate_params( assert ( key_cache.size() == value_cache.size() ), f"Key cache and value cache must have same size but got {key_cache.size()} and {value_cache.size()}" + assert query.size(0) == key_cache.size(0) == value_cache.size(0), ( + "Query, key cache, and value cache must have matching batch sizes but got " + f"{query.size(0)}, {key_cache.size(0)}, and {value_cache.size(0)}" + ) + assert query.size(3) == key_cache.size(3) == value_cache.size(3), ( + "Query, key cache, and value cache must have matching head dimensions " + f"but got {query.size(3)}, {key_cache.size(3)}, and {value_cache.size(3)}" + ) + assert key_cache.size(2) > 0, "Expected key/value cache to have at least one head." + assert query.size(2) % key_cache.size(2) == 0, ( + "Query head count must match the KV head count or be divisible by it, " + f"but got {query.size(2)} and {key_cache.size(2)}" + ) # These asserts are real but they require me to add constrain_as_size/value calls to the model and I dont want to do that right now # assert start_pos < key_cache.size( @@ -265,35 +326,6 @@ def quantized_moe_ffn_meta( return torch.empty((x.size(0), dim), dtype=torch.float32, device=x.device) -@impl(custom_ops_lib, "custom_sdpa", "Meta") -def custom_sdpa( - query, - key_cache, - value_cache, - start_pos, - attn_mask=None, - drpout_p=0.0, - is_causal=False, - scale=None, -): - seq_len = query.size(1) - _validate_params( - query, - key_cache, - value_cache, - key_cache, - value_cache, - start_pos, - seq_len, - attn_mask, - drpout_p, - is_causal, - scale, - ) - - return torch.empty(query.shape, dtype=query.dtype, device=query.device) - - def _validate_update_cache_params( value, cache, diff --git a/extension/llm/custom_ops/custom_sdpa_schema.py b/extension/llm/custom_ops/custom_sdpa_schema.py new file mode 100644 index 00000000000..93298e073fe --- /dev/null +++ b/extension/llm/custom_ops/custom_sdpa_schema.py @@ -0,0 +1,105 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + + +class CustomSDPANativeOpUnavailable(ImportError): + """Raised when the compiled custom SDPA op is unavailable.""" + + +def _validate_custom_sdpa_output_shape_inputs(query, key, value): + assert query.dim() == key.dim() == value.dim() == 4, ( + "custom_sdpa requires 4D Q/K/V tensors but got " + f"{query.dim()}D, {key.dim()}D, and {value.dim()}D" + ) + supported_dtypes = (torch.float32, torch.bfloat16, torch.float16) + assert query.dtype in supported_dtypes, ( + "custom_sdpa requires query to use float32, bfloat16, or float16 " + f"but got {query.dtype}" + ) + assert key.dtype == query.dtype == value.dtype, ( + "custom_sdpa requires matching Q/K/V dtypes but got " + f"{query.dtype}, {key.dtype}, and {value.dtype}" + ) + assert ( + tuple(query.shape[:-3]) == tuple(key.shape[:-3]) == tuple(value.shape[:-3]) + ), ( + "custom_sdpa does not support batch broadcasting; expected query, key, " + f"and value to have matching batch shape but got {query.shape}, " + f"{key.shape}, and {value.shape}" + ) + assert query.shape[-1] == key.shape[-1] == value.shape[-1], ( + "custom_sdpa requires matching Q/K/V head dimensions but got " + f"{query.shape[-1]}, {key.shape[-1]}, and {value.shape[-1]}" + ) + assert key.shape[-2] == value.shape[-2], ( + "custom_sdpa requires key and value to have matching KV head counts " + f"but got {key.shape[-2]} and {value.shape[-2]}" + ) + assert key.shape[-3] == value.shape[-3], ( + "custom_sdpa requires key and value to have matching sequence lengths " + f"but got {key.shape[-3]} and {value.shape[-3]}" + ) + assert key.shape[-2] > 0, "custom_sdpa requires at least one KV head" + assert query.shape[-2] % key.shape[-2] == 0, ( + "custom_sdpa requires the query head count to match the KV head count " + f"or be divisible by it, but got {query.shape[-2]} and {key.shape[-2]}" + ) + + +def validate_custom_sdpa_args( + query, + key, + value, + start_pos, + attn_mask=None, + drpout_p=0.0, + is_causal=False, + scale=None, +): + del start_pos, drpout_p, scale + _validate_custom_sdpa_output_shape_inputs(query, key, value) + assert not ( + attn_mask is not None and is_causal + ), "attn_mask and is_causal cannot be set at the same time" + if attn_mask is not None: + assert attn_mask.dim() == 2, ( + "custom_sdpa requires a 2D attention mask but got " f"{attn_mask.dim()}D" + ) + assert attn_mask.dtype == torch.float32, ( + "custom_sdpa requires attn_mask to use float32 but got " + f"{attn_mask.dtype}" + ) + assert attn_mask.shape == (query.shape[-3], key.shape[-3]), ( + "custom_sdpa requires attn_mask to match [q_len, kv_len], but got " + f"{tuple(attn_mask.shape)} for q_len={query.shape[-3]} " + f"and kv_len={key.shape[-3]}" + ) + + +def custom_sdpa_fake( + query, + key, + value, + start_pos, + attn_mask=None, + drpout_p=0.0, + is_causal=False, + scale=None, +): + validate_custom_sdpa_args( + query, + key, + value, + start_pos, + attn_mask, + drpout_p, + is_causal, + scale, + ) + return torch.empty(query.shape, dtype=query.dtype, device=query.device) diff --git a/extension/llm/custom_ops/targets.bzl b/extension/llm/custom_ops/targets.bzl index 4c2efb9adaa..8251ac98349 100644 --- a/extension/llm/custom_ops/targets.bzl +++ b/extension/llm/custom_ops/targets.bzl @@ -1,3 +1,7 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. load("@fbsource//xplat/executorch/build:build_variables.bzl", "EXTENSION_LLM_CUSTOM_OPS_BUCK_SRCS") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "is_xplat", "runtime") load( @@ -157,6 +161,7 @@ def define_common_targets(): runtime.python_library( name = "custom_ops_aot_py", srcs = [ + "custom_sdpa_schema.py", "custom_ops.py", ], visibility = ["PUBLIC"], diff --git a/extension/llm/export/export_passes.py b/extension/llm/export/export_passes.py index a3cb2cdd98a..d1fd4ea22a7 100644 --- a/extension/llm/export/export_passes.py +++ b/extension/llm/export/export_passes.py @@ -1,9 +1,14 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + import logging import torch - from executorch.exir.pass_base import ExportPass from torch._subclasses import FakeTensor +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.passes.infra.pass_base import PassResult @@ -111,31 +116,26 @@ def __init__(self, assume_causal_mask=False): self.assume_causal_mask = assume_causal_mask def call_operator(self, op, args, kwargs, meta): - from executorch.extension.llm.custom_ops import custom_ops # noqa - if op != torch.ops.aten.scaled_dot_product_attention.default: return super().call_operator(op, args, kwargs, meta) - q, k, v, mask, dropout, is_causal, scale = self._extract_args(args, kwargs) + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 - qT = self._transpose(q, meta) - kT = self._transpose(k, meta) - vT = self._transpose(v, meta) - - if not ( - q.node.meta["val"].dim() - == k.node.meta["val"].dim() - == v.node.meta["val"].dim() - == 4 - ): - logging.info("ReplaceSDPAWithCustomSDPAPass only supports 4D QKV inputs.") - return super().call_operator(op, args, kwargs, meta) + q, k, v, mask, dropout, is_causal, scale, enable_gqa = self._extract_args( + args, kwargs + ) if self.assume_causal_mask: - # Ignore specified mask simply set the is_causal flag. mask = None is_causal = True + if not self._can_rewrite_sdpa(q, k, v, mask, dropout, is_causal, enable_gqa): + return super().call_operator(op, args, kwargs, meta) + + qT = self._transpose(q, meta) + kT = self._transpose(k, meta) + vT = self._transpose(v, meta) + if mask is not None: mask_fake_tensor = mask.node.meta["val"] if mask_fake_tensor.dim() > 2: @@ -169,12 +169,66 @@ def call_operator(self, op, args, kwargs, meta): ) return self._transpose(custom_sdpa, meta) + def _can_rewrite_sdpa(self, q, k, v, mask, dropout, is_causal, enable_gqa): + if not ( + q.node.meta["val"].dim() + == k.node.meta["val"].dim() + == v.node.meta["val"].dim() + == 4 + ): + logging.info("ReplaceSDPAWithCustomSDPAPass only supports 4D QKV inputs.") + return False + + if not self._supports_head_layout(q, k, v, enable_gqa): + logging.info( + "ReplaceSDPAWithCustomSDPAPass requires matching heads or " + "enable_gqa=True with compatible query/KV heads." + ) + return False + if not self._supports_native_shape_contract(q, k, v): + logging.info( + "ReplaceSDPAWithCustomSDPAPass requires matching batch shapes " + "and Q/K/V head dimensions." + ) + return False + if dropout != 0.0: + logging.info( + "ReplaceSDPAWithCustomSDPAPass does not rewrite SDPA with " + "non-zero dropout." + ) + return False + if q.node.meta["val"].dtype not in ( + torch.float32, + torch.bfloat16, + torch.float16, + ): + logging.info( + "ReplaceSDPAWithCustomSDPAPass only rewrites float32, " + "bfloat16, or float16 SDPA inputs." + ) + return False + if not self._supports_unmasked_attention_contract(q, k, v, mask, is_causal): + logging.info( + "ReplaceSDPAWithCustomSDPAPass does not rewrite unmasked SDPA " + "when the native custom-op KV-prefix contract would change " + "aten semantics." + ) + return False + if not self._supports_custom_mask_contract(q, k, mask): + logging.info( + "ReplaceSDPAWithCustomSDPAPass only rewrites attention masks " + "that already match the custom SDPA contract." + ) + return False + return True + def _extract_args(self, args, kwargs): q, k, v, *rest = args mask = None dropout = 0.0 is_causal = False scale = None + enable_gqa = False if len(rest) > 0: mask = rest[0] if len(rest) > 1: @@ -183,8 +237,98 @@ def _extract_args(self, args, kwargs): is_causal = rest[2] if "scale" in kwargs: scale = kwargs["scale"] + if "enable_gqa" in kwargs: + enable_gqa = kwargs["enable_gqa"] + + return q, k, v, mask, dropout, is_causal, scale, enable_gqa + + @staticmethod + def _known_true(expr) -> bool: + try: + return statically_known_true(expr) + except Exception: + return False + + @staticmethod + def _fake_val(x): + return x.node.meta["val"] + + @classmethod + def _supports_head_layout(cls, q, k, v, enable_gqa): + q_heads = cls._fake_val(q).shape[1] + k_heads = cls._fake_val(k).shape[1] + v_heads = cls._fake_val(v).shape[1] + + if not cls._known_true(k_heads == v_heads): + return False + if cls._known_true(q_heads == k_heads): + return True + if cls._known_true(k_heads == 1): + return True + if not enable_gqa: + return False + return cls._known_true((q_heads % k_heads) == 0) + + @classmethod + def _supports_native_shape_contract(cls, q, k, v): + q_shape = cls._fake_val(q).shape + k_shape = cls._fake_val(k).shape + v_shape = cls._fake_val(v).shape + + batch_dims_match = all( + cls._known_true(q_dim == k_dim) and cls._known_true(k_dim == v_dim) + for q_dim, k_dim, v_dim in zip(q_shape[:-3], k_shape[:-3], v_shape[:-3]) + ) + return ( + batch_dims_match + and cls._known_true(q_shape[-1] == k_shape[-1]) + and cls._known_true(k_shape[-1] == v_shape[-1]) + and cls._known_true(k_shape[-2] == v_shape[-2]) + and cls._known_true(k_shape[-3] == v_shape[-3]) + ) + + @classmethod + def _supports_unmasked_attention_contract(cls, q, k, v, mask, is_causal): + if mask is not None: + return True + + q_shape = cls._fake_val(q).shape + k_shape = cls._fake_val(k).shape + v_shape = cls._fake_val(v).shape + q_seq_len = q_shape[2] + k_seq_len = k_shape[2] + v_seq_len = v_shape[2] + + if is_causal: + return cls._known_true(q_seq_len <= k_seq_len) and cls._known_true( + q_seq_len <= v_seq_len + ) + + return cls._known_true(q_seq_len == k_seq_len) and cls._known_true( + k_seq_len == v_seq_len + ) - return q, k, v, mask, dropout, is_causal, scale + @classmethod + def _supports_custom_mask_contract(cls, q, k, mask): + if mask is None: + return True + + mask_fake_tensor = cls._fake_val(mask) + if mask_fake_tensor.dim() < 2: + return False + if mask_fake_tensor.dim() > 2 and any( + not cls._known_true(d == 1) for d in mask_fake_tensor.size()[:-2] + ): + return False + + if mask_fake_tensor.dtype not in (torch.bool, torch.float32): + return False + + q_seq_len = cls._fake_val(q).shape[2] + k_seq_len = cls._fake_val(k).shape[2] + return cls._known_true(mask_fake_tensor.shape[-2] == q_seq_len) and ( + cls._known_true(mask_fake_tensor.shape[-1] == k_seq_len) + ) def _transpose(self, x, meta): transpose = super().call_operator( diff --git a/extension/llm/export/test_export_passes.py b/extension/llm/export/test_export_passes.py index 1b59892ff88..20eb261898e 100644 --- a/extension/llm/export/test_export_passes.py +++ b/extension/llm/export/test_export_passes.py @@ -1,7 +1,16 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + import unittest import torch +from executorch.extension.llm.custom_ops.custom_sdpa_schema import ( + custom_sdpa_fake, + CustomSDPANativeOpUnavailable, +) from executorch.extension.llm.export.export_passes import ( RemoveRedundantTransposes, ReplaceSDPAWithCustomSDPAPass, @@ -172,6 +181,16 @@ def forward(self, x, mask, is_causal): x, x, x, attn_mask=mask, is_causal=is_causal ) + @classmethod + def setUpClass(cls): + super().setUpClass() + try: + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + except CustomSDPANativeOpUnavailable as exc: + raise unittest.SkipTest( + f"llama.custom_sdpa native op is unavailable: {exc}" + ) from exc + def setUp(self): torch.manual_seed(0) @@ -190,6 +209,30 @@ def _test(self, args, assume_causal_mask=False): y2 = gm(*args) self.assertTrue(torch.allclose(y1, y2)) + def _assert_not_rewritten(self, model, args, *, compare_outputs=True): + gm = export(model, args, strict=True).module() + + sdpa_key = "torch.ops.aten.scaled_dot_product_attention.default" + custom_sdpa_key = "torch.ops.llama.custom_sdpa.default" + gm = ReplaceSDPAWithCustomSDPAPass()(gm).graph_module + FileCheck().check_count(sdpa_key, 1, exactly=True).run(gm.code) + FileCheck().check_count(custom_sdpa_key, 0, exactly=True).run(gm.code) + if compare_outputs: + expected = model(*args) + actual = gm(*args) + torch.testing.assert_close(expected, actual) + + def _assert_rewritten(self, model, args, assume_causal_mask=False): + gm = export(model, args, strict=True).module() + + sdpa_key = "torch.ops.aten.scaled_dot_product_attention.default" + custom_sdpa_key = "torch.ops.llama.custom_sdpa.default" + FileCheck().check_count(sdpa_key, 1, exactly=True).run(gm.code) + gm = ReplaceSDPAWithCustomSDPAPass(assume_causal_mask)(gm).graph_module + FileCheck().check_count(sdpa_key, 0, exactly=True).run(gm.code) + FileCheck().check_count(custom_sdpa_key, 1, exactly=True).run(gm.code) + return gm + def test_causal_mask(self): self._test((torch.rand(1, 4, 32, 64), None, True)) @@ -207,3 +250,279 @@ def test_squeezable_mask(self): m2 = torch.tril(torch.ones(32, 32, dtype=torch.bool), diagonal=-16) m = torch.logical_xor(m1, m2).view(1, 1, 32, 32) self._test((torch.rand(1, 4, 32, 64), m, False)) + + def test_float16_mask_is_not_rewritten(self): + mask = torch.zeros(32, 32, dtype=torch.float16) + self._assert_not_rewritten( + self.TestModule(), + (torch.rand(1, 4, 32, 64, dtype=torch.float16), mask, False), + ) + + def test_broadcastable_2d_mask_is_not_rewritten(self): + mask = torch.zeros(1, 32, dtype=torch.float32) + self._assert_not_rewritten( + self.TestModule(), + (torch.rand(1, 4, 32, 64), mask, False), + ) + + def test_dropout_is_not_rewritten(self): + class DropoutModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.5, + is_causal=False, + ) + + q = torch.rand(1, 4, 32, 64) + k = torch.rand(1, 4, 32, 64) + v = torch.rand(1, 4, 32, 64) + self._assert_not_rewritten(DropoutModule(), (q, k, v), compare_outputs=False) + + def test_unsupported_dtype_is_not_rewritten(self): + class Float64Module(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(1, 4, 32, 64, dtype=torch.float64) + k = torch.rand(1, 4, 32, 64, dtype=torch.float64) + v = torch.rand(1, 4, 32, 64, dtype=torch.float64) + self._assert_not_rewritten(Float64Module(), (q, k, v)) + + def test_gqa_masked_decode_exercises_attention_math(self): + class MaskedGQADecodeModule(torch.nn.Module): + def forward(self, q, k, v, mask): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + enable_gqa=True, + ) + + q = torch.rand(1, 32, 1, 16) + k = torch.rand(1, 8, 4, 16) + v = torch.rand(1, 8, 4, 16) + mask = torch.tensor([[0.0, -4.0, float("-inf"), -1.5]], dtype=torch.float32) + + gm = self._assert_rewritten(MaskedGQADecodeModule(), (q, k, v, mask)) + expected = MaskedGQADecodeModule()(q, k, v, mask) + actual = gm(q, k, v, mask) + torch.testing.assert_close(expected, actual) + + def test_gqa_causal_prefill_exercises_attention_math(self): + class GQACausalPrefillModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + enable_gqa=True, + ) + + q = torch.rand(1, 32, 4, 16) + k = torch.rand(1, 8, 4, 16) + v = torch.rand(1, 8, 4, 16) + + gm = self._assert_rewritten(GQACausalPrefillModule(), (q, k, v)) + expected = GQACausalPrefillModule()(q, k, v) + actual = gm(q, k, v) + torch.testing.assert_close(expected, actual) + + def test_broadcast_mqa_without_enable_gqa_is_rewritten(self): + class BroadcastMQAModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + ) + + q = torch.rand(1, 4, 3, 8) + k = torch.rand(1, 1, 3, 8) + v = torch.rand(1, 1, 3, 8) + + gm = self._assert_rewritten(BroadcastMQAModule(), (q, k, v)) + expected = BroadcastMQAModule()(q, k, v) + actual = gm(q, k, v) + self.assertTrue(torch.allclose(expected, actual)) + + def test_value_embedding_dim_mismatch_is_not_supported(self): + class ValueDimMismatchModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(2, 4, 3, 8) + k = torch.rand(2, 4, 3, 8) + v = torch.rand(2, 4, 3, 6) + + self._assert_not_rewritten(ValueDimMismatchModule(), (q, k, v)) + + def test_batch_broadcast_is_not_supported(self): + class BatchBroadcastModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(2, 4, 3, 8) + k = torch.rand(1, 4, 3, 8) + v = torch.rand(1, 4, 3, 8) + + self._assert_not_rewritten(BatchBroadcastModule(), (q, k, v)) + + def test_head_mismatch_without_enable_gqa_is_not_supported(self): + class HeadMismatchModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(1, 1, 3, 4) + k = torch.rand(1, 4, 3, 4) + v = torch.rand(1, 4, 3, 4) + self._assert_not_rewritten(HeadMismatchModule(), (q, k, v)) + + def test_unmasked_non_causal_unequal_seq_len_is_not_supported(self): + class NonCausalUnequalSeqLenModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(1, 4, 3, 8) + k = torch.rand(1, 4, 5, 8) + v = torch.rand(1, 4, 5, 8) + + self._assert_not_rewritten(NonCausalUnequalSeqLenModule(), (q, k, v)) + + def test_unmasked_causal_query_longer_than_kv_is_not_supported(self): + class CausalQueryLongerThanKVModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + ) + + q = torch.rand(1, 4, 5, 8) + k = torch.rand(1, 4, 3, 8) + v = torch.rand(1, 4, 3, 8) + + self._assert_not_rewritten(CausalQueryLongerThanKVModule(), (q, k, v)) + + def test_assume_causal_mask_rewrites_unmasked_prefix_attention(self): + class PrefixAttentionModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(1, 4, 3, 8) + k = torch.rand(1, 4, 5, 8) + v = torch.rand(1, 4, 5, 8) + + gm = self._assert_rewritten( + PrefixAttentionModule(), + (q, k, v), + assume_causal_mask=True, + ) + expected = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + ) + actual = gm(q, k, v) + self.assertTrue(torch.allclose(expected, actual)) + + def test_dynamic_unmasked_non_causal_seq_lens_are_not_rewritten(self): + class DynamicNonCausalModule(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + + q = torch.rand(1, 4, 4, 8) + k = torch.rand(1, 4, 4, 8) + v = torch.rand(1, 4, 4, 8) + q_seq = torch.export.Dim("q_seq", min=1, max=8) + kv_seq = torch.export.Dim("kv_seq", min=1, max=8) + gm = export( + DynamicNonCausalModule(), + (q, k, v), + dynamic_shapes=({2: q_seq}, {2: kv_seq}, {2: kv_seq}), + strict=True, + ).module() + + custom_sdpa_key = "torch.ops.llama.custom_sdpa.default" + gm = ReplaceSDPAWithCustomSDPAPass()(gm).graph_module + FileCheck().check_count(custom_sdpa_key, 0, exactly=True).run(gm.code) + + +class CustomSDPASchemaTest(unittest.TestCase): + def test_custom_sdpa_fake_rejects_mismatched_kv_seq_len(self): + q = torch.rand(1, 3, 4, 8) + k = torch.rand(1, 5, 4, 8) + v = torch.rand(1, 4, 4, 8) + + with self.assertRaisesRegex(AssertionError, "sequence lengths"): + custom_sdpa_fake(q, k, v, 0, None, 0.0, False, None) + + def test_native_unavailable_is_an_import_error(self): + self.assertTrue(issubclass(CustomSDPANativeOpUnavailable, ImportError))