fix(build): separate linked operator allowlist

docs/build.md

@@ -27,11 +27,17 @@ entry is `python -m pip install` with CMake options passed through
 | `INFINI_RT_ROOT` | InfiniRT install prefix containing `include/` and `lib/`. | `$INFINI_RT_ROOT` |
 | `INFINI_OPS_SMOKE_BUILD` | Build only the smoke-test operator subset. | `OFF` |
 | `INFINI_OPS_OPS` | Comma- or semicolon-separated operator allowlist. | empty |
+| `INFINI_OPS_LINKED_OPS` | Comma- or semicolon-separated linked implementation allowlist. | `INFINI_OPS_OPS` |
 | `INFINI_OPS_TORCH_OPS` | Comma- or semicolon-separated ATen operator allowlist. | empty |
 
 Only one GPU backend should be enabled in a build. CPU may be enabled with the
 selected accelerator backend.
 
+When `WITH_LINKED=ON`, `INFINI_OPS_LINKED_OPS` can select a narrower linked
+provider subset without removing native or ATen implementations for the same
+operators. If `INFINI_OPS_OPS` is set, every linked entry must also appear in
+that main operator allowlist.
+
 ## Python Wheel Build
 
 Using CPU as the smallest backend:

docs/linked-operators.md

@@ -83,9 +83,15 @@ cmake -S . -B build \
   -DWITH_METAX=ON \
   -DWITH_LINKED=ON \
   -DWITH_TORCH=OFF \
-  -DINFINI_OPS_OPS=silu_and_mul
+  -DINFINI_OPS_OPS=silu_and_mul \
+  -DINFINI_OPS_LINKED_OPS=silu_and_mul
 ```
 
+`INFINI_OPS_LINKED_OPS` defaults to `INFINI_OPS_OPS`. Set it explicitly when
+only a subset of the generated operators should resolve linked providers. This
+does not remove native or ATen implementations for the same operator. When the
+main allowlist is set, the linked allowlist must be its subset.
+
 The `torch` transport uses the installed PyTorch C++ headers and libraries for
 `at::Tensor`, but it does not enable the standard `src/torch` operator backend.
 Provider and PyTorch C++ ABIs must match. Configuration fails before compilation

scripts/generate_wrappers.py

@@ -1751,6 +1751,28 @@ def _filter_ops(ops, op_allowlist, *, strict=False):
     return {op_name: ops[op_name] for op_name in op_allowlist if op_name in ops}
 
 
+def _filter_linked_implementations(ops, linked_op_allowlist):
+    if linked_op_allowlist is None:
+        return ops
+
+    linked_root = (_SRC_DIR / "linked").resolve()
+    allowed = set(linked_op_allowlist)
+    filtered = {}
+
+    for op_name, impl_paths in ops.items():
+        selected = [
+            impl_path
+            for impl_path in impl_paths
+            if op_name in allowed
+            or not pathlib.Path(impl_path).resolve().is_relative_to(linked_root)
+        ]
+
+        if selected:
+            filtered[op_name] = selected
+
+    return filtered
+
+
 def _get_all_ops(devices, with_torch=False, with_ninetoothed=False, with_linked=False):
     scan_dirs = set(devices)
 
@@ -1964,6 +1986,13 @@ def _dispatch_gen_batch_size():
         action="store_true",
         help="Include linked third-party backend implementations.",
     )
+    parser.add_argument(
+        "--linked-ops",
+        nargs="+",
+        default=None,
+        type=str,
+        help="Linked implementation allowlist. Accepts names separated by spaces or commas.",
+    )
     parser.add_argument(
         "--ops",
         nargs="+",
@@ -1994,6 +2023,10 @@ def _dispatch_gen_batch_size():
             with_linked=args.with_linked,
         )
 
+    linked_op_allowlist = (
+        None if args.linked_ops is None else _normalize_op_allowlist(args.linked_ops)
+    )
+    ops = _filter_linked_implementations(ops, linked_op_allowlist)
     ops = _filter_ops(
         ops,
         _normalize_op_allowlist(args.ops),

src/CMakeLists.txt

@@ -479,6 +479,8 @@ endif()
 
 set(INFINI_OPS_OPS "" CACHE STRING
     "Semicolon- or comma-separated operator allowlist for generated wrappers and bindings")
+set(INFINI_OPS_LINKED_OPS "" CACHE STRING
+    "Semicolon- or comma-separated linked operator allowlist; defaults to INFINI_OPS_OPS")
 set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL
     "Build only the smoke-test operator subset")
 set(_infini_ops_smoke_ops
@@ -510,6 +512,24 @@ if(INFINI_OPS_OPS)
     message(STATUS "Wrapper op allowlist: ${_infini_ops_op_allowlist}")
 endif()
 
+if(INFINI_OPS_LINKED_OPS)
+    string(REPLACE "," ";" _infini_ops_linked_op_allowlist
+        "${INFINI_OPS_LINKED_OPS}")
+    message(STATUS "Linked op allowlist: ${_infini_ops_linked_op_allowlist}")
+
+    if(INFINI_OPS_OPS)
+        foreach(_linked_op IN LISTS _infini_ops_linked_op_allowlist)
+            if(NOT "${_linked_op}" IN_LIST _infini_ops_op_allowlist)
+                message(FATAL_ERROR
+                    "`INFINI_OPS_LINKED_OPS` entry `${_linked_op}` is not "
+                    "present in `INFINI_OPS_OPS`.")
+            endif()
+        endforeach()
+    endif()
+elseif(INFINI_OPS_OPS)
+    set(_infini_ops_linked_op_allowlist "${_infini_ops_op_allowlist}")
+endif()
+
 set(INFINI_OPS_LINKED_SOURCES "")
 set(INFINI_OPS_LINKED_LIBRARIES "")
 set(INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES "")
@@ -554,8 +574,8 @@ if(WITH_LINKED)
         --nm "${_linked_nm}"
         --readelf "${_linked_readelf}"
         --cxxfilt "${_linked_cxxfilt}")
-    if(INFINI_OPS_OPS)
-        list(APPEND _linked_resolver_args --ops ${_infini_ops_op_allowlist})
+    if(_infini_ops_linked_op_allowlist)
+        list(APPEND _linked_resolver_args --ops ${_infini_ops_linked_op_allowlist})
     endif()
 
     execute_process(
@@ -969,6 +989,10 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
     endif()
     if(WITH_LINKED)
         list(APPEND GENERATOR_ARGS --with-linked)
+        if(_infini_ops_linked_op_allowlist)
+            list(APPEND GENERATOR_ARGS
+                --linked-ops ${_infini_ops_linked_op_allowlist})
+        endif()
     endif()
     if(WITH_NINETOOTHED)
         list(APPEND GENERATOR_ARGS --with-ninetoothed)

tests/test_generate_wrappers.py

@@ -582,6 +582,17 @@ def test_torch_system_compiler_receives_host_range_profile_definition():
     ) in cmake
 
 
+def test_cmake_uses_an_independent_linked_operator_allowlist():
+    cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text(
+        encoding="utf-8"
+    )
+
+    assert 'set(INFINI_OPS_LINKED_OPS "" CACHE STRING' in cmake
+    assert 'set(_infini_ops_linked_op_allowlist "${_infini_ops_op_allowlist}")' in cmake
+    assert "--ops ${_infini_ops_linked_op_allowlist}" in cmake
+    assert "--linked-ops ${_infini_ops_linked_op_allowlist}" in cmake
+
+
 def test_generated_dispatch_calls_start_with_dispatch_profile_scope(
     tmp_path, monkeypatch
 ):
@@ -686,6 +697,28 @@ def test_filter_ops_strict_rejects_unavailable_ops():
         raise AssertionError("strict unknown ops should fail")
 
 
+def test_filter_linked_implementations_preserves_other_backends(monkeypatch, tmp_path):
+    module = _load_generator_module()
+    src_dir = tmp_path / "src"
+    native_topk = src_dir / "native" / "cuda" / "nvidia" / "topk_softmax.h"
+    linked_topk = src_dir / "linked" / "torch" / "nvidia" / "topk_softmax.h"
+    linked_flash = src_dir / "linked" / "torch" / "nvidia" / "flash_attn_with_kvcache.h"
+    ops = {
+        "topk_softmax": [native_topk, linked_topk],
+        "flash_attn_with_kvcache": [linked_flash],
+    }
+    monkeypatch.setattr(module, "_SRC_DIR", src_dir)
+
+    assert module._filter_linked_implementations(ops, None) == ops
+    assert module._filter_linked_implementations(ops, ["flash_attn_with_kvcache"]) == {
+        "topk_softmax": [native_topk],
+        "flash_attn_with_kvcache": [linked_flash],
+    }
+    assert module._filter_linked_implementations(ops, ["topk_softmax"]) == {
+        "topk_softmax": [native_topk, linked_topk]
+    }
+
+
 def test_linked_implementations_require_explicit_scan_flag(monkeypatch, tmp_path):
     module = _load_generator_module()
     src_dir = tmp_path / "moore" / "src"