revert: restore public backend extension integration

cmake/ggml.cmake

@@ -9,6 +9,7 @@ if(NOT TARGET ggml)
     add_library(ggml ALIAS ggml::ggml)
 endif()
 
+if(SD_USE_UPSTREAM_GGML)
     get_target_property(sd_ggml_imported ggml IMPORTED)
     if(sd_ggml_imported)
         set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src")
@@ -20,18 +21,9 @@ if(NOT EXISTS "${sd_ggml_private_include}/ggml-impl.h")
     endif()
     target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}")
     set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}")
+endif()
 
 if(SD_USE_UPSTREAM_GGML)
     target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML)
     message(WARNING "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are converted to F16 at load time. Some operators may be unsupported and performance may be lower than with patched GGML.")
 endif()
-
-if(NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
-    set(SD_GGML_CUDA_INCLUDE_DIR "${sd_ggml_private_include}/ggml-cuda")
-    if(NOT EXISTS "${SD_GGML_CUDA_INCLUDE_DIR}/vendors/cuda.h" OR
-       NOT EXISTS "${SD_GGML_CUDA_INCLUDE_DIR}/vendors/hip.h" OR
-       NOT EXISTS "${SD_GGML_CUDA_INCLUDE_DIR}/vendors/musa.h")
-        message(FATAL_ERROR "Set SD_GGML_SOURCE_DIR to the matching ggml source tree for CUDA/HIP/MUSA vendor headers.")
-    endif()
-    target_include_directories(${SD_LIB} PRIVATE "${SD_GGML_CUDA_INCLUDE_DIR}")
-endif()

docs/build.md

@@ -32,9 +32,10 @@ cmake -S . -B build-upstream -DSD_USE_UPSTREAM_GGML=ON -DSD_GGML_SOURCE_DIR=../g
 cmake --build build-upstream --config Release
 ```
 
-Both patched and upstream builds use the selected library and matching private
+Upstream compatibility mode uses the selected library and matching private
-headers for backend integration. Backend options such as `-DSD_CUDA=ON` apply
+headers. The patched GGML build uses public headers, including `ggml-ext.h` and
-to the selected tree as usual.
+`ggml-graph.h`. Backend options such as `-DSD_CUDA=ON` apply to the selected tree
+as usual.
 
 `SD_USE_UPSTREAM_GGML` defaults to `OFF`, which enables the patched GGML
 extensions. Set it to `ON` when using upstream GGML; it selects the compatibility
@@ -52,9 +53,8 @@ quantization types remain available, subject to backend operator support.
 
 `SD_USE_SYSTEM_GGML=ON` instead links an installed GGML CMake package, located
 with `ggml_DIR` or `CMAKE_PREFIX_PATH`. An installed patched package must include
-the extension and graph APIs. Set `SD_GGML_SOURCE_DIR` to its matching source
+the extension and graph APIs; its source tree is not required for private
-tree for private backend headers, including CUDA/HIP/MUSA vendor headers.
+headers. The installed library must use the same ABI settings as sd.cpp,
-The installed library must use the same ABI settings as sd.cpp,
 including `GGML_MAX_NAME`. For an upstream package, also set
 `SD_USE_UPSTREAM_GGML=ON` and point `SD_GGML_SOURCE_DIR` to its matching source
 tree for the compatibility layer's private headers.

docs/ggml_extensions.md

@@ -104,21 +104,22 @@ and backend `graph_plan_*` interfaces do not dispatch these extensions.
 
 ## Backend services
 
-GPU services are private backend integration details declared in
+Services are discovered through `ggml_backend_reg_get_proc_address`:
-`ggml/src/ggml-ext-gpu.h` and discovered through
-`ggml_backend_reg_get_proc_address`. They are not installed public API.
-CPU callbacks use the native CPU backend entry point and its configured thread
-pool directly from `src/core/cpu/registration.cpp`.
 
 | Service | Purpose |
 | --- | --- |
-| `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, INT8 BLAS capability, compute capability and SM limits from the existing device cache. |
+| `ggml_ext_cpu_parallel_v1` | Run an application callback on the CPU backend's configured thread pool, with its abort callback. |
+| `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. The
+or destroy it, or change its stream, pointer mode, math mode, or workspace. This
-same services support CUDA, HIP, and MUSA with the same ownership rules. Extensions
+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.
 
@@ -153,11 +154,9 @@ 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. Both sd.cpp and ggml use the same
+`sd_musa`, so implementations can coexist. `runtime.h` handles the small runtime
-existing `ggml/src/ggml-cuda/vendors/{cuda,hip,musa}.h` headers for runtime and BLAS API
+differences, including logical 32-lane shuffles on HIP. The application defines
-compatibility. Kernels call CUDA-style APIs directly; `backend.h` only selects
+and launches its kernels directly; it includes only public ggml headers.
-the vendor/namespace and queries ggml's private device and stream services.
-CUDA-only PTX/CUTLASS kernels remain restricted to NVIDIA builds.
 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.
@@ -170,7 +169,7 @@ separate weight scale tensor.
 CUTLASS headers are configured. It registers an override of `FLASH_ATTN_EXT`,
 with the existing sequence-length preference and native fallback. It also
 registers Q preparation (`RMSNorm -> Mul -> RoPE`), output scaling, and packed
-ConvRot output fusion. These kernels use the extension registration API, borrowed
+ConvRot output fusion. These kernels use only the public extension API, borrowed
 CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or
 private conversion helpers. The existing CUDA fusion and CuTe disable environment
 variables are honored.
@@ -407,9 +406,9 @@ application-owned implementations. This interface does not add dynamic dtypes.
 The Vulkan compute metadata service exposes shared-memory limits, subgroup
 arithmetic, and cooperative-matrix properties; shape selection stays in sd.cpp.
 Shader variants preserve the former tiling and compilation policy, including
-SwiGLU, residual, row-concatenation, and modulation paths. Kernels include the
+SwiGLU, residual, row-concatenation, and modulation paths. Kernels use public
-private Vulkan service header for dispatch, buffer, copy and workspace access.
+Vulkan dispatch, buffer, copy, and workspace services without private backend
-Vulkan currently handles packed group-256 activations with separate
+headers. Vulkan currently handles packed group-256 activations with separate
 weight scales; unsupported custom nodes fall back to CPU.
 
 ## SageAttention
@@ -422,7 +421,7 @@ backend dispatch branch.
 `src/core/cuda/sage_attention.cu` and its `sageattention/` device headers own the
 CUDA implementation. Quantization and temporary attention output use scheduler
 workspace, rather than an oversized output tensor. The implementation obtains
-the backend stream through the private GPU service and converts the final FP16
+the backend stream through the public GPU service and converts the final FP16
 output with a local CUDA kernel. It retains the SM80 compatibility path and the
 SageAttention2/2++ toolkit and architecture requirements. Capability checks use
 the application's compiled CUDA target as well as the physical device.
@@ -442,7 +441,7 @@ CLI path calls this entry point; the migration preserves its low-level behavior.
 `src/core/cuda/sol_attention.cu` owns packing, block summaries, threshold
 selection, and the BF16 tensor-core kernel. Temporary storage is separate
 scheduler workspace, so the output allocation contains only the output tensor.
-The kernel uses the private backend stream service and direct CUDA/PTX operations.
+The kernel uses the public backend stream service and direct CUDA/PTX operations.
 Registration requires CUDA 12 or newer, an SM80-or-newer device, and an application
 kernel compiled for SM80 or newer. CPU, HIP, MUSA, and Vulkan do not advertise
 support. ggml no longer contains a Sol-Attn opcode, constructor, allocation rule,
@@ -459,18 +458,8 @@ leave the destination graph unchanged. Range execution uses already allocated
 tensors and retains the original graph's consumer counts.
 
 Workspace measurement, graph cuts, and evaluation callbacks use these APIs.
-Backend extensions require matching private headers. Compatibility with an
+Normal sd.cpp builds require only ggml's public headers. Compatibility with an
 upstream ggml lacking this header is isolated in `src/core/ggml_compat.cpp`; that
 optional build still requires matching upstream private headers. The compatibility
 layer also retains the workaround for old Metal backends wrapping failed host
 mappings. The patched Metal backend returns null directly when mapping fails.
-
-Vulkan command recording services are declared in the private
-`ggml/src/ggml-ext-vulkan.h`. Pipeline ownership, submission and synchronization
-remain inside the backend, where they share the native execution machinery.
-
-CUTLASS is pinned by sd.cpp at `thirdparty/cutlass` and configured through
-`SD_CUDA_CUTLASS_PATH`. ggml itself has no CUTLASS dependency or configuration.
-CUDA/HIP/MUSA extensions include the existing vendor headers directly from the
-ggml source tree. They are not installed as a public ggml API. System ggml builds
-must also provide the matching `SD_GGML_SOURCE_DIR`.

src/core/cpu/int8.cpp

@@ -101,18 +101,24 @@ namespace sd_cpu {
         return GGML_STATUS_SUCCESS;
     }
     static void destroy(void*, void*) {}
+    static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
+        return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
+    }
     static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
         auto* dst     = nodes[0];
         I8Execution e = {};
         if (!sd_ggml_i8_describe(dst, &e.op))
             return GGML_STATUS_FAILED;
         e.rows   = ggml_nrows(dst);
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
         char* scratch           = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
         const size_t quant_size = quant_workspace(dst, e.op);
         if (e.op.input->type == GGML_TYPE_F32) {
             e.quant     = reinterpret_cast<int8_t*>(scratch);
             e.scales    = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
-            auto status = parallel(context->backend, dst, quantize, &e);
+            auto status = run(context->backend, dst, quantize, &e);
             if (status != GGML_STATUS_SUCCESS)
                 return status;
         } else {
@@ -123,7 +129,7 @@ namespace sd_cpu {
         for (auto* input : dst->src)
             alias |= overlaps(dst, input);
         e.output    = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
-        auto status = parallel(context->backend, dst, multiply, &e);
+        auto status = run(context->backend, dst, multiply, &e);
         if (status == GGML_STATUS_SUCCESS && alias)
             std::memcpy(dst->data, e.output, ggml_nbytes(dst));
         return status;
@@ -133,13 +139,16 @@ namespace sd_cpu {
         I8Execution e = {};
         if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
             return GGML_STATUS_FAILED;
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
         e.op.input      = dst->src[0];
         e.rows          = ggml_nrows(e.op.input);
         bool alias      = overlaps(dst, e.op.input);
         e.quant         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
         const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
         e.scales    = reinterpret_cast<float*>(e.quant + k * padded);
-        auto status = parallel(context->backend, dst, quantize, &e);
+        auto status = run(context->backend, dst, quantize, &e);
         if (status == GGML_STATUS_SUCCESS) {
             std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
             if (alias)

src/core/cpu/registration.cpp

@@ -1,29 +1,8 @@
 #include "core/cpu/registration.h"
 
 #ifndef SD_USE_UPSTREAM_GGML
-#include <cstring>
-#include "ggml-backend-impl.h"
-#include "ggml-cpu.h"
-#include "ggml-impl.h"
 namespace sd_cpu {
 
-    ggml_status parallel(ggml_backend_t backend, ggml_tensor* output, ggml_custom_op_t callback, void* user_data) {
-        if (!ggml_backend_is_cpu(backend) || !output || !callback) {
-            return GGML_STATUS_FAILED;
-        }
-        ggml_tensor node             = *output;
-        ggml_custom_op_params params = {callback, GGML_N_TASKS_MAX, user_data};
-        node.op                      = GGML_OP_CUSTOM;
-        node.flags |= GGML_TENSOR_FLAG_COMPUTE;
-        memcpy(node.op_params, &params, sizeof(params));
-        ggml_tensor* nodes[] = {&node};
-        ggml_cgraph graph    = {};
-        graph.n_nodes        = 1;
-        graph.nodes          = nodes;
-        // Reuse the configured thread pool without dispatching this synthetic node through extensions.
-        return backend->iface.graph_compute(backend, &graph);
-    }
-
     void register_extensions(ggml_backend_dev_t device) {
         register_i8(device);
         register_rope(device);

src/core/cpu/registration.h

@@ -5,7 +5,6 @@
 
 #ifndef SD_USE_UPSTREAM_GGML
 namespace sd_cpu {
-    ggml_status parallel(ggml_backend_t backend, ggml_tensor* output, ggml_custom_op_t callback, void* user_data);
     void register_extensions(ggml_backend_dev_t device);
     void register_rope(ggml_backend_dev_t device);
     void register_i8(ggml_backend_dev_t device);

src/core/cpu/rope.cpp

@@ -103,12 +103,16 @@ namespace sd_cpu {
     }
 
     static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
+        auto device = ggml_backend_get_device(context->backend);
+        auto reg    = ggml_backend_dev_backend_reg(device);
+        auto run    = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        GGML_ASSERT(run);
         RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
         if (count == 3) {
             execution.norm   = nodes[0];
             execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
         }
-        return parallel(context->backend, nodes[count - 1], rope_compute, &execution);
+        return run(context->backend, nodes[count - 1], rope_compute, &execution);
     }
 
     void register_rope(ggml_backend_dev_t device) {

src/core/cuda/CMakeLists.txt

@@ -11,7 +11,7 @@ function(sd_cuda_sources VENDOR LANGUAGE)
     endforeach()
     set(TARGET "sd-${VENDOR}")
     add_library(${TARGET} OBJECT ${SOURCES})
-    target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src" "${SD_GGML_CUDA_INCLUDE_DIR}" "${sd_ggml_private_include}")
+    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)

src/core/cuda/affine.cu

@@ -1,8 +1,8 @@
 #include <climits>
 #include <cstdlib>
 #include <cstring>
-#include "backend.h"
 #include "ggml-ext.h"
+#include "runtime.h"
 
 namespace SD_CUDA_NAMESPACE {
     namespace sd_affine_cuda {
@@ -129,7 +129,7 @@ namespace SD_CUDA_NAMESPACE {
             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];
@@ -154,10 +154,10 @@ namespace SD_CUDA_NAMESPACE {
                 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();
+            auto error = cuda_last_error();
-            if (error == cudaSuccess && alias)
+            if (error == cuda_success && alias)
-                error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
-            return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
         }
     }
 

src/core/cuda/backend.h

@@ -1,41 +0,0 @@
-#ifndef __SD_CORE_CUDA_BACKEND_H__
-#define __SD_CORE_CUDA_BACKEND_H__
-
-#include "ggml-ext-gpu.h"
-#include "ggml-impl.h"
-#if defined(SD_CUDA_VENDOR_HIP)
-#include "vendors/hip.h"
-#define SD_CUDA_NAMESPACE sd_hip
-#define SD_CUDA_NAME "hip"
-#elif defined(SD_CUDA_VENDOR_MUSA)
-#include "vendors/musa.h"
-#define SD_CUDA_NAMESPACE sd_musa
-#define SD_CUDA_NAME "musa"
-#else
-#include "vendors/cuda.h"
-#define SD_CUDA_NAMESPACE sd_cuda
-#define SD_CUDA_NAME "cuda"
-#endif
-
-namespace SD_CUDA_NAMESPACE {
-#if defined(SD_CUDA_VENDOR_HIP)
-    static constexpr auto cuda_kind = GGML_EXT_GPU_HIP;
-#elif defined(SD_CUDA_VENDOR_MUSA)
-    static constexpr auto cuda_kind = GGML_EXT_GPU_MUSA;
-#else
-    static constexpr auto cuda_kind = GGML_EXT_GPU_CUDA;
-#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/cuda/convrot.cu

@@ -3,9 +3,9 @@
 #include <cstdlib>
 #include <cstring>
 #include <string>
-#include "backend.h"
 #include "convrot_kernel.h"
 #include "core/ggml_extensions.h"
+#include "runtime.h"
 
 namespace SD_CUDA_NAMESPACE {
     namespace sd_convrot {
@@ -14,6 +14,11 @@ namespace SD_CUDA_NAMESPACE {
             static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
             return disabled;
         }
+        static float param_f32(const ggml_tensor* tensor, int index) {
+            float value;
+            std::memcpy(&value, tensor->op_params + index, sizeof(value));
+            return value;
+        }
         static bool f32(const ggml_tensor* tensor) {
             return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
         }
@@ -52,12 +57,12 @@ namespace SD_CUDA_NAMESPACE {
             op.up_stride = op.k;
             if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
                 op.kind  = SCALE;
-                op.scale = ggml_get_op_params_f32(first, 0);
+                op.scale = param_f32(first, 0);
-                op.bias  = ggml_get_op_params_f32(first, 1);
+                op.bias  = param_f32(first, 1);
             } else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
                 op.kind       = NORM;
                 op.modulation = nodes[1]->src[1];
-                op.eps        = ggml_get_op_params_f32(first, 0);
+                op.eps        = param_f32(first, 0);
                 if (op.k < 1024 || !vector(op.modulation, op.k))
                     return false;
             } else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
@@ -68,7 +73,7 @@ namespace SD_CUDA_NAMESPACE {
                 op.weight     = nodes[1]->src[1];
                 op.modulation = nodes[2]->src[1];
                 op.shift      = nodes[4]->src[1];
-                op.eps        = ggml_get_op_params_f32(first, 0);
+                op.eps        = param_f32(first, 0);
                 if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !vector(op.shift, op.k))
                     return false;
             } else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
@@ -206,16 +211,15 @@ namespace SD_CUDA_NAMESPACE {
             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) {
-#if !defined(SD_CUDA_VENDOR_HIP)
+                auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
-                auto status = cudaFuncSetAttribute(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared_memory - extra));
+                                                     int(p.device.shared_memory - extra));
-                if (status != cudaSuccess)
+                if (status != cuda_success)
                     return status;
-#endif
                 p.configured = true;
             }
             const int threads    = Normalize ? 1024 : RMS                               ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
@@ -228,7 +232,7 @@ namespace SD_CUDA_NAMESPACE {
             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);
@@ -237,15 +241,15 @@ namespace SD_CUDA_NAMESPACE {
             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);
@@ -266,7 +270,7 @@ namespace SD_CUDA_NAMESPACE {
                 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;
         }

src/core/cuda/convrot_kernel.h

@@ -1,28 +1,28 @@
 #ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
 #define __SD_CORE_CUDA_CONVROT_KERNEL_H__
 #include <cstdint>
-#include "backend.h"
+#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.x += cuda_shuffle_xor(x.x, offset, 32);
-                x.y += __shfl_xor_sync(0xffffffff, x.y, offset, 32);
+                x.y += cuda_shuffle_xor(x.y, offset, 32);
             }
             return x;
         }
-        template <int Width = 32>
+        template <int Width>
         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) {
@@ -130,19 +130,19 @@ namespace SD_CUDA_NAMESPACE {
                         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 a = cuda_shuffle(v[j], base, 32);
-                            const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
-                            const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
-                            const float d = __shfl_sync(0xffffffff, v[j], base + 3 * 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 a = cuda_shuffle(v[j], lane % 16, 32);
-                        const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
-                        const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
-                        const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 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);
                     }

src/core/cuda/cute_attention.cu

@@ -1,4 +1,3 @@
-#include "backend.h"
 /*
  * Adapted from FlashAttention v2.8.3:
  * https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
@@ -461,23 +460,51 @@ static __global__ void copy_from_half(const half* src, float* dst, int64_t count
         dst[i] = scaled ? __fmaf_rn(__half2float(src[i]), scale, bias) : __half2float(src[i]);
 }
 
+static float param_f32(const ggml_tensor* node, int index) {
+    float value;
+    std::memcpy(&value, reinterpret_cast<const char*>(node->op_params) + index * sizeof(value), sizeof(value));
+    return value;
+}
 
 static int padded_dim(int64_t dim) {
     return dim == 80 || dim == 96 || dim == 112 ? int(dim) : int((dim + 63) / 64 * 64);
 }
 
+struct CUDADevice {
+    int ordinal;
+    cudaDeviceProp properties;
+};
+
+static const CUDADevice* cuda_device(ggml_backend_dev_t device) {
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, CUDADevice> devices;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto found = devices.find(device);
+    if (found != devices.end())
+        return &found->second;
+    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"));
+    if (!query)
+        return nullptr;
+    CUDADevice info = {};
+    info.ordinal    = query(device);
+    if (info.ordinal < 0 || cudaGetDeviceProperties(&info.properties, info.ordinal) != cudaSuccess)
+        return nullptr;
+    return &devices.emplace(device, info).first->second;
+}
+
 static bool attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
     static const bool disabled = std::getenv("GGML_CUDA_DISABLE_CUTE_FATTN") || std::getenv("GGML_CUDA_DISABLE_CUTLASS_FATTN");
-    const auto info            = sd_cuda::cuda_device_info(device);
+    const auto* info           = cuda_device(device);
     const auto* q              = node->src[0];
     const auto* k              = node->src[1];
     const auto* v              = node->src[2];
-    if (disabled || info.device < 0 || info.cc < 800 || !q || !k || !v ||
+    if (disabled || !info || info->properties.major < 8 || !q || !k || !v ||
         q->type != GGML_TYPE_F32 || node->type != GGML_TYPE_F32 ||
         q->ne[0] < 32 || q->ne[0] > 256 || q->ne[0] % 8 ||
-        size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info.shared_memory ||
+        size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info->properties.sharedMemPerBlockOptin ||
         q->ne[1] < 1 || k->ne[1] < 1 || k->ne[2] < 1 ||
-        node->src[3] || node->src[4] || ggml_get_op_params_f32(node, 1) != 0.0f || ggml_get_op_params_f32(node, 2) != 0.0f ||
+        node->src[3] || node->src[4] || param_f32(node, 1) != 0.0f || param_f32(node, 2) != 0.0f ||
         q->ne[2] % k->ne[2] || q->ne[3] != k->ne[3] || !ggml_are_same_shape(k, v) ||
         q->nb[0] != sizeof(float) || !ggml_is_contiguous(node) ||
         (k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_F32) ||
@@ -535,7 +562,7 @@ static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, in
             scaled = true;
         } else if (sd_ggml_convrot_describe(node, &convrot_group)) {
             if (convrot_group == 256 && previous->ne[0] % 256 == 0 &&
-                size_t(previous->ne[0]) * sizeof(float) + 128 <= sd_cuda::cuda_device_info(device).shared_memory)
+                size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
                 return count + 1;
             break;
         } else
@@ -617,7 +644,7 @@ static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes,
 }
 
 template <int D, bool FullDim>
-static cudaError_t attention_launch(const ggml_ext_gpu_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
     auto kernel                = sd_fattn_cute::kernel<64, 64, D, FullDim>;
     constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
     static std::mutex mutex;
@@ -636,7 +663,7 @@ static cudaError_t attention_launch(const ggml_ext_gpu_context& context, const s
 }
 
 template <int D>
-static cudaError_t attention_launch(const ggml_ext_gpu_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
     return p.dim == D ? attention_launch<D, true>(context, p, batches) : attention_launch<D, false>(context, p, batches);
 }
 
@@ -644,8 +671,10 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
     const auto& plan = *static_cast<AttentionPlan*>(state);
     const auto* attn = nodes[plan.index];
     auto device      = ggml_backend_get_device(execution->backend);
-    ggml_ext_gpu_context context;
+    auto reg         = ggml_backend_dev_backend_reg(device);
-    if (!sd_cuda::cuda_context(execution->backend, &context))
+    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))
         return GGML_STATUS_FAILED;
     auto stream = static_cast<cudaStream_t>(context.stream);
     auto* base  = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
@@ -655,7 +684,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
     const half* inputs[3] = {work[0], work[1], work[2]};
     for (int i = 0; i < 3; ++i) {
         const auto* src   = attn->src[i];
-        const float scale = i == 0 ? ggml_get_op_params_f32(attn, 0) : 1.0f;
+        const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
         if (i == 0 && plan.index == 5) {
             if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
                 return GGML_STATUS_FAILED;
@@ -706,8 +735,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
     float scale = 1.0f, bias = 0.0f;
     for (int i = plan.index + 1; i < count; ++i) {
         if (nodes[i]->op == GGML_OP_SCALE) {
-            scale = ggml_get_op_params_f32(nodes[i], 0);
+            scale = param_f32(nodes[i], 0);
-            bias  = ggml_get_op_params_f32(nodes[i], 1);
+            bias  = param_f32(nodes[i], 1);
         }
     }
     auto* dst = nodes[count - 1];
@@ -718,7 +747,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
         if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
             return GGML_STATUS_FAILED;
         if (cudaFuncSetAttribute(quantize_attention, cudaFuncAttributeMaxDynamicSharedMemorySize,
-                                 int(sd_cuda::cuda_device_info(device).shared_memory - 128)) != cudaSuccess)
+                                 int(cuda_device(device)->properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
             return GGML_STATUS_FAILED;
         quantize_attention<<<rows, 256, width * sizeof(float), stream>>>(work[3], packed,
                                                                          reinterpret_cast<float*>(packed + width * padded_rows), width, rows, scale, bias);
@@ -733,7 +762,7 @@ static void attention_destroy(void* state, void*) {
 }
 
 void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
-    if (sd_cuda::cuda_device_info(device).device < 0)
+    if (!cuda_device(device))
         return;
     ggml_ext_kernel_desc desc = {};
     desc.abi_version          = GGML_EXT_ABI_VERSION;

src/core/cuda/i8_gated_residual.cu

@@ -7,11 +7,71 @@
 #include <cstdlib>
 #include <cstring>
 #include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
-#include "backend.h"
+#include <mutex>
+#include <unordered_map>
 #include "core/ggml_extensions.h"
-#include "i8_gemm.cuh"
 
 namespace sd_i8_gated {
+    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+        using Base::Base;
+
+        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+            const auto shape = tiled_shape();
+            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+                constexpr int band_width = 16;
+                if (index >= shape.m() * shape.n()) {
+                    return {shape.m(), shape.n(), get_batch_idx()};
+                }
+                const int band   = index / (shape.m() * band_width);
+                const int width  = min(band_width, shape.n() - band * band_width);
+                const int within = index - band * shape.m() * band_width;
+                return {within / width, band * band_width + within % width, get_batch_idx()};
+            }
+#endif
+            return Base::get_tile_offset(index);
+        }
+    };
+
+    template <int M, int N, int WM>
+    struct TileConfig {
+        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+    };
+
+    enum Tile {
+        TILE_128_256,
+        TILE_128_128,
+        TILE_64_128,
+        TILE_16_128,
+    };
+
+    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+    }
+
+    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+        if (rows < 32) {
+            return TILE_16_128;
+        }
+        const int64_t columns = (n + 127) / 128;
+        if (((rows + 15) / 16) * columns <= nsm) {
+            return TILE_16_128;
+        }
+        if (((rows + 63) / 64) * columns <= nsm) {
+            return TILE_64_128;
+        }
+        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
+        const int64_t square_tiles = ((rows + 127) / 128) * columns;
+        // A narrower tile costs more per output element, but can waste less of the last wave.
+        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+    }
+
     template <class T>
     struct MulRound {
         CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
@@ -40,7 +100,7 @@ namespace sd_i8_gated {
     struct KernelTypes {
         using Tile         = typename Config::Tile;
         using Warp         = typename Config::Warp;
-        using Swizzle      = std::conditional_t<stream_k, sd_i8_cutlass::BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Swizzle      = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
         using Map          = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
         using Col          = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
         using Row          = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
@@ -84,6 +144,35 @@ namespace sd_i8_gated {
         using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
 
+    struct DeviceInfo {
+        int ordinal        = -1;
+        int major          = 0;
+        int sms            = 0;
+        int threads_per_sm = 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.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+        cache.emplace(device, info);
+        return info;
+    }
+    static float param_f32(const ggml_tensor* tensor, int index) {
+        float value;
+        std::memcpy(&value, tensor->op_params + index, sizeof(value));
+        return value;
+    }
     static bool aligned(const ggml_tensor* tensor) {
         if (!tensor)
             return false;
@@ -111,7 +200,7 @@ namespace sd_i8_gated {
         op.gate           = mul->src[1];
         op.residual       = add->src[0] == mul ? add->src[1] : add->src[0];
         op.prefix         = concat ? nodes[count - 1]->src[0] : nullptr;
-        op.scale          = nodes[1]->op == GGML_OP_SCALE ? ggml_get_op_params_f32(nodes[1], 0) : 1.f;
+        op.scale          = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f;
         op.rows           = ggml_nrows(nodes[0]);
         op.n              = nodes[0]->ne[0];
         op.k              = op.mat.weight->ne[0];
@@ -136,7 +225,7 @@ namespace sd_i8_gated {
         return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0};
     }
     struct Plan {
-        ggml_ext_gpu_device_info device;
+        DeviceInfo device;
         size_t kernel_workspace                                                     = 0;
         size_t output_offset                                                        = 0;
         int occupancy                                                               = 0;
@@ -159,7 +248,7 @@ namespace sd_i8_gated {
                 return GGML_STATUS_FAILED;
             plan.occupancy = occupancy;
         }
-        typename Kernel::Params params(args, plan.device.nsm, plan.occupancy);
+        typename Kernel::Params params(args, plan.device.sms, plan.occupancy);
         if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess)
             return GGML_STATUS_FAILED;
         cutlass::arch::synclog_setup();
@@ -170,15 +259,15 @@ namespace sd_i8_gated {
     static bool configure(const Operation& op, Plan& plan) {
         using Types  = KernelTypes<Config, Scaled, StreamK>;
         using Kernel = typename Types::Kernel;
-        if (sizeof(typename Kernel::SharedStorage) > plan.device.shared_memory)
+        if (sizeof(typename Kernel::SharedStorage) > plan.device.shared)
             return false;
         auto args = arguments<Types>(op, nullptr, true);
         if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
             return false;
         size_t workspace = 0;
         // Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement.
-        for (int occupancy = 1; occupancy <= plan.device.max_threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+        for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
-            typename Kernel::Params params(args, plan.device.nsm, occupancy);
+            typename Kernel::Params params(args, plan.device.sms, occupancy);
             const auto grid = params.get_grid_dims();
             if (grid.y > 65535 || grid.z > 65535)
                 return false;
@@ -196,23 +285,23 @@ namespace sd_i8_gated {
         return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
     }
     static bool configure(const Operation& op, Plan& plan) {
-        auto try_tile = [&](sd_i8_cutlass::Tile tile) {
+        auto try_tile = [&](Tile tile) {
             switch (tile) {
-                case sd_i8_cutlass::TILE_128_256:
+                case TILE_128_256:
-                    return configure<sd_i8_cutlass::TileConfig<128, 256, 64>>(op, plan);
+                    return configure<TileConfig<128, 256, 64>>(op, plan);
-                case sd_i8_cutlass::TILE_128_128:
+                case TILE_128_128:
-                    return configure<sd_i8_cutlass::TileConfig<128, 128, 64>>(op, plan);
+                    return configure<TileConfig<128, 128, 64>>(op, plan);
-                case sd_i8_cutlass::TILE_64_128:
+                case TILE_64_128:
-                    return configure<sd_i8_cutlass::TileConfig<64, 128, 32>>(op, plan);
+                    return configure<TileConfig<64, 128, 32>>(op, plan);
-                case sd_i8_cutlass::TILE_16_128:
+                case TILE_16_128:
-                    return configure<sd_i8_cutlass::TileConfig<16, 128, 16>>(op, plan);
+                    return configure<TileConfig<16, 128, 16>>(op, plan);
             }
             return false;
         };
-        const auto tile = sd_i8_cutlass::select_tile(op.rows, op.n, plan.device.nsm);
+        const auto tile = select_tile(op.rows, op.n, plan.device.sms);
         if (try_tile(tile))
             return true;
-        for (auto fallback : {sd_i8_cutlass::TILE_128_256, sd_i8_cutlass::TILE_128_128, sd_i8_cutlass::TILE_64_128, sd_i8_cutlass::TILE_16_128})
+        for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
             if (fallback != tile && try_tile(fallback))
                 return true;
         return false;
@@ -268,7 +357,7 @@ namespace sd_i8_gated {
         if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias)
             return 0;
         const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0];
-        if (rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
+        if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
             ((rows + 15) / 16) > INT_MAX / ((n + 127) / 128))
             return 0;
         for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale})
@@ -276,7 +365,7 @@ namespace sd_i8_gated {
                 return 0;
         const bool scaled = nodes[1]->op == GGML_OP_SCALE;
         int count         = scaled ? 4 : 3;
-        if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || ggml_get_op_params_f32(nodes[1], 1) != 0.f)))
+        if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f)))
             return 0;
         const auto* mul = nodes[count - 2];
         const auto* add = nodes[count - 1];
@@ -302,8 +391,8 @@ namespace sd_i8_gated {
             if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix)
                 return 0;
         Plan plan;
-        plan.device = sd_cuda::cuda_device_info(device);
+        plan.device = device_info(device);
-        if (plan.device.cc < 800 || plan.device.nsm <= 0)
+        if (plan.device.major < 8 || plan.device.sms <= 0)
             return 0;
         return configure(describe(nodes, count), plan) ? count : 0;
 #endif
@@ -311,7 +400,7 @@ namespace sd_i8_gated {
     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 = sd_cuda::cuda_device_info(ggml_backend_get_device(backend));
+        plan->device = device_info(ggml_backend_get_device(backend));
         if (!configure(describe(nodes, count), *plan))
             return GGML_STATUS_FAILED;
         *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]);
@@ -330,8 +419,10 @@ namespace sd_i8_gated {
         auto& plan    = *static_cast<Plan*>(state);
         const auto op = describe(nodes, count);
         auto* dst     = nodes[count - 1];
-        ggml_ext_gpu_context context;
+        auto reg      = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
-        if (!sd_cuda::cuda_context(execution->backend, &context) || context.device != plan.device.device)
+        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)
             return GGML_STATUS_FAILED;
         auto stream   = static_cast<cudaStream_t>(context.stream);
         auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));

src/core/cuda/i8_gemm.cu

@@ -5,11 +5,95 @@
 #include <climits>
 #include <cstdlib>
 #include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
-#include "backend.h"
+#include <mutex>
-#include "i8_gemm.cuh"
+#include <unordered_map>
 #include "i8_gemm.h"
 
 namespace sd_i8_cutlass {
+    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+        using Base::Base;
+
+        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+            const auto shape = tiled_shape();
+            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+                constexpr int band_width = 16;
+                if (index >= shape.m() * shape.n()) {
+                    return {shape.m(), shape.n(), get_batch_idx()};
+                }
+                const int band   = index / (shape.m() * band_width);
+                const int width  = min(band_width, shape.n() - band * band_width);
+                const int within = index - band * shape.m() * band_width;
+                return {within / width, band * band_width + within % width, get_batch_idx()};
+            }
+#endif
+            return Base::get_tile_offset(index);
+        }
+    };
+
+    template <int M, int N, int WM>
+    struct TileConfig {
+        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+    };
+
+    enum Tile {
+        TILE_128_256,
+        TILE_128_128,
+        TILE_64_128,
+        TILE_16_128,
+    };
+
+    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+    }
+
+    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+        if (rows < 32) {
+            return TILE_16_128;
+        }
+        const int64_t columns = (n + 127) / 128;
+        if (((rows + 15) / 16) * columns <= nsm) {
+            return TILE_16_128;
+        }
+        if (((rows + 63) / 64) * columns <= nsm) {
+            return TILE_64_128;
+        }
+        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
+        const int64_t square_tiles = ((rows + 127) / 128) * columns;
+        // A narrower tile costs more per output element, but can waste less of the last wave.
+        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+    }
+
+    struct DeviceInfo {
+        int ordinal        = -1;
+        int major          = 0;
+        int sms            = 0;
+        int threads_per_sm = 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.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+        cache.emplace(device, info);
+        return info;
+    }
     // The iterator pairs gate/up columns while reading the original checkpoint layout.
     struct WeightLayout : cutlass::layout::PermuteBase {
         int stride, half_n;
@@ -118,7 +202,7 @@ namespace sd_i8_cutlass {
         using Gemm            = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
     struct Plan {
-        ggml_ext_gpu_device_info device;
+        DeviceInfo device;
         int64_t rows, n, k;
         bool swiglu;
         size_t kernel_bytes = 0, scratch_offset = 0;
@@ -151,7 +235,7 @@ namespace sd_i8_cutlass {
                 p.occupancy <= 0)
                 return GGML_STATUS_FAILED;
         }
-        typename Kernel::Params params(args, p.device.nsm, p.occupancy);
+        typename Kernel::Params params(args, p.device.sms, p.occupancy);
         if (params.get_workspace_size() > p.kernel_bytes)
             return GGML_STATUS_FAILED;
         if (p.swiglu) {
@@ -168,14 +252,14 @@ namespace sd_i8_cutlass {
     template <class Types>
     static bool configure_kernel(Plan& p) {
         using Kernel = typename Types::Kernel;
-        if (sizeof(typename Kernel::SharedStorage) > p.device.shared_memory)
+        if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
             return false;
         auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
         if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
             return false;
         size_t bytes = 0;
-        for (int occupancy = 1; occupancy <= p.device.max_threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+        for (int occupancy = 1; occupancy <= p.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
-            typename Kernel::Params params(args, p.device.nsm, occupancy);
+            typename Kernel::Params params(args, p.device.sms, occupancy);
             const auto grid = params.get_grid_dims();
             if (grid.y > 65535 || grid.z > 65535)
                 return false;
@@ -206,7 +290,7 @@ namespace sd_i8_cutlass {
             }
             return false;
         };
-        const auto tile = select_tile(p.rows, p.n, p.device.nsm);
+        const auto tile = select_tile(p.rows, p.n, p.device.sms);
         if (try_tile(tile))
             return true;
         for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
@@ -220,8 +304,8 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row
 #ifdef SD_CUDA_FORCE_CUBLAS
     return nullptr;
 #else
-    const auto info = sd_cuda::cuda_device_info(device);
+    const auto info = sd_i8_cutlass::device_info(device);
-    if (info.device < 0 || info.cc < 800 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
+    if (info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
         rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
         return nullptr;
     auto* p   = new SDI8CudaGemmPlan;

src/core/cuda/i8_gemm.cuh

@@ -1,70 +0,0 @@
-#ifndef __SD_CORE_CUDA_I8_GEMM_CUH__
-#define __SD_CORE_CUDA_I8_GEMM_CUH__
-
-#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
-#include <cstdlib>
-
-namespace sd_i8_cutlass {
-    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
-        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
-        using Base::Base;
-
-        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
-            const auto shape = tiled_shape();
-            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
-                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
-                constexpr int band_width = 16;
-                if (index >= shape.m() * shape.n()) {
-                    return {shape.m(), shape.n(), get_batch_idx()};
-                }
-                const int band   = index / (shape.m() * band_width);
-                const int width  = min(band_width, shape.n() - band * band_width);
-                const int within = index - band * shape.m() * band_width;
-                return {within / width, band * band_width + within % width, get_batch_idx()};
-            }
-#endif
-            return Base::get_tile_offset(index);
-        }
-    };
-
-    template <int M, int N, int WM>
-    struct TileConfig {
-        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
-        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
-    };
-
-    enum Tile {
-        TILE_128_256,
-        TILE_128_128,
-        TILE_64_128,
-        TILE_16_128,
-    };
-
-    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
-        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
-        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
-    }
-
-    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
-        if (rows < 32) {
-            return TILE_16_128;
-        }
-        const int64_t columns = (n + 127) / 128;
-        if (((rows + 15) / 16) * columns <= nsm) {
-            return TILE_16_128;
-        }
-        if (((rows + 63) / 64) * columns <= nsm) {
-            return TILE_64_128;
-        }
-        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
-        const int64_t square_tiles = ((rows + 127) / 128) * columns;
-        // A narrower tile costs more per output element, but can waste less of the last wave.
-        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
-        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
-        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
-    }
-
-}
-
-#endif

src/core/cuda/i8_projection.cu

@@ -1,4 +1,3 @@
-#include "backend.h"
 #include "core/cuda/i8_projection.h"
 #include "core/cuda/rope.h"
 #include "core/ggml_extensions.h"
@@ -11,7 +10,9 @@
 #include <cstdlib>
 #include <cstring>
 #include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
 #include <type_traits>
+#include <unordered_map>
 
 static __device__ __forceinline__ float warp_sum8(float value) {
 #pragma unroll
@@ -149,16 +150,32 @@ struct ProjectionKernel {
     using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
 };
 
+static float param_f32(const ggml_tensor* tensor, int index) {
+    float result;
+    std::memcpy(&result, tensor->op_params + index, sizeof(result));
+    return result;
+}
 
 static bool device_supports(ggml_backend_dev_t device) {
 #ifdef SD_CUDA_FORCE_CUBLAS
     (void)device;
     return false;
 #else
-    const auto info      = sd_cuda::cuda_device_info(device);
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, bool> cache;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto found = cache.find(device);
+    if (found != cache.end())
+        return found->second;
+    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 = {};
     size_t shared        = std::max({sizeof(typename ProjectionKernel<float>::Kernel::SharedStorage),
                                      sizeof(typename ProjectionKernel<half>::Kernel::SharedStorage), sizeof(typename ProjectionKernel<half, false>::Kernel::SharedStorage)});
-    return info.device >= 0 && info.cc >= 800 && info.shared_memory >= shared;
+    bool supported       = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared;
+    cache.emplace(device, supported);
+    return supported;
 #endif
 }
 
@@ -222,7 +239,7 @@ static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat
         reinterpret_cast<uintptr_t>(output) % 8)
         return cudaErrorInvalidValue;
     typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}},
-                                      {{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? ggml_get_op_params_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
+                                      {{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
     typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr,
                                   int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0);
     Gemm gemm;
@@ -284,7 +301,7 @@ static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
             const auto* rope  = node_at(graph, first + 4);
             const auto* scale = node_at(graph, first + 5);
             const auto* copy  = node_at(graph, first + 6);
-            if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && ggml_get_op_params_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
+            if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
                 copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
                 return 7;
         }
@@ -330,8 +347,10 @@ static ggml_status projection_execute(const ggml_ext_execution_context* executio
     const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2];
     const bool joined  = concat->op == GGML_OP_CONCAT;
     auto* dst          = nodes[count - 1];
-    ggml_ext_gpu_context context;
+    auto reg           = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
-    if (!sd_cuda::cuda_context(execution->backend, &context))
+    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))
         return GGML_STATUS_FAILED;
     auto stream = static_cast<cudaStream_t>(context.stream);
     bool alias  = reinterpret_cast<uintptr_t>(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0;
@@ -340,7 +359,7 @@ static ggml_status projection_execute(const ggml_ext_execution_context* executio
             alias |= overlaps(dst, input);
     void* output   = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
     int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0;
-    float scale    = key && count == 7 && !joined ? ggml_get_op_params_f32(nodes[5], 0) : 1.f;
+    float scale    = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f;
     cudaError_t error;
     if (dst->type == GGML_TYPE_F32)
         error = projection_launch<float>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0);

src/core/cuda/int8.cu

@@ -7,17 +7,26 @@
 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>
+#endif
 #include <algorithm>
 #include <climits>
 #include <mutex>
 #include <unordered_map>
-#include "backend.h"
 
 namespace SD_CUDA_NAMESPACE {
     namespace sd_i8_cuda {
-        using sd_convrot::convrot_h4_cuda;
+        template <int Width = 32>
-        using sd_convrot::warp_reduce_max;
+        static __device__ __forceinline__ float warp_reduce_max(float value) {
-
+#pragma unroll
+            for (int offset = Width / 2; offset > 0; offset >>= 1)
+                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) {
             const int64_t row = blockIdx.x;
             const int tid     = threadIdx.x;
@@ -60,6 +69,13 @@ namespace SD_CUDA_NAMESPACE {
                 dst_row[i] = (int8_t)value;
             }
         }
+        static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+            a = index == 3 ? -a : a;
+            b = index == 2 ? -b : b;
+            c = index == 1 ? -c : c;
+            d = index == 0 ? -d : d;
+            return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+        }
         static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) {
             constexpr int group_size = 256;
             const int64_t row        = blockIdx.x;
@@ -86,19 +102,19 @@ namespace SD_CUDA_NAMESPACE {
                         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 a = cuda_shuffle(v[j], base, 32);
-                            const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
-                            const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
-                            const float d = __shfl_sync(0xffffffff, v[j], base + 3 * 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 a = cuda_shuffle(v[j], lane % 16, 32);
-                        const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
-                        const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
-                        const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 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);
                     }
@@ -483,10 +499,10 @@ namespace SD_CUDA_NAMESPACE {
             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);
+                auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
-                if (error != cudaSuccess)
+                if (error != cuda_success)
                     return error;
             }
             const size_t shared = size_t(p.k) * sizeof(float);
@@ -496,11 +512,9 @@ namespace SD_CUDA_NAMESPACE {
                 {
                     std::lock_guard<std::mutex> lock(mutex);
                     if (!configured[p.device.device]) {
-#if !defined(SD_CUDA_VENDOR_HIP)
+                        auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
-                        auto error = cudaFuncSetAttribute(quantize_convrot, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared_memory - 128));
+                        if (error != cuda_success)
-                        if (error != cudaSuccess)
                             return error;
-#endif
                         configured[p.device.device] = true;
                     }
                 }
@@ -513,7 +527,7 @@ namespace SD_CUDA_NAMESPACE {
                 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) {
@@ -535,7 +549,7 @@ namespace SD_CUDA_NAMESPACE {
             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;
@@ -543,17 +557,15 @@ namespace SD_CUDA_NAMESPACE {
             {
                 std::lock_guard<std::mutex> lock(mutex);
                 if (!initialized[p.device.device]) {
-#if !defined(SD_CUDA_VENDOR_HIP)
+                    auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
-                    auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared_memory - 128));
+                    if (error != cuda_success)
-                    if (error != cudaSuccess)
                         return error;
-#endif
                     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);
+                auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
-                if (error != cudaSuccess)
+                if (error != cuda_success)
                     return error;
             }
             const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
@@ -562,7 +574,7 @@ namespace SD_CUDA_NAMESPACE {
             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);
@@ -570,7 +582,7 @@ namespace SD_CUDA_NAMESPACE {
             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;
@@ -588,7 +600,7 @@ namespace SD_CUDA_NAMESPACE {
             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);
@@ -603,8 +615,8 @@ namespace SD_CUDA_NAMESPACE {
                         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);
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
-                        if (error == cudaSuccess)
+                        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 {
@@ -613,12 +625,12 @@ namespace SD_CUDA_NAMESPACE {
                     if (op.packed || p.rows == p.padded)
                         quant = data;
                     else {
-                        error = cudaMemcpyAsync(quant, data, size_t(p.k) * p.rows, cudaMemcpyDeviceToDevice, stream);
+                        error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
-                        if (error == cudaSuccess)
+                        if (error == cuda_success)
-                            error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                            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;
@@ -643,11 +655,24 @@ namespace SD_CUDA_NAMESPACE {
                     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
                 }
@@ -665,7 +690,7 @@ namespace SD_CUDA_NAMESPACE {
                         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) {
@@ -674,14 +699,14 @@ namespace SD_CUDA_NAMESPACE {
                     } 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)
+                if (error == cuda_success && prefix)
-                    error = cudaMemcpyAsync(output, dst->src[0]->data, prefix, cudaMemcpyDeviceToDevice, stream);
+                    error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
             }
-            if (error == cudaSuccess && alias)
+            if (error == cuda_success && alias)
-                error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
-            return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
         }
 
     }

src/core/cuda/registration.cu

@@ -1,4 +1,4 @@
-#include "backend.h"
+#include "runtime.h"
 
 #ifdef SD_CUDA_CUTE_ATTENTION
 void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);

src/core/cuda/rope.cu

@@ -9,7 +9,7 @@
 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, 32));
+            v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
         return v;
     }
     template <int D, bool interleaved, typename T>
@@ -54,7 +54,7 @@ namespace SD_CUDA_NAMESPACE {
             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]);
@@ -136,6 +136,11 @@ namespace SD_CUDA_NAMESPACE {
         return value != 0;
     }
 
+    static float param_f32(const ggml_tensor* node, int index) {
+        float value;
+        std::memcpy(&value, node->op_params + index, sizeof(value));
+        return value;
+    }
 
     static bool rope_is_node(const ggml_tensor* node) {
         return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
@@ -174,7 +179,7 @@ namespace SD_CUDA_NAMESPACE {
     }
 
     template <typename T, bool Interleaved>
-    static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) {
+    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);
@@ -183,7 +188,7 @@ namespace SD_CUDA_NAMESPACE {
         rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
             static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),                               \
             static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows,                                               \
-            x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), ggml_get_op_params_f32(nodes[0], 0), scale);               \
+            x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale);                            \
         break
         switch (x->ne[0]) {
             SD_ROPE(32);
@@ -195,13 +200,13 @@ namespace SD_CUDA_NAMESPACE {
             SD_ROPE(224);
             SD_ROPE(256);
             default:
-                return cudaErrorInvalidValue;
+                return cuda_invalid_value;
         }
 #undef SD_ROPE
-        return cudaGetLastError();
+        return cuda_last_error();
     }
 
-    cudaError_t 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);
@@ -223,7 +228,7 @@ namespace SD_CUDA_NAMESPACE {
             const auto* scale = ggml_graph_node(g, first + 3);
             const auto* copy  = ggml_graph_node(g, first + 4);
             if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
-                scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && ggml_get_op_params_f32(scale, 1) == 0.f &&
+                scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
                 ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
                 return 5;
         }
@@ -280,15 +285,15 @@ namespace SD_CUDA_NAMESPACE {
         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 = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? ggml_get_op_params_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];
@@ -300,11 +305,11 @@ namespace SD_CUDA_NAMESPACE {
             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 == cudaSuccess && alias)
+        if (error == cuda_success && alias)
-            error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+            error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
-        return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
     }
 
     void register_rope(ggml_backend_dev_t device) {

src/core/cuda/rope.h

@@ -1,10 +1,11 @@
 #ifndef __SD_CORE_CUDA_ROPE_H__
 #define __SD_CORE_CUDA_ROPE_H__
 
-#include "backend.h"
+#include "runtime.h"
 
 namespace SD_CUDA_NAMESPACE {
     bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
-    cudaError_t 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);
+
 }
 #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/cuda/sage_attention.cu

@@ -1,5 +1,5 @@
-#include "backend.h"
 #include "core/ggml_extensions.h"
+#include "runtime.h"
 
 #if __CUDACC_VER_MAJOR__ >= 12
 
@@ -353,8 +353,9 @@ namespace sd_cuda {
         if (info.device < 0)
             return;
         int previous;
+        cudaDeviceProp properties;
         cudaFuncAttributes attributes;
-        if (cudaGetDevice(&previous) != cudaSuccess ||
+        if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess ||
             cudaSetDevice(info.device) != cudaSuccess)
             return;
         const auto error = cudaFuncGetAttributes(&attributes, sage_output);
@@ -363,7 +364,7 @@ namespace sd_cuda {
         // PTX compiled for an older target retains that target's instruction guards after JIT.
         static std::map<int, SageDevice> devices;
         auto& data                = devices[info.device];
-        data                      = {info.cc / 10, std::min(attributes.binaryVersion, attributes.ptxVersion)};
+        data                      = {properties.major * 10 + properties.minor, std::min(attributes.binaryVersion, attributes.ptxVersion)};
         ggml_ext_kernel_desc desc = {};
         desc.abi_version          = GGML_EXT_ABI_VERSION;
         desc.name                 = "sd.sage_attn.cuda.v1";

src/core/cuda/sol_attention.cu

@@ -1,7 +1,7 @@
 // Sol-Attn diagonal-threshold algorithm:
 // https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
-#include "backend.h"
 #include "core/ggml_extensions.h"
+#include "runtime.h"
 
 #if __CUDACC_VER_MAJOR__ >= 12
 
@@ -379,8 +379,9 @@ namespace sd_cuda {
         if (info.device < 0)
             return;
         int previous;
+        cudaDeviceProp properties;
         cudaFuncAttributes attributes;
-        if (info.cc < 800 ||
+        if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
             cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
             return;
         const auto error = cudaFuncGetAttributes(&attributes, sol_forward);

src/core/ggml_extensions.cpp

@@ -2,7 +2,6 @@
 
 #ifndef SD_USE_UPSTREAM_GGML
 #include "core/cpu/registration.h"
-#include "ggml-cpu.h"
 #include "ggml-ext.h"
 
 #include <algorithm>
@@ -335,7 +334,8 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
         return;
     }
     auto device    = ggml_backend_get_device(backend);
-    const bool cpu = ggml_backend_is_cpu(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;
     static std::mutex mutex;
     static std::unordered_set<ggml_backend_dev_t> registered;
     std::lock_guard<std::mutex> lock(mutex);