refactor(build): select InfiniOps implementations from JSON

Author: voltjia

csrc/infinicore/src/ops/random_sample/random_sample_infiniops.cc

@@ -57,6 +57,14 @@ void calculate(
     infini::ops::Handle handle;
     handle.set_stream(context::getStream());
     infini::ops::Config config;
+    infini::ops::Config sampling_config;
+    static const auto implementation_index = []() {
+        const auto indices = infini::ops::TopKTopPSamplingFromLogits::active_implementation_indices(
+            infini::ops::Device::Type::kNvidia);
+        INFINICORE_ASSERT(!indices.empty());
+        return indices.front();
+    }();
+    sampling_config.set_implementation_index(implementation_index);
 
     auto logits_2d = logits->view({1, logits->numel()});
     Tensor sampled_logits = logits_2d;
@@ -112,7 +120,7 @@ void calculate(
     const std::optional<int64_t> offset{0};
     infini::ops::TopKTopPSamplingFromLogits::Call(
         handle,
-        config,
+        sampling_config,
         TensorMeta(sampled_logits).tensor(sampled_logits),
         top_k_tensor,
         top_p_tensor,

docs/superpowers/specs/2026-07-16-infini-stack-repository-boundaries-design.md

@@ -82,7 +82,7 @@ InfiniLM will own these files:
 
 ```text
 scripts/build_infini_stack.py
-scripts/configs/infiniops_ops.txt
+scripts/configs/infiniops_ops.json
 test/scripts/test_build_infini_stack.py
 ```
 

scripts/build_infini_stack.py

@@ -9,28 +9,41 @@
 from typing import Dict, List, Mapping, Optional, Sequence
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
-OPERATOR_SET = PROJECT_ROOT / "scripts/configs/infiniops_ops.txt"
-LINKED_OPERATORS = (
-    "flash_attn_varlen_func",
-    "flash_attn_with_kvcache",
-)
+OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops.json"
 SUBMODULES = {
     "InfiniRT": Path("submodules/InfiniRT"),
     "InfiniOps": Path("submodules/InfiniOps"),
     "InfiniCCL": Path("submodules/InfiniCCL"),
 }
 
 
-def read_operator_set(path: Path = OPERATOR_SET) -> List[str]:
-    operators = [line.strip() for line in path.read_text(encoding="utf-8").splitlines()]
-    operators = [operator for operator in operators if operator]
-    if not operators:
-        raise ValueError(f"Operator set is empty: {path}")
-    if len(operators) != len(set(operators)):
-        raise ValueError(f"Operator set contains duplicates: {path}")
-    if operators != sorted(operators):
-        raise ValueError(f"Operator set must be sorted: {path}")
-    return operators
+def read_operator_config(path: Path = OPERATOR_CONFIG) -> Dict[str, Dict[str, object]]:
+    config = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(config, dict) or not config:
+        raise ValueError(f"Operator config must be a non-empty object: {path}")
+    if list(config) != sorted(config):
+        raise ValueError(f"Operator config must be sorted: {path}")
+    for operator, selection in config.items():
+        if not isinstance(operator, str) or not operator:
+            raise ValueError(f"Operator names must be non-empty strings: {path}")
+        if not isinstance(selection, dict) or set(selection) != {"implementations"}:
+            raise ValueError(
+                f"Operator {operator!r} must define only 'implementations': {path}"
+            )
+        implementations = selection["implementations"]
+        if (
+            not isinstance(implementations, list)
+            or not implementations
+            or any(
+                type(slot) is not int or not 0 <= slot < 32 for slot in implementations
+            )
+            or len(implementations) != len(set(implementations))
+        ):
+            raise ValueError(
+                f"Operator {operator!r} implementations must be unique "
+                f"integers between 0 and 31: {path}"
+            )
+    return config
 
 
 def parse_gitlink(output: str, relative_path: Path) -> str:
@@ -152,7 +165,7 @@ def build_infiniops_commands(
     build_type: str,
     jobs: int,
     cuda_arch: Optional[str],
-    operators: Sequence[str],
+    operator_config: Path,
 ) -> List[List[str]]:
     configure = [
         "cmake",
@@ -163,14 +176,12 @@ def build_infiniops_commands(
         "-DWITH_CPU=ON",
         "-DWITH_NVIDIA=ON",
         "-DWITH_LINKED=ON",
-        f"-DINFINI_OPS_LINKED_OPS={','.join(LINKED_OPERATORS)}",
         "-DWITH_TORCH=ON",
-        "-DINFINI_OPS_TORCH_OPS=argmax",
         "-DAUTO_DETECT_DEVICES=OFF",
         "-DAUTO_DETECT_BACKENDS=OFF",
         "-DGENERATE_PYTHON_BINDINGS=OFF",
         f"-DINFINI_RT_ROOT={prefix}",
-        f"-DINFINI_OPS_OPS={','.join(operators)}",
+        f"-DINFINI_OPS_OPS={operator_config}",
         f"-DCMAKE_BUILD_TYPE={build_type}",
         f"-DCMAKE_INSTALL_PREFIX={prefix}",
     ]
@@ -263,7 +274,7 @@ def write_manifest(
     infinilm_revision: str,
     infinicore_revision: str,
     revisions: Mapping[str, str],
-    operators: Sequence[str],
+    operator_config: Mapping[str, Mapping[str, object]],
     prefix: Path,
     build_type: str,
     cuda_arch: Optional[str],
@@ -278,7 +289,11 @@ def write_manifest(
         "infinilm": infinilm_revision,
         "install_prefix": str(prefix),
         "jobs": jobs,
-        "operators": list(operators),
+        "operators": list(operator_config),
+        "operator_implementations": {
+            operator: selection["implementations"]
+            for operator, selection in operator_config.items()
+        },
         "submodules": dict(revisions),
         "test": test,
     }
@@ -362,7 +377,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
     build_root = resolve_from_project(args.build_root)
     prefix = build_root / "prefix"
     manifest_path = build_root / "manifest.json"
-    operators = read_operator_set()
+    operator_config = read_operator_config()
     revisions = {
         name: validate_submodule(infinicore_root, relative_path)
         for name, relative_path in SUBMODULES.items()
@@ -376,7 +391,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
     print(f"InfiniRT: {revisions['InfiniRT']}")
     print(f"InfiniOps: {revisions['InfiniOps']}")
     print(f"InfiniCCL: {revisions['InfiniCCL']}")
-    print(f"Operators ({len(operators)}): {','.join(operators)}")
+    print(f"Operators ({len(operator_config)}): {','.join(operator_config)}")
 
     if not args.dry_run:
         manifest_path.unlink(missing_ok=True)
@@ -402,7 +417,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
         args.build_type,
         args.jobs,
         args.cuda_arch,
-        operators,
+        OPERATOR_CONFIG,
     ):
         run(command, PROJECT_ROOT, integration_env, args.dry_run)
 
@@ -423,7 +438,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
             infinilm_revision,
             infinicore_revision,
             revisions,
-            operators,
+            operator_config,
             prefix,
             args.build_type,
             args.cuda_arch,

scripts/configs/infiniops_ops.json

@@ -0,0 +1,25 @@
+{
+  "add": {"implementations": [0]},
+  "argmax": {"implementations": [8]},
+  "causal_softmax": {"implementations": [0]},
+  "convolution": {"implementations": [0]},
+  "copy": {"implementations": [0]},
+  "embedding": {"implementations": [0]},
+  "fill": {"implementations": [0]},
+  "flash_attn_varlen_func": {"implementations": [16]},
+  "flash_attn_with_kvcache": {"implementations": [16]},
+  "fused_add_rms_norm": {"implementations": [0]},
+  "gelu": {"implementations": [0]},
+  "gemm": {"implementations": [0]},
+  "mul": {"implementations": [0]},
+  "relu": {"implementations": [0]},
+  "reshape_and_cache_flash": {"implementations": [0]},
+  "rms_norm": {"implementations": [0]},
+  "rotary_embedding": {"implementations": [0]},
+  "sigmoid": {"implementations": [0]},
+  "silu": {"implementations": [0]},
+  "silu_and_mul": {"implementations": [0]},
+  "softmax": {"implementations": [0]},
+  "top_k_top_p_sampling_from_logits": {"implementations": [16]},
+  "topk_softmax": {"implementations": [0]}
+}

scripts/configs/infiniops_ops.txt

@@ -53,11 +53,11 @@ def _patch_main_dependencies(self, stack, args):
         )
         return run, write_manifest
 
-    def test_operator_set_is_stable(self):
-        operators = build_infini_stack.read_operator_set()
+    def test_operator_config_is_stable(self):
+        config = build_infini_stack.read_operator_config()
 
         self.assertEqual(
-            operators,
+            list(config),
             [
                 "add",
                 "argmax",
@@ -84,15 +84,23 @@ def test_operator_set_is_stable(self):
                 "topk_softmax",
             ],
         )
-        self.assertTrue(all("_infinilm" not in operator for operator in operators))
-
-    def test_invalid_operator_set_is_rejected(self):
+        self.assertEqual(
+            config["top_k_top_p_sampling_from_logits"]["implementations"], [16]
+        )
+        self.assertEqual(config["argmax"]["implementations"], [8])
+        self.assertEqual(config["flash_attn_varlen_func"]["implementations"], [16])
+        self.assertEqual(config["flash_attn_with_kvcache"]["implementations"], [16])
+        self.assertEqual(config["topk_softmax"]["implementations"], [0])
+        self.assertEqual(config["add"]["implementations"], [0])
+        self.assertTrue(all("_infinilm" not in operator for operator in config))
+
+    def test_invalid_operator_config_is_rejected(self):
         with tempfile.TemporaryDirectory() as directory:
-            path = Path(directory) / "ops.txt"
-            path.write_text("rms_norm\nadd\nadd\n", encoding="utf-8")
+            path = Path(directory) / "ops.json"
+            path.write_text('{"add": {"implementations": [32]}}', encoding="utf-8")
 
-            with self.assertRaisesRegex(ValueError, "duplicates"):
-                build_infini_stack.read_operator_set(path)
+            with self.assertRaisesRegex(ValueError, "between 0 and 31"):
+                build_infini_stack.read_operator_config(path)
 
     def test_parse_gitlink(self):
         output = "160000 commit abcdef1234567890\tsubmodules/InfiniRT\n"
@@ -266,28 +274,20 @@ def test_infiniops_commands_use_pinned_runtime_and_operator_set(self):
             "Release",
             8,
             "sm_80",
-            [
-                "add",
-                "flash_attn_varlen_func",
-                "flash_attn_with_kvcache",
-                "rms_norm",
-            ],
+            Path("ops.json"),
         )
 
         configure = commands[0]
         self.assertIn(f"-DINFINI_RT_ROOT={Path('build/prefix')}", configure)
-        self.assertIn(
-            "-DINFINI_OPS_OPS="
-            "add,flash_attn_varlen_func,flash_attn_with_kvcache,rms_norm",
-            configure,
-        )
+        self.assertIn(f"-DINFINI_OPS_OPS={Path('ops.json')}", configure)
         self.assertIn("-DWITH_LINKED=ON", configure)
-        self.assertIn(
-            "-DINFINI_OPS_LINKED_OPS=flash_attn_varlen_func,flash_attn_with_kvcache",
-            configure,
-        )
         self.assertIn("-DWITH_TORCH=ON", configure)
-        self.assertIn("-DINFINI_OPS_TORCH_OPS=argmax", configure)
+        self.assertFalse(
+            any(option.startswith("-DINFINI_OPS_LINKED_OPS=") for option in configure)
+        )
+        self.assertFalse(
+            any(option.startswith("-DINFINI_OPS_TORCH_OPS=") for option in configure)
+        )
         self.assertIn("-DCMAKE_CUDA_ARCHITECTURES=80", configure)
         self.assertEqual(commands[1][3:5], ["--target", "infiniops"])
         self.assertEqual(commands[-1], ["cmake", "--install", str(Path("build/ops"))])
@@ -478,7 +478,10 @@ def test_manifest_records_revisions_operators_and_options(self):
                     "InfiniOps": "ops-sha",
                     "InfiniCCL": "ccl-sha",
                 },
-                ["add", "rms_norm"],
+                {
+                    "add": {"implementations": [0]},
+                    "rms_norm": {"implementations": [0]},
+                },
                 Path("prefix"),
                 "Release",
                 "sm_80,sm_90a",
@@ -493,6 +496,10 @@ def test_manifest_records_revisions_operators_and_options(self):
             self.assertEqual(manifest["submodules"]["InfiniRT"], "rt-sha")
             self.assertEqual(manifest["submodules"]["InfiniCCL"], "ccl-sha")
             self.assertEqual(manifest["operators"], ["add", "rms_norm"])
+            self.assertEqual(
+                manifest["operator_implementations"],
+                {"add": [0], "rms_norm": [0]},
+            )
             self.assertEqual(manifest["install_prefix"], str(Path("prefix")))
             self.assertEqual(manifest["build_type"], "Release")
             self.assertEqual(manifest["cuda_arch"], "sm_80,sm_90a")

test/scripts/test_build_infini_stack.py

@@ -107,8 +107,8 @@ def test_infiniops_adapters_do_not_use_infinilm_suffixed_apis(self) -> None:
                     source, re.compile(r"infini::ops::[A-Za-z0-9_]*Infinilm")
                 )
 
-        allowlist = read_source("scripts/configs/infiniops_ops.txt")
-        self.assertNotIn("_infinilm", allowlist)
+        operator_config = read_source("scripts/configs/infiniops_ops.json")
+        self.assertNotIn("_infinilm", operator_config)
 
     def test_infiniops_adapter_temporaries_keep_owning_tensors(self) -> None:
         cases = {
@@ -623,6 +623,18 @@ def test_random_sampling_uses_only_canonical_infiniops(self) -> None:
             self.assertIn(sentinel, wrapper)
         self.assertIn("base/top_k_top_p_sampling_from_logits.h", adapter)
         self.assertIn("infini::ops::TopKTopPSamplingFromLogits::Call", adapter)
+        self.assertIn("active_implementation_indices", adapter)
+        self.assertIn("INFINICORE_ASSERT(!indices.empty())", adapter)
+        self.assertIn("sampling_config.set_implementation_index", adapter)
+        self.assertRegex(adapter, re.compile(r"Fill::Call\(\s*handle,\s*config,"))
+        self.assertRegex(adapter, re.compile(r"Mul::Call\(\s*handle,\s*config,"))
+        self.assertRegex(
+            adapter,
+            re.compile(
+                r"TopKTopPSamplingFromLogits::Call\(\s*handle,"
+                r"\s*sampling_config,"
+            ),
+        )
         self.assertIn("infini::ops::Fill::Call", adapter)
         self.assertIn("infini::ops::Mul::Call", adapter)
         self.assertIn("1.0f / temperature", adapter)