[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>

sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu

@@ -652,7 +652,11 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch:
   TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "Only BFloat16 output dtype is supported")
 
   auto const sm = getSMVersion();
+#ifndef USE_MUSA
   TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");
+#else
+  TORCH_CHECK(sm >= 22, "required MUSA ARCH >= MP_22");
+#endif
 
   auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
   if (num_tokens <= 8) {

sgl-kernel/csrc/gemm/dsv3_router_gemm_entry.cu

@@ -121,7 +121,11 @@ void dsv3_router_gemm(
       output.dtype() == torch::kFloat32 || output.dtype() == torch::kBFloat16, "output must be float32 or bf16");
 
   auto const sm = getSMVersion();
+#ifndef USE_MUSA
   TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");
+#else
+  TORCH_CHECK(sm >= 22, "required MUSA ARCH >= MP_22");
+#endif
 
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 

sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu

@@ -90,16 +90,25 @@ __forceinline__ __device__ OUT_DTYPE_T extract_required_scale_format(float value
 }
 
 __device__ __forceinline__ void st_global(const int4* ptr, const int4& value) {
+#ifndef USE_MUSA
   asm volatile(
       "st.global.v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w));
+#else
+  int4* p = const_cast<int4*>(ptr);
+  *p = value;
+#endif
 }
 
 __device__ __forceinline__ int4 ld_global_nc(const int4* ptr) {
+#ifndef USE_MUSA
   int4 ret;
   asm volatile("ld.global.nc.v4.s32 {%0, %1, %2, %3}, [%4];"
                : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
                : "l"(ptr));
   return ret;
+#else
+  return *ptr;
+#endif
 }
 
 template <typename T>