From b2dd99f4f453186feb3b42922d7fe36713175921 Mon Sep 17 00:00:00 2001 From: Baris Demir Date: Sat, 10 Oct 2026 23:21:45 +0100 Subject: [PATCH] Harden custom SDPA export path Tighten the LLM custom SDPA export rewrite so it only rewrites GQA decode shapes whose query and KV head layout matches the native custom op contract. Keep unsupported cases on the original SDPA path instead of rewriting them. Add export-pass coverage for a representative decode-time GQA shape, unsupported head layouts, mask constraints, dtype constraints, and dropout handling. Preserve the compatibility Python wrapper for llama.custom_sdpa so downstream imports keep working with the native custom op path. Test: PYTHONPATH=/private/tmp/executorch_patch1_sdpa/src \ PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ /Users/bardem01/devkit7/executorch/.venv-fvp/bin/pytest -q \ extension/llm/export/test_export_passes.py Authored with assistance from OpenAI Codex. Change-Id: I880b2dbd201218467254db2eb5a981f00254623b Signed-off-by: Baris Demir --- extension/llm/custom_ops/custom_ops.py | 124 ++++--- .../llm/custom_ops/custom_sdpa_schema.py | 105 ++++++ extension/llm/custom_ops/targets.bzl | 5 + extension/llm/export/export_passes.py | 180 +++++++++- extension/llm/export/test_export_passes.py | 319 ++++++++++++++++++ 5 files changed, 669 insertions(+), 64 deletions(-) create mode 100644 extension/llm/custom_ops/custom_sdpa_schema.py 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))