fix(npu): avoid false gather index bounds failures on 910B, fix ACLNN executor lifecycle

src/candle/_backends/npu/aclnn.py

@@ -4287,7 +4287,14 @@ def _destroy_deferred_executor(executor):
     handle = _executor_handle(executor)
     if handle == 0:
         return
-    _run_deferred_executor_cleanup(handle)
+    cleanup = _DEFERRED_EXECUTOR_CLEANUP.pop(handle, None)
+    if _ffi is None or not _ffi.is_initialized():
+        _apply_deferred_cleanup(cleanup)
+        return
+    try:
+        _ffi.destroy_executor(handle)
+    finally:
+        _apply_deferred_cleanup(cleanup)
 
 
 def _cleanup_aclnn():

src/candle/_backends/npu/ops_soc.py

@@ -112,6 +112,7 @@
     "910b": {
         "use_smallop_arange_1d": False,
         "use_smallop_linspace": True,
+        "use_safe_int64_index_compare": True,
     },
     "310b": {
         "use_smallop_arange_1d": True,

src/candle/_cython/_aclnn_ffi.pyx

@@ -1020,6 +1020,8 @@ def tensor_scalar_op_with_alpha(
                         NULL, 0, executor, <void*>stream)
                 if ret != 0:
                     raise RuntimeError(f"Execute failed: {ret}")
+                _release_executor_cleanup(<uintptr_t>executor)
+                executor = NULL
             except Exception:
                 destroy_executor(<uintptr_t>executor)
                 executor = NULL
@@ -1099,6 +1101,8 @@ def tensor_scalar_op_no_alpha(
                         NULL, 0, executor, <void*>stream)
                 if ret != 0:
                     raise RuntimeError(f"Execute failed: {ret}")
+                _release_executor_cleanup(<uintptr_t>executor)
+                executor = NULL
             except Exception:
                 destroy_executor(<uintptr_t>executor)
                 executor = NULL
@@ -1174,6 +1178,8 @@ def unary_op(
                     ret = (<aclnnExec_t>exec_ptr)(NULL, 0, executor, <void*>stream)
                 if ret != 0:
                     raise RuntimeError(f"Execute failed: {ret}")
+                _release_executor_cleanup(<uintptr_t>executor)
+                executor = NULL
             except Exception:
                 destroy_executor(<uintptr_t>executor)
                 executor = NULL
@@ -1328,6 +1334,8 @@ def reduce_sum_op(
                     ret = (<aclnnExec_t>exec_ptr)(NULL, 0, executor, <void*>stream)
                 if ret != 0:
                     raise RuntimeError(f"Execute failed: {ret}")
+                _release_executor_cleanup(<uintptr_t>executor)
+                executor = NULL
             except Exception:
                 destroy_executor(<uintptr_t>executor)
                 executor = NULL
@@ -1484,6 +1492,8 @@ def arg_reduce_op(
                     ret = (<aclnnExec_t>exec_ptr)(NULL, 0, executor, <void*>stream)
                 if ret != 0:
                     raise RuntimeError(f"Execute failed: {ret}")
+                _release_executor_cleanup(<uintptr_t>executor)
+                executor = NULL
             except Exception:
                 destroy_executor(<uintptr_t>executor)
                 executor = NULL

src/candle/_dispatch/dispatcher.py

@@ -1,4 +1,5 @@
 import inspect
+import functools
 import numpy as np
 
 from .registry import registry
@@ -49,6 +50,7 @@ def _pop_dispatch_context():
         stack.pop()
 
 
+@functools.lru_cache(maxsize=None)
 def _accepts_device(func):
     try:
         sig = inspect.signature(func)

tests/contract/test_npu_no_fallback_contract.py

@@ -6117,7 +6117,31 @@ def test_inplace_unary_op_fast_path_releases_cleanup_and_returns_null_executor(t
     assert state['destroy_executor_after'] == 0
 
 
-def test_flush_deferred_executors_releases_cleanup_without_destroying_executor(monkeypatch):
+def test_unary_op_fast_path_releases_cleanup_and_returns_null_executor(tmp_path):
+    state = _run_compiled_executor_cleanup_contract(
+        tmp_path,
+        lib_stem='fake_aclnn_neg_fastpath',
+        c_source=_tensor_cleanup_contract_source(
+            'aclnnNegGetWorkspaceSize(void* self, void* out, uint64_t* workspace_size, void** executor)',
+            'aclnnNegGetWorkspaceSize',
+            'aclnnNeg',
+            workspace_size=0,
+        ),
+        op_name='Neg',
+        ffi_call='ws_size, executor = ffi.unary_op(getws_ptr, exec_ptr, (3,), (1,), (3,), (1,), 9, 9, 2, 1, 2, 0)',
+        destroy_returned_executor=False,
+    )
+
+    assert state['ws_size'] == 0
+    assert state['executor'] == 0
+    assert state['create'] == 2
+    assert state['destroy_before'] == 2
+    assert state['destroy_after'] == 2
+    assert state['destroy_executor_before'] == 0
+    assert state['destroy_executor_after'] == 0
+
+
+def test_flush_deferred_executors_releases_cleanup_and_destroys_executor(monkeypatch):
     calls = []
 
     class _FakeFfi:
@@ -6141,7 +6165,9 @@ def destroy_tensor(self, handle):
 
     assert ("destroy_scalar", 0x2000) in calls
     assert ("destroy_tensor", 0x3000) in calls
-    assert not [entry for entry in calls if entry[0] == "destroy_executor"]
+    assert ("destroy_executor", 0xBEEF) in calls
+    assert calls.index(("destroy_executor", 0xBEEF)) < calls.index(("destroy_scalar", 0x2000))
+    assert calls.index(("destroy_executor", 0xBEEF)) < calls.index(("destroy_tensor", 0x3000))
     assert aclnn._DEFERRED_EXECUTORS == []
     assert aclnn._DEFERRED_EXECUTOR_CLEANUP == {}
 

tests/npu/test_npu_soc_policy.py

@@ -32,10 +32,10 @@ def test_soc_capability_table_routes_smallop_linspace_for_310b_910a_and_910b():
     assert not ops_soc.use_smallop_linspace(profile="310p")
 
 
-def test_soc_capability_table_routes_safe_int64_index_compare_for_310b_only():
+def test_soc_capability_table_routes_safe_int64_index_compare_for_310b_and_910b():
     assert ops_soc.capability("use_safe_int64_index_compare", profile="310b")
     assert not ops_soc.capability("use_safe_int64_index_compare", profile="910a")
-    assert not ops_soc.capability("use_safe_int64_index_compare", profile="910b")
+    assert ops_soc.capability("use_safe_int64_index_compare", profile="910b")
     assert not ops_soc.capability("use_safe_int64_index_compare", profile="310p")