[Inductor] Fix cuda_kernel typing

mlazos · mlazos · commit aba636d50bd5 · 2025-04-27T14:05:26.000-07:00
ghstack-source-id: 1853112 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
@@ -520,7 +521,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[
