[MUSA][8/N] Port CUDA kernels that are compatible with MUSA (#17946)

Signed-off-by: yafeng.li <yafeng.li@mthreads.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
MARATRIX
2026-04-23 18:04:58 -07:00
committed by GitHub
co-authored by Alex Nails
parent c0166355ae
commit 74c2e5bacd
15 changed files with 1905 additions and 66 deletions
+2 -2
View File
@@ -7,7 +7,7 @@
#include <limits>
#include <vector>
#ifndef USE_ROCM
#if !defined(USE_ROCM) && !defined(USE_MUSA)
#include <dlfcn.h>
#define WARP_SIZE 32
#include "pytorch_extension_utils.h"
@@ -24,7 +24,7 @@ transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_
#pragma unroll
for (int j = lane_id; j < total_chunks; j += WARP_SIZE) {
#ifndef USE_ROCM
#if !defined(USE_ROCM) && !defined(USE_MUSA)
uint64_t tmp;
asm volatile("ld.global.nc.b64 %0,[%1];" : "=l"(tmp) : "l"(src + j) : "memory");
asm volatile("st.global.cg.b64 [%0],%1;" ::"l"(dst + j), "l"(tmp) : "memory");