已合并
fix(inductor): support Flex Attention mask_out cppwrapper #45565
Xuan Peng创建于 25 天前
fix(inductor): support Flex Attention mask_out cppwrapper #45565
已合并
共 2 个文件变更+76-3
| @@ -489,10 +489,21 @@ class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper): | |||
| 489 | enable_simt = npu_config.is_ascend950 and ( | 489 | enable_simt = npu_config.is_ascend950 and ( |
| 490 | "simt" in params["parallel_mode"] or params["is_pure_simt"] | 490 | "simt" in params["parallel_mode"] or params["is_pure_simt"] |
| 491 | ) | 491 | ) |
| 492 | - enable_auto_blockify = not params.get("has_auto_blockify_blacklist_op", False) and triton_support_auto_blockify() | 492 | + enable_auto_blockify = not params.get( |
| 493 | + "has_auto_blockify_blacklist_op", False | ||
| 494 | + ) and triton_support_auto_blockify() | ||
| 495 | + args_decl = wrapper.generate_args_decl( | ||
| 496 | + prefix, | ||
| 497 | + call_args, | ||
| 498 | + arg_types, | ||
| 499 | + arg_signatures, | ||
| 500 | + is_triton_kernel=True, | ||
| 501 | + is_pure_simt=is_pure_simt, | ||
| 502 | + kernel_params=params, | ||
| 503 | + ) | ||
| 493 | prefix.splice(f""" | 504 | prefix.splice(f""" |
| 494 | auto launch_call = [=]() {{ | 505 | auto launch_call = [=]() {{ |
| 495 | - {wrapper.generate_args_decl(prefix, call_args, arg_types, arg_signatures, True, is_pure_simt)} | 506 | + {args_decl} |
| 496 | {wrapper.generate_launch_preparation(kernel_var_name, params, enable_simt, enable_auto_blockify)} | 507 | {wrapper.generate_launch_preparation(kernel_var_name, params, enable_simt, enable_auto_blockify)} |
| 497 | }}; | 508 | }}; |
| 498 | """) | 509 | """) |
| @@ -656,6 +667,7 @@ class CppWrapperNpu(CppWrapperGpu): | |||
| 656 | #include <acl/acl_rt.h> | 667 | #include <acl/acl_rt.h> |
| 657 | #include <runtime/runtime/rt.h> | 668 | #include <runtime/runtime/rt.h> |
| 658 | #include <torch_npu/csrc/core/npu/NPUStream.h> | 669 | #include <torch_npu/csrc/core/npu/NPUStream.h> |
| 670 | + #include <torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h> | ||
| 659 | #include <torch_npu/csrc/framework/OpCommand.h> | 671 | #include <torch_npu/csrc/framework/OpCommand.h> |
| 660 | """) | 672 | """) |
| 661 | super().add_device_include(device) | 673 | super().add_device_include(device) |
| @@ -845,6 +857,7 @@ static inline void load_{kernel_name}() {{ | |||
| 845 | arg_signatures, | 857 | arg_signatures, |
| 846 | is_triton_kernel=True, | 858 | is_triton_kernel=True, |
| 847 | is_pure_simt=False, | 859 | is_pure_simt=False, |
| 860 | + kernel_params: Optional[dict[str, Any]] = None, | ||
| 848 | ): | 861 | ): |
| 849 | """ | 862 | """ |
| 850 | Generates any declarations of args to pass into a kernel call, and then returns the arg names. | 863 | Generates any declarations of args to pass into a kernel call, and then returns the arg names. |
| @@ -884,6 +897,10 @@ static inline void load_{kernel_name}() {{ | |||
| 884 | struct_arg_body = "" | 897 | struct_arg_body = "" |
| 885 | 898 | ||
| 886 | target_support_ffts = triton_support_ffts() | 899 | target_support_ffts = triton_support_ffts() |
| 900 | + kernel_params = kernel_params or {} | ||
| 901 | + lock_num = int(kernel_params.get("lock_num", 0) or 0) | ||
| 902 | + lock_init_val = int(kernel_params.get("lock_init_val", 0) or 0) | ||
| 903 | + workspace_size = int(kernel_params.get("workspace_size", 0) or 0) | ||
| 887 | 904 | ||
| 888 | def process_args(arg, arg_type, arg_signature=None): | 905 | def process_args(arg, arg_type, arg_signature=None): |
| 889 | var_name = f"var_{next(self.arg_var_id)}" | 906 | var_name = f"var_{next(self.arg_var_id)}" |
| @@ -945,11 +962,47 @@ static inline void load_{kernel_name}() {{ | |||
| 945 | } | 962 | } |
| 946 | """ | 963 | """ |
| 947 | 964 | ||
| 965 | + workspace_str = "" | ||
| 966 | + if not is_pure_simt and workspace_size > 0: | ||
| 967 | + workspace_str = f""" | ||
| 968 | + uint64_t workspace_size = static_cast<uint64_t>({workspace_size}) | ||
| 969 | + * grid_0 * grid_1 * grid_2; | ||
| 970 | + auto workspace_tensor = at_npu::native::allocate_workspace( | ||
| 971 | + workspace_size, stream_); | ||
| 972 | + workspace_addr = const_cast<void *>(workspace_tensor.storage().data()); | ||
| 973 | + """ | ||
| 974 | + | ||
| 975 | + sync_block_lock_str = "" | ||
| 976 | + if not is_pure_simt and lock_num > 0: | ||
| 977 | + sync_block_lock_str = f""" | ||
| 978 | + uint64_t sync_block_lock_size = static_cast<uint64_t>({lock_num}) | ||
| 979 | + * sizeof(int64_t); | ||
| 980 | + auto sync_block_lock_tensor = at_npu::native::allocate_workspace( | ||
| 981 | + sync_block_lock_size, stream_); | ||
| 982 | + sync_block_lock = const_cast<void *>( | ||
| 983 | + sync_block_lock_tensor.storage().data()); | ||
| 984 | + std::vector<int64_t> sync_block_lock_init( | ||
| 985 | + {lock_num}, static_cast<int64_t>({lock_init_val})); | ||
| 986 | + ret = aclrtMemcpy( | ||
| 987 | + sync_block_lock, | ||
| 988 | + sync_block_lock_size, | ||
| 989 | + sync_block_lock_init.data(), | ||
| 990 | + sync_block_lock_size, | ||
| 991 | + ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 992 | + if (ret != ACL_SUCCESS) {{ | ||
| 993 | + throw std::runtime_error( | ||
| 994 | + std::string("initialize Triton sync block lock failed, 0x") | ||
| 995 | + + std::to_string(ret)); | ||
| 996 | + }} | ||
| 997 | + """ | ||
| 998 | + | ||
| 948 | args_str = f""" | 999 | args_str = f""" |
| 949 | aclError ret; | 1000 | aclError ret; |
| 950 | {ffts_str if target_support_ffts else ""} | 1001 | {ffts_str if target_support_ffts else ""} |
| 951 | {"void* workspace_addr = NULL;" if not is_pure_simt else ""} | 1002 | {"void* workspace_addr = NULL;" if not is_pure_simt else ""} |
| 952 | {"void* sync_block_lock = NULL;" if not is_pure_simt else ""} | 1003 | {"void* sync_block_lock = NULL;" if not is_pure_simt else ""} |
| 1004 | + {workspace_str} | ||
| 1005 | + {sync_block_lock_str} | ||
| 953 | struct __attribute__((packed)) {{ | 1006 | struct __attribute__((packed)) {{ |
| 954 | {"void* ffts_addr __attribute__((aligned(8)));" if target_support_ffts else ""} | 1007 | {"void* ffts_addr __attribute__((aligned(8)));" if target_support_ffts else ""} |
| 955 | {"void* sync_block_lock __attribute__((aligned(8)));" if not is_pure_simt else ""} | 1008 | {"void* sync_block_lock __attribute__((aligned(8)));" if not is_pure_simt else ""} |
| @@ -1426,7 +1426,20 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1426 | "mix_mode": input_launcher.bin.metadata.mix_mode, | 1426 | "mix_mode": input_launcher.bin.metadata.mix_mode, |
| 1427 | "parallel_mode": input_launcher.bin.metadata.parallel_mode, | 1427 | "parallel_mode": input_launcher.bin.metadata.parallel_mode, |
| 1428 | "is_pure_simt": input_launcher.bin.metadata.is_pure_simt, | 1428 | "is_pure_simt": input_launcher.bin.metadata.is_pure_simt, |
| 1429 | - "has_auto_blockify_blacklist_op": getattr(input_launcher.bin.metadata, "has_auto_blockify_blacklist_op", False) | 1429 | + "has_auto_blockify_blacklist_op": getattr( |
| 1430 | + input_launcher.bin.metadata, | ||
| 1431 | + "has_auto_blockify_blacklist_op", | ||
| 1432 | + False, | ||
| 1433 | + ), | ||
| 1434 | + "lock_num": int( | ||
| 1435 | + getattr(input_launcher.bin.metadata, "lock_num", 0) or 0 | ||
| 1436 | + ), | ||
| 1437 | + "lock_init_val": int( | ||
| 1438 | + getattr(input_launcher.bin.metadata, "lock_init_val", 0) or 0 | ||
| 1439 | + ), | ||
| 1440 | + "workspace_size": int( | ||
| 1441 | + getattr(input_launcher.bin.metadata, "workspace_size", 0) or 0 | ||
| 1442 | + ), | ||
| 1430 | } | 1443 | } |
| 1431 | enable_simt = ("simt" in params["parallel_mode"]) or params["is_pure_simt"] | 1444 | enable_simt = ("simt" in params["parallel_mode"]) or params["is_pure_simt"] |
| 1432 | if npu_config.is_ascend950 and enable_simt: | 1445 | if npu_config.is_ascend950 and enable_simt: |
| @@ -2189,6 +2202,13 @@ class NPUSymbolicGroupedAutotuner(NPUCachingAutotuner): | |||
| 2189 | "has_auto_blockify_blacklist_op": getattr( | 2202 | "has_auto_blockify_blacklist_op": getattr( |
| 2190 | metadata, "has_auto_blockify_blacklist_op", False | 2203 | metadata, "has_auto_blockify_blacklist_op", False |
| 2191 | ), | 2204 | ), |
| 2205 | + "lock_num": int(getattr(metadata, "lock_num", 0) or 0), | ||
| 2206 | + "lock_init_val": int( | ||
| 2207 | + getattr(metadata, "lock_init_val", 0) or 0 | ||
| 2208 | + ), | ||
| 2209 | + "workspace_size": int( | ||
| 2210 | + getattr(metadata, "workspace_size", 0) or 0 | ||
| 2211 | + ), | ||
| 2192 | "cubin_path": binary_path, | 2212 | "cubin_path": binary_path, |
| 2193 | } | 2213 | } |
| 2194 | 2214 | ||