[Bugfix][CPU] Skip unsupported custom op register on CPU (#25534)

Signed-off-by: jiang1.li <jiang1.li@intel.com>
This commit is contained in:
Li, Jiang 2025-09-24 14:21:51 +08:00 committed by GitHub
parent fed8a9b107
commit 1cbcfb94de
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -130,12 +130,14 @@ def _w8a8_triton_block_scaled_mm_fake(
device=qx.device) device=qx.device)
direct_register_custom_op( # Note: the check can be removed when CPU torch > 2.7
"w8a8_triton_block_scaled_mm_func", if not current_platform.is_cpu():
_w8a8_triton_block_scaled_mm_func, direct_register_custom_op(
fake_impl=_w8a8_triton_block_scaled_mm_fake, "w8a8_triton_block_scaled_mm_func",
dispatch_key="CUDA", _w8a8_triton_block_scaled_mm_func,
) fake_impl=_w8a8_triton_block_scaled_mm_fake,
dispatch_key="CUDA",
)
# TODO fix ROCm->Triton custom path: # TODO fix ROCm->Triton custom path: