已合并
catlass adapt 2.13.0 #43293
catlass adapt 2.13.0 #43293
已合并
shi-yufeng99创建于 7月29日
3 个文件变更+14-2
@@ -28,6 +28,18 @@ def patch_tuning_process():
28 28 
29 autotune_process.CUDA_VISIBLE_DEVICES = ASCEND_VISIBLE_DEVICES29 autotune_process.CUDA_VISIBLE_DEVICES = ASCEND_VISIBLE_DEVICES
30 30 
31+ # Replace GPUDeviceBenchmarkMixin with NPUDeviceBenchmarkMixin in
32+ # ExternKernelGPUBenchmarkRequest so that extern kernel benchmarks
33+ # (aten::mm, aten::addmm, etc.) use torch.npu instead of torch.cuda.
34+ from torch._inductor.autotune_process import (
35+ ExternKernelGPUBenchmarkRequest,
36+ ExternKernelBenchmarkRequest,
37+ )
38+ ExternKernelGPUBenchmarkRequest.__bases__ = (
39+ NPUDeviceBenchmarkMixin,
40+ ExternKernelBenchmarkRequest,
41+ )
42+ 
31 43 
32class NPUDeviceBenchmarkMixin:44class NPUDeviceBenchmarkMixin:
33 def do_bench(45 def do_bench(
@@ -898,7 +898,7 @@ class CATLASS1xGemmTemplate(CATLASSGemmTemplate):
898 examples = create_example_tensors(898 examples = create_example_tensors(
899 var_name_to_buffer_name,899 var_name_to_buffer_name,
900 name_to_buffer, # type: ignore[arg-type]900 name_to_buffer, # type: ignore[arg-type]
901- V.graph.sizevars.size_hint,901+ V.graph.sizevars.optimization_hint,
902 )902 )
903 903 
904 evg_config = {904 evg_config = {
@@ -45,7 +45,7 @@ def _finalize_stride_collection(symbol_stride_map, axis_vars):
45 def sort_key(item):45 def sort_key(item):
46 var, stride = item46 var, stride = item
47 try:47 try:
48- hint = V.graph.sizevars.size_hint(stride)48+ hint = V.graph.sizevars.optimization_hint(stride)
49 except TypeError:49 except TypeError:
50 hint = 050 hint = 0
51 return (51 return (