diff --git a/README.md b/README.md index 02786c280..d4e7535fc 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,10 @@ export MOJO_RUN_MODE="EAGER" source code: mojo_opset/backends/ttx/kernels +### UC NPU + +UC NPU is a Unified Compiler based backend for Ascend NPU kernels. It is selected with `MOJO_BACKEND="uc"`; the support matrix labels it as `uc_npu`. + ### Ixformer Ixformer is a backend powered by the [ixformer](https://github.com/AIsoftwareFactory/ixformer) library, providing high-performance fused operator implementations for Iluvatar GPUs. @@ -49,6 +53,7 @@ You can control the backend you want to use via the `MOJO_BACKEND` environment v - "ixformer" - "ttx" +- "uc" - "torch_npu" - "torch" @@ -58,94 +63,94 @@ When multiple backends are added, Mojo Opset selects the backend implementation ### Core Mojo Operator List -| Op Category | Op Name | torch native | torch_npu | ttx | ixformer | -| :---------- | :------ | :----------- | :-------- | :-- | :------- | -| Activation | `MojoGelu` | ✅ | ✅ | ✅ | TBD | -| Activation | `MojoSilu` | ✅ | ✅ | ✅ | TBD | -| Activation | `MojoSwiGLU` | ✅ | ✅ | ✅ | TBD | -| Activation | `MojoRotateActivation` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPrefillGQA` | ✅ | ✅ | TBD | TBD | -| Attention | `MojoPagedPrefillGQA` | ✅ | ✅ | ✅ | ✅ | -| Attention | `MojoPrefillMLA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPagedPrefillMLA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPrefillNSA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPagedPrefillNSA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoDecodeGQA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPagedDecodeGQA` | ✅ | ✅ | ✅ | ✅ | -| Attention | `MojoDecodeMLA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPagedDecodeMLA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoDecodeNSA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoPagedDecodeNSA` | ✅ | TBD | TBD | TBD | -| Attention | `MojoSdpa` | ✅ | TBD | ✅ | TBD | -| Attention | `MojoPagedPrefillSWA` | ✅ | TBD | ✅ | TBD | -| Attention | `MojoPagedDecodeSWA` | ✅ | TBD | ✅ | TBD | -| Attention | `MojoSWA` | ✅ | TBD | ✅ | TBD | -| KVCache | `MojoStorePagedKVCache` | ✅ | TBD | ✅ | TBD | -| KVCache | `MojoStoreMLAKVCache` | ✅ | TBD | TBD | TBD | -| KVCache | `MojoStorePagedMLAKVCache` | ✅ | TBD | TBD | TBD | -| Gemm | `MojoGemm` | ✅ | TBD | TBD | TBD | -| Gemm | `MojoQuantGemm` | ✅ | ✅ | ✅ | TBD | -| Gemm | `MojoGroupGemm` | ✅ | ✅ | ✅ | TBD | -| ComputeComm | `MojoGemmAll2All` | ✅ | TBD | TBD | TBD | -| ComputeComm | `MojoAllGatherGemm` | ✅ | TBD | TBD | TBD | -| ComputeComm | `MojoGemmAllReduce` | ✅ | TBD | TBD | TBD | -| ComputeComm | `MojoGemmReduceScatter` | ✅ | TBD | TBD | TBD | -| ComputeComm | `MojoQuantGemmAll2All` | ✅ | TBD | TBD | TBD | -| ComputeComm | `MojoAll2AllQuantGemm` | ✅ | TBD | TBD | TBD | -| Embedding | `MojoEmbedding` | ✅ | TBD | TBD | TBD | -| Embedding | `MojoParallelEmbedding` | ✅ | TBD | TBD | TBD | -| OverEncoding | `MojoOverEncoding` | ✅ | TBD | ✅ | TBD | -| OverEncoding | `MojoOverEncodingNGram` | ✅ | TBD | ✅ | TBD | -| Quantize | `MojoStaticQuant` | ✅ | TBD | ✅ | TBD | -| Quantize | `MojoDequant` | ✅ | TBD | TBD | TBD | -| Quantize | `MojoDynamicQuant` | ✅ | ✅ | ✅ | TBD | -| Quantize | `MojoMoEDynamicQuant` | ✅ | ✅ | ✅ | TBD | -| Quantize | `MojoDequantSwiGLUQuant` | ✅ | ✅ | TBD | TBD | -| MoE | `MojoMoE` | ✅ | TBD | TBD | TBD | -| MoE | `MojoMoEGating` | ✅ | TBD | TBD | TBD | -| MoE | `MojoMoEDispatch` | ✅ | TBD | TBD | TBD | -| MoE | `MojoMoEInitRoutingDynamicQuant` | ✅ | TBD | TBD | TBD | -| MoE | `MojoFusedSwiGLUMoEScaleDynamicQuantize` | ✅ | TBD | TBD | TBD | -| MoE | `MojoExperts` | ✅ | TBD | TBD | TBD | -| MoE | `MojoMoECombine` | ✅ | TBD | TBD | TBD | -| MoE | `MojoQuantExperts` | ✅ | TBD | TBD | TBD | -| MoE | `MojoQuantMoE` | ✅ | TBD | TBD | TBD | -| Norm | `MojoLayerNorm` | ✅ | TBD | ✅ | ✅ | -| Norm | `MojoRMSNorm` | ✅ | ✅ | ✅ | ✅ | -| Norm | `MojoGroupLayerNorm` | ✅ | TBD | TBD | TBD | -| Norm | `MojoGroupRMSNorm` | ✅ | TBD | ✅ | TBD | -| Norm | `MojoChannelRMSNorm` | ✅ | TBD | TBD | TBD | -| Norm | `MojoRMSNormQuant` | ✅ | ✅ | TBD | TBD | -| Norm | `MojoLayerNormQuant` | ✅ | ✅ | TBD | TBD | -| Norm | `MojoResidualAddRMSNorm` | ✅ | ✅ | ✅ | ✅ | -| Norm | `MojoResidualAddLayerNorm` | ✅ | TBD | ✅ | ✅ | -| Norm | `MojoResidualAddRMSNormQuant` | ✅ | ✅ | TBD | TBD | -| Norm | `MojoResidualAddLayerNormQuant` | ✅ | ✅ | TBD | TBD | -| Norm | `MojoResidualAddNormCast` | ✅ | TBD | TBD | TBD | -| PositionEmb | `MojoRotaryEmbedding` | ✅ | TBD | ✅ | TBD | -| PositionEmb | `MojoRelativeEmbedding` | ✅ | TBD | ✅ | TBD | -| PositionEmb | `MojoApplyRoPE` | ✅ | ✅ | ✅ | TBD | -| PositionEmb | `MojoRoPEStoreKV` | ✅ | TBD | TBD | TBD | -| PositionEmb | `MojoNormRoPE` | ✅ | TBD | TBD | TBD | -| PositionEmb | `MojoNormRoPEStoreKV` | ✅ | TBD | TBD | TBD | -| PositionEmb | `MojoGridRoPE` | ✅ | TBD | TBD | TBD | -| Sampling | `MojoTopPSampling` | ✅ | TBD | ✅ | TBD | -| Sampling | `MojoTopKSampling` | ✅ | TBD | ✅ | TBD | -| Sampling | `MojoRejectSampling` | ✅ | TBD | ✅ | TBD | -| Sampling | `MojoJoinProbRejectSampling` | ✅ | TBD | ✅ | TBD | -| Sampling | `MojoApplyPenaltiesTempurate` | ✅ | TBD | ✅ | TBD | -| Sampling | `MojoTopPFilter` | ✅ | TBD | ✅ | TBD | -| Convolution | `MojoCausalConv1dUpdateState` | ✅ | TBD | ✅ | TBD | -| MLP | `MojoSwiGLUMLP` | ✅ | TBD | TBD | TBD | -| Indexer | `MojoLightningIndexer` | ✅ | TBD | ✅ | TBD | +| Op Category | Op Name | torch native | torch_npu | ttx | ixformer | uc_npu | +| :---------- | :------ | :----------- | :-------- | :-- | :------- | :----- | +| Activation | `MojoGelu` | ✅ | ✅ | ✅ | TBD | ✅ | +| Activation | `MojoSilu` | ✅ | ✅ | ✅ | TBD | ✅ | +| Activation | `MojoSwiGLU` | ✅ | ✅ | ✅ | TBD | ✅ | +| Activation | `MojoRotateActivation` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPrefillGQA` | ✅ | ✅ | TBD | TBD | TBD | +| Attention | `MojoPagedPrefillGQA` | ✅ | ✅ | ✅ | ✅ | TBD | +| Attention | `MojoPrefillMLA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPagedPrefillMLA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPrefillNSA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPagedPrefillNSA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoDecodeGQA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPagedDecodeGQA` | ✅ | ✅ | ✅ | ✅ | TBD | +| Attention | `MojoDecodeMLA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPagedDecodeMLA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoDecodeNSA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoPagedDecodeNSA` | ✅ | TBD | TBD | TBD | TBD | +| Attention | `MojoSdpa` | ✅ | TBD | ✅ | TBD | ✅ | +| Attention | `MojoPagedPrefillSWA` | ✅ | TBD | ✅ | TBD | TBD | +| Attention | `MojoPagedDecodeSWA` | ✅ | TBD | ✅ | TBD | TBD | +| Attention | `MojoSWA` | ✅ | TBD | ✅ | TBD | TBD | +| KVCache | `MojoStorePagedKVCache` | ✅ | TBD | ✅ | TBD | TBD | +| KVCache | `MojoStoreMLAKVCache` | ✅ | TBD | TBD | TBD | TBD | +| KVCache | `MojoStorePagedMLAKVCache` | ✅ | TBD | TBD | TBD | TBD | +| Gemm | `MojoGemm` | ✅ | TBD | TBD | TBD | TBD | +| Gemm | `MojoQuantGemm` | ✅ | ✅ | ✅ | TBD | ✅ | +| Gemm | `MojoGroupGemm` | ✅ | ✅ | ✅ | TBD | TBD | +| ComputeComm | `MojoGemmAll2All` | ✅ | TBD | TBD | TBD | TBD | +| ComputeComm | `MojoAllGatherGemm` | ✅ | TBD | TBD | TBD | TBD | +| ComputeComm | `MojoGemmAllReduce` | ✅ | TBD | TBD | TBD | TBD | +| ComputeComm | `MojoGemmReduceScatter` | ✅ | TBD | TBD | TBD | TBD | +| ComputeComm | `MojoQuantGemmAll2All` | ✅ | TBD | TBD | TBD | TBD | +| ComputeComm | `MojoAll2AllQuantGemm` | ✅ | TBD | TBD | TBD | TBD | +| Embedding | `MojoEmbedding` | ✅ | TBD | TBD | TBD | TBD | +| Embedding | `MojoParallelEmbedding` | ✅ | TBD | TBD | TBD | TBD | +| OverEncoding | `MojoOverEncoding` | ✅ | TBD | ✅ | TBD | TBD | +| OverEncoding | `MojoOverEncodingNGram` | ✅ | TBD | ✅ | TBD | TBD | +| Quantize | `MojoStaticQuant` | ✅ | TBD | ✅ | TBD | ✅ | +| Quantize | `MojoDequant` | ✅ | TBD | TBD | TBD | TBD | +| Quantize | `MojoDynamicQuant` | ✅ | ✅ | ✅ | TBD | ✅ | +| Quantize | `MojoMoEDynamicQuant` | ✅ | ✅ | ✅ | TBD | ✅ | +| Quantize | `MojoDequantSwiGLUQuant` | ✅ | ✅ | TBD | TBD | TBD | +| MoE | `MojoMoE` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoMoEGating` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoMoEDispatch` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoMoEInitRoutingDynamicQuant` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoFusedSwiGLUMoEScaleDynamicQuantize` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoExperts` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoMoECombine` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoQuantExperts` | ✅ | TBD | TBD | TBD | TBD | +| MoE | `MojoQuantMoE` | ✅ | TBD | TBD | TBD | TBD | +| Norm | `MojoLayerNorm` | ✅ | TBD | ✅ | ✅ | ✅ | +| Norm | `MojoRMSNorm` | ✅ | ✅ | ✅ | ✅ | ✅ | +| Norm | `MojoGroupLayerNorm` | ✅ | TBD | TBD | TBD | TBD | +| Norm | `MojoGroupRMSNorm` | ✅ | TBD | ✅ | TBD | TBD | +| Norm | `MojoChannelRMSNorm` | ✅ | TBD | TBD | TBD | TBD | +| Norm | `MojoRMSNormQuant` | ✅ | ✅ | TBD | TBD | TBD | +| Norm | `MojoLayerNormQuant` | ✅ | ✅ | TBD | TBD | TBD | +| Norm | `MojoResidualAddRMSNorm` | ✅ | ✅ | ✅ | ✅ | ✅ | +| Norm | `MojoResidualAddLayerNorm` | ✅ | TBD | ✅ | ✅ | ✅ | +| Norm | `MojoResidualAddRMSNormQuant` | ✅ | ✅ | TBD | TBD | TBD | +| Norm | `MojoResidualAddLayerNormQuant` | ✅ | ✅ | TBD | TBD | TBD | +| Norm | `MojoResidualAddNormCast` | ✅ | TBD | TBD | TBD | TBD | +| PositionEmb | `MojoRotaryEmbedding` | ✅ | TBD | ✅ | TBD | ✅ | +| PositionEmb | `MojoRelativeEmbedding` | ✅ | TBD | ✅ | TBD | TBD | +| PositionEmb | `MojoApplyRoPE` | ✅ | ✅ | ✅ | TBD | ✅ | +| PositionEmb | `MojoRoPEStoreKV` | ✅ | TBD | TBD | TBD | TBD | +| PositionEmb | `MojoNormRoPE` | ✅ | TBD | TBD | TBD | TBD | +| PositionEmb | `MojoNormRoPEStoreKV` | ✅ | TBD | TBD | TBD | TBD | +| PositionEmb | `MojoGridRoPE` | ✅ | TBD | TBD | TBD | TBD | +| Sampling | `MojoTopPSampling` | ✅ | TBD | ✅ | TBD | TBD | +| Sampling | `MojoTopKSampling` | ✅ | TBD | ✅ | TBD | TBD | +| Sampling | `MojoRejectSampling` | ✅ | TBD | ✅ | TBD | TBD | +| Sampling | `MojoJoinProbRejectSampling` | ✅ | TBD | ✅ | TBD | TBD | +| Sampling | `MojoApplyPenaltiesTempurate` | ✅ | TBD | ✅ | TBD | TBD | +| Sampling | `MojoTopPFilter` | ✅ | TBD | ✅ | TBD | TBD | +| Convolution | `MojoCausalConv1dUpdateState` | ✅ | TBD | ✅ | TBD | TBD | +| MLP | `MojoSwiGLUMLP` | ✅ | TBD | TBD | TBD | TBD | +| Indexer | `MojoLightningIndexer` | ✅ | TBD | ✅ | TBD | TBD | ### Experimental Mojo Operator List -| Op Category | Op Name | torch native | torch_npu | ttx | ixformer | -| :---------- | :------ | :----------- | :-------- | :-- | :------- | -| Experimental | `MojoQuantBatchGemmReduceSum` | ✅ | ✅ | ✅ | TBD | -| Experimental | `MojoIndexer` | ✅ | TBD | ✅ | TBD | -| Experimental | `MojoStoreLowrank` | ✅ | TBD | ✅ | TBD | +| Op Category | Op Name | torch native | torch_npu | ttx | ixformer | uc_npu | +| :---------- | :------ | :----------- | :-------- | :-- | :------- | :----- | +| Experimental | `MojoQuantBatchGemmReduceSum` | ✅ | ✅ | ✅ | TBD | TBD | +| Experimental | `MojoIndexer` | ✅ | TBD | ✅ | TBD | TBD | +| Experimental | `MojoStoreLowrank` | ✅ | TBD | ✅ | TBD | TBD | ### Core Mojo Function List diff --git a/mojo_opset/backends/__init__.py b/mojo_opset/backends/__init__.py index 4911b60da..6181deb29 100644 --- a/mojo_opset/backends/__init__.py +++ b/mojo_opset/backends/__init__.py @@ -12,12 +12,13 @@ if platform == "npu": import os - # special setting for npu deterministic matmul - os.environ["CLOSE_MATMUL_K_SHIFT"] = "1" + # Default for deterministic NPU matmul; respect explicit user/CANN settings. + os.environ.setdefault("CLOSE_MATMUL_K_SHIFT", "1") configure_torch_deterministic() _SUPPORT_TTX_PLATFROM = ["npu", "ilu", "mlu"] _SUPPORT_TORCH_NPU_PLATFROM = ["npu"] +_SUPPORT_UC_PLATFORM = ["npu"] _SUPPORT_IXFORMER_PLATFORM = ["ilu"] if platform in _SUPPORT_IXFORMER_PLATFORM: @@ -27,7 +28,16 @@ logger.warning("Skipping ixformer backend (import failed): %s", e) if platform in _SUPPORT_TTX_PLATFROM: - from .ttx import * + try: + from .ttx import * + except Exception as e: + logger.warning("Skipping ttx backend (import failed): %s", e) if platform in _SUPPORT_TORCH_NPU_PLATFROM: from .torch_npu import * + +if platform in _SUPPORT_UC_PLATFORM: + try: + from .uc import * + except ImportError as e: + logger.warning("Skipping uc backend (import failed): %s", e) diff --git a/mojo_opset/backends/uc/__init__.py b/mojo_opset/backends/uc/__init__.py new file mode 100644 index 000000000..060fa70b9 --- /dev/null +++ b/mojo_opset/backends/uc/__init__.py @@ -0,0 +1,5 @@ +from mojo_opset.utils.platform import get_impl_by_platform + +_op_map = get_impl_by_platform() +globals().update(_op_map) +__all__ = list(_op_map.keys()) diff --git a/mojo_opset/backends/uc/operators/__init__.py b/mojo_opset/backends/uc/operators/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/mojo_opset/backends/uc/operators/__init__.py @@ -0,0 +1 @@ + diff --git a/mojo_opset/backends/uc/operators/_utils.py b/mojo_opset/backends/uc/operators/_utils.py new file mode 100644 index 000000000..dd6019c22 --- /dev/null +++ b/mojo_opset/backends/uc/operators/_utils.py @@ -0,0 +1,68 @@ +from functools import lru_cache + +import torch + + +_DTYPE_API_SUFFIX = { + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.float32: "fp32", +} + + +@lru_cache(maxsize=1) +def _uc_kernels(): + import uc_kernel + + return uc_kernel.load() + + +def _matrix_shape(tensor: torch.Tensor) -> tuple[int, int]: + if tensor.dim() == 0: + return 1, 1 + if tensor.dim() == 1: + return 1, tensor.numel() + return tensor.numel() // tensor.shape[-1], tensor.shape[-1] + + +def _typed_api(api: str, dtype: torch.dtype) -> str: + suffix = _DTYPE_API_SUFFIX.get(dtype) + if suffix is None: + raise NotImplementedError(f"UC backend {api} does not support dtype {dtype}.") + + kernels = _uc_kernels() + typed_api = f"{api}_{suffix}" + if typed_api in kernels.keys(): + return typed_api + raise NotImplementedError(f"UC backend {api} does not provide a {suffix} kernel artifact.") + + +def run_unary_kernel(api: str, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + return torch.empty_like(x) + + kernel_input = x.contiguous() + kernel_output = torch.empty_like(kernel_input) + rows, cols = _matrix_shape(kernel_input) + _uc_kernels()[_typed_api(api, kernel_input.dtype)](kernel_input, kernel_output, rows, cols) + return kernel_output.reshape(x.shape) + + +def run_binary_kernel(api: str, lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor: + if lhs.shape != rhs.shape: + raise ValueError(f"UC backend {api} expects matching input shapes, got {lhs.shape} and {rhs.shape}.") + if lhs.dtype != rhs.dtype: + raise ValueError(f"UC backend {api} expects matching input dtypes, got {lhs.dtype} and {rhs.dtype}.") + if lhs.numel() == 0: + return torch.empty_like(lhs) + + kernel_lhs = lhs.contiguous() + kernel_rhs = rhs.contiguous() + kernel_output = torch.empty_like(kernel_lhs) + rows, cols = _matrix_shape(kernel_lhs) + _uc_kernels()[_typed_api(api, kernel_lhs.dtype)](kernel_lhs, kernel_rhs, kernel_output, rows, cols) + return kernel_output.reshape(lhs.shape) + + +def run_kernel(api: str, dtype: torch.dtype, *args) -> None: + _uc_kernels()[_typed_api(api, dtype)](*args) diff --git a/mojo_opset/backends/uc/operators/activation.py b/mojo_opset/backends/uc/operators/activation.py new file mode 100644 index 000000000..0aa78a816 --- /dev/null +++ b/mojo_opset/backends/uc/operators/activation.py @@ -0,0 +1,31 @@ +import torch + +from mojo_opset.core import MojoGelu +from mojo_opset.core import MojoSilu +from mojo_opset.core import MojoSwiGLU + +from ._utils import run_binary_kernel +from ._utils import run_unary_kernel + + +class UCGelu(MojoGelu): + supported_platforms_list = ["npu"] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return run_unary_kernel("mojo_gelu", x) + + +class UCSilu(MojoSilu): + supported_platforms_list = ["npu"] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return run_unary_kernel("mojo_silu", x) + + +class UCSwiGLU(MojoSwiGLU): + supported_platforms_list = ["npu"] + + def forward(self, gate_out: torch.Tensor, up_out: torch.Tensor) -> torch.Tensor: + if self.swiglu_limit > 0: + return super().forward(gate_out, up_out) + return run_binary_kernel("mojo_swiglu", gate_out, up_out) diff --git a/mojo_opset/backends/uc/operators/attention.py b/mojo_opset/backends/uc/operators/attention.py new file mode 100644 index 000000000..e78fe9fa2 --- /dev/null +++ b/mojo_opset/backends/uc/operators/attention.py @@ -0,0 +1,104 @@ +import math +from typing import Optional + +import torch + +from mojo_opset.core import MojoSdpa +from mojo_opset.utils.logging import get_logger + +from ._utils import _typed_api +from ._utils import _uc_kernels + + +logger = get_logger(__name__) + +_SUPPORTED_DTYPES = (torch.bfloat16,) +_KERNEL_BY_SHAPE = { + (1, 5, 1, 4096, 128): "mojo_sdpa_b1_qh5_kvh1_s4096_d128", +} + + +def _is_default_scale(scale: Optional[float], head_dim: int) -> bool: + if scale is None: + return True + default_scale = 1.0 / math.sqrt(head_dim) + return math.isclose(float(scale), default_scale, rel_tol=1e-6, abs_tol=1e-12) + + +def _assert_sdpa_contract( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + enable_gqa: bool, +) -> None: + assert query.dim() == 4 and key.dim() == 4 and value.dim() == 4 + assert query.shape[0] == key.shape[0] == value.shape[0] + assert query.shape[-1] == key.shape[-1] and key.shape[-1] == value.shape[-1] + head_dim = query.shape[-1] + assert head_dim in {64, 128} + assert query.shape[-2] == key.shape[-2] and key.shape[-2] == value.shape[-2] + + if not enable_gqa: + assert query.shape[1] == key.shape[1] and key.shape[1] == value.shape[1] + else: + assert key.shape[1] == value.shape[1] and query.shape[1] % key.shape[1] == 0 + + assert query.dtype == key.dtype == value.dtype + + +def _assert_sdpa_mask(attn_mask: torch.Tensor, seq_length: int) -> None: + assert len(attn_mask.shape) == 2 and attn_mask.shape[0] == seq_length and attn_mask.shape[1] == seq_length + assert attn_mask.dtype == torch.bool + + +def _assert_uc_static_kernel_contract(query: torch.Tensor, scale: Optional[float]) -> None: + head_dim = query.shape[-1] + assert query.dtype in _SUPPORTED_DTYPES + assert _is_default_scale(scale, head_dim) + + +def _static_kernel_api(query: torch.Tensor, key: torch.Tensor) -> str: + batch, q_heads, q_seq, head_dim = query.shape + kv_heads = key.shape[1] + api = _KERNEL_BY_SHAPE.get((batch, q_heads, kv_heads, q_seq, head_dim)) + if api is None: + raise NotImplementedError( + "UC SDPA only supports static shape " + "(batch=1, q_heads=5, kv_heads=1, seq=4096, head_dim=128) for now, " + f"got {(batch, q_heads, kv_heads, q_seq, head_dim)}." + ) + return api + + +class UCSdpa(MojoSdpa): + supported_platforms_list = ["npu"] + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + ): + if query.numel() == 0: + return torch.empty_like(query) + + _assert_sdpa_contract(query, key, value, self.enable_gqa) + if attn_mask is not None: + _assert_sdpa_mask(attn_mask, query.shape[-2]) + logger.warning_once( + "UC SDPA does not support attention masks yet; falling back to torch implementation." + ) + return super().forward(query, key, value, attn_mask) + + _assert_uc_static_kernel_contract(query, self.scale) + api = _static_kernel_api(query, key) + + q = query.contiguous() + k = key.contiguous() + v = value.contiguous() + out = torch.empty_like(q) + + kernel = _uc_kernels()[_typed_api(api, q.dtype)] + kernel(q, k, v, out) + return out.reshape(query.shape) diff --git a/mojo_opset/backends/uc/operators/gemm.py b/mojo_opset/backends/uc/operators/gemm.py new file mode 100644 index 000000000..78dbe0ab6 --- /dev/null +++ b/mojo_opset/backends/uc/operators/gemm.py @@ -0,0 +1,60 @@ +import torch + +from mojo_opset.core import MojoQuantGemm + +from ._utils import _uc_kernels + + +_OUTPUT_DTYPE_SUFFIX = { + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.float32: "fp32", +} + + +def _require_kernel(api: str): + kernels = _uc_kernels() + if api not in kernels.keys(): + raise NotImplementedError( + f"UC backend kernel {api!r} is not available. Rebuild uc-kernel after adding the TileLang source." + ) + return kernels[api] + + +class UCQuantGemm(MojoQuantGemm): + supported_platforms_list = ["npu"] + + def forward(self, input: torch.Tensor, input_scale: torch.Tensor) -> torch.Tensor: + if input.dim() != 2: + raise ValueError(f"input must be 2D, got shape {tuple(input.shape)}.") + if input.dtype != torch.int8: + raise NotImplementedError(f"UC QuantGemm supports int8 input, got {input.dtype}.") + if self.trans_weight: + weight = self.weight.t().contiguous() + else: + weight = self.weight + input_scale = input_scale.flatten().float().contiguous() + weight_scale = self.weight_scale.flatten().float().contiguous() + if not input.is_contiguous(): + input = input.contiguous() + if not weight.is_contiguous(): + weight = weight.contiguous() + M, K = input.shape + K_w, N = weight.shape + if K_w != K: + raise ValueError(f"input K {K} must match weight K {K_w}.") + if input_scale.numel() != M: + raise ValueError(f"input_scale length {input_scale.numel()} must equal M {M}.") + if weight_scale.numel() != N: + raise ValueError(f"weight_scale length {weight_scale.numel()} must equal N {N}.") + + output = torch.empty((M, N), device=input.device, dtype=self.output_dtype) + if output.numel() == 0: + return output + suffix = _OUTPUT_DTYPE_SUFFIX.get(self.output_dtype) + if suffix is None: + raise NotImplementedError(f"UC QuantGemm does not support output dtype {self.output_dtype}.") + + kernel = _require_kernel(f"mojo_quant_gemm_{suffix}") + kernel(input, weight, input_scale, weight_scale, output, M, K, N) + return output diff --git a/mojo_opset/backends/uc/operators/normalization.py b/mojo_opset/backends/uc/operators/normalization.py new file mode 100644 index 000000000..2f11dc54c --- /dev/null +++ b/mojo_opset/backends/uc/operators/normalization.py @@ -0,0 +1,159 @@ +import torch + +from mojo_opset.core import MojoLayerNorm +from mojo_opset.core import MojoResidualAddLayerNorm +from mojo_opset.core import MojoResidualAddRMSNorm +from mojo_opset.core import MojoRMSNorm + +from ._utils import _matrix_shape +from ._utils import run_kernel + + +class UCRMSNorm(MojoRMSNorm): + supported_platforms_list = ["npu"] + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + if hidden_state.numel() == 0: + return torch.empty_like(hidden_state) + + kernel_input = hidden_state.contiguous() + rows, cols = _matrix_shape(kernel_input) + kernel_output = torch.empty_like(kernel_input) + eps = float(self.variance_epsilon) + + run_kernel( + "mojo_rmsnorm", + kernel_input.dtype, + kernel_input, + self.weight.contiguous(), + kernel_output, + rows, + cols, + eps, + ) + return kernel_output.reshape(hidden_state.shape) + + +class UCResidualAddRMSNorm(MojoResidualAddRMSNorm): + supported_platforms_list = ["npu"] + + def forward(self, hidden_state: torch.Tensor, residual: torch.Tensor = None): + if residual is None: + raise ValueError("UC backend MojoResidualAddRMSNorm requires residual.") + if hidden_state.shape != residual.shape: + raise ValueError( + f"UC backend MojoResidualAddRMSNorm expects matching shapes, " + f"got {hidden_state.shape} and {residual.shape}." + ) + if hidden_state.dtype != residual.dtype: + raise ValueError( + f"UC backend MojoResidualAddRMSNorm expects matching dtypes, " + f"got {hidden_state.dtype} and {residual.dtype}." + ) + if hidden_state.numel() == 0: + empty = torch.empty_like(hidden_state) + return empty, empty + + kernel_input = hidden_state.contiguous() + kernel_residual = residual.contiguous() + rows, cols = _matrix_shape(kernel_input) + kernel_output = torch.empty_like(kernel_input) + kernel_residual_output = torch.empty_like(kernel_input) + eps = float(self.variance_epsilon) + + run_kernel( + "mojo_residual_add_rmsnorm", + kernel_input.dtype, + kernel_input, + kernel_residual, + self.weight.contiguous(), + kernel_output, + kernel_residual_output, + rows, + cols, + eps, + ) + + output = kernel_output.reshape(hidden_state.shape) + updated_residual = kernel_residual_output.reshape(hidden_state.shape) + if self.norm_pos == "pre": + return output, updated_residual + return output, output + + +class UCResidualAddLayerNorm(MojoResidualAddLayerNorm): + supported_platforms_list = ["npu"] + + def forward(self, hidden_state: torch.Tensor, residual: torch.Tensor = None): + if residual is None: + raise ValueError("UC backend MojoResidualAddLayerNorm requires residual.") + if hidden_state.shape != residual.shape: + raise ValueError( + f"UC backend MojoResidualAddLayerNorm expects matching shapes, " + f"got {hidden_state.shape} and {residual.shape}." + ) + if hidden_state.dtype != residual.dtype: + raise ValueError( + f"UC backend MojoResidualAddLayerNorm expects matching dtypes, " + f"got {hidden_state.dtype} and {residual.dtype}." + ) + if self.weight is None or self.bias is None: + raise NotImplementedError("UC backend mojo_residual_add_layernorm requires weight and bias.") + if hidden_state.numel() == 0: + empty = torch.empty_like(hidden_state) + return empty, empty + + kernel_input = hidden_state.contiguous() + kernel_residual = residual.contiguous() + rows, cols = _matrix_shape(kernel_input) + kernel_output = torch.empty_like(kernel_input) + kernel_residual_output = torch.empty_like(kernel_input) + eps = float(self.variance_epsilon) + + run_kernel( + "mojo_residual_add_layernorm", + kernel_input.dtype, + kernel_input, + kernel_residual, + self.weight.contiguous(), + self.bias.contiguous(), + kernel_output, + kernel_residual_output, + rows, + cols, + eps, + ) + + output = kernel_output.reshape(hidden_state.shape) + updated_residual = kernel_residual_output.reshape(hidden_state.shape) + if self.norm_pos == "pre": + return output, updated_residual + return output, output + + +class UCLayerNorm(MojoLayerNorm): + supported_platforms_list = ["npu"] + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + if hidden_state.numel() == 0: + return torch.empty_like(hidden_state) + if self.weight is None or self.bias is None: + raise NotImplementedError("UC backend mojo_layernorm requires weight and bias.") + + kernel_input = hidden_state.contiguous() + rows, cols = _matrix_shape(kernel_input) + kernel_output = torch.empty_like(kernel_input) + eps = float(self.variance_epsilon) + + run_kernel( + "mojo_layernorm", + kernel_input.dtype, + kernel_input, + self.weight.contiguous(), + self.bias.contiguous(), + kernel_output, + rows, + cols, + eps, + ) + return kernel_output.reshape(hidden_state.shape) diff --git a/mojo_opset/backends/uc/operators/position_embedding.py b/mojo_opset/backends/uc/operators/position_embedding.py new file mode 100644 index 000000000..4faf6970c --- /dev/null +++ b/mojo_opset/backends/uc/operators/position_embedding.py @@ -0,0 +1,237 @@ +from typing import Optional +from typing import Tuple + +import torch + +from mojo_opset.core import MojoApplyRoPE +from mojo_opset.core import MojoRotaryEmbedding +from mojo_opset.utils.logging import get_logger + +from ._utils import run_kernel + + +logger = get_logger(__name__) + + +class UCRotaryEmbedding(MojoRotaryEmbedding): + supported_platforms_list = ["npu"] + + def __init__( + self, + rope_theta, + rope_dim, + attention_scaling: float = 1.0, + init_max_length: Optional[int] = None, + **kwargs, + ): + super().__init__(rope_theta, rope_dim, attention_scaling, init_max_length, **kwargs) + if init_max_length is None: + raise ValueError("init_max_length must be provided for UCRotaryEmbedding") + + def forward( + self, + x: torch.Tensor, + cu_q_lens: Optional[torch.Tensor] = None, + total_seq_lens: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cu_q_lens is not None: + assert cu_q_lens.dtype == torch.int32 + if total_seq_lens is not None: + assert total_seq_lens.dtype == torch.int32 + if position_ids is not None: + assert position_ids.dtype == torch.int32 + assert position_ids is None or cu_q_lens is None, "At most one of cu_q_lens or position_ids should be provided" + + if cu_q_lens is not None: + logger.warning_once( + "UC rotary embedding does not support varlen mode yet; falling back to torch implementation." + ) + return super().forward(x, cu_q_lens, total_seq_lens, position_ids) + elif position_ids is not None: + assert position_ids.shape == x.shape[:-1], "position_ids must have the same shape as x except the hidden dimension" + position_ids = position_ids.contiguous() + else: + return self._arange_cache(x) + + return self._position_ids_cache(position_ids) + + def _arange_cache(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + if x.dim() < 2: + raise AssertionError("x must have at least two dimensions for padded prefill rotary embedding") + + seq_len = x.shape[1] + rope_dim = self.cos.shape[-1] + if seq_len > self.cos.shape[0]: + raise ValueError(f"seq_len {seq_len} exceeds rotary cache length {self.cos.shape[0]}") + + cos_out = torch.empty((seq_len, rope_dim), device=self.cos.device, dtype=self.cos.dtype) + sin_out = torch.empty((seq_len, rope_dim), device=self.sin.device, dtype=self.sin.dtype) + if seq_len == 0 or rope_dim == 0: + return cos_out, sin_out + + run_kernel( + "mojo_rotary_embedding_arange", + self.cos.dtype, + self.cos.contiguous(), + self.sin.contiguous(), + cos_out, + sin_out, + self.cos.shape[0], + rope_dim, + seq_len, + ) + return cos_out, sin_out + + def _position_ids_cache(self, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + rope_dim = self.cos.shape[-1] + rows = position_ids.numel() + output_shape = tuple(position_ids.shape) + (rope_dim,) + cos_out = torch.empty(output_shape, device=self.cos.device, dtype=self.cos.dtype) + sin_out = torch.empty(output_shape, device=self.sin.device, dtype=self.sin.dtype) + if rows == 0 or rope_dim == 0: + return cos_out, sin_out + + run_kernel( + "mojo_rotary_embedding_position_ids", + self.cos.dtype, + self.cos.contiguous(), + self.sin.contiguous(), + position_ids.reshape(-1).contiguous(), + cos_out.reshape(rows, rope_dim), + sin_out.reshape(rows, rope_dim), + self.cos.shape[0], + rope_dim, + rows, + ) + return cos_out, sin_out + + +class UCApplyRoPE(MojoApplyRoPE): + supported_platforms_list = ["npu"] + _STATIC_APPLY_ROPE_KERNELS = frozenset( + { + (96, 96, torch.float16), + (96, 32, torch.bfloat16), + (88, 88, torch.bfloat16), + (128, 48, torch.float16), + (128, 128, torch.float16), + } + ) + + @staticmethod + def _normalize_to_token_head( + q: torch.Tensor, + k: torch.Tensor, + head_first: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[int], Optional[int]]: + if q.ndim == 3: + if head_first: + q = q.transpose(0, 1).contiguous() + k = k.transpose(0, 1).contiguous() + else: + q = q.contiguous() + k = k.contiguous() + return q, k, None, None + + if head_first: + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + else: + q = q.contiguous() + k = k.contiguous() + + batch_size, seq_len = q.shape[0], q.shape[1] + q = q.reshape(batch_size * seq_len, q.shape[2], q.shape[3]).contiguous() + k = k.reshape(k.shape[0] * k.shape[1], k.shape[2], k.shape[3]).contiguous() + return q, k, batch_size, seq_len + + @staticmethod + def _restore_from_token_head( + x: torch.Tensor, + original_shape: torch.Size, + head_first: bool, + batch_size: Optional[int], + seq_len: Optional[int], + ) -> torch.Tensor: + if len(original_shape) == 3: + if head_first: + x = x.transpose(0, 1).contiguous() + return x.reshape(original_shape) + + x = x.reshape(batch_size, seq_len, x.shape[1], x.shape[2]) + if head_first: + x = x.transpose(1, 2).contiguous() + return x.reshape(original_shape) + + @staticmethod + def _run_static_token_head_kernel( + api: str, + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + *shape_args: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + q_out = q.clone(memory_format=torch.contiguous_format) + k_out = k.clone(memory_format=torch.contiguous_format) + run_kernel(api, q.dtype, q, k, cos, sin, q_out, k_out, *shape_args) + return q_out, k_out + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + head_first: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert q.ndim == k.ndim, "q and k must have the same dimension" + assert q.ndim == 3 or q.ndim == 4, "q and k must be 3D or 4D" + assert cos.shape == sin.shape, "cos and sin must have the same shape" + if q.ndim == 3: + assert cos.ndim == 2, "3D q/k inputs expect 2D cos/sin" + elif cos.ndim not in (2, 3): + raise ValueError("4D q/k inputs expect 2D or 3D cos/sin") + + if q.dtype != k.dtype: + raise ValueError(f"q and k must have the same dtype, got {q.dtype} and {k.dtype}.") + if cos.dtype != torch.float32 or sin.dtype != torch.float32: + raise NotImplementedError("UC backend MojoApplyRoPE expects float32 cos/sin tensors.") + + if q.numel() == 0 or k.numel() == 0: + return torch.empty_like(q), torch.empty_like(k) + + rope_dim = cos.shape[-1] + q_norm, k_norm, batch_size, seq_len = self._normalize_to_token_head(q, k, head_first) + rows, q_heads, head_dim = q_norm.shape + k_rows, k_heads, k_head_dim = k_norm.shape + if rows != k_rows or head_dim != k_head_dim: + raise ValueError("q and k must have matching token count and head dimension") + + config_key = (head_dim, rope_dim, q_norm.dtype) + if config_key not in self._STATIC_APPLY_ROPE_KERNELS: + raise NotImplementedError( + "UC backend MojoApplyRoPE does not provide an aligned static kernel for " + f"head_dim={head_dim}, rope_dim={rope_dim}, dtype={q_norm.dtype}." + ) + + if cos.ndim == 2: + cos_kind = "cos2d" + cos_kernel = cos.contiguous() + sin_kernel = sin.contiguous() + shape_args = (rows, q_heads, k_heads, cos.shape[0]) + else: + cos_kind = "costoken" + cos_kernel = cos.reshape(rows, rope_dim).contiguous() + sin_kernel = sin.reshape(rows, rope_dim).contiguous() + shape_args = (rows, q_heads, k_heads) + + api = f"mojo_apply_rope_tnh_d{head_dim}_r{rope_dim}_{cos_kind}" + q_out, k_out = self._run_static_token_head_kernel( + api, q_norm, k_norm, cos_kernel, sin_kernel, *shape_args + ) + return ( + self._restore_from_token_head(q_out, q.shape, head_first, batch_size, seq_len), + self._restore_from_token_head(k_out, k.shape, head_first, batch_size, seq_len), + ) diff --git a/mojo_opset/backends/uc/operators/quant.py b/mojo_opset/backends/uc/operators/quant.py new file mode 100644 index 000000000..596288d94 --- /dev/null +++ b/mojo_opset/backends/uc/operators/quant.py @@ -0,0 +1,133 @@ +import torch + +from mojo_opset.core import MojoDynamicQuant +from mojo_opset.core import MojoStaticQuant +from mojo_opset.core import MojoMoEDynamicQuant + +from ._utils import _matrix_shape +from ._utils import run_kernel + + +class UCStaticQuant(MojoStaticQuant): + supported_platforms_list = ["npu"] + + def forward(self, input: torch.Tensor): + if self.quant_dtype != torch.int8: + raise NotImplementedError(f"UCStaticQuant only supports torch.int8, got {self.quant_dtype}.") + if input.dim() < len(self.input_size): + raise ValueError( + f"input must have at least {len(self.input_size)} dims for scale shape " + f"{self.input_size}, got {tuple(input.shape)}." + ) + if tuple(input.shape[-len(self.input_size):]) != self.input_size: + raise ValueError( + f"input trailing dims {tuple(input.shape[-len(self.input_size):])} must " + f"match scale shape {self.input_size}." + ) + if input.numel() == 0: + return torch.empty_like(input, dtype=self.quant_dtype), self.scale + + kernel_input = input.contiguous() + scale = self.scale.to(device=kernel_input.device, dtype=torch.float32).contiguous() + cols = scale.numel() + rows = kernel_input.numel() // cols + kernel_input_2d = kernel_input.reshape(rows, cols) + scale_1d = scale.reshape(cols) + kernel_output = torch.empty_like(kernel_input_2d, dtype=self.quant_dtype) + + run_kernel( + "mojo_static_quant", + kernel_input.dtype, + kernel_input_2d, + scale_1d, + kernel_output, + rows, + cols, + ) + return kernel_output.reshape(input.shape), self.scale + + +class UCDynamicQuant(MojoDynamicQuant): + supported_platforms_list = ["npu"] + + def forward(self, input: torch.Tensor): + if input.dim() < 1: + raise ValueError("input must have at least one dimension.") + if input.numel() == 0: + output = torch.empty_like(input, dtype=torch.int8) + scale = torch.empty((*input.shape[:-1], 1), dtype=torch.float32, device=input.device) + return output, scale + + kernel_input = input.contiguous() + rows, cols = _matrix_shape(kernel_input) + kernel_input_2d = kernel_input.reshape(rows, cols) + kernel_output = torch.empty_like(kernel_input_2d, dtype=torch.int8) + kernel_scale = torch.empty((rows,), dtype=torch.float32, device=kernel_input.device) + if self.inv_smooth_scale is None: + inv_smooth_scale = torch.ones((cols,), dtype=torch.float32, device=kernel_input.device) + else: + inv_smooth_scale = self.inv_smooth_scale.to(device=kernel_input.device, dtype=torch.float32).contiguous() + + run_kernel( + "mojo_dynamic_quant", + kernel_input.dtype, + kernel_input_2d, + inv_smooth_scale, + kernel_output, + kernel_scale, + rows, + cols, + ) + return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) + + +class UCMoEDynamicQuant(MojoMoEDynamicQuant): + supported_platforms_list = ["npu"] + + def forward(self, input: torch.Tensor, token_count: torch.Tensor): + if input.dim() < 2: + raise ValueError(f"input must have at least 2 dimensions for MoE dynamic quant, got {input.dim()}.") + if token_count.dim() != 1: + raise ValueError(f"token_count must be 1D, got shape {tuple(token_count.shape)}.") + if token_count.dtype not in (torch.int32, torch.int64): + raise TypeError(f"token_count must be int32 or int64, got {token_count.dtype}.") + if token_count.numel() != self.expert_num: + raise ValueError(f"token_count length must equal expert_num {self.expert_num}, got {token_count.numel()}.") + if bool(torch.any(token_count < 0).item()): + raise ValueError("token_count must be non-negative.") + + kernel_input = input.contiguous() + rows, cols = _matrix_shape(kernel_input) + if cols != self.input_size: + raise ValueError(f"input last dimension must equal input_size {self.input_size}, got {cols}.") + + token_total = int(token_count.sum().item()) + if token_total != rows: + raise ValueError(f"token_count sum must equal flattened row count {rows}, got {token_total}.") + + if input.numel() == 0: + output = torch.empty_like(input, dtype=torch.int8) + scale = torch.empty((*input.shape[:-1], 1), dtype=torch.float32, device=input.device) + return output, scale + + inv_smooth_scale = self.inv_smooth_scale.to(device=kernel_input.device, dtype=torch.float32).contiguous() + token_count_device = token_count.to(device=kernel_input.device, dtype=torch.int64) + expanded_inv_smooth_scale = inv_smooth_scale.repeat_interleave(token_count_device, dim=0).contiguous() + scaled_input = (kernel_input.float().reshape(rows, cols) * expanded_inv_smooth_scale).to(kernel_input.dtype) + + kernel_input_2d = scaled_input.contiguous() + kernel_output = torch.empty_like(kernel_input_2d, dtype=torch.int8) + kernel_scale = torch.empty((rows,), dtype=torch.float32, device=kernel_input.device) + dynamic_inv_smooth_scale = torch.ones((cols,), dtype=torch.float32, device=kernel_input.device) + + run_kernel( + "mojo_dynamic_quant", + kernel_input_2d.dtype, + kernel_input_2d, + dynamic_inv_smooth_scale, + kernel_output, + kernel_scale, + rows, + cols, + ) + return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) diff --git a/mojo_opset/core/backend_registry.py b/mojo_opset/core/backend_registry.py index 6f4627540..0ef31c6e2 100644 --- a/mojo_opset/core/backend_registry.py +++ b/mojo_opset/core/backend_registry.py @@ -11,7 +11,7 @@ logger = get_logger(__name__) PLATFORM_BACKEND_PRIORITY = { - "npu": ["ttx", "torch_npu", "torch", "xops"], + "npu": ["ttx", "torch_npu", "torch", "xops", "uc"], "ilu": ["ixformer", "ttx", "torch"], "mlu": ["ttx", "torch"], "meta_device": ["torch"],