Commit f9af9be21 for llama.cpp

commit f9af9be219ca647a59106f6201bf0d85fab00224
Author: R0CKSTAR <yeahdongcn@gmail.com>
Date:   Fri Sep 25 15:16:29 2026 +0800

    musa: fix PH1 (MTT S5000) operator failures and build issues (#29193)

    * musa: use 16-byte copies for MUSA like sm_70+

    ggml_cuda_get_max_cpy_bytes() derives the copy width from __CUDA_ARCH__. mcc
    never defines it, so MUSA fell into the generic branch and returned 8 bytes
    instead of the 16 bytes that every sm_70+ target gets. The value sizes the
    per-thread copy unit of the FlashAttention K/V staging code (fattn-common,
    fattn-vec, fattn-tile, fattn-mma-f16 shared-memory loads) and of mmq-vec-dot,
    so every MUSA FlashAttention kernel moved half as many bytes per instruction.

    On an MTT S5000 (mp_31, MUSA SDK 5.2.0) with Qwen3.8-27B-UD-Q4_K_M, -ngl 999,
    -p 512 -n 64, -fa on: 751.15 -> 794.73 t/s prefill and 15.59 -> 15.69 t/s
    decode. -fa off is unchanged (1050.05 -> 1052.86 t/s prefill), FLASH_ATTN_EXT
    is unchanged (3984 ok / 0 fail / 1323 unsupported) and perplexity is
    unchanged.

    * musa: enable the CUB paths on MUSA

    GGML_CUDA_USE_CUB and USE_CUB are selected by "CUDART_VERSION >= 11070", which
    the MUSA SDK never satisfies: CUDART_VERSION is not defined anywhere under
    /usr/local/musa/include, so the condition is always false and every CUB-based
    path stayed compiled out on MUSA even though the SDK ships CUB and the kernels
    build for mp_31.  Select them from GGML_USE_MUSA as well.  The device-wide
    algorithms are usable too: cub::DeviceSegmentedSort compiles and produces
    correct results on mp_31.

    This lifts the ne[0] <= 1024 limit that ggml_backend_cuda_device_supports_op
    applied to ARGSORT and TOP_K on MUSA.  On an MTT S5000 (S5000, mcc 5.2.0):
    ARGSORT 48 ok / 52 not supported -> 100 ok / 0 (CUDA parity), TOP_K 0 ok /
    354 not supported -> 527 ok / 0.  The other 20 per-op suites are unchanged, the
    Qwen3-0.6B f16 (14.4679) and Qwen3.8-27B iq4_nl (5.1724) perplexities are
    unchanged, and the 0.6B graph keeps the same nodes and splits (18 CPU + 18
    MUSA0, SET_ROWS 1008) as before.

    * musa: take the upstream code path where the toolkit supports it

    Several guards were written for an older MUSA toolkit. Verified against MUSA SDK
    5.2.0 and on an MTT S5000 (mp_31):

    - device init: query cudaDevAttrCooperativeLaunch instead of hardcoding false.
      The device reports cooperativeLaunch=1 and musaLaunchCooperativeKernel works
      (verified with a kernel whose result was checked).
    - device init: keep prop.warpSize instead of overriding it with 32. The device
      reports 32 anyway, so this only removes the divergence.
    - CUDA_SET_SHARED_MEMORY_LIMIT and the FA shared-memory raise: musaFuncSetAttribute
      returns success and sharedMemPerBlockOptin is 192 KiB, so the kernels can use
      more than the default 48 KiB.
    - vendors/musa.h: add the cudaDeviceGetAttribute and cudaDevAttrCooperativeLaunch
      mappings the device-init change needs.

    Measured on one S5000 with Qwen3.8-27B Q4_K_M (-ngl 999, -r 3): pp512 968.27 ->
    957.09 t/s, tg64 10.09 -> 10.23 t/s, FLASH_ATTN_EXT sweep identical (3975/3982
    both), perplexity identical (80.2841 +/- 7.26772 both).

    * musa: drop compile-time guards that MUSA's runtime gates already cover

    mcc never defines __CUDA_ARCH__, so the arch-gated fallbacks in this group
    were already taken on MUSA and the GGML_USE_MUSA guards on top of them only
    kept the upstream text from being compiled:

      - wkv.cu: the "#pragma unroll" suppression has no effect on the generated
        code that is not already covered by the surrounding guards
      - common.cuh: the MUSA-only __builtin_unreachable() in no_device_code() is
        not needed to silence the compiler
      - ssm-scan.cu: the SSD (Mamba-2 prefill) block and its dispatch are gated at
        runtime by GGML_CUDA_CC_IS_NVIDIA(cc) and turing_mma_available(cc), which
        are both false for PH1 (cc 0x100310), so compiling them changes nothing
      - common.cuh: warp_reduce_max(half2) is guarded the same way as
        warp_reduce_sum(half2) (FP16_AVAILABLE); the MUSA-only guard left the
        function with no return statement. It has no caller today.

    MTT S5000 (mp_31, MUSA SDK 5.2.0), MUSA_ARCHITECTURES=31: build rc=0. Against
    an unmodified build of the same tree on the same card, FLASH_ATTN_EXT
    (3984 ok / 0 fail / 1323 unsupported), SSM_SCAN (15/0), RWKV_WKV6 (6/0),
    GATED_DELTA_NET (38/0) and MUL_MAT (1299/0/385 unsupported) are identical, and
    perplexity with -fa on is bit-identical (5.1639 +/- 0.36673, 4 chunks).

    * musa: do not use MMQ on PH1

    test-backend-ops on an MTT S5000 (mp_31, MUSA SDK 5.2.0) fails 260 cases and every
    one of them goes through the MMQ path:

      - MUL_MAT with a batched src1 (any bs/nr != [1,1]): 109 cases across all
        quantized types, e.g. 12 of 13 cases at n=16, while the plain [1,1] layout
        passes
      - every quantized MUL_MAT_ID: 147 cases, while the f16/f32 variants of the same
        shapes pass
      - MUL_MAT with more than ~512 tokens: 4 cases (n=509..4096); the small-n cases pass

    The cuBLAS/dequant path is correct for all of them and the MMVQ path used for
    small batches is unaffected, so quantized matmuls now take that path on PH1
    instead of returning wrong values. 27B perplexity with default flags goes from
    nan to finite, and the full suite reports 0 failures out of 22237 cases.

    The MMQ defect itself (fastdiv, __umulhi, uint3 kernel parameters and
    __CUDA_ARCH__-based MMA availability were all checked and are correct on this
    part) is not addressed here.

    * musa: keep the block barrier of the fused TOPK_MOE kernel reachable

    topk_moe_cuda returns early for the rows past the end of the graph, but one block
    covers TOPK_MOE_ROWS_PER_BLOCK (8) rows, so the last block is only partially filled
    whenever n_rows is not a multiple of 8.  On MUSA a warp that has already returned
    blocks the block wide __syncthreads() below, which makes the kernel hang and the
    launch time out.  CUDA tolerates the exited warps, which is why the CUDA numbers
    never showed it.

    For MUSA, clamp the row index of those warps to the last row so that every warp of
    the block reaches the barrier; they recompute the last row and write the same
    values.  The CUDA code path is unchanged.

    On an MTT S5000 (mp_31) the fused TOPK_MOE cases change from a launch timeout with
    no completed case to 418 ok / 0 not supported / 0 failed, i.e. the CUDA result, and
    the other 101 per op suites are unchanged (0 failed, no count changes).

    * musa: enable GATED_DELTA_NET

    The op was turned off for every MUSA target because mcc could not build the kernel
    at the time. The current toolkit builds it: with mp_31 and MUSA SDK 5.2.0 the file
    compiles with zero errors and all 36 test-backend-ops GATED_DELTA_NET cases pass
    against the CPU reference. 27B perplexity is unchanged.

    While the op is refused, the scheduler has no choice but to run it on the CPU: 48
    GATED_DELTA_NET nodes per forward pass. On an MTT S5000 (Qwen3.8-27B Q4_K_M, -ngl
    999, one container, -r 3):

        pp512 (FA off)   964.51 -> 2119.26 t/s
        tg64  (FA off)    10.15 ->   15.50 t/s

    * musa: name the stream capture query API for the graph aware kernels

    argsort.cu and mean.cu call cudaStreamCaptureStatus, cudaStreamIsCapturing and
    cudaStreamCaptureStatusNone inside their USE_CUDA_GRAPH blocks, but the MUSA
    compatibility headers do not alias those names, so building with the experimental
    GGML_MUSA_GRAPHS option fails with 7 errors in those two files.  Map the three
    names to their musa* counterparts, under the same guard that enables the graph
    code, so the default build is untouched.

    The option stays off by default: on an MTT S5000 the captured path measured
    slower (pp512 693 vs 772 t/s, tg128 15.20 vs 15.39 t/s over two sessions) and the
    borderline MUL_MAT cases are not reproducible between runs.

    * musa: build the CI and docs for PH1 (MTT S5000)

    The MUSA CI job and the documented default still targeted the first generation
    (MTT S80, MUSA_ARCHITECTURES=21) while the current MUSA SDK targets PH1
    (MTT S5000, 31).  Move the job, ci/run.sh's default and the build docs to 31,
    and run the job in the PH1 MUSA SDK devel image:

        registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64

    That image needs two things the previous one did not: python3-venv for the
    ccache-buckets step, which builds a virtual environment for the Hugging Face
    CLI, and no time prefix on the build command, because container jobs run their
    steps with sh and the image ships no time binary.

diff --git a/.github/workflows/build-cuda-ubuntu.yml b/.github/workflows/build-cuda-ubuntu.yml
index 30029887c..3dc9255be 100644
--- a/.github/workflows/build-cuda-ubuntu.yml
+++ b/.github/workflows/build-cuda-ubuntu.yml
@@ -145,7 +145,7 @@ jobs:

   musa:
     runs-on: ubuntu-22.04
-    container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
+    container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64

     steps:
       - name: Clone
@@ -156,7 +156,7 @@ jobs:
         id: depends
         run: |
           apt-get update
-          apt-get install -y build-essential git cmake libssl-dev jq
+          apt-get install -y build-essential git cmake libssl-dev jq python3-venv

       - name: ccache
         uses: ggml-org/ccache-action@v1.2.24
@@ -178,8 +178,8 @@ jobs:
         run: |
           cmake -B build -S . \
             -DGGML_MUSA=ON \
-            -DMUSA_ARCHITECTURES=21
-          time cmake --build build --config Release -j $(nproc)
+            -DMUSA_ARCHITECTURES=31
+          cmake --build build --config Release -j $(nproc)

       - name: ccache-buckets-save
         if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
diff --git a/ci/README-MUSA.md b/ci/README-MUSA.md
index c5e24c5d9..2101dfa2c 100644
--- a/ci/README-MUSA.md
+++ b/ci/README-MUSA.md
@@ -21,7 +21,7 @@ docker run --privileged -it \
     -v $HOME/llama.cpp/ci-cache:/ci-cache \
     -v $HOME/llama.cpp/ci-results:/ci-results \
     -v $PWD:/ws -w /ws \
-    mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
+    registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64
 ```

 Inside the container, execute the following commands:
diff --git a/ci/run.sh b/ci/run.sh
index 0595fac5a..fc661f976 100755
--- a/ci/run.sh
+++ b/ci/run.sh
@@ -158,8 +158,8 @@ if [ ! -z ${GG_BUILD_WEBGPU} ]; then
 fi

 if [ ! -z ${GG_BUILD_MUSA} ]; then
-    # Use qy1 by default (MTT S80)
-    MUSA_ARCH=${MUSA_ARCH:-21}
+    # Use ph1 by default (MTT S5000)
+    MUSA_ARCH=${MUSA_ARCH:-31}
     CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_MUSA=ON -DMUSA_ARCHITECTURES=${MUSA_ARCH}"
 fi

diff --git a/docs/build.md b/docs/build.md
index 70fc17af2..bd666c145 100644
--- a/docs/build.md
+++ b/docs/build.md
@@ -323,11 +323,11 @@ cmake --build build --config Release
 By default, all supported compute capabilities are enabled. To customize this behavior, you can specify the `MUSA_ARCHITECTURES` option in the CMake command:

 ```bash
-cmake -B build -DGGML_MUSA=ON -DMUSA_ARCHITECTURES="21"
+cmake -B build -DGGML_MUSA=ON -DMUSA_ARCHITECTURES="31"
 cmake --build build --config Release
 ```

-This configuration enables only compute capability `2.1` (MTT S80) during compilation, which can help reduce compilation time.
+This configuration enables only compute capability `3.1` (MTT S5000) during compilation, which can help reduce compilation time.

 #### Compilation options

diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
index 2e78ae4fa..5363dd001 100644
--- a/ggml/src/ggml-cuda/common.cuh
+++ b/ggml/src/ggml-cuda/common.cuh
@@ -111,9 +111,9 @@
 #define GGML_CUDA_CC_IS_QY2(cc)      (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_PH1)
 #define GGML_CUDA_CC_IS_PH1(cc)      (cc >= GGML_CUDA_CC_PH1)

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#if !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070)
 #    define GGML_CUDA_USE_CUB
-#endif  // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#endif  // !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070)

 // PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8.
 // However, this has been bugged in CTK < 12.3 for MSVC builds, see
@@ -237,7 +237,7 @@ static const char * cu_get_error_str(CUresult err) {
 #define CU_CHECK(err) CUDA_CHECK_GEN(err, CUDA_SUCCESS, cu_get_error_str)
 #endif

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_HIP)
 #    define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes)                                                       \
         do {                                                                                                   \
             static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = { false };                         \
@@ -252,7 +252,7 @@ static const char * cu_get_error_str(CUresult err) {
         do {                                             \
             GGML_UNUSED(nbytes);                         \
         } while (0)
-#endif // !(defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_HIP)

 #if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA)
 #define GGML_CUDA_ASSUME(x) __builtin_assume(x)
@@ -397,7 +397,7 @@ static constexpr __device__ int ggml_cuda_get_physical_warp_size() {

 // Maximum number of bytes that can be copied in a single instruction.
 static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() {
-#ifdef GGML_USE_HIP
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
     return 16;
 #else
 #if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
@@ -405,7 +405,7 @@ static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() {
 #else
     return 8;
 #endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
-#endif // GGML_USE_HIP
+#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
 }


@@ -424,10 +424,6 @@ static __device__ void no_device_code(
     __trap();

     GGML_UNUSED(no_device_code); // suppress unused function warning
-
-#if defined(GGML_USE_MUSA)
-    __builtin_unreachable();
-#endif // defined(GGML_USE_MUSA)
 }

 #ifdef __CUDA_ARCH__
@@ -696,16 +692,11 @@ static __device__ __forceinline__ half2 ggml_cuda_hmax2(const half2 a, const hal

 template<int width = WARP_SIZE>
 static __device__ __forceinline__ half2 warp_reduce_max(half2 x) {
-#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL || defined(GGML_USE_HIP)
 #pragma unroll
    for (int offset = width/2; offset > 0; offset >>= 1) {
        x = ggml_cuda_hmax2(x, __shfl_xor_sync(0xffffffff, x, offset, width));
    }
    return x;
-#else
-   GGML_UNUSED(x);
-   NO_DEVICE_CODE;
-#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL || defined(GGML_USE_HIP)
 }

 #if (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || \
diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index 449a77c5b..abb99a354 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -2091,26 +2091,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
             constexpr bool use_sparse_kernel = false;
             fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;

-#if !defined(GGML_USE_MUSA)
             static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
             if (!shared_memory_limit_raised[id]) {
                 CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
                 shared_memory_limit_raised[id] = true;
             }
-#endif // !defined(GGML_USE_MUSA)
         }
     } else {
         constexpr bool use_logit_softcap = true;
         constexpr bool use_sparse_kernel = false;
         fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;

-#if !defined(GGML_USE_MUSA)
         static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
         if (!shared_memory_limit_raised[id]) {
             CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
             shared_memory_limit_raised[id] = true;
         }
-#endif // !defined(GGML_USE_MUSA)
     }

     launch_fattn<DV, ncols1, ncols2>
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 36c6a93b1..e76ff3128 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -310,13 +310,9 @@ static ggml_cuda_device_info ggml_cuda_init() {
         info.devices[id].smpb       = prop.sharedMemPerBlock;
         info.devices[id].warp_size  = prop.warpSize;

-#ifndef GGML_USE_MUSA
         int supports_coop_launch = 0;
         CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, physical_id));
         info.devices[id].supports_cooperative_launch = !!supports_coop_launch;
-#else
-        info.devices[id].supports_cooperative_launch = false;
-#endif // !(GGML_USE_MUSA)

 #if defined(GGML_USE_HIP)
         info.devices[id].smpbo = prop.sharedMemPerBlock;
@@ -337,8 +333,6 @@ static ggml_cuda_device_info ggml_cuda_init() {
                       device_vmm ? "yes" : "no", prop.warpSize,
                       device_vram_mib);
 #elif defined(GGML_USE_MUSA)
-        // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs.
-        info.devices[id].warp_size = 32;
         info.devices[id].smpbo = prop.sharedMemPerBlockOptin;
         info.devices[id].cc = GGML_CUDA_CC_OFFSET_MTHREADS + prop.major * 0x100;
         info.devices[id].cc += prop.minor * 0x10;
@@ -5581,12 +5575,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_RWKV_WKV7:
             return true;
         case GGML_OP_GATED_DELTA_NET:
-            //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327
-#ifdef GGML_USE_MUSA
-            return false;
-#else
             return true;
-#endif // GGML_USE_MUSA
         case GGML_OP_DSV4_HC_COMB:
             return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
                 op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu
index b13b34ee9..f68b3df18 100644
--- a/ggml/src/ggml-cuda/mmq.cu
+++ b/ggml/src/ggml-cuda/mmq.cu
@@ -389,5 +389,10 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
         return n_experts > 0;
     }

+    // MUSA: the MMQ kernels compute wrong values on PH1 (MTT S5000).
+    if (cc == GGML_CUDA_CC_PH1) {
+        return false;
+    }
+
     return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
 }
diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu
index 40cb38dee..e9a1043fb 100644
--- a/ggml/src/ggml-cuda/ssm-scan.cu
+++ b/ggml/src/ggml-cuda/ssm-scan.cu
@@ -1,6 +1,6 @@
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#if !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070)
 #define USE_CUB
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#endif // !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070)

 #ifdef USE_CUB
 #include <cub/cub.cuh>
@@ -342,7 +342,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
     }
 }

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_HIP)
 // ============================================================================
 // SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS)
 //
@@ -821,7 +821,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX);
     GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX);

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_HIP)
     // Mamba-2 with scalar A per head: use SSD matmul path for long sequences.
     // Requires NVIDIA Turing+ otherwise fallback to scan.
     const bool is_mamba2 = (src3->nb[1] == sizeof(float));
diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu
index dadcd601c..ee903bf21 100644
--- a/ggml/src/ggml-cuda/topk-moe.cu
+++ b/ggml/src/ggml-cuda/topk-moe.cu
@@ -98,7 +98,12 @@ __global__ void topk_moe_cuda(const float *         logits,
                               const float           clamp_val,
                               const float           scale_val,
                               const topk_moe_config config) {
+#if defined(GGML_USE_MUSA)
+    // MUSA: every warp of a partially filled block must reach the barrier below.
+    const int row = MIN(blockIdx.x * blockDim.y + threadIdx.y, n_rows - 1);
+#else
     const int row = blockIdx.x * blockDim.y + threadIdx.y;
+#endif // defined(GGML_USE_MUSA)
     if (row >= n_rows) {
         return;
     }
diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h
index 6d725c7ec..4243caab4 100644
--- a/ggml/src/ggml-cuda/vendors/musa.h
+++ b/ggml/src/ggml-cuda/vendors/musa.h
@@ -44,6 +44,7 @@
 #define cudaDeviceGetPCIBusId musaDeviceGetPCIBusId
 #define cudaDeviceProp musaDeviceProp
 #define cudaDeviceSynchronize musaDeviceSynchronize
+#define cudaDeviceGetAttribute musaDeviceGetAttribute
 #define cudaError_t musaError_t
 #define cudaErrorMemoryAllocation musaErrorMemoryAllocation
 #define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled
@@ -114,6 +115,7 @@
 #define cuMemRelease muMemRelease
 #define cuMemSetAccess muMemSetAccess
 #define cuMemUnmap muMemUnmap
+#define cudaDevAttrCooperativeLaunch musaDevAttrCooperativeLaunch
 #define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize
 #define cudaFuncSetAttribute musaFuncSetAttribute
 #define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms
@@ -145,6 +147,9 @@
 #define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed
 #define cudaStreamBeginCapture musaStreamBeginCapture
 #define cudaStreamEndCapture musaStreamEndCapture
+#define cudaStreamCaptureStatus musaStreamCaptureStatus
+#define cudaStreamCaptureStatusNone musaStreamCaptureStatusNone
+#define cudaStreamIsCapturing musaStreamIsCapturing
 #define cudaOccupancyMaxActiveBlocksPerMultiprocessor musaOccupancyMaxActiveBlocksPerMultiprocessor

 typedef __mt_bfloat16 nv_bfloat16;
diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu
index 236111212..0bf997760 100644
--- a/ggml/src/ggml-cuda/wkv.cu
+++ b/ggml/src/ggml-cuda/wkv.cu
@@ -79,9 +79,7 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons
     float state[head_size];
     __shared__ float _r[head_size], _w[head_size], _k[head_size], _a[head_size], _b[head_size];

-#ifndef GGML_USE_MUSA
     #pragma unroll
-#endif
     for (int i = 0; i < head_size; i++) {
         state[i] = s[batch_i * state_size + head_i * head_size * head_size + tid * head_size + i];
     }