[bdemirb]

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.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/25032

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b2dd99f with merge base d3369c4 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[bdemirb]

@pytorchbot label "partner: arm"

[bdemirb]

@pytorchbot label "release notes: arm"

[shoumikhin]

@claude review effort=max

[bdemirb]

The previous Claude review failed during repository setup, before Claude could start. Retrying the review.

@claude review effort=max

[claude]

Claude finished @bdemirb's task in 42m 1s —— View job

---

Request Changes: the new Python fallback for llama::custom_sdpa can abort the process if the compiled library loads after it, silently hides broken installs and drops the op's export-time validation, two new pass guards regress broadcast MQA and miss the assume_causal_mask path, and no test exercises the real fallback path.

Detailed review

PR Review: #25032

Summary

The PR adds a pure-Python schema, fake and CompositeExplicitAutograd fallback for llama::custom_sdpa. ReplaceSDPAWithCustomSDPAPass uses it when the compiled custom-ops library cannot be loaded, and the pass gains head-layout, shape and unmasked-attention guards.

Verdict: request changes.

• The fallback's registration can collide with the compiled op.
• It hides broken installs and does not follow the kernel's contract.
• Two of the new guards regress or miss cases.
• The tests do not cover the new paths.

All 8 changed files were reviewed against merge base 8c77fe1; the base ref could not be fetched in this job. Nothing was executed, because running Python or pytest was not permitted here. The runtime outcomes below are derived from reading the code.

Code Quality

• The Python schema can collide with the compiled op and abort the process.
  • Loading custom_sdpa_fallback.py:14,115 defines llama::custom_sdpa (via custom_sdpa_schema.py:105-113) and registers an impl as a side effect.
  • If custom_ops_aot_lib is loaded later in the same process, its static TORCH_LIBRARY_FRAGMENT(llama, m) calls def on custom_sdpa a second time (op_sdpa_aot.cpp:470,479-482).
  • The dispatcher rejects duplicate defs, and an exception that escapes a static initializer during dlopen ends in std::terminate. This is the expected outcome; I did not reproduce it.
  • Trigger: run pytest extension/llm/custom_ops/test_custom_sdpa_fallback.py extension/llm/export/test_export_passes.py with the library built.
    • The first file loads the fallback module before anything loads the library.
    • The pass then loads custom_ops. Its probe checks sdpa_with_kv_cache, fast_hadamard_transform and quantized_moe_ffn, not custom_sdpa (custom_ops.py:30-36), so it reaches torch.ops.load_library (:53).
  • A full-directory run hides this, because test_sdpa_with_kv_cache.py:14 and its siblings load the library at collection.
  • The Python schema string is also a second copy of the C++ one, and the two can drift.
  • Fix: make the module side-effect free. Define and register the fallback only from one explicit call made after loading the compiled module fails. Fix this →
• The catch-all fallback hides broken installs.
  • custom_sdpa_schema.py:131-135 silently falls back on (AssertionError, ImportError, OSError, RuntimeError), and nothing is logged. That covers:
    • assert len(libs) == 1 when two copies are on the path, the case described at custom_ops.py:41-43,51;
    • ABI and dlopen OSErrors;
    • portable_lib failures (custom_ops.py:39).
  • Before this PR, the lazy load in call_operator failed loudly.
  • Now export continues on a fallback that has no custom_sdpa.out (C++ defines it at op_sdpa_aot.cpp:483-486). to_executorch() then fails much later with Missing out variants (exir/passes/__init__.py:462).
  • Fix: keep the hard dependency by default and make the fallback an explicit opt-in, or at least log a warning that includes the original exception. Fix this →
• Validation was lost, and the fallback does not follow the kernel's contract.
  • The removed custom_sdpa Meta impl ran _validate_params: 4D inputs, supported and matching dtypes, equal K/V sizes, and a 2D f32/f16 mask.
  • The new fake (custom_sdpa_schema.py:20-35,53-68) checks only batch, head dim and KV-head equality. It does not check Hq % Hkv == 0.
  • The asserts this PR adds to _validate_params (custom_ops.py:139-151) only reach sdpa_with_kv_cache, whose Meta impl is the function's sole caller (custom_ops.py:184).
  • The fallback (custom_sdpa_fallback.py:94-102) also accepts any mask and applies drpout_p. The kernel requires a 2D Float mask of exactly [q_len, kv_len] (op_sdpa.cpp:77-84, op_sdpa_impl.h:836-848) and ignores dropout (op_sdpa_impl.h:775).
  • These inputs pass the export-time checks on the fallback but fail or differ on device:
    • a [1, 1, 1, S] padding mask with L > 1, which export_passes.py:167-176 squeezes to [1, S] and F.sdpa then broadcasts;
    • a bf16 mask in a bf16 model, which the old Meta impl rejected;
    • dropout_p > 0.
  • Fix: restore these checks in the shared fake. Make the fallback ignore dropout and assert the kernel's mask contract. Fix this →
• Dead and trivial code.
  • custom_ops.custom_sdpa (custom_ops.py:65-85) has no callers. It reuses the name of the removed Meta impl, and its ensure_custom_sdpa_registered() call just imports custom_ops again.
  • custom_sdpa_output_shape (custom_sdpa_schema.py:38-40) is a single-use two-liner, which CLAUDE.md asks to avoid.
  • Fix this →
• The shape helper is duplicated, with test-only branches.
  • export_passes.py:222-225,241-244,259-262 define the same nested _heads/_shape helper three times.
  • Each copy has a hasattr(x, "node") branch that exists only so test_export_passes.py:289-300 can pass eager tensors.
  • call_operator already reads node.meta["val"] directly (:129-133).
  • Fix this →
• The DEF → FRAGMENT change in op_fallback.py:13 is unexplained.
  • This Qualcomm-owned file is loaded by exir/passes/spec_prop_pass.py:23-29 and extension/llm/custom_ops/model_sharding.py:10.
  • No other Library("llama", "DEF") exists, and neither the C++ TORCH_LIBRARY_FRAGMENT nor the new Python FRAGMENT library conflicts with a DEF.
  • Revert it, or document the conflict it fixes. Fix this →
• (Low confidence) register_fake is used on a C++-defined op.
  • custom_sdpa_schema.py:80 departs from the @impl(custom_ops_lib, ..., "Meta") pattern used for every other op in custom_ops.py (:170,201,276,357,375,452,475,570).
  • For a C++-defined op, torch.library.register_fake wraps the fake in a pystub check.
  • That check raises on first use if REQUIRES_SET_PYTHON_MODULE is enabled and the library has no m.set_python_module(...). There is none in this repo.
  • OSS disables the flag. I could not check fbcode.
• (Low confidence) The guards compare symbolic dims.
  • export_passes.py:249-251,271-274 compare meta["val"] sizes with ==/<=.
  • For dynamic seq-len exports, the example-input hints decide those comparisons after export has finished, so the rewrite decision can be specialized to the example.
  • statically_known_true would make the pass skip the rewrite instead.
• Copyright headers.
  • export_passes.py, test_export_passes.py and targets.bzl had no header before and now carry an Arm-only one.
  • custom_ops.py and op_fallback.py keep their original owner's line next to Arm's.
  • Maintainers should confirm the intended header.

Infrastructure

• extension/llm/custom_ops/BUCK has no runtime.python_test target for test_custom_sdpa_fallback.py, while every other test there has one (BUCK:23-134), so internal CI never runs it. Fix this →

Testing

• The fallback tests do not validate the fallback.
  • test_fallback_respects_start_pos_for_causal_decode builds its expected value with the module's own _build_offset_causal_mask. It also uses start_pos + L == S (3 + 2 = 5), so the KV slice does nothing.
  • test_fallback_keeps_full_kv_for_non_causal_unmasked_sdpa (test_custom_sdpa_fallback.py:76) asserts that KV is truncated to L (:96-97), the opposite of its name.
  • No test compares the fallback against the compiled torch.ops.llama.custom_sdpa, the pattern test_sdpa_with_kv_cache.py uses.
  • The registration tests (:111-171) only exercise mocks of __import__ and private flags.
  • No test runs the pass → fallback path. OSS pytest loads the library at collection, and Buck export_passes_test preloads custom_ops_aot_lib (extension/llm/export/BUCK:244-246). Only a subprocess test with the library unavailable can cover it.
  • Fix this →
• test_gqa_decode_shape is degenerate (test_export_passes.py:223-251).
  • With is_causal=True and L=1, aten's top-left causal mask attends only key 0, so both sides reduce to v[:, h // 4, 0, :].
  • Add a masked GQA decode test ([1, 1, 1, 512] float mask, is_causal=False) and a GQA prefill test (L == S > 1). Fix this →

Thread Safety

• The registration flags in custom_sdpa_schema.py are set before the work they guard is done.
  • _CUSTOM_SDPA_SCHEMA_REGISTERED = True is set at :97/:115, but register_custom_sdpa_fake() runs afterwards and outside the lock (:98/:117).
  • Its check-then-set on _CUSTOM_SDPA_FAKE_REGISTERED (:74-81) is unsynchronized.
  • A concurrent caller can therefore either return early and trace before the fake exists, or call register_fake a second time, which raises.
  • _has_custom_sdpa_meta_impl() only becomes true after this module's own register_fake has run, because C++ registers just CompositeExplicitAutograd (op_sdpa_aot.cpp:530,537).
  • :137 is redundant with custom_ops.py:62.
  • Fix: register once when the owning module loads; Python's module-loading lock already serializes that. This removes all three flags and both locks. Fix this →

Backward Compatibility

• Broadcast MQA is no longer rewritten.
  • export_passes.py:235-236 rejects every layout where q_heads != k_heads unless enable_gqa=True.
  • aten SDPA accepts Hkv == 1 without enable_gqa by broadcasting. The pass rewrote that case before this PR, and the kernel handles it via num_reps (op_sdpa_impl.h:825). Such models now silently stay on aten SDPA.
  • The new test (test_export_passes.py:289-300) calls the static method on eager Hq=4, Hkv=2 tensors. That layout fails at export without enable_gqa, so the test covers a graph that cannot exist and misses the one that regressed.
  • Fix: accept k_heads == 1 without enable_gqa, and test it on an exported graph. Fix this →
• The guards run before the assume_causal_mask override. _supports_unmasked_attention_contract (export_passes.py:150) checks the original mask/is_causal, which :162-165 then replaces.
  • Regression: with assume_causal_mask=True, an unmasked non-causal call with L < S was rewritten before this PR, to top-left causal, which is what the flag asserts. It is now skipped.
  • Gap: an explicit mask with L > S passes the guard. It is then rewritten to unmasked causal with start_pos=0 and fails num_keys_for_causal_attention <= kvSize on device (op_sdpa_impl.h:806-810). That is exactly the case the guard was added to prevent.
  • Fix: apply the override before the guards. Fix this →

Performance

• custom_sdpa_fallback.py:82-89 calls repeat_interleave on the whole KV cache before slicing it to start_pos + L. Every eager call therefore materializes Hq / Hkv copies of the full cache.
• Fix: slice first, or pass enable_gqa=True to F.scaled_dot_product_attention. Fix this →

Recommendation

Request Changes

These block approval:

• the process abort from the duplicate def when the fallback and the compiled library meet in one process;
• the silent catch-all fallback;
• the lost validation, and the fallback not following the kernel's contract;
• the two guard problems: the MQA regression and the guards running before the assume_causal_mask override;
• no test compares the fallback with the compiled op, and none covers the fallback path.

---

[bdemirb]

@claude review effort=max

[claude]

Claude finished @bdemirb's task in 40m 44s —— View job

---

Request Changes: the automatic Python fallback for llama::custom_sdpa can abort the process if the native library loads later, and its graphs fail at to_executorch(); the PR also breaks callers that catch ImportError and drops the check that K and V have the same length from the op's fake.

Detailed review

PR Review: #25032

Summary

This PR adds semantic guards to ReplaceSDPAWithCustomSDPAPass, moves the llama::custom_sdpa fake into a new custom_sdpa_schema.py, and adds a Python schema and implementation that the pass registers when the native custom-ops library is unavailable.

The earlier points about import-time registration, error filtering, mask checks and dropout are resolved. The fallback design is still the main problem, along with the issues below.

How this review was done:

• origin/main could not be fetched, so the diff was taken against merge base 3e85cae.
• Python and pytest could not run in this job, so the findings come from reading the code and none was reproduced.

Code Quality

• The automatic fallback (custom_sdpa_schema.py:194-203) defines llama::custom_sdpa in Python (171-176).
  • Process abort if the native library loads afterwards.
    • op_sdpa_aot.cpp:479-482 defines the same op with def inside a TORCH_LIBRARY_FRAGMENT static initializer. c10 rejects a duplicate def, and a C++ exception thrown from a static initializer during dlopen ends the process.
    • Trigger, with the native library built: pytest extension/llm/custom_ops/test_custom_sdpa_fallback.py extension/llm/export/test_export_passes.py.
    • Sequence: test_fallback_impl_is_not_registered_when_compiled_impl_exists reaches ensure_custom_sdpa_schema() through custom_sdpa_fallback.py:114 and defines the Python schema. Each tearDown restores _CUSTOM_SDPA_IMPLEMENTATION_REGISTERED to False, so the first pass test imports custom_ops, and torch.ops.load_library runs the C++ def.
    • The full suite hides this because test_gated_delta.py and other modules import custom_ops at collection time.
  • Broken installs fall back instead of failing. custom_ops.py:63-68 turns every ImportError from portable_lib into CustomSDPANativeOpUnavailable, including an undefined symbol or ABI mismatch from a stale build. The pass then logs a warning and continues.
  • Fallback graphs can't be serialized. The Python schema has no custom_sdpa.out. Any graph where the op isn't delegated fails much later in ToOutVarPass with Missing out variants (exir/passes/__init__.py:461-462).
  • The only in-tree caller of the pass is test_export_passes.py. The two locks, three flags and two new Library objects mostly let those tests run without the native library. Keeping the hard dependency and skipping those tests when the op can't load is simpler. If a Python reference is wanted, test code can register it explicitly. Fix this →
• The new fake no longer checks that K and V have the same sequence length.
  • The old custom_sdpa Meta asserted key_cache.size() == value_cache.size(). That check is still at custom_ops.py:169-171, but now only for sdpa_with_kv_cache.
  • _validate_custom_sdpa_output_shape_inputs (custom_sdpa_schema.py:26-59) never compares shape[-3], and the mask check at line 90 uses K's length.
  • The native kernel takes kvSize from V (op_sdpa_impl.h:796,803), and validate_flash_attention_args (op_sdpa.cpp:45-106) doesn't compare the two lengths.
  • So mismatched K/V now export, then fail the mask check at runtime. If K is the shorter one, the kernel probably reads past the end of K (suspected).
  • The head and batch asserts the PR adds to _validate_params (custom_ops.py:172-184) overlap with this validator and only reach sdpa_with_kv_cache. One shared validator would keep them in sync. Fix this →
• Contract gaps in the pass:
  • dropout_p is still forwarded (export_passes.py:199) although the kernel ignores it, so SDPA with dropout_p > 0 is rewritten with different semantics.
  • Dtypes outside f32/bf16/f16 pass the guards and then hit the fake's assertion (custom_sdpa_schema.py:31-35), instead of being left as aten SDPA. Fix this →
• Dynamic shapes in the guards (low confidence):
  • The guards compare SymInts with plain operators (export_passes.py:254-257, 276-279, 297-299). For dynamic dims, bool() decides using the example hints.
  • Example: a non-causal unmasked SDPA whose q and k lengths are different symbols with equal hints would be rewritten. At runtime the kernel would then attend only to the first q_len keys.
  • statically_known_true would rewrite only when the contract provably holds. Fix this →
• Duplication and dead code:
  • _has_native_custom_sdpa_kernel (custom_ops.py:32-39) and _has_custom_sdpa_composite_impl (custom_sdpa_fallback.py:18-25) are identical. _has_custom_sdpa_meta_impl (custom_sdpa_schema.py:97-104) differs only in the dispatch key.
  • export_passes.py has three nested _heads/_shape helpers (227-230, 246-249, 264-267). Their hasattr(x, "node") branch exists only for an eager-tensor test, while _supports_custom_mask_contract reads .node.meta directly.
  • The else branch at export_passes.py:182-186 can't be reached after the check at 289-292.
  • The new public custom_ops.custom_sdpa wrapper (custom_ops.py:98-118) has no callers.
  • custom_sdpa_output_shape (custom_sdpa_schema.py:62-64) is a single-use helper that repeats the validation custom_sdpa_fake just ran.
  • Line 77 (_ = (...)) does nothing useful.
  • Lines 204-205 repeat what custom_ops.py:95 already does.
  • ensure_custom_sdpa_schema() runs twice on the fallback path (197 and custom_sdpa_fallback.py:114). Fix this →
• Unrelated changes:
  • op_fallback.py:13 switches Library("llama", "DEF") to "FRAGMENT", but nothing in-tree creates another llama DEF library. Drop it, or explain it in the description.
  • Arm-only copyright headers were added to Meta-authored files that had none (export_passes.py, test_export_passes.py, targets.bzl). Maintainers should confirm that is the intended header.

Infrastructure

• extension/llm/custom_ops/BUCK has no python_test for test_custom_sdpa_fallback.py, so the new tests don't run in Buck CI. export_passes_test preloads custom_ops_aot_lib (extension/llm/export/BUCK:244-246), so the fallback isn't exercised there either. Fix this →

Testing

• test_gqa_decode_shape (test_export_passes.py:248-276) can't catch errors in Q·K or in the K-head mapping.
  • With is_causal=True and q_len=1, both aten and the kernel attend only to key 0, so the softmax weight is 1 whatever Q and K hold.
  • Add GQA cases that attend to many keys: a masked decode over the 512-key cache and a causal prefill with q_len == kv_len. Reuse _assert_rewritten instead of copying it. Fix this →
• The fallback tests don't check the fallback independently:
  • test_fallback_respects_start_pos_for_causal_decode (test_custom_sdpa_fallback.py:19-55) builds its expected mask with _build_offset_causal_mask, the helper under test. It also uses start_pos + q_len == kv_len (3 + 2 = 5), so the prefix slice changes nothing.
  • Nothing compares the fallback with the native op for masks, scale, GQA, non-causal start_pos > 0 or bf16/fp16.
  • Nothing checks that the schema string (custom_sdpa_schema.py:173-175) matches op_sdpa_aot.cpp:479-482.
  • No test runs the pass end to end through the fallback.
  • test_native_load_errors_are_not_silently_swallowed (189-202) covers OSError, not the ImportError conversion that actually hides errors.
  • test_fallback_registration_is_explicit (174-187) patches impl only after importlib.reload, so it can't detect registration at import. The reload also drops the old custom_sdpa_fallback_lib, which may deregister a fallback kernel registered earlier in the same process (suspected). Fix this →

Thread Safety

• ensure_custom_sdpa_schema sets _CUSTOM_SDPA_SCHEMA_REGISTERED (custom_sdpa_schema.py:160, 178) before register_custom_sdpa_fake() runs, and that call happens outside the lock (161, 180).
  • A second thread can return at 152-153 and start tracing before the fake exists.
  • Two threads can both pass the unlocked check in register_custom_sdpa_fake (137-143). The second torch.library.register_fake call then raises.
  • Removing the fallback and its flags removes this problem too. Fix this →

Backward Compatibility

• CustomSDPANativeOpUnavailable subclasses RuntimeError (custom_sdpa_schema.py:14).
  • Importing custom_ops without portable_lib used to raise ImportError.
  • Callers that catch only ImportError now error instead of skipping: examples/models/phonon/tests/test_phonon.py:290-296, 322-328 and examples/models/llama/source_transformation/test_quantized_sdpa.py:123-128. Out-of-tree callers that catch ImportError break the same way.
  • Subclassing ImportError keeps them working. Fix this →
• Broadcast MQA is no longer rewritten.
  • _supports_head_layout (export_passes.py:236-242) rejects q_heads != k_heads without enable_gqa, including a single KV head.
  • aten broadcasts a single KV head, and the old pass rewrote it correctly because the kernel maps it via num_reps. These graphs now silently stay as aten SDPA. export_lib is visible to targets outside ExecuTorch (extension/llm/export/BUCK:33-39), so they are affected too.
  • test_head_mismatch_without_enable_gqa_is_not_supported (test_export_passes.py:314-326) calls the static method on eager tensors with Hq=4 and Hkv=2, a layout aten can't run. It doesn't cover this case. Fix this →

Performance

• The fallback runs repeat_interleave on the whole KV cache (custom_sdpa_fallback.py:89-93) before slicing the start_pos + q_len prefix (96). Each eager decode step copies the full cache num_reps times. Fix this →

Recommendation

Request Changes

These block approval:

• the automatic fallback (process abort, hidden broken installs, graphs that can't be serialized);
• the switch from ImportError to RuntimeError;
• the dropped K/V length check in the fake.

The other items are smaller and can be fixed in the same pass.

[bdemirb]

@claude review effort=max

[claude]

Claude finished @bdemirb's task in 37m 40s —— View job

---

Request Changes: the new rewrite guards compare symbolic shapes with plain operators, so dynamic-shape exports are decided from the example sizes. custom_sdpa is also the only op that moves to a lock-and-flag register_fake path instead of the file's @impl(custom_ops_lib, ..., "Meta") pattern, and the GQA decode test only ever attends key 0.

Detailed review

PR Review: #25032

Summary

This PR adds eligibility guards to ReplaceSDPAWithCustomSDPAPass, moves the llama::custom_sdpa fake into a new custom_sdpa_schema.py that registers it with register_fake, and replaces the assertion in native-op loading with typed errors. Three problems remain:

• The guards are unsound under dynamic shapes.
• The registration path is much heavier than the pattern custom_ops.py already uses.
• The new GQA test never exercises the attention math the guards are meant to protect.

The diff is against origin/main (merge base a532d10), and all 6 changed files were reviewed. Python and pytest could not run in this job, so every finding comes from reading the code. The dynamic-shape finding was traced through the code but not reproduced.

Code Quality

• Guards use example sizes under dynamic shapes (major). These guards compare meta["val"] sizes with plain ==, !=, <=, % and tuple equality:

  • _supports_head_layout (export_passes.py:257-265)
  • _supports_native_shape_contract (277-280)
  • _supports_unmasked_attention_contract (299-302)
  • _supports_custom_mask_contract (322)

  For backed SymInts, bool() evaluates on the example hints and records a guard, and nothing enforces that guard after export() returns. This causes two failures:

  • Unmasked, non-causal SDPA with independent Dims for the q and kv sequence lengths and equal example sizes gets rewritten. At runtime with q_len < kv_len, the kernel attends only the first start_pos + seq_len keys (op_sdpa.cpp:506-507, op_sdpa_impl.h:806-811), so the output is silently wrong.
  • In the causal case, a runtime q_len > kv_len trips ET_CHECK_MSG(num_keys_for_causal_attention <= kvSize).

  For unbacked sizes, the same bool() raises GuardOnDataDependentSymNode and the export fails, when the node should simply stay on aten. Use torch.fx.experimental.symbolic_shapes.statically_known_true(...) and add a dynamic-shape test.
  Fix this →

• custom_sdpa registration breaks the file's pattern and adds unneeded machinery (major). Every other C++-defined op in custom_ops.py registers its fake with @impl(custom_ops_lib, "<op>", "Meta") (custom_ops.py:201, 232, …), and custom_sdpa did too on main. This PR replaces that with:

  • a module lock and a global flag (custom_sdpa_schema.py:15-16)
  • a probe for an existing Meta kernel (89-96)
  • a bool return value that nothing reads (122-134)
  • ensure_custom_sdpa_registered (137-140), whose register_custom_sdpa_fake() call always returns early because importing custom_ops already ran it (custom_ops.py:94)

  The import lock already serializes module-level registration. There are smaller leftovers too:

  • _ = (start_pos, drpout_p, scale) (custom_sdpa_schema.py:69) only exists to silence a lint.
  • The module is named custom_sdpa_schema but holds no schema.
  • _has_native_custom_sdpa_kernel (custom_ops.py:31-38) ties the Python import to the C++ dispatch key. Adding torch.ops.llama.custom_sdpa.default to the tuple in _native_custom_ops_are_registered is enough.

  Fix this →

• New _validate_params asserts only reach sdpa_with_kv_cache (minor). The custom_sdpa fake no longer calls _validate_params. So the asserts added at custom_ops.py:170-182 (batch, head dim, KV heads > 0, Hq % Hkv) only affect sdpa_with_kv_cache_meta (custom_ops.py:215). That is an untested behavior change the PR description doesn't mention.

  • The SDPA contract now lives in three places.
  • The Q/K/V dtype tuple is repeated at export_passes.py:201-205, custom_sdpa_schema.py:24 and custom_ops.py:142.

  Fix this →

• Production branches that only an eager-tensor test needs (minor). The hasattr(x, "node") fallbacks (export_passes.py:248-251, 269-272, 287-290) exist only because test_head_mismatch_without_enable_gqa_is_not_supported (test_export_passes.py:384-395) calls a private static method with eager tensors.

  • aten cannot trace that test's layout (4 query heads, 2 KV heads, no enable_gqa).
  • _supports_custom_mask_contract already reads mask.node.meta directly.
  • An exportable case reaches the same branch, for example 1 query head with 4 KV heads.

  Fix this →

• Unreachable branch (minor). _supports_custom_mask_contract already rejects masks whose leading dims are not 1 (export_passes.py:312-315), so the else at 150-154 can never run. If it were reached, it would return after the q/k/v transposes at 136-138 had already been emitted.
  Fix this →

• fp16/bf16 masks are rejected instead of cast (minor). export_passes.py:317-318 accepts only bool and float32 masks, so fp16/bf16 models whose mask has the same dtype lose the rewrite. A lossless cast to float32, placed next to the bool→float conversion at 157-163, would satisfy the kernel's Float-mask check (op_sdpa.cpp:82-83). test_float16_mask_is_not_rewritten (test_export_passes.py:249) locks the rejection in.
  Fix this →

• Unexplained namespace change in a module that core exir imports (question). op_fallback.py:13 switches Library("llama", "DEF") to "FRAGMENT". A single DEF conflicts with neither register_fake (which uses a FRAGMENT library) nor the C++ TORCH_LIBRARY_FRAGMENT (op_sdpa_aot.cpp:470). exir/passes/spec_prop_pass.py:25 imports this module. Which conflict does the change fix? If none, please split it out.
  Fix this →

• Arm-only copyright headers on files that had none (question for maintainers). export_passes.py:1-4, test_export_passes.py:1-4 and custom_ops/targets.bzl:1-4 had no header on main and now carry only an Arm notice. In the same directory, builder.py:1-6 keeps the Meta notice and adds the Arm line after it.
  Fix this →

Testing

• test_gqa_decode_shape never exercises the attention math (major). The test (test_export_passes.py:297-325) uses is_causal=True with q_len 1 and start_pos=0. Under those settings, aten (top-left causal) and custom_sdpa (kvSize = 0 + 1) both attend only key 0. Softmax over a single logit is 1, so the output is just v[..., 0, :]. The test therefore checks the V-head mapping and none of Q·K, scale, softmax, or the K-head mapping. Missing tests:

  • a masked GQA decode with a [1, 512] float mask containing -inf, plus a bool variant
  • a GQA causal prefill with q_len == kv_len
  • positive fp16/bf16 rewrites
  • a not-rewritten case for enable_gqa=True with mismatched K/V heads (8/2/4)

  Fix this →

• Negative tests don't isolate the guard they're named after (minor).

  • test_value_embedding_dim_mismatch_is_not_supported (test_export_passes.py:348) and test_batch_broadcast_is_not_supported (366) use unmasked, non-causal attention with q_len 3 and kv_len 5. _supports_unmasked_attention_contract rejects that, so both tests still pass with _supports_native_shape_contract deleted.
  • _assert_not_rewritten (213-218) only checks that there are zero custom_sdpa nodes. It doesn't check that aten SDPA survives or that the outputs match.
  • test_gqa_decode_shape re-implements _assert_rewritten inline (314-321).

  Fix this →

• A missing native op now becomes a skip (minor).

  • setUpClass (test_export_passes.py:185-193) skips the whole class, including the four tests that predate this PR.
  • On main, a missing custom_ops_aot_lib raised AssertionError. It now raises an ImportError subclass (custom_ops.py:78-81). The existing except ImportError skips in examples/models/llama/source_transformation/test_quantized_sdpa.py:123-127 and examples/models/phonon/tests/test_phonon.py:353-359/385-391 will now hide a broken build instead of failing.
  • test_native_unavailable_is_an_import_error (475-476) only pins the class hierarchy.
  • Buck preloads the library (extension/llm/export/BUCK:244-246), so the skip can only trigger in OSS.

  Fix this →

Recommendation

Request Changes. Three things block approval:

• The guards must be sound for dynamic and unbacked shapes.
• custom_sdpa should go back to the file's @impl(custom_ops_lib, ..., "Meta") registration, without the lock, flag and probe.
• GQA needs tests that exercise the attention math (masked decode and causal prefill).

The other findings are minor.

• Read PR context and diffed against origin/main
• Ran /pr-review with three independent reviewers
• Consolidated and fact-checked findings against the code
• Posted review

[bdemirb]

@claude review effort=max

[claude]

Claude finished @bdemirb's task in 31m 38s —— View job

---

Request Changes: the class-wide SkipTest from the previous review is still there, and combined with the new ImportError subclass it turns a missing custom_ops_aot_lib into skipped tests in three files; everything else is minor.

Detailed review

PR Review: #25032

Summary

The PR guards ReplaceSDPAWithCustomSDPAPass so it only rewrites SDPA that provably fits the native llama.custom_sdpa contract. It also moves the custom_sdpa fake into a new custom_sdpa_schema.py and makes the native-library loader in custom_ops.py raise typed errors.

One item blocks approval: the test skip under Testing. The rest are minor: cleanups (some still open from earlier reviews), coverage gaps, and one low-confidence backward-compatibility question. All five changed files were reviewed by reading the code only, because this job could not run Python or pytest.

Code Quality

• The custom SDPA contract is still written out in three places (minor, still open from earlier reviews).
  • The (float32, bfloat16, float16) tuple appears at custom_sdpa_schema.py:20, export_passes.py:200-204 and custom_ops.py:130.
  • The asserts added to _validate_params (custom_ops.py:158-170) repeat custom_sdpa_schema.py:29-52, but custom_sdpa no longer goes through _validate_params.
  • Their only caller is sdpa_with_kv_cache_meta (custom_ops.py:203), an op this PR doesn't otherwise change or test.
  • Define the tuple once and drop those asserts, or land them separately with sdpa_with_kv_cache tests. Fix this →
• Redundant and unreachable guards in the pass (minor; the else is still open from earlier reviews).
  • In the aten layout, export_passes.py:287 repeats the K/V head check at :262.
  • :286 already requires equal K/V sequence lengths, so the v_seq_len comparisons at :303-309 add nothing.
  • _supports_custom_mask_contract (:319-322) already rejects masks whose leading dims aren't statically 1. So the all(...) at :142 is always true, and the else at :149-153 can never run.
  • Fix this →

Testing

• The class-wide skip hides a missing native library (major, still open from the previous review).
  • setUpClass (test_export_passes.py:184-192) turns CustomSDPANativeOpUnavailable into SkipTest for the whole ReplaceSDPAWithCustomSDPAPassTest class. That includes test_causal_mask, test_explicit_causal_mask, test_custom_mask and test_squeezable_mask, which run unconditionally on main.
  • Buck preloads custom_ops_aot_lib (extension/llm/export/BUCK:244-246). OSS pytest collects extension/ (pytest.ini:126), and those four tests already need the library there. So the skip only fires on a broken build, and it reports that build as green.
  • On main, a missing library raised AssertionError; it now raises an ImportError subclass (custom_ops.py:66-69). That also flips the existing except ImportError checks from failing to skipping, in examples/models/llama/source_transformation/test_quantized_sdpa.py:123-128 and examples/models/phonon/tests/test_phonon.py:353-359 / 385-391.
  • Fix: drop the skip and raise RuntimeError for the missing library, as :71 and :76 already do. Keep the missing-portable_lib case a plain ImportError as on main, since that is the case the earlier automated suggestion to subclass ImportError was meant to protect.
  • CustomSDPANativeOpUnavailable and test_native_unavailable_is_an_import_error (:527-528) can then be deleted.
  • After that, custom_sdpa_schema.py and its targets.bzl:164 entry would exist only so CustomSDPASchemaTest can import the fake without the library. custom_sdpa_fake could then live next to the other Meta impls in custom_ops.py. Fix this →
• Coverage gaps in the new guards (minor).
  • The only dynamic-shape test (test_export_passes.py:489-515) is negative. Nothing checks that the statically_known_true guards still allow the rewrite in the common LLM export case: causal SDPA with one shared sequence Dim. A regression there would silently disable the rewrite.
  • CustomSDPASchemaTest (:518-525) covers only a K/V sequence mismatch. Direct torch.ops.llama.custom_sdpa callers now hit the other new fake asserts at export, and none of them are tested: mask with is_causal, mask rank/dtype/shape (custom_sdpa_schema.py:67-82), and head divisibility (:49). Fix this →

Backward Compatibility

• The fake now rejects float16 masks (low confidence: no CI flow I found hits this).
  • custom_sdpa_schema.py:74-77 accepts only float32. The removed Meta impl also accepted float16, through _validate_params (custom_ops.py:184-186).
  • The new check matches the portable CPU kernel (op_sdpa.cpp:81-84). But the native ggml engine also lowers llama.custom_sdpa, and it casts any mask to F16 (backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp:136).
  • SDPACustom casts fp16 Q/K/V to float32 but passes the mask through unchanged (examples/models/llama/source_transformation/sdpa.py:46-61). So an fp16 masked export that ran on that backend would now fail at export.
  • Consider keeping float16 in the fake's mask check. The CPU kernel still rejects it at runtime, as on main. Fix this →

Recommendation

Request Changes

Remove the class-wide skip and the ImportError change behind it before merging. Together they take failure signal away from existing tests in three files. Everything else is minor.

• Read the PR context and diff it against the merge base d3369c4 (the checkout is shallow)
• Run the /pr-review skill with three independent reviewers (all 5 changed files covered)
• Check whether findings from earlier reviews were addressed
• Consolidate and fact-check findings against the code
• Post the review