refactor: share CUDA-style extension kernels across backends

CMakeLists.txt

@@ -345,8 +345,7 @@ if (SD_CUDA)
     set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA)
     if (NOT SD_USE_UPSTREAM_GGML)
         enable_language(CUDA)
-        target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu src/core/cuda/affine.cu)
-        set_source_files_properties(src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu src/core/cuda/affine.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+        set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
         if (GGML_CUDA_CUTLASS_PATH)
             target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu)
             set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu PROPERTIES
@@ -358,17 +357,28 @@ if (SD_CUDA)
             if (GGML_CUDA_FORCE_CUBLAS)
                 set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS)
             endif()
-            set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
-            set_property(SOURCE src/core/cuda/int8.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
         endif()
-        set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
-        target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
-        set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_EXTENSIONS)
     endif()
 endif()
+set(SD_BACKEND_REGISTRATIONS)
+if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
+    add_subdirectory(src/core/cuda)
+    foreach(VENDOR CUDA HIP MUSA)
+        if(SD_${VENDOR} OR (VENDOR STREQUAL "HIP" AND SD_HIPBLAS))
+            string(TOLOWER "${VENDOR}" NAMESPACE)
+            list(APPEND SD_BACKEND_REGISTRATIONS "sd_${NAMESPACE}::register_extensions")
+        endif()
+    endforeach()
+endif()
 if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
     add_subdirectory(src/core/vulkan)
-    set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_VULKAN_EXTENSIONS)
+    list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope)
+endif()
+if(SD_BACKEND_REGISTRATIONS)
+    string(REPLACE ";" ",\n    " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")
+    configure_file(cmake/backend_extensions.h.in generated/sd_backend_extensions.h @ONLY)
+    set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY INCLUDE_DIRECTORIES "${CMAKE_CURRENT_BINARY_DIR}/generated")
+    set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_BACKEND_EXTENSIONS)
 endif()
 if (SD_HIPBLAS)
     find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm)

cmake/backend_extensions.h.in

@@ -0,0 +1,12 @@
+#ifndef __SD_BACKEND_EXTENSIONS_H__
+#define __SD_BACKEND_EXTENSIONS_H__
+
+#include "core/cuda/registration.h"
+
+void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
+
+static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = {
+    @SD_BACKEND_REGISTRATIONS@
+};
+
+#endif

docs/ggml_extensions.md

@@ -112,14 +112,20 @@ Services are discovered through `ggml_backend_reg_get_proc_address`:
 | `ggml_ext_cuda_context_v1` | Borrow the current CUDA device and stream, ordered with native work. |
 | `ggml_ext_cuda_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. |
 | `ggml_ext_cuda_blas_v1` | Borrow the current stream's cuBLAS handle during execution, with host scalar pointers. |
+| `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, and INT8 BLAS capability. |
+| `ggml_ext_gpu_context_v1` | Borrow the matching runtime's device and execution stream. |
+| `ggml_ext_gpu_blas_v1` | Borrow the matching cuBLAS, hipBLAS, or muBLAS handle on that stream. |
 
 The borrowed cuBLAS handle remains owned by the backend. Extensions must not retain
 or destroy it, or change its stream, pointer mode, math mode, or workspace. This
-service is exposed by the CUDA backend, not HIP/MUSA.
+CUDA-specific BLAS service is retained for existing callers; the runtime-aware
+services support CUDA, HIP, and MUSA with the same ownership rules. Extensions
+call vendor runtime and BLAS APIs directly. These services do not execute GEMM,
+quantization, or model fusion on the application's behalf.
 
 Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers
-must use the representation appropriate to their backend. A Vulkan command
-recording service is not implemented yet.
+must use the representation appropriate to their backend. Vulkan extensions use
+the command recording services described below.
 
 ## sd.cpp integration
 
@@ -127,8 +133,8 @@ recording service is not implemented yet.
 `sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a
 `RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
 vector. Both interleaved and split-half RoPE, partial rotation, and strided input
-rows are supported. Model construction uses the custom operator on CPU and CUDA
-and retains existing backend implementations elsewhere.
+rows are supported. Model construction uses the custom operator on CPU, CUDA,
+HIP, MUSA, and Vulkan, and retains existing backend implementations elsewhere.
 
 `src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override,
 and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
@@ -136,8 +142,16 @@ its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
 position-encoding strides are supported by the standalone kernel; RMSNorm fusion
 requires contiguous position encoding and weights. Copy fusion uses workspace
 when the destination overlaps an input. These registrations do not require
-CUTLASS. The old CUDA RoPE dispatch is removed; HIP and MUSA retain their existing
-implementations.
+CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed.
+
+Shared CUDA-style kernels live in `src/core/cuda/` and are compiled separately
+for each enabled runtime. `SD_CUDA_NAMESPACE` expands to `sd_cuda`, `sd_hip`, or
+`sd_musa`, so implementations can coexist. `runtime.h` handles the small runtime
+differences, including logical 32-lane shuffles on HIP. The application defines
+and launches its kernels directly; it includes only public ggml headers.
+CMake generates the registration table from enabled backend options. The shared
+entry point uses one `SD_BACKEND_EXTENSIONS` macro; backend registrars verify the
+device kind, and disabled backends add no symbol references.
 
 `SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
 API is unchanged. INT8 ConvRot retains its existing packed activation layout and
@@ -216,10 +230,12 @@ Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to
 preserve their previous accumulation behavior. Fusion disable and force-cuBLAS
 settings are honored. No CUTLASS headers are needed for the cuBLAS path.
 
-The native CUDA CUTLASS GEMM implementation and its SwiGLU/row-concat matchers
-have been removed from ggml. Native cuBLAS GEMM and standalone ConvRot remain
-compatibility implementations for callers without these registrations. HIP
-retains its native INT8/activation fusion paths.
+The native CUDA/HIP INT8 GEMM, standalone ConvRot, and associated fusion paths
+have been removed from ggml. The shared application implementation calls cuBLAS
+or hipBLAS directly. HIP retains the existing architecture eligibility policy;
+MUSA does not advertise INT8 BLAS support and uses scheduler fallback for GEMM.
+Standalone ConvRot kernels are compiled for all three runtimes. CUTLASS
+implementations remain CUDA-specific.
 
 `src/core/cuda/i8_gated_residual.cu` implements
 `INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered
@@ -255,8 +271,8 @@ overlapping outputs. Long rows or external intermediate consumers retain
 ordinary execution. These fusions do not require CUTLASS and honor
 `GGML_CUDA_DISABLE_FUSION`.
 
-The corresponding scale, normalization, and unary-activation fusion rules in
-ggml now only serve HIP. CUDA GEMM/SwiGLU/ConvRot chains are selected by the
+The corresponding scale, normalization, and unary-activation fusion rules have
+been removed from ggml for CUDA and HIP. GEMM/SwiGLU/ConvRot chains are selected by the
 application INT8 extension before the standalone activation fusion.
 
 `src/core/cuda/affine.cu` registers FP32 `Mul(gate) -> Add(residual)`,
@@ -268,8 +284,8 @@ Add inputs, and fusion disabling retain ordinary operators. Internal in-place
 Add outputs are accepted when they overwrite the complete fused Mul allocation. Exact in-place residual
 or input addresses are safe for the pointwise kernel; partial overlap, broadcast
 gate overlap, and FP16 copy overlap use scheduler workspace. These registrations
-require no CUTLASS and support reusable CUDA Graph replay. Their former native
-fusion implementations now only serve HIP/MUSA.
+require no CUTLASS, compile for CUDA/HIP/MUSA, and support reusable CUDA Graph
+replay. Their former native fusion implementations have been removed from ggml.
 
 CUDA kernels can opt into device replay with `ggml_ext_register_kernel_ex` and
 `GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its

src/core/compute_workspace.cpp

@@ -147,6 +147,9 @@ namespace sd {
                         append(usage);
                         append(type);
                     }
+                } else {
+                    const auto alignment_offset = reinterpret_cast<uintptr_t>(tensor->data) % ggml_backend_get_alignment(external);
+                    append(alignment_offset);
                 }
                 const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
                 const int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
@@ -187,7 +190,8 @@ namespace sd {
                     ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
                 }
                 copy->buffer = buffer;
-                copy->data   = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
+                // Capability checks inspect alignment even before weights are resident.
+                copy->data   = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
                 copy->extra  = nullptr;
             }
         }

src/core/cuda/CMakeLists.txt

@@ -0,0 +1,90 @@
+cmake_policy(SET CMP0079 NEW)
+
+function(sd_cuda_sources VENDOR LANGUAGE)
+    set(SOURCES)
+    foreach(OP rope int8 convrot affine registration)
+        set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
+        string(TOUPPER "${VENDOR}" UPPER_VENDOR)
+        file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
+        set_source_files_properties("${SOURCE}" PROPERTIES LANGUAGE ${LANGUAGE})
+        list(APPEND SOURCES "${SOURCE}")
+    endforeach()
+    set(TARGET "sd-${VENDOR}")
+    add_library(${TARGET} OBJECT ${SOURCES})
+    target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src")
+    target_link_libraries(${TARGET} PRIVATE ggml)
+    set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
+    target_compile_features(${TARGET} PRIVATE cxx_std_17)
+    target_link_libraries(${SD_LIB} PRIVATE ${TARGET})
+    set(SD_CUDA_TARGET ${TARGET} PARENT_SCOPE)
+    set(SD_CUDA_SOURCES ${SOURCES} PARENT_SCOPE)
+endfunction()
+
+if(SD_CUDA)
+    enable_language(CUDA)
+    sd_cuda_sources(cuda CUDA)
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+    if(GGML_CUDA_CUTLASS_PATH)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${PROJECT_SOURCE_DIR}/src/core/cuda")
+    endif()
+    set_target_properties(${SD_CUDA_TARGET} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
+    target_link_libraries(${SD_CUDA_TARGET} PRIVATE CUDA::cudart CUDA::cublas)
+    target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
+endif()
+
+if(SD_HIPBLAS)
+    list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}" /opt/rocm /usr)
+    find_package(hip REQUIRED)
+    find_package(hipblas REQUIRED)
+    if(WIN32 OR CMAKE_CXX_COMPILER MATCHES "hipcc")
+        sd_cuda_sources(hip CXX)
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::device)
+    else()
+        if(NOT CMAKE_HIP_ARCHITECTURES)
+            if(GPU_TARGETS)
+                set(CMAKE_HIP_ARCHITECTURES ${GPU_TARGETS})
+            elseif(AMDGPU_TARGETS)
+                set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS})
+            endif()
+        endif()
+        enable_language(HIP)
+        sd_cuda_sources(hip HIP)
+        set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
+    endif()
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
+    target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
+    target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
+endif()
+
+if(SD_MUSA)
+    if(NOT MUSA_PATH)
+        if(DEFINED ENV{MUSA_PATH})
+            set(MUSA_PATH "$ENV{MUSA_PATH}")
+        elseif(EXISTS /opt/musa)
+            set(MUSA_PATH /opt/musa)
+        else()
+            set(MUSA_PATH /usr/local/musa)
+        endif()
+    endif()
+    list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
+    find_package(MUSAToolkit REQUIRED)
+    set(CMAKE_CXX_COMPILER "${MUSA_PATH}/bin/clang++")
+    sd_cuda_sources(musa CXX)
+    set(FLAGS -Od3 -fno-strict-aliasing -ffast-math -fsigned-char -x musa -mtgpu -fmusa-flush-denormals-to-zero)
+    if(NOT DEFINED MUSA_ARCHITECTURES)
+        set(MUSA_ARCHITECTURES "22;31")
+    endif()
+    foreach(ARCH ${MUSA_ARCHITECTURES})
+        list(APPEND FLAGS "--cuda-gpu-arch=mp_${ARCH}")
+    endforeach()
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "${FLAGS}")
+    if(GGML_STATIC)
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart_static MUSA::mublas_static)
+        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart_static MUSA::mublas_static)
+    else()
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart MUSA::mublas)
+        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart MUSA::mublas)
+    endif()
+endif()

src/core/cuda/affine.cu

@@ -1,11 +1,11 @@
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
 #include <climits>
 #include <cstdlib>
 #include <cstring>
 #include "ggml-ext.h"
+#include "runtime.h"
 
-namespace sd_affine_cuda {
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_affine_cuda {
         static bool disabled() {
             static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
             return value;
@@ -126,12 +126,10 @@ namespace sd_affine_cuda {
         }
         static void destroy(void*, void*) {}
         static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
-        auto reg   = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
-        auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
-        ggml_ext_cuda_context context;
-        if (!query || !query(execution->backend, &context))
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context))
                 return GGML_STATUS_FAILED;
-        auto stream          = static_cast<cudaStream_t>(context.stream);
+            auto stream          = static_cast<CUDAStream>(context.stream);
             const auto* node     = nodes[0];
             const auto* a        = node->src[0];
             const auto* b        = node->op == GGML_OP_CONT ? nullptr : node->src[1];
@@ -156,14 +154,14 @@ namespace sd_affine_cuda {
                 scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
                                                           a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
             }
-        auto error = cudaGetLastError();
-        if (error == cudaSuccess && alias)
-            error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
-        return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+            auto error = cuda_last_error();
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
     }
-}
 
-void sd_ggml_register_cuda_affine(ggml_backend_dev_t device) {
+    void register_affine(ggml_backend_dev_t device) {
         ggml_ext_kernel_desc desc = {};
         desc.abi_version          = GGML_EXT_ABI_VERSION;
         desc.priority             = 200;
@@ -171,13 +169,15 @@ void sd_ggml_register_cuda_affine(ggml_backend_dev_t device) {
         desc.prepare              = sd_affine_cuda::prepare;
         desc.destroy              = sd_affine_cuda::destroy;
         desc.execute              = sd_affine_cuda::execute;
-    desc.name                 = "sd.gated_residual.cuda.v1";
+        desc.name                 = "sd.gated_residual." SD_CUDA_NAME ".v1";
         desc.op                   = GGML_OP_MUL;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name = "sd.concat_cast.cuda.v1";
+        desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
         desc.op   = GGML_OP_CONCAT;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name = "sd.cont_scale_permute_cast.cuda.v1";
+        desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
         desc.op   = GGML_OP_CONT;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
 }

src/core/cuda/convrot.cu

@@ -2,33 +2,14 @@
 #include <climits>
 #include <cstdlib>
 #include <cstring>
-#include <mutex>
 #include <string>
-#include <unordered_map>
 #include "convrot_kernel.h"
 #include "core/int8_ops.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_convrot {
 
-namespace sd_convrot {
-    struct DeviceInfo {
-        int ordinal   = -1;
-        size_t shared = 0;
-    };
-    static DeviceInfo device_info(ggml_backend_dev_t device) {
-        static std::mutex mutex;
-        static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
-        std::lock_guard<std::mutex> lock(mutex);
-        auto found = cache.find(device);
-        if (found != cache.end())
-            return found->second;
-        DeviceInfo info;
-        auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_cuda_device_v1"));
-        int ordinal          = query ? query(device) : -1;
-        cudaDeviceProp props = {};
-        if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
-            info = {ordinal, size_t(props.sharedMemPerBlockOptin)};
-        cache.emplace(device, info);
-        return info;
-    }
         static bool fusion_disabled() {
             static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
             return disabled;
@@ -127,10 +108,10 @@ namespace sd_convrot {
             }
             return true;
         }
-    static bool supported(const Operation& op, const DeviceInfo& device) {
+        static bool supported(const Operation& op, const ggml_ext_gpu_device_info& device) {
             const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256
                                                                             : 128;
-        return device.ordinal >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared;
+            return device.device >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared_memory;
         }
         static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
             if (fusion_disabled())
@@ -146,7 +127,7 @@ namespace sd_convrot {
             for (int j = 0; j < count; ++j)
                 nodes[j] = ggml_graph_node(g, first + j);
             Operation op;
-        return describe(nodes, count, op) && supported(op, device_info(device)) ? count : 0;
+            return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0;
         }
         static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
             if (fusion_disabled())
@@ -182,7 +163,7 @@ namespace sd_convrot {
                         }
                 }
                 Operation op;
-            if (!count || !describe(chain, count, op) || !supported(op, device_info(device)))
+                if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device)))
                     continue;
                 int begin = -1;
                 for (int j = std::max(0, i - 64); j < i; ++j)
@@ -207,14 +188,14 @@ namespace sd_convrot {
             }
         }
         struct Plan {
-        DeviceInfo device;
+            ggml_ext_gpu_device_info device;
             Operation op;
             bool configured = false;
         };
         static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
             auto* plan   = new Plan{};
             *state       = plan;
-        plan->device = device_info(ggml_backend_get_device(backend));
+            plan->device = cuda_device_info(ggml_backend_get_device(backend));
             if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
                 return GGML_STATUS_FAILED;
             *workspace = ggml_nbytes(nodes[count - 1]);
@@ -230,14 +211,14 @@ namespace sd_convrot {
             return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
         }
         template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
-    static cudaError_t launch(Plan& p, const Operation& op, int8_t* output, cudaStream_t stream) {
+        static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
             constexpr bool warp = Scaled || RMS;
             constexpr int extra = Normalize ? 384 : RMS ? 256
                                                         : 128;
             if (!p.configured) {
-            auto status = cudaFuncSetAttribute(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
-                                               cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - extra));
-            if (status != cudaSuccess)
+                auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
+                                                     int(p.device.shared_memory - extra));
+                if (status != cuda_success)
                     return status;
                 p.configured = true;
             }
@@ -251,26 +232,24 @@ namespace sd_convrot {
             quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
                 data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
                 op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
-        return cudaGetLastError();
+            return cuda_last_error();
         }
         static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
             auto& plan     = *static_cast<Plan*>(state);
             const auto& op = plan.op;
             auto* dst      = nodes[count - 1];
-        auto reg       = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
-        auto query     = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
-        ggml_ext_cuda_context context;
-        if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != plan.device.device)
                 return GGML_STATUS_FAILED;
-        auto stream = static_cast<cudaStream_t>(context.stream);
+            auto stream = static_cast<CUDAStream>(context.stream);
             bool alias  = false;
             for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
                 alias |= overlaps(dst, input);
             auto* output         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
             const int64_t padded = (op.rows + 3) / 4 * 4;
-        if (padded > op.rows && cudaMemsetAsync(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cudaSuccess)
+            if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success)
                 return GGML_STATUS_FAILED;
-        cudaError_t error;
+            CUDAError error;
             switch (op.kind) {
                 case SCALE:
                     error = launch<false, false, true>(plan, op, output, stream);
@@ -291,17 +270,17 @@ namespace sd_convrot {
                 default:
                     return GGML_STATUS_FAILED;
             }
-        if (error != cudaSuccess || (alias && cudaMemcpyAsync(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), cudaMemcpyDeviceToDevice, stream) != cudaSuccess))
+            if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success))
                 return GGML_STATUS_FAILED;
             return GGML_STATUS_SUCCESS;
         }
-}
+    }
 
-void sd_ggml_register_cuda_convrot(ggml_backend_dev_t device) {
+    void register_convrot(ggml_backend_dev_t device) {
         for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
             ggml_ext_kernel_desc desc = {};
             desc.abi_version          = GGML_EXT_ABI_VERSION;
-        const std::string name    = std::string("sd.convrot_fusion.") + ggml_op_name(op) + ".cuda.v1";
+            const std::string name    = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
             desc.name                 = name.c_str();
             desc.op                   = op;
             desc.priority             = 225;
@@ -312,4 +291,6 @@ void sd_ggml_register_cuda_convrot(ggml_backend_dev_t device) {
             desc.destroy              = sd_convrot::destroy;
             GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
         }
+    }
+
 }

src/core/cuda/convrot_kernel.h

@@ -1,19 +1,20 @@
 #ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
 #define __SD_CORE_CUDA_CONVROT_KERNEL_H__
-#include <cuda_runtime.h>
 #include <cstdint>
-namespace sd_convrot {
+#include "runtime.h"
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_convrot {
         static __device__ __forceinline__ float warp_reduce_sum(float x) {
 #pragma unroll
             for (int offset = 16; offset > 0; offset >>= 1)
-            x += __shfl_xor_sync(0xffffffff, x, offset, 32);
+                x += cuda_shuffle_xor(x, offset, 32);
             return x;
         }
         static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
 #pragma unroll
             for (int offset = 16; offset > 0; offset >>= 1) {
-            x.x += __shfl_xor_sync(0xffffffff, x.x, offset, 32);
-            x.y += __shfl_xor_sync(0xffffffff, x.y, offset, 32);
+                x.x += cuda_shuffle_xor(x.x, offset, 32);
+                x.y += cuda_shuffle_xor(x.y, offset, 32);
             }
             return x;
         }
@@ -21,7 +22,7 @@ namespace sd_convrot {
         static __device__ __forceinline__ float warp_reduce_max(float x) {
 #pragma unroll
             for (int offset = Width / 2; offset > 0; offset >>= 1)
-            x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, Width));
+                x = fmaxf(x, cuda_shuffle_xor(x, offset, Width));
             return x;
         }
         static __device__ float2 block_sum(float2 x, float2* shared) {
@@ -129,19 +130,19 @@ namespace sd_convrot {
                         const int base = lane & ~(3 * stride);
 #pragma unroll
                         for (int j = 0; j < 8; ++j) {
-                        const float a = __shfl_sync(0xffffffff, v[j], base, 32);
-                        const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
-                        const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
-                        const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
+                            const float a = cuda_shuffle(v[j], base, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+                            const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
                             v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
                         }
                     }
 #pragma unroll
                     for (int j = 0; j < 8; j += 2) {
-                    const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
-                    const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
-                    const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
-                    const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
+                        const float a = cuda_shuffle(v[j], lane % 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+                        const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
                         v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
                         v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
                     }
@@ -245,5 +246,6 @@ namespace sd_convrot {
             }
         }
 
+    }
 }
 #endif

src/core/cuda/cute_attention.cu

@@ -588,7 +588,7 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
         offset = 5;
     } else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
         if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
-            !sd_cuda_rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
+            !sd_cuda::rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
             graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
             return 0;
         offset = 3;
@@ -608,7 +608,7 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
             continue;
         for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
             if (attn->src[0] == graph_node(graph, j + 2) &&
-                sd_cuda_rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
+                sd_cuda::rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
                 if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) &&
                     ggml_ext_graph_delay(graph, j - 2, 5, i))
                     break;
@@ -687,7 +687,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
             if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
                 return GGML_STATUS_FAILED;
         } else if (i == 0 && plan.index == 3) {
-            if (sd_cuda_rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
+            if (sd_cuda::rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
                 return GGML_STATUS_FAILED;
         } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
             inputs[i] = static_cast<const half*>(src->data);

src/core/cuda/i8_projection.cu

@@ -208,7 +208,7 @@ static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* ma
 
 bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
     if (!projection_supports(device, mat, view) || norm->src[0] != view ||
-        !sd_cuda_rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
+        !sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
         return false;
     int32_t interleaved;
     std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));

src/core/cuda/int8.cu

@@ -1,21 +1,30 @@
 #include <cstdlib>
 #include "convrot_kernel.h"
 #include "core/int8_ops.h"
-#include "i8_gemm.h"
+#ifdef SD_CUDA_I8_CUTLASS
+#include "core/cuda/i8_gemm.h"
+#else
+struct SDI8CudaGemmPlan;
+#endif
 
+#include "runtime.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hipblas/hipblas.h>
+#elif !defined(SD_CUDA_VENDOR_MUSA)
 #include <cublas_v2.h>
-#include <cuda_runtime.h>
+#endif
 #include <algorithm>
 #include <climits>
 #include <mutex>
 #include <unordered_map>
 
-namespace sd_i8_cuda {
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_i8_cuda {
         template <int Width = 32>
         static __device__ __forceinline__ float warp_reduce_max(float value) {
 #pragma unroll
             for (int offset = Width / 2; offset > 0; offset >>= 1)
-            value = fmaxf(value, __shfl_xor_sync(0xffffffff, value, offset, Width));
+                value = fmaxf(value, cuda_shuffle_xor(value, offset, Width));
             return value;
         }
         static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) {
@@ -93,19 +102,19 @@ namespace sd_i8_cuda {
                         const int base = lane & ~(3 * stride);
 #pragma unroll
                         for (int j = 0; j < 8; ++j) {
-                        const float a = __shfl_sync(0xffffffff, v[j], base, 32);
-                        const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
-                        const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
-                        const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
+                            const float a = cuda_shuffle(v[j], base, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+                            const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
                             v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
                         }
                     }
 #pragma unroll
                     for (int j = 0; j < 8; j += 2) {
-                    const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
-                    const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
-                    const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
-                    const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
+                        const float a = cuda_shuffle(v[j], lane % 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+                        const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
                         v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
                         v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
                     }
@@ -353,29 +362,6 @@ namespace sd_i8_cuda {
             dst[index] = value;
         }
 
-    struct DeviceInfo {
-        int ordinal   = -1;
-        int major     = 0;
-        int minor     = 0;
-        size_t shared = 0;
-    };
-    static DeviceInfo device_info(ggml_backend_dev_t device) {
-        static std::mutex mutex;
-        static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
-        std::lock_guard<std::mutex> lock(mutex);
-        auto found = cache.find(device);
-        if (found != cache.end())
-            return found->second;
-        DeviceInfo info;
-        auto reg             = ggml_backend_dev_backend_reg(device);
-        auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
-        int ordinal          = query ? query(device) : -1;
-        cudaDeviceProp props = {};
-        if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
-            info = {ordinal, props.major, props.minor, size_t(props.sharedMemPerBlockOptin)};
-        cache.emplace(device, info);
-        return info;
-    }
         static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) {
             if (!tensor)
                 return true;
@@ -396,16 +382,16 @@ namespace sd_i8_cuda {
             for (auto* input : node->src)
                 if (!aligned(input))
                     return false;
-        const auto info = device_info(device);
-        return (info.major > 7 || (info.major == 7 && info.minor >= 5)) &&
-               ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_cuda_blas_v1");
+            const auto info = cuda_device_info(device);
+            return info.int8_blas &&
+                   ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
         }
         static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
             int group;
             if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0]))
                 return false;
             const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]);
-        return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && device_info(device).ordinal >= 0;
+            return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0;
         }
         enum Kind { MAT,
                     QUANTIZE,
@@ -436,13 +422,13 @@ namespace sd_i8_cuda {
             if (first + 2 < ggml_graph_n_nodes(g)) {
                 const auto* quant = ggml_graph_node(g, first + 2);
                 if (quant->src[0] == next && convrot_supports(device, quant, nullptr) &&
-                size_t(next->ne[0]) * sizeof(float) + 128 <= device_info(device).shared)
+                    size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
                     return 3;
             }
             return 2;
         }
         struct Plan {
-        DeviceInfo device;
+            ggml_ext_gpu_device_info device;
             Kind kind;
             int64_t k, n, rows, padded;
             size_t quant, scales, accum, partial, output, gemm_offset;
@@ -461,7 +447,7 @@ namespace sd_i8_cuda {
             auto* p        = new Plan{};
             *state         = p;
             auto device    = ggml_backend_get_device(backend);
-        p->device      = device_info(device);
+            p->device      = cuda_device_info(device);
             p->kind        = !mat ? QUANTIZE : count == 3              ? SWIGLU_QUANTIZE
                                            : count == 1                ? MAT
                                            : dst->op == GGML_OP_CONCAT ? CONCAT
@@ -493,7 +479,7 @@ namespace sd_i8_cuda {
                 const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu);
                 p->accum                = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
             }
-        const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared;
+            const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
             p->partial                = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
             p->output                 = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
             *workspace                = total;
@@ -513,23 +499,23 @@ namespace sd_i8_cuda {
             const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
             return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
         }
-    static cudaError_t convrot_launch(const Plan& p, cudaStream_t stream, const float* src, int8_t* dst, float* scales, float* partial) {
+        static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) {
             if (p.padded > p.rows) {
-            auto error = cudaMemsetAsync(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
-            if (error != cudaSuccess)
+                auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                if (error != cuda_success)
                     return error;
             }
             const size_t shared = size_t(p.k) * sizeof(float);
-        if (shared + 128 <= p.device.shared) {
+            if (shared + 128 <= p.device.shared_memory) {
                 static std::mutex mutex;
                 static std::unordered_map<int, bool> configured;
                 {
                     std::lock_guard<std::mutex> lock(mutex);
-                if (!configured[p.device.ordinal]) {
-                    auto error = cudaFuncSetAttribute(quantize_convrot, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - 128));
-                    if (error != cudaSuccess)
+                    if (!configured[p.device.device]) {
+                        auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
+                        if (error != cuda_success)
                             return error;
-                    configured[p.device.ordinal] = true;
+                        configured[p.device.device] = true;
                     }
                 }
                 int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256
@@ -541,7 +527,7 @@ namespace sd_i8_cuda {
                 reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
                 convrot_group_quantize_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, dst, scales, p.k, groups, total_groups, p.k);
             }
-        return cudaGetLastError();
+            return cuda_last_error();
         }
         template <bool Dequantize>
         static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
@@ -563,23 +549,23 @@ namespace sd_i8_cuda {
             dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
         }
         template <bool Swiglu, bool Dequantize, bool Rotated>
-    static cudaError_t quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, cudaStream_t stream) {
+        static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
             const int64_t hidden = p.n / 2;
             auto kernel          = sd_convrot::quantize<Swiglu, Dequantize, false, false, Rotated>;
             static std::mutex mutex;
             static std::unordered_map<int, bool> initialized;
             {
                 std::lock_guard<std::mutex> lock(mutex);
-            if (!initialized[p.device.ordinal]) {
-                auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - 128));
-                if (error != cudaSuccess)
+                if (!initialized[p.device.device]) {
+                    auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
+                    if (error != cuda_success)
                         return error;
-                initialized[p.device.ordinal] = true;
+                    initialized[p.device.device] = true;
                 }
             }
             if (p.padded > p.rows) {
-            auto error = cudaMemsetAsync(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
-            if (error != cudaSuccess)
+                auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
+                if (error != cuda_success)
                     return error;
             }
             const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
@@ -588,16 +574,15 @@ namespace sd_i8_cuda {
             kernel<<<unsigned(p.rows), threads, size_t(hidden) * sizeof(float), stream>>>(src, output,
                                                                                           reinterpret_cast<float*>(output + hidden * p.padded), hidden, p.rows, hidden,
                                                                                           Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr);
-        return cudaGetLastError();
+            return cuda_last_error();
         }
         static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
             const auto& p = *static_cast<const Plan*>(state);
             auto reg      = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
-        auto query    = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
-        ggml_ext_cuda_context context;
-        if (!query || !query(execution->backend, &context) || context.device != p.device.ordinal)
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
                 return GGML_STATUS_FAILED;
-        auto stream   = static_cast<cudaStream_t>(context.stream);
+            auto stream   = static_cast<CUDAStream>(context.stream);
             auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
             auto* dst     = nodes[count - 1];
             SDI8Operation op;
@@ -615,7 +600,7 @@ namespace sd_i8_cuda {
             if (p.gemm && p.kind == SWIGLU_QUANTIZE)
                 alias = false;
             void* output    = alias ? scratch + p.output : dst->data;
-        cudaError_t error = cudaSuccess;
+            CUDAError error = cuda_success;
             if (!mat) {
                 auto* quant  = static_cast<int8_t*>(output);
                 auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
@@ -630,8 +615,8 @@ namespace sd_i8_cuda {
                         error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
                     else {
                         if (p.padded > p.rows)
-                        error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
-                    if (error == cudaSuccess)
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                        if (error == cuda_success)
                             quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(static_cast<const float*>(op.input->data), quant, temporary_scales, p.k, p.rows);
                     }
                 } else {
@@ -640,12 +625,12 @@ namespace sd_i8_cuda {
                     if (op.packed || p.rows == p.padded)
                         quant = data;
                     else {
-                    error = cudaMemcpyAsync(quant, data, size_t(p.k) * p.rows, cudaMemcpyDeviceToDevice, stream);
-                    if (error == cudaSuccess)
-                        error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                        error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
+                        if (error == cuda_success)
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
                     }
                 }
-            if (error != cudaSuccess || cudaGetLastError() != cudaSuccess)
+                if (error != cuda_success || cuda_last_error() != cuda_success)
                     return GGML_STATUS_FAILED;
                 const auto* ws          = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
                 const auto* bias        = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
@@ -662,15 +647,34 @@ namespace sd_i8_cuda {
                 } else
 #endif
                 {
-                auto blas_query = reinterpret_cast<ggml_ext_cuda_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_blas_v1"));
-                auto handle     = blas_query ? static_cast<cublasHandle_t>(blas_query(execution->backend)) : nullptr;
+                    auto blas_query = reinterpret_cast<ggml_ext_gpu_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1"));
+                    auto handle     = blas_query ? blas_query(execution->backend) : nullptr;
                     if (!handle)
                         return GGML_STATUS_FAILED;
-                const int32_t alpha = 1, beta = 0;
-                if (cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, int(p.n), int(p.padded), int(p.k), &alpha,
-                                 op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k), &beta, accum, CUDA_R_32I,
-                                 int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP) != CUBLAS_STATUS_SUCCESS)
+#if defined(SD_CUDA_VENDOR_MUSA)
                     return GGML_STATUS_FAILED;
+#else
+                    const int32_t alpha = 1, beta = 0;
+#if defined(SD_CUDA_VENDOR_HIP)
+#if HIP_VERSION >= 60500000
+                    constexpr auto compute_type = HIPBLAS_COMPUTE_32I;
+#else
+                    constexpr auto compute_type = HIPBLAS_R_32I;
+#endif
+                    const auto status           = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
+                                                                int(p.n), int(p.padded), int(p.k), &alpha,
+                                                                op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
+                                                                &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
+                    if (status != HIPBLAS_STATUS_SUCCESS)
+#else
+                    const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
+                                                     int(p.n), int(p.padded), int(p.k), &alpha,
+                                                     op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k),
+                                                     &beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+                    if (status != CUBLAS_STATUS_SUCCESS)
+#endif
+                        return GGML_STATUS_FAILED;
+#endif
                 }
                 const auto* values = reinterpret_cast<const float*>(accum);
                 if (p.kind == SWIGLU_QUANTIZE) {
@@ -686,7 +690,7 @@ namespace sd_i8_cuda {
                         swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
                     else
                         swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
-                error = cudaGetLastError();
+                    error = cuda_last_error();
                 } else if (!p.gemm) {
                     const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
                     if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(gemm_output)) % 16 == 0) {
@@ -695,22 +699,22 @@ namespace sd_i8_cuda {
                     } else {
                         dequantize_i32_rows_cuda<<<dim3(unsigned((p.n + 255) / 256), gy, gz), 256, 0, stream>>>(accum, scales, ws, bias, gemm_output, p.n, p.rows);
                     }
-                error = cudaGetLastError();
+                    error = cuda_last_error();
                 }
-            if (error == cudaSuccess && prefix)
-                error = cudaMemcpyAsync(output, dst->src[0]->data, prefix, cudaMemcpyDeviceToDevice, stream);
+                if (error == cuda_success && prefix)
+                    error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
             }
-        if (error == cudaSuccess && alias)
-            error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
-        return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
         }
 
-}
+    }
 
-void sd_ggml_register_cuda_i8(ggml_backend_dev_t device) {
+    void register_i8(ggml_backend_dev_t device) {
         ggml_ext_kernel_desc desc = {};
         desc.abi_version          = GGML_EXT_ABI_VERSION;
-    desc.name                 = "sd.linear_i8.cuda.v1";
+        desc.name                 = "sd.linear_i8." SD_CUDA_NAME ".v1";
         desc.op                   = GGML_OP_CUSTOM;
         desc.custom_op            = sd_ggml_i8_op();
         desc.priority             = 100;
@@ -719,27 +723,29 @@ void sd_ggml_register_cuda_i8(ggml_backend_dev_t device) {
         desc.execute              = sd_i8_cuda::execute;
         desc.destroy              = sd_i8_cuda::destroy;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name      = "sd.linear_i8.cuda.override.v1";
+        desc.name      = "sd.linear_i8." SD_CUDA_NAME ".override.v1";
         desc.op        = GGML_OP_MUL_MAT;
         desc.custom_op = nullptr;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name     = "sd.linear_i8.epilogue.cuda.v1";
+        desc.name     = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1";
         desc.priority = 240;
         desc.supports = nullptr;
         desc.match    = sd_i8_cuda::match;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name      = "sd.linear_i8.epilogue.cuda.custom.v1";
+        desc.name      = "sd.linear_i8.epilogue." SD_CUDA_NAME ".custom.v1";
         desc.op        = GGML_OP_CUSTOM;
         desc.custom_op = sd_ggml_i8_op();
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
         desc.priority  = 100;
         desc.match     = nullptr;
-    desc.name      = "sd.quantize_i8_convrot.cuda.v1";
+        desc.name      = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
         desc.custom_op = sd_ggml_convrot_op();
         desc.supports  = sd_i8_cuda::convrot_supports;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name      = "sd.quantize_i8_convrot.cuda.override.v1";
+        desc.name      = "sd.quantize_i8_convrot." SD_CUDA_NAME ".override.v1";
         desc.op        = GGML_OP_QUANTIZE_I8_CONVROT;
         desc.custom_op = nullptr;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
 }

src/core/cuda/registration.cu

@@ -0,0 +1,28 @@
+#include "runtime.h"
+
+#ifdef SD_CUDA_CUTE_ATTENTION
+void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+    void register_rope(ggml_backend_dev_t device);
+    void register_i8(ggml_backend_dev_t device);
+    void register_convrot(ggml_backend_dev_t device);
+    void register_affine(ggml_backend_dev_t device);
+
+    void register_extensions(ggml_backend_dev_t device) {
+        if (cuda_device_info(device).device < 0)
+            return;
+        register_affine(device);
+        register_i8(device);
+        register_convrot(device);
+        register_rope(device);
+#ifdef SD_CUDA_CUTE_ATTENTION
+        sd_ggml_register_cuda_i8_projection(device);
+        sd_ggml_register_cuda_i8_gated_residual(device);
+        sd_ggml_register_cuda_attention(device);
+#endif
+    }
+}

src/core/cuda/registration.h

@@ -0,0 +1,13 @@
+#ifndef __SD_CORE_CUDA_REGISTRATION_H__
+#define __SD_CORE_CUDA_REGISTRATION_H__
+#include "ggml-backend.h"
+namespace sd_cuda {
+    void register_extensions(ggml_backend_dev_t device);
+}
+namespace sd_hip {
+    void register_extensions(ggml_backend_dev_t device);
+}
+namespace sd_musa {
+    void register_extensions(ggml_backend_dev_t device);
+}
+#endif

src/core/cuda/rope.cu

@@ -1,19 +1,19 @@
 #include "core/cuda/rope.h"
 #include "core/ggml_extensions.h"
 
-#include <cuda_fp16.h>
 #include <climits>
 #include <cstdlib>
 #include <cstring>
 #include <type_traits>
 
-static __device__ __forceinline__ float sd_warp_sum(float v) {
+namespace SD_CUDA_NAMESPACE {
+    static __device__ __forceinline__ float sd_warp_sum(float v) {
         for (int mask = 16; mask > 0; mask >>= 1)
-        v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask));
+            v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
         return v;
-}
-template <int D, bool interleaved, typename T>
-static __global__ void rms_norm_mul_rope_apply_fused(
+    }
+    template <int D, bool interleaved, typename T>
+    static __global__ void rms_norm_mul_rope_apply_fused(
         const float* x,
         const float* weight,
         const float* pe,
@@ -54,7 +54,7 @@ static __global__ void rms_norm_mul_rope_apply_fused(
             const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
             float rotated     = value;
             if constexpr (interleaved) {
-            const float other = __shfl_xor_sync(0xffffffff, value, 1, 32);
+                const float other = cuda_shuffle_xor(value, 1, 32);
                 if (col < 2 * rot_half) {
                     const float* m      = pe + 4 * (token * rot_half + col / 2);
                     const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
@@ -79,15 +79,15 @@ static __global__ void rms_norm_mul_rope_apply_fused(
                 out[col] = rotated;
             }
         }
-}
+    }
 
-struct RopeApplyStrides {
+    struct RopeApplyStrides {
         int64_t x[4];
         int64_t pe[4];
-};
+    };
 
-template <bool interleaved>
-static __global__ void rope_apply_f32(
+    template <bool interleaved>
+    static __global__ void rope_apply_f32(
         const float* x,
         const float* pe,
         float* dst,
@@ -128,25 +128,25 @@ static __global__ void rope_apply_f32(
             out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
             out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
         }
-}
+    }
 
-static bool rope_interleaved(const ggml_tensor* rope) {
+    static bool rope_interleaved(const ggml_tensor* rope) {
         int32_t value;
         std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value));
         return value != 0;
-}
+    }
 
-static float param_f32(const ggml_tensor* node, int index) {
+    static float param_f32(const ggml_tensor* node, int index) {
         float value;
         std::memcpy(&value, node->op_params + index, sizeof(value));
         return value;
-}
+    }
 
-bool sd_cuda_rope_is_node(const ggml_tensor* node) {
+    static bool rope_is_node(const ggml_tensor* node) {
         return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
-}
+    }
 
-static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
+    static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
         const auto* x  = rope->src[0];
         const auto* pe = rope->src[1];
         if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
@@ -161,10 +161,10 @@ static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
                 return false;
         }
         return true;
-}
+    }
 
-bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
-    if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) ||
+    bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+        if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !rope_is_node(rope) ||
             mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
             return false;
         const auto* x      = norm->src[0];
@@ -176,10 +176,10 @@ bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mu
                x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
                x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
                (ggml_nrows(x) + 3) / 4 <= INT_MAX;
-}
+    }
 
-template <typename T, bool Interleaved>
-static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) {
+    template <typename T, bool Interleaved>
+    static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
         const auto* x      = nodes[0]->src[0];
         const auto* pe     = nodes[2]->src[1];
         const int64_t rows = ggml_nrows(x);
@@ -200,27 +200,27 @@ static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* node
             SD_ROPE(224);
             SD_ROPE(256);
             default:
-            return cudaErrorInvalidValue;
+                return cuda_invalid_value;
         }
 #undef SD_ROPE
-    return cudaGetLastError();
-}
+        return cuda_last_error();
+    }
 
-cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
+    CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
         if (fp16) {
             return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
                                               : rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
         }
         return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
                                           : rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
-}
+    }
 
-static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+    static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
         static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
         if (disabled)
             return 0;
         auto* g = const_cast<ggml_cgraph*>(graph);
-    if (first + 2 >= ggml_graph_n_nodes(g) || !sd_cuda_rope_fusion_supports(
+        if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
                                                       ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
             return 0;
         if (first + 4 < ggml_graph_n_nodes(g)) {
@@ -233,9 +233,9 @@ static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, v
                 return 5;
         }
         return 3;
-}
+    }
 
-static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+    static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
         static const bool disabled = [] {
             const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
             return value && std::atoi(value) != 0;
@@ -246,7 +246,7 @@ static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
         for (int i = 2; i + 1 < total; ++i) {
             const auto* mul  = ggml_graph_node(graph, i);
             const auto* rope = ggml_graph_node(graph, i + 1);
-        if (mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) || rope->src[0] != mul ||
+            if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
                 !mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
                 continue;
             const auto* norm = mul->src[0];
@@ -267,35 +267,33 @@ static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
                 break;
             }
         }
-}
+    }
 
-static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+    static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
         *workspace = ggml_nbytes(nodes[count - 1]);
         return GGML_STATUS_SUCCESS;
-}
+    }
 
-static void rope_destroy(void*, void*) {}
+    static void rope_destroy(void*, void*) {}
 
-static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+    static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
         const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
         return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
-}
+    }
 
-static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
-    auto reg   = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
-    auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
-    ggml_ext_cuda_context context;
-    if (!query || !query(execution->backend, &context))
+    static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
             return GGML_STATUS_FAILED;
-    auto stream      = static_cast<cudaStream_t>(context.stream);
+        auto stream      = static_cast<CUDAStream>(context.stream);
         const auto* rope = nodes[count == 1 ? 0 : 2];
         const auto* x    = count == 1 ? rope->src[0] : nodes[0]->src[0];
         auto* dst        = nodes[count - 1];
         const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
         void* output     = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
-    cudaError_t error;
+        CUDAError error;
         if (count > 1) {
-        error = sd_cuda_rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
+            error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
         } else {
             RopeApplyStrides strides;
             const auto* pe = rope->src[1];
@@ -307,17 +305,17 @@ static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggm
             auto kernel         = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>;
             kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
                                                                                        static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
-        error = cudaGetLastError();
+            error = cuda_last_error();
+        }
+        if (error == cuda_success && alias)
+            error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+        return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
     }
-    if (error == cudaSuccess && alias)
-        error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
-    return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
-}
 
-void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) {
+    void register_rope(ggml_backend_dev_t device) {
         ggml_ext_kernel_desc desc = {};
         desc.abi_version          = GGML_EXT_ABI_VERSION;
-    desc.name                 = "sd.rope_apply.cuda.v1";
+        desc.name                 = "sd.rope_apply." SD_CUDA_NAME ".v1";
         desc.op                   = GGML_OP_CUSTOM;
         desc.custom_op            = sd_ggml_rope_op();
         desc.priority             = 100;
@@ -326,15 +324,17 @@ void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) {
         desc.destroy              = rope_destroy;
         desc.execute              = rope_execute;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name      = "sd.rope_apply.cuda.override.v1";
+        desc.name      = "sd.rope_apply." SD_CUDA_NAME ".override.v1";
         desc.op        = GGML_OP_ROPE_APPLY;
         desc.custom_op = nullptr;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name     = "sd.rms_norm_mul_rope.cuda.v1";
+        desc.name     = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
         desc.op       = GGML_OP_RMS_NORM;
         desc.priority = 250;
         desc.supports = nullptr;
         desc.match    = rope_match;
         desc.optimize = rope_optimize;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
 }

src/core/cuda/rope.h

@@ -1,12 +1,11 @@
 #ifndef __SD_CORE_CUDA_ROPE_H__
 #define __SD_CORE_CUDA_ROPE_H__
 
-#include <cuda_runtime.h>
-#include "ggml-ext.h"
+#include "runtime.h"
 
-bool sd_cuda_rope_is_node(const ggml_tensor* node);
-bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
-cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
-void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
+namespace SD_CUDA_NAMESPACE {
+    bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+    CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
 
+}
 #endif

src/core/cuda/runtime.h

@@ -0,0 +1,111 @@
+#ifndef __SD_CORE_CUDA_RUNTIME_H__
+#define __SD_CORE_CUDA_RUNTIME_H__
+
+#include "ggml-ext.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#define SD_CUDA_NAMESPACE sd_hip
+#define SD_CUDA_NAME "hip"
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <musa_fp16.h>
+#include <musa_runtime.h>
+#define SD_CUDA_NAMESPACE sd_musa
+#define SD_CUDA_NAME "musa"
+#else
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#define SD_CUDA_NAMESPACE sd_cuda
+#define SD_CUDA_NAME "cuda"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+#if defined(SD_CUDA_VENDOR_HIP)
+    using CUDAError                          = hipError_t;
+    using CUDAStream                         = hipStream_t;
+    static constexpr auto cuda_success       = hipSuccess;
+    static constexpr auto cuda_invalid_value = hipErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_HIP;
+    static CUDAError cuda_last_error() {
+        return hipGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return hipMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel, int) {
+        // HIP uses the device's shared-memory limit without CUDA's opt-in attribute.
+        return cuda_success;
+    }
+#elif defined(SD_CUDA_VENDOR_MUSA)
+    using CUDAError                          = musaError_t;
+    using CUDAStream                         = musaStream_t;
+    static constexpr auto cuda_success       = musaSuccess;
+    static constexpr auto cuda_invalid_value = musaErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_MUSA;
+    static CUDAError cuda_last_error() {
+        return musaGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return musaMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+        return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+    }
+#else
+    using CUDAError                          = cudaError_t;
+    using CUDAStream                         = cudaStream_t;
+    static constexpr auto cuda_success       = cudaSuccess;
+    static constexpr auto cuda_invalid_value = cudaErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_CUDA;
+    static CUDAError cuda_last_error() {
+        return cudaGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return cudaMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+        return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+    }
+#endif
+    template <typename T>
+    static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+        return __shfl(value, lane, width);
+#else
+        return __shfl_sync(0xffffffff, value, lane, width);
+#endif
+    }
+    template <typename T>
+    static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+        return __shfl_xor(value, lane, width);
+#else
+        return __shfl_xor_sync(0xffffffff, value, lane, width);
+#endif
+    }
+    static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
+        auto reg   = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
+        return query && query(backend, context) && context->kind == cuda_kind;
+    }
+    static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
+        ggml_ext_gpu_device_info info = {};
+        auto query                    = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
+        if (!query || !query(device, &info) || info.kind != cuda_kind)
+            return {cuda_kind, -1, 0, 0, false};
+        return info;
+    }
+}
+#endif

src/core/ggml_extensions.cpp

@@ -9,19 +9,8 @@
 #include <mutex>
 #include <unordered_set>
 
-#ifdef SD_VULKAN_EXTENSIONS
-void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
-#endif
-#ifdef SD_CUDA_EXTENSIONS
-void sd_ggml_register_cuda_affine(ggml_backend_dev_t device);
-void sd_ggml_register_cuda_i8(ggml_backend_dev_t device);
-void sd_ggml_register_cuda_convrot(ggml_backend_dev_t device);
-void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
-#endif
-#ifdef SD_CUDA_CUTE_ATTENTION
-void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
-void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
-void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
+#ifdef SD_BACKEND_EXTENSIONS
+#include "sd_backend_extensions.h"
 #endif
 
 static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
@@ -174,34 +163,16 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
     auto device    = ggml_backend_get_device(backend);
     auto reg       = ggml_backend_dev_backend_reg(device);
     const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
-    const bool cuda   = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") != nullptr;
-    const bool vulkan = ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1") != nullptr;
-    if (!cpu && !cuda && !vulkan) {
-        return;
-    }
     static std::mutex mutex;
     static std::unordered_set<ggml_backend_dev_t> registered;
     std::lock_guard<std::mutex> lock(mutex);
     if (!registered.insert(device).second) {
         return;
     }
-    if (vulkan) {
-#ifdef SD_VULKAN_EXTENSIONS
-        sd_ggml_register_vulkan_rope(device);
-#endif
-        return;
-    }
-    if (cuda) {
-#ifdef SD_CUDA_EXTENSIONS
-        sd_ggml_register_cuda_affine(device);
-        sd_ggml_register_cuda_i8(device);
-        sd_ggml_register_cuda_convrot(device);
-        sd_ggml_register_cuda_rope(device);
-#endif
-#ifdef SD_CUDA_CUTE_ATTENTION
-        sd_ggml_register_cuda_i8_projection(device);
-        sd_ggml_register_cuda_i8_gated_residual(device);
-        sd_ggml_register_cuda_attention(device);
+    if (!cpu) {
+#ifdef SD_BACKEND_EXTENSIONS
+        for (auto register_backend : sd_backend_extensions)
+            register_backend(device);
 #endif
         return;
     }
@@ -234,11 +205,8 @@ ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_
     if (backend) {
         auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
         if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")
-#ifdef SD_CUDA_EXTENSIONS
-            || ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1")
-#endif
-#ifdef SD_VULKAN_EXTENSIONS
-            || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
+#ifdef SD_BACKEND_EXTENSIONS
+            || ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1") || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
 #endif
         ) {
             ggml_tensor* inputs[] = {x, pe};

src/core/vulkan/rope.cpp

@@ -200,6 +200,8 @@ namespace sd_vulkan_rope {
 }
 
 void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) {
+    if (!sd_vulkan_rope::api(device))
+        return;
     ggml_ext_kernel_desc desc = {};
     desc.abi_version          = GGML_EXT_ABI_VERSION;
     desc.name                 = "sd.rope_apply.vulkan.v1";