[Inductor] Fix cuda_kernel typing

mlazos · mlazos · commit b70bd22fd00a · 2025-04-09T12:43:32.000-07:00
ghstack-source-id: 317fe01 Pull Request resolved: #150908
diff --git a/torch/_inductor/codegen/cuda/cuda_kernel.py b/torch/_inductor/codegen/cuda/cuda_kernel.py
@@ -1,4 +1,5 @@
 # mypy: allow-untyped-defs
+import functools
 import logging
 from dataclasses import dataclass
 from typing import Any, Callable, Literal, Optional, TYPE_CHECKING, Union
@@ -504,7 +505,10 @@ def __init__(
         category: str,
         input_nodes: list[Buffer],
         layout: Layout,
-        make_kernel_render: Callable[[CUDATemplateBuffer, Optional[list[IRNode]]], str],
+        make_kernel_render: Callable[
+            [CUDATemplateBuffer, Optional[list[IRNode]]],
+            tuple[CUDATemplateKernel, functools.partial[str]],
+        ],
         bmreq: CUDABenchmarkRequest,
         template: "CUDATemplate",  # type: ignore[name-defined]
         info_kwargs: Optional[
