Harden custom SDPA export path

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,

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)

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"],