[DAS] Adapt code for dcu

csrc/custom_marlin/gptq_marlin/gptq_marlin.cu

@@ -1765,7 +1765,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
         if (has_act_order) {
             // Permute A columns
             int block_rows = div_ceil(prob_m, blocks);
-            permute_cols_kernel << <blocks, default_threads, 0, stream >> > (
+            permute_cols_kernel <<<blocks, default_threads, 0, stream >>> (
                 A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, block_rows);
             A_ptr = a_tmp_ptr;
         }

csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh

@@ -52,24 +52,24 @@ template <> class ScalarType<nv_bfloat16> {
     using FragC = Vec<float, 4>;
     using FragS = Vec<nv_bfloat162, 1>;
 
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
-    static __device__ float inline num2float(const nv_bfloat16 x) {
+//#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+static __device__ float inline num2float(const nv_bfloat16 x) {
     return __bfloat162float(x);
-    }
+}
 
-    static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
+static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
     return __bfloat162bfloat162(x);
-    }
+}
 
-    static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
+static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
                                                 const nv_bfloat16 x2) {
     return __halves2bfloat162(x1, x2);
-    }
+}
 
-    static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
+static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
     return __float2bfloat16(x);
-    }
-#endif
+}
+//#endif
 };
 
 } // namespace gptq_marlin

csrc/ktransformers_ext/CMakeLists.txt

@@ -118,7 +118,8 @@ elseif (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LW
     message(STATUS "x86 detected")
     set(HOST_IS_X86 TRUE)
     set(HAS_AVX512 TRUE)
-    set(__HAS_AMX__ TRUE)
+    set(__HAS_AMX__ False)
+    #set(__HAS_AMX__ TRUE)
     add_compile_definitions(__x86_64__)
     # check AVX512
     execute_process(
@@ -141,12 +142,12 @@ elseif (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LW
     # check AMX
     string(FIND "${LSCPU_OUTPUT}" "amx" COMPILER_SUPPORTS_AMX)
     
-    if(COMPILER_SUPPORTS_AMX GREATER -1)
-        message(STATUS "Compiler supports AMX")
-        add_compile_definitions(__HAS_AMX__)
-    else()
+    #if(COMPILER_SUPPORTS_AMX GREATER -1)
+    #   message(STATUS "Compiler supports AMX")
+    #    add_compile_definitions(__HAS_AMX__)
+    #else()
     message(STATUS "Compiler does NOT support AMX")
-    endif()
+    #endif()
     if (MSVC)
         # instruction set detection for MSVC only
         if (LLAMA_NATIVE)

ktransformers/operators/attention.py

@@ -704,11 +704,7 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention):
                 cache_position,
                 **kwargs,
             )
-        elif (os.name == 'nt'  or get_compute_capability()<8
-              or hidden_states.device.type == 'cpu'
-              or device_manager.gpu_vendor != GPUVendor.NVIDIA) 
-              or ("Z100" in get_device_name()) 
-              or ("Z100L" in get_device_name()) or ("K100" in get_device_name()):
+        elif (os.name == 'nt'  or get_compute_capability()<8 or hidden_states.device.type == 'cpu' or device_manager.gpu_vendor != GPUVendor.NVIDIA or ("Z100" in get_device_name()) or ("Z100L" in get_device_name()) or ("K100" in get_device_name())):
             print("for Windows or GPU before ampere or Z100/Z100L or K100, use forward_windows")
             return self.forward_windows(
                 hidden_states,

ktransformers/operators/models.py

@@ -660,12 +660,7 @@ class KDeepseekV2Model(BaseInjectedModule):
         else:
             #if os.name == 'nt' or get_compute_capability()<8:
              #   print("for Windows or GPU before ampere, use forward_windows")
-            if os.name == 'nt' or get_compute_capability()<8  
-                or (self.transfer_map is not None and 'cpu' in self.transfer_map.values())
-                or device_manager.gpu_vendor != GPUVendor.NVIDIA
-                or (self.transfer_map is not None and 'cpu' in self.transfer_map.values())
-                or device_manager.gpu_vendor != GPUVendor.NVIDIA
-                or ("Z100" in get_device_name()) or ("Z100L" in get_device_name()) or ("K100" in get_device_name()):
+            if os.name == 'nt' or get_compute_capability()<8  or (self.transfer_map is not None and 'cpu' in self.transfer_map.values()) or device_manager.gpu_vendor != GPUVendor.NVIDIA  or ("Z100" in get_device_name()) or ("Z100L" in get_device_name()) or ("K100" in get_device_name()):
                 print("for Windows or GPU before ampere or Z100/Z100L or K100, use forward_windows")
                 # only use mask in forward windows or can't flash attn
                 causal_mask = self._update_causal_mask(