From 7de7b02bad6a3e998c2eef2a2ac4eda409ca21f4 Mon Sep 17 00:00:00 2001 From: Bob Ren Date: Thu, 22 Oct 2026 13:49:46 -0700 Subject: [PATCH] [executorch][cuda] Find the tuned INT4 kernel's def under its module-level name too `test_a_tuned_shape_traces_only_its_picked_config` finds the user kernel in Inductor's wrapper by the exact text `def _int4_w4a16_direct_kernel(`. That only matches the wrapper's `async_compile.triton('''...''')` string form, which keeps the kernel's own name. pytorch/pytorch#200502 makes Inductor's python wrapper define each Triton kernel as module-level code by default. There, the def takes Inductor's kernel name, `_int4_w4a16_direct_kernel_0`. The test failed on its DiffTrain import, D126740082, which was reverted partly for this. The test now matches `^def \w*\(` and slices back to the `user_autotune(` before that match. The decorator's `configs=[...]` still comes right before the def in both forms, so the check that only the picked config is traced is unchanged. This lands ahead of the #200502 reland. ___ landed-with-radar-review Differential Revision: [D126992211](https://our.internmc.facebook.com/intern/diff/D126992211/) ghstack-source-id: 454871348 Pull-Request: https://github.com/pytorch/executorch/pull/24096 --- backends/cuda/tests/test_int4_quantized_gemm.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/backends/cuda/tests/test_int4_quantized_gemm.py b/backends/cuda/tests/test_int4_quantized_gemm.py index 89ab1c62015..816cd719c08 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/backends/cuda/tests/test_int4_quantized_gemm.py @@ -20,6 +20,7 @@ import contextlib import os import pickle +import re import statistics import tempfile import types @@ -887,10 +888,10 @@ def linear(x): ) source = "\n".join(code) name = autotuner.fn.fn.__name__ - self.assertIn(f"def {name}(", source) - block = source[ - source.rindex("user_autotune(", 0, source.index(f"def {name}(")) : - ] + # A module-level kernel's def is renamed to Inductor's kernel name: name plus a suffix. + kernel_def = re.search(rf"^def {re.escape(name)}\w*\(", source, re.MULTILINE) + self.assertIsNotNone(kernel_def, source) + block = source[source.rindex("user_autotune(", 0, kernel_def.start()) :] configs = block[: block.index("],")] self.assertEqual(configs.count("'num_warps'"), 1) _check_close(self, compiled, eager)