Repository: vllm-project/vllm
File: vllm/kernels/helion/register.py (Python)

Code before:
        tree_spec: pytree.TreeSpec,
    ) -> dict[str, Any]:
        leaf_specs: list[dict[str, Any]] = []
        for out in flat_outputs:
            if isinstance(out, torch.Tensor):
                leaf_specs.append(
                    {
                        "type": "tensor",
                        "shape": list(out.shape),
                        "stride": list(out.stride()),
                        "dtype": out.dtype,
                        "device": str(out.device),
                    }
                )
            elif isinstance(out, (torch.SymInt, int, float, bool)):
                leaf_specs.append({"type": "scalar", "scalar_value": out})
            else:
                raise TypeError(
                    f"Unsupported output type {type(out).__name__} in "
                    f"kernel output. Only tensors, int, float, and bool "
                    f"are supported."
                )
        return {
            "leaf_specs": leaf_specs,
            "tree_spec_str": pytree.treespec_dumps(tree_spec),
        }

    def register_config_picker(
        self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
    ) -> Callable[[tuple[Any, ...], list[str]], str | None]:
        self._config_picker = picker_func
        return picker_func

    def register_input_generator(
        self, generator_func: Callable[[], dict[str, tuple[Any, ...]]]
    ) -> Callable[[], dict[str, tuple[Any, ...]]]:
        """
        Register a function to generate inputs for autotuning and benchmarking.

        Args:
            generator_func: Function that returns dict[str, tuple] where:
                - key: Configuration identifier (e.g., "4096", "hidden_4096")
                - value: Tuple of arguments to pass to the kernel

        Returns:
            The registered function (for decorator usage)

        Example:
            @kernel_wrapper.register_input_generator
            def generate_inputs():
                return {

Review comment (suggestion):
wonder if it's better to expose this as util function from Helion, so that internal and external use the same logic

Code after:
            else:
                reconstructed.append(next(hop_iter))
        return pytree.tree_unflatten(reconstructed, tree_spec)

    @staticmethod
    def _partition_args(
        kernel: Any,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        constant_args: dict[str, Any] = {}
        tensor_args: dict[str, Any] = {}
        params = list(kernel.signature.parameters.keys())
        for i, val in enumerate(args):
            name = params[i]
            if isinstance(val, torch.Tensor):
                tensor_args[name] = val
            else:
                constant_args[name] = val
        for name, val in kwargs.items():
            if isinstance(val, torch.Tensor):
                tensor_args[name] = val
            else:
                constant_args[name] = val
        return constant_args, tensor_args

    def register_config_picker(
        self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
    ) -> Callable[[tuple[Any, ...], list[str]], str | None]:
        self._config_picker = picker_func
        return picker_func

    def register_input_generator(
        self, generator_func: Callable[[], dict[str, tuple[Any, ...]]]
    ) -> Callable[[], dict[str, tuple[Any, ...]]]:
        """
        Register a function to generate inputs for autotuning and benchmarking.

        Args:
            generator_func: Function that returns dict[str, tuple] where:
                - key: Configuration identifier (e.g., "4096", "hidden_4096")
                - value: Tuple of arguments to pass to the kernel

        Returns:
            The registered function (for decorator usage)

        Example:
            @kernel_wrapper.register_input_generator
            def generate_inputs():
                return {
                    "4096": (torch.randn(4096, device="cuda"), 0.5),